Transformer+Mamba黃金組合！長文推理效能飆升3倍，效能還更強

新智元報道

編輯：LRS

【新智元導讀】Nemotron-H模型混合了Transformer和Mamba架構，使長文字推理速度提升3倍，同時還能保持高效能，開源版本包括8B和56B尺寸。訓練過程採用FP8訓練和壓縮技術，進一步提高了20%推理速度

過去幾年，Transformer雖穩坐AI架構「鐵王座」，但其二次方複雜度帶來的算力消耗和長序列處理瓶頸，限制了大模型在推理階段處理長文字。

Mamba憑藉「線性複雜度」異軍突起，非常適合長序列任務，有望成為Transformer架構的替代品，但在處理全域性關係上偏弱。

Mamba+Transformer混合架構可以將二者的優勢互補，實現「效率」和「效能」的雙豐收。

最近英偉達釋出了Nemotron-H系列模型，模型尺寸為8B和56B（蒸餾版本47B），用Mamba-2層替換了Transformer中的自注意力層，關鍵創新在於對Transformer和Mamba的平衡，實現了高效處理長上下文的同時，還不犧牲模型效能，顯著提高了推理速度，並且記憶體佔用更少。

論文連結：https://arxiv.org/pdf/2504.03624

實驗結果表明，Nemotron-H模型在準確度上優於同尺寸的開源Transformer模型（例如Qwen-2.5-7B/72B和Llama-3.1-8B/70B），同時在推理速度上提速3倍。

為了進一步提高推理速度並降低推理時所需的記憶體量，研究人員使用一種剪枝和蒸餾壓縮技術MiniPuzzle，將56B尺寸的模型蒸餾為NemotronH-47B-Base，在基準的準確率上與56B模型相當，同時推理速度提升20%

論文中還提出了一種基於FP8的訓練方案，使56B模型實現了與BF16訓練相當的效能。

Nemotron-H架構

Nemotron-H模型由Mamba-2、自注意力層和前饋神經網路（FFN）層組成，其中總層數的8%為自注意力層，均勻分佈在模型中。

即，Nemotron-H-8B模型包含52層，其中4層為注意力層；Nemotron-H-56B模型包含118層，其中10層為注意力層，其餘層平均分配為FFN層和Mamba-2層。

為了與標準Transformer模組的結構一致，研究人員提出三條設計準則：模型的第一層必須是Mamba-2層，最後一層必須是FFN層，並且自注意力層總是位於FFN層之前。

模型的具體引數

需要注意的是，8B和56B模型在FFN隱藏維度、注意力查詢頭和狀態維度設定上存在差異；

對於Mamba-2層，保持預設的頭維度（64）、擴充套件因子（2）和卷機視窗（4）；
使用RMSNorm進行歸一化；
不適用位置嵌入；
模型的嵌入層和輸出層使用了獨立的權重，沒有使用線性層的偏置權重，也沒有使用dropout；
在每個Mamba-2層、自注意力層和FFN層周圍都加入了殘差跳躍連線。

訓練過程

資料準備

訓練資料從來源上大體上可以分為多語言、網路爬取、學術、程式碼、維基百科和數學資料，這種資料組合可以全面覆蓋通用知識，同時在程式設計和數學等領域培養強大的專業能力。

其中多語言資料涵蓋了九種語言：德語、西班牙語、法語、義大利語、葡萄牙語、中文、日語、韓語和俄語。

研究人員設計資料組合時，確保所有相同質量的資料來源權重相似，而高質量的資料來源權重會高於低質量的資料來源。

在訓練56B尺寸的模型時，使用了大約20萬億個token的資料，其中，網頁爬取資料佔比最大，達到了59%，其次是程式碼資料，佔20%，學術內容佔8.8%

在預訓練Nemotron-H基礎模型時，研究人員採用了分階段的資料混合方法：

第一階段，使用一種促進資料多樣性的資料組合；在第二和第三階段，主要使用高質量的資料集（例如維基百科），其中第二階段為訓練進度達到60%時，第三階段為訓練進度達到80%時；第四階段則使用最後3800億個訓練token

在後訓練階段，研究人員調整了資料的分佈，更加註重有監督的微調（SFT）樣本。

FP8訓練策略

Nemotron-H訓練的一個創新在於使用8位浮點數（FP8），在降低記憶體需求和計算成本的同時，還能保持模型的質量，主要包括以下關鍵點：

