Multi-Rollout On-Policy Distillation via Peer Successes and Failures

TL;DR

Proposes Multi-Rollout On-Policy Distillation (MOPD), leveraging peer successes and failures to enhance model reasoning by 20% on math tasks.

cs.LG 🔴 Advanced 2026-05-13 70 views
Weichen Yu Xiaomin Li Yizhou Zhao Xiaoze Liu Ruowang Zhang Haixin Wang Yinyi Luo Chen Henry Wu Gaurav Mittal Matt Fredrikson Yu Hu
deep learning policy distillation multi-rollout contrastive learning reasoning

Key Findings

Methodology

MOPD introduces peer-conditioned distillation by sampling multiple trajectories per prompt, scoring via verifier, and partitioning into success/failure sets. It constructs peer contexts using successful and failed trajectories, applying divergence-based losses (inverse KL or Jensen-Shannon). This approach enables the teacher to compare current trajectories against peer attempts, providing sharper, instance-adaptive supervision. Experiments across reasoning, scientific QA, coding, and tool use show consistent improvements over baseline methods, with analyses confirming better alignment of teacher signals with verifier rewards, especially when combining success and failure trajectories.

Key Results

  • On math benchmarks AIME2024/2025, MOPD improves mean@8 by over 20% compared to SDPO, demonstrating superior reasoning capabilities.
  • In scientific QA and tool use tasks, performance gains of 10-15% are observed, with the contrastive peer context outperforming success-only conditioning.
  • Ablation studies confirm that integrating both successful and failed peer trajectories yields the most faithful supervision, leading to faster convergence and robustness.

Significance

This work advances the field by effectively utilizing local, instance-specific information from multiple rollouts, addressing the limitations of traditional single-trajectory distillation. It demonstrates that leveraging peer success and failure signals enhances the model’s ability to distinguish valid reasoning paths from misleading ones, significantly boosting performance in complex reasoning tasks. The approach offers a new paradigm for instance-adaptive supervision, with broad implications for improving autonomous learning, robustness, and generalization in large language models and decision-making systems.

Technical Contribution

The paper introduces a novel peer-conditioned distillation framework that integrates multiple trajectories within the same problem instance, employing divergence-based contrastive losses. It innovatively combines success and failure trajectories to produce instance-specific, sharper supervision signals, surpassing traditional methods that treat trajectories independently. The approach provides theoretical guarantees for improved alignment with verifier rewards and practical benefits in diverse reasoning tasks, opening new avenues for instance-adaptive learning and multi-trajectory exploitation.

Novelty

This is the first systematic integration of peer success and failure trajectories into on-policy policy distillation, transforming the process from independent imitation to a local contrastive learning paradigm. The dual strategies—positive peer imitation and contrastive success-failure conditioning—offer a new way to utilize local structure information, setting a new standard for instance-specific supervision in large-scale language models. This innovation significantly extends prior work that only used success signals or external advantage estimates.

Limitations

  • The effectiveness depends on the quality and diversity of sampled trajectories; insufficient or biased sampling may limit gains.
  • Computational overhead increases with the number of trajectories, posing challenges for large-scale deployment.
  • In environments with extremely sparse rewards, peer trajectories may lack informative structure, reducing supervision quality.

Future Work

Future research could focus on adaptive sampling strategies to improve trajectory diversity and quality, integrating multi-modal data for richer context, and exploring reinforcement learning techniques to optimize trajectory generation. Extending the framework to real-world applications such as robotics and autonomous systems, and reducing computational costs, are promising directions.

AI Executive Summary

In recent years, large language models have demonstrated remarkable capabilities in reasoning and decision-making tasks. However, their training often relies on sparse reward signals, which only indicate success or failure without providing detailed guidance on the reasoning process. Traditional on-policy distillation methods improve supervision density but treat each sampled trajectory independently, missing the opportunity to leverage the rich local structure within multiple attempts.

