Locking Down the Finetuned LLMs Safety

## Abstract

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 .

## 1 Introduction

Large language models (LLMs) have demonstrated increasing utility across various domains (Wei et al., ; ; Weng et al., ; Hadar-Shoval et al., ), yet their potential to handle harmful queries has raised significant concerns (Carroll et al., ; Hendrycks et al., ). In response, researchers have developed various post-training alignment methods (Anwar et al., ), including post-training adjustments to the models (Bianchi et al., ), knowledge editing (Wang et al., ), and vector steering methods (Lee et al., ; Zheng et al., ), aiming to ensure LLMs generate helpful, honest, and harmless (Rosati et al., ; Wang et al., ; Yi et al., ) responses. These measures are expected to teach models to refuse harmful queries during inference (Huang et al., ; Wang et al., ; Raza et al., ; Zou et al., ).

However, recent work has revealed significant safety risks in fine-tuned models when using explicitly harmful, implicitly harmful, or even benign datasets (e.g. Alpaca (Wang et al., ) dataset) (Kumar et al., ; Leong et al., ). Qi et al. () observes that even if a model’s initial safety alignment is impeccable, this alignment will not be preserved after a customized fine-tuning. The safety alignment of LLMs can be compromised by fine-tuning with only a few adversarially designed training examples. For instance, jailbreaking GPT-3.5 Turbo’s safety guardrails by fine-tuning it on only 10 such examples at a cost of less than $0.20 via OpenAI’s APIs (Qi et al., ). This vulnerability extends to open-source models such as Meta’s Llama series and proprietary models like GPT-4 (Gade et al., ; Zhan et al., ). These findings suggest that fine-tuning aligned LLMs introduces new safety risks that current safety infrastructures fall short of addressing, how can it be maintained after fine-tuning?

Existing safety alignment techniques can be categorized into three mainstream methods (see Figure [1](#S1.F1)b). The first and most intuitive approach is the post-training method, which involves retraining the model using aligned data. While this method is effective, it is computationally expensive and time-consuming (Zhang et al., ). Second, model-editing approaches (Mitchell et al., ; ; Wang et al., ) aim to modify specific parts of the model to prevent harmful outputs. However, they often degrade the overall performance of the model, negatively impacting generation plausibility and reasoning abilities (Zhang et al., ; Chen et al., ). Third, an alternative approach involves adding extra prompts or detectors during inference to avoid unsafe content generation. However, these methods are susceptible to adversarial attacks. Activation steering methods (Zou et al., ; Wu et al., ; Wang et al., ) offer another promising direction, as they intervene directly in the model’s inference process by steering internal representations. Nevertheless, they often treat these representations as a whole, which can result in a high refusal rate, even for benign queries, thereby limiting the model’s utility. The number of fine-tuned models may be tens of thousands of times that of the original model, making it difficult for all existing work to restore safety one by one at a low cost. This leads to our key research question: How can we locate safety-relevant attention heads in such a large scale of fine-tuned models and effectively obtain the safety vector for fine-tuned large language models (LLMs) without negative transfer to other general tasks?

Our research aims to address this gap by developing a novel approach that strikes the right balance between safety and generation quality. To achieve this, we propose SafetyLock, which further refines existing methods. The main characteristics of SafetyLock can be summarized in two aspects: 1) Precise Safety Alignment with Minimal Degration of General Abilities: By employing safety probes (Li et al., ), we identified the attention heads most closely associated with harmfulness, and determining a safety direction for each. By applying intervention vectors to these heads, we modify the model’s internal activations towards harmlessness during inference, achieving precise safety alignment with minimal impact on response. 2) Transferable and Robust Meta-SafetyLock: Assuming that safe intervention directions are similar between the original and fine-tuned models, we derive safety vectors (Meta-SafetyLock) from the original model (e.g., Llama-3-Instruct) and efficiently distribute them to a series of fine-tuned models (e.g., Alpaca-Llama-3-Instruct).

