Improving Dictionary Learning with Gated Sparse Autoencoders

I believe that equation (10) giving the analytical solution to the optimization problem defining the relative reconstruction bias is incorrect. I believe the correct expression should be $\gamma =E_{x\sim D}\left[\frac{\hat{x}\cdot x}{\parallel x{\parallel }_2^2}\right]$.

You could compute this by differentiating equation (9), setting it equal to 0 and solving for $\gamma$. But here's a more geometrical argument.

By definition, $\gamma x$ is the multiple of $x$ closest to $\hat{x}$. Equivalently, this closest such vector can be described as the projection ${proj}_x\left(\hat{x}\right)=\frac{\hat{x}\cdot x}{\parallel x{\parallel }_2^2}x$. Setting these equal, we get the claimed expression for $\gamma$.

As a sanity check, when our vectors are 1-dimensional, $x=1$, and $\hat{x}=\frac{1}{2}$, we my expression gives $\gamma =\frac{1}{2}$ (which is correct), but equation (10) in the paper gives $\frac{1}{\sqrt{3}}$.

↑ comment by Arthur Conmy (arthur-conmy) · 2024-04-25T21:29:38.307Z · LW(p) · GW(p)

Oh oops, thanks so much. We'll update the paper accordingly. Nit: it's actually 

$\frac{E_{x\sim D}[\dot{x}\cdot x]}{E_{x\sim D}\left[||x|{|}_2^2\right]}$

(it's just minimizing a quadratic)

ETA: the reason we have complicated equations is that we didn't compute $E_{x\sim D}[\dot{x}\cdot x]$ during training (this quantity is kinda weird). However, you can compute $\gamma$ from quantities that are usually tracked in SAE training. Specifically, $\gamma =\frac{1}{2}(1+\frac{E\left[||xˆ|{|}_2^2\right]-E[||x-xˆ|{|}_2^2]}{E\left[||x|{|}_2^2\right]})$ and all terms here are clearly helpful to track in SAE training.

Replies from: samuel-marks

↑ comment by Sam Marks (samuel-marks) · 2024-04-25T22:13:00.663Z · LW(p) · GW(p)

Oh, one other issue relating to this: in the paper it's claimed that if $\gamma$ is the argmin of $E[\parallel \hat{x}-{\gamma }^{\prime }x{\parallel }_2^2]$ then $1/\gamma$ is the argmin of $E[\parallel {\gamma }^{\prime }\hat{x}-x\parallel ]$. However, this is not actually true: the argmin of the latter expression is $\frac{E[x\cdot \hat{x}]}{E[\parallel \hat{x}{\parallel }_2^2]}\ne {\left(\frac{E[x\cdot \hat{x}]}{E[\parallel x{\parallel }_2]}\right)}^{-1}$. To get an intuition here, consider the case where $\hat{x}$ and $x$ are very nearly perpendicular, with the angle between them just slightly less than ${90}^{\circ }$. Then you should be able to convince yourself that the best factor to scale either $x$ or $\hat{x}$ by in order to minimize the distance to the other will be just slightly greater than 0. Thus the optimal scaling factors cannot be reciprocals of each other.

ETA: Thinking on this a bit more, this might actually reflect a general issue with the way we think about feature shrinkage; namely, that whenever there is a nonzero angle between two vectors of the same length, the best way to make either vector close to the other will be by shrinking it. I'll need to think about whether this makes me less convinced that the usual measures of feature shrinkage are capturing a real thing.

ETA2: In fact, now I'm a bit confused why your figure 6 shows no shrinkage. Based on what I wrote above in this comment, we should generally expect to see shrinkage (according to the definition given in equation (9)) whenever the autoencoder isn't perfect. I guess the answer must somehow be "equation (10) actually is a good measure of shrinkage, in fact a better measure of shrinkage than the 'corrected' version of equation (10)." That's pretty cool and surprising, because I don't really have a great intuition for what equation (10) is actually capturing.

Replies from: rohinmshah, SenR

↑ comment by Rohin Shah (rohinmshah) · 2024-04-25T22:54:11.997Z · LW(p) · GW(p)

Thinking on this a bit more, this might actually reflect a general issue with the way we think about feature shrinkage; namely, that whenever there is a nonzero angle between two vectors of the same length, the best way to make either vector close to the other will be by shrinking it.

This was actually the key motivation for building this metric in the first place, instead of just looking at the ratio $\frac{E\left[||\hat{x}|{|}^2\right]}{E\left[||x|{|}^2\right]}$. Looking at the $\gamma$ that would optimize the reconstruction loss ensures that we're capturing only bias from the L1 regularization, and not capturing the "inherent" need to shrink the vector given these nonzero angles. (In particular, if we computed $\frac{E\left[||\hat{x}|{|}^2\right]}{E\left[||x|{|}^2\right]}$ for Gated SAEs, I expect that would be below 1.)

I think the main thing we got wrong is that we accidentally treated $E[||\hat{x}-x|{|}^2]$ as though it were $E[||\hat{x}-\gamma x|{|}^2]$. To the extent that was the main mistake, I think it explains why our results still look how we expected them to -- usually $\gamma$ is going to be close to 1 (and should be almost exactly 1 if shrinkage is solved), so in practice the error introduced from this mistake is going to be extremely small.

We're going to take a closer look at this tomorrow, check everything more carefully, and post an update after doing that. I think it's probably worth waiting for that -- I expect we'll provide much more detailed derivations that make everything a lot clearer.

↑ comment by Senthooran Rajamanoharan (SenR) · 2024-04-25T22:57:22.598Z · LW(p) · GW(p)

Hey Sam, thanks - you're right. The definition of reconstruction bias is actually the argmin of

$$E[|\hat{x}/{\gamma }^{\prime }-x{|}^2]$$

which I'd (incorrectly) rearranged as the expression in the paper. As a result, the optimum is

$${\gamma }^{-1}=E[\hat{x}\cdot x]/E\left[|\hat{x}{|}^2\right]$$

That being said, the derivation we gave was not quite right, as I'd incorrectly substituted the optimised loss rather than the original reconstruction loss, which makes equation (10) incorrect. However the difference between the two is small exactly when gamma is close to one (and indeed vanishes when there is no shrinkage), which is probably why we didn't pick this up. Anyway, we plan to correct these two equations and update the graphs, and will submit a revised version.

Replies from: SenR

↑ comment by Senthooran Rajamanoharan (SenR) · 2024-04-26T12:38:44.850Z · LW(p) · GW(p)

UPDATE: we've corrected equations 9 and 10 in the paper (screenshot of the draft below) and also added a footnote that hopefully helps clarify the derivation. I've also attached a revised figure 6, showing that this doesn't change the overall story (for the mathematical reasons I mentioned in my previous comment). These will go up on arXiv, along with some other minor changes (like remembering to mention SAEs' widths), likely some point next week. Thanks again Sam for pointing this out!

