Language models learn better on long tasks with step by step rewards
Long-Horizon Language Model Reinforcement Learning via Progressive Point Matching
Machine Learning
Summary
Training language models to complete very long or complex tasks often fails because they get feedback only at the very end, making learning slow and difficult. The authors show that giving rewards not just at the end but also throughout the task, by checking progress at segments, helps the model learn much faster and better. They call this method progressive point matching and prove it works well both in theory and in tests on simple environments. When used on difficult math problems, this approach allowed language models to improve while traditional methods could not.
language modelreinforcement learningsparse rewardsdense rewardslong-horizon tasksprogressive point matchingtrajectorysegment-level rewardssuccess ratepass@k
Authors
Preston Fu, Kevin Frans, Oleh Rybkin, Sergey Levine, Aviral Kumar
Abstract
Current paradigms for training language models via reinforcement learning rely heavily on sparse outcome rewards. However, as we pursue tasks that require longer and more complicated trajectories, such strategies result in slow learning. Prior work has attempted to address this problem by rewarding partial progress; however, naive formulations are often biased and converge to suboptimal policies. We show that a simple and unbiased dense reward formulation, which we term progressive point matching, scales exponentially more efficiently to long-horizon tasks by rewarding partial progress on a segment level, both theoretically and empirically via synthetic environments. We then show how progressive point matching can be practically instantiated using a single reference trajectory per task. On extremely hard math reasoning problems, sparse outcome rewards cannot make any progress, whereas segment-level rewards enable improvements at larger test-time token budgets when measured by success rate or pass@k.