HyperAIHyperAI

Command Palette

Search for a command to run...

Distillation sur politique multi-tours avec rejeu de préfixes

Baohao Liao Hanze Dong Christof Monz Xinxing Xu Li Dong Furu Wei

Résumé

Nous étudions la distillation sur politique (OPD) pour des tâches agentiques, où un agent LLM interagit avec un environnement sur plusieurs tours et un étudiant imite un enseignant à travers ces historiques d'interactions multi-tours. L'OPD entièrement en ligne est coûteuse car chaque mise à jour nécessite de nouveaux déploiements de l'étudiant dans l'environnement et des requêtes à l'enseignant sur les historiques visités. Nous proposons Replayed-Prefix On-Policy Distillation (ReOPD), une alternative hors environnement qui réutilise des trajectoires d'enseignant pré-collectées comme préfixes rejoués : l'étudiant agit à des étapes sélectionnées, tandis que l'enseignant fournit une supervision dense par étape sans exécuter de nouvelles interactions avec l'environnement. Nous montrons que l'OPD multi-tours introduit un piège de préfixe : rendre les historiques plus conformes à la politique de l'étudiant améliore la pertinence pour celui-ci, mais peut interroger l'enseignant sur des historiques où sa cible n'est pas fiable. Cela crée un décalage de distribution bilatéral entre l'occupation de l'étudiant et la fiabilité de l'enseignant. ReOPD résout ce problème en traitant l'OPD multi-tours comme une conception de distribution de préfixes sensible à la fiabilité et le met en œuvre avec un schéma d'échantillonnage à décroissance par étapes qui privilégie les préfixes précoces à faible décalage. À travers des environnements de raisonnement mathématique avec Python et de recherche sur plusieurs échelles de modèles enseignant et étudiant, ReOPD préserve ou améliore la précision de niveau OPD, n'utilise aucun appel d'outil pendant l'entraînement de l'étudiant, et est au moins 4 fois plus rapide par étape d'entraînement que l'OPD. ReOPD transforme donc l'interaction coûteuse agent-environnement en une ressource hors ligne réutilisable, permettant une distillation évolutive à travers outils, tâches et environnements.

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×4\times4× 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\omega(t; \kappa) = \kappa^tω(t;κ)=κt, where κ(0,1]\kappa \in (0, 1]κ(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 κ\kappaκ 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.


Créer de l'IA avec l'IA

De l'idée au lancement — accélérez votre développement IA avec le co-codage IA gratuit, un environnement prêt à l'emploi et le meilleur prix pour les GPU.

Codage assisté par IA
GPU prêts à l’emploi
Tarifs les plus avantageux

HyperAI Newsletters

Abonnez-vous à nos dernières mises à jour
Nous vous enverrons les dernières mises à jour de la semaine dans votre boîte de réception à neuf heures chaque lundi matin
Propulsé par MailChimp