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; Xipeng Qiu

Source-linked AI summary

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, Xipeng Qiu

arXiv:2604.11297v1Apr 13, 2026cs.LGcs.AIcs.CL

TL;DR

LLM reinforcement learning can reduce sampling diversity by repeatedly generating similar errors, while current-policy entropy regularization does not explicitly target recurrent failures. MEDS stores lightweight representations of past rollouts, clusters recurring error patterns, and penalizes prevalent clusters; across benchmarks and models, it improves performance and exploration diversity. The study’s main limitation is that its logit utilization uses relatively simple aggregation strategies.

Problem
On-policy LLM reinforcement learning can collapse into repetitive erroneous behaviors, and entropy regularization does not explicitly discourage recurrent failure patterns.
Method
MEDS reuses layer-wise logits as reasoning representations, clusters historical responses, and dynamically increases penalties for more recurrent error patterns.
Results
MEDS improves pass@1 and pass@128 across three models, with gains of up to 4.13 points and 4.37 points, respectively, while increasing exploration diversity.
Takeaways & Limitations
Historical trajectory information can be incorporated into reward design to suppress repeated error modes while improving reasoning performance and exploration diversity.
Takeaways & Limitations
The study uses relatively simple logit aggregation and leaves the benefits of more sophisticated aggregation strategies unclear.

Abstract

from arXiv · show

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.

1 Introduction

MEDS addresses recurrent reasoning failures in LLM reinforcement learning by using historical error patterns to shape rewards, rather than relying only on current-policy randomness. It combines lightweight logit representations, online clustering, and dynamic penalties, with experiments reporting improved performance and exploration diversity.

On-policy optimization can collapse policies into narrow, repetitive error patterns that weaken exploration and reinforce erroneous reasoning trajectories.
MEDS dynamically records historical error patterns and applies stronger penalties to recurrent failure paths, encouraging escape from local optima.
MEDS improves pass@1 and pass@128 across three models, with gains of up to 4.13 points and 4.37 points, respectively.
Analyses show increased exploration diversity, agreement between logit-based and LLM-based clustering, and a positive correlation between clustering quality and downstream performance.
MEDS explicitly incorporates historical error patterns into reward modeling to recognize and avoid recurrent failure behaviors.
Layer-wise logits provide a compact response-similarity representation with little additional computational overhead.

2 Related Work

Prior approaches improve reasoning diversity through regularization, current-batch clustering, historical outcome frequencies, or external reward models. MEDS instead uses layer-wise logits to track reasoning patterns and cluster them without requiring extra models.

RLVR training can over-optimize dominant solution patterns, reducing reasoning diversity; KL and entropy regularization constrain policy shifts or encourage randomness.
Some methods cluster current rollouts or reward rare correct trajectories, but they do not track solution patterns over long-term training history.
Historical outcome-frequency rewards add exploration information but lack fine-grained cross-step reasoning behavior tracking.
External models can adapt rewards from historical performance, but their additional computational cost limits scalability.
MEDS clusters layer-wise logits to capture shared reasoning patterns without requiring extra models.
MEDS uses historical responses, HDBSCAN, and cluster-based penalties to adjust rewards according to recurring failure patterns.

3 Method

MEDS augments reward-based optimization with a memory of historically repeated errors, using response representations to cluster failures and shape penalties. The method reuses layer-wise logits, applies density-based clustering, and adjusts rewards according to cluster prevalence.

Theoretical Benefits of Penalizing Repeated Errors: MEDS augments the task-specific reward with a repetition penalty based on historically sampled errors.The indicator increases with repeated-error frequency, while the penalty coefficient controls its effect on the updated policy.
Theoretical Benefits of Penalizing Repeated Errors: The indicator function acts as reward memory by recording historical sampling behavior and adjusting scores for recurrent failures.Repeated errors are penalized because they indicate higher probability mass under the current policy.
Logic Feature Extraction: MEDS reuses layer-wise logits from final-answer token positions as compact representations of model reasoning trajectories.The method selects the relevant token logit at each layer and aggregates logits, using later-layer features to represent responses.
Cluster-based Reward Shaping: For each prompt, MEDS stores historical response representations and applies HDBSCAN to obtain dynamically determined error clusters.The historical representation set is maintained across training epochs, and the number of clusters is determined dynamically.
Cluster-based Reward Shaping: The cluster-size indicator c(ỹ) = log(|C_k| + 1) increases monotonically with cluster prevalence, while α and β control penalty strength and its upper bound.This transformation preserves the ordering used in the theoretical analysis and makes more prevalent clusters receive stronger shaping.

4 Experiments

Experiments across three models and five mathematical benchmarks show that MEDS consistently achieves strong performance while improving exploration diversity and clustering quality with modest overhead.

Experimental Setup: Experiments use three base models, five mathematical benchmarks, and comparisons against GRPO, DAPO, entropy regularization, and the untuned base model.The benchmarks are AIME24, AMC23, MATH500, Minerva, and OlympiadBench.
Main Results: MEDS consistently achieves the best overall performance across model scales and outperforms the base model and DAPO in most pass@k settings.Figure 4 reports pass@k for k ∈ {1, 2, 4, 8, 16, 32, 64, 128}.
Main Results: MEDS introduces only a small computational overhead compared with DAPO.On Qwen2.5-Math-7B, MEDS takes 8.46 minutes for 50 steps and 9.73 minutes for 100 steps, versus DAPO’s 8.00 and 8.95 minutes.
Impact on Exploration Behavior: MEDS produces higher Within-Step and Across-Step Diversity scores and lower Top-1 Eigen Ratios than the baseline during training.The lower eigen ratio indicates a more uniform spread of logits across representation-space directions.
Logits can Reflect Reasoning Patterns: Layer-wise logit trajectories tend to be similar for responses with similar reasoning patterns and distinguish responses following different reasoning paths.The case study indicates that later-layer logits capture reasoning-pattern differences beyond final answer tokens.
Logits Provide Correct Clustering Signals: Logit-based clustering has a 61.2% correlation with Claude annotations, and using the last 14 layers yields the best downstream performance among feature constructions.Every clustering variant outperforms DAPO, whereas excluding logits performs worse than DAPO.

