A distributional simplicity bias in the learning dynamics of transformers

Authors: riccardo  rende
federica  gerace
alessandro  laio
sebastian  goldt
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:16:04

Abstract: The remarkable capability of over-parameterised neural networks to generalise effectively has been explained by invoking a ``simplicity bias'': neural networks prevent overfitting by initially learning simple classifiers before progressing to more complex, non-linear functions. While simplicity biases have been described theoretically and experimentally in feed-forward networks for supervised learning, the extent to which they also explain the remarkable success of transformers trained with self-supervised techniques remains unclear. In our study, we demonstrate that transformers, trained on natural language data, also display a simplicity bias. Specifically, they sequentially learn many-body interactions among input tokens, reaching a saturation point in the prediction error for low-degree interactions while continuing to learn high-degree interactions.  To conduct this analysis, we develop a procedure to generate \textit{clones} of a given natural language data set, which rigorously capture the interactions between tokens up to a specified order. This approach opens up the possibilities of studying how interactions of different orders in the data affect learning, in natural language processing and beyond.

Main paper: 
 introduction: 
 The success of neural networks has been explained using various "simplicity biases", which postulate that neural networks trained using stochastic gradient descent (SGD) learn increasingly complex functions of their data, going from generalised linear models to more complex, non-linear functions [1] [2] [3] [4] [5] . These simplicity biases have been analysed theoretically in linear networks [6] [7] [8] , in non-linear shallow neural networks [1, 9] and in the NTK regime [10] [11] [12] . The viability of such a scenario has also been demonstrated experimentally, for example in image classification with convolutional neural networks [4] . More recently, a complementary distributional simplicity bias was demonstrated theoretically and experimentally [13, 14] : deep neural networks learn increasingly complex approximations of the distribution of their training data. A natural question to ask is then: Does this distributional simplicity bias extend to the realm of natural language processing (NLP), where transformers are typically pre-trained via Masked Language Modeling (MLM) or Next-Token Prediction? As a first step in this direction, Belrose et al. [15] recently showed that deep transformers of the Pythia suite [16] first learn the token unigram and then the bigram distribution. However, estimating the higher-order interactions is prohibitively expensive on large corpora with many tokens due to the curse of dimensionality.In this work, we design a novel framework based on MLM to build principled approximations of a NLP data set that captures increasingly higher-order interactions between tokens. Using these approximations or "clones", we show that deep transformers trained on natural language also exhibit a simplicity bias: they sequentially learn increasingly higher-order interactions among tokens. We illustrate the idea of a statistical "clone" of a data set, which approximates the underlying data distribution by keeping only interactions between tokens up to a fixed degree (in this case, three-body interactions). We introduce a principled approach to create clones by training a transformer with multiple layers of factored self-attention [17] with x 2 activation function between layers. The depth of the architecture controls the degree of the approximation. Clones can then be sampled from these models. Right: Test loss of a standard BERT-like transformer encoder [18, 19] with four attention blocks trained on the WikiText-103 [20] data set and tested on clones of this data set with a truncated maximum degree of many-body interactions between tokens. We show the average over five training runs starting from the same initial condition. The shaded area indicates one standard deviation.The key idea of our framework is shown on the left of fig. 1 . Given a data set containing natural language (we will use WikiText-103 [20] as a running example), we first create a number of clones.Each clone approximates the underlying data distribution of WikiText-103, but only including interactions between tokens up to a specified order (in the sketch, up to order three). We build these clones by using MLM to train a new type of transformer that consists of multiple layers of factored attention [17] with quadratic activation functions. In section 2, we show analytically that the depth of this architecture precisely controls the degree of interactions among tokens in each clone. We then sample the trained models using state-of-the-art Monte Carlo techniques [21] to obtain the cloned data sets. The clones thus form a hierarchical family of data sets with increasing order of interactions among tokens that approximate the true distribution of WikiText-103 with increasing fidelity.Using this set of clones, we investigate the sequential learning of higher-order interactions between tokens in a standard BERT-like transformer encoder [18, 19] with four attention blocks. On the right of fig. 1 , we show the test loss of this architecture during MLM training on WikiText-103, but testing on various clones of WikiText-103 with a truncated maximum degree of many-body interactions between tokens, as well as on a clone sampled directly from the trained BERT model. The test losses of the transformer on the clones are close in the beginning of training up to ≈ 20 epochs. Beyond this threshold, BERT exhibits a different generalisation behaviour on clones that capture different orders of interaction: while the test loss on the clones containing only low-degree interactions among input tokens (blue and green curves) saturates after about 60 SGD epochs, the test loss evaluated on clones having high-order interactions (red curve) keeps improving during training. In other words, the accuracy of the transformer in the beginning of training derives from having learnt lower-order interactions, and only later during training the transformer also exploits higher-order interactions between tokens. In section 4, we extend our experimental analysis to an alternative dataset, TinyStories [22] , and further include experiments involving GPT models trained on next-token prediction tasks.The main contributions of this work are the following:• We introduce generative models based on transformers with factored attention, and show that the degree of interactions they capture is rigorously controlled by their depth (section 2); • We study sequential learning of higher-order interactions both numerically and analytically in a controlled setting where we regulate the degree of interactions between tokens (section 3); • We train these generative models on benchmark NLP datasets, such as TinyStories [22] , and sample them using Monte Carlo techniques to demonstrate the sequential learning of higher-order interactions by transformers in MLM and next-token prediction (section 4).Our work shows that the many-body interactions between tokens in natural text are key for learning, and we note that our approach naturally extends to other data modalities.Further related work on simplicity biases in transformers Bhattamishra et al. [23] found a simplicity bias in the ability of transformers to learn sparse boolean functions; here, we focus on the MLM task on natural language. Abbe et al. [24] showed that the difference between trained and initial weights progressively increases in rank throughout training, hence increasing in complexity as the task requires. Vasudeva et al. [25] investigated a simplicity bias in transformers in terms of their sensitivity to random changes in the input. Other recent experimental work showed that RoBERTa achieves high performance on most linguistic benchmarks early in pre-training, while more complex tasks require longer pre-training time [26] , and that transformers learn different linguistic phenomena in a prototypical way [27] . To the best of our knowledge, no prior research has investigated the progressive increase in complexity of the distribution learned by a transformer from its data over time.More recently, Cagnetta et al. [28, 29] studied probabilistic context-free grammars and showed that increasing the training set size allows transformers to learn correlations at increasingly large distances sequentially, corresponding to their resolving increasingly higher layers of the hierarchical structure of the tree. While we were preparing the camera-ready version of this paper, Garnier-Brun et al. [30] demonstrated a similar phenomenon in the training dynamics of transformers.Reproducibility We share the scripts to train and sample the networks and reproduce our results at https://doi.org/10.5281/zenodo.13992398. 
