深度學習架構有很多，但近些年最成功的莫過于 Transformer，其已經(jīng)在多個應用領域確立了自己的主導地位。

如此成功的一大關鍵推動力是注意力機制，這能讓基于 Transformer 的模型關注與輸入序列相關的部分，實現(xiàn)更好的上下文理解。但是，注意力機制的缺點是計算開銷大，會隨輸入規(guī)模而二次增長，也因此就難以處理非常長的文本。

好在前段時間誕生了一種頗具潛力的新架構：結構化的狀態(tài)空間序列模型（SSM）。該架構能高效地捕獲序列數(shù)據(jù)中的復雜依賴關系，并由此成為 Transformer 的一大強勁對手。

這類模型的設計靈感來自經(jīng)典的狀態(tài)空間模型 —— 我們可以將其看作是循環(huán)神經(jīng)網(wǎng)絡和卷積神經(jīng)網(wǎng)絡的融合模型。它們可使用循環(huán)或卷積運算進行高效地計算，從而讓計算開銷隨序列長度而線性或近線性地變化，由此大幅降低計算成本。

更具體而言，SSM 最成功的變體之一 Mamba 的建模能力已經(jīng)可以比肩 Transformer，同時還能維持隨序列長度的線性可擴展性。

Mamba 首先引入了一個簡單卻有效選擇機制，其可根據(jù)輸入對 SSM 進行重新參數(shù)化，從而可讓模型在濾除不相關信息的同時無限期地保留必要和相關的數(shù)據(jù)。然后，Mamba 還包含一種硬件感知型算法，可使用掃描（scan）而非卷積來循環(huán)地計算模型，這在 A100 GPU 上能讓計算速度提升 3 倍。

如圖 1 所示，憑借強大的建模復雜長序列數(shù)據(jù)的能力和近乎線性的可擴展性，Mamba 已經(jīng)崛起成為一種基礎模型，并有望變革計算機視覺、自然語言處理和醫(yī)療等多個研究和應用領域。

因此，研究和應用 Mamba 的文獻迅速增長，讓人目不暇接，一篇全面的綜述報告必定大有裨益。近日，香港理工大學的一個研究團隊在 arXiv 上發(fā)布了他們的貢獻。

• 論文標題：A Survey of Mamba
• 論文地址：https://arxiv.org/pdf/2408.01129

這份綜述報告從多個角度對 Mamba 進行了總結，既能幫助初學者學習 Mamba 的基礎工作機制，也能助力經(jīng)驗豐富的實踐者了解最新進展。

Mamba 是一個熱門研究方向，也因此有多個團隊都在嘗試編寫綜述報告，除了本文介紹的這一篇，還有另一些關注狀態(tài)空間模型或視覺 Mamba 的綜述，詳情請參閱相應論文：

• Mamba-360: Survey of state space models as transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv:2404.16112
• State space model for new-generation network alternative to transformers: A survey. arXiv:2404.09516
• Vision Mamba: A Comprehensive Survey and Taxonomy. arXiv:2405.04404
• A survey on vision mamba: Models, applications and challenges. arXiv:2404.18861
• A survey on visual mamba. arXiv:2404.15956

預備知識

Mamba 集中了循環(huán)神經(jīng)網(wǎng)絡（RNN）的循環(huán)框架、Transformer 的并行計算和注意力機制、狀態(tài)空間模型（SSM）的線性特性。因此，為了透徹地理解 Mamba，就必需先理解這三種架構。

循環(huán)神經(jīng)網(wǎng)絡

循環(huán)神經(jīng)網(wǎng)絡（RNN）具有保留內(nèi)部記憶的能力，因此很擅長處理序列數(shù)據(jù)。

具體來說，在每個離散時間步驟 k，標準 RNN 在處理一個向量時會連同前一時間步驟的隱藏狀態(tài)一起處理，之后輸出另一個向量并更新隱藏狀態(tài)。這個隱藏狀態(tài)就可作為 RNN 的記憶，其能保留過去已見過的輸入的信息。這種動態(tài)記憶讓 RNN 可處理不同長度的序列。

