Video summary

Training Agents 2: Live tutorial on model distillation for training custom agents.

Main summary

Key takeaways

Educational

Main ideas & lessons

  • Goal of the series: Learn how to train custom agents, with emphasis on how “training an agent” differs from other model-training approaches.
  • Where this session fits: This is post-training, specifically model distillation, to train smaller (cheaper/faster) models for targeted tasks.
  • Why distillation matters for agents:
    • Enables compression (smaller student model).
    • Enables capability transfer (teacher knowledge → student behavior).
    • Can support synthetic data generation.
  • Core conceptual framing: Distillation can be described along multiple axes; the most important ones here are:
    • Off-policy vs on-policy (what trajectories the student learns from relative to its own behavior)
    • Plus related axes (e.g., signal type, timing, teacher identity)

Training background (recap)

Common phases of LLM training:

  1. Pre-training: general language/world knowledge.
  2. Mid-training: code knowledge.
  3. Post-training: targeted behaviors (formatting, correct actions, etc.).

This stream focuses on post-training methods.

Distillation: key definitions and axes

1) Signal: “hard” vs “soft” learning targets

  • Hard-label distillation (SFT-like):
    • Student trains to imitate the teacher’s produced tokens (sequence-level text).
  • Soft-label distillation (logit/KD-like):
    • Student trains using the teacher’s logits (full output distribution).
    • Benefit: richer supervision because it reflects alternatives the teacher considered.

2) Accessibility: black-box vs white-box teacher

  • Black-box teacher: no logits access; only outputs/strings are available (possibly produced by an agentic harness).
  • White-box teacher: logits are available and the student can use them directly.

3) Timing: offline vs online

  • Offline: teacher data/trajectories are prepared before student training begins.
  • Online: teacher is consulted during/within training (teacher may be “in the loop”).

4) Teacher identity: separate teacher vs same model (self-distillation)

  • Separate teacher: teacher and student are different models.
  • Self-distillation: teacher is an earlier checkpoint/version of the same model family.

Off-policy distillation (Sergio + Ben overview)

Concept

  • Teacher generates data; student learns from it without generating its own trajectories during training.
  • In an agent/coding scenario:
    • Student sees prompts/tasks.
    • A frozen teacher (could be larger) generates solutions/trajectories.
    • Student trains to imitate those teacher outputs/logits.

Two ways to do off-policy distillation

  • Sequence KD (token/text-level imitation):
    • Student trains on the teacher-generated text/tokens, not the logits.
  • Logit KD (distribution-level imitation):
    • Student trains to match the teacher’s logits distribution, which is richer.

Implementation in TRL (Generalized Knowledge Distillation Trainer)

They describe using TRL’s generalized knowledge distillation trainer (configurable for off-policy vs on-policy) and tuning two parameters (“knobs”):

  • lambda: controls source of training data / policy behavior
  • beta: controls KL direction / how strictly the distribution is matched

For off-policy distillation, they set:

  • lambda = 0 → student learns from teacher’s fixed (offline) data
  • beta = 0 → uses the teacher distribution corresponding to forward KL (described as matching more of “everything the teacher considers”)

They note these parameters can be adjusted between extremes (0 and 1) to switch behaviors (forward vs reverse KL; off-policy vs on-policy).

Experiment they ran (reproducing prior class idea)

  • Task/domain: coding agent
  • Method compared: distilling a student coding agent from a teacher using knowledge distillation
  • Models:
    • Teacher: ~4B parameters
    • Student: ~0.6B parameters
  • Data: same dataset as prior session: coding traces (from an agentic coding session; traces available on the hub)
  • Procedure:
    • Train three student variants using different learning rates
    • Track losses/evaluation progress in TrackAI
  • Result framing: losses decrease for the student models, indicating distillation is working; presented as a fast, short illustrative run.

On-policy distillation (Ben)

Concept

  • Student generates its own trajectories during training, then a frozen teacher evaluates them.

Schema:

  1. Dataset provides prompts
  2. Student (smaller model) generates a completion (may make mistakes)
  3. Teacher performs a single forward pass (prefill stage, not full decode) to obtain distributions
  4. Compute per-token loss using reverse KL (student vs teacher distributions)
  5. Update student to reduce mismatch where it matters most for the student’s own generated attempts

Why on-policy is beneficial

  • Supervision is more relevant to the student’s actual mistakes because the training loop depends on student-generated trajectories.
  • Contrast with off-policy: a teacher trajectory might be “too advanced” and miss what matters for the student’s failures.

Key references/lineage

  • On-policy distillation attributed to DAGGER (framed as iterative SFT with teacher-in-the-loop evaluation).
  • Also mentioned as used in modern systems (e.g., GLM 5.2, Gwen, DeepSeek).
  • Suggested reading:
    • Hugging Face blog post(s)
    • Thinking Machines blog post (detailed)
    • Additional papers listed in description (not enumerated fully in subtitles)

Loss details: reverse KL and “mode-seeking”

  • Emphasized points:
    • Loss is per token, conditioned on the student’s trajectory context.
    • Reverse KL behaves more mode-seeking vs forward KL being more mean-seeking.
    • Intuition: reverse KL encourages focus on the teacher’s most relevant probability mass/gradient peaks, improving student behavior on its own trajectory.

