A distributional simplicity bias in the learning dynamics of transformers

We want to thank you for your positive comments and your enthusiam for our work. We really appreciated the series of language evolution studies on simplicity biases that you suggested, and we decided to include them in the revised version of our manuscript.

Below we reply to your questions and concerns point-by-point; we hope that these points alleviate your concerns regarding the architecture, the time complexity of sampling, and the need for additional empirical evidence. If so, we would appreciate it if you could reconsider your rating; if not, we're looking forward to discussing further during the discussion period.

• Title: We agree with your suggestion for the title, which would be more specific.

• Activation function: In our framework, we only use the $x^2$ activation function in the architectures that are subsequently used to generate the clones of a real NLP dataset, in which the maximum degree of interaction among tokens can be controlled by tuning the number of layers. We use these clones to demonstrate the simplicity bias of a standard transformer encoder (BERT-like) with standard activation function trained under MLM on a real NLP task. Thus, the main result of our paper is obtained for a standard architecture, trained with a standard task on a standard dataset.

• Decoder-only transformers are indeed popular language models. The procedure we present in this manuscript, and in particular the analytic derivations, are however designed for BERT-style transformers, and validated only on such models. Extending it to causal models and other NLP tasks is for sure possible and interesting, but non-trivial, and we plan to investigate this topic in our future research.

• Time complexity of sampling The reviewer is right that the sampling of the clones is not trivial; it is difficult to fully decorrelate the samples, and the complexity of the method we use is not linear in the number of samples. However, the modified transformer architectures that model the data are computationally much more easy to handle than, say, estimating the corresponding $n$-grams. The clone models require only a couple of layers of factored self-attention, which is easier to compute than vanilla self-attention layers, and they do not need fully-connected networks in-between attention layers. This makes sampling of our clones, while challenging, also way cheaper than sampling the distribution generated by an ordinary transformer.

• Evidence for simplicity bias We appreciate the need for additional experimental verification, and have repeated our experiments on another data set, WikiText101. As we show in our plot in the additional one-side PDF, we can see sequential learning of increasingly higher-order interactions also on this data set. In our manuscript, we further show sequential learning in a controlled setting with synthetic data (cf. Section 2) and we see the same behaviour in our analytical treatment of Sec. 3.1. We hope that this combination of experimental verification on tinystories, WikiText101, synthetic data, combined with our analytical results, is convincing evidence for our claims.

We thank you for your careful reading of the manuscript, your detailed review, and the many references. Below we reply point-by-point; we hope that our replies alleviate your concerns. If so, we would appreciate it if you could revisit your rating; if not, we look forward to discussing further during the discussion period.

• Related works We thank the referee for the references. Indeed, the connection with n-gram models is extremely important, since our clones implement n-grams using a transformer architecture. We will explicitly mention this in the revised version. On the other hand, we are aware that the NLP community has made enormous steps forward in understanding learning dynamics, with paradigms that go well beyond n-grams. In the revised version of the manuscript we will explicitly mention ref https://aclanthology.org/2022.acl-long.568.pdf and https://aclanthology.org/2022.acl-long.568.pdf as examples of advanced quantitative analysis in this direction. We will also mention https://aclanthology.org/N19-1329/ as a technique which allow probing the learning dynamics in language. Finally, we discuss the relation of our work to the very recent paper https://arxiv.org/pdf/2402.04362. They show that deep transformers first learn the token unigram and then the bigram distribution. However, estimating the higher-order interactions with the approach described in this work is prohibitively expensive on large corpora with many tokens due to the curse of dimensionality. Moreover, we would like to highlight that none of these papers, including arXiv:2402.04362, introduces an architecture which produces distributions in which the simplicity bias can be controlled by tuning the interactions between tokens that are captured by the model, at an arbitrary order. This analytically guaranteed control, as also recognized by the reviewer, is an important and novel contribution of our work.

• The nature of simplicity bias Thank you for sharing these references on various ways in which learning simple rules might interfere with learning more complex rules. These papers demonstrate "simplicity biases" by performing experiments on appropriately fine-tuned training tasks, where simple rules based on spurious correlations hinder the network from learning the generalisable input features. These behaviours are interesting and important, but have to be induced by an external intervention (a fine-tuned task). If trained by a standard task (say text completion), transformers seem instead to pick up the most relevant underlying rules of language rather well, producing flawless chunks of text. It is this process that we want to study through the lens of another line of research. We hope that our paper convincingly shows that during a single training run, a model often seems to be learning increasingly complex functions or distributions that explain its training data. We discuss previous theoretical and experimental work in the introduction, in particular in lines 18-26, many of which have been published at NeurIPS/ICML/ICLR in the last couple of years (Refs 2,3,4,5,6,10,12,13). We will mention this separate line of work on simplicity biases in the revised version of the manuscript.

Regarding your questions:

Can you give a simple explanation of what a k-body interaction would look like? is it essentially stating something about ngram models, but not restricted to the most proximate tokens?

A $k$-body interaction is any interaction between $k$ tokens anywhere along the sequence that can not be written as a product of lower order interactions. In the energy-based models we use to construct the clones, each such interaction corresponds to one term in the energy function involving $k$ tokens. Indeed, the tokens do not have to be proximate. N-gram models instead simply define the number of previous tokens on which a given token is conditionally dependent, i.e. how many conditions there are in $p(x_n | x_{n-1}, x_{n-2}, …)$.

Can you guarantee that the difference provably limited models actually differ by nothing but their limitations? What about random variation?

