Exploring Unsupervised Methods for Recovering the Refusal Direction in LLMs

Introduction

This is the first part of what we plan to be a research exploration into the geometries of LLMs. We start with building off existing work by Arditi et al focused on the refusal direction in LLMs.

Background

We recommend reading the Arditi et al. Lesswrong post and paper for more details. They explore the idea of the behavior categorized as Refusal being mediated by a single direction. Refusal is basically when an LLM refuses to do something, so examples like

"I cannot fulfill..."

"I cannot complete..."

etc.

A direction is defined by a vector. Here, it means a vector in an LLM's activation space.

An LLM's activation space is the space where an LLM's internal vectors live. Think of the intermediary vectors that are used in the calculation of the final next-token-prediction process.

Mediated means that if this direction were to be removed, then the behavior categorized as Refusal would stop. Arditi et al. performs directional ablation as their method of removal.

They identify the refusal direction using difference-in-means.

Methodology

We apply Principal Component Analysis and Isomap alongside diff-in-means to an experimental setup of refusal. We test the Llama 8b 3.1 Instruct and Qwen2.5-7B-Instruct models.

For the fit harmless dataset, we utilize the same Alpaca dataset used in Arditi et al.

For the fit harmful dataset, we utilize the Beavertails dataset.

We also amalgamate a held out harmful dataset from many of the same sources Arditi et al. procured its harmful dataset from, including AdvBench, JailbreakBench, HarmBench, and MaliciousInstruct. We also use CatQA, StrongREJECT, SorryBench, TDC2023, and HEx-Phi.

*by deduped, we make sure the prompts are not too similar within Beavertails itself as well as with the held out harmful set

Explaining PCA, Isomap, and Diff-in-means

PCA (Principal Component Analysis) reduces high-dimensional data into lower dimensions whilst preserving as much variance of the data as possible. PCA does this by finding the eigenvectors of the data's covariance matrix (though in practice a technique called SVD is most often used). Applying Multidimensional Scaling (MDS) to the Euclidean distances between the data points produces the PCA scores.

Isomap (Isometric Mapping) performs MDS with Geodesic distances. A Geodesic distance is the shortest path between two points on a Manifold. A Manifold is a shape that for each point appears to be flat within its neighborhood. Isomap approximates this Manifold using the kNN (or ε-ball) algorithm.

Diff-in-means takes the mean of vectors within classes and subtracts one from another, producing another vector. In Arditi et al the classes are harmful and harmless prompts, and the resulting vector is labeled the refusal direction. d = mean harmful - mean harmless

Results

Cosine Similarity

How similar are the directions (refusal direction computed from diff-in-means) and ?

Llama

Qwen

Ablation Efficacy

How effective is performing causal ablation in suppressing refusal behavior in harmful prompts?

Llama

Qwen

KL Divergence

How similar are the output log probabilities between the ablated model and the non-ablated model on harmless prompts? If the KL Divergence is too high, the model stops performing like it should. Arditi et al. sets a threshold of < 0.1 for KL Divergence, but that figure is arbitrary. Qwen has no layers that meet the threshold of <0.1 for KL and >0.80 for bypass, so we chose L20 as it combined the highest causal ablation efficacy (0.99 for , 0.97 for ) with a reasonable KL divergence (0.133 for , 0.197 for ).

Llama

Qwen

Validation

For Validation, we project the held out Harmless and held out Harmful activations onto and at the layers of L14 for Llama and L20 for Qwen.

Projecting means taking the dot product between the direction vector ( or ) and the activation vector . We sweep a threshold where a prompt is classified as harmful if its projection ( or ) exceeds , and harmless otherwise. For each we measure the TPR and FPR, and plot them as an ROC curve.

Llama

Qwen

Isomap

How the Isomap algorithm works, in more detail

1) For each of the points, perform kNN, and build a sparse distance matrix with neighbor distances populated at each point . If you're curious, Compressed Sparse Row is the format used to store the sparse matrix.

2) Using M, perform Dijkstra's Shortest Path Algorithm from each of the points to the other points. This will take for each point, thus being in total, setting to be constant. Construct a matrix of geodesic distances between all points.

3) Apply (classical) MDS. (Much of this is paraphrased from Wikipedia) Apply double centering with , where is the identity matrix, is a Matrix of all 1s, and is our geodesic distance matrix. Then we determine the largest eigenvalues and corresponding eigenvectors of . .

This is concretely what it means when Isomap "approximates" a Manifold. It is assuming the shape is flat locally but not necessarily globally; kNN utilizes euclidean distances, but the shortest path from one point to another follows the geodesic distance.

Note that is of size , whereas for PCA the matrix was of size . If Isomap were performed on euclidean distances, PCA's covariance Matrix and Isomap's would share the same nonzero eigenvalues where is the activation matrix. We could also recover from and vice versa using and where is the largest eigenvalue of and .

In the figures below we see the Isomap being plotted against each prompt's projection onto . Let's take a moment to think about what actually means. Because is an matrix of the geodesic distances, and the eigenvectors can be thought of to represent the vectors that best capture the essence of in dimensions, we can think of (where ) as the one dimensional embedding of the n points that best preserves their pairwise geodesic distances.

Our plots show a positive linear correlation between and projection onto , and the concentration of harmful and harmless points are clustered in their own regions. This indicates that the dominant axis of the points' geodesic distances is closely aligned with the refusal direction.

The Pearson Correlation Coefficient is given by the equation , where quantifies how much two variables and change together, and and are the standard deviations of and .

Llama

Qwen

Sweeping n

How does varying the size n affect how well and Isomap recover the refusal direction? We show that even a modest is sufficient for and Isomap to recover the refusal direction . is also computed with samples. We also show that Isomap consistently performs worse than PCA, indicating that Isomap provides no evidence that the structure is nonlinear.

Llama

Qwen

Conclusion

In this project we show that the refusal direction as computed by diff-in-means can also be recovered by , and is shown to be both directionally and causally similar. We show that we can recover via unsupervised PCA and utilize both and as a classifier for harmful and harmless prompts. We show that the sample size required is modest and doesn't require a large dataset of harmful/harmless prompts.

We also show that Isomap provides no indication that the activation geometry of refusal is supported by a curved Manifold. We argue that the strong linear correlation between Isomap and indicates that Isomap is trying to recover the underlying structure of refusal being a linear direction.

Limitations/Future Work

We were surprised at how well PCA was able to recover the refusal direction. We thought maybe it was due to it being the nature of having balanced class sizes between two distinct classes, and were a bit worried that this result would be trivial. We have since performed PCA on imbalanced classes, and it could reliably produce a with high cosine similarity with the refusal direction as well as high classification AUROC when harmful prompts were between 0.20 and 0.90 of the total sample. Isomap however performed well throughout all class distributions. These results can be seen in our code repo.

Future questions are then: for what types of binary LLM behaviors could the method of applying PCA to reliably classify them? Are there any causal implications of recovering such a direction as there is with refusal? Are there any unsupervised methods (besides SAEs in the Weekdays case) that can reliably recover the activation space geometries of curved concepts such as weekdays and months, as shown in Not All Language Model Features Are One-Dimensionally Linear by Engels et al.?

Code

https://github.com/ostenloo/unsupervised_refusal_direction