PaperlayerView exact revision on arXiv

Source-linked AI summary

MetaOPD: Meta-Learned Token Weighting for On-Policy Distillation

Zipeng Wang, Xinpeng Dong, Yuefan Wang, Pingchen Lu, Xian Wei, Kun Kuang, Fei Wu, Zhongxiang Dai, Min Zhang

arXiv:2610.11989v1cs.AI

TL;DR

Uniform OPD overlooks differences in token learning value, and existing proxy-based weighting rules are predefined rather than adapted through student-update feedback. MetaOPD learns token weights through validation loss after a virtual weighted update; across six mathematical-reasoning benchmarks, it achieves the highest Avg@8 and Pass@8 averages at both student scales, with gains over OPD at each scale.

  • Problem

    Uniform OPD overlooks differences in token learning value, while proxy-based weighting uses predefined mappings rather than learning from the effectiveness of student updates.

    Source: §1 · p. 1

  • Method

    MetaOPD jointly learns the student through weighted OPD and a lightweight token-weighting network trained by validation loss after a virtual student update.

    Source: §1 · p. 2

  • Results

    MetaOPD has the highest six-benchmark Avg@8 and Pass@8 averages at both student scales, improving over OPD by 2.25/6.41 percentage points for the 1.7B student and 1.99/5.97 points for the 0.6B student.

    Source: §4.2 · p. 8

  • Takeaways & Limitations

    Local uncertainty or disagreement alone does not determine whether increasing a token’s OPD weight will improve student learning.

    Source: §5 · p. 9

    The work uses existing models and datasets, and improved benchmark performance does not establish safety; downstream use needs appropriate safety evaluation and must comply with underlying licenses and terms.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student on its own generated responses using token-level teacher supervision. However, uniform weighting overlooks differences in token learning value, while existing weighting methods rely on predefined mappings from prediction signals to token weights. These mappings are not learned from the effectiveness of the resulting student updates, limiting their ability to adapt to evolving learning needs. In this paper, we propose MetaOPD, a bilevel optimization framework that jointly learns the student model and a lightweight token-weighting network. The inner objective updates the student through weighted OPD, while the outer objective optimizes the weighting network using validation loss on reference solutions after a virtual student update. Differentiating through this update connects weighting decisions to their effects on post-update performance, allowing the mapping from prediction signals to token weights to evolve alongside the student. Experiments on six mathematical reasoning and three out-of-domain datasets, covering two student scales and seven baselines, demonstrate the effectiveness of MetaOPD, with Avg@8/Pass@8 gains over OPD of 1.99/5.97 percentage points for the 0.6B student and 2.25/6.41 points for the 1.7B student.

1 INTRODUCTION

The introduction argues that fixed proxy rules cannot adapt token weights to changing learning needs, and proposes MetaOPD, which learns weights from student post-update validation performance. Its motivating experiment reports a larger validation-loss reduction for MetaOPD than for the compared weighting strategies.

1 INTRODUCTION

  • Standard OPD uniformly weights tokens despite differences in learning value, while fixed proxy mappings may not adapt as student capabilities and contexts evolve.Existing approaches use signals such as predictive entropy and teacher–student disagreement, but their mappings are predefined.Source: §1 · p. 1
  • MetaOPD produces the largest validation-loss reduction relative to Uniform OPD, while TIP and random weighting perform close to the uniform baseline.EOPD yields only a modest improvement; the comparison uses validation-loss feedback to learn weights.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • MetaOPD jointly optimizes the student through weighted OPD and a lightweight network that learns token weights from post-update validation loss.A virtual student update lets validation loss on reference solutions guide the weighting network.Source: §1 · p. 2

2 RELATED WORK

Related work covers OPD objectives and training improvements, static token-weighting rules, and bilevel methods that adapt training weights using feedback on model updates.

On-policy distillation

  • OPD research improves distillation objectives and optimization, and also selects prompts, improves rollouts, or refines teacher supervision.Token weighting is presented as another way to improve OPD because tokens differ in learning value.Source: §2 · pp. 2–3

Static proxy weighting OPD

  • Static weighting methods use signals such as entropy, teacher–student disagreement, prefix discrepancy, or token position, but retain predefined mappings to supervision allocation.These mappings do not learn from post-update feedback how to adapt to evolving student needs.Source: §2 · p. 3

