MEDS: Memory-Enhanced Dynamic Reward Shaping for LLM Reinforcement Learning

Summary (Overview)

• Core Problem: On-policy reinforcement learning (RL) for large language models (LLMs) often suffers from "error collapse," where the policy degenerates into generating highly repetitive erroneous reasoning patterns.
• Proposed Solution: Introduces MEDS (Memory-Enhanced Dynamic reward Shaping), a framework that penalizes repeated historical error patterns to encourage broader exploration and reduce recurrent mistakes.
• Key Mechanism: Uses layer-wise logits as lightweight, low-overhead representations of the model's internal reasoning process. These are clustered (using HDBSCAN) to identify and track recurring error patterns across training epochs.
• Reward Shaping: Dynamically adjusts the reward for a new response based on the size of its assigned error cluster: $\tilde{r}(x, \tilde{y}) = r(x, \tilde{y}) - \min(\alpha \log(|C_k| + 1), \beta)$.
• Empirical Results: MEDS consistently outperforms baselines (GRPO, DAPO) across five mathematical reasoning benchmarks and three base models, achieving gains of up to 4.13 pass@1 points and 4.37 pass@128 points. It also demonstrably increases behavioral diversity during sampling.

Introduction and Theoretical Foundation

Reinforcement learning has become a key method for improving LLM performance. However, a persistent failure mode is error collapse: as training progresses, the policy often converges to a narrow set of behaviors, repeatedly generating similar erroneous reasoning trajectories. This reduces exploration and entrenches the model in self-reinforcing failure modes.

While entropy regularization encourages randomness under the current policy distribution, it does not explicitly discourage recurrent error patterns that appear across different rollouts and training steps. The vast action space of LLMs makes distribution-level stochasticity insufficient to escape these stable error basins.

Inspired by dynamic reward mechanisms in human learning (where stronger penalties are imposed on recurring mistakes), this paper proposes MEDS. The core idea is to dynamically record historical error patterns and impose incremental penalties on repetitive failure paths. This provides more informative supervision than static rewards, helping the model escape local optima.

Theoretical Analysis: The paper provides a proof (Theorem 1) that penalizing repeated errors improves the expected return. Given two reward signals:

• $\mu_1 = r(x, y)$ (standard reward)
• $\mu_2 = r(x, y) - \lambda c(y)$ (reward with penalty for repetition)

Let $q_1$ and $q_2$ be the corresponding updated policies, and $J(q) = \mathbb{E}_{x,y \sim q}[r(x, y)]$ be the expected return. The theorem states:

$J(q_2) \geq J(q_1)$

where $\lambda$ is the penalty coefficient and $c(y)$ is an indicator function that increases monotonically with the number of past occurrences of error $y$. The proof leverages the Gibbs form of the KL-regularized policy update and Chebyshev's rearrangement inequality to show that the covariance between the reward $r$ and the penalty weight $w(y) = \exp(-\eta\lambda c(y))$ is non-negative, leading to the inequality.

Methodology

MEDS operates in three stages, as illustrated in Figure 2 of the paper.

1. Preliminaries: The standard RL objective for an LLM policy $p_\theta$ is:

$$\max_\theta \mathbb{E}_{x \sim D}\left[ \mathbb{E}_{y \sim p_\theta(\cdot|x)}[r(x, y)] - \frac{1}{\eta} \text{KL}(p_\theta(\cdot | x) \| p_{\text{ref}}(\cdot | x)) \right]$$

where $p_{\text{ref}}$ is a reference policy and $\eta$ controls the regularization strength.

2. Logic Feature Extraction: To implement the indicator function $c(y)$ efficiently, MEDS reuses the model's internal layer-wise logits. For a response $y$, let $y^* = y_{L-t}$ be the first token of the final answer. The logit for this token at layer $n$ is $l^{(n)}_{y_{L-t}} \in \mathbb{R}^V$. The scalar value for the chosen token is denoted $l^{*(n)} \in \mathbb{R}$.

Since earlier layers model simpler semantics, the feature vector is constructed from the latter half of the Transformer layers:

$$f(y) = \text{concat}\left( l^{*(n)} \ \middle| \ n = \frac{N}{2}, \dots, N \right) \in \mathbb{R}^{\frac{N}{2}}$$

This vector serves as a compact proxy for the model's reasoning trajectory.

3. Cluster-based Reward Shaping: For each prompt $x$, a memory $G_x$ is maintained containing feature vectors of all historical responses. HDBSCAN clustering is applied to $G_x$ to obtain clusters $\\{C_1, C_2, \dots, C_K, C_{\text{noise}}\\}$.

For a new response $\tilde{y}$ assigned to cluster $C_k$, the penalty is defined as a monotonic function of the cluster size:

$$c(\tilde{y}) = \log(|C_k| + 1)$$

The shaped reward is then:

$$\tilde{r}(x, \tilde{y}) = r(x, \tilde{y}) - \min\left(\alpha \log(|C_k| + 1), \beta\right), \quad \text{s.t.} \ f(\tilde{y}) \in C_k$$

where $\alpha$ and $\beta$ are hyperparameters controlling penalty strength and its upper bound. This directly implements the theoretical penalty on repeated errors.

Empirical Validation / Results

Experimental Setup:

• Models: Qwen3-1.7B, Qwen2.5-Math-7B, Qwen3-8B.
• Baselines: Base model, GRPO, DAPO, GRPO with Entropy Advantage.
• Benchmarks: Five mathematical reasoning datasets (AIME24, AMC23, MATH500, Minerva, OlympiadBench).
• Evaluation: pass@1 and pass@128.

Main Results

Table 1: Main results on five mathematical reasoning benchmarks. MEDS consistently achieves the best overall performance.

• Key Findings:
  • MEDS achieves the best average performance across all three model scales.
  • The improvement is most pronounced on the larger Qwen3-8B model, suggesting benefits scale with model capability.
  • On Qwen3-8B, MEDS improves pass@128 from 70.81 to 82.67 on OlympiadBench, a 17% relative gain.
  • Pass@k curves (Figure 4) show MEDS consistently matches or outperforms baselines across all k.

Impact on Exploration Behavior:

• Diversity Metrics: Using LLM-based annotation, MEDS shows higher Within-Step Diversity (diversity of rollouts at the same step) and Across-Step Diversity (novelty of later rollouts compared to earlier ones) than DAPO.
• Representation Diversity: The Top-1 Eigen Ratio of the logit covariance matrix is lower for MEDS, indicating a more uniform spread of representations and less concentration on a single error direction.

  Top-1 Eigen Ratio $= \frac{\lambda_1}{\sum_{j=1}^d \lambda_j}$, where $\lambda_i$ are eigenvalues of the logit covariance matrix.

Logits Reflect Reasoning Patterns:

• Qualitative Case Study (Figure 6): Responses with similar underlying reasoning patterns (even if producing different final answers) exhibit correlated layer-wise logit trajectories, especially in later layers.
• Quantitative Validation: Logit-based clustering shows a 61.2% correlation with error categories annotated by Claude-Haiku-4.5, confirming logits serve as an acceptable proxy for reasoning structure.

Clustering Quality Matters:

• Table 2: Comparison of different clustering feature constructions on Qwen2.5-Math-7B. Using the last 14 layers (MEDS-14) yields the best performance.

• Key Findings:
  • Any variant using clustering information outperforms the DAPO baseline.
  • The single cluster degenerate baseline performs worse, highlighting the importance of meaningful clustering.
  • The ranking of downstream performance closely mirrors the ranking of clustering consistency with LLM annotations, showing that better clustering quality leads to larger performance gains.

Theoretical and Practical Implications

Theoretical Implications:

• Provides a formal justification for incorporating historical error information into reward shaping, proving it can non-decrease the expected return.
• Establishes layer-wise logits as a valid, efficient proxy for capturing the model's internal reasoning state, connecting representation-level analysis with RL optimization.

Practical Implications:

• Efficient Exploration: MEDS offers a low-overhead method (reusing existing logits) to significantly improve exploration diversity and mitigate error collapse in LLM RL.
• Reward Design: Demonstrates the value of dynamic, memory-aware reward functions over static ones, opening a new direction for RL reward engineering.
• Model Analysis: The success of logit-based clustering validates the use of internal model representations for understanding and steering model behavior beyond surface-level text.

Conclusion

MEDS addresses the critical problem of error collapse in LLM reinforcement learning by introducing a memory-enhanced dynamic reward shaping framework. Key takeaways:

1. Core Innovation: MEDS is the first method to explicitly incorporate historical error patterns into reward modeling, penalizing repetitive failures to encourage broader exploration.
2. Efficient Implementation: By leveraging layer-wise logits as lightweight reasoning representations, it achieves significant gains with minimal computational overhead.
3. Consistent Improvements: Empirical results across multiple models and benchmarks show MEDS consistently improves both pass@1 and pass@128 performance and increases behavioral diversity.
4. Validated Mechanism: Theoretical analysis and empirical correlations confirm that logit-based clustering effectively captures reasoning patterns and that better clustering leads to better downstream performance.

Limitations and Future Work: The current logit aggregation strategies are relatively simple. Future work could explore more sophisticated feature extraction methods and apply the MEDS framework to other domains beyond mathematical reasoning.

Code: Available at https://github.com/Linxi000/MEDS