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

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

圖片來(lái)源：由無(wú)界 AI生成

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

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

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

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

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

MoE-Mamba 能同時(shí)提升 SSM 和 MoE 的效率,。而且該團(tuán)隊(duì)還發(fā)現(xiàn),，當(dāng)專(zhuān)家的數(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）是一類(lèi)用于序列建模的架構(gòu)。這些模型的思想源自控制論領(lǐng)域,，可被看作是 RNN 和 CNN 的組合,。盡管它們具有相當(dāng)大的優(yōu)勢(shì)，但也有一些問(wèn)題,，因此難以成為語(yǔ)言建模任務(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ì)上下文長(zhǎng)度而言），并且其還通過(guò)硬件感知型設(shè)計(jì)實(shí)現(xiàn)了高效的訓(xùn)練流程,。Mamba 采用了一種工作高效型的并行掃描方法,，可以減輕循環(huán)的序列性的影響，而融合 GPU 操作則可無(wú)需實(shí)現(xiàn)擴(kuò)展?fàn)顟B(tài),。反向傳播所必需的中間狀態(tài)不會(huì)被保存下來(lái),，而是會(huì)在反向通過(guò)過(guò)程中被重新計(jì)算，由此可以降低內(nèi)存需求,。Mamba 優(yōu)于注意力機(jī)制的優(yōu)勢(shì)在推理階段尤其顯著,，因?yàn)槠洳粌H能降低計(jì)算復(fù)雜度，而且內(nèi)存使用量還不會(huì)取決于上下文長(zhǎng)度,。

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

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

混合專(zhuān)家

混合專(zhuān)家（MoE）這類(lèi)技術(shù)能極大提升模型的參數(shù)數(shù)量,，同時(shí)不會(huì)影響模型推理和訓(xùn)練所需的 FLOPs。MoE 最早由 Jacobs et al. 于 1991 年提出,，并在 2017 年由 Shazeer et al. 開(kāi)始用于 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 的核心問(wèn)題，即將 token 分配給專(zhuān)家的過(guò)程,，也稱(chēng)路由（routing）過(guò)程,。目前有兩種基本的路由算法：Token Choice 和 Expert Choice。其中前者是將每個(gè) token 路由到一定數(shù)量（K）的專(zhuān)家,，至于后者則是路由到每個(gè)專(zhuān)家的 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è)專(zhuān)家（K=1）,，而他們使用該方法將 Transformer 的參數(shù)規(guī)模成功擴(kuò)增至了 1.6 萬(wàn)億,。波蘭的這個(gè)團(tuán)隊(duì)在實(shí)驗(yàn)中也采用了這種 MoE 設(shè)計(jì)。

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

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

尤其值得一提的是 Mistral 開(kāi)源的 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）中，見(jiàn)圖 2,。之后,，這個(gè)殘差流的最終值會(huì)被用于預(yù)測(cè)語(yǔ)言建模任務(wù)的下一個(gè) token。

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

不過(guò)該團(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 分開(kāi)了 Mamba 層執(zhí)行的每個(gè) token 的無(wú)條件處理和 MoE 層執(zhí)行的有條件處理,；其中的無(wú)條件處理可高效地將序列的整個(gè)上下文整合到一個(gè)內(nèi)部表征中,，而有條件處理可為每個(gè) token 使用最相關(guān)的專(zhuān)家,。這種將有條件處理和無(wú)條件處理交替起來(lái)的思路在一些基于 MoE 的模型中已經(jīng)得到了應(yīng)用，不過(guò)它們通常是交替基本的和 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è)注意力層加起來(lái)差不多。為了讓 Mamba 和新模型中每個(gè) token 的活動(dòng)參數(shù)數(shù)量大致一樣,，該團(tuán)隊(duì)將每個(gè)專(zhuān)家前向?qū)拥拇笮】s小到了 6dm2,。除了嵌入層和解除嵌入（unembedding）層，所有模型都是每個(gè) token 使用大約 2600 萬(wàn)參數(shù),。訓(xùn)練過(guò)程使用了 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 是否能隨專(zhuān)家數(shù)量的增長(zhǎng)而很好地?cái)U(kuò)展，研究者比較了使用不同數(shù)量專(zhuān)家的模型,。

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

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

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