Towards Context-Invariant Safety Alignment for Large Language Models

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.

When aligning Large Language Models (LLMs) with human values, safety behaviors are often fragile. Robust AI requires context-invariance: the model's behavior should be driven by the user's underlying intent, regardless of how the prompt is formatted or phrased. However, teaching this is uniquely challenging due to a gap in supervision. Standard machine learning techniques that try to force consistency across all formats often fail; they inadvertently drag the model's reliable performance down to match the noisy evaluations. To solve this, we introduce Anchor Invariance Regularization (AIR). AIR treats verifiable, highly reliable prompts as fixed "anchors." During training, it applies an asymmetric constraint: it forces the model's behavior on tricky, open-ended prompts to improve and match the anchor standard, while mathematically preventing the anchor performance from degrading. Essentially, it ties noisy generation to rigorous, checkable constraints. Across safety, moral reasoning, and mathematics, our method dramatically increases the model's consistency and resilience against adversarial prompts.