Bilevel Optimization

  • Prior bilevel methods use outer-loop feedback to learn example or token weights, and MetaOPD applies this update-aware principle to token weighting on student-generated trajectories.The related work frames training-signal value as dependent on the induced update and the reference objective used to assess it.Source: §2 · p. 3

3 METHODOLOGY

The methodology first defines token-weighted OPD, then learns weights through a virtual student update and validation objective before applying an actual student update. It also interprets the token-level meta-gradient and describes its computation.

3.1 TOKEN-WEIGHTED ON-POLICY DISTILLATION

  • Student-generated responses supply visited prefixes and valid completion tokens for a clipped, sampled OPD objective; setting valid-token weights to one recovers uniform sampled OPD.The teacher is frozen, and the student-to-teacher reverse-KL gradient interpretation applies at the rollout policy; away from it, clipping and cached log-probabilities define a local surrogate.Source: §3.1 · pp. 3–4
  • A shared network maps token descriptors to scores, which are centered and normalized to positive weights with mean one per response before weighting the OPD loss.Descriptors include student and teacher log-probabilities, entropy, prediction agreement, position, and ranked top-32 log-probability profiles.Source: §3.1 · pp. 3–4

3.2 BILEVEL OPTIMIZATION OF METAOPD

  • For each minibatch, a virtual weighted-OPD update yields reference-solution validation loss that trains the weighting network; its updated weights are then recomputed for the actual student update.The virtual parameters are discarded, and the actual update starts from the original student parameters and optimizer state.Source: §3.2 · pp. 4–5

3.3 TOKEN-LEVEL INTERPRETATION AND GRADIENT COMPUTATION

  • With other normalized-weight coordinates fixed, a positive token contribution locally lowers validation loss, but centering, shared parameters, and regularization affect the learned weight.This derivative describes a local effect and does not guarantee finite-step or long-term improvement.Source: §3.3 · pp. 6–7
  • The meta-gradient propagates the outer validation gradient through clipping and AdamW, then uses vector–Jacobian and Jacobian–vector products to compute token contributions without materializing a token-by-parameter Jacobian.Holding the computed token contributions fixed yields the complete meta-gradient at the iterate where they were computed.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

MetaOPD achieves the highest six-benchmark Avg@8 and Pass@8 averages at both student scales, with gains over OPD also reported on out-of-domain tasks. The controls show that direct SFT supervision does not reproduce its gains in both metrics, while combining token descriptors outperforms the strongest single-group variants.

4.1 EXPERIMENTAL SETUP

  • The experiments cover six mathematical-reasoning and three out-of-domain datasets, two Qwen3 student scales, and seven baselines under common evaluation settings.Students are Qwen3-0.6B-Base and Qwen3-1.7B-Base, with Qwen3-8B as teacher; evaluation reports Avg@8 and Pass@8 across eight responses per problem.Source: §4.1 · pp. 7–8

4.2 MAIN RESULTS

  • MetaOPD leads six-benchmark Avg@8 and Pass@8 at both scales; versus OPD, gains are 1.99/5.97 points (0.6B) and 2.25/6.41 (1.7B).The 1.7B model also exceeds TIP by 1.13/1.56 points in Avg@8/Pass@8.Source: §4.2 · p. 8
  • On three tasks evaluated without additional training, MetaOPD leads both GPQA-Diamond metrics and MMLU-Pro Pass@1; its AlpacaEval gains over OPD are smaller.On GPQA-Diamond, MetaOPD scores 27.27 Avg@8 and 78.28 Pass@8, versus OPD’s 26.20 and 77.78.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 ABLATIONS AND CONTROLS

  • Across three coefficients, SFT on OPD raises Pass@8 but lowers Avg@8; at γ⋆, adding it to MetaOPD lowers both by 5.87/6.72 points.MetaOPD improves both metrics when SFT is added to OPD, unlike uniform OPD.Source: §4.3 · p. 9, Figure 3 · p. 9
  • The full descriptor outperforms the strongest single-group variants by 0.57 Avg@8 and 3.20 Pass@8 points; entropy leads single-group Avg@8, while log-probability gap leads Pass@8.The feature ablation reports Avg. scores in Table 3.Source: §4.3 · p. 9

