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 indomain, 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 pretrained LMs.1

heading: 1 Introduction
text: A central question in machine learning is to understand how models generalize from data that supports different possible solutions. McCoy et al. (2020) investigated this question in the context of pretrained language models (LMs) like BERT (Devlin et al., 2019) trained on natural language inference (NLI; Williams et al., 2018), showing that models trained with different random seeds perform very similarly on in-domain (ID) evaluation
1We release our code publicly at https://github.com/ princeton-nlp/Heuristic-Core.
sets but generalize very differently to adversarial out-of-domain (OOD) evaluation sets. However, we have a limited understanding of how any individual model learns to generalize.
In this work, we use subnetwork analysis to better understand how pretrained LMs generalize on NLP tasks. Specifically, we use structured pruning (Wang et al., 2020; Xia et al., 2022) 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
14351
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 (Barak et al., 2022; Nanda et al., 2023; Varma et al., 2023). We refer to this as the Competing Subnetworks explanation, following Merrill et al. (2023). Our setting bears a similarity to grokking, originally documented by Tu et al. (2020): 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 sentences— supporting 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 sparsity2 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.

heading: 2 Problem Setup
text: We focus on two sentence-pair classification tasks. MNLI (Williams et al., 2018) is a natural language inference (NLI) dataset, which involves predicting whether a premise sentence logically entails a hypothesis sentence. QQP (Iyer et al., 2017) 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, McCoy et al. (2019) 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 Table 1 for examples and the list of abbreviations used for HANS subcases. Each subcase of HANS has 1000 validation examples. Similarly, the QQPPAWS dataset (Zhang et al., 2019) tests whether paraphrase identification models are susceptible to a word-overlap heuristic (Appendix Table 2). PAWS-QQP has 12, 663 examples in all, of which 8, 696 bear the verdict “not equivalent”. We work with the latter subset in this paper.
Models. We mainly focus on BERT (Devlin et al., 2019) throughout the paper, but our findings also generalize to RoBERTa (Liu et al., 2019) and GPT2 (Radford et al., 2019) (Section 6). We fine-tune BERT models on MNLI and QQP (see Appendix A for more training details). Our results reproduce the observations of Tu et al. (2020): in-domain accuracy saturates early, and out-of-domain accuracy increases much later. We focus on the four subcases in Table 1 for HANS, as the model gener-
2In this paper, sparsity refers to the fraction (or percentage) of attention heads and MLPs that have been pruned away.
alizes well on them. Other subcases are discussed in Appendix C. The effect is smaller, but still consistent, on PAWS-QQP (Table 2).
Finding subnetworks via pruning. For a Transformer model, we define a subnetwork as a subset of the attention heads and MLP (feed-forward) layers. An important motif in our analysis will be the search for subnetworks that preserve model performance. We choose structured pruning towards this end. While circuit-finding methods (Conmy et al., 2023) would also be useful, pruning is faster and allows injecting randomness more readily. We use structured rather than unstructured pruning for two reasons. First, unstructured pruning can be expressive enough to find subnetworks with completely different behavior from the original model (Wen et al., 2023). Second, structured pruning lets us compare subnetworks in terms of semantically meaningful components: attention heads.
We prune via optimization of a binary mask, where 0 and 1 indicate dropping and retention, respectively. Our runs prune attention heads and entire MLP layers, each corresponding to one bit in this mask. We use mean ablation per recommendations from work in Mechanistic Interpretability (Zhang and Nanda, 2024)—pruned layers’ (or heads’) activations are replaced by the mean activation over the training data. Following CoFi Pruning (Xia et al., 2022), we relax binary masks to floats
in [0, 1] and then perform gradient descent to optimize them. The objective optimized is the KL loss between the predictions of the full model and those of the pruned subnetwork. Finally, the mask entries are discretized to {0, 1} based on a threshold. L0 regularization is used to enforce a target sparsity. We adopt Louizos et al. (2018)’s recipe to model the masks and furnish the details in Appendix B. More details and hyperparameters, along with a more detailed discussion of the method, are provided in Appendix A. We freeze the model after fine-tuning and only optimize the pruning masks to preserve faithfulness to the full model.

