Entropy-Regularized Process Reward Model

2024-12-15

Trans. Mach. Learn. Res.

AI Summary

To address systematic errors in large language models (LLMs) during mathematical reasoning, this paper proposes a process-reward-based reinforcement learning optimization framework. Methodologically, it introduces entropy regularization into process reward modeling for the first time and rigorously derives an optimal reward construction scheme—balancing policy optimization and initial distribution constraints—within the theoretical framework of KL-regularized Markov decision processes. By assigning stepwise scores to explicitly guide reasoning-path generation, the approach enhances controllability and robustness of the inference process. Empirically, on the GSM8K and MATH benchmarks, the method achieves 1–3% improvement in Best-of-N evaluation and over 1% absolute accuracy gain after RLHF fine-tuning. These results demonstrate substantial improvements in both reliability and generalization capability of LLMs for mathematical reasoning.

Technology Category

• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Large language models (LLMs) have shown promise in performing complex multi-step reasoning, yet they continue to struggle with mathematical reasoning, often making systematic errors. A promising solution is reinforcement learning (RL) guided by reward models, particularly those focusing on process rewards, which score each intermediate step rather than solely evaluating the final outcome. This approach is more effective at guiding policy models towards correct reasoning trajectories. In this work, we propose an entropy-regularized process reward model (ER-PRM) that integrates KL-regularized Markov Decision Processes (MDP) to balance policy optimization with the need to prevent the policy from shifting too far from its initial distribution. We derive a novel reward construction method based on the theoretical results. Our theoretical analysis shows that we could derive the optimal reward model from the initial policy sampling. Our empirical experiments on the MATH and GSM8K benchmarks demonstrate that ER-PRM consistently outperforms existing process reward models, achieving 1% improvement on GSM8K and 2-3% improvement on MATH under best-of-N evaluation, and more than 1% improvement under RLHF. These results highlight the efficacy of entropy-regularization in enhancing LLMs' reasoning capabilities.

Problem

• Improving mathematical reasoning in LLMs with process rewards
• Balancing policy optimization via entropy-regularized reward modeling
• Reducing systematic errors in multi-step reasoning trajectories

Innovation

• Entropy-regularized process reward model balances optimization
• KL-regularized MDP prevents excessive policy distribution shift
• Novel reward construction derived from initial policy sampling