Source-linked AI summary

Open-MOPD: Diagnosing and Fixing Capability Imbalance in Multi-Teacher On-Policy Distillation

Huan-ang Gao, Haohan Chi, Yong Yan, Shiyuan Feng, Hanlin Wu, Zheng Jiang, Bingxiang He, Wei-Ying Ma, Ya-Qin Zhang, Hao Zhou

arXiv:2608.19098v1cs.LGcs.AIcs.CL

TL;DR

Multi-teacher on-policy distillation lacks a rigorously reproducible account of how specialized capabilities integrate in one student. Open-MOPD addresses this with token-share balancing, gap-aware allocation, and reward refresh, raising recovery relative to RouteRL from 35.6% to 83.4%.

  • Problem

    The optimization dynamics of multi-teacher capability integration remain poorly understood, and the field lacks an open, rigorously reproducible recipe.

  • Method

    Open-MOPD combines token-share balancing, gap-aware dynamic budget allocation, and student reward refresh to correct distinct token-level optimization imbalances.

  • Results

    Recovery relative to RouteRL rises from 35.6% to 83.4% in one student, while the integration gap falls from 3.50 points to 0.31 points.

  • Takeaways & Limitations

    The released recipe and mechanism implementations provide an open baseline for reproducible multi-teacher on-policy distillation research.

  • Takeaways & Limitations

    Reward refresh removes only staleness caused by student-dependent rewards; sampled student states still come from rollout time.

Abstract

from arXiv · show

Multi-teacher on-policy distillation (M-OPD) has emerged as a promising paradigm for consolidating domain-specialized reinforcement learning (RL) experts into a single generalist student via dense, token-level reward supervision. Despite its practical success, the optimization dynamics governing multi-teacher capability integration remain poorly understood, and open, rigorously reproducible recipes are conspicuously lacking. In this work, we establish a controlled M-OPD benchmark on SmolLM3-3B-Base with oracle routing, isolating capability integration from routing ambiguity. Our investigation reveals a pronounced capability integration gap: standard M-OPD captures only 35.6% of the available headroom relative to a domain-routed oracle ensemble, with concise tasks such as instruction following suffering severe degradation and premature stagnation. Crucially, we show that this failure stems not from gradient conflict, but from a severe misallocation of the token-level optimization budget. This pathology is driven by three orthogonal factors: structural sequence-length disparities across domains, dynamic convergence drift due to non-uniform learning rates, and multi-step reward staleness from asynchronous policy updates. To resolve these imbalances, we introduce Open-MOPD, a principled framework incorporating token-share balancing, gap-aware dynamic budget allocation, and student reward refresh. Together, these mechanisms systematically restore cross-domain balance, elevating headroom recovery from 35.6% to 83.4% in a single deployable student. We fully open-source our end-to-end post-training recipe, training trajectories, and evaluation suites on an academically accessible hardware budget.

1 Introduction

The introduction frames multi-teacher on-policy distillation as an underexplored capability-integration problem and diagnoses severe cross-domain imbalance despite oracle routing. Open-MOPD addresses this imbalance through three budget-allocation mechanisms, raising recovery from 35.6% to 83.4% in one student model.

  • Controlled benchmark: Oracle routing on SmolLM3-3B-Base isolates capability integration from routing errors across mathematics, coding, and instruction following.The controlled framework uses mixed-SFT initialization, three domain-expert RL teachers, and ground-truth domain labels.
  • Empirical diagnosis: 28.05 versus 31.55 average score exposes a capability-integration gap, with concise instruction-following tasks 6.16 points below RouteOPD.The instruction-following degradation is 3.3× that observed in mathematics, indicating pronounced cross-domain asymmetry.
  • Mechanistic decomposition: Token-level optimization budget distortion, rather than catastrophic gradient conflict, explains the gap through sequence-length disparities, convergence-rate differences, and reward staleness.Gradient allocation follows gradient token volume rather than prompt frequency, causing concise instruction-following responses to receive a negligible share.
  • Open-MOPD: Open-MOPD combines token-share balancing, gap-aware dynamic budget allocation, and student reward refresh to resolve domain imbalances.Student reward refresh recomputes student log-probabilities before each gradient step while caching teacher states, eliminating sample staleness with negligible overhead.
  • Results and reproducibility: 35.6% to 83.4% recovery demonstrates that Open-MOPD restores most available headroom in a single student model.The recipe, model checkpoints, and evaluation suites are fully open-sourced on a reproducible academic compute budget.

