PaperlayerView exact revision on arXiv

Source-linked AI summary

AsyncOPD: How Stale Can On-Policy Distillation Be?

Wonjun Kang, Kevin Galim, Seunghyuk Oh, Minjun Kang, Sanghyun Park, Donghoon Kim, Minjae Lee, Minseo Kim, Rishabh Tiwari, Yuchen Zeng, Hyung Il Koo, Kangwook Lee

arXiv:2606.24143v1cs.LG

TL;DR

Asynchronous OPD can reduce rollout bottlenecks but introduces stale-policy data, whose effects and mitigation in OPD remain underexplored. This paper studies staleness, develops estimator choices including learner-time reverse-KL recomputation and multi-sample MC, and presents AsyncOPD, which improves throughput 1.6× to 3.8× over strict synchronous training while maintaining comparable accuracy.

  • Problem

    Asynchronous OPD alleviates rollout bottlenecks but introduces stale-policy data, whose effects and suitable mitigation remain underexplored in practical teacher-cache settings.

  • Method

    The paper studies KL-direction effects, evaluates reverse-KL surrogates, analyzes sparse and sampled teacher-cache estimators, and builds AsyncOPD with learner-time recomputation and multi-sample MC.

  • Results

    1.6× to 3.8× throughput improvement over strict synchronous training was achieved with comparable accuracy on Qwen3-Base models.

  • Takeaways & Limitations

    Forward KL is more robust to stale rollouts, while reverse KL benefits from learner-time recomputation and multi-sample MC under finite teacher-score caches.

  • Takeaways & Limitations

    The study covers sparse and Monte Carlo estimators rather than dense asynchronous full-vocabulary KL, and experiments are limited to a single 8-GPU node.

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student on its own rollouts guided by teacher feedback and is becoming increasingly important for large language model (LLM) post-training. Like reinforcement learning (RL), however, OPD faces an on-policy systems bottleneck, as rollouts can dominate training time for reasoning workloads. Asynchronous training pipelines can alleviate this bottleneck by decoupling rollout generation from learner updates, but doing so introduces stale-policy data. While prior work has studied stale data in asynchronous RL, its effects in OPD remain underexplored. We present the first systematic study of staleness in asynchronous OPD, focusing on a practical setting where teacher feedback is implemented through local KL losses and full-vocabulary teacher logits are too expensive to store or transfer, necessitating finite teacher-score caches. We first show that KL direction changes the stale-data problem: teacher-weighted forward KL is more robust to stale rollouts, whereas student-weighted reverse KL is vulnerable. Second, for this vulnerable reverse-KL case, we study whether methods designed to stabilize asynchronous RL can mitigate OPD staleness. In our experiments, they do not improve over a simpler OPD-specific surrogate: recomputing the reverse-KL signal under the current student at learner time. Third, we analyze how finite teacher-score caches create a bias-variance tradeoff for sparse and sampled reverse-KL OPD estimators. This motivates multi-sample Monte Carlo (MC), which preserves MC correctability while reducing one-sample variance. Finally, we present and open-source AsyncOPD, a fully asynchronous OPD training pipeline built from these estimator choices. Experiments show that AsyncOPD improves training throughput by $1.6\times$ to $3.8\times$ over strict synchronous training while reaching comparable accuracy.

1 Introduction

This paper systematically studies stale-policy data in asynchronous on-policy distillation under practical finite teacher-score caches. It shows that KL direction and estimator design determine robustness, motivating AsyncOPD, which improves throughput while maintaining comparable accuracy.

  • Motivation: Asynchronous OPD decouples rollout generation from learner updates, addressing expensive reasoning-task rollouts that leave synchronous learners underutilized.The trade-off is that asynchronous execution introduces policy lag through stale-policy data, which can degrade model quality.
  • KL Direction and Staleness: Forward KL is more robust to stale rollouts, whereas reverse KL is vulnerable because current-student actions may fall outside the cached teacher-scored support.This difference arises because forward KL is teacher-weighted, while reverse KL is student-weighted.
  • Reverse-KL Surrogates: PPO-style clipping, decoupled PPO, and M2PO do not improve reverse-KL OPD over recomputing the token-level advantage under the current student at learner time without clipping.The comparison is motivated by reverse KL's RL-style policy-gradient surrogate, with the teacher-student log-ratio acting as a token-level advantage.
  • Teacher-Cache Estimators: Sparse reverse-KL implementations suffer support mismatch, while one-sample Monte Carlo avoids fixed-support mismatch through importance-correctable sampling but has high variance.Stale student top-k supports may omit actions required by the current top-k objective, and reweighting within the stale support cannot recover missing teacher scores.
  • AsyncOPD: AsyncOPD overlaps student rollout, teacher scoring, and learner updates in a fully asynchronous OPD pipeline.The pipeline instantiates the paper's estimator and staleness findings under the finite teacher-cache constraint.
  • AsyncOPD: 1.6× to 3.8×: AsyncOPD improves training throughput over strict synchronous training while maintaining comparable accuracy on Qwen3-Base models.This result is reported for the fully asynchronous pipeline on Qwen3-Base models.

