Produced as part of the ML Alignment Theory Scholars Program - Winter 2023-24 Cohort as part of Lee Sharkey's stream.

TL;DR

Sparse autoencoders are a method of resolving superposition by recovering linearly encoded “features” inside activations. Unfortunately, despite the great recent success of SAEs at extracting human interpretable features, they fail to perfectly reconstruct the activations. For instance, Cunningham et al. (2023) note that replacing the residual stream of layer 2 of Pythia-70m with the reconstructed output of an SAE increased the perplexity of the model on the Pile from 25 to 40. It is important for interpretability that the features we extract accurately represent what the model is doing.

In this post, I show how and why SAEs have a reconstruction gap due to ‘feature suppression’. Then, I look at a few ways to fix this while maintaining SAEs interpretability. By modifying and fine-tuning a pre-trained SAE, we achieve a $9 %$ decrease in mean square error and a $24 %$ reduction in the perplexity increase upon patching activations into the LLM.

Finally, I compare a theoretical example to the observed amounts of feature suppression in Pythia 70m, showing that features are suppressed based on both the strength of their activations and their frequency of activation.

Feature Suppression

The architecture of an SAE is:

$f (x) = ReLU (W_{e} x + b_{e})$

$y = W_{d} f (x) + b_{d}$

The loss function usually combines a MSE reconstruction loss with a sparsity term, like $L (x, f (x), y) = | | y - x | |^{2} / d + c | f (x) |$ , where $d$ is the dimension of $x$ . When training the SAE on this loss, the decoder’s weight matrix is fixed to have unit norm for each feature (column).

The reason for feature suppression is simple: The training loss has two terms, only one of which is reconstruction. Therefore, reconstruction isn’t perfect. In particular, the loss function pushes for smaller $f (x)$ values, leading to suppressed features and worse reconstruction.

An illustrative example of feature suppression

As an example, consider the trivial case where there is only one binary feature in one dimension. That is, $x = 1$ with probability $p$ and $x = 0$ otherwise. Then, ideally the optimal SAE would extract feature activations of $f (x) \in {0, 1}$ and have a decoder with $W_{d} = 1$ .

However, if we were to train an SAE optimizing the loss function $L (x, f (x), y) = |$ ...

Posts

Wikitag Contributions

Comments