expressing many-body distributions with factored attention: 
 In this section, we present generative models utilising deep networks composed of multiple attention layers. We demonstrate analytically that these models, when trained with MLM, are capable of learning increasingly higher-order interactions among tokens. In the MLM task, transformers are trained to predict missing tokens in sentences of tokenized text. We model sentences as sequences of random variables (s 1 , . . . , s L ), where each token s i is an element taking values in a discrete vocabulary:s i ∈ {1, . . . , |V|}.When masking the token in position i, the sequence (s 1 , . . . , s i = 0, . . . , s L ) is given as input to the transformer, which embeds the input sequence into a collection of vectors (x 1 , . . . , x L ), each of them living in a d-dimensional space. Then, this sequence of vectors is processed by the self-attention layers and transformed into a sequence of context-dependent vectors (y 1 , . . . , y L ), where y j ∈ R d for all j. At the end, the output y i , corresponding to the masked input token, is mapped into a vector of the same size of the vocabulary: The key ingredient to the transformer architecture [31] is the self-attention mechanism, which recombines input tokens into output tokens by weighing them with an input-dependent attention matrix. Rende et al. [17] recently considered a simplified attention mechanism, the so-called factored attention, where the attention weights are input-independent and the input sequence is processed as:h i ∈ R |V| . The prediction p mlm (s i |s ̸ =i ) = softmax(h i ) cany iα = L j=1 A ij |V| β=1 V αβ x jβ ,where A ∈ R L×L and V ∈ R d×d are matrices of trainable parameters. In this way, the prediction of a single factored attention layer for the masked token is expressed in terms of the following conditional distribution (assuming the embedding matrix to be the identity):EQUATIONIt can be shown analytically that, when training a single layer of factored attention via MLM, the model is capable of exactly inferring the two-body interactions among input tokens [17] , when sampled according to the joint distribution:EQUATIONwhere J ij describes the interactions among positions along the sequence, U αβ describes similarities between tokens and Z normalises the distribution. In the following, we introduce the convention of using Latin indices to run over the length of the sequences L and the Greek indices to run over the size of the vocabulary |V|. As can be easily checked, the expression in eq. ( 1) exactly matches the conditionals associated with the two-body joint distribution in eq. ( 2). Remarkably, a single layer of standard self-attention, where the attention weights depend explicitly on the inputs, struggles to learn these same conditionals [17] . As a consequence, factored attention has been applied successfully in various areas of science, like protein structure prediction [32] or the approximation of ground states of frustrated quantum systems [33] [34] [35] [36] .We now demonstrate how we can easily express distributions with many-body interactions by stacking layers of factored self-attention. In particular, we introduce an architecture that can learn interactions of a specified order by tuning the number of layers. To this end, we start considering the simplest case where we stack just two layers of factored attention, with f (x) = x 2 as non-linearity between them and skip connections, see fig. 2a ) for a sketch of this architecture. By writing down the conditional distribution corresponding to this two-layer model:EQUATIONEQUATIONjk V(1)EQUATIONβµ s lµ ,(3) we notice that the second layer expresses a three-body interaction, parameterised by a specific combination of the attention weights of the first and second layer through the tensor Jikl =j A (2) ij A (1) jk A (1)jl ( Ũ is similarly expressed in terms of V ). In particular, the first element of the sum in eq. ( 3) accounts for the contribution of the skip connections, while the second one matches the functional form of the exact conditionals of a three-body joint distribution:EQUATIONGeneralising these observations to a generic number of factored attention layers, we conclude that by stacking n layers of factored attention with x 2 activation function we can infer interactions up to order 2 n-1 + 1. In this perspective, if the activation function is quadratic, the number of layers precisely controls the degree of interactions which can be inferred by a deep factored transformer. 
