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

2026-04-13

AI Summary

This work addresses the limited sampling diversity of large language models in reinforcement learning, which often leads to repetitive generation of similar errors. To mitigate this issue, the authors propose the MEDS framework, which uniquely integrates historical behavior memory with density-based clustering to explicitly suppress high-frequency error patterns through dynamic reward shaping. Specifically, MEDS stores past behaviors using intermediate representations, identifies recurrent errors via density clustering, and applies stronger penalties to frequent mistakes to encourage exploration. Experimental results demonstrate that MEDS consistently outperforms existing methods across five datasets and three base models, achieving gains of up to 4.13 (pass@1) and 4.37 (pass@128), while effectively enhancing behavioral diversity.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Multiagent Learning
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Metrics for user behavior and evaluating success
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

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.

Problem

• reinforcement learning
• sampling diversity
• repeated errors
• reward shaping
• large language models

Innovation

• reward shaping
• memory-enhanced learning
• behavioral diversity
• error pattern clustering
• reinforcement learning for LLMs