論文標題

《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》

作者

Albert Gu 和 Tri Dao

初讀

摘要

• Transformer 架構及其核心注意力模塊

  • 地位：目前深度學習領域普遍的基礎模型。

  • 為了解決 Transformers 在長序列上的計算效率低下的問題，已經開發了許多次二次時間架構（subquadratic-time architectures），如：

    • 線性注意力

    • 門控卷積

    • 遞歸模型

    • 結構化狀態空間模型（SSM）

  • 仍存在的問題及原因：

    • 它們在語言等重要模態上的表現不如關注

    • 一個關鍵弱點是它們無法執行基于內容的推理。

• 本文改進：

  • 讓 SSM 參數成為輸入的函數

    • 能解決離散模態的弱點，

    • 使模型能夠根據當前標記有選擇地沿序列長度維度傳播或遺忘信息。

  • 設計了一種硬件感知的并行遞歸模式算法：

    • 為了解決上述變化對高效卷積使用的阻礙，
    • 將這些選擇性 SSM 集成到一個簡化的端到端神經網絡架構中，該架構不含注意力，甚至不含 MLP 塊（Mamba）。

• Mamba 的特點：

  • 快速推理（吞吐量比 Transformers 高 5 倍）
  • 序列長度線性伸縮，其性能在實際數據中可提高到百萬長度序列。

• Mamba 的性能：

  • 作為通用序列模型的骨干，Mamba 在語言、音頻和基因組學等多種模式中都達到了最先進的性能。
  • 在語言建模方面，無論是預訓練還是下游評估，Mamba-3B 模型都優于同等規模的 Transformers，并能與兩倍于其規模的 Transformers 相媲美。

結論

• 我們為結構化狀態空間模型引入了一種選擇機制，使其能夠執行與上下文相關的推理，同時按序列長度線性擴展。
• 當把 Mamba 納入一個簡單的無注意力架構時，它在一系列不同的領域取得了最先進的結果，其性能達到或超過了強 Transformer 模型。
• 我們對選擇性狀態空間模型在不同領域建立基礎模型的廣泛應用感到興奮，尤其是在基因組學、音頻和視頻等需要長語境的新興模式中。我們的研究結果表明，Mamba 是通用序列模型骨干的有力候選者。

再讀

Section 1 Introduction

• 第一段：基于 Transformer 及其核心注意力層的基礎模型

  • 現代機器學習的有效范式：

    • 基礎模型（Foundation models，FM），即在海量數據上進行預訓練，然后針對下游任務進行調整的大型模型。

    • 這些基礎模型的骨干通常是序列模型，可在語言、圖像、語音、音頻、時間序列和基因組學等各種領域的任意輸入序列上運行。

  • 現代 FM 主要基于一種單一類型的序列模型：Transformer 及其核心注意力層。

    • 優點：自注意力的功效歸功于其在上下文窗口中密集路由信息的能力，使其能夠對復雜數據進行建模。
    • 缺陷：無法對有限窗口之外的任何事物進行建模，以及與窗口長度相關的二次縮放。

  • 當前進展：

    • 為了克服這些缺點，人們對注意力的更有效變體進行了大量研究，但這些研究往往以犧牲注意力的有效特性為代價。
    • 到目前為止，這些變體還沒有一個在跨領域的規模化應用中顯示出經驗上的有效性。

• 第二段：結構化狀態空間序列模型（SSMs）

  • 特點：

    • 已成為序列建模領域的一類有前景的架構。
    • 靈感來自經典狀態空間模型。
    • 這些模型可以被解釋為循環神經網絡（RNN）和卷積神經網絡（CNN）的組合，
    • 這類模型可以非常有效地進行遞歸或卷積計算，序列長度呈線性或近線性縮放。

  • 優點：

    在某些數據形式中具有建模長程依賴性的原理機制，并主導了諸如長程競技場等基準測試。許多 SSMs 在涉及連續信號數據（如音頻和視覺）的領域取得了成功。

  • 缺陷：

    在建模離散和信息密集數據（如文本）方面效果不佳。

• 本文提出了一類新的選擇性狀態空間模型，該模型在多個軸上改進了先前的工作，從而在序列長度線性縮放的同時實現了 Transformer 的建模能力。

  • 選擇機制

    • 先前模型的一個關鍵局限：以依賴輸入的方式有效選擇數據的能力（即關注或忽略特定輸入）。
    • 基于選擇性復制和感應頭等重要合成任務的直覺，本文設計了一種簡單的選擇機制，根據輸入對 SSM 參數進行參數化。這樣，模型就能過濾掉無關信息，并無限期地記住相關信息。

  • 硬件感知算法

    • 上述簡單的變化給模型的計算帶來了技術挑戰；事實上，所有先前的 SSM 模型都必須是時間和輸入不變的，這樣才能提高計算效率。
    • 本文采用了一種硬件感知算法來克服這一難題，該算法通過掃描而不是卷積來計算模型，但不會將擴展狀態具體化，以避免在 GPU 存儲器層次結構的不同級別之間進行 I/O 訪問。
    • 由此產生的實現方法無論是在理論上（與所有基于卷積的 SSM 的偽線性相比，隨序列長度線性擴展）還是在現代硬件上（在 A100 GPU 上可快達 3 倍）都比以前的方法更快。

  • 體系結構

    • 我們將先前的 SSM 架構設計與 Transformers 的 MLP 模塊合并為一個模塊，從而簡化了先前的深度序列模型架構，形成了一種包含選擇性狀態空間的簡單、同質的架構設計（Mamba）。
    • 選擇性 SSM 以及 Mamba 架構是完全遞歸模型，其關鍵特性使其適合作為在序列上運行的通用基礎模型的骨干。
      • 高質量：選擇性為語言和基因組學等密集模式帶來了強大的性能。
      • 快速訓練和推理：在訓練過程中，計算和內存與序列長度成線性關系，而在推理過程中，由于不需要緩存以前的元素，因此自回歸展開模型每一步只需要恒定的時間。
      • 長上下文：質量和效率的共同作用使實際數據的性能提高到序列長度 100 萬。

• Mamba 的性能：

  本文通過經驗驗證了 Mamba 作為通用序列 FM 骨干的潛力，無論是在預訓練質量還是特定領域的任務性能方面，它都能在幾種模式和環境中發揮作用：

  • 合成：

    在一些重要的合成任務上，如被認為是大型語言模型關鍵的復制和歸納頭，Mamba 不僅能輕松解決，還能推導出無限長（>100 萬個詞組）的解決方案。

  • 音頻和基因組學：

    在音頻波形和 DNA 序列建模方面，Mamba 在預訓練質量和下游指標（例如，在具有挑戰性的語音生成數據集上，FID 降低了一半以上）方面都優于 SaShiMi、Hyena 和 Transformers 等先前的一流模型。在這兩種情況下，它的性能都隨著上下文長度的增加而提高，最高可達百萬長度的序列。

  • 語言建模：

    Mamba 是第一個線性時間序列模型，無論是在預訓練復雜度還是下游評估方面，都真正達到了 Transformer 質量的性能。通過多達 1B 個參數的縮放規律，我們發現 Mamba 的性能超過了大量基線模型，包括基于 LLaMa 的非常強大的現代變換器訓練配方。我們的 Mamba 語言模型與類似規模的 Transformer 相比，具有 5 倍的生成吞吐量，而且 Mamba-3B 的質量與兩倍于其規模的 Transformer 相當（例如，與 Pythia-3B 相比，常識推理的平均值高出 4 分，甚至超過 Pythia-7B）。

Section 2 State Space Models

狀態空間模型

