Towards Understanding Safety Alignment: A Mechanistic Perspective from Safety Neurons

Thank you for your constructive and valuable comments. Here is our response to the weaknesses and questions.

W1 The presentation of this paper can be substantially improved

Thank you for your suggestion, and we apologize for any unclear statements. We have revised the corresponding sections in the newly submitted version to improve clarity.

W2 Safety neurons in related work

Please refer to our general response. We also added a discussion on this aspect in the related work section of our newly uploaded paper.

W3 What is the alignment in 4.1？

We agree that the released chat- or instruct- versions of current LLMs are generally safe, but to conduct experiments in a more controllable setting, we chose to begin alignment from the pre-trained base models. As mentioned at the end of Section 3, we refer to "the pre-trained LLMs before SFT (denoted as Base)." Additionally, the names used in Section 4.1, such as Mistral-7b-v0.1, refer to the official model names, not the abbreviated versions of the aligned chat models. Sorry for the potential misunderstandings. The results in Table 2 show that the base models used in this paper are not safe enough and our alignment improves safety.

W4 It is necessary to include some advanced jailbreaking attacks for evaluation

We believe our practice is appropriate since using red teaming benchmarks to evaluate model safety is a very common setting used by existing works, such as the two papers mentioned in your review. Moreover, similar to the misunderstanding in Weakness #3, we agree that current LLMs can generally refuse harmful questions, but this is the case for chat- and instruct- models. For the based models used in our experiments, they are clearly not safe enough in the adopted benchmarks (shown in Table 2).

W5 Can the conclusion generalize to larger models?

In general, we believe it is a common practice to use this model size in interpretability research [1-5], but we agree this is a valid concern and are happy to add more experiments. Due to the more computation and limited resources, more experiments are still running. We provide the results for Llama2-13B below (the format is similar to Table 2), which show similar trends to the original experiments in the paper.

[1] Zhenhong Zhou, Haiyang Yu, Xinghua Zhang, Rongwu Xu, Fei Huang, Kun Wang, Yang Liu, Junfeng Fang, & Yongbin Li. (2024). On the Role of Attention Heads in Large Language Model Safety.

[2] Shen Li, Liuyi Yao, Lan Zhang, & Yaliang Li. (2024). Safety Layers in Aligned Large Language Models: The Key to LLM Security.

[3] Zeping Yu, & Sophia Ananiadou. (2024). Neuron-Level Knowledge Attribution in Large Language Models.

[4] Ameen Ali, Lior Wolf, & Ivan Titov. (2024). Mitigating Copy Bias in In-Context Learning through Neuron Pruning.

[5] Alessandro Stolfo, Ben Wu, Wes Gurnee, Yonatan Belinkov, Xingyi Song, Mrinmaya Sachan, & Neel Nanda. (2024). Confidence Regulation Neurons in Language Models.

Thank you for your constructive and valuable comments. Here is our response to the weaknesses and questions.

W1 & Q1 How does our method differ from prior ones that target other types of critical neurons?

Please refer to Section 7 of the revised paper (Lines 499-516) for discussions on related neuron-based works.

Also please refer to our general response for discussions on related works on interpreting safety. Thanks for the question and we will add the discussions if more space is permitted.

W2 Limited Discussion

Please refer to our general response.

Q2 Can the "dynamic activation patching" method be generalized

Dynamic activation patching is task-agnostic. At the very least, Figure 4(b) in our paper demonstrates its effectiveness in altering the model’s helpfulness. As for whether it can be extended to other aspects of value alignment, we believe it is possible and this is an interesting direction for future work.

Q3 Mechanistic insight from safety neuron

Our current mechanistic insight suggests that safety and helpfulness may share the same set of neurons but exhibit different activation patterns on these neurons, which could potentially explain the alignment tax phenomenon. We think your suggestion about exploring patterns among the "safety neurons" related to specific types of safety risks to be a very interesting perspective. We plan to investigate this further in future work and sincerely thank you for proposing this valuable direction.

Q4 Overhead of safeguard applications

The overhead of our safeguard mechanism primarily comes from a logistic regression classifier. When using activations from only 1,500 neurons, this requires merely computing the inner product of 1,500-dimensional vectors, which is negligible compared to the billions of parameters in an LLM. In fact, if certain outputs can be rejected early, the process could even accelerate generation. Based on our measurements, the classification step takes less than 0.01 seconds, accounting for less than 1/2500 of the total inference time.

[1] Xiaozhi Wang, Kaiyue Wen, Zhengyan Zhang, Lei Hou, Zhiyuan Liu, & Juanzi Li. (2022). Finding Skill Neurons in Pre-Trained Transformer-Based Language Models.

[2] Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, & Furu Wei. (2022). Knowledge Neurons in Pretrained Transformers.

