文章來源：機(jī)器之心

性能與 Mamba 一樣,，但所需訓(xùn)練步驟數(shù)卻少 2.2 倍,。

圖片來源：由無界 AI生成

狀態(tài)空間模型（SSM）是近來一種備受關(guān)注的 Transformer 替代技術(shù),，其優(yōu)勢是能在長上下文任務(wù)上實(shí)現(xiàn)線性時間的推理,、并行化訓(xùn)練和強(qiáng)大的性能。而基于選擇性 SSM 和硬件感知型設(shè)計的 Mamba 更是表現(xiàn)出色,，成為了基于注意力的 Transformer 架構(gòu)的一大有力替代架構(gòu),。

近期也有一些研究者在探索將 SSM 和 Mamba 與其它方法組合起來創(chuàng)造更強(qiáng)大的架構(gòu)，比如機(jī)器之心曾報告過《Mamba 可以替代 Transformer,，但它們也能組合起來使用》,。

近日，波蘭一個研究團(tuán)隊發(fā)現(xiàn),，如果將 SSM 與混合專家系統(tǒng)（MoE/Mixture of Experts）組合起來,，可望讓 SSM 實(shí)現(xiàn)大規(guī)模擴(kuò)展。MoE 是目前常用于擴(kuò)展 Transformer 的技術(shù),，比如近期的 Mixtral 模型就使用了這一技術(shù),，參閱機(jī)器之心文章。

這個波蘭研究團(tuán)隊給出的研究成果是 MoE-Mamba,，即將 Mamba 和混合專家層組合起來的模型,。

論文地址：https://arxiv.org/pdf/2401.04081.pdf

MoE-Mamba 能同時提升 SSM 和 MoE 的效率,。而且該團(tuán)隊還發(fā)現(xiàn),，當(dāng)專家的數(shù)量發(fā)生變化時，MoE-Mamba 的行為是可預(yù)測的,。

該團(tuán)隊也進(jìn)行了實(shí)驗(yàn)論證,，如圖 1 所示,，結(jié)果表明：相比于 Mamba，MoE-Mamba 達(dá)到同等性能時所需的訓(xùn)練步驟數(shù)少 2.2 倍,，這彰顯了新方法相較于 Transformer 和 Transformer-MoE 的潛在優(yōu)勢,。這些初步結(jié)果也指出了一個頗具潛力的研究方向：SSM 也許可以擴(kuò)展到數(shù)百億參數(shù)！

相關(guān)研究

狀態(tài)空間模型

狀態(tài)空間模型（SSM）是一類用于序列建模的架構(gòu),。這些模型的思想源自控制論領(lǐng)域,，可被看作是 RNN 和 CNN 的組合。盡管它們具有相當(dāng)大的優(yōu)勢,，但也有一些問題，因此難以成為語言建模任務(wù)的主導(dǎo)架構(gòu)。但是，近期的一些研究突破卻讓深度 SSM 可以擴(kuò)展到數(shù)十億參數(shù),，同時還能維持計算效率和強(qiáng)大的性能表現(xiàn),。

Mamba

Mamba 是基于 SSM 構(gòu)建的模型,，能實(shí)現(xiàn)線性時間的推理速度（對上下文長度而言），并且其還通過硬件感知型設(shè)計實(shí)現(xiàn)了高效的訓(xùn)練流程。Mamba 采用了一種工作高效型的并行掃描方法,，可以減輕循環(huán)的序列性的影響,，而融合 GPU 操作則可無需實(shí)現(xiàn)擴(kuò)展?fàn)顟B(tài),。反向傳播所必需的中間狀態(tài)不會被保存下來，而是會在反向通過過程中被重新計算,，由此可以降低內(nèi)存需求。Mamba 優(yōu)于注意力機(jī)制的優(yōu)勢在推理階段尤其顯著，因?yàn)槠洳粌H能降低計算復(fù)雜度，而且內(nèi)存使用量還不會取決于上下文長度。

Mamba 能解決序列模型的效率和效果之間的根本性權(quán)衡,，這就凸顯了狀態(tài)壓縮的重要性,。高效的模型必需要小狀態(tài),，而有效的模型所需的狀態(tài)應(yīng)當(dāng)包含上下文的所有關(guān)鍵信息,。不同于其它 SSM 對時間和輸入不變性的需求，Mamba 引入了一種選擇機(jī)制,，可以控制信息沿序列維度傳播的方式。這一設(shè)計選擇的靈感來自對選擇性復(fù)制和歸納頭等合成任務(wù)的直觀理解，讓模型可以分辨和保留關(guān)鍵信息,，同時濾除無關(guān)信息。

研究發(fā)現(xiàn),，Mamba 有能力高效地利用更長的上下文（長達(dá) 1M token）,，并且隨著上下文長度增長，預(yù)訓(xùn)練困惑度也會得到改善,。Mamba 模型是由堆疊的 Mamba 塊構(gòu)成的,，在 NLP、基因組學(xué),、音頻等多個不同領(lǐng)域都取得了非常好的結(jié)果，其性能可以媲美和超越已有的 Transformer 模型。因此，Mamba 成為了通用序列建模骨干模型的一個有力候選模型,，參閱《五倍吞吐量,，性能全面包圍 Transformer：新架構(gòu) Mamba 引爆 AI 圈》,。

混合專家

混合專家（MoE）這類技術(shù)能極大提升模型的參數(shù)數(shù)量,，同時不會影響模型推理和訓(xùn)練所需的 FLOPs,。MoE 最早由 Jacobs et al. 于 1991 年提出,，并在 2017 年由 Shazeer et al. 開始用于 NLP 任務(wù)。

