SOPD: Step-Level On-Policy Distillation Between SFT and OPD
Reading notes on:
Supervised fine-tuning and on-policy distillation fail in opposite ways.
- SFT gives the student coherent long-range targets, but all those targets live on the teacher’s clean trajectory distribution.
- OPD trains on the student’s own states, but its token-by-token corrections are often too fragmented to form a usable recovery path once the student has already drifted off course.
SOPD sits between those extremes. It keeps the student’s visited prefixes, but asks the teacher to generate a full step of correction at each prefix. That single change is enough to make the target locally coherent without giving up on-policy state coverage.
1. The Original Trade-Off
Let $x$ be a prompt, $\pi_\theta$ the student, and $\pi_T$ the teacher.
Sequence-level SFT uses teacher forcing on a complete teacher trajectory $y^T$:
\[\mathcal{L}_{SFT} = \mathbb{E}_{(x, y^T) \sim \mathcal{D}_T} \left[ -\frac{1}{L} \sum_{t=1}^{L} \log \pi_\theta(y_t^T \mid x, y_{<t}^T) \right].\]This is coherent, but off-policy. At inference time the student will visit prefixes that do not exist in $\mathcal{D}_T$.
Standard OPD fixes the state-distribution mismatch by sampling student rollouts $y^S \sim \pi_\theta(\cdot \mid x)$ and matching the teacher on those prefixes:
\[\mathcal{L}_{OPD} = \mathbb{E}_{x, y^S \sim \pi_\theta} \left[ \frac{1}{|y^S|} \sum_t \mathbb{D}_{KL} \left( \pi_\theta(\cdot \mid h_t^S) \parallel \pi_T(\cdot \mid h_t^S) \right) \right].\]The trouble is that each teacher target is only local to one token. When the student drifts into a bad region, the teacher is forced to answer inside that bad prefix. The position-wise signals often do not combine into a globally sensible repair strategy. That failure mode lines up well with the brittleness documented in Revisiting On-Policy Distillation.
2. SOPD’s Core Move: Freeze the Student Prefix, Distill One Whole Step
SOPD first rolls out the student policy $\pi_{\theta^r}$ to obtain a step sequence
\[\tau_S = (s_1, s_2, \dots, s_K), \qquad s_k \sim \pi_{\theta^r}(\cdot \mid h_k),\]where $h_k$ is the prefix before step $k$.
Then, for every recorded prefix, the teacher is asked to generate exactly one step:
\[\tilde{s}_k \sim q_T(\cdot \mid h_k).\]Two details make this more subtle than “teacher forcing, but shorter”:
- The student rollout is frozen first. The teacher cannot alter the visited state distribution after the fact.
- Each teacher query is detached. The teacher sees the student prefix $h_k$, not previous teacher corrections $\tilde{s}_{<k}$.
That is what preserves the on-policy state occupancy while still making each local target coherent.
The loss is step-balanced:
\[\mathcal{L}_{SOPD}(\mathcal{B}) = \frac{1}{|\mathcal{S}_\mathcal{B}|} \sum_{(h_k, \tilde{s}_k) \in \mathcal{S}_\mathcal{B}} \frac{1}{|\tilde{s}_k|} \sum_{j=1}^{|\tilde{s}_k|} -\log \pi_\theta(\tilde{s}_{k,j} \mid h_k, \tilde{s}_{k,<j}).\]The inner average over tokens is critical. Without it, long verbose steps would dominate short but decisive actions.
3. Why Step-Level Is the Right Granularity
SOPD is nice mathematically because it really does interpolate between the two older regimes.
3.1 One-step limit: recover SFT
If the whole response is treated as one step, then $h_1 = x$ and SOPD becomes
\[\mathcal{L}_{K=1} = -\frac{1}{|\tilde{y}|} \sum_{j=1}^{|\tilde{y}|} \log \pi_\theta(\tilde{y}_j \mid x, \tilde{y}_{<j}),\]which is ordinary sequence-level SFT.
3.2 One-token limit: recover token-wise distillation on student states
If every token is its own step, then the expected one-token cross-entropy becomes
\[\mathbb{E}_{\tilde{y}_t \sim q_T} [-\log \pi_\theta(\tilde{y}_t \mid h_t^S)] = \mathcal{H}(q_T(\cdot \mid h_t^S)) + \mathbb{D}_{KL}\!\left( q_T(\cdot \mid h_t^S) \parallel \pi_\theta(\cdot \mid h_t^S) \right).\]Up to the teacher entropy constant, that is forward-KL distribution matching on the student’s visited states. So the one-token limit lands back in the token-wise OPD regime; the step-level version simply chooses a more useful granularity.
That is why SOPD feels less like a brand-new algorithm than like the right point on a continuum already sketched by The Distributional Lens of Post-Training.
4. The Packed Attention Mask Is a Quietly Important Systems Trick
Naively, step-level distillation would duplicate prefixes many times: once per step target. SOPD avoids that by packing the prompt, student steps, and teacher targets into one sequence:
[Prompt] [Student s1] [Student s2] ... [Teacher Target 1] [Teacher Target 2] ...
The attention mask enforces three rules:
- student steps attend to the prompt and prior student steps;
- teacher target $k$ attends to the prompt, the student prefix $(s_1, \dots, s_{k-1})$, and its own earlier teacher tokens;
- teacher target $k$ is blocked from future student steps and from all other teacher targets.
This makes all step-level losses trainable in parallel in a single forward pass. It also makes SOPD friendly to black-box APIs: the teacher only needs to return text for each step, not logits or gradients, and the teacher calls can be issued asynchronously after the rollout ends.
5. The Results Are Better Than “Fixing OPD”
5.1 Agentic tasks: ALFWorld
Distilling a Qwen2.5-7B-RL teacher into a Qwen2.5-3B-Instruct student yields:
| Method | Valid Unseen SR (%) | Valid Unseen Rounds |
|---|---|---|
| Teacher | 76.87 | 13.06 |
| SFT | 25.37 | 24.16 |
| Vanilla OPD | 60.45 | 16.21 |
| SOPD | 82.09 | 11.88 |
That is a 21.64-point jump over vanilla OPD on unseen tasks, while cutting the interaction rounds close to teacher level.
5.2 Mathematical reasoning
On AIME and HMMT-style evaluation, SOPD does even better:
| Method | Avg4 |
|---|---|
| Teacher | 46.0 |
| Student base | 15.4 |
| Vanilla OPD | 47.7 |
| SOPD | 57.7 |
SOPD not only beats token-wise OPD by 10.0 points; it also beats the teacher by 11.7 points. My read is that the student is not simply copying the teacher’s trajectory distribution. It is using coherent local step repairs to learn a more stable reasoning structure than either SFT or token-wise OPD can induce.
5.3 The structure of reasoning changes
The paper measures how responses evolve during training:
- mean student step count grows from roughly 120 to over 500;
- average student step length stays around 26 tokens;
- teacher supervision per step grows from 29 to 45 tokens.
That is a strong sign that SOPD is teaching modular reasoning structure, not merely producing longer outputs.
6. Where SOPD Sits in the Bigger Post-Training Landscape
SOPD solves a different problem from Open-MOPD. Open-MOPD asks how to allocate budget across multiple domains and teachers. SOPD asks what the shape of a good teacher target should be once you are already training on student-visited states.
The answer here is persuasive:
- Sequence targets are too coarse.
- Token targets are too fragmented.
- Natural steps are the sweet spot.
That is the right abstraction for long-horizon reasoning and agents. A student that has already made a mistake does not need one more next-token nudge. It needs a locally coherent recovery move.