2 Related Works

Related work spans token- and sequence-level on-policy distillation, asynchronous reinforcement-learning systems, and emerging asynchronous OPD schedulers. Prior asynchronous OPD systems establish feasibility but leave estimator behavior under stale teacher-scored caches open.

  • On-Policy Distillation: On-policy distillation uses student rollouts and dense teacher feedback, with prior work covering token-level KL, sequence-level reverse KL, training dynamics, and token importance.GKD introduced token-level KL; MiniLLM studied sequence-level reverse KL; Li et al. analyzed unstable configurations; TIP characterized importance using entropy and teacher-student divergence.
  • Asynchronous RL: Asynchronous RL decouples rollout generation from learner updates to improve hardware utilization, including overlapped and streaming RLHF pipelines.Synchronous pipelines can leave learner resources idle while waiting for the longest rollout in a batch.
  • Asynchronous OPD: Asynchronous OPD schedulers overlap student rollout, teacher scoring, and learner updates, but estimator behavior under stale teacher-scored caches remains open.VeRL fixes rollout lag to one or two learner steps, while KDFlow decouples teacher inference from learner training and transmits teacher hidden states.

3 Preliminaries: On-Policy Distillation

This section formulates token-level on-policy distillation at fixed prefix states, where forward and reverse KL use teacher- and student-weighted objectives, respectively. It also connects reverse-KL OPD to policy-gradient methods and motivates sparse or sampled estimators as computation–coverage–variance tradeoffs.

  • OPD setup: At each decoding timestep, OPD treats the visited prefix s as the state and the next token a as the action, applying local losses to generated tokens.The teacher and student policies are q(a | s) and pθ(a | s), respectively.
  • Forward-KL OPD: Forward-KL OPD weights the local loss by the teacher policy q(a | s).The objective is evaluated at a fixed prefix over the vocabulary.
  • Reverse-KL OPD: Reverse-KL OPD weights the local loss by the student policy pθ(a | s), with its gradient expressed through student-policy log-probability terms.The reverse-KL gradient is given in the supplied equations as a vocabulary sum involving pθ(a | s), log pθ(a | s), and log q(a | s).
  • Reverse-KL OPD: Viewing reverse-KL OPD as policy-gradient estimation identifies A = log q(a | s) − log p(a | s) as an advantage term and motivates PPO-style surrogates with behavior-policy ratios and clipping.The practical surrogate uses ρθ(a, s), clipped ¯ρθ(a, s), and a frozen behavior-time signal Abeh(a, s).
  • Sparse and sampled implementations: Sparse and sampled OPD implementations trade computation against support coverage and estimator variance instead of evaluating the full-vocabulary objectives.Sparse top-k methods restrict and renormalize distributions on S(s), while MC methods sample actions from a proposal distribution.

4 Forward- and Reverse-KL OPD Under Staleness

