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}}$.

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?)

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.

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.