科技

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

04月20日 12:36 新浪網 tech-auto-hilite

新智元報導

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

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