KL direction determines how asynchronous OPD responds to rollout staleness: forward KL is more robust, while reverse KL degrades faster as staleness grows. This difference reflects teacher-weighted versus student-weighted action terms under scored-cache restrictions.

  • Scope of analysis: Action-level staleness is the focus because rollout-fixed prefixes cannot be changed by learner-time estimator design.Asynchronous OPD has both prefix-level and action-level staleness, but only action-level staleness is directly addressable by estimator design.
  • Cached-data pipeline: Asynchronous OPD separates rollout, teacher scoring, and learner updates, allowing visited prefixes, cached actions, teacher scores, and update policies to come from different student versions.The pipeline uses cached data, and teacher logits are available only on the teacher-scored action cache; the learner cannot recover missing teacher signals without additional teacher access.
  • KL direction and staleness: Reverse KL starts higher at zero staleness but drops faster as staleness increases, eventually falling below forward KL.Figure 2 compares sparse top-k forward-KL and PPO-style reverse-KL implementations.
  • KL direction and staleness: Forward KL is teacher-weighted and robust to rollout staleness, whereas reverse KL is student-weighted and vulnerable.Under scored-cache restrictions, forward KL avoids converting stale student-sampled actions into a current-student expectation, while reverse KL depends on student-weighted action terms.

5 Reverse-KL: Policy-Gradient Surrogates Under Staleness

For stale reverse-KL OPD, recomputing the current-student advantage at learner time without clipping is the most effective surrogate. Advanced asynchronous RL surrogates do not improve over this OPD/IS choice, while PPO-style clipping remains a strong baseline when paired with the stale advantage.

  • Exact importance-sampling identity: The OPD/IS surrogate uses the current advantage Aθ as a stop-gradient weight with an unclipped old-to-current importance-sampling ratio.Its validity follows from the exact importance-sampling identity, assuming pold has support wherever pθ does.
  • Surrogate ablation: Recomputing Aθ at learner time without clipping is the best-performing reverse-KL surrogate under staleness.This OPD/IS choice is more effective than the alternatives evaluated.
  • Surrogate ablation: Clipping helps Aold by limiting stale, large-ratio updates but hurts Aθ because recomputing it already reduces the high-percentile ratio tail.At staleness 64, recomputing Aθ reduces the p99 ρθ tail, so clipping removes useful signal.
  • Advanced asynchronous RL surrogates: Decoupled PPO and M2PO do not improve over the simpler OPD/IS surrogate that recomputes Aθ without clipping.These methods were designed to improve robustness to stale-policy updates but do not consistently help OPD under staleness.

6 Reverse-KL: Cached Supports Under Staleness

Reverse-KL cached-support estimators behave differently under staleness: sparse top-k cannot recover missing current-support teacher scores, while one-sample MC remains correctable but high-variance. Multi-sample MC reduces that variance while preserving MC correctability and improves performance at large staleness.

  • Sparse top-k: Sparse top-k remains biased under stale rollouts because reweighting cannot recover teacher scores for actions missing from the rollout-time support.The stale support is TopK(p_old(· | s), k), which may omit actions in the current support TopK(p_θ(· | s), k).
  • One-sample MC: One-sample MC is unbiased in expectation under old-to-current importance sampling when the behavior policy covers the current support, but its estimator can have high variance.It caches a sampled action and its rollout log probability, enabling action-level correction at learner time.
  • Multi-sample MC: Multi-sample MC averages IS-corrected gradients from multiple local behavior-policy samples, reducing one-sample variance while preserving the same expectation.Independent samples are drawn at each decoding timestep without generating additional full trajectories.
  • Sparse top-k vs. one-sample MC: IS substantially improves one-sample MC robustness as staleness increases, whereas sparse-top-k reweighting does not; one-sample MC with IS is strongest among the four methods.Sparse-top-k reweighting fails because it cannot recover missing current-support actions.
  • One-sample MC vs. multi-sample MC: m = 4 already gives a clear jump over one-sample MC at large staleness, while m ∈{4, 16, 64} performs similarly.This supports multi-sample MC as the preferred cached-support estimator for asynchronous OPD.

7 AsyncOPD: Fully Asynchronous OPD

