English

Revisiting Robustness for LLM Safety Alignment via Selective Geometry Control

Machine Learning 2026-05-22 v2

Abstract

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.

Keywords

Cite

@article{arxiv.2602.07340,
  title  = {Revisiting Robustness for LLM Safety Alignment via Selective Geometry Control},
  author = {Yonghui Yang and Wenjian Tao and Jilong Liu and Xingyu Zhu and Junfeng Fang and Weibiao Huang and Le Wu and Richang Hong and Tat-Sent Chua},
  journal= {arXiv preprint arXiv:2602.07340},
  year   = {2026}
}
R2 v1 2026-07-01T10:25:38.881Z