Mamba LLM 架構簡介:機器學習的新范式
探索 Mamba LLM 的強大功能,Mamba LLM 是來自一流大學的變革性架構,重新定義了 AI 中的序列處理。語言模型是一種經過訓練的機器學習模型,用于在自然語言上執行概率分布。它們的架構主要由多層神經網絡組成,例如遞歸層、前饋層、嵌入層和注意力層。這些層組合在一起以處理給定的輸入文本并生成輸出預測。
2023 年底,卡內基梅隆大學和普林斯頓大學的研究人員發表了一篇研究論文,揭示了一種名為 Mamba 的大型語言模型 (LLM) 的新架構。Mamba 是一種與序列建模有關的新狀態空間模型架構。它是為了解決 transformer 模型的一些限制而開發的,尤其是在處理長序列時,并且已經顯示出有希望的性能。
Mamba
Mamba 是一種新的 LLM 架構,它集成了結構化狀態空間序列 (S4) 模型來管理冗長的數據序列。S4 結合了遞歸、卷積和連續時間模型的最佳功能,可以有效且高效地模擬長期依賴關系。這使它能夠處理不規則采樣的數據,具有無限的上下文,并在整個訓練和測試過程中保持計算效率。
Mamba 擴展了 S4 范式,帶來了幾項值得注意的改進,尤其是在處理時變作方面。它的架構圍繞著一種特殊的選擇機制展開,該機制根據輸入修改結構化狀態空間模型 (SSM) 參數。
因此,Mamba 可以通過只關注序列中的關鍵信息來成功過濾掉不太重要的數據。根據 Wikipedia 的說法,“該模型從時不變框架過渡到時變框架,這會影響系統的計算和效率。
主要特點和創新
Mamba 偏離了傳統的 attention 和 MLP 塊,使其與眾不同。這種簡化使模型更輕、更快,并且與序列的長度呈線性縮放,這是其前輩都無法實現的壯舉。
Mamba 的關鍵組件包括:
選擇性狀態空間 (SSM):根據當前輸入有選擇地處理信息的遞歸模型是 Mamba SSM 的基礎。這使他們能夠過濾掉無關數據并專注于相關信息,這可能會帶來更高效的處理。
簡化的架構: Mamba 用一個單一的、有凝聚力的 SSM 塊取代了 Transformers 錯綜復雜的注意力和 MLP 塊。這旨在加速推理并降低計算復雜性。
硬件感知并行性:Mamba 的性能可能會更好,因為它使用循環模式和專為硬件效率而創建的并行算法。
另一個關鍵組成部分是線性時間不變性 (LTI);LTI 是 S4 模型的核心功能之一。這一特性表明,模型的參數在所有時間步中保持不變,從而保持模型動力學的一致性。使用 LTI 構建序列模型更容易、更有效,LTI 是遞歸和卷積的基礎。
Mamba LLM 架構詳細信息
Mamba 的架構進一步強調了機器學習進步的重要性。它通過引入選定的狀態空間模型 (SSM) 層來修改模型處理序列的方式。這使 Mamba 能夠做兩件極其重要的事情:
關注相關信息 – Mamba 可以通過為每個輸入分配不同的權重來為任務確定更多預測數據的優先級。
動態適應輸入 – 由于模型能夠適應輸入,Mamba 可以輕松處理各種序列建模工作。
因此,Mamba 可以以前所未有的效率處理序列,這使其成為涉及冗長數據序列的任務的完美選擇。
Mamba 的設計理念基于對現代硬件功能的認識。它旨在充分利用 GPU 計算能力,保證:
**優化的內存使用:**通過設計 Mamba 的狀態擴展以適應 GPU 的高帶寬內存 (HBM),可以縮短數據傳輸時間并加快處理速度。
**最大化并行處理:**Mamba 通過協調其計算與 GPU 計算的并行性質,達到了為序列模型建立新基準的性能水平。
Mamba VS Transformer
GPT-4 等 Transformer 的引入進入了自然語言處理 (NLP) 領域,并為多項自然語言任務建立了基準。長期以來,較長的序列一直是變壓器的眼中釘,因為它們會嚴重阻礙其效率。
這個缺陷正是 Mamba 擅長的地方。也就是說,mamba 可以比變壓器更快地處理冗長的序列,并且由于其獨特的架構而更加簡單。
Transformer 架構
Transformer 非常擅長處理數據序列,例如語言模型的文本。它們同時處理完整的序列,這與早期按順序處理數據的模型不同。這種固有的功能使他們能夠捕獲數據中錯綜復雜的關系。它們使用注意力機制,使模型能夠在生成預測時專注于各種序列段。使用三組權重來計算此關注度:從輸入數據獲取的值、鍵和查詢。
序列中的每個元素都相對于其他每個元素進行加權,以指示它應該有多少權重(或 “注意力”)來預測序列中的下一個元素。Transformer 由兩個主要塊組成:創建輸出的 decoder 和處理輸入數據的 encoder。編碼器由幾層組成 - 每層都有兩個子層:一個基本的、按位置的、完全連接的前饋網絡和一個多頭自注意力機制。為了幫助訓練深度網絡,每個子層都使用殘差連接和歸一化。
與編碼器一樣,解碼器由兩層和兩個子層組成,但它也增加了第三個子層,用于處理編碼器輸出上的多頭注意力。由于解碼器的 sequential 性質,解碼器的 autoregressive 屬性被保留下來,這將對位置的預測限制為僅考慮較早的位置。
因此,Transformers 試圖通過利用更復雜的注意力過程來解決冗長序列的問題,但 Mamba 采取了不同的方法。
Mamba 架構
Mamba 利用選擇性狀態空間。這種方法解決了Transformers 在長序列下計算效率低下的問題。Mamba 的架構使更快的推理和線性序列長度擴展成為可能,為序列建模創造了一種新的范式,隨著序列越來越長,這種范式可能會被證明會更加有效。由于我們在上面深入探討了 Mamba 的架構,因此我們不會在這里深入討論。
特征 | 變壓器 | 曼巴 |
---|---|---|
建筑 | 基于注意力 | 基于 SSM |
復雜性 | 高 | 降低 |
推理速度 | O(n) | O(1) |
訓練速度 | O(注2) | O(n) |
值得注意的是,盡管 SSM 與 Transformer 相比具有許多優勢,但后者可以處理比 SSM 存儲在內存中更長的序列,需要更少的數據來學習類似的任務,并且在需要從輸入上下文檢索或復制的任務中優于 SSM,即使參數更少。
開始使用 Mamba
如果您有興趣使用 Mamba 或在項目中利用它,則必須具備以下條件:
- Linux的
- 英偉達圖形處理器
- PyTorch 1.12+ 版本
- CUDA 11.6+ 的
要從 Mamba 存儲庫安裝所需的軟件包,請使用一些簡單的 pip 說明:
- [Option] : 在 Mamba 塊內部使用的簡單因果 Conv1d 層的有效實現。
pip install causal-conv1d>=1.2.0
pip install mamba-ssm
:核心 Mamba 包。它也可以通過 from this repository從源代碼構建。
pip install .
如果 PyTorch 版本導致兼容性問題,可以使用與 switch 一起使用來幫助。這些模型是在 Pile 和 SlimPajama 數據集等大型數據集上訓練的,旨在滿足各種計算要求和性能基準。
pip``--no-build-isolation
Mamba 模型具有多個接口級別,但主模塊是包裝選擇性 SSM 的 Mamba 架構塊。
# Source: Mamba Repository
import torch
from mamba_ssm import Mambabatch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(# This module uses roughly 3 * expand * d_model^2 parametersd_model=dim, # Model dimension d_modeld_state=16, # SSM state expansion factord_conv=4, # Local convolution widthexpand=2, # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape
Mamba 的應用
Mamba LLM 的推出是 LLM 架構領域的一個重大潛在轉變。Mamba 更快、更高效且可擴展,可以毫不費力地以高性能標準處理長序列,這解釋了為什么它將在塑造復雜 AI 系統的未來方面發揮關鍵作用。
也就是說,下一波 AI 創新可能由 Mamba 的有效性和性能帶來,它為創建越來越復雜的模型和應用程序鋪平了道路。它的潛在影響力是巨大的,包括音頻和語音處理應用程序、長篇文本分析、內容創建、實時語言翻譯等。
這可能會帶來革命性的行業包括:
醫療: Mamba 可以通過快速分析遺傳數據來加快開發個性化健康藥物的過程。
**金融:**可以部署 Mamba 來分析長期市場趨勢,從而獲得更準確的庫存預測。
顧客服務: Mamba 能夠為監控長篇討論的聊天機器人提供支持,從而改善客戶溝通
結論
Mamba 不僅為當前的序列模型提供了增量改進;它重新定義了什么是可能的。隨著它的推出,人工智能的歷史將翻開新的篇章,計算效率低下和序列長度限制最終變得過時。在過去的幾年里,我們看到了 AI 從 RNN 到Transformers,再到現在的 Mamba,每一步都離實現能夠與人類相媲美的 AI 進行深度思考和信息處理更近一步。Mamba 體現了革命性的精神,通過其選定的狀態空間方法和線性時間縮放推動 AI 領域向前發展。Mamba 標志著人工智能前景廣闊的開始。這是一種為未來設計的范式,并將以其無限的潛力對 AI 產生重大影響。