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

1 Introduction

• On-policy optimization can collapse into repetitive error patterns, and entropy regularization may be insufficient to prevent repeated mistakes in LLM reasoning.Once the policy enters an error pattern, it can keep reproducing the mistake through slightly varied reasoning expressions despite continued optimization.Source: §1 · pp. 1–2
• MEDS dynamically records historical errors and imposes stronger penalties on recurrent failure paths, rather than relying only on randomness under the current policy.The paper also provides theoretical analysis that additional penalties on repeated errors can improve performance.Source: §1 · p. 2
• MEDS reuses layer-wise logits as a low-overhead proxy for reasoning trajectories and applies online HDBSCAN clustering to adapt penalties to recurring error patterns.The representation-level similarity criterion uses logits already produced during the forward pass, adding almost no computational overhead.Source: §1 · pp. 2–3
• MEDS improves pass@1 and pass@128 across all three models, with gains up to 4.13 points in pass@1 and 4.37 points in pass@128.On Qwen3-8B, pass@128 rises from 70.81 to 82.67 on a single dataset, a relative gain of 17%.Source: §1 · pp. 2–3
• Analyses show that MEDS increases exploration diversity, logit-based clustering aligns with LLM-based semantic clustering, and clustering quality positively correlates with downstream performance.The paper reports these findings across multiple benchmarks and models.Source: §1 · pp. 2–3

2.1 Reinforcement Learning with Verifiable Rewards

• RLVR verifiers use exact-match [15] [1], sandbox execution [16], video QA [17], image similarity [18], or nDCG@10 [19], but score current samples independently.Source: §2.1 · p. 3

2.2 Reward Shaping for Diversity Exploration

• RLVR can collapse reasoning diversity [20] [21]; prior remedies use KL regularization [4] [22], entropy [10], or rare-behavior rewards [23] [11], but lack long-term tracking.Hu et al. [23] cluster rollouts by strategy, while Wan et al. [11] promote diversity among correct trajectories with token-level entropy regularization.Source: §2.2 · p. 3
• Historical UCB rewards outcome frequencies [24] but lacks fine-grained tracking; external-model reward adaptation [25] [26] [27] [28] [29] adds cost, whereas MEDS uses layer-wise logits.Source: §2.2 · p. 3

2.3 Inner Thoughts of LLMs

• Logit-based studies localize factual knowledge [30], analyze reasoning with sparse autoencoders and interventions [32] [33] [34], detect errors [35], and provide fine-grained supervision [36].Source: §2.3 · p. 3
• MEDS clusters logits to identify shared patterns across reasoning trajectories without requiring extra models.Source: §2.3 · p. 3

3.1 Preliminaries

• The policy generates responses and maximizes expected reward under a KL-regularized objective, with a reference policy constraining updates and a coefficient controlling constraint strength.Source: §3.1 · p. 4

3.2 Theoretical Benefits of Penalizing Repeated Errors

• Repeated-error penalties increase with past occurrences; the analysis argues that penalizing errors with higher policy probability mass reallocates probability and makes expected return non-decreasing under optimization.The repeated-error indicator acts as memory for the reward model and adjusts scores using historical policy sampling.Source: §3.2 · pp. 4–5

3.3 Logic Feature Extraction

• MEDS aggregates final-answer-token logits from latter-half layers into response features, using internal activations to capture reasoning similarity beyond surface-level string matching.The description uses direct aggregation and notes that alternative aggregation approaches are compared in Section 4.6.Source: §3.3 · p. 5

3.4 Cluster-based Reward Shaping

• For each prompt, MEDS clusters historical response features with HDBSCAN, then assigns responses a log cluster-size penalty and adjusts reward with strength and upper-bound hyperparameters.The cluster count is determined dynamically, and the indicator is c(ỹ) = log(|C_k| + 1).Source: §3.4 · pp. 5–6

4.1 Experimental Setup

• MEDS is trained on Qwen3-1.7B, Qwen3-8B, and Qwen2.5-Math-7B using DAPO-Math-17K merged with MATH levels 3–5, and compared with GRPO, DAPO, an entropy-regularization exploration approach, and untuned base models.Evaluation uses AIME24, AMC23, MATH500, Minerva, and OlympiadBench.Source: §4.1 · pp. 6–7

4.2 Main Results

• Across three base-model scales and five mathematical-reasoning benchmarks, the method achieves the best overall results; its pass@k curves match or outperform baselines across settings and perform best in most cases.Training takes 9.73 minutes for 100 steps, compared with DAPO times of 8.00 and 8.95 minutes.Source: §4.2 · pp. 7–8, Figure 4 · p. 7

4.3 Impact on Exploration Behavior

• MEDS scores higher than the baseline on Within-Step Diversity and Across-Step Diversity, and has a lower Top-1 Eigen Ratio; diversity generally declines as training progresses.A lower Top-1 Eigen Ratio indicates a more uniform spread of logits across representation-space directions.Source: §4.3 · p. 8

4.4 Logits can Reflect Reasoning Patterns

• Responses with similar reasoning patterns tend to have similar layer-wise logit trajectories, while different reasoning paths remain distinguishable, including responses that share a final answer.The analysis reports limited logit variation in the first 20 layers and stronger pattern-correlated changes in later layers.Source: §4.4 · p. 9

