現今大多數深度學習的基礎模型（foundation models）都仰賴 Transformer 架構及其核心的注意力機制來推動各種強大應用，例如自然語言處理、語音辨識與基因序列分析。然而，Transformer 在計算效率上仍有明顯瓶頸，尤其面對長序列時，注意力機制的計算複雜度呈二次方成長，導致訓練和推論都相當耗時。為克服此缺陷，研究社群先後提出許多子二次方時間複雜度的架構，包括線性注意力（linear attention）、門控卷積、循環模型與結構化狀態空間模型（SSMs）。不過，這些方法在重要的序列資料類型（如自然語言）上，效能尚未能全面超越傳統的 Transformer。
本文由 Gu Dao 等人發表於 ICLR 2024，榮獲 Outstanding Paper 獎的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》對上述不足提出了深入診斷與關鍵改進，開創了一種效率與效能兼具的新型序列模型。本文將依序探討本研究的背景動機、核心技術創新、實驗驗證，以及其對 AI 領域的深遠影響。
研究背景與動機
Transformer 雖具備優越的內容基礎推理能力（content-based reasoning），但其注意力機制的計算成本高昂，特別是在長序列任務中顯得不夠可擴展。儘管 SSM 與類似技術透過動態系統理論或遞迴結構，達到線性甚至更低的時間與空間複雜度，卻因缺乏足夠的「內容選擇」能力，在需要根據輸入序列內容動態調整信息流時表現不佳。換句話說，SSM 難以判斷哪些序列片段需要強化或遺忘，這阻礙了它在語言等離散模態上的競爭力。
基於此關鍵觀察，作者希望解決兩大挑戰：
- 如何讓 SSM 能夠依據輸入序列的當前 token 動態調整參數與狀態轉移，強化對離散內容的處理能力。
- 如何維持線性時間處理長序列的能力，同時克服因參數變動所失去的快速卷積計算優勢，保持高效推論。
核心方法與技術創新
本研究提出了一種名為 Mamba 的新型架構，其核心在於「選擇性狀態空間模型」（Selective State Spaces, SSS）。具體而言，Mamba 讓原先的狀態空間模型參數由靜態變為能動態依據當前輸入內容改變，類似具「內容感知」(content-aware) 的機制，這是解決離散與多樣化序列處理的核心突破。
這個改變打破了傳統 SSM 可用於高效傅立葉變換或快速卷積的條件，使傳統快速運算方法無法直接應用。為此，作者設計了一套硬體友好的平行遞迴演算法，利用硬體並行與流水線特性來實現加速推論，使得即便是在無法用快速卷積處理的前提下，仍能實現近似線性時間的序列處理效率。
除此之外，Mamba 簡化了神經網路整體架構，不使用任何 Transformer 中的注意力層甚至多層感知機（MLP）塊，取而代之的是一個純淨的端到端序列建模框架。這樣不僅降低了系統複雜度，也加強了模型專注於高效且動態資訊流動的能力。
主要實驗結果
作者以多模態、多長度的序列資料驗證 Mamba 表現。首先，在自然語言建模中，他們訓練了一個擁有 30 億參數的 Mamba-3B 模型。該模型在預訓練和下游任務中，都超越了同等規模的 Transformer，並與兩倍規模的 Transformer 相當，顯現出其優異的效能與擴展性。
除了語言，Mamba 在音訊處理與基因序列分析等多個領域也展現競爭力，且序列長度擴展到百萬級別，其效能表現依然保持穩定，證明模型的確能夠處理極長序列，且推論速度比傳統 Transformer 快 5 倍以上。
這些實驗結果不僅證明了選擇性狀態空間模型的實用價值，也證明了在無需傳統注意力和 MLP 的簡化架構中可以達成更高效且優異的序列建模效果。
對 AI 領域的深遠影響
Mamba 的提出為序列建模領域開啟了一條全新途徑，特別是在高效且可擴展的基礎模型設計上具有重大啟示。其成功整合動態參數調整與平行遞迴演算法，不僅克服了 Transformer 的時間複雜度瓶頸，亦解決了傳統 SSM 在離散模態上無法有效內容推理的弱點。
未來，這樣的技術有潛力推動超長序列的深度學習應用，涵蓋多種關鍵領域例如長文本生成、高解析度音訊理解、及大規模基因資料分析。同時，Mamba 精簡的模型結構也為硬體友好型 AI 系統設計提供了範例，有助於降低推論延遲與能耗，進一步促進 AI 在資源受限環境（如手機邊緣計算或實時系統）的部署。
此外，Mamba 強調內容感知的狀態空間機制，激發學術界重新思考如何在保持計算效率下，提升模型對序列內容的動態適應與選擇性記憶，這將成為未來序列模型發展的重要方向之一。
總結
綜合而言，Mamba 不僅提出了一種全新的選擇性狀態空間建模方式，突破傳統限制，實現在長序列上線性時間且高效的序列建模，更在多模態任務中展現了卓越的效能。它成功平衡了效率與內容推理能力，為 Transformer 後的序列模型研究提供了具有開創性的解決方案，必將深刻影響未來 AI 基礎模型架構的設計與應用。
論文資訊
📄 Mamba: Linear-Time Sequence Modeling with Selective State Spaces
👥 Gu, Dao
🏆 ICLR 2024 · Outstanding Paper
🔗 arxiv.org/abs/2312.00752
沒有留言:
張貼留言