Small-Model Distillation — Part 5: On-Policy Self-Distillation for a 0.8B SQL Agent

TL;DR

I had a 0.8B SQL agent scoring 69/220 on the fixed eval split, built through four stages of distillation from larger teacher models. The question was whether it could improve further by teaching itself. No larger teacher, no corrected actions, just the model sampling its own next action and training on probability feedback from a frozen copy of itself.

The answer came in two parts. When I stacked self-distillation on top of the 69/220 model, every variant degraded the parent. The self-teacher at 69/220 was too similar to the student to provide useful signal. But when I ran self-distillation from the same 67/220 checkpoint that Blog 4 used as its starting point for large-teacher on-policy distillation, five steps of self-OPD matched the 35B teacher’s best result: both reached 69/220. The self-teacher can match the large teacher, but its signal is shorter-lived. At ten steps, the 35B teacher still gained +2 while the self-teacher degraded to 63/220.

The practical finding: you can skip the large teacher for on-policy distillation, but the step budget matters more. The large teacher’s advantage is signal durability, not signal quality.

What I Wanted To Test

At this point in the series, the small model was no longer a base model pretending to be an agent. It had learned the SQL action schema, could inspect schemas, run queries, read observations, and submit final SQL through the fixed harness. Four rounds of distillation had taken it from 1/220 to 69/220. The natural next question was: if the model already knows how to act, can it learn from its own current behavior?

I wanted this post to be strict about the method. Earlier posts used hard-token targets: a teacher wrote a corrected action, I verified it in the environment, and the student learned to imitate that action with supervised fine-tuning. Here I wanted probability feedback only. The current student samples tokens, and a teacher scores those exact tokens. If the teacher writes a replacement action, the method is hard-token correction SFT, not on-policy distillation.

That boundary matters because it changes what the model can learn. Hard-token correction can show the model a new successful path. Token-level on-policy distillation can mostly reshape the model’s probability distribution around behavior it already samples. That can clean up local behavior, but it may not invent a missing multi-step SQL strategy.

I also wanted to answer a second question that the series had not addressed directly: how much does the large teacher actually buy you in on-policy distillation? Blog 4 used a Qwen3.5-35B-A3B teacher to go from 67/220 to 69/220. Could self-distillation from the same 67/220 starting point reach a similar level? If yes, the practical implication is significant: you can do on-policy distillation without serving a large teacher.

The SQL-Agent Problem

The task is a small tool-use SQL agent. The model receives a natural-language question, interacts with a SQLite database, and must eventually submit SQL. The harness owns the action schema, parsing, SQL execution, observations, stop rules, and scoring. I kept the benchmark, train/eval split, parser, scorer, action schema, max turns, and eval data fixed.

BAML, or Boundary Markup Language, is the structured-output layer around the model call. In this project it defines the action objects, renders the prompt, parses the model response, and gives the harness something executable. I did not optimize BAML in this experiment. It is part of the fixed task interface.

A trace is not just one answer. The model may first inspect the schema, then run a draft query, then use the observation to submit a final query. A simplified action shape looks like this:

action = {
    "draft": "I need to inspect the relevant tables first.",
    "output": {"action": "inspect_schema"},
}

The harness turns that action into a new state. If the model emits malformed JSON, the parser fails. If it repeats the same action, the harness can stop it. If it submits SQL that executes but answers the wrong question, the scorer marks the task failed. This is why final eval success matters more than training loss.

The Training Idea

On-policy distillation means the training signal comes from the current policy, or from a very fresh version of it. In this setup, the current policy is the 0.8B SQL agent. The core loop is:

student sees current SQL-agent state
student samples next assistant/action tokens
self-teacher scores those exact sampled tokens
student updates toward the teacher probabilities

The teacher is a frozen copy of the same checkpoint. What makes it probability distillation is that the target is a distribution over the sampled tokens, not a teacher-written replacement action.

The most important training knob is lmbda. In TRL’s GKD-style trainers, lmbda=1.0 means the trainer samples completions from the current student during training. With lmbda=0.0, the run becomes fixed-buffer knowledge distillation on stored completions. That can be useful, but I do not count it here as strict on-policy self-distillation.

For the privileged-context variant, the student still sees the normal prompt. Only the teacher side sees train-only extra information, such as the reference answer for that train task. The teacher is not allowed to write a new answer. It only scores the student’s sampled tokens under a richer context.

