Paper: Interpretability Illusions in the Generalization of Simplified Models (ICLR-2024)

Reviewer comment (Experiments):
The focus of the experimental analysis is too narrow. The introduction section does a good job in outlining the research goals, but this aim is then overly specialized to a specific class of models trained on a toy problem.

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

Reviewer comment (Experiments):
Can you include the analysis of other datasets and architectures? I do appreciate your work on toy-languages. However, I would like to see these insights to be applied to models trained on naturalistic data.

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

Reviewer comment (Evaluation):
In Figure 6, it seems like for the longer key depths, the predictions diverge more. Would the authors be able to emphasise more whether it is always the case and if there are any solid reasons behind this particular behaviour?

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

Reviewer comment (Evaluation):
The analysis is not paired with proposed improvements or solutions, which I find to be a weakness of the paper.

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

Reviewer comment (Writing):
The phrase “simplified models” can be misleading in this paper’s context. I was thinking of methods such as knowledge distillation or network pruning when I first see the phrase “simplified models”. However, what the paper mainly focuses on are dimensionality reduction and clustering methods. It is encouraged to use a more precise word other than “simplified”.

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