heading: 3 Comparing Subnetwork Generalization
text: 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).
Different subnetworks generalize differently. Figures 2a and 2c plot the results for a target sparsity of 50% and 30% in the case of MNLI and QQP, respectively. The pruned models perform close to each other in-domain, and are all within 2- 3% of the full model’s accuracy on both the MNLI and QQP validation splits. However, their accuracies on the out-of-domain evaluation splits varies widely. In the most extreme case, the accuracies on the Embed-If (Constituent) subcase of HANS vary from just 4.1% to 50.2%. Moreover, different subnetworks generalize to different subcases, as Figure 3 illustrates. Subnetwork #8 achieves 41.4% at SO-Swap (Lexical Overlap) but only 18.0% at Embed-If (Constituent). Subnetwork #5, on the other hand, achieves only 19.5% at the former but
50.2% at the latter. Thus, pruning different parts of the model seems to sacrifice generalization on the OOD subcases to varying degrees. The results in the case of QQP and PAWS-QQP are similar.
Sparser subnetworks generalize worse. We also observe that sparser subnetworks consistently generalize worse. First, we prune the model to a higher sparsity (70% for MNLI and 60% for QQP) with 12 random seeds. As Figures 2b and 2d show, virtually all subnetworks at this sparsity show a complete lack of generalization on almost every OOD subcase. The ID accuracy stays at a respectable 74- 76%, indicating that a large portion of model performance is explained by these sparse subnetworks. In Figure 4, we plot the generalization accuracy of subnetworks pruned at different sparsity levels,
500 12500 24500 36500 48500 Steps
0
20
40
60
80
100
Ef fe
ct iv
e Si
ze (%
)
ID + OOD ID OOD
0 10 20 30 40 50 60 70 Sparsity (%)
0
20
40
60
80 100 Av g. A cc ur ac y (% )
(a) MNLI/HANS
0 10 20 30 40 50 60 Sparsity (%)
0
20
40
60
80
(b) QQP/PAWS-QQP
Figure 4: The OOD accuracy decreases fairly smoothly with sparsity (3 seeds). The drop in ID accuracy is slow and has low variance. We find no subnetworks sparser than 30% generalizing as well as the full model on either dataset.
0 1 2 3 4 5 6 7 8 9 10 11 12 Frequency, 70% subnetworks
12 11 10 9 8 7 6 5 4 3 2 1
0Fr eq
ue nc
y, 5
0% su
bn et
wo rk
s 1 2 1 1 4 9 3 1 1 2 2 3
1 1 1 2 1 1 1 1 1 1 1 1 2 7 1 1 1 1 1 1 1 1 5 2 2 1 1 4 1 10 1 2 1 8 14 32 0 5 10 15 20 25
30
Figure 5: This heatmap quantifies the frequencies of attention heads in the 50% and 70% sparsity subnetworks (MNLI). Entry (i, j) corresponds to the number of heads appearing in i/12 50% subnetworks and j/12 70% subnetworks. In particular, we note that nine attention heads appear in all of the subnetworks.
for 3 random seeds. The generalization accuracy generally drops steadily at higher sparsity levels and is devoid of any drastic phase transitions.

heading: 4 Competing Subnetworks Hypothesis
text: 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 (Merrill et al., 2023). Additionally, as Tu et al. (2020) 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 Merrill et al. (2023). 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 size— the 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-value = 1.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 ones— that 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−5 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 Merrill et al. (2023) 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.