Updated equations (draft):

Updated figure 6 (shrinkage comparison for GELU-1L):

I'm a bit perplexed by the choice of loss function for training GSAEs (given by equation (8) in the paper). The intuitive (to me) thing to do here would be would be to have the $L_{\text{reconstruct}}$ and $L_{\text{sparsity}}$ terms, but not the $L_{\text{aux}}$ term, since the point of ${\pi }_{\text{gate}}$ is to tell you which features should be active, not to itself provide good feature coefficients for reconstructing $x$. I can sort of see how not including this term might result in the coordinates of ${\pi }_{\text{gate}}$ all being extremely small (but barely positive when it's appropriate to use a feature), such that the sparsity term doesn't contribute much to the loss. Is that what goes wrong? Are there ablation experiments you can report for this? If so, including this $L_{\text{aux}}$ term still currently seems to me like a pretty unprincipled way to deal with this -- can the authors provide any flavor here?

Here are two ways that I've come up with for thinking about this loss function -- let me know if either of these are on the right track. Let $f_{\text{gate},\text{ReLU}}$ denote the gated encoder, but with a ReLU activation instead of Heaviside. Note then that $f_{\text{gate},\text{ReLU}}$ is just the standard SAE encoder from Towards Monosemanticity.

Perspective 1: The usual loss from Towards Monosemanticity for training SAEs is $\parallel x-\hat{x}\left(f_{\text{gate},\text{ReLU}}\right(x\left)\right){\parallel }_2^2+\lambda \parallel f_{\text{gate},\text{ReLU}}\left(x\right){\parallel }_1$ (this is the same as your $L_{\text{sparsity}}$ and $L_{\text{aux}}$ up to the detaching thing). But now you have this magnitude network which needs to get a gradient signal. Let's do that by adding an additional term $\parallel x-\hat{x}\left(\tilde{f}\right(x\left)\right){\parallel }_2^2$ -- your $L_{\text{reconstruction}}$. So under this perspective, it's the reconstruction term which is new, with the sparsity and auxiliary terms being carried over from the usual way of doing things.

