Beyond Teacher Likelihood: Group-Calibrated On-Policy Distillation for Long-Context Reasoning

AI summary

Overview

Research area: Machine learning, specifically knowledge distillation and post-training of compact large language models for long-context reasoning.

Technical level: Intermediate. The paper assumes familiarity with policy-gradient style objectives, log-probability-based advantages, and group-relative advantage normalization, though its core idea (a correction term built from the gap between two scores) is explainable without that background.

Scope: This paper diagnoses a mismatch between token-level teacher guidance and response-level task verification in on-policy distillation for long-context tasks, and proposes a group-relative residual correction, GC-OPD, to reconcile the two.

What This Paper Is About

On-policy distillation (OPD) trains a student model on its own generated responses, using dense per-token guidance from a stronger teacher. The problem is that a teacher's token-level preference can favor locally plausible responses that omit evidence scattered across a long input or violate global task constraints, while a task-specific verifier judges the response as a whole and may return graded rewards reflecting partial success. The paper names this mismatch "teacher–verifier disagreement," shows it grows with input length in two long-context evidence-aggregation tasks, and proposes Group-Calibrated On-Policy Distillation (GC-OPD) to correct dense teacher guidance with response-level verifier outcomes without discarding either signal.

Key Contributions

1. Diagnosis of teacher–verifier disagreement. Across two distributed-evidence long-context tasks (Multi-Table Extraction and High-Recall Retrieval in GoLongRL), the paper shows that trajectory-level OPD scores become progressively less aligned with verifier rewards as prompts get longer, measured by pairwise disagreement rate and OPD preference gap.

2. The GC-OPD method. GC-OPD separately z-score-normalizes verifier rewards and trajectory-level OPD scores within each rollout group, takes their difference as a signed teacher–verifier disagreement residual, and adds a token-scaled correction to the original OPD advantage. The original dense OPD advantage is retained as the base term.

3. Relative-advantage-based credit assignment (RACA). A bounded, monotonic mapping converts each token's OPD advantage relative to its response mean into a positive credit in (0, 2), which scales the trajectory-level residual across tokens. Setting the residual coefficient to zero recovers vanilla OPD.

4. Empirical validation at two model scales plus ablations. GC-OPD achieves the highest five-benchmark average among the compared methods for both Qwen3-4B and Qwen3-8B, and controlled ablations show the signed residual outperforms both an additional OPD-derived term and direct group-normalized verifier reward addition, while RACA outperforms uniform token allocation.

Main Findings

• Teacher–verifier disagreement grows with input length. For Multi-Table Extraction, the pairwise disagreement rate rises from 40.6% below 8K tokens to 64.0% at 32–64K tokens, and the OPD preference gap declines from +0.35 to −0.37. For High-Recall Retrieval, the values change from 35.2% to 60.2% and from +0.65 to −0.35. The diagnostic covers 751 Multi-Table prompts and 2,908 High-Recall prompts across three prompt-length ranges below 64K tokens, using Qwen3-8B as student and Qwen3-30B-A3B-Thinking-2507 as teacher with eight responses per prompt.

• GC-OPD gives the best five-benchmark average at both scales. Post-training raises the five-benchmark average of the official Qwen3-4B checkpoint from 29.08 to 40.47 and of the official Qwen3-8B checkpoint from 35.12 to 44.65. Vanilla OPD reaches 39.31 and 43.56 under the same setup.

• Shared-setup implementations of other verifier-aware methods do not match the aggregate. For Qwen3-4B, ExOPD reaches 38.22, Uni-OPD 38.53, PowerOPD 38.88, and FiRe-OPD 39.50 against GC-OPD's 40.47. For Qwen3-8B, PowerOPD reaches 41.53, Uni-OPD 43.41, ExOPD 43.49, and FiRe-OPD 44.01 against GC-OPD's 44.65. The paper notes these methods often lead on individual benchmarks but do not maintain the strongest aggregate performance.

• Gains concentrate on structured reasoning and evidence aggregation. Relative to vanilla OPD, GC-OPD improves DocMath, MRCR, and CorpusQA for both students, with the largest gains on CorpusQA. Frames and LBv1QA show smaller or model-dependent changes.

