Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models

Thank you for the helpful comments!

The found parameters are mainly on the first layer which reduces the strength of the claims, it makes it less clear wthat the model uses it at the beginning when it overfits bigrams (with all parameters) and such issues.

Yes, we find that bigram subnetwork parameters are consistently concentrated in the first layer. Early in training when the model overfits to bigrams, it may be that the rest of the model is still essentially untrained, and the noise of the later layer parameters washes out, leaving the early layer bigram predictions unchanged even if the later layer parameters are included. This result is closely related to model distillation and capacity measurement, as it indicates that for un-fully-trained models, a substantial proportion of parameters can be dropped without changing model behavior.

Minor: l26 wouldn't it be better for the citations to be near the topic they refer to, instead of all together?

Good point, we will update this!

In Fig. 1 why not write e^x instead of just x? Isn't it an easier way of conveying the number and that it is log scale.

We used x instead of 10^x to avoid having the exponents too small to read, but we can update this in the revised manuscript.

l127 it is shown in the left or center not on the right figure.

Raw surprisal correlations between bigram subnetworks and the bigram distribution are reported in the table on the right of Figure 1. We agree that the left and center also show these effects, so we will update this mention in the revised manuscript.

Thank you for pointing out these areas to clarify!

Which parameters are the target of masking in Eq. 1? Does it also include parameters in the language model head? I’m curious about whether the first layer parameters indeed have bigram knowledge or if it just triggers some bigram knowledge in the language model head... Language model prediction head (and its interplay with activations) should also be an important analysis target, given the fact that it handles primitive language statistics. However, it’s not explicitly analyzed or discussed.

We learn the subnetwork mask over all model parameters except the token embedding, token unembedding (language model head), and layernorm parameters (Section 3.1, Appendix A.2). We agree that the language model head contains important information about primitive language statistics (e.g. token frequencies and bigram probabilities). Indeed, in Appendix A.4, we find bigram subnetworks in randomly-initialized models as long as trained token embeddings and unembeddings are patched in. This indicates that, as the reviewer suggests, the first layer parameters might primarily trigger some bigram knowledge in the embedding and unembedding parameters; still, our ablation results indicate that this triggering of bigram knowledge is a key part of Transformer language model processing, and thus these non-embedding subnetwork parameters are still a central component of the bigram subnetwork.

We also note that we did not include embedding and unembedding (language model head) parameters in the mask because we hypothesized that embedding parameters for rare tokens might be highly likely to get masked. Specifically, when optimizing for bigram cross-entropy loss and encouraging sparsity, the continuous sparsification method might be likely to find it optimal to simply drop many embedding and unembedding parameters for rarer tokens, because they contribute less to the loss.

How did you estimate the gold bigram distribution in Eq.1? Are you targeting bigram at the subword level or word level?

Thank you for pointing out that we omitted this detail! Bigram distributions were estimated at the subword (token) level for each model, using 10M sequences of 128 tokens (1.28B tokens total) from the OSCAR English corpus of web text. Estimates were computed by counting bigram frequencies in the corpus. Because the corpus was large enough that all tokens appeared at least once (i.e. all possible bigram "prefixes"), there were no undefined bigram probabilities in our experiments. Thus, we did not smooth the bigram distribution estimate. We will include these details in the updated manuscript.

Is there any concern that the findings are biased toward the bigram estimation method? (e.g., reporting bias, smoothing...)

Reported above, we did not use any smoothing for the bigram distribution estimate. However, we agree that there may be reporting bias: we estimated the bigram distributions from OSCAR English web text, but the models were not all trained on exactly this dataset. Still, both the Pythia and GPT-2 models are primarily trained on English web text, so we expect the bigram distributions in their training datasets to be roughly similar to our estimates from OSCAR.

What type of statistical test yielded p<0.0001 in Section 6.2?

When considering the parameter overlap between two subnetworks A and B, we generate 10,000 random subnetwork pairs (pairs of parameter masks) that randomly select the same number of parameters as A and B respectively within each parameter block. Then, we count the number of random pairs whose parameter overlap is greater than or equal to the parameter overlap of A and B. This gives us the probability that two random subnetworks would have equal or greater overlap than A and B, after accounting for the number of parameters kept per block in A and B. In all cases, none of the 10,000 randomly generated pairs had greater parameter overlap than A and B, i.e. $p < 0.0001$. We will include these details in the updated manuscript.

Thank you for the helpful comments!

The work may have some implications for practice (say for pruning), and it would be helpful and make the paper more impactful to include explicit discussion of them. This would give the work some practical relevance on top of being a nice scientific work... What practical implications may exist as a result of identifying bigram subnetworks, that don't exist when identifying the optimal subnetworks (using the definitions in the paper)?

We hope that the bigram subnetworks would be useful as starting points for interpretable circuits that trigger behaviors that diverge from bigram predictions (Section 7). The bigram subnetworks allow researchers to disentangle the subnetwork for basic next token prediction (the bigram subnetwork) from an actual behavior of interest. In particular, unlike an optimal subnetwork, the bigram subnetwork's predictions are highly interpretable, so we can be more confident that the bigram subnetworks are not performing opaque operations that might confound results for other circuits. Finding these more precise and interpretable circuits could facilitate model pruning methods that better target specific behaviors rather than "optimal" behavior as a whole.

Thank you for the thorough review!

L174 suggests (if I understand correctly?) that the concentration of subnetwork parameters in early layers is not necessarily due to an actual computation localized in these layers, and may be an artifact of the architecture+objective+optimization setting.

Yes, we agree that the concentration of bigram subnetwork parameters in early layers in the randomly-initialized networks suggests that such a property might be an artifact of the architecture+objective+optimization setting. Other subsets of parameters in other layers might recreate similar properties, and we do not necessarily claim that the bigram subnetworks we find are unique (l. 265-268; 552-557), despite their apparent importance to language model predictions.

Still, we speculate that as models compress more information into their parameters throughout training, the bias towards early layers for bigram computations might make the models likely to compress any redundant bigram-like computations into those early layers. This could compound with the fact that later model computations could scaffold off of bigram computations (e.g. refining the bigram predictions), making it more optimal for bigram computations to appear in earlier layers. Future work might investigate whether these types of suboperation dependencies and parameter location biases explain the tendencies for particular suboperations to appear in specific layers (e.g. to better understand the learning mechanisms driving classic results such as "BERT Rediscovers the Classical NLP Pipeline"; Tenney et al., 2019).

Together, these findings cut down the (mostly implicit?) intuition that these bigram subnetworks are performing a more meaningful computation, e.g. as discussed in Appendix A.5. I would appreciate if the authors could add at least a bit of summary speculation late in the paper about what these findings mean for possible computational functions of the "bigram subnetwork."

We speculate that early in training (or for low-performing models), bigram heuristics near-optimally decrease the language modeling loss overall, so the optimizer pushes parameters towards recreating a similar distribution. Thus, early in training (e.g. when the model overall approximates the bigram distribution), the bigram subnetworks are likely indeed serving a computational function similar to bigram prediction. Later in training, we speculate that the model scaffolds off of these bigram predictions and subnetworks, adapting the early layers to perform generally useful early-layer computations (e.g. transforming to next token space) that still maintain some similarity to bigram predictions (even if not exactly an "intermediate bigram prediction", as discussed in Appendix A.5). It may be that this persisting similarity to bigram predictions in early layers is what makes it possible to extract bigram subnetworks even late in training, when these subnetworks have adapted to scaffold other computational roles as well.