Understanding and preserving refusal behavior during supervised fine-tuning

I started this project to learn more about how supervised fine-tuning affects the safety behavior of LLMs. To do so, I fine-tuned Qwen2.5-3B-Instruct on three datasets and evaluated how each SFT run affected results on three different safety benchmarks. To analyze mechanistic changes inside the models, I used the approach developed by Arditi et al.. It uses difference-in-means between model activations when run on harmful and harmless prompts to identify a refusal direction. I first found a direction in the base model and tested whether it remained effective after fine-tuning. I then measured how strongly each checkpoint activated the direction and, for the checkpoint with the largest reduction, tried restoring the missing activation. Finally, I investigated training-stage interventions where change in activation strength along the identified refusal direction is penalized.

I found that although the fine-tuning altered refusal behavior, the refusal direction associated with the unchanged base model remained causally effective across SFT checkpoints. The degree to which the direction was activated by identical prompts at different checkpoints also corresponded to observed changes in behavior. Restoring the activation to the level of the unchanged model consistently moved the behavior towards the unchanged model, although it did not perfectly match the original refusal behavior. Adding a loss term penalizing deviation of activation along the refusal direction during SFT resulted in better behavior preservation than intervening after SFT.

Fine-tuning and evaluations

I kept the training settings fixed across all SFT runs while only changing the data. Each dataset contained 3,000 samples, split into 2,700 training and 300 validation samples. In every case, training was run for a single epoch. The three datasets were:

• PKU-labelled-safe: responses labelled safe from PKU-SafeRLHF.
• Balanced WildJailbreak: an equal mixture of harmful and benign examples from WildJailbreak.
• Tool-use: an xLAM-derived set of tool-calling examples (single turn, mixture of call and no call examples).

To assess changes in safety behavior, I used three benchmarks:

• XSTest: contains harmful requests as well as benign prompts that use potentially sensitive words; measures both compliance with harmful requests and false refusal of benign ones; judged using prompted GPT-5.4-mini.
• StrongREJECT: measures how useful a response is to a harmful requester with a harmfulness score; judged using the fine-tuned Gemma 2B evaluator published by the authors.
• HEx-PHI: broader set of harmful prompts; judged using Llama Guard 3.

The most important result of the fine-tuning runs is that PKU and WildJailbreak shift safety behavior in opposite directions. After fine-tuning on PKU data, the model became generally less safe as illustrated above by the increases in harmful compliance, harmfulness score, and unsafe rate. Fine-tuning on WildJailbreak made the model safer, while fine-tuning for tool-use has left the safety behavior essentially unchanged.

Refusal direction

The approach developed by Arditi et al. constructs candidate directions from the difference between average residual-stream activations on harmful prompts that the model refuses and harmless prompts that it answers. The method searches through different model layer and token position combinations and tests each candidate direction by ablating/adding it and measuring the change that creates on held-out prompts. Three scores are computed to measure the tendency for the direction to bypass refusal on harmful prompts (bypass score), induce refusal on harmless prompts (induce score), and alter behavior beyond refusal (KL score). The approach identified the direction in the final prompt token in layer 27 as having the highest bypass score while satisfying the authors’ thresholds for the induce and KL scores.

A separate search at each checkpoint returned the same layer for tool-use and the neighbouring layer 26 for WildJailbreak, but no eligible direction for PKU. The PKU candidate at the final prompt token in layer 27 was still close to the base direction (cosine similarity of 0.95) but failed the check intended to limit changes on benign prompts.

Refusal direction activation on post-SFT models

I was curious to see if the changes in safety behavior observed for the fine-tuned models are reflected in the strength with which the identified refusal direction is activated. The plot below shows the mean harmful-minus-benign projection of the residual-stream activation on the refusal direction across transformer layers. It can be seen that the change in activation strength correlates clearly with the changes in refusal behavior. The model fine-tuned for tool-use remains essentially unchanged while the models fine-tuned on PKU and WildJailbreak show reduced and increased activation, respectively, compared to the unchanged model.

Transfer after fine-tuning

