mechanistic interpretability

whirlwind tour, initial exploration, glossary

The subfield of alignment that delves into reverse engineering of a neural network, especially LLMs

To attack the curse of dimensionality, the question remains: how do we hope to understand a function over such a large space, without an exponential amount of time? ¹

steering

refers to the process of manually modifying certain activations and hidden state of the neural net to influence its outputs

For example, the following is a toy example of how a decoder-only transformers (i.e: GPT-2) generate text given the prompt “The weather in California is”

flowchart LR
  A[The weather in California is] --> B[H0] --> D[H1] --> E[H2] --> C[... hot]

To steer to model, we modify $H_2$ layers with certain features amplifier with scale 20 (called it $H_3$)²

flowchart LR
  A[The weather in California is] --> B[H0] --> D[H1] --> E[H3] --> C[... cold]

One usually use techniques such as sparse autoencoders to decompose model activations into a set of interpretable features.

For feature ablation, we observe that manipulation of features activation can be strengthened or weakened to directly influence the model’s outputs

A few examples where (Panickssery et al., 2024) uses contrastive activation additions to steer Llama 2

contrastive activation additions

intuition: using a contrast pair for steering vector additions at certain activations layers

Uses mean difference which produce difference vector similar to PCA:

Given a dataset $\mathcal{D}$ of prompt $p$ with positive completion $c_p$ and negative completion $c_n$, we calculate mean-difference $v_{\text{MD}}$ at layer $L$ as follow:

$$v_{\text{MD}}=\frac{1}{\mid \mathcal{D}\mid }\sum _{p,c_p,c_n\in \mathcal{D}}{a}_L(p,c_p)-a_L(p,c_n)$$

implication

by steering existing learned representations of behaviors, CAA results in better out-of-distribution generalization than basic supervised finetuning of the entire model.

sparse autoencoders

abbrev: SAE

see also: landspace

Often contains one layers of MLP with few linear ReLU that is trained on a subset of datasets the main LLMs is trained on.

empirical example: if we wish to interpret all features related to the author Camus, we might want to train an SAEs based on all given text of Camus to interpret “similar” features from Llama-3.1

definition

We wish to decompose a models’ activitation $x\in {\mathbb{R}}^n$ into sparse, linear combination of feature directions:

$$\begin{gathered} x\sim x_0+\sum _{i=1}^M{f}_i(x)d_i \\ \because \begin{array}{rl}{d}_{i}M\gg n& :\text{ latent unit-norm feature direction}\\ f_i(x)\ge 0& :\text{ corresponding feature activation for }x\end{array} \end{gathered}$$

Thus, the baseline architecture of SAEs is a linear autoencoder with L1 penalty on the activations:

$$\begin{array}{rl}f(x)& :=\text{ReLU}(W_{\text{enc}}(x-b_{\text{dec}})+b_{\text{enc}})\\ \hat{x}(f)& :=W_{\text{dec}}f(x)+b_{\text{dec}}\end{array}$$

training it to reconstruct a large dataset of model activations $x\sim \mathcal{D}$, constraining hidden representation $f$ to be sparse

L1 norm with coefficient $\lambda$ to construct loss during training:

$$\begin{gathered} \mathcal{L}(x):=\Vert x-\hat{x}(f(x)){\Vert }_2^2+\lambda \Vert f(x){\Vert }_1 \\ \because \Vert x-\hat{x}(f(x)){\Vert }_2^2:\text{ reconstruction loss} \end{gathered}$$

intuition

We need to reconstruction fidelity at a given sparsity level, as measured by L0 via a mixture of reconstruction fidelity and L1 regularization.

We can reduce sparsity loss term without affecting reconstruction by scaling up norm of decoder weights, or constraining norms of columns $W_{\text{dec}}$ during training

Ideas: output of decoder $f(x)$ has two roles

• detects what features acre active ⇐ L1 is crucial to ensure sparsity in decomposition
• estimates magnitudes of active features ⇐ L1 is unwanted bias

Gated SAE

see also: paper

uses Pareto improvement over training to reduce L1 penalty (Rajamanoharan et al., 2024)

Clear consequence of the bias during training is shrinkage (Sharkey, 2024) ¹

