Interpretability Illusions in the Generalization of Simplified Models

Thank you for your helpful feedback. We are glad that you find our paper addresses an important problem, the experiments are detailed and rigorous, the writing is clear, and the paper presents interesting insights.

The focus of the experimental analysis is too narrow

As described in the overall response, we have extended our experiments to include a more realistic setting of language modeling in code datasets. We originally chose to focus on the Dyck languages because they capture a fundamental aspect of natural languages–hierarchical syntactic structure–but are simple enough to allow for some mechanistic understanding. The Dyck languages are considered a canonical instance of context free languages, because any context free language can be formed by the intersection of Dyck and a regular language [1]. For this reason, a variety of prior work has also focused exclusively on this setting as a controlled test bed for studying questions about expressivity and interpretability (e.g. [2-4]).

However, as noted in the general response, we have extended our experiments to incorporate an additional setting: language modeling on code datasets, training models with more layers and attention heads. These experiments show that generalization gaps also occur in more naturalistic settings and provide some additional insight into the relationship between generalization gaps and properties of the task. To summarize:

• We trained character-level language models on Java functions and evaluated generalization to functions with deeper nesting depths and unseen programming languages.
• These experiments revealed a consistent generalization gap: simplified proxy models are more faithful to the original model on in-distribution data and worse on OOD data.
• We also broke down these results according to different subsets of the code completion task, and found that the generalization gap is notably larger on some types of predictions. In particular, the generalization gap is larger on subtasks that can be seen as more “algorithmic”, including predicting closing brackets and copying variable names from earlier in the sequence. The gap is smaller when predicting whitespace characters and keywords, perhaps because these predictions rely more on local context than precise attention.
• We were able to attribute some of these generalization gaps to an “induction head” mechanism, which copies words from earlier in the sequence. Simplifying this head leads to large gaps in approximation quality between in-domain and out-of-domain samples, meaning that the low-dimensional approximation underestimates the extent to which the induction head mechanism will generalize to unseen distributions.

Overall, these experiments provide evidence that generalization gaps also occur in naturalistic settings, and they provide additional insight into how these gaps are related to different properties of sequence modeling tasks. Specifically, generalization gaps might be most pronounced in “algorithmic” settings where the model must use a particular feature in a precise, context-dependent way, and there are fewer or no other proxies of that feature available in the data. The effect is diminished in settings where a variety of local features contribute to the model’s prediction. We hope that these additional experiments help to address concerns about the scope of our analysis, and, if so, you might consider reevaluating your score.

Specialized to a specific class of models

We note that our additional experiments extend our analysis to models with more layers and attention heads. We do focus on a specific class of models–Transformer language models. However, this is now the predominant class of models across AI, and we are joining a considerable body of research in investigating how we can better understand these systems.

References

[1] Kambites, M. (2009). Formal languages and groups as memory. Communications in Algebra, 37(1), 193-208.

[2] Yao, S., Peng, B., Papadimitriou, C., & Narasimhan, K. (2021). Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115.

[3] Hewitt, J., Hahn, M., Ganguli, S., Liang, P., & Manning, C. D. (2020). RNNs can generate bounded hierarchical languages with optimal memory. arXiv preprint arXiv:2010.07515.

[4] Wen, K., Li, Y., Liu, B., & Risteski, A. (2023). (Un) interpretability of Transformers: a case study with Dyck grammars.

Many thanks! I had a look at the experiments and responses, and it appears to improve the paper and address the reviewers' concerns. I am saying 'appears to improve' because I would like to read it more thoroughly, and there has been unfortunately only a short period of time between the rebuttal submission and the deadline. I understand this is due to the very limited time constraints for the rebuttal preparation, and that it took lots of effort to prepare the revision and I would like to thank the authors for it.

Therefore, I will read it carefully again and take these comments and revision into account when submitting post-rebuttal recommendation. My recommendation remains unchanged for now.

Thank you for your helpful feedback! We’re glad you see the paper as addressing an important problem, found the experiments to be original and thoroughly reproducible, and found the paper to be clearly written.

Limited scope

As noted in the general response, we have extended our experiments to incorporate an additional setting: language modeling on code datasets, training models with more layers and attention heads. These experiments show that generalization gaps also occur in more naturalistic settings and provide some additional insight into the relationship between generalization gaps and properties of the task. To summarize:

• We trained character-level language models on Java functions and evaluated generalization to functions with deeper nesting depths and unseen programming languages.
• These experiments revealed a consistent generalization gap: simplified proxy models are more faithful to the original model on in-distribution data and worse on OOD data.
• We also broke down these results according to different subsets of the code completion task, and found that the generalization gap is notably larger on some types of predictions. In particular, the generalization gap is larger on subtasks that can be seen as more “algorithmic”, including predicting closing brackets and copying variable names from earlier in the sequence. The gap is smaller when predicting whitespace characters and keywords, perhaps because these predictions rely more on local context than precise attention.
• We were able to attribute some of these generalization gaps to an “induction head” mechanism, which copies words from earlier in the sequence. Simplifying this head leads to large gaps in approximation quality between in-domain and out-of-domain samples, meaning that the low-dimensional approximation underestimates the extent to which the induction head mechanism will generalize to unseen distributions.

Overall, these experiments provide evidence that generalization gaps also occur in naturalistic settings, and they provide additional insight into how these gaps are related to different properties of sequence modeling tasks. Specifically, generalization gaps might be most pronounced in “algorithmic” settings where the model must use a particular feature in a precise, context-dependent way, and there are fewer or no other proxies of that feature available in the data. The effect is diminished in settings where a variety of local features contribute to the model’s prediction.

Why do predictions diverge more on longer key depths?

While we don’t have a precise mechanistic explanation for this result, we think our observations suggest some hypotheses (in particular, see Appendix Figure 10 in our revision). Specifically, the success of the Transformer depends on the mechanism for representing nesting depth. To succeed on the depth generalization split, this mechanism must also extrapolate to unseen depths. Figure 10 suggests that the original model does extrapolate to some extent, but the simplified models fail to fully capture the mechanism for extrapolating to deeper depths. This could lead to two kinds of error. First, even if the query token appears at a seen depth, there might be deeper key tokens earlier in the sequence, and if these keys are not represented well, they might act as “distractors” for shallower queries. Second, if the query token appears at an unseen depth, the model might fail to match it to the correct bracket if the representations for these depths are not well-separated. In the error patterns in Figure 4 (Figure 6 in our original draft) we see evidence that both types of errors occur, and the divergences are greater when the query token is at an unseen depth.

Are these results consistent across training trajectories?

In the appendix, we have added a figure showing the error patterns for another model trained with a different random initialization (Figure 11). The general trends are similar, with predictions diverging more on queries appearing at unseen depths. We would also note that, in our additional experiments on code completion, we reported the results from models trained with multiple random initializations and found that different training runs lead to similar trends in terms of generalization gaps.

Description of Figure 1 (accuracy on depth generalization)

Thank you for pointing this out! This was an oversight and we have corrected the description of the figure in our revision.

I find quite concerning that the reviewer is proposing a rejection based on its own lack of understanding. If you do not feel qualified to review a paper, you should communicate it to the AC and ask to be dispensed.

Thank you for your feedback.

The analysis is not paired with proposed improvements or solutions

We would like to emphasize that our aim in this paper isn’t to present new methods for analyzing Transformer language models. Our goal is to better understand the limitations of interpretability methods that are commonly used, which try to understand models by way of simplified representations. It is important to note that practitioners may use this approach without recognizing the implicit assumption that “the simplified and original model respond similarly to distribution shifts.” Our main point is to illustrate how simplified proxy models may be faithful to the original model on in-distribution data but fail to explain the model’s behavior on out-of-distribution data, giving rise to interpretability illusions. We conduct further analysis to characterize how the faithfulness of these explanations is influenced by the choice of simplification, the underlying task, the required features for solving the problem, the form of distribution shift, and the model's performance. Our findings also provide a counterpoint to a common assumption about the relationship between simplification and generalization in machine learning, highlighting settings where simplified versions of a model incur a larger performance drop on tests of systematic generalization.

We agree with all reviewers (iDw5, oEsb, H66k, 5xMH) that this is an important question. We argue that extensive analyses are needed to disentangle the role of these variables, investigate and explain the underlying reasons, and effectively convey the message. We believe this is an important contribution, and the best way to share these results with the community is through an “insight and analysis paper”. While a “method paper” could be a valuable future step, it serves a different purpose and may not be the most suitable platform for sharing this particular insight.

