Codebook Features: Sparse and Discrete Interpretability for Neural Networks

We thank the reviewer for their review! We are glad they thought highly of the soundness, presentation, and contribution of the work.

Generalizability to other forms of neural networks

It is certainly true that we focus on Transformer neural networks in this work due to how widely used the architecture is. We present results applying codebook features to different parts of the transformer (attention and MLP layers), different size transformers (from 1 to 24 layers) and different datasets (synthetic data to natural language). We agree that an exciting direction for future work is in extending our method to other modalities and network architectures.

Other ways to control the network such as sentiment, style, and logical flow

We certainly agree; as we note in our paper, this is an exciting opportunity for future work.

How robust is the method of generating hypotheses and then steering the network?

In our experiments, we were able to find topic codes across a wide range of inputs, including a wide range of different domains within Wikipedia as well as the TinyStories dataset. However we agree that studying the relationship between codebooks and adversarial robustness would be interesting more broadly; we discuss this direction in a bit more detail in Appendix G3.

Tradeoffs of codebook features and generalization to other data distributions

Yes, there are a few trade-offs to using codebooks. First, codebooks require additional computations to be performed on the network in exchange for their interpretability benefits. Second, the performance of the codebook model is often slightly lower than the original model, although the difference is small in practice (Table 2). In general, the loss of a language model is a good measure of its performance across distributions it is evaluated on, so the small difference in loss between the codebook and base models aligns with our experience that the model continues to perform well across distributions.

We thank the reviewer for their review! We are glad they find the results interesting and applicable in a wide range of use cases.

Computational efficiency and sparsity

This is a good question. We have provided some information about computational efficiency in Table 9. Indeed, leveraging sparsity through fast maximum inner product search does increase the throughput of the model. We suspect that custom sparse kernels would continue to increase the efficiency of codebook features further.

How do we associate codes with topics?

We follow a two-stage process. First, we take the code and retrieve examples where that code is activated. A person (or in the future, a language model) then looks at the examples and develops a hypothesis for the role of the code (e.g. “dragon,” if many dragon-related passages are retrieved). We then verify this role by activating the code during generation and seeing if the model’s behavior changes as we expected. Section 2.2 describes this process in more detail, and Sections 3 and 4 describe our automated evaluations for verifying the role of codes.

How are dead codes avoided?

While we see some number of dead codes during training, they do not significantly impact the success of the method. For example, in the 410M parameter model we train, fewer than 2% of codes were dead across the 384 codebooks we train on all the attention heads. Moreover, these dead codes do not appear to negatively affect the performance of the model, and they can be discarded after training.

Are some codes activated all the time?

We do see a small number of codes that are activated very frequently (see Figure 7). However, the majority of codes are activated far more rarely. The codes that occur most frequently tend to be associated with very general topics or very common words (e.g. articles like “the”).

We thank the reviewer for their review! We are glad they found the results surprising and promising.

Can multiple concepts be associated with a single code?

It is certainly possible that a code could be associated with multiple distinct concepts. We study this in Appendix D.4 and Figure 8, where we find that almost every code holds a single concept at the first layer. In later layers, we find that more concepts share a code, however we observe a good reason for this—namely, that these concepts are related to one another and tend to share the same next token (Figure 9). We agree that further investigation of this question is a fruitful direction for future work and have noted this in our revision.

Why we use attention codes in one model and MLP codes in another

This is a good question. The reason we use different types of codes here is that we are trying to control different aspects of the sequence/text in each model. In the TokFSM environment, we are trying to alter the prediction of an individual state/token. We find codes in the MLP layers are most associated with these single tokens. For the language modeling experiments, we are trying to alter the global topic of a generation. Topics typically manifest across many tokens, rather than a single token, and we find the attention layers are most associated with these features. However, we believe it is likely that, e.g., for more local linguistic features (such as word choice) editing the MLP codes in a language model may prove to be the best choice.

”In Table 3, the authors have presented PriViT and MPCViT in distinct hyper-parameter configurations, giving rise to concerns regarding the fairness of the comparisons.”

We believe this comment may be about another paper, as we do not present or discuss results on methods called PriViT or MPCViT.

