英偉達韓松團隊新作：具有后神經架構搜索的高效語言模型

時令 發自 凹非寺

量子位 | 公眾號 QbitAI

英偉達開源又放大招了！

韓松團隊推出了一款全新的基於后神經架構搜索的高效語言模型——Jet-Nemotron。

該模型在一系列基準測試中，不僅表現出與Qwen3、Qwen2.5、Gemma 3和Llama 3.2相當甚至更優的準確率，還在生成吞吐量上實現最高53.6倍加速，在預填充階段達到6.1倍加速。

值得一提的是，在MMLU、MMLU-Pro和BBH基準上，Jet-Nemotron-2B相比Qwen3-1.7B-Base吞吐量提高了47倍，緩存大小縮小至1/47。

同時，它還實現了比DeepSeek-V3-Small和Moonlight（共150億參數，22億激活參數）更高的準確率。

代碼和預訓練模型都將開源，我們先來看看Jet-Nemotron是如何構建的。

Jet-Nemotron：基於后神經架構搜索構建

首先，Jet-Nemotron是在后神經架構搜索（Post Neural Architecture Search，PostNAS）的基礎上構建的。

其中，后神經架構搜索（PostNAS）模型是一種「站在大模型肩膀上做改造」的架構搜索方法。

它從一個預訓練的全注意力模型出發，並直接繼承其多層感知機權重，且在整個過程中保持這些權重被凍結（不再更新）。

Jet-Nemotron就是將PostNAS通過以下4個步驟優化得到的：

全注意力層的放置和消除

在模型中保留少數幾個全注意力層，對於在檢索等高難度任務上保持高準確率至關重要。

然而，這些層的最佳放置位置一直不明確。

因此，研究團隊引入了一種新方法，通過訓練一個「一次性」超級網絡 (once-for-all super network)，自動學習應該在哪些位置使用全注意力層。

實驗結果表明，與常用的均勻放置策略相比，這種學習到的放置方式在MMLU基準上的準確率有顯著提升。

選擇線性注意力模塊

在確定了全注意力層的放置方案后，研究團隊進行注意力模塊搜索，以確定最優的線性注意力模塊。

在實驗中，他們評估了6個最先進的線性注意力模塊（RWKV7由於訓練吞吐量過低排除在外），結果如下。

由上表觀察到，Gated DeltaNet實現了最優的整體準確率。因此，研究團隊在后續實驗中都採用Gated DeltaNet。

設計新型注意力模塊

添加捲積是增強線性注意力能力的一種常用策略。然而，以往的方法僅僅依賴於靜態卷積核 ，缺乏動態適應卷積核特徵提取模式的能力。

於是，研究團隊引入一種名為JetBlock的新型線性注意力模塊。

此模塊使用一個卷積核生成器 (kernel generator)，能夠根據輸入內容動態地生成因果卷積核 (dynamic causal convolution kernels)，然后將這些卷積核應用於 V (值) 詞元上。此外，它還移除了在 Q (查詢) 和 K (鍵) 上的冗余靜態卷積，從而簡化了計算流程。

執行硬件感知架構搜索

傳統上，參數量被用作語言模型效率的代理指標。然而，參數數量與硬件效率並不直接相關。

基於「KV 緩存大小是影響長上下文和長生成吞吐量的最關鍵因素」的發現。

研究團隊將KV緩存大小固定為原始設計的規格，並對key維度、value維度以及注意力頭數進行了小規模的網格搜索。

這種硬件感知搜索能夠在保持相似生成吞吐量的同時，利用更多參數以實現更高準確率。

好消息是，研究團隊計劃在GitHub上公開代碼和模型，目前正等待法律合規審覈。

顯著的效率提升

Jet-Nemotron-2B和Jet-Nemotron-4B分別基於Qwen2.5-1.5B和Qwen2.5-3B模型構建。

爲了全面評估模型性能，研究團隊在數學、常識、檢索、編碼以及長上下文中都進行了測試。

