Score centering improves stability of reinforcement learning under mismatch
Score Centering Stabilizes Off-policy Reinforcement Learning
Machine Learning
Summary
When computers learn to do tasks by trial and error, small differences between how they train and how they act can cause problems. The authors show that these issues come from a biased shift that builds up during training. They introduce a simple fix, called score centering, that cancels out this bias and helps the learning stay stable. Their method works well for models of many sizes and works even better than some previous solutions when the differences are bigger. Combining their fix with older techniques makes the learning even more reliable.
What this means in practice
- •For machine learning engineers: Improve training stability of large language models when there are small differences between training and deployment settings.
- •For software engineers deploying ai: Reduce failures in AI applications caused by mismatch between training simulations and real-world inference.
Authors
Martin Marek, Max Ryabinin
Abstract
Reinforcement learning (RL) of large language models is notoriously sensitive to small differences between training and inference engines, often referred to as the training-inference mismatch (TIM). However, completely eliminating TIM is impractical, as it would come at a major cost to rollout efficiency. In this paper, we show that the instability of RL under TIM is primarily caused by drift: a persistent bias between training and inference engines that accumulates with every training step. We derive an additive "score centering" correction term that stabilizes RL under TIM by canceling drift. When training models from 0.6B to 30B parameters, score centering alone matches or outperforms methods based on importance sampling under quantization, with the gap growing as the mismatch becomes more severe. Because the correction is additive, score centering also composes with importance sampling -- their composition outperforms pure importance-sampling baselines in our staleness experiments.