一文看懂Mamba，Transformer最強競爭者

机器之心發表於2024-08-19
ORM

Mamba 雖好，但發展尚早。


深度學習架構有很多，但近些年最成功的莫過於 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）。

相關文章