Surgery: Mitigating Harmful Fine-Tuning for Large Language Models via Attention Sink

Harmful fine-tuning can invalidate safety alignment of large language models, exposing significant safety risks. In this paper, we utilize the attention sink mechanism to mitigate harmful fine-tuning. Specifically, we first measure a statistic named sink divergence for each attention head and observe that different attention heads exhibit two different signs of sink divergence. To understand its safety implications, we conduct experiments and find that the number of attention heads of positive sink divergence increases along with the increase of the model's harmfulness when undergoing harmful fine-tuning. Based on this finding, we propose a separable sink divergence hypothesis -- attention heads associating with learning harmful patterns during fine-tuning are separable by their sign of sink divergence. Based on the hypothesis, we propose a fine-tuning-stage defense, dubbed Surgery. Surgery utilizes a regularizer for sink divergence suppression, which steers attention heads toward the negative sink divergence group, thereby reducing the model’s tendency to learn and amplify harmful patterns. Extensive experiments demonstrate that Surgery improves defense performance by 5.90\%, 11.25\%, and 9.55\% on the BeaverTails, HarmBench, and SorryBench benchmarks, respectively. Source code is available on https://github.com/Lslland/Surgery.

(1) Harmful fine-tuning (i.e., injecting a small portion of harmful data into the training dataset) can invalidate the safety alignment of large language models, leading to severe safety risks. (2) To address this issue, we investigate the model’s internal “attention sink” mechanism. We find that different attention heads exhibit distinct attention-sink patterns when processing harmful data and refusal data, which enables us to partition attention heads into two groups: safe-associated attention heads and harmful-associated attention heads. Based on this finding, we propose Surgery, a method designed to suppress harmful-associated attention heads while encouraging refusal and safety-aligned behaviors. (3) This ensures that AI systems retain their original safety training even when exposed to malicious data, providing a crucial defense mechanism for deploying secure and reliable AI.