結構化狀態空間序列模型（Structured state space sequence models，S4）是最近一類用于深度學習的序列模型，與 RNN、CNN 和經典狀態空間模型廣泛相關。它們受到一個特定連續系統 (1) 的啟發，該系統通過一個隱含的潛在狀態 $h(t)\in {\mathbb{R}}^N$ 映射一個一維函數或序列 $x(t)\in \mathbb{R}\to y(t)\in \mathbb{R}$。

具體來說，S4 模型由四個參數 $(\Delta ,A,B,C)$ 定義，它們分兩個階段定義序列到序列的轉換
$$\begin{array}{rl}{h}^{\prime }(t)& =Ah(t)+Bx(t)(1a)\\ y(t)& =Ch(t)(1b)\end{array}\text{\hspace{1em}?}$$

$$\begin{array}{rl}{h}_t& =\overline{A}{h}_{t?1}+\overline{B}{x}_t(2a)\\ y_t& =C{h}_t(2b)\end{array}\text{\hspace{1em}?}$$

$$\begin{array}{rl}\overline{K}& =(C\overline{B},C\overline{AB},\dots ,C\overline{AkB},\dots )(1a)\\ y& =x?\overline{K}(1b)\end{array}\text{\hspace{1em}?}$$

• 離散化

  第一階段通過固定公式 $\overline{A}=f_A(\Delta ,A)$ 和 $\overline{B}=f_B(\Delta ,A,B)$ 將 “連續參數” $(\Delta ,A,B)$ 轉換為 “離散參數” $(\overline{A},\overline{B})$ ，其中一對 $(f_A,f_B)$ 稱為離散化規則。可以使用各種規則，例如等式 (4) 中定義的零階保持 (ZOH)。
  $$\overline{A}=\exp (\Delta A)\overline{B}=(\Delta A{)}^{?1}(\exp (\Delta A)?I)?\Delta B$$
  離散化與連續時間系統有很深的聯系，可以賦予連續時間系統更多的特性，如分辨率不變性和自動確保模型正確歸一化。它還與 RNN 的門控機制有關，我們將在第 3.5 節中再次討論。不過，從機械的角度來看，離散化可以簡單地看作是 SSM 前向傳遞中計算圖的第一步。另一種 SSM 可以繞過離散化步驟，直接對 $(\overline{A},\overline{B})$ 進行參數化，這可能更容易推理。

• 計算

  在參數從 $(\Delta ,A,B,C)\to (\overline{A},\overline{B},C)$ 轉化之后，模型可以通過兩種方式計算，一種是線性遞推（2），另一種是全局卷積（3）。

  通常，該模型使用卷積模式 (3) 進行高效的并行訓練（提前看到整個輸入序列），并切換到遞歸模式 (2) 進行高效的自回歸推理（每次看到一個時間步的輸入）。

• 線性時間不變性（LTI）

  方程 (1) 至 (3) 的一個重要特性是，模型的動態在時間上是恒定的。換句話說， $(\Delta ,A,B,C)$ 以及 $(\overline{A},\overline{B})$ 在所有時間步中都是固定不變的。這一特性被稱為線性時間不變性（LTI），它與遞推和卷積有著深刻的聯系。非正式地，我們認為 LTI SSM 等同于任何線性遞推（2a）或卷積（3b），并用 LTI 作為這些模型類別的總稱。

  迄今為止，所有結構化 SSM 都是 LTI 模型（例如，以卷積方式計算），因為存在基本的效率限制，這將在第 3.3 節中討論。然而，這項工作的一個核心觀點是，LTI 模型在對某些類型的數據進行建模時具有根本性的局限性，我們的技術貢獻在于消除 LTI 限制，同時克服效率瓶頸。

• 結構和尺寸

  結構化 SSM 之所以被命名為結構化 SSM，是因為高效計算 SSM 還需要對 A 矩陣施加結構。最常用的結構形式是對角線結構，我們也使用這種結構。

  在這種情況下，$A\in {\mathbb{R}}^{N\times N}$，$B\in {\mathbb{R}}^{N\times 1}$，$C\in {\mathbb{R}}^{1\times N}$ 矩陣都可以用數字 $N$ 來表示。要對批量大小為 $B$、長度為 $L$ 、且具有 $D$ 個通道的輸入序列 $x$ 進行操作，則 SSM 將獨立應用于每個通道。請注意，在這種情況下，每個輸入的總隱藏狀態維度為 $DN$，在序列長度上計算它需要 $O(BLDN)$ 時間和內存；這就是第 3.3 節中提到的基本效率瓶頸的根源。

• 一般狀態空間模型

  "狀態空間模型 "一詞的含義非常廣泛，它簡單地代表了任何具有潛在狀態的循環過程的概念。它被用來指代不同學科中的許多不同概念，包括馬爾可夫決策過程（MDP）（強化學習）、動態因果建模（DCM）（計算神經科學）、卡爾曼濾波器（控制）、隱馬爾可夫模型（HMM）和線性動力系統（LDS）（機器學習），以及大量的遞歸（有時是卷積）模型（深度學習）。

  本文使用 "SSM "一詞專指結構化 SSM 或 S4 模型，這些術語可以互換使用。為方便起見，我們還可以包括這些模型的衍生模型，如那些側重于線性遞推或全局演化觀點的模型，并在必要時澄清細微差別。

