Low-Rank Adapting Models for Sparse Autoencoders

Thank you for your feedback and time! We are especially glad to hear you found our experiments and results to be thorough and convincing.

Can you explain 5.2 in more detail?

Certainly! We apologize for not being clearer and hope the below will explain in more detail how we’re computing our metric:

For a given dataset D, for both the base and LoRA model:

• For a given text sample $i \in D$, we compute the average per token log likelihood $a_i$.
• On the same text sample $i \in D$, we compute the average per token log likelihood $b_i$ when we steer with an SAE feature $f$.
• We repeat this over a dataset of text samples.
• We normalize all the log likelihoods (both the ones with and without steering) by subtracting a constant and scaling factor such that the non-steering log likelihoods range from 0 to 100. The constant and scaling factor are computed based on the non-steered log-likelihoods. We normalize this way so that the change in log likelihood after steering loosely represents a % change in likelihood after steering. We keep track of all the changes in LL after steering as $\delta_i$.

We record $\Delta_D$ as the mean change $\delta_{i, \text{LoRA}} - \delta_{i, \text{base}}$, with standard error estimates.

On a dataset consisting of text that exhibit the SAE feature (which we denote as “positive” texts), a model that steers more effectively would have a higher change in LL after steering (that is $\Delta_{\text{positive}} > 0$) because it would indicate the positive text sample is more likely after steering.

However, on a dataset of texts that do not exhibit the SAE feature (which we denote as “negative” texts), we want the change in LL to remain the same or slightly decrease (that is, $\Delta_{\text{negative}} \le 0$) after steering, as this indicates that steering with an SAE feature f does not significantly affect the likelihood of unrelated texts.

the method is somewhat a simple premise: SAE + LORA. There isn’t really a novel contribution in making LORA or SAE more amenable to solving the problem.

We agree the method is simple! We believe that’s a strong appeal of our results, as practitioners prefer implementing simpler techniques. We also study the performance of our technique in many settings and analyze why the method works (Section 6), which we believe goes beyond merely combining SAEs and LoRA.

What is “nats”? (e.g., line 088 “0.17 nats”)

nats is a unit measuring information and is the unit for cross entropy loss (it is the base e analog of bits)

captions above tables is a little unusual

We agree and usually prefer having them on the bottom, but having captions above the table is required as per the ICML style guide.

We are also excited to share we ran new evaluations on just our adapted model (no SAE inserted) and found that it outperforms the original model on many of the capability benchmarks across different model sizes (see table). We find it compelling that the models created using our method not only demonstrate stronger evidence of actually using the interpretable SAE features downstream, but also match or even exceed the general performance capabilities of the original pretrained model.

Italicized metrics indicate better performing between models with SAE inserted

Bolded metrics indicate best performing method overall.

Gemma-2-2b

GEMMA-2-9b

GEMMA-2-27b

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? And are there any further clarification or modifications we could make to improve your score?

We are greatly thankful for your time and feedback, especially related to the limitations of our work. We were glad to hear you found the evaluations comprehensive.

Additionally, I am concerned that far too much focus is on improving simple SAE metrics… At absolute minimum, the authors should be much more open about the limitations of their work.

This is a fair criticism, thank you! We agree that our multi-SAE section does not directly show improvements to SAE circuits. We have added the following to the limitations section, please let us know what you think!

While our approach develops models that more effectively leverage the interpretable features learned by SAEs in downstream tasks, we do not yet show that these can be leveraged into additional insights as to how the original model internally computes using these features. We leave this important direction for future work.

In addition to those worries about circuit analysis, I don’t think you can use error terms with your approach? But often attaching lots of SAEs is very lossy, so error terms are essential.

This is a good question! It is true that one cannot use error terms with our method to obtain the activations of the original model, but in return we get the benefit that our method directly attacks the issue that “attaching lots of SAEs is very lossy.” For example, note that in Figure 6, even inserting 7 SAEs with our method results in a reasonable cross entropy loss of 2.78, and inserting 3 SAEs into the LoRA model has better loss than the pretrained model with 1 SAE.

We also note that one could use error terms to recover the performance of the adapted model, which we show below is similar (and in some cases even better) than the original model.

Your paper requires a very expensive training procedure that will need to be redone if different SAEs need to be attached

While it’s true the training procedure needs to be redone if different SAEs are attached, the training procedure is actually very cheap. Taking off the shelf SAEs, training LoRA adapters takes only a few minutes or hours (depending on your model size), which is many times cheaper than the original SAE training process.

Why is this a promising method for applications of probing? …if the way to monitor requires editing weights before the SAE then the latents will not be able to used on their own as probes

This is a great point and something we overlooked. We have added a discussion of this to Section 5.1.

However, as you note someone using our method for probing could run the adapted model without the SAE (because they would just be using it for probing), and we are excited to show that in new experiments, the adapted model with no SAE is actually as good or better than the original model on many capability benchmarks:

Italicized indicates best performing between models with SAE inserted

Bolded indicates best overall.

Gemma-2-2b

GEMMA-2-9b

GEMMA-2-27b