也就是說，RNN 是一種非線性的循環(huán)模型，可通過使用存儲在隱藏狀態(tài)中歷史知識來有效地捕獲時間模式。

Transformer

Transformer 的自注意力機制有助于捕獲輸入之中的全局依賴。其實現(xiàn)方式是基于每個位置相對于其它位置的重要程度為它們分配權重。更具體而言，首先對原始輸入進行線性變換，將輸入向量的序列 x 轉換成三類向量：查詢 Q、鍵 K 和值 V。

然后計算歸一化的注意力分數(shù) S 并計算注意力權重。

除了可以執(zhí)行單個注意力函數(shù)，我們還可以執(zhí)行多頭注意力。這讓模型可以捕獲不同類型的關系，并從多個視角理解輸入序列。多頭注意力會使用多組自注意力模塊并行地處理輸入序列。其中每個頭都獨立運作，執(zhí)行的計算與標準自注意力機制一樣。

之后，將每個頭的注意力權重匯聚組合，得到值向量的加權和。這個聚合步驟可讓模型使用來自多個頭的信息并捕獲輸入序列中的多種不同模式和關系。

狀態(tài)空間

狀態(tài)空間模型（SSM）是一種傳統(tǒng)的數(shù)學框架，可用于描述系統(tǒng)隨時間變化的動態(tài)行為。近些年來，人們已將 SSM 廣泛應用于控制論、機器人學和經(jīng)濟學等多個不同領域。

究其核心，SSM 是通過一組名為「狀態(tài)」的隱藏變量來體現(xiàn)系統(tǒng)的行為，使其能有效捕獲時間數(shù)據(jù)的依賴關系。不同于 RNN，SSM 是一種具有關聯(lián)（associative）屬性的線性模型。具體來說，經(jīng)典的狀態(tài)空間模型會構建兩個關鍵方程（狀態(tài)方程和觀察方程），以通過一個 N 維的隱藏狀態(tài) h (t) 建模當前時間 t 時輸入 x 與輸出 y 之間的關系。

• 離散化

為了滿足機器學習的需求，SSM 必需經(jīng)歷一個離散化過程 —— 將連續(xù)參數(shù)轉變成離散參數(shù)。通常來說，離散化方法的目標是將連續(xù)時間劃分為具有盡可能相等積分面積的 K 個離散區(qū)間。為了實現(xiàn)這一目標，SSM 采用的最具代表性的解決方案之一是 Zero-Order Hold（ZOH），其假設區(qū)間 Δ = [??_{???1}, ??_?? ] 上的函數(shù)值保持不變。離散 SSM 與循環(huán)神經(jīng)網(wǎng)絡結構相似，因此離散 SSM 能比基于 Transformer 的模型更高效地執(zhí)行推理過程。

• 卷積計算

離散 SSM 是一個具有結合屬性的線性系統(tǒng)，因此可以與卷積計算無縫整合。

RNN、Transformer 和 SSM 之間的關系

圖 2 展示了 RNN、Transformer 和 SSM 的計算算法。

一方面，常規(guī) RNN 的運作基于一種非線性的循環(huán)框架，其中每個計算都僅依賴于之前的隱藏狀態(tài)和當前輸入。

盡管這種形式可讓 RNN 在自回歸推理時快速生成輸出，但它也讓 RNN 難以充分利用 GPU 的并行計算能力，導致模型訓練速度變慢。

另一方面，Transformer 架構是在多個「查詢 - 鍵」對上并行執(zhí)行矩陣乘法，而矩陣乘法可以高效地分配給硬件資源，從而更快地訓練基于注意力的模型。但是，如果要讓基于 Transformer 的模型生成響應或預測，則推理過程會非常耗時。