數學任務上，Jet-Nemotron-2B取得了49.6的平均準確率，比Qwen3-1.7B-Base高6.3，同時速度快47倍。

相比之下，之前的線性注意力和混合模型在數學任務上遠遠落后於Qwen3-1.7B-Base。

常識推理任務上，Jet-Nemotron-2B平均準確率達到62.0，超越所有基線模型。

檢索任務上，Jet-Nemotron-2B的表現優於除 Qwen3-1.7B-Base之外的所有基線模型。

當擴展到4B時，Jet-Nemotron-4B達到了76.2的最佳平均準確率，同時與Qwen3相比仍保持21倍的速度提升。

編碼任務上，Jet-Nemotron-2B的平均準確率高於所有基線模型。

同時，Jet-Nemotron-4B在所有編碼任務中都實現了更高的準確率。

長下文任務上，可以看出Jet-Nemotron-2B雖然只有兩個全注意力層，但性能堪比擁有更多全注意力層的Qwen2.5-1.5B和Gemma3n-E2B等領先模型。

綜合來看，Jet-Nemotron-2B和Jet-Nemotron-4B在這些領域的表現均與Qwen3-1.7B-Base相當，甚至更勝一籌。

而由於全注意力層顯著減少且KV緩存規模更小，Jet-Nemotron與Qwen3相比有明顯優勢。

團隊介紹

值得一提的是，此研究團隊全為華人。

Yuxian Gu，本科與博士均就讀於清華大學計算機科學與技術系，導師為黃民烈教授。

此前，他還在微軟亞洲研究院實習，導師為研究員董力。

他的研究興趣主要集中在語言模型的全生命周期，包括預訓練、下游任務適配以及推理階段的高效方法。

最近，他的研究重點是面向預訓練大語言模型的數據構建理論與算法（如PDS、指令預訓練、Learning Law），以及利用知識蒸餾進行語言模型壓縮（如MiniLLM、MiniPLM）。

胡擎昊，本科畢業於浙江大學，碩士畢業於新加坡國立大學，現為麻省理工學院韓松教授的博士后研究員。

尚揚，現為麻省理工學院電子工程學院的一年級博士生，導師為韓松教授。在此之前，他以最高榮譽獲得了清華大學電子工程系的理學學士學位。

Haochen Xi，本科畢業於清華大學姚班，導師為姚期智院士，目前博士就讀於美國加州大學伯克利分校計算機科學專業，現為加州大學伯克利分校MLsys研究員。

Junyu Chen，現為清華大學姚班的一名本科生。曾在麻省理工學院HAN實驗室擔任研究實習生，導師為韓松教授。此前，還曾在清華大學與李毅教授合作研究3D視覺感知和人機交互。

韓松，本科畢業於清華大學電子工程系，在斯坦福大學獲得博士學位，目前是麻省理工學院電子工程學院副教授。

他提出了被廣泛用於高效人工智能計算的「深度壓縮」技術，並且首次給現代人工智能芯片帶來權重稀疏性的「高效推理機」，這些技術影響了NVIDIA的安培GPU架構等。

韓松還是TinyML研究的先驅，這項研究將深度學習帶到物聯網設備上，使邊緣端機器學習成為可能。

2023年，韓松創辦的專注邊緣設備機器學習優化的OmniML被英偉達收購，他也因此加入英偉達成為傑出科學家，其公司的CEO吳迪和CTO毛慧子同樣也入職英偉達。

蔡涵，NVIDIA研究院研究科學家。在上海交通大學獲得碩士和學士學位，在麻省理工學院電子工程與計算機科學系獲得博士學位。

參考鏈接：

[1]https://arxiv.org/abs/2508.15884

[2]https://github.com/NVlabs/Jet-Nemotron

[3]https://x.com/iScienceLuvr/status/1959832287073403137

（聲明：本文僅代表作者觀點，不代表新浪網立場。）