Falcon Mamba 7B 版本說明

Falcon Mamba 7B 是由科技創新研究院(Technology Innovation Institute,簡稱 TII)推出的全新開放存取模型,證明了狀態空間語言模型(State Space Language Models,SSLM)在效能上可以與最先進(SoTA)的 Transformer 競爭,同時不會因注意力機制而產生線性增加的計算與記憶體成本。

純粹 Mamba 架構與序列擴展性

Falcon Mamba 基於原始的 Mamba 架構,利用選擇性的狀態空間克服 Transformer 在序列擴展性上的限制。為了確保大規模訓練的穩定性,TII 在基礎 Mamba 設計上額外加入了 RMS 正規化層。

此架構相較於基於注意力的模型提供了兩大技術優勢:

  • 任意序列長度:模型能處理任意長度的序列,且不會導致記憶體需求同步增加,因而能在單張 A10 24GB GPU 上運行。
  • 恆定生成時間:產生新 token 所需的時間不受上下文大小影響,保持恆定。

模型訓練與資料

Falcon Mamba 使用約 5,500GT 的資料進行訓練。資料集主要由 RefinedWeb 資料構成,並輔以來自公開來源的高品質技術資料與程式碼。訓練過程在大部分時間採用恆定學習率,最後以短暫的學習率衰減階段結束,該階段加入少量高品質精選資料以提升最終效能。

效能基準

透過 lm-evaluation-harness(使用 Hugging Face 分數正規化)與 lighteval 進行的評測顯示,Falcon Mamba-7B 是迄今為止最強的純 SSM 模型,常常能與或超越混合 SSM‑attention 模型以及多種基於 Transformer 的模型。

新版排行榜結果

模型名稱 IFEval BBH MATH LvL5 GPQA MUSR MMLU-PRO 平均
Falcon Mamba-7B 33.36 19.88 3.63 8.05 10.86 14.47 15.04
recurrentgemma-9b 30.76 14.80 4.83 4.70 6.60 17.88 13.20
Falcon2-11B 32.61 21.94 2.34 2.80 7.53 15.44 13.78
Meta-Llama-3.1-8B 12.70 25.29 4.61 6.15 8.98 24.95 13.78
gemma-7B 26.59 21.12 6.42 4.92 10.98 21.64 15.28

LLM 排行榜(第一版)結果

模型名稱 ARC HellaSwag MMLU Winogrande TruthfulQA GSM8K 平均
Falcon Mamba-7B 62.03 80.82 62.11 73.64 53.42 52.54 64.09
recurrentgemma-9b 52.00 80.40 60.50 73.60 38.60 42.60 57.95
Falcon2-11B 59.73 82.91 58.37 78.30 52.56 53.83 64.28
Meta-Llama-3-8B 60.24 82.23 66.70 78.45 42.93 45.19 62.62

記憶體與吞吐量分析

Falcon Mamba 在序列處理的預填(prefill)與解碼(decode)階段相較於 Transformer 模型展現出顯著的效率提升。

預填效率

處理提示(prefill)主要有兩種方法:

  1. 平行預填:將整段提示平行處理,以最大化 GPU 使用率。此模式下記憶體需求會隨提示長度而成比例增加,因需儲存隱藏狀態。但即使如此,Falcon Mamba 仍能容納比 Transformer 更長的序列。
  2. 序列預填:逐 token 或分塊處理提示。此方式允許 SSM 在不受記憶體擴展限制的情況下處理任意長度的提示。

生成吞吐量

在使用 H100 GPU、提示長度為 1、最多產生 130k token 的測試中,Falcon Mamba 能保持恆定的吞吐量與穩定的 CUDA 峰值記憶體。相較之下,Transformer 模型的峰值記憶體會隨序列長度上升而增加,且生成速度會下降。

實作與可取得性

Falcon Mamba 已整合至 Hugging Face transformers 套件(版本 >4.45.0),可與 AutoModelForCausalLM 以及 pipeline API 直接搭配使用。

模型變體

  • 基礎模型:標準的預訓練模型。
  • 指令模型:經過 50 億 token 監督式微調(SFT)以提升指令式任務表現的版本。
  • 量化版本:提供 4 位元轉換的基礎與指令模型,適用於支援 bitsandbytes 的 GPU。

最佳化

使用者可在載入模型後執行 torch.compile(model) 以獲得更快的推論速度。

Sources