英偉達發布 Jet-Nemotron 系列小模型，理論最大加速比 56 倍

注意力機制或成為小模型之爭的最前沿。

作者丨梁丙鑒

編輯丨馬曉寧

小模型 SOTA 之爭持續升溫，英偉達也再次下場。

在最新論文中，其最新發布的 Jet-Nemotron 系列混合架構語言模型在多項基準測試中超越或追平了 Qwen3，Qwen2.5，Gemma3 和 Llama3.2 的精度，同時實現了高達 53.6 倍的吞吐量加速和 6.1 倍的預填充加速。在與 DeepSeek-V3-Small 和 Moonlight 等最新 MoE 全注意力模型的比較中，Jet-Nemotron-2B 也在 MMLU 和 MMLU- pro 等數學任務的處理上表現出了更高的精度。

如上圖所示，在 NVIDIA H100 GPU 上，以 64K token 的上下文長度測量生成吞吐量。Jet-Nemotron-2B 在 MMLU-Pro 上比 Qwen3-1.7B-Base 具有更高的準確性，同時生成吞吐量提高了 47 倍。Jet-Nemotron-4B 盡管模型規模更大，但生成吞吐量仍比所有參數小于 2B 的全注意力模型更高。

一句話來說，就是以更短的時間實現了更高的準確率。

有意思的是，Jet-Nemotron 系列模型的背后，還是小模型中的老朋友 Qwen2.5-1.5B 和 Qwen2.5-3B。

由于預訓練的成本高企不下，設計新的語言模型架構具有具有極大的挑戰和風險。在此背景下，Jet-Nemotron 的訓練采用了 Post Neural Architecture Search （PostNAS）方案，這是種以完成預訓練的全注意力模型為基礎，尋找最優注意力模塊配置的方法。簡單來說，英偉達選擇對既有的模型“開刀”進行針對性優化，而從結果來看，手術很成功。

論文鏈接：https://arxiv.org/pdf/2508.15884

01

站在大模型的肩膀上

Qwen2.5 系列本身就是同尺寸模型中的佼佼者，提供了一個極高的性能起點。PostNAS 的目標不是從零學習知識，而是在保留其核心知識的前提下，對其進行“架構手術”，優化其推理效率。

與主流的語言模型架構設計方法不同，PostNAS 從預訓練的全注意力模型開始，繼承其 MLP （多層感知器）權重并使其保持凍結，僅對注意力機制部分進行優化，這種策略顯著降低了訓練成本。

PostNAS 的后續流程包括四個關鍵步驟：

1.全注意力機制的位置選擇與消除

在模型中保留少數全注意力層，對于在檢索等任務中保持高精度至關重要，然而這些層的最佳位置并不明確。解決該問題的常見策略之一是將少數全注意力層整合，在全注意力模型的部分層上均勻應用，其余層使用線性注意力。

這種均勻策略并非最優，英偉達團隊提出了一種自動方法，用于高效確定全注意力層的位置。

其核心思路是為預訓練的全注意力模型增加替代線性注意力路徑，構建一個一次性超網絡。在訓練過程中，研究人員會在每一步隨機采樣一條活躍路徑，形成一個子網絡，并對該子網絡使用特征蒸餾損失進行訓練。訓練完成后，使用束搜索就可以確定全注意力層在給定約束條件下的最佳位置。其中搜索目標與任務相關：對于 MMLU，研究人員選擇在正確答案上損失最低的配置（即最大化?），而對于數學和檢索任務，則選擇準確率最高的配置。

上圖展示了 Qwen2.5-1.5B 的搜索結果。通過依次將每一層配置為全注意力機制，同時將所有剩余層設置為線性注意力機制，研究人員從超網絡中提取到了相應的子網絡。每個子網絡的準確率和損失可以在給定任務下進行評估，并通過熱力圖可視化結果。英偉達團隊的分析揭示了三個關鍵發現：