What conclusions can be drawn from Deep Learning practitioners? Are transformers reliable for learning languages?

The main takeaways for practitioners are to be aware of limitations to an approach to model interpretation: simplified proxy models might be faithful to a model on some distributions but fail to predict how the model will generalize to different distributions. Our analysis is not focused on the question of whether or not Transformers are reliable for learning languages (although we think our findings might provide some information about the extent to which we can expect Transformers to generalize to different kinds of distribution shifts, such as to deeper hierarchical structures).

Does the analysis extend also to similar models or other datasets?

As noted in the general response, we have extended our experiments to incorporate an additional setting: language modeling on code datasets, training models with more layers and attention heads. These experiments show that generalization gaps also occur in more naturalistic settings and provide some additional insight into the relationship between generalization gaps and properties of the task. To summarize:

• We trained character-level language models on Java functions and evaluated generalization to functions with deeper nesting depths and unseen programming languages.
• These experiments revealed a consistent generalization gap: simplified proxy models are more faithful to the original model on in-distribution data and worse on OOD data.
• We also broke down these results according to different subsets of the code completion task, and found that the generalization gap is notably larger on some types of predictions. In particular, the generalization gap is larger on subtasks that can be seen as more “algorithmic”, including predicting closing brackets and copying variable names from earlier in the sequence. The gap is smaller when predicting whitespace characters and keywords, perhaps because these predictions rely more on local context than precise attention.
• We were able to attribute some of these generalization gaps to an “induction head” mechanism, which copies words from earlier in the sequence. Simplifying this head leads to large gaps in approximation quality between in-domain and out-of-domain samples, meaning that the low-dimensional approximation underestimates the extent to which the induction head mechanism will generalize to unseen distributions.

Overall, these experiments provide evidence that generalization gaps also occur in naturalistic settings, and they provide additional insight into how these gaps are related to different properties of sequence modeling tasks. Specifically, generalization gaps might be most pronounced in “algorithmic” settings where the model must use a particular feature in a precise, context-dependent way, and there are fewer or no other proxies of that feature available in the data. The effect is diminished in settings where a variety of local features contribute to the model’s prediction.

Thank you for your helpful comments! We are glad you find that our paper addresses an important question about the faithfulness of explanations. And thank you for pointing out areas where the presentation of the paper can be improved. We respond to your questions below and have incorporated your suggestions into our updated draft.

Introducing the Dyck languages

Thank you for the suggestion! We have revised section 2.1 and Appendix A.1 to present the Dyck languages more clearly. In particular, please see Table 1 in the appendix, which provides toy examples illustrating the different generalization sets.

The phrase “simplified models” can be misleading

Thank you for pointing this out. We agree that the phrase “simplified models” could be more precise. In this paper, we simplify the representations of a model, using tools like dimensionality reduction, clustering, and discretization, and then analyze these simplified representations. In other words, we are essentially replacing the original model with a simplified proxy model which uses a more limited—and thus easier to interpret—set of features. We have revised our draft to communicate this more precisely in the introduction.

Which of the testing datasets is considered in-distribution / IID vs. Unseen struct

We apologize for any confusion in our presentation of the different testing splits. We constructed these splits to test how the model generalized to different kinds of unseen data. In each split, some aspect of the data can be considered “out-of-distribution” and some can be considered “in-distribution”–for example, in the “Unseen structure” datasets, the structures were not seen during training, but the nesting depth was seen.

Regarding the IID split: The subset was constructed by sampling from the same distribution used to construct the training set but rejecting any strings that appeared in the training set. It turns out that almost all sequences in this subset have unseen structures. (The number of bracket structures at length $n$ is given by the n-th Catalan number, so the chance of sampling the same bracket structure twice is very low.) We have updated Section 2.1 to clarify the sampling strategy of different generalization sets. We have updated Figures 1-3 to focus on generalization splits more clearly. Additionally, we have added more details to the Appendix A.1 to further clarify the IID split, and positioning of our experiments with respect to the prior work by Murty et al. [1], who introduced the structural generalization split we adopt here.

[1] Murty et al., 2023. Grokking of Hierarchical Structure in Vanilla Transformers.

How to interpret the results in Figure 4. How can one conclude that there is a generalization gap between the simplified model and the original model? How can one conclude that the simplified model underestimate/overestimate the generalization ability of the original model? Is it more appropriate to compare the prediction accuracy of the simplified model with that of the original model on both in-distribution and out-of-distribution data?