learning many-body interactions with factored attention: 
 In this section we show how, in a controlled setup, the multi-layer factored attention model introduced in section 2 sequentially learns increasingly higher-order interactions among tokens during training.To this end, we perform a numerical experiment where we sample the input sequences from a joint distribution with tokens interacting only up to the 4th-order, that is:EQUATIONHere, we ensure the couplings to be symmetric tensors by considering a set of vectors {j p } P p=1 and setting J ijkl = P p=1 j p i j p j j p k j p l . We choose P = 200 and we sample j p i iid ∼ Bernoulli(0.8). The same holds for U . The exact conditionals corresponding to this joint distribution can be computed analytically, precisely as in the three-body case exemplified in eq. 4. From the practical point of view, we use these conditionals to sample a set of M sequences {s µ } M µ=1 via Gibbs sampling, which will constitute the training set. We then train via MLM an increasing number of factored self-attention layers (from one to three) on these data, monitoring the test loss as a function of the number of epochs. The outcome of this experiment is shown in fig. 2b ). The dashed orange line corresponds to the lowest possible value of the test loss that a learning model would achieve in case it would be able to perfectly reconstruct the four-body tensors J and U of the generative model. The dashed blue and green lines correspond instead to the lowest test loss achievable by the best two and three body approximations of the input distribution, respectively. These last two bounds can be obtained by training via MLM a model whose predictions correspond to the exact conditional of a two and three-body distribution, where the learnable parameters are the tensors in eq. ( 1) and eq. ( 4). As shown by the red curve in fig. 2b ), if we train a single layer of factored attention on the four-body interacting sequences, the best we can do is to exactly recover the best two-body approximation [17] . By adding one extra layer of factored attention with x 2 non-linearity, we can reconstruct the best three-body approximation (violet curve). Then, the reconstruction of the full four-body interactions can be achieved by adding one further layer of factored attention (grey curve). This is consistent with the observation that by stacking n layers of factored attention with x 2 activation function we can infer interactions up to order 2 n-1 + 1.) also reveals that, in this controlled setup, the learning of higher-order interactions is preceded by distinct plateaus in the test loss. For example, in the three layers case, an initial plateau aligns with the best two-body approximation, followed by a subsequent plateau corresponding to the best three-body approximation. Interestingly, this sequential learning behaviour further implies that the layers are activated sequentially in time. This is illustrated in Figure 2c ), where we display the Mean Square Displacement (MSD) of the weights for three layers of factored attention across training epochs, that is:EQUATIONwith A (l) (t) and V (l) (t) being the attention and the value matrix of the l-th layer at time t. As we can see, despite the simultaneous update of all network parameters during training, a distinctive sequential learning pattern emerges. In the early epochs, only the last layer (orange curve) activates concurrently to the appearance of the first plateau, thus suggesting that the last layer is the one responsible for learning the best two-body fit of the training data. Subsequently, the weights of the second layer come into play (green curve), and combined with the weights of the last layer can reconstruct the best three-body approximation of the underlying four-body data distribution. Close to the end of the training, also the weights of the first layer start changing and, combined with the other layers, effectively infer the ground-truth four-body interaction among tokens in the training data. A hint of the reason why the weights of the last layer find the best two-body interaction is that, due to the presence of skip connections, they appear linearly combined with the input tokens in the output prediction. In the next section we provide a more formal analysis of this phenomenon.To gain theoretical insights on the sequential learning behaviour observed in the previous numerical experiments, we analyse the gradient flow of two-layers of factored attention with x 2 activation function when trained via MLM. For the sake of simplicity, we restrict to the case d = 1, that is when the tokens in the sequence (x 0 , x 1 , . . . , x L ) are real numbers rather than vectors. We consider the learning framework where the model is trained to always predict the first element of each sequence, using the square loss L = 1 2 (x 0y 0 ) 2 . Here, y 0 represents the transformer prediction about the masked token which, in this setting, acquires the following functional form:EQUATIONjk A(1)EQUATIONThe gradient flow equations can then be written explicitly for the attention weights of both factored attention layers as:EQUATIONwhereEQUATIONkl x l . From these equations, we can immediately see that, if we initialise the network with zero weights A (1) = A (2) = 0, the first layer attention weights do not change in time, Ȧ(1) kl = 0, while the second-layer weights do, Ȧ(2) 0j ̸ = 0, consistently with what we observe in fig. 2c ). We can gain some further intuition into which degree of interactions each layer learns by simplifying the data model further. In particular, we can restrict to a Gaussian design where tokens are now drawn from a zero-centered Gaussian distribution. By taking the expectation over the data, the correlations of an odd number of random variables vanish and the gradient flow equations can be written as:EQUATIONwhere we have defined h j = E[x 0 x j ] for j > 0 and C ij = E[x i x j ] for i, j > 0. In appendix B we further show that, if data are Gaussian distributed, the expectation value of the loss is bounded from below by the value it takes for A (1) = 0. Therefore, if one adds an ℓ 2 regularisation of arbitrarily small strength, A (1) = 0 is the only stable fixed point of the gradient flow equations. This is consistent with the numerical experiments, which show that for two-body distributed data the weights of the first layer vanish during the optimisation. In this case, only the second layer weights contribute to the learning dynamics:EQUATIONwhose solution describes an exponential convergence in time toA (2) 0j (t → ∞) = (C -1 h) j .This stationary point reproduces the exact conditional distribution of a Gaussian entry given all the others. Indeed, considering a multivariate Gaussian distribution with zero mean, indicating the first row of the covariance matrix (except the 0, 0 element) by h and the block obtained with i, j > 0 by C we have:p(x 0 |x >0 ) = N (x T >0 C -1 h, Σ). This means that the second layer is learning the covariance matrix, and thus the interactions up to the second order.If data contains higher-order interactions among tokens, we need also the first layer to unveil the underlying data distributions. Once again, we can see this from the gradient flow equations. In the general case of higher-order interactions (thus for non-Gaussian data distributions), the stationary point A(t) = 0 is generally unstable. For example, considering the case of a diagonal perturbation A(1) kh = ϵδ kh around A (1) = 0, eq. ( 9) becomes (neglecting terms of order ϵ 2 ):EQUATIONwhere we have taken the expectation with respect to a generic data distribution and definedΓ ijk = E [x i x j x k ].In this case A (1) (t) = 0 is a stable solution with respect to a diagonal perturbation only if the matrixEQUATION0j Γ jlk is negative definite (implying that the perturbation decays to zero as t goes to infinity). Therefore, the structure of the three-body correlations Γ ijk determines the stability of the stationary point A (1) = 0. Indeed, if, for example, one takes Γ jlk = γδ ij δ jk , we haveG lk = γδ lk A (2) 0k δ 0k -A (2)0k , which is not negative definite for a generic choice of A(2) 0 . This analysis allows concluding that there are conditions for which the solution A (1) = 0 is unstable when data contains interactions of order higher than second, which is precisely in line with what we observe in fig. 2b ).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes]Answer: [Yes] Justification: The code has been shared on zenodo.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
