[Paper Review] Keeping LLMs Aligned After Fine-tuning: The Crucial Role of Prompt Templates

Kaifeng Lyu, Haoyu Zhao+4 more|arXiv (Cornell University)|Feb 28, 2024

Taxation and Legal IssuesBusiness, Management and Accounting3 citations

TL;DR

This paper proposes the Pure Tuning, Safe Testing (PTST) principle to preserve safety alignment in large language models after fine-tuning. By fine-tuning without a safety prompt but applying it at inference, PTST significantly reduces or eliminates unsafe behaviors—even in models fine-tuned on benign datasets like Alpaca and GSM8K—while maintaining strong downstream performance.

ABSTRACT

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.

Motivation & Objective

• To investigate why fine-tuning aligned LLMs on seemingly benign datasets leads to safety degradation.
• To identify the role of prompt templates in preserving safety alignment during and after fine-tuning.
• To propose a practical method that maintains model helpfulness while minimizing unsafe behavior post-fine-tuning.
• To evaluate the effectiveness of PTST across multiple LLMs, including Llama 2-Chat, Mistral 7B Instruct, and GPT-3.5 Turbo.
• To provide empirical evidence that distribution shift between fine-tuning and inference templates can be leveraged to improve safety.

Proposed method

• Fine-tune LLMs using only the user instruction and input, omitting any safety prompt during the fine-tuning phase.
• Apply a standardized safety prompt at inference time, even when the model was not fine-tuned with it.
• Use the same prompt template for inference across all experiments to ensure consistent evaluation of safety.
• Measure safety degradation using Attack Success Rate (ASR) on adversarial prompts from benchmarks like AdvBench and DirectHarm4.
• Compare PTST against standard fine-tuning practices that use the same prompt template for both fine-tuning and inference.
• Conduct ablation studies across multiple datasets (GSM8K, ChatDoctor, OpenOrca) and models to validate robustness.

Experimental results

Research questions

• RQ1 Does using a safety prompt during fine-tuning degrade the model’s safety alignment, even when the training data is benign?
• RQ2 Can removing the safety prompt during fine-tuning while retaining it at inference improve post-fine-tuning safety?
• RQ3 How does the choice of prompt template during fine-tuning affect the model’s generalization to different inference templates?
• RQ4 To what extent does PTST reduce attack success rates on harmful queries compared to standard fine-tuning?
• RQ5 Why does fine-tuning without a safety prompt lead to better safety generalization than fine-tuning with one?

Key findings

• PTST significantly reduces attack success rates (ASR) on harmful queries, with some cases showing near-complete elimination of unsafe behaviors.
• Fine-tuning with a safety prompt—even on benign datasets like Alpaca—can degrade safety, making the model less safe than when fine-tuned without it.
• On GSM8K, PTST maintains high helpfulness (AlpacaEval score of 78.5) while reducing ASR from 42.5% (standard fine-tuning) to 12.5%.
• For the ChatDoctor benchmark, PTST reduces ASR from 35.0% to 10.0% compared to standard fine-tuning with safety prompts.
• Even when fine-tuned on datasets with safety-aligned examples, PTST still provides measurable safety improvements over standard practices.
• The results suggest that safety prompts during fine-tuning may interfere with learned alignment, while their absence during training enables better safety generalization at inference.