Is Token Weighting in OPD Really Reasonable? Maybe Tokens Are the Key, and Uniform Weighting Is Wrong from the Start!

Original Link: https://yannx1e.github.io/IW-OPD/
Paper Authors: Yan Xie, Sijie Zhu, Tiansheng Wen, Bo Chen, Yifei Wan

Abstract

On-Policy Distillation (OPD) has gained attention because it combines two useful elements: the student model trains on its own rollouts, while the teacher model provides dense token-level feedback on those rollouts. This also makes OPD a natural alternative to sparse-reward reinforcement learning in the post-training of language models.

In this paper, we study a less verbose problem: once dense teacher feedback is obtained, which positions of the response are actually useful? The effectiveness of OPD depends not only on how many student rollouts we collect, but also on whether the teacher signal can continue to guide the student toward better reasoning positions.

We find a strong positional effect. Tokens early in the student rollout tend to provide more useful supervision signals than later tokens. As the student continues to generate, its tail may drift away from the teacher’s high-probability reasoning region. The teacher is still queried, but now generating it requires scoring tokens under what may not be its own neighborhood. Therefore, dense supervision does not automatically equate to equally useful supervision.

We call this phenomenon position bias in OPD. To further diagnose this issue, we conduct controlled studies and discover two interesting phenomena:

• Early token supervision drives OPD performance;
• The teacher-student gap largely persists after OPD training.

To understand this phenomenon, we re-derive OPD from a constrained optimization perspective. We then derive the optimal policy under this constraint and use this optimal solution to explain why neighborhoods compatible with the teacher should receive more weight. IW-OPD is designed based on this principle: retain the dense token-level OPD signal, but strengthen the emphasis on neighborhoods that still appear compatible with the teacher.

1. Where Is the Supervision Signal Spent?

Let denote the student model and the teacher model. For a prompt and a response sampled by the student, standard OPD minimizes a reverse KL-style objective over the distribution of the student’s own rollouts:

In on-policy fine-tuning implementations, this yields a token-level advantage:

OPD applies token-level loss to each token of the sampled response; a natural implementation averages these losses uniformly. This is simple, but assumes that teacher signals are equally useful at every position. Before introducing a weighting rule, we ask a more fundamental importance question: If we train on different parts of the same student response, would the effect be the same?

2. Position Bias

We start with a simple operation: keeping the model, data, and training settings unchanged, we train OPD variants using tokens from different parts of the response.

• Prefix-30: Apply OPD only to the first 30% of tokens in the response
• Suffix-30: Apply OPD only to the last 30% of tokens in the response
• Standard OPD: Use all valid response tokens

This experiment is deliberately coarse; the goal is to test whether position itself has a tangible impact.

Same quantity. In this controlled experiment, Prefix-30 matches or exceeds the full OPD, while Suffix-30 learns poorly. This shows that the quantity of dense supervision is not everything—the location of supervision matters equally.

This result reveals a position-dependent effect in OPD. Early and late tokens are not interchangeable: they are sampled under different conditional distributions, and training on them can lead to different outcomes. We refer to this phenomenon as position bias.

This result is not a simple gem that “early tokens are always better.” A better explanation: early tokens are more likely to be scored under a context where the teacher distribution still has meaningful guidance for improving the student’s reasoning path.

Late tokens are conditioned on the student’s entire generated history. If the student has already left the region compatible with the teacher, the suffix may be locally predictable, but the correction has less global impact on earlier decisions.

This is loosely connected with recent analyses of reasoning updates, which find disproportionate influence of certain tokens or positions, such as work on high-entropy branch tokens, critical token training, and separating reasoning tokens from boilerplate tokens.

The next question is: why does this positional effect occur? The following diagnostics probe whether late positions are simply “later” or because they are conditioned on more incompatible neighborhoods.

3. Diagnosing the Bias

Diagnosis 1: Can a Strong Teacher Recover from the Student?

