英偉達韓松團隊新作：具有后神經架構搜索的高效語言模型
英偉達開源又放大招了！
韓松團隊推出了一款全新的基于后神經架構搜索的高效語言模型——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研究院研究科學家。在上海交通大學獲得碩士和學士學位，在麻省理工學院電子工程與計算機科學系獲得博士學位。