…surely this is supposed to say un-frozen?

We think this is written correctly. To clarify, in LoRA finetuning, the weight matrices W of the original model are frozen, and only low rank matrices A and B are learned. The forward pass through the model is then (W + AB)x instead of Wx. We hope this clarifies things; we apologize for not being more clear.

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? If not, what further clarification or modifications could we make to improve your score?

We are grateful for your time and thorough feedback, especially with regards to our methodology. We were glad to hear you found the implementation to be well executed.

How do you reconcile your approach of modifying the model with the traditional goal of mechanistic interpretability?

This is an important point, thank you! As we state in our limitations (Section 7.1), our method is not appropriate when the model cannot be modified. However, while much interpretability research focuses on fixed models, other work builds interpretable models from scratch or modifies existing models to be interpretable. As cited in our related work, this includes Lai & Heimersheim, 2024; Lai & Huang, 2024; Elhage et al., 2022b; Liu et al., 2023; Liu et al., 2024; Heimersheim 2024. We consider our paper well situated in this latter category and thus an important research contribution.

Moreover, because pretrained models perform poorly with SAEs inserted, one may be skeptical if models are actually using the more interpretable SAE latents. An interpretable basis is unhelpful if the model does not actually use it. Because our method optimizes the model to use SAE latents downstream, we can be more confident they are really using them.

Have you analyzed whether the adapted model still exhibits the same internal computational patterns?

Great question! In Figure 11, we compare the distance between the downstream layer activations of the pretrained model (gemma-2-2b) with the downstream layer activations of 1) the pretrained model with the SAE, and 2) the adapted model with the SAE. We find that the adapted model with the SAE has a higher cosine similarity with the pretrained model’s original activations than the pretrained model with the SAE. Thus, our adapted model + SAE is more similar to the original model than the pretrained model + SAE.

We also ran a quick experiment comparing just the adapted and pretrained model (no SAE involved). We found the average per-layer cosine similarity between pretrained and adapted model activations on the Pile was consistently greater than 0.99 (for comparison, the average cosine similarity with the inserted SAE is ~0.95).

How do you ensure that the modified model still maintains the properties that made the original model worth studying?

Great question! It is true that by modifying the model, we may be losing capabilities of the original model. To allay this concern, we ran our adapted model on the benchmarks from Table 5 and found that not only does the modified model still perform competitively with the pretrained model, but that in many cases it outperforms the original pretrained model.

Italicized metrics indicate best performing between models with SAE inserted

Bolded metrics indicate best method overall.

Gemma-2-2b

GEMMA-2-9b

GEMMA-2-27b

Do you have evidence that the features discovered after model adaptation are more semantically interpretable to humans?

Thank you for this question; we apologize for not being clearer! We are not trying to find more interpretable features. Indeed, in sections 4.1 & 4.2 (where we recover up to 55% of the loss gap), we adapt only after the SAE, so the features (and their interpretability) are unchanged. Rather, we are creating an alternative model that better uses these interpretable features downstream.

What does the f in Eq. 5 (line 158) stand for?

SAE dictionary features.

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? If not, what further clarification or modifications could we make to improve your score?

We thank you for your time and help, especially in regards to your suggestions for additional experiments, which we have now run. We are especially glad that you appreciated the novelty of modifying the model itself to better fit a sparse autoencoder, as we agree that this idea is the major impact of our work.

More diverse model architectures could be tested, e.g., Mistral and Qwen (where there are already some SAEs and there's no need to train SAE from zero).

Thank you for this suggestion! We ran additional experiments on the layer 16 Mistral 7B SAE from SAELens and found that our method still works (we did not see a Qwen SAE on SAELens). The Mistral 7B base model has a CE loss of 1.5485 on our val set, the model with the SAE inserted has a CE loss of 1.8007, and the LoRA model with SAE inserted has a CE loss of 1.6548, reducing the cross entropy loss gap by 57.8%.

Different SAE sizes and capabilities could also be considered to prove that this method is constantly surpassing SAE rather than needing a strong SAE to start.

We show that our method works on a large number of SAE sizes and sparsities in Section 4; for all of them it surpasses the performance of the original SAE. We also show that even for partially trained SAEs our method works well; see Figure 1. Please let us know if there is anything we could add to the main text (experiments or wording) to make this stronger, thank you!

Could you replicate the performance experiment on standard NLP benchmarks on the 27B model? I'm concerned about whether this method could scale up to bigger models (and also more complicated features and capabilities).

Thank you for this suggestion! Below, we include a version of Table 5 with a 27B model. Our method seems to scale up well to bigger models. Interestingly, the adapted model alone outperforms the original pretrained model on many of the capability benchmarks across all model sizes; we believe this is evidence our method scales well, and that the adapted models from our method could be of interest for real world use cases.

Italicized metrics indicate best performing between models with SAE inserted Bolded metrics indicate best performing method overall.

Gemma-2-2b

GEMMA-2-9b

GEMMA-2-27b

Thank you again for taking the time to review the paper and providing helpful feedback! Do the above actions address your concerns with the paper? If not, what further clarification or modifications could we make to improve your score?