🤖 AI Summary
Existing exploration methods in reinforcement learning for large language models struggle to distinguish whether behavioral diversity stems from genuine reasoning or mere memorized pattern replay, often biasing exploration toward memory recall at the expense of reasoning capability improvement. To address this, this work proposes Direction-aware Reinforcement Learning (DiRL), a novel framework that introduces, for the first time, the concept of a “reasoning–memory direction.” By extracting this direction from internal model representations, DiRL constructs direction-weighted gradient features and leverages them for reward shaping, thereby steering exploration toward authentic reasoning processes. Integrated within the Group Relative Policy Optimization framework, DiRL demonstrates significant performance gains over existing approaches on both mathematical and general reasoning benchmarks, effectively suppressing memory-driven updates while genuinely enhancing the model’s reasoning abilities.
📝 Abstract
Reinforcement learning has become a key paradigm for eliciting reasoning abilities in large language models, where exploration is crucial for discovering effective solution trajectories. Existing exploration methods typically encourage diversity in semantic or gradient spaces, without distinguishing what drives this diversity. A trajectory may appear novel because it follows a new reasoning process, or because it varies memorized patterns and shortcuts. Rewarding both cases equally may steer exploration toward memorization rather than genuine reasoning improvement. In this paper, we propose DiRL, a Direction-Aware Reinforcement Learning framework that anchors exploration to an internal reasoning-memorization direction of the policy. Specifically, DiRL extracts this direction from model representations, constructs direction-weighted gradient features to characterize rollout updates, and shapes rewards to amplify reasoning-aligned exploration while suppressing memorization-aligned variations. DiRL integrates seamlessly into standard Group Relative Policy Optimization (GRPO). Extensive experiments on mathematical and general reasoning benchmarks demonstrate the effectiveness of DiRL, showing significant improvements over various existing exploration methods.