AsyncOPD is a fully asynchronous OPD system that overlaps rollout, teacher scoring, and learner updates while streaming scored examples to the learner. Across matched Qwen3-Base experiments, it achieves the highest throughput and overlap, reaching up to 2.7× strict-sync throughput in MC64 and 3.3× in MC1 while maintaining strong final accuracy.

  • System: AsyncOPD overlaps rollout, teacher scoring, and learner updates in a fully asynchronous OPD pipeline.The system follows AReaL and decouples these OPD stages to improve pipeline utilization.
  • Scheduler: The streaming scheduler pauses workers only for weight synchronization, preserves in-flight prefixes, and updates the learner when a scored batch is ready.Unlike k-step-off scheduling, it streams completed examples instead of waiting for complete rollout batches.
  • Experimental setup: Experiments compare strict sync, two-step-off, and AsyncOPD on Qwen3-1.7B, 4B, and 8B students using a Qwen3-30B-A3B-Instruct-2507 teacher for 100 iterations.All runs use the same reverse-KL estimator, with MC64 or MC1, on the same 8-GPU node; AsyncOPD uses τ = 4.
  • Experimental Results: 2.7× strict-sync throughput: AsyncOPD reaches this maximum in MC64 while achieving the best or tied-best final AIME24 Avg@32 accuracy.AsyncOPD also achieves the highest throughput and pipeline overlap in every matched MC64 comparison.
  • Experimental Results: 3.3× strict-sync throughput: AsyncOPD reaches this maximum in MC1 while delivering the highest throughput and overlap for every student.MC1 shows the same overall trend as MC64 across the matched student comparisons.

8 Conclusion … A.2 Monte Carlo Reverse-KL OPD

The paper studies stale data in asynchronous OPD and finds that forward KL is more robust than student-weighted reverse KL, for which learner-time advantage recomputation is the most effective surrogate. It also formalizes sparse and Monte Carlo reverse-KL estimators while noting that dense asynchronous KL remains an implementation challenge.

  • 8 Conclusion: Forward KL remains robust to stale rollouts, whereas student-weighted reverse KL is more vulnerable.The conclusion identifies KL direction as the factor shaping the stale-data problem.
  • 8 Conclusion: In reverse-KL OPD, recomputing the current advantage at learner time without clipping is the most effective policy-gradient surrogate.Advanced asynchronous RL surrogates do not improve over this choice.
  • 8 Conclusion: The study evaluates sparse and Monte Carlo OPD estimators rather than dense full-vocabulary KL in the asynchronous setting.Dense KL avoids cached-support mismatch but is difficult to implement when rollout, teacher scoring, and learner updates are decoupled.
  • A.1 Sparse Top-k Reverse-KL OPD: Sparse top-k reverse KL evaluates the objective on a finite student support instead of summing over the full vocabulary.The dense reverse-KL objective sums over the full vocabulary, while sparse implementation restricts evaluation to finite support.
  • A.1 Sparse Top-k Reverse-KL OPD: For support S, restricted normalizers define the sparse reverse-KL distribution and objective.The construction uses support-restricted normalization for the student and teacher distributions.
  • A.1 Sparse Top-k Reverse-KL OPD: When S = Sθ(s), the selected top-k support is treated as fixed during the local update.This is the practical assumption used for sparse reverse-KL optimization.
  • A.2 Monte Carlo Reverse-KL OPD: The dense reverse-KL gradient is an expectation over current-policy actions of the advantage-weighted log-policy gradient.The advantage is Aθ(a, s) = log q(a | s) − log pθ(a | s).
  • A.2 Monte Carlo Reverse-KL OPD: A one-sample Monte Carlo estimator uses one action from the current policy, while a multi-sample estimator averages the same local term over m independent samples.Both estimators use bgMC(s, a) = −Aθ(a, s)∇θ log pθ(a | s).

B Experimental Details · C Datasets and Metrics · D Existing Asset Licenses

Experiments use a standardized single-node setup with controlled train-batch staleness, while evaluation relies on Avg@32 accuracy across three math benchmarks. The sections also document the filtered training data and reused-asset licensing information.

  • B Experimental Details: Experiments use vLLM for rollouts and teacher scoring, PyTorch FSDP for learner training, and a single 8×B200 node.Individual experiments take roughly 1–12 hours, depending on the setting.
  • D Existing Asset Licenses: Reused assets are documented with their URLs, license names, and versions in Table 5.Table 5 lists the reused assets.
  • B Experimental Details: Staleness is controlled in train-batch steps, where k denotes how many consumed cache steps trail the current learner.The controlled sweep uses k ∈ {0, 1, 2, 4, 8, 16, 32, 64, 128}.
  • B Experimental Details: A 64-batch target cache depth corresponds to plotted staleness 64, and each train-batch step contains four optimizer updates in the common setup.With B = 256 and Bmini = 64, M = B/Bmini = 4.
  • C Datasets and Metrics: The training set contains 57,630 DeepMath problems filtered to difficulty level greater than or equal to 6.This filtered subset is used as the training data.
  • C Datasets and Metrics: Evaluation uses AIME 2024, AIME 2025, and AMC 2023; AIME24 is evaluated every 20 steps, while other datasets are evaluated only at the final checkpoint.These datasets are listed in Table 4.
  • C Datasets and Metrics: Avg@32 is the mean per-problem pass rate from 32 sampled responses, and it is the default metric in tables and plots.The metric counts correct sampled responses for each problem before averaging across the dataset.