不同于僅支持一類計算的 RNN 和 Transformer，離散 SSM 靈活性很高；得益于其線性性質(zhì)，它既能支持循環(huán)計算，也可支持卷積計算。這種特性讓 SSM 不僅能實現(xiàn)高效推理，也能實現(xiàn)并行訓練。但是，需要指出，最常規(guī)的 SSM 是時不變的，也就是說其 A、B、C 和 Δ 與模型輸入 x 無關。這會限制其上下文感知型建模的能力，導致 SSM 在選擇性復制等一些特定任務上表現(xiàn)不佳。

Mamba

為了解決上述傳統(tǒng) SSM 的缺點，實現(xiàn)上下文感知型建模，Albert Gu 和 Tri Dao 提出了可用作通用序列基礎模型主干網(wǎng)絡的 Mamba，參閱機器之心報道《五倍吞吐量，性能全面包圍 Transformer：新架構 Mamba 引爆 AI 圈》。

之后，他們倆又進一步提出了 Mamba-2，其中的結構化空間狀態(tài)對偶（SSD/Structured Space-State Duality）構建了一個將結構化 SSM 與多種形式的注意力連接起來的穩(wěn)健的理論框架，讓我們可將原本為 Transformer 開發(fā)的算法和系統(tǒng)優(yōu)化技術遷移用于 SSM，也可參閱機器之心報道《再戰(zhàn) Transformer！原作者帶隊的 Mamba 2 來了，新架構訓練效率大幅提升》。

Mamba-1：使用硬件感知型算法的選擇式狀態(tài)空間模型

Mamba-1 基于結構化狀態(tài)空間模型引入了三大創(chuàng)新技術，即基于高階多項式投影算子（HiPPO）的內(nèi)存初始化、選擇機制和硬件感知型計算。如圖 3 所示。這些技術的目標是提升 SSM 的長程線性時間序列建模能力。

具體來說，其中的初始化策略可構建一個連貫的隱藏狀態(tài)矩陣，以有效地促進長程記憶。

然后，選擇機制可讓 SSM 有能力獲取可感知內(nèi)容的表征。

最后，為了提升訓練效率，Mamba 還包含兩種硬件感知型計算算法：Parallel Associative Scan（并行關聯(lián)掃描）和 Memory Recomputation（內(nèi)存重新計算）。

Mamba-2：狀態(tài)空間對偶

Transformer 啟發(fā)了多種不同技術的發(fā)展，比如參數(shù)高效型微調(diào)、災難性遺忘緩解、模型量化。為了讓狀態(tài)空間模型也能受益于這些原本為 Transformer 開發(fā)的技術，Mamba-2 引入了一個新框架：結構化狀態(tài)空間對偶（SSD）。該框架在理論上將 SSM 和不同形式的注意力連接到了一起。

本質(zhì)上講，SSD 表明，Transformer 使用的注意力機制和 SSM 中使用的線性時不變系統(tǒng)都可被視為半可分離的矩陣變換。

此外，Albert Gu 和 Tri Dao 還證明選擇式 SSM 等價于使用一種半可分離掩碼矩陣實現(xiàn)的結構化線性注意力機制。

Mamba-2 基于 SSD 設計了一種能更高效使用硬件的計算方法，這要用到一種塊分解矩陣乘法算法。

具體來說，通過這種矩陣變換將狀態(tài)空間模型視為半可分離矩陣，Mamba-2 能將該計算分解為矩陣塊，其中對角塊表示塊內(nèi)計算。而非對角塊則表示通過 SSM 的隱藏狀態(tài)分解的塊間計算。該方法可讓 Mamba-2 的訓練速度超過 Mamba-1 的并行關聯(lián)掃描的 2-8 倍，同時性能還能媲美 Transformer。

Mamba 塊

下面來看看 Mamba-1 和 Mamba-2 的塊設計。圖 4 比較了這兩種架構。

