PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kaifeng Lyu, Haoyu Zhao, Xinran Gu, Dingli Yu, Anirudh Goyal, Sanjeev Arora

arXiv:2402.18540v2cs.LGcs.AIcs.CL

TL;DR

Benign fine-tuning can weaken an aligned LLM, raising the question of how to preserve safety while improving downstream utility. The paper studies prompt-template choices and proposes PTST: fine-tune without a safety prompt, then include it at test time. Across the reported tasks and models, PTST reduces safety degradation while retaining downstream helpfulness improvements.

  • Problem

    Even fine-tuning on clearly benign datasets can raise unsafe behavior, motivating the question of how to preserve alignment while improving a model for downstream use.

    Source: §1 · pp. 1–2

  • Method

    PTST fine-tunes without a safety prompt and uses one at inference, creating a prompt-template discrepancy between training and testing.

    Source: §1 · p. 2

  • Results

    Across several models and fine-tuning tasks, PTST reduces safety loss while maintaining substantial downstream helpfulness improvements.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The paper presents PTST as a simple amendment to fine-tuning practice for preserving safety alignment.

    Source: §6 · p. 10

    The study does not aim to defend against prompt-engineered jailbreak attacks, and most safety tests use queries the original model already handles with low ASR.

    Source: §2 · p. 3

Abstract

from arXiv · show