heading: 5 Understanding the Heuristic Core
text: Instead of finding evidence for distinct, competing subnetworks, we have found that all subnetworks share a set of common components. In particular, nine attention heads appear in all MNLI subnetworks, including subnetworks that do not generalize at all. We find the same nine attention heads
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 Table 3. None of the heads’ ablation affects in-domain accuracy greatly.
Surprisingly, however, many heads seem to be critical to performing well on HANS: for example, ablating a single head (L4.H0) reduces performance on SO-Swap by more than 40%. Moreover, the drops observed on ablating these heads is superadditive: ablating just three of them leads to the OOD accuracy plummeting. This indicates that heads outside the heuristic core rely on the heuristic core to implement more complex behavior.
Table 4 paints a similar picture for QQP. Perhaps most interestingly, the heads 2.5 and 4.0 that had the maximum ablation effect for MNLI also stand out here. Along with the fact that 5 of these 8 heads also exist in the MNLI model’s core, this highlights the possibility that these heads already extract word overlap features in the BERT model prior to finetuning. We leave this question for future work.

heading: 6 Findings Generalize to RoBERTa and GPT-2
text: We show in Appendix D that our results extend to RoBERTa (Liu et al., 2019) models. Appendix E finds that they also transfer to GPT-2 (Radford et al., 2019) models, albeit with some novel behavior. Specifically, we observe that the accuracy on the non-adversarial OOD (‘entailment’) cases slightly drops late into training, and is accompanied by a mild increase in effective size. The performance on these cases also varies across different subnetworks of the same sparsity. Explaining these additional features requires further research and makes for future work. We also note that it is possible that models much larger (such as those with 7B+ parameters) take a different route to generalization. They might thus show qualitatively different behavior, and we urge caution in applying the results of this paper to them.

heading: 7 Related Work
text: Pruning. Pruning (LeCun et al., 1989) removes layers or parameters from a network for space efficiency and speed-ups. Structured Pruning (Wang et al., 2020; Ma et al., 2023; Xia et al., 2022, 2024) preserves some structure in the model, e.g. equal fraction of surviving parameters in each attention head, accepting lower compression in exchange for higher speed-ups. Unstructured pruning (Sun et al., 2024) approaches like Magnitude Pruning (LeCun et al., 1989; Hassibi and Stork, 1992) achieve better compression at the cost of lower potential speedup. Proposals such as Layer dropping (Zhang and He,
2020) and Block Pruning (Lagunas et al., 2021) have explored pruning at coarser granularities.
Pruning for interpretability. Mechanistic Interpretability seeks to understand models through their components, e.g., neurons, attention heads, and MLPs. Such analysis uncovers intricate interactions between different atomic components, which form circuits (Elhage et al., 2021; Olah et al., 2020) implementing various tasks. Pruning has occasionally been used as part of a larger interpretability effort (Lepori et al., 2023a), though mostly in an ad-hoc manner. Jain et al. (2023) use pruning and probing techniques to show that fine-tuning learns a thin “wrapper” around model capabilities. Lepori et al. (2023b) use pruning to identify subnetworks which they then show form modular building blocks of model behavior. In a similar vein, (De Cao et al., 2020) learn differentiable masks over the input to investigate model behavior. While most of these approaches replace missing layers with an identity function, work in Mechanistic Interpetability (Zhang and Nanda, 2024; Conmy et al., 2023) has shown that adding their mean activations might preserve faithfulness better. Prior work has examined attention patterns to decipher model behavior (e.g. Clark et al., 2019), although attention patterns have also been shown to be misleading in some cases (Jain and Wallace, 2019).
Generalization. The question of how a neural network generalizes has remained a compelling research direction for the NLP community (Yang et al., 2023). McCoy et al. (2020) find that fine-tuning BERT with different random seeds leads to disparate generalization despite similar in-domain accuracies. Other works have identified tasks where OOD accuracies increase long after in-domain accuracy saturates and attributed them to hard minority examples (Tu et al., 2020; Bhargava et al., 2021). Recently, a phenomenon called Grokking (Power et al., 2022) was discovered, where test accuracy on a toy algorithmic task increases abruptly long after overfitting to the training data. It has been explained in terms of competing subnetworks (Merrill et al., 2023; Varma et al., 2023). Grokking-like phenomena have also been identified in other settings (Liu et al., 2023) and for small transformers (Murty et al., 2023) on hierarchical depth generalization. Simultaneously, questions about memorizing in pretrained LMs have been asked in other settings (Tänzer et al., 2022). Friedman et al. (2024) document cases where sim-
plifying a Transformer’s representations preserves in-domain performance but impairs systematic generalization; we discover a similar “generalization gap” for another simplification method, pruning.
Subnetworks as modular components. An emerging body of work (Lepori et al., 2023b; Choenni et al., 2023; Hupkes et al., 2020) suggests that neural networks and language models may consist of modular subnetworks each performing different subtasks. Bayazit et al. (2023) find a subnetwork in GPT-2 responsible for most of its factual knowledge – removing it does not lead to loss of performance on non-knowledge-intensive tasks. Li et al. (2023) report the possibility of ablating away undesired behavior in a model (e.g. toxicity), while Wei et al. (2024) identify regions that when pruned away compromise the safety alignment of a Language Model. In concurrent work, Zhang et al. (2024) find a subset of less than 1% of an LM’s parameters vital for linguistic competence across 30 languages—which they call the “Linguistic Core”. Prakash et al. (2024) report that entitytracking circuits in fine-tuned models exist even before fine-tuning, providing a potential explanation for why the heuristic core heads we found for MNLI and QQP had a substantial overlap. Panigrahi et al. (2023) find that grafting ∼ 0.01% of model parameters from a fine-tuned version can recover most of the gains from fine-tuning.