We condition both models on two neighborhoods generated by the student and measure the probability of eventually reaching the correct answer. This does not directly train the model; rather, it probes whether that neighborhood is still a useful state for teacher guidance.

Teacher-student gap. The teacher starts with a higher correction rate, but its recovery ability decreases as the student’s generation grows. The student’s tail may drag it downstream, making it harder even for the teacher to reach the region containing the correct answer.

This figure is central to our explanation. OPD queries the teacher on the student’s rollout, not on the teacher’s own ideal rollout. In the early part of the response, the teacher can usually still identify a path to the correct solution.

Later, after the student has made enough local choices, the same teacher is asked to continue from a context that may encode a different plan. The supervision is still dense and updated, but it is less relevant to the corrections we wish to make.

Diagnosis 2: Can Teacher-Student KL Self-Repair During Training?

If the teacher-student KL divergence could be reduced to a sufficiently low value, OPD might gradually fix the neighborhood problem: as training progresses, late student neighborhoods become more compatible with the teacher. Therefore, this diagnosis tests whether OPD can repair this mismatch on its own.

If OPD simply moves the student all the way to the teacher distribution, we would expect KL to shrink continuously toward zero. However, the average KL flattens out after an initial decrease.

KL reduction is local. Token-level KL decreases early in training but leaves a non-zero residual. In our diagnostics, this reduction is limited even after extensive training.

This plateau suggests that OPD is better viewed as a local update strategy, not unconstrained teacher matching. This is important for position bias: when updates can only move the student a limited distance, different token positions can have different leverage.

Diagnosis 3: Does the Residual Mismatch Concentrate in the First Part of the Response?

Not really. Token-level KL remains non-zero across a wide range of relative positions even after training.

Residual mismatch across positions. OPD improves alignment with the teacher, but it does not eliminate the teacher-student mismatch throughout the entire rollout. Therefore, uniform weighting continues to average over positions with different rollout qualities.

Diagnosis 4: Visualization of Neighborhood Likelihood

As the student’s rollout unfolds, the likelihood under the student model naturally remains relatively high (by construction of the sampling procedure). But under the teacher model, the same likelihood drops faster. This is a finer-grained version of the above result.

Another party in log-likelihood space. The teacher-student likelihood ratio tends to shrink as the student’s generation grows. This motivates using neighborhood space (rather than token position) as the unit of reweighting.

Taken together, these diagnostics illustrate why position bias is important for rollout quality. We are not simply dealing with noisy supervision at random positions, but a reconditioning mismatch: late tokens are more likely conditioned on less compatible neighborhoods, and uniform OPD continues to assign the same token weight to them.

4. Limited Altitude

These diagnostics motivate a limited-horizon perspective, which we use as an analytical tool to explain the observed position bias. We do not ask: if the student can move freely, what distribution would match the teacher? Instead, we ask: within a small KL ball around the current student, what distribution is closest to the teacher?

Here is the effective local update layout. This perspective is related to trust region and proximal policy optimization ideas, such as TRPO and PPO: a single update can improve the policy, but cannot replace the student distribution with the teacher distribution.

Under the non-vacuous condition that the teacher lies outside the student’s local KL ball, under common support, the constrained optimal solution has a closed form:

The optimal solution under constraint provides an explanation for the observed position bias. The optimal local target is not the teacher itself, but the student distribution tilted toward the teacher. This tilt is the teacher-student likelihood ratio.

This does not mean suffix tokens are useless. It means that uniform token weighting is not the allocation rule implied by the limited budget. As the student improves and the other branch of the tree becomes compatible with the teacher, late tokens can once again receive useful supervision.

5. From Projection to IW-OPD

The local projection identifies the teacher’s tilt distribution. But the data we actually observe during OPD is sampled from the current student. This mismatch is the position where priority allocation enters.

Similar distribution-correction ideas appear in policy optimization and preference optimization works, such as GSPO and DPO, though our usage here is specifically for the neighborhood-level allocation rule in OPD.

