Keeping LLMs Aligned After Fine-tuning: The Crucial Role of Prompt Templates

We sincerely thank the reviewer for the positive and constructive feedback. Below, we address the reviewer’s questions.

Q1 (rephrased): It would be appreciated if the authors could provide an explanation for why using different templates does not simultaneously lower the helpfulness while the harmfulness is lowered. Is it because longer prompts such as CL can elicit better helpfulness performance?

A1: Thanks for the reviewer’s suggestion. We respond to this question in our general response. Our hypothesis is that there is some compositional generalization ability arising from the pretraining and alignment phases (prior to the fine-tuning phase we study). Understanding this would be an interesting future direction.

The length of the prompt may also be a factor, but Figure 3(b) shows that a shorter version of CL (llama-short) can achieve an even better helpfulness than CL.

---

Q2: The VLGuard paper (Zong et al., 2024) also observed that mixing safety data can reduce ASR. It would be nice if the authors could mention and discuss it in Section 4.

A2: Yes, the VLGuard paper curated a vision-language safety instruction-following data and showed that mixing this data during fine-tuning can mitigate safety issues. In our NeurIPS submission, we have already included a discussion on this work in Section 5 (Related Works). We will also cite this paper in the first paragraph of Section 4.

Meanwhile, we want to clarify that the main point of our Section 4 is not to repeat the point that mixing some safety data mitigates the safety issue, but to show that adding the PTST strategy on top of mixing the safety data can further enhance the safety performance.

---

Q3 (rephrased): Table 2 includes too many prompt templates, making it hard for readers to see which ones contain safety prompts.

A3: Below, we respond to the review assuming the reviewer was talking about Table 1, since Table 2 only contains 3 templates.

We apologize for any inconvenience caused by the presentation of Table 1. In the next version, we will add a “*” to the prompt template (CL) that contains a safety prompt. In the current version, blue cells stand for training and testing with the same templates and orange cells stand for PTST. We will revise the caption to explain the meaning of these colors and the “*” symbol to be added.

---

Q4: The literature review seems to be comprehensive but still some related works are missing.

A4: Thanks for listing so many related works before and concurrent to our work. We will definitely cite and discuss them in the next version of our paper!

---

Q5: I suggest the authors add one or two baselines for comparison, e.g., Vaccine (Huang et al., 2024).

A5: Thanks for the suggestion. We added Self-Reminder (SR) and In-context Defense (ICD) as two lightweight baselines. Please see the general response above for the results. Comparing with Vaccine would also be interesting, but we want to note that this method is orthogonal to our PTST. Vaccine is a robust alignment method that works in the alignment phase (prior to the fine-tuning phase we study). It can make the model’s safety more robust to fine-tuning, but even if Vaccine works well, one can still add PTST during the fine-tuning phase to enhance the model’s safety even further. In our new experiments, we choose to focus on comparing with baselines that work in the fine-tuning phase.

---

Q6: In Table 1(b)(c), why does changing CL->CV increase harmfulness while changing CL->CA decreases it?

A6: Thank the reviewer for carefully reviewing our results. First, we would like to clarify our main point in Table 1:

1. There exist many cases where using different training and testing templates leads to lower ASR than using the same one, while still improving the model on helpfulness. This can be seen by comparing diagonal entries and off-diagonal entries in Table 1.
2. However, not all off-diagonal entries are safer than diagonal entries, just as the reviewer points out.
3. Among the off-diagonal entries, the PTST strategy consistently succeeds. That is, when the training template does not have a safety prompt (e.g., TV, TA, CV, CA) and the testing template has a safety prompt (e.g., CL in Table 1 and three other such templates we tried in Figure 3), the ASR is consistently lower than training with the same prompt template and the helpfulness consistently improves. This can be seen by comparing the diagonal entries and the orange-colored off-diagonal entries in Table 1.

Regarding the reviewer’s question on the results of CL->CV and CL->CA, we would like to point out that these strategies are not PTST, and we do not claim that they will consistently reduce or increase the ASR. Besides PTST, there may be other strategies that consistently work, and exploring these other strategies can be an interesting future direction.

For Q3, I am referring to Table 2. After one month, I again forgot how to interpret Table 2. I think it is generally not a good idea to use the table showing training/testing combination. Dumb people (like me) may not easily interpret your results (a lot of abberviation for prompt template make situation worse). Perhaps it will better to provide at least one easier to interpret table to show your performance?