heading: 8 Conclusion
text: 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.
Limitations
Our approaches are not without limitations. Our analysis is based on a top-down approach of identifying subnetworks. Thus, we do not look for or discover “circuits” that are atomically interpretable. Our experiments only use text classification tasks, although we believe that such a task is ideal due to clear definitions of heuristics. Additionally, we only work with BERT models—it is possible that larger transformers, or those pretrained differently, exhibit different behavior. Finally, although we have analyzed the mechanisms of generalization in this paper, the question of why they are undertaken remains elusive for the moment. Our experiments also raised many novel questions, which require further research to answer.
Ethical Considerations
We do not foresee any ethical considerations raised due to our work. We work with pretrained language models, which can at times hallucinate or demonstrate other undesirable behavior. Our work, however, does not interact with these facets of model behavior. Additionally, we only work with English datasets in our paper.

heading: Acknowledgements
text: We thank Michael Tang for his detailed suggestions and alacritous insights, without which this project would be incomplete. We are grateful to Alexander Wettig and Zirui Wang for their feedback on paper drafts. Finally, we thank everyone at the Princeton NLP group—especially Mengzhou Xia and Tianyu Gao, for fruitful discussions regarding our experiments. This research is funded by the National Science Foundation (IIS-2211779), a Sloan Research Fellowship, and a Hisashi and Masae Kobayashi *67 Fellowship.

heading: A Model and Training Details
text: BERT Architecture Our experiments use finetuned BERT models (Devlin et al., 2019). In the simplest setting, a Transformer Model (Vaswani et al., 2017) consists of L blocks, where each block is in turn composed of a Multi-Head Attention layer followed by an MLP layer. Given an input X , the Attention Layer outputs
X +
Nh∑
i=1
Attn(W (i)Q ,W (i) K ,W (i) V , X)W (i) O
where Nh is the number of attention heads, and W
(i) K , W (i) Q and W (i) V denote the Key, Query, and Value matrices, respectively. W (i)O performs a linear output projection, before adding the result to the residual stream via a skip connection. On the other hand, an MLP layer first up-projects its input X by multiplying with WU ∈ Rd×dh , then passes the result through an activation function such as GeLU, then down-projects the result back via WD ∈ Rd×dh . Here d is the hidden dimension and dh typically equals 4d. This is then added to the residual stream to produce
X + GeLU(XWU )WD
Layer Normalization precedes both layers.
BERT Fine-Tuning We fine-tuned BERT (bert_base_cased, 125M parameters) on MNLI for 5 epochs (61360 steps with batch size = 32), with a learning rate of 2 · 10−5, a and a maximum sequence length of 128. Checkpointing and evaluation were performed every 500 steps. For QQP, we opted for 40000 steps of fine-tuning with an effective batch size of 256 and a learning rate of 2 · 10−5. Evaluation and checkpoint were once again performed every 500 steps. All runs used the Adam optimizer (Kingma and Ba, 2015), with β = (0.9, 0.999) and ϵ = 10−8.
Each run takes approximately 3 hours to complete on one NVIDIA A5000 GPU.