採用逐張量（per-tensor）的當前縮放技術，以提高訓練的穩定性。

將模型中最初的四個和最後四個矩陣乘法（GEMM）操作保持在BF16精度，以確保關鍵部分的高精度處理；

在訓練過程中，FP8訓練逐漸與BF16訓練收斂，最終達到類似的效能水平。

實驗結果表明，FP8訓練在多種基準測試中能夠達到或超過BF16訓練的效能，不僅提高了訓練效率，還透過在MMLU、常識理解、程式碼生成和GSM8K等基準測試中的表現，證明了其在保持或提升模型質量方面的有效性

模型壓縮

為了進一步提高模型部署的效率，研究人員開發了一種新型壓縮框架MiniPuzzle的，結合了剪枝、神經架構搜尋和知識蒸餾技術。

MiniPuzzle壓縮框架的工作流程如下圖所示，展示了從預訓練模型到壓縮模型的轉變過程，包括重要性評估、神經架構搜尋和蒸餾等步驟。

MiniPuzzle方法包含五個階段：

1. 重要度估計：分析每一層對模型效能的貢獻。



defimportance_estimation(model, dataset):# Compute per-layer importance scores    scores = []for layer in model.layers:# Zero out layer outputs and measure impact on loss        scores.append(measure_impact_on_loss(model, layer, dataset))return scores

2. 層重要度分析：研究人員需要分析出哪些層對模型效能的貢獻最大。

3. 條件神經架構搜尋：探索不同壓縮架構方案，在每個候選壓縮模型中保留不同的層。

4. 記憶體與效能權衡：根據記憶體使用量和準確度對模型進行評估，對候選架構的記憶體負載與基準效能進行權衡。

5. 知識蒸餾：透過訓練，使壓縮後的模型能夠匹配甚至超越原始模型的能力。

透過這一過程，Nemotron-H-56B模型成功被壓縮為Nemotron-H-47B模型，引數減少了16%，同時保持了相當的準確度，並將推理吞吐量提高了20%。

實驗結果

Nemotron-H模型在效能和效率方面相較於其他基於Transformer的模型取得了顯著進步。

推理吞吐量

混合架構使得推理速度大幅提升，尤其是在處理長序列時：

Nemotron-H-56B的推理吞吐量比Qwen-2.5-72B和Llama-3.1-70B高出多達3倍。
Nemotron-H-8B在類似準確度水平下，比Qwen-2.5-7B的吞吐量高出1.8倍。

效率提升在處理長序列（例如65,536個token）時尤為明顯，突顯了Mamba層在輸出token時計算複雜度固定的優勢。

多基準測試中的準確度

儘管架構發生了變化，但Nemotron-H模型在廣泛的基準測試中仍保持了強勁的效能表現。

在評估的17項任務中，Nemotron-H-56B在16項任務中的表現優於Llama-3.1-70B，在數學推理任務上表現尤為出色。

應用與多功能性

Nemotron-H模型可以進行擴充套件，以適應各種應用場景。

視覺-語言能力：基礎模型透過NVLM-D架構擴充套件，建立了視覺-語言模型（VLM），在VQAv2、GQA和VizWiz等基準測試中表現出色，顯示出混合架構對多模態任務的適應性。
程式碼生成：模型在與程式碼相關的任務上表現尤為出色。訓練資料中包含大量程式碼資料（佔比20%），使得模型能夠理解和生成多種程式語言的高質量程式碼。
長文字處理：混合架構的一個顯著優勢是能夠高效處理長文字。Nemotron-H-8B模型經過專門的長文字處理能力微調，在RULER基準測試和其他長文字評估任務中表現出色。
針對不同能力的資料分佈：研究人員針對不同的訓練階段精心調整了資料分佈，以培養特定的能力，透過調整不同資料型別（網頁爬取、程式碼、數學、學術等）的比例，可以在不需要架構變更的情況下增強模型的特定能力。比如針對STEM能力最佳化時，訓練資料增加了數學和程式碼內容的比重。

Mamba架構簡介

原版Mamba架構

Mamba是一種新型的序列建模架構，透過選擇性狀態空間模型（Selective State Space Model, SSM）和硬體最佳化演算法，將計算複雜度降低到線性級別O(L），能夠高效處理長達百萬級的序列長度，推理速度比Transformer快5倍，在短序列任務中也實現了超越Transformer的效能。