The Heuristic Core: Understanding Subnetwork Generalization in Pretrained Language Models (2024)

Abstract: Prior work has found that pretrained language models (LMs) fine-tuned with different random seeds can achieve similar in-domain performance but generalize differently on tests of syntactic generalization. In this work, we show that, even within a single model, we can find multiple subnetworks that perform similarly in-domain, but generalize vastly differently. To better understand these phenomena, we investigate if they can be understood in terms of “competing subnetworks”: the model initially represents a variety of distinct algorithms, corresponding to different subnetworks, and generalization occurs when it ultimately converges to one. This explanation has been used to account for generalization in simple algorithmic tasks (“grokking”). Instead of finding competing subnetworks, we find that all subnetworks—whether they generalize or not—share a set of attention heads, which we refer to as the _heuristic core_. Further analysis suggests that these attention heads emerge early in training and compute shallow, non-generalizing features. The model learns to generalize by incorporating additional attention heads, which depend on the outputs of the “heuristic” heads to compute higher-level features. Overall, our results offer a more detailed picture of the mechanisms for syntactic generalization in pre-trained LMs.

## introduction a central question in machine learning is to understand how models generalize from data that supports different possible solutions. @xcite investigated this question in the context of pretrained language models (lms) like bert @xcite trained on natural language inference (nli; @xcite , showing that models trained with different random seeds perform very similarly on in-domain (id) evaluation id ood id ood id ood id ood id ood id ood full model competing subnetworks heuristic core subnetwork #1 subnetwork #2 subnetwork #1 subnetwork #2 heuristic core accuracy figure 1 : we find different subnetworks in a pretrained lm that achieve similar in-domain performance but generalize differently. prior work has explained similar generalization phenomena in synthetic tasks in terms of distinct subnetworks that compete during training. we instead find evidence of a heuristic core: a set of attention heads that appear in all generalizing subnetworks but, on their own, do not generalize. in this work, we use subnetwork analysis to better understand how pretrained lms generalize on nlp tasks. specifically, we use structured pruning @xcite to isolate various subnetworks-subsets of attention heads and mlp layers-that approximate the behavior of the full model. our main finding is, that even within a single model, there exist multiple subnetworks that all match the model's performance closely on in-domain evaluation sets, but exhibit vastly different generalization. this result has meaningful consequences for practitioners-for example, it underscores the importance of ood evaluation of pruning methods. it also raises several questions about the mechanisms underlying generalization, which we proceed to investigate. one possible explanation for our results is that the model initially consists of a variety of disjoint subnetworks, corresponding to distinct possible solutions for the task; the model ultimately generalizes if it converges to a subnetwork that generalizes. this model of generalization has been used to explain the "grokking" phenomenon on toy algorithmic tasks, whereby a neural network initially overfits a training set-but, after continued training, suddenly undergoes a phase shift to perfect generalization @xcite @xcite . we refer to this as the competing subnetworks explanation, following @xcite . our setting bears a similarity to grokking, originally documented by @xcite : models converge early on in-domain evaluations and only start to generalize after training for additional epochs. therefore, we test if generalization in this case also arises from competition between disjoint subnetworks. first, we check if these behaviorally different subnetworks consist of disjoint subsets of model components (attention heads). to the contrary, we find a small set of nine attention heads that occur in all subnetworks-even the ones that do not generalize at all. furthermore, we find that this same set of attention heads also appears when we prune earlier checkpoints of the model before it starts to generalize. in the grokking setting, generalization is marked by a decrease in effective size-defined as the size of the smallest subnetwork that matches the full network's performance on the generalization sets-as the model switches from a dense, memorizing subnetwork to a sparse, generalizing subnetwork. in contrast, we find that generalization in our case is accompanied by a sharp increase in effective size. together with the set of common attention heads, our results suggest that, instead of selecting between competing subnetworks, the model initially learns a "core" of attention heads that implement simple heuristics. the model ultimately generalizes by learning additional attention heads that interact with these heads. we refer to this core set of attention heads as the heuristic core, and conduct further analysis to better understand what roles they play in the model. we find that these attention heads are associated with entailment heuristics-namely, attending to words that are repeated across sentencessupporting our characterization of these components as heuristic core. on the other hand, ablating these attention heads from the original model leads to a dramatic decrease in performance on the generalization set. this result further supports the idea that generalization arises from additional components building off simple features extracted by heuristic components. at intermediate sparsity foot_1 levels, subnetworks contain different "counterheuristic" components, leading to different degrees of partial generalization to different subcases. these results have important practical implications. for instance, they suggest that we cannot make the model more robust by ablating the "heuristic" components. more broadly, our experiments paint a more detailed picture of the mechanisms underlying generalization in natural language understanding tasks. in doing so, they highlight interesting phenomena and open new avenues to future research into language models' internals. ## problem setup we focus on two sentence-pair classification tasks. mnli @xcite is a natural language inference (nli) dataset, which involves predicting whether a premise sentence logically entails a hypothesis sentence. qqp @xcite is a paraphrase identification dataset. prior work has noted that these datasets admit "heuristics": solutions that achieve high in-domain accuracy but do not generalize well. in particular, @xcite introduced the hans dataset, which tests whether nli models use specific heuristics, such as lexical overlap (predicting that sentences entail each other if they share many words); see appendix ## comparing subnetwork generalization we start by examining how different subnetworks generalize on mnli/hans and paws/qqp. specifically, we compare the generalization behavior of subnetworks with the same sparsity level, but pruned with different seeds; and we compare subnetworks pruned with different sparsity levels. we prune the bert models fine-tuned on mnli and qqp with 12 random seeds for different target sparsities (more details in appendix a). the resulting subnetworks are evaluated on the id (mnli and qqp validation sets) and ood (hans subcases, and paws-qqp's adversarial subset) evaluation sets. in our runs, we observe that the mlp layers are never pruned (although they are allowed to be). mnli (matched) mnli (mis.) so-swap [lo] prep [lo] embed-if [c] embed-verb [c] 0 20 40 60 80 100 accuracy (%) in domain out of domain (a) mnli, 50% mnli (matched) mnli (mis.) so-swap [lo] prep [lo] embed-if [c] embed-verb [c] in domain out of domain (b) mnli, 70% qqp paws-qqp (only adversarial) in domain out of domain (c) qqp, 30% qqp paws-qqp (only adversarial) in domain out of domain (d) qqp, 60% figure 2 : pruning a bert model with different random seeds results in subnetworks that perform similarly indomain but generalize differently. the dots refer to the accuracy of the pruned subnetworks, while solid lines indicate full model performance. mnli/hans: at 50% sparsity, the subnetworks perform within 3% of the model on mnli but show varying generalization. at 70% sparsity, the subnetworks behave as pure heuristics despite respectable mnli accuracy. the trend also holds for qqp/paws, with sparsities of 30% and 60%. figure 12 in appendix c shows the plot for all subcases of hans. full model #1 #2 #3 #4 #5 #6 #7 #8 #9 #10 #11 #12 0 20 40 60 80 accuracy (%) so-swap [lo] prep [lo] embed-if [c] embed-verb [c] figure 3: different subnetworks at 50% sparsity (found by pruning with different random seeds) generalize partially to different subcases of hans. subnetwork #5 generalizes to the constituent subcases, whereas #8 generalizes to the lexical overlap subcases. subnetwork #2 does well on the embed-if subcase of constituent, but not embed-verb. ## competing subnetworks hypothesis in the previous section, we observed that different subnetworks perform similarly to the full model in-domain but generalize differently. moreover, no sparse subnetwork generalizes as well as the full model. we now seek to understand how these subnetworks emerge during training and interact to give rise to generalization. specifically, we explore if our findings can be explained in terms of competing subnetworks: the model initially represents multiple solutions as dif-ferent subnetworks, and it ultimately converges to one. whether it generalizes is dependent on which of the subnetworks it converges to. this hypothesis has been found to explain the curious phenomenon of grokking @xcite . additionally, as @xcite note, the id performance of our bert models saturates early in training (figure 6 ) but the ood accuracy starts to rise much later. motivated by the similarity to grokking, we test if generalization in our case also arises from competition between disjoint subnetworks. we consider the competing subnetwork hypothesis to consist of two main predictions, mirroring @xcite . first, the model can be composed into subnetworks that are disjoint, representing distinct algorithms consistent with the training data. second, a rise in generalization is accompanied by a reduction in the model's effective sizethe smallest subnetwork whose accuracy matches the model's within a given threshold (e.g., 3%) on a target dataset-which occurs when the model switches from a dense, non-generalizing subnetwork to a sparse, generalizing subnetwork. all subnetworks share a common set of components. we start by measuring how many model components are shared between different subnetworks. in particular, we compute the frequencies of each attention head in the 50% (partially generalizing) subnetworks and in the 70% (nongeneralizing) subnetworks of the mnli model. these frequencies are visualized in figure 5 (a corresponding plot for qqp is presented in figure 11 of appendix c). the spearman's rho correlation between the two frequencies if 0.82 (p-@xmath0.6 • 10 -36 ), i.e., very strong agreement. (the corresponding value for qqp is 0.74, corresponding to strong agreement; see figure 14 in appendix c.) in particular, we observe that nine attention heads appear in all subnetworks, even the sparsest onesthat is, the subnetworks that behave most consistently with the heuristics. in other words, instead of finding disjoint subnetworks, we find a high degree of overlap (the expected intersection of 12 random 70% sparsity subnetworks is 7.7 • 10 - foot_2 heads) between the subnetworks that partially generalize and the subnetworks that do not generalize at all. effective size increases as the model generalizes. next, we investigate how these subnetworks emerge over training. we conduct a similar analysis to @xcite and plot the effective size of the model, defined as the size of the smallest subnetwork that is faithful to the model on a particular evaluation dataset d. we say that a subnetwork is faithful to a model on dataset d if their accuracies are within 3%. we report faithfulness using different choices of d-either the in-domain validation set, the out-of-domain validation sets, or the equally weighted combination of the two (mixed data). we prune each checkpoint to sparsities of multiples of 5% and pick the highest one that yields a faithful subnetwork. the reported accuracies at each sparsity are averages over three seeds. the results are in figure 6 . the effective size needed to approximate the model on the in-domain evaluation remains relatively flat over the course of training. however, on out-of-domain and mixed data, generalization is accompanied by an increase in effective size. we also examine the set of attention heads that appear in the mnli subnetworks at the checkpoint immediately before the model starts to generalize. we find that the same nine attention heads that appear in all subnetworks at the end of training appear in these subnetworks as well. this suggests that, rather than switching between competing subnetworks, the model first learns a set of attention heads that compute simple, non-generalizing features, and then generalizes by incorporating more components into this core. the president advised the doctor . the president advised the doctor . figure 7 : many attention heads common to the mnli subnetworks attend to tokens co-occurring in the premise and context. since the order of, e.g., "doctor" and "president" is reversed, the heads seem to identify lexical overlap. in subnetworks when we prune checkpoints early in training, before the model generalizes. we refer to these nine attention heads as the heuristic core of the model. in this section, we inspect these attention heads in more detail, to understand what features they compute, and how they interact with the rest of the model to enable generalization. how do the heuristic core attention heads behave? first, we look at the attention patterns exhib-ited by these attention heads. we inspect example attention patterns and observe that most heads exhibit a simple attention pattern, and we quantify these in figure 9 . four out of nine heads (mnli) attend to tokens that repeat across the premise and hypothesis, suggesting that they extract token overlap features. figure 7 shows example attention patterns. notably, we find that eight attention heads also appear in all qqp subnetworks. figure 8 presents an example indicating that these heads play a similar role for qqp, attending between repeated words. other low-layer heads attend to the previous or subsequent token. the heads in the higher layers are more difficult to characterize but often attend to the special separator token, perhaps aggregating information from across the sequence. overall, these findings are consistent with the notion that these attention heads calculate simple, shallow features. how do the heuristic core attention heads interact with the rest of the model? next, we ablate each heuristic core attention head from the full model and observe how the in-domain and out-ofdomain accuracies change, in ## conclusion we showed that pruning a pretrained lm with different random seeds can yield subnetworks that perform similarly to the original model in-domain but generalize differently. moreover, sparser subnetworks generalize worse. to explain this result, we first considered the competing subnetworks hypothesis, which has been used to explain grokking on algorithmic tasks; it suggests that the model consists of disjoint subnetworks representing different solutions, some which generalize and some which do not. we found that our results are instead more consistent with the existence of a heuristic core: a subset of attention heads that appear in all subnetworks and calculate shallow, non-generalizing features. the model generalizes by incorporating additional heads that interact with the heuristic core. our results have practical implications about the effect of pruning on generalization and raise intriguing questions about how language models learn to generalize by composing simple components.