Jet-Nemotron
A hybrid-architecture LM family built by adapting after pretraining. Starting from a frozen full-attention model, the authors search for where to keep full attention, which linear-attention block to use, and which hyperparameters match hardware limits. The result, Jet-Nemotron-2B/4B, matches or surpasses popular full-attention baselines while massively increasing throughput on long contexts.
Ask this paper
PostNAS pipeline — Begins with a pre-trained full-attention model and freezes MLPs, then proceeds in four steps: learn optimal placement or removal of full-attention layers, select a linear-attention block, design a new attention block, and run a hardware-aware hyperparameter search.
Learning where full attention actually matters — A once-for-all super-network plus beam search identifies only a few layers as critical, and the important layers differ by task.
JetBlock: linear attention with dynamic convolution — The new block adds a kernel generator that produces input-conditioned causal convolutions applied to V tokens and removes static convolutions on Q/K. They report higher math and retrieval accuracy vs. prior linear blocks at similar training and inference speed.
Hardware-aware design insight — Through grid search at fixed KV cache size, they show generation speed tracks KV cache more than parameter count. The work shares head/dimension settings that hold throughput roughly constant while improving accuracy. This leads to comparable tokens/s with more parameters and better scores.
Results at a glance — Jet-Nemotron-2B outperforms or matches small full-attention models on MMLU, MMLU-Pro, BBH, math, commonsense, retrieval, coding, and long-context tasks, while delivering up to 47x decoding throughput at 64K and as high as 53.6x decoding and 6.14x prefilling speedup at 256K on H100.