Revisiting Robustness for LLM Safety Alignment via Selective Geometry Control
ShaPO enhances LLM safety alignment robustness via selective geometry control, outperforming benchmarks.
Key Findings
Methodology
ShaPO is a geometry-aware preference optimization framework enforcing worst-case alignment objectives via selective geometry control over alignment-critical parameter subspace. It avoids uniform geometry constraints, mitigating over-regularization under distribution shifts. ShaPO is instantiated at two levels: token-level ShaPO stabilizes likelihood-based surrogate optimization, while reward-level ShaPO enforces reward-consistent optimization under noisy supervision.
Key Results
- ShaPO achieved an 85.88% WinRate on the PKU-SafeRLHF-30K dataset, significantly outperforming DPO's 73.70%.
- In cross-domain safety benchmarks, ShaPO-T's MD score was as low as 0.43%, compared to rDPO's 1.17%.
- ShaPO maintained high robustness under noisy preference supervision, outperforming other data-robust methods.
Significance
ShaPO offers a new perspective for LLM safety alignment by focusing on optimization geometry rather than solely relying on data robustness, addressing long-standing issues of fragility under distribution shifts and noisy supervision. This approach has significant academic impact and provides a more robust deployment strategy for the industry.
Technical Contribution
ShaPO introduces a new optimization framework through selective geometry control, offering new theoretical guarantees and engineering possibilities compared to existing SOTA methods. It enhances alignment robustness by restricting worst-case parameter perturbations without affecting capability-related parameters.
Novelty
ShaPO is the first to apply geometry control to LLM safety alignment, differing from traditional data-robust methods by offering a new optimization perspective. Compared to related work, ShaPO selectively controls alignment-critical parameters, avoiding over-constraining.
Limitations
- ShaPO may still fail under extreme noise scenarios as selective control cannot eliminate all uncertainties.
- The method relies on initial probing signals for parameter subspace selection, which may lead to incomplete safety parameter identification.
Future Work
Future work could explore finer geometry control strategies and integration with other data-robust methods. Further research on automating the identification of safety-critical subspaces is also important.
AI Executive Summary
Safety alignment of large language models (LLMs) remains brittle under domain shifts and noisy preference supervision. Existing robust alignment methods focus on uncertainty in alignment data, overlooking optimization-induced fragility in preference-based objectives. In this study, we revisit robustness for LLM safety alignment from an optimization geometry perspective and propose the ShaPO framework. ShaPO enforces worst-case alignment objectives via selective geometry control over alignment-critical parameter subspace, avoiding uniform geometry constraints that can harm robustness under distribution shifts. Experimental results show that ShaPO consistently improves safety robustness across diverse benchmarks and noisy preference settings, outperforming 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 on GitHub for researchers to use.
Deep Analysis
Background
Large language models (LLMs) have made significant progress in natural language processing, but their safety alignment remains challenging. Traditional methods rely heavily on reinforcement learning from human feedback (RLHF) to align model behavior with human safety preferences. However, these methods exhibit fragility under domain shifts and noisy supervision, leading to potentially unsafe responses during deployment.
Core Problem
There is a gap between training-time and deployment-time robustness in LLM safety alignment. Models performing well on curated safety benchmarks may still produce unsafe responses under domain shifts or noisy preference supervision. This fragility is primarily attributed to uncertainty in alignment data and sensitivity introduced by the optimization process.
Innovation
ShaPO achieves worst-case alignment objectives through selective geometry control over alignment-critical parameter subspace, avoiding over-regularization issues of uniform geometry constraints. It is instantiated at token and reward levels, stabilizing likelihood-based surrogate optimization and enforcing reward-consistent optimization under noisy supervision.
Methodology
- �� ShaPO enforces worst-case alignment objectives via selective geometry control. • At the token level, ShaPO stabilizes likelihood-based surrogate optimization. • At the reward level, ShaPO enforces reward-consistent optimization under noisy supervision. • ShaPO integrates with data-robust objectives for enhanced robustness.
Experiments
Experiments were conducted on the PKU-SafeRLHF-30K dataset and evaluated ShaPO's robustness across five safety benchmarks, including HH-RLHF-Safety and HarmBench. Controlled label flips were introduced to simulate noisy preference supervision, assessing ShaPO's alignment robustness under noisy conditions.
Results
ShaPO achieved an 85.88% WinRate on the PKU-SafeRLHF-30K dataset, significantly outperforming DPO's 73.70%. In cross-domain safety benchmarks, ShaPO-T's MD score was as low as 0.43%, compared to rDPO's 1.17%. ShaPO maintained high robustness under noisy preference supervision, outperforming other data-robust methods.
Applications
ShaPO can be directly applied to LLM deployment scenarios requiring high safety and robustness, such as automated customer service and medical diagnostics. Its selective geometry control strategy excels under domain shifts and noisy supervision.
Limitations & Outlook
ShaPO may still fail under extreme noise scenarios as selective control cannot eliminate all uncertainties. The method relies on initial probing signals for parameter subspace selection, which may lead to incomplete safety parameter identification.
Plain Language Accessible to non-experts
Imagine a kitchen where ShaPO acts like a smart chef who knows which ingredients are key and which can be flexible. During cooking, the chef pays special attention to those key ingredients, ensuring their quality and safety, rather than treating all ingredients equally. This selective focus ensures that the final dish remains delicious and safe across different tables.
ELI14 Explained like you're 14
Hey, imagine you're playing a super complex game with lots of levels and challenges. ShaPO is like a smart game assistant that knows which levels are the hardest and need special attention. It helps you stay robust on these key levels, rather than spending tons of energy on every level. This way, you can stay ahead in the game and tackle various challenges with ease!
Glossary
ShaPO (Selective Geometry Control)
A framework that achieves worst-case alignment objectives via selective geometry control over alignment-critical parameter subspace.
Used to enhance LLM safety alignment robustness.
LLM (Large Language Model)
A large-scale model capable of processing and generating natural language, typically with billions of parameters.
The subject of study requiring safety alignment.
RLHF (Reinforcement Learning from Human Feedback)
A method of training models through human feedback to achieve safety alignment.
Traditional method for safety alignment.
Geometry Control
A strategy to achieve optimization objectives by selectively constraining parameter subspaces.
Core innovation of ShaPO.
Safety Alignment
The process of ensuring model outputs meet safety standards.
Primary goal of the study.
Open Questions Unanswered questions from this research
- 1 How to automatically identify safety-critical subspaces remains an open question, as current methods rely on initial probing signals.
- 2 ShaPO's performance under extreme noise conditions needs further study to ensure robustness.
Applications
Immediate Applications
Automated Customer Service
ShaPO can enhance the safety of automated customer service systems, ensuring robustness against diverse user inputs.
Long-term Vision
Medical Diagnostic Systems
ShaPO can be used in medical diagnostic systems to ensure safety and accuracy when handling diverse patient data.
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.