Idea is to use gated ReLU encoder (Dauphin et al., 2017; Shazeer, 2020):

$$\tilde{f}(\mathbf{x}):=\underset{f_{\text{gate}}(\mathbf{x})}{\underbrace{1[\underset{{\pi }_{\text{gate}}(\mathbf{x})}{\underbrace{({\mathbf{W}}_{\text{gate}}(\mathbf{x}-{\mathbf{b}}_{\text{dec}})+{\mathbf{b}}_{\text{gate}})>0}}]}}\odot \underset{f_{\text{mag}}(\mathbf{x})}{\underbrace{\text{ReLU}({\mathbf{W}}_{\text{mag}}(\mathbf{x}-{\mathbf{b}}_{\text{dec}})+{\mathbf{b}}_{\text{mag}})}}$$

where $1[\bullet >0]$ is the (pointwise) Heaviside step function and $\odot$ denotes elementwise multiplication.

term	annotations
$f_{\text{gate}}$	which features are deemed to be active
$f_{\text{mag}}$	feature activation magnitudes (for features that have been deemed to be active)
${\pi }_{\text{gate}}(x)$	$f_{\text{gate}}$ sub-layer’s pre-activations

to negate the increases in parameters, use weight sharing:

Scale $W_{\text{mag}}$ in terms of $W_{\text{gate}}$ with a vector-valued rescaling parameter $r_{\text{mag}}\in {\mathbb{R}}^M$:

$$(W_{\text{mag}}{)}_{ij}:=(\exp (r_{\text{mag}}){)}_i\cdot (W_{\text{gate}}{)}_{ij}$$

Figure 3: Gated SAE with weight sharing between gating and magnitude paths

Figure 4: A gated encoder become a single layer linear encoder with Jump ReLU (Erichson et al., 2019) activation function ${\sigma }_{\theta }$

feature suppression

See also: link

Loss function of SAEs combines a MSE reconstruction loss with sparsity term:

$$\begin{gathered} L(x,f(x),y)=\Vert y-x{\Vert }^2/d+c\mid f(x)\mid \\ \because d:\text{ dimensionality of }x \end{gathered}$$

the reconstruction is not perfect, given that only one is reconstruction. For smaller value of $f(x)$, features will be suppressed

illustrated example

consider one binary feature in one dimension $x=1$ with probability $p$ and $x=0$ otherwise. Ideally, optimal SAE would extract feature activation of $f(x)\in \{0,1\}$ and have decoder $W_d=1$

However, if we train SAE optimizing loss function $L(x,f(x),y)$, let say encoder outputs feature activation $a$ if $x=1$ and 0 otherwise, ignore bias term, the optimization problem becomes:

$$\begin{gathered} \begin{array}{rl}a& =\mathrm{argmin}p\ast L(1,a,a)+(1-p)\ast L(0,0,0)\\ & =\mathrm{argmin}(1-a{)}^2+\mid a\mid \ast c\\ & =\mathrm{argmin}{a}^2+(c-2)\ast a+1\end{array} \\ ⟹\boxed{a=1-\frac{c}{2}} \end{gathered}$$

How do we fix feature suppression in training SAEs?

introduce element-wise scaling factor per feature in-between encoder and decoder, represented by vector $s$:

$$\begin{array}{rl}f(x)& =\text{ReLU}(W_{e}x+b_e)\\ f_s(x)& =s\odot f(x)\\ y& =W_d{f}_s(x)+b_d\end{array}$$

Dauphin, Y. N., Fan, A., Auli, M., & Grangier, D. (2017). Language Modeling with Gated Convolutional Networks. arXiv preprint arXiv:1612.08083 arxiv

Erichson, N. B., Yao, Z., & Mahoney, M. W. (2019). JumpReLU: A Retrofit Defense Strategy for Adversarial Attacks. arXiv preprint arXiv:1904.03750 arxiv

Rajamanoharan, S., Conmy, A., Smith, L., Lieberum, T., Varma, V., Kramár, J., Shah, R., & Nanda, N. (2024). Improving Dictionary Learning with Gated Sparse Autoencoders. arXiv preprint arXiv:2404.16014 arxiv