sequential learning in nlp: 
 In this section we use the tools introduced in the previous sections to investigate which features of the input data influence the training dynamics of transformers pretrained on real NLP datasets. Specifically, we analyze the TinyStories dataset [22] , training a standard BERT encoder [18] using the masked language modeling objective, as well as a standard GPT model for the next-token prediction task. Our focus lies on understanding whether a sequential learning of many-body interactions among input tokens, akin to the observed phenomenon in the synthetic scenario, is also prevalent in a real-world context. Since we do not know the underlying joint distribution over words, our approach will be to first train the attention-based models of section 2 with varying number of layers on TinyStories, that can be sampled to generate "clones" of the original data with a truncated maximum degree of interaction. We now discuss our experimental methods and analyse the results of our experiments; for full details on the methods, see appendix A.The clones We exploit the generative models described in section 2 to generate clones of the original TinyStories data set (see appendix A.1 for details on how we prepare the data set). The clones are meant to approximate the underlying TinyStories distribution with increasing fidelity. We use n = 2, 4 and 6 layers of factored attention to generate clones that include effective interactions up to degree 3, 9, and 33, respectively. Since we use a vocabulary of size |V| = 10000, we perform a finite-rank approximation of the value matrix with rank 400 (see appendix A.2 for more details). InOnce upon a time there lived a cat named Max. He was a restless cat, never staying in one place for too long. One day, Max demanded something funny. He wanted a friend.Once upon a time there was a boy named Timmy. He was very hungry lion the with not could play the ball to get wet.One day, Lily saw something else. He shook his head.Once upon a time there was a boy named Tom. He saw a small cat, could not play in the park with flowers pretty.One day, Lily smelled and nice. They wanted the ball.Once upon a time there was a puppy named Max. Max was a small puppy who was buzzing in a basket for very long.One day, Max saw something amazing. He wanted a puppy. fig. 3a ), we show the test loss curves when training models with two, four and six layers of factored attention featuring the x 2 activation function on TinyStories. Notably, the final value of the test loss decreases systematically with the order of the interaction modelled by the architecture; for comparison, the test loss curve associated to the training of a standard BERT architecture [18] is also shown.The bidirectional architectures we consider here are more expressive, yet also more challenging to sample than auto-regressive models such as GPT [37] , as we discuss with an example in appendix A.3. Here, we adopt the sampling approach of Goyal et al. [21] which is based on the Metropolis-Hastings (MH) algorithm. We first select 620 examples of the test set at random. For each input sequence (s 1 , . . . , s L ), we define its score in terms of the masked conditionals used for training the masked language models, E(s 1 , . . . , s L ) = -L i=1 h i • s i , with h i being the logits as defined in section 2. Then, a MH sampler is run on the joint distribution defined as p(s 1 , . . . , s L ) ∝ exp [-E(s 1 , . . . , s L )]. Starting from an initial sequence (s t=0 1 , . . . , s t=0 L ), at each step a local move is proposed, in which the i-th token is masked and replaced with s t+1 i ∼ p mlm (s i |s t ̸ =i ). The move is then accepted according to the standard MH criterion. In this framework, the conditional distributions obtained after training the transformer with MLM are used as the transition kernel of the MH algorithm. In fig. 3b ), we show the evolution of the mean score after each sweep of the sampling procedure, in the case of the standard BERT transformer and the four layers factored architecture. In table 1, we finally show how one randomly selected part of a story from the test set of TinyStories is modified after 20 sweeps of Monte Carlo sampling associated to two and four layers of factored attention, and four layers BERT architecture.Validating the clones As a sanity check of the statistical consistency of the clones, we trained a standard four-layer transformer encoder on TinyStories, but tested it on the clones generated after 20 and 70 MH sweeps, see panel c) of fig. 3 . For both the clones generated from a four-layer factored architecture and from a standard four-layer transformer encoder, we can check that the test loss after 20 and 70 sweeps is really similar, thus we decided to consider the samples obtained after 20 sweeps reliable clones of the original data set.Sequential learning of many-body interactions by BERT As already anticipated, we use these clones to investigate how a standard transformer encoder trained on MLM learns the different degrees of interactions among input words. Our results in fig. 1 already suggested a sequential learning of the various degrees of interaction among words over time for the WikiText-103 dataset. In fig. 4 , we show an additional experiment where we trained the same transformer (for more details about the architecture see appendix A.4) on a different data set, Tinystories [22] , finding the same sequential learning dynamics. We show this effect in the left panel of fig. 4 , where we plot the test loss of the four-layer BERT model of fig. 1 trained on TinyStories, while tested on clones with a fixed maximum Figure 4 : BERT models trained on masked-language modelling learn increasingly higher-order interactions during training. Left panel: In an experiment analogous to the one shown in fig. 1 , we show the test loss of a standard BERT-like transformer encoder trained trained on the TinyStories data set [22] and tested on clones of this data set with a truncated maximum degree of many-body interactions between tokens. The inset shows the corresponding test accuracy. We show the average over five different training runs, all starting from the same initial condition. The shaded area indicates one standard deviation. Right panel: An alternative way to visualise the data from the left panel is to plot the test loss at steps 10 4 , 3 × 10 4 , and 10 5 (blue, green and orange points respectively). This visualisation highlights the sequential learning of higher-order interactions, showing that for the clones derived from two-and four-layer factored architectures the loss saturates after 3 × 10 4 training steps, while on the clones derived from a six-layer architecture, as well as for the clone sampled from a BERT model, the test loss continues to decrease, as indicated by the black arrows.degree of interactions that we generated. After the initial plateau, which finishes at ∼ 5000 training steps, the loss curves on the various clones continue to overlap up to ∼ 6000 steps, despite the quick improvement in test error. This behaviour suggests that in the early stages of training, the model primarily learns low-order interactions without yet capturing high-order ones. Furthermore, in the final phase of training, the test error on the three-body clone, derived from a two-layer factored architecture, shows little change after 3 × 10 4 steps (blue curve). A similar behaviour happens for the clone obtained from a four-layer architecture (green curve). Instead, on the clone obtained from a six-layer factored architecture and from a BERT architecture, the loss continues to decrease (orange and red curves respectively) until the end of the training process, highlighting the sequential learning of the higher-order interactions by BERT. In the right panel of fig. 4 , we present an alternative visualization of the same test loss values on the clones at steps 10 4 , 3 × 10 4 , and 10 5 (blue, green and orange points respectively). This visualization highlights that, for the clone derived from a two-layer factored architecture, training beyond 10 4 steps does not change the test error significantly. In contrast, for the clone based on a four-layer factored architecture, the test loss shows significant improvement between 10 4 and 3 × 10 4 steps, but effectively saturates after 3 × 10 4 steps. Notably, for the clone derived from a six-layer architecture, as well as for the clone sampled from a BERT model, the test loss continues to decrease throughout the entire training process, as indicated by the black arrows.Sequential learning of many-body interactions by GPT Finally, we extend the analysis to the learning dynamics of an autoregressive transformer model trained for the next-token prediction task. Specifically, we examine a two-layer GPT-Neo model, as also investigated in the TinyStories paper [22] , training it on the original dataset and evaluating its performance on the different clones (see appendix A.4 for full experimental details). While the loss curve shows a single decay of the loss, the sequential learning is again clearly apparent, as shown in the left panel of fig. 5 . An alternative visualization of the same data presented in the left panel of the same figure reveals that the test loss on clones with low-order interactions saturates within the first 500 training steps, while the test loss on clones describing higher-order interactions continues to decrease throughout the entire training period. Crucially, these findings demonstrate the applicability of our framework also to autoregressive models, which are the foundation of state-of-the-art large language models.  