5 CONCLUSION

MetaOPD learns token weights using validation loss after a virtual distillation update, and achieves the highest six-benchmark Avg@8 and Pass@8 averages among compared methods at both Qwen3 student scales. Broader model and task coverage, long-horizon learning effects, and meta-update cost remain open directions.

5 CONCLUSION

  • Across two Qwen3 student scales, MetaOPD achieves the highest six-benchmark Avg@8 and Pass@8 averages among compared methods.Its one-step bilevel framework uses validation loss after a virtual distillation update to learn token-level teacher supervision weights.Source: §5 · p. 9
  • Broader models and tasks, links between local update quality and long-horizon learning, and lower meta-update cost remain open directions.Source: §5 · p. 9
  • Benchmark gains do not establish safety; downstream use should include safety evaluation and follow underlying model and dataset licenses and terms.Source: ETHICS STATEMENT · p. 10

A.1 SAMPLED-OPD GRADIENT AND CLIPPING

The appendices specify the sampled-OPD gradient, response-level weight normalization, and one-step virtual-update derivatives, alongside descriptor and implementation details.

A.1 SAMPLED-OPD GRADIENT AND CLIPPING

  • At θ = θold, a behavior-policy action is an unbiased single-sample estimator conditional on the visited prefix; away from that policy, the cached-action objective is a local surrogate.The conditional statement does not differentiate the trajectory distribution.Source: §A.1 · p. 13

A.2 SCORE CENTERING AND RESPONSE NORMALIZATION

  • Response-level normalization couples valid token weights; with ρ = 0.25, normalized weights lie in (0.6, 5/3), and a one-token response has weight one.The score derivative includes the mean-score derivative; detaching that mean changes the meta-gradient.Source: §A.2 · p. 14

A.3 ADAMW VIRTUAL UPDATE

  • The virtual AdamW update reads the current optimizer state, while its differentiated map retains the new second-moment effect and the forward influence of decoupled weight decay.Global gradient clipping uses one norm across all trainable student parameters, not separate norms per microbatch or tensor.Source: §A.3 · pp. 14–15

A.4 ONE-STEP HYPERGRADIENT AND VJP–JVP FORM

  • The one-step hypergradient differentiates through normalized token weights and the specified optimizer map while holding rollout inputs, student parameters, and optimizer state fixed.Its exactness is for that single-step map, including new moments and the clipping coefficient.Source: §A.4 · pp. 15–16

B IMPLEMENTATION AND REPRODUCIBILITY DETAILS

  • The inner rollout and reference-solution streams both draw from the MATH training split, may overlap, and do not use a separate validation pool for meta-updates.Evaluation examples update neither the student nor the weighting network.Source: B IMPLEMENTATION AND REPRODUCIBILITY… · p. 16

B.1 TOKEN DESCRIPTOR

  • The 74-dimensional descriptor includes top-32 overlap features, but top-k token identities are used only to compute match and overlap, not as network inputs.The weighting network has widths 74–256–64–1 and is initialized so valid-token weights start at one.Source: §B.1 · pp. 16–17, Table 4 · p. 16

B.2 TRAINING AND EVALUATION CONFIGURATION

The configuration distinguishes settings shared across experiments from settings that apply only to particular variants, and specifies selected baseline setups.

B.2 TRAINING AND EVALUATION CONFIGURATION

  • Training and evaluation configurations are shared across the main and ablation experiments, while weighting-network, virtual-update, and reference-solution settings apply only to variants that use them.Learned-weight variants use response-centered scores and the listed schedules, except where an ablation explicitly changes a factor.Source: §B.2 · p. 17
  • KD distills on teacher-generated sequences, while GRPO uses answer-correctness rewards with group size eight.Source: §B.2 · p. 17
  • EOPD uses the method-specific settings reported in its paper; the appendix details how FiRe-OPD, REOPOLD, and TIP allocation rules are implemented within the shared sampled-OPD framework.Source: §B.2 · p. 17

B.2.1 TOKEN-ALLOCATION BASELINE IMPLEMENTATIONS

The appendix describes how allocation baselines are implemented within a shared sampled-OPD setup, then specifies reduction, virtual-state, and EOS conventions.

