PaperlayerView exact revision on arXiv

Source-linked AI summary

OmniOPD: Logit-Free On-Policy Distillation via Speculative Verification

Yuhang Zhou, Lizhu Zhang, Yifan Wu, Mingyi Wang, Bo Peng, Jiayi Liu, Xiangjun Fan, Zhuokai Zhao

arXiv:2606.01476v2cs.LGcs.CL

TL;DR

Standard on-policy distillation requires teacher logits and relies on brittle token-level matching. OmniOPD uses logit-free chunk-level semantic verification with selective uncertainty-based auditing, outperforming white-box OPD by up to +28.64% on mathematical reasoning.

  • Problem

    Standard OPD requires direct teacher-logit access, while token-level matching remains brittle even when logits are available.

  • Method

    OmniOPD estimates semantic chunk-level supervision from sampled text-only teacher rollouts and concentrates auditing at high-uncertainty reasoning forks.

  • Results

    +28.64% on mathematical reasoning over white-box OPD, while OmniOPD also surpasses SFT and self-exploratory GRPO across evaluated benchmarks.

  • Takeaways & Limitations

    OmniOPD provides a logit-free distillation route that enables capable proprietary models to serve as teachers and can outperform standard white-box OPD.

  • Takeaways & Limitations

    Chunk-level invariance helps less on competitive programming, where syntactic rigidity makes token-level matching more directly meaningful.

Abstract

from arXiv · show

On-Policy Distillation (OPD) trains a student model on its own generative trajectories under dense token-level feedback from a stronger teacher, mitigating both the off-policy distribution shift of Supervised Fine-Tuning (SFT) and the sparse credit assignment of Reinforcement Learning (RL). However, standard OPD faces two coupled limitations. First, it requires direct access to the teacher's token-level logits, excluding a broad class of capable proprietary models from serving as teachers. Second, the token-level logit signal itself is brittle, depending on a narrow overlap of plausible next tokens between teacher and student, and prone to amplifying degenerate patterns such as repetition loops. In this paper, we introduce OmniOPD, a novel framework that addresses both limitations through a logit-free, chunk-level supervision signal. OmniOPD replaces deterministic logit matching with Monte Carlo rollouts that approximate the teacher's local preferences through a continuous semantic similarity metric over multi-token chunks, and concentrates this supervision via a peak-entropy scheduler that audits the student only at its high-uncertainty reasoning forks. A Dirichlet-Multinomial Bayesian prior and a base-model KL anchor further bound the variance of discrete sampling and prevent policy collapse across unaudited tokens. Across competitive benchmarks, OmniOPD surpasses the standard OPD approach by up to +28.64% on math, confirming that chunk-level semantic verification extracts a more reliable learning signal than token-level logit matching, whose high information density is offset by significant noise and brittleness. Furthermore, when paired with stronger black-box teachers such as Claude-4.5-Haiku and Gemini-2.5-Flash, OmniOPD achieves an additional +9.54% relative on math over its open-weight teacher counterpart, advancing the student past the performance of self-exploratory RL.

1 Introduction

OmniOPD addresses standard OPD’s teacher-logit access requirement and brittle token-level supervision with logit-free, chunk-level semantic verification. It combines selective Monte Carlo teacher auditing with Bayesian stabilization and KL anchoring, achieving strong gains over SFT and white-box OPD.

  • Limitations: Standard OPD requires exact teacher next-token logits, excluding capable proprietary models, while token-level matching relies on a narrow and fragile teacher–student overlap.The signal can degrade under stylistic differences and become harmful on degenerate prefixes such as repetition loops.
  • OmniOPD: OmniOPD replaces deterministic logit matching with logit-free, multi-token semantic supervision estimated from sampled rollouts and their similarity to the student’s generation.This estimates the teacher’s local preferences from text alone rather than requiring direct probability distributions.
  • Method: OmniOPD concentrates verification at high-uncertainty reasoning forks and stabilizes learning with a Dirichlet-Multinomial prior plus a trust-region KL penalty on unaudited tokens.The components target computational tractability, discrete-estimate stability, and prevention of policy collapse across unsupervised gaps.
  • Results: +28.64% on mathematical reasoning is OmniOPD’s maximum improvement over white-box OPD, and stronger black-box teachers advance the student past self-exploratory GRPO.The introduction also reports that OmniOPD frequently outperforms white-box OPD despite full teacher-logit availability.
  • Results: +45.31% relative on mathematical reasoning and +18.52% on competitive programming are OmniOPD’s gains over SFT.These results are reported across mathematical reasoning and competitive programming benchmarks.