limitations: 
 The biggest limitation of our methodology is associated to the sampling procedure we used to generate clones. Although we used state-of-the-art methods to sample bidirectional models [21] , the presence of residual correlations cannot be excluded, even if we carefully tested the convergence of the results shown in the figures. Improving current sampling methods is an important objective for future work, but it is beyond the goals of this paper. On the theoretical side, the analytical computations of section 3 assume the embedding dimension to be d = 1, which is a simplification, and they do not give a quantitative estimate of the timescales over which the interactions of different orders are learnt in the numerical experiments of fig. 2 , or the number of samples ("sample complexities") required to learn them. Estimating such sample complexities is an active research topic even in the case of vanilla single-and two-layer neural networks [38] [39] [40] [41] [42] [43] 14] ; adapting these techniques to attention-based neural networks is an intriguing challenge for future work. 
concluding perspectives: 
 We have shown that transformer architectures trained on masked language modeling and next-token prediction tasks are able to learn the many-body interactions that govern natural text in a sequential way, going from lower to higher orders as training progresses. We showed this using a new approach to assess simplicity biases by first generating clones of a NLP data set, where the maximum level of interaction among tokens is limited to a specific order. Testing a transformer on these clones while training it on the original data set provides insights into the sequential learning patterns of the transformer. We conducted numerical simulations to show that, when training a transformer on the WikiText-103 and TinyStories data sets, the test loss curves for clones with low-degree interactions (e.g., three-body) exhibit plateaus after a fraction of the total training steps. In contrast, higher-order interactions exhibit continuous test loss reduction throughout the entire training procedure. The key element for generating the clones is the use of architectures combining factored attention layers [17] with the square activation function; the depth of these architectures controls the degree of the manybody interactions captured. In the future, it will be intriguing to extend our approach for constructing clones and for analysing the learning dynamics of transformers to different data modalities, like protein sequences [44, 32, 45] . On the theoretical side, developing a quantitative treatment of the sequential learning dynamics of (deep) transformers looms as an intriguing challenge.A.1 The data sets TinyStories is a synthetic data set available on Huggingface [46] and generated by GPT-3.5 and GPT-4, which contains a collection of stories that typical 3 to 4-year-olds usually understand [22] . This data set was recently introduced to facilitate the development, analysis and research of LLMs.In particular, it can be used to train and evaluate LLMs much smaller than state-of-the-art models, yet still being able to produce fluent stories with consistent meaning, thus demonstrating reasoning capabilities. This has been shown even for models with one layer and below 10 million total parameters [22] . In the same spirit of Ref. [19] , we use the full-sentences training scheme for MLM settings. Therefore, we first concatenate all the stories of the data set, tokenising the resulting text with the ByteLevelBPETokenizer (as implemented in Huggingface) and setting the vocabulary size to |V| = 10000. We then extract sequences of tokens of length L = 128 [47] ; notice that, in the full-sentences scheme, the sequences may contain sentences that are unfinished or that cross the boundaries of different stories. In the training phase, we adopt a strategy wherein, for each sequence, we randomly mask 15% of the tokens individually, as opposed to a simultaneous masking approach.WikiText We also used the WikiText-103 data set [20] , which is another standard NLP data set. Similar to tinystories, we tokenized using the ByteLevelBPETokenizer, setting the vocabulary size to |V| = 500 and considering sequences of length L = 64 tokens.We exploit the generative model described in section 2 to generate clones of the original TinyStories data set (see appendix A.1 for details on how we prepare the data set). The clones are meant to approximate the underlying TinyStories distribution with increasing fidelity. In particular, as already discussed, by stacking n layers of factored attention, we can generate copies of TinyStories whose maximum degree of interaction is n + 1. To enhance training stability, we normalise the activations through Layer Normalisation. Each layer of factored attention involves two matrices of parameters, A ∈ R L×L and V ∈ R |V|×|V| . In our setup, the size of the vocabulary is |V| = 10000, thus the matrix V would involve 100 millions of parameters. To reduce this number, we perform a finite-rank approximation of the value matrix, writing V = AB, with A ∈ R |V|×r and B ∈ R r×|V| . In practice, we choose r = 400. In this setting, we train two, four and six layers of factored attention via MLM. These models can then be sampled to generate three distinct clones of the original TinyStories data set, with a different degree of interaction among tokens.Auto-regressive architectures [37] are trained to learn the set of distributions {p(s i |s <i )} L i=1 and can be efficiently sampled. These models are associated to a joint over the tokens through the chain rule. On the other hand, bidirectional architectures model are trained to learn the set of conditionals {p(s i |s ̸ =i )} L i=1 . One may be tempted to sample these models with a Gibbs sampler, however this can result in an inconsistent sampler, where the conditionals may not be associated to any valid joint distribution. To understand why this could be the case, consider two discrete random variables X and Y , taking respectively m and n different possible values. The conditional distributions are defined in terms of two matrices A ij = P (X = x i |Y = y j ) and B ij = P (Y = y j |X = x i ), subject to the constraints i A ij = 1 for every j and j B ij = 1 for every i. The conditionals A and B are compatible if and only if two vectors τ i = P (X = x i ) and η j = P (Y = y j ), defining the marginals, exist such that τ i B ij = η j A ij for every i and j. The latter is a linear system, which does not always have a solution (see Ref. [48] for more details and for the generalisation of the problem to more than two random variables). A possible solution is proposed in Ref. [21] , where the authors develop a sampling scheme based on the Metropolis Hastings (MH) algorithm, starting from the evidence that BERT can be used to reliably score a set of sequences [49] . In this approach, for each given input sequence (s 1 , . . . , s L ), its score is defined as E(s 1 , . . . , s L ) = -L i=1 h i • s i , with h i being the logits as defined in section 2. Then, the MH sampler is run on the joint distribution defined as p(s 1 , . . . , s L ) ∝ exp (-E(s 1 , . . . , s L )). Starting from an initial sequence (s t=0 1 , . . . , s t=0 L ), at each step a local move is proposed, in which the i-th token is masked and replaced with s t+1 i ∼ p mlm (s i |s t ̸ =i ). The move is then accepted according to the standard MH criterion. Hence, in this framework, the conditional distributions obtained after training the transformer with MLM are used as the transition kernel of the MH algorithm. In order to generate the clones of TinyStories, we choose randomly 620 examples of the test set and evolve them for 20 MH sweeps (2560 steps).A.4 Architecture and the training procedure BERT We use a standard BERT-like transformer encoder [18, 19] with input-dependent attention weights (thus using queries and keys) and MLP after each self-attention layer. We use skip connections and put LayerNorm inside each residual block, thus adopting the successful Pre-LN scheme introduced in Ref. [50] . We choose the number of layers equal to 4, number of heads equal to 16, embedding dimension equal to 768 and size of the hidden layer of the MLPs equal to 1536. Similar hyper-parameters are also used for the models trained in Ref. [22] . We choose SGD for the optimiser, setting the batch size to 1024. We start with a learning rate of 0.03, annealing it with a cosine decay scheduler. We trained the model for 10 hours on eight A100 GPUs.GPT For our experiments with next-token prediction, we trained a two-layer GPT-Neo model, as also investigated in the TinyStories paper [22] . We set the hyperparameters of the network to be: number of layers equal to 2, number of heads equal to 16 and embedding dimension equal to 768. We used the AdamW optimiser, with the same hyperparameters used in Ref. [22] , namely learning rate equal to 5e -4, exponential decay β 1 = 0.9 and β 2 = 0.95. The weight decay was set to 0.1 and we used a batch size equal to 128.Solving for A (2) in the case A (1) (t) = 0We aim to solve the dynamical equations presented in Eqs. 9 of the main text, specifically in the scenario where the first layer is skipped, which corresponds to setting A (1) (t) = 0. Under this condition, the dynamics of the second layer are described by:EQUATION.To solve this differential equation, we first write C = U T DU with D a diagonal matrix with diagonal elements D j and U U T = I since the covariance matrix is symmetric. Thus we haveȦ(2) 0j + D j Ã(2) 0j = hj ,where we have definedEQUATION0j and hj = U h j . In this new basis, the equations are decoupled. With the initial conditions Ã(2) 0j (t = 0) = Ã(2) 0j (0), the solution is:Ã(2) 0j (t) = hj D j + Ã(2) 0j (0) - hj D j e -Dj t .It is easy to verify that, in the original basis, the stationary solution is A(2) 0j (t → ∞) = (C -1 h) j as discussed in the main text.Consider the analytical setup presented in section 3 of the main text. The mean square error loss associated with the model in Eq. 7 is:L(A (1) , A (2) 0 ) =   x 0 - j A (2) 0j x j - jkl A (2) 0j A (1) jk A (1) jl x k x l   2where b i = B 0i . We want to show that, in the case of Gaussian distributed data, for any values of the parameters b i it is bounded from below by the value it has in A (1) = 0. Taking the expectation, we get:EQUATIONjl A(1)EQUATIONEQUATION0 alone, and we defined Γ = A (1) CA (1) . For any possible choice of the vector A(2) 0 , we have:L(A (1) | A (2) 0 ) ≥ L(A (1) = 0 | A (2) 0 ) = f (A (2) 0 ) . Indeed, j A (2) 0j Γ jj 2≥ 0 for all choices of A (1) . Moreover, jj ′ b j ′ A(2) 0j Γ 2 jj ′ ≥ 0, since all the matrix Γ 2 jj ′ are non-negative, therefore all its eigenvalues are non-negative.NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide technical information to reproduce the results in the caption of the figures and in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: See captions of the figures.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We don't see any possible societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: We don't develop any model which have risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification:Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: This paper studies the simplicity bias in BERT-like models, by asking how transformers learn high-order interactions between tokens. The reviewers all appreciated this work, in particular the importance of the relatively underexplored problem, the novel (theoretical) insights, and the clarity. Main concerns regarding framing and some claims not being supported, overhead, missing related work, and the quality of the results have largely been addressed in the rebuttal.

