GeoAlign: Geometric Rollout Curation for Robust LLM Reinforcement Learning
Abstract
Online reinforcement learning is widely used to align large language models (LLMs) with reward signals, yet training can be unstable under noisy or misspecified rewards. We identify a failure mode we call directional inconsistency: within a batch, a small set of high-reward rollouts induces representation-space preference directions that sharply disagree with the batch majority, resulting in high-variance and destabilizing updates. We propose GeoAlign, a lightweight plug-in for rollout curation in iterative policy optimization. GeoAlign (i) forms within-prompt preference pairs, (ii) learns an online projector on per-rollout hidden states to concentrate reward-ordered displacement directions, and (iii) detects directionally inconsistent rollouts via their angular deviation from a batch consensus prototype and rectifies them with within-prompt stable alternatives. GeoAlign is forward-pass only and adds negligible overhead. Across dialogue alignment with a learned reward model and mathematical reasoning with binary verified rewards, GeoAlign improves final performance and reduces training oscillation, outperforming PF-PPO, PAR, PODS, and Seed-GRPO. These results suggest latent directional consensus as an effective reliability signal for online LLM RL.
Lay Summary
Large language models are often fine-tuned with the help of a reward model that scores which answers are better. In practice, a perfect reward model is hard to build, and the scores it gives always carry some noise. This noise can interfere with the optimization process and steer learning away from the direction we actually want. We found a useful warning sign for this. Looking at how each answer would nudge the model inside its own representation space, most answers push in a similar direction, like a crowd walking together. A small number point sharply away from the crowd, even when their reward scores look good. These off-direction cases are often the ones driven by noisy rewards rather than by genuinely better answers. Our method spots these cases on the fly and swaps them for safer alternatives from the same prompt, with negligible added cost. On dialogue and math reasoning tasks, it reaches better downstream evaluation than baseline methods under standard training, and shows stronger noise resistance when reward noise is deliberately injected.