隨著深度學習在自然語言處理、語音辨識、基因組學等多種序列資料上的廣泛應用，序列建模技術的效能與效率成為研究的核心焦點。當前最成功且最具代表性的基礎模型，多是基於 Transformer 架構，其核心的自注意力（Self-Attention）機制擅長捕捉長距離依賴關係，然而自注意力機制在序列長度增加時面臨計算和記憶成本呈二次方成長的挑戰，極大限制了其在超長序列（例如百萬級別長度）的應用。
針對此瓶頸，過去許多研究嘗試提出次二次時間複雜度（subquadratic complexity）的替代方案，包括線性注意力（linear attention）、門控卷積（gated convolution）、遞迴模型以及結構化狀態空間模型（Structured State Space Models, SSMs）等。這些方法雖然在理論上能夠降低計算成本，但在關鍵的離散模態（例如語言）表現上常常不敵 Transformer。本文作者 Dao Gu 等人發現其中的原因，主要是這些模型缺乏「基於內容的推理能力（content-based reasoning）」，即無法根據序列中目前的輸入內容動態且選擇性地傳遞或遺忘資訊。
核心貢獻與技術創新
本文提出的 Mamba 模型，在結構化狀態空間模型基礎上引入了「選擇性狀態空間（Selective State Spaces）」的概念，即讓狀態空間模型的參數隨輸入而動態調整，藉此大幅提升模型對離散序列（像是語言）中內容變化的敏感度和適應性。這一點打破了以往 SSM 參數固定不變，透過線性卷積高效計算的限制，讓模型能根據當前 token 決定資訊的繼續傳遞或遺忘，實現更有效的內容理解與推理。
不過動態參數化通常會帶來計算上的瓶頸，失去原本可並行計算的優勢。論文中針對此設計了一套硬體感知的「迴圈（Recurrent）平行演算法」，同時保留了選擇性動態參數的彈性，並優化計算流程以平衡序列長度的線性延展性與 GPU 計算效率。這使得 Mamba 不僅在理論上擁有線性時間複雜度，更實際在硬體上實現了 5 倍較 Transformer 更高的推理吞吐量（throughput）。
值得一提的是，Mamba 模型架構設計剔除了傳統 Transformer 所依賴的注意力機制（attention）及多層感知機（MLP）模塊，僅透過精簡的選擇性狀態空間層架構就能完成端對端的序列建模任務。這種極度簡化卻高效的架構，是該研究的一大亮點，也讓模型更易於擴展至超長序列。
實驗結果與應用
作者在多個序列任務上驗證 Mamba 的性能，包括語言模型、語音分析與基因序列處理，展示其優於現有有效率序列模型的方法的表現。而在語言模型領域，Mamba-3B（約三十億參數規模）不僅優於相同參數量級的 Transformer，甚至能媲美兩倍規模 Transformer 的表現，這包括在預訓練階段及下游任務評估上的一致提升。
此外，Mamba 在處理極長序列（長度可達百萬級別）時，仍維持良好的效能和穩定性，這是現有 Transformer 及其變體難以做到的。其線性時間複雜度使得在硬體與能耗限制環境中，具備明顯的實務價值。
對 AI 領域的深遠影響
Mamba 的提出對序列建模領域具有多方面的深遠意義：
- 突破 Transformer 計算瓶頸：以選擇性動態的狀態空間模型取代自注意力機制，解決了 Transformer 在長序列上的資源龐大難題，讓基礎模型可應對百萬級序列長度。
- 內容感知推理能力的提升：透過讓模型參數依內容變化，可動態選擇資訊傳遞策略，這為序列模型在語言等離散模態中帶來全新推理途徑，可能影響未來語言模型結構設計。
- 硬體加速與效率優化的典範：巧妙整合硬體敏感演算法與動態模型設計，為高效能神經網路提供一條既非犧牲性能又非犧牲效率的可行路徑。
- 簡化架構的挑戰與可能：捨棄傳統 attention 與 MLP 模塊，利用純粹狀態空間結構進行序列建模，展現了模型簡化與性能優化可同步完成的潛力，為未來架構設計帶來不同思考角度。
總結來說，Mamba 代表了一種突破性且實用性的序列建模新範式。它不只是針對 Transformer 的一種替代方法，更是一種對「如何讓模型既敏銳於內容又高效處理長序列」這一核心難題的新解，未來有望推動大型基礎模型及多模態應用邁向更高效、更深度的發展。
對研究者與工程師而言，深入理解 Mamba 的選擇性狀態空間設計理念與硬體感知算法實作，有助於拓展序列模型的設計視野，也激發出跨領域理論與系統層面優化的潛力。
論文資訊
📄 Mamba: Linear-Time Sequence Modeling with Selective State Spaces
👥 Gu, Dao
🏆 ICLR 2024 · Outstanding Paper
🔗 arxiv.org/abs/2312.00752
沒有留言:
張貼留言