Measuring Progress in Dictionary Learning for Language Model Interpretability with Board Game Models

Summary

This paper proposes two measures to reflect the quality of the trained sparse auto-encoder when decomposing superposition features. In addition, this paper achieves approximate L0 optimization through dynamic adjustment of the p-norm and testing it on two types of board games.

Strengths

The research direction of this article is very valuable. Previous work using SAE for feature detection and circuit discovery did not directly measure the interpretability of the model decomposition results but relied on indirect measures such as reconstruction and sparsity levels of SAE. This work proposes using a few rules (low or high level) to automatically analyze the results of SAE decomposition under different hyperparameters, which is helpful for subsequent selection of SAE used for feature detection. Additionally, dynamically adjusting the p-norm is a straightforward method that intuitively approximates optimization of L0.

Weaknesses

- The experimental objectives of the article are not clear: The article lacks of comparing the advantages of the two proposed measurement methods compared to reconstruction loss and sparsity. For example, demonstrating cases where both reconstruction loss and sparsity perform well but coverage and reconstruction are low, resulting in biases in downstream tasks (such as circuit discovery). - The rationality of the two proposed measurement methods in the article is questionable: The purpose of using SAE is to decompose the features of superposition into monosemantic features, with the expected decomposition results leaning towards several low-level features. However, based on the experimental results in Figure 2, there is no significant difference between low-level and high-level coverage results. This phenomenon is likely due to the unreasonable selection of rules. - Similarly, due to the fact that both coverage and reconstruction are ultimately detected using linear probing on individual features, some high-level features require linear combinations of multiple low-level features to be expressed [1,2,3]. If the measurement methods described in the article are used, it will not be effective in detecting these features. - Typos (not limited to): Line 154 "than than" should be corrected. - The image captions are incorrect: Figure 2 should be labeled "coverage" instead of "reconstruction" for the bottom row. [1] Templeton, Adly, et al. "Scaling monosemanticity: Extracting interpretable features from claude 3 sonnet." Transformer Circuits Thread (2024). [2] He, Zhengfu, et al. "Dictionary Learning Improves Patch-Free Circuit Discovery in Mechanistic Interpretability: A Case Study on Othello-GPT." arXiv preprint arXiv:2402.12201 (2024). [3] Rajamanoharan, Senthooran, et al. "Improving dictionary learning with gated sparse autoencoders." arXiv preprint arXiv:2404.16014 (2024).

Questions

- Is there a separate training of SAE for MLPs, Attention outputs, etc., which write into the residual stream instead of activations between attention blocks? - Is there any relevant reference for the assumption about *High Precision* in line 130? - How were the hyperparameters adjusted in Figures 2, 3, 4, and 5 to obtain several different SAEs? (Causality between hyperparameters and SAE results needs further clarification [1].) - In [2], it's mentioned that Gated SAE can solve the shrinkage problem mentioned in line 167. Have you compared the relative reconstruction bias (gamma) proposed in [2] when using gated SAE w/wo p-annealing? [1] Kramár, János, et al. "AtP*: An efficient and scalable method for localizing LLM behaviour to components." arXiv preprint arXiv:2403.00745 (2024). [2] Rajamanoharan, Senthooran, et al. "Improving dictionary learning with gated sparse autoencoders." arXiv preprint arXiv:2404.16014 (2024).

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

3

Limitations

I suggest the authors expand the scope of the research to include more domains and datasets, and validate the effectiveness of the proposed methods in other downstream interpretability tasks.

Summary

This manuscript applies sparse autoencoders (SAE) to detect interpretable features from autoregressive transformer models trained on Othello and chess. These controlled scenarios provide suitable testbeds, in the sense that we can extract ground truth features to measure progress in dictionary learning. Based on the observations that existing SAE optimizations may indeed be suboptimal, the authors propose *p-annealing*, a warm-start technique that iteratively optimize a sequence of (increasingly more non-convex) SAE objectives, in order to approximate the intractable $\ell_0$ objective.

Strengths

* The presentation of this paper is very clear. I really enjoyed reading this paper overall. * The topic of this paper is of high relevance NeurIPS community. SAEs are shown to be a promising approach for mechanistic interpretability. Providing a quick-to-iterate testbed as well as a set of improvement techniques would surely be valuable to this research. * The $p$-annealing technique is reasonable, is easy to implement, and has rich history in sequential and/or global optimization. It also demonstrate superior performance compared to standard SAE objectives. * Experiments are thorough, and demonstrate (1) the applicability of SAEs to learn interpretable features for Othello and chess tasks, and (2) the benefits of improved techniques on these tasks.

Weaknesses

* While it is reasonable to operate in a controlled setup (i.e. Othello and chess) in this paper, there lacks concrete evidence on how these improved techniques can potentially improve SAEs in LLMs and natural language domains. Including a discussion and/or additional experiments on this could further strengthen the paper.

Questions

* How sensitive is the SAE training to the choice of $p$-annealing schedule and other choice of hyperparameters?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have adequately addressed limitations and potential societal impacts of their work

Summary

The paper introduces a setting for evaluating dictionary learning methods (and in particular sparse autoencoders) for language model interpretability. The setting proposed is that of interpreting features learned by language models trained on data representing Chess and Othello games. This setting should allow testing based on a more easily enumerable set of ground truth features (compared to natural language) which we know are relevant to the task and are captured by the underlying model. The authors propose two new F1 based metrics to evaluate the SAEs ability to recover features learned by the model, namely a "coverage" metric and a "board reconstruction" metric that measures the quality of high precision features in the SAE. Both these metrics are computed by treating SAE features as classifiers for binary board state properties. The paper also introduces a new training technique for SAEs called p-annealing that modifies the L1 loss to more closely approximate optimizing the L0 norm for which L1 loss acts as a proxy.