Reviews: review: paper_summary: paper_summary: This paper shows that transformers sequentially learn high-order interactions between input tokens during the training process, which echoes the simplicity bias prevalent in neural networks. Specifically, the paper proposes a method to extract certain orders of interactions from the original training set, and construct a new dataset that contains purely low-order interactions. Then, the new dataset is used to evaluate how well a transformer learns certain orders of interactions.
strength_weakness: strength_weakness: 1.	The main claim of the paper “sequential learning increasingly higher-order interactions by BERT” is not well supported by the results in Figure 4. The term “sequential learning” implies that (i) at the early stage of training, the model first learns low-order interactions *without learning high-order interactions*, (ii) at the later stage of training, the model learns high-order interactions. However, although Figure 4 can support the above point (ii), but it cannot support point (i). The reason is as follows: at epoch 3, the testing loss on 7-body clones is already significantly lower than the testing loss on 3-body clones, which means that the model already learns part of the 7-order interactions at the early training stage. This further indicates that the model learns 3-order, 5-order, and 7-order interactions at the same time in the early stage of training.

2. Figure 4 is only tested on one instance, thus lacking statistical significance. 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.

3.  The paper is hard to follow due to inconsistent and undefined notations.

(1) I suggest adding a preliminary section with a formal definition of what “factored attention” means. There should be at least an equation showing how factored attention operates given a sequence of input tokens, what parameters it contains, etc.

