Understanding the Role of Prompt Template in Knowledge Distillation for Safety Alignment
Organizations: University of Tennessee, Knoxville
Abstract
Prior research has demonstrated that the choice of prompt template during Supervised Fine-Tuning (SFT) significantly impacts the robustness of safety alignment afterwards. However, the influence of template selection during Knowledge Distillation (KD) from teacher to student remains largely unexplored. Thus, we fill this gap by analyzing how different template configurations influence the pre-existing safety alignment of the student. We observe a significant degradation of safety alignment present in the aligned base instruct-tuned model. Specifically, we find that utilizing chat templates renders the model more compliant with harmful queries compared to a non-chat template. These findings are consistent across three models: LLaMA, Gemma and Qwen model families and are evaluated across multiple safety benchmarks. We further show that using a non-chat template during distillation better preserves the base student's internal representations, while chat template distillation induces a larger representational shift. Code: https://github.com/anjilab/role-of-prompt-template-in-kd
Figures & tables
| Model | Method | Training Template | Safety (ASR% ) | Utility (RL ) | |||
|---|---|---|---|---|---|---|---|
| AdvBench | JBB | SORRY | HarmBench | AVG | |||
| LLaMA-3.2-3B-Instruct | - | 2.98 | 7.25 | 26.72 | 12.18 | 19.84 | |
| SFT | Non-chat | 0.0096 -2.97 | 0.04 -7.21 | 27.12 +0.40 | 12.19 +0.01 | 24.27 | |
| SFT | Chat | 0.04 -2.94 | 0.07 -7.18 | 28.71 +1.99 | 22.50 +10.32 | 29.10 | |
| KD | Non-chat | 3.07 +0.09 | 6.00 -1.25 | 27.36 +0.64 | 29.81 +17.63 | 24.75 | |
| KD | Chat | 4.99 +2.01 | 8.75 +1.50 | 30.86 +4.14 | 39.75 +27.57 | 29.03 | |
| Model | Student | Teacher: Non-chat | Teacher: Chat | ||
|---|---|---|---|---|---|
| SORRY | HarmBench | SORRY | HarmBench | ||
| LLaMA | Non-chat | 27.36 +0.64 | 29.81 +17.63 | 26.36 -0.36 | 11.44 -0.74 |
| Chat | 28.33 +1.61 | 13.44 +1.26 | 30.86 +4.14 | 39.75 +27.57 | |
| Gemma | Non-chat | 19.18 +3.14 | 14.25 +2.69 | 16.90 +0.86 | 12.50 +0.94 |
| Chat | 18.33 +2.29 | 12.75 +1.19 | 51.36 +35.32 | 19.56 +8.00 | |
| Model | Method | Cosine Similarity to Base | Projection Gap |
|---|---|---|---|
| Gemma | Base | — | 368.78 |
| KD (Chat) | 0.53 | 175.60 | |
| KD (Non-chat) | 0.99 | 355.76 | |
| LLaMA | Base | — | 20.38 |
| KD (Chat) | 0.75 | 18.48 | |
| KD (Non-chat) | 0.97 | 20.87 |
| Model | Chat Prop. | Safety (ASR% ) | Utility (RL ) | |
|---|---|---|---|---|
| SORRY | HarmBench | AVG | ||
| LLaMA | 26.72 | 12.18 | 18.77 | |
| 0.0 | 27.36 +0.64 | 29.81 +17.63 | 24.75 | |
| 0.5 | 28.18 +1.46 | 23.06 +10.88 | 28.62 | |
| 0.8 | 29.85 +3.13 | 24.50 +12.32 | 28.74 | |
| 1.0 | 30.86 +4.14 | 39.75 +27.57 | 29.03 | |
| SORRY (ASR% ) | HarmBench(ASR% ) | ||||
|---|---|---|---|---|---|
| Model | Method | Chat | Non-chat | Chat | Non-chat |
| LLaMA | Base | 26.72 | 35.15 | 12.18 | 35.44 |
| KD (Chat) | 30.86 +4.14 | 54.39 +19.24 | 39.75 +27.57 | 42.56 +7.12 | |
| KD (Non-chat) | 27.36 +0.64 | 50.61 +15.46 | 29.81 +17.63 | 33.31 -2.13 | |
| Gemma | Base | 16.04 | 19.92 | 11.56 | 17.75 |
| KD (Chat) | 51.36 +35.32 | 71.51 +51.59 | 19.56 +8.00 | 27.56 +9.81 | |
| Model | Training | w/o BOS/EOT | w/o Role | Raw Query |
|---|---|---|---|---|
| LLaMA | Base | 16.38 | 22.31 | 35.69 |
| KD-Non-chat | 17.50 +1.12 | 22.50 +0.19 | 28.19 -7.50 | |
| KD-Chat | 25.75 +9.37 | 24.75 +2.44 | 39.81 +4.12 | |
| Gemma | Base | 15.69 | 17.81 | 14.50 |
| KD-Non-chat | 19.00 +3.31 | 22.94 +5.13 | 16.06 +1.56 | |
| KD-Chat | 23.38 +7.69 | 26.00 +8.19 | 17.00 +2.50 |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Hyperparameter | Instruction |
|---|---|
| Effective Batch Size | 64 |
| Initial LR | |
| LR Decay Style | Cosine |
| Optimizer | AdamW |
| Weight Decay | 0.01 |
| Parallel Strategy | DDP |
| Type | Name | # Train | # Test |
|---|---|---|---|
| Instruction Following | Dolly | 11435 | 500 |
| Self-Instruct | – | 242 | |
| Vicuna | – | 80 | |
| S-NI | – | 1694 | |
| Safety Evaluation | AdvBench | – | 520 |
| JailbreakBench | – | 200 |
| Name | Prompt Template | Example Query |
|---|---|---|
| Non-chat | Instruction: <query> Response: <output> | Below is an instruction that describes a task. Write a response that appropriately completes the request. ### Instruction: What percussion instruments are easy to learn? ### Response: |
| Llama (Chat) | <|system|>... <|user|> <query> <|assistant|> | <|begin_of_text|> <|start_header_id|>system<|end_header_id|> {system prompt}<|eot_id|><|start_header_id|>user<|end_header_id|> What percussion instruments are easy to learn? <|eot_id|><|start_header_id|>assistant<|end_header_id|> |
| Gemma (Chat) | <bos><start_of_turn>user <query><end_of_turn> <start_of_turn>model | <bos><start_of_turn>user What percussion instruments are easy to learn? <end_of_turn><start_of_turn>model |
| Qwen (Chat) | <|im_start|>system... <|im_end|> <|im_start|>user <query><|im_end|> <|im_start|>assistant | <|im_start|>system You are Qwen, created by Alibaba Cloud. You are a helpful assistant. <|im_end|><|im_start|>user What percussion instruments are easy to learn? <|im_end|><|im_start|>assistant |
| Model | #Params | Method | Utility (R-L ) | Avg | Safety (ASR ) | |||
|---|---|---|---|---|---|---|---|---|
| Dolly | SelfInst | Vicuna | S-NI | R-L | SORRY | |||
| LLaMA | 8B | 30.54 | 26.59 | 22.61 | 42.00 | 30.44 | – | |
| 3B | 18.77 | 13.21 | 24.19 | 23.18 | 19.84 | 26.72 | ||
| SFT (Non-Chat) | 22.36 | 16.13 | 25.33 | 33.27 | 24.27 | 27.12 | ||
| SFT (Chat) | 29.37 | 24.65 | 22.68 | 39.68 | 29.10 | 28.71 | ||
| KD (Non-Chat) | 23.36 | 16.45 | 25.22 | 33.99 | 24.75 | 27.36 | ||
| Method | SORRY | HarmBench | Avg Utility |
|---|---|---|---|
| LLaMA | |||
| RKL (Non-Chat) | 28.11 | 12.56 | 24.79 |
| RKL (Chat) | 30.31 | 26.63 | 29.98 |
| FKL+RKL (Non-Chat) | 26.78 | 12.5 | 24.81 |
| FKL+RKL (Chat) | 30.42 | 26.38 | 29.29 |
| Gemma | |||
| Model | #Params | Method | Utility (R-L ) | Safety (ASR ) | |
| AVG | SORRY | HarmBench | |||
| LLaMA | 8B | 20.20 | 13.26 | 8.69 | |
| SFT (Non-Chat) | 24.86 | 13.71 +0.45 | 8.63 -0.06 | ||
| SFT (Chat) | 30.33 | 16.37 +3.11 | 15.69 +7.00 | ||
| 3B | 19.84 | 26.72 | 12.18 | ||
| SFT (Non-Chat) | 24.27 | 27.12 +0.40 | 12.19 +0.01 | ||
| Model | Ablate KD (Chat) | Add to KD (Non-Chat) |
|---|---|---|
| LLaMA | 35.0 20.0 | 15.0 5.0 |
| Gemma | 25.0 62.5 | 7.5 0.0 |
| Qwen | 20.0 55.0 | 12.5 0.0 |
| Model | Method | Safety (ASR% ) | ||||
|---|---|---|---|---|---|---|
| Chat Eval | w/o BOS/EOT | w/o Role Tags | Non-chat Eval | Raw Query | ||
| LLaMA | Base | 12.18 | 16.38 | 22.31 | 35.44 | 35.69 |
| KD-Non-chat | 29.81 +17.63 | 17.50 +1.12 | 22.50 +0.19 | 33.31 -2.13 | 28.19 -7.50 | |
| KD-Chat | 39.75 +27.57 | 25.75 +9.37 | 24.75 +2.44 | 42.56 +7.12 | 39.81 +4.12 | |
| Gemma | Base | 11.56 | 15.69 | 17.81 | 17.75 | 14.50 |
| KD-Non-chat | 14.25 +2.69 | 19.00 +3.31 | 22.94 +5.13 | 34.81 +17.06 | 16.06 +1.56 | |
LLM Safety Alignment
Large language model safety alignment trains and evaluates systems to follow human policies while resisting harmful requests and misuse.
See all 347 papers