再戰(zhàn)Transformer！原作者帶隊的Mamba 2來了，新架構(gòu)訓練效率大幅提升
自 2017 年被提出以來，Transformer 已經(jīng)成為 AI 大模型的主流架構(gòu)，一直穩(wěn)居語言建模方面 C 位。
但隨著模型規(guī)模的擴展和需要處理的序列不斷變長，Transformer 的局限性也逐漸凸顯。一個很明顯的缺陷是：Transformer 模型中自注意力機制的計算量會隨著上下文長度的增加呈平方級增長。
幾個月前，Mamba 的出現(xiàn)打破了這一局面，它可以隨上下文長度的增加實現(xiàn)線性擴展。隨著 Mamba 的發(fā)布，這些狀態(tài)空間模型 (SSM) 在中小型規(guī)模上已經(jīng)實現(xiàn)了與 Transformers 匹敵，甚至超越 Transformers。
Mamba 的作者只有兩位，一位是卡內(nèi)基梅隆大學機器學習系助理教授 Albert Gu，另一位是 Together.AI 首席科學家、普林斯頓大學計算機科學助理教授 Tri Dao。
Mamba 面世之后的這段時間里，社區(qū)反應熱烈?？上У氖?，Mamba 的論文卻慘遭 ICLR 拒稿，讓一眾研究者頗感意外。
僅僅六個月后，原作者帶隊，更強大的 Mamba 2 正式發(fā)布了。
- 論文地址：???https://arxiv.org/pdf/2405.21060???
- GitHub 地址：???https://github.com/state-spaces/mamba???
- 論文標題：Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
總體而言，本文提出了 SSD（state space duality）框架，基于此，研究者設(shè)計了一個新的體系架構(gòu) Mamba-2，其核心層是對 Mamba 的選擇性 SSM 的改進，速度提高了 2-8 倍，同時在語言建模方面繼續(xù)與 Transformers 競爭。
Tri Dao 表示，他們構(gòu)建了一個豐富的 SSD 理論框架，許多線性注意力變體和 SSM 是等效的，由此產(chǎn)生的模型 Mamba-2 比 Mamba-1 更好、更快。
Mamba-2 的新算法使其能夠利用更大的狀態(tài)維度 (16 → 256)，同時訓練速度更快。在需要更大狀態(tài)容量的任務上，例如 MQAR 任務，它比 Mamba-1 有了顯著的改進。
此外研究者還發(fā)現(xiàn)，最近新出的混合模型（Jamba、Zamba）增加了一些注意力層來提高模型質(zhì)量?；谶@些發(fā)現(xiàn)，研究者將 4-6 個注意力層與 Mamba-2 層混合，其表現(xiàn)優(yōu)于 Transformer++ 和純 Mamba-2，因而得出注意力和 SSM 是互補的。
這項研究的貢獻概括為：
本文展示了狀態(tài)空間模型與一類稱為半可分矩陣的結(jié)構(gòu)化矩陣族之間的等價性。這一聯(lián)系是 Mamba-2 框架的核心，揭示了狀態(tài)空間模型的新屬性和算法。
本文顯著改進了線性注意力理論，首先通過張量收縮的語言對其循環(huán)形式提供了一個明確的證明，然后將其推廣到一種新的結(jié)構(gòu)化掩碼注意力（SMA）家族。
本文將 SSM（狀態(tài)空間模型）和 SMA（結(jié)構(gòu)化掩碼注意力）聯(lián)系起來，顯示它們有一個很大的交集，彼此是對偶的，同時具有 SSM 式的線性形式和類似注意力的二次方形式。本文還證明了任何具有快速循環(huán)形式的核注意方法都是 SSM。
除了內(nèi)在的理論價值外，研究者所提出的框架為理解和改進序列模型開辟了廣闊的方向。
在算法層面。所提框架為計算 SSM 提供了新的高效且易于實現(xiàn)的算法。本文提出了一種基于半可分離矩陣塊分解的 SSD 算法，該算法利用了 SSM 線性遞推和二次對偶形式，在所有主要效率軸上獲得了最優(yōu)的權(quán)衡。基于 SSD 的實現(xiàn)比 Mamba 的優(yōu)化選擇性掃描實現(xiàn)快 2 到 8 倍，同時允許使用更大的循環(huán)狀態(tài)大?。ㄊ?Mamba 的 8 倍甚至更高，且?guī)缀醪挥绊懰俣龋?。SSD 與優(yōu)化過的 softmax 注意力實現(xiàn)（FlashAttention-2）具有高度競爭力，在序列長度 2k 時性能相當，在序列長度 16K 時速度快 6 倍。
架構(gòu)設(shè)計。采用 SSM 等新架構(gòu)的一個主要障礙是針對 Transformers 量身定制的生態(tài)系統(tǒng)，例如用于大規(guī)模訓練的硬件高效優(yōu)化和并行技術(shù)。本文框架允許使用已建立的慣例和技術(shù)來構(gòu)建 SSM 的架構(gòu)設(shè)計選擇詞匯表，并進一步改進它們。
本文還對 Mamba 塊做了一些修改，這些修改允許實現(xiàn)張量并行，其主要思想包括引入分組值注意力 (GVA，grouped-value attention) 頭結(jié)構(gòu)。
將修改后的并行 Mamba 塊與作為內(nèi)部 SSM 層的 SSD 結(jié)合使用，形成了 Mamba-2 架構(gòu)。研究者在與 Mamba 相同的設(shè)置中研究了 Mamba-2 的 Chinchilla 擴展法則，發(fā)現(xiàn)它在困惑度和實際運行時間方面均優(yōu)于 Mamba 和 Transformer++。研究者還在 Pile 數(shù)據(jù)集上訓練了一系列 Mamba-2 模型，結(jié)果顯示 Mamba-2 在標準下游評估中匹配或超過 Mamba 和開源的 Transformers。例如，在 Pile 上訓練了 3000 億 token 的 2.7B 參數(shù)的 Mamba-2 在性能上超過了在同一數(shù)據(jù)集上訓練的 2.8B 參數(shù)的 Mamba 和 Pythia 以及 6.9B 參數(shù)的 Pythia。
系統(tǒng)優(yōu)化：SSD 框架連接 SSM 和 transformer，允許利用為 transformer 開發(fā)的豐富的系統(tǒng)優(yōu)化工作。
SSD 層
Mamba-2 的核心貢獻是新的 SSD（state space dual）層。SSD 層可以被定義為選擇性 SSM 的特例。與 Mamba 相比，Mamba-2 的改動會略微降低表達能力，但卻顯著提高了訓練效率，特別是允許在現(xiàn)代加速器上使用矩陣乘法單元。
SSD 層的對偶注意力：
除了最新的 SSD 層，研究者也對 Mamba 的神經(jīng)網(wǎng)絡(luò)架構(gòu)做了一些小的改變，Mamba-2 架構(gòu)如下所示。
Mamba-2 在網(wǎng)絡(luò)架構(gòu)上的主要變化是從順序生成變?yōu)椴⑿猩?SSM 參數(shù)，并且 Mamba-2 更適合張量并行等擴展方法。
通過提供狀態(tài)空間模型的顯式矩陣變換形式，研究團隊揭示了理解和使用它們的新方法。從計算的角度來看，任何計算狀態(tài)空間模型前向傳播的方法都可以看作是半可分離矩陣上的矩陣乘法算法。半可分離矩陣視角為 SSD 提供了一個視角，其中雙重模式分別指的是線性時間半可分離矩陣乘法算法和二次時間樸素矩陣乘法。
研究團隊定義了結(jié)構(gòu)化狀態(tài)空間模型和結(jié)構(gòu)化注意力，討論了它們的屬性，并表明它們都有二次算法和線性算法。
自最初的 Mamba 論文研究了合成任務 —— 如：合成復制和歸納 Head 以來，許多后續(xù)工作開始研究更難的關(guān)聯(lián)回憶任務。由 Zoology 和 Based 系列工作引入的 MQAR（multi-query associative recall）任務已成為事實上的標準。
通過運行一個比文獻中通常報告的版本要難得多的任務，該團隊發(fā)現(xiàn) Mamba-2 明顯優(yōu)于 Mamba-1，而改善性能的一個原因是狀態(tài)大小（比 Mamba-1 大約 16 倍）。
在這篇文章中，作者深入探討了模型背后的理論。
從兩個完全不同的角度推導出 SSD 的「對偶性」：
- 一個從 SSM 的角度出發(fā)；
- 另一個從注意力機制的角度出發(fā)。
SSD 框架提供了狀態(tài)空間模型、注意力機制和結(jié)構(gòu)化矩陣之間豐富的聯(lián)系。
雖然 SSD 模型可以被視為框架內(nèi)每個分支的具體實例，但 SSD 框架本身更加通用，為未來的工作開辟了許多方向。
SSD 框架（紅色，藍色）：狀態(tài)空間模型（即半可分矩陣）和結(jié)構(gòu)化掩碼注意力機制包含了大量高效的序列模型。它們的交集是 SSD 模型（紫色）。
SSD 算法
通常，矩陣乘法（matmul）的 FLOPs 速度要比非矩陣乘法 FLOPs 快得多（高達 16 倍）：A100 GPU 具有 312 TFLOPS 的 BF16 矩陣乘法性能，但只有 19 TFLOPS 的 FP32 算術(shù)性能，而 H100 具有 989 TFLOPS 的 BF16 矩陣乘法性能，但只有 67 TFLOPS 的 FP32 算術(shù)性能。
Mamba-2 的主要目標之一是「利用張量核心加速 SSM」。
在綁定參數(shù)并引入 Head 結(jié)構(gòu)后，Mamba-1 中的 SSM 變成了 SSD，這是一種更具限制性的形式，具有類似注意力的公式。并且由于 SSD 連接 SSM 和結(jié)構(gòu)化矩陣，計算 SSM 的高效算法直接對應于「token-mixing」或「sequence-mixing」矩陣 M 的不同分解。
因此，可以通過尋找替代的矩陣乘法方式，例如通過各種方式對其進行分解，從而創(chuàng)建計算 SSM 的新算法。
通過精心選擇塊大小，對這個矩陣進行簡單塊分解，就可以集 SSD 線性遞歸和二次注意力對偶形式的兩種優(yōu)勢于一身。
而這也就是 SSD 算法的起源，它有 4 個步驟，并且對于這個算法有兩種完全不同的詮釋。
SSD 算法：分塊矩陣分解
首先將半可分 SSM 矩陣劃分為大小為 Q×Q 的塊，然后，利用半分矩陣的性質(zhì)來分解每個低秩的非對角塊：
- （橙色）每個對角塊是一個更小的半可分矩陣，可以以喜歡的方式計算這個乘法，特別是使用 SSD 的二次（類似注意力機制）形式。
- （綠色）總共有 T/Q 個不同的綠色塊，通過批處理矩陣乘法來計算。
- （黃色）注意，黃色項本身是一個 1 - 半可分矩陣，這一步等價于對某些修改后的 A 因子的 SSM 掃描。
- （藍色）與綠色類似，通過批處理矩陣乘法來計算。
SSD 算法：分塊和狀態(tài)傳遞
該算法的另一種詮釋涉及「推理 SSM 如何在實際序列上進行操作」。
首先將輸入序列分割成大小為 Q 的塊，步驟可以分為：
- 分塊內(nèi)部輸出：計算每個塊的局部輸出（假設(shè)初始狀態(tài)（對于塊）為 0，則每個塊的輸出是多少？）
- 塊狀態(tài)：計算每個塊的最終狀態(tài)（假設(shè)初始狀態(tài)（對于塊）為 0，則每個塊的最終狀態(tài)是多少？）
- 傳遞狀態(tài)：計算所有塊的最終狀態(tài)的遞歸 - 使用任何所需的算法，例如并行或順序掃描（考慮到所有先前輸入，每個塊的實際最終狀態(tài)是多少？）
- 輸出狀態(tài)：對于每個塊，根據(jù)其真實的初始狀態(tài)（在步驟 3 中計算），僅從初始狀態(tài)得出的輸出計算貢獻
可以看到，大部分算法（步驟 1、2 和 4）利用了矩陣乘法（因此利用了張量核心），而且可以并行計算。
只有步驟 3 需要掃描，但它只操作一個非常短的序列，通常只需要很少時間。
系統(tǒng)及擴展優(yōu)化
張量并行
使用張量并行對 Mamba-1 進行大規(guī)模訓練的一項困難是，每層都需要 2 次 all-reduce，而在 Transformer 中，每個注意力或 MLP 層只需 1 次 all-reduce。這是因為 SSM 的一些參數(shù)是內(nèi)部激活的函數(shù)，而不是層的輸入函數(shù)。在 Mamba-2 中，由于采用了「并行投影」結(jié)構(gòu)，所有 SSM 參數(shù)都是層輸入的函數(shù)，因此可以輕松地將張量并行應用于輸入投影：將輸入投影和輸出投影矩陣分割成 2、4、8 個碎片，具體取決于張量并行度。研究者使用 grouped norm，分組數(shù)除以張量并行度，這樣每個 GPU 都能單獨完成歸一化。這些變化導致每層只需 1 次 all-reduce，而不是 2 次。
序列并行
在對超長序列進行訓練時，可能需要沿著序列長度進行分割，并將不同部分分配給不同的設(shè)備。序列并行主要有兩種形式：對于殘差和歸一化操作，用 reduce-scatter、殘差 + 歸一化、然后 all-gather，取代張量并行中的 all-reduce。由于 Mamba-2 使用與 Transformer 相同的殘差和歸一化結(jié)構(gòu)，因此這種形式的序列并行無需修改即可直接應用。對于注意力或 SSM 操作，又稱上下文并行（CP）。對于注意力，可以使用環(huán)形注意力沿序列維度進行分割。對于 Mamba-2，SSD 框架再次提供了幫助：使用相同的蒯分解，可以讓每個 GPU 計算其本地輸出和最終狀態(tài)，然后在更新每個 GPU 的最終輸出之前，在 GPU 之間傳遞狀態(tài)（使用發(fā)送 / 接收通信原語）。
實驗結(jié)果
該研究在 MQAR 的一種具有挑戰(zhàn)性的版本上，使用更難的任務、更長的序列和更小的模型進行了對比實驗?；€包括標準的多頭 softmax 注意力以及 Based 架構(gòu)，實驗結(jié)果如圖 8 所示。
下表顯示了 Mamba-2 在一系列下游零樣本評估任務上的性能：
感興趣的讀者可以閱讀論文原文，了解更多研究內(nèi)容。
本文轉(zhuǎn)自 機器之心 ，作者：機器之心
