Improving LLM Safety Alignment with Dual-Objective Optimization

We thank the reviewer for the insightful feedback and thoughtful questions. We address each of your concerns in detail below:

1. Justification for Combining Token-Level Refusal Training and Sequence-Level Unlearning

This is a great question. Our design reflects the intuition that refusal and unlearning operate at different granularities: Refusal behaviors are typically localized in specific tokens or phrases (e.g., “Sorry,” “I cannot help with that.”). These early “transition” tokens are pivotal in shifting the model’s trajectory from harmful continuation to safe refusal. Hence, we apply token-level weighting to reinforce them.

Harmful generations, on the other hand, often unfold over longer sequences and require sequence-level treatment to be effectively unlearned, particularly since the model may learn harmful behaviors in aggregate.

While we do not yet provide a formal theoretical framework for this combination, empirical results show that W-DOOR consistently outperforms DOOR, demonstrating that token-level granularity enhances safety without disrupting the effect of sequence-level unlearning. We will revise the paper to better clarify this intuition and interaction.

2. Lack of Theoretical Guarantees for Gradient Analysis

You're right to note that the gradient analysis in Section 3 is illustrative rather than formal. Our aim was to provide a mechanistic understanding rather than formal proof. Formal convergence or divergence speed analyses, such as those in Negative Preference Optimization (Zhang et al., 2024a), require strong assumptions about model distributions which may not hold in real-world LLMs. As such, we chose to emphasize empirical validation across multiple attack types and scenarios, which supports our claims more robustly in practical settings.

3. Lack of Ablations and Hyperparameter Sensitivity Analysis

We appreciate this suggestion. While we did not initially label them as ablations, our experiments inherently contain component-level comparisons, including:

• Data Augmentation: Comparing DPO vs. DPO w/o Aug isolates the effect of harmful prefix augmentation.

• Robust Refusal Training (SFT component): Comparing NPO (only unlearning) vs. DOOR (NPO + refusal SFT) isolates the benefit of adding the refusal training objective.

• Targeted Unlearning (NPO component): Comparing SFT (only refusal training) vs. DOOR (SFT + NPO unlearning) isolates the benefit of adding NPO.

• Token Weighting: Comparing DOOR vs. W-DOOR isolates the effect of the reward-based token weighting.

Furthermore, to directly address hyperparameter sensitivity for W-DOOR, we conducted additional experiments varying the calculation of the token weight $\beta_t$. We tested the default setting (exponential weighting, $\tau=5$, as in paper), $\tau=1$ for the exponential weighting, and replacing the exponential function with a sigmoid normalization. The results (Table below) show that the method performs robustly across these variations, consistently outperforming baselines.

W-DOOR Hyperparameter Sensitivity (Llama-3-8B)

(Note: Gemma results show similar robustness, omitted for brevity)

We also have further results in this figure: https://anonymous.4open.science/r/icml25-8B51/safety_align.pdf

4. How Does Token-Level Weighting Help Prevent Harmful Continuations?

Excellent question. Token-level weighting helps the model learn to transition away from unsafe continuations, especially in the augmented training setup where harmful prefixes are inserted.

Empirically, we observe that refusal tokens (e.g., “Sorry,” “I cannot...”) exhibit the highest divergence between the reference and target policies (Figures 18, 19, 20). W-DOOR amplifies gradients at those tokens, which steers the model to recognize harmful contexts and pivot to refusal. This is especially useful in prefix attacks where the model is “in the middle” of an unsafe sequence and needs to course-correct.

---

We appreciate the reviewer’s thoughtful critique and hope our clarifications strengthen your confidence in our methodology. We will revise the paper to better connect our design decisions, highlight ablation results more clearly, and clarify the interaction between granular objectives.

We thank the reviewer for the thoughtful comments and constructive suggestions. Below, we address each concern in detail:

However, better justification is needed for using part of the benchmark data as training data, and the small number of data used in the fine-tuning.

We appreciate your concern. Our design choice reflects two priorities:

• Minimizing distribution drift: We intentionally use a small subset of benchmark data to prevent large shifts away from the original model’s capabilities, which often occurs in post-alignment tuning with large datasets.

• Avoiding evaluation leakage: We ensure strict separation between training and evaluation splits within each benchmark. For example, we use a held-out subset of SORRY-Bench for training and evaluate on different, unseen prompts from the same benchmark.

Additionally, we include out-of-distribution (OOD) evaluations, such as HarmBench and XSTest, to demonstrate generalization beyond the training set.