Mamba-1 的設計是以 SSM 為中心，其中選擇式 SSM 層的任務是執(zhí)行從輸入序列 X 到 Y 的映射。在這種設計中，經(jīng)過了初始的創(chuàng)建 X 的線性投射之后，會使用 (A, B, C) 的線性投射。然后，輸入 token 和狀態(tài)矩陣會通過選擇式 SSM 單元，利用并行關聯(lián)掃描，從而得到輸出 Y。之后，Mamba-1 采用了一個 skip 連接，以鼓勵特征復用和緩解常在模型訓練過程中發(fā)生的性能下降問題。最后，通過交錯地堆疊該模塊與標準歸一化和殘差連接，便可構建出 Mamba 模型。

至于 Mamba-2，則是引入了 SSD 層來創(chuàng)建從 [X, A, B, C] 到 Y 的映射。其實現(xiàn)方式是在塊的起點處使用單個投射來同時處理 [X, A, B, C]，這類似于標準注意力架構以并行方式生成 Q、K、V 投射的方式。

也就是說，通過移除序列線性投射，Mamba-2 塊是在 Mamba-1 塊的基礎上進行了簡化。這能讓 SSD 結構的計算速度超過 Mamba-1 的并行選擇式掃描。此外，為了提升訓練穩(wěn)定性，Mamba-2 還在 skip 連接之后添加了一個歸一化層。

Mamba 模型正在發(fā)展進步

狀態(tài)空間模型和 Mamba 近來發(fā)展迅猛，已經(jīng)成為了一大極具潛力的基礎模型骨干網(wǎng)絡選擇。盡管 Mamba 在自然語言處理任務上表現(xiàn)不俗，但也仍具有一些難題，比如記憶丟失、難以泛化到不同任務、在復雜模式方面的表現(xiàn)不及基于 Transformer 的語言模型。為了解決這些難題，研究社區(qū)為 Mamba 架構提出了諸多改進方案?，F(xiàn)有的研究主要集中于修改塊設計、掃描模式和記憶管理。表 1 分類總結了相關研究。

塊設計

Mamba 塊的設計和結構對 Mamba 模型的總體性能有很大的影響，也因此這成為了一大研究熱點。

如圖 5 所示，基于構建新 Mamba 模塊的不同方法，現(xiàn)有研究可以分為三類：

• 集成方法：將 Mamba 塊與其它模型集成到一起，實現(xiàn)效果與效率的平衡；
• 替換方法：用 Mamba 塊替換其它模型框架中的主要層；
• 修改方法：修改經(jīng)典 Mamba 塊內(nèi)的組件。

掃描模式

并行關聯(lián)掃描是 Mamba 模型內(nèi)的一大關鍵組件，其目標是解決由選擇機制導致的計算問題、提升訓練過程速度以及降低內(nèi)存需求。其實現(xiàn)方式是利用時變的 SSM 的線性性質(zhì)來在硬件層級上設計核融合和重新計算。但是，Mamba 的單向序列建模范式不利于全面學習多樣化的數(shù)據(jù)，比如圖像和視頻。

為緩解這一問題，一些研究者探索了新的高效掃描方法，以提升 Mamba 模型的性能以及促進其訓練過程。如圖 6 所示，在開發(fā)掃描模式方面，現(xiàn)有的研究成果可以分為兩類：

• 展平式掃描方法：以展平的視角看待 token 序列，并基于此處理模型輸入；
• 立體式掃描方法：跨維度、通道或尺度掃描模型輸入，這又可進一步分為三類：分層掃描、時空掃描、混合掃描。

記憶管理

類似于 RNN，在狀態(tài)空間模型內(nèi)，隱藏狀態(tài)的記憶有效地存儲了之前步驟的信息，因此對 SSM 的整體性能有著至關重要的影響。盡管 Mamba 引入了基于 HiPPO 的方法來進行記憶初始化，但管理 SSM 單元中的記憶依然難度很大，其中包括在層之前轉移隱藏信息以及實現(xiàn)無損記憶壓縮。

