Paper: Codebook Features: Sparse and Discrete Interpretability for Neural Networks (ICLR-2024)

Reviewer comment (Experiments):
While Transformers are prevalent, there are many architecture differences between different models (e.g. novel layers, group-query attention, etc.). In this sense the study is limited in scope by focusing only on Transformer neural networks and examining their performance on a singular algorithmic dataset and two natural language datasets. This leaves unanswered questions about the generalizability of codebook features to other neural network architectures or different types of data, such as visual information.

Author response:
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.

Reviewer comment (Experiments):
While the paper demonstrates the capability of codebook features in topic manipulation for language models, it does not explore other linguistic features like sentiment, style, or logical flow. This limitation narrows the understanding of how versatile and broadly applicable codebook features might be for controlling various aspects of neural network behavior.

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

Reviewer comment (Experiments):
In authors' two-phase method for understanding and controlling the network's behavior, you focus on generating hypotheses for the role of codes and then steering the network by activating these codes. How robust is this method to the presence of adversarial or noisy input?

Author response:
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.

Reviewer comment (Evaluation):
Authors mention that codebook features reduce the complexity of a neural network’s hidden state, making it easier to control the network’s behavior. Could you provide more details on the trade-offs involved? Specifically, how does the use of codebooks affect the model's capacity for generalization across different tasks or data distributions?

Author response:
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.

Reviewer comment (Evaluation):
How do you find the activated codes and their corresponding topics? How do you know for example 'code 123' leads to the topic of 'dragon'?

Author response:
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.

Reviewer comment (Evaluation):
Will the model lead to better computational time?

Author response:
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.

Reviewer comment (Experiments):
One shortcoming of the experiment in table 1 is that it has two dependent variables: both the quantization level k and the codebook size C. It is important to ablate changes to these two variables separately to better understand if they both independently provide benefits. There is a similar issue of two dependent variables being tested at the same time in Table 2 where both k is modified as well the features being quantized (i.e., attention vs mlp features).

Author response:
**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.

Reviewer comment (Experiments):
The benefit of sparsity in the combination of codebook entries has not been clearly articulated. In fact, in section F.1.1 of the manuscript, it is argued that the continuous combination of different atoms reduces interpretability. Selecting codebook entries via top-k cosine similarity is a naive approach that has not been compared to more recent sparse coding techniques (e.g., variational sparse coding methods like in Tonolini et al. 2020 and Fallah et al. 2022). Except for the potential increase in modeling capacity, the work does not demonstrate the benefit of quantizing with multiple codebook entries per token.

Author response:
**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.

Reviewer comment (Theory):
Can the authors clarify the benefit from increasing k? Since the authors use cosine similarity between the activations and the codebook to pick the top-k entries, what would the difference be between each of these k codebook entries? It seems that taking k entries contradicts the viewpoint of “features-as-points”, and may even lead to what the authors refer to as “smuggling of information”. I would expect that increasing k may improve performance of the model, at potential cost to interpretability.

Author response:
**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.

Reviewer comment (Writing):
Out of curiosity, have the authors considered quantizing each vector along the feature components (i.e., divide the N features of each token into k blocks)? Is this what the authors refer to as “grouped codebooks” in E3? If so, I believe this warrants discussion and attention in the main text. This would be an alternative to quantizing each token with k codebook entries that are very close in cosine similarity and be closer to the VQ-VAE setting.

Author response:
**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.

Reviewer comment (Experiments):
**Experiments:** As mentioned in the paper that a code can be related to some specific semantics; however, the results supporting such claim appears insufficient. The authors may consider conducting comparative analyses of the distribution disparities between code associated with similar and dissimilar semantics to substantiate this claim. In addition, is it possible that a (some) certain code(s) may correspond to a multitude of distinct semantic contexts?

Author response:
**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.

Reviewer comment (Experiments):
**Experiments:** It is interesting that model outputs is controlled by the codes. However, based on the results on TokFSM dataset, it appears that the code following the MLP layer plays a more significant role. Nevertheless, the experiments conducted by the authors on WikiText-103 dataset only involve discrete attention layers (in Table 2 (b)). Does this incongruity potentially render it challenging to control the model?

Author response:
**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.

Reviewer comment (Evaluation):
**Experiments:** In Section 2, the authors sum top-k codes weighted by the same value (specifically, 1). How will each code influence to the model's decision? (For example, does codes having semantics related to the target task contributes most to the model's decision evaluated by attribution methods?) In particular, the authors can utilize feature attribution methods [3-5] to present quantitative and qualitative analyses.

Author response:
**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.

Reviewer comment (Reproducibility):
The authors do not provide codes for reproducibility check.

Author response:
**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