雷峰網(wǎng)訊 小模型 SOTA 之爭持續(xù)升溫，英偉達(dá)也再次下場。

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

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

一句話來說，就是以更短的時(shí)間實(shí)現(xiàn)了更高的準(zhǔn)確率。

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

由于預(yù)訓(xùn)練的成本高企不下，設(shè)計(jì)新的語言模型架構(gòu)具有具有極大的挑戰(zhàn)和風(fēng)險(xiǎn)。在此背景下，Jet-Nemotron 的訓(xùn)練采用了 Post Neural Architecture Search （PostNAS）方案，這是種以完成預(yù)訓(xùn)練的全注意力模型為基礎(chǔ)，尋找最優(yōu)注意力模塊配置的方法。簡單來說，英偉達(dá)選擇對既有的模型 " 開刀 " 進(jìn)行針對性優(yōu)化，而從結(jié)果來看，手術(shù)很成功。

站在大模型的肩膀上

Qwen2.5 系列本身就是同尺寸模型中的佼佼者，提供了一個(gè)極高的性能起點(diǎn)。PostNAS 的目標(biāo)不是從零學(xué)習(xí)知識(shí)，而是在保留其核心知識(shí)的前提下，對其進(jìn)行 " 架構(gòu)手術(shù) "，優(yōu)化其推理效率。

與主流的語言模型架構(gòu)設(shè)計(jì)方法不同，PostNAS 從預(yù)訓(xùn)練的全注意力模型開始，繼承其 MLP （多層感知器）權(quán)重并使其保持凍結(jié)，僅對注意力機(jī)制部分進(jìn)行優(yōu)化，這種策略顯著降低了訓(xùn)練成本。

PostNAS 的后續(xù)流程包括四個(gè)關(guān)鍵步驟：

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

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

這種均勻策略并非最優(yōu)，英偉達(dá)團(tuán)隊(duì)提出了一種自動(dòng)方法，用于高效確定全注意力層的位置。

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

通過依次將每一層配置為全注意力機(jī)制，同時(shí)將所有剩余層設(shè)置為線性注意力機(jī)制，研究人員從超網(wǎng)絡(luò)中提取到了相應(yīng)的子網(wǎng)絡(luò)。每個(gè)子網(wǎng)絡(luò)的準(zhǔn)確率和損失可以在給定任務(wù)下進(jìn)行評估，并通過熱力圖可視化結(jié)果。英偉達(dá)團(tuán)隊(duì)的分析揭示了三個(gè)關(guān)鍵發(fā)現(xiàn)：

在預(yù)訓(xùn)練的全注意力模型中，并非所有注意力層都同等重要。對于 MMLU 任務(wù)，只有兩層表現(xiàn)出關(guān)鍵重要性，而對于檢索任務(wù)，僅有兩到三層特別關(guān)鍵。

不同的注意力層對不同的能力有所貢獻(xiàn)。對 MMLU 準(zhǔn)確率起關(guān)鍵作用的注意力層，不一定對檢索任務(wù)重要。

對于數(shù)學(xué)推理等復(fù)雜任務(wù)，注意力重要性的模式變得更加復(fù)雜。幸運(yùn)的是，為 MMLU 和檢索任務(wù)識(shí)別出的關(guān)鍵層集合已經(jīng)涵蓋了數(shù)學(xué)所需的大部分關(guān)鍵層。

2. 線性注意力模塊選擇

在全注意力層的最優(yōu)位置之后，還需要確定最適合的線性注意力模塊。英偉達(dá)團(tuán)隊(duì)評估了六種 SOTA 線性注意力模塊，其中包括 RWKV7, RetNet, Mamba2, GLA, Deltanet 和 Gated DeltaNet。分析發(fā)現(xiàn)，RWKV7 的訓(xùn)練吞吐量與其它線性注意力模塊相比顯著較低，訓(xùn)練吞吐量顯著較低，而 Gated DeltaNet 則實(shí)現(xiàn)了最佳的整體準(zhǔn)確率。

這歸因于兩個(gè)因素的結(jié)合：數(shù)據(jù)相關(guān)門控機(jī)制可以實(shí)現(xiàn)模型在關(guān)注當(dāng)前 token 和歷史狀態(tài)之間的動(dòng)態(tài)控制，Delta 規(guī)則通過當(dāng)前 token 的信息增量更新歷史狀態(tài)，在此基礎(chǔ)上節(jié)省有限的狀態(tài)內(nèi)存。

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

3. 新的注意力模塊設(shè)計(jì)

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