5 Conclusion

The paper introduces MEDS to identify and penalize recurrent failure patterns during LLM reinforcement learning. It reports improved reasoning performance and exploration diversity, supporting historical trajectory information as a reward-design signal.

Conclusion: MEDS dynamically penalizes recurrent failure patterns by clustering historical responses represented with layer-wise logits.The framework is presented as memory-aware reward shaping for LLM reasoning.
Conclusion: Empirical results show that MEDS consistently improves reasoning performance and exploration diversity across multiple benchmarks.The conclusion highlights historical trajectory information as a direction for reward design in LLM reinforcement learning.

6 Limitations

The study’s main limitation is that its logit-utilization methods use relatively simple aggregation strategies, leaving the value of more sophisticated alternatives unresolved.

Limitations: The explored logit methods use relatively simple aggregation functions, so the benefits of more complex aggregation strategies remain unclear.This is identified as the study’s main limitation.

A Implementation Details

Training and evaluation problems use the standardized Qwen-Math prompt format, which instructs step-by-step reasoning and boxed final answers.

Prompt Format: All training and evaluation problems are processed using the standardized Qwen-Math template.The template formats mathematical problem-solving prompts consistently.
Prompt Format: The template instructs the assistant to reason step by step and place the final answer within boxed notation.The system prompt contains the explicit formatting instruction.

A.2 Training Hyperparameters and Objectives

The section specifies baseline configurations and states that training objectives are defined for comparison methods.

Training configurations: GRPO follows the official DeepSeekMath configuration with clipping ratio 𝜖= 0.2 and KL penalty coefficient 𝛽= 10−3.
Training configurations: DAPO uses decoupled clipping with asymmetric thresholds 𝜀low = 0.2 and 𝜀high = 0.28, while removing the KL penalty term.
Training configurations: GRPO w/ Entropy Adv augments the group-relative advantage with a per-token entropy bonus weighted by 𝛼.
Training objectives: The section introduces the training objectives and definitions for the probability ratio, group-relative advantage, entropy-augmented advantage, and per-token entropy.

B Theoratical Proof

The proof analyzes a reward modified by a repetition penalty and concludes that the updated policy has no lower expected task reward under the stated assumptions.

Setup: Theorem 2 compares the task reward r(x, y) with the penalized reward r(x, y) −𝜆c(y), where c(y) measures repeated errors.
Setup: The penalty indicator c(y) is positive when a response corresponds to an error repeatedly sampled in the past.
Assumption: The proof assumes responses with larger repetition counts are not higher-scoring under the task reward.
Policy updates: Using KL-regularized one-step updates, the proof defines q1 under r(x, y) and q2 under r(x, y) −𝜆c(y), with 𝜆> 0.
Covariance argument: Because both the task reward and the weighting term decrease with repetition count, the proof obtains Cov𝑞1(𝑟(x, Y), w(Y)) ≥0.
Conclusion: The resulting inequality is Jx(q2) ≥ Jx(q1) for all x.

C Full Results of Case Study

The case study presents a good-number problem alongside a heatmap visualization whose logit values are displayed.

Problem: The case study asks for the smallest positive integer having three distinct positive divisors summing to 2022.
Visualization: The displayed labels include A1, A2, B1, B2, and B3.
Visualization: Figure 7 is described as a heatmap from Figure 6 with logit values.

C.2 Custer A Responses

Cluster A responses include reasoning about numbers with exactly three divisors, computational search code, and two reported candidate answers for the divisor-sum problem.

Mathematical reasoning: The divisor condition is reduced to 1 + p + p^2 = 2022, equivalently p^2 + p −2021 = 0.
Verification: The response checks nearby values including 89, 83, 43, and 41, each of which is reported not to satisfy the divisor-sum equation.
Computational verification: The accompanying code iterates over candidate values, tests primality, and searches for the smallest number satisfying the target sum.
Reported results: One reported result gives 1342 as the smallest good number.
Mathematical reasoning: A number with exactly three distinct positive divisors must be the square of a prime, with divisors 1, p, and p^2.
Mathematical reasoning: The solution discusses solving the quadratic using the positive root and checking whether the resulting value of p is an integer prime.
Reported results: Another reported result gives 1351 as the smallest number with exactly three distinct positive divisors summing to 2022.

C.3 Cluster B Responses

The responses explore divisor structures and systematic computation for identifying the smallest integer meeting the target sum, with conflicting reported results.

Divisor-structure analysis: The solutions examine prime-square and composite forms to characterize numbers with three distinct divisors summing to 2022.They consider n = p^2 and n = p^2·q, alongside direct divisor enumeration.
Systematic search: A brute-force approach iterates through candidate integers, computes their divisors, and checks triples whose sum equals 2022.The implementation removes 1 where appropriate and tests combinations of distinct divisors.
Diversity evaluation: The accompanying prompts evaluate answer diversity by comparing reasoning strategies, conclusions, correctness, and quality within or across groups.One scoring rule treats substantively different reasoning ideas as new paths while excluding wording or formatting changes.