Reproduction of Refusal Direction in LLMs

TL;DR

I reproduced the main experiment from the paper “Refusal in Language Models Is Mediated by a Single Direction” (arXiv:2406.11717).

The paper claims that refusal behavior corresponds to a single direction in the residual stream of a transformer model. By adding this direction, the model can be induced to refuse harmless prompts, while removing it suppresses refusal on harmful prompts.

Using Gemma 2B, I reproduced the main experiment and observed similar trends: injecting the refusal direction increases refusal on harmless prompts, while ablating it reduces refusal on harmful prompts.

Implemented repository is available here.

Motivation of the paper

Large language models are trained to refuse harmful requests through post-training techniques such as safety fine-tuning. However, these safeguards are not always reliable, and models may still produce harmful responses in certain situations.

Meanwhile, recent progress in mechanistic interpretability suggests that some concepts can be represented as a single direction in the residual stream of a transformer model.

If this is the case, it raises an interesting possibility: rather than relying solely on post-training alignment, we might be able to induce or suppress refusal behaviors by directly manipulating the model’s internal representations.

Refusal in Language Models Is Mediated by a Single Direction explores this idea. In this post, I reproduce its experiment to see how the method works in practice.

Experimental Setup

The experiment is conducted on the Gemma 2B model. The dataset is directly imported from the codebase of the paper.

1. Extracting the residual stream representations from harmful and harmless prompts

Residual stream representations are extracted from harmful and harmless prompts. The prompts are constructed by inserting harmful or harmless instructions into the chosen model’s instruction prompt format. After tokenization, the prompts are passed to the model to obtain the residual stream activations. These representations are collected from all layers of the model at post-instruction token positions (Appendix C.3).

2. Computing candidate refusal directions from the difference between mean representations

Candidate refusal directions are computed from the difference between the mean representations. The representations at the same layer and position index are averaged over the batch dimension. The difference between the averaged representations is taken to compute the candidate refusal directions. Thus, the final candidate refusal directions have shape (num_layers, num_post_instruction_tokens, hidden_size).

3. Selecting the optimal direction using bypass/induce scores

The optimal direction is selected following the selection algorithm described in the paper (Appendix C). However, evaluating the scores via generation is computationally expensive. Therefore, the proxy metric proposed in the paper is instead used to reduce the cost (Appendix B).

Implementation

The implementation is available here. The implementation relies on the TransformerLense library for convenient activation caching and interventions. Due to limited computational resources, the experiments are conducted only on google/gemma-2b-it.

Results

Paper Results

Ablating the refusal direction (Figure 1 in the paper)

As shown in the paper, ablating the refusal direction reduces refusal behavior on harmful prompts.

Adding the refusal direction (Figure 3 in the paper)

Adding the refusal direction induces refusal behavior even for harmless prompts.

Reproduced Results

Our reproduction produces the following result:

The reproduced results show similar trends: adding the refusal direction increases refusal on harmless prompts, while removing the direction reduces refusal on harmful prompts.

Interpretation

What the result suggests

The reproduced results show that manipulating a single direction in the residual stream changes the model’s refusal behavior. This suggests that the model’s refusal behavior is strongly associated with a specific representation in the residual stream and intervening on this representation can alter the model’s decision to refuse or comply with a prompt.

Reflection

Reproducing the results of Refusal in Language Models Is Mediated by a Single Direction provided a useful perspective on how safety behaviors can emerge from relatively simple internal representations. In particular, the experiment suggests that refusal behavior can be manipulated through a single direction in the residual stream, which supports the idea that certain high-level behaviors in language models may correspond to relatively localized representations.

One interesting aspect of the implementation is that the intervention is applied across all layers and token positions when ablating the refusal direction. This design choice appears to ensure that the target direction is fully removed from the model’s computation. However, this raises an additional question: it remains unclear at which layer the refusal signal is most strongly represented. A more detailed layer-wise analysis could provide further insight into where refusal behavior is actually formed in the network.

Several potential extensions remain for future exploration. The rank-one weight editing method proposed in the paper was not reproduced in this experiment due to time constraints. Investigating this approach could provide further insight into how internal representations can be modified through weight-level interventions.