Agent Distillation in Practice: Replicating Large Model Intelligence in Small Models

In the context of large model training, the word "distillation" has recently taken on a negative connotation due to so-called "distillation attacks." But as the Hugging Face team emphasized in this livestream on Agent training: distillation is standard practice in machine learning, dating back to Hinton's classic 2015 paper, and is used in some form by virtually every machine learning lab as a practical engineering tool.

This article is based on the second session of that livestream series, systematically covering the core concepts of Agent distillation, its four key dimensions, and how to implement both off-policy and on-policy distillation in the TRL framework with minimal code.

Where Distillation Fits in the Training Pipeline

Modern large model training typically consists of three stages: pre-training, mid-training, and post-training. The pre-training stage gives the model general capabilities, language understanding, and world knowledge; mid-training covers specialized knowledge like coding; and post-training delivers targeted refinements, such as using the correct output format.

All the methods discussed in this series — SFT (Supervised Fine-Tuning) from the previous session, distillation in this session, and reinforcement learning to be covered later — all fall under the post-training umbrella.

The core value of distillation lies in taking a smaller, cheaper, faster, and potentially locally-runnable model and improving it for specific tasks with the help of a larger or domain-specific "teacher model." These tasks can range from mathematical reasoning and clinical information extraction to dialogue improvement or, as this session focuses on, agentic tasks.

Interestingly, the teacher model doesn't have to be larger than the student — it can also be a domain expert model of comparable size. Distillation can achieve model compression, capability transfer, and synthetic data generation.

Understanding the logic behind choosing a "teacher model" helps clarify the practical boundaries of distillation. In terms of scale, the teacher doesn't need to be larger — a point with deep engineering implications: a same-size model specifically trained in a particular domain (e.g., medical reports, legal contracts, Python tool calls) often has a domain distribution far better suited as a distillation target than a general large model, yielding more precise transfer effects. The distinction between "white-box" and "black-box" is crucial in practice: white-box teachers (with access to weights and logits) support soft-label distillation with high information transfer efficiency; black-box teachers (such as closed-source models accessed via API) can only provide text output, limiting distillation to hard-label offline approaches — essentially an indirect path through synthetic data generation followed by SFT. This is precisely why the open-source model ecosystem matters so much to distillation research — access to logits directly determines which training methods are available.

Breaking Down the Four Core Dimensions of Distillation

The livestream introduced a very clear analytical framework that breaks down various distillation methods into four dimensions. Understanding these four dimensions makes it easy to decode nearly any combination of terms in paper titles.

Dimension 1: Hard Labels vs. Soft Labels

The first dimension is the source of the training signal. Hard labels (sequence knowledge distillation) are what was done in the previous session's SFT — obtaining tokens in trajectory form and training the student model to imitate those tokens. Soft labels (logits knowledge distillation) involve obtaining the teacher model's logits and training the student model based on probability distributions.

Soft labels provide richer information: not just the top-ranked word, but also the second and third most plausible alternatives, as well as information about words the teacher model considers completely inappropriate. This gives the student model a much richer signal to learn from.

Understanding the difference between the two from an information-theoretic perspective is more intuitive. Hard labels are essentially one-hot distributions — probability 1 for the correct answer, 0 for everything else — which means a large amount of "dark knowledge" is discarded. The core insight of Hinton's 2015 paper was precisely that the probability distribution output by a teacher model itself carries structural information. For example, given the input "cat," the teacher model might assign 0.05 probability to "dog" and 0.001 to "car" — this gap reveals that "cats and dogs are semantically more similar than cats and cars," information that hard labels cannot convey at all. Soft labels accelerate training and improve generalization precisely because the effective information content per sample is far higher than with hard labels. In practice, soft-label distillation also typically requires a "temperature" hyperparameter T: dividing logits by T before applying softmax makes the distribution smoother as T increases, enriching the dark knowledge but also introducing more noise.

Dimensions 2–4: Data Source, Timing, and Teacher Identity

The second dimension is data source — whether training data comes from the teacher or the student. The third dimension is timing — offline (data prepared separately before training) or online (the teacher model participates in real-time during the training loop). The fourth dimension is teacher identity — whether it's a separate model or the same model (the latter being self-distillation).

