Improving LLM Safety Alignment with Dual-Objective Optimization

TL;DR

This work tackles the limitations of Direct Preference Optimization (DPO) for safety alignment in large language models, particularly its poor refusal learning and limited out-of-distribution generalization. It introduces Dual-Objective Optimization for Refusal (DOOR), which combines robust refusal training with targeted unlearning via Negative Preference Optimization (NPO), and extends it with Weighted DOOR (W-DOOR) that emphasizes critical refusal tokens through token-level reward weighting. The approach yields significantly improved resilience to jailbreak attacks (prefilling, suffix, and multi-turn) while retaining general capabilities, and reveals that robustness correlates with token-distribution shifts and internal representations of safety tokens. These findings advance practical safety-alignment strategies by focusing on explicit refusal reinforcement, harmful-knowledge unlearning, and fine-grained token-level optimization, offering a path toward safer, more trustworthy LLM deployments.

Abstract

Existing training-time safety alignment techniques for large language models (LLMs) remain vulnerable to jailbreak attacks. Direct preference optimization (DPO), a widely deployed alignment method, exhibits limitations in both experimental and theoretical contexts as its loss function proves suboptimal for refusal learning. Through gradient-based analysis, we identify these shortcomings and propose an improved safety alignment that disentangles DPO objectives into two components: (1) robust refusal training, which encourages refusal even when partial unsafe generations are produced, and (2) targeted unlearning of harmful knowledge. This approach significantly increases LLM robustness against a wide range of jailbreak attacks, including prefilling, suffix, and multi-turn attacks across both in-distribution and out-of-distribution scenarios. Furthermore, we introduce a method to emphasize critical refusal tokens by incorporating a reward-based token-level weighting mechanism for refusal learning, which further improves the robustness against adversarial exploits. Our research also suggests that robustness to jailbreak attacks is correlated with token distribution shifts in the training process and internal representations of refusal and harmful tokens, offering valuable directions for future research in LLM safety alignment. The code is available at https://github.com/wicai24/DOOR-Alignment

Figures (15)

• Figure 1: Weighted Dual-Objective Optimization for Refusal (W-DOOR) framework. W-DOOR integrates token-weighted refusal training with harmful response unlearning.
• Figure 2: Prefilling Attack Success Rate (ASR) for different alignment methods across various prefilling lengths in 3, 4, 5, 7, 10, 15, 20, 25. The attack is conducted by attaching a harmful prefix before the model generates a response to the harmful prompt. The ASR is calculated based on the LLM-Judge evaluation of the generated response. DOOR and W-DOOR consistently achieve low ASR across all prefilling attacks.
• Figure 3: Accuracy of different alignment methods evaluated on Hellaswag benchmark.
• Figure 4: Attack success rate (ASR) on Multi-turn SORRY-Bench across different alignment methods. The number of turns ranges from 2 to 10 and is not uniformly distributed. Details on the experimental setup are provided in Appendix \ref{['sec:app_exp']}.
• Figure 5: Prefilling attack success rate vs. Hellaswag accuracy over 10 epochs of training.