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.
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.?
We were surprised at how well PCA was able to recover the refusal direction.
Maybe I just need more explanation, but I'm not sure what the result here is? My understanding is that you should expect this for any strong linear feature in any dataset.
"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"
This, to my mind, says more about the within class variation between prompt classes than anything else. You should expect total inter-class variation to scale on the order of \sqrt(p) where p is the probability of a given class if I'm not getting my maths incorrect.
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.
Model
Layers
Dimensions
Llama 8b 3.1 Instruct
32
4096
Qwen2.5-7B-Instruct
28
3584
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.
Dataset
n
Role
tatsu-lab/alpaca
14355 (sampled, seed = 42)
Fit Harmless
PKU-Alignment/Beavertails
14355 (deduped*)
Fit Harmful
*by deduped, we make sure the prompts are not too similar within Beavertails itself as well as with the held out harmful set
Dataset
n
Role
JailbreakBench/JBB-Behaviors
100
Held out Harmful
walledai/MaliciousInstruct
100
Held out Harmful
walledai/HarmBench
200
Held out Harmful
walledai/AdvBench
520
Held out Harmful
walledai/CatHarmfulQA
500
Held out Harmful
walledai/StrongREJECT
313
Held out Harmful
sorry-bench/sorry-bench-202503
450
Held out Harmful
walledai/TDC23-RedTeaming
100
Held out Harmful
LLM-Tuning-Safety/HEx-Phi
300
Held out Harmful
tatsu-lab/alpaca
2000 (sampled, seed = 43)
Held out Harmless
walledai/XSTest
450 (250 Harmful/200 Harmless)
Held out Harmful + Harmless
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
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
Llama
Qwen
Validation
For Validation, we project the held out Harmless and held out Harmful activations onto
Projecting means taking the dot product between the direction vector (
Llama
Qwen
Isomap
How the Isomap algorithm works, in more detail
1) For each of the
2) Using M, perform Dijkstra's Shortest Path Algorithm from each of the
3) Apply (classical) MDS. (Much of this is paraphrased from Wikipedia) Apply double centering
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
In the figures below we see the Isomap
Our plots show a positive linear correlation between
The Pearson Correlation Coefficient
Llama
Qwen
Sweeping n
How does varying the size n affect how well
Llama
Qwen
Conclusion
In this project we show that the refusal direction
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
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
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