We sincerely thank the reviewer for their valuable comments and insights. Below, we address the points raised in the review.

Explanation Needed for How P(⋅∣w<t−1) is Derived.

We thank the reviewer for carefully reading our method and raising this important question. To address the concern, we would like to guide the reviewer to the implementation details provided in Line 210 of the paper. In this section, we explain that we treat P as a one-hot vector of the next token as present in the safety related dataset. We greatly appreciate the reviewer’s feedback and will revise this section to improve its readability.

Increasing the Finetuning Dataset Size

We conducted an additional evaluation of our models trained with a larger dataset to explore the impact of increased data size on performance. Limited by the available instruction-following data in the Alpaca dataset, we expanded the fine-tuning dataset to 2.5 times its original size, resulting in a total of 50k data samples from Alpaca and 2,500 unsafe (toxic prompts and unsafe responses) pairs.

As shown in Table 2 of the main paper, our TA-SFT method with EMD loss already achieves peak performance on the four safety evaluation benchmarks (with 0 harmful responses). To further challenge our method, we evaluated its performance on the larger, more diverse and newly released OR-Bench-Toxic benchmark. This benchmark includes 655 toxic prompts across 10 toxic types, providing broader coverage and a more rigorous evaluation. We use the following metric to evaluate the performance.

• Safety Level: Measured by the number of harmful responses (lower is better).
• Response Quality: Measured using the same method and settings on the AlpacaEval benchmark as in the main paper (higher is better).

As shown in the table below, models trained with the larger dataset achieved:

1. Slightly better safety levels on the OR-Bench-Toxic benchmark.
2. Comparable response quality to models trained on the original dataset.

These results demonstrate that increasing the dataset size can further enhance the model’s safety levels without compromising response quality. This finding suggests that scaling the fine-tuning dataset is a promising approach for improving safety in large language models.

We will update the PDF with our new results and typo fixes after performing all experiments asked by all reviewers.

We sincerely thank the reviewer for their valuable comments and insights. Below, we address the points raised in the review.

Calculating EMD on Part of the Unsafe Responses

The idea of focusing on specific tokens in harmful responses when calculating the EMD loss is indeed attractive. However, defining what constitutes a "harmful word" is quite challenging. For instance, while the word "kill" may generally be considered harmful, in a context such as "how to kill a python process," it is not harmful. This context-dependence complicates the process of reliably identifying harmful keywords.

Yet, taking inspiration from the reviewer’s suggestion, we consider penalizing only parts of a response. To address this idea in a simpler and more targeted way, we focused on the beginning of the response. Intuitively, the initial words of an LLM's response often set the tone or "attitude" of the entire reply, determining whether it refuses to answer or engages with the prompt. In this additional experiment, we only penalize the first 5 words of the response, aiming to influence the LLM's overall attitude toward toxic prompts while minimizing unintended effects on its original capabilities.

We evaluated the safety levels of LLaMA-7B and LLaMA-13B using the same fine-tuning dataset described in the main paper, measured on the I-CoNa benchmark, and recorded the corresponding response quality. Our findings indicate that achieving comparable safety levels in the fine-tuned models leads to a significant decrease in response quality if only penalize the first 5 words. During the experiments, we observed that when penalizing only the first 5 words of the response, the weight of the EMD term must be adjusted to a much higher value to reach a comparable safety level. Intuitively, this is because penalizing only the first 5 words involves a smaller subset of tokens compared to penalizing the entire response. But a high penalty weight would cause the next token prediction distribution of some tokens to change significantly. We believe that increasing the diversity and quantity of toxic prompts and their corresponding harmful responses in the fine-tuning dataset could mitigate this issue. This represents an interesting direction for future exploration.

Upscaling the Loss for Safety Related QA Pairs

We would like to clarify that in the standard SFT approach, the model is trained only on safety-unrelated datasets, which do not include toxic prompts or harmful responses. In contrast, our proposed TA-SFT method is trained on a mixed dataset comprising both safety-unrelated datasets and safety-related datasets. To that end, in our TA-SFT method, besides the SFT term in the loss, there is an additional EMD term (or NLCL term) to penalize the harmful responses. A weight $\\lambda$ is multiplied with the EMD or NLCL term. Adjusting this weight $\\lambda$ is equivalent to upscaling the loss for safety related QA pairs. We have tuned these weights for both the EMD term and NLCL term to achieve extremely high safety levels, as demonstrated by the results on four benchmarks in Table 2 and Figure 2 in the paper.