[3] Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, & Dimitris Bertsimas. (2023). Finding Neurons in a Haystack: Case Studies with Sparse Probing.

[4] Zeping Yu, & Sophia Ananiadou. (2024). Neuron-Level Knowledge Attribution in Large Language Models.

[5] Ameen Ali, Lior Wolf, & Ivan Titov. (2024). Mitigating Copy Bias in In-Context Learning through Neuron Pruning.

[6] Alessandro Stolfo, Ben Wu, Wes Gurnee, Yonatan Belinkov, Xingyi Song, Mrinmaya Sachan, & Neel Nanda. (2024). Confidence Regulation Neurons in Language Models.

[7] Wes Gurnee, Theo Horsley, Zifan Carl Guo, Tara Rezaei Kheirkhah, Qinyi Sun, Will Hathaway, Neel Nanda, & Dimitris Bertsimas. (2024). Universal Neurons in GPT2 Language Models.

Thank you for your acknowledgment and for giving us the chance to explain. We are so sorry for the potential confusion caused by the original response, which is inaccurate and could be explained in a dangerous way like we are criticizing the related works. Here are our clarifications regarding your questions. To avoid spreading confusion, we have revised the last comment (the old version is visible in the revision history if needed) and the PDF.

Question 1: Clarification of "Properties"

• Causal Effect: In this paper, the causal effect is measured by the formula in Equation 4.
• Sparsity: The sparsity in this work is measured by the proportion of neurons required to achieve a decent level of causal effect (e.g., 0.9) on model safety. In our main results, safety neurons constitute approximately 5% of all neurons.
• Transferability: Transferability refers to how well the safety neurons identified in one dataset also work on others. In our experiments, we identified safety neurons in Beavertails, and these neurons also work well on other red-teaming benchmarks, such as in Table 2.
• Stability on Training: This evaluates the consistency of safety neurons identified across models trained with different random seeds. We measured this using both neuron overlap and Spearman rank correlation, achieving values above 0.95 across the three model families. Additionally, the narrow error bars in Figure 2 also exhibit stability.

Question 2: Controlling Properties

We apologize for the imprecise expression. The properties here do not refer to the properties of safety neurons in Section 4. Here we use “properties” to refer to the phenomena or properties of LLMs that we want to interpret by identifying neurons. For instance, the property corresponding to safety neurons in this context is model safety. In [7], the authors identified universal neurons that are responsible for “properties” like alphabet, position, suppression, etc. However, the method described in [7] does not allow specifying the “properties” to be interpreted before identifying neurons, and thus cannot be directly applied to the goal of this work, i.e., interpreting model safety.

Question 3: Comparison with [6]

First, we need to clarify that the “reproducibility issue” definitely does not mean there are difficulties in reproducing the original results from our side. It is a bad expression for “the method in [6] cannot be directly reused in our work”, and we deeply apologize for the horrible implications. In [6], the authors identified entropy neurons by searching for neurons with a high weight norm and minimal impact on the logits. The underlying assumption is that these neurons act as a near-constant addition to all logits before the softmax, resulting in a minimal effect on output probabilities while increasing their entropy.

For model safety, however, we lack such a clear mechanistic intuition about how safety-related neurons should work. Therefore, our only assumption is that safety neurons should work in different ways between safety-aligned and unaligned models, and we design the inference-time activation contrasting to identify them.

Thank you for your constructive and valuable comments. Here is our response to the weaknesses and questions.

W1 Lack of comprehensive discussion

Please refer to our general response.

W2 & Q1 Why focus on neurons?

1. From the perspective of research objectives: Our goal is to develop a mechanistic understanding of LLM safety. The residual stream studied in RepE represents the combined effects of attention heads and MLPs. Understanding how these effects are formed requires digging deeper into the MI perspective. Since MLP neurons account for approximately two-thirds of the model’s parameters and serve as the basic functional units, we chose neurons as the focus of our research.
2. From a technical perspective: Compared to attention heads, identifying specific neurons presents a greater technical challenge due to their vastly larger quantity. For instance, in LLaMA2-7B, the number of neurons is about 340 times that of attention heads, and neurons often function in combinations, making it difficult to pinpoint them using DFA methods like those in [1]. These challenges motivated us to prioritize studying neuron findings. Of course, we do not claim that neurons alone provide a complete understanding of the safety mechanism. Given the complexity of safety, it likely requires the joint participation of neurons and attention heads. Additionally, the methods proposed in this paper can also be applied to identify attention heads, facilitating further exploration of their interactions, which we consider an interesting direction for future work.
3. Regarding related work on attention heads: [2] was published after our submission, while [1] identifies attention heads correlated with refusal directions from the RepE perspective but does not verify their causal effects on safety. This focus differs from the scope of our study.

W3 & Q2 Evidence provided in favor of the alignment tax hypothesis is limited

