TL;DR
I started this post from the best hard-token correction-family SQL agent I had so far: unsloth/Qwen3.5-0.8B, evaluated at 67/220 on the fixed eval set. That parent model came from supervised fine-tuning, or SFT, on hard assistant/action tokens, first from teacher traces, then from DAgger-style expert corrections, and finally from retention-heavy anti-forgetting consolidation. Here I wanted to test a different signal: keep the same 0.8B student, keep the same SQL harness, but use a larger same-family Qwen/Qwen3.5-35B-A3B teacher to score the student's own sampled action tokens with log probabilities.
The clean result was small but real. The best on-policy probability distillation run reached 69/220, a +2 task gain over the 67/220 parent. It used TRL's standard DistillationTrainer, lmbda=1.0, a forward-Kullback-Leibler-style setting, top-20 teacher probabilities, 10 optimizer steps, and full 220-task evaluation. The negative result is just as important: longer runs, a lower learning rate, a second fresh-state round, and a different random seed all got worse. Probability distillation helped as a careful nudge, not as a magic new capability source.
One honest caveat up front: the teacher here scores tokens the student already samples. If the student never samples a good SQL strategy from a given state, token-level teacher probabilities cannot invent that missing behavior. The gain came from cleaning up the student's local action distribution, not from teaching it entirely new SQL.
What I Wanted To Test
The question was simple: can a larger same-family teacher improve a small SQL tool-use agent by scoring the tokens the student already produces? This is different from asking the teacher to fix the task. In the earlier hard-token correction work in Blog 3, the teacher wrote a target action and the student learned to imitate it. In this post, the teacher does not write the action. The current student samples an assistant/action span, and the larger teacher scores those exact tokens.
That distinction matters because the learning signal is softer. A hard target says, "write this next action." A probability target says, "given this state and this sampled action, here is how the larger model distributes probability over the action tokens." For a tiny SQL agent, I wanted to see whether this dense token-level feedback could clean up behavior without changing the benchmark, parser, scorer, harness, action schema, train split, eval split, max turns, or eval budget.
This post is the fourth in the series, and each post has taught the student from a different kind of signal. Blog 1 used offline hard-token teacher traces. Blog 2 added off-policy top-k soft labels over those frozen traces. Blog 3 moved the states on-policy (student-created) but kept the targets hard (teacher-written corrections). This post moves the targets on-policy too: the student samples the tokens, and the teacher only scores them. That is the clean progression the series was building toward.
The SQL-Agent Problem
The task is still the same SQL-agent benchmark from the earlier posts. The model receives a natural-language SQL debugging or query-writing task, can inspect the schema, can run SQL against a temporary SQLite database, reads the environment observation, and eventually submits final SQL. The fixed eval set has 220 tasks. A task is counted as solved only when the submitted SQL passes the scorer.
BAML, or Boundary Markup Language, is the structured-output layer around the harness. In this project it renders the model request, defines the action schema, parses the model response, and gives the harness an executable object. I did not tune BAML or the harness in this post. A typical action still looks like this:
{"draft":"Inspect the schema before writing SQL.","output":{"action":"inspect_schema"}}
{"draft":"Test the join and filter.","output":{"action":"run_sql_query","sql":"SELECT ..."}}
{"draft":"Submit the corrected query.","output":{"action":"submit_sql","sql":["SELECT ...;"]}}
The failure categories are the same four terminal outcomes tracked across this series. A parse failure means the model did not produce a valid structured action. A repeated-action stop means the model repeated a previous action and the harness stopped the loop. A SQL execution error means SQLite rejected a query the model ran or submitted. A submitted-but-failed task means the model produced final SQL, but the scorer rejected it. These categories matter because an agent can become cleaner without becoming more correct. A model can submit more often, repeat itself less, and still submit the wrong query. I track all four categories throughout this post, not just the success count.
The Training Idea: Off-Policy Versus On-Policy
The simplest way to understand this post is to place it next to the earlier ones along two axes: where do the training tokens come from, and what does the student learn from at each position?
In Blogs 1 and 2, both the prompt state and the assistant target came from the teacher. The teacher ran the harness, left successful traces, and each teacher action became a supervised row. Blog 2 added a richer target, top-k teacher probabilities at each position, but the rows were still frozen teacher traces. The student never acted while those targets were being scored. That is off-policy distillation: the data comes from another policy, and it does not change between training steps.
In Blog 3, the prompt state moved on-policy. The student rolled out, reached its own states, and the teacher wrote a continuation from there. But the target was still a hard token: the teacher's chosen action text. So the state distribution was the student's, while the target was the teacher's.
This post moves the target on-policy too. During training, the current student samples the assistant/action span. There is no stored target in the training rows at all. The prompt/state pool is collected from the harness before training, but the completion is sampled fresh by the current student inside the trainer, at every step. This is on-policy probability distillation.
Off-policy (Blog 1): teacher state -> teacher hard target
Off-policy (Blog 2): teacher state -> teacher soft-label target
On-policy state (Blog 3): student state -> teacher hard correction
On-policy target (Blog 4): student state -> student-sampled tokens, teacher-scored
The exact nuance is important. The prompt/state pool is collected from the SQL harness before training. Those rows are static during a training run. The on-policy part is the completion span sampled by the current student inside the trainer, which changes as the LoRA adapter updates. I also tested one refreshed-state second round, where I collected a new prompt/state pool from the 69/220 OPD model and trained again from that model. That second round did not help.
This is the conceptual loss shape:
prompt = sql_agent_state_before_next_action
sampled_tokens = current_student.sample(prompt)
teacher_logprobs = qwen_35b_teacher.score(prompt, sampled_tokens)
student_logprobs = current_student.score(prompt, sampled_tokens)
loss = distillation_loss(student_logprobs, teacher_logprobs)
The important phrase is "same sampled tokens." The student generates a completion, and the teacher must score that exact completion under the same prefix. If the teacher writes a replacement action, the method becomes hard-token correction SFT, which belongs to Blog 3. If the student trains on a frozen rollout buffer for many epochs without refreshing or sampling from the current policy, the method becomes more like the off-policy soft-label distillation of Blog 2. Here the teacher is a scorer, not a repair model.
Forward-KL Versus Reverse-KL, And The Three Trainer Knobs
The TRL DistillationTrainer exposes three knobs that decide the shape of the on-policy update. This post is mostly a study of how those knobs move a small SQL agent, so they deserve a real explanation.
lmbda controls whether the completion tokens come from the current student or from a stored buffer. lmbda=1.0 means the trainer samples completions from the current student policy at every step. That is what makes this on-policy. A lower lmbda would mix in stored completions, which moves the method back toward off-policy. Every reportable run in this post used lmbda=1.0.
beta controls the direction of the Kullback-Leibler, or KL, divergence in the distillation loss. This is the knob that decides whether the update is mass-covering or mode-seeking. Forward KL, which corresponds to beta=0.0 in this trainer, measures the expected divergence from the teacher distribution to the student distribution. It penalizes the student whenever the teacher has probability mass that the student does not cover, so the student is pushed to spread mass across all of the teacher's alternatives. That is why forward-KL runs can use top-20 teacher probabilities with tail mass: the student is being asked to respect the full shape of the teacher distribution, not just its peak.
Reverse KL, which corresponds to beta=1.0 here, measures the divergence from the student distribution to the teacher distribution. It penalizes the student for putting mass where the teacher has little. The student is pushed to concentrate on the teacher's single most likely token, because any student mass outside the teacher's peak is costly. That is why reverse-KL-style runs in this post used top-1 teacher scoring: the target collapses to the teacher's argmax, and the update becomes a soft version of "copy the teacher's pick."
The practical difference for a SQL agent is real. Forward-KL nudges the student to consider the same alternatives the teacher considers, which can make the action distribution richer but also noisier. Reverse-KL sharpens the student toward one token, which can make it more decisive but also more brittle. The results section shows that the forward-KL top-20 setting was the only one that beat the parent.
loss_top_k controls how many teacher alternatives survive in the sparse probability target. Full-vocabulary logits would be the most complete signal, but they are expensive to move across the network and expensive to train against. With an external teacher served over HTTP, the trainer asks the teacher for the top-k token logprobs at each sampled position. loss_top_k=20 keeps the top 20 plus a residual tail mass for everything else, which is the same sparse-plus-tail compromise I used in Blog 2. loss_top_k=1 keeps only the teacher's single most likely token, which pairs naturally with the reverse-KL setting.
Educational Code Walkthrough
The training row is a prompt/state row, not a supervised example with an assistant target. The row ends right before the next assistant action and contains no stored completion. A simplified row looks like this:
row = {
"id": "TRAIN_123_turn_2_difficult_state",
"messages": [
{"role": "system", "content": sql_agent_instructions},
{"role": "user", "content": task_and_current_observation},
{"role": "assistant", "content": previous_action_if_any},
{"role": "user", "content": previous_tool_observation_if_any},
],
"metadata": {
"task_id": "TRAIN_123",
"db_id": "books",
"task_category": "bugfix",
"state_role": "difficult_state",
"turn": 2,
"source": "blog3_final_student_harness_state",
"row_contract": "prompt_state_only_no_assistant_target",
"prompt_tokens": 1881,
},
}
Notice what is absent. There is no teacher_action, no target, no token_scores. The row is only the state. The trainer creates the sampled assistant completion during training, and the teacher scores it then. This is the hard boundary that separates this post from Blog 3: Blog 3 rows carried a verified teacher-written target, while these rows carry only the state.
The mask is what keeps the training objective honest. System text, user task text, previous tool observations, and previous environment outputs are context. The loss belongs only on the sampled assistant/action tokens that the current student generates during training:
system/user/tool-observation tokens -> context only -> loss mask = 0
current sampled assistant/action tokens -> trainable span -> loss mask = 1
teacher logprobs -> scores for the same sampled assistant/action tokens
During training, the trainer does three things per row. First, it renders the prompt and samples a completion from the current student. Second, it sends the prompt plus sampled token IDs to the external teacher server and receives teacher logprobs over the sampled positions. Third, it computes the distillation loss on the sampled positions only.
# Simplified on-policy sampling and scoring inside the trainer.
prompt_ids = tokenize(render_chat_template(row.messages, add_generation_prompt=True))
sampled_ids = current_student.generate(
prompt_ids,
max_new_tokens=128,
temperature=1.0,
top_p=0.95,
)
full_ids = concat(prompt_ids, sampled_ids)
# Ask the external teacher to score the exact sampled tokens.
teacher_response = teacher_server.post("/get_sequence_logprobs/", json={
"sequences": [full_ids.tolist()],
"prompt_lengths": [len(prompt_ids)],
"top_logprobs": 20,
"temperature": 1.0,
})
teacher_top_ids = teacher_response["top_token_ids"]
teacher_top_logprobs = teacher_response["top_logprobs"]
# Student scores the same sampled positions.
student_logits = current_student.forward(full_ids).logits
student_logprobs = log_softmax(student_logits, dim=-1)
# Loss only on sampled positions, not on the prompt.
loss = distillation_loss(
student_logprobs[len(prompt_ids):],
teacher_top_ids,
teacher_top_logprobs,
beta=0.0,
loss_top_k=20,
)
That snippet is the shape of the method, not production code. The real TRL DistillationTrainer handles batching, padding, KV-cache reuse, gradient accumulation, and the LoRA backward pass. But the contract is the same: the student samples, the teacher scores the same tokens, and the loss applies only to sampled positions.
The teacher and student are from the same model family, which keeps token scoring auditable. Both use the same Qwen vocabulary, so the token IDs the student samples are the same IDs the teacher scores. The trainer still checks tokenization compatibility before training and saves a one-batch proof that the teacher scored the same token span the student sampled. I will come back to that proof in the infrastructure section.
Algorithm Walkthrough
I like writing this as an algorithm because otherwise it is too easy to mix it up with hard-token correction. The teacher is not fixing the trajectory. The teacher is scoring the student's sampled tokens.
Algorithm: Larger-Teacher On-Policy Probability Distillation for a SQL Agent
Input:
M0: Qwen3.5 0.8B SQL agent from hard-token correction SFT, evaluated at 67/220
T: Qwen/Qwen3.5-35B-A3B probability teacher, served as mlx-community/Qwen3.5-35B-A3B-8bit
train_tasks: fixed 879-task SQL-agent train split
eval_tasks: unchanged 220-task SQL-agent eval split
harness: unchanged parser, action schema, SQL tools, scorer, and stop rules
trainer: TRL DistillationTrainer with lmbda=1.0 (current-student sampling enabled)
Output:
M1: Qwen3.5 0.8B adapter trained with larger-teacher token probabilities
eval_1: full 220-task eval report
Data:
Prompt/state rows collected from M0 running in the unchanged SQL harness
Rows end before the next assistant action and contain no assistant target
Runtime flow:
1. Run M0 on train_tasks in the unchanged SQL harness.
2. Save prompt/state rows before assistant turns.
3. Mix difficult states with successful-anchor states from the same parent.
4. Start TRL probability-distillation training from M0.
5. For each prompt/state row, sample assistant/action tokens from the current student.
6. Ask T to score the exact sampled tokens under the same prefix.
7. Mask system, user, and environment-observation tokens.
8. Update only sampled assistant/action tokens with the distillation loss.
9. Save the adapter, trainer config, teacher-request proof, and token-alignment proof.
10. Evaluate the adapter on eval_tasks with the unchanged harness.
Acceptance rule:
Count only full 220-task evals where the trainer proof shows current-student sampling
and larger-teacher scoring of the same sampled tokens.
Method boundary:
This is probability distillation, not teacher-written correction, RL, DPO,
adapter mixing, or a harness change.
The hardware flow also matters, and it is different from the earlier posts. In Blog 2, the teacher scoring was a batch pass over frozen traces on the Mac. In Blog 3, the teacher was a hosted API. Here the teacher is a live scorer that must respond during training, one request per microbatch. The student trained on a CUDA GPU, while the 35B teacher ran on my Mac through MLX and served logprobs over HTTP. No model weights or gradients crossed the network. The GPU sent prompt plus sampled token IDs to the Mac, and the Mac returned teacher log probabilities.
Dataset And Filtering
I collected the first prompt/state pool by running the 67/220 parent over all 879 train tasks in the unchanged SQL harness. The parent solved 280 train tasks during that rollout and produced 2,847 usable prompt/state rows. The mix had 1,983 difficult-state rows and 864 successful-anchor rows, at roughly a 2:1 difficult-to-anchor ratio.
The two row types do different things. Difficult states are prompt/state rows taken from trajectories where the parent eventually failed: wrong submissions, SQL execution errors, repeated-action stops, and near-max-turn states. These expose the student's current failure distribution, which is exactly where a teacher nudge could help. Successful-anchor rows are prompt/state rows taken from trajectories the parent solved. They do not contain the parent's actions as targets. They are just states from good trajectories, included so the on-policy update does not drift away from behavior that already works. This is the same retention intuition from Blog 3, adapted to the prompt/state-only contract: the anchor rows carry no hard target, but their presence in the pool keeps the sampled completions close to states the student already handles well.
The rows were filtered by rendered prompt length, not by target length, because there is no stored target. The maximum sequence length was 8,192 tokens. Most runs, including the selected run, used a 128-token assistant budget during OPD training, and Run A tested a smaller 64-token cap. The final parent-state pool had prompt-token statistics of min 580, median 1,938, p90 3,016, p95 3,276, and max 5,778. No rows were dropped for being malformed, having assistant targets, or being too long.
| Pool source | Train tasks | Parent train successes | Total rows | Difficult-state rows | Successful-anchor rows |
| --- | ---: | ---: | ---: | ---: | ---: |
| First pool from 67/220 parent | 879 | 280 | 2,847 | 1,983 | 864 |
| Second fresh-state pool from 69/220 OPD model | 879 | 289 | 2,854 | 1,956 | 898 |
I also collected a second fresh-state pool from the best 69/220 OPD model to test whether a second on-policy round would compound. That pool had 2,854 rows: 1,956 difficult states and 898 successful anchors. These were also prompt/state-only rows, not stored action targets. The pool was useful as an experiment, but the trained second-round model fell to 64/220.
Training Setup
The student is unsloth/Qwen3.5-0.8B, initialized from the Blog 3 hard-token correction-family checkpoint that scored 67/220. The teacher is Qwen/Qwen3.5-35B-A3B, served as mlx-community/Qwen3.5-35B-A3B-8bit on the Mac for logprob inference. The training backend uses TRL 1.6.0 experimental DistillationTrainer, PEFT LoRA, Transformers 5.12.1, PyTorch 2.10.0, and bf16 on the GPU.
The best run used LoRA, or Low-Rank Adaptation, instead of updating every student weight. LoRA rank was 32, alpha was 32, dropout was 0.0, and the trained modules were the attention and MLP projections: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, and down_proj. The sequence length was 8,192, max completion length during OPD training was 128, batch size was 1, gradient accumulation was 8, learning rate was 5e-6, and the selected run trained for 10 optimizer steps. The full stack was TRL 1.6.0, Transformers 5.12.1, PyTorch 2.10.0, PEFT 0.19.1, and bf16 precision on the GPU. Evaluation used the same max turns, 512-token generation budget, action schema, parser, scorer, and 220-task eval split as the parent.
| Setting | Value |
| --- | --- |
| Student | unsloth/Qwen3.5-0.8B |
| Starting checkpoint | Blog 3 hard-token correction-family adapter at 67/220 |
| Probability teacher | Qwen/Qwen3.5-35B-A3B, served as mlx-community/Qwen3.5-35B-A3B-8bit |
| Trainer | TRL 1.6.0 trl.experimental.distillation.DistillationTrainer |
| lmbda | 1.0 (current-student sampled completions) |
| beta | 0.0 (forward-KL-style update) |
| loss_top_k | 20 (top-20 teacher probabilities plus tail mass) |
| LoRA rank / alpha / dropout | 32 / 32 / 0.0 |
| LoRA target modules | q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj |
| Batch / grad accumulation | 1 / 8 |
| Learning rate | 5e-6 |
| Optimizer steps | 10 (selected run) |
| Max sequence length | 8,192 |
| Max completion length | 128 (Run A used 64) |
| Precision | bf16 |
| Transform / trainer stack | Transformers 5.12.1, PyTorch 2.10.0, PEFT 0.19.1 |
The three distillation knobs are the ones I defined earlier, and the sweep tested each of them. The selected run used lmbda=1.0, beta=0.0, and loss_top_k=20, which is the forward-KL-style top-20 setting. The reverse-KL-style runs used beta=1.0 and loss_top_k=1, collapsing the target to the teacher's single most likely token. That is the cleanest way to see the forward-versus-reverse distinction in the results: same student, same teacher, same prompt/state pool, different divergence direction and different target sparsity.
Results As Research Questions
Did Larger-Teacher Probability Distillation Improve The 67/220 Parent?
Yes, but only by a small amount. The best run reached 69/220. The pattern was sharper than I expected: reverse-KL-style top-1 updates underperformed, a 5-step forward-KL-style top-20 update tied the parent, a 10-step forward-KL-style top-20 update improved to 69, and longer updates fell back down.
This is the exact table behind the chart. I include the behavior counters because the success score alone hides how the policy moved. Submitted means the agent reached submit_sql; repeated, parse, and SQL-error counts are terminal failure modes from the unchanged harness.
| Run | Method setting | Eval success | Submitted | Repeated stop | Parse failures | SQL errors | Avg turns | | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | | Starting parent | Hard-token correction-family SFT checkpoint | 67/220 | 174 | 40 | 5 | 8 | 3.21 | | A | Reverse-KL-style, top-1 teacher token, 5 steps, 64-token cap | 65/220 | 171 | 42 | 7 | 7 | 3.25 | | B | Reverse-KL-style, top-1 teacher token, 5 steps, 128-token cap | 66/220 | 174 | 39 | 6 | 6 | 3.23 | | C | Reverse-KL-style, top-1 teacher token, lower learning rate, 5 steps | 62/220 | 170 | 43 | 7 | 5 | 3.16 | | D | Forward-KL-style, top-20 teacher distribution, 5 steps | 67/220 | 171 | 42 | 6 | 8 | 3.21 | | E | Forward-KL-style, top-20 teacher distribution, 10 steps | 69/220 | 174 | 38 | 7 | 8 | 3.22 | | F | Forward-KL-style, top-20 teacher distribution, 15 steps | 63/220 | 175 | 39 | 6 | 9 | 3.28 | | G | Forward-KL-style, top-20 teacher distribution, 12 steps | 65/220 | 172 | 41 | 7 | 11 | 3.29 | | H | Iterative fresh-state second OPD round from the 69/220 model, 5 steps | 64/220 | 176 | 38 | 6 | 14 | 3.23 | | I | Forward-KL-style, top-20 teacher distribution, 10 steps, seed 7 | 61/220 | 170 | 43 | 6 | 10 | 3.30 |
The selected model is Run E. I would not call the method broadly robust from this evidence. The same 10-step recipe with seed 7 scored 61/220, and a second round from the 69/220 model scored 64/220. My read is that this kind of probability distillation can nudge a good agent, but the update budget is fragile. It is easy to move the policy away from useful behavior.
The KL-direction pattern is the cleanest signal in the sweep. All three reverse-KL-style runs (A, B, C) scored below the parent, ranging from 62 to 66. All five forward-KL-style runs that used the same 128-token cap and the standard learning rate landed between 63 and 69. The forward-KL setting gave the student a richer target, the teacher's full top-20 distribution, and that richer signal was what produced the only gain over the parent. The reverse-KL top-1 setting collapsed the target to one token and consistently pulled the student slightly below where it started. That is consistent with the intuition: sharpening a small agent toward a single teacher token can overwrite fragile skills faster than it adds useful ones.
The step-count pattern is the other clean signal. Five steps of forward-KL tied the parent at 67. Ten steps improved to 69. Twelve steps fell to 65, and fifteen steps fell to 63. There is a narrow window where the update helps, and past that window the policy drifts. This is the same fragility Blog 3 saw with correction-only training: more updates did not keep helping.
These are the completed reportable full-eval OPD runs for this post. I also had earlier smoke and proof runs and older probability-distillation scaffolding from pre-final checkpoints, but I am not mixing those into the public result story because they do not start from the final 67/220 parent or do not match this post's larger-teacher OPD contract.
Did The Failure Shape Improve?
The best run solved two more tasks than the parent, but the failure shape did not transform. The parent had 107 submitted-but-failed tasks, 40 repeated-action stops, 5 parse failures, and 8 SQL execution errors. The best OPD run had 105 submitted-but-failed tasks, 38 repeated-action stops, 7 parse failures, and 8 SQL execution errors.
That is a small local cleanup, not a new SQL reasoning breakthrough. Repeated-action stops went down by two and submitted-but-failed tasks went down by two, but parse failures rose slightly from 5 to 7. SQL execution errors held flat at 8. The model became marginally cleaner in loop control without becoming more accurate in SQL judgment. This is the same pattern that recurred across the series: a small score gain can come from fewer repeated actions and slightly better submit decisions, while the dominant submitted-but-failed failure mode barely moves.
Did The Best OPD Model Solve Different Tasks?
Yes, a little. The best OPD model and the starting parent overlapped on 61 solved eval tasks. The OPD model solved 69 tasks total and the parent solved 67 tasks total. The OPD model gained 8 tasks that the parent missed, but lost 6 tasks the parent had solved. The oracle union would be 75, which is useful analysis but not a deployable single-model score.
| Overlap | Tasks | | --- | ---: | | Both solved | 61 | | Only best OPD model solved | 8 | | Only starting parent solved | 6 | | Neither solved | 145 | | Oracle union | 75 |
This overlap is the part that keeps showing up across the series. In Blog 3, the 64/220 and 67/220 checkpoints also traded tasks rather than simply accumulating them. Small training changes often move the solved-task set rather than simply adding solved tasks. For this post, I am not turning that into adapter mixing or selector logic, because that would be a different method. The lesson here is narrower: probability distillation changed the policy enough to trade a few tasks, and the net trade was +2.
Failure Analysis
If I only looked at the best run, the story would be too clean: forward-KL top-20 at 10 steps wins, full stop. The sweep shows the real, messier shape, and it matches what the other posts found. The headline lesson is the fragile one.
The gain was narrow and seed-sensitive. The selected run reached 69/220, but the same 10-step forward-KL top-20 recipe with seed 7 reached 61/220. That is an 8-task swing from a seed change alone, on a 220-task eval. I would not claim the method is robust from a single winning run. What I can say is that, under the right settings, the forward-KL top-20 update produced a small real gain, and the reverse-KL top-1 setting never did.
Submitted-but-wrong SQL is still the dominant failure. The best OPD run submitted 174 times and solved 69, which means 105 submitted queries still failed the hidden tests. That is only two fewer than the parent's 107. The probability distillation update changed how the student distributes mass over action tokens, but it did not meaningfully 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 second fresh-state round is especially instructive. Run H started from the 69/220 model, collected a new prompt/state pool from it, and trained another 5 steps. That model scored 64/220. It kept repeated-action stops low at 38, but SQL execution errors rose to 14, the highest in the sweep. More confident action was not automatically better action. The second round pushed the policy further from the parent, and some of that push turned into faster, wronger SQL. This mirrors the Blog 3 finding that correction-only iteration did not compound: once the parent is already strong, a second round of the same recipe can do more harm than good.
The update-budget cliff is the clearest negative result. Forward-KL top-20 improved monotonically from 5 steps (67) to 10 steps (69), then fell off a cliff: 12 steps reached 65 and 15 steps reached 63. The policy moved past the useful update window and into territory where the LoRA adapter was overwriting fragile skills. The on-policy samples at step 12 and 15 were coming from a student that had already drifted, so the 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.
My interpretation ties back to the method boundary. Hard-token correction (Blog 3) can show the student a better action it never would have sampled. Probability distillation (this post) scores only the actions the student already samples. If the student's current policy never puts meaningful mass on the right SQL strategy in a given state, token-level teacher probabilities may not be enough to invent that missing behavior. The forward-KL top-20 setting helps because it spreads the student's mass toward teacher-likely alternatives, but it still cannot steer the student toward a token the student essentially never considers. That is why the gain was +2 and not +10.
Hardware And Infrastructure Lessons
The Mac-plus-GPU split was the practical path, and it is different from the infrastructure pattern in the earlier posts. In Blog 2, the teacher scoring was a one-time batch pass over frozen traces on the Mac, done before any GPU training. Here the teacher is a live scorer that must respond during training, one request per microbatch, because the completions are sampled on-policy and are not known ahead of time.
The Mac has enough unified memory to serve the 35B 8-bit teacher with MLX. The remote NVIDIA GPU has much faster CUDA kernels for LoRA training, but not enough VRAM to comfortably host both student training and the 35B teacher. Splitting them let the GPU train the 0.8B student while the Mac only returned teacher logprobs. The GPU sent raw token IDs, not text, so the network payload was small. No model weights or gradients crossed the network.
This split added its own failure modes. MLX needed to run from a normal macOS Terminal session so it could see the Metal device. When I launched it from a headless or sandboxed process, MLX reported No Metal device available and the server could not start. The remote GPU needed a stable reverse tunnel back to the Mac teacher server, and that tunnel had to stay up for the entire training run. If the tunnel dropped mid-step, the trainer would hang waiting for a teacher response.
The trainer also needed to prove that it was actually doing on-policy probability distillation, not silently falling back to a custom path. The best run's trainer proof recorded several things that together confirm the method contract:
trainer: trl.experimental.distillation.DistillationTrainer
lmbda: 1.0
beta: 0.0
loss_top_k: 20
use_teacher_server: true
teacher model: mlx-community/Qwen3.5-35B-A3B-8bit
on_policy_loss present in log history: yes
completions/on_policy_mean_length present: yes
teacher_request_delta during training: 80 real requests
tokenizer vocab_sha256 match: yes
trainer_or_loss_modified: false
The on_policy_loss and completions/on_policy_mean_length metrics in the trainer log history are the key evidence. They only appear when TRL is sampling completions from the current student. The teacher_request_delta of 80 means the teacher server received 80 real scoring requests during the 10-step run, which matches the expected number of microbatches. The tokenizer check confirms that the student and teacher share the same vocabulary, so the token IDs the student samples are the same IDs the teacher scores.
The runtime also shaped the experiments. The best 10-step run took about 571 seconds of training time. That sounds short, and the optimizer steps themselves are fast. But each reportable experiment also needed a full 220-task harness eval, and those evals are much slower than the optimizer steps because each eval task is a multi-turn agent episode with SQLite execution. This is why I kept the ladder small: first prove the path with a smoke run, then test the update budget and divergence shape, then stop once the method's behavior was clear. The cost of a full sweep is dominated by eval, not by training.
What I Learned
The clean answer is that larger-teacher probability distillation helped, but only slightly. Starting from a strong 67/220 hard-token SFT parent, the best same-family 35B teacher scoring run reached 69/220. That is a real result, but it is not strong enough to claim that OPD solved the remaining bottleneck.
The more useful lesson is about what the method can and cannot teach. Hard-token correction can show the student a better action. Probability distillation scores the student's own action. If the student never samples the right SQL strategy in a given state, token-level teacher probabilities may not be enough to invent that missing behavior. It can make the local action distribution more teacher-like, but the final task still depends on exploration, schema understanding, SQL judgment, and not overwriting fragile skills.
The forward-KL-versus-reverse-KL distinction was the clearest method-level signal. Forward-KL top-20, which asks the student to cover the teacher's full top-20 distribution, produced the only gain. Reverse-KL top-1, which sharpens the student toward the teacher's single pick, consistently underperformed. For a small agent that already has fragile but useful behavior, spreading mass toward alternatives was safer than collapsing onto one token.
I also learned that more OPD is not automatically better. Ten steps beat five steps, but twelve and fifteen steps got worse. A second fresh-state round also got worse. The method is sensitive enough that I would treat it as a short, carefully evaluated nudge, not a long training recipe to run blindly. The compounding-drift risk is real: each step trains on samples from the current policy, so a few bad steps can pull the policy into a worse region and stay there.
Conclusion
This post answered the narrow question I cared about: can a larger same-family logprob teacher improve the best 0.8B SQL agent by scoring the student's own sampled tokens? Yes, but modestly. The best run moved from 67/220 to 69/220 on the unchanged 220-task eval.
The result is valuable because it is clean. Same student, same harness, same eval, no teacher corrections, no RL, no adapter mixing, and a standard TRL probability-distillation trainer. My takeaway is that larger-teacher OPD is useful as a precise post-training tool, but for this SQL agent it is not enough by itself. The remaining gains probably need methods that add missing successful behavior while preserving what the agent already knows. Probability distillation can clean up the action distribution the student already has, but it cannot teach the student a strategy it never samples.
That question sets up the next post in the series: instead of a larger teacher, can the 0.8B agent improve by teaching itself, scoring its own sampled tokens with on-policy probability feedback from its own checkpoints?