2 Related Work

Related work progresses from offline sequence-level imitation toward fine-grained white-box distillation and on-policy training. Offline methods can miss token-level reasoning dynamics, while OPD evaluates student-generated trajectories to reduce exposure bias and align training with inference distributions.

  • Offline policy distillation of LLMs: Offline knowledge distillation transfers frontier-model reasoning to smaller models through generated Chain-of-Thought trajectories and student SFT.This paradigm includes sequence-level distillation from black-box APIs.
  • Offline policy distillation of LLMs: Sequence-level imitation often fails to capture the teacher’s nuanced, token-by-token reasoning dynamics.
  • Offline policy distillation of LLMs: As capable open-source models emerged, distillation increasingly optimized students against continuous teacher distributions through Forward KL divergence or token-level logit matching.
  • On-policy distillation and its limitations: On-policy distillation addresses offline exposure bias by training on the student’s own generated distribution.
  • On-policy distillation and its limitations: In standard OPD, students generate reasoning trajectories that teachers evaluate and score, bridging training and inference distributions.

3 Methodology

OmniOPD replaces token-level logit matching with chunk-level semantic verification from black-box teacher rollouts, targeting high-uncertainty reasoning forks. Bayesian smoothing and a trust-region KL anchor stabilize sparse feedback and constrain policy drift on unaudited tokens.

  • 3 Methodology: OmniOPD combines a peak-entropy scheduler, Monte Carlo chunk-level teacher critic with semantic similarity, and a trust-region anchor.The scheduler selects high-uncertainty chunks, while the critic estimates teacher preference from rollouts and the anchor prevents drift on unaudited tokens.
  • Chunk-Level Teacher Critic: The teacher critic evaluates student chunks against black-box rollouts using continuous semantic similarity, avoiding tokenizer-dependent exact probability matching.This verifies conceptual milestones across architectures without penalizing stylistic or vocabulary differences.
  • Bayesian Stabilization: The Dirichlet-Multinomial prior bounds estimation error and guarantees non-zero gradients when semantic matches are zero, eliminating supervision collapse.The prior uses the student’s per-chunk normalized probability and stabilizes sparse teacher feedback.
  • Peak-Entropy Scheduler: The scheduler anchors M contiguous C-token chunks to the trajectory’s highest-entropy positions, where the student faces complex decision boundaries and needs correction.Overlapping dynamically formed chunks are merged or resampled before auditing.
  • Trust-Region Anchor: For unaudited tokens, a KL penalty against the frozen initial student policy constrains total-variation drift, preserving stability and coherence between audited chunks.The trust region prevents structural shortcuts, incoherence, or entropy collapse caused by sparse localized supervision.

4 Theoretical Analysis of OmniOPD

OmniOPD’s theoretical analysis links its chunk-level loss, Bayesian estimator, and KL anchor to distinct failure modes: gradient explosion, supervision collapse, and drift on un-audited tokens. It further establishes chunk-level invariance to tokenizer, paraphrase, and stylistic variation, while characterizing a bias-variance tradeoff in rollout-based estimation.

  • Gradient stability: The chunk-level loss bounds each per-chunk gradient because the teacher estimator is a multiplier in [0, 1], avoiding standard OPD’s reverse-KL gradient explosion.Reverse-KL gradients become unbounded when the teacher assigns near-zero probability to a sampled token.
  • Gradient stability: The Bayesian prior prevents supervision collapse by keeping the estimator strictly positive even when no teacher rollout semantically matches the student chunk.The empirical estimator becomes zero when k(c)_sem = 0, eliminating that chunk’s gradient contribution.
  • Bayesian estimation: The prior strength α controls a bias-variance tradeoff: larger α reduces noise but increases bias, whereas smaller α retains high variance under small N.Because the estimator multiplies the chunk gradient, this tradeoff directly affects gradient stability; accuracy and gradient noise improve at O(1/N).
  • Trust-region anchoring: The KL anchor constrains policy drift on un-audited tokens, preventing sparse supervision from inducing entropy collapse, shortened trajectories, or globally incoherent behavior.Pinsker’s inequality connects reduced per-token KL to bounded total-variation drift from the frozen base policy.