B.2.1 TOKEN-ALLOCATION BASELINE IMPLEMENTATIONS

  • The baseline comparison places FiRe-OPD, REOPOLD, and TIP on a shared sampled-OPD surrogate and controls the underlying objective, optimizer, and schedule to focus comparison on supervision allocation.Selection statistics are detached rollout-time predictions; these baselines use no learned weighting network, meta-update, or reference-solution loss.Source: §B.2.1 · p. 17
  • FiRe-OPD removes the lowest-scoring 20% of responses by mean teacher log-probability and assigns retained-token weights; TIP retains each response’s highest-scoring half, while REOPOLD and TIP use unit weights on retained tokens.FiRe-OPD weights are normalized to mean one within each retained response; the baselines share the surrogate, optimizer, and training schedule.Source: §B.2.1 · pp. 17–19

B.3 GLOBAL TOKEN-MEAN REDUCTION

  • Weighted OPD losses are divided by the global valid-token count across the complete optimizer minibatch, rather than by separately normalized microbatch losses.Reference-solution loss uses its own global valid-token count, which can differ from the rollout count.Source: §B.3 · p. 19

B.4 VIRTUAL STATE AND LEARNING-RATE CONVENTION

  • Training-time virtual steps read the current student AdamW state and scheduled learning rate without committing parameter or optimizer-state changes.The weighting network is updated at each student step throughout all 696 updates.Source: §B.4 · p. 19

B.5 END-OF-SEQUENCE MAPPING

  • The student EOS identity remains the generation stop token, while the teacher-view mapping converts the first valid student EOS occurrence to the teacher EOS identity.Reference-solution targets retain the student EOS token.Source: §B.5 · p. 19

B.6 ABLATION AND RESOURCE DETAILS

The appendix details controls for direct SFT supervision, feature-ablation comparisons, training cost, and the shared protocol for one-step strategy comparisons.

B.6 ABLATION AND RESOURCE DETAILS

  • The SFT coefficient is calibrated once at initialization to approximately match gradient scales, and the resulting γ⋆ is about 2.16.The coefficient applies only to the SFT loss in the student objective; matching gradient norms does not imply matching gradient directions or optimizer updates.Source: §B.6 · p. 20
  • The feature ablation compares four named descriptor subsets with the full input while holding the 74-dimensional architecture and 35,713 parameters fixed.This tests input subsets rather than removing one feature at a time.Source: §B.6 · p. 20, Table 8 · p. 21
  • The SFT controls compare SFT, a uniform-OPD SFT loss-weight sweep, and an SFT-plus-MetaOPD interaction control under shared training conventions.The MetaOPD variants retain the shared virtual-update rule; SFT and the uniform-OPD sweep have no meta-update.Source: §B.6 · p. 20, Table 9 · p. 21
  • MetaOPD’s complete 1.7B training procedure costs 88.5 GPU-hours versus 45.0 for OPD, a 1.97× ratio.The accounting includes all eight allocated GPUs, including paired teacher GPUs, and excludes evaluation and diagnostic probes.Source: §B.6 · p. 21, Table 10 · p. 21

C.1 SHARED SETTING

  • The shared comparison uses four fixed prompt folds and three response-sampling seeds, with two shared reference sets assigned distinct roles.Set A drives the weighting-network meta-update after a virtual weighted OPD update; set B evaluates the actual student update and does not enter weight learning.Source: §C.1 · pp. 21–22, §C.2 · p. 22

C.2 COMPARED STRATEGIES

  • Uniform, MetaOPD, EOPD, TIP, and Random strategies are compared using shared responses, learning rate, and clipping threshold.Random weighting serves as an allocation reference without validation-loss feedback.Source: §C.1 · p. 21, §C.2 · p. 22

C.3 EVALUATION AND AGGREGATION

The motivating comparison favors MetaOPD over the proxy and random weighting strategies in a one-step setting. The diagnostic protocol uses fixed checkpoints, responses, and reference banks to examine derivative agreement and update quality.

C.3 EVALUATION AND AGGREGATION

  • The reported intervals condition on fixed prompts and reference sets; response-sampling seeds are not independent full training runs.The comparison averages paired reductions across 12 panels and bootstraps panels.Source: §C.3 · p. 22
  • MetaOPD achieves the largest observed post-update validation-loss reduction; EOPD improves less, while TIP and Random remain close to Uniform OPD.This supports using post-update performance to learn token weights in this one-step setting, but does not establish the effects of repeatedly applying the strategies over a full training run.Source: §C.3 · p. 22

