Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Abstract

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.

DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Abstract

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- 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.
Paper Structure (35 sections, 11 equations, 3 figures, 5 tables, 1 algorithm)

This paper contains 35 sections, 11 equations, 3 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. RL Backbone: GRPO
  4. Forward KL Divergence
  5. Jensen-Shannon Divergence
  6. Method
  7. Dynamic Reference Construction
  8. Dynamic Replay Buffer
  9. Bias-Aware Adaptive Data Selection
  10. Mitigating Early Diversity Collapse
  11. JS Regularization Implementation
  12. Joint Optimization Objective
  13. Experiments
  14. Setting
  15. Baseline
  16. ...and 20 more sections

Figures (3)

  • Figure 1: We track the evolution of the Qwen3-4B-Base model throughout the training process. The recorded metrics include the entropy of the output distribution, the average probability of rank-$k$ tokens ($\bar{P}_{\text{rank-}k}$), and the performance fluctuations on the evaluation datasets.
  • Figure 2: Comparison of rank-$k$ trends across the transition from RLEP to DyJR.
  • Figure 3: (a) Pass@$k$ performance on the Beyond AIME dataset across varying sampling budgets $k$. (b) and (c) Computational latency comparison between GRPO (w/o KL) and DyJR during the Rollout phase and the Actor Update phase (encompassing log-probability calculation, forward passes, and gradient updates), respectively.