How does each code influence a model’s decisions?

This is a good question. To evaluate this, we choose several families of topic codes and compute how varying the number of topic codes we activate changes the fraction of generations that follow that topic. As expected, we find that increasing the number of topic codes changes the fraction of decisions (Figure 10). Qualitatively, when we look at the activated tokens, we note that there is significant overlap in the tokens activated by each code; however, there are still many tokens which are activated by some codes and not others, suggesting that the multiple codes together may help the network form a more robust version of the topic concept.

Computer vision papers

We thank the reviewer for the pointers to related work and have included them in our most recent revision.

Code for reproducibility

We have included an anonymized codebase in the newest supplementary materials of our paper. We have also uploaded the same codebase along with our model weights here: https://drive.proton.me/urls/C205JG7GQM#l7CfinlSZUHM

Vector graphics

The newest version of our paper now uses vector graphics for most of the figures, which should result in improved readability. We thank the reviewer for the suggestion.

We thank the reviewer for their review!

Untangling dependent variables in Table 1

This is a good point. To address this issue, we reran our experiments to hold $C$ constant while changing $k$. We include these results in Table 1, which mirror the results from our previous version. We also include an experiment with an MLP codebook with $k=8$ and $16$ groups in Table 2; this enables a direct comparison to the attention $k=8$ model because our model has $16$ attention heads.

Reason for using multiple codes per layer

Yes, the primary benefit of working with $k>1$ codebooks is expressivity; because models likely represent more than 1 feature per layer, our interpretability method needs to be able to model multiple features. While we chose a simple method for our initial experiments, we agree that more sophisticated methods, including variational sparse coding approaches, are promising. We have mentioned this, including Tonolini et al. 2020 and Fallah et al. 2022, in our newest revision.

What happens with increasing k?

Yes, increasing $k$ does improve the performance of the codebooks model. In our updated Table 1, we show several runs where we only change $k$, demonstrating steady improvements in performance. In terms of interpretability, we agree that in the limit increasing $k$ could make the hidden state harder to understand. However, as Table 5 shows, even large values of $k$ transmit many fewer bits of information than a standard neural network layer, limiting the complexity of the hidden state. We will include this discussion in our revision.

Additional motivation for top-k, and whether the top-k codes have different roles

High dimensional spaces enable code vectors containing different information to be selected through top-k cosine similarity. As an example, consider the vector [$a$; $b$], the concatenation of two smaller vectors $a$ and $b$. This vector would have relatively high cosine similarity with [$a$; 0] and [0; $b$], despite the two vectors being potentially very distant from each other. Thus, we believe the $k>1$ models are fully compatible with the features-as-points viewpoint. However, there are other cases we have seen where several of the $k$ codes appear to activate on similar tokens. This may indicate that the model needs to represent fewer features than $k$, or that the redundancy is helpful for the network computation.

Is interpretation manual?

Yes, currently interpreting the codebook features requires manual interpretation. However, other work (e.g. [1, 2]) has demonstrated that language models can be used to interpret neurons or features as well, which we expect to be a more scalable solution.

Does interpretation require forward-passes?

Yes, currently interpretation does require a forward pass over a subset of the data. However, forward passes on this subset of data can be computed once and then the resulting data can be used to interpret the entire network. In practice, we expect the amount of compute needed to do this will be much smaller than that used to finetune the network to produce the codebook features.

Grouped codebooks

Yes, we believe the reviewer’s description is the same as our grouped codebooks method. We have included more discussion of this in the paper, however we note (per the example given in the previous paragraph) the $k>1$ codebooks can implement something quite similar to grouped codebooks. Thus, it is possible the primary benefit of grouped codebooks may be in improving optimization of the network, rather than in the expressivity or interpretability of the codebook itself.

Notation in 2.1

Yes, for the reconstruction loss, $x$ refers to the continuous activations. We changed this to $a$ in the most recent revision. The MSE is taken with respect to the output of the codebook (i.e., the sum of the codes).

[1] Natural Language Descriptions of Deep Visual Features, Hernandez et al 2022

[2] Language models can explain neurons in language models, Bills et al 2023