5 Experiments

Experiments evaluate OmniOPD across mathematical reasoning and competitive programming with Qwen students, open-weight and proprietary teachers, and teacher-access-matched baselines. OmniOPD improves over offline SFT and standard white-box OPD, while its advantage over GRPO depends on teacher capacity and task structure.

  • Experimental Setup: OmniOPD is evaluated with Qwen3-1.7B and Qwen3-4B students distilled from open-weight or proprietary teachers, including Qwen3-32B, Qwen3-30B-A3B-Instruct, Claude-4.5-Haiku, and Gemini-2.5-Flash.Mathematical reasoning uses DAPO-Math-17K with five evaluation benchmarks, while programming uses PRIME-RL datasets including APPS, CodeContests, TACO, and Codeforces.
  • Mathematical Reasoning: 69.08% average mathematical accuracy with Qwen3-4B and Qwen3-32B exceeds base inference by +15.07% and offline SFT by +5.28%.The corresponding base inference and SFT baselines are 54.01% and 63.80%, respectively.
  • Mathematical Reasoning: 74.92% with Claude-4.5-Haiku and 75.67% with Gemini-2.5-Flash exceed their SFT baselines by +7.40% and +2.16%, respectively.The corresponding SFT baselines are 67.52% for Claude-4.5-Haiku and 73.51% for Gemini-2.5-Flash.
  • Comparison with White-Box OPD: 69.08% with Qwen3-4B and Qwen3-32B exceeds standard white-box OPD’s 64.16%, despite OPD having access to teacher logits.The paper attributes this advantage to OmniOPD’s chunk-level semantic loss being invariant to stylistic surface-form differences that make token-level matching brittle.
  • Teacher Scaling and RL: 60.67% with Qwen3-32B teaching Qwen3-1.7B slightly trails GRPO’s 62.41%, indicating that limited teacher capability margins can favor self-exploration.The paper also reports that chunk-level invariance helps less in programming than in mathematical prose because syntactic matching is more meaningful there.

6 Analysis and Ablation Studies

The analysis validates OmniOPD’s Bayesian estimation and chunk-level design, while identifying chunk size and KL anchoring as especially important for performance. Sensitivity experiments show diminishing returns from additional rollouts and robustness to audited-chunk counts and semantic metrics.

  • Estimation accuracy: The Bayesian estimator consistently reduces MSE and improves correlation with ground-truth trajectory probabilities relative to the frequentist baseline.Ground truth is computed from Qwen3-32B’s internal distributions, with comparisons spanning semantic metrics and Monte Carlo budgets.
  • Estimation accuracy: MSE > 0.24 for Exact Match, while Edit Distance reaches MSE ≈0.148 and Pearson correlation 0.72.Exact Match suffers sparsity, whereas Edit Distance better preserves the ranking of reasoning paths despite underestimating absolute magnitudes.
  • Hyperparameter sensitivity: 71.58% is achieved with chunk size C = 100, while reducing C from 50 to 25 drops performance from 69.08% to 24.48%.The results indicate that larger chunks provide more structural context, whereas 25-token windows may be too narrow for meaningful deductions.
  • Hyperparameter sensitivity: +0.49% is the gain from increasing rollouts from N = 10 to N = 20 at C = 50, with negligible or negative benefit under C = 100.This diminishing-returns pattern matches the estimator’s O(1/N) accuracy improvement rate.
  • Hyperparameter sensitivity: 69.08% to 68.38% is the performance change when audited chunks decrease from M = 10 to M = 5, while M = 20 raises the average to 70.96%.OmniOPD therefore gains from denser supervision but remains resilient to substantial verification-budget reductions.
  • Component ablations: 69.08% to 8.28% is the collapse caused by removing the base-model KL anchor, whereas raw frequentist estimation reduces average accuracy from 69.08% to 68.63%.These ablations support the KL trust region and Bayesian smoothing as safeguards against sparse-verification collapse and limited-sampling error.