It is also interesting to see that the method can be combined with alignment stage solution, e.g., Vaccine. Could you discuss a little bit about this potential combination in the next version of the paper? (perhaps in the last section when you talk about conclusion and future direction)?

We sincerely thank the reviewer for acknowledging that the paper is clearly written and that the proposed PTST strategy is novel. Below, we address the reviewer’s questions.

Q1: This paper argues that training and testing on the same prompt template makes attacking easier. However, the trainTA-testCV entry suffers the same ASR as trainCV-testCV in Table 1(b), and the trainCV-testTV entry suffers an even higher ASR than trainTV-testTV in 1(d).

A1: Thank the reviewer for carefully reviewing our results. In fact, we would like to clarify that we do NOT claim that using different training and testing templates is always safer than training with the same template, but there exist many such cases, and PTST is a particular case where such a phenomenon consistently happens. More specifically, our claim is the following:

1. There exist many cases where using different training and testing templates leads to lower ASR than using the same one, while still improving the model on helpfulness. This can be seen by comparing diagonal entries and off-diagonal entries in Table 1.
2. Among these cases, the PTST strategy (described in the gray box on Page 2) consistently succeeds. That is, when the training template does not have a safety prompt (e.g., TV, TA, CV, CA) and the testing template has a safety prompt (e.g., CL in Table 1 and three other such templates we tried in Figure 3), the ASR is consistently lower than training with the same prompt template and the helpfulness consistently improves. This can be seen by comparing the diagonal entries and the orange-colored off-diagonal entries in Table 1.

We will make this point more clear in the next version of our paper.

---

Q2: Is there a particular reason that TV and TA is not included in the further experiments like GPT-3.5

A2: Yes, the main reason is that we are fine-tuning GPT-3.5 with OpenAI’s APIs, but these APIs do not support TV and TA and only accept chat-mode data represented as a list of system, user and assistant messages.

---

Q3: Throughout experiments in this paper, it seems like there is always a tradeoff between helpfulness and ASR. Philosophically, is that really the case? I wonder whether a sufficiently intelligence machine still needs to give up either helpfulness or safety? The paper is also lacking in discussion about the implications of PTST, and how the PTST method might inspire future papers to explore the direction of safety-aligned fine tuning.

A3: Thank the reviewer’s suggestion on adding more high-level discussion. We will definitely do in the next version of the paper. For existing models, Table 2 and Figure 2(b) showed results on GPT-3.5 Turbo, which is more intelligent than Llama 2-Chat, but still GPT-3.5 Turbo suffers from safety issues after custom fine-tuning with benign data. In this situation, PTST is helpful to mitigate such issues. Although we cannot perfectly predict the future, we believe that a model that is good at both helpfulness and safety won’t appear just by adding more data, as discussed in many existing papers [1, 2]. We indeed need better designs of pretraining and alignment methods to achieve this goal. For designing better pretraining and alignment methods, it would be helpful if one could know better how the model is going to be used after fine-tuning. Our paper points out the PTST principle for fine-tuning the model after the alignment phase, and an interesting future direction is to identify important factors in pretraining and alignment methods that make PTST work for current models, and improve these methods to make PTST work better in future models. Please also see our general response for a more detailed discussion.

[1]. Jailbroken: How Does LLM Safety Training Fail? https://arxiv.org/abs/2307.02483

[2]. Sleeper Agents: Training Deceptive LLMs that Persist Through Safety Training https://arxiv.org/abs/2401.05566

We sincerely thank the reviewer for acknowledging the novelty and promising performance of PTST. Below, we address the reviewer’s questions.

Q1: I feel it necessary for the authors to at least propose some hypotheses on the underlying mechanism of PTST and try to verify them with concrete experiments. I also wonder whether this principle can help us design better prompt templates for fine-tuning.

A1: Thanks for the reviewer’s suggestion. We have responded to this question in our general response.

• Our hypothesis is that there is some compositional generalization ability arising from the pretraining and alignment phases (prior to the fine-tuning phase we study). Understanding this would be an interesting future direction.
• The PTST principle may be helpful for designing better prompt templates. It recommends searching for new training templates without mentioning safety and for new testing templates among those emphasizing safety.

---

Q2: There are cases when the helpfulness of models is notably decreased if we adopt the PTST rule, such as (TV, CL) and (TA, CL) for Llama-7B.