Sharkey, L. (2024). Addressing Feature Suppression in SAEs. AI Alignment Forum. [post]

Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv preprint arXiv:2002.05202 arxiv

Footnotes

1. If we hold $\hat{x}(\bullet )$ fixed, thus L1 pushes $f(x)\to 0$, while reconstruction loss pushes $f(x)$ high enough to produce accurate reconstruction.
   An optimal value is somewhere between.
   However, rescaling the shrink feature activations (Sharkey, 2024) is not necessarily enough to overcome bias induced by L1: a SAE might learnt sub-optimal encoder and decoder directions that is not improved by the fixed. ↩

Lien vers l'original

superposition hypothesis

tl/dr

phenomena when a neural network represents more than $n$ features in a $n$-dimensional space

Linear representation of neurons can represent more features than dimensions. As sparsity increases, model use superposition to represent more features than dimensions.

neural networks “want to represent more features than they have neurons”.

When features are sparsed, superposition allows compression beyond what linear model can do, at a cost of interference that requires non-linear filtering.

reasoning: “noisy simulation”, where small neural networks exploit feature sparsity and properties of high-dimensional spaces to approximately simulate much larger much sparser neural networks

In a sense, superposition is a form of lossy compression

importance

• sparsity: how frequently is it in the input?

• importance: how useful is it for lowering loss?

overcomplete basis

reasoning for the set of $n$ directions ³

features

A property of an input to the model

When we talk about features (Elhage et al., 2022, p. see “Empirical Phenomena”), the theory building around several observed empirical phenomena:

1. Word Embeddings: have direction which corresponding to semantic properties (Mikolov et al., 2013). For example:

   V(king) - V(man) = V(monarch)

2. Latent space: similar vector arithmetics and interpretable directions have also been found in generative adversarial network.

We can define features as properties of inputs which a sufficiently large neural network will reliably dedicate a neuron to represent (Elhage et al., 2022, p. see “Features as Direction”)

ablation

refers to the process of removing a subset of a model’s parameters to evaluate its predictions outcome.

idea: deletes one activation of the network to see how performance on a task changes.

• zero ablation or pruning: Deletion by setting activations to zero
• mean ablation: Deletion by setting activations to the mean of the dataset
• random ablation or resampling

residual stream

flowchart LR
  A[Token] --> B[Embeddings] --> C[x0]
  C[x0] --> E[H] --> D[x1]
  C[x0] --> D
  D --> F[MLP] --> G[x2]
  D --> G[x2]
  G --> I[...] --> J[unembed] --> X[logits]

residual stream $x_0$ has dimension $(C,E)$ where

• $C$: the number of tokens in context windows and
• $E$: embedding dimension.

Attention mechanism $H$ process given residual stream $x_0$ as the result is added back to $x_1$:

$$x_1=H(x_0)+x_0$$

grokking

See also: writeup, code, circuit threads

A phenomena discovered by (Power et al., 2022) where small algorithmic tasks like modular addition will initially memorise training data, but after a long time ti will suddenly learn to generalise to unseen data

empirical claims

related to phase change

Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., & Olah, C. (2022). Toy Models of Superposition. Transformer Circuits Thread. [link]

Mikolov, T., Yih, W., & Zweig, G. (2013). Linguistic Regularities in Continuous Space Word Representations. In L. Vanderwende, H. Daumé III, & K. Kirchhoff (Eds.), Proceedings of the 2013 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (pp. 746–751). Association for Computational Linguistics. https://aclanthology.org/N13-1090

Panickssery, N., Gabrieli, N., Schulz, J., Tong, M., Hubinger, E., & Turner, A. M. (2024). Steering Llama 2 via Contrastive Activation Addition. arXiv preprint arXiv:2312.06681 arxiv

Power, A., Burda, Y., Edwards, H., Babuschkin, I., & Misra, V. (2022). Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv preprint arXiv:2201.02177 arxiv

Footnotes

1. good read from Lawrence C for ambitious mech interp. ↩

2. An example steering function can be:

   $$H_3=H_2+\text{steering\_strength}\ast \text{SAE}.W_{\text{dec}}[20]\ast \text{max\_activation}$$ ↩

3. Even though features still correspond to directions, the set of interpretable direction is larger than the number of dimensions ↩