DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay
作者: Long Li, Zhijian Zhou, Tianyi Wang, Weidi Xu, Zuming Huang, Wei Chu, Zhe Wang, Shirui Pan, Chao Qu, Yuan Qi
分类: cs.LG, cs.AI
发布日期: 2026-03-17
备注: 14 pages, 3 figures
💡 一句话要点
DyJR通过动态Jensen-Shannon回放,在强化学习中保持多样性并提升大语言模型推理能力。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 强化学习 大语言模型 经验回放 多样性保持 Jensen-Shannon散度
📋 核心要点
- 现有强化学习方法在增强大语言模型推理能力时,存在样本效率低和易模式崩溃的问题。
- DyJR通过动态Jensen-Shannon回放,利用时间敏感的动态缓冲区和分布约束来保持多样性。
- 实验表明,DyJR在数学推理和Text-to-SQL任务上显著优于现有方法,并保持了训练效率。
📝 摘要(中文)
本文提出了一种名为动态Jensen-Shannon回放(DyJR)的正则化框架,旨在解决强化学习增强大语言模型推理能力时,on-policy算法样本效率低和经验回放方法易导致模式崩溃的问题。DyJR使用来自最近轨迹的动态参考分布,包含两个创新点:一是时间敏感的动态缓冲区,采用FIFO和自适应大小来保留时间上接近的样本,与模型演化同步;二是Jensen-Shannon散度正则化,用分布约束代替直接梯度更新,防止多样性崩溃。在数学推理和Text-to-SQL基准测试上的实验表明,DyJR显著优于GRPO以及RLEP和Ex-GRPO等基线,同时保持与原始GRPO相当的训练效率。通过Rank-$k$ token概率演化的角度,展示了DyJR增强了多样性并减轻了对Rank-1 token的过度依赖。
🔬 方法详解
问题定义:现有基于强化学习的大语言模型训练方法,如GRPO,是on-policy算法,样本效率低,丢弃了大量的历史轨迹数据。虽然经验回放可以重复利用历史数据,但直接使用历史数据进行策略更新容易导致过拟合,从而造成模式崩溃,降低生成结果的多样性。因此,如何在利用历史数据的同时,保持生成结果的多样性是一个关键问题。
核心思路:DyJR的核心思路是,与其直接利用历史数据进行梯度更新以追求准确性,不如利用历史数据来维持生成结果的多样性。具体来说,DyJR使用一个动态的参考分布,该分布来自最近的轨迹,并利用Jensen-Shannon散度来约束策略的更新,从而防止多样性崩溃。
技术框架:DyJR包含两个主要模块:时间敏感的动态缓冲区和Jensen-Shannon散度正则化。时间敏感的动态缓冲区使用FIFO策略和自适应大小来保留时间上接近的样本,确保缓冲区中的数据与当前模型的演化状态相匹配。Jensen-Shannon散度正则化则通过约束策略的输出分布与动态参考分布之间的距离,来防止策略过度拟合历史数据,从而保持生成结果的多样性。整体流程是,在每个训练迭代中,首先从环境中收集新的轨迹数据,然后将这些数据添加到动态缓冲区中。接着,从缓冲区中采样一部分数据,并计算当前策略的输出分布与动态参考分布之间的Jensen-Shannon散度。最后,将该散度作为正则化项添加到损失函数中,用于更新策略。
关键创新:DyJR的关键创新在于,它将经验回放的重点从追求准确性转移到维持多样性。传统的经验回放方法通常直接使用历史数据进行梯度更新,而DyJR则使用历史数据来构建一个动态的参考分布,并利用Jensen-Shannon散度来约束策略的更新。这种方法可以有效地防止策略过度拟合历史数据,从而保持生成结果的多样性。与现有方法的本质区别在于,DyJR不是直接利用历史数据来提升性能,而是利用历史数据来正则化策略,从而提升性能和多样性。
关键设计:时间敏感的动态缓冲区的大小是自适应的,它会根据模型的训练进度动态调整。Jensen-Shannon散度正则化的系数是一个超参数,需要根据具体的任务进行调整。损失函数包含两部分:一个是标准的强化学习损失函数,用于提升性能;另一个是Jensen-Shannon散度正则化项,用于保持多样性。网络结构与GRPO保持一致,没有引入额外的复杂性。
🖼️ 关键图片
📊 实验亮点
实验结果表明,DyJR在数学推理和Text-to-SQL任务上显著优于GRPO以及RLEP和Ex-GRPO等基线。例如,在数学推理任务上,DyJR的性能比GRPO提升了超过10%。此外,DyJR还保持了与原始GRPO相当的训练效率,表明其具有良好的实用性。通过Rank-$k$ token概率演化的分析,进一步验证了DyJR能够有效增强生成结果的多样性。
🎯 应用场景
DyJR可应用于各种需要大语言模型进行推理和生成的任务,例如数学问题求解、代码生成、文本到SQL查询等。通过保持生成结果的多样性,DyJR可以提升模型的泛化能力和鲁棒性,使其在实际应用中更加可靠。未来,DyJR可以进一步扩展到其他类型的强化学习算法和任务中,例如多智能体强化学习和机器人控制。
📄 摘要(原文)
While Reinforcement Learning (RL) enhances Large Language Model reasoning, on-policy algorithms like GRPO are sample-inefficient as they discard past rollouts. Existing experience replay methods address this by reusing accurate samples for direct policy updates, but this often incurs high computational costs and causes mode collapse via overfitting. We argue that historical data should prioritize sustaining diversity rather than simply reinforcing accuracy. To this end, we propose Dynamic Jensen-Shannon Replay (DyJR), a simple yet effective regularization framework using a dynamic reference distribution from recent trajectories. DyJR introduces two innovations: (1) A Time-Sensitive Dynamic Buffer that uses FIFO and adaptive sizing to retain only temporally proximal samples, synchronizing with model evolution; and (2) Jensen-Shannon Divergence Regularization, which replaces direct gradient updates with a distributional constraint to prevent diversity collapse. Experiments on mathematical reasoning and Text-to-SQL benchmarks demonstrate that DyJR significantly outperforms GRPO as well as baselines such as RLEP and Ex-GRPO, while maintaining training efficiency comparable to the original GRPO. Furthermore, from the perspective of Rank-$k$ token probability evolution, we show that DyJR enhances diversity and mitigates over-reliance on Rank-1 tokens, elucidating how specific sub-modules of DyJR influence the training dynamics.