E Multi-Sample MC Variance at Large Staleness · F Importance-Sampling Ablation

At large staleness, multi-sample Monte Carlo consistently lowers estimator variance, with sequence-level reductions less extreme than local next-token reductions. Importance sampling remains the mechanism for correcting old-policy samples toward the current reverse-KL gradient, distinct from increasing the sample count.

  • E Multi-Sample MC Variance at Large Staleness: Multi-sample MC draws iid local actions with replacement from the old policy at each fixed prefix, allowing duplicate actions.The learner evaluates these local scorer queries rather than separate rollout branches.
  • E Multi-Sample MC Variance at Large Staleness: Variance ratios are normalized to the corresponding m = 1 estimator within the same old-to-current policy pair.The fixed-prefix and sequence-level ratios use Qwen3-4B-Base staleness-128 runs.
  • E Multi-Sample MC Variance at Large Staleness: The generated prefix path is fixed during variance computation, so Monte Carlo samples are local scorer queries at each prefix.They are not separate rollout branches.
  • E Multi-Sample MC Variance at Large Staleness: 1.49%: with m = 64, fixed-prefix variance falls to 1.49% of one-sample local variance, near the 1/64 = 1.56% reference.This supports the expected 1/m reduction for local action-sampling variance.
  • E Multi-Sample MC Variance at Large Staleness: 11.2%: after timestep aggregation, m = 64 leaves 11.2% of one-sample sequence-level variance, showing a less extreme but practical reduction.Larger m consistently reduces variance, while sequence-level averaging moderates the local effect.
  • F Importance-Sampling Ablation: Increasing m changes Monte Carlo variance but does not change the target distribution; old-to-current importance sampling performs the correction.Figure 10 compares MC1 and MC16 with and without IS to separate these effects.

G AsyncOPD Scheduler Details

AsyncOPD implements a fully asynchronous OPD pipeline with separate rollout, teacher-scoring, and learner stages, using queued OPD cache items rather than RL trajectories. Compared with strict synchronization and a controlled two-update offset, it achieves higher throughput and overlap while preserving comparable accuracy and earlier wall-clock progress.

  • Scheduler architecture: AsyncOPD uses three long-running stages—rollout generation, teacher scoring, and learner training—with a queue of OPD cache items.Rollout workers cache MC actions, rollout log probabilities, and student versions for teacher scoring and learner updates.
  • Scheduler architecture: During keep-mode synchronization, in-flight prefixes are retained, student weights are updated, and prefix caches are reset before generation resumes.This reuses tokens sampled before synchronization while rebuilding attention state under the new weights.
  • Experimental protocol: Each scheduler run uses one 8-GPU node, with one GPU for teacher scoring and the remaining seven split between rollout and training in asynchronous runs.Strict synchronization time-shares all seven pool GPUs, whereas asynchronous runs allocate four to rollout and three to training.
  • Experimental protocol: The comparison uses strict sync, two-step-off, and AsyncOPD with the same teacher, data, metrics, and reverse-KL estimator, including current-policy Aθ and old-to-current IS correction.Two-step-off imposes a fixed two-update rollout offset to control stale reuse independently of queue timing.
  • Results: Up to 3.8× strict-sync throughput on MC64 and up to 3.2× on MC1, while maintaining comparable final accuracy and the highest throughput and overlap in matched Qwen3 settings.AsyncOPD also reaches later checkpoints sooner, so accuracy improves earlier across student sizes and MC settings.
Loading 2606.24143v1…