What is that neural network actually doing?
by Yours Truly
Why would one read the weights?
A trained network is a pile of floats that happens to work. Mechanistic interpretability is the attempt to read that as a program. That looks as instead of
“the model pays attention to token five”
saying
“these two heads implement a lookup. Here is the algorithm, and here is what breaks when I delete it” .
That is reminiscent of reverse engineering. The binary is written by gradient descent, there are no symbols, and the calling convention is linear algebra :D
Features are directions, and there are far too many of them
The residual stream of a transformer is a vector space \(\mathbb{R}^d\) that every layer reads from and adds to. The working hypothesis is that a concept is a direction within it, and that concepts compose by vector addition.
A serious model tracks far more than \(d\) concepts though, hence they get packed.
Johnson-Lindenstrauss claims that \(\mathbb{R}^d\) holds \(\exp(O(\varepsilon^2 d))\) unit vectors that are pairwise almost orthogonal, so \(|\langle u, v \rangle| \le \varepsilon\) - exponentially many. So if features are sparse (== most are off for most inputs), the network can store them in an overcomplete dictionary and eat the interference. That is superposition, and it shows that a single neuron will fire for academic citations, English dialogue, HTTP requests and Korean text all at once. Unpacking it is dictionary learning where you fit a wide, sparse autoencoder to the activations and hope the atoms come out one concept at a time.
The residual stream has no privileged basis speaking more of the linear algebra.
If you rotate it by any orthogonal \(Q\), and push \(Q^{\top}\) into the next weight matrix, then the network computes the identical function. So “coordinate 1337 of the residual stream” doesn’t mean anything. The MLP hidden layer is privileged due to \(\mathrm{ReLU}\) acting coordinate-wise, and does not commute with rotations. Symmetry decides where the coordinates mean anything.
The network that rediscovered the DFT
If you train a small transformer on \(a + b \bmod p\) and read the weights afterwards then the embedding rows are sines and cosines at a handful of frequencies. The thing implements a discrete Fourier transform out of the blue :) The phenomenon is called grokking.
With a bit of group theory jargon, here is how that goes. The irreducible representations of the cyclic group \(\mathbb{Z}/p\mathbb{Z}\) are exactly the characters \(\chi_k(a) = e^{2\pi i k a/p}\), and they are the basis in which group convolution becomes pointwise multiplication. “Add mod \(p\)” is convolution on that group. Here, the gradient descent found the change of basis.
Trace it by hand
Let \(p = 5\), \(\omega_k = 2\pi k/5\), and ask for \(3 + 4\).
-
Embed each input onto the unit circle, \(a \mapsto (\cos \omega_k a,\ \sin \omega_k a)\). At \(k = 1\) that is \(3 \mapsto (-0.809, -0.588)\) and \(4 \mapsto (0.309, -0.951)\).
-
Multiply the coordinates pairwise. That is what an attention head followed by a roughly-quadratic nonlinearity hands us. Collecting terms gives:
Numerically, \((-0.809)(0.309) - (-0.588)(-0.951) = -0.809\), which is \(\cos 7\omega_1\). The sum now lives in a phase, and a phase is mod \(p\) for free. Hence, no carry, comparisons, and branching.
- Unembed into one logit per candidate answer \(c\):
- Read it off. Characters are orthogonal so their sum \(\sum_{k=1}^{p-1} \cos(2\pi k m / p)\) is \(p - 1\) when \(m \equiv 0\) and \(-1\) otherwise. Two frequencies suffice here, since \(\cos\) is even and so \(k \in \{1, 2\}\) covers all four nonzero ones:
| c | 0 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|---|
| logit | -0.5 | -0.5 | 2.0 | -0.5 | -0.5 |
A discrete delta at \(c = 2 = 7 \bmod 5\). This is exact, shows what people found by staring at a trained network, then confirming by ablating individual frequencies and watching accuracy fall over.
Layers are time
A residual block is \(x_{l+1} = x_l + f_l(x_l)\), which is one forward Euler step of \(\dot{x} = f(x, t)\). One might imagine a deep network as a discretized flow and depth as time, so interpretability gets to ask what this flow transports, and where it “mixes”.
Training is a flow too. Gradient descent integrates \(\dot{\theta} = -\nabla L(\theta)\); put the minibatch noise back in and the distribution of parameters obeys Fokker-Planck:
\[\partial_t \rho = \nabla \cdot (\rho \nabla L) + \tfrac{\beta}{2} \Delta \rho\]with the Gibbs steady state \(\rho_\infty \propto e^{-2L/\beta}\). If the width is infinity, then density of neurons obeys a continuity equation (!), namely \(\partial_t \rho + \nabla \cdot (\rho v) = 0\). Looks like a transport.
Perhaps grokking is a phase transition in that flow, and maybe because of that it looks so abrupt from the loss curve.
It is interesting to think how a Green’s function fits into the picture.
Am I using that?
Not at work (yet?). The things I debug there still have source, and I can attach a debugger. At home, sure. The modular addition network fits on one screen, trains in a couple of minutes on a laptop, and you can print the embedding matrix and see the sine waves sitting in it.
Much the same feeling as the first time a hex dump turned into a file format.
What I have been reading
- Olah et al., Zoom In: An Introduction to Circuits
- Elhage et al., A Mathematical Framework for Transformer Circuits
- Elhage et al., Toy Models of Superposition
- Nanda et al., Progress Measures for Grokking via Mechanistic Interpretability
- Chughtai, Chan, Nanda, A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations
- Templeton et al., Scaling Monosemanticity