Perspective 2 (h/t Jannik Brinkmann): let's just add together the usual Towards Monosemanticity loss function for both the usual architecture and the new modified archiecture: $L=L_{\text{reconstruction}}\left(\tilde{f}\right)+L_{\text{reconstruction}}\left(\tilde{f}\right)+L_{\text{sparsity}}\left(f_{\text{gate},\text{ReLU}}\right)+L_{\text{sparsity}}\left(f_{\text{gate},\text{ReLU}}\right)$.

However, the gradients with respect to the second term in this sum vanish because of the use of the Heaviside, so the gradient with respect to this loss is the same as the gradient with respect to the loss you actually used.

(The question in this comment is more narrow and probably not interesting to most people.)

The limitations section includes this paragraph:

One worry about increasing the expressivity of sparse autoencoders is that they will overfit when
reconstructing activations (Olah et al., 2023, Dictionary Learning Worries), since the underlying
model only uses simple MLPs and attention heads, and in particular lacks discontinuities such as step
functions. Overall we do not see evidence for this. Our evaluations use held-out test data and we
check for interpretability manually. But these evaluations are not totally comprehensive: for example,
they do not test that the dictionaries learned contain causally meaningful intermediate variables in the
model’s computation. The discontinuity in particular introduces issues with methods like integrated
gradients (Sundararajan et al., 2017) that discretely approximate a path integral, as applied to SAEs
by Marks et al. (2024).

I'm not sure I understand the point about integrated gradients here. I understand this sentence as meaning: since model outputs are a discontinuous function of feature activations, integrated gradients will do a bad job of estimating the effect of patching feature activations to counterfactual values.

If that interpretation is correct, then I guess I'm confused because I think IG actually handles this sort of thing pretty gracefully. As long as the number of intermediate points you're using is large enough that you're sampling points pretty close to the discontinuity on both sides, then your error won't be too large. This is in contrast to attribution patching which will have a pretty rough time here (but not really that much worse than with the normal ReLU encoders, I guess). (And maybe you also meant for this point to apply to attribution patching?)

↑ comment by Neel Nanda (neel-nanda-1) · 2024-04-26T00:14:38.219Z · LW(p) · GW(p)

I haven't fully worked through the maths, but I think both IG and attribution patching break down here? The fundamental problem is that the discontinuity is invisible to IG because it only takes derivatives. Eg the ReLU and Jump ReLU below look identical from the perspective of IG, but not from the perspective of activation patching, I think.

Replies from: samuel-marks

↑ comment by Sam Marks (samuel-marks) · 2024-04-26T00:36:58.112Z · LW(p) · GW(p)

Yep, you're totally right -- thanks!

Great work! Obviously the results here speak for themselves, but I especially wanted to complement the authors on the writing. I thought this paper was a pleasure to read, and easily a top 5% exemplar of clear technical writing. Thanks for putting in the effort on that.

<3 Thanks so much, that's extremely kind. Credit entirely goes to Sen and Arthur, which is even more impressive given that they somehow took this from a blog post to a paper in a two week sprint! (including re-running all the experiments!!)

Thanks for the feedback, we will put up an update to the paper with all these numbers in tables, tomorrow night. For now I have sent you them (and can send anyone else them who wants them in the next 24H)

We use learning rate 0.0003 for all Gated SAE experiments, and also the GELU-1L baseline experiment. We swept for optimal baseline learning rates on GELU-1L for the baseline SAE to generate this value. 

