Apriel-H1: 透過 Mamba 混合架構蒸餾高效能推理模型
ServiceNow AI 開發了 Apriel-H1,這是一個 15B 推理模型系列,將全注意力(full-attention)教師模型轉換為 Mamba 混合架構。這種方法在推理品質損失極小的情況下,實現了高達 2.1 倍的吞吐量提升,證明了透過針對性的蒸餾,可以為現有的強大模型進行效率改造。
核心洞察:能力匹配蒸餾
蒸餾推理模型需要使用與要保留的特定能力相匹配的數據,而不是通用的預訓練數據。最初嘗試使用預訓練數據或預訓練與 SFT 數據混合進行蒸餾,導致推理品質大幅下降。
成功的方法是利用來自教師模型 SFT 數據集的高品質推理軌跡(reasoning traces),包括:
- 多步驟數學證明
- 具有清晰邏輯依賴關係的編碼任務
- 具有詳細解釋鏈的科學分析
這是必要的,因為用 Mamba 的線性遞歸取代注意力機制(例如檢索與歸納頭)會破壞用於推理的計算路徑。高品質的推理軌跡提供了混合模型發現通往相同結果之新路徑所需的顯式結構。此外,團隊在蒸餾時使用了 reverse KL divergence(溫度為 1),因為其尋求模式(mode-seeking)的特性會鼓勵學生模型致力於教師模型的高置信度、結構化預測。
Apriel-H1 性能與基準測試
旗艦模型 Apriel-H1-15b-Thinker-SFT 與全注意力教師模型相比,實現了 2.1 倍的吞吐量。對推理基準測試的影響如下:
| 基準測試 | 教師模型分數 | Apriel-H1-15b-Thinker-SFT 分數 |
|---|---|---|
| MATH500 | 0.90 $\rightarrow$ 0.92 | 進步 |
| MTBench | 8.30 $\rightarrow$ 8.58 | 進步 |
| GSM8k | 0.97 $\rightarrow$ 0.95 | 輕微下降 |
| GPQA | 0.59 $\rightarrow$ 0.55 | 輕微下降 |
| AIME24 | 0.70 $\rightarrow$ 0.65 | 輕微下降 |
在 Apriel-H1 系列中,檢查點(checkpoints)的 Mamba 層數範圍為 25 到 40 層(總共 50 層)。雖然 H-30 變體提供了最佳平衡,但 H-40 變體為延遲敏感型工作負載實現了最高的吞吐量(高達 3.4 倍)。
分階段蒸餾流程
將全注意力模型轉換為混合模型無法一蹴而就。ServiceNow AI 採用了三個階段的程序:
階段 1:初始層替換
透過對 MMLU 進行留一法(Leave-One-Out, LOO)分析,團隊識別出 25 個最不重要的層。這些層被替換為使用 Mamba-in-Llama (MIL) 初始化的混合器(mixers),並進行端到端蒸餾。
階段 2:漸進式轉換
由於 LOO 分析在超過 25 層後會變得不可靠,團隊實施了 MIL-Mamba-Replacement (MMR)。這種動態啟發式方法會初始化一個 Mamba mixer,運行 100 個訓練步驟並記錄蒸餾損失。收斂到較低損失的層會被優先替換。此過程是增量進行的(25 $\rightarrow$ 27 $\rightarrow$ 30 $\rightarrow$ 34 $\rightarrow$ 37 $\rightarrow$ 40 層)。
階段 3:最終 SFT
在達到目標 Mamba 層數後,進行最後一次 SFT,直到推理性能穩定。最終的 Apriel-H1-15b-Thinker-SFT 模型是在經過 55.9B 蒸餾 token 和 20.9B SFT token 後產生的。
實作與可重現性
Apriel-H1 是使用 Fast-LLM 構建的,這是一個開源(Apache 2.0)訓練框架,將注意力與 Mamba 視為模組化混合介面。這允許透過配置中的 pattern 欄位指定層順序,從而實現區塊類型的自由交換。
生產部署
Apriel-H1 已與 Hugging Face Transformers 和 vLLM 整合。vLLM 整合利用 Mamba 快取操作進行連續批處理(continuous batching)、前綴快取(prefix caching)和分塊預填充(chunked prefill)。作者指出,雖然吞吐量增益顯著,但由於工具鏈尚在成熟中,目前部署混合模型需要自定義代碼和仔細的數值驗證。
技術規格
- Mamba 超參數: 狀態大小 16,DT 秩 16,內部維度 4096。
- 架構: Mamba-1 與 Attention 層的混合。
- 訓練數據: 旗艦模型總計 76.8B tokens。