OPD can only optimize a local region around the current student policy. Under the constrained optimization perspective, the best policy that OPD can achieve is one that uses the likelihood ratio (the simple teacher-student gap) to reweight the base student policy.

From our earlier analysis, only tokens with a relatively similar likelihood ratio provide a meaningful learning signal for OPD, because other tokens are sampled with very low probability. Based on this, we use the projected teacher-aligned policy as the learning target, emphasizing samples with high probability and compatibility with the teacher:

Since the projection lies on the boundary of the local KL ball, the positive KL term decomposes into a fixed local penalty plus an expectation term. This is a three-step derivation that transforms the limited-horizon perspective into an importance-weighted objective:

The problem arises because during OPD we sample from , not from . The measure change rewrites the expectation with intensity weights. In the main paper, the likelihood ratio is , and the density ratio induced by the projected target is:

For token-level OPD, the relevant causal object is the continuation. Following the notation in the main paper, the corresponding normalized continuation ratio is:

This gives the token-level objective starting from sampled rollouts:

Its semi-gradient has the same policy gradient form as OPD, but with priority weights on the advantage:

This is why the OPD advantage reappears in IW-OPD. The derivation does not introduce a new token loss; it changes the coefficient multiplying the usual OPD policy gradient signal.

Interpretation: standard OPD treats the observed student distribution as the desired local target. IW-OPD asks how to tilt the observed distribution toward the teacher’s distribution, which is the preset local target. Tokens with a high teacher-student likelihood ratio receive more weight; tokens that accumulate teacher-student divergence receive lighter weight.

6. Practical Surrogate Objective

The first ratio objective provides a clear allocation principle, but using it directly in limited online training is unstable. Practical IW-OPD weights follow a four-step modification outlined in the paper: scaling the ratio in log space, modifying sign cancellation, normalizing within a rollout, and interpolating the result with standard OPD. Among these modifications, ablation shows why simple sharpening is not enough.

First, direct ratio weighting is fragile because probabilities multiply across tokens. Even with an exponent, the coefficients across a rollout can become aggressive. The first ablation asks whether the adjusted ratio alone is sufficient.

Visualizing the token-level behavior shows the same problem. It motivates the first transformation below: the log-space view suggests the downstream trend and local rebound that we want our practical surrogate to avoid.

Problem: Does the raw inverse ratio with appropriate exponent work? Large values make weights adversarial and harm training; very small values make weights flat, behavior closer to standard OPD. This suggests the ratio is useful as an allocation principle but fragile as a practical training signal.

Problem: What does the signed token-level weight do? The downstream trend aligns with position bias intuition, but there are local sharp increases, and signed probability gaps can compensate each other. This makes the literal outer ratio unstable as an update coefficient.

I. Log Scaling

Instead of using the product of probability ratios, it is better to work in log space. Let . The log-scaled tilt score is:

The digital-space form is cleaner for scores, but introduces another problem. Cumulated values are signed: positive and negative token gaps can cancel each other out, so a single score might rebound after the teacher has accumulated several mismatches.

II. Modify Positive Advantage Cancellation

Since rollout tokens are sampled from the student, the teacher-student deviation on these tokens is often negative. Sometimes the teacher assigns a higher probability than the student, yielding a positive deviation term. We reduce this cancellation by converting the per-token inconsistency into a non-positive metric:

Problem: Should the mutual difference be signed or unsigned? The unsigned version performs better in this ablation. The result supports the idea that inconsistencies in both directions should count, rather than allowing positive and negative gaps to cancel each other out.

III. Normalization

The signed metric is monotonic across the rollout, but its magnitude can vary widely across different samples. Therefore, we normalize within each rollout, mapping the beginning of the sequence close to 1, and positions that have accumulated larger mismatches close to 0:

IV. Interpolation with OPD

