The Past Is Not Past: Memory-Enhanced Dynamic Reward Shaping

Yang Liu; Enxi Wang; Yufei Gao; Weixin Zhang; Bo Wang; Zhiyuan Zeng; Yikai Zhang; Yining Zheng

cs.LGApr 13, 2026

The Past Is Not Past: Memory-Enhanced Dynamic Reward Shaping

Authors: Yang Liu, Enxi Wang, Yufei Gao, Weixin Zhang, Bo Wang, Zhiyuan Zeng, Yikai Zhang, Yining Zheng,

Organizations: Fudan University · Shanghai Innovation Institute

Abstract

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

Panel of Figure 1 — 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 $\tilde{y}$ and extract a logic feature vector $f(\tilde{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 $G_{x}$ and apply HDBSCAN to group historical responses with similar logic features. (C) Reward Shaping: For a response assigned to cluster $C_{k}$ , we compute a penalty $\min(\alpha\log(|C_{k}|+1),\beta)$ proportional to the cluster size, and adjust the reward as $\tilde{r}(x,\tilde{y})=r(x,\tilde{y})-\text{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
Method	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\in\{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 $\mathbf{A_{1},A_{2}}$ and Cluster B containing $\mathbf{B_{1},B_{2},B_{3}}$ by our clustering method. Although $\mathbf{A_{1},B_{1}}$ and $\mathbf{B_{2}}$ 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
Method	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.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Heatmap in Figure 6 with logit values.

RL for Language Models

Reinforcement learning for language models uses feedback-driven policy updates to improve reasoning, alignment, and interactive behavior.

See all 410 papers

DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay
Mar 17, 2026Long Li, Zhijian Zhou, Tianyi Wang +7Behavioral Diversity in RL Reinforcement Learning
Griffith University, Australia · Fudan University, China · Shanghai Innovation Institute, China +3
DSDR: Dual-Scale Diversity Regularization for Exploration in LLM Reasoning
Feb 23, 2026Zhongwei Wan, Yun Shen, Zhihao Dou +9Reinforcement Learning RL for Language Model Reasoning
The Ohio State University, USA · Case Western Reserve University, USA · The Chinese University of Hong Kong, Hong Kong +4
Reasoning or Memorization? Direction-Aware Diversity Exploration in LLM Reinforcement Learning
Jun 9, 2026Jiangnan Xia, Yucheng Shi, Yu Yang +3Memorization in Language Models RL for Language Model Reasoning
University of Georgia · Tencent AI Lab · The Education University of Hong Kong +1

Mar 17, 2026J/K move · Enter open · S save

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

Long Li, Zhijian Zhou, Tianyi Wang, Weidi Xu+6

Griffith University, Australia · Fudan University, China · Shanghai Innovation Institute, China+3

Behavioral Diversity in RL Reinforcement Learning RL for Language Model Reasoning Experience Replay Reinforcement Learning with Verifiable Rewards

Open paper AskPDF arXiv

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.

Abstract

RL for Language Models

Explore similar work

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

DSDR: Dual-Scale Diversity Regularization for Exploration in LLM Reasoning

Reasoning or Memorization? Direction-Aware Diversity Exploration in LLM Reinforcement Learning