• The signed residual beats both control signals. On Qwen3-8B, with RACA and β = 0.10 held fixed, an additional OPD-derived term moves the average only from 43.56 to 43.60, direct group-normalized verifier reward addition reaches 44.19, and the signed residual reaches 44.65.

• RACA beats uniform and sign-discarding allocation. With the same signed residual and β = 0.10, Absolute OPD allocation reaches 43.93, uniform allocation 44.28, and RACA 44.65, against the vanilla OPD anchor of 43.56.

• Disagreement persists across task families but varies. The appendix reports pairwise disagreement and top-1 mismatch for four GoLongRL task families containing at least 500 prompts, plus a prompt-macro aggregate over all nine families, showing the phenomenon extends beyond the two tasks analyzed in detail while remaining task dependent.

• The residual coefficient was selected on a held-out set. The first 256 examples of the ordered GoLongRL shards were reserved, and applying the same 32K-token limit left 231 validation examples with no overlap with the 9,527 training examples. Both the step-100 reward and the mean over steps 60, 70, 80, 90, and 100 select β = 0.10 for both Qwen3-4B and Qwen3-8B. The holdout contains only the High-Recall Retrieval task family and is used exclusively for coefficient selection.

Methodology in Plain English

The researchers start from a standard distillation setup: a student model generates several responses to the same long prompt, and a stronger teacher provides a per-token preference signal by comparing its own token probabilities to the student's. Averaging that per-token signal over a response gives one number summarizing how much the teacher likes that response.

Separately, a task-specific verifier scores each response for whether the task was actually completed, sometimes with a graded score for partial success.

To find out whether these two signals agree, the researchers take a fixed set of generated responses from two long-context aggregation tasks, sort responses within each group by verifier reward, and check whether the teacher-based score puts them in the same order. They measure how often the orderings conflict and how large the teacher's preference is along the verifier's ordering, then break the results down by prompt length.

The fix, GC-OPD, works as follows. Within each group of responses to the same prompt, both the verifier rewards and the teacher-based scores are rescaled to have zero mean and unit variance, so they become comparable despite different raw ranges. Their difference is a signed residual: positive when the verifier likes a response more than the teacher's score does, negative when the reverse. Crucially, the method uses this difference rather than the verifier reward directly, so that calibration vanishes when the two assessments already agree and grows with their disagreement. A simple argument shows that when the verifier prefers response i and the teacher prefers response j, the residual assigns the larger value to response i.

A single residual per response is not enough, because the correction has to be spread across the response's tokens. RACA measures each token's teacher advantage relative to the mean advantage of its own response, maps that through a bounded monotonic function to a positive credit between 0 and 2, and multiplies the residual by that credit. Tokens the teacher favors more get a larger share of the correction; because the credit is always positive it changes the size but not the direction of the residual. The corrected advantage is the original OPD advantage plus the scaled residual, and it is substituted into the existing clipped policy objective unchanged.

Implementation details include numerical guards: if either group-level signal has negligible variation, the residual is set to zero for that group and the update reduces to vanilla OPD; token credits fall back to unit credit when a response has fewer than two valid tokens or negligible token-level variation. The final corrected advantage is clipped before the policy update; the paper states that the threshold and clipping values are reported in Table 6, which is not included in the content provided.

The experimental setup uses Qwen3-4B and Qwen3-8B as students and Qwen3-30B-A3B-Thinking-2507 as teacher. The training set is a GoLongRL subset of 9,527 prompts of at most 32K tokens across nine task families; three families use binary rewards and six use graded rewards, and the two largest families (precise long-range retrieval with 4,693 samples and evidence-grounded reasoning with 3,204 samples) account for 82.9% of the training set. Training uses 100 steps, a batch of 32 prompts, eight responses per prompt, a 32,768-token maximum prompt length, and a 10,240-token response cap. Evaluation uses DocMath, Frames, MRCR, CorpusQA, and LBv1QA, with at most 120,000 input tokens and at most 8,192 generated tokens within a 131,072-token serving context, and a YaRN scaling factor of 4.

Why This Matters