This paper introduces Multi-Rollout On-Policy Distillation (MOPD), a novel framework that exploits the local context of multiple peer trajectories generated for the same problem instance. By conditioning the teacher model on both successful and failed peer attempts, MOPD provides more targeted, instance-specific supervision. The core idea is to compare the current trajectory against peer successes, which serve as positive evidence, and peer failures, which act as structured negative evidence. This contrastive approach sharpens the supervision signals, enabling the student model to better distinguish valid reasoning paths from misleading ones.

The methodology involves sampling multiple trajectories per prompt, scoring them with a verifier, and constructing peer contexts using success and failure sets. The teacher then uses divergence-based losses—such as inverse KL or Jensen-Shannon—to guide the student. Two main strategies are explored: positive peer imitation, which conditions only on successful peers, and contrastive success-failure conditioning, which incorporates both. Experimental results across tasks like mathematical reasoning, scientific QA, programming, and tool use show that MOPD consistently outperforms baseline methods, with improvements exceeding 20% in some benchmarks.

Analysis reveals that the mixed success-failure peer context aligns the teacher’s signals more closely with verifier rewards, leading to more faithful supervision. This approach effectively utilizes the local trial-and-error behavior of the student, transforming the distillation process into a more adaptive, instance-aware learning paradigm. The findings suggest that leveraging multi-trajectory information is crucial for advancing reasoning capabilities in large language models, with broad implications for AI research and practical applications.

Deep Analysis

Background

深度学习中的策略蒸馏技术经过多年发展,逐渐成为提升模型推理能力的关键手段。早期工作如Hinton等提出的知识蒸馏,主要关注模型输出的概率匹配,后续研究引入强化学习和自我蒸馏,增强模型的泛化和鲁棒性。近年来,基于轨迹采样的策略蒸馏(如AGARWAL等的OPD)逐步成为主流,强调在模型自身行为中进行学习,减少分布偏差。然而,现有方法多忽视多轨迹组的局部结构信息,未能充分利用同行轨迹中的正负证据。与此同时,强化学习中的多轮次采样与验证器引导的改进(如Shao等的多轮强化学习)虽在奖励密度上有所突破,但在知识迁移中仍未融合多轨迹对比信息,限制了模型的推理深度和鲁棒性。

Core Problem

当前深度模型在复杂推理任务中表现有限,部分原因在于训练过程中未能充分利用多次采样的局部结构信息。单一轨迹的知识蒸馏忽略了同行轨迹中的正负证据,导致模型难以区分正确路径与误导路径。尤其在推理任务中,成功轨迹提供有效的解题策略,而失败轨迹揭示潜在的误区和错误模式。如何设计一种机制,将多轨迹的局部差异作为指导信号,提升模型的实例适应性和推理能力,成为亟待解决的问题。

Innovation

本文的核心创新在于引入同行轨迹的对比信息到策略蒸馏中,提出多轮次策略蒸馏(MOPD)。具体包括:1)利用成功轨迹作为正向同行证据,2)结合失败轨迹作为结构化负证据,3)设计基于逆KL和Jensen-Shannon散度的对比损失,4)构建多轨迹条件化教师信号,增强模型对局部推理路径的识别能力。这一机制突破了传统单轨迹蒸馏的局限,为模型提供了更丰富、更具针对性的监督信息,有效提升推理任务中的表现。

Methodology

  • �� 在每个问题实例中采样多条轨迹(N条),每条轨迹由学生策略πθ生成。
  • �� 利用验证器对轨迹进行评分,将轨迹划分为成功(Y+)和失败(Y−)两组。
  • �� 构建同行轨迹上下文:成功轨迹作为正证据,失败轨迹作为负证据,形成条件化教师信号。
  • �� 采用逆KL或Jensen-Shannon散度作为轨迹对比损失,优化学生模型。
  • �� 训练流程包括采样、评分、构建上下文、计算损失、模型更新,循环进行。

Experiments

