GEOALIGN: Geometric Rollout Curation for Robust LLM Reinforcement Learning

Abstract: Online reinforcement learning is widely used to align 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.

• The paper introduces a geometric rollout curation method that uses a reward-sensitive projector to distill consistent directional signals from noisy RL updates.
• The paper demonstrates improved evaluation scores and smoother learning trajectories in both dialogue alignment and mathematical reasoning tasks.
• The paper shows that leveraging a geometric consensus prototype mitigates reward corruption and preserves diverse, robust policy updates.

GeoAlign: Geometric Rollout Curation for Robust LLM Reinforcement Learning

Motivation and Geometric Failure Modes in Online RL

This paper identifies and formalizes directional inconsistency as a critical instability in online RL for LLM alignment. In RLHF and reasoning tasks, updates are guided by rollouts with high reward or advantage scores, but the update directions in latent representation space are not always concordant. Even when reward is used as a scalar reliability signal, high-reward rollouts may induce steeply conflicting preference directions, generating angular outliers that destabilize training. These inconsistencies are exacerbated by reward noise, model exploitation, or misspecification, causing learning curves to oscillate and policies to regress despite apparently favorable updates.

Figure 1: Geometric turbulence in preference latent space: within one update step, most preference-induced directions align, while a small fraction forms angular outliers that can destabilize training.

Unlike prior stabilizers that operate solely on reward magnitude—via clipping, shaping, entropy filtering, or uncertainty weighting—this paper demonstrates that robust policy optimization must also account for the vector geometry of learning signals. Even rollouts with similarly high reward can push the policy in diverging directions, so scalar-based filtering is insufficient for preventing high-variance, contradictory updates. Empirical evidence (see Figure 1) shows most preference directions align, but a nontrivial fraction forms outliers responsible for instability.

GeoAlign Method: Geometric Rollout Curation

GeoAlign is introduced as a lightweight, forward-pass-only plug-in for rollout curation during iterative policy optimization. GeoAlign operates by:

1. Forming within-prompt preference pairs from the RL batch, mapping low-reward to high-reward rollouts.
2. Learning a reward-sensitive projector (small MLP) to distill reward-ordered directions in hidden state space into a concentrated improvement manifold.
3. Constructing a batch-wise consensus prototype direction.
4. Scoring each rollout via its angular deviation from this prototype and computing a per-rollout Geometric Deviation Index (GDI).
5. Rectifying the experience buffer by replacing directionally inconsistent rollouts with stable alternatives from the same prompt, thus neutralizing anomalous signals without shrinking per-prompt batch density.

   Figure 2: GeoAlign overview. At each iteration, within-prompt preference displacements are projected onto a reward-sensitive manifold, consensus prototype is constructed, rollouts are scored and rectified, before running the RL update.

Unlike methods requiring per-sample gradients or backward passes, GeoAlign uses only forward-pass representations and adds negligible overhead. The process preserves rollout diversity and batch density, intervening only on the most severe outliers as detected by adaptive density thresholding in the GDI score distribution.

Empirical Evaluation: Alignment and Reasoning Robustness

GeoAlign is validated across dialogue alignment (HH-RLHF with learned reward model) and mathematical reasoning (DAPO-Math with binary verified rewards), using Qwen3-1.7B and Qwen3-4B LLMs. Experiments cover competitive baselines, including Base-GRPO, PF-PPO, PAR, PODS, and Seed-GRPO. GeoAlign consistently improves final evaluation scores and yields smoother learning trajectories with reduced training oscillation, both for continuous and binary reward regimes.

Figure 3: Qwen3-1.7B on Math

GeoAlign achieves the highest accuracy on Math and the highest mean score on HH-RLHF relative to baselines. Performance gains range from 1.6–2.1% in reasoning, and up to 6.4% relative improvement in alignment. Evaluation trajectories show GeoAlign’s advantage in both median and stability, outperforming strong robust-RL alternatives.

Moreover, robustness to reward corruption is demonstrated: when rewards are intentionally corrupted (flipped or assigned to extremes), GeoAlign’s GDI metric identifies and rectifies these anomalies, maintaining the highest overall performance under noise. Visualization shows injected reward errors concentrate among high-GDI rollouts, confirming the directional metric’s reliability.

Figure 4: Injected reward errors concentrate among high-GDI rollouts for Math and HH-RLHF, confirming GeoAlign's anomaly detection.

Geometric Analysis and Ablations

A crucial step is learning the improvement manifold via a simple projector. Without projection, latent preference directions remain scattered; after projection, most directions align strongly with the geometric prototype.

Figure 5: Under a continuous reward model proxy.

Ablations indicate the projector achieves validation accuracy >85% (HH-RLHF) and >97% (Math) and is stable across depth variations (2–5 layers), supporting computational efficiency and generalizability.

The cumulative form of GDI score strongly enhances contrast between disrupting rollouts and normal samples, outclassing normalization alternatives in both boundary sharpness and task-level performance.

Training Stability and Diverse Solution Preservation

GeoAlign does not introduce additional variance, as shown by tightly clustered training curves across seeds and reductions in PPO trust-region diagnostics (KL divergence and clipping fraction). Rectification preserves per-prompt batch density; replacement outperforms zero-weighting and dropping, avoiding artificial variance and neutrality loss.

Pass@$k$ analysis confirms GeoAlign maintains or increases slope versus baseline, proving solution diversity is preserved—correct but creative solutions serve as positive anchors in preference pairs, contributing to consensus direction as long as they are not conflicting outliers.

Practical and Theoretical Implications

Practically, GeoAlign is a modular algorithm that can be incorporated into any group-normalized policy optimization pipeline for LLM RL, with negligible computational and memory overhead. It targets instability due to reward artifacts, misspecification, and reward-model exploitation, all endemic to RLHF and reasoning tasks. The geometric approach provides an orthogonal reliability signal to reward magnitude, allowing principled targeting of training inefficiency and variance inflation.

Theoretically, the directional consensus prototype connects directly to the variance of stochastic policy gradients, promising future formal links with batch-level update statistics and principled integration with other alignment objectives.

Speculation and Future Directions

The latent directional consensus approach opens avenues for integrating geometric signals with semantic uncertainty measures, reward-model ensembles, and dual-objective RL. Further research may develop online consensus estimation across multiple tasks, layer-wise geometric calibration, and transfer across domains with varying reward structures. Scaling GeoAlign to larger models and broader alignment objectives (multi-objective RL, agentic behaviors) is anticipated.

Conclusion

GeoAlign addresses geometric instability in online RL for LLM alignment by leveraging preference-induced directions and batch-level consensus prototypes. The method stabilizes learning and improves final performance, showing resilience to reward corruption and preserving solution diversity. Latent directional consensus is a robust reliability signal, motivating future work on geometric-statistical alignment mechanisms and integration into broader RL pipelines.