D.1 DIAGNOSTIC SETTING

  • Checkpoint probes use student-generated responses for candidate updates and reference solutions for validation loss, with 64 responses split into two groups per checkpoint.The probes examine three MetaOPD training checkpoints and do not commit candidate updates to the model.Source: §D.1 · pp. 22–23
  • Figure 4 tracks weight dispersion, median within-response Spearman correlation with c_i,t, and paired Learned−Uniform validation-loss differences across three checkpoints.Negative loss differences favor Learned; intervals are reference-bank bootstrap intervals with checkpoints and rollouts fixed.Source: Figure 4 · p. 23
  • Derivative-consistency analyses compare token-level scores across reference batches using rank correlation, sign agreement, and top-20% Jaccard overlap.The analyses use median summaries over batch pairs and responses; derivative aggregation weights reference banks by valid reference-token counts.Source: §D.1 · p. 23

D.2 LEARNED WEIGHTS AND UPDATE QUALITY

MetaOPD’s learned weights become positively associated with validation-loss derivatives, and Learned lowers validation loss relative to Uniform at all three checkpoints. These one-step comparisons do not show that either allocation improves on the unchanged checkpoint or throughout full training.

D.2 LEARNED WEIGHTS AND UPDATE QUALITY

  • Weight dispersion increases from uniform initialization and then levels off, while the online correlation with c_i,t turns positive and reaches about 0.3 later in the recorded first epoch.The online correlation uses the current training reference batch and is not interchangeable with checkpoint alignment.Source: §D.2 · pp. 23–24
  • Learned weights correlate positively with c◦_i,t at all three training checkpoints; the final correlation is 0.267 [0.246, 0.278], with 95% bootstrap intervals above zero.These correlations do not establish convergence of the weighting rule.Source: §D.2 · p. 24, Table 11 · p. 24

D.3 ONE-STEP UPDATE COMPARISONS

  • Reweighted yields lower validation loss than Learned at each checkpoint, although all three allocations increase loss at the final checkpoint.The local derivative-ranked finite-step ordering does not guarantee descent or improvement over the full training trajectory.Source: §D.3 · pp. 24–25
  • −1.95 × 10^-3: at the diagnostic reference learning rate, Learned’s final paired validation-loss difference versus Uniform is approximately this value.Learned has lower validation loss than Uniform at all three checkpoints, although both allocations increase mean loss relative to the unchanged checkpoint at every stage.Source: Table 12 · p. 25, §D.3 · p. 25
  • All four restricted-input variants show consistent derivative estimates and negative Reweighted–Learned contrasts at student step 696.These comparisons assess local update quality rather than benchmark accuracy.Source: §D.3 · p. 25, Table 13 · p. 25

E OUT-OF-DOMAIN GENERALIZATION

On out-of-domain evaluations, MetaOPD leads on GPQA-Diamond and MMLU-Pro but trails TIP on AlpacaEval preference scores. It also has the highest descriptive average across the five reported metrics.

E OUT-OF-DOMAIN GENERALIZATION

  • The evaluations use Qwen3-1.7B-Base checkpoints without additional training on the three out-of-domain tasks.Distillation methods use Qwen3-8B as the teacher and a common local evaluation protocol for each benchmark.Source: E OUT-OF-DOMAIN GENERALIZATION · p. 26
  • MetaOPD has the highest reported scores on both GPQA-Diamond metrics and MMLU-Pro Pass@1, while its AlpacaEval scores are close to OPD’s and below TIP’s.The observed advantage varies across the evaluated tasks.Source: E OUT-OF-DOMAIN GENERALIZATION · p. 26
  • MetaOPD has the highest descriptive average across the five reported out-of-domain metrics.The unweighted average combines two GPQA-Diamond metrics, MMLU-Pro Pass@1, and two AlpacaEval win-rate metrics; individual metrics remain necessary to interpret transfer.Source: E OUT-OF-DOMAIN GENERALIZATION · p. 26

Related papers

Loading 2610.11989v1…