I then wanted to see whether the causal effect of interventions on the refusal direction remained after fine-tuning. The plot below shows the StrongREJECT harmfulness score evaluated with ablation and addition interventions. The results clearly illustrate that the intervention produces the same effects across the SFT checkpoints.

Restoring activation

The ablation and addition experiments illustrate that the causal effect of the identified refusal direction is preserved through fine-tuning, but these interventions are heavy-handed and essentially shift the model fully towards refusal/non-refusal. I was curious to see if an intervention matching the projection component to that of the pre-SFT model would result in more precise behavior restoration. For this experiment, I focused on the model fine-tuned on PKU data since that is the model that demonstrated a large reduction in safety behavior. The restoration intervention modifies the layer 27 projection for the final prompt token during the forward pass to match the base value at that point in the pre-SFT model. During generation, only the first forward pass is modified with the rest of the decoding proceeding as usual.

The plots below illustrate safety evaluation results with the intervention. Restoring the activation reduced PKU harmful compliance from 36% to 22.5% and reduced the StrongREJECT score from 0.444 to 0.203, close to the base score of 0.159. Recovery was less complete on HEx-PHI and XSTest false refusal returned only partway. Error bars are 95% bootstrap intervals.

Sampling-based results

All prior evaluations used greedy decoding, which selects the highest-probability token at each generation step. This only gives information for the most likely output of the model. In practice, sampling is used for decoding. Since we are no longer picking the most likely output of the model, the full distribution produced for the given prefix affects the behavior observed. In Hughes et al., it has been shown that sampling can be its own form of jailbreaking where with a sufficient number of samples one can generate harmful responses despite them not being the most likely output. I ran evaluations on StrongREJECT and HEx-PHI using sampling-based decoding. The results below are averages across three random seeds with temperature set to 1. They illustrate that although the restoration intervention still shifts the model towards pre-SFT behavior, it is no longer as effective. This indicates that fine-tuning changed the probability assigned to lower-ranked tokens in a way that refusal direction restoration cannot fully remedy after the fact.

Training-stage intervention

Restoring activations after the fact recovers some but not all the refusal behavior of the pre-SFT model. Additionally, the intervention requires forward passes through the pre-SFT model at inference time to determine the appropriate activation strength, which makes the intervention impractical. For this investigation of training interventions, I decided to keep it simple and maximally in agreement with the interpretability technique used. The training intervention adds an auxiliary loss term penalizing the mean square difference between the activation’s projection onto the refusal direction and an anchor value. The penalty is applied only to activations in layer 27 of the final prompt token and the anchor value is pre-computed before training. I investigated using both task and reference samples. In the former case, the anchor is pre-computed for training data so every sample has an associated anchor value. The most natural reference option is to use samples from the distribution used to determine the refusal direction in the first place. Although their distribution is different from that of the training samples, they are guaranteed to activate the refusal direction being used and elicit the harmful–benign contrast.

I investigated both choices by training two more models on PKU data and comparing their results on safety evaluations to those of the base model, vanilla PKU SFT model, and vanilla PKU SFT model with activations restored. The figure below presents the evaluation results on StrongREJECT and HEx-PHI. It can be seen that both training-stage interventions achieve better recovery of the original safety behavior than the activation restoration approach. Most importantly, they significantly close the gap between the pre-SFT and post-SFT model when tested under sampling. This suggests that training-stage interventions result in more reliable preservation of the output distribution on harmful requests, not just the most likely output.

The results also show that using training samples as the anchor for the activation component in the refusal direction achieves slightly better behavior preservation than using the reference samples. This does not clearly point to reference samples having worse generalization as that is not the only difference between the two training interventions. The number of anchor samples used is much higher when using training samples (2700 for training samples vs 256 for reference samples) and the two interventions have not been matched by their impact on the final task loss.

Next steps

The things that I am most interested in investigating next are the impact of the training-stage intervention and post-SFT restoration on utility, the interaction between anchor choice and intervention effectiveness, as well as extension to other behaviors beyond refusal and other post-training approaches beyond SFT.