After normalization, later tokens in the sequence may receive very low weights. Instead of removing dense supervision, IW-OPD interpolates the normalized tail share with standard OPD:

In simple terms: tokens that have accumulated more teacher-student divergence receive a smaller multiplier, while those that have diverged less remain close to standard OPD. This design choice is important.

IW-OPD does not remove dense supervision from late tokens; it uses standard OPD as a floor, separates weights, and allocates additional temporary budget to more compatible tokens. In our main experiments, we use , while exactly recovers standard OPD.

Practical Sketch

• Sample prompts and generate student rollouts on-policy
• Evaluate student and teacher log probabilities on the sampled response tokens
• Compute standard OPD advantages from teacher-student log probability gap
• Compute the cumulative unsigned discrepancy and convert to
• Train with a PPO-style objective, replacing with

This method requires no additional teacher evaluations beyond standard OPD. It changes the allocation of already computed signals.

7. Experiments

We tune IW-OPD on the same family and topology. Student models: Qwen3-4B, Qwen3-1.7B, and Qwen3-0.6B. Teacher models include Qwen3-4B-Instruct-2507 (larger overlap setting) and Qwen3-30B-A3B-Instruct-2507 (smaller overlap setting). We also report experiments with Qwen3-235B-A22B to Qwen3-30B-A3B.

Training uses DeepMath problems (difficulty at least 6, about 57K prompts) and Eurus-RL-Code (about 25K prompts). Math evaluated on AIME 2024, AIME 2025, and HMMT 2025 with mean@32 accuracy. Code evaluated on HumanEval+ and MBPP+. Results averaged over three random seeds.

Step-10 Results

Step-10 checkpoints test sample efficiency: under the same training settings and early update budget, does reweighting better utilize the teacher signal? In both teacher settings, IW-OPD corresponds to standard OPD at step 10 across all student sizes.

Table: Early checkpoint at training step 10. Avg is the aggregate over AIME24, AIME25, HMMT25, HumanEval+, and MBPP+.

The largest gains are in the Qwen3-30B-A3B to Qwen3-4B setting, with AIME25 improving by 6.9 points at step 10. We interpret this as evidence that IW-OPD changes learning efficiency, not just the final checkpoint. Results are based on the tested model families and tasks and do not imply universal step reduction.

TeacherStudentOPD10 AIME25IW-OPD10 AIME25Delta
Qwen3-30B-A3B Qwen3-4B 42.4 49.3 +6.9
Qwen3-30B-A3B Qwen3-1.7B 20.2 23.2 +3.0
Qwen3-30B-A3B Qwen3-0.6B 14.1 15.8 +1.7
Qwen3-4B Qwen3-4B 45.7 46.7 +1.0
Qwen3-4B Qwen3-1.7B 24.7 25.9 +1.2
Qwen3-4B Qwen3-0.6B 17.1 19.0 +1.9

Qwen3-30B-A3B Teacher

The 30B-A3B teacher is the less overlapping setting in our main experiments. In this case, students more often leave regions compatible with the teacher, so if the limited budget explanation is correct, additional weighting should be useful.

Table: Results with the Qwen3-30B-A3B-Instruct-2507 teacher. Math results are mean@32 accuracy (%); code results use HumanEval+ and MBPP+.

StudentMethodAIME24AIME25HMMT25HE+MBPP+Avg
Qwen3-4B OPD 55.3 48.0 27.1 77.2 69.1 55.3
Qwen3-4B IW-OPD 57.5 49.7 28.7 78.7 70.9 57.1
Qwen3-1.7B OPD 34.6 28.7 15.5 64.6 53.7 39.4
Qwen3-1.7B IW-OPD 35.5 29.5 16.4 65.2 55.0 40.3
Qwen3-0.6B OPD 11.0 17.8 7.1 29.6 28.7 18.8
Qwen3-0.6B IW-OPD 11.5 19.3 8.0 32.5 31.9 20.2

