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

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

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

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

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

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

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

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

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

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

Mamba 是一個熱門研究方向，也因此有多個團隊都在嘗試編寫綜述報告，除了本文介紹的這一篇，還有另一些關注狀態空間模型或視覺 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 集中了循環神經網絡（RNN）的循環框架、Transformer 的并行計算和注意力機制、狀態空間模型（SSM）的線性特性。因此，為了透徹地理解 Mamba，就必需先理解這三種架構。

循環神經網絡

循環神經網絡（RNN）具有保留內部記憶的能力，因此很擅長處理序列數據。

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

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

Transformer

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

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

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

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

狀態空間

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

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

• 離散化

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

• 卷積計算

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

RNN、Transformer 和 SSM 之間的關系

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

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

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

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

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

Mamba

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

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

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

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

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

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

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

Mamba-2：狀態空間對偶

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

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

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

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

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

Mamba 塊

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

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

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

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

Mamba 模型正在發展進步

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

塊設計

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

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

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

掃描模式

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

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

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

記憶管理

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

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

讓 Mamba 適應多樣化的數據

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

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

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

序列數據

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

非序列數據

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

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

多模態數據

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

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

應用

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

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

挑戰與機遇

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

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