Our experiments (e.g. the curves in Figure 4(b)) rely on testing the causal effects of different neurons (e.g., safety, helpfulness, reasoning, etc.) on model safety and helpfulness, and do not involve transferring other abilities. The effectiveness of dynamic activation patching in transferring safety and helpfulness has been validated by previous experiments such as Figure 2 and Table 2. Therefore, we cannot see another alternative interpretation of our experimental results. We are willing to address your concerns but are unsure about the reasoning behind your interpretation or the specific experiments you propose. We would really appreciate it if you could elaborate more, like describing the suggested experiments in detail, and we are more than happy to verify it.

W4 & Q3 The classification results in Section 6 are very misleading

Thank you for your valuable suggestions. We agree that the margin of the original results is not clear enough. To more comprehensively verify whether safety neurons encode more safety-related information compared to random neurons, we conducted additional experiments:

1. For the datasets used in the experiments, we selected one dataset as the training set and merged the others as a single test set at a time, averaging the results across all rotations.
2. We excluded safety neurons from the randomly sampled set.
3. We added a group of random neurons sampled from all layers, as following the layer distribution of safety neurons may inherently carry safety-related information.

The updated results have been added to the appendix, and we revised the corresponding descriptions in the main paper to be more precise. Below is a brief summary of the results and our explanations:

From the results, we observe that safety neurons are indeed more effective than random neurons in predictions. Additionally, random neurons with the same layer distribution as safety neurons are more effective than those sampled from other layers, which indicates the layer distribution of safety neurons may also encode safety information. This may partially explain the results in Appendix E. We sincerely apologize for our oversight and thank you for pointing this out.

Lastly, we would like to note that the differences in prediction performance are not very significant, which may be due to the following reasons:

1. Safety neurons may not directly encode information about whether harmful content will be generated but instead exert their effects through subsequent components.
2. Random neurons may still receive information from safety neurons.

We plan to further investigate these aspects in future work.

Regarding W1 & W2

I acknowledge the engineering challenges involved in interpreting neurons compared to more coarse-grained units like layers and attention heads. The authors correctly note that previous works have examined safety behavior at both higher granularity (parameters) and lower granularity (attention heads/layers). However, my concern about novelty remains: the central question is not whether one can study safety behaviors at different granularities, but whether doing so provides meaningful new insights.

Even if the proposed probing method were significantly better on safety neurons versus random neurons, I'd still be unsure what the actionable insight would be. Researchers building latent space defenses (in the RepE camp) aren't probing at random neurons right now, so it feels like an artificial comparison. A more convincing demonstration would have been to show that neuron-level probing outperforms probing at the residual stream (representation engineering), which would have provided stronger validation for the need for additional mechanistic interpretability approaches.

I question whether papers that merely localize the same behaviors at different levels of granularity, without demonstrating novel insights or superior practical utility, represent sufficiently substantial contributions to warrant publication.

Regarding W3

I think there are two related but distinct claims here

1. There is some sort of fundamental conflict between the neurons the model tends to use for helpful behavior, versus the neurons that the model tends to use for safety behavior.
2. Finetuning a base model on just helpfulness allows it to repurpose some of its safety neurons. Similarly, fine-tuning a base model on just safety allows it to repurpose some of its helpful neurons.

The results in the paper suggest 2, but do not prove 1. I think proving 1 will be a challenge, that will require a more careful methodology. For example, fine-tuning a model on both helpful and harmless behavior, and then patching its safety mechanism over to a model just trained on helpfulness, and measuring how safe it is. I think there is a lot of subtlety that could be discussed here.

Regarding W4

I think the authors should move Figure 10b, into Section 6. Without any additional context, it is hard to understand how good an accuracy of 76.2% is. I would have appreciated additional baselines here, such as probing directly on the residual stream. I would also appreciate error bars for the random neuron lines, considering that the margins are so small.

While I appreciate the authors' effort to strengthen the paper, particularly with the addition of a more comprehensive related works section, my core concerns remain. Therefore, I maintain my score of 5.

We agree that patching over the residual activations is possible, but this paper is not about discussing which one of the neuron and representation interpretations is better. This paper has always been about interpreting model safety from the very beginning, and we are just saying that in this topic, our neuron-level interpretation demonstrates an advantage in understanding the mechanism (of alignment tax) over the representation-engineering related works. To avoid misunderstanding again, we are not saying that representation engineering has no way to gain a better mechanistic interpretation of model safety, but this hasn't been achieved (or demonstrated by a clear case like the alignment tax interpretation in our paper), and thus we believe our findings have unique values for now. We understand that one can have the favor of a technical direction over another, but we do not think the possibility of another direction should influence the value of actual findings.

Also, thanks for acknowledging our new experiments. Although in different positions, we respect your opinion.