Locking Down the Finetuned LLMs Safety

Authors: Minjun Zhu, Yixuan Weng, Linyi Yang, Yifan Wei, Ningyu Zhang, Yue Zhang
Primary area: generative models
Keywords: large language models, Safety, Finetuned, fast

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.

Decision: reject
Reviewer scores: 5, 3, 6, 5

Reviews and discussion:
<thread>
creation date:
1730630356294

thread_type:
official_review

<thread_object>
creator: reviewer_sw5b

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730647394970

thread_type:
official_review

<thread_object>
creator: reviewer_zum2

<content>
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.

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.

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.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708677717

thread_type:
official_review

<thread_object>
creator: reviewer_uk5c

<content>
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.

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

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730741859146

thread_type:
official_review

<thread_object>
creator: reviewer_etzc

<content>
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.

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.

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)

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732735233142

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Responses and Summary of Revisions

comment:
We sincerely thank all reviewers for their careful and constructive feedback, which has helped us significantly improve our work. We are encouraged by the reviewers' recognition of several key strengths:

- "The paper proposes a new idea to prevent safety compromise owing to downstream fine-tuning of the model" (*Review eTzc*)
- "Experiments show that SafetyLock significantly reduces harmful behavior" (*Review ZuM2*)
- "The paper provides insights into the inner mechanism of LLMs on safety" (*Review ZuM2*)
- "The experiments are sound, demonstrating the method's effectiveness in mitigating safety issues across various risk levels" (*Review eTzc*)

Based on your constructive comments, we have made substantial improvements to strengthen our paper's contributions and address the limitations. The major enhancements include:

1. **Comparison with Circuit Breakers (NeurIPS 2024)**: We have added comprehensive comparison with Circuit Breakers, including performance evaluation across three risk levels (Level-1: ASR from 84.62% to 0.19%, Level-2: 27.12% to 5.19%, Level-3: 94.04% to 0.19%), efficiency analysis (SafetyLock: 5min + 0.1s/model vs. Circuit Breakers: 22min/model), superior performance demonstration on their original benchmarks (AutoDAN, PAIR, GCG), SafetyLock achieved superior safety performance while maintaining 0% ASR across all scenarios. And detailed analysis in Appendix D.6.

1. **Citation and Related Work**: We have thoroughly expanded our literature review and comparisons with closely related work, added comprehensive comparisons with safety vector approaches (R1-R4 in Review eTzc), included recent works on safety mechanisms and neuron-level analysis, and removed claims of being "first" while properly positioning our contributions (Section 1, 2).
2. **Domain-Specific Experiments**: We have substantially expanded our experimental validation by adding GSM8K mathematical reasoning case study, including detailed performance analysis across different domains, and comparing domain-specific capabilities with and without SafetyLock (Appendix D.4).
3. **Activation Analysis**: We have enhanced our technical analysis by adding ablation studies on standard deviation term effects, including multi-head activation pattern analysis, and expanding intermediate layer analysis beyond the final layer (Section 3.2, Appendix D.3).
4. **Technical Clarifications**: We have added detailed analysis of learning rate impacts (1e-6 vs 2e-5), clarified token window size selection (r parameter), and expanded explanation of intervention degree (K) selection (Appendix D.1, D.2).
5. **Experimental Robustness**: We have added XSTest benchmark evaluation, included comprehensive baselines comparison, and enhanced analysis of safety direction consistency (Section 4.4, Table 2).
6. **Implementation Details**: We have added ablation studies on activation normalization, clarified the choice of preference-style safety dataset, and detailed the impact of token alignment in safety direction calculation (Appendix C, D).
7. **Open-Source Considerations**: We have added detailed discussion on deployment strategies, included hybrid approach combining transparency with controlled access, and expanded practical guidelines for implementation (Appendix E).
8. **Presentation Improvements**: We have fixed citation formatting and duplicates, improved figure captions and spacing, and enhanced overall clarity and readability throughout the paper.

We believe these comprehensive improvements address the key concerns raised while strengthening the paper's technical depth and practical impact. The additions substantially enhance the paper's contribution to the field of LLM safety alignment. We would be honored if the reviewers find these enhancements worthy of a more favorable evaluation. We welcome any additional feedback or questions from the reviewers.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.