Combined, these four dimensions give rise to terms like "on-policy distillation," "off-policy distillation," and "self-distillation," and also explain the distinction between white-box (logits accessible) and black-box (only string outputs available) approaches.

Off-Policy vs. On-Policy: The Chess Analogy

The livestream used a chess analogy to brilliantly illustrate the essential difference between the two strategies.

Reinforcement learning rewards the model after completing a task (e.g., winning a game), but the signal is very sparse — feedback only comes at the end of the game. If the model can't even tell a knight from a rook and can't make a legal first move, it will never receive any reward and has no way to learn.

Off-policy distillation (including SFT from the previous session) provides the complete game trajectories of a teacher model, yielding very dense signals. But its limitation is this: if the teacher is a grandmaster who wins every game in a few moves and never encounters a situation with the king exposed, then the student model lacks relevant information to learn from at its own foundational level of challenge — the trajectories the teacher takes are completely different from what the student would play.

On-policy distillation fills exactly this gap: the student model generates its own trajectories (plays the game), and the teacher model scores those moves. If the student consistently moves the rook incorrectly, it receives low scores. This gives the student dense gradients that are also relevant to its own performance.

A two-dimensional chart summarizes this well: the vertical axis is off-policy/on-policy, and the horizontal axis is dense/sparse. Reinforcement learning is on-policy but with sparse signals; SFT is off-policy but with dense signals; on-policy distillation enjoys both density and on-policy relevance simultaneously.

TRL in Practice: One Trainer, Two Knobs

The most impressive aspect of the hands-on portion was the elegant simplicity of the TRL framework. The Generalized Knowledge Distillation Trainer (GKD Trainer) lets you freely switch between off-policy and on-policy distillation by adjusting just two parameters.

• lambda: Controls the data source. Set to 0 for off-policy (training on teacher data); set to 1 for on-policy (generating data with the student model).
• beta: Controls the loss type. Set to 0 to use forward KL divergence (considering the teacher's full distribution); set to 1 to use reverse KL divergence (focusing only on the teacher's hard choices/top options).

Off-Policy Distillation Experiment

The experimental setup continues from the first session: the student model is Qwen3 0.6B, the teacher model is Qwen3 4B, and the dataset consists of agent trajectories extracted from Python programming-based drawing tasks. The team ran a coding agent with hfjobs across three different learning rates, logging results to Trackio.

Simply setting lambda=0, beta=0, four lines of code are enough to instantiate the trainer and complete off-policy training — essentially dense supervision with forward KL divergence. The entire demo runs on a single H200 GPU and completes in a few minutes, with an extremely low barrier to reproduction.

On-Policy Distillation Experiment

The complete on-policy distillation pipeline works as follows: input prompt → student model samples to generate completions (expected to make mistakes) → frozen teacher model performs one forward pass over these tokens (prefill only, low overhead) → compute reverse KL divergence between student and teacher at each token → update the student model based on the gap.

Why use reverse KL instead of forward KL? The key is that reverse KL is "mode-seeking" while forward KL is "mean-seeking." Forward KL might push the student toward a mean point that doesn't actually represent any specific behavior of the teacher; reverse KL pushes the distribution toward the parts that best characterize the teacher's behavior, making the student imitate the teacher in the most action-relevant areas.

The code is almost identical — just set both lambda and beta to 1. This "change two numbers" design makes switching between experiments virtually frictionless.

On-policy distillation draws from the DAGGER-related literature and can be viewed as iterative SFT that incorporates teacher evaluation in a closed loop. It has been applied recently in GLM-4.5, Qwen3, and DeepSeek, making it one of the most popular distillation approaches today.

