TL;DR: We use SGD to find sparse connections between features; additionally a large fraction of features between the residual stream & MLP can be modeled as linearly computed despite the non-linearity in the MLP. See linear feature section for examples.

Special thanks to fellow AISST member, Adam Kaufman, who originally thought of the idea of learning sparse connections between features & to Jannik Brinkmann for training these SAE’s.

Sparse AutoEncoders (SAE)’s are able to turn the activations of an LLM into interpretable features. To define circuits, we would like to find how these features connect to each other. SAE’s allowed us to scalably find interpretable features using SGD, so why not use SGD to find the connections too?

A layer in Pythia-70M, with SAE's before & after the MLP. Note, this is different than the typical Transformer Architecture, Pythia models have parallel MLP & Attn modules as opposed to sequential (this was to copy SOTA models which parallelized for more GPU utilization AFAIK)

We have a set of features before the MLP, F1, and a set of features after the MLP, F2. These features were learned by training SAE’s on the activations at these layers.

Ideally, we learn a linear function such that F2 = W(F1), & W is sparse (ie L1 penalty on weights of W). So then we can look at a feature in F2, and say "Oh, it's just a sparse linear combination of features of F1 e.g. 0.8*(however feature) + 0.6*(but feature)", which would be quite interpretable!

However, we’re trying to replicate an MLP’s computation, which surely can’t be all linear!^[1] So, what’s the simplest computation from F1 to F2 that gets the lowest loss (ignoring L1 weight sparsity penalty for now)?

Training on only MSE between F1 & F2, we plot the MSE throughout training across 5 layers in Pythia-70m-deduped in 4 settings:

Linear: $y = W x$

Nonlinear: $y = R e l u (W x)$

MLP: $y = W_{2} R e L U (W_{1} x)$

Two Nonlinear: $R e L U (W_{2} R e L U (W_{1} x))$

For all layers, training loss clusters along (MLP & two nonlinear) and (linear & nonlinear). Since MLP & linear are the simplest of these two clusters, the rest of the analysis will only look at those two.

[I also looked at bias vs no-bias: adding a bias didn’t positively improve loss, so it was excluded]

Interestingly enough, the relative linear-MLP difference is huge in the last layer (and layer 2). The...

Posts

Wikitag Contributions

Comments