時(shí)令 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

英偉達(dá)開源又放大招了！

韓松團(tuán)隊(duì)推出了一款全新的基于后神經(jīng)架構(gòu)搜索的高效語言模型——Jet-Nemotron。

該模型在一系列基準(zhǔn)測試中，不僅表現(xiàn)出與Qwen3、Qwen2.5、Gemma 3和Llama 3.2相當(dāng)甚至更優(yōu)的準(zhǔn)確率，還在生成吞吐量上實(shí)現(xiàn)最高53.6倍加速，在預(yù)填充階段達(dá)到6.1倍加速。

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

同時(shí)，它還實(shí)現(xiàn)了比DeepSeek-V3-Small和Moonlight（共150億參數(shù)，22億激活參數(shù)）更高的準(zhǔn)確率。

代碼和預(yù)訓(xùn)練模型都將開源，我們先來看看Jet-Nemotron是如何構(gòu)建的。

Jet-Nemotron：基于后神經(jīng)架構(gòu)搜索構(gòu)建

首先，Jet-Nemotron是在后神經(jīng)架構(gòu)搜索（Post Neural Architecture Search，PostNAS）的基礎(chǔ)上構(gòu)建的。

其中，后神經(jīng)架構(gòu)搜索（PostNAS）模型是一種“站在大模型肩膀上做改造”的架構(gòu)搜索方法。

它從一個(gè)預(yù)訓(xùn)練的全注意力模型出發(fā)，并直接繼承其多層感知機(jī)權(quán)重，且在整個(gè)過程中保持這些權(quán)重被凍結(jié)（不再更新）。

Jet-Nemotron就是將PostNAS通過以下4個(gè)步驟優(yōu)化得到的：

全注意力層的放置和消除

在模型中保留少數(shù)幾個(gè)全注意力層，對于在檢索等高難度任務(wù)上保持高準(zhǔn)確率至關(guān)重要。

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

因此，研究團(tuán)隊(duì)引入了一種新方法，通過訓(xùn)練一個(gè)“一次性”超級網(wǎng)絡(luò) (once-for-all super network)，自動學(xué)習(xí)應(yīng)該在哪些位置使用全注意力層。

實(shí)驗(yàn)結(jié)果表明，與常用的均勻放置策略相比，這種學(xué)習(xí)到的放置方式在MMLU基準(zhǔn)上的準(zhǔn)確率有顯著提升。

選擇線性注意力模塊

在確定了全注意力層的放置方案后，研究團(tuán)隊(duì)進(jìn)行注意力模塊搜索，以確定最優(yōu)的線性注意力模塊。

在實(shí)驗(yàn)中，他們評估了6個(gè)最先進(jìn)的線性注意力模塊（RWKV7由于訓(xùn)練吞吐量過低排除在外），結(jié)果如下。

由上表觀察到，Gated DeltaNet實(shí)現(xiàn)了最優(yōu)的整體準(zhǔn)確率。因此，研究團(tuán)隊(duì)在后續(xù)實(shí)驗(yàn)中都采用Gated DeltaNet。

設(shè)計(jì)新型注意力模塊

添加卷積是增強(qiáng)線性注意力能力的一種常用策略。然而，以往的方法僅僅依賴于靜態(tài)卷積核，缺乏動態(tài)適應(yīng)卷積核特征提取模式的能力。

于是，研究團(tuán)隊(duì)引入一種名為JetBlock的新型線性注意力模塊。

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

執(zhí)行硬件感知架構(gòu)搜索

傳統(tǒng)上，參數(shù)量被用作語言模型效率的代理指標(biāo)。然而，參數(shù)數(shù)量與硬件效率并不直接相關(guān)。

基于“KV 緩存大小是影響長上下文和長生成吞吐量的最關(guān)鍵因素”的發(fā)現(xiàn)。

研究團(tuán)隊(duì)將KV緩存大小固定為原始設(shè)計(jì)的規(guī)格，并對key維度、value維度以及注意力頭數(shù)進(jìn)行了小規(guī)模的網(wǎng)格搜索。

這種硬件感知搜索能夠在保持相似生成吞吐量的同時(shí)，利用更多參數(shù)以實(shí)現(xiàn)更高準(zhǔn)確率。

好消息是，研究團(tuán)隊(duì)計(jì)劃在GitHub上公開代碼和模型，目前正等待法律合規(guī)審核。

顯著的效率提升

Jet-Nemotron-2B和Jet-Nemotron-4B分別基于Qwen2.5-1.5B和Qwen2.5-3B模型構(gòu)建。