We will update the PDF with our new results and typo fixes after performing all experiments asked by all reviewers.

We sincerely thank the reviewer for their valuable comments and insights. Below, we address the points raised in the review.

Formatting and Expression Issues

We sincerely thank the reviewer for their valuable comments and insights. Below, we address the points raised in the review.

We have conducted a comprehensive review of the manuscript to address all formatting and expression issues, ensuring a polished and professional presentation. We will update the manuscript once we complete all experiments asked by all reviewers.

Complexity of EMD Computation

The additional training time required for our TA-SFT method is minimal, only 1–2% longer than that of Original SFT. Therefore, TA-SFT can scale up to a very large dataset. The computational requirements of TA-SFT are modest, as detailed below:

In addition to the original SFT loss, TA-SFT involves calculating the Earth Mover’s Distance (EMD), which requires two matrix multiplications and one squared Euclidean distance calculation. These operations are efficiently performed on GPUs. Once the EMD calculation is complete, the backpropagation in TA-SFT is identical to that of Original SFT. Furthermore, the forward pass in TA-SFT remains the same as in Original SFT. Thus, the additional computational workload introduced by TA-SFT is negligible.

We conducted experiments to measure the training time for both SFT and TA-SFT using consistent hardware and configurations:

• LLaMA-7B: Trained on NVIDIA L40 GPUs.
• LLaMA-13B: Trained on NVIDIA H100 96GB GPUs.

All experiments used the same batch size, gradient accumulation steps, and training for 3 epochs. The key difference is that SFT was trained on 20k Alpaca instruction-following data, while TA-SFT included an additional 1k unsafe (toxic prompt, harmful response) pairs, leading to slightly longer total training steps for TA-SFT. As shown in the table below, the Average Training Time per Step indicates that TA-SFT is only 1.17% slower for LLaMA-7B and 2.36% slower for LLaMA-13B compared to SFT.

Limited Over-Alignment Solution

We wish to point out that our work's goal is not to solve the overalignment issue. Instead, our goal is to show that safety can be improved in the SFT stage with a very limited number of safety-related data points, which we demonstrate successfully. Indeed, overalignment remains an open issue that researchers are still struggling to solve.

We sincerely thank the reviewer for their valuable comments and insights. Below, we address the points raised in the review.

Evaluation with Jailbreaking Attacks

First, we would like to clarify our approach focuses on incorporating safety alignment during the supervised fine-tuning (SFT) stage of large language models (LLMs). This is distinct from traditional defense methods which aim to make LLMs safer after SFT or RLHF. Importantly, the LLMs trained using our method can be further combined with existing defense mechanisms to enhance robustness against jailbreaking attacks.

We appreciate the reviewer's suggestion to use jailbreaking to evaluate the generalization and robustness of our method. Following this recommendation, we selected the jailbreaking method outlined in reference $3$ provided by the reviewer: "Jailbreaking Black Box Large Language Models in Twenty Queries"

We compared the performance of our method against the baseline approach, Safety-Tuned LLaMAs (STL), across various data sizes of unsafe examples (these are not jailbreaking examples). Key differences between the two methods include:

• TA-SFT with EMD (ours): Trained using instruction-following data mixed with unsafe examples.
• STL: Trained using instruction-following data mixed with safe examples.

Jailbreaking method from $3$ requires the following three components:

• Attacker LLM: GPT-4O generates jailbreaking prompts based on
  $3$

• Target LLM: Models trained with our method.

• Judging LLM: GPT-4O evaluates responses and scores their harmfulness on a scale of 1 to 10, where a score of 10 indicates a successful attack. Lower scores signify less harmfulness and better robustness against jailbreaking.

Due to time constraints, we conducted the performance evaluation against jailbreaking examples only for the LLaMA-7B model.. The results, summarized in the table below, demonstrate that our method significantly outperforms STL in both Attack Success Rate, ASR (lower is better) and Mean Judge Score (lower is better). Notably, our method achieves a 19.57% ASR and a 6.02 Mean Judge Score, indicating superior robustness. Additionally, the robustness improves as the number of unsafe examples used during training increases.

Experiments With Various Sample Sizes for Over Alignment Issue.