Public LLMs such as the Llama 2-Chat underwent alignment training and were considered safe. Recently Qi et al. [2024] reported that even benign fine-tuning on seemingly safe datasets can give rise to unsafe behaviors in the models. The current paper is about methods and best practices to mitigate such loss of alignment. We focus on the setting where a public model is fine-tuned before serving users for specific usage, where the model should improve on the downstream task while maintaining alignment. Through extensive experiments on several chat models (Meta's Llama 2-Chat, Mistral AI's Mistral 7B Instruct v0.2, and OpenAI's GPT-3.5 Turbo), this paper uncovers that the prompt templates used during fine-tuning and inference play a crucial role in preserving safety alignment, and proposes the ``Pure Tuning, Safe Testing'' (PTST) strategy -- fine-tune models without a safety prompt, but include it at test time. This seemingly counterintuitive strategy incorporates an intended distribution shift to encourage alignment preservation. Fine-tuning experiments on GSM8K, ChatDoctor, and OpenOrca show that PTST significantly reduces the rise of unsafe behaviors.

1 Introduction

Benign fine-tuning can weaken an aligned model, and the paper finds that prompt-template choices during fine-tuning and inference shape this risk. It proposes PTST and reports that it reduces safety degradation across tasks while retaining downstream helpfulness.

1 Introduction

  • Safety degradation depends on input format: models are less safe on test inputs formatted like fine-tuning inputs, while a discrepancy between formats can preserve safety.The paper creates this discrepancy by changing prompt templates.Source: §1 · p. 2
  • The study addresses whether prompt-template mismatch can improve safety without sacrificing downstream utility, against evidence that even clearly benign fine-tuning can raise unsafe behavior.Prior work cited in the paper reported such increases even with datasets lacking harmful data (Qi et al. [2024]).Source: §1 · pp. 1–2
  • Across Llama 2-Chat, Mistral 7B Instruct v0.2, and GPT-3.5 Turbo, PTST reduces safety loss while maintaining substantial downstream helpfulness improvements.PTST fine-tunes without a safety prompt and includes one at test time.Source: §1 · p. 2
  • Using the same prompt template for GSM8K fine-tuning and inference breaks safety alignment to a large extent, while PTST is the most effective tested strategy in many cases.The findings are also confirmed on ChatDoctor and OpenOrca.Source: §1 · p. 3
  • With added safety examples, benchmark-style harmful requests can bring ASR close to 0%, but GSM8K-mixed harmful queries still produce high ASR after fine-tuning; PTST significantly reduces this degradation.The authors attribute the difference to whether safety examples cover the test-time distribution of harmful queries.Source: §1 · p. 3

2 Threat Model and Safety Evaluation

The evaluation models an owner fine-tuning an aligned model before online deployment, with a black-box attacker submitting harmful queries. Safety is assessed using harmfulness judgments, while the study excludes jailbreak resistance as a general goal.

2 Threat Model and Safety Evaluation

  • The threat model has an owner choose training and test templates and a helpfulness metric, while a black-box attacker submits harmful queries using the enforced test template.The attacker has no access to model weights or exact fine-tuning or pretraining data.Source: §2 · p. 3
  • A GPT-4 judge scores harmfulness on a five-point scale, and ASR is the percentage of harmful queries receiving a score of 5.The evaluation uses harmful-query datasets to calculate ASR.Source: §2 · p. 3
  • Defending against prompt-engineered jailbreaks is outside the study's scope; most tests use harmful queries the original model already handles with low ASR.The paper separately reports a GCG attack evaluation in Table 1d.Source: §2 · pp. 3–4
  • The study tests AdvBench's 520 direct harmful requests and introduces DirectHarm4, a 400-query set spanning four categories associated with higher ASRs.DirectHarm4 queries are also phrased as direct requests in imperative tone.Source: §2 · p. 4

3 Role of Prompt Templates

Across the experiments, safety outcomes depend on how training and inference prompt templates are paired: PTST omits the safety prompt during fine-tuning and uses it at inference. This approach generally preserves safety while retaining downstream helpfulness gains across models, datasets, and safety prompts.

3.1 Case Study: Fine-tuning on GSM8K

  • On GSM8K, chat:vanilla raises exact match from 20.32% to 33.39% as DirectHarm4 ASR rises from 2.75% to 11.00%; AdvBench reports only 0.26% ASR.With chat:llama, DirectHarm4 ASR rises from 0.00% to 18.08%.Source: §3.1 · pp. 4–5
  • With chat:vanilla training, PTST cuts DirectHarm4 ASR from 18.08% to 1.08%, while helpfulness moves from 33.51% to 30.00%.With Self-Reminder testing after chat:vanilla training, DirectHarm4 ASR is 0.17% and helpfulness is 31.26%.Source: Table 1 · p. 5, §3.1 · p. 5
  • Across training steps, PTST maintains low ASR and a better helpfulness-safety balance than same-template training, including when early stopping is considered.The plotted trajectories connect points in training-step order.Source: §3.1 · p. 6, Figure 2 · p. 6

3.2 Experiments on Other Models: GPT-3.5 and Mistral

  • On GPT-3.5 Turbo, switching chat:vanilla inference to chat:llama lowers DirectHarm4 ASR from 22.75% to 4.50%; Mistral experiments show the same PTST safety-helpfulness pattern.Source: §3.2 · pp. 6–7
  • On Orca-Math, PTST maintains lower ASR with similar helpfulness across 10,000-, 20,000-, and 40,000-example training horizons.Source: §3.2 · p. 6

3.3 Experiments on Other Datasets: ChatDoctor and OpenOrca

  • On ChatDoctor and OpenOrca, the same-template safety degradation recurs, while chat:vanilla fine-tuning followed by chat:llama inference preserves safety with decent helpfulness gains.Source: §3.3 · p. 7

3.4 Experiments on Other Safety Prompts

  • With alternative safety prompts, PTST preserves safety and much helpfulness, whereas training with chat:llama and testing with another safety prompt still degrades safety.Source: §3.4 · pp. 7–8, Figure 3 · p. 8

4 Effects of Mixing Safety Data

The section tests whether adding safety examples can preserve alignment during fine-tuning, including on harmful queries unlike the safety examples. It finds that PTST can further reduce unsafe responses on these out-of-distribution queries.

4.1 Adding Safety Examples Can Reduce the ASR on Similar Queries Without PTST

  • Adding safety examples reduced ASRs on AdvBench and DirectHarm4 to nearly 0%, and PTST could reduce them further, though sometimes only marginally.The examples comprise 2,483 harmful queries with safe responses and resemble the style and format of those two evaluation sets.Source: §4.1 · p. 8

4 Effects of Mixing Safety Data

  • Because safety data may not cover all harmful-query distributions, its protection can be limited on out-of-distribution queries.The section explores combining safety examples with PTST during fine-tuning.Source: §4.2 · pp. 8–9, §4 · p. 8

4.2 PTST Leads to More Substantial Improvements on OOD Queries

  • Even with safety data, matching training and test templates increased GSM-Danger ASR by more than 10% for chat:vanilla, chat:alpaca, and chat:llama.GSM-Danger queries begin like GSM8K math problems and end with a harmful request, making them unlike direct safety examples.Source: §4.2 · p. 9, Table 6 · p. 9
  • 5% without safety data and 4% with it were the GSM-Danger ASRs for chat:vanilla training and chat:llama testing; chat:llama training and testing yielded 12% even with safety data.Changing the training template to chat:alpaca produced 1% ASR with or without safety data.Source: §4.2 · p. 9, Table 6 · p. 9

5 Related Works

Prior work studies prompt-based alignment and safety degradation from fine-tuning, while approaches to preserve safety include data and model defenses. This paper’s related-work discussion identifies prompt-template differences between fine-tuning and inference as underexplored.

Prompting for LLM alignment

  • Prompt engineering can improve safety, including through system prompts for Llama 2-Chat [Touvron et al., 2023], but prompt differences between fine-tuning and inference remain underexplored.Other approaches include instruction and in-context-example prompts [Askell et al., 2021], Directed Representation Optimization [Zheng et al., 2024], and self-reflection prompts [Ganguli et al., 2023, Wu et al., 2023].Source: §5 · p. 9

Removing safety guardrails via fine-tuning

  • Prior studies found that benign fine-tuning can degrade safety, but often did not evaluate the corresponding downstream task performance [Qi et al., 2024, Pelrine et al., 2023].The paper reports that its reproduction of fine-tuning Llama-2-7B-chat on Alpaca did not improve instruction-following ability measured by AlpacaEval.Source: §5 · p. 10

Preserving safety during fine-tuning

  • Proposed defenses include robust alignment training, safety-data mixtures, and projecting LoRA weights into a safe subspace; the authors say these can be combined with PTST.Examples include Vaccine [Huang et al., 2024b], safety-data approaches [Wang et al., 2024; Zong et al., 2024], and the training-free method of Hsu et al. [2024].Source: §5 · p. 10

6 Conclusions

The paper reports that fine-tuning can degrade safety across several aligned models and proposes PTST: fine-tune without a safety prompt, then include one at test time. It also notes that the mechanism behind PTST’s template generalization remains unclear.

6 Conclusions

  • Fine-tuning can degrade safety in Llama, Mistral, and GPT-3.5 Turbo, and PTST fine-tunes without a safety prompt but includes one at test time.The paper presents PTST as an empirical amendment to current fine-tuning practice.Source: §6 · p. 10
  • The authors suggest PTST’s success may reflect compositional generalization across templates, but its underlying mechanisms are not well understood.They identify further empirical and theoretical exploration as warranted.Source: §6 · p. 10
  • The authors call for safety to be considered throughout training and suggest more robust alignment and prompt-template augmentation as future directions.They propose exploring alignment algorithms that make PTST more effective and help models handle template changes while attending to safety prompts.Source: §6 · p. 10

A Limitation

The appendix notes that experimental costs constrained hyperparameter sweeps and repeated runs, alongside ethics discussion, prompt-use context, and additional related work.

A Limitation

  • Compute and API-evaluation costs limited hyperparameter sweeps and repeated seeds; OpenOrca fine-tuning for one epoch exceeded 100 A100 GPU hours, and safety evaluation exceeded $5 per test template.Repeated runs were conducted for Llama GSM8K main experiments and ChatDoctor sampling; the authors believe the findings are robust.Source: A Limitation · p. 15

B Ethics and Broader Impact

  • The authors frame the work as reducing risks from harmful content while acknowledging that its examples of fine-tuning-induced degradation could also be misused.They argue the findings’ safety benefits outweigh those potential concerns.Source: B Ethics and Broader Impact · p. 15

C Current Practice of Using Safety Prompts

  • Llama 2-Chat training distills safety prompts into the model, while its evaluations and released code have used or changed the default safety prompt over time.Mistral’s report describes a safety prompt and its codebase exposes a boolean option to enable it; MPT chat tokenizers enforce a default prompt unless overridden.Source: C Current Practice of Using Safety… · p. 15

D Addtional Related Works

  • Jailbreak research manipulates prompts or generation, while defenses include self-reminders, perplexity filtering, goal prioritization, and Llama Guard [Zou et al., 2023; Chao et al., 2023; Yong et al., 2023; Deng et al., 2023; Zeng et al., 2024; Huang et al., 2023; Zhang et al., 2023b; Xie et al., 2023; Jain et al., 2023; Zhang et al., 2023a; Inan et al., 2023].The cited approaches include adversarial suffixes, iterative prompt refinement, low-resource-language translation, persuasion-based paraphrases, and decoding manipulation.Source: D Addtional Related Works · pp. 15–16
  • For Mistral, testing without CL yields high ASR even before fine-tuning, and fine-tuning with any template while omitting CL also yields very high ASR.The passage identifies CL as the relevant test condition but does not expand the abbreviation.Source: Table 7 · p. 16

E Additional Experiments: Fine-tuning Mistral on GSM8K

On Mistral GSM8K fine-tuning, safety on AdvBench depends on the prompt-template relationship between training and testing, and retaining a system prompt at test time matters. PTST produced lower reported ASRs than using the safety template throughout.

E Additional Experiments: Fine-tuning Mistral on GSM8K

F.1 Models and Fine-tuning Tasks

The experiments cover three aligned chat models and downstream tasks in math, medical consultation, and general reasoning. This appendix also reports AlpacaEval results and identifies the prompt-template tables used across tasks.

F.1 Models and Fine-tuning Tasks

  • The study evaluated Llama-2-7B-chat, Mistral 7B Instruct v0.2, and GPT-3.5 Turbo across its case studies.Source: §F.1 · p. 16
  • Fine-tuning Llama-2-7B-chat on Alpaca or Alpaca-GPT4 significantly decreased AlpacaEval instruction-following win rate, with Alpaca-GPT4 results of 70.72% for LoRA and 73.98% for Full fine-tuning.The paper motivates using datasets that improve the models on their corresponding task metrics instead.Source: §F.1 · p. 17
  • For math, GSM8K and Orca-Math fine-tuning was assessed on GSM8K test-set zero-shot performance, using answer extraction suited to each dataset’s format.GSM8K answers were extracted with regular expressions; Orca-Math answers were extracted by GPT-4 and compared with gold answers.Source: §F.1 · p. 17
  • ChatDoctor fine-tuning simulated medical consultation using 100,000 patient-physician conversations, with helpfulness evaluated by BERTScore on 1,000 held-out queries.The model was fine-tuned for three epochs with LoRA and a cosine learning-rate schedule.Source: §F.1 · p. 17
  • OpenOrca fine-tuning used a random 600,000-example subset of the 4.2-million-entry dataset, with reasoning and comprehension evaluated on ARC-easy and ARC-challenge.Llama-7B-chat was trained for one epoch, and evaluation used exact-match scores on generated answers.Source: §F.1 · p. 17

F.2 Prompt Templates

  • The appendix lists prompt templates for Llama 2 Chat on GSM-8K, ChatDoctor, and OpenOrca, and for GPT-3.5 Turbo on GSM-8K.Source: §F.2 · p. 18

F.3 Harmful Query Datasets

The appendix describes harmful-query evaluation and the construction of DirectHarm4 and GSM-Danger. DirectHarm4 targets several harmful-instruction categories, while its requests are deliberately direct rather than jailbreaking prompts.

F.3 Harmful Query Datasets

  • AdvBench safety evaluation uses 520 harmful instructions spanning areas including misinformation, discrimination, cybercrime, and dangerous or illegal suggestions.Source: §F.3 · p. 18
  • DirectHarm4 was created after observing that fine-tuned models could pass AdvBench yet still produce harmful responses to other instructions.The authors selected four of about ten reviewed categories because they elicited harmful responses more easily.Source: §F.3 · p. 18
  • DirectHarm4 contains 100 direct harmful instructions per category, designed without jailbreak tricks as a minimal safety requirement for aligned models.The four categories are malware, drugs, phishing, and disinformation.Source: §F.3 · p. 18
  • For malware and phishing, GPT-4 generated candidate instructions that were manually reviewed; drug and disinformation examples were iteratively generated from selected examples and quality-filtered.The final subset for each category contained 100 examples.Source: §F.3 · p. 18
  • GSM-Danger prompts were generated by sampling three GSM8K test prompts and asking GPT-4 to create social-media-writing prompts encouraging dangerous behaviors.About 200 prompts were collected, manually filtered, and reduced to 100 final examples.Source: §F.3 · pp. 18, 21

F.4 Experimental Details on the GCG Attack from the JailbreakBench

The GCG experiment tests PTST against optimized adversarial suffixes from JailbreakBench. The reported AdvBench attack success rate exceeded 20% when fine-tuning and inference used the same prompt template.

F.4 Experimental Details on the GCG Attack from the JailbreakBench

  • With GCG adversarial suffixes, ASR exceeded 20% when the same prompt templates were used for fine-tuning and inference.The suffixes were paired with 100 JBB-Behaviors harmful queries, whose attack suffixes had been optimized on Llama 2-Chat 7B.Source: §F.4 · p. 21
  • The GCG strings came from JailbreakBench’s JBB-Behaviors set of 100 harmful queries and were optimized for Llama 2-Chat 7B.Source: §F.4 · p. 21
  • Using suffixes optimized for Llama 2-Chat was considered plausible because GCG transfers well and the assumed attacker has only black-box access.Source: §F.4 · p. 21

F.5 Helpfulness Evaluation

This appendix checklist directs readers to the paper’s claims, evaluation and reproducibility details, and supporting disclosures. The helpfulness evaluation uses exact-match scoring for generated GSM8K and ARC answers.

F.5 Helpfulness Evaluation

  • For GSM8K, Llama and Mistral generation stops at <s>, while GPT-3.5 Turbo uses the full API output; extracted answers are scored by exact match.The evaluation identifies the last numerical value in generated text as the response.Source: §F.5 · p. 23
  • For ARC-Easy and ARC-Challenge, models generate answers to transformed multiple-choice prompts, which are scored by exact match against the correct answer.The prompt asks for a selected choice and a final “The answer is” response.Source: §F.5 · p. 23

1. Claims

  • The paper identifies prompt templates as crucial to preserving safety alignment after fine-tuning and proposes PTST to encourage alignment preservation.The abstract and introduction are described as closely reflecting these contributions.Source: §1 · p. 24

2. Limitations

  • The authors direct readers to Appendix A for limitations and state that the paper presents no theorems requiring assumptions or proof.Source: §2 · p. 24, §3 · p. 25

4. Experimental Result Reproducibility

  • Experimental reproducibility details are in Section 3; the fine-tuning code is released, and the listed fine-tuning datasets are publicly available.The paper also points to Appendix F for experimental settings and Appendix F.1 for dataset licenses.Source: §4 · p. 25, §5 · p. 26, §6 · p. 26

7. Experiment Statistical Significance

  • Standard deviations are reported for the main GSM8K and ChatDoctor experiments, but repeating every experiment is described as expensive.A single GPT-3.5-turbo-0613 run on 40k Orca-Math samples costs approximately 150 USD.Source: §7 · p. 26

9. Code Of Ethics

  • The checklist states that the authors adhere to the NeurIPS Code of Ethics, refer to Appendix B for broader impacts, and mark safeguards as not applicable.Source: §9 · p. 27, §10 · p. 27, §11 · p. 28

12. Licenses for existing assets

  • The paper cites all public datasets used and explicitly mentions their licenses in Appendix F.1; crowdsourcing and human-subjects research, including IRB approval or equivalent review, are marked not applicable.Source: §12 · p. 28, §14 · p. 29
Loading 2402.18540v2…