• 在預訓練的全注意力模型中，并非所有注意力層都同等重要。對于 MMLU 任務，只有兩層表現出關鍵重要性，而對于檢索任務，僅有兩到三層特別關鍵。

• 不同的注意力層對不同的能力有所貢獻。對 MMLU 準確率起關鍵作用的注意力層，不一定對檢索任務重要。

• 對于數學推理等復雜任務，注意力重要性的模式變得更加復雜。幸運的是，為 MMLU 和檢索任務識別出的關鍵層集合已經涵蓋了數學所需的大部分關鍵層。

2.線性注意力模塊選擇

在全注意力層的最優位置之后，還需要確定最適合的線性注意力模塊。英偉達團隊評估了六種 SOTA 線性注意力模塊，其中包括 RWKV7, RetNet, Mamba2, GLA, Deltanet和 Gated DeltaNet。分析發現，RWKV7 的訓練吞吐量與其它線性注意力模塊相比顯著較低，訓練吞吐量顯著較低，而 Gated DeltaNet 則實現了最佳的整體準確率。

這歸因于兩個因素的結合：數據相關門控機制可以實現模型在關注當前 token 和歷史狀態之間的動態控制，Delta 規則通過當前 token 的信息增量更新歷史狀態，在此基礎上節省有限的狀態內存。

值得注意的是，得益于 PostNAS 框架的低訓練成本，系統評估既有線性注意力模塊在不同任務上的準確性、訓練效率和推理速度成為可能。因此隨著新的線性注意力模塊出現，當他們表現出超越 Gated DeltaNet 的性能時，也可以被快速替換。

3.新的注意力模塊設計

Jet-Nemotron 系列模型還應用了一種全新的注意力模塊 JetBlock 以擴展線性注意力容量。

不同的此前依賴靜態卷積核的方案，英偉達團隊引入了一個卷積核生成模塊，該模塊根據輸入特征動態生成卷積核。如下圖所示，該模塊與 Q/K/V 投影層共享相同輸入，首先通過線性縮減層提高效率，縮減比例為 8，然后應用 SiLU 激活函數，最后通過一個最終線性層輸出卷積核權重。

研究人員還發現，一旦對值（V）標記應用動態卷積，對查詢（Q）和鍵（K）標記的靜態卷積就可以被移除，而對最終模型精度的影響可以忽略不計。最終實驗采用了這種設計，JetBlock 在與之前線性注意力模塊的比較中，在數學推理和檢索任務上比 Gated DeltaNet 表現出了更高的精度，同時保持了相似的效率。

4.硬件感知架構搜索

PostNAS 還應用了一種硬件感知架構搜索方案，以識別最優的架構超參數，包括鍵/值維度和注意力頭的數量。

雖然參數數量通常被用作衡量語言模型效率的指標，但它實際上與硬件的生成效率并不直接相關。英偉達團隊將生成吞吐量作為選擇架構超參數的直接目標，通過硬件感知架構搜索，發現了能夠實現相似生成吞吐量、同時使用更多參數以獲得更高精度的架構超參數。

這一原理在于，KV 緩存大小是影響長上下文和長生成吞吐量的最關鍵因素。當 KV 緩存大小保持不變時，不同參數數量的模型表現出相似的生成吞吐量。研究人員認為，其原因是解碼階段通常受內存帶寬限制而非計算限制。在長上下文場景中，KV 緩存往往比模型權重占用更多內存。減小其大小可以減少每步解碼的內存傳輸時間，并支持更大的批量大小，從而提高生成吞吐量。

02

理論最大加速比 56 倍

Jet-Nemotron 系列模型包含兩個不同參數規模的版本：Jet-Nemotron-2B 和 Jet-Nemotron-4B，分別基于 Qwen2.5-1.5B 和 Qwen2.5-3B 構建。

基準測試結果亦于同期發布。兩者在所有六個評估領域中的表現均與前沿完整注意力模型（Qwen3-1.7B-Base）相當甚至更好，具有顯著更少的完整注意力層和更小的 KV 緩存大小。在生成吞吐量上，Jet-Nemotron-2B 和 Jet-Nemotron-4B 分別比 Qwen3-1.7B-Base 高出 47 倍和 21 倍。