Results provided in Figure 3 seem too noisy to infer any meaningful conclusion. Also, Figure 3 seems to be not discussed anywhere in the text.

Thank you for pointing this out. Apologies for this oversight and the lack of clarity.. We intended to discuss this figure in the paragraph "Robustness Against Stronger Multi-Turn Attacks" (Lines 321-329). The key takeaway, despite some noise due to limited samples at higher turn counts, is that W-DOOR generally maintains lower ASR compared to other methods, especially as the turn number increases, whereas methods like DPO show less robustness in these longer multi-turn interactions. This suggests W-DOOR's alignment is somewhat more resilient in conversational attack scenarios.

There is a mismatch between the W-DOOR objective in (2) and in Figure 1, compared to the expression at the beginning of Section 3.4.

Great catch — thank you. The discrepancy between Equation (2), Figure 1, and the start of Section 3.4 arose due to a simplification: we use $x \oplus y^h_{<k}$​ as the new input $x$ during robust refusal training. We will clarify this notation to ensure consistency and prevent confusion in the revised version.

Statement "..we show that the robustness demonstrated in the prefilling attack generalizes to other forms of adversarial attacks (Figures 1)" is not correct (second column lines 222-224)

You're absolutely right — this should refer to Table 1, not Figure 1. We appreciate you catching this, and will correct it in the revised manuscript.

The results in Section B in Supplementary do not seem to accompany any analysis/discussion of the results.

We acknowledge the lack of accompanying analysis in the current draft. Due to time constraints before the deadline, we were unable to include our full commentary. We now have an updated version with complete discussions and improved clarity, and we will ensure this is reflected in the final camera-ready version.

It is recommended to refer to a particular Figure/Table when some result is discussed (e.g. second column lines 257-274), so the reader can better understand the context.

Thank you — this is a valuable suggestion. We have revised the experimental section to explicitly reference specific figures and tables (e.g., lines 257–274), which greatly improves readability and traceability of claims to evidence.

How can DOOR/W-DOOR perform well on the XSTest benchmark, even when the objectives only focus on enhancing safety refusal without any consideration for preventing over-refusal?

Excellent question. Our alignment objectives (DOOR and W-DOOR) do not directly optimize for minimizing over-refusal. However, our training pipeline includes a retain loss that helps preserve general capabilities and reduces over-conservatism.

While W-DOOR shows slightly higher refusal rates on XSTest compared to the original model, we believe this is a reasonable trade-off given the substantial gains in robustness. We include a Pareto frontier analysis in Appendix Figure 14, plotting ASR vs. over-refusal rate. W-DOOR consistently dominates DPO and other baselines, achieving lower ASR with minimal over-refusal, placing it closer to the optimal trade-off.

We will clarify this trade-off more explicitly in the revised version.

---

We hope these clarifications address your concerns and help strengthen our paper. Thank you again for your constructive feedback and insightful questions.

We sincerely thank the reviewer for the detailed and insightful questions. Below, we clarify key aspects of our methodology, contributions, and evaluation setup:

1. Novel Contributions Relative to Prior Work

Thank you for prompting us to make the contributions more explicit. Our work builds on prior safety alignment techniques (e.g., DPO and NPO) and introduces the following innovations:

• Dual-Objective Integration: We combine robust refusal training via prefix augmentation with harmful content unlearning using NPO into a unified training objective (DOOR). To the best of our knowledge, this is the first work to effectively integrate these two complementary strategies in a cohesive framework.

• Token-Level Adaptive Weighting: We propose W-DOOR, which enhances refusal learning by emphasizing critical refusal tokens via a KL-based reward weighting mechanism. This token-level gradient refinement is novel and leads to improved safety across adversarial attacks.

• Diagnostic Analysis of DPO Limitations: We provide a gradient-based critique of DPO’s underperformance in safety contexts and show, both empirically and analytically, how DOOR and W-DOOR address these issues.

We will revise the paper to more clearly distinguish these novel elements from prior work.

2. Threat Model and Scope of Evaluation

Our training objectives are primarily motivated by prefilling attacks, which exploit partial unsafe generations. However, the robustness imparted by DOOR and W-DOOR generalizes well to other jailbreak scenarios, including suffix-based (AutoDAN, GCG) and multi-turn attacks (Crescendo), as shown in Table 1 and Appendix Figures 2–4.

While prefilling motivates our data augmentation design, the dual-objective formulation is general-purpose and can be applied to other jailbreak modes, given appropriate adversarial data.

