Research map / Architectures and efficiency / Neural architecture search
Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Jet-Nemotron builds hybrid linear/full-attention LMs from pretrained full-attention models via PostNAS (frozen MLPs, staged architecture search) and a new JetBlock linear attention with dynamic convolution, reaching Qwen3-level accuracy with up to 47x higher generation throughput.
Problem
Efficient language models with linear or hybrid attention typically lag full-attention SOTA models in accuracy, and designing new LM architectures from scratch is prohibitively expensive and risky.
Approach
Reuse pretrained full-attention models: freeze the MLPs and search attention designs post hoc (full-attention placement via a once-for-all super network with beam search, linear block selection, new JetBlock with input-conditioned dynamic causal convolution on V, and hardware-aware search that holds KV cache size fixed).
Key results
- Jet-Nemotron-2B reaches 47x higher generation throughput than Qwen3-1.7B-Base (2,885 vs 61 token/s at 64K) with higher MMLU-Pro (39.0 vs 37.8) and average math accuracy 49.6 vs 42.3.
- Jet-Nemotron-2B achieves up to 53.6x decoding speedup and 6.14x prefilling speedup at 256K context versus Qwen3-1.7B-Base.
- Jet-Nemotron-4B achieves the best average on retrieval (76.2) and commonsense (64.7) among compared models while keeping 21x throughput advantage over Qwen3-1.7B-Base.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use learned full-attention layer placement instead of uniform placement (ablation) | PostNAS · MMLU | accuracy | Helped |
| select Gated DeltaNet rather than other evaluated linear attention blocks (ablation) | PostNAS linear attention block selection · Math tasks | accuracy: 32.3 | Helped |
| add JetBlock in place of Gated DeltaNet (ablation) | JetBlock · Math reasoning | accuracy: 32.3 -> 32.8 | Helped |
| add JetBlock in place of Gated DeltaNet (ablation) | JetBlock · Retrieval | accuracy: 69.3 -> 69.9 | Helped |
| apply hardware-aware architecture search to JetBlock (ablation) | JetBlock · MMLU | accuracy: 56.3 -> 58.1 | Helped |
| apply hardware-aware architecture search to JetBlock (ablation) | JetBlock · Math | accuracy: 32.8 -> 34.9 | Helped |
| apply hardware-aware architecture search to JetBlock (ablation) | JetBlock · Retrieval | accuracy: 69.9 -> 70.4 | Helped |
| apply hardware-aware architecture search to JetBlock (ablation) | JetBlock · Commonsense | accuracy: 58.5 -> 59.5 | Helped |
| apply dynamic convolution to V tokens rather than Q or K (ablation) | JetBlock · final model accuracy | accuracy | No clear effect |
| remove static convolutions on Q and K after applying dynamic convolution to V (ablation) | JetBlock · final model accuracy | accuracy | No clear effect |
| increase parameter count while keeping KV cache size constant through hardware-aware hyperparameter search (ablation) | Jet-Nemotron architecture · generation throughput | throughput | No clear effect |
| use Jet-Nemotron-2B instead of Qwen3-1.7B-Base | Jet-Nemotron-2B · MMLU-Pro | accuracy: 37.8 -> 39.0 | Helped |
Surprising findings
- KV cache size, rather than parameter count, is the dominant factor for long-context generation throughput; models with equal cache size but different parameter counts have similar throughput.
- Only two full-attention layers are critical for MMLU and two to three for retrieval in a pretrained full-attention model; placement importance differs by task.
Limitations
- Architectures are constrained to the pretrained base model's MLPs and are not optimal when trained from scratch, as the authors acknowledge.
- Linear attention kernels are less optimized (e.g., RWKV7 training throughput low due to kernel implementation), and prefill gains at short contexts are modest (1.14-1.15x at 4K-8K).
Experimental setup
Datasets: Nemotron-CC Redstone-QA math corpora code corpora LongBench MMLU MMLU-Pro GSM8K MATH GPQA
Benchmarks: MMLU MMLU-Pro BBH GSM8K MATH MathQA MMLU-Stem GPQA ARC-c ARC-e PIQA Winogrande
Metrics: accuracy generation throughput (token/s) prefilling speedup decoding speedup cache size (MB)
Baselines: Qwen2.5-1.5B Qwen3-1.7B-Base Llama3.2-3B MiniCPM-2B-128K SmolLM2-1.7B Mamba2-2.7B RWKV7-1.5B Gemma3n-E2B Hymba-1.5B Falcon-H1-1.5B
Builds on: Gated DeltaNet GLA once-for-all super network FlashAttention Flash-Linear-Attention
Closest papers by meaning
- Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models (NeurIPS 2025 · similarity 0.87)
- Puzzle: Distillation-Based NAS for Inference-Optimized LLMs (ICML 2025 · similarity 0.85)
- On Data Engineering for Scaling LLM Terminal Capabilities (NeurIPS 2026 · similarity 0.85)
- RZ-NAS: Enhancing LLM-guided Neural Architecture Search via Reflective Zero-Cost Strategy (ICML 2025 · similarity 0.85)
- Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.84)
- M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.