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

Thank you for your thoughtful review. We’re glad that you think the research direction in our paper is very valuable.

Comparison of the advantages of the two proposed measurement methods compared to reconstruction loss and sparsity: Thank you for giving us the opportunity to clarify this. In our paper, we compare three SAE training methodologies: Gated, Gated with p-annealing, and Standard with p-annealing. All three show Pareto improvements over the Standard approach in terms of reconstruction loss and sparsity but cannot be distinguished using existing unsupervised metrics (L0 and fraction of loss recovered). However, as you can see in the pdf attached to our global comment, the metrics we introduce clearly distinguish the SAEs. In other words, our metrics reveal differences in SAE quality which are invisible to prior unsupervised metrics. (As an especially clear example, the parallel curves that appear in our plots correspond to SAEs trained with the same algorithm but different hidden dimensions; our metrics clearly demonstrate that the SAEs with larger hidden dimensions are better than those with smaller hidden dimensions, as expected.)

No significant difference between low-level and high-level coverage, unreasonable selection of rules: Thank you for pointing this out. In our original submission there was a mismatch between graphs and captions, and our actual high-level coverage graph was not included. We have provided an updated selection of graphs in the PDF.

We also agree that our initial split into low vs. high seemed ambiguous and subjective. To address this, we moved to a more principled division of BSPs into (1) board state BSPs, which correspond to the presence of a particular piece at a particular board space (giving 8 x 8 x 12 board state BSPs for chess and 8 x 8 x 2 for Othello), and (2) researcher-selected strategy BSPs for chess. Our goal was to confine researcher choice to the strategy BSPs as much as possible, while keeping the class of board state BSPs natural and principled.

Some high-level features require linear combinations of multiple low-level features to be expressed: We would like to clarify that linear probing is not used in either the coverage or reconstruction metrics, which instead measure whether individual features found using SAEs correspond to individual BSPs. From our perspective, neither our work nor prior work suggests that interpretable high-level properties are necessarily expressed using an ensemble of low-level features. For example, in our work, we find individual features corresponding to high-level properties such as “an en passant capture is available”. Similarly, in Anthropic’s Scaling Monosemanticity [4] they find “a diversity of highly abstract features”. We hope this addresses your concern, but we also acknowledge that we might be misinterpreting your question. If this is the case, could you please provide further clarification?

Training of SAEs on various locations: We only train SAEs on the residual stream between layers 6 and 7, and do not train SAEs on other locations. We chose this location because it achieved the best performance when reconstructing the board state with linear probes, which is consistent with prior work [1, 2, 3].

Assumption of High Precision: We had several motivations for choosing high precision features. The first was that during manual inspection of SAE features, we found that features often seemed to activate on specific board configurations, rather than individual board square states. We also noted studies on chess players showing that chess experts excel at remembering realistic board configurations, but not random piece placements[4]. This suggests experts (and potentially AI models) encode board states as meaningful patterns rather than individual square occupancies.

Concurrently, Anthropic has also evaluated features on their precision in Scaling Monosemanticity, which they called “specificity” [5].

Hyperparameter adjustment: We sweep over the hyperparameters of sparsity penalty, learning rate, and expansion factor. A lower sparsity penalty tends to increase L0, while a larger expansion factor tends to increase Loss Recovered.

Relative Reconstruction Bias (Gamma from GDM paper [6]) Thank you for this suggestion. We have performed additional experiments measuring relative reconstruction bias. Our results show p-annealing achieves similar relative reconstruction bias to gated and gated w/ p-annealing, all improving over the standard SAE. The results can be seen in Figure 4 of our global pdf. We note that the y-axis scale for relative reconstruction bias differs between Chess and Othello SAEs. Chess shows a narrower range (minimum γ ≈ 0.98) compared to Othello (minimum γ ≈ 0.80), which more closely resembles the original GDM plot [6]. This may reflect differences in the underlying models or data. In the Othello plot, we found erratic relative reconstruction bias values for SAEs with L0s very near 0; these are very bad SAEs outside of the sparsity range usually considered (including in [6]).

[1] K. Li, A. K. Hopkins, D. Bau, F. Viégas, H. Pfister, and M. Wattenberg, ‘Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task’, arXiv [cs.LG]. 2024.

[2] N. Nanda, A. Lee, and M. Wattenberg, ‘Emergent Linear Representations in World Models of Self-Supervised Sequence Models’, arXiv [cs.LG]. 2023.

[3] A. Karvonen, ‘Emergent World Models and Latent Variable Estimation in Chess-Playing Language Models’, arXiv [cs.LG]. 2024.

[4] Frey PW, Adesman P. ‘Recall memory for visually presented chess positions.’ Mem Cognit. 1976.

[5] A. Templeton et al., ‘Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet’, Transformer Circuits Thread, 2024.

[6] Rajamanoharan, Senthooran, et al. "Improving dictionary learning with gated sparse autoencoders." arXiv preprint arXiv:2404.16014 (2024).

Thank you for your review. We’re glad that you enjoyed reading the paper and think it’s of high relevance to the NeurIPS community.

Demonstrating transfer to natural language: Yes, we agree that this is a limitation in our paper and future work should study this. However, note that our metrics reveal progress in SAE training which is invisible to existing proxy metrics. Specifically, SAEs with larger hidden dimensions perform better on our metrics, e.g. for the Standard architecture this is reflected in the parallel lines of purple diamonds (see attached PDF). The hope would be that training techniques which work better in the board game setting also work better in the natural language setting, though we don't validate that in this paper.