• SSM 架構

  SSM 是獨立的序列變換，可被納入端到端神經網絡架構。(有時也稱 SSM 架構為 SSNN，它與 SSM 層的關系就像 CNN 與線性卷積層的關系）。我們將討論一些最著名的 SSM 架構，其中許多也將作為我們的主要基線。

  • 線性注意是自我注意的一種近似方法，它涉及一種遞歸，可被視為一種退化的線性 SSM。
  • H3 將這種遞歸概括為使用 S4；它可以被看作是一種由兩個門控連接夾著一個 SSM 的架構（圖 3）。H3 還在主 SSM 層之前插入了一個標準的局部卷積，并將其定格為移位 SSM。
  • Hyena 使用與 H3 相同的架構，但用 MLP 參數化全局卷積取代了 S4 層。
  • RetNet 在架構中增加了一個額外的門，并使用更簡單的 SSM，允許另一種可并行計算的路徑，使用多頭注意力（MHA）的變體來代替卷積。
  • RWKV 是最近基于另一種線性注意近似（無注意轉換器）設計用于語言建模的 RNN。它的主要 "WKV "機制涉及 LTI 循環，可視為兩個 SSM 的比率。

  其他密切相關的 SSM 和架構將在擴展的相關工作（附錄 B）中進一步討論。我們特別強調 S5、QRNN 和 SRU，我們認為它們是與我們的核心選擇性 SSM 關系最密切的方法。

Section 3 Selective State Space Models

選擇性狀態空間模型

• 第 3.1 節：利用合成任務的直覺來激發我們的選擇機制，
• 第 3.2 節：解釋如何將這一機制納入狀態空間模型。
• 第 3.3 節：由此產生的時變 SSM 不能使用卷積，這就提出了如何高效計算的技術難題。本文采用一種硬件感知算法，利用現代硬件的內存層次結構來克服這一難題。
• 第 3.4 節：描述了一個簡單的 SSM 架構，它不需要注意力，甚至不需要 MLP 塊。
• 第 3.5 節：討論選擇機制的一些其他特性。

3.1 Motivation: Selection as a Means of Compression

動機：選擇作為一種壓縮手段

• 序列建模的一個基本問題：將上下文壓縮到更小的狀態中。

• 從這個基本問題的角度看待流行的序列模型的取舍：

  • 注意力既有效又低效，因為它根本沒有明確壓縮上下文。這一點可以從自回歸推理需要明確存儲整個上下文（即 KV 緩存）這一事實中看出，這直接導致了 Transformers 緩慢的線性時間推理和二次時間訓練。
  • 遞歸模型的效率很高，因為它們的狀態是有限的，這意味著恒定時間推理和線性時間訓練。然而，它們的效率受限于這種狀態對上下文的壓縮程度。

• 為了理解這一原理，重點討論兩個合成任務的運行示例：

  

  • （左）復制任務的標準版本涉及輸入和輸出元素之間的恒定間距，可通過線性遞歸和全局卷積等時變模型輕松解決。
  • (右上）選擇性復制任務的輸入間距是隨機的，需要時變模型能夠根據輸入內容選擇性地記住或忽略輸入。
  • (右下）"誘導頭 "任務是聯想記憶的一個例子，要求根據上下文檢索答案，這是 LLM 的一項關鍵能力。
  • 選擇性復制任務修改了流行的復制任務，改變了要記憶的標記的位置。它要求內容感知推理能夠記憶相關標記（彩色）并過濾掉不相關的標記（白色）。
  • 誘導頭任務是一種眾所周知的機制，據推測，它可以解釋大多數 LLMs 的情境學習能力。該任務要求進行情境感知推理，以了解何時在適當的情境中產生正確的輸出（黑色）。

• LTI 的缺點：

  這些任務揭示了 LTI 模型的失敗模式。從遞歸的角度來看，它們的恒定動態（例如（2）中的 $(\overline{A},\overline{B})$ 轉換）無法讓它們從上下文中選擇正確的信息，也無法以依賴輸入的方式影響沿序列傳遞的隱藏狀態。從卷積的角度來看，眾所周知，全局卷積可以解決初始復制任務（ the vanilla Copying task），因為它只需要時間感知，但由于缺乏內容感知，全局卷積很難解決選擇性復制任務（the Selective Copying task）（圖 2）。更具體地說，輸入到輸出之間的間距是變化的，無法用靜態卷積核建模。

• 總結：

  序列模型的效率與效果權衡的特點在于它們如何很好地壓縮其狀態：**高效的模型必須有一個較小的狀態，而有效的模型必須有一個包含上下文所有必要信息的狀態。**反過來，我們提出建立序列模型的一個基本原則是選擇性：或者說是將輸入集中或過濾到序列狀態中的上下文感知能力。特別是，選擇機制可以控制信息如何沿著序列維度傳播或交互（更多討論見第 3.5 節）。

3.2 Improving SSMs with Selection

通過選擇改進SSM

• 核心思想：讓影響序列交互的參數（如 RNN 的遞歸動力學或 CNN 的卷積核）取決于輸入，從而實現將選擇機制納入模型。

算法 1 和 2 展示了我們使用的主要選擇機制。主要區別在于，我們只需將幾個參數 $\Delta 、B、C$ 變為輸入的函數，同時改變張量的形狀。我們特別強調，這些參數現在都有一個長度維度 $L$，這意味著模型已經從時間不變變為時間可變。(請注意，形狀注釋已在第 2 節中進行了描述）。這就失去了與卷積（3）的等價性，對其效率產生了影響，將在下文討論。

我們特別選擇 $s_B(x)={\mathrm{Linear}}_N(x)$、$s_C(x)={\mathrm{Linear}}_N(x)$、$s_{\Delta }(x)={\mathrm{Broadcast}}_D({\mathrm{Linear}}_1(x))$ 和 ${\tau }_{\Delta }(x)=\mathrm{softplus}$，其中 ${\mathrm{Linear}}_d$ 是維度 $d$ 的參數化投影。選擇 $s_{\Delta }$ 和 ${\tau }_{\Delta }$ 是因為第 3.5 節中解釋的 RNN 門控機制。

3.3 Ecient Implementation of Selective SSMs

高效實施選擇性 SSM

諸如卷積和 Transformer 等硬件友好型架構得到了廣泛應用。在這里，我們的目標是讓選擇性 SSM 在現代硬件（GPU）上也能高效運行。選擇機制是非常自然的，早期的研究曾試圖將選擇的特殊情況納入其中，例如讓 $\Delta$ 在循環 SSM 中隨時間變化。然而，如前所述，使用 SSM 的一個核心限制是其計算效率，這也是 S4 及其所有衍生產品使用 LTI（非選擇性）模型（最常見的形式是全局卷積）的原因。

3.3.1 Motivation of Prior Models

先行模型的動機

• 在高層次上，SSM 等循環模型總是在表現力和速度之間進行權衡：如第 3.1 節所述，隱藏狀態維度越大的模型越有效，但速度越慢。因此 我們希望在不付出速度和內存代價的情況下最大化隱藏狀態維度。
• 請注意，循環模式比卷積模式更靈活，因為后者（3）是由前者（2）擴展而來的。然而，這需要計算和物化形狀為 $(B,L,D,N)$ 的隱藏狀態矩陣 $h$，該矩陣比輸入 $x$ 和形狀為 $(B,L,D)$ 的輸出矩陣 $y$ 要大得多（大 $N$ 倍，即 SSM 狀態維度）。因此，我們引入了更高效的卷積模式，該模式可以繞過狀態計算，實現僅為 $(B,L,D)$ 的卷積核 (3a)。
• 先前的 LTI SSM 利用雙遞歸卷積形式，將有效狀態維度提高了 $N(\approx 10?100)$，遠大于傳統的 RNN，而且不會降低效率。

3.3.2 Overview of Selective Scan: Hardware-Aware State Expansion

選擇性掃描概述： 硬件感知狀態擴展

選擇機制旨在克服 LTI 模型的局限性，因此我們需要重新審視 SSM 的計算問題。我們用三種經典技術來解決這個問題：內核融合（kernel fusion）、并行掃描（parallel scan）和重新計算（recomputation）。我們提出了兩個主要觀點：

• 初級遞歸計算使用了 $O(BLDN)$ FLOPs，而卷積計算使用了 $O(BLD\log (L))$ 而前者的常數因子更低。因此，對于長序列和不太大的狀態維度 $N$，遞歸模式使用的 FLOP 實際上更少。
• 這兩個挑戰是遞歸的順序性和大量的內存使用。為了解決后者，就像卷積模式一樣，我們可以嘗試不實際實現完整的狀態 $h$。

其主要思路是利用現代加速器（GPU）的特性，僅在更高效的內存層次結構中實現狀態 “$h$”。特別是，大多數操作（矩陣乘法除外）都受到內存帶寬的限制。這包括我們的掃描操作，我們使用內核融合來減少內存 I/O 的數量，從而比標準實現顯著提速。

• 具體來說，我們不在 GPU HBM（高帶寬存儲器）中準備大小為 $(B,L,D,N)$ 的掃描輸入 $(\overline{A},\overline{B})$ ，而是直接從慢速 HBM 將 SSM 參數 $(\Delta ,A,B,C)$ 直接從慢速 HBM 加載到快速 SRAM，在 SRAM 中執行離散化和遞歸，然后將大小為 $(B,L,D)$ 的最終輸出寫回 HBM。
• 為了避免順序遞歸，我們注意到，盡管它不是線性的，但仍然可以用一種工作效率高的并行掃描算法將其并行化。
• 最后，我們還必須避免保存反向傳播所必需的中間狀態。我們謹慎地應用了經典的重新計算技術來降低內存需求：當輸入從 HBM 加載到 SRAM 時，中間狀態不會被存儲，而是在后向傳遞中重新計算。因此，融合選擇性掃描層的內存需求與使用 FlashAttention 的優化 transformer 實現相同。

有關融合內核和重新計算的詳情見附錄 D。圖 1 展示了完整的選擇性 SSM 層和算法。

圖片概述：結構化 SSM 通過更高維的潛在狀態 $h$（例如 $N=4$），將輸入 $x$ 的每個通道（例如 $D=5$）獨立映射到輸出 $y$。先前的 SSM 通過巧妙的交替計算路徑避免了這一大型有效狀態（ $DN$，乘以批量大小 $B$ 和序列長度 $L$）的具體化，這些路徑需要時間不變性： $(\Delta ,A,B,C)$ 參數在整個時間內保持不變。我們的選擇機制又增加了依賴于輸入的動態，這也需要一種謹慎的硬件感知算法，以便只在 GPU 存儲器層次結構中更高效的層級中實現擴展狀態。

3.4 A Simplied SSM Architecture

簡化的 SSM 架構

與結構化 SSM 一樣，選擇性 SSM 也是獨立的序列變換，可以靈活地融入神經網絡。H3 架構是最著名的 SSM 架構（第 2 節）的基礎，一般由線性注意啟發的區塊和 MLP（多層感知器）區塊交錯組成。我們簡化了這一架構，將這兩個部分合二為一，均勻堆疊（圖 3）。這是受門控注意力單元（GAU）的啟發，該單元也對注意力做了類似的處理。

這種結構是通過一個可控的擴展因子 $E$ 來擴展模型維度 $D$ 。對于每個區塊，大部分參數（$3E{D}^2$）都在線性投影中（輸入投影為 $2E{D}^2$，輸出投影為 $E{D}^2$），而內部 SSM 的貢獻較小。相比之下，SSM 參數（$\Delta$、$B$、$C$ 的投影和矩陣 $A$）的數量要少得多。我們重復這一模塊，并交錯使用標準歸一化和殘差連接，形成 Mamba 架構。在實驗中，我們始終將 $E=2$ 固定下來，并使用該模塊的兩個堆棧來匹配 Transformer 交錯 MHA（多頭注意力）和 MLP 模塊的 $12D^2$ 參數。我們使用 SiLU/Swish 激活函數，從而使門控 MLP 成為流行的 "SwiGLU "變體。最后，我們還使用了一個可選的歸一化層（我們選擇了 LayerNorm），其動機是 RetNet 在類似位置使用了歸一化層。

圖片概述：我們的簡化塊設計結合了 H3 塊（大多數 SSM 架構的基礎）和現代神經網絡中無處不在的 MLP 塊。我們只是簡單地重復 Mamba 模塊，而不是交錯使用這兩個模塊。與 H3 模塊相比，Mamba 用激活函數取代了第一個乘法門。與 MLP 程序塊相比，Mamba 在主分支上增加了一個 SSM。對于 $\sigma$，我們使用 SiLU / Swish 激活。

3.5 Properties of Selection Mechanisms

選擇機制的特性

選擇機制是一個更寬泛的概念，可以以不同的方式應用，例如應用于更傳統的 RNN 或 CNN，應用于不同的參數（例如算法 2 中的 $A$），或使用不同的變換 $s(x)$。

3.5.1 Connection to Gating Mechanisms

與門控機制的聯系

我們強調最重要的聯系：RNN 的經典門控機制是我們 SSM 選擇機制的一個實例。我們注意到，RNN 門控與連續時間系統離散化之間的聯系已經得到證實。事實上，定理 1 是 Gu、Johnson、Goel 等人對 ZOH 離散化和輸入依賴門的改進（證明見附錄 C）。更廣義地說，SSM 中的 $\Delta$ 可以看作是 RNN 門控機制的一種廣義作用。與之前的研究一致，我們認為 SSM 的離散化是啟發式門控機制的原則基礎。

定理 1. 當 $N=1$，$A=?1$，$B=1$，$s_{\Delta }=\mathrm{Linear}(x)$ ，且 ${\tau }_{\Delta }=\mathrm{softplus}$ 時，選擇性 SSM 遞推算法（算法 2）的形式為：
$$\begin{array}{rl}{g}_t& =\sigma (\mathrm{Linear}(x_t))\\ h_t& =(1?g_t)h_{t?1}+g_t{x}_t\end{array}\text{\hspace{1em}?}$$
如第 3.2 節所述，我們對 $s_{\Delta }$ 和 ${\tau }_{\Delta }$ 的具體選擇就是基于這種聯系。特別要注意的是，如果給定的輸入 $x_t$ 應該被完全忽略（這在合成任務中是必要的），那么所有 $D$ 通道都應該忽略它，因此我們在用 $\Delta$ 重復/廣播之前，會將輸入投影到 1 維。

3.5.2 Interpretation of Selection Mechanisms

選擇機制的解釋

• 可變間距

  選擇性可以過濾掉相關輸入之間可能出現的無關噪聲標記。選擇性復制任務（the Selective Copying task）就是一個例子，但在常見的數據模式中，尤其是離散數據中，例如存在 "嗯 "等語言填充物時，也會出現這種情況。這一特性的產生是因為模型可以機械地過濾掉任何特定的輸入 $x_t$，例如在門控 RNN 的情況下（定理 1），當 $g_t\to 0$ 時。

• 篩選上下文

  根據經驗觀察，許多序列模型的性能并沒有隨著上下文時間的延長而提高，盡管從嚴格意義上講，更多的上下文應該會帶來更好的性能。一種解釋是，許多序列模型在必要時無法有效忽略無關上下文；一個直觀的例子是全局卷積（和一般的 LTI 模型）。另一方面，選擇性模型可以在任何時候簡單地重置它們的狀態以去除無關的歷史，因此它們的性能原則上隨著上下文長度的增加而單調提高。

• 邊界重置

  在多個獨立序列拼接在一起的情況下，Transformer 可以通過實例化特定的注意力掩碼將它們分開，而 LTI 模型則會在序列之間滲入信息。選擇性 SSM 還可以在邊界重置狀態（例如，當 ${\Delta }_t\to \infty$ 或定理 1 $g_t\to 1$ 時）。這些設置可能是人為的（例如將文檔打包在一起以提高硬件利用率），也可能是自然的（例如強化學習中的情節邊界）。

  此外，我們還闡述了每個選擇性參數的影響。

• 對 $\Delta$ 的解釋。

  一般來說，$\Delta$ 控制著對當前輸入 $x_t$ 的關注或忽略程度。它概括了 RNN 邏輯門（例如定理 1 中的 $g_t$），從機制上講，大的 $\Delta$ 會重置狀態 $h$ 并關注當前輸入 $x$，而小的 $\Delta$ 則會保持狀態并忽略當前輸入。SSM (1)-(2) 可以解釋為一個被時間步長 $\Delta$ 離散化的連續系統，直觀地說，大 $\Delta \to \infty$ 表示系統更長時間地關注當前輸入（從而 "選擇 "當前輸入并忘記當前狀態），而小 $\Delta \to 0$ 則表示忽略瞬時輸入。

• 對 $A$ 的解釋

  我們注意到，雖然 $A$ 參數也可能是選擇性的，但它最終只是通過 $\overline{A}=\exp (\Delta A)$ （離散化 (4)）與 $\Delta$ 的相互作用來影響模型。因此，$\Delta$ 的選擇性足以確保 $(\overline{A},\overline{B})$ 的選擇性，這也是改進的主要原因。我們假設，除了 $\Delta$ 之外（或代替 $\Delta$），讓 $A$ 也具有選擇性也會有類似的效果，為了簡單起見，我們不考慮它。

• 對 $B$ 和 $C$ 的解釋

  如第 3.1 節所述，選擇性最重要的特性是過濾掉無關信息，從而將序列模型的上下文壓縮成一個有效的狀態。在 SSM 中，修改 $B$ 和 $C$ 使其具有選擇性，可以更精細地控制是讓輸入 $x_t$ 進入狀態 $t$，還是讓狀態進入輸出 $y_t$。這可以解釋為，模型可以分別根據內容（輸入）和情境（隱藏狀態）來調節循環動態。

3.6 Additional Model Details

其他模型細節

• 實數與復數

  大多數先前的 SSM 在其狀態迭代中使用復數，這在許多任務中都是提高性能的必要條件。然而，根據經驗觀察，完全實值的 SSM 似乎也能正常工作，在某些情況下甚至可能更好。我們將實值作為默認值，除一項任務外，其他任務都能很好地使用實值；我們假設，復數與實值的權衡與數據模態的連續-離散光譜有關，復數有助于連續模態（如音頻、視頻），但對離散模態（如文本、DNA）沒有幫助。

• 初始化

  大多數先前的 SSM 還建議進行特殊的初始化，尤其是在復值情況下，這在低數據量等多種情況下都有幫助。我們對復值情況的默認初始化是 S4D-Lin，對實值情況的默認初始化是 S4D-Real，它基于 HIPPO 理論。它們將 $A$ 的 $n$-th 元素分別定義為 $?1/2+ni$ 和 $?(n+1)$。不過，我們預計許多初始化都能順利進行，尤其是在大數據和實值 SSM 環境中；第 4.6 節將考慮一些消融問題。

• $\Delta$ 的參數化

  我們將 $\Delta$ 的選擇性調整定義為 $s_{\Delta }(x)={\mathrm{Broadcast}}_D({\mathrm{Linear}}_1(x))$，這是由 $\Delta$ 的力學原理（第 3.5 節）激發的。） 我們發現它可以從維度 1 推廣到更大的維度 $R$。我們將其設定為 $D$ 的一小部分，與區塊中的主要線性投影相比，使用的參數數量可以忽略不計。我們還注意到，廣播操作可以被看作是另一種線性投影，初始化為 1 和 0 的特定模式； 如果這個投影是可訓練的，那么就會產生另一種 $s_{\Delta }(x)={\mathrm{Linear}}_D({\mathrm{Linear}}_1(x))$，這可以看作是一種低階投影。

  在我們的實驗中，根據先前關于 SSM 的研究，將 $\Delta$ 參數（可視為偏置項）初始化為 ${\tau }_{\Delta }^{?1}(\mathrm{Uniform}(0.001,0.1))$。

  備注 3.1. 在我們的實驗結果中，為了簡潔起見，我們有時把選擇性 SSM 簡寫為 S6 模型，因為它們是帶有選擇機制的 S4 模型，是用掃描計算的。

Section 4 Empirical Evaluation

經驗評估

• 第 4.1 節中，我們測試了 Mamba 解決第 3.1 節中提出的兩個合成任務的能力。然后，我們對三個領域進行了評估，每個領域都對自回歸預訓練和下游任務進行了評估。
• 第 4.2 節：語言模型預訓練（縮放規律）和零點下游評估。
• 第 4.3 節：DNA 序列預訓練 DNA 序列預訓練和長序列分類任務的微調。
• 第 4.4 節：音頻波形預訓練和自回歸生成語音片段的質量。
• 第 4.5 節展示了 Mamba 在訓練和推理時的計算效率，第 4.6 節介紹了架構的各個組件和選擇性 SSM。

4.1 Synthetic Tasks

合成任務

4.1.1 Selective Copying

選擇性復制

復制任務是序列建模中研究最深入的合成任務之一，最初是為了測試遞歸模型的記憶能力而設計的。正如第 3.1 節所討論的，LTI SSM（線性遞歸和全局卷積）只需跟蹤時間而不是推理數據，就能輕松解決這個任務；例如，通過構建一個長度完全正確的卷積核（圖 2）。這一點在早期的全局卷積研究中得到了明確驗證。選擇性復制任務通過隨機化標記間的間距來防止這種捷徑。請注意，這項任務之前已作為去噪任務引入。

需要注意的是，許多前人的研究認為，增加架構門控（乘法交互）可以賦予模型以 "數據依賴性 "并解決相關任務。然而，我們發現這種解釋不夠直觀，因為這種門控并不沿著序列軸相互作用，也不能影響標記之間的間距。尤其是架構門控并不是選擇機制的實例（附錄 A）。

表 1 證實，H3 和 Mamba 等門控架構只能部分提高性能，而選擇機制（將 S4 修改為 S6）可以輕松解決這一問題，尤其是在與這些功能更強大的架構相結合時。

4.1.2 Induction Heads

感應頭

誘導頭是機械可解釋性視角中的一項簡單任務，它對 LLM 的語境學習能力具有驚人的預測能力。它要求模型進行聯想回憶和復制：例如，如果模型在序列中見過 "哈利-波特 "這樣的大構詞，那么下一次 "哈利 "出現在同一序列中時，模型應該能夠通過復制歷史來預測 “波特”。

• 數據集

  我們在序列長度為 256、詞匯量為 16 的誘導頭任務中訓練了一個雙層模型，這與之前關于該任務的研究相當，但是序列更長。此外，我們還通過在測試時對從 $2^6=64$ 到 $2^{20}=1048576$ 的一系列序列長度進行評估，研究了泛化和外推能力。

• 模型

  根據有關誘導頭的既有研究，我們使用 2 層模型，使注意力能夠機械地解決誘導頭任務。我們測試了多頭注意力（8 個頭，有不同的位置編碼）和 SSM 變體。我們對 Mamba 使用的模型維度 $D$ 為 64，對其他模型使用的模型維度 $D$ 為 128。

• 結果

  表 2 顯示，Mamba–或者更準確地說，它的選擇性 SSM 層–有能力完美地解決任務，因為它能夠選擇性地記住相關標記，而忽略中間的所有其他標記。它能完美地泛化到百萬長度的序列，即比訓練時看到的長度長 4000 倍，而其他方法都不能超過 2 倍。

在注意力模型的位置編碼變體中，xPos（專為長度外推法設計）略優于其他變體；還需注意的是，由于內存限制，所有注意力模型只測試到序列長度 214 = 16384。在其他 SSM 中，H3 和 Hyena 相似，這與 Poli 等人（2023 年）的研究結果相反。

4.2 Language Modeling

語言建模

我們在標準自回歸語言建模上評估了 Mamba 架構與其他架構的預訓練指標（perplexity）和零點評估。我們將模型大小（深度和寬度）設置為 GPT3 規格。我們使用 Pile 數據集，并遵循 Brown 等人（2020 年）中描述的訓練方法。所有訓練細節見附錄 E.2。

4.2.1 Scaling Laws

縮放定律

在基線方面，我們與標準 Transformer 架構（GPT3 架構）以及我們所知的最強 Transformer 配方（此處稱為 Transformer++）進行了比較，后者基于 PaLM 和 LLaMa 架構（例如旋轉嵌入、SwiGLU MLP、RMSNorm 而非 LayerNorm、無線性偏差和更高的學習率）。我們還與其他最新的亞二次方架構進行了比較（圖 4）。所有模型細節見附錄 E.2。

圖 4 顯示了在標準 Chinchilla 協議下，從 $\approx 125M$ 到 $\approx 1.3B$ 參數模型的縮放規律。Mamba 是第一個與非常強大的 Transformer 配方（Transformer++）性能相匹配的無注意力模型，后者現已成為標準配置，尤其是當序列長度增加時。我們注意到，RWKV 和 RetNet 基線在上下文長度為 $8k$ 的情況下，由于缺乏有效的實現，導致內存不足或不切實際的計算要求，因此缺少完整的結果。

4.2.2 Downstream Evaluations

下游評估

表 3 顯示了 Mamba 在一系列流行的下游零點評估任務中的表現。我們將這些模型與最著名的開源模型進行了比較，其中最重要的是 Pythia 和 RWKV，它們采用了與我們的模型相同的標記符、數據集和訓練長度（300B 標記符）進行訓練。(請注意，Mamba 和 Pythia 是在上下文長度為 2048 的情況下訓練的，而 RWKV 是在上下文長度為 1024 的情況下訓練的）。

4.3 DNA Modeling

DNA 建模

在大型語言模型取得成功的推動下，人們最近開始探索在基因組學中使用基礎模型范例。DNA 被比作語言，因為它由具有有限詞匯量的離散標記序列組成。眾所周知，它還需要長程依賴關系來建模。我們將 Mamba 作為預訓練和微調的調頻骨干進行研究，其環境與最近關于 DNA 長序列模型的研究相同。特別是，我們重點探討了模型大小和序列長度的縮放規律（圖 5），以及需要長上下文的下游合成分類任務（圖 6）。

在預訓練方面，我們主要采用標準的因果語言建模（下一個標記預測）設置來進行訓練和模型細節（另見附錄 E.2）。在數據集方面，我們基本沿用了 HyenaDNA 的設置，它使用 HG38 數據集進行預訓練，該數據集由單個人類基因組組成，在訓練分割中包含約 45 億個標記（DNA 堿基對）。

4.3.1 Scaling: Model Size

縮放： 模型尺寸

• 訓練

  為了使基線更具優勢，我們在 1024 的短序列長度上進行訓練；如第 4.3.2 節所示，我們預計在更長的序列長度上，結果會更有利于 Mamba。我們將全局批次大小固定為 1024，每批共有 $2^{20}\approx 1M$ 個標記。對模型進行了 $10$ 個梯度訓練，共訓練了 $10B$ 個詞塊。

• 結果

  圖 5（左）顯示，Mamba 的預訓練復雜度隨模型大小的變化而平滑提高，Mamba 的擴展能力優于 HyenaDNA 和 Transformer++。例如，在參數 $\approx 40M$ 的最大模型規模下，曲線顯示 Mamba 可以用大約少 3 倍到 4 倍的參數與 Transformer++ 和 HyenaDNA 模型相媲美。

4.3.2 Scaling: Context Length

縮放： 上下文長度

在下一個 DNA 實驗中，我們研究了模型隨序列長度的縮放特性。我們只對 HyenaDNA 和 Mamba 模型進行比較，因為在較長的序列長度上，二次注意的成本過高。我們在序列長度為 $2^{10}=1024$、$2^{12}=4096$、$2^{14}=16384$、$2^{16}=65536$、$2^{18}=262144$、$2^{20}=1048576$ 時對模型進行預訓練。我們將模型大小固定為 6 層 128 寬（約 $1.3M?1.4M$ 個參數）。對模型進行了 $20$ 個梯度的訓練，共訓練了 $\approx 330B$ 標記。較長的序列長度使用了類似于的序列長度熱身。

• 結果

  圖 5（右）顯示，即使是長度為 100 萬的超長序列，Mamba 也能利用較長的上下文。另一方面，HyenaDNA 模型隨著序列長度的增加而變差。從第 3.5 節關于選擇機制特性的討論中可以直觀地看出這一點。特別是，LTI 模型不能選擇性地忽略信息；從卷積的角度來看，一個很長的卷積核是在一個很長的序列中匯總所有信息，而這個序列可能是非常嘈雜的。需要注意的是，雖然 HyenaDNA 聲稱在較長的上下文中會有所改進，但他們的結果并沒有控制計算時間。

4.3.3 Synthetic Species Classication

合成物種分類

我們在一項下游任務中對模型進行了評估，該任務是通過隨機取樣 5 個不同物種的 DNA 連續片段來對其進行分類。這項任務改編自 HyenaDNA，后者使用的物種是 $人類、狐猴、小鼠、豬、河馬$。我們對這項任務進行了修改，使其更具挑戰性，即對已知共享 99% DNA 的五個類人猿物種 $人類、黑猩猩、大猩猩、猩猩、倭黑猩猩$ 進行分類。

4.4 Audio Modeling and Generation

音頻建模和生成

對于音頻波形模式，我們主要與 SaShiMi 架構和訓練協議進行比較。該模型包括：

1. 一個 U-Net 主干網，兩級池化系數𝑝使每級模型維度𝐷翻倍；
2. 每級交替使用 S4 和 MLP 塊。

我們考慮用 Mamba 模塊取代 S4+MLP 模塊。實驗詳情見附錄 E.4。

4.4.1 Long-Context Autoregressive Pretraining

長語境自回歸預訓練

我們在 YouTubeMix（DeepSound 2017）上對預訓練質量（自回歸下一樣本預測）進行了評估，YouTubeMix 是之前工作中使用的標準鋼琴音樂數據集，由 4 小時的鋼琴獨奏音樂組成，采樣率為 16000 Hz。 預訓練細節基本遵循標準語言建模設置（第 4.2 節）。圖 7 評估了在計算量保持不變的情況下，將訓練序列長度從 213 = 8192 增加到 220 ≈ 106 的效果。(數據整理方式存在一些細微的邊緣情況，可能會導致縮放曲線出現扭結。例如，只有分鐘長的片段可用，因此最大序列長度實際上以 60𝑠 ? 16000𝐻𝑧 = 960000 為界）。

Mamba 和 SaShiMi（S4+MLP）基線都隨著上下文長度的延長而不斷改進；Mamba 在整個過程中都更勝一籌，而在更長的上下文長度下，差距就會拉大。主要指標是每字節比特數 (BPB)，它是預訓練其他模式時標準負對數似然 (NLL) 損失的一個常數因子 $log(2)$。

我們注意到一個重要細節：這是本文中唯一一次將實參數化轉換為復參數化的實驗（第 3.6 節）。我們在附錄 E.4 中展示了其他消融情況。

4.4.2 Autoregressive Speech Generation

自回歸語音生成

SC09 是一個基準語音生成數據集，由 1 秒鐘的片段組成，采樣頻率為 16000 Hz，包含數字 "0 "到 “9”，具有高度可變的特征。我們基本沿用了 Goel 等人（2022 年）的自回歸訓練設置和生成協議。

表 4 顯示了 Mamba-UNet 模型與 Goel 等人（2022 年）的各種基線相比的自動化指標： WaveNet、SampleRNN、WaveGAN、DiffWave 和 SaShiMi。小型 Mamba 模型的性能優于最先進的（大得多的）基于 GAN 和擴散的模型。與基線模型參數匹配的更大模型則進一步顯著提高了保真度指標。

表 5 采用了小型 Mamba 模型，并研究了外層和中心階段的不同架構組合。結果表明，在外圍區塊，Mamba 一直優于 S4+MLP，而在中心區塊，Mamba > S4+MLP > MHA+MLP。

4.5 Speed and Memory Benchmarks

速度和內存基準測試

我們在圖 8 中對 SSM 掃描操作（狀態擴展 $N=16$）的速度以及 Mamba 的端到端推理吞吐量進行了基準測試。我們的高效 SSM 掃描比目前已知的最佳注意力實現（FlashAttention-2）更快，超過序列長度 2K，比 PyTorch 中的標準掃描實現快 20-40倍。Mamba 的推理吞吐量比類似規模的 Transformer 高出 4-5 倍，因為沒有 KV 緩存，它可以使用更高的批處理規模。例如，Mamba-6.9B（未經訓練）的推理吞吐量比 Transformer-1.3B 小 5 倍。詳情見附錄 E.5，其中還包括內存消耗基準。

4.6 Model Ablations

模型消融

4.6.1 Architecture

結構

表 6 調查了結構（塊）及其內部 SSM 層（圖 3）的影響。我們發現

• 在以前的非選擇性（LTI）SSM（相當于全局卷積）中，性能非常相似。
• 用實值 SSM 取代之前工作中的復值 S4 變體對性能影響不大，這表明（至少對 LM 來說）在考慮硬件效率時，實值 SSM 可能是更好的選擇。
• 用選擇性 SSM（S6）取代其中任何一個都能顯著提高性能，這也驗證了第 3 節的動機。
• Mamba 架構的性能與 H3 架構類似（使用選擇層時似乎略勝一籌）。

在附錄 E.2.2 中，我們還研究了 Mamba 模塊與其他模塊的交錯，如 MLP（傳統架構）、MHA（混合注意力架構）。

4.6.2 Selective SSM

選擇性 SSM

表 7 通過考慮選擇性 $\Delta$、$B$ 和 $C$ 參數的不同組合（算法 2），對選擇性 SSM 層進行了分析，結果表明 $\Delta$ 是最重要的參數，因為它與 RNN 門控有關（定理 1）。

表 8 考慮了 SSM 的不同初始化，這些初始化已被證明在某些數據模式和設置中會產生巨大差異。在語言建模方面，我們發現更簡單的實值對角初始化（S4D-Real，第 3 行）比更標準的復值參數化（S4D-Lin，第 1 行）效果更好。隨機初始化的效果也很好，這與之前的研究結果一致。

表 9 和表 10 分別考慮了改變 $\Delta$ 和 $(B,C)$ 投影的維度。將它們從靜態變為選擇性帶來的好處最大，而進一步增加維度通常會在參數數少量增加的情況下適度提高性能。

特別值得注意的是，當狀態大小 $N$ 增加時，選擇性 SSM 得到了顯著改善，只需增加 1% 的參數，就能提高超過 1.0 的困惑度。這驗證了我們在第 3.1 節和第 3.3 節中提出的核心動機。

Section 5 Discussion

討論

• 相關工作

  附錄 Ad 討論了選擇機制與類似概念的關系。附錄 B 是 SSM 和其他相關模型的擴展相關工作。

• 沒有免費的午餐 連續-離散光譜

  結構化 SSM 最初被定義為連續系統的離散化(1)，對連續時間數據模態（如感知信號，如音頻、視頻）有很強的歸納傾向。正如第 3.1 節和第 3.5 節所討論的，選擇機制克服了它們在離散模式（如文本和 DNA）上的弱點；但反過來，這也會妨礙它們在 LTI SSM 擅長的數據上的表現。我們對音頻波形的消融更詳細地研究了這種權衡。

• 下游能力

  基于轉換器的基礎模型（尤其是 LLM）具有豐富的生態系統特性和與預訓練模型的交互模式，如微調、適應、提示、上下文學習、指令調整、RLHF、量化等。我們尤其感興趣的是，變形器的替代品（如 SSM）是否也具有類似的特性和能力。

• 規模

  我們的實證評估僅限于較小的模型規模，低于大多數強大的開源 LLM 以及其他循環模型和 RetNet 的閾值，這些模型已在 7B 參數規模及以上進行了評估。我們仍需評估 Mamba 在這些更大的規模下是否仍有優勢。我們還注意到，擴展 SSM 可能涉及進一步的工程挑戰和對模型的調整，本文對此未作討論。

附錄

A Discussion: Selection Mechanism

討論： 選擇機制

我們的選擇機制受到門控、超網絡和數據依賴等概念的啟發，并與之相關。它也可以被視為與 "快速權重"相關，后者將經典 RNN 與線性注意機制聯系起來。然而，我們認為這是一個值得澄清的獨特概念。

• 門控

  門控最初指的是 RNN 的門控機制，如 LSTM 和 GRU，或門控方程 (5)n 定理 1。這被解釋為一種控制是否讓輸入進入 RNN 隱藏狀態的特殊機制。特別是，這會影響信號在時間中的傳播，并導致輸入沿著序列長度維度進行交互。

  不過，門控的概念在流行的用法中已被放寬，僅指任何乘法相互作用（通常與激活函數有關）。例如，神經網絡架構中的元素相乘成分（不沿序列長度發生相互作用）現在通常被稱為門控架構，盡管其含義與最初的 RNN 意義截然不同。因此，我們認為 RNN 門控的原始概念與乘法門控的流行用法實際上具有非常不同的語義。

• 超網絡

  超網絡指的是其參數本身由更小的神經網絡生成的神經網絡。最初的想法從狹義上將其用于定義一個大型 RNN，其遞歸參數由一個較小的 RNN 生成。

• 數據依賴性

  與超網絡類似，數據依賴性可指模型的某些參數取決于數據的任何概念。

• 舉例說明： GLU 激活

  為了說明這些概念的問題，請考慮一個簡單的對角線線性層 $y=Dx$，其中 $D$ 是一個對角線權重參數。現在假設 $D$ 本身是由 $x$ 的線性變換生成的，帶有可選的非線性：$D=\sigma (Wx)$。由于它是對角線，乘法就變成了元素乘積：$y=\sigma (Wx)°x$ 。

  這是一個相當瑣碎的變換，但它在技術上卻滿足了門控（因為它有一個乘法 “分支”）、超網絡（因為參數 $D$ 是由另一層產生的）和數據依賴（因為 $D$ 取決于數據 $x$）的常見含義。然而，這實際上只是定義了一個 GLU 函數，它是如此簡單，以至于經常被認為只是一個激活函數，而不是一個有意義的層。

• 選擇

  因此，雖然選擇機制可以被視為結構門控、超網絡或數據依賴等觀點的特例，但其他大量的構造也可以被視為特例–基本上任何有乘法的構造都可以，包括標準的注意力機制。

  相反，我們認為它與傳統 RNN 的門控機制關系最為密切，后者是一個特例（定理 1），而且通過 $\Delta$ 的變量（依賴于輸入）離散化，與 SSM 也有著更深的聯系。我們也摒棄了 “門控”（gating）一詞，轉而使用選擇（selection），以澄清前者過多的使用。更狹義地說，我們使用選擇來指代模型選擇或忽略輸入的機制作用，并促進序列長度上的數據交互（第 3.1 節）。除了選擇性 SSM 和門控 RNN 之外，其他例子還包括依賴輸入的卷積，甚至注意力。

B Related Work

B.1 S4 Variants and Derivatives

我們簡要介紹了過去工作中的一些結構化 SSM，特別是那些與我們的方法有關的 SSM。

• S4 引入了第一個結構化 SSM，描述了對角線結構和對角線加低階（DPLR）。由于與連續時間在線記憶（HIPPO）有關，該研究重點關注 DPLR SSM 的高效卷積算法。

• DSS 通過近似 HIPPO 初始化首次發現了對角結構 SSM 的經驗有效性。這在 S4D 中得到了理論上的拓展。

• S5 獨立地發現了對角 SSM 近似，并且是第一個利用并行掃描反復計算的 S4 模型。然而，這需要降低有效狀態維度，他們通過將 SSM 維度從 SISO（單輸入單輸出）公式轉換為 MIMO（多輸入多輸出）公式實現了這一目標。我們提出的 S6 與掃描相同，但不同之處在于：

  • 保留了 SISO 維度，從而提供了更大的有效循環狀態；
  • 使用硬件感知算法來克服計算問題；
  • 增加了選擇機制。

  Lu 等人（2023 年）將 S5 應用于元 RL，以處理情節軌跡之間 SSM 狀態的重置。他們的機制可以被看作是選擇機制的一個特殊硬編碼實例，其中 A 被手動設置為 0，而不是我們這種依賴于輸入的可學習機制。如果能將選擇性 SSM 通用地應用到這種設置中，并探究模型是否學會了在情節邊界自動重置其狀態，那將會非常有趣。

• Mega 將 S4 簡化為實值而非復值，并將其解釋為指數移動平均（EMA）。此外，他們還將 SSM 的離散化步驟與 EMA 阻尼項進行了有趣的聯系。與最初的 S4 論文中的發現相反，這是第一個模型顯示實值 SSM 在某些情況下或與不同的架構組件相結合時是有效的。

• Liquid S4 也是以依賴輸入的狀態轉換來增強 S4。從這個角度看，它與選擇機制有相似之處，盡管其形式有限，仍然是卷積計算，接近 LTI。

• SGConv、Hyena、LongConv 、MultiresConv 和 Toeplitz Neural Network 都側重于 S4 的卷積表示，并以不同的參數化創建全局或長卷積核。然而，這些方法無法直接進行快速自回歸推斷。

值得注意的是，所有這些方法以及我們所知的所有其他結構化 SSM 都是非選擇性的，而且通常是嚴格的 LTI（線性時間不變）。

B.2 SSM Architectures

我們使用 SSM 架構或狀態空間神經網絡（SSNN）來指代將之前的 SSM 之一作為黑盒層的深度神經網絡架構。

• GSS 是第一個包含 SSM 的門控神經網絡架構。它受到 Hua 等人（2022 年）的門控注意單元（GAU）的啟發，看起來與我們的區塊非常相似，只是多了一些投影。最重要的是，基于第 3.1 節中的動機，它的投影收縮了模型維度以減少 SSM 的狀態大小，而我們的投影則擴大了模型維度以增加狀態大小。

• Mega 將上述 S4 的 EMA 簡化與使用高效注意力近似的混合架構相結合。

• H3 的動機是將 S4 與線性注意相結合。它首次將線性注意的這一表述推廣到更一般的遞歸中，這也是后來架構的基礎。

• 選擇性 S4 將 S4 作為一個黑盒來生成二進制掩碼，并與輸入相乘。雖然與 "選擇 "同名，但我們認為這是一種架構修改，更接近于架構門控而非選擇機制（附錄 A）。例如，我們假設它無法解決 "選擇性復制 "任務，因為簡單地屏蔽掉不相關的輸入并不會影響相關輸入之間的間距（事實上，如果噪聲標記被嵌入為 0，"選擇性復制 "任務甚至可以被視為預先屏蔽）。

• RetNet 也是基于線性注意（Linear Attention），與 H3 非常相似，但將內部 S4 層簡化為狀態維度為 $N=1$ 的特殊情況。雖然其框架并非如此，但其遞推可視為線性 SSM 的特例。

  它的主要改進來源于使用大頭部維度的線性注意，這可以看作是另一種進行輸入依賴性狀態擴展的方法。在線性注意變體中使用較大的頭部維度是由 H3 首次實現的，但由于這需要相應的額外計算量，因此并未得到廣泛應用。RetNet 采用標準多頭注意力變體而不是卷積來并行計算，從而避免了這一問題。

• RWKV 是另一種最近為語言建模而設計的 RNN。它基于 AFT（無注意力轉換器），是線性注意力的另一種變體。它的主要 "WKV "機制涉及 LTI 循環，可以看作是兩個 SSM 的比率。

我們還重點介紹了 Hua 等人（2022 年）的門控注意力單元（GAU），該單元的設計靈感來自于將 Transformer 的 MHA 和 MLP 模塊組合在一起，也是我們將 H3 和 MLP 模塊組合在一起的架構（第 3.4 節）的靈感來源。

B.3 Relationship to RNNs

RNN 和 SSM 有著廣泛的聯系，因為它們都涉及潛在狀態的遞歸概念。

一些較早的 RNN，如強類型 RNN、準 RNN（QRNN）和簡單遞歸單元（SRU）涉及無時間非線性的門控 RNN 形式。由于門控機制和選擇機制的聯系，這些機制可被視為選擇性 SSM，因此在某種意義上比上述 LTI 結構 SSM 系列更為強大。其主要區別在于：

• 它們沒有使用狀態擴展（$N=1$）或選擇性 $B$、$C$ 參數，而這兩個參數對性能都很重要（第 4.6 節）。
• 他們使用啟發式門控機制，我們將其概括為選擇機制+離散化的結果（定理 1）。與原則性 SSM 理論的聯系提供了更好的參數化和初始化（第 3.6 節）。

此外，較早的 RNNs 還存在著名的效率問題和梯度消失問題，這兩個問題都是由其順序性引起的。對于上述一些 RNN，后者可以通過利用并行掃描來解決，但如果沒有后來為 SSM 開發的理論，前者就很難解決。例如，現代結構化 SSM 與經典 SSM 理論啟發下的遞歸動力學參數化（例如通過離散化）或直接分析有所不同。

我們還注意到，關于正交 RNN 的研究已有很長一段時間，其動機是約束 A 過渡矩陣為正交或單元矩陣，以控制其特征值并防止梯度消失問題。然而，這些研究還有其他局限性；我們認為，這些局限性源于正交/單元 RNN 也是 LTI 的這一事實。例如，它們幾乎總是在 "復制 "任務中得到評估，并能完美地解決該任務，但在 "選擇性復制 "任務中卻顯得力不從心。

B.4 Linear Attention

線性注意（LA）框架是普及核注意并展示其與遞歸自回歸模型關系的重要成果。許多變體提出了替代核和其他修改。隨機特征注意（RFA）利用高斯核的隨機傅里葉特征近似，選擇核特征圖來近似軟最大注意（即 exp 特征圖）。Performer 則找到了指數核的近似值，該近似值只涉及正特征，同時也允許使用 softmax 歸一化項。TransNormer 表明，LA 的分母項可能不穩定，因此建議用 LayerNorm 取而代之。cosFormer 用余弦重權機制增強了 RFA，該機制結合了位置信息以強調位置性。Linear Randomized Attention 從重要性采樣的角度對 RFA 進行了概括，并將其推廣到對整個 softmax 內核（而不僅僅是指數變換后的分子）提供更好的估計。

除了內核注意力，高效注意力還有許多其他變體；Tay、Dehghani、Bahri 等人（2022 年）的調查報告對其中許多變體進行了廣泛分類。

B.5 Long Context Models

長語境已成為一個熱門話題，最近有幾個模型聲稱可以擴展到越來越長的序列。不過，這些模型通常都是從計算角度出發，尚未經過廣泛驗證。這些模型包括：

• Recurrent Memory Transformer 是 Transformer 骨干的輕量級封裝。它顯示了泛化多達 100 萬個序列的能力，但僅限于合成記憶任務；他們的主要結果與我們的誘導頭外推法實驗類似（表 2）。
• LongNet 號稱可擴展至 1B 長度，但在實際任務中只對長度 < 100 的情況進行了評估。
• Hyena 和 HyenaDNA，它們聲稱最多可利用 100 萬上下文。然而，他們的實驗在更長的上下文中對更多的數據進行了訓練，因此很難斷定在 100 萬上下文中質量的提高是由于上下文長度還是由于更多的數據和計算。
• Sparse Transformer 展示了一個概念驗證，即使用分步稀疏注意力變換器對長度為 $2^{20}=1048576$ 的音頻波形進行建模，但并未討論控制計算和模型大小時的性能權衡。

與此相反，我們認為，這項工作首次提出了一種方法，有意義地證明了隨著語境時間的延長，性能也在不斷提高。