heading: B Details of Pruning and L0 Regularization
text: Our formulation of the masks is the same as CoFi Pruning (Xia et al., 2022)’s, and is built on top of Louizos et al. (2018)’s work. Specifically, the masks z are modeled as hard concrete distributions:
u ∼ Uniform(ϵ, 1− ϵ)
s = σ ( 1
β · u 1− u + logα
)
s̃ = s× (r − l) + l z = min(1,max(0, s̃))
Here, σ denotes the sigmoid function, and ϵ = 10−6 is chosen to avoid division by 0 errors. The temperature is fixed at 1β = 2 3 . The result is then stretched to the interval [l, r] = [−0.1, 1.1], with the probability mass from [−0.1, 0] and [1, 1.1] accumulated to 0 and 1 in the end. This ensures that the mask is incentivized to contain values close to 0 or 1. Here, the log alphas logα are the main learnable parameters.
Following Wang et al. (2020) and Xia et al. (2022), we enforce a target sparsity via a Lagrangian Term. Supposing the target and current sparsity to be s and t, we add
Ls = λ1(t− s) + λ2(t− s)2
to the loss (CoFi uses s − t instead of t − s, but the two are functionally the same; consider the transform λ1 −→ −λ1). Along with training the model parameters, a gradient ascent is performed on λ1 and λ2, so that the updates of the lambdas always keep adjust the contribution of Ls to keep the sparsity in check.
Pruning - General Commentary We found that the hyperparameters were very sensitive to the choice of learning rate. Deviation even by a factor of 3 would cause instability in training. Specifically, if the learning rate for the ascent on the Lagrangian Parameters (λ1 and λ2) was too low, the model would stay under a sparsity of 10−5 until the target sparsity climbed to 0.15 or so, and would then jump to a value beyond the target, sacrificing a lot of performance in the process. On the other hand, a large value also leads to insufficient time for “settling into” the intermediate sparsities. The exact choice of the number of optimization steps was also important, although not as sensitive as the learning rate.
BERT Pruning We used a maximum sequence length of 128 as before, with an effective batch size of 128 including gradient accumulation. The learning rates for log alphas and lambdas were 0.1 and 1, respectively. We use the formulae
warmup_steps = 6500 + 50× sparsity_pct
total_steps = warmup_steps+6×sparsity_pct
to compute the number of steps taken to linearly warm up the target sparsity, and the total number of steps, respectively. For sparsity targets of 50%, these translate to 9300 training steps, of which the first 9000 warm up the target sparsity. For 70% target, these change to 10500 and 10000, respectively. Evaluation and checkpointing were done every 64 steps but we always used the final checkpoint to report results. The thresholds were chosen with a grid search on [0, 1] with stride 0.05 to match the target sparsity as closely as possible. They were found to always be 0.5 and 0.35 for target sparsities of 0.5 and 0.7.
Each pruning run takes approximately 90 minutes to complete on an A5000 GPU. We estimate our total computational budget to be 750 GPU hours.

heading: C More results
text: In this section, we provide further exposition on some of the results from the main text. For these results, we use the full names of the HANS subcases instead of abbreviations.
Figure 10 showcases the in-domain v/s out-ofdomain training dynamics on MNLI/HANS for more subcases. The trend remains the same—OOD accuracy starts increasing much after ID accuracy saturates. For the other subcases not shown here, the model never gets off the ground.
Figures 13 and 14 show that the frequency of the heads is strongly correlated between the 50% and 70% sparsity subnetworks for MNLI/HANS, and the 30% and 60% sparsity subnetworks for
QQP. This point is further elaborated by the frequency heatmap of Figure 11, the counterpart of Figure 5. Essentially, the heads most important to the heuristic subnetworks also appear often in the generalizing subnetworks, once again showing that the heuristic subnetworks are subsets of generalizing subnetworks.

