Mamba (深度學習架構)

Mamba是一種用於序列建模任務的深度學習架構。它是由卡內基梅隆大學和普林斯頓大學的研究人員開發的，用於解決Transformer架構的一些局限性，尤其是在處理長序列方面。它基於結構化狀態空間序列（S4）模型。 ^[1][2][3]

架構

為了能夠處理長資料序列，Mamba採用了結構化狀態空間序列模型（S4）。^[1]S4可以通過結合連續時間、迴圈和卷積模型有效且高效地對長序列進行建模。這使得它能夠處理不規則採樣的資料和長上下文，並在訓練和推理期間保持計算效率。^[4]

Mamba對S4模型在時序運算處理方面實現了突破性改進，其創新性的動態參數選擇機制能夠根據輸入特徵自適應調整結構化狀態空間模型（SSM）的參數組態。^[5][1]這種機制使模型具備序列資訊的選擇性注意力能力，通過動態篩選關鍵特徵實現對冗餘資料的智慧型過濾。相較於傳統時不變系統框架，Mamba採用的時變架構不僅提升了計算效率，更最佳化了複雜序列建模的能量消耗比，實現了模型效能的階躍式提升。^[1][6]

Mamba採用一種硬體感知演算法。該實現利用GPU去進行核心融合、並列掃描和重新計算來避免在主記憶體密集型層中實現擴充狀態^[1]，從而提高效能和主記憶體使用率。與Transformer相比，該方法在處理長序列時效率明顯更高。^[1][6]

此外，Mamba通過將SSM設計與MLP塊相整合來簡化其架構，從而形成相同品質但精簡的結構，進一步增強了該模型進行跨文字、音訊和基因組等資料類型進行通用序列建模的能力，同時保持了訓練和推理的高效率。^[1]

關鍵組件

• 選擇性狀態空間（SSM）：SSM是Mamba的核心，它是根據當前輸入選擇性處理資訊的迴圈模型。這使得能夠專注於相關資訊，丟棄不相關的資料。^[1]
• 簡化架構：Mamba用統一的SSM塊取代了Transformers的複雜注意力和MLP塊。目的是降低計算複雜度並提高推理速度。^[1]
• 硬體感知並列：Mamba採用迴圈模式，並採用專為提高硬體效率而設計的並列演算法，從而進一步提升其效能。^[1]

與 Transformer 的比較

特徵	Transfromer	Mamba
架構	基於注意力	基於SSM
複雜	高	低
推理時間複雜度	O(n)	O(1)
訓練時間複雜度	O(n 2 )	O(n)

變種模型

無token的語言模型：MambaByte

由於要對每個位元組大小的token進行操作，Transformer的擴充性較差，因為每個token都必須「關注」其他每個token，從而導致O(n 2 )的計算複雜度，因此，Transformer 選擇使用子詞分詞器來減少文字中的token數量，然而，這會導致詞彙表和詞嵌入非常大。

這項研究探討了一種新穎的語言建模方法：MambaByte，它不同於標準的基於token的方法。與依賴於將文字分解為離散單元的傳統模型不同，MambaByte 直接處理原始位元組序列。這消除了token化的需要，可能帶來以下幾個優點：^[7]

• 語言獨立性：tokenization通常依賴於特定於語言的規則和詞彙，從而限制了其在不同語言中的適用性。 MambaByte 的位元組級表示使其能夠處理不同的語言，而無需特定於語言的適應。
• 消除子詞分詞器帶來的偏見：常見子詞被過度代表，而罕見詞或新詞被低估或被分成意義較小的單元。這會影響模型的理解和生成能力，特別是對於形態豐富的語言或在訓練資料中沒有很好表現的token。
• 預處理的簡化：通過消除對複雜token和詞彙管理的需求，簡化了預處理流程，減少了預處理步驟和潛在錯誤。

子詞分詞器在LLM中引入了許多奇怪的問題，例如LLM無法拼寫單詞、反轉某些單詞、處理罕見token，而這些在位元組級token化中是不存在的。^[8]

MOE與Mamba模型的結合（Mamba Mixture of Experts，MOE）

MoE Mamba代表了混合專家（MoE）技術與Mamba架構的開創性結合，增強了狀態空間模型（SSM）在語言建模中的效率和可延伸性。該模型充分利用了MoE和SSM的優勢，顯著提高了訓練效率——所需的訓練時間比其前身Mamba減少了2.2倍，同時保持了與其相匹配的效能。MoE Mamba通過將選擇性狀態空間建模與基於混合專家技術的處理相結合，展示了更高的效率和效能，為未來擴充SSM來進行數百億級別參數的模型研究提供了有潛力的途徑。該模型的設計涉及互相交替的Mamba層和MoE層，使其能夠有效地整合所有的序列上下文，並為每個token應用最相關的專家模型。^[9][10]

Mamba在視覺上的使用（Vision Mamba，Vim）

Vision Mamba（Vim）將SSM用於視覺資料處理。它採用雙向Mamba塊進行視覺序列編碼並於此減少了視覺任務中通常與self-attention機制相關的計算需求。經過ImageNet分類資料集、COCO對象檢測和ADE20k語意分割的測試，Vim展示了更好的效能，並且能夠以較低的計算資源處理高解析度圖像。這使得Vim成為未來視覺表徵學習進步的可延伸模型。^[11]

Jamba

Jamba是一種將Transformer和Mamba SSM架構相結合的新型架構，由AI21 Labs開發，擁有520億個參數，是迄今為止建立的參數最多的Mamba變種。它有一個包含256k token的上下文窗口。 ^[12]

影響和未來方向

Mamba模型顯著影響了深度學習領域，通過其選擇性狀態空間模型（Selective State Space Model）提高了長序列處理的效率，在自然語言處理中用於長上下文語言建模和文字摘要，提升了推理速度 ^[12][13]。在電腦視覺中，它應用於醫學圖像分割和疾病診斷，顯著減少了GPU主記憶體使用^[14][15]。

未來方向包括探索硬體感知計算（如Parallel Associative Scan）最佳化資源利用^[16][17]。此外，增強模型的可信賴性（如安全性與公平性）是新興研究重點^[18][19]。

參見

• 語言建模
• Transformer（機器學習模型）
• 狀態空間模型
• 迴圈神經網路

注釋