2 Open-MOPD: An Open Recipe for Multi-Teacher Capability Integration

Open-MOPD provides a reproducible three-stage recipe for integrating math, code, and instruction-following experts into one student on SmolLM3-3B-Base. Its on-policy distillation uses student-generated responses, routed domain teachers, dense top-k token rewards, and a controlled evaluation protocol.

  • Recipe setup: The recipe integrates three domain experts—math, code, and instruction-following—into one student, with each labeled sample routed to its corresponding teacher.The student is πθ, while domain teachers are πϕd; domain labels select supervision across the three tasks.
  • On-policy distillation: Student-generated responses define the on-policy distribution, while a single teacher prefill produces routed teacher distributions for dense token-level supervision.This avoids the fixed data distribution and potential distribution shift associated with offline teacher trajectories.
  • Three-stage recipe: The three stages are mixed-domain SFT, separate per-domain RL teacher training, and multi-teacher on-policy distillation from the shared mixed-domain SFT checkpoint.Stage I runs four SFT epochs after balancing domains by response-token count; Stage II trains math, code, and IF teachers on their domain rewards.
  • Training configuration: 16K for math and code and 2K for IF: the recipe sets domain-specific response length limits while computing dense rewards on student top-k tokens.Training prompts are sampled from the union of the three domains, and domain labels route each prompt to its corresponding teacher.
  • On-policy distillation: k = 16 by default: Open-MOPD computes dense rewards over the student’s top-k tokens, with tokens outside that set contributing zero.The teacher–student log-probability difference is aggregated with softmax weights and placed directly in PPO’s advantage slot without a critic.
  • Evaluation protocol: The evaluation averages dataset scores within each domain and then averages the three domain means, using AIME24/AIME25, IFEval/IFBenchtest, and code benchmarks under a shared verifier.The protocol reports math accuracy mean@64 with n = 64 and temperature = 0.6, and instruction-following accuracy mean@1 with n = 1.

3 Diagnosing the Multi-Teacher Integration Gap

Naive M-OPD improves over mixed-domain SFT but remains substantially below routed distillation, with the deficit concentrated in instruction following and worsening during training. The evidence points away from teacher conflict and toward three training-signal imbalances: token budget, reward magnitude, and reward freshness.

  • Integration gap: Naive M-OPD raises πmixsft from 25.67 to 28.05 but remains 3.50 points below RouteOPD at 31.55.The result shows multi-teacher supervision helps, yet combining teacher signals in one student loses capability relative to routed distillation.
  • Uneven domain degradation: IF falls 6.16 points below its RouteOPD reference, versus a 1.89-point math gap, making it the main source of integration loss.Over [200, 300] steps, math closes 64% of its gap, code 40%, and IF only 26%; IF also falls 11% during [100, 200].
  • Teacher conflict: Teacher conflict is unlikely to be the dominant bottleneck: disagreement averages 0.126 nat, stays below 0.27 nat, and exceeds 1 nat for only 0.62% of tokens.Removing or replacing high-conflict tokens reduces total score by 0.52, 0.73, 0.75, and 0.83 points across the tested interventions.
  • Token budget: Token shares are severely length-skewed: math/code receive 49.7%/49.3% of tokens, while IF receives 0.99%, despite prompt shares of 39.8%/39.8%/20.3%.Math and code responses average about 10,500 tokens, versus 409 for IF; IF’s token share remains between 0.44% and 1.65% across training.
  • Reward magnitude: Reward magnitudes create dynamic budget drift: IF’s training share drops from 48.7% to around 9% within 25 steps and ends at 11.4%, while code rises to 63.8%.Early average rewards are 0.019 for math, 0.063 for code, and 0.091 for IF; their largest-to-smallest ratio is 4.9.
  • Reward freshness: Reward staleness grows with inner updates: policy-shift KL rises from 0 at K=1 to 0.059 at K=4 and 0.216 at K=32.Later updates reuse rollout-time student probabilities, making the student-dependent loss inconsistent with the current policy.

4 From Diagnosis to Method: Open-MOPD