7 Conclusion · Appendix · A Prompt Templates

OmniOPD is presented as a logit-free, chunk-level OPD framework that combines scheduled speculative verification, Bayesian smoothing, and KL anchoring to address standard OPD’s access and brittleness limitations. The appendix specifies a shared prompt template for training and evaluation, including final-answer formatting for automated parsing.

  • 7 Conclusion: OmniOPD removes standard OPD’s white-box teacher-access requirement through logit-free, chunk-level supervision.The framework is designed to address both the access requirement and brittleness of token-level supervision.
  • 7 Conclusion: Peak-entropy chunk scheduling concentrates supervision within the OmniOPD framework.This mechanism is unified with Bayesian-smoothed speculative verification and base-model KL anchoring.
  • 7 Conclusion: Bayesian-smoothed speculative verification extracts variance-reduced gradient signals from teacher models.The conclusion identifies this component as part of OmniOPD’s unified framework.
  • 7 Conclusion: Base-model KL anchoring mathematically bounds policy degeneration.KL anchoring is presented as the framework’s mechanism for constraining policy degeneration.
  • A Prompt Templates: The same prompt template is used during training on DAPO-Math-17K and across all evaluation benchmarks.This shared template is intended to ensure consistency between learning and testing distributions.
  • A Prompt Templates: The prompt template includes a directive for final-answer formatting to support reliable automated parsing via regular expressions.The {Question} placeholder is replaced with the specific problem text during inference.

B Viability and Teacher Inference Analysis · C Supplementary Details of OmniOPD

OmniOPD reduces teacher-inference burden by shifting supervision from expensive full-trajectory decoding to parallel prefill plus short, targeted rollout bursts. Its configurations provide a tunable effort–accuracy frontier, from below standard SFT cost to higher-cost maximum accuracy.

  • B Viability and Teacher Inference Analysis: Teacher inference dominates on-policy distillation cost in both self-hosted and hosted deployments, determining GPU-hours, wall-clock training time, and hosted-teacher expense.Modern inference separates parallelizable, compute-bound prefill from sequential, memory-bandwidth-bound decode.
  • B Viability and Teacher Inference Analysis: Decode is roughly an order of magnitude more demanding per token than prefill across typical hardware and batching regimes.The passage attributes this gap to decode’s sequential execution and memory-bandwidth bottleneck.
  • B Viability and Teacher Inference Analysis: Standard offline SFT generates the full reasoning trajectory, whereas OmniOPD uses the student trajectory as context and requests short teacher decode bursts only at audited chunks.DAPO-Math-17K uses 8,000 decode tokens per sample under standard offline SFT.
  • B Viability and Teacher Inference Analysis: OmniOPD’s rollout decode consumption is M × N × C tokens per sample, totaling exactly 10,000 decode tokens for M = 10, N = 20, and C = 50.Prefill additionally conditions on the task prompt and the student’s trajectory prefix at each forking point.
  • B Viability and Teacher Inference Analysis: 1.75× standard SFT is the default OmniOPD cost for M = 10 and C = 50 while retaining dense on-policy correction at each audited chunk.This overhead is contrasted with full-trajectory teacher regeneration required by typical on-policy supervision.
  • B Viability and Teacher Inference Analysis: 0.88x SFT is the sparse M = 5 configuration’s cost, while its baseline accuracy remains 99%: 68.38% vs 69.08%.This makes the sparse configuration strictly more efficient than offline distillation.
  • B Viability and Teacher Inference Analysis: 3.00× SFT is the cost of the highest-accuracy C = 100 configuration, defining a controllable effort–accuracy frontier across training budgets.The framework scales across this range without changing its core mechanism.

C.1 Full Algorithm