heading: D Results on RoBERTa
text: In this section, we verify that our findings extend to a RoBERTa (Liu et al., 2019) (roberta-base) model. We follow the same recipes as the main text, using the same hyperparameters as Appendices A and B for fine-tuning and pruning. Once again, we evaluate on the subcases SO-Swap, Prep, Embed-If, and Embed-Verb in the MNLI-HANS setting.
Table 5 demonstrates that different subnetworks at 50% sparsity still generalize differently despite
performing similarly to the full model in-domain. Figure 15 tells us that the effective size again goes up along with generalization. The 12 subnetworks share 11 attention heads that also appear in the effective size subnetworks before generalization. We list these heads along with the effect of ablating them in Table 6. The effects are smaller but consistent with the findings on the BERT model: these heads have much larger ablation effects than random other heads despite not generalizing on their own.

heading: E Results on GPT-2
text: In this subsection, we ask if our findings and explanations can extend to decoder-only models. We work with a GPT-2 small (117M) model (Radford et al., 2019). Most of the hyperparameters for both fine-tuning and pruning are the same as those listed in Appendices A and B. However, we found that the GPT-2 model, even as a whole, generalizes very poorly to the HANS subcases. After training for 10 epochs (12270 steps), constituent_cn_embed_under_if is the only “contradiction” subcase with an accuracy above 30%. Hence, we select this setting, and Embed-If
as the only adversarial out-of-domain subcase. Interestingly, however, we observed that the accuracies of some of the heuristic-friendly subcases started dropping later into the training. We select subsequence_se_relative_clause_on_obj (Rel-Clause [SE]) and constituent_ce_embedded_under_since (Embed-Since [CE]) as two representative subcases and group them under the split OOD-E (OOD-Entailment). We perform the experiments for both the OOD and OOD-E subcases.
We evaluate 12 pruned subnetworks at 40% and 60% sparsity in Table 7. Like before, the subnetworks perform similarly in-domain but generalize differently. However, one striking feature stands out. As we increase the sparsity, the average performance on the ‘contradiction’ OOD cases increases, whereas that of the ‘entailment’ cases decreases.
One way to reconcile these results is for some heuristics to be contradiction heuristics, and for the model to develop further non-heuristic heads around a contradiction heuristic core. Indeed, in Figure 16, we note two interesting trends: (1) the model accuracy decreases slightly on the OOD-E subcases late into training and (2) a mild increase
in the OOD-E effective size accompanies this decrease in accuracy. A possible explanation for the latter could be that some heuristics for spotting ‘contradiction’ are also learned, and play an analogous role in the heuristic core.
We find the heads common to all 40% and 60% subnetworks and tabulate their ablation effects in Table 8. We note the following: (1) On Embed-If, some heads show a strong increase in accuracy on ablation, and some a strong decrease. This would be consistent with some heads implementing entailment heuristics and some contradiction heuristics. On the other hand, (2) None of the heads show large ablation effects on Rel-Clause. On EmbedSince—the other entailment subcase—some heads show small but significant effects.
Therefore, our results on the OOD subcases are consistent with the model implementing both entailment and contradiction heuristics. On the other hand, some questions remain regarding the OOD-E subcases in GPT-2—namely, why the performance of the model tapers off late into the training, and why there is a variance in performance at high sparsities (in contrast to the results on BERT) despite none of the common heads having significant ablation effects. We leave it to future work to answer these questions.
0 20000 40000 60000 80000 100000 120000
Steps
0
Steps

heading: Future Work
text: We leave it to future work to answer these questions. 0 20000 40000 60000 80000 100000 120000
Steps
0
Steps