Educational Code Walkthrough

The prompt/state row for strict OPD is deliberately boring. It contains the SQL-agent state before the next model action. It does not contain the target action.

row = {
    "messages": [
        {"role": "system", "content": sql_agent_instructions},
        {"role": "user", "content": current_question_and_observation},
    ],
    "metadata": {
        "row_contract": "prompt_state_only_no_assistant_target",
        "prompt_tokens": 1945,
    },
}

During training, the trainer samples the action tokens. The teacher scores those same tokens.

prompt = render_sql_agent_state(row)
sampled_tokens = student.generate(prompt)
teacher_logprobs = teacher.score(prompt, sampled_tokens)
student_logprobs = student.score(prompt, sampled_tokens)

loss = distillation_loss(student_logprobs, teacher_logprobs)

This runs through the nablo package’s train_opd function, which wraps TRL’s GKDTrainer with lmbda=1.0. The teacher is loaded as a frozen LoRA adapter over the same base model using PeftModel.from_pretrained(model, adapter, is_trainable=False). Both models fit on a single 16GB GPU because the 0.8B base is small and only the student’s LoRA adapter is trainable.

For privileged self-distillation, the student still sees the normal prompt. Only the teacher side sees train-only extra information.

student_prompt = normal_sql_agent_state
teacher_prompt = normal_sql_agent_state + train_only_reference_answer
sampled_tokens = student.generate(student_prompt)
teacher_logprobs = teacher.score(teacher_prompt, sampled_tokens)
student_logprobs = student.score(student_prompt, sampled_tokens)
loss = distillation_loss(student_logprobs, teacher_logprobs)

This runs through TRL’s experimental SDPOTrainer with distillation_weight=1.0, which keeps the loss as probability distillation rather than reinforcement learning.

Algorithm Walkthroughs

Algorithm 1: Exact Frozen-Self OPD

The student and teacher start from the same checkpoint, but the teacher is frozen while the student updates.

Algorithm: Exact Frozen-Self On-Policy Distillation
Input:
  M0: the SQL-agent checkpoint (69/220 or 67/220)
  train_tasks: fixed 879-task train split
  eval_tasks: fixed 220-task eval split
  harness: fixed SQL-agent parser, tools, scorer, and stop rules
Output:
  M1: self-distilled student adapter
  eval_1: full eval report
Data:
  prompt_state_rows: one row before each assistant/action span
Steps:
  1. Run M0 through the unchanged train harness.
  2. Store prompt/state rows before assistant actions.
  3. Load M0 as the trainable student.
  4. Load a frozen copy of M0 as the self-teacher.
  5. Train with standard TRL GKD and lmbda = 1.0.
  6. Let the trainer sample current-student action tokens.
  7. Score those same tokens with the frozen self-teacher.
  8. Update only the student LoRA adapter.
  9. Evaluate the produced adapter on all 220 eval tasks.
Method boundary:
  This is probability distillation, not hard-token correction, DPO, or RL.

Algorithm 2: Privileged Self-Teacher SDPO

Exact self-distillation has a weak signal because the teacher and student begin almost identical. The privileged variant gives the teacher side more train-only information while keeping the student side normal.

Algorithm: Privileged-Context Self-Distillation
Input:
  M0: the 69/220 SQL-agent checkpoint
  train_tasks: fixed train split with reference answers available
  eval_tasks: fixed eval split with no privileged context
Output:
  M1: student adapter trained from privileged self-teacher probabilities
Data:
  student_prompt: normal SQL-agent prompt/state
  teacher_prompt: normal prompt/state plus train-only reference answer
Steps:
  1. Build rows from normal train prompt/states.
  2. Add train-only reference answer to the teacher context only.
  3. Keep rows under the context budget.
  4. Train with standard TRL SDPO and distillation_weight = 1.0.
  5. Let the student sample completions from the normal prompt.
  6. Append the same sampled completion tokens to the teacher prompt.
  7. Score those tokens with the teacher-side context.
  8. Update the student under the normal prompt.
  9. Evaluate with the normal harness and no privileged information.
Method boundary:
  This is not hard-token SFT because the reference answer is never used as
  the student's target text.

Dataset And Filtering

I used two sets of prompt/state rows, both collected from the fixed 879-task train split:

• From the 69/220 parent: 2,854 rows (1,956 difficult-state + 898 successful-anchor) from 289/879 parent train successes. Used for the self-OPD-from-69 experiments.
• From the 67/220 parent: 2,847 rows (1,983 difficult-state + 864 successful-anchor) from 280/879 parent train successes. Used for the parallel self-OPD-from-67 experiments. These are the exact same rows Blog 4 used for its 35B-teacher OPD runs.

The prompt lengths were wide: min 580, p50 1,945, p90 3,027, p95 3,280, and max 5,780 tokens. The training context budget was 8,192 tokens. For privileged SDPO, the teacher prompt was longer because it included train-only reference context, so I filtered to shorter prompt states.

I did not filter eval tasks, and I did not change the eval harness. Every reportable score below is a full run on the same 220 held-out tasks.

Training Setup

The student was unsloth/Qwen3.5-0.8B trained as a LoRA adapter (rank 32, alpha 32). The exact and reverse-KL self-OPD controls used standard TRL GKDTrainer from trl.experimental. The privileged-context runs used standard TRL SDPOTrainer from trl.experimental. I used distillation_weight=1.0 for SDPO and lmbda=1.0 for OPD-style runs so the runs stayed probability distillation, not reinforcement learning or preference optimization.

I tested two KL configurations:

• Forward-KL (beta=0.0, loss_top_k=20): the teacher’s top-20 token probabilities plus tail mass. This is the mass-covering objective that Blog 4 found most effective with the 35B teacher.
• Reverse-KL (beta=1.0, loss_top_k=1): only the teacher’s top-1 token. This is the mode-seeking objective that applies the gentlest update.

And two learning rate regimes:

• Standard (lr=5e-6, 128-token completions): matching Blog 4’s exact best configuration.
• Gentle (lr=5e-7, 32-token completions): a 10x lower learning rate with shorter completions.

All runs used full eval after training, with gc.collect() and torch.cuda.empty_cache() between training and eval to free GPU memory on the 16GB RTX 4080.

Results As Research Questions

1. Did any self-distillation variant improve the 69/220 parent?

No. Every variant degraded the model.

The best result from the 69/220 parent was 67/220 (5-step forward-KL at lr=5e-6), still a -2 regression. The 10-step run at the same learning rate collapsed to 59/220, losing ten parent wins and finding zero new ones. Lowering the learning rate to 5e-7 with 32-token completions helped marginally (63 vs 59 at 10 steps) but still degraded. The reverse-KL 10-step at 5e-7 was the most stable variant (66/220, only -3), because the mode-seeking objective applies a narrower update.

The privileged SDPO runs gave the self-teacher reference answers that the student could not see. The 5-step runs reached 65/220 with four unique wins, but still lost eight parent wins. The 10-step run collapsed to 35/220: with a neutral reward function (all zeros, since this is pure distillation not RL), the SDPO trainer’s DPO component has no signal, and the flat-reward warnings that TRL emits became reality at 10 steps.

2. Can self-distillation match the 35B teacher from the same starting point?

Yes, with the right step budget. The 5-step forward-KL self-OPD run from 67/220 reached 69/220, exactly matching Blog 4’s best result with the 35B teacher.

Config	Self-OPD	35B-OPD (Blog 4)
5 steps, lr=5e-6, 128 tok, fwd-KL	69/220 (+2)	67/220 (0)
10 steps, lr=5e-6, 128 tok, fwd-KL	63/220 (-4)	69/220 (+2)
10 steps, lr=5e-7, 32 tok, fwd-KL	64/220 (-3)	-
10 steps, lr=5e-7, 32 tok, rev-KL	66/220 (-1)	-

But at 10 steps, the self-teacher degraded to 63/220 while the 35B teacher reached 69/220. The self-teacher’s signal runs out after about five steps. The 35B teacher, by contrast, has genuine knowledge the student lacks (it scores 96/220 vs the student’s 67/220), so its signal does not exhaust as quickly.

3. Did the variants find new wins or just lose old ones?

Both. The pattern across all experiments was consistent: self-distillation moves the policy into new territory but cannot hold onto the parent’s wins.

The five-step run from 67/220 was the exception: it found four new wins and lost only two, for a net +2. Every other variant lost more than it gained. The oracle union (tasks solved by either the parent or the candidate) was consistently higher than either model alone, reaching 71-74/220 depending on the variant. That means there is complementary behavior to capture, but a single model cannot hold it all.