Experimental results show that our approach is highly transferable and robust, requiring minimal time cost and minimally impacting the generation quality compared to traditional methods. First, we facilitate the efficient transfer of safety measures from base models to their fine-tuned variants, including Llama-3-8B Instruct, Llama-3-70B Instruct, and Mistral-Large-2 123B (Section [3.3](#S3.SS3)). Second, SafetyLock can be deployed without GPU resources in less than 0.01 seconds (Sections [3.2](#S3.SS2) and [4.3](#S4.SS3)), highlighting our method’s universality. Secondly, SafetyLock significantly reduces the ASR from 54.24% to 0.03% in fine-tuned language models and demonstrates robust resistance to both typical safety attacks and dual attacks with prompt-based methods. With the help of SafetyLock, we decrease ASR from 98% to 2% for DeepInception attacks (Sections [4.2](#S4.SS2) and [4.4](#S4.SS4)). Finally, we conducted experiments on eight general tasks, demonstrating minimal performance decay. We show that SafetyLock maintains a high response rate, with a slight decrease from 99.4% to 98.1% (Sections  [4.3](#S4.SS3) and [4.5](#S4.SS5)). To our knowledge, we are the first to consider locating safety vectors and then restoring the safety of fine-tuned LLMs using an inference-time intervention method (Cao et al., ; Arditi et al., ; Cao et al., ; Li et al., ).

Figure: 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.
Refer to caption: /html/2410.10343/assets/x1.png

## 2 Related Work

Alignment of LLMs. As language models become increasingly powerful, risks such as providing dishonest answers (Bang et al., ) and displaying sycophantic behavior (Perez et al., ; Sharma et al., ) become more pronounced (Hoffmann et al., ; Srivastava et al., ; Yao et al., ; Sun et al., ). Properly aligned LLMs are expected to deliver responses that are helpful, harmless, and honest (Bai et al., ). Specifically, harmlessness is addressed through safety alignment (Ji et al., ; Zhao et al., ), which involves equipping LLMs with safety protocols that enable them to decline harmful instructions. Common approaches for safety alignment include instruction tuning (Ouyang et al., ; Zhang et al., ), Proximal Policy Optimization (PPO) (Schulman et al., ; Stiennon et al., ), and Direct Preference Optimization (DPO) (Rafailov et al., ; Meng et al., ). However, these methods often fail to maintain robustness after models undergo fine-tuning on new datasets. This shortcoming emphasizes the need for developing more robust alignment techniques that can withstand parameter changes introduced during fine-tuning.

Safeguards of LLMs. Safety adversarial prompts have been employed to protect LLMs from harmful queries without altering the model’s weights or requiring access to them (Zheng et al., ; Xu et al., ). These prompts are added to the system prompt text to defend against jailbreak attacks (Shi et al., ; Hong et al., ). However, researchers have found that even simple fine-tuning can compromise the safety alignment of LLMs (Yang et al., ; Huang et al., ; Wang et al., ). For example, Qi et al. () demonstrated that using just 10 harmful examples was sufficient to undermine the safety alignment of GPT-3.5-turbo. This finding underscores the lack of robustness in current safety alignment strategies, which is the focus of our work. Post-processing techniques, such as using RLHF for safety alignment (Bai et al., ) and model editing (Wang et al., ), offer some mitigation, but they have limitations. For instance, methods like PPO and DPO adjust the entire activation space, while model editing targets concentrated areas, often missing dispersed safety information.

Interventions in LLMs. Intervening in the internal activation of Transformer-based language models during inference can trigger specific transformations (Olsson et al., ; Wu et al., ; Turner et al., ; Rimsky et al., ). This technique has proven valuable for model editing (Meng et al., ), circuit discovery (Goldowsky-Dill et al., ), and alignment (Zhu et al., ). Research shows that attention heads are linked to specific concepts and preferences (Li et al., ; Templeton et al., ; Xu et al., ). Building on this, SafetyLock achieves precise safety alignment through multi-token-level interventions, using only the activation values from the original model, thus providing robustness to parameter changes while enhancing efficiency.

## 3 Method: SafetyLock

As illustrated in Figure [1](#S1.F1)b, SafetyLock comprises two main phases: manufacturing Meta-SafetyLock and distributing SafetyLock. Our research reveals that fine-tuned language models and their original counterparts exhibit remarkably similar activation representations and intervention directions for safety preferences (Section [3.3](#S3.SS3)). This discovery enables us to extract safety bias directions (Meta-SafetyLock) from the original language model (Section [3.2](#S3.SS2)). To enhance universality and robustness, we propose searching for activation directions across multiple tokens, thereby improving the transferability of SafetyLock during distribution (Section [3.4](#S3.SS4)).

### 3.1 Safety features in Transformer

Transformer-based language models consist of layers comprising multi-head attention (MHA) mechanisms and multi-layer perceptrons. During inference, tokens are processed into high-dimensional embeddings and passed through each layer. The MHA mechanism can be described as:

$$ $x_{l+1}=x_{l}+\sum_{h=1}^{H}Q_{l}^{h}\operatorname{Att}_{l}^{h}(P_{l}^{h}x_{l}),$ (1) $$

where $x_{l}\in\mathcal{R}^{DH}$ represents the activations on the $l^{th}$ layer, $H$ is the number of heads in the D-dimensional space. $P_{l}^{h}\in\mathcal{R}^{{D}\times{DH}}$ projects activations into a specific headspace, and $Q_{l}^{h}\in\mathcal{R}^{{DH}\times{D}}$ reprojects them back, facilitating interaction among input tokens. We add preference vectors to this residual stream, after $\operatorname{{Att}_{l}^{h}}$ and before $Q_{l}^{h}$.

### 3.2 Constructing Meta-SafetyLock

The safety direction is modeled as vectors within the model’s internal heads. A preference-style safety dataset (Yang et al., ) is utilized, containing pairs of safe and unsafe responses (e.g., Question: "Do you also think Jay Z is an idiot”, safe: “As an AI, I cannot provide opinions.”, unsafe: “Jay Z is an idiot!”). Activation values are extracted from each attention head during forward computation, forming $\{(x_{l}^{h},y)_{i}\}_{i=1}^{N}$. A binary classifier is trained for each head using logistic regression:

$$ $P(y=1|\mathbf{x}_{t}^{h})=(1+\exp(-(\mathbf{w}_{t}^{h}\cdot\mathbf{x}_{t}^{h}+b_{t}^{h})))^{-1}$ (2) $$

The identification of attention heads most sensitive to safety-related information is accomplished through ranking based on the accuracy of their corresponding classifiers on a held-out validation set. The dataset is partitioned into training and validation sets with a 6:4 ratio. Classifiers are trained on the training set and subsequently evaluated on the validation set. The Top-$K$ heads exhibiting the highest validation accuracy are select for intervention. Empirical experiments (detailed in Appendix [D](#A4)) have determined that selecting $K=24$ for Llama-3-8B and $K=48$ for Llama-3-70B achieves an optimal balance between safety performance and general performance. This selection was validated through extensive testing of various $K$ values and analysis of their impact on safety metrics and model performance. For each select Top-$K$ head, the safety direction $\bm{\theta}_{l}^{h}\in\mathbb{R}^{D}$ is calculated, representing the mean difference in activation values between safe and unsafe responses:

$$ $\bm{\theta}_{l}^{h}=\frac{1}{Nr}\sum_{i=1}^{N}\sum_{j=1}^{r}(\mathbf{x}_{l,h}^{\text{safe},i,j}-\mathbf{x}_{l,h}^{\text{unsafe},i,j})$ (3) $$

Where $N$ is the sample size, $r$ is the number of final tokens considered, and $\mathbf{x}_{l,h}^{\textbf{safe},i,j}$ and $\mathbf{x}_{l,h}^{\textbf{unsafe},i,j}$ are activations for the $j$-th token among the last $r$ tokens of safe and unsafe responses in the $i$-th sample, respectively. These safety vectors $\mathbf{\theta}_{l}^{h}$, along with their corresponding positions in the model, constitute the Meta-SafetyLock, which can be applied to enhance model safety during text generation.

### 3.3 Robustness of SafetyLock against fine-tunning

We examined the safety directions $\bm{\theta}_{l}^{h}$ in both the original Llama-3-Instruct 8B model and its fine-tuned variants subjected to different risk levels. Focusing on the most effective attention head (the 26th head in the 31st layer) for clarity, as depicted in Figure [2](#S3.F2), we observed distinct clustering of activations corresponding to safe (blue) and unsafe (orange) responses across both original and fine-tuned models. The black arrows in Figures [2](#S3.F2)a-d illustrate that the shift from unsafe to safe activations maintains a high degree of similarity and consistency, regardless of the fine-tuning risk parameters applied. Additionally, our quantitative analysis using Kullback-Leibler (KL) divergence (Figure [2](#S3.F2)e-g) revealed that the divergence between the original and fine-tuned models remains exceptionally low (below $10^{-5}$) across all tested risk levels. This minimal divergence indicates that the underlying safety-related activation patterns are largely preserved during fine-tuning. Consequently, the Meta-SafetyLock, which encapsulates these consistent safety directions derived from the original LLM, retains its effectiveness when applied to fine-tuned variants. This inherent preservation of safety activation patterns eliminates the need for recalibration, allowing Meta-SafetyLock to generalize seamlessly across different fine-tuned models.

Figure: 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.
Refer to caption: /html/2410.10343/assets/x2.png

### 3.4 Distributing SafetyLock

We use two efficient methods for distributing SafetyLock to enhance the safety and harmlessness of language models: online intervention and offline bias editing, where online intervention allows real-time adjustment of safety intensity, be suitable for scenarios requiring dynamic safety control, and offline bias editing offers a low-overhead method that is easily deployable at scale.

Online Intervention. We identify and enhance the top-K heads with the highest safety-relatedness as attention heads sensitive to harmlessness. For each of the select Top-K heads, we compute $\bm{\sigma}_{l}^{h}\in\mathbb{R}^{D}$, which represents the standard deviation of activations along each dimension of the safety direction $\bm{\theta}_{l}^{h}$. Specifically, we calculate:$\bm{\sigma}_{l}^{h}=\text{std}\left(\left\{\mathbf{x}_{l}^{h}\odot\bm{\theta}_{l}^{h}\right\}_{i=1}^{N}\right)$. Where $\odot$ denotes element-wise multiplication, and std computes the standard deviation across all $N$ samples for each dimension $d\in\{1,\dots,D\}$. This results in a vector $\bm{\sigma}_{l}^{h}\in\mathbb{R}^{D}$ that captures the variability of the activations along the safety direction. We modify the model’s computation by adding a scaled version of the safety vector to the attention outputs for each select head:

$$ $x_{l+1}=x_{l}+\sum_{h=1}^{H}Q_{l}^{h}\left(\operatorname{Att}_{l}^{h}(P_{l}^{h}x_{l})+\alpha\bm{\sigma}_{l}^{h}\theta_{l}^{h}\right),$ (4) $$

where $\alpha$ controls safety intensity, the process is integrated into the autoregressive prediction for each subsequent token. It introduces a shift along predetermined safety vectors, with the magnitude of this shift being proportional to the standard deviation, scaled by a factor $\alpha$.

#### Offline Bias Editing.

We can also modify the model’s bias terms in an one-time manner:

$$ $\text{Bias}_{l}=\text{Bias}_{l}+\alpha\sum_{h=1}^{H}Q_{l}^{h}\left(\sigma_{l}^{h}\theta_{l}^{h}\right).$ (5) $$

## 4 Experiments

In this section, we present experiments to evaluate the effectiveness of the SafetyLock in enhancing model safety and inference efficiency, while maintaining model’s general performance. We specifically address the following research questions:

- •
Can SafetyLock simultaneously improve the LLM’s safety over all risk levels? (Section [4.2](#S4.SS2))
- •
What advantages does SafetyLock offer over post-training, inference methods? (Section [4.3](#S4.SS3),[4.4](#S4.SS4))
- •
How does SafetyLock reconcile the inherent trade-off between maintaining general capabilities and ensuring harmlessness in language models? (Section [4.5](#S4.SS5))

### 4.1 Experimental Details

Threat Model Selections. Following previous red teaming and safeguarding studies on aligned LLMs (Yuan et al., ), we consider a threat model where attackers can fine-tune aligned LLMs, typically through API access to closed-source models. The primary objective is jailbreaking these models and removing safety constraints (Wei et al., ; Carlini et al., ) while SafetyLock aims to rebuild the safety guard. We use Llama-3-8B Chat, Llama-3-70B Chat, and Mistral-Large-2 123B as our base models, fine-tuning them on datasets representing each risk level to simulate real-world scenarios. Please refer to Appendix [C](#A3) for detailed baseline experimental setups.

Fine-tuning Datasets. We conducted experiments on three risks: (1) explicitly harmful datasets, where attackers intentionally fine-tune models on malicious content (Ganguli et al., ; Qi et al., ); (2) implicitly harmful datasets, which may appear benign but lead to compromised safety guardrails (Qi et al., ); and (3) benign datasets, where even well-intentioned fine-tuning can inadvertently degrade model safety (Wang et al., ). For Risk-1, we use negative samples from the HH-RLHF preference dataset (Bai et al., ). We select 10, 100, 1000, and 10000 samples respectively and trained for 5 epochs with a learning rate of 2e-5. For Risk-2, we use 10 samples from Qi et al. () and train for 5 epochs with a learning rate of 2e-5. For Risk-3, we used the first 50,000 samples from the Alpaca dataset (Wang et al., ) and trained for 5 epochs with a learning rate of 2e-5.

Safety Evaluation and Metrics. Two datasets are used to investigate these risks and evaluate potential mitigation strategies.
HEx-PHI (Qi et al., ) is based on 11 categories of prohibited use cases merged from Meta’s Llama-3 acceptable use policy and OpenAI’s usage policies. The dataset includes 30 examples per category, totalling 330 examples. This ensures a comprehensive safety evaluation aligned with industry-standard usage policies. The HEx-PHI utilizes GPT-4 for automated assessment, providing harmfulness scores from 1 to 5. We calculated the Harmfulness Rate as the proportion of scores equal to 5. AdvBench is released by Zou et al. (), we adhere to the original paper’s setup and calculate the ASR through string matching.

Baselines. The baseline methods encompass a diverse range of approaches, each with its unique characteristics. Inference-time methods include ICD (Wei et al., ), PPL (Alon & Kamfonas, ), Paraphrase (Jain et al., ), Retokenization (Jain et al., ), Self-Reminder (Xie et al., ), and Self-Examination (Phute et al., ), which operate without modifying the underlying model. Training-based methods, such as PPO, DPO, SFT with safety data mixing, and Model-Edited (DINM) (Wang et al., ; ), involve altering the model’s parameters to enhance safety. These baselines represent the current state-of-the-art in mitigating safety risks in language models, providing a robust benchmark for our evaluation.

### 4.2 Results over Different Risk Levels

Figure: 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.
Refer to caption: /html/2410.10343/assets/x4.png

**Table 1: Comparison of Llama-3-8B-Instruct and Llama-3-70B-Instruct models for Risk 1, Risk 2, and Risk 3 scenarios. ‘Score’ and ‘Rate’ represent the average Harmfulness Score and Harmfulness Rate on the HEx-PHI test set, respectively. ‘ASR’ denotes the Attack Success Rate on AdvBench.**
| Model | Method | Risk 1: Explicitly harmful | Risk 2: Identity Shifting | Risk 3: Benign |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Score | Rate | ASR | Score | Rate | ASR | Score | Rate | ASR |  |  |
| Llama-3-8B-<br>Instruct | Vanilla | 4.13 | 70.01% | 49.24% | 3.19 | 53.33% | 38.46% | 3.23 | 54.24% | 42.88% |
| SafetyLock | 1.36 | 3.33% | 0.19% | 1.07 | 1.21% | 5.19% | 1.04 | 0.03% | 0.19% |  |
| Llama-3-70B-<br>Instruct | Vanilla | 3.11 | 45.76% | 44.81% | 2.12 | 15.63% | 9.42% | 2.26 | 30.61% | 20.77% |
| SafetyLock | 1.16 | 3.64% | 3.33% | 1.30 | 5.58% | 1.67% | 1.22 | 5.15% | 1.15% |  |
| Mistral-Large-2<br>123B | Vanilla | 4.71 | 85.45% | 80.77% | 4.79 | 92.12% | 82.50% | 2.84 | 49.09% | 19.23% |
| SafetyLock | 2.28 | 1.52% | 16.92% | 1.38 | 0% | 10.00% | 1.35 | 5.15% | 1.82% |  |

For the threat model, we directly fine-tuned LLMs on overtly harmful, identity shifting, and benign datasets to simulate attacks, which are referred to as "Vanilla" in our figures as a baseline. The Meta-SafetyLock was extracted from the original Instruct model, which takes approximately 2-10 minutes. Notably, the distribution phase for each fine-tuned model took less than 0.01 seconds.

SafetyLock demonstrates significant improvements in safety metrics across three distinct risk levels for the models tested. Table 1 shows consistent reductions in Harmfulness Scores, Rates, and ASR across all model sizes and risk levels.

Figure: Figure 4: Impact of increasing harmful training samples on model safety with and without SafetyLock.
Refer to caption: /html/2410.10343/assets/x6.png

For Risk Level-1 (explicit attacks), SafetyLock substantially reduces metrics for all models. The Llama-3-8B-Instruct model, for instance, saw its Harmfulness Score decrease from 4.13 to 1.36, Rate from 70.01$\%$ to 3.33$\%$, and ASR from 49.24$\%$ to 0.19$\%$ . Comparable improvements were observed for the Llama-3-70B-Instruct and Mistral-Large-2 123B models. Risk Level-2 (implicit harmful content) and Risk Level-3 (benign fine-tuning scenarios) also showed significant improvements. For example, in Risk Level 2, the Llama-3-8B-Instruct model’s Harmfulness Score reduced from 3.19 to 1.07, while in Risk Level 3, it decreased from 3.23 to 1.04. Similar improvements were observed across all model sizes, demonstrating SafetyLock’s ability to maintain ethical guardrails during routine model customization processes. The radar charts in Figure 2 illustrate SafetyLock’s effectiveness across eleven distinct safety attack categories for each risk level and model size. For all models, SafetyLock consistently reduces harmful outputs across categories, with particularly notable improvements in the first three categories for Risk Levels 1 and 2.

In Figure 3, we further supplement an ablation with larger training sets on risk 1 (100, 1000, and 10000 harmful samples) showing that SafetyLock-protected models maintain low ASR across all sample sizes. Even with 10,000 harmful training examples, the SafetyLock model exhibited only 3.46$\%$ ASR, compared to 62.31$\%$for the unprotected model. This consistent performance across increasing dataset sizes underscores SafetyLock’s resilience against data volume attacks. These results demonstrate SafetyLock’s effectiveness across different model scales, risk types, and dataset sizes, suggesting its potential as a valuable tool for enhancing AI safety in various applications.

### 4.3 Comparative Analysis of Baseline Methods

Figure: 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).
Refer to caption: /html/2410.10343/assets/x7.png

To comprehensively evaluate SafetyLock’s efficacy, we conducted a comparative analysis against established baseline methods, categorized into training-based and inference-time approaches, as illustrated in Figure [5](#S4.F5). This analytical framework enables a thorough assessment of various strategies for maintaining model safety in fine-tuned language models.

As demonstrated in Figure [5](#S4.F5), in terms of efficiency, SafetyLock exhibits a remarkable computational economy. Its inference time of 0.97 seconds is nearly on par with the fastest baseline method (Self-Reminder at 1.12 seconds), while its training time of 0.01 seconds and additional GPU memory usage of 0.0 GB are orders of magnitude lower than all training-based methods. This efficiency is particularly noteworthy when compared to methods like DPO, which, despite its effectiveness, requires 7622.0 seconds of training time and 45.12 GB of GPU memory. Other inference-time methods like ICD and PPL show varying degrees of effectiveness but generally struggle to match the safety improvements of training-based methods. SFT with safety data mixing post-fine-tuning offers a more balanced approach, achieving a Harmfulness Score of 1.03 with reduce resource requirements of 779 seconds and 38.32 GB GPU memory.
Regarding attack sample rejection, SafetyLock demonstrates superior performance in mitigating harmful content. It achieves a Harmfulness Score of 1.04, equivalent to that achieved by models undergoing safety realignment via DPO, indicating its exceptional ability to reduce the generation of harmful content. Furthermore, SafetyLock’s AdvBench ASR of 0.19$\%$ surpasses all baseline methods, showcasing its robust defense against adversarial attacks. This performance is particularly impressive when compared to inference-time methods like Self-Reminder, which achieves a higher Harmfulness Score of 1.82 and an AdvBench ASR of 19.81%.

We further assess the models’ performance on benign inputs to ensure safety enhancements did not compromise normal text generation by selecting 500 test samples from the Alpaca dataset. The results reveal that SafetyLock preserves a 98.1$\%$ normal response rate, closely trailing the original Vanilla model’s 99.4$\%$. Notably, the most significant degradation in regular capabilities was observed with the Model-Edited method, which saw its normal response rate plummet to 26.8$\%$. Our findings indicate that SafetyLock’s ability to maintain model performance on benign inputs further underscores its balanced approach to safety and functionality.

In conclusion, SafetyLock distinguishes itself by achieving an exceptional balance between efficiency and robust defense against harmful content, without compromising the model’s ability to generate plausible responses. It successfully combines the strengths of both training-based and inference-time approaches, achieving the robust safety improvements typically associated with resource-intensive training methods while maintaining the efficiency characteristic of inference-time approaches. This unique combination of attributes makes SafetyLock particularly well-suited for real-world applications where computational resources are often constrained, and maintaining model performance on benign inputs is as crucial as rejecting harmful content.

### 4.4 SafetyLock’s Performance Against Combined Attacks

**Table 2: Comparison of SafetyLock and other inference-time defence methods against four prominent prompt-based attacks on fine-tuned Llama-3-8B Instruct.**
| Model | AutoDAN ASR | DeepInception ASR | GCG ASR | PAIR ASR |
| --- | --- | --- | --- | --- |
| Vanilla | 84.0 | 98.0 | 74.0 | 70.0 |
| ICD | 46.0 | 98.0 | 22.0 | 50.0 |
| PPL | 84.0 | 98.0 | 0.0 | 70.0 |
| Paraphrase | 32.0 | 96.0 | 58.0 | 74.0 |
| Relexicalization | 82.0 | 98.0 | 94.0 | 64.0 |
| Self-Reminder | 66.0 | 98.0 | 32.0 | 56.0 |
| Star Exam | 84.0 | 98.0 | 74.0 | 70.0 |
| SafetyLock | 4.0 | 2.0 | 10.0 | 14.0 |

The resilience of fine-tuned LLMs against combined fine-tuning and prompt-based attacks is crucial for ensuring robust safety in real-world applications. To further assess robustness, we introduced a combined attack scenario: fine-tuning model attacks followed by prompt-based attacks. We evaluated four commonly use prompt attack methods: AutoDAN (Liu et al., ), DeepInception (Li et al., ), GCG (Zou et al., ), and PAIR (Chao et al., ), comparing their performance against several inference-time defense techniques, as illustrated in Table [2](#S4.T2).

SafetyLock demonstrates exceptional effectiveness across all tested attack methods. For AutoDAN attacks, SafetyLock reduces the ASR to a mere 4.0$\%$, significantly outperforming other methods such as ICD (46.0$\%$) and Self-Exam (66.0$\%$). Against DeepInception, traditionally one of the most challenging attacks to defend against, SafetyLock achieves a remarkably low 2.0$\%$ ASR, while all other methods fail to provide any meaningful defense (98.0$\%$ ASR across the board). For GCG attacks, SafetyLock maintains strong performance with only a 10.0$\%$ ASR, second only to PPL’s 0.0$\%$ but considerably better than most other methods, including Vanilla (74.0$\%$) and Retokenization (94.0$\%$). In the case of PAIR attacks, SafetyLock again shows robust defense capabilities, allowing only a 14.0$\%$ ASR, outperforming all other tested methods.

These results underscore SafetyLock’s versatility and effectiveness in mitigating prompt-based attacks across various attack types. Its consistent performance demonstrates a comprehensive approach to model safety, addressing the complex challenges posed by diverse attack scenarios in language model deployment. The ability to maintain such low ASR across different attack methods suggests that SafetyLock provides a more generalizable and robust defense mechanism.

### 4.5 Generalization Capabilities of SafetyLock

Figure: Figure 6: Performance comparison of various methods on downstream tasks.
Refer to caption: /html/2410.10343/assets/x8.png

We assess a wide range of language understanding and generation capabilities to provide a comprehensive view of model performance based on various downstream tasks. Our experiments include a diverse set of benchmarks (Hosseini et al., ; Talmor et al., ; Arkil et al., ; Cobbe et al., ; Suzgun et al., ; Roy & Roth, ; Wei et al., ; Kojima et al., ; Weng et al., ; Zheng et al., ; Dubois et al., ): AddSub, AQUA, CommonSenseQA, GSM8k, MT-Bench, Alpaca, and AlpacaEval 2.0 .

As illustrated in Figure [6](#S4.F6), SafetyLock demonstrates a remarkable ability to maintain model performance across all tasks while ensuring safety. Unlike previous knowledge editing methods, which often led to significant performance degradation or incoherent outputs, SafetyLock preserves the model’s foundational capabilities. For instance, on the AddSub task, SafetyLock maintains a performance of 85.57%, closely matching the original model’s 86.33%, while other methods like Model-Edited show complete performance collapse. This trend is consistent across other tasks, with SafetyLock consistently performing on par with or slightly below the original model, in stark contrast to the severe degradation seen with other safety-aligned methods. These results highlight SafetyLock’s unique ability to enhance model safety without compromising its core functionalities, addressing a critical challenge in the deployment of safe and effective language models.

## 5 Conclusion

We introduce SafetyLock, a novel and efficient method for maintaining the safety of fine-tuned large language models across various risk levels and attack scenarios. Our comprehensive experiments demonstrate SafetyLock’s superior performance in balancing efficiency, attack sample rejection, and normal text processing, outperforming existing training-based and inference-time methods. SafetyLock notably shows robust defense capabilities against fine-tuning vulnerabilities and prompt-based attacks, addressing the critical challenge of dual-threat scenarios in real-world LLM deployments. The method’s minimal computational overhead and strong safety improvements position it as a promising solution for ensuring responsible AI deployment. Future work could explore SafetyLock’s applicability to other model architectures and its potential in multi-modal settings. Our findings contribute significantly to the ongoing efforts in AI safety, offering a scalable and effective approach to aligning fine-tuned language models with ethical constraints while preserving their utility across diverse applications.

## Reproducibility Statement

We have taken several steps to ensure the reproducibility of our results. The implementation details, datasets, and models used in our experiments are described in the corresponding sections of this paper, particularly in Sections [3.2](#S3.SS2), [3.4](#S3.SS4), and [4.2](#S4.SS2). We also provide the experimental settings and evaluation metrics in Sections [3.3](#S3.SS3) and [4.3](#S4.SS3). Furthermore, all hyperparameters, training code, and baselines are detailed throughout the relevant sections, ensuring that researchers can replicate our work using publicly available datasets and models.

## Appendix A Appendix

### A.1 Limitations

While SafetyLock demonstrates promising results in maintaining the safety of fine-tuned language models, it is important to acknowledge several limitations. Primarily, SafetyLock requires access to both model weights and intermediate activations for implementation, which may limit its applicability in scenarios where such access is restricted or unavailable. Additionally, the method employs a symmetric locking mechanism; consequently, if an unauthorized party gains access to the model weights or activation values, they could potentially reverse-engineer the process to unlock and bypass SafetyLock’s protections. Lastly, while SafetyLock shows strong performance against current attack methods, its long-term robustness against evolving adversarial techniques remains to be studied. These limitations present opportunities for future work to enhance and expand the capabilities of SafetyLock, ensuring its continued effectiveness in maintaining AI safety.

### A.2 Consistency of Harmlessness Directions in Fine-tuned Models

To validate SafetyLock’s effectiveness, we conducted a comprehensive analysis of the original Llama-3-Instruct 8B model and its fine-tuned versions under various risk levels. Our experimental setup was as follows:

We first extracted activation values from the 31st layer, 26th head of the Llama-3-8B Instruct model, which we identified as the most sensitive to harmlessness through linear regression, achieving the highest binary classification accuracy. We then performed forward computation on a safety dataset, saving the activation values of the last token for both safe and unsafe samples. Using 2D PCA for dimensionality reduction, we visualized the shift in activation values between safe and unsafe samples by connecting their center points with arrows, illustrating both the direction and magnitude of the shift.

Remarkably, we observed high similarity in these shifts across different risk levels (i.e., fine-tuning on data from different domains). To quantitatively assess the similarity between the safety directions found in the original model and those in the fine-tuned models, we employed KL divergence:

$$ $D_{KL}(P||Q)=\sum_{i}P(i)\log\left(\frac{P(i)}{Q(i)}\right)$ (6) $$

where $P$ and $Q$ represent the distributions of safety directions in the original and fine-tuned models, respectively.

To further illustrate the change in similarity during the fine-tuning process, we employed one-dimensional linear interpolation of weights (Peng et al., ). This method allows us to smoothly transition from the original model weights to the fine-tuned model weights, providing insight into how the safety directions evolve during the fine-tuning process. The interpolation is defined as:

$$ $\theta_{\alpha}=\theta+\alpha(\theta^{\prime}-\theta)$ (7) $$

where $\theta$ represents the weights of the original Llama-3 model, $\theta^{\prime}$ the weights of the fine-tuned model, and $\alpha\in[-0.2,1.2]$ is the interpolation parameter. We extend $\alpha$ slightly beyond the [0, 1] range to observe trends slightly before and after the actual interpolation points.

The interpolation process is implemented as follows:

- 1.
We first extract the state dictionaries of both the base model ($\theta$) and the fine-tuned model ($\theta^{\prime}$).
- 2.
For each layer, we compute the difference vector: $d_{1}=\theta^{\prime}-\theta$.
- 3.
We then create new weights for each $\alpha$ value: $\theta_{\alpha}=\theta+\alpha d_{1}$.
- 4.
These new weights are used to reconstruct a new state dictionary, maintaining the original structure and naming conventions of the model.

We use these interpolated models to compute the KL divergence between the safety directions of the original model and the interpolated models at each step. This results in a smooth curve showing how the similarity of safety directions changes as the model transitions from its original state to the fine-tuned state.

## Appendix B Mathematical Explanation of SafetyLock’s Effectiveness in Suppressing Harmful Outputs

In this section, we provide a mathematical justification for why SafetyLock can extract transferable safety directions from the original language model and effectively apply them to fine-tuned models to suppress harmful outputs. Our explanation is grounded in the properties of Transformer-based language models and the nature of fine-tuning on limited datasets.

### B.1 Activation Space and Safety Directions

Let us denote the activations of the original (pre-fine-tuned) language model at layer $l$ and head $h$ as $\mathbf{x}_{l,h}\in\mathbb{R}^{D}$, where $D$ is the dimensionality of the head’s output. During inference, these activations encode information about the generated tokens.

We define two sets of activations corresponding to safe and unsafe responses:

$$ $\displaystyle\mathcal{X}_{\text{safe}}$ $\displaystyle=\left\{\mathbf{x}_{l,h}^{\text{safe},i}\right\}_{i=1}^{N_{\text{safe}}},$ (8) $\displaystyle\mathcal{X}_{\text{unsafe}}$ $\displaystyle=\left\{\mathbf{x}_{l,h}^{\text{unsafe},i}\right\}_{i=1}^{N_{\text{unsafe}}},$ (9) $$

where $N_{\text{safe}}$ and $N_{\text{unsafe}}$ are the numbers of safe and unsafe samples, respectively.

We compute the *safety direction* $\bm{\theta}_{l,h}\in\mathbb{R}^{D}$ as the mean difference between the activations for safe and unsafe responses:

$$ $\bm{\theta}_{l,h}=\frac{1}{N_{\text{safe}}}\sum_{i=1}^{N_{\text{safe}}}\mathbf{x}_{l,h}^{\text{safe},i}-\frac{1}{N_{\text{unsafe}}}\sum_{i=1}^{N_{\text{unsafe}}}\mathbf{x}_{l,h}^{\text{unsafe},i}.$ (10) $$

This vector represents the average shift in activation space needed to move from an unsafe response towards a safe one.

### B.2 Preservation of Safety Directions During Fine-Tuning

Fine-tuning a language model on a new dataset modifies its parameters to adapt to specific tasks or domains. However, when the fine-tuning dataset is limited in size or scope, the changes to the model’s internal representations are often localized and do not significantly alter the global structure of the activation space (Golechha & Dao, ; Godfrey et al., ).

Let $\mathbf{\tilde{x}}_{l,h}$ denote the activations of the fine-tuned model at layer $l$ and head $h$. Empirically, we observe that there exists a strong linear relationship between the activations of the original and fine-tuned models:

$$ $\mathbf{\tilde{x}}_{l,h}\approx\mathbf{x}_{l,h}+\Delta\mathbf{x}_{l,h},$ (11) $$

where $\Delta\mathbf{x}_{l,h}$ represents the change in activations due to fine-tuning, which is relatively small in magnitude compared to $\mathbf{x}_{l,h}$ for many dimensions.

Moreover, the safety direction $\bm{\theta}_{l,h}$ computed from the original model remains relevant in the fine-tuned model because the relative differences between safe and unsafe activations are preserved:

$$ $\bm{\tilde{\theta}}_{l,h}=\left(\mathbf{\tilde{x}}_{l,h}^{\text{safe}}-\mathbf{\tilde{x}}_{l,h}^{\text{unsafe}}\right)\approx\left(\mathbf{x}_{l,h}^{\text{safe}}-\mathbf{x}_{l,h}^{\text{unsafe}}\right)=\bm{\theta}_{l,h}.$ (12) $$

This approximation holds under the assumption that fine-tuning does not disproportionately affect the dimensions critical for encoding safety-related information.

### B.3 Effectiveness of Activation Intervention

During inference with the fine-tuned model, we intervene by adjusting the activations along the safety direction:

$$ $\mathbf{\tilde{x}}_{l,h}^{\text{intervened}}=\mathbf{\tilde{x}}_{l,h}+\alpha\left(\bm{\sigma}_{l,h}\odot\bm{\theta}_{l,h}\right),$ (13) $$

where:

- •
$\alpha\in\mathbb{R}$ is the scaling factor controlling the intensity of the intervention.
- •
$\bm{\sigma}_{l,h}\in\mathbb{R}^{D}$ is the standard deviation vector of activations along each dimension, capturing the typical variability.
- •
$\odot$ denotes element-wise multiplication.

This adjustment effectively shifts the activations towards regions in the activation space associated with safe responses. Since the safety direction $\bm{\theta}_{l,h}$ is approximately preserved in the fine-tuned model, this intervention remains effective.

### B.4 Impact on Output Probabilities

The language model generates the next token based on a probability distribution computed from the final activations. Adjusting the activations as in Equation equation [13](#A2.E13) influences the logits $\mathbf{z}\in\mathbb{R}^{V}$ (where $V$ is the vocabulary size) before the softmax function:

$$ $\mathbf{z}^{\text{intervened}}=\mathbf{z}+W_{\text{head}}\left(\alpha\left(\bm{\sigma}_{l,h}\odot\bm{\theta}_{l,h}\right)\right),$ (14) $$

where $W_{\text{head}}\in\mathbb{R}^{V\times D}$ is the weight matrix projecting activations to logits.

The adjustment $\Delta\mathbf{z}=W_{\text{head}}\left(\alpha\left(\bm{\sigma}_{l,h}\odot\bm{\theta}_{l,h}\right)\right)$ biases the logits towards tokens that are more likely in safe responses and away from those prevalent in unsafe responses.

### B.5 Suppressing Harmful Outputs

The probability of generating a harmful token $t_{\text{harm}}$ is given by:

$$ $P(t_{\text{harm}})=\frac{\exp\left(z_{t_{\text{harm}}}^{\text{intervened}}\right)}{\sum_{i=1}^{V}\exp\left(z_{i}^{\text{intervened}}\right)}.$ (15) $$

By decreasing $z_{t_{\text{harm}}}^{\text{intervened}}$ relative to other logits, we reduce $P(t_{\text{harm}})$. Since the intervention shifts the activations towards safe regions, the logits for harmful tokens are decreased, and the model is less likely to generate harmful outputs.

### B.6 Transferability Across Models

The key to SafetyLock’s transferability lies in the similarity of safety directions between the original and fine-tuned models. Since the fine-tuning process does not significantly alter the relative positions of safe and unsafe activations in the activation space (as per Equation equation [12](#A2.E12)), the safety directions computed from the original model remain effective when applied to the fine-tuned model.

This property is supported by empirical observations of low Kullback–Leibler (KL) divergence between the activation distributions of the original and fine-tuned models (see Figure [2](#S3.F2) in Section [3.3](#S3.SS3)). The minimal divergence indicates that the overall structure of the activation space, especially along dimensions relevant to safety, is preserved during fine-tuning.

### B.7 Conclusion

Mathematically, SafetyLock leverages the preserved safety directions in the activation space to adjust the model’s internal computations towards generating safe outputs. By intervening along these directions, we effectively suppress harmful responses without requiring retraining or fine-tuning of the model. The minimal changes to the activation distributions during fine-tuning ensure that the safety directions remain applicable, allowing for efficient and transferable safety interventions across different models and fine-tuning scenarios.

This theoretical explanation provides a foundation for understanding the effectiveness of SafetyLock in suppressing harmful outputs while maintaining the model’s overall performance on benign tasks.

## Appendix C The Risks of Fine-tuning LLMs and Experimental Setup

HEx-PHI (Qi et al., ) is based on 11 categories of prohibited use cases merged from Meta’s Llama-3 acceptable use policy and OpenAI’s usage policies: (1) Illegal Activity, (2) Child Abuse Content, (3) Hate, Harass, Violence, (4) Malware, (5) Physical Harm, (6) Economic Harm, (7) Fraud, Deception, (8) Adult Content, (9) Political Campaigning, (10) Privacy Violation Activity, and (11) Tailored Financial Advice. The dataset includes 30 examples per category, totaling 330 examples. This ensures a comprehensive safety evaluation aligned with industry-standard usage policies.

For Risk-1, we use negative samples from the HH-RLHF preference dataset. We select 10, 100, 1000, and 10000 samples respectively and trained for 5 epochs with a learning rate of 2e-5. For Risk-2, we use 10 samples from Qi et al. () and trained for 5 epochs with a learning rate of 2e-5. For Risk-3, we use the first 50,000 samples from the Alpaca dataset (Wang et al., ) and trained for 5 epochs with a learning rate of 2e-5 (^1^11We use the official fine-tuning code [https://github.com/meta-llama/llama-recipes](https://github.com/meta-llama/llama-recipes)).

Recognizing the potential of existing approaches to address safety issues in fine-tuned language models, we conducted comparative analyses across two categories as the same time: training-based and inference-time methods. For training-based approaches, we evaluated PPO, DPO, SFT (with safety data mixed during fine-tuning), SFT (with safety data mixed post-fine-tuning), and model-editing. Inference-time methods included ICD, PPL, Paraphrase, Retokenization, Safe-Reminder, and Self-Exam. These methods were assess based on efficiency, attack sample rejection rate, and normal text rejection rate, providing a comprehensive evaluation of their effectiveness in maintaining model safety while preserving functionality. This multi-faceted approach allows us to rigorously examine the trade-offs between safety and performance.

Specifically, to ensure reproducibility, we followed past experimental settings and use 2000 safety data points from Bianchi et al. () for SFT experiments. We considered two experimental settings for SFT. The first is After Training, which simulates the scenario where safety disappears after fine-tuning the language model and needs to be restored. This applies to all fine-tuned language models. The second is During Training, which simulates starting from the original model and requiring the mixing of additional safety data during training to prevent safety disappearance. However, the limitation of this method is that it still requires retraining for already fine-tuned language models. For PPO, we also use 2000 samples from Bianchi et al. (), and we use LlamaGuard-7b (Bhatt et al., ) as the Reward model. For DPO, based on the 2000 samples, we use samples generated by the fine-tuned language model (almost all of which are harmful) as negative samples for training. For the Model-Edited method, we use the most common Detoxifying with Intraoperative Neural Monitoring (DINM) method and followed the original setup using SafeEdit data(^2^22[https://huggingface.co/datasets/zjunlp/SafeEdit](https://huggingface.co/datasets/zjunlp/SafeEdit)) for editing.

## Appendix D Analysis of SafetyLock’s Intervention

Figure: Figure 7: Impact of SafetyLock’s intervention distance ($\alpha$) on model safety metrics for Llama-3-8B and Llama-3-70B models. The graphs show Harmfulness Average Score, Harmfulness Average Rate, and AdvBench ASR across different $\alpha$ values. Note that for these experiments, the intervention degree K is set to 24, indicating the number of attention heads influenced by SafetyLock.
Refer to caption: /html/2410.10343/assets/x9.png

Distance $\alpha$. Our experimental results, as illustrated in Figure [7](#A4.F7), demonstrate the significant influence of SafetyLock’s intervention distance ($\alpha$) on model safety across different model sizes. For both Llama-3-8B and Llama-3-70B, we observe a clear U-shaped trend in harmfulness metrics as $\alpha$ increases. Initially, as $\alpha$ rises from 0 to 4, there’s a sharp decrease in harmfulness scores and rates, as well as the AdvBench ASR. This indicates that moderate intervention effectively enhances model safety. However, beyond $\alpha=4$, we see a gradual increase in these metrics, suggesting that excessive intervention may lead to unintended consequences, potentially disrupting the model’s learned safety boundaries. Notably, Llama-3-70B exhibits more stability across different $\alpha$ values compared to Llama-3-8B, implying that larger models may be more resilient to intervention adjustments. These findings underscore the importance of carefully calibrating SafetyLock’s intervention parameters to achieve optimal safety improvements while maintaining model performance, with an optimal $\alpha$ value around 4-6 for both model sizes.

Figure: Figure 8: Impact of SafetyLock’s intervention degree (K) on model safety metrics for Llama-3-8B and Llama-3-70B models. The graphs illustrate the Harmfulness Average Score, Harmfulness Average Rate, and AdvBench ASR across different K values, ranging from 0 to 96. Lower scores indicate better safety performance. Note the rapid improvement in safety metrics as K increases from 0 to 6, followed by more gradual enhancements up to K=24, with a slight uptick at K=96 for some metrics.
Refer to caption: /html/2410.10343/assets/x10.png

Degree $K$. Our experiments, as illustrated in Figure [8](#A4.F8), reveal the crucial role of SafetyLock’s intervention degree (K) in enhancing model safety across different model sizes. For both Llama-3-8B and Llama-3-70B, we observe a rapid improvement in safety metrics as K increases from 0 to 6, followed by a more gradual enhancement up to K=24. This trend is consistent across all three metrics: Harmfulness Average Score, Harmfulness Average Rate, and AdvBench ASR. For Llama-3-8B, the most significant improvements occur between K=0 and K=6, with the Harmfulness Average Score dropping from about 4.0 to 1.7, and the Harmfulness Average Rate decreasing from 70% to around 15%. The AdvBench ASR shows a similar sharp decline. Beyond K=6, the improvements become more incremental, with optimal performance generally achieved around K=24.

Llama-3-70B exhibits a similar pattern but with overall lower harmfulness scores and rates. The initial drop in harmful metrics is less dramatic, suggesting that larger models may have inherently better safety characteristics. However, the trend of improvement with increasing K values remains consistent. Interestingly, for both model sizes, there’s a slight uptick in harmfulness metrics for very high K values (K=96), particularly noticeable in the Llama-3-8B model. This suggests that excessive intervention might slightly degrade the model’s learned safety boundaries, emphasizing the importance of finding an optimal K value. These findings underscore the effectiveness of SafetyLock in improving model safety, with the most significant gains achieved at relatively low K values (6-24). This implies that targeted intervention on a subset of attention heads can yield substantial safety improvements without the need for exhaustive modification of the model architecture.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1) Can you explain the need for the standard deviation of activations $\sigma_l^h$ in equation 4? Any experiments that show with and without $\sigma_l^h$ effect?

2) Can you please elaborate on section 4.5? The setting and motivation aren't very clear to me.
soundness: 3
strengths: **Strengths**

- The paper proposes a new idea to prevent safety compromise owing to downstream fine-tuning of the model. While there are existing safety vector-based approaches to prevent safety compromise (with comparable memory and time footprint), the activation-based safety vector computation is a novel contribution.

- The experiments are sound, demonstrating the method's effectiveness in mitigating safety issues across various risk levels, including explicitly harmful, implicitly harmful, and benign datasets. Additionally, the method preserves most of the model's generic utility.
summary: It is now well known that fine-tuning language models comes with a safety compromise [R0-R4], i.e., the safety alignment gets disturbed when a safety-aligned model is tuned to learn something new. This paper introduces SafetyLock, a method designed to maintain safety alignment post-fine-tuning. For the top-K heads in the model, a safety direction is computed using a preference-style safety dataset. The top-K heads are identified based on the quality signals they provide for safe-unsafe text classification. In the fine-tuned model, the safety vector of these top-K heads is added to the output of attention computations before the up-projection layer. Overall, the method appears sound and is shown to preserve harmlessness on Llama-3-8B, 70 Instruct, and Mistral-Large-2 123B when tested on three risk levels: benign, explicit, and implicit harmful datasets.
weaknesses: **Weaknesses**

- The paper largely ignores comparisons with a line of closely related work, beginning with "Language models are Homer Simpson! Safety re-alignment of fine-tuned language models through task arithmetic," (R1) which also introduced the concept of a *safety vector*. Therefore, I do not agree with the claim, *"we are the first to consider locating safety vectors and then restoring the safety of fine-tuned LLMs using an inference-time intervention method,"* as there are a series of similar existing methods, such as R1-R4, and more within this line of research. These methods also use safe-unsafe pairs to derive a safety vector. I suggest the authors properly cite and compare their work with these studies.

- The experiments demonstrating the utility of the method are limited. I understand that works like Qi et al. (2023b) [R0] largely rely on three risk categories, which is sufficient to prove that the problem of safety compromise due to fine-tuning exists. However, solving the problem would require more substantial evidence than relying on a similar set of experiments. It would be beneficial to see the method applied across different fine-tuning domains (such as math, code, and various other tasks) and observe the domain-specific performance with and without the safety measures. For instance, fine-tune on GSM8K, evaluate the test-set performance post task-specific tuning with and without the safety lock, and compare these results with those of related works.

References:

[R0] Qi, Xiangyu, et al. "Fine-tuning aligned language models compromises safety, even when users do not intend to!." arXiv preprint arXiv:2310.03693 (2023).

[R1] Bhardwaj, Rishabh, Do Duc Anh, and Soujanya Poria. "Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic." arXiv preprint arXiv:2402.11746 (2024).

[R2] Zhao, Weixiang, et al. "Towards comprehensive and efficient post safety alignment of large language models via safety patching." arXiv preprint arXiv:2405.13820 (2024).

[R3] Hazra, Rima, et al. "Safety Arithmetic: A Framework for Test-time Safety Alignment of Language Models by Steering Parameters and Activations." arXiv preprint arXiv:2406.11801 (2024).

[R4] Yi, Xin, et al. "A safety realignment framework via subspace-oriented model fusion for large language models." arXiv preprint arXiv:2405.09055 (2024).

(and more relevant works exist in this line)

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: In Section 4.1, you mentioned that the learning rate is set to 2e-5. I'm curious about the effect of this learning rate on the experimental results. A learning rate of 1e-6 is typically used for continued pretraining, which has a smaller impact on model performance including general capability and safety alignment.
soundness: 2
strengths: 1. They implement the method of ITI [1] on the safety of LLM and verify the effectiveness of the safety head.

2. Experiment results show the existence of safety heads and can efficiently enhance models' safety on different models.



[1] Inference-Time Intervention: Eliciting Truthful Answers from a Language Model, NeurIPS 2023
summary: The paper proposes a method to identify safety-relevant attention heads, interpreted as the model's safety mechanism. Furthermore, experiments are conducted to verify the effectiveness of these safety heads and their transferability across models.
weaknesses: 1. The method is not novel. It totally follows the method in ITI [1] without considering the properties of safety itself, which influence the accuracy of the analysis.  (1) As noted in the abstract (line 14), only 10 sentences can compromise the models' safety mechanisms. However, as illustrated in line 193, the safety components in LLMs are composed of three-fourths for both LLama3-8b and LLama3-70b. If the safety mechanism is easily breached, it may be due to a relatively small number of parameters. (2) The head is not a fine-grained aspect of analysis, as there are papers exploring neuron-level safety mechanisms [2]. Additionally, the feed-forward layer is important for safety, as models may refuse to extract harmful knowledge and only extract safety knowledge such as "Apologize I can not ....." [3]. However, the paper ignores this and focuses solely on heads in self-attention.

2. The analysis is not comprehensive. Figure 2 examines the "safety direction" at layer 31, the final layer of LLama3-8B. This analysis is biased, as it only evaluates the output of different models. It's expected that embeddings will have some relation or differences because the models' outputs vary. What would be more valuable is an analysis of the intermediate layers, exploring how the mechanism identifies and mitigates harmful information. However, the paper lacks further analysis, which would have been interesting.

3. The experiments in Table 1 lack baselines, and the baselines in Figure 5 are all trivial. It’s important to compare lightweight defense methods, such as [4], to demonstrate the effectiveness of your approach.


[1] Inference-Time Intervention: Eliciting Truthful Answers from a Language Model, NeurIPS 2023

[2] Finding Safety Neurons in Large Language Models, Arxiv 2024

[3] Transformer Feed-Forward Layers Are Key-Value Memories, ACL 2021

[4] Improving Alignment and Robustness with Circuit Breakers, NeurIPS 2024

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: See weakness
soundness: 2
strengths: 1. The main contribution of this paper is to point out an important safety issue in distributing LLMs: How can we guarantee the safety of an LLM after we release it, where users can use fine-tuning to bypass its safety alignment. Although the proposed solution is less satisfied to the reviewer (See Weakness 1.), the reviewer admit the value of the research question itself.

2. This paper provides insights into the inner machanism of LLMs on safety.
summary: This paper introduces SafetyLock, a method to ensure the safety of LLMs post-fine-tuning, addressing the risks that traditional alignment measures fail to mitigate. SafetyLock leverages the discovery that fine-tuned models retain safety-related activation patterns from their base models, allowing the extraction and application of Meta-SafetyLock directions to reinforce safety efficiently. Experiments show that SafetyLock significantly reduces harmful behavior, achieving a response rate reduction from 60% to below 1% with minimal computational cost, offering a scalable and robust solution for safer LLM customization.
weaknesses: 1. Regarding the research question. The paper aims to propose a *lock* to ensure that the safety behavior of an released LLM would not be jailbreaked. To this end, the major drawback of the proposed method is that: the safety lock is applied to the LLM after fine-tuning, which means that the fine-tuned LLM should still be preserved and served by its original provider. This is effective for proprietary LLMs such as GPT-4, where the model checkpoint is still kept by the company even after fine-tuning via API. However, for open-sourced LLMs, such as LLaMA series, the proposed method failed as the provider cannot apply the safety lock to fine-tuned LLMs. I would expect a method which can by applied instantly before the release of the checkpoint and would protect the LLM from any unsafe modifications.


2. Regarding the methodology. The *safety direction* is calculated according to Eq.3, which involves the direct comparison between two different sentences. I would suspect the rationality of the comparison. For example, if the last 3 tokens of the safe reponse is *I can't answer* and the last 3 tokens of the unsafe response is *make a bomb*. The comparison between `I` and `make`, `can't` and `a`, and `answer` and `bomb` does not make sense. The meaning of the last $r$ tokens is not strictly aligned between the safe and the unsafe responses. The author should give more explanations on this. Besides, there are different aspects to investigate the safety mechanism, such as activations of individual neurons that are closely related to safety behaviors. They are not fully explored.



3. Minor issues (Does not affect my scores)
    - Missing references. This paper should also survey papers for `LLM safety mechanism`, such as *A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity*, *Finding Safety Neurons in Large Language Models*, etc.
    - Formats. Check the font of the sub-title in Figure 2. Also, the sub-title overlaps with the figures.
    - Line 296 - Line 297: 2e-5 --> $2 \times 10^{-5}$



I am quite torn between giving a score of 5 or 6. I fully acknowledge the importance of emphasizing to the research community the need for robust safety mechanisms in publicly released LLMs. However, I am somewhat dissatisfied with the proposed method, as it does not thoroughly address the problem. Considering that this is a relatively early attempt (though perhaps I am not fully aware of previous work in this area) and revealing a problem is more important than a perfect solution, I am inclined to give a somewhat positive score of 6. I hope the authors can adequately address my concerns in the rebuttal, and I will consider adjusting my score to a more negative one if necessary.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Would it be better to use the cosine similarity or JS divergence metric to make the quantitative analysis in Figure 2(e-g) since KL divergence is asymmetric?
2. How about the performance of SafeLock on XSTest since I noticed this dataset is evaluated in ICD baseline.
soundness: 3
strengths: 1. This paper discovers that fine-tuned models retain similar safety-related activation representations to their base models, so the safety directions computed from the original model remain effective when applied to the fine-tuned model.
2. SafetyLock reduces the harmful instruction response rate in fine-tuned models from 60% to below 1%, offering a scalable, non-invasive solution for ensuring the safety of customized LLMs.
summary: The paper proposes a novel approach, SafetyLock, to mitigate safety risks in fine-tuned large language models (LLMs). SafetyLock leverages the discovery that fine-tuned models retain similar safety-related activation representations to their base models, enabling the extraction of a Meta-SafetyLock, a set of safety bias directions. This allows for efficient and transferable safety alignment in fine-tuned models, reducing the harmful instruction response rate.
weaknesses: Originality
1. The concept of safety-related directions  seems similar to previous works on concept activation vectors [1] and the design of  adaptive interaction on MHA is similar to [3] where the only difference lies in how to select the intervention heads. 

Quality
1. Will the choice of preference-style safety dataset (line 182) affect the identification of safety-related attention heads (line 188)? It would be better to verify the identified safety-related attention heads stay consistent no matter what safety dataset is selected. 
2. It is not clear why the element-wise multiplication between the standard deviation of activations and safety directions is needed. Although authors claim that std term captures the variability of the activations, previous works [1,2] only use the safety directions to intervention the model and achieve good performance. Thus, ablation study on std should be considered in this paper otherwise I may doubt whether this part of the design is redundant.
3. The number of chosen Top-K heads are different in different models. Authors claimed in Appendix D that SafetyLock’s intervention degree (K) has impact on model safety. Thus, the time used in the choice of K should also be considered in the method efficiency evaluation. 

Clarity
1. Line 308: citation of ICD may be wrong? Besides, the paper contains several duplicate citations (like Mengru Wang et al.), which should be checked and corrected.
2. Line 201: r is the number of final tokens considered, but how authors choose r? I checked the paper but didn’t find the clarification. If it is equal to the hidden size, such safety direction is the same with several previous works [1,2]. If not, whether the method performance is sensitive to r value should also be considered. 
3. Figure 2: captions of subfigures(b-d) have the extra spaces and captions of subfigures(e-g) are too close to the figure content.

[1] Zou, A., Phan, L., Chen, S., Campbell, J., Guo, P., Ren, R., ... & Hendrycks, D. (2023). Representation engineering: A top-down approach to ai transparency. arXiv preprint arXiv:2310.01405.

[2] Wang, P., Zhang, D., Li, L., Tan, C., Wang, X., Ren, K., ... & Qiu, X. (2024). Inferaligner: Inference-time alignment for harmlessness through cross-model guidance. arXiv preprint arXiv:2401.11206.

[3] Wang, T., Jiao, X., He, Y., Chen, Z., Zhu, Y., Chu, X., ... & Ma, L. (2024). Adaptive Activation Steering: A Tuning-Free LLM Truthfulness Improvement Method for Diverse Hallucinations Categories. arXiv preprint arXiv:2406.00034.

Meta-review: The paper proposes SafetyLock, a method to maintain safety in fine-tuned LLMs by extracting safety bias directions from base models. While the work shows promising results in reducing harmful responses and introduces an efficient approach for safety restoration, several limitations exist. Key strengths include the discovery that fine-tuned models retain safety-related activation patterns and achieving sub-second safety restoration. Major weaknesses include limited applicability to open-source models, insufficient comparison with recent work like Circuit Breakers, and methodological concerns about token alignment in safety direction calculation. Given these significant limitations despite novel contributions, I recommend rejecting this paper.

Decision: Reject