4.5 Logits Provide Correct Clustering Signals

• Logit-based clustering has a 61.2% correlation with Claude-Haiku-4.5 annotations of logical similarity.The evaluation samples 800 negative responses and uses 11 error categories distilled through manual inspection.Source: §4.5 · pp. 9–10

4.6 Clustering Quality Matters

• Using logits from the last 14 layers yields the best downstream performance; every clustering variant beats DAPO, while excluding logits performs worse than DAPO.The ranking of clustering consistency with Claude annotations closely mirrors the ranking of downstream performance.Source: §4.6 · p. 10, Table 2 · p. 10

5 Conclusion

• MEDS reuses layer-wise logits as lightweight trajectory representations, clustering historical responses to identify and suppress repeated error modes.Its dynamic reward shaping penalizes recurrent failure patterns.Source: §5 · p. 10
• The paper provides theoretical analysis demonstrating the effectiveness of MEDS.Source: §5 · p. 10
• Across multiple benchmarks, MEDS consistently improves reasoning performance and exploration diversity.The conclusion emphasizes incorporating historical trajectory information into reward design as a direction for improving exploration in LLM reinforcement learning.Source: §5 · p. 10

6 Limitations

• MEDS explores only relatively simple logit aggregation, and whether more complex strategies could improve performance remains unclear.Source: §6 · p. 11

A.1 Qwen-Math Template

• All training and evaluation problems are processed using the Qwen-Math template for standardized mathematical problem-solving prompts.Source: §A.1 · p. 15
• The template instructs the model to reason step by step and put its final answer within \boxed{}.Source: §A.1 · p. 15

• All methods use Adam at 1 × 10−6, weight decay 0.1, gradient clipping 1.0, 10 warmup steps, and PPO mini-batches of 16; GRPO uses ε=0.2 and β=10−3.GRPO follows the official DeepSeekMath configuration [1].Source: §A.2 · p. 15
• DAPO uses decoupled clipping with εlow = 0.2 and εhigh = 0.28, and removes the KL penalty term.Source: §A.2 · p. 15
• GRPO w/ Entropy Adv augments the group-relative advantage with a per-token entropy bonus weighted by α, following Cheng et al. [41].Source: §A.2 · p. 15

B Theoratical Proof

• The proof assumes a finite response space (with sums replaced by integrals in the continuous case) and uses aligned monotonicity in repetition count to establish the reward guarantee.The repetition weight and conditional mean task reward are both non-increasing in the count, enabling the covariance argument via Chebyshev’s rearrangement inequality.Source: B Theoratical Proof · pp. 16–17
• Theorem 2 concludes that adding a positive repetition penalty for repeatedly sampled errors yields J_x(q2) ≥ J_x(q1) for every x.Here, q1 uses the task reward alone, while q2 uses the task reward minus λc(y), where c(y) indicates repeated errors.Source: B Theoratical Proof · pp. 16–17

Question

• The case-study question asks for the smallest positive integer with three distinct positive divisors summing to 2022.The passage defines such an integer as “good.”Source: Question · p. 18

C.1 Logit Values

• Figure 7 presents the heatmap from Figure 6 with logit values.The supplied passages identify the visualization but do not describe or interpret its values.Source: Figure 7 · p. 18, §C.1 · p. 18

C.2 Custer A Responses

• The responses report different smallest good numbers: 1342 in one answer and 1351 in another.The supplied passages do not reconcile these two results.Source: §C.2 · pp. 20–21
• The responses reason that a number with exactly three distinct divisors must be p^2 for prime p, giving 1 + p + p^2 = 2022.One response then proposes checking candidate primes computationally.Source: §C.2 · pp. 19–21

C.3 Cluster B Responses

• Cluster B responses try prime-square and other divisor structures before proposing direct enumeration, with one response rejecting 44 because it is not prime.Source: §C.3 · pp. 21–23
• Cluster B answers report 1342 as the smallest n, once for exactly three factors and once for three divisors excluding 1.The two passages phrase the divisor condition differently.Source: §C.3 · pp. 22–23
• Another Cluster B answer reports 1349 as the smallest integer with three distinct positive divisors summing to 2022.Source: §C.3 · p. 24

Prompt used in Figure 1, Within-Step Diversity

• The within-step prompt rates each answer group’s internal diversity from 1 to 5 based on shared conclusions, methods, and variation in correctness or reasoning quality.Source: Prompt used in Figure 1, Within-Step… · p. 25

Prompt used in Figure 1, Across-Step Diversity

• The across-step prompt scores whether a later group adds new reasoning paths relative to an earlier group on a 1–3 scale, disregarding wording changes and correctness unless they reflect a new path.Source: Prompt used in Figure 1, Across-Step… · p. 26

Prompt used in 4.5

• The error-analysis prompt asks for a wrong answer’s reasoning errors and assigns a category, with options including misinterpretation, modeling, calculation, and implementation errors.Source: Prompt used in 4.5 · p. 26

E Top-1 Eigen Ratio

• The Top-1 Eigen Ratio section orders covariance-matrix eigenvalues and explains that the largest eigenvalue measures variance along the dominant logit-space direction, while their sum is total variance.Source: E Top-1 Eigen Ratio · p. 27