Open-MOPD preserves teacher routing and the on-policy distillation objective while reallocating the optimization budget and refreshing student-dependent rewards. Its three mechanisms balance token shares, follow remaining capability gaps, and reduce reward staleness during inner updates.

  • Token-share balancing: Within 25 steps, instruction-following’s token share falls to about 9% while code rises to 63.8%, showing that gap-driven drift persists after token-share flattening.Flattening token share does not lock the effective budget in place.
  • Gap-following allocation: Over the first 75 steps, instruction-following reward magnitude decreases 35.3-fold, while α = 0.5 makes its weight rise from 26.7 to 90.7 and code’s fall from 0.44 to 0.27.These dynamics motivate avoiding positive-exponent inverse normalization, which creates an unstable feedback loop.
  • Gap-following allocation: Gap-following allocation assigns more budget to domains with larger remaining teacher–student gaps, rather than inversely normalizing reward magnitudes.It clips the relative reward factor to [0.05, 20] and preserves the batch’s total loss scale after normalization.
  • Reward refresh: Reward refresh recomputes the student-dependent reward before each of K inner updates, correcting mismatch between pre-update rollout trajectories and the changing student.It reuses the PPO student forward pass, adds no teacher forward, and at K = 1 reduces to the original objective.

5 Ablation Studies

The ablations show that token-share balancing removes severe token-allocation imbalance, while gap-following reallocates budget toward larger teacher–student gaps. Reward refresh preserves throughput, and each mechanism improves the corresponding training behavior or domain score.

  • Token-share balancing: 33.33% weighted token share across domains removes the imbalance that otherwise sends 99% of gradient tokens to math and code and about 1% to IF.Token-share balancing produces a +1.17-point gain, almost entirely from IF, while math and code barely change.
  • Gap-following allocation: 55.4% average code budget and 17.1% IF budget reflect allocation toward domains with larger remaining teacher–student gaps.Baseline middle-window gaps are 0.028 for code, 0.008 for math, and 0.003 for IF; code receives more resources while quickly converging domains lose budget.
  • Gap-following allocation: +0.72 points result when gap-following allocation is added without reward refresh.The mechanism shifts budget dynamically: math rises from 15.1% to 32.6%, IF falls from 34.5% to 17.1%, and code averages 55.4% while peaking at 87.6%.
  • Gap-following allocation: 32.5× IF-gap shrinkage under the inverse formulation causes its weight to rise from 24.4 to 80.9 and training to stop at step 74.The reversed factor favors domains the student already handles well, creating a positive feedback loop; the sign of α is therefore fixed by gap direction.
  • Reward refresh: 27.3 s with refresh versus 27.8 s without refresh shows negligible dense-reward runtime difference at K=4.The full outer step takes 1298 s with refresh and 1313 s without it, so reward refresh adds no measurable extra runtime cost.

6 Related Work

Related work spans fixed-corpus distillation, on-policy distillation, and staleness in reinforcement-learning updates. It also covers multi-teacher OPD and alternative strategies for training or combining domain specialists across data, weight, and module spaces.

  • Distillation on a fixed corpus: Fixed-corpus distillation trains sequence-model students on teacher-generated prefixes, creating exposure bias because inference uses student-generated prefixes.The student matches the teacher’s output distribution using responses generated in advance.
  • On-policy distillation: On-policy distillation trains students on their own rollouts and applies teacher feedback to the prefixes they actually visit.Recent work varies divergences, sampling rules, and token-level signals while studying stability and unreliable feedback on long or drifted prefixes.
  • Staleness under batch reuse: PPO reuses rollout data across minibatch updates, while asynchronous reinforcement learning introduces policy lag when rollouts come from older policies.Prior work studies stale-data tolerance, and recent research examines staleness specifically in on-policy distillation.
  • From Single-Teacher to Multi-Teacher OPD: Multi-teacher OPD trains domain specialists, routes student rollouts to domain teachers, and uses token-level teacher feedback for consolidation.Public examples include Nemotron-Cascade 2, Agents-A1, DeepSeek-V4, and Kimi K3, which use multiple teachers or domain-specific teacher definitions.
  • Integration strategies: Domain integration can occur in data space through mixed or sequential training, or in weight and module space through merging, sparsification, and expert combination.Examples include Qwen3, Nemotron-Cascade, Task Arithmetic, TIES-Merging, DARE, and BTM.

7 Conclusion

