首頁 > 文章 > 科技週邊 > MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模

MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模

王林轉載: 2024-01-23 18:00:14729瀏覽

狀態空間模型（SSM）是一種備受關注的技術，它被認為是Transformer的替代方案。相較於Transformer，SSM在處理長上下文任務時能夠實現線性時間的推理，同時具備平行化訓練和出色的效能。特別是基於選擇性SSM和硬體感知型設計的Mamba，更是展現出了卓越的表現，成為了基於注意力的Transformer架構的有力替代之一。

最近，研究人員也正在探索將SSM和Mamba與其他方法結合起來，以創造更強大的架構。例如，機器之心曾報道《Mamba可以取代Transformer，但它們也可以結合使用》。

近日，波蘭一個研究團隊發現，如果將 SSM 與混合專家系統（MoE/Mixture of Experts）組合起來，可望讓 SSM 實現大規模擴展。 MoE 是目前常用於擴展 Transformer 的技術，例如近期的 Mixtral 模型就使用了這項技術，請參閱機器之心文章。

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

MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模

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

MoE -Mamba 能同時提升SSM 和MoE 的效率。而團隊也發現，當專家的數量改變時，MoE-Mamba 的行為是可預測的。

該團隊進行了實驗論證，結果顯示，MoE-Mamba相較於Mamba，以相同性能需求，所需的訓練步驟減少了2.2倍，顯示了新方法相較於Transformer和Transformer-MoE的潛在優勢。這些初步結果也指出了一個有潛力的研究方向：SSM可能可以擴展到數百億參數。

MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模

相關研究

#狀態空間模型

#狀態空間模型（SSM）是一類用於序列建模的架構。這些模型的想法源自於控制論領域，可被視為 RNN 和 CNN 的組合。儘管它們具有相當大的優勢，但也有一些問題，因此難以成為語言建模任務的主導架構。但是，近期的一些研究突破卻讓深度 SSM 可以擴展到數十億參數，同時也能維持運算效率和強大的效能表現。

Mamba

#Mamba 是一個基於SSM 建構的模型，能實現線性時間的推理速度（對上下文長度而言），並且其也透過硬體感知型設計實現了高效的訓練流程。 Mamba 採用了一種工作高效型的平行掃描方法，可減輕循環的序列性的影響，而融合 GPU 操作則可無需實現擴展狀態。反向傳播所必需的中間狀態不會被保存下來，而是會在反向通過過程中被重新計算，由此可以降低記憶體需求。 Mamba 優於注意力機制的優勢在推理階段尤其顯著，因為其不僅能降低計算複雜度，而且記憶體使用量還不會取決於上下文長度。

Mamba 能解決序列模型的效率和效果之間的根本性權衡，這就凸顯了狀態壓縮的重要性。高效率的模型必需要小狀態，而有效的模型所需的狀態應包含上下文的所有關鍵資訊。不同於其它 SSM 對時間和輸入不變性的需求，Mamba 引入了一種選擇機制，可以控制訊息沿著序列維度傳播的方式。這項設計選擇的靈感來自對選擇性複製和歸納頭等合成任務的直觀理解，讓模型可以分辨和保留關鍵訊息，同時濾除無關訊息。

研究發現，Mamba 有能力有效率地利用更長的脈絡（長達 1M token），並且隨著情境長度增長，預訓練困惑度也會得到改善。 Mamba 模型是由堆疊的 Mamba 區塊構成的，在 NLP、基因組學、音訊等多個不同領域都取得了非常好的結果，其性能可以媲美和超越現有的 Transformer 模型。因此，Mamba 成為了通用序列建模骨幹模型的一個有力候選模型，請參閱《五倍吞吐量，性能全麵包圍Transformer：新架構Mamba 引爆AI 圈#》。

混合專家

#混合專家（MoE）這類技術能大幅提昇模型的參數數量，同時不會影響模型推理和訓練所需的FLOPs。 MoE 最早由 Jacobs et al. 於 1991 年提出，並在 2017 年由 Shazeer et al. 開始用於 NLP 任務。

MoE 有一個優點：活化很稀疏 —— 對於處理的每個 token，只會用到模型的一小部分參數。由於其計算需求，Transformer 中的前向層已經變成了多種 MoE 技術的標準目標。

研究社群已經提出了多種方法用於解決 MoE 的核心問題，即將 token 分配給專家的過程，也稱為路由（routing）過程。目前有兩種基本的路由演算法：Token Choice 和 Expert Choice。其中前者是將每個 token 路由到一定數量（K）的專家，至於後者則是路由到每個專家的 token 數量是固定的。

Fedus et al. 在2022 年的論文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》中提出的Switch 是一種Token Choice 架構，其是將每個token 路由到單一專家（K=1），而他們使用該方法將Transformer 的參數規模成功擴增至了1.6 兆。波蘭的這個團隊在實驗中也採用了這種 MoE 設計。

最近，MoE 也開始進入開源社區，例如 OpenMoE。

專案網址：https://github.com/XueFuzhao/OpenMoE

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

模型架構

儘管Mamba 的主要底層機制與Transformer 中使用的注意力機制大不相同，但Mamba 保留了Transformer 模型的高層級、基於模組的結構。使用這個範式，由相同模組構成的一層或多層會彼此堆疊在一起，而每一層的輸出都會被加入到殘差流（residual stream）中，見圖 2。之後，這個殘差流的最終值會被用來預測語言建模任務的下一個 token。

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

MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模

不過團隊也注意到這個設計和《Mamba: Linear-time sequence modeling with selective state spaces》的設計有些相似；後者交替堆疊了Mamba 層和前饋層，但得到的模型相比於單純的Mamba 還略有不及。此設計在圖 1 中被記為 Mamba-MLP。

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

主要結果

訓練設定

##該團隊比較了5 種不同設定：基本Transformer、Mamba、Mamba-MLP、MoE 和MoE-Mamba。

在大多数 Transformers 中，前馈层包含 8dm² 个参数，而 Mamba 论文中则让 Mamba 更小一些（约 6dm²），这样两个 Mamba 层的参数数量与一个前馈层和一个注意力层加起来差不多。为了让 Mamba 和新模型中每个 token 的活动参数数量大致一样，该团队将每个专家前向层的大小缩小到了 6dm²。除了嵌入层和解除嵌入（unembedding）层，所有模型都是每个 token 使用大约 2600 万参数。训练过程使用了 65 亿个 token，训练步骤数为 100k。

训练使用的数据集是 English C4 数据集，任务是预测下一 token。文本的 token 化使用了 GPT2 tokenizer。表 3 给出了超参数的完整列表。

MoE與Mamba攜手合作，將狀態空間模型推廣至數十億參數規模