近年來,Transformer 架構憑藉其強大的建模能力,已成為深度學習領域中的基礎模型核心,推動語言、音訊、基因組等多種序列資料的突破性發展。然而,Transformer 在面對長序列時的計算複雜度高達二次方,使其在資源有限或長序列應用中效率不彰,催生了大量針對加速和降低運算負擔的架構改良方案。
本次於 ICLR 2024 獲得 Outstanding Paper 的《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》一文,由 Gu Dao 等人提出了一種基於選擇性狀態空間模型(Selective State Space Models, S3M)的新型序列建模架構——Mamba,成功在保持甚至超越 Transformer 性能的同時,以線性時間的推理效率解決長序列處理的核心挑戰,為序列建模領域開闢一條新的技術路徑。
研究背景與動機
Transformer 的自注意力機制對序列中任意位置的內容均可直接關聯,為語言模型等提供了極強的上下文理解能力。但其計算代價隨序列長度平方級增長,限制了在超長序列或低延遲場景的應用。為解決此問題,近年研究社群嘗試以線性注意力、門控卷積、循環神經網路及結構化狀態空間模型(SSMs)等架構降低計算成本。
其中,結構化狀態空間模型因長距依賴建模能力佳且可實現線性時間複雜度,被看作取代 Transformer 的有力競爭者。但現有 SSM 架構在離散符號(如語言)等模態表現未能達到 Transformer 水準,研究者發現其關鍵瓶頸在於模型缺乏以內容為依據的推理能力,即無法根據當前輸入動態調節狀態的更新與遺忘策略,導致上下文捕捉不夠彈性與精確。
核心方法與創新
Mamba 論文的最重要貢獻即在於設計出具備「選擇性參數」的狀態空間模型。具體而言,作者提出將 SSM 的關鍵參數表示為當前輸入的函數,讓模型能動態根據每個標記(token)來調整狀態轉移與輸出過程,類似於賦予模型內容敏感的「控制門」,實現資訊的有選擇性保存或遺忘。這種設計大幅提升 SSM 在離散模態上的建模能力,使其兼具內容感知與長距依賴捕捉。
引入動態參數化後,傳統基於卷積的快速計算路徑被打破,為解決這一效率瓶頸,作者設計了一套「硬體感知(hardware-aware)」的並行演算法,在循環 (recurrent) 模式下有效利用現代 GPU 的平行計算能力,達成線性時間複雜度且保持高吞吐量。
此外,Mamba 架構設計上極致簡化,捨棄了經典 Transformer 的注意力機制與多層感知器 (MLP) 區塊,以純粹且統一的選擇性狀態空間層組成端到端神經網路骨幹,這不僅降低了架構的複雜度,也強化了模型的理論解釋與實際應用便利性。
主要實驗結果
經由大規模的多模態評測,Mamba 展現出驚人的表現:在自然語言建模任務中,3B 參數版本的 Mamba,無論是在大規模預訓練或下游任務上,都超越了相同規模 Transformer 模型的效能,甚至追平了兩倍參數尺寸 Transformer 的表現,顯示出優異的參數效率和泛化能力。
音訊和基因組資料上的測試同樣驗證了 Mamba 的廣泛適用性和領先性能,使其成為跨領域序列資料分析的新興主力架構。
在推理速度方面,Mamba 擁有遠高於 Transformer 約 5 倍的吞吐量,且在序列長度接近百萬 token 時仍保持線性運算複雜度,為超長序列實時處理提供了可行解決方案。
對 AI 領域的深遠影響
Mamba 的成功不僅是對傳統 Transformer 架構的一次強力挑戰,更重要的是對序列建模范式的創新拓展。其選擇性狀態空間策略突破了以往 SSM 模型無法靈活響應輸入內容的短板,填補了內容敏感與計算高效之間的鴻溝。
在實務角度,Mamba 為開發面臨長序列、大規模資料的系統工程師帶來了新的利器,可應用於語言理解、語音識別、基因序列分析等多種領域,尤其適合需在有限資源下進行高效推理的嵌入式或邊緣應用。
學術層面,Mamba 提供了可解釋的、結構化且硬體友好的序列處理架構設計範例,激發未來研究者探索結合動態參數化與高效計算的深度模型新方向,可能引領序列模型從單純的注意力機制走向更靈活且高效的混合建模體系。
總體而言,Mamba 不僅完成了效率與效能之間的雙贏,也為長序列處理帶來了理論與應用上的雙重革命,預示著未來大型基礎模型的架構創新將更加多元與實用。
論文資訊
📄 Mamba: Linear-Time Sequence Modeling with Selective State Spaces
👥 Gu, Dao
🏆 ICLR 2024 · Outstanding Paper
🔗 arxiv.org/abs/2312.00752

沒有留言:
張貼留言