Compact Language Models via Pruning and Knowledge Distillation

Instead of training an 8B and a 4B model from scratch, NVIDIA takes its trained Nemotron-4 15B and cuts it down. Structured pruning removes whole units: attention heads, MLP neurons, embedding channels (the model's width) or entire layers (its depth). The question is which to remove and how to recover the lost accuracy cheaply.

Ranking what to cut. Run 1,024 calibration samples through the model, forward passes only, and score each head, neuron and channel by the size of its activations, aggregated across the batch and sequence (L2 over the batch, mean over the sequence worked best). Layers are scored by how much perplexity rises when each is removed, or by how little a layer changes its input.

Choosing a shape. Enumerate architectures that fit the parameter budget, trim the original to each, retrain each for about 1.8B tokens to let rankings settle, and keep the best.

Recovering accuracy. Retrain the chosen model by distillation: minimize the KL divergence between its next-token distribution and the original model's, instead of training on the text's next tokens. Add intermediate-state matching only when many layers were removed. Then prune the 8B again to make the 4B.

The resulting Minitron 8B and 4B used 94B tokens each. The 8B beats the previous-generation Nemotron-3 8B and Llama-2 7B and is close to Mistral 7B, Gemma 7B and Llama-3 8B on most benchmarks; the 4B leads Gemma-2B on most tasks. Both beat other published pruning methods by wide margins. The paper distills its findings into ten best practices, including: prefer width over depth at these sizes, retrain only with distillation, and prune from the model closest to the target size.

Customers often settle on a model that works and then need a smaller one: for a cheaper GPU, an edge device, or lower latency at high volume. Training a smaller model from scratch is out of reach for almost everyone. Pruning plus distillation from the model they already trust is the practical path, and this paper is the most detailed public recipe for it.

The data cost is small. 94B tokens is under 3% of the original training run, and the retraining data does not need labels: the teacher's own distributions are the target. Distillation, not plain retraining, made the difference: at equal compute, the pruned 4B trained on ordinary next-token labels scored at chance on MMLU (24.6), while the distilled one scored 37.8.

Width cuts keep more than depth cuts. At these sizes, removing heads, neurons and channels beat removing layers once the model was retrained, even though depth pruning looked better immediately after cutting.

Check what the cut actually saves at serving time. All three Nemotron and Minitron models keep 32 layers and 8 key-value groups (Table 5). Parameters drop two to four times, but unless the head size also changed, which the paper does not say, the KV cache per token does not shrink. For long-context serving, memory per session may barely move.