Mitigating the Safety Alignment Tax with Null-Space Constrained Policy Optimization

TL;DR

This paper introduces Null-Space constrained Policy Optimization (NSPO), a reinforcement learning framework that enhances safety alignment in large language models while preserving their core abilities, using geometric projections to mitigate the safety alignment tax.

Contribution

The paper proposes NSPO, a novel RL method that preserves language models' abilities during safety alignment by projecting safety policy gradients into the null space of general tasks.

Findings

NSPO outperforms existing safety alignment methods.

Achieves state-of-the-art safety performance on multiple tasks.

Requires only 40% of safety data for effective alignment.

Abstract

As Large Language Models (LLMs) are increasingly deployed in real-world applications, it is important to ensure their behaviors align with human values, societal norms, and ethical principles. However, safety alignment under Reinforcement Learning (RL) often suffers from forgetting learned general abilities, which is also known as the alignment tax. To address this issue, we introduce Null-Space constrained Policy Optimization (NSPO), a novel RL framework for LLM safety alignment while preserving their core abilities. The safety policy gradients are geometrically projected into the null space of general tasks, thereby mitigating the safety alignment tax. In addition, we theoretically prove that NSPO preserves the model's original core capabilities, while still guaranteeing a descent direction for effective safety alignment. Extensive experiments demonstrate that NSPO outperforms existing methods by a large margin, achieving state-of-the-art safety performance without sacrificing accuracy on general tasks, including math, code, and instruction-following tasks. Notably, NSPO is data-efficient and only requires 40% of public human-annotated safety data from PKU-SafeRLHF to achieve promising safety performance, without a large amount of mixed general tasks data in existing alignment methods.