Sensitivity of SAE training to p-annealing schedule and other hyperparameters: We did not conduct a systematic evaluation of the p-annealing schedule and other hyperparameters. From our experience it did not seem that sensitive to the schedule, but is a bit sensitive to the p_end value, i.e. the final p value at which we would stop the training. If p_end was below 0.15, the training started to destabilize.

We also noticed an effect on the initial lambda coefficient (strength of the sparsity penalty). Specifically, when using p-annealing, a small change to lambda could lead to larger differences in L0 than when training without it. This means it would often tighten the range of lambdas that we would explore. However, this only meant that we would explore over less orders of magnitude of lambda coefficients and get the same L0 spread.

Thank you for your review. We’re glad that you think that our metrics intuitively make sense and that the motivation for p-annealing is well thought out and presented.

Mismatch between figures and caption: Thank you; this is a mistake and we will correct this in the paper. We have included the updated plots in the attached PDF. Further updates are outlined in our global response.

Less sparse SAEs have worse coverage: Here are two related perspectives on the connection between sparsity and coverage:

• BSPs are typically sparse (e.g. only a sparse subset of chess boards have a knight on E3). Thus, if the SAE features are dense (as is the case for high L0 SAEs), then they have no chance of corresponding to sparse BSPs; thus sufficiently dense SAEs must have low coverage scores.
• Consistent with prior work, we find that sparser SAEs have more interpretable features. Since coverage is a measure of whether our SAEs have features belonging to a certain class of interpretable properties, more interpretable SAEs will also tend to have better coverage scores.

(These two perspectives are really two sides of the same coin: the connection between sparsity and interpretable properties is a key motivator for applying SAEs for interpretability [1, 2].)

Separate plots for evaluating p-annealing: In the three variable plots, we wanted to demonstrate that p-annealing is a Pareto improvement in L0 / Loss recovered over the standard SAE and at the same time showcase how our new metrics can help to distinguish improvements invisible to existing metrics. Specifically, it becomes difficult to distinguish between the gated and p-annealing SAEs using existing metrics, while their differences are still visible in plots using our new metrics. However, we agree that it is difficult to distinguish between the gated and p-anneal SAEs in these plots and will make sure to improve the presentation in the final version.

Efficacy of p-annealing: Thank you for pointing this out. The final sentence in the conclusion should be updated to better reflect our results. Specifically, we find that standard SAEs trained using p-annealing consistently perform better than those trained with constant L1 penalty by existing proxy metrics and in terms of coverage. However, Gated SAEs are on par with them in terms of coverage. We will include a more nuanced discussion in the results section.

Qualitative analysis of trained SAEs: We did perform qualitative analysis of the trained SAEs, e.g. Figure 1 includes examples of game states that correspond to interpretable SAE features that we found. However, we found many more interpretable features that did not make it into the paper, such as an “en passant” feature that only activates when an en passant capture is available. We will provide additional qualitative results of trained SAEs alongside the new metrics in the appendix.

F_low vs F_high distinction: You're right, our split into low vs. high seemed ambiguous and subjective. To address this, we moved to a more principled division of BSPs into (1) board state BSPs, which correspond to the presence of a particular piece at a particular board space, and (2) researcher-selected strategy BSPs for chess This is described in more detail in the global rebuttal. Please find our revised figures in the attached PDF.

Board reconstruction without high precision features: Only high precision features are used in calculating the metric. If there are no high precision features for a BSP, the board reconstruction score for that BSP would be 0. We will clarify this in the text.

Batch of train data: We apologize for the confusion caused by our terminology. To clarify, we use a consistent dataset of 1000 games as our training set for identifying high-precision features across all Board State Properties (BSPs). An additional, separate set of 1000 games serves as our test set. We will clarify this in our updated paper.

Qualitative analysis of metrics: In qualitative analysis of high level BSP features, the features are inherently interpretable because we screen for 95% precision. Thus, for example, at least 95% of a pin feature’s activations will be when there is a pin on the board. We also noticed that high level features tend to activate in specific scenarios. For example, a pin feature may only activate on a more specific type of pin, such as “a pin by a bishop in the bottom left corner of the board”.

Number of SAEs trained The plots contain 40 SAEs from each SAE training methodology for each board game. We had additional sweeps of 40 Standard SAEs on layer 1 of each trained model, and on layers 1 and 6 of the randomly initialized models.

Line 220, optimal L0 range We try to make the same point as in our comment on Less sparse SAEs have worse coverage: above. Thank you for pointing this out, the word optimal is too strong here. We will replace “optimal L0 range” with “expected L0 range in line with previous work”.

Correlation between proposed metrics and final SAE reconstruction loss While there's some ambiguity in the term "final SAE reconstruction loss”, we interpret this question as comparing Loss Recovered with our new metrics. Our metrics generally perform best at the elbow of the L0 / Loss Recovered plot, balancing sparsity and reconstruction quality. A higher Loss Recovered doesn't necessarily mean better performance on our metrics, as it often comes with higher L0 (less sparsity). Conversely, very low L0 tends to have low Loss Recovered. Thus, the relationship between our metrics and reconstruction loss isn't straightforward, as we're optimizing for a balance rather than maximizing either variable.

References

[1] A. Templeton et al., ‘Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet’, Transformer Circuits Thread, 2024.

[2] Elhage, et al., "Toy Models of Superposition", Transformer Circuits Thread, 2022.