0
雷峰網(wǎng)訊 小模型 SOTA 之爭持續(xù)升溫，英偉達也再次下場。
在最新論文中，其最新發(fā)布的 Jet-Nemotron 系列混合架構(gòu)語言模型在多項基準測試中超越或追平了 Qwen3，Qwen2.5，Gemma3 和 Llama3.2 的精度，同時實現(xiàn)了高達 53.6 倍的吞吐量加速和 6.1 倍的預填充加速。在與 DeepSeek-V3-Small 和 Moonlight 等最新 MoE 全注意力模型的比較中，Jet-Nemotron-2B 也在 MMLU 和 MMLU- pro 等數(shù)學任務的處理上表現(xiàn)出了更高的精度。
在 NVIDIA H100 GPU 上，以 64K token 的上下文長度測量生成吞吐量。Jet-Nemotron-2B 在 MMLU-Pro 上比 Qwen3-1.7B-Base 具有更高的準確性，同時生成吞吐量提高了 47 倍。Jet-Nemotron-4B 盡管模型規(guī)模更大，但生成吞吐量仍比所有參數(shù)小于 2B 的全注意力模型更高。
一句話來說，就是以更短的時間實現(xiàn)了更高的準確率。
有意思的是，Jet-Nemotron 系列模型的背后，還是小模型中的老朋友 Qwen2.5-1.5B 和 Qwen2.5-3B。
由于預訓練的成本高企不下，設計新的語言模型架構(gòu)具有具有極大的挑戰(zhàn)和風險。在此背景下，Jet-Nemotron 的訓練采用了 Post Neural Architecture Search （PostNAS）方案，這是種以完成預訓練的全注意力模型為基礎，尋找最優(yōu)注意力模塊配置的方法。簡單來說，英偉達選擇對既有的模型“開刀”進行針對性優(yōu)化，而從結(jié)果來看，手術很成功。
論文鏈接：https://arxiv.org/pdf/2508.15884
站在大模型的肩膀上
Qwen2.5 系列本身就是同尺寸模型中的佼佼者，提供了一個極高的性能起點。PostNAS 的目標不是從零學習知識，而是在保留其核心知識的前提下，對其進行“架構(gòu)手術”，優(yōu)化其推理效率。
與主流的語言模型架構(gòu)設計方法不同，PostNAS 從預訓練的全注意力模型開始，繼承其 MLP （多層感知器）權(quán)重并使其保持凍結(jié)，僅對注意力機制部分進行優(yōu)化，這種策略顯著降低了訓練成本。
PostNAS 的后續(xù)流程包括四個關鍵步驟：
1.全注意力機制的位置選擇與消除
在模型中保留少數(shù)全注意力層，對于在檢索等任務中保持高精度至關重要，然而這些層的最佳位置并不明確。解決該問題的常見策略之一是將少數(shù)全注意力層整合，在全注意力模型的部分層上均勻應用，其余層使用線性注意力。
這種均勻策略并非最優(yōu)，英偉達團隊提出了一種自動方法，用于高效確定全注意力層的位置。
其核心思路是為預訓練的全注意力模型增加替代線性注意力路徑，構(gòu)建一個一次性超網(wǎng)絡。在訓練過程中，研究人員會在每一步隨機采樣一條活躍路徑，形成一個子網(wǎng)絡，并對該子網(wǎng)絡使用特征蒸餾損失進行訓練。訓練完成后，使用束搜索就可以確定全注意力層在給定約束條件下的最佳位置。其中搜索目標與任務相關：對于 MMLU，研究人員選擇在正確答案上損失最低的配置（即最大化?????），而對于數(shù)學和檢索任務，則選擇準確率最高的配置。
通過依次將每一層配置為全注意力機制，同時將所有剩余層設置為線性注意力機制，研究人員從超網(wǎng)絡中提取到了相應的子網(wǎng)絡。每個子網(wǎng)絡的準確率和損失可以在給定任務下進行評估，并通過熱力圖可視化結(jié)果。英偉達團隊的分析揭示了三個關鍵發(fā)現(xiàn)：
在預訓練的全注意力模型中，并非所有注意力層都同等重要。對于 MMLU 任務，只有兩層表現(xiàn)出關鍵重要性，而對于檢索任務，僅有兩到三層特別關鍵。
不同的注意力層對不同的能力有所貢獻。對 MMLU 準確率起關鍵作用的注意力層，不一定對檢索任務重要。
對于數(shù)學推理等復雜任務，注意力重要性的模式變得更加復雜。幸運的是，為 MMLU 和檢索任務識別出的關鍵層集合已經(jīng)涵蓋了數(shù)學所需的大部分關鍵層。
2.線性注意力模塊選擇
在全注意力層的最優(yōu)位置之后，還需要確定最適合的線性注意力模塊。英偉達團隊評估了六種 SOTA 線性注意力模塊，其中包括 RWKV7, RetNet, Mamba2, GLA, Deltanet和 Gated DeltaNet。分析發(fā)現(xiàn)，RWKV7 的訓練吞吐量與其它線性注意力模塊相比顯著較低，訓練吞吐量顯著較低，而 Gated DeltaNet 則實現(xiàn)了最佳的整體準確率。
這歸因于兩個因素的結(jié)合：數(shù)據(jù)相關門控機制可以實現(xiàn)模型在關注當前 token 和歷史狀態(tài)之間的動態(tài)控制，Delta 規(guī)則通過當前 token 的信息增量更新歷史狀態(tài)，在此基礎上節(jié)省有限的狀態(tài)內(nèi)存。
值得注意的是，得益于 PostNAS 框架的低訓練成本，系統(tǒng)評估既有線性注意力模塊在不同任務上的準確性、訓練效率和推理速度成為可能。因此隨著新的線性注意力模塊出現(xiàn)，當他們表現(xiàn)出超越 Gated DeltaNet 的性能時，也可以被快速替換。
3.新的注意力模塊設計
Jet-Nemotron 系列模型還應用了一種全新的注意力模塊 JetBlock 以擴展線性注意力容量。
不同的此前依賴靜態(tài)卷積核的方案，英偉達團隊引入了一個卷積核生成模塊，該模塊根據(jù)輸入特征動態(tài)生成卷積核。如下圖所示，該模塊與 Q/K/V 投影層共享相同輸入，首先通過線性縮減層提高效率，縮減比例為 8，然后應用 SiLU 激活函數(shù)，最后通過一個最終線性層輸出卷積核權(quán)重。
研究人員還發(fā)現(xiàn)，一旦對值（V）標記應用動態(tài)卷積，對查詢（Q）和鍵（K）標記的靜態(tài)卷積就可以被移除，而對最終模型精度的影響可以忽略不計。最終實驗采用了這種設計，JetBlock 在與之前線性注意力模塊的比較中，在數(shù)學推理和檢索任務上比 Gated DeltaNet 表現(xiàn)出了更高的精度，同時保持了相似的效率。
4.硬件感知架構(gòu)搜索
PostNAS 還應用了一種硬件感知架構(gòu)搜索方案，以識別最優(yōu)的架構(gòu)超參數(shù)，包括鍵/值維度和注意力頭的數(shù)量。
雖然參數(shù)數(shù)量通常被用作衡量語言模型效率的指標，但它實際上與硬件的生成效率并不直接相關。英偉達團隊將生成吞吐量作為選擇架構(gòu)超參數(shù)的直接目標，通過硬件感知架構(gòu)搜索，發(fā)現(xiàn)了能夠?qū)崿F(xiàn)相似生成吞吐量、同時使用更多參數(shù)以獲得更高精度的架構(gòu)超參數(shù)。
這一原理在于，KV 緩存大小是影響長上下文和長生成吞吐量的最關鍵因素。當 KV 緩存大小保持不變時，不同參數(shù)數(shù)量的模型表現(xiàn)出相似的生成吞吐量。研究人員認為，其原因是解碼階段通常受內(nèi)存帶寬限制而非計算限制。在長上下文場景中，KV 緩存往往比模型權(quán)重占用更多內(nèi)存。減小其大小可以減少每步解碼的內(nèi)存?zhèn)鬏敃r間，并支持更大的批量大小，從而提高生成吞吐量。
理論最大加速比 56 倍
Jet-Nemotron 系列模型包含兩個不同參數(shù)規(guī)模的版本：Jet-Nemotron-2B 和 Jet-Nemotron-4B，分別基于 Qwen2.5-1.5B 和 Qwen2.5-3B 構(gòu)建。
基準測試結(jié)果亦于同期發(fā)布。兩者在所有六個評估領域中的表現(xiàn)均與前沿完整注意力模型（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 甚至超過了參數(shù)量更大的 MoE 模型如 DeepSeek-V3-Small 和 Moonlight，兩者分別有 2.2B 激活參數(shù)和 15B 總參數(shù)。當擴展到 4B 參數(shù)時，Jet-Nemotron-4B 仍然比 Qwen3-1.7B-Base 在吞吐量上高 21 倍。
數(shù)學任務上，Jet-Nemotron-2B 實現(xiàn)了平均 49.6 的準確率，比 Qwen3-1.7B-Base 高出 6.3，同時速度提升了 47 倍。相比之下，之前的線性注意力機制和混合模型在數(shù)學任務上則遠遠落后于 Qwen3。
作為 Jet-Nemotron 系列模型基礎的 Qwen2.5 和 Qwen3 在常識推理任務上的性能相對較弱，但 Jet-Nemotron-2B 仍然表現(xiàn)出色，平均準確率達到了 62.0，超越了所有基線模型。此外在檢索和編程等任務中，Jet-Nemotron-2B 和 Jet-Nemotron-4B 仍在保持吞吐量優(yōu)勢的前提下，表現(xiàn)出全面超越基線模型的性能。
英偉達團隊量化了此種吞吐量優(yōu)勢，下圖展示了 Qwen3-1.7B-Base 和 Jet-Nemotron-2B 在不同上下文長度下的吞吐量對比。在預填充階段，Jet-Nemotron-2B 在較短的上下文長度（4K 和 8K）下初始速度比 Qwen3-1.7B-Base 快 1.14 倍和 1.15 倍。通過設計更優(yōu)化的 JetBlock 內(nèi)核實現(xiàn)，這還可以進一步改進。隨著上下文長度的增加，線性注意力機制變得突出，使得 Jet-Nemotron-2B 在 256K 上下文長度下實現(xiàn)了 6.14 倍的加速。
在解碼階段，Jet-Nemotron-2B 始終大幅優(yōu)于 Qwen3-1.7B-Base。由于 Jet-Nemotron-2B 包含 2 個全注意力層，每組鍵值狀態(tài)有 2 組，而 Qwen3-1.7B-Base 包含 28 個全注意力層，每組鍵值狀態(tài)有 8 組，因此 Jet-Nemotron-2B 的理論最大加速比是 14 × 4 = 56 倍。據(jù)英偉達介紹，在其吞吐量測試平臺中，Jet-Nemotron-2B 在 4K 上下文長度下實現(xiàn)了 15.6 倍的加速，在 256K 上下文長度下實現(xiàn)了高達 53.6 倍的加速，幾乎達到了理論上的上界。
One more thing
從Jet-Nemotron 的經(jīng)驗來看，在日漸升溫的小模型之爭中，對現(xiàn)有模型進行深度優(yōu)化、剪枝、架構(gòu)改造的能力也將成為多方爭奪的前沿，而注意力機制只是起點。
這種基于成熟預訓練模型進行架構(gòu)優(yōu)化的策略，不僅為英偉達節(jié)省了天價開發(fā)成本，而且借助 Qwen2.5 和 Qwen3 經(jīng)過驗證的性能保障了下限。其成功之處不在于讓英偉達在基準測試榜單上占據(jù)幾個名次，更深遠的意義是 Jet-Nemotron 驗證了一條可行的技術路徑，即頂級的開源基座模型、創(chuàng)新的高效架構(gòu)加上低成本的訓練后架構(gòu)搜索，就可以實現(xiàn)新一代高性能、高效率的模型。
這場競賽的贏家，將是那些在“算法創(chuàng)新”和“工程極致”結(jié)合得最好的團隊。 注意力機制的創(chuàng)新是發(fā)令槍，而整個優(yōu)化技術體系的馬拉松，才剛剛開始。
雷峰網(wǎng)(公眾號：雷峰網(wǎng))文章
雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