4. Why does the self-teacher match the 35B teacher at 5 steps but not at 10?

At five steps from 67/220, the self-teacher has a small but real advantage over the student: the student has already updated its LoRA adapter, while the teacher is still the original frozen checkpoint. That gap is the signal. The student moves toward the teacher’s distribution, which sharpens its behavior on tasks where it was uncertain.

By step ten, the student has converged toward the teacher on the easy cases. The remaining signal is from hard cases where the self-teacher itself is wrong. Training on those cases moves the student in random directions. The 35B teacher, by contrast, can still teach new things at step ten because it has higher capability. Its signal is a knowledge-transfer signal that is more durable. The self-teacher’s signal is a sharpening signal that exhausts itself quickly.

Exact Numbers Behind The Charts

Self-OPD and SDPO from 69/220:

Method	Full eval	Submitted	SQL errors	Parse fails	Repeated action	Max-turn	Avg turns
Parent (69/220)	69/220	174	8	7	38	1	3.22
fwd-KL 5-step, lr=5e-6, 128 tok	67/220	171	7	5	44	0	3.24
fwd-KL 10-step, lr=5e-6, 128 tok	59/220	168	13	5	46	1	3.33
rev-KL 10-step, lr=5e-7, 32 tok	66/220	171	13	8	40	1	3.23
SDPO solution 5-step	65/220	171	13	6	42	1	3.27

Self-OPD from 67/220 (parallel comparison):

Method	Full eval	Submitted	SQL errors	Parse fails	Repeated action	Max-turn	Avg turns
Parent (67/220)	67/220	174	8	5	40	1	3.21
fwd-KL 5-step, lr=5e-6, 128 tok	69/220	176	6	8	35	1	3.19
fwd-KL 10-step, lr=5e-6, 128 tok	63/220	172	6	6	40	2	3.24

The 5-step run from 67/220 is notable: it submitted more often (176 vs 174), had fewer repeated-action stops (35 vs 40), and fewer SQL execution errors (6 vs 8). The training cleaned up the action loop and the student submitted more confidently. That confidence translated into correct SQL often enough to gain two tasks.

Failure Analysis

The dominant failure did not disappear. Submitted-but-wrong SQL stayed high across all variants. The best run from 67/220 submitted 176 times and solved 69, meaning 107 submitted queries still failed the hidden tests. The probability distillation update changed how the student distributes mass over action tokens, but it did not fundamentally change whether the student’s final SQL was correct. This is the same wall every post in the series has hit: the student knows how to run the loop, but its SQL judgment is not strong enough to pass the scorer on the hard tasks.

The step-budget cliff is the clearest negative result from the self-OPD-from-69 experiments. Forward-KL at lr=5e-6 degraded monotonically: 5 steps lost 2, 10 steps lost 10, 15 steps lost 8. The on-policy samples at step 10 and 15 were coming from a student that had already drifted, so the self-teacher was scoring worse completions than it was at step 5. This is the compounding-drift risk of on-policy methods: each step’s data depends on the current policy, so a few bad steps can cascade.

The SDPO 10-step collapse to 35/220 is a separate failure mode. With distillation_weight=1.0 and a neutral reward function, the SDPO trainer’s DPO component has no signal (all rewards are zero). TRL warns about “flat SDPO rewards” at every step. At 5 steps, the distillation loss alone is enough to move the policy usefully. At 10 steps, the degenerate reward signal compounds and the policy collapses. This is an inherent limitation of pure distillation with SDPO: without a real reward signal, the DPO component is dead weight that can destabilize training.

The deeper pattern is forgetting. The best self-distillation variants had useful unique wins, but they lost parent wins at the same time. The five-step run from 67/220 was the only variant where new wins exceeded lost wins (4 vs 2). Every other variant lost at least as many as it gained. That means the next useful method is probably not “more self-OPD steps.” It is a retention-aware consolidation method, or a different training stage, that can absorb complementary wins without erasing the stable parent behavior.

Hardware And Infrastructure Lessons

These runs used a single remote NVIDIA RTX 4080 (16GB VRAM) for both training and eval. Unlike Blog 4, which split the Mac (35B teacher server) and GPU (student training), this self-distillation work needed no Mac at all. The teacher was either the same 0.8B model loaded as a frozen PEFT adapter (for GKD-based OPD) or the same model with privileged context inside the SDPO trainer. The entire pipeline ran on one GPU.