為了全面評估模型性能，研究團(tuán)隊(duì)在數(shù)學(xué)、常識、檢索、編碼以及長上下文中都進(jìn)行了測試。

數(shù)學(xué)任務(wù)上，Jet-Nemotron-2B取得了49.6的平均準(zhǔn)確率，比Qwen3-1.7B-Base高6.3，同時(shí)速度快47倍。

相比之下，之前的線性注意力和混合模型在數(shù)學(xué)任務(wù)上遠(yuǎn)遠(yuǎn)落后于Qwen3-1.7B-Base。

常識推理任務(wù)上，Jet-Nemotron-2B平均準(zhǔn)確率達(dá)到62.0，超越所有基線模型。

檢索任務(wù)上，Jet-Nemotron-2B的表現(xiàn)優(yōu)于除 Qwen3-1.7B-Base之外的所有基線模型。

當(dāng)擴(kuò)展到4B時(shí)，Jet-Nemotron-4B達(dá)到了76.2的最佳平均準(zhǔn)確率，同時(shí)與Qwen3相比仍保持21倍的速度提升。

編碼任務(wù)上，Jet-Nemotron-2B的平均準(zhǔn)確率高于所有基線模型。

同時(shí)，Jet-Nemotron-4B在所有編碼任務(wù)中都實(shí)現(xiàn)了更高的準(zhǔn)確率。

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

綜合來看，Jet-Nemotron-2B和Jet-Nemotron-4B在這些領(lǐng)域的表現(xiàn)均與Qwen3-1.7B-Base相當(dāng)，甚至更勝一籌。

而由于全注意力層顯著減少且KV緩存規(guī)模更小，Jet-Nemotron與Qwen3相比有明顯優(yōu)勢。

團(tuán)隊(duì)介紹

值得一提的是，此研究團(tuán)隊(duì)全為華人。

Yuxian Gu，本科與博士均就讀于清華大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)系，導(dǎo)師為黃民烈教授。

此前，他還在微軟亞洲研究院實(shí)習(xí)，導(dǎo)師為研究員董力。

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

最近，他的研究重點(diǎn)是面向預(yù)訓(xùn)練大語言模型的數(shù)據(jù)構(gòu)建理論與算法（如PDS、指令預(yù)訓(xùn)練、Learning Law），以及利用知識蒸餾進(jìn)行語言模型壓縮（如MiniLLM、MiniPLM）。

胡擎昊，本科畢業(yè)于浙江大學(xué)，碩士畢業(yè)于新加坡國立大學(xué)，現(xiàn)為麻省理工學(xué)院韓松教授的博士后研究員。

Shang Yang，現(xiàn)為麻省理工學(xué)院三年級博士生，導(dǎo)師為韓松教授。在此之前，他本科畢業(yè)于清華大學(xué)電子工程系。

Haochen Xi，本科畢業(yè)于清華大學(xué)姚班，導(dǎo)師為姚期智院士，目前博士就讀于美國加州大學(xué)伯克利分校計(jì)算機(jī)科學(xué)專業(yè)，現(xiàn)為加州大學(xué)伯克利分校MLsys研究員。

Junyu Chen，現(xiàn)為清華大學(xué)姚班的一名本科生。曾在麻省理工學(xué)院HAN實(shí)驗(yàn)室擔(dān)任研究實(shí)習(xí)生，導(dǎo)師為韓松教授。此前，還曾在清華大學(xué)與李毅教授合作研究3D視覺感知和人機(jī)交互。

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

他提出了被廣泛用于高效人工智能計(jì)算的“深度壓縮”技術(shù)，并且首次給現(xiàn)代人工智能芯片帶來權(quán)重稀疏性的“高效推理機(jī)”，這些技術(shù)影響了NVIDIA的安培GPU架構(gòu)等。

韓松還是TinyML研究的先驅(qū)，這項(xiàng)研究將深度學(xué)習(xí)帶到物聯(lián)網(wǎng)設(shè)備上，使邊緣端機(jī)器學(xué)習(xí)成為可能。

2023年，韓松創(chuàng)辦的專注邊緣設(shè)備機(jī)器學(xué)習(xí)優(yōu)化的OmniML被英偉達(dá)收購，他也因此加入英偉達(dá)成為杰出科學(xué)家，其公司的CEO吳迪和CTO毛慧子同樣也入職英偉達(dá)。

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

參考鏈接：

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

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

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