MoE 有一個優(yōu)勢：激活很稀疏 —— 對于處理的每個 token，只會用到模型的一小部分參數(shù),。由于其計算需求，Transformer 中的前向?qū)右呀?jīng)變成了多種 MoE 技術(shù)的標(biāo)準(zhǔn)目標(biāo)。

研究社區(qū)已經(jīng)提出了多種方法用于解決 MoE 的核心問題，即將 token 分配給專家的過程,，也稱路由（routing）過程,。目前有兩種基本的路由算法：Token Choice 和 Expert Choice,。其中前者是將每個 token 路由到一定數(shù)量（K）的專家，至于后者則是路由到每個專家的 token 數(shù)量是固定的,。

Fedus et al. 在 2022 年的論文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》中提出的 Switch 是一種 Token Choice 架構(gòu)，其是將每個 token 路由到單個專家（K=1），而他們使用該方法將 Transformer 的參數(shù)規(guī)模成功擴(kuò)增至了 1.6 萬億,。波蘭的這個團(tuán)隊在實(shí)驗(yàn)中也采用了這種 MoE 設(shè)計,。

最近,，MoE 也開始進(jìn)入開源社區(qū)，比如 OpenMoE,。

項(xiàng)目地址：https://github.com/XueFuzhao/OpenMoE

尤其值得一提的是 Mistral 開源的 Mixtral 8×7B,，其性能可比肩 LLaMa 2 70B,，同時所需的推理計算預(yù)算只有后者的約六分之一。

模型架構(gòu)

盡管 Mamba 的主要底層機(jī)制與 Transformer 中使用的注意力機(jī)制大不相同,，但 Mamba 保留了 Transformer 模型的高層級,、基于模塊的結(jié)構(gòu)。使用這一范式,，由相同模塊構(gòu)成的一層或多層會彼此堆疊在一起,，而每一層的輸出都會被添加到殘差流（residual stream）中，見圖 2,。之后,，這個殘差流的最終值會被用于預(yù)測語言建模任務(wù)的下一個 token。

MoE-Mamba 利用了這兩種架構(gòu)的兼容能力,。如圖 2 所示，在 MoE-Mamba 中,，每間隔一個 Mamba 層就會被替換成一個基于 Switch 的 MoE 前饋層,。

不過該團(tuán)隊也注意到這一設(shè)計和《Mamba: Linear-time sequence modeling with selective state spaces》的設(shè)計有些相似,；后者交替堆疊了 Mamba 層和前饋層,，但得到的模型相比于單純的 Mamba 還略有不及。該設(shè)計在圖 1 中被記為 Mamba-MLP,。

MoE-Mamba 分開了 Mamba 層執(zhí)行的每個 token 的無條件處理和 MoE 層執(zhí)行的有條件處理,；其中的無條件處理可高效地將序列的整個上下文整合到一個內(nèi)部表征中，而有條件處理可為每個 token 使用最相關(guān)的專家,。這種將有條件處理和無條件處理交替起來的思路在一些基于 MoE 的模型中已經(jīng)得到了應(yīng)用,，不過它們通常是交替基本的和 MoE 的前饋層。

主要結(jié)果

訓(xùn)練設(shè)置

該團(tuán)隊比較了 5 種不同設(shè)置：基本 Transformer,、Mamba,、Mamba-MLP、MoE 和 MoE-Mamba,。

在大多數(shù) Transformers 中,，前饋層包含 8dm2 個參數(shù)，而 Mamba 論文中則讓 Mamba 更小一些（約 6dm2）,，這樣兩個 Mamba 層的參數(shù)數(shù)量與一個前饋層和一個注意力層加起來差不多,。為了讓 Mamba 和新模型中每個 token 的活動參數(shù)數(shù)量大致一樣，該團(tuán)隊將每個專家前向?qū)拥拇笮】s小到了 6dm2,。除了嵌入層和解除嵌入（unembedding）層,，所有模型都是每個 token 使用大約 2600 萬參數(shù)。訓(xùn)練過程使用了 65 億個 token,，訓(xùn)練步驟數(shù)為 100k,。

訓(xùn)練使用的數(shù)據(jù)集是 English C4 數(shù)據(jù)集,，任務(wù)是預(yù)測下一 token。文本的 token 化使用了 GPT2 tokenizer,。表 3 給出了超參數(shù)的完整列表,。

結(jié)果

表 1 給出了訓(xùn)練結(jié)果,。MoE-Mamba 的表現(xiàn)顯著優(yōu)于普通 Mamba 模型,。

值得注意的是,，MoE-Mamba 只用僅僅 46% 的訓(xùn)練步驟就達(dá)到了與普通 Mamba 同等的結(jié)果水平,。由于學(xué)習(xí)率是針對普通 Mamba 進(jìn)行調(diào)整的，因此可以預(yù)計,，如果再針對 MoE-Mamba 對訓(xùn)練流程進(jìn)行一番優(yōu)化,，MoE-Mamba 的表現(xiàn)還會更好。

消融研究

為了評估 Mamba 是否能隨專家數(shù)量的增長而很好地擴(kuò)展,，研究者比較了使用不同數(shù)量專家的模型,。

圖 3 展示了使用不同數(shù)量的專家時的訓(xùn)練運(yùn)行步驟情況。

表 2 給出了 100k 步驟后的結(jié)果,。

這些結(jié)果表明新提出的方法能隨專家數(shù)量而很好地擴(kuò)展,。如果專家的數(shù)量為 8 或更多,，新模型的最終性能優(yōu)于普通 Mamba。由于 Mamba-MLP 比普通 Mamba 差,，可以預(yù)見使用少量專家的 MoE-Mamba 的性能表現(xiàn)會比 Mamba 差,。當(dāng)專家數(shù)為 32 時，新方法得到了最佳結(jié)果,。