The 16GB VRAM was enough for same-process GKD with a frozen teacher. The 0.8B base model is small, and only the student’s LoRA adapter is trainable. The frozen teacher adapter adds negligible memory overhead because it shares the same base model weights. The main memory consumers were the training batches with 128-token completions and the full-vocabulary teacher logits for the top-20 forward-KL loss.

The practical bottleneck was not training but eval. Each 220-task eval took about 27 minutes because every eval task is a multi-turn agent episode with SQLite execution. A 10-step training run took about 24 minutes, so training and eval were roughly equal in wall time. I kept the experiment ladder small and evaluated after every candidate because the cost of a full sweep is dominated by eval, not by training.

One implementation detail mattered enough to record. The Unsloth model loader creates a Qwen3VLProcessor wrapper that is incompatible with TRL’s GKDTrainer, which expects a standard HuggingFace tokenizer with a backend_tokenizer attribute. When the parent adapter was saved by Unsloth (indicated by the unsloth_fixed flag in adapter_config.json), I used the Unsloth loader. When it was a standard HuggingFace PEFT adapter, I fell back to AutoModelForCausalLM.from_pretrained with standard HuggingFace loading. The nablo package’s train_opd function handles this automatically by checking the adapter format.

I also needed to explicitly release GPU memory between training and eval. Without gc.collect() and torch.cuda.empty_cache(), the eval subprocess would OOM because the training process held GPU memory. On one occasion, a zombie training process from an earlier run held 7.2GB of VRAM, causing the eval to fail with an out-of-memory error. Killing the zombie and adding the cleanup calls fixed this.

What I Learned

The clean answer is that self-distillation can match the large teacher, but only with the right step budget. Five steps of forward-KL self-OPD from 67/220 reached 69/220, exactly matching Blog 4’s best result with the 35B teacher. That is a real result with a practical implication: you can do on-policy distillation without serving a large teacher model, as long as you keep the step count low and evaluate after every candidate.

The more useful lesson is about signal durability. The self-teacher’s signal is a sharpening signal: it pushes the student’s distribution toward the frozen teacher’s distribution on actions the student already samples. That signal is real but short-lived. By step five, the student has absorbed the useful sharpening. By step ten, the remaining signal is from hard cases where the self-teacher itself is wrong, and training on those cases moves the student in random directions. The 35B teacher’s signal is a knowledge-transfer signal: it can teach the student things the student does not know. That signal is more durable because the 35B teacher has higher capability (96/220 vs the student’s 67/220) and can still provide useful guidance at step ten.

Stacking self-distillation on top of the large-teacher result (from 69/220) did not help. Every variant degraded the parent. The self-teacher at 69/220 is even more similar to the student than at 67/220, so the signal-to-noise ratio is worse. The practical takeaway: self-distillation and large-teacher distillation are alternatives, not complements. You pick one or the other, and you tune the step budget accordingly.

The forward-KL-versus-reverse-KL distinction mattered here too, just as it did in Blog 4. Forward-KL top-20, which spreads the student’s mass toward teacher-likely alternatives, produced the only gain from 67/220. Reverse-KL top-1, which sharpens toward the teacher’s single pick, was more stable at 10 steps (only -1 from 67/220) but never improved the parent. For a small agent that already has fragile but useful behavior, spreading mass toward alternatives was the safer update when the step budget was right, and collapsing onto one token was the safer update when the step budget was too high.

Conclusion

The current best single model remains 69/220, reachable either by a 35B teacher with 10 steps or by self-distillation with 5 steps from the 67/220 checkpoint. Self-distillation is a viable replacement for the large teacher in on-policy distillation, but only with careful step budget control. The large teacher’s advantage is signal durability, not signal quality: it can teach for more steps before the signal runs out.

For practical use, this means a small-model-only training pipeline is possible. You do not need to serve a 35B model to do on-policy distillation on a 0.8B agent. You do need to keep the step count low and evaluate after every candidate, because the self-teacher’s signal exhausts quickly and extra steps actively damage the model.

The next direction is not more self-OPD steps. It is a retention-aware consolidation method that can absorb the complementary wins without erasing stable parent behavior, or a training stage that can optimize task success more directly.