在数学、科学问答、编程和工具使用任务中,使用AIME、SciKnowEval、LiveCodeBench和ToolAlpaca等数据集进行训练。对比基线包括传统OPD、GRPO和SDPO,指标涵盖平均正确率、pass@8等。实验中,采样轨迹数保持一致,验证器评分统一,进行消融验证不同同行轨迹构建策略的效果。超参数如采样次数、奖励阈值等保持一致,确保公平。

Results

MOPD在所有任务中均优于对比方法,数学推理提升超过20%,科学问答和工具任务提升10-15%。结合成功与失败轨迹的对比条件化显著优于仅用成功轨迹,验证了结构化负证据的有效性。分析显示,混合轨迹条件化能更好匹配验证器奖励,模型在复杂推理中表现更鲁棒,训练收敛速度更快。

Applications

该方法适用于多步推理、结构化输出和复杂决策场景,如自动编程、科学研究、智能问答和机器人控制。通过同行轨迹对比,模型能更有效识别潜在错误,提升自主学习和鲁棒性。未来结合多模态信息和强化采样,有望在更复杂环境中实现自主推理与决策。

Limitations & Outlook

当前方法对轨迹采样数量敏感,采样不足或偏差会影响效果。在高复杂度任务中,计算成本较高,硬件资源要求大。稀疏奖励环境中,同行轨迹缺乏结构信息,影响监督效果。未来需优化采样策略和结构信息利用,以适应更复杂应用。

Plain Language Accessible to non-experts

想象你在学习做菜,你尝试了几次不同的方法。有时候你做得很好,味道很棒;有时候则失败了。这些尝试就像模型的轨迹,有成功的,也有失败的。现在,你的老师不仅告诉你哪次做得好,还会告诉你哪次做错了,为什么错了。这样,你就能更清楚地知道哪些步骤是正确的,哪些是需要改正的。通过比较这些不同的尝试,你可以学得更快、更好。这就像MOPD的方法,用多次尝试的成功和失败来指导学习,而不是只看最终的结果。这样,学习变得更有效,也更贴近实际的试错过程。

ELI14 Explained like you're 14

想象你在玩一款游戏,你每次都试不同的策略。有时候你赢了,有时候输了。你的哥哥会告诉你:‘这个策略不错,下次可以试试这个!’或者:‘这个方法不行,要注意这个陷阱!’通过不断比较成功和失败的经验,你学得更快、更聪明。这就像模型在学习时,不只是看最终答案,而是用多次尝试的结果来判断哪些做法是对的,哪些要避免。这样,模型就能像你一样,逐渐变得更厉害,能应对各种难题。这个方法叫做多轮次策略蒸馏,它让模型在试错中不断改进,就像你在游戏中变得越来越厉害一样。

Abstract

Large language models are often post-trained with sparse verifier rewards, which indicate whether a sampled trajectory succeeds but provide limited guidance about where reasoning succeeds or fails. On-policy distillation (OPD) offers denser token-level supervision by training on student-generated trajectories, yet existing methods typically distill each rollout independently and ignore the other attempts sampled for the same prompt. We introduce Multi-Rollout On-Policy Distillation (MOPD), a peer-conditioned distillation framework that uses the student's local rollout group to construct more informative teacher signals. MOPD conditions the teacher on both successful and failed peer rollouts: successes provide positive evidence for valid reasoning patterns, while failures provide structured negative evidence about plausible mistakes to avoid. We study two peer-context constructions: positive peer imitation and contrastive success-failure conditioning. Experiments on competitive programming, mathematical reasoning, scientific question answering, and tool-use benchmarks show that MOPD consistently improves over standard on-policy baselines. Further teacher-signal analysis shows that mixed success-failure contexts better align teacher scores with verifier rewards, indicating that the gains arise from more faithful, instance-adaptive supervision. These results indicate that effective on-policy distillation should exploit the student's multi-rollout trial-and-error behavior rather than treating rollouts as isolated samples.

cs.LG cs.AI