(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.

(3) According to Line 84 “each token $s_i$ is an element taking values in a discrete vocabulary $s_i\in \\{1,...,|\mathbb{V}|\\}$”, the equation $s_{i \alpha}^\mu=1$ in Line 93 means that we are always measuring the probability of generating the first token in the vocabulary. This seems unreasonable.

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

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

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

4.	Figure 1 is misleading. The input sentence is shown as a folded tape, which highly resembles the structure of amino acid. However, this paper has no connection with amino acid, nor does it discuss about the spatial position of each token. Thus, I think this design is quite redundant. Why not simply write the input sentence on a horizontal line?
questions: questions: Figure 2(c) shows the mean squared displacement (MSD) of the weight for a 3-layer transformer with *factored attention* during training. It shows that the weight of the last layer first begins to change, then the second layer, and at last the first layer. I wonder if this phenomenon can also be observed on a transformer with *standard attention*.
limitations: limitations: Limitations are properly stated in Section 5.
paper_summary: paper_summary: The paper investigates the simplicity bias in BERT-style Transformers trained with MLM. The study reveals that these models initially learn simpler interactions between tokens and gradually learn higher-order interactions. This finding aligns with simplicity biases observed in many neural network architectures. The authors develop a method to create dataset "clones" that capture token interactions up to specified orders -- which they called many-body interactions. This allows detailed analysis of the learning dynamics in Transformers.
strength_weakness: strength_weakness: **Computational overhead**. The method to create dataset clones and analyze the learning dynamics is **computationally intensive**. It will potentially limit its applicability to larger datasets or more complex models. I would like to hear the authors proposals on how to deal with this issue.

**Generality** While the study focuses on BERT-style Transformers and the TinyStories dataset, the generalizability of the findings to other Transformer architectures or more diverse NLP tasks remains to be demonstrated. One thing to note, that Transformers are also applied to vision tasks, speech tasks, etc. Would like to hear how this technique can generalize to studying transformers on non-language tasks and if the conclusion would be similar. Similarly, how this would extend to decoder-only models such as GPT, remains unknown.
questions: questions: It's a very clearly written and presented paper, I don't have further questions to ask.
limitations: limitations: n/a
paper_summary: paper_summary: This paper demonstrates that masked language models like BERT approximate distributions with increasing complexity in terms of the number of interactions, as tracked over the course of training. They do this by approximating lower order interactions with provably limited models.
strength_weakness: strength_weakness: *Post rebuttals: The authors have expressed a willingness to rewrite the misleading framing and expand their literature review, so I have raised my score from 4 to 6.*

The paper should be updated to reflect related literature in language modeling or NLP. Here are a few other examples:
- Language models start out by behaving like ngram models in their grammatical capabilities and eventually begin to act more like humans https://aclanthology.org/2022.acl-long.568.pdf
- there are well documented dependencies including the famous checkmate in one emergence example from https://arxiv.org/pdf/2206.04615
- RoBERTa achieves high performance on most linguistic benchmarks early in pre-training, but more complex tasks require longer pre-training time. https://aclanthology.org/2021.findings-emnlp.71/
- LSTMs start by learning representations that correspond to simple groupings like part of speech and eventually learn more complex language modeling features https://aclanthology.org/N19-1329/
- Section 4 explores increasingly complex ngram learning in https://arxiv.org/pdf/2402.04362
in context learning approximates shorter ngrams before longer ngram distributions https://arxiv.org/abs/2402.11004

The core issue with this paper is that at no point does it describe a simplicity bias. While conceptual dependencies and progressively complex learned distributions are clearly related to simplicity bias, these things are not the same. A simplicity bias requires one to show that the model fails to learn something difficult because it has learned something easier, but arguably what is described here might actually be a sequence of strict dependencies between interaction levels. In order to demonstrate a simplicity bias, the authors would have to show that the model fails to learn higher order interactions because of the availability of lower order interactions, but the lower order interactions are available regardless of whether the higher ordinary interactions are present so there are no conditions to contrast in this way. 

Some ways of demonstrating simplicity bias are:
- demonstrating that introducing a simple rule prevents the model from learning a more complex rule https://arxiv.org/abs/2011.09468 https://arxiv.org/abs/2006.07710
- demonstrating that suppressing a simple rule allows the model to get better at learning a more complex rule https://arxiv.org/abs/2309.07311 (specifically is focused on MLMs as well)
- Correlating the accessibility of a rule and the degree to which the model ultimately captures a different rule https://arxiv.org/abs/2006.12433
- analytic proofs of a preference towards simplicity https://arxiv.org/abs/1805.08522
It is not sufficient to simply demonstrate that the model happens to learn simple things and then more complex things later on.

I want to be clear that this paper still has value, as it demonstrates yet another way in which these models learn increasingly complex distributions, which is a different but related thread of the literature. However, the framing is misleading and I do not believe it is ready for publication without a substantial rewrite.
questions: 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?

Can you guarantee that the difference provably limited models actually differ by nothing but their limitations? What about random variation?
limitations: limitations: nothing obvious
paper_summary: paper_summary: The paper investigates whether transformers trained on natural language data using masked language modelling (MLM) exhibit a "simplicity bias" - learning increasingly complex interactions between input tokens over the course of training. The work provides both theoretical and empirical evidence for a simplicity bias in transformer models trained on text data.​​​​​​​​​​​​​​​​

The key contributions of the paper are listed below:

1) Introduction of a novel framework to study the learning dynamics of transformers:
- They develop a method to create "clones" of natural language datasets that capture token interactions up to a specified order.
- This is done by using transformers with factored attention layers and quadratic $x^2$ activation functions.

2) Demonstration of sequential learning in transformers: By using their cloning method on the TinyStories dataset, the authors show that a BERT-like transformer sequentially learns increasingly higher-order interactions between tokens. More specifically, lower-order interactions (e.g. 3-body) are learned earlier, while higher-order interactions (e.g. 7-body) continue to be learned later in the training.

3) Theoretical analysis: the authors provide analytical insights into the learning dynamics of multi-layer factored attention models, showing how different layers contribute to learning interactions of increasing complexity.

