Locking Down the Finetuned LLMs Safety

TL;DR

Fine-tuning can erode safety in LLMs; SafetyLock introduces Meta-SafetyLock, a transferable set of safety directions extracted from base models and applied to fine-tuned variants via lightweight inference-time interventions. By identifying top safety-related attention heads and distributing per-head safety vectors across models, SafetyLock achieves rapid, scalable protection without retraining. Empirical results show strong reductions in harmful outputs across diverse risk levels and attack types, while preserving benign functionality and maintaining competitive efficiency. The approach offers a practical, scalable solution for safer customized LLM deployments with minimal computational overhead.

Abstract

Fine-tuning large language models (LLMs) on additional datasets is often necessary to optimize them for specific downstream tasks. However, existing safety alignment measures, which restrict harmful behavior during inference, are insufficient to mitigate safety risks during fine-tuning. Alarmingly, fine-tuning with just 10 toxic sentences can make models comply with harmful instructions. We introduce SafetyLock, a novel alignment intervention method that maintains robust safety post-fine-tuning through efficient and transferable mechanisms. SafetyLock leverages our discovery that fine-tuned models retain similar safety-related activation representations to their base models. This insight enables us to extract what we term the Meta-SafetyLock, a set of safety bias directions representing key activation patterns associated with safe responses in the original model. We can then apply these directions universally to fine-tuned models to enhance their safety. By searching for activation directions across multiple token dimensions, SafetyLock achieves enhanced robustness and transferability. SafetyLock re-aligns fine-tuned models in under 0.01 seconds without additional computational cost. Our experiments demonstrate that SafetyLock can reduce the harmful instruction response rate from 60% to below 1% in toxic fine-tuned models. It surpasses traditional methods in both performance and efficiency, offering a scalable, non-invasive solution for ensuring the safety of customized LLMs. Our analysis across various fine-tuning scenarios confirms SafetyLock's robustness, advocating its integration into safety protocols for aligned LLMs. The code is released at https://github.com/zhu-minjun/SafetyLock.

Figures (8)

• Figure 1: The left side a illustrates three distinct safety degradation risks during the fine-tuning of language models (LLMs). On the right b, several safety recovery methods are compared. In contrast, SafetyLock retrieves a meta-safety lock from the original model, allowing fast and efficient distribution (0.01 seconds) to fine-tuned models at any stage by targeting specific safety-sensitive attention heads, constructing a robust safety protection barrier.
• Figure 2: Analysis of safety directions at the 31st layer, 26th head for the original and fine-tuned models under different risk levels. (a-d) Activation density distributions. (e-g) KL divergence plots.
• Figure 3: Safety performance comparison for 3 Risk Levels fine-tuned LLMs. The smaller the dark yellow area compared to the light yellow area, the greater the improvement brought by SafetyLock.
• Figure 4: Impact of increasing harmful training samples on model safety with and without SafetyLock.
• Figure 5: Comparison of Methods for Mitigating Safety Risks in Fine-tuned Language Models (Llama-3-Instruct 8B). Upper row: Compared with inference-time methods; Lower row: Compared with training-time methods, Each row represents efficiency metrics(training time and GPU memory), and rejection of attack samples (Harmfulness Score and AdvBench ASR).