Despite the success of reinforcement learning for large language models, a common failure mode is reduced sampling diversity, where the policy repeatedly generates similar erroneous behaviors. Classical entropy regularization encourages randomness under the current policy, but does not explicitly discourage recurrent failure patterns across rollouts. We propose MEDS, a Memory-Enhanced Dynamic reward Shaping framework that incorporates historical behavioral signals into reward design. By storing and leveraging intermediate model representations, we capture features of past rollouts and use density-based clustering to identify frequently recurring error patterns. Rollouts assigned to more prevalent error clusters are penalized more heavily, encouraging broader exploration while reducing repeated mistakes. Across five datasets and three base models, MEDS consistently improves average performance over existing baselines, achieving gains of up to 4.13 pass@1 points and 4.37 pass@128 points. Additional analyses using both LLM-based annotations and quantitative diversity metrics show that MEDS increases behavioral diversity during sampling.
Figures & tables
Figure 1: Illustration of error collapse in on-policy optimization and its effect on output diversity. (a) During on-policy optimization, the policy evolves from a high-entropy distribution into a narrow, stable error basin. (b) Different errors share the same failure mode and become progressively reinforced during training. (c) Comparison of output diversity between DAPO and our method, with detailed description provided in Section 4.3 .
Figure 2: Overview of MEDS . The process consists of three stages: (A) Logic Feature Extraction: For an input x , we sample a response y~ and extract a logic feature vector f(y~) from the layer-wise logits of the first token y∗ in the final answer, using the latter half of the Transformer layers to capture reasoning patterns. (B) Memory-based Clustering: We maintain a per-prompt error memory Gx and apply HDBSCAN to group historical responses with similar logic features. (C) Reward Shaping: For a response assigned to cluster Ck , we compute a penalty min(αlog(∣Ck∣+1),β) proportional to the cluster size, and adjust the reward as r~(x,y~)=r(x,y~)−penalty , which is then used to update the policy.
Figure 3: The logits aggregation process, exemplified by a 3-layer model feature graph. Each layer contributes a scalar logit at the final answer token position, whose value is determined by the activated computational circuits. The left and middle graphs activate nearly identical circuits, resulting in nearly identical 3-layer logit vectors. In contrast, the right graph exhibits a substantially different circuit structure, yielding divergent logits.
Method
AIME24
AMC23
MATH500
Minerva
OlympiadBench
Average
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
Qwen3-1.7B
22.58
66.67
58.57
87.50
80.86
95.20
37.82
63.6
41.97
64.15
48.36
75.42
+GRPO
18.20
46.67
51.58
92.50
79.58
95.60
38.62
65.81
40.19
70.37
45.63
74.19
+DAPO
13.46
63.33
57.65
95.00
80.25
96.60
39.45
69.49
42.20
77.04
46.6
80.29
+GRPO w/ Entropy Adv
12.50
50.00
51.80
95.00
78.72
94.00
38.91
62.87
38.47
66.96
44.08
73.77
+ MEDS
23.65
70.00
65.51
95.00
84.66
97.40
41.77
68.75
46.86
77.63
52.49
81.76
Table 1: Main results on five mathematical reasoning benchmarks. Best and second-best are marked with bold and underline ; MEDS rows are highlighted.
Figure 4: Pass@ k across k∈{1,2,4,8,16,32,64,128} on five math benchmarks for three model scales. Each row corresponds to one model; each column to one dataset. Ours consistently matches or outperforms Base, GRPO, DAPO, and GRPO with entropy advantage across settings.
Figure 5: (Left) Top-1 eigen ratio over training steps for DAPO and Ours. Our method consistently yields lower ratios at later steps, indicating less concentration on a single dominant error direction. (Right) Correlation of different logit aggregation strategies on Qwen2.5-Math-7B. Single cluster assigns all logits to one cluster; 28 and 14 use direct aggregation over all 28 layers or the last 14 layers; 28–diff and 14–diff use aggregation over layer-wise differences between adjacent logits.
Figure 6: A case study. We manually annotate the reasoning structures of different responses. The responses are grouped into Cluster A containing A1,A2 and Cluster B containing B1,B2,B3 by our clustering method. Although A1,B1 and B2 produce the same incorrect answer 1342, their reasoning structures are different. A similar phenomenon can also be observed for the responses within Cluster B. We further visualize the logits using a heatmap. See Appendix C for the complete case study and the corresponding logit values.
Method
AIME24
AMC
MATH500
Minerva
OlympiadBench
Average
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
pass@1
pass@128
DAPO
32.27
70.00
72.71
97.50
85.61
96.00
43.01
75.74
44.43
75.00
55.61
82.85
MEDS -single cluster
29.70
70.00
73.60
100.00
86.60
96.80
41.20
71.69
44.60
74.52
55.14
82.60
MEDS -28-diff
31.88
70.00
72.34
97.50
85.80
96.00
43.38
73.53
45.48
74.37
55.78
82.28
MEDS -14-diff
32.40
73.33
71.64
100.00
86.40
96.60
43.01
72.43
44.15
74.52
55.52
83.38
MEDS -28
35.63
70.00
73.13
97.50
86.20
96.40
41.90
73.16
44.00
73.78
56.17
82.17
Table 2: Comparison of different clustering feature constructions on Qwen2.5-Math-7B, including assigning all logits to a single cluster ( single cluster ); using adjacent logit differences over 28 layers ( 28-diff ) or the last 14 layers ( 14-diff ); using logits from all 28 layers; and using only the last 14 layers. The first row corresponds to the vanilla DAPO baseline.
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 can incur additional storage costs and concentrate the policy on a narrow set of solutions. We argue that historical data should prioritize sustaining local policy plasticity rather than simply reinforcing accuracy. We define local policy plasticity as a persistent ability to remain backtrackable and switchable with respect to recent successful trajectories: throughout training, the policy maintains non-negligible support on recently reachable alternatives, so it can smoothly reallocate probability mass when needed without being locked into a single dominant path. 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 replay policy-gradient updates with a regularization signal that discourages premature over-commitment to a single Rank-1 trajectory. Experiments on mathematical reasoning, Text-to-SQL, and coding benchmarks demonstrate that DyJR improves over GRPO as well as baselines such as RLEP and Ex-GRPO, while maintaining training efficiency comparable to the original GRPO. Furthermore, Rank-k token probability evolution and semantic-mode analyses show reduced concentration and greater retention of alternative continuations. Matched admission-objective comparisons and repeated runs provide additional evidence for the roles of DyJR's components.
Long Li, Zhijian Zhou, Tianyi Wang +7
Griffith University, Australia · Fudan University, China · Shanghai Innovation Institute, China +3
Reinforcement learning with verifiers (RLVR) is a central paradigm for improving large language model (LLM) reasoning, yet existing methods often suffer from limited exploration. Policies tend to collapse onto a few reasoning patterns and prematurely stop deep exploration, while conventional entropy regularization introduces only local stochasticity and fails to induce meaningful path-level diversity, leading to weak and unstable learning signals in group-based policy optimization. We propose DSDR, a Dual-Scale Diversity Regularization reinforcement learning framework that decomposes diversity in LLM reasoning into global and coupling components. Globally, DSDR promotes diversity among correct reasoning trajectories to explore distinct solution modes. Locally, it applies a length-invariant, token-level entropy regularization restricted to correct trajectories, preventing entropy collapse within each mode while preserving correctness. The two scales are coupled through a global-to-local allocation mechanism that emphasizes local regularization for more distinctive correct trajectories. We provide theoretical support showing that DSDR preserves optimal correctness under bounded regularization, sustains informative learning signals in group-based optimization, and yields a principled global-to-local coupling rule. Experiments on multiple reasoning benchmarks demonstrate consistent improvements in accuracy and pass@k, highlighting the importance of dual-scale diversity for deep exploration in RLVR. Code is available at https://github.com/SUSTechBruce/DSDR.
Zhongwei Wan, Yun Shen, Zhihao Dou +9
The Ohio State University, USA · Case Western Reserve University, USA · The Chinese University of Hong Kong, Hong Kong +4
Reinforcement learning has become a key paradigm for eliciting reasoning abilities in large language models, where exploration is crucial for discovering effective solution trajectories. Existing exploration methods typically encourage diversity in semantic or gradient spaces, without distinguishing what drives this diversity. A trajectory may appear novel because it follows a new reasoning process, or because it varies memorized patterns and shortcuts. Rewarding both cases equally may steer exploration toward memorization rather than genuine reasoning improvement. In this paper, we propose DiRL, a Direction-Aware Reinforcement Learning framework that anchors exploration to an internal reasoning-memorization direction of the policy. Specifically, DiRL extracts this direction from model representations, constructs direction-weighted gradient features to characterize rollout updates, and shapes rewards to amplify reasoning-aligned exploration while suppressing memorization-aligned variations. DiRL integrates seamlessly into standard Group Relative Policy Optimization (GRPO). Extensive experiments on mathematical and general reasoning benchmarks demonstrate the effectiveness of DiRL, showing significant improvements over various existing exploration methods.
Jiangnan Xia, Yucheng Shi, Yu Yang +3
University of Georgia · Tencent AI Lab · The Education University of Hong Kong +1