文章來源：機(jī)器之心
性能與 Mamba 一樣，但所需訓(xùn)練步驟數(shù)卻少 2.2 倍。
圖片來源：由無界 AI生成
狀態(tài)空間模型（SSM）是近來一種備受關(guān)注的 Transformer 替代技術(shù)，其優(yōu)勢(shì)是能在長上下文任務(wù)上實(shí)現(xiàn)線性時(shí)間的推理、并行化訓(xùn)練和強(qiáng)大的性能。而基于選擇性 SSM 和硬件感知型設(shè)計(jì)的 Mamba 更是表現(xiàn)出色，成為了基于注意力的 Transformer 架構(gòu)的一大有力替代架構(gòu)。
近期也有一些研究者在探索將 SSM 和 Mamba 與其它方法組合起來創(chuàng)造更強(qiáng)大的架構(gòu)，比如機(jī)器之心曾報(bào)告過《Mamba 可以替代 Transformer，但它們也能組合起來使用》。
近日，波蘭一個(gè)研究團(tuán)隊(duì)發(fā)現(xiàn)，如果將 SSM 與混合專家系統(tǒng)（MoE/Mixture of Experts）組合起來，可望讓 SSM 實(shí)現(xiàn)大規(guī)模擴(kuò)展。MoE 是目前常用于擴(kuò)展 Transformer 的技術(shù)，比如近期的 Mixtral 模型就使用了這一技術(shù)，參閱機(jī)器之心文章。
這個(gè)波蘭研究團(tuán)隊(duì)給出的研究成果是 MoE-Mamba，即將 Mamba 和混合專家層組合起來的模型。
論文地址：https://arxiv.org/pdf/2401.04081.pdf
MoE-Mamba 能同時(shí)提升 SSM 和 MoE 的效率。而且該團(tuán)隊(duì)還發(fā)現(xiàn)，當(dāng)專家的數(shù)量發(fā)生變化時(shí)，MoE-Mamba 的行為是可預(yù)測(cè)的。
該團(tuán)隊(duì)也進(jìn)行了實(shí)驗(yàn)論證，如圖 1 所示，結(jié)果表明：相比于 Mamba，MoE-Mamba 達(dá)到同等性能時(shí)所需的訓(xùn)練步驟數(shù)少 2.2 倍，這彰顯了新方法相較于 Transformer 和 Transformer-MoE 的潛在優(yōu)勢(shì)。這些初步結(jié)果也指出了一個(gè)頗具潛力的研究方向：SSM 也許可以擴(kuò)展到數(shù)百億參數(shù)！
相關(guān)研究
狀態(tài)空間模型
狀態(tài)空間模型（SSM）是一類用于序列建模的架構(gòu)。這些模型的思想源自控制論領(lǐng)域，可被看作是 RNN 和 CNN 的組合。盡管它們具有相當(dāng)大的優(yōu)勢(shì)，但也有一些問題，因此難以成為語言建模任務(wù)的主導(dǎo)架構(gòu)。但是，近期的一些研究突破卻讓深度 SSM 可以擴(kuò)展到數(shù)十億參數(shù)，同時(shí)還能維持計(jì)算效率和強(qiáng)大的性能表現(xiàn)。
Mamba
Mamba 是基于 SSM 構(gòu)建的模型，能實(shí)現(xiàn)線性時(shí)間的推理速度（對(duì)上下文長度而言），并且其還通過硬件感知型設(shè)計(jì)實(shí)現(xiàn)了高效的訓(xùn)練流程。Mamba 采用了一種工作高效型的并行掃描方法，可以減輕循環(huán)的序列性的影響，而融合 GPU 操作則可無需實(shí)現(xiàn)擴(kuò)展?fàn)顟B(tài)。反向傳播所必需的中間狀態(tài)不會(huì)被保存下來，而是會(huì)在反向通過過程中被重新計(jì)算，由此可以降低內(nèi)存需求。Mamba 優(yōu)于注意力機(jī)制的優(yōu)勢(shì)在推理階段尤其顯著，因?yàn)槠洳粌H能降低計(jì)算復(fù)雜度，而且內(nèi)存使用量還不會(huì)取決于上下文長度。
Mamba 能解決序列模型的效率和效果之間的根本性權(quán)衡，這就凸顯了狀態(tài)壓縮的重要性。高效的模型必需要小狀態(tài)，而有效的模型所需的狀態(tài)應(yīng)當(dāng)包含上下文的所有關(guān)鍵信息。不同于其它 SSM 對(duì)時(shí)間和輸入不變性的需求，Mamba 引入了一種選擇機(jī)制，可以控制信息沿序列維度傳播的方式。這一設(shè)計(jì)選擇的靈感來自對(duì)選擇性復(fù)制和歸納頭等合成任務(wù)的直觀理解，讓模型可以分辨和保留關(guān)鍵信息，同時(shí)濾除無關(guān)信息。
研究發(fā)現(xiàn)，Mamba 有能力高效地利用更長的上下文（長達(dá) 1M token），并且隨著上下文長度增長，預(yù)訓(xùn)練困惑度也會(huì)得到改善。Mamba 模型是由堆疊的 Mamba 塊構(gòu)成的，在 NLP、基因組學(xué)、音頻等多個(gè)不同領(lǐng)域都取得了非常好的結(jié)果，其性能可以媲美和超越已有的 Transformer 模型。因此，Mamba 成為了通用序列建模骨干模型的一個(gè)有力候選模型，參閱《五倍吞吐量，性能全面包圍 Transformer：新架構(gòu) Mamba 引爆 AI 圈》。
混合專家
混合專家（MoE）這類技術(shù)能極大提升模型的參數(shù)數(shù)量，同時(shí)不會(huì)影響模型推理和訓(xùn)練所需的 FLOPs。MoE 最早由 Jacobs et al. 于 1991 年提出，并在 2017 年由 Shazeer et al. 開始用于 NLP 任務(wù)。
MoE 有一個(gè)優(yōu)勢(shì)：激活很稀疏 —— 對(duì)于處理的每個(gè) token，只會(huì)用到模型的一小部分參數(shù)。由于其計(jì)算需求，Transformer 中的前向?qū)右呀?jīng)變成了多種 MoE 技術(shù)的標(biāo)準(zhǔn)目標(biāo)。
研究社區(qū)已經(jīng)提出了多種方法用于解決 MoE 的核心問題，即將 token 分配給專家的過程，也稱路由（routing）過程。目前有兩種基本的路由算法：Token Choice 和 Expert Choice。其中前者是將每個(gè) token 路由到一定數(shù)量（K）的專家，至于后者則是路由到每個(gè)專家的 token 數(shù)量是固定的。
Fedus et al. 在 2022 年的論文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》中提出的 Switch 是一種 Token Choice 架構(gòu)，其是將每個(gè) token 路由到單個(gè)專家（K=1），而他們使用該方法將 Transformer 的參數(shù)規(guī)模成功擴(kuò)增至了 1.6 萬億。波蘭的這個(gè)團(tuán)隊(duì)在實(shí)驗(yàn)中也采用了這種 MoE 設(shè)計(jì)。
最近，MoE 也開始進(jìn)入開源社區(qū)，比如 OpenMoE。
項(xiàng)目地址：https://github.com/XueFuzhao/OpenMoE
尤其值得一提的是 Mistral 開源的 Mixtral 8×7B，其性能可比肩 LLaMa 2 70B，同時(shí)所需的推理計(jì)算預(yù)算只有后者的約六分之一。
模型架構(gòu)
盡管 Mamba 的主要底層機(jī)制與 Transformer 中使用的注意力機(jī)制大不相同，但 Mamba 保留了 Transformer 模型的高層級(jí)、基于模塊的結(jié)構(gòu)。使用這一范式，由相同模塊構(gòu)成的一層或多層會(huì)彼此堆疊在一起，而每一層的輸出都會(huì)被添加到殘差流（residual stream）中，見圖 2。之后，這個(gè)殘差流的最終值會(huì)被用于預(yù)測(cè)語言建模任務(wù)的下一個(gè) token。
MoE-Mamba 利用了這兩種架構(gòu)的兼容能力。如圖 2 所示，在 MoE-Mamba 中，每間隔一個(gè) Mamba 層就會(huì)被替換成一個(gè)基于 Switch 的 MoE 前饋層。
不過該團(tuán)隊(duì)也注意到這一設(shè)計(jì)和《Mamba: Linear-time sequence modeling with selective state spaces》的設(shè)計(jì)有些相似；后者交替堆疊了 Mamba 層和前饋層，但得到的模型相比于單純的 Mamba 還略有不及。該設(shè)計(jì)在圖 1 中被記為 Mamba-MLP。
MoE-Mamba 分開了 Mamba 層執(zhí)行的每個(gè) token 的無條件處理和 MoE 層執(zhí)行的有條件處理；其中的無條件處理可高效地將序列的整個(gè)上下文整合到一個(gè)內(nèi)部表征中，而有條件處理可為每個(gè) token 使用最相關(guān)的專家。這種將有條件處理和無條件處理交替起來的思路在一些基于 MoE 的模型中已經(jīng)得到了應(yīng)用，不過它們通常是交替基本的和 MoE 的前饋層。
主要結(jié)果
訓(xùn)練設(shè)置
該團(tuán)隊(duì)比較了 5 種不同設(shè)置：基本 Transformer、Mamba、Mamba-MLP、MoE 和 MoE-Mamba。
在大多數(shù) Transformers 中，前饋層包含 8dm2 個(gè)參數(shù)，而 Mamba 論文中則讓 Mamba 更小一些（約 6dm2），這樣兩個(gè) Mamba 層的參數(shù)數(shù)量與一個(gè)前饋層和一個(gè)注意力層加起來差不多。為了讓 Mamba 和新模型中每個(gè) token 的活動(dòng)參數(shù)數(shù)量大致一樣，該團(tuán)隊(duì)將每個(gè)專家前向?qū)拥拇笮】s小到了 6dm2。除了嵌入層和解除嵌入（unembedding）層，所有模型都是每個(gè) token 使用大約 2600 萬參數(shù)。訓(xùn)練過程使用了 65 億個(gè) token，訓(xùn)練步驟數(shù)為 100k。
訓(xùn)練使用的數(shù)據(jù)集是 English C4 數(shù)據(jù)集，任務(wù)是預(yù)測(cè)下一 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í)率是針對(duì)普通 Mamba 進(jìn)行調(diào)整的，因此可以預(yù)計(jì)，如果再針對(duì) MoE-Mamba 對(duì)訓(xùn)練流程進(jìn)行一番優(yōu)化，MoE-Mamba 的表現(xiàn)還會(huì)更好。
消融研究
為了評(píng)估 Mamba 是否能隨專家數(shù)量的增長而很好地?cái)U(kuò)展，研究者比較了使用不同數(shù)量專家的模型。
圖 3 展示了使用不同數(shù)量的專家時(shí)的訓(xùn)練運(yùn)行步驟情況。
表 2 給出了 100k 步驟后的結(jié)果。
這些結(jié)果表明新提出的方法能隨專家數(shù)量而很好地?cái)U(kuò)展。如果專家的數(shù)量為 8 或更多，新模型的最終性能優(yōu)于普通 Mamba。由于 Mamba-MLP 比普通 Mamba 差，可以預(yù)見使用少量專家的 MoE-Mamba 的性能表現(xiàn)會(huì)比 Mamba 差。當(dāng)專家數(shù)為 32 時(shí)，新方法得到了最佳結(jié)果。