For the Pythia-2.8B and Gemma-7B baseline SAE experiments, we divided the L2 loss by $E||x|{|}_2$, motivated by wanting better hyperparameter transfer, and so changed learning rate to 0.001 or 0.00075 for all the runs (currently in Figure 1, only attention output pre-linear uses 0.00075. In the rerelease we'll state all the values used). We didn't see noticable difference in the Pareto frontier changing between 0.001 and 0.00075 so did not sweep the baseline hyperparameter further than this.

Re dictionary width, 2**17 (~131K) for most Gated SAEs, 3*(2**16) for baseline SAEs, except for the (Pythia-2.8B, Residual Stream) sites we used 2**15 for Gated and 3*(2**14) for baseline since early runs of these had lots of feature death. (This'll be added to the paper soon, sorry!). I'll leave the other Qs for my co-authors

On $b_{\text{mag}}$, it's unclear what a "natural" choice would be for setting this parameter in order to simplify the architecture further. One natural reference point is to set it to $e^{r_{\text{mag}}}\odot b_{\text{gate}}$, but this corresponds to getting rid of the discontinuity in the Jump ReLU (turning the magnitude encoder into a ReLU on multiplicatively rescaled gate encoder preactivations). Effectively (removing the now unnecessary auxiliary task), this would give results similar to the "baseline + rescale & shift" benchmark in section 5.2 of the paper, although probably worse, as we wouldn't have the shift.

Nah I think it's pretty sketchy. I personally prefer mean ablation, especially for residual stream SAEs where zero ablation is super damaging. But even there I agree. Compute efficiency hit would be nice, though it's a pain to get the scaling laws precise enough

For our paper this is irrelevant though IMO because we're comparing gated and normal SAEs, and I think this is just scaling by a constant? It's at least monotonic in CE loss degradation

I don't think zero ablation is that great a baseline. We're mostly using it for continuity's sake with Anthropic's prior work (and also it's a bit easier to explain than a mean ablation baseline which requires specifying where the mean is calculated from). In the updated paper https://arxiv.org/pdf/2404.16014v2 (up in a few hours) we show all the CE loss numbers for anyone to scale how they wish.

I don't think compute efficiency hit^[1] is ideal. It's really expensive to compute, since you can't just calculate it from an SAE alone as you need to know facts about smaller LLMs. It also doesn't transfer as well between sites (splicing in an attention layer SAE doesn't impact loss much, splicing in an MLP SAE impacts loss more, and residual stream SAEs impact loss the most). Overall I expect it's a useful expensive alternative to loss recovered, not a replacement.

EDIT: on consideration of Leo's reply, I think my point about transfer is wrong; a metric like "compute efficiency recovered" could always be created by rescaling the compute efficiency number.

1. ^{^}

   What I understand "compute efficiency hit" to mean is: for a given (SAE, $L{M}_1$) pair, how much less compute you'd need (as a multiplier) to train a different LM, $L{M}_2$ such that $L{M}_2$ gets the same loss as $L{M}_1$-with-the-SAE-spliced-in.

Good question - we're planning to post an update on this point about combining the new sparsity penalty from Anthropic with Gated SAEs. The TL;DR is that you can replace the L1 term in the Gated SAE loss with the analogous (gated feature magnitudes dotted with decoder magnitudes) sparsity term introduced by Anthropic and thereby do away with the decoder norms constraint and resampling. If you're going to do this, you also need to either unfreeze the decoder in the auxiliary task, or freeze the decoder weights where they appear in the sparsity penalty; both attain reasonably similar performance, and are definitely better than having the decoder weights frozen in one place but not the other. Put together, this seems to a marginal hit (versus the original Gated loss with L1 penalty and resampling) when comparing Pareto curves, but may be worth to the extent it simplifies training (with this loss function, the SAE training loop just becomes a vanilla neural network training loop).

PS With either the original (L1-based) loss or the modified loss of the previous paragraph, some of the other improvements suggested in the Anthropic post -- in particular, initializing the encoder weights to the transpose of the decoder weights (only at initialisation, not tying them thereafter), and warming up lambda. My point about the new loss not being Pareto better than L1 applies only if you compare like with like -- i.e. apply these other improvements in both cases.

We found that exactly that form of sparsity penalty did improve shrinkage with standard (ungated) SAEs, and provide a decent boost to loss recovered at low L0. (We didn't evaluate interpretability though.) But then we hit upon Gated SAEs which looked even better, and for which modifying the sparsity penalty in this way feels less necessary, so we haven't experimented with combining the two.

This suggestion seems less expressive than (but similar in spirit to) the "rescale & shift" baseline we compare to in Figure 9. The rescale & shift baseline is sufficient to resolve shrinkage, but it doesn't capture all the benefits of Gated SAEs.

The core point is that L1 regularization adds lots of biases, of which shrinkage is just one example, so you want to localize the effect of L1 as much as possible. In our setup L1 applies to $\text{ReLU}\left({\pi }_{\text{gate}}\right(x\left)\right)$, so you might think of ${\pi }_{\text{gate}}$ as "tainted", and want to use it as little as possible. The only thing you really need L1 for is to deter the model from setting too many features active, i.e. you need it to apply to one bit per feature (whether that feature is on / off). The Heaviside step function makes sure we are extracting just that one bit, and relying on ${\text{f}}_{\text{mag}}$ for everything else.