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

作者丨梁丙鑒

編輯丨馬曉寧

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

在最新論文中，其最新發(fā)布的 Jet-Nemotron 系列混合架構(gòu)語言模型在多項基準(zhǔn)測試中超越或追平了 Qwen3，Qwen2.5，Gemma3 和 Llama3.2 的精度，同時實現(xiàn)了高達 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)確性，同時生成吞吐量提高了 47 倍。Jet-Nemotron-4B 盡管模型規(guī)模更大，但生成吞吐量仍比所有參數(shù)小于 2B 的全注意力模型更高。

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

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

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

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

01

站在大模型的肩膀上

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

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

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

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

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

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

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

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

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

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

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

2.線性注意力模塊選擇

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

這歸因于兩個因素的結(jié)合：數(shù)據(jù)相關(guān)門控機制可以實現(xiàn)模型在關(guān)注當(dāng)前 token 和歷史狀態(tài)之間的動態(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 的性能時，也可以被快速替換。

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

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

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

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

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

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

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

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

02

理論最大加速比 56 倍

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

基準(zhǔn)測試結(jié)果亦于同期發(fā)布。兩者在所有六個評估領(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 倍，同時在 MMLU、MMLU-Pro 和 BBH 上顯著提高了準(zhǔn)確率。JetNemotron-2B 甚至超過了參數(shù)量更大的 MoE 模型如 DeepSeek-V3-Small 和 Moonlight，兩者分別有 2.2B 激活參數(shù)和 15B 總參數(shù)。當(dāng)擴展到 4B 參數(shù)時，Jet-Nemotron-4B 仍然比 Qwen3-1.7B-Base 在吞吐量上高 21 倍。

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

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

英偉達團隊量化了此種吞吐量優(yōu)勢，下圖展示了 Qwen3-1.7B-Base 和 Jet-Nemotron-2B 在不同上下文長度下的吞吐量對比。在預(yù)填充階段，Jet-Nemotron-2B 在較短的上下文長度（4K 和 8K）下初始速度比 Qwen3-1.7B-Base 快 1.14 倍和 1.15 倍。通過設(shè)計更優(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 倍的加速，幾乎達到了理論上的上界。

03

One more thing

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

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

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

未經(jīng)「AI科技評論」授權(quán)，嚴禁以任何方式在網(wǎng)頁、論壇、社區(qū)進行轉(zhuǎn)載！

公眾號轉(zhuǎn)載請先在「AI科技評論」后臺留言取得授權(quán)，轉(zhuǎn)載時需標(biāo)注來源并插入本公眾號名片。