Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search

TL;DR

Jet-Nemotron is an NVIDIA hybrid-architecture LM family (2B and 4B) built by Post Neural Architecture Search (PostNAS) — a four-stage NAS pipeline that starts from a frozen-MLP pretrained full-attention model (Qwen2.5) and only searches over the attention stack. The four stages: learn where to keep full-attention layers via a once-for-all super-network, select a linear-attention block (Gated DeltaNet wins), introduce a new JetBlock that uses an input-conditioned dynamic causal convolution on V tokens, and run a hardware-aware hyperparameter search whose target is KV-cache size (not parameter count). Jet-Nemotron-2B matches or beats Qwen3-1.7B-Base on MMLU/MMLU-Pro/BBH/math/retrieval/coding while delivering up to 47× higher generation throughput and 47× smaller KV cache at 64K context; throughput gains grow with context length (53.6× at 256K).

Key claims

Starting NAS from a pretrained full-attention model with frozen MLPs lets the search be done in 50B distillation tokens + 350B full-training tokens — orders of magnitude cheaper than from-scratch architecture search [§2.1, §3.1].
Optimal placement of full-attention layers is task-dependent and learned, not uniform: PostNAS-discovered placements significantly outperform the standard “every-k-layers” baseline on MMLU [§2.2, Fig. 5].
A once-for-all super network with feature-distillation training, plus beam search under a budget of N full-attention layers, identifies high-value attention positions per task (MMLU vs math vs retrieval each prefer different layers) [§2.2, Fig. 4].
Among six SOTA linear-attention blocks under the same PostNAS protocol, Gated DeltaNet ranks first on average — attributed to combining data-dependent gating with the Delta Rule [§2.3, Table 1].
JetBlock improves over Gated DeltaNet by replacing static Q/K convolutions with an input-conditioned dynamic causal convolution applied to V tokens only: +0.7 MMLU, +0.5 math, +0.6 retrieval at comparable throughput [§2.4, Table 1].
KV-cache size, not parameter count, is the right efficiency target — for fixed cache size, more parameters can be allocated to improve accuracy without hurting generation throughput because decoding is memory-bandwidth-bound [§2.5, Table 2, Finding 4].
Hardware-aware search adds another +1.8 MMLU / +2.1 math / +0.5 retrieval on top of JetBlock at fixed cache size and matched training/inference throughput [§2.5, Table 1].
Jet-Nemotron-2B beats Qwen3-1.7B-Base on MMLU (60.8 vs 60.3), MMLU-Pro (39.0 vs 37.8), and BBH (58.3 vs 54.2), while using a 154 MB KV cache vs Qwen3’s 7,168 MB at 64K — a 47× cache reduction and 47× throughput multiplier [§3.2, Table 3].
Jet-Nemotron-2B also exceeds MoE models with much larger total/active parameters (DeepSeek-V3-Small 15B/2.2B-active and Moonlight 15B/2.2B-active) on MMLU and MMLU-Pro [§3.2, Table 3].
Math accuracy gap is large: Jet-Nemotron-2B averages 49.6 across GSM8K/MATH/MathQA/MMLU-Stem/GPQA vs Qwen3-1.7B-Base 42.3, while being 47× faster [§3.2, Table 4].
Prior linear-attention and hybrid models (Mamba2-2.7B, RWKV7-1.5B, Rec.Gemma-2B, Hymba-1.5B, Zamba2-1.2B) trail Jet-Nemotron-2B by 30+ points on coding and math benchmarks — i.e. the throughput advantage of pure linear models came at quality cost that the PostNAS-derived hybrid avoids [§3.2, Tables 4, 7].
Decoding throughput advantage widens with context length: 21× at 4K → 47× at 64K → 53.6× at 256K (vs Qwen3-1.7B-Base) [§1, Fig. 6].

Method

PostNAS is a four-stage pipeline (Fig. 2) that takes a pretrained full-attention model (here Qwen2.5-1.5B for 2B; Qwen2.5-3B for 4B), freezes its MLP weights, and only searches the attention stack.

(1) Full-attention placement & elimination. A once-for-all super network is built where every original attention block has two paths (full vs linear); during training a path is sampled per step and optimized with a feature-distillation loss from the original teacher. After training, beam search finds an optimal placement of N full- attention layers under a task-specific objective (P(answer) for MMLU; accuracy for math/retrieval). For Jet-Nemotron-2B this picks two full-attention layers (15, 20) plus two sliding-window-attention layers (21, 22) for MMLU-style pattern-matching, with the remaining layers as JetBlock.