In Figure 4 (Figure 2 in the revision), we concluded that there is a generalization gap because the approximation quality metrics are better on in-distribution data (e.g. comparing seen structures vs. unseen structures, or seen depths vs. unseen depths). This implies that the simplified models are better approximations of the original model on the in-distribution data than the out-of-distribution data. We concluded that the simplified model underestimates generalization because the original model gets near-perfect accuracy on all generalization splits (with the exception of depth generalization). Therefore, the fact that the simplified model makes a different prediction than the original model implies that it makes the wrong prediction, which would lead to an underestimate of the generalization ability of the original model. To conclude that one-hot attention overestimates generalization, we do compare the prediction accuracy in Figure 5b (Figure 3b in the revision): the one-hot attention model achieves higher accuracy on the depth generalization split than the original model.

Thank you for your helpful feedback! We are glad that you find our results to be interesting, surprising, and relevant, and you find our paper well-written and clear.

Limited scope

As noted in the general response, we have extended our experiments to incorporate an additional setting — language modeling on code datasets. These experiments show that generalization gaps also occur in more naturalistic settings and provide some additional insight into the relationship between generalization gaps and properties of the task. To summarize:

• We trained character-level language models on Java functions and evaluated generalization to functions with deeper nesting depths and unseen programming languages.
• These experiments revealed a consistent generalization gap: simplified proxy models are more faithful to the original model on in-distribution data and worse on OOD data.
• We also broke down these results according to different subsets of the code completion task, and found that the generalization gap is notably larger on some types of predictions. In particular, the generalization gap is larger on subtasks that can be seen as more “algorithmic”, including predicting closing brackets and copying variable names from earlier in the sequence. The gap is smaller when predicting whitespace characters and keywords, perhaps because these predictions rely more on local context than precise attention.
• We were able to attribute some of these generalization gaps to an “induction head” mechanism, which copies words from earlier in the sequence. Simplifying this head leads to large gaps in approximation quality between in-domain and out-of-domain samples, meaning that the low-dimensional approximation underestimates the extent to which the induction head mechanism will generalize to unseen distributions.

Overall, these experiments provide evidence that generalization gaps also occur in naturalistic settings, and they provide additional insight into how these gaps are related to different properties of sequence modeling tasks. Specifically, generalization gaps might be most pronounced in “algorithmic” settings where the model must use a particular feature in a precise, context-dependent way, and there are fewer or no other proxies of that feature available in the data. The effect is diminished in settings where a variety of local features contribute to the model’s prediction.

Why does the simplified model generalize less well to the depth generalization split?

We don’t have a precise mechanistic explanation for this result but we think our observations suggest some hypotheses (in particular, see Appendix Figure 10 in our revision). Specifically, the success of the Transformer depends on the mechanism for representing nesting depth. Each depth has to be represented as a separate direction, to be used in the second attention layer. To succeed on the structural generalization split, the simplified model just needs to capture the in-distribution behavior of this mechanism–embedding seen depths. But to succeed on the depth generalization split, this mechanism must also extrapolate to unseen depths. Figure 10 suggests that the model does extrapolate to some extent, but the simplified models fail to fully capture this behavior. For example, perhaps the model uses additional directions for encoding deeper depths, but these directions are dropped when we fit SVD on data containing only lower depths. We have updated Appendix B.3 to include this discussion.

What if we don’t prepend a START token?

In preliminary experiments, we found that the model learned a similar attention pattern even when we didn’t prepend a START token – that is, the first layer attention head attended uniformly to all positions but placed higher attention on the first token in the sequence. This pattern resembles a construction from [1]: the first token is used to absorb extra attention so that the normalizing constant in the attention output is consistent across sequence lengths. We have noted this in Appendix B.1 in our revision.

[1] Liu et al., 2022. Transformers Learn Shortcuts to Automata.

What extra information is encoded in the value embeddings?

From plotting the value embeddings (Figure 7c in the revision), we can see that the first components of the value embeddings encode position. The value embeddings also encode bracket type – each cluster of value embeddings corresponds to either opening or closing brackets for a single bracket type. In contrast, in the human-written Transformer algorithm, the value embeddings only encode whether the bracket is an opening or closing bracket and are invariant to position and bracket type. We have updated section B.1 to explain this point.