The behavioral difference between forward KL divergence and reverse KL divergence warrants further elaboration. Forward KL $D_{KL}(P_{teacher} | P_{student})$ requires the student to cover all regions where the teacher's distribution has non-zero probability, imposing infinite penalty for any gaps — this causes the student to tend toward "mass spreading," distributing probability across all modes the teacher considers reasonable, resulting in conservative and averaged generations. Reverse KL $D_{KL}(P_{student} | P_{teacher})$ works in the opposite direction: it penalizes the student for assigning probability to regions the teacher considers impossible, but allows the student to ignore some of the teacher's modes — this leads the student toward "mode collapse," focusing on reproducing the teacher's most prominent behavioral characteristics. In agentic tasks, this means the student more precisely imitates the teacher's most common, highest-confidence decision chains rather than attempting to average-fit all possible teacher behaviors. This also explains why reverse KL tends to work better in on-policy distillation: the trajectories the student generates already reflect its current policy, and using reverse KL to drive it toward the teacher's dominant modes is more efficient than using forward KL to require coverage of the teacher's entire distribution.

Self-Distillation: Being Your Own Teacher

Self-distillation is an advanced topic worth attention — the teacher and student are different versions of the same model. The livestream introduced three main approaches:

Self-Distillation Fine-Tuning (SDFT): The model trains on high-quality samples it generated itself. A typical use case is capability recovery — for instance, after performing distillation to improve coding ability, you might find that instruction-following capability (measured by IFEval) has degraded. In this case, you can extract an early checkpoint of the model and perform self-distillation to recover instruction-following ability. This makes distillation a practical engineering tool for checkpoint capability management.

Self-Distillation Preference Optimization (SDPO): Two trajectories are created, with the second giving the model some kind of privileged information or hint to better complete the task. For example, for an agent coding problem that took five hours to solve, feeding the final solution directly into the starting trajectory produces a high-quality, more efficient reference trajectory. The loss is then computed from both trajectories.

On-Policy Self-Distillation: The best trajectory is selected from a set of runs as the superior sample, and the difference between each generation and the best run is computed.

There is an important mechanistic explanation for why self-distillation can recover or improve capabilities in certain scenarios: neural networks often suffer from "catastrophic forgetting" after fine-tuning for specific tasks — gradient updates for new tasks overwrite weights encoding old capabilities. By continuing to train on data generated from the model's own earlier checkpoints, self-distillation implicitly introduces a regularization constraint on the original capability distribution in the loss function, functioning similarly to EWC (Elastic Weight Consolidation) and other continual learning methods, but with simpler implementation. SDPO (Self-Distillation Preference Optimization) is highly similar in form to DPO (Direct Preference Optimization), with the key difference being the source of the "preference pairs": DPO relies on human-annotated preference data, while SDPO automatically constructs high-quality positive trajectory samples by injecting privileged information (e.g., final answers, key hints) into the model, avoiding expensive human annotation and making it more suitable for automated training pipelines.

The Limits of Distillation and Frontier Convergence

The main limitation of distillation is the performance ceiling imposed by the teacher model — under on-policy distillation, the student struggles to surpass the teacher. Self-distillation can break through in specific capabilities using privileged information, but still has an overall upper bound. Reinforcement learning's limit, on the other hand, lies in environment and reward design. In practice, "once you've exhausted the potential of distillation, you transition to reinforcement learning."

The more important insight is that frontier model training is no longer the domain of any single algorithm. Take NVIDIA's Nemotron paper as an example: the final model underwent a multi-stage training pipeline, using a series of domain expert models for distillation, mixing early checkpoints, multiple distillation strategies, and reinforcement learning strategies. Some checkpoints excel at STEM, others at code, and engineers combine different capabilities according to the use case.

The livestream also addressed several common practical questions:

• Can closed-source models be used as teachers? Only offline off-policy distillation (i.e., SFT) is possible, since closed-source models don't provide logits, making on-policy distillation — which requires scoring information — infeasible. However, synthetic data generation pipelines (such as Distilabel, Wisdom, etc.) are a mature approach for this kind of scenario.
• Can TRL use multiple teacher models? Multi-teacher on-policy distillation isn't directly supported, but it can be implemented through staged training or by extending the trainer — which is actually part of how frontier models are trained today.

Post-training has essentially become the process of managing different loss functions and signal sources at this stage — you need to decide what to use at each phase and where to get the signal. Distillation and reinforcement learning are increasingly converging on dimensions like process rewards. For developers looking to get started, the livestream team's advice is straightforward: start with lambda and beta at 0 or 1, train a model, observe its behavior, and dissect the data — that's the best way to truly understand what these parameters do.