Improving Safety Alignment via Balanced Direct Preference Optimization

Shiji Zhao; Mengyang Wang; Shukun Xiong; Fangzhou Chen; Qihui Zhu; Shouwei Ruan; Yisong Xiao; Ranjie Duan

cs.AIMar 24, 2026

Improving Safety Alignment via Balanced Direct Preference Optimization

Authors: Shiji Zhao, Mengyang Wang, Shukun Xiong, Fangzhou Chen, Qihui Zhu, Shouwei Ruan, Yisong Xiao, Ranjie Duan,

Organizations: Institute of Artificial Intelligence, Beihang University, Beijing, China.

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 & tables

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.

Panel of Figure 2 — 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).

Panel of Figure 3 — 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.

Panel of Figure 4 — 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.

	General				Safety				Overall
	SimpleQA	GSM8K	AdvGLUE	HHH	StrongReject	XsTest	GCG	PAIR	Avg.
Qwen2-7B-Instruct
Base	5.30%	71.19%	64.91%	88.69%	80.51%	82.50%	97.37%	58.00%	68.56%
DPO	5.00%	72.25%	65.85%	86.88%	84.03%	81.50%	93.54%	76.00%	70.63%
CPO	5.50%	72.63%	65.04%	88.69%	83.71%	81.00%	89.09%	60.00%	68.21%
SimPO	5.70%	71.42%	64.50%	88.24%	78.85%	80.50%	85.05%	52.00%	65.78%

Table 1: Performance on diverse benchmarks reflects both safety and general performance. “Avg.” is the average performance of all benchmarks. For all reported metrics, the best results are marked in bold and the second best results are marked by underline.

Method	StrongReject	XsTest	AdvGLUE
Base	80.51%	82.50%	64.91%
DPO	84.03%	81.50%	65.85%
DPO+BW	85.30%	82.50%	65.18%
DPO+SF	86.58%	83.00%	65.31%
B-DPO(ours)	92.01%	84.50%	65.72%

Table 2: Ablation study for the components of B-DPO. “Base” denotes the basic LLM. “DPO” denotes the fine-tuned version based on “Base” model, “DPO+BW” denotes B-DPO version that adds the balanced weight but does not include the scaling factor, “DPO+SF” denotes only adding the scaling factor but without balanced weight, “B-DPO” is the final version of our method.

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.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: The responses towards harmful prompts of LLMs trained by DPO and B-DPO.

Figure 7: The responses towards jailbreak instructions of LLMs trained by DPO and B-DPO.

LLM Safety Alignment

Large language model safety alignment trains and evaluates systems to follow human policies while resisting harmful requests and misuse.

See all 341 papers

CompassDPO: Dynamics-Controlled Direct Preference Optimization for Robust Safety Alignment
Mar 7, 2026Jilong Liu, Yonghui Yang, Pengyang Shao +5Direct Preference Optimization LLM Safety Alignment
Hefei University of Technology · National University of Singapore
Cat-DPO: Category-Adaptive Safety Alignment
Apr 19, 2026Tiankai Yang, Yi Nian, Xinyuan Li +3LLM Alignment Direct Preference Optimization
University of Southern California · 2Northwestern University
Suan: Rectifying Direct Preference Safety Alignment in Large Language Models
Sep 8, 2026Oleksandr Cherednichenko, Roman KlypaLLM Alignment LLM Safety Alignment
Department of Mathematics and Mathematical Statistics Integrated Science Lab, Umeå University Umeå, Sweden · Univ. Grenoble Alpes, CNRS, Grenoble INP, LJK 38000 Grenoble, France

Mar 7, 2026J/K move · Enter open · S save

CompassDPO: Dynamics-Controlled Direct Preference Optimization for Robust Safety Alignment

Jilong Liu, Yonghui Yang, Pengyang Shao, Wenjian Tao+4

Hefei University of Technology · National University of Singapore

Direct Preference Optimization LLM Safety Alignment Preference-Based RL

Open paper AskPDF arXiv

Direct Preference Optimization (DPO) has become a standard framework for safety alignment, but its reliance on pairwise preference updates makes training sensitive to imperfect supervision. Existing robust DPO methods often address this sensitivity through global loss corrections or external data-level interventions, while largely overlooking how unreliable comparisons distort batch-level optimization dynamics. We propose CompassDPO, a reward-free DPO framework that stabilizes preference optimization through dynamics control. Using the implicit DPO reward margin as a training-time compass, CompassDPO regulates sample influence along two complementary axes: update direction and update magnitude. For directional control, it applies sparse, budgeted, and warm-up delayed loss mixing to attenuate update components that conflict with the emerging preference direction. For magnitude control, it adaptively soft-winsorizes high-loss tail contributions, reducing tail dominance while preserving useful gradients from hard examples. Both mechanisms use only signals available during standard DPO training and require no external reward model or additional supervision. Experiments on PKU-SafeRLHF across four backbones and multiple out-of-distribution safety benchmarks show that CompassDPO consistently improves robustness over vanilla DPO and strong DPO-family baselines, especially under controlled label-flip noise. Code is available at https://anonymous.4open.science/r/CompassDPO-4D00

Abstract

LLM Safety Alignment

Explore similar work

CompassDPO: Dynamics-Controlled Direct Preference Optimization for Robust Safety Alignment

Cat-DPO: Category-Adaptive Safety Alignment

Suan: Rectifying Direct Preference Safety Alignment in Large Language Models