6 Knowledge recovery through distillation

Knowledge recovery after structural pruning is the central challenge in rebuilding a smaller LLM, because removing depth or width inevitably strips away knowledge and capability. The chapter argues that recovery begins before training, when you decide what to prune: data-driven layer importance analysis tends to preserve more recoverable knowledge than simple heuristics like dropping the last blocks or blindly protecting the final layers. In the experiments, pruning four of Gemma-3-270M’s 18 blocks with cosine-similarity-based importance scoring produced the strongest starting point for distillation, showing that the choice of removed blocks strongly shapes both the immediate damage and the eventual recovery ceiling.

To restore performance, the chapter moves beyond classic logits-only distillation and builds a compound training setup that combines hard labels, soft labels, and hidden-state alignment. The student and teacher are aligned through mapping strategies that handle unequal depth, including uniform, last-block, and original-index mappings, and the loss blends task loss, KL-based logit loss, and feature alignment. The results show that logits distillation alone can recover a pruned model effectively when enough data is available, while feature alignment becomes more useful when the hidden states themselves matter or when pruning is more aggressive. The chapter also introduces advanced recovery techniques: Skew KLD smooths overly sharp mismatches between teacher and student distributions, and Feature Dynamics Distillation aligns how representations change from one layer to the next rather than only matching static activations.

Across the recovery experiments, more data consistently improved the outcome, with 2K samples giving only modest gains, 15K enabling strong recovery, and 40K allowing the student to match or even surpass the teacher on perplexity while keeping most benchmark capability. Advanced methods delivered their best value at larger data scales, where they slightly improved perplexity at a small cost in compute and, sometimes, a tiny capability tradeoff. The overall guidance is practical: use data-driven pruning, start with a balanced compound loss, keep feature losses low or off for purely generative models unless needed, and reserve advanced techniques for cases with limited data, aggressive pruning, or tasks where intermediate representations are the end product. The chapter closes by emphasizing that this pruning-plus-distillation pipeline is not just theoretical, but reflects real industry practice for building efficient small language models.

Comparison of layer removal strategies. All experiments remove 4 blocks (striped) from the original 18-block model. EXP01-02 use data-driven importance analysis (Listing 6.5), while EXP03-04 follow heuristic rules.

The three mapping strategies: Uniform mapping distributes proportionally, last mapping aligns with the deepest layers, and original mapping preserves the original architectural connections, leaving gaps where the removed layers were.

The compound loss combines three signals: Hard Labels from the dataset (the ground truth "France"), Soft Labels capturing the teacher's probability distribution across candidates, and Hidden States alignment between corresponding Transformer blocks. The weights α, β, γ control each component's influence on the total loss.

Temperature softens the teacher's probability distribution. At T=1.0, the model assigns 94.5% to "hard". At T=2.0, probabilities spread across alternatives (23.9%, 4.7%, 1.3%, 1.2%), exposing the teacher's reasoning process to the student.

Token shifting aligns predictions with labels. Each logit position predicts the next token in the sequence. The first label (WORK) is discarded because no logit predicts it, and the last logit is discarded because there's no future token to predict after <EOS>. The arrows show how logits[0] compares against labels[1], logits[1] against labels[2], and so on.

Convergence analysis showing perplexity (bars, left axis) and capabilities retention (lines with markers, right axis) across three training data scales. Different bar patterns (diagonal hatching, dots, solid) and line markers (circles, squares, triangles) distinguish the three distillation approaches tested at each scale.

Target distribution comparison. Left panel: Standard KLD with teacher targets (solid) and student predictions (diagonal hatching). Right panel: Skew KLD adds hybrid targets (dotted) computed as 90% teacher + 10% student.

The rearchitecting pipeline is flexible and modular. You can apply Width Pruning to reduce MLP dimensions, Depth Pruning to remove blocks, or combine both as shown here for maximum efficiency. Regardless of the structural changes chosen, the final step uses the original Base Model as the Teacher to recover lost knowledge through Distillation.

• Efficient model recovery doesn't start at training, but in the design phase, and that intelligently selecting which blocks to keep (instead of pruning blindly) drastically reduces the computational cost later on.
• Compound loss functions combine three complementary signals: hard labels from the dataset (task loss), soft probability distributions from the teacher (logits loss), and intermediate representations (hidden loss). The balance between these signals should align your final objective. Generative models prioritize logits, while embedding models need higher hidden loss weights.
• You can design Layer Mapping strategies to align a large Teacher with a small Student, addressing the challenge of mismatched architectural depth.
• You've seen how to implement specialized loss functions like Skew KLD and Feature Dynamics, designed to improve convergence when data is scarce or pruning has been aggressive.
• This approach isn't theoretical, but the production standard used by Nvidia (Minitron) and Mistral (Ministral), which allows you to apply top-tier engineering techniques at an accessible scale.

FAQ