為此，一些開創(chuàng)性研究提出了一些不同的解決方案，包括記憶的初始化、壓縮和連接。

讓 Mamba 適應多樣化的數(shù)據(jù)

Mamba 架構是選擇式狀態(tài)空間模型的一種擴展，其具備循環(huán)模型的基本特性，因而非常適合作為處理文本、時間序列、語音等序列數(shù)據(jù)的通用基礎模型。

不僅如此，近期一些開創(chuàng)性研究更是擴展了 Mamba 架構的應用場景，使其不僅能處理序列數(shù)據(jù)，還能用于圖像和圖譜等領域，如圖 7 所示。

這些研究的目標是既充分利用 Mamba 能獲取長程依賴關系的出色能力，也讓其發(fā)揮學習和推理過程中的效率優(yōu)勢。表 2 簡單總結了這些研究成果。

序列數(shù)據(jù)

序列數(shù)據(jù)是指以特定順序收集和整理的數(shù)據(jù)，其中數(shù)據(jù)點的順序具有重要意義。這份綜述報告全面總結了 Mamba 在多種序列數(shù)據(jù)上的應用，包括自然語言、視頻、時間序列、語音和人體運動數(shù)據(jù)。詳見原論文。

非序列數(shù)據(jù)

不同于序列數(shù)據(jù)，非序列數(shù)據(jù)并不遵循特定的順序。其數(shù)據(jù)點可以任意順序進行組織而不會對數(shù)據(jù)的含義造成顯著影響。對于專門設計用于捕獲數(shù)據(jù)中時間依賴關系的循環(huán)模型（RNN 和 SSM 等）來說，這種缺乏固有順序的數(shù)據(jù)會很難處理。

令人驚訝的是，近期的一些研究成功讓 Mamba（代表性的 SSM）實現(xiàn)了對非序列數(shù)據(jù)的高效處理，包括圖像、圖譜和點云數(shù)據(jù)。

多模態(tài)數(shù)據(jù)

為了提升 AI 的感知和場景理解能力，可以整合多個模態(tài)的數(shù)據(jù)，比如語言（序列數(shù)據(jù)）和圖像（非序列數(shù)據(jù)）。這樣的整合能提供非常有價值和補充性的信息。

近段時間來，多模態(tài)大型語言模型（MLLM）是最受關注的研究熱點；這類模型繼承了大型語言模型（LLM）的強大能力，包括強大的語言表達和邏輯推理能力。盡管 Transformer 已經(jīng)成為該領域的主導方法，但 Mamba 也正在崛起成為一大強勁競爭者，其在對齊混合源數(shù)據(jù)和實現(xiàn)序列長度的線性復雜度擴展方面表現(xiàn)出色，這使 Mamba 有望在多模態(tài)學習方面替代 Transformer。

應用

下面介紹基于 Mamba 的模型的一些值得注意的應用。該團隊將這些應用分為了以下類別：自然語言處理、計算機視覺、語音分析、藥物發(fā)現(xiàn)、推薦系統(tǒng)以及機器人和自主系統(tǒng)。

這里我們不再過多介紹，詳見原論文。

挑戰(zhàn)與機遇

Mamba 雖然已經(jīng)在一些領域取得了出色表現(xiàn)，但總體而言，Mamba 研究仍還處于起步階段，前方仍還有一些挑戰(zhàn)有待克服。當然，這些挑戰(zhàn)同時也是機遇。

• 如何開發(fā)和改進基于 Mamba 的基礎模型；
• 如何充分實現(xiàn)硬件感知型計算，以盡可能利用 GPU 和 TPU 等硬件，提升模型效率；
• 如何提升 Mamba 模型的可信度，這需要安全和穩(wěn)健性、公平性、可解釋性以及隱私方面的進一步研究；
• 如何將 Transformer 領域的新技術用于 Mamba，如參數(shù)高效型微調(diào)、災難性遺忘緩解、檢索增強式生成（RAG）。