A2: We would like to thank the reviewer for carefully reviewing our results. It is true that the helpfulness of the models with training and testing templates TV -> CL and TA -> CL are notably lower than those with the same training and testing template. However, we want to note the following two points:

• Still, the helpfulness is improved after fine-tuning according to the PTST rule. E.g., GSM8K accuracy is 15.31 under No FT -> TV and 6.52 under No FT -> CL, but TV -> CL gives 23.76. This suggests that PTST consistently improves the helpfulness upon the model without fine-tuning, though in some cases PTST may lead to worse helpfulness than the model fine-tuned and tested with the same template.
• Whether PTST can lead to comparable helpfulness as using the same training and testing templates can depend on the model. For example, on Mistral-7b-Instruct-v0.2, both TV -> CL and TA -> CL give comparable helpfulness improvement as using the same training and test templates (see Table 6). Thus, an interesting future direction is to explore whether we can improve the aligned model in the alignment or even the pretraining phase to make the helpfulness improvement more robust to the change of template. See also the general response above for a detailed discussion.

---

Q3: Some lightweight defenses such as Self-Reminder [1] and ICD [2] can be incorporated into the (CL, CL) training scheme, which will serve as good baselines for PTST. Comparison with safeguarding algorithms can help readers better understand the significance of PTST.

A3: Thanks for the suggestion. We added Self-Reminder (SR) and In-context Defense (ICD) as two other lightweight baselines. Please see the general response above for the results.

We sincerely thank the reviewer for reviewing our paper and for acknowledging our experiments as extensive. However, we’d like to point out that the reviewer’s main comment (1, 2 below) under “weakness” (as well as Question 2, Question 3 and Limitation 2) suggests they may not have absorbed the core idea of our paper.

1. Why fine-tune on math datasets (gsm8k, Orca-Math) to verify the model's safety? How does the performance compare when fine-tuned on safety-specific datasets, such as Anthropic/hh-rlhf?
2. The experiments on PTST are insufficient, as they do not adequately compare the effectiveness of the approach with current alignment algorithms such as PPO, DPO, KTO, among others.

The starting point of our paper is that often end-users try to improve the capabilities of aligned models (such as llama chat) on specialized tasks like GSM8K. This fine-tuning leads to a loss of safety alignment (i.e., a rise in unsafe responses) if one follows the common practice that uses the same training and testing prompt templates, as discussed in the introduction (Section 1; see also Appendix C). Our paper proposes a lightweight trick (called PTST) that avoids such a drastic loss of alignment during post-alignment fine-tuning (see, e.g., Table 1 for safety with/without our method).

Clarification of Misunderstanding

The reviewer asks for the comparison of PTST and alignment algorithms in Weakness 2, Question 2, Question 3 and Limitation 2 and for experiments on the human preference dataset hh-rlhf in Weakness 1 and Question 1. However, PTST and alignment algorithms are used in different stages of LLM training and deployment. As highlighted in Sections 1 and 2, our paper studies the scenario where a model owner fine-tunes an already aligned LLM to enhance a certain helpfulness metric, e.g., fine-tune on GSM8k to improve math capabilities, measured by the accuracy on the GSM8k test set. PTST helps the fine-tuned LLM retain the safety attribute previously acquired during the alignment training. By contrast, alignment algorithms, e.g., PPO, DPO, KTO, are used to align base LLMs, e.g., Llama 2, with human preferences prior to the fine-tuning process we investigate.

We believe this clarification already resolves the reviewer’s questions raised in Weaknesses 1, 2, Questions 1, 2, 3, and Limitation 2. Below we respond to the other questions.

Other Concerns

W3: This paper proposes the PTST algorithm, but it is a training technique and lacks a certain level of innovation.

A: We argue that PTST is a non-trivial training technique. A common practice for fine-tuning is to use the same prompt templates for training and inference to maximize the downstream performance. By contrast, PTST encourages safety by creating a distribution shift in training and inference while still maintaining the downstream performance. Furthermore, PTST is lightweight and easy to implement, allowing for further refinement of aligned LLMs with significantly reduced safety concerns.

L1: Is the PTST algorithm still effective with an increasing amount of data or the introduction of mixed datasets?

A: Yes. As detailed in Section 3.3 and Appendix F, besides GSM8k, which contains 8k samples, we also conducted experiments with larger datasets, including ChatDoctor with 100k samples and a 600k-sample subset of OpenOrca. I. In Section 4, we fine-tuned the models on a mixture of GSM8k and safety data. PTST consistently maintains safety across all these settings.