Revisiting Robustness for LLM Safety Alignment via Selective Geometry Control

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 \url{https://github.com/liujilong0116/ShaPO}.

AI language models are trained to avoid harmful outputs, but this safety is often fragile; a model that passes safety tests may still fail in new real-world contexts or when trained on imperfect data. Most existing fixes focus on improving training data, but we show that the problem also lies in how models are trained. We find that a tiny fraction of a model's internal components are disproportionately responsible for safe behavior, yet standard training makes them unnecessarily sensitive. Our method, ShaPO, identifies and stabilizes exactly these components, leaving the rest of the model free to adapt. This makes safety significantly more robust: it holds under distribution shifts, resists noisy supervision, and generalizes to unseen scenarios. As AI is deployed in increasingly unpredictable environments, robust safety guarantees are essential for responsible AI development.