3. On "Removing" Harmful Knowledge vs. Making It Harder to Access

We appreciate the reviewer referencing recent literature (e.g., Wu et al., 2024) and agree with the critique. Like NPO, our approach does not guarantee full removal of harmful knowledge but rather reduces its accessibility under typical prompting. We will revise the language in the paper to make this distinction explicit and avoid overstating unlearning effects. Thank you for catching this.

4. Confidence Intervals and Variance Analysis

We conducted additional experiments over 5 random seeds and report the mean ± standard deviation below:

Gemma-2B Results:

LLaMA-3-8B Results

These results confirm that our reported improvements are robust and statistically significant. We will include these confidence intervals in the updated manuscript.

5. Details on KL Divergence Computation (Figure 5)

We compute forward KL divergence, $D_\text{KL}(\pi_\theta \| \pi_\text{base})$, measuring how the aligned model diverges from the base model. This is computed on the evaluation subset of harmful prompts, which mirrors the training distribution but excludes any training samples.

We will make the KL direction and evaluation set details explicit in the revised paper.

6. Why Not Use the Oracle Policy $\pi^*$ for Inference Directly?

Thank you for raising this subtle point. In theory, an ideal oracle policy $\pi^*$ would be used directly. However, in practice: $\pi^*$ is an idealized abstraction and not available in most scenarios. We instead approximate $\pi^*$ using a stronger model aligned via DPO, serving as a proxy to compute token-level importance scores for weighting. Our goal is to train a smaller or aligned model that mimics the safety behavior of $\pi^*$ without requiring access to it at inference time. We’ll revise our explanation of this to improve clarity.

---

We hope these clarifications address your concerns and better highlight the novelty, scope, and rigor of our work. Thank you again for the thoughtful and constructive feedback.

We sincerely thank the reviewer for the thoughtful and constructive feedback, as well as for recognizing the contributions of our work. Below, we address the specific concerns raised:

Change of loss function in LLM also needs to test the performance of other capabilities (e.g., math, coding). How does the proposed approach work with the other capabilities? How do we use the proposed method and combine it with other training data of different capabilities?

Thank you for this important point regarding the impact on general capabilities beyond standard instruction following and safety. To address this, we conducted additional experiments evaluating the performance of our aligned models on standard math (MATH [1]) and coding (HumanEval [2]) benchmarks. The results for the Gemma-2-2B model are presented below:

The results show that alignment methods generally cause a reduction in performance on these specialized tasks compared to the original model. This is likely because our general utility dataset (Alpaca) does not contain significant amounts of math or coding data, and the alignment process (including ours) shifts the model's distribution. However, W-DOOR demonstrates the best retention of these capabilities among the evaluated alignment methods (SFT, DPO, DOOR, W-DOOR). This suggests that our weighted approach, by focusing gradients more precisely, helps mitigate excessive capability degradation compared to other techniques. Moreover, incorporating coding and math data into the utility dataset could potentially alleviate the performance drop.

The proposed method is sensitive/dependent on the reference model performance using the NPO loss. What if the reference model contains harmful responses already?

Thank you for raising this point. To clarify: in our implementation, the NPO loss is only applied to harmful responses generated by the reference model, i.e., cases where the reference itself outputs unsafe content. The intent of NPO in this context is not to imitate the reference but to unlearn these harmful generations by penalizing the target model's likelihood of reproducing them. Therefore, even if the reference model contains harmful outputs, NPO is used to actively reduce their influence in the aligned model. We will clarify this more explicitly in the final version.

In Table 3, W-DOOR increase over refusal evaluation? Can you provide more explanation on this?

Thank you for noticing this. While W-DOOR does show a slightly higher refusal rate on XSTest than DOOR, it still maintains a significantly better trade-off compared to DPO, achieving lower attack success rates (ASR) and lower over-refusal than baseline methods overall. Moreover, we include a Pareto analysis in Appendix Figure 14 (Gemma: Prefill ASR vs. XSTest Refusal Rate), where each model's robustness is plotted against its over-refusal rate. DOOR/W-DOOR consistently lies on a better Pareto frontier—achieving strong safety without excessive conservatism. This indicates that the slight increase in over-refusal is a reasonable and effective trade-off for significantly improved robustness.

We will add a clearer explanation of this in the revised version, highlighting the nuanced balance between robustness and utility.

---

We are grateful for your detailed feedback and hope our clarifications address your concerns. We are committed to further improving our paper based on this helpful input.