Strengths

- Having a better grasp on how to evaluate SAE's training methods would be very useful as they are generally fairly difficult objects to evaluate in terms of how interpretable they are (without a lot of manual inspection). The setting used is complex enough to be interesting while maintaining an ability to know ground truth concepts that should be present and separately probe that the models the SAEs are trying to give insight into actually have representations of those concepts. - The metrics proposed intuitively make sense and appear to relate well to established metrics like precision, recall and F1. - The motivation for p-annealing is well thought out and presented. And the method appears to have a strong effect on lowering overall sparsity.

Weaknesses

- The figures and caption for figure 2 don't seem to match up (there may be a mixup with subfigures for figure 3). From reading the caption I would have expected the bottom row to show "Coverage" for high level BSPs in a similar pattern to the top row, however it shows "Reconstruction" of high level BSPs and reconstruction of low-level BSPs (which seem to appear as subfigures in figure 3). Could the authors correct this/post the appropriate charts? - I can't easily understand the comparison of the new metrics compared to the existing proxy metrics, the figures (2-5) seem to attempt to both validate the new metrics as well as compare the proposed p-annealing technique and thus get fairly busy making them hard to parse visually: - The 3 variable plots on the left hand column of the figures are somewhat hard to read, especially since the proposed metric is encoded using color. Having a plot comparing each proxy metric (loss recovered and sparsity) to each proposed metrics would be easier to read. I do understand that the authors are trying to demonstrate that there is a region on the Pareto frontier where their metrics perform best, its not so apparent to me that sparser SAEs are always more "interpretable" SAEs so it is an interesting finding, but I still think it would be helpful to more easily compare each metric to the existing proxy metrics (for example I wouldn't necessarily expect that less sparse SAEs should have worse coverage—but maybe the authors have some thoughts on why that should be the case?) - Having separate plots that focus on evaluating the p-annealing technique using the various metrics (proxy and proposed) when also be helpful (i.e. similar to Fig 5 in Rajamanoharan et al 2024). The overplotting of the various shapes for the different conditions (particularly in the left hand column) makes it a bit hard to understand the the trends and understand how well p-annealing works actually works. - From my read of the charts, the results for p-annealing seem quite different between Othello and Chess. I'm trying to map the final sentence about the efficacy p-annealing in the conclusion to the charts and i'm not sure if i'm able to see how the conclusion follows from the results. Could the authors comment on this? - There isn't much description of qualitative/manual analysis of the trained SAEs alongside the new metrics. Because measurement is relatively nascent in this space, it would be helpful to have some qualitative evaluation/results from manual inspection demonstrated.

Questions

- What is functionally/representationally different about BSPs in F_low vs F_high? "Any square being threatened" seems as high level as "has a queen on the board". I'm just trying to understand if there is a critical difference between these two classes of BSP from the perspective of the SAE or BSP classifier. It is nice that these concepts come from different sources, just wondering about the low-level vs high level implication. - What is the effect on the board reconstruction metric if there are BSPs for which there are no high precision features (line 131)? Are these low precision features still used in calculating the final metric? It was not super clear to me from the notation in the equation below line 134 if **all** features are used used in this calculation, _but_ φfi,t(x) = 1 is only calculated for high precision features, or if **only** high precision features are included in this computation at all? - Why use a "batch of train data" to identify the high precision features? Why not use the whole dataset? How big should this batch be relative to the whole dataset? - Have the authors done any qualitative analysis of their metrics? - For example board reconstruction seems to overall be lower for high-level vs low-level BSPs. Have the authors observed that this translates into more difficulty interpreting board states retrieves by featured for the high level BSPs? - How many SAEs are trained for each of the 4 SAE training conditions {Gated, Standard} x {p-annealing, no p-annealing}? - Could you elaborate on line 220 where you say "in the region with L0 below 100, which is the optimal range according to proxy metrics"? Is this a reference to an existing finding, or how do we know that this is the optimal range for sparsity? A more speculative question: - Do the authors know if there is any correlation between their proposed metrics and the final reconstruction loss of the SAE? Particuarly since the linear probe performance on predicting BSPs seems to be high, might we expect to see a correlation to just straight reconstruction loss as another signal to validate these metrics or am I off base?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Any final questions?

Dear reviewers, The discussion period is ending soon. If you have additional questions for the authors, please post them now. One of the questions that two of the reviewers raised was the paper does not address enough domains, especially not natural language and so the utility of the approach may be harder to judge. Reviewer pBfq, your rating differs significantly from the others -- have the author's responses addressed any of your concerns ? Thanks.

Reply for Authors

I appreciate the authors' response. My main question focuses on the first point. Let me briefly explain my understanding of the first question. > The author believes that the existing L0 and reconstruction loss cannot accurately measure the training results of SAE (**I also agree**), Therefore, the author proposed two new measurement methods and observed that the extreme value area of ​​the new measurement method is at the Pareto front of the original measurement method, so the two new measures are considered to be better than the original method. But this only shows that the new measurement method considers L0 and reconstruction loss at the same time (similar to harmonic averaging), but does not reflect the loopholes of the original measurement (for example, different training settings obtained two groups of L0 and reconstruction loss with the same SAE, but the subsequent performance of these two groups of SAE is very different). And this is also my main concern with this article. Therefore, I maintain my rating.
 And I am open to revising my rating based on discussions with other reviewers/ AC.