IW-OPD improves final outcomes for all three students in this setting. Gains are not uniform across scales, consistent with our view that the relevant variable is effective overlap, not model size.

Qwen3-4B Teacher

The 4B teacher setting has larger student-teacher overlap. This is a useful test because if IW-OPD only works when students are far from the teacher, gains might disappear. Gains are smaller in some cases, but remain positive for the reported students.

Table: Results with the Qwen3-4B-Instruct-2507 teacher. The teacher and some students share more scale/family overlap, but IW-OPD still improves the reported tasks.

StudentMethodAIME24AIME25HMMT25HE+MBPP+Avg
Qwen3-4B OPD 56.5 46.3 24.4 76.3 67.8 54.3
Qwen3-4B IW-OPD 58.7 46.7 25.0 77.9 68.2 55.3
Qwen3-1.7B OPD 34.0 26.4 13.7 61.5 53.7 37.9
Qwen3-1.7B IW-OPD 35.2 27.1 15.3 62.8 54.9 39.1
Qwen3-0.6B OPD 11.8 17.1 2.8 29.8 33.3 19.2
Qwen3-0.6B IW-OPD 13.6 19.0 6.1 31.6 35.7 21.2

Qwen3-235B-A22B Teacher

We also evaluate distillation from the stronger Qwen3-235B-A22B teacher to Qwen3-30B-A3B. This setting is not meant to exhaustively characterize very large teachers, but to check whether the same tilt weighting idea remains useful outside the smaller student setting.

Table: Distillation from Qwen3-235B-A22B-Instruct-2507 to Qwen3-30B-A3B. IW-OPD improves all reported benchmarks in this table.

StudentMethodAIME24AIME25HMMT25HE+MBPP+Avg
Qwen3-30B-A3B Base 28.4 23.4 15.2 77.8 69.5 42.9
Qwen3-30B-A3B OPD 69.5 56.7 38.4 82.1 71.3 63.6
Qwen3-30B-A3B IW-OPD 70.8 58.9 40.5 83.5 73.7 65.5

Our claims in these settings are cautious: IW-OPD consistently improves standard OPD in our evaluated environments, consistent with the limited budget interpretation. We do not claim that the same gains must hold for every model family, task, or training scheme.

8. Weighting and Reward Design Are Orthogonal

IW-OPD changes how the already computed OPD signal is distributed across positions; it does not depend on the specific advantage or reward design. For example, ExOPD (Yang et al., 2026) reframes OPD as a reinforcement learning problem with KL constraints, separates the reward term, and improves exploration by scaling this reward with a fixed hyperparameter.

Since IW-OPD provides token-level importance weights, we can combine it with ExOPD by making the reward scale depend on the neighborhood, e.g., modulating with IW-OPD weights. We call this combination IW-ExOPD. The purpose of this experiment is broader than a specific ExOPD variant: it tests whether token-level weighting can complement another advantage or modification method.

Table: Distillation results with Qwen3-30B-A3B-Instruct-2507 as teacher in the smaller student-teacher overlap setting. Math results reported as mean@32 accuracy (%). Methods with subscript evaluated at training step 10. Bold indicates best result within each student group.

StudentMethodAIME24AIME25HMMT25HE+MBPP+Avg
Teacher—74.7 62.8 44.2 86.6 75.1 68.7
Qwen3-4B Base 23.1 21.4 10.0 75.3 64.5 38.9
Qwen3-4B OPD 55.3 48.0 27.1 77.2 69.1 55.3
Qwen3-4B ExOPD 57.9 50.1 31.7 78.9 70.2 57.8
Qwen3-4B IW-ExOPD 59.4 51.7 32.0 80.1 71.0 58.8
Qwen3-1.7B Base 13.4 11.0 6.8 59.6 52.5 28.7
Qwen3-1.7B OPD 34.6 28.7 15.5 64.6 53.7 39.4
Qwen3-1.7B ExOPD 37.6 31.8 16.8 67.2 55.0 41.7
Qwen3-1.7B IW-ExOPD 38.9 33.2 18.3 68.5 57.4 43.2

