Towards Context-Invariant Safety Alignment for Large Language Models

TL;DR

This paper introduces Anchor Invariance Regularization (AIR), a novel method to improve context-invariant safety alignment in large language models by anchoring open-ended prompt responses to verifiable ones.

Contribution

The paper proposes AIR, a plug-in auxiliary loss that enhances safety robustness by enforcing invariance using verifiable prompts as anchors, combined with group-based preference optimization.

Findings

AIR improves in-distribution group accuracy by 12.71%.

AIR boosts out-of-distribution consistency by 33.49%.

Enhances safety robustness against adversarial framings.

Abstract

Preference-based post-training aligns LLMs with human intent, yet safety behavior often remains brittle. A model may refuse a harmful request in a standard prompt but comply when the same intent is wrapped in adversarial wording. We suggest that robust safety requires context-invariant alignment, where behavior depends on the underlying intent rather than surface form. Enforcing invariance is difficult in alignment because not all training signals are equally trustworthy; for some prompt variants we can obtain verifiable feedback (e.g., multiple-choice), while for open-ended variants we typically rely on noisy, gameable reward proxies (e.g., learned judges). As a result, standard symmetric invariance regularizers can reduce cross-context discrepancies by lowering performance on reliable variants instead of improving open-ended robustness. To address this, we introduce Anchor Invariance Regularization (AIR), which treats verifiable prompts as anchors and uses a stop-gradient target to regularize only the open-ended variants toward the anchor performance. AIR is implemented as a plug-in auxiliary loss and combined with group-based preference optimization (e.g., GRPO) via heterogeneous prompt grouping. Across Safety, Moral Reasoning, and Math, AIR improves context invariance, boosting in-distribution group accuracy by 12.71% and out-of-distribution consistency by 33.49%, making safety constraints robust to adversarial framings.