準確性方面，Jet-Nemotron-2B 在 MMLU(-Pro)和 BBH 上的吞吐量比 Qwen3-1.7B-Base 高 47 倍，緩存大小則小 47 倍，同時在 MMLU、MMLU-Pro 和 BBH 上顯著提高了準確率。JetNemotron-2B 甚至超過了參數量更大的 MoE 模型如 DeepSeek-V3-Small 和 Moonlight，兩者分別有 2.2B 激活參數和 15B 總參數。當擴展到 4B 參數時，Jet-Nemotron-4B 仍然比 Qwen3-1.7B-Base 在吞吐量上高 21 倍。

數學任務上，Jet-Nemotron-2B 實現了平均 49.6 的準確率，比 Qwen3-1.7B-Base 高出 6.3，同時速度提升了 47 倍。相比之下，之前的線性注意力機制和混合模型在數學任務上則遠遠落后于 Qwen3。

作為 Jet-Nemotron 系列模型基礎的 Qwen2.5 和 Qwen3 在常識推理任務上的性能相對較弱，但 Jet-Nemotron-2B 仍然表現出色，平均準確率達到了 62.0，超越了所有基線模型。此外在檢索和編程等任務中，Jet-Nemotron-2B 和 Jet-Nemotron-4B 仍在保持吞吐量優勢的前提下，表現出全面超越基線模型的性能。

英偉達團隊量化了此種吞吐量優勢，下圖展示了 Qwen3-1.7B-Base 和 Jet-Nemotron-2B 在不同上下文長度下的吞吐量對比。在預填充階段，Jet-Nemotron-2B 在較短的上下文長度（4K 和 8K）下初始速度比 Qwen3-1.7B-Base 快 1.14 倍和 1.15 倍。通過設計更優化的 JetBlock 內核實現，這還可以進一步改進。隨著上下文長度的增加，線性注意力機制變得突出，使得 Jet-Nemotron-2B 在 256K 上下文長度下實現了 6.14 倍的加速。

在解碼階段，Jet-Nemotron-2B 始終大幅優于 Qwen3-1.7B-Base。由于 Jet-Nemotron-2B 包含 2 個全注意力層，每組鍵值狀態有 2 組，而 Qwen3-1.7B-Base 包含 28 個全注意力層，每組鍵值狀態有 8 組，因此 Jet-Nemotron-2B 的理論最大加速比是 14 × 4 = 56 倍。據英偉達介紹，在其吞吐量測試平臺中，Jet-Nemotron-2B 在 4K 上下文長度下實現了 15.6 倍的加速，在 256K 上下文長度下實現了高達 53.6 倍的加速，幾乎達到了理論上的上界。

03

One more thing

從Jet-Nemotron 的經驗來看，在日漸升溫的小模型之爭中，對現有模型進行深度優化、剪枝、架構改造的能力也將成為多方爭奪的前沿，而注意力機制只是起點。

這種基于成熟預訓練模型進行架構優化的策略，不僅為英偉達節省了天價開發成本，而且借助 Qwen2.5 和 Qwen3 經過驗證的性能保障了下限。其成功之處不在于讓英偉達在基準測試榜單上占據幾個名次，更深遠的意義是 Jet-Nemotron 驗證了一條可行的技術路徑，即頂級的開源基座模型、創新的高效架構加上低成本的訓練后架構搜索，就可以實現新一代高性能、高效率的模型。

這場競賽的贏家，將是那些在“算法創新”和“工程極致”結合得最好的團隊。 注意力機制的創新是發令槍，而整個優化技術體系的馬拉松，才剛剛開始。

未經「AI科技評論」授權，嚴禁以任何方式在網頁、論壇、社區進行轉載！

公眾號轉載請先在「AI科技評論」后臺留言取得授權，轉載時需標注來源并插入本公眾號名片。