Implementation in TRL via generalized trainer (knob changes)

Reusing the generalized KD trainer configuration, switching knobs for on-policy:

  • lambda = 1 → on-policy behavior (student generates trajectories)
  • beta = 1 → reverse KL behavior instead of forward KL

Practical experimentation notes:

  • Vary lambda to blend on-policy/off-policy (control exposure to student vs teacher trajectories).
  • Vary beta to adjust from mode-seeking toward less mode-seeking, depending on observed training/evals.

Experiment they ran for on-policy

  • Same coding-agent distillation setup as off-policy, but with on-policy settings
  • Models (example):
    • Teacher: ~4B
    • Student: ~0.6B
  • Data: same coding traces dataset from hub
  • Procedure:
    • Train three variants with different learning rates
    • Evaluate and plot losses via TrackAI
  • Takeaway: with lambda/beta at the on-policy extreme (1,1), the run with the lowest eval loss is highlighted (illustrative proof-of-concept).

Self-distillation (Ben’s overview; no major experiment shown)

What it is

  • Using an earlier checkpoint/version of the same model as a teacher, or using self-generated privileged improvements.
  • Presented as especially practical for checkpoint-based ability management (restore/transfer capabilities across phases).

Three self-distillation approaches

1) Simple self-distillation (SFT-like self improvement + checkpoint restoration)

  • Idea:
    • After domain-focused training improves one capability, you may lose another (e.g., instruction following).
    • Use an earlier checkpoint (better instruction following) as the teacher to restore that capability.
  • Practical pipeline purpose:
    • Save checkpoints, then reuse them to transfer/restore abilities across training phases.

2) Self-distillation via policy optimization (SDPO-style)

  • Idea:
    • Create two traces:
      • A baseline/less optimal generation trace
      • Another trace with privileged information or hints enabling better performance
    • Train using the improved trace’s advantage, teaching the model to adopt better behavior conditioned on the hint/extra context.
  • Examples of “privileged information”:
    • Final solution inserted back into the beginning context
    • Environment/tool feedback from a coding system
  • Mentioned reference:
    • SDPO trainer in TRL; also recommended videos/blog posts.

3) Self-distilled fine-tuning (SDFT-style)

  • Idea:
    • Generate a set (“group”) of candidate traces.
    • Select the best traces.
    • Train by encouraging the model to move its generations toward the best traces (by contrasting each generation with the best run).
  • Relationship to RL:
    • Self-distillation and reinforcement learning can blur when “process rewards” or environment-derived signals appear.
    • Post-training becomes a mix of loss functions and reward-signal construction.

Limitations / “ceilings” of distillation

  • Distillation performance is capped by the teacher’s quality:
    • On-policy: student generally can’t exceed teacher abilities beyond what the teacher can evaluate/teach.
    • Self-distillation: can surpass via privileged information or checkpoint reuse/restoration.
  • RL ceilings differ:
    • RL limits are tied to environment interaction and reward design.

Methodology / instruction-like summary (how to reproduce the presented workflow)

Use TRL generalized knowledge distillation trainer

Configure lambda and beta to choose behavior:

  • Off-policy distillation
    • lambda = 0
    • beta = 0 (forward-KL-like distribution coverage)
  • On-policy distillation
    • lambda = 1
    • beta = 1 (reverse-KL-like distribution matching)

Prepare teacher/student and data

  • Select:
    • a teacher model (often larger) and freeze it
    • a student model (smaller)
    • a dataset of prompts or agent traces (e.g., coding traces)

Training loop behavior

  • Off-policy: teacher generates trajectories ahead of student training; student trains on that fixed dataset.
  • On-policy: student generates completions; teacher evaluates via a forward pass (prefill) and per-token loss is computed.

Run hyperparameter comparisons

  • Train multiple students with different learning rates
  • Log/evaluate losses via TrackAI
  • Pick the best learning-rate variant using eval loss trends

Self-distillation extensions (optional)

  • Use prior checkpoints as teachers to restore lost capabilities
  • Optionally use SDPO/SDFT trainers in TRL with:
    • privileged-hint trace construction, or
    • best-trace selection

Speakers / sources featured

  • Ben (main presenter): distillation overview, on-policy distillation, self-distillation, and Q&A explanations
  • Sergio (co-presenter): off-policy distillation, implementation notes, and experiment walkthrough

Sources referenced (papers/blogs/systems mentioned)

  • Hinton (2015) — distillation / model extraction (historical framing)
  • DAGGER — iterative on-policy distillation framing
  • GLM 5.2, Gwen, DeepSeek — mentioned as using on-policy distillation ideas
  • SDPO (Sasha Rush video mentioned)
  • Thinking Machines blog post (recommended)
  • Neumotron 3 (paper mentioned; post-training pipeline overview)
  • Data“D E I T A” (synthetic data generation paper mentioned; abbreviated as “data d e i t a” in subtitles)
  • WizardLM (synthetic data pipelines mentioned)
  • HF course page on inference (prefill vs decode) at hf.co/learn

Software/documentation

  • TRL (Transformer Reinforcement Learning) and its distillation trainers/documentation
  • TrackAI (experiment tracking/plots)
  • HF hub (datasets/collections/traces referenced as accessible there)

Original video