Command Palette
Search for a command to run...
プレフィックスリプレイを用いたマルチターンオンポリシー蒸留
プレフィックスリプレイを用いたマルチターンオンポリシー蒸留
Baohao Liao Hanze Dong Christof Monz Xinxing Xu Li Dong Furu Wei
概要
本研究では、LLMエージェントが複数ターンにわたり環境と対話し、生徒モデルがその対話履歴を通じて教師モデルを模倣するエージェント的タスクにおけるオンポリシー蒸留(OPD)を検討する。完全オンラインのOPDは、更新ごとに新たな生徒のロールアウトと訪問履歴での教師クエリが必要となるためコストが高い。我々は、事前収集した教師の軌跡をリプレイ用プレフィックスとして再利用するオフ環境代替手法、Replayed-Prefix On-Policy Distillation(ReOPD)を提案する。生徒は選択されたステップで行動し、教師は新たな環境対話を実行せずに密なステップ単位の監督を提供する。マルチターンOPDにはプレフィックストラップが存在することを示す。履歴をより生徒オンポリシーにすることは生徒への関連性を向上させるが、教師の目標が信頼できない履歴で教師にクエリを行う可能性がある。これにより、生徒の占有度と教師の信頼性の間に両面の分布シフトが生じる。ReOPDは、マルチターンOPDを信頼性を考慮したプレフィックス分布設計として扱い、シフトの少ない初期のプレフィックスを重視する単純なステップ減衰サンプリングスケジュールで実装することでこれに対処する。Pythonを用いた数学的推論および検索環境において、複数の教師・生徒モデル規模で評価した結果、ReOPDはOPDレベルの精度を維持または向上させ、生徒の訓練中にツール呼び出しを一切使用せず、訓練ステップあたり少なくとも4倍高速である。したがってReOPDは、高コストなエージェントと環境の対話を再利用可能なオフラインリソースに変換し、ツール、タスク、環境を横断するスケーラブルな蒸留を可能にする。
One-sentence Summary
Researchers from Microsoft Research and the University of Amsterdam propose Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment method that reuses teacher trajectories as replayed prefixes and employs a step-decaying sampling schedule to mitigate the prefix trap in multi-turn agentic tasks, thereby preserving or improving accuracy while eliminating tool calls and achieving at least 4× faster training per step than standard on-policy distillation.
Key Contributions
- The paper proposes Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment distillation method that reuses pre-collected teacher trajectories as replayed prefixes, enabling the student to act at selected steps while the teacher provides dense per-step supervision without live environment interactions.
- The paper identifies and formalizes the prefix trap in multi-turn on-policy distillation, a two-sided distribution shift where making histories more student-on-policy can lead to teacher queries on unreliable histories, and frames the problem as reliability-aware prefix distribution design with a bound that decomposes the gap into student occupancy mismatch and teacher reliability terms.
- The paper proposes a step-decaying sampling schedule to implement reliability-aware prefix distribution, and demonstrates that ReOPD preserves or improves OPD-level accuracy across mathematical reasoning with Python and search environments, uses zero tool calls during student training, and achieves at least 4× faster per-step training than OPD.
Introduction
Reinforcement learning (RL) from verifiable rewards drives strong reasoning in large language models, but its sparse scalar signal is sample-inefficient, especially for smaller models. Knowledge distillation provides dense supervision by imitating a teacher, yet standard off-policy distillation trains on teacher prefixes, causing the student to drift and compound errors when it makes mistakes. On-policy distillation (OPD) samples from the student while still using the teacher’s per-token distribution, but extending OPD to multi-turn agentic tasks is expensive because it requires live environment interaction and fresh teacher queries at every step. Replaying pre-collected teacher trajectories as prefixes avoids these costs, but introduces a prefix trap: a two-sided distribution shift where student-relevant prefixes may land in regions where the teacher’s guidance is unreliable, so fully student-on-policy roll-ins are not always optimal. The authors formulate multi-turn OPD as a reliability-aware prefix distribution design, derive a bound that decomposes the gap into student occupancy mismatch and teacher reliability, and propose ReOPD, a practical method that uses a step-decay schedule to sample prefixes from an offline teacher pool. ReOPD matches on-policy OPD when the teacher is reliable on student histories, and outperforms it when the teacher-student gap is large, while eliminating environment interaction during distillation.
Method
The authors formulate multi-turn on-policy distillation as an interactive process where the goal is to optimize a student policy on the interaction histories it will actually encounter. However, relying solely on student-generated histories can lead to unreliable teacher supervision if those histories drift far from the teacher's reliable interaction support. To address this, the authors propose a reliability-aware interaction-prefix distribution that balances student relevance against teacher reliability. This is formalized through a bridge objective that interpolates between the student's occupancy and the teacher's occupancy using a geometric bridge distribution.
To realize this principle efficiently without requiring continuous environment interaction, the authors introduce Replayed-Prefix On-Policy Distillation. This algorithm operates off-environment using a fixed pool of pre-collected teacher trajectories. The method relies on two core components: a teacher-forced prefix with student action and teacher supervision, and a reliability-aware step schedule. For any supervised step, the interaction prefix is replayed verbatim from a teacher trajectory. The student then generates its own action autoregressively on this prefix, and the teacher's recorded conditional distribution provides the distillation target. This design ensures that the student's evaluated step is genuinely on-policy while keeping the roll-in strictly within the teacher's occupancy, completely eliminating the need for environment queries.
The exact bridge weight required to reshape the teacher pool toward the ideal distribution involves computing a per-history student-to-teacher likelihood ratio along the prefix. The authors observe that this ratio naturally decays as the prefix lengthens, because the student is increasingly unlikely to follow the teacher's recorded actions at deeper steps. As shown in the figure below, the logarithm of the likelihood ratio exhibits a strong, monotonic decay with respect to the normalized step index, demonstrating that the step index serves as a highly effective proxy for the expensive prefix-level ratio.
Consequently, the authors replace the exact density-ratio weight with a stable, one-parameter step-decay schedule defined as ω(t;κ)=κt, where κ∈(0,1] controls the steepness of the decay.
The final weighted distillation objective minimizes the token-level KL divergence between the student and teacher distributions, reweighted by this step-decay schedule. By sampling training positions in proportion to the schedule, the method concentrates supervision on early, low-shift positions where the teacher signal remains highly reliable, while gracefully de-emphasizing later steps where the teacher's conditional distribution becomes a poor proxy for the ideal improvement target. The steepness parameter κ is directly tied to the average per-step KL divergence between the teacher and the current student, allowing the schedule to automatically adapt to the teacher-student gap.
Experiment
ReOPD replaces online student-environment interaction with replay of teacher-recorded prefixes, making distillation at least four times faster per training step while preserving or improving student accuracy. On mathematical reasoning tasks, where teacher–student capability gaps can reduce teacher reliability on student-induced histories, ReOPD consistently outperforms fully online OPD; on search tasks, it matches OPD accuracy. The gains come from a teacher-reliability-aware prefix selection that favors early teacher-forced steps, and the method effectively combines multiple environments without requiring live tools during student training.
When the teacher–student capability gap is large, the teacher’s own GRPO rollouts provide a more reliable prefix pool than student-on-policy rollouts, leading to substantial gains on mathematical reasoning tasks. On search tasks where the teacher is already reliable on student histories, student-on-policy OPD and ReOPD perform similarly, confirming that the optimal roll-in depends on the reliability regime. ReOPD reuses the teacher’s GRPO trajectories, avoiding any extra collection cost while matching a stationary teacher pool. ReOPD reuses the teacher’s GRPO trajectories as the prefix pool, eliminating the need for dedicated collection. On math reasoning, ReOPD improves AIME24 from 28.3 to 36.7 under a Qwen3-8B teacher and lifts the average across all teacher settings. On search tasks, student-on-policy OPD and ReOPD yield nearly identical averages (40.5 vs. 40.6), showing the teacher-anchored roll-in ties when the teacher is already reliable. Teacher-anchored prefixes outperform those from a stronger model, confirming that reliability with respect to the teacher, not standalone capability, is what matters. A single fixed κ=0.6 across tasks already achieves regime-appropriate behavior, with gap-adaptive schedules suggested as a natural extension.
Both environments use a temperature of 1.0. Math tasks allow up to 16,384 tokens for generation, while search tasks are capped at 8,192 tokens. Accuracy is averaged over 8 samples for the hardest math benchmarks (AIME24, AIME25, AMC23), 4 samples for other math benchmarks, and a single sample for search. Challenging competition math benchmarks use 8-fold averaging to reduce variance in pass@1 estimates. Search evaluation uses n=1, reflecting the deterministic nature of the retrieval environment.
On mathematical reasoning, ReOPD consistently improves student performance over fully online OPD, with the largest gains when the teacher-student gap is wider. The method anchors roll-ins to teacher-supported histories, sidestepping student-induced drift where the teacher becomes less reliable. For the Qwen3-4B student learning from the same-scale Qwen3-4B teacher, ReOPD raised the average accuracy across six math benchmarks from 55.1 to 57.2. ReOPD's advantage over OPD grows as the teacher-student gap widens, aligning with the teacher-reliability regime where student-on-policy roll-ins drift into less trustworthy teacher targets.
In search distillation, ReOPD performs on par with online student-on-policy distillation (OPD) when the teacher remains reliable on student trajectories. The student achieves an average accuracy of 40.5 with ReOPD versus 40.6 with OPD across general and multi-hop QA benchmarks, while the base model obtains 35.6. This confirms that teacher-anchored replay can replace live environment interaction without sacrificing accuracy in this regime. ReOPD matches OPD accuracy (40.5 vs 40.6 average) on search tasks, eliminating the need for online rollouts when the teacher is reliable. Student distillation (OPD) lifts average accuracy from 35.6 (base) to 40.6, slightly surpassing the teacher's own GRPO result of 40.4. ReOPD maintains OPD-level performance while making training over 4x faster per step by using zero tool calls during student updates.
A single Qwen3-4B student distilled jointly from separate math and search teachers with ReOPD matches the performance of online student-on-policy distillation (OPD) on math benchmarks. The joint training removes the need for online environment interaction during student updates, showing that offline replay from heterogeneous environments can be combined effectively. ReOPD’s reliability-aware schedule carries over to the multi-domain setting, keeping the student on par with OPD. On math, the joint ReOPD student outperforms OPD on AIME24 (36.3 vs. 32.5) and Minerva (44.5 vs. 43.5) while tying on Olympiad (57.4), and is only slightly behind on AIME25 (30.8 vs. 35.4) and AMC23 (75.5 vs. 75.9). Joint training with ReOPD avoids online tool calls during student updates, enabling a single student to be distilled from two domains without keeping either environment active.
The experiments evaluate ReOPD, a method that reuses the teacher’s GRPO trajectories as a prefix pool, on mathematical reasoning and search tasks under varying teacher-student capability gaps. When the gap is large, student-on-policy rollouts become unreliable and ReOPD substantially improves math performance by anchoring to teacher-supported histories; on search tasks where the teacher is already reliable, ReOPD matches online distillation without live environment interaction. Joint training across multiple domains with ReOPD also preserves accuracy while avoiding online tool calls during student updates. The findings confirm that teacher-anchored replay provides a robust, cost-efficient alternative to online distillation by adapting to the teacher’s reliability regime.