SPTA: Teaching Teachers How to Guide by Learning from Student Prefixes

Introduction: When the Teacher Can't Keep Up with the Student's Reasoning

In knowledge distillation for large language models, On-Policy Distillation (OPD) has become the mainstream approach to mitigating student model exposure bias. Its core idea is to have the student generate its own trajectories during inference, and then have the teacher provide supervision signals for each token along those trajectories. This sounds reasonable—the student follows its own path, and the teacher corrects along the way. But reality is often more complicated: the prefixes generated by the student may contain reasoning errors or follow patterns completely unfamiliar to the teacher. When the teacher faces these "unfamiliar road conditions," the reliability of its token-level supervision drops significantly. Worse still, the longer the prefix, the more severe this distribution mismatch becomes.

A recent paper proposes Student-Prefix Teacher Adaptation (SPTA) , cleverly shifting the perspective: rather than forcing the teacher to adapt to the student's "wrong paths," it's better to have the teacher learn in advance how to provide effective guidance on these paths. This approach outperforms both standard OPD and task-level RLVR-adapted teachers on math reasoning benchmarks.

Knowledge distillation was originally proposed by Hinton et al., referring to transferring knowledge from a large teacher model to a small student model. The traditional approach has the student imitate the teacher's output distribution on a fixed dataset (off-policy), but the distribution the student encounters during inference differs from the training data, leading to exposure bias—one wrong step leads to compounding errors. On-Policy Distillation has the model receive teacher supervision on its own generated trajectories, thereby alleviating this problem, with supervision typically at the token level: the teacher provides a probability distribution or score for each token generated by the student, and the student learns by minimizing the divergence from the teacher's distribution.

Problem Diagnosis: Why the Teacher Becomes "Maladapted"

OPD's supervision signal depends on the teacher's judgment of "the student's current state." The standard approach has the teacher score based on reasoning patterns it is familiar with, but the student may reach states the teacher has never considered due to insufficient knowledge or strategic deviation. For example, the student may continue deriving from an incorrect intermediate conclusion, and if the teacher forcibly corrects according to its own reasoning, it may instead provide incoherent supervision.

The paper points out that this teacher-side distribution mismatch worsens as prefix length accumulates. Intuitively: with short prefixes, the teacher can still "tolerate" deviations; but long prefixes may deviate entirely from the teacher's reasoning manifold, causing the supervision signal to become noisy. This explains why OPD's gains on complex reasoning tasks are limited—the further the student errs, the less the teacher can help.

Distribution mismatch is a common problem in the teacher-student framework: the teacher provides supervision based on its own reasoning distribution, but the state distribution actually visited by the student deviates from it. In sequence generation, this deviation accumulates with prefix length—a small deviation at each step can cause subsequent states to drift away from the teacher's familiar region, ultimately making the teacher's token-level signals inconsistent with the optimal guidance direction. This is essentially the same as covariate shift in imitation learning: the input distribution during training differs from that during inference, causing policy performance to degrade over long sequences.

SPTA's Core Mechanism: Letting the Teacher "Preview" the Student's States

SPTA's solution consists of two stages and does not change the distillation objective itself.

Stage 1: Collect Student Prefixes of Random Lengths

Before formal OPD begins, the student model freely generates a series of reasoning trajectories, and prefixes of random lengths are truncated. These prefixes are samples of the states the student "might visit," including correct ones, incorrect ones, and various strange intermediate steps.

Stage 2: Train the Teacher with RLVR to Adapt to These Prefixes

For each prefix, it is used as context, and the teacher continues generating subsequent reasoning. Then Reinforcement Learning with Verifiable Rewards (RLVR) is used to train the teacher: if the final answer is correct, a reward is given. The key design is that the teacher does not imitate the student's reasoning process, but learns to "recover" to the correct answer from student-induced states. In other words, the teacher is trained to be an expert who can "pull the reasoning back on track no matter how absurd the starting point."

After adaptation is complete, the teacher then participates in the standard OPD process. At this point, when facing prefixes generated by the student, the teacher has already seen similar states, and the reliability of its token-level supervision is significantly improved.

RLVR (Reinforcement Learning with Verifiable Rewards) refers to reinforcement learning using automatically verifiable reward signals. Unlike RLHF, which relies on human preference models, RLVR's rewards come directly from rules or programs: for example, whether the answer to a math problem matches the standard answer, or whether code passes unit tests. This reward signal is objective, cheap, and scalable, making it particularly suitable for reasoning tasks. In SPTA, after the teacher model continues reasoning from each student prefix, the system only checks whether the final answer is correct and gives a reward, thereby using reinforcement learning to optimize the teacher's policy so that it learns to recover to the correct reasoning path from any prefix state.

Experimental Results: Comprehensive Outperformance on Math Reasoning

The paper evaluates on math reasoning benchmarks. Results show that SPTA yields better average student performance than two baselines: standard OPD, and OPD with a task-level RLVR-adapted teacher (i.e., directly performing RLVR on task data rather than on student prefixes).

This comparison reveals an important insight: the teacher must not only be good at solving problems independently, but also good at providing guidance on the states the student actually visits. Although task-level RLVR improves the teacher's problem-solving ability, it does not solve the distribution mismatch problem; SPTA, by matching the student prefix distribution, makes the teacher's guidance more targeted.

Implications: "Teaching According to Aptitude" in Distillation

SPTA's value lies in redefining the direction of teacher adaptation. The traditional approach is to have the student move toward the teacher (e.g., sequence-level distillation), or to have the teacher generalize to more tasks (e.g., task-level RLVR). SPTA instead proposes that the teacher should proactively adapt to the student's "behavior distribution," even if those behaviors contain errors. This is similar to human teaching, where excellent teachers adjust their explanations based on students' typical mistakes rather than simply demonstrating standard solutions.

For AI programming and AI applications, this idea also has reference value. In code generation or tool invocation scenarios, the student model may produce intermediate states that do not match the teacher's expectations. If the teacher can be trained on these states in advance, distillation may be more stable. In addition, SPTA's two-stage process is relatively lightweight—it only requires an additional round of prefix collection and RLVR training before OPD—making it engineering-feasible.

Of course, the paper has not yet discussed more complex open-domain tasks or strategies for designing the prefix length distribution. But as an improvement targeting teacher-side distribution mismatch, SPTA provides a new optimization dimension for on-policy distillation.

One-line Summary

SPTA improves the reliability of on-policy distillation by having student prefixes and RLVR train the teacher to adapt to student states.

Paragraph Summary

This article introduces an improved method for addressing teacher-side distribution mismatch in on-policy distillation—Student-Prefix Teacher Adaptation (SPTA). Standard OPD has the student generate trajectories and the teacher supervise token by token, but when student prefixes deviate from the teacher's familiar reasoning patterns, the reliability of teacher supervision declines as prefixes grow longer. SPTA has two stages: first, the student freely generates and truncates prefixes of random lengths; then, reinforcement learning with verifiable rewards (RLVR) trains the teacher to continue reasoning from these prefixes and recover the correct answer. The adapted teacher then participates in standard OPD, outperforming both standard OPD and task-level RLVR-adapted teachers on average across math reasoning benchmarks. The article points out that the teacher must not only be good at solving problems independently but also match the state distribution actually visited by the student, and discusses the implications of this idea for scenarios such as code generation.