Bellman policy optimization improves reasoning in language models
Bellman Policy Optimization
Machine LearningComputation and Language
Summary
Reinforcement learning helps language models learn better by getting rewards for good answers. The authors introduce a new method called Bellman Policy Optimization (BPO) that simplifies how the model improves its decision-making during text generation. BPO uses a mathematical trick to avoid hard calculations in the middle steps, focusing only on the overall outcome. Their tests show BPO helps models do better on math reasoning tasks.
What this means in practice
- •For language model engineers: Train language models to improve final task performance without estimating intermediate steps using the BPO approach.
- •For automated reasoning teams: Use BPO to enhance mathematical reasoning skills in language-based AI assistants.
Authors
Zhuoqing Song, Haotian Xu, Xikun Zhang, Lidong Bing
Abstract
Reinforcement learning with verifiable rewards (RLVR) improves the reasoning capabilities of large language models (LLMs). We introduce Bellman Policy Optimization (BPO), a critic-free method derived from Policy Mirror Descent (PMD). For autoregressive generation with terminal rewards, BPO uses the Bellman equations to reformulate PMD as a trajectory-level objective. The reformulation avoids estimating state values at intermediate states. We prove that it has the same unique optimal solution as the original PMD objective. We derive the practical BPO loss by approximating this objective. Its mismatch-correction weight is a smoothed ratio of complementary token probabilities. Experiments on mathematical reasoning benchmarks demonstrate the effectiveness of BPO.