4) Evidence on language data: the authors demonstrate that interactions between words in natural text (specifically TinyStories) are highly many-body, up to at least 7th degree.
strength_weakness: strength_weakness: 1) *Uncommonly used architecture*.

As a first step towards understanding the learning dynamics of Transformer-based language models, I understand the model choice of the authors. However, this also limits the contributions and impact of this work. From my perspective, the "Transformers" in the title might better be replaced by "BERT-like Language Models", as the work investigates only a particular implementation of transformers on only language modelling. So, "learning dynamics of transformers" might be far bigger than the research scope of this work.

Moreover, for language modelling, decoder-only Transformers are wider applied these days. The $x^2$ activation function studied in this work is not common for language modelling either. These two factors further limit the contributions and impact of the work.

2) *Higher orders are not straightforward to sample.*

As mentioned by the authors in Section 5, the time complexity of the sampling algorithm for creating clones is not linear with respect to the number of samples. Thus, it would become increasingly harder to sample clones of higher-orders. This is a major bottleneck to apply the work at larger scales.

3) *The empirical evidence is not strong enough.*

To me, only one model on one specific dataset is not very convincing to empirically verify the proposed simplicity bias hypothesis. It would be much more convincing if the authors can show empirical evidence on more diverse language tasks, e.g. machine translation, and etc.
questions: questions: 1) What does the "last layer" in line 163 mean?

It's unclear to me how is the "last layer" in line 163 defined. I kind of assume that it means the (bottom) layer that's closest to the inputs, a.k.a. the secondly activated layer in the model. The authors need to clarify the meaning of it to avoid unnecessaty confusion.
limitations: limitations: Please refer to my weaknesses section.

Comments: title: Reframing
Reply
Response
comment: Thank you for the clarification. The main issue lies in the inconsistent label of the x-axis in Fig.1 and Fig.4, which causes my confusion. I'm happy to see this being resolved in the future version of the manuscript, and I'm raising my score.
Thank you for engaging with our rebuttal, we appreciate your effort.

> Regarding the notation

Thanks again for spotting these typos; we have already changed the indices in those sums for clarity.

> Further questions regarding Fig. 4 and Fig. 1

Thank you for following up on this. 
- We drew Fig 4 using the same data as Fig 1. In this experiment, 1 epoch = 3000 steps of SGD. Fig 1 shows that the test losses of the transformer on the different clones are different after 9k steps (=3 epochs). Hence the blue dots in Fig 4, which represent the performance of the transformer on the different clones (from left to right) after three epochs, have different values. 
- Sequential learning in sense (ii) of your definition (learning higher-order interactions at a later stage of training) is borne out by Fig 4 since the performance on the three-body clone does not improve at later epochs (left-most column), while the performance on the five-body clones continues to improve until ~epoch 10, and improves on the seven-body clone for even longer. 
- Sequential learning in sense (i) of your definition (learning low-order interactions without learning high-order interactions) can be seen in Fig 1 during the first two epochs, where the loss curves on the various clones continue to overlap despite the quick improvement in test error.

Does this answer your questions? We will add epochs to the $x$-axis of Fig 1 for increased readability, and explicitly state the number of steps per epoch for the experiments. Please let us know if there are other changes that could help the reader, or if you have any other questions.
Thank you for your point-by-point response. I appreciate the new experiment on the transformer with standard attention. Regarding the notations, I encourage the authors to clarify all the mentioned issues in the future manuscript, especially replacing the dummy index $\mu$ in Eq. (3) and $p$ in Line 134 with other appropriate symbols.

I have further questions regarding Fig. 4 and Fig. 1. The authors respond to my concern about the insufficient empirical support of Fig.4 for the main claim by referring to Fig.1. However, I find that the results in Fig.1 and Fig.4 seem to conflict with each other. In Fig.1, the test losses on 3-body, 5-body, and 7-body clones are almost the same in the early stage of training (before Step 5000), while in Fig.4, the test losses on 3-body, 5-body, and 7-body clones significantly differ in the early stage of training (epoch 3). Could you provide a further explanation?
Thank you first of all for your continued engagement, and for increasing your score, we really appreciate it. Your final suggestion of avoiding the term bias altogether is a good one; we will change the title in that sense.

All the best, The Authors
I would perhaps use a framing that's not oriented around the term bias, which carries implications around the simpler solution being an actual attractor (eg, you could use snowclone boilerplate title of "Transformers learn feature interactions of increasing complexity"). However, given that you're willing to change the misleading and confusing framing, I will be raising my score, as I think that the results and method are overall valuable.
Thank you for engaging with our rebuttal so quickly, we appreciate it. 

> Re: framing

We now realise that the term simplicity bias can be misleading; since we strive for maximum clarity of our paper, we are considering changing the framing of the paper by changing its title to “A bias towards low-order interactions in the early learning dynamics of BERT-like language models”, and generally changing “simplicity bias” to the more descriptive “bias towards low-order interactions” throughout the text. Would that address your criticism?

> Is it empirically true that models trained under different initial conditions with the same set of limitations exhibit all of the same behavior that you measure?

Yes; in the course of our work, we found that if we generate clones from networks with the same interaction-limited architecture trained from different initial conditions, the performance of a given transformer tested on the different clones will be the same to within the variation in performance of the same transformer trained from different initial conditions.
> The nature of simplicity bias 

I do not feel that this response addresses my criticism, which is that none of the results justify framing the paper as studying simplicity bias. Although training dynamics of increasing complexity are clearly related to simplicity bias, they are not evidence of simplicity bias and therefore the paper should not be framed around describing them as such. This criticism is fairly mild, but I consider it to be crucial: I believe that the title as written is making a claim that is unrelated to the results and not supported by evidence in the paper. I don't think that this paper should be published with a misleading title or framing, although the results themselves are interesting. This issue is easily fixed; I would easily raise my score for a version of the paper that had a different title and discussed simplicity bias as related work rather than as a description of the results here.

> 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.

Is it empirically true that models trained under different initial conditions with the same set of limitations exhibit all of the same behavior that you measure?