The conclusion identifies optimization-budget allocation as the main bottleneck preventing error-free multi-teacher OPD from integrating three experts, and presents Open-MOPD as a corresponding three-mechanism remedy. Together, its mechanisms reduce the integration gap and substantially increase recovery relative to RouteRL.

  • Reproducibility: The work releases the complete recipe and mechanism implementations to support reproducible follow-up.The pipeline is built from scratch and includes mixed-domain SFT, three domain RL teachers, and multi-teacher OPD.
  • Core finding: Optimization-budget allocation is the main bottleneck when error-free routing must integrate three experts into shared parameters.The identified imbalance sources are unequal token counts, unequal reward magnitudes, and outdated rewards during repeated inner updates.
  • Open-MOPD mechanisms: Open-MOPD addresses the three imbalance sources with token-share balancing, gap-following allocation, and reward refresh.Each mechanism corresponds one-to-one with a separate source of imbalance and is validated separately in ablation experiments.
  • Results: 3.50 points to 0.31 points: the three mechanisms reduce the integration gap.The reduction is reported for the combined Open-MOPD system.
  • Results: 35.6% to 83.4%: recovery relative to RouteRL rises under Open-MOPD.The conclusion reports this as the combined recovery-rate improvement.

Appendix · A Constructing the Open-MOPD Recipe

The Open-MOPD recipe first selects a base model through math SFT for stable reasoning trajectories, then extends it with mixed-domain SFT across math, code, and instruction following.

  • A Constructing the Open-MOPD Recipe: The recipe screens base models with math SFT for stable reasoning trajectories.This screening determines the base model used in later recipe construction.
  • A Constructing the Open-MOPD Recipe: The selected base model is then trained with mixed-domain SFT on math, code, and instruction following.The mixed-domain stage broadens the model beyond the initial math-focused screening.

A.1 Base Model Choice

The base-model study selects SmolLM3-3B-Base as a practical compromise: Qwen3-1.7B-Base frequently truncates before final answers, while Qwen2.5-7B-Base is stronger but exceeds the resource boundary for the full recipe and ablations.

  • Candidate-model screening: 69.17–80.42% of Qwen3-1.7B-Base samples exhaust the 31K budget, while only 20.00–31.25% close </think> across five checkpoints.AIME24 remains below 7.1%, and additional epochs do not yield monotone improvement.
  • Candidate-model screening: Qwen3-1.7B-Base leaves many trajectories as reasoning prefixes without final answers, confounding attribution of later RL or OPD performance.The passage attributes this ambiguity to the base model and data mixture not reliably learning trajectory closure.
  • Candidate-model screening: AIME24 rises from 16.25% to 31.25% on Qwen2.5-7B-Base as truncation falls from 42.92% to 12.92%, showing that greater capacity improves trajectory closure.This contradicts the explanation that OpenR1 trajectories are inherently too long for any base model.
  • Final base-model choice: SmolLM3-3B-Base is chosen as the compromise between 1.7B’s unclosed reasoning and 7B’s stronger performance with an end-to-end recipe beyond the resource boundary.The final recipe uses four epochs of mixed-domain SFT with response-token shares of roughly 37.3%, 28.1%, and 34.6%.

A.2 Stage Hyperparameters

Stage Hyperparameters reports the configurations actually run across SFT, domain RL teachers, single-domain OPD, and multi-teacher OPD. The recipe uses shared teacher settings with domain-specific adjustments, filtering for discriminative groups, and distinct multi-teacher balancing mechanisms.

  • Stage configurations: The recipe documents hyperparameters from executed training configurations across mixed-domain SFT, three domain RL teachers, single-domain OPD, and multi-teacher OPD.Unlisted items use verl defaults.
  • Domain RL teachers: The three GRPO teachers share one configuration but vary data, sequence length, rollout group size, and training steps by domain.Instruction-following answers use shorter responses, smaller groups, and more steps, while math and code allow 30,000-token responses for full reasoning chains.
  • Domain RL teachers: All teachers disable KL regularization and filter groups whose samples are uniformly correct or uniformly wrong, resampling up to 8 generation batches.This dynamic sampling setup ensures gradients come from groups that discriminate.
  • OPD configurations: Single-domain and multi-teacher OPD share optimizer and distillation settings, isolating their difference to teacher count and domain mixture.Dense rewards use the student top-k with k=16, and the off-policy depth is K=train batch/mini batch=4.
  • OPD configurations: The multi-teacher OPD configuration additionally enables domain sampling and token-share balancing, gap-following allocation, and reward refresh.Table 8 identifies these switch values as the final recipe’s mechanism settings.
Loading 2608.19098v1…