Distillation Variants › Limitations and Practical Considerations

Each distillation variant introduces computational and design overhead beyond standard KD, and understanding these tradeoffs is essential before committing to one approach. None of these methods is universally best; the right choice depends on your task, your architecture, and your compute budget.

Feature distillation requires designing the projection layer $W_r$ and choosing which teacher layers to align with which student layers. This layer correspondence problem has no universal solution. Incorrect alignments can harm the student, forcing it to mimic teacher representations that are architecturally incompatible with its capacity. For transformer models, a useful heuristic is to match layers by relative depth, aligning the student's $i$-th layer to the teacher's $\lfloor i \cdot N_T / N_S \rfloor$-th layer. But this is domain-dependent: in vision transformers, early layers specialize for local texture, and matching them strictly to the corresponding teacher layers may not benefit semantic tasks. A more principled approach is to search over layer correspondences using a small validation set, selecting the alignment that maximizes student validation performance, but this adds another hyperparameter search loop to an already expensive training process.

There is also the risk of over-regularization from feature distillation. The student has less representational capacity than the teacher. Forcing it to reproduce the teacher's high-dimensional geometry can cause the student to allocate its limited capacity to mimicking teacher structure rather than solving the task. In practice, researchers report that the feature loss coefficient $\gamma$ needs careful tuning: too high and the student becomes a teacher-geometry-replicator that fails at the task; too low and the feature alignment provides no benefit over standard KD.

Attention transfer faces a related correspondence problem when teacher and student have different numbers of heads. Averaging across heads before computing the loss loses head specialization information: some teacher heads specialize for syntactic dependencies while others handle positional patterns. Matching individual student heads to teacher heads preserves this, but requires a head alignment procedure that may be brittle. The benefit of attention transfer also diminishes for tasks where attention patterns are diffuse or task-agnostic, making it primarily valuable for structured prediction problems. For generative language modeling where attention patterns are inherently noisy and variable across sequences, the signal from attention transfer may add less than the computational overhead it introduces.

Progressive distillation multiplies the training cost by the number of stages. A three-stage pipeline requires training three models, each for multiple epochs, compared to a single training run for direct distillation. The intermediate models add memory overhead. In practice, this cost is offset by the improved final student quality, but for resource-constrained settings, it may be preferable to use a single-step approach with a more carefully tuned loss. Another subtlety is model selection at each stage: if an intermediate model is poorly trained, errors propagate forward through the chain, potentially making the final student worse than a direct one-step distillation would have been.

On-policy distillation is the most computationally expensive variant because it requires autoregressive generation at every training step. For long sequences, this is prohibitively slow without significant engineering effort. The standard remedy is to generate shorter prefixes and limit the on-policy component to a fraction of training batches. Mixed training, alternating between on-policy and offline steps, often achieves most of the benefit at lower cost. The choice between forward and reverse KL also has significant practical implications: forward KL produces more diverse generations but can result in overly hedged outputs, while reverse KL produces more focused, confident text but may be overconfident on inputs outside the student's learned distribution.

A recurring practical limitation across all variants is that they assume the teacher is a reliable guide. If the teacher is itself imperfect, biased, or overfit to the training distribution, the student faithfully replicates those flaws. This is particularly relevant for attention transfer: if the teacher's attention patterns encode spurious correlations, teaching the student to replicate them transfers the flaw. It is also relevant for on-policy distillation in a subtle way: the teacher is queried on student-generated sequences, which may be very different from the sequences the teacher was trained on. If the teacher's distribution is poorly calibrated on out-of-distribution inputs, the on-policy KL divergence signal becomes unreliable precisely in the regime where the student most needs guidance. Auditing teacher behavior before distillation, and monitoring the quality of teacher outputs on student-generated prefixes throughout on-policy training, is sound engineering practice.

Given the distinct costs and benefits, a practical decision guide:

• Use feature distillation when your student architecture mirrors the teacher's (e.g., both are transformers with comparable layer structure) and you have sufficient compute to handle the projection overhead. It provides the richest training signal for structural tasks.
• Use attention transfer when your task requires precise relational understanding and the teacher has developed well-defined structural attention patterns. It is most valuable for encoder models on parsing, coreference, and structured prediction tasks.
• Use progressive distillation when your target compression ratio exceeds roughly 8-10x. Below that threshold, single-step distillation usually suffices. Above it, the quality gains from staged compression typically justify the training cost.
• Use on-policy distillation for any generative task where inference quality on long sequences matters. The distribution shift problem is severe enough for autoregressive generation that on-policy training, even a mixed on-policy/offline schedule, almost always improves over purely offline distillation.

In production, these methods are routinely combined. A common pipeline for compressing a large instruction-tuned LLM into a deployable small model uses progressive staging to reach the target size, feature distillation and attention transfer within each stage for richer signal, and on-policy training in the final stage to align the student's generation distribution with the teacher's.