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

作者：新智元 2025-04-21 09:07:00

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. 重要度估計：分析每一層對模型性能的貢獻。

def importance_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

def importance_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的性能。