CoDistill-GRPO: A Co-Distillation Recipe for Efficient Group Relative Policy Optimization
Proposes CoDistill-GRPO, jointly training large and small models, boosting small model accuracy by 11.6%, with 18% speedup.
Key Findings
Methodology
CoDistill-GRPO integrates group relative policy optimization (GRPO) with co-distillation. The small model learns from the large model's distribution via on-policy KD rewards, while the large model is updated using rollouts generated solely by the small model with importance reweighting. The two models mutually teach each other, optimizing combined objectives. The approach involves multi-round rollout generation, filtering based on combined rewards, and decoupled gradient updates, significantly reducing computational overhead while maintaining performance. The framework enables simultaneous training, leveraging many inexpensive small-model rollouts for effective learning.
Key Results
- On Qwen2.5-Math-1.5B, accuracy improved from 20.37% to 31.99%, a 11.6 percentage point increase, surpassing standard GRPO by 6.0 points. The large model (Qwen2.5-Math-7B) trained with CoDistill-GRPO nearly matches traditional GRPO performance, with an 18% speedup, reducing training time.
- In Llama models, performance on math benchmarks (e.g., MATH500, AMC2024) improved by over 15%, demonstrating broad applicability. Theoretical analysis confirms unbiased gradient estimates, and multi-round filtering balances exploration and exploitation, ensuring stable training.
- The method effectively enhances small model reasoning capabilities, reduces training costs, and accelerates convergence, making it suitable for resource-constrained environments and large-scale deployment scenarios.
Significance
This work addresses the longstanding challenge of sparse rewards in small models during complex reasoning tasks. By jointly training large and small models with co-distillation, it achieves significant performance gains while reducing computational costs. The approach bridges the gap between model capacity and training efficiency, offering a scalable solution for deploying high-performance AI systems in resource-limited settings. Its theoretical guarantees and empirical success open new avenues for multi-model collaborative learning, promising broader impacts across NLP and AI fields.
Technical Contribution
The paper introduces CoDistill-GRPO, a novel algorithm combining multi-model joint training, importance reweighting, and multi-round rollout filtering. It provides rigorous theoretical guarantees, including unbiased gradient estimation, and demonstrates that integrating on-policy KD rewards guides small models toward large model distributions effectively. The framework reduces rollout computation by leveraging inexpensive small-model generations, enabling near-parity performance of large models trained on small rollouts, with an 18% speedup. This advances policy optimization methods by enabling cost-effective, scalable training of large models with high accuracy.
Novelty
This is the first work to jointly train large and small models via a co-distillation framework integrated with GRPO, using multi-round rollout filtering and importance reweighting. Unlike prior static teacher-student setups, this dynamic mutual learning approach improves efficiency and performance simultaneously. The method innovatively combines policy optimization with on-policy KD rewards, enabling small models to learn effectively from limited rollouts, and achieves near-large-model performance at a fraction of the cost, setting a new standard for resource-efficient training.
Limitations
- Despite reducing rollout costs, the multi-round filtering process still requires considerable computation, especially for very large models.
- Sensitivity to hyperparameters like α affects training stability; improper tuning may lead to suboptimal results.
- Effectiveness on extremely complex or multi-task scenarios remains to be validated; further adaptation may be necessary.
Future Work
Future directions include refining rollout filtering strategies, automating hyperparameter tuning, and extending the framework to multi-task and multi-modal settings. Exploring adaptive reward functions and integrating self-supervised pretraining could further enhance robustness and generalization, making the approach more versatile for diverse AI applications.
AI Executive Summary
This paper introduces CoDistill-GRPO, an innovative framework that jointly trains large and small language models by integrating group relative policy optimization (GRPO) with co-distillation. The core challenge addressed is the sparse reward problem faced by small models in complex reasoning tasks, which hampers learning efficiency. Traditional approaches rely heavily on large static teachers, incurring high computational costs and limited flexibility. In contrast, CoDistill-GRPO employs a dynamic mutual learning process where the small model learns from the large model's distribution via on-policy knowledge distillation, while the large model benefits from rollouts generated solely by the small model, utilizing importance reweighting to reduce computational overhead. This bidirectional training setup enables the models to improve simultaneously, leveraging many inexpensive small-model rollouts for effective learning.
The methodology involves multi-stage sampling, reward-based filtering, and gradient decoupling. During training, the large model provides initial hints (when T > 0), which the small model completes, and both are evaluated using a combined reward that includes an on-policy KD component. The most informative rollouts are selected through downsampling, balancing exploration and exploitation. Theoretical analysis confirms that the gradient estimates are unbiased, ensuring training stability. Empirical results on Qwen and Llama models across multiple mathematical benchmarks demonstrate that CoDistill-GRPO achieves over 11% accuracy improvement and reduces training time by 18%. These findings highlight the method's effectiveness in enhancing small model reasoning while maintaining computational efficiency. Overall, this work offers a scalable, cost-effective approach for high-performance language model training, with promising implications for resource-constrained AI deployment and future multi-model collaborative learning.
Deep Analysis
Background
近年来，深度学习中的大规模语言模型（如GPT、LLaMA）在自然语言理解和推理任务中取得巨大突破，但其训练成本极高。为降低成本，知识蒸馏（KD）和强化学习（RL）策略逐渐普及，尤其在数学推理等复杂任务中，模型表现受限于奖励稀疏和训练效率。传统方法多依赖预训练或静态大模型作为教师，存在效率低、成本高的问题。近年来，群相对策略优化（GRPO）作为一种无需训练值函数的策略优化算法，因其简洁高效而受到关注，但在小模型上仍面临奖励稀疏的瓶颈。为解决这一难题，本文提出了协同蒸馏机制，结合多轮rollout筛选，提升小模型推理能力，推动大模型训练的高效化。
Core Problem
核心问题在于小模型在复杂推理任务中因奖励稀疏，学习缓慢甚至停滞，导致训练效果不理想。传统方法依赖大模型作为教师，存在训练时间长、模型差异大带来的性能瓶颈。此外，rollout生成成本高，难以在大规模训练中普遍应用。如何在保证模型性能的同时，降低训练成本、提升效率，成为亟待解决的关键问题。
Innovation
本研究的创新点包括：1）提出协同蒸馏（CoDistill）机制，联合训练大模型与小模型，避免单向依赖；2）引入多轮rollout筛选策略，通过奖励重加权筛选出最具代表性样本，提升训练效率；3）结合GRPO与on-policy知识蒸馏，设计有效奖励机制，增强小模型学习能力；4）理论上证明梯度无偏，确保训练稳定性。这些创新共同推动了小模型在复杂推理任务中的性能提升，显著降低训练成本。
Methodology
- �� 初始化大模型和小模型参数，设定训练数据和超参数。• 在每轮训练中，从数据集中采样prompt，先由大模型生成部分hint（T个tokens），作为提示输入小模型。• 小模型完成剩余tokens，生成完整序列。• 计算奖励：大模型用标准奖励（如准确率），小模型用结合KD的有效奖励（加入对大模型分布的匹配项）。• 利用多轮rollout，筛选出奖励最高和最低的样本，进行下采样。• 通过梯度反向传播，分别更新大模型和小模型参数，优化目标包括GRPO损失和KD奖励。• 重复上述步骤，逐步提升模型性能，直至收敛。
Experiments
在Qwen和Llama模型上，使用数学推理基准（Minerva、MATH500、AMC2024、OlympiadBench）进行评估。对比基线包括标准GRPO和静态大模型蒸馏。超参数包括rollout数量（M=14）、筛选后样本数（G=8）、奖励系数（α=1或2）等。训练过程中，采用多轮rollout生成、多次筛选和梯度优化，确保模型在复杂推理任务中的表现提升。
Results
在Qwen2.5-Math-1.5B模型上，使用CoDistill-GRPO后，准确率从20.37%提升至31.99%，比标准GRPO提升6.0个百分点（25.88%到31.99%）。在Minerva数据集上，性能提升显著，训练速度提升18%。在Llama模型上，平均性能提升超过15%，在多个基准任务中表现优异。理论分析验证了梯度无偏性，筛选机制有效平衡探索与利用，整体提升训练效率和模型性能。
Applications
该算法适用于需要高效推理能力的场景，如自动问答、数学题解、代码生成等。在资源有限的环境中，能显著降低训练成本，加快模型部署速度。未来，结合多任务学习和自监督预训练，有望推动大模型在工业界的广泛应用，尤其是在边缘设备和实时系统中实现高性能推理。
Limitations & Outlook
尽管算法降低了训练成本，但多轮rollout筛选仍需一定计算资源，尤其在超大模型上可能存在瓶颈。参数调优（如α值）对性能影响较大，且在极端复杂任务中效果尚待验证；未来需优化奖励机制和筛选策略。
Plain Language Accessible to non-experts
想象你在做一道难题，比如拼装一台复杂的机械装置。你有两个助手：一个经验丰富、懂得很多（大模型），另一个新手（小模型）。大助手可以给你一些提示，但每次提示都很费时间。小助手可以自己尝试拼装，但容易迷路。为了快点成功，你让两个助手合作：大助手先给出部分提示，小助手根据提示继续拼装，然后你根据拼装的效果给他们评分。你还让他们多试几次，挑出最好的方案。这样，两个助手互相学习，既节省时间，又能拼装得更好。这个过程就像论文中的算法，让小模型学得更快、更准，同时减少了繁琐的计算。
ELI14 Explained like you're 14
想象你在学校里参加拼图比赛。有两个队友：一个是经验丰富的老师（大模型），另一个是新手学生（小模型）。老师可以给你一些提示，但每次都要花很多时间。学生可以自己试着拼，但有时候拼错了。为了赢得比赛，你让老师先帮忙画出一部分拼图（提示），学生根据提示继续拼。然后，你根据拼图的完整度给他们打分，鼓励学生多试几次，挑出最好的拼图。老师和学生不断互相学习，学生变得越来越厉害，老师也能更快地指导。这个方法就像论文里的CoDistill-GRPO，让小模型学得更快、更准，还能节省时间和精力。
Abstract
Group Relative Policy Optimization (GRPO) has emerged as a powerful algorithm for improving the reasoning capabilities of language models, but often fails to improve small models due to sparse rewards on difficult tasks. Existing works mitigate this issue by leveraging a larger model, either to provide hints for rollouts or to provide dense reward signals through knowledge distillation (KD). However, this assumes the existence of such an oracle, and training one can significantly increase total training time. In this work, we propose CoDistill-GRPO, a co-distillation algorithm that simultaneously trains a large and a small model by maximizing carefully designed GRPO objectives. The two models learn from each other: the small model uses an on-policy KD reward to learn from the large model's distribution, while the large model is updated using rollouts generated by the small model with importance reweighting, reducing the computational overhead of rollout generation. We show that CoDistill-GRPO substantially improves small model performance over standard GRPO on mathematical benchmarks across both Qwen and Llama models. Specifically, with Qwen2.5-Math-1.5B, we observe an accuracy increase of over 11.6 percentage points over the base model and an additional 6.0 percentage points over GRPO on the Minerva dataset. Interestingly, the larger model (Qwen2.5-Math-7B) trained with CoDistill-GRPO nearly matches standard GRPO performance despite training on small-model rollouts. This highlights CoDistill-GRPO as a cost-effective alternative to GRPO for larger models, yielding an approximate 18% speedup, which may be of independent interest.