(2) Linear-attention block selection. Six candidates (RWKV7, RetNet, Mamba2, GLA, DeltaNet, Gated DeltaNet) are evaluated under the same training budget; Gated DeltaNet wins on average. RWKV7 is dropped due to a 2× lower training throughput from suboptimal kernels.

(3) JetBlock. A new linear-attention block based on Gated DeltaNet time-mixing, but with a kernel-generator module that takes the same input as the Q/K/V projection, reduces it 8×, applies SiLU, and emits weights of a causal convolution kernel applied to V tokens only. Static Q/K convolutions are removed (no quality loss once V has dynamic conv).

(4) Hardware-aware architecture search. A grid over key dim, value dim, and number of heads under a fixed KV-cache size constraint (set to match the original model). Configurations with more parameters but the same cache size deliver near-identical generation throughput yet better accuracy.

Training is two-stage: stage 1 freezes MLPs and uses distillation loss for 50B tokens on Nemotron-CC + Redstone-QA; stage 2 unfreezes and trains on 350B tokens with math + code data mixed in.

Results

Cache size: 154 MB for Jet-Nemotron-2B vs 1,792 MB Qwen2.5-1.5B / 7,168 MB Qwen3-1.7B-Base at 64K context [Table 3].
Throughput: 2,885 tok/s vs Qwen3-1.7B-Base 61 tok/s (47×) at 64K; 6.14× prefill speedup [§1, Table 3].
MMLU-Pro: 39.0 (2B) and 44.2 (4B) — vs Qwen3-1.7B-Base 37.8 [Table 3].
Math average across GSM8K/MATH/MathQA/MMLU-Stem/GPQA: 49.6 (2B) / 51.3 (4B) vs Qwen3-1.7B-Base 42.3 [Table 4].
Retrieval average: 74.2 (2B) / 76.2 (4B) — 4B beats Qwen3-1.7B-Base 76.1 [Table 6].
Coding average (EvalPlus/CRUXEval-I-cot/CRUXEval-O-cot): 59.5 (2B) / 63.5 (4B) vs Qwen3-1.7B-Base 58.9 [Table 7].
4B model retains 21× throughput advantage over Qwen3-1.7B-Base while improving accuracy across all benchmark suites [Table 3].
Linear-attention-only baselines (Mamba2-2.7B, RWKV7-1.5B) reach Jet-Nemotron-2B’s throughput but lose 35+ points on math and coding [Tables 4, 7].

Why it’s interesting

Two ideas in this paper are directly portable to other domains the wiki tracks. First, PostNAS’s core trick — freeze MLPs from a pretrained full-attention checkpoint and only search the attention stack with distillation — is the same template as ZAYA1-8B-Diffusion-Preview: Converting an AR MoE to Diffusion on AMD (converting an AR MoE to diffusion by keeping FFNs and retraining only the attention/sampling stack) and SlimQwen: Exploring the Pruning and Distillation in Large MoE Model Pre-training (Qwen pruning a pretrained 80A3B MoE to 23A2B via continued-pretraining distillation rather than scratch training). All three claim that pretrained-then-modify dominates from-scratch under matched token budget; PostNAS extends the recipe to architecture search rather than just compression. Second, the “KV-cache size is the right efficiency target, not parameter count” finding sits alongside the systematic SD study in Speculative Decoding: Performance or Illusion? (verification dominates 42–95% of SD wall-clock because it’s memory-bandwidth-bound) and WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference (prefix cacheability is the right metric for KV-cached serving) — a converging 2025–2026 thesis that in long-context generation, memory traffic, not FLOPs is the binding constraint. Jet-Nemotron is the architecture-side expression of that thesis: trade attention compute for less KV cache, then spend the freed cache budget on parameters that improve accuracy. The JetBlock design (dynamic V-side convolution kernel) is also a concrete follow-up to AlphaGo Moment for Model Architecture Discovery‘s ASI-ARCH observation that “gating layers + small convolutions dominate the winners” — Jet-Nemotron arrives at the same conclusion via human- directed NAS rather than autonomous loops.