Mechanistic Interpretability

This is the master note for Mechanistic Interpretability

---

what

• reverse engineering the computational mechanism and representation learnt by neural networks to human-understandable algorithms and concepts¹
• a paradigm shift in interpretability: surface-level analysis (input/output relations) → inner interpretability (internal mechanisms) ¹
  • similar to the shift from behaviourism to cognitive neuroscience in psychology
  • Mech Interp is an approach toward inner interpretability

    

    Link to original

---

some terminologies

• Features:
  • independent yet repeatable units that a neural network representation can decompose into ²
  • fundamental units of neural network representations ¹
  • are different from neurons
• Neurons ¹
  • are computational units
  • potentially representing individual features
  • forming privileged bases
• Circuit
  • sub-graphs of the network, consisting of features and the weights connecting them
• monosemantic and polysemantic neurons
  • monosemantic: neurons corresponding to a single semantic concept
  • polysemantic: neurons associated with multiple, unrelated concepts
  • if neurons were the fundamental primitives of model representations
    • all neurons would be monosemantic
    • implying a 1-to-1 relation between neurons and features
    • however, empirical studies observed that neurons are polysemantic

---

• Privileged bases
  • the standard bases of the representation space

    

    Link to original

---

hypotheses

---

Linear representation hypothesis

• Features / concepts are represented by orthogonal directions in activation space
  • but need not be aligned with the privileged bases
• The network represents features as linear combinations of neurons

  

  Link to original

---

Superposition Hypothesis

• The representation may encode features not with the $n$ basis directions (neurons) but with $\propto exp(n)$ possible almost orthogonal directions.

---

• Neural networks represent more features than they have neurons by encoding features in ovelapping combination of neurons ¹

  

  Link to original

---

• Non-orthogonality means that features interfere with one another.
• Sparsity means that a feature rarely occurs. The assumption is most features are sparse.

---

Toy model of superposition

---

Link to original

---

A Hierarchy of Feature Properties³

4 progressively more strict properties that neural network representations might have:

• Decomposability: Neural network activations which are decomposable can be decomposed into features, the meaning of which is not dependent on the value of other features.
• Linearity: Features correspond to directions. Each feature $f_i$​ has a corresponding representation direction $W_i$. The presence of multiple features $f_1,f_2\dots$ activating with values $x_{f_1},x_{f_2}\dots$ is represented by $x_{f_1}{W}_{f_1}+x{f}_{2}W{f}_2\cdots$
• Superposition vs Non-Superposition: A linear representation exhibits superposition if $W{W}^T$ is not invertible. If $W{W}^T$ is invertible, it does not exhibit superposition.
• Basis-Aligned: A representation is basis aligned if all $W_i$​ are one-hot basis vectors. A representation is partially basis aligned if all $W_i$ are sparse. This requires a privileged basis.

The first two are hypothesized to be widespread, while the latter are believed to be occured only sometimes.

---

Universality hypothesis

• Analogous features and circuits form across model and tasks.

---

A mechanistic view on LLM

---

Virtual Weights and the Residual Stream as a Communication Channel ⁴

• View the residual stream as the main object that accumulate information

---

• MLP and Attention are branches that write to the stream

---

Observations about Attention ⁴

• Applying attention can be described as

$$\begin{array}{rl}h(x)& =\underset{\begin{array}{l}\text{Project result vectors}\\ \text{out for each token}\\ (h(x{)}_i=W_O{r}_i)\end{array}}{\underbrace{(I\otimes W_O)}}\cdot \underset{\begin{array}{l}\text{Mix value vectors across}\\ \text{tokens to compute}\\ \text{result vectors}\\ (r_i={\sum }_j{A}_{i,j}{v}_j)\end{array}}{\underbrace{(A\otimes I)}}\cdot \underset{\begin{array}{l}\text{Compute value vector}\\ \text{for each token}\\ (v_i=W_V{x}_i)\end{array}}{\underbrace{(I\otimes W_V)}}\\ & =\underset{\begin{array}{l}\text{A mixes across tokens while}\\ W_O{W}_V\text{ acts on each vector}\\ \text{independently}\end{array}}{\underbrace{(A\otimes W_O{W}_V)\cdot }}x\end{array}$$

• And the attention pattern is

$$\begin{array}{rl}{k}_i& =W_K{x}_i\\ q_i& =W_Q{x}_i\\ & \\ A& =softmax(q^{T}k)\\ & =softmax(x^T{W}_Q^T{W}_{K}x)\end{array}$$

• $W_Q$​ and $W_K$​ always operate together. They’re never independent. Similarly, $W_O$​ and $W_V$​ always operate together as well.

---

• An attention head is really applying two linear operations, $A$ and $W_O{W}_V$​, which operate on different dimensions and act independently.
  • $A$ governs which token’s information is moved from and to.
  • $W_O{W}_V$​ governs which information is read from the source token and how it is written to the destination token.
• Products of attention heads behave much like attention heads themselves. By the distributive property

$$(A^{h_2}\otimes W_{OV}^{h_2}\cdot A^{h_1}\otimes W_{OV}^{h_1}=(A^{h_2}{A}^{h_1})\otimes (W_{OV}^{h_2}{W}_{OV}^{h_1}))$$

• These are called virtual attention heads

---

MLP layers are KV memories ⁵

Link to original

---

• MLP layers are Unnormalized Key-Value Memories
  • MLP layers ⁶ : $FF(x)=f(x\cdot K^T)\cdot V$
  • Neural Memory ⁷ : $$\begin{array}{rl}p(k_i|x)& \propto exp(x\cdot k_i)\\ & \\ MN(x)& =\sum ^{d_m}i=1p(k_i|x)v_i\\ & \\ & =softmax(x\cdot K^T)\cdot V\end{array}$$
• MLP layers are almost identical to KV neural memories. The only different is
  • neural memory uses $softmax(\cdot )$ as the non-linearity $f(\cdot )$
  • while MLPs in transformer doesn’t use a normalizing function

---

• Intra-layer and inter-layer memory composition

  • “the layer-level prediction is typically not the result of a single dominant memory cell, but a composition of multiple memories.”

    Link to original

  • “the model uses the sequential composition apparatus as a means to refine its prediction from layer to layer, often deciding what the prediction will be at one of the lower layers.”

    Link to original

    

    Link to original

---

my observations, hypothesis and findings

• Theory of activation space
• The effects of neural net layers on activation space

literature

papers

• Refusal in Language Models Is Mediated by a Single Direction
• Toy Models of Superposition
• The Linear Representation Hypothesis and the Geometry of Large Language Models
• Representation Engineering: A Top-Down Approach to AI Transparency
• Steering Language Models With Activation Engineering
• A Language Model’s Guide Through Latent Space
• Linear Representations of Sentiment in Large Language Models
• Universal and Transferable Adversarial Attacks on Aligned Language Models
• Transformer Feed-Forward Layers Are Key-Value Memories
• Mechanistic Interpretability for AI Safety — A Review
• A Mathematical Framework for Transformer Circuits

articles

study

• Towards Monosemanticity: Decomposing Language Models With Dictionary Learning
• Transformer Circuits Thread
• Thread: Circuits

references

Footnotes

1. Mechanistic Interpretability for AI Safety — A Review ↩ ↩² ↩³ ↩⁴ ↩⁵

2. Distributed Representations: Composition & Superposition ↩

3. Toy Models of Superposition ↩

4. A Mathematical Framework for Transformer Circuits ↩ ↩²

5. Transformer Feed-Forward Layers Are Key-Value Memories ↩

6. 

   Link to original

   ↩

7. 

   Link to original

   ↩