Impact on research. The paper reframes the role of verifier feedback in distillation. Rather than routing objectives, conditioning the teacher on outcomes, gating tokens, or weighting trajectories, it treats the verifier as a reference point against which the teacher's trajectory-level assessment is calibrated. The residual formulation makes the correction vanish when teacher and verifier already agree, which is a different design principle from simply adding outcome reward. The diagnostic metrics (pairwise disagreement rate and OPD preference gap) also give the field a way to measure teacher–verifier misalignment as a function of input length rather than only reporting end-task scores.

Real-world applications:

• Long-document question answering and corpus aggregation, where answers depend on evidence spread across a 100K-token-plus input and a locally fluent response can omit required evidence.
• Numerical reasoning over structured documents such as financial filings, corresponding to the DocMath evaluation.
• Multi-round co-reference and multi-hop synthesis over long conversations or document sets, corresponding to MRCR and Frames.
• High-recall retrieval and multi-table extraction pipelines, where graded metrics such as set F1, IoU, NDCG, and ROUGE-L rather than binary correctness are the natural success criteria.

Industry relevance. Compressing strong teacher behavior into 4B and 8B parameter students is directly relevant to cost and latency budgets for serving long-context models. The method requires no additional teacher or student forward pass beyond the existing pipeline: given the OPD advantages and verifier rewards already produced, it adds only group-level aggregation, normalization, and elementwise token transformations. The paper releases code at https://github.com/SolereZhang/GC-OPD.

Future Directions

• Understanding why disagreement grows with length. The diagnostic shows the trend in two tasks and task-dependent prevalence in four families, but the paper explicitly notes the fixed-response analysis characterizes signal disagreement rather than causal effects of task type or context length. Isolating the mechanism remains open.

• Removing the need for per-setup coefficient tuning. The residual coefficient β = 0.10 was selected on a 231-example holdout drawn from a single task family (High-Recall Retrieval). Whether that value transfers to other data mixtures, longer contexts, or larger students is not established.

• Ablating the numerical guards and clipping. The paper describes thresholds for negligible group-level and token-level variation plus a final advantage-clipping bound, but the values themselves are deferred to Table 6, which is not part of the provided content. Their sensitivity is not reported.

• Extending to settings without a strong verifier or with binary-only feedback. The design retains graded within-group spacing when rewards are graded, but with binary rewards the residual carries less outcome information. The paper also states that complete optimization and implementation details are provided in the supplementary material, which is not included in the content provided.

Target Audience

Researchers and engineers working on knowledge distillation, reinforcement-learning-style post-training, and long-context language models. It is most useful to readers who already understand on-policy distillation and group-relative advantage estimation and want a concrete, low-overhead way to incorporate task verification into dense token-level supervision. Practitioners training compact long-context models with an existing task verifier will find the method directly applicable; readers new to policy-gradient methods will need to work through Section 3 first.

Authors’ abstract

On-policy distillation (OPD) trains a student on its own responses using dense token-level guidance from a stronger teacher. In long-context tasks, however, token-level teacher support can favor locally plausible responses that omit evidence distributed across the input or violate global task constraints. Task-specific verifiers, in contrast, evaluate task completion at the response level and may return graded rewards that reflect partial success. We diagnose this mismatch on fixed responses from two representative long-context evidence-aggregation tasks. Across longer input ranges, trajectory-level OPD scores become progressively less aligned with verifier rewards, indicating teacher-verifier disagreement. Motivated by this observation, we introduce Group-Calibrated On-Policy Distillation (GC-OPD). GC-OPD separately normalizes verifier rewards and trajectory-level OPD scores within each rollout group and uses their difference as a signed teacher-verifier disagreement residual. Relative-advantage-based credit assignment (RACA) distributes this trajectory-level residual across tokens according to their relative OPD advantages while preserving the original OPD signal. Across five long-context benchmarks, post-training with GC-OPD raises the five-benchmark averages of the official Qwen3-4B and Qwen3-8B checkpoints from 29.08 to 40.47 and from 35.12 to 44.65, respectively. Vanilla OPD reaches 39.31 and 43.56 under the same setup. Controlled ablations show that the signed residual is more effective than either an additional OPD-derived term or direct group-normalized verifier reward addition, while RACA further improves over uniform token allocation. Together, these results demonstrate that group-relative residual calibration can incorporate verifier outcomes without discarding dense token-level guidance. Code is available at https://github.com/SolereZhang/GC-OPD.