The architecture of the "limited models" that we use to obtain the clones only differ by their "limitations", namely by the degree of the interactions between tokens that they cover. After training, the values of the parameters of these models could have some random variations in their values due to different initial conditions; however, this will not change the degree of the interactions between tokens represented by these models.

We thank you for your detailed questions and your feedback, which led us to conduct additional experiments. Below we reply point-by-point; we hope that our replies alleviate your concerns. If so, we would appreciate if you revisit your rating; if not, we look forward to discussing further during the discussion period.

• Computational cost: applying our method to train a many-body approximation of a data set, or "clone", is less expensive than training a standard BERT-like model, since the architecture is much simplified. Likewise, the sampling procedure that we use from Goyal et al. [ICLR '22, arXiv:2106.02736] can sample the standard BERT models. Hence we are confident that the method can be applied to any data set on which BERT was applied to.

• Generality: The procedure, as presented in the manuscript, and in particular the analytic derivations, are designed for BERT-style transformers, and validated only on such models. Extending it to causal models and other NLP tasks is for sure possible and interesting, but non-trivial, as the probability structure induced by causal training is analytically (and empirically) different. We plan to investigate these topics in our future research. To further support the empirical evidence of sequential learning in BERT-like architectures trained with MLM on NLP datasets, we replicated the analysis on the WikiText-2 dataset, observing a similar hierarchy in the learning process, and with a vanilla transformer on the synthetic data set we created (see plots in the attached one-page PDF). In both cases, we see clear sequential learning by the transformers.

We thank the referee for their detailed questions and their feedback. Below we reply point-by-point; we hope that our replies alleviate your concerns. If so, we would appreciate if you revisit your rating; if not, we look forward to discussing further during the discussion period.

Support for our main claim in Figures 1 and 4:

We're glad the reviewer agrees that our results show that transformers learn higher-order interactions between tokens in the later stages of training, point (ii). However, we think our results support also point (i) both in experiments on real data and in the controlled setting with synthetic data. The left panel of Figure 1 shows that the test loss curves for the different clones overlap during the initial stage of training. Specifically, even after the initial plateau, between 5000 and 6000 training steps, the loss curves on the various clones continue to overlap despite the quick improvement in test error. The crossover between the two-body and the higher-order interactions is at step 15000, and only at 22000 the models start differing significantly. These results show that, at the early stage of training, the model first learns low-order interactions without learning high-order interactions. Furthermore, point (i) is clearly supported by the experiments on the synthetic dataset (Section 2) and from the analytical computation (Section 3.1).

I encourage the authors show the test loss on more individual instances as well as the average test loss across all instances to consolidate this conclusion.

We have performed our experiments using various random seeds, as specified in the figure captions, and found the results to be consistent across different seeds. We have updated the plots by adding the corresponding error bars; we show an example in the rebuttal figure 1, where we redrew Fig 1 with error bars. We will update all figures in this way in the revised version.

Clarifying the nature of “factored attention”

We define factored attention, where the attention matrix depends only on the position of the tokens, by writing the corresponding probability for the masked token $p_{\mathrm{mlm}}$ in Eq (2). To clarify this we will first define only the factored self-attention layer, and then the equation for $p_{\mathrm{mlm}}$. The trainable parameters in a factored self-attention layer are the input-independent attention weights $A$ (an $L\times L$ matrix, with $L$ the length of the input sequence) and the value matrix $V$.

(2) In Line 93, $\mu$ and $\alpha$ are not defined. I guess $\mu$ means the $\mu$-th sample in the training set, but I’m not sure what $\alpha$ refers to.

You are right: $\mu$ is an index that runs over training examples; $\alpha$ indexes positions along the sequence. We will clarify this in the revised manuscript.

(3) According to Line 84 “each token is an element taking values in a discrete vocabulary ”, the equation in Line 93 means that we are always measuring the probability of generating the first token in the vocabulary. This seems unreasonable.

We wrote the equation in line 93 for the first token to keep the equation readable; this equation (and any other) naturally extends to predicting any other token.

(4) In Equation (2), $A$ and $V$ are not defined.

Thank you for pointing this out -- $A$ and $V$ are the attention and the value matrix, respectively. We will add a comment.

(5) The meaning of the subscript $\mu$ in Equation (3) seems to conflict with the meaning of the superscript in Line 93.

In both equations, $\mu$ is simply a "dummy" index that is summed over.

(6) In Line 134, $p$ is not defined. Moreover, it is not clear whether the superscript is an exponent, or simply an index.

In this equation, $p$ is also a dummy index. We now see that this might clash with the notation for probabilities, and use a different dummy index for this sum.

Regarding the design of Figure 1

We have chosen the folded tape shape because it allows highlighting groups of words which are far away, but (qualitatively) form a group of interactors. Importantly, the factored attention architecture learns exactly the strength of the many-body interactions along the sequence. This is the message that we would like to convey with the figure. These days we attempted to follow the reviewer’s suggestion and write the sentence on a horizontal line, but we did not manage to find a graphically satisfactory alternative.

Repeating the experiments on synthetic data with a transformer with standard attention.

This is a great question, and inspired us to perform an additional analysis. We trained a standard three-layer transformer encoder on clones of a synthetic datasets characterized by interactions involving up to four bodies. Being a standard transformer, this architecture also had several layers of Layer Normalization, which mitigate the occurrence of plateaus even on synthetic datasets. However, the plot for this standard transformer in Fig 1 of the rebuttal figures clearly shows sequential learning also in this case.