Improving Safety Alignment via Balanced Direct Preference Optimization

Abstract

With the rapid development and widespread application of Large Language Models (LLMs), their potential safety risks have attracted widespread attention. Reinforcement Learning from Human Feedback (RLHF) has been adopted to enhance the safety performance of LLMs. As a simple and effective alternative to RLHF, Direct Preference Optimization (DPO) is widely used for safety alignment. However, safety alignment still suffers from severe overfitting, which limits its actual performance. This paper revisits the overfitting phenomenon from the perspective of the model's comprehension of the training data. We find that the Imbalanced Preference Comprehension phenomenon exists between responses in preference pairs, which compromises the model's safety performance. To address this, we propose Balanced Direct Preference Optimization (B-DPO), which adaptively modulates optimization strength between preferred and dispreferred responses based on mutual information. A series of experimental results show that B-DPO can enhance the safety capability while maintaining the competitive general capabilities of LLMs on various mainstream benchmarks compared to state-of-the-art methods. \color{red}{Warning: This paper contains examples of harmful texts, and reader discretion is recommended.

Figures (7)

• Figure 1: The Illustration of B-DPO. We find Imbalanced Preference Comprehension in safety preference pairs, which produces a negative impact on safety alignment. Our B-DPO adjusts optimization strengths towards different responses based on the comprehension level, finally improving the safety performance of LLMs.
• Figure 2: The statistical analysis of Mutual Information for preference pairs on Qwen-2-7B-Instruct and Mistral-7B-Instruct-v0.3. The x and y axes represent the Mutual Information between queries and dispreferred vs. preferred responses, respectively. We can find that both the LLMs exist Imbalanced Preference Comprehension phenomenon for both the safe and unsafe query types (For a portion of the queries, preferred responses exhibit higher mutual information than dispreferred ones, whereas the reverse holds for the remaining queries).
• Figure 3: The curve of the reward margin between preferred and dispreferred responses trained by the relative balanced dataset $\mathcal{D}_{\text{balanced}}$ and the imbalanced dataset $\mathcal{D}_{\text{imbalanced}}$. The reward margin of the $\mathcal{D}_{\text{imbalanced}}$ is higher than the $\mathcal{D}_{\text{balanced}}$, which indicates LLMs fit the imbalanced data better than the balanced data.
• Figure 4: The safety and general performance of LLMs trained by a balanced dataset $\mathcal{D}_{\text{balanced}}$ and an imbalanced dataset $\mathcal{D}_{\text{imbalanced}}$. We can observe that LLMs trained by the balanced dataset $\mathcal{D}_{\text{balanced}}$ achieve better safety and general performance.
• Figure 5: Ablation study on hyper-parameter $\alpha$ selection. The safety performance is evaluated based on StrongReject. And setting $\alpha$ value to 1.5 can yield the highest safety result.