OmniOPD optimizes a student policy from on-policy reasoning trajectories by selecting peak-entropy anchor tokens, auditing resulting chunks with black-box teacher rollouts, and applying trust-region-anchored gradient updates. The framework uses a reference policy, student policy, black-box teacher, and hyperparameters governing chunks, rollouts, Bayesian prior strength, KL penalty, and learning rate.

  • Inputs and outputs: The algorithm takes a base reference policy, student policy, and black-box teacher as inputs, then returns an optimized student policy.Its hyperparameters include chunk number M, chunk size C, Monte Carlo rollouts N, prior strength α, KL penalty weight β, and learning rate η.
  • Trajectory generation and scheduling: It generates an on-policy student reasoning trajectory y = (y_1, . . . , y_T) ∼ π_θ(· | x) before scheduling supervision.Vocabulary entropy is computed across trajectory positions, and M peak-entropy anchor tokens are selected.
  • Chunk extraction: The selected anchors define a chunk set C = {c_1, . . . , c_M}, while tokens outside audited chunks are identified as the unaudited set U.This concentrates verification on selected chunks rather than the full trajectory.
  • Speculative verification and Bayesian smoothing: For each chunk, OmniOPD queries the black-box teacher with the student prefix and performs N Monte Carlo rollouts for speculative verification and Bayesian smoothing.The supplied algorithm specifies teacher queries for each chunk but does not expose the subsequent rollout aggregation equations in the provided passage.
  • Trust-region optimization: The method computes a chunk-level policy loss and an unaudited trust-region penalty, combines them as L_OmniOPD(θ) = L_chunk + L_KL, and updates θ by gradient descent.The optimization loop returns the updated student policy π_θ.

C.2 Training Dynamics

OmniOPD training converges rapidly and stably under both open-weight and black-box teachers, while KL divergence plateaus rather than growing unboundedly. This stable optimization accompanies substantial improvement on AIME-2025, with Qwen3-32B-teacher accuracy rising from below 30% to above 50%.

  • Training Dynamics: The chunk-level on-policy loss converges rapidly and stably across Qwen3-32B and Gemini-2.5-Flash teacher setups.With Qwen3-32B, loss falls from approximately 0.33 to 0.24 within the first 100 steps before stabilizing.
  • Training Dynamics: KL divergence initially rises as the student departs from base behavior, then safely plateaus within the theoretical bounds.The trajectory supports the trust-region penalty’s role in constraining policy shifts across unaudited tokens.
  • Training Dynamics: AIME-2025 pass-rate accuracy rises alongside decreasing on-policy loss, increasing from below 30% to above 50% under Qwen3-32B.Intermediate validation results may differ from Table 1 because of maximum token truncation during evaluation.

D Supplementary Results · D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks

OmniOPD’s peak-entropy scheduling concentrates teacher verification on high-uncertainty reasoning forks rather than auditing entire trajectories. Across geometry, combinatorics, and algebra, selected chunks target theorem choice, constraint validation, state tracking, and solution verification.

  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: Peak-entropy scheduling concentrates teacher verification on high-uncertainty “forks in the road” instead of auditing entire trajectories.The examples use 10 selected chunks per problem, with || marking each verification chunk.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: Entropy spikes at strategic transitions where the model must select a specific theorem or tool, such as Law of Cosines versus area formulas.These transitions occur after setup phases, when the model commits to a particular solution route.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: Entropy peaks at algebraic branching points involving absolute-value cases and checks that candidate solutions satisfy the original domain constraints.The analysis contrasts relatively low-entropy calculation with high-entropy logical validity checks.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: Entropy rises during multi-step counting when the model must track previously assigned states, creating a memory bottleneck where teacher correction is most vital.This pattern is described as combinatorial state tracking, exemplified by vertex-color assignments.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: In the geometry case, entropy highlights uncertainty over geometric-theorem memory, enabling teacher correction of potentially hallucinated formulas.The selected forks involve choosing the Law of Cosines and checking Stewart’s theorem against the angle-bisector formula.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: In the hexagon-coloring case, entropy selection captures cyclic-constraint struggles and identifies the wrap-around check between vertex 6 and vertex 1 as the highest-entropy region.The analysis identifies this region as the most common point of failure for smaller models.
  • D.1 Case Studies: Qualitative Analysis of Entropy-Selected Chunks: In the algebra case, selection focuses on verifying case-split logic, including whether x = 1 and x = 3 satisfy the relevant conditions.The selected chunks check non-negativity requirements and whether roots of the transformed equations solve the original equation.
Loading 2606.01476v2…