We sincerely appreciate the reviewer’s recognition that the overalignment issue identified in the paper is a valuable starting point. To further evaluate this issue, we conducted additional experiments using LLaMA-7B trained with our TA-SFT approach, leveraging 1000, 500, 300, and 100 unsafe examples. We also examined LLaMA-7B models trained with Safety-Tuned-LLaMA (STL) using 10k, 4k, 2k, 1k, and 300 safe examples. However, consistent with the findings presented in Figure 7 of the main paper, we observed that the Toxic-Rejection-Rate versus Seemingly-Toxic-Rejection-Rate results remain on the same curve. Based on these findings, we conclude that the number of safe or unsafe examples used in training does not explain the over alignment issue.

While the over alignment issue is a remaining noteworthy challenge, solving it is not the primary focus of this paper. Identifying the root cause of overalignment remains an exciting avenue for future research.

We will update the PDF with our new results and typo fixes after performing all experiments asked by all reviewers.

The proposed experiment sounds good. I might be nitpicking but I am not sure the comparison between (training with unsafe answers from LLaMA) and (training with safe answers from GPT4o) is completely fair. Nevertheless, in a perfectly ideal scenario we need answers with similar information content but one safe and other unsafe, which is hard and time-taking to design.

The main question here is

Q) Which is better in safety training: (a) pushing towards safe answers and/or (b) pushing away from unsafe answers.

We could not have asked this question without this work, hence this is a great contribution to the community. I would love to look at any preliminary results you have by the end of rebuttal.

Thank you for pointing out the fair comparison issue in our experimental plan. Your insightful comment has further refined our understanding and approach.

While collecting human-annotated data is indeed the most ideal solution, it is very expensive and time-consuming. To address this, we still seek to utilize an existing dataset and identified the PKU-SafeRLHF dataset as particularly suitable for our purpose. This dataset, available on HuggingFace, includes toxic prompts. Each prompt is paired with two responses, accompanied by labels indicating which response is better and whether it is safe.

To ensure the dataset meets our requirements for fair comparison, we applied the following filtering steps:

1. Keep entries where both safe and unsafe responses are present for the same prompt.
2. Exclude entries where the safe response begins with a direct rejection (e.g., "I'm sorry, I cannot..."), focusing on actual safe responses, where we use a dictionary of direct rejection responses from the prior work $1$ .
3. Segregate the data based on the source of the responses, as these are generated by Alpaca, Alpaca2-7B, and Alpaca3-8B. For simplicity, we refer to the cleaned subsets as safety_1, safety_2, and safety_3, respectively.

After cleaning, we set up three datasets with approximately 3,000 toxic prompts in each of them, each paired with an actual safe response and a corresponding unsafe response.

Using these datasets, we plan to conduct a fair comparison by training LLaMA-7B with the same prompts and safe/unsafe responses from the same LLMs using two approaches:

1. TA-SFT with EMD (ours): Trained using instruction-following data mixed with unsafe examples.
2. STL: Trained using instruction-following data mixed with safe examples.

However, it is important to note a key difference between this experimental setup and the pipeline described in our main paper. In our main paper, unsafe responses were generated by the target LLM itself, whereas in this experiment, unsafe responses are sourced from external LLMs (Alpaca models). This difference may result in a performance decline when applying our TA-SFT method.

To explore the effect of data size, we will train LLaMA-7B on a mix of safety-unrelated data and varying amounts of safety-related data sampled from safety_1 and Safety_3. Due to limited time, we don’t test on Safety_2. Additionally, since we observed that oversampling safety-related data improves STL performance when training with safe responses (pointed out by the third Reviewer hhQn), we will include oversampling in our comparison for a more comprehensive analysis. In the ‘Finetuning Data’ column of the following table, ‘Alpaca+Safety_1-xxx-y’ indicate that the oversampling rate of y times oversampling for the number of xxx safety-related data sampled from Safety_1.

We report the number of harmful responses on the four safety evaluation benchmarks used in our main paper. Surprisingly, we find that training with actual safe responses using the STL method does not improve the safety level at all. While our method, trained with unsafe responses obtained from external LLMs (different from the dataset pipeline in the main paper), experiences a slight performance decrease, it still achieves a satisfying safety level. This finding suggests that our method is robust even when using datasets sourced from external LLMs.

Due to limited time, we have not explored training with actual safe response and unsafe response simultaneously using our method. However, based on our preliminary findings, we believe that pushing away from unsafe answers appears to be a more robust approach to improving safety. The underlying reason why training with actual safe responses fails to improve safety remains unclear and needs further investigation. We consider this an important direction for future work.

$1$ Röttger P, Kirk H R, Vidgen B, et al. Xstest: A test suite for identifying exaggerated safety behaviours in large language models $J$ . arXiv preprint arXiv:2308.01263, 2023.