GEOALIGN: Geometric Rollout Curation for Robust LLM Reinforcement Learning

2026-06-25

AI Summary

Online reinforcement learning often suffers from training instability due to reward noise or misspecification, particularly when influenced by high-reward trajectories within a batch that exhibit inconsistent directional signals. This work proposes GeoAlign, a lightweight plug-in for replay trajectory filtering that introduces, for the first time, the consistency of reward-induced directional alignment in the latent space of hidden states as a reliability signal. GeoAlign employs an online projector to align the displacement directions of hidden states across trajectories sharing the same prompt, identifies and discards outlier trajectories whose angles significantly deviate from the batch-wise consensus prototype, and replaces them with aligned alternatives. Requiring no additional training overhead, GeoAlign substantially improves final performance and mitigates training oscillations on dialogue alignment and mathematical reasoning tasks, outperforming existing methods such as PF-PPO, PAR, PODS, and Seed-GRPO.

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.

Problem

• reinforcement learning
• large language models
• directional inconsistency
• reward noise
• training instability

Innovation

• directional inconsistency
• geometric rollout curation
• online reinforcement learning
• latent directional consensus
• LLM alignment