不同的此前依賴靜態(tài)卷積核的方案，英偉達(dá)團(tuán)隊(duì)引入了一個(gè)卷積核生成模塊，該模塊根據(jù)輸入特征動(dòng)態(tài)生成卷積核。如下圖所示，該模塊與 Q/K/V 投影層共享相同輸入，首先通過線性縮減層提高效率，縮減比例為 8，然后應(yīng)用 SiLU 激活函數(shù)，最后通過一個(gè)最終線性層輸出卷積核權(quán)重。

研究人員還發(fā)現(xiàn)，一旦對值（V）標(biāo)記應(yīng)用動(dòng)態(tài)卷積，對查詢（Q）和鍵（K）標(biāo)記的靜態(tài)卷積就可以被移除，而對最終模型精度的影響可以忽略不計(jì)。最終實(shí)驗(yàn)采用了這種設(shè)計(jì)，JetBlock 在與之前線性注意力模塊的比較中，在數(shù)學(xué)推理和檢索任務(wù)上比 Gated DeltaNet 表現(xiàn)出了更高的精度，同時(shí)保持了相似的效率。

4. 硬件感知架構(gòu)搜索

PostNAS 還應(yīng)用了一種硬件感知架構(gòu)搜索方案，以識(shí)別最優(yōu)的架構(gòu)超參數(shù)，包括鍵 / 值維度和注意力頭的數(shù)量。

雖然參數(shù)數(shù)量通常被用作衡量語言模型效率的指標(biāo)，但它實(shí)際上與硬件的生成效率并不直接相關(guān)。英偉達(dá)團(tuán)隊(duì)將生成吞吐量作為選擇架構(gòu)超參數(shù)的直接目標(biāo)，通過硬件感知架構(gòu)搜索，發(fā)現(xiàn)了能夠?qū)崿F(xiàn)相似生成吞吐量、同時(shí)使用更多參數(shù)以獲得更高精度的架構(gòu)超參數(shù)。

這一原理在于，KV 緩存大小是影響長上下文和長生成吞吐量的最關(guān)鍵因素。當(dāng) KV 緩存大小保持不變時(shí)，不同參數(shù)數(shù)量的模型表現(xiàn)出相似的生成吞吐量。研究人員認(rèn)為，其原因是解碼階段通常受內(nèi)存帶寬限制而非計(jì)算限制。在長上下文場景中，KV 緩存往往比模型權(quán)重占用更多內(nèi)存。減小其大小可以減少每步解碼的內(nèi)存?zhèn)鬏敃r(shí)間，并支持更大的批量大小，從而提高生成吞吐量。

理論最大加速比 56 倍

Jet-Nemotron 系列模型包含兩個(gè)不同參數(shù)規(guī)模的版本：Jet-Nemotron-2B 和 Jet-Nemotron-4B，分別基于 Qwen2.5-1.5B 和 Qwen2.5-3B 構(gòu)建。

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

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

數(shù)學(xué)任務(wù)上，Jet-Nemotron-2B 實(shí)現(xiàn)了平均 49.6 的準(zhǔn)確率，比 Qwen3-1.7B-Base 高出 6.3，同時(shí)速度提升了 47 倍。相比之下，之前的線性注意力機(jī)制和混合模型在數(shù)學(xué)任務(wù)上則遠(yuǎn)遠(yuǎn)落后于 Qwen3。

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

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

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

One more thing

從 Jet-Nemotron 的經(jīng)驗(yàn)來看，在日漸升溫的小模型之爭中，對現(xiàn)有模型進(jìn)行深度優(yōu)化、剪枝、架構(gòu)改造的能力也將成為多方爭奪的前沿，而注意力機(jī)制只是起點(diǎn)。

這種基于成熟預(yù)訓(xùn)練模型進(jìn)行架構(gòu)優(yōu)化的策略，不僅為英偉達(dá)節(jié)省了天價(jià)開發(fā)成本，而且借助 Qwen2.5 和 Qwen3 經(jīng)過驗(yàn)證的性能保障了下限。其成功之處不在于讓英偉達(dá)在基準(zhǔn)測試榜單上占據(jù)幾個(gè)名次，更深遠(yuǎn)的意義是 Jet-Nemotron 驗(yàn)證了一條可行的技術(shù)路徑，即頂級的開源基座模型、創(chuàng)新的高效架構(gòu)加上低成本的訓(xùn)練后架構(gòu)搜索，就可以實(shí)現(xiàn)新一代高性能、高效率的模型。

這場競賽的贏家，將是那些在 " 算法創(chuàng)新 " 和 " 工程極致 " 結(jié)合得最好的團(tuán)隊(duì)。 注意力機(jī)制的創(chuàng)新是發(fā)令槍，而整個(gè)優(yōu)化技術(shù)體系的馬拉松，才剛剛開始。

雷峰網(wǎng)文章