Revisiting Robustness for LLM Safety Alignment via Selective Geometry Control
For practitioners of LLM safety alignment, this work provides a new optimization perspective and method to improve robustness under distribution shift and noisy supervision.
The paper identifies that robustness failures in LLM safety alignment stem from optimization geometry issues, not just data noise, and proposes ShaPO, a geometry-aware preference optimization framework that selectively controls alignment-critical parameters. ShaPO consistently improves safety robustness over existing methods across diverse benchmarks and noisy settings.
Safety alignment of large language models remains brittle under domain shift and noisy preference supervision. Most existing robust alignment methods focus on uncertainty in alignment data, while overlooking optimization-induced fragility in preference-based objectives. In this work, we revisit robustness for LLM safety alignment from an optimization geometry perspective, and argue that robustness failures cannot be addressed by data-centric methods alone. We propose \textit{ShaPO}, a geometry-aware preference optimization framework that enforces worst-case alignment objectives via selective geometry control over alignment-critical parameter subspace. By avoiding uniform geometry constraints, ShaPO mitigates the over-regularization that can harm robustness under distribution shift. We instantiate ShaPO at two levels: token-level ShaPO stabilizes likelihood-based surrogate optimization, while reward-level ShaPO enforces reward-consistent optimization under noisy supervision. Across diverse safety benchmarks and noisy preference settings, ShaPO consistently improves safety robustness over popular preference optimization methods. Moreover, ShaPO composes cleanly with data-robust objectives, yielding additional gains and empirically supporting the proposed optimization-geometry perspective. The code is available at https://github.com/liujilong0116/ShaPO.