IW-ExOPD improves ExOPD for both Qwen3-4B and Qwen3-1.7B students. This suggests that token-level weighting is complementary to ExOPD’s novel scaling mechanism, not only bound to standard OPD.

9. Ablation Experiments

Ablations isolate three design choices: whether the tilt selection should adapt per rollout, whether the discrepancy should be signed or unsigned, and whether the weighting term should replace OPD or be allocated on top of it.

Table: AIME25 improvement over OPD for different variants.

VariantAIME25Delta vs. OPD
Standard OPD 43.3 0.0
Amplify fixed prefix 43.7 +0.4
Manual curriculum 44.8 +1.5
Cumulative-share (ours) 48.9 +5.6
Signed accumulated discrepancy 45.9 +2.6
Unsigned accumulated discrepancy 48.9 +5.6
Ideal weight only 42.1 -1.2
Ideal weight with OPD blend 43.9 +0.6
Surrogate weight only 46.2 +2.9
Surrogate with OPD blend 48.9 +5.6

Fixed prefix shows that IW-OPD is not merely a hardcoded preference for early tokens. Hand-designed curricula help only marginally. Stronger results come from adapting the tilt to each rollout’s own variance.

The signed vs. unsigned comparison shows why cancellation matters. If the student and teacher disagree in both directions along the neighborhood, the signed metric could make the neighborhood look less divergent. The unsigned statistic treats each disagreement as evidence of deviation from the shared region.

The final group shows why we keep OPD as a floor. The surrogate weight works best as an extra tilt on top of standard OPD, rather than a replacement for dense supervision.

10. Discussion

Our work relates to OPD methods such as GKD and MiniLLM, and to practical schemes for supervising student rollout states (e.g., Thinking Machines Lab post). It also relates to token-level curriculum training and token-level curriculum learning, which ask which examples or tokens deserve more weight.

We specifically focus on the imprecise position bias distinction within the student’s own long reasoning rollout: even within a single response, early and late tokens may not be equally useful supervision targets.

This creates an imprecise position bias: the teacher’s supervision at rollout positions near the start tends to be more useful than near the end. This bias is not emphasized in OPD.

It may also appear in many sequences involving two autoregressive models, such as on-policy reinforcement learning that samples from the current policy but updates toward a new policy. This problem is particularly salient in OPD because the teacher-student distribution gap is unknown in advance, so we cannot predefine where they will remain compatible.

This also highlights a broader tension in many sequence optimization problems. Theoretical objectives defined at the sequence level, including the constrained projection view used in the paper, are natural. However, directly optimizing sequence-level objectives in autoregressive models tends to be unstable. Sequence-level objectives for long sequences may have lower reliability than possible.

Standard OPD already practices ignoring the gradient’s return part in token-level proxies.

IW-OPD follows the same philosophy. Since it needs to estimate a sequence-level discrepancy, it uses a smoother cumulative discrepancy proxy instead of the exact likelihood ratio. How to stably estimate and optimize more precise sequence-level objectives is an important open problem.

11. Conclusion

The main lesson is that dense supervision in OPD still has an allocation problem. When the student samples its own long rollouts, later tokens may drift toward the teacher distribution. Then uniform token weighting spends part of the limited update budget in contexts where teacher feedback is less actionable.

Our limited budget perspective transforms this into a local projection problem. The near-teacher alignment target reweights student samples by the teacher-student likelihood ratio, which naturally leads to external importance weighting. IW-OPD implements a stable version of this idea using cumulative unsigned discrepancies and separate extra budget multipliers.

The method is deliberately simple: it uses the same teacher log probabilities as standard OPD, only changing token advantage scales. In our experiments, this allocation changes early learning and final performance in the reported teacher-student settings.

https://qingkeai.online/archives/IW-OPD