A distributional simplicity bias in the learning dynamics of transformers
Authors: Riccardo Rende
Federica Gerace
Alessandro Laio
Sebastian Goldt
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformers
Natural Language Processing
Sequential Learning
Simplicity Bias
Primary area: natural_language_processing

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.

A distributional simplicity bias in the learning dynamics of transformers
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt
International School for Advanced Studies Trieste, Italy
(rrende, fgerace, laio, sgoldt]@sisse
Abstract
The remarkable capability of weerparaaeterised neural networks to generalise effectively has been explained by invoking a "simplicity bias": neural networks prevent overfitting by initially learning simple classifiers before progressing more complex. non-linear functions. While simplicity biases have been described theoretically and experimentally in fced-forward networks for supervised learning, the extent to which they also explain the remarkable success of transformers trained with supe-sup 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 learn high-degree interactions. To conduct this analysis, we develop a procedure to generate clones of a given natural language data set, which rigorously capture the interactions between tokens up to specified order. This approach opens up the possibilities of studying how interactions of different orders in the data affect learning, natural language processing and beyond.
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 func- tions [[55.. These simplicity biases have been analysed theoretically in linear networks (6-3), in non-linear shallow neural networks I 9 and in the NTK regime (1-H2). The viability of such 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 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 "clones", we show that deep transformers trained natural language also exhibit a simplicity bias: they sequentially learn increasingly higher-order interactions among tokens.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

WikiText- 103
Data: many- body interactions
Clone: up 3-body
fustcrod losss faztored, lavee
Test loss on clones
Factorec oftectttttnn pftection tisi
b O + P
Factor aod ottemtion
BERT
was Hey
say LEanu
place
10e
Epoch
Figure 1: Transformers learn increasingly higher-order interactions from their data. Left: 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 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² 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 [T8] 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 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
Monte Carlo techniques [21] to obtain the cloned data sets. The clones thus form 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 [8] 19] with four attention blocks. On the right of fig. [ we show the test loss of this architecture during MLM training on WikiTexx- 103, but testing on various clones of WikiText- 103 with a truncated maximum degree of many-body interactions between tokens, as well as 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 Rv 20 epochs. Beyond this threshold, BERT exhibits different generalisation behaviour 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 benchmark NL.P 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
Our work shows that the many-body interactions between tokens natural text are key for learning, and we note that approach naturally extends to other data modalities.

---

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 learn- ing 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 ($1..... .s....... sL). where each token is an element taking values in discrete vocabulary: 8, .......III.. When masking the token in position i, the sequence (S1,..., S; O.... .......... sL) is given as input to the transformer, which embeds the input sequence into a collection of vectors (Zi...... Tr)) each of them living in a d-dimensional space. Then, this sequence of vectors pro- cessed by the self-attention layers and transformed into sequence of context-dependent vectors (y1,.... V2), where Vj E Rdd for all j. At the end, the output yi. corresponding to the masked input token, is mapped into a vector of the same size of the vocabulary: h; E RI The prediction Further related work on simplicity biases in transformers Bhattamishra et al. [23] found 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 ef al. [25] investigated 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, prior research has investigated the progressive increase in complexity of the distribution learned by a transformer from its data over time. More recently, Cagnetta 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-Brn et al. (30 demonstrated 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//ddb.orr/115522///22 23329///onodoo3199923888 13992398
2 Expressing many-body distributions with factored attention
Pmim ($||$ = softmax(h.) can be finally interpreted as the probability distribution over all possible values the masked word can assume within the vocabulary. The transformer is trained to minimise
the MLM cross-entropy loss, L = - Em=1 Eta (mm) log (Pmlm(#1a (+(rr) (142)) (5 ) where M denotes the total number of examples in the dataset, and m indexes individual samples. The index i refers to the masked token position, while represents the feature index. During training, the spins are cyclically masked one by one in each sequence.
The key ingredient 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 al. (17) recently considered a simplified attention mechanism, the so-called factored attention, where the attention weights are input-nndependent and the input sequence is processed as: V/o Σラ= Aij EM Va33ja, where A R and V R are matrices of trainable parameters. In this way, the prediction of 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):
M V
Punm
(Sia Pulm(Siaa 1|84s) softmax
Σ jz
44
V888
(1)
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 [7] when sampled according to the joint distribution:
1 p(s) = - exp
-
-E i.j
JigU (aUB888883
(2)
o
3

---

where Jıy describes the interactions among positions along the sequence, Uas 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 I. and the Greek indices to run over the size of the vocabulary [V]. As can be easily checked, the expression in eq. [ exactly matches the conditionals associated with the two-hody joint distribution in eq. (2) Remarkably, 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 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 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² as non-linearity between them and skip connections, see fig- (1) for sketch of this architecture. By writing down the conditional distribution corresponding this two-layer model:
p(s/a 111))t) softmax
Σ jB
( 0D
08
Σ by
A By
jk
Sky
Aw Ji Σ 4p
Sta
j9
(3)
notice that the second layer expresses three-body interaction, parameterised by specific combination of the attention weights of the first and second layer through the tensor Jık (2) f A] (Ü similarly expressed in terms of V). In particular, the first element of the sum (1)
Σ, A, 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:
p(sia= 1|5)4r) softmax
3
3 jk
Σ BY
JaykU
Generalising these observations to a generic number of factored attention layers, we conclude that by stacking layers of factored attention with x² activation function can infer interactions up to order 2n-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.
3 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:
1-EE ijki a.qu
(8)
7 exp
JujuUaShy 8188818
(5)
Here, we ensure the couplings to be symmetric tensors by considering set of vectors {1p} and setting Jishs = Ep- J'P We choose P = 200 and we sample of Bernoulli 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" J==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. [p). 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 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 model whose predictions correspond to the exact conditional of two and

---

s)
2.30
layer
layer
softmax
2.95
0.55
layer
2.20
3 byyers
layer
2.15
0.20
Self. Attention
2.10
bodies
MSD
0.15
2.00
0.30
1.90
bodies budies
0.06
Self-Attention
0.00
1.85
100
250
200
Ippor Emboodding
Epoch
Epoch
Figure 2: a) Multi-layer factored self-attention architecture with x:2 activation function. b) Test loss learning curves of one, two and three factored self-attention layers with x?² activation function. The models were trained a synthetic data generated from four-body Hamiltonian. The dashed horizontal lines correspond to the convergence value of the loss for two, three and four bodies energy based models trained on the same data set. c) Mean Square Displacement of the weights across different layers in three-layers factored attention architecture. In these experiments, the size of the vocabulary was [V] = 10 and the sequence length
20. We used training set of M = 25600 samples, training the models with SGD, choosing mini-batch size of 256. The initial learning rate chosen be 0.1.
three-body distribution, where the learnable parameters are the tensors in eq. ( and eq. (4) As shown by the red curve in fig. 2D), if we train single layer of factored attention on the four-body interacting sequences, the best we can do to exactly recover the best two-body approximation [17]. By adding one extra layer of factored attention with x² non-linearity, 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² activation function we can infer interactions up order 2"
reconstruct the best
Figure 2b) 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, initial plateau aligns with the best two-body approximation, followed by 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 illustrated in Figure where we display the Mean Square Displacement (MSD) of the weights for three layers of factored attention across training epochs, that is:
MtD 27 (t) - A
(t) (0)12 Ivp
(6)
with (t) and / (t) being the attention and the value matrix of the (-th layer at time t. As we can see, despite the simultaneous update of all network parameters during training, 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 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 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 the presence of skip connections, they appear linearly combined with the input tokens in the output prediction. In the next section we provide more formal analysis of this phenomenon.
3.1 Analysing the gradient flow for two-layer factored attention
To gain theoretical insights on the sequential learning behaviour observed in the previous numerical experiments, analyse the gradient flow of two-layers of factored attention with z2 activation
5

---

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 (zo. .1...... .t). 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 (ro mn)² Here, yo represents the transformer prediction about the masked token which, in this setting, acquires the following functional form:
Σ
+EE
A (1) (1) TA-
The gradient flow equations can then be written explicitly for the attention weights of both factored attention layers as:
A
To
(2) 2j -
Σ
Aoj 13
(2)
(x, + 13)
(8)
3
4 =
To - Σ
(2)) 53j -
Σ 1 x3
(1 (2) mi)a)))
(9)
where A = Σ, 4
Tj. From these equations, we can immediately see that, if we initialise the network with zero weights A = A = O. the first layer attention weights do not change in time, A 0, while the second-layer weights do, Ab # O, consistently with what we observe in fig. 2
(1)
We can gain some further intuition into which degree of interactions each layer learns by simplifying the data model further. In particular, from a zero-centered Gaussian distribution. By taking the expectation over the data, the correlations of odd number of random variables vanish and the gradient flow equations can be written as:
restrict to Gaussian design where tokens are now drawn
1
h2 - (CAb
(2) UK
[AAAR]
(10)
A
-2
(2) (2) Σ A0j E De
(11)
where we have defined h = B[ for 0 and Ct2 E[r,,y] for i, 0. In appendix B we further show that, if data are Gaussian distributed, the expectation value of the loss bounded from below by the value it takes for A O. Therefore, if one adds an l2 regularisation of arbitrarily small strength, A11) 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 the learning dynamics:
A = h, (CA(
(12)
whose solution describes an exponential convergence in time to (22 + oo) = (C- This stationary point reproduces the exact conditional distribution of 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, element) by and the block obtained with i, by C have: p(z0 (o)) = Ni h- E). 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- -an-Gussiann data distributions). the stationary point A(t) 0 is generally unstable. For example, considering the case of diagonal perturbation A(1) cōga around (1)) = 0. eq. 9 becomes (neglecting terms of order e?): Ahh
t6 = (a 0s
(2)
Took
(2) T
(13)
where we have taken the expectation with respect to generic data distribution and defined I Piya- E [x]z) In this case A (1) 0 is stable solution with respect to diagonal perturbation only

---

a) Train factored models on Tinystories
6) Sample factored models
c) Check generated clones
3100
BERT
BERT, 20 sweeps BERTT; 70 sweeps factored layers, 20 sweeps
factored 2 layers
3200
factored 4 layers
Test loss on clones
factored 6 layors
3000
Tost loss
factured layers, 70 sweeps
2800
BERT
factored layers
3
2600
2
2100
2-
2200
20000
401000
63000
20000
Step
Swoep
Step
Figure 3: Three steps for cloning a data set using a) Train factored-attention models on TinyStories. Test loss curves of different factored-attention based architectures trained on Ttorito and tested on TinyStories. Specifically, we consider architectures with two, four and six factored self self-attention layers with x² activation function. For comparison, also the loss of four-layers BERT is shown. b) Sample factored models. Mean score of a batch of sentences taken from the test set of the TinyStories data set and evolved with the Metropolis- Metropolis-Hasting sampling scheme described in appendix AA.3, c) Check generated clones. Test loss curves of standard four layers transformer encoder, trained on TinyStories and tested on clones generated after 20 and 70 Setropolis-Hasting sweeps. The clones were generated from a four layers standard BERT and from architecture with four layers of factored self-attention and z² activation function (associated with a nine bodies approximation of Tiny-Stories).
based generative models.
if the matrix Gik ( (Tou- Ej Titk T is negative definite (implying that the perturbation decays to zero as t goes to infinity). Therefore, the structure of the three-body correlations Fijk determines the stability of the stationary point A = O. Indeed, if, for example, one takes I' Ju-- yd 15zō we have G1 y0u A (Sue -1 which is not negative definite for a generic choice
of (2) This analysis allows concluding that there are conditions for which the solution A(") = is unstable when data contains interactions of order higher than second, which precisely in line with what we observe in fig. D)
4 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 real NLP datasets. Specifically, we analyze the TinyStories dataset [22], training a standard BERT encoder [8 using the masked language modeling objective, well 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- Stoyis, that can be sampled to generate "clones" of the original data with a truncated maximum degree of interaction. We now discuss experimental methods and analyse the results of our experiments; for full details on the methods, see appendix A
models of section 2 with varying number of layers on
The clones We exploit the generative models described in section 2 to generate clones of the original TinyStories data set (see appendix A. for details on how we prepare the data set). The clones are meant to approximate the underlying TinyStories distribution with increasing fidelity. We use 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). In
7

---

Once upon time there lived 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.
Original text
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.
Factored 2 layers
Once upon time there was boy named Tom. He saw small cat, could not play in the park with flowers pretty. One day, Lily smelled and nice. They wanted the ball.
Factored 4 layers
Once upon 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 puppy.
BERT
Table I: Sampling the clones. In the first row, we show part of sentence taken from the test set of TinyStories. The second, third and fourth rows show how the original text is modified after 20 sweeps of Monte Carlo sampling associated to two and four layers factored architectures and BERT architectures, respectively.
fig- 31), we show the test loss curves when training models with two, four and six layers of factored attention featuring the x² 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 standard BERT architecture [18] is also shown.
Sampling procedure 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-
(MH) algorithm. We first select 620 examples of the test set at random. For each input sequence ($1...., .L), we define its score in terms of the masked conditionals used for training the masked language models, F(s1,...., 81.) ((81,....88L)
Σ h; 8i, with h; being the logits as defined in section Then, a MH sampler is run on the joint distribution defined as p(#1,-.... &1. xx exp ( $ )] Starting from an initial sequence (s!=0
st=0). each step local move proposed, in which the i-th token is masked and replaced with of Pmim (st s5))). The move is then accepted according to the standard MH criterion. In this framework, the conditional distributions obtained after training the transformer with MIM are used as the transition kernel of the MH algorithm. In fig- B)), 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 ta 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 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 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 consider the samples obtained after 20 sweeps reliable clones of the original data set.
Sequential learning of mony-body interactions by BERT As already anticipated, we use these clones to investigate how standard transformer encoder trained on MLM learns the different degrees of interactions among input words. Our results in fig. [ alrrady suggested 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) 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. Trained on StnyStories, while tested on clones with a fixed maximum
8

---

10
factored, layers factored, layers
Test loss clones
8
factored, layers
Test. loss clones
BERT
6
1.76
step 10
step
Step
step 10"
10
10a
10
10
10
factored laayers
factored
factored 6 layers
BERT
Step
4 layers
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. D we show the test loss of 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 3 x 10*, and 10 (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 x 10 training steps, while on the clones derived from six-lay layer architecture, as well as for the clone sampled from a BERT model, the test loss continues to decrease, indicated by the black arrows.
degree of interactions that 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 two-layer factored architecture, shows little change after 3 x 10a steps (blue curve). A similar behaviour happens for the clone obtained from four-layer architecture (green curve). Instead, the clone obtained from 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 3 10 and 10t (blue, green and orange points respectively). This visualization highlights that, for the clone derived from two-layer factored architecture, training beyond 10 steps does not change the test error significantly. In contrast, for the clone based on four-layer factored architecture, the test loss shows significant improvement between 10s and 3 x 10d steps, but effectively saturates after 3 x 10a steps. Notably, for the clone derived from six-layer architecture, as well as for the clone sampled from 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 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 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 tfate-ff-the-ar large language models.

---

factored, layers factored, layers
Tost loss clones
4.00
factord,, layers
Test loss on clones
3.75
BERT
5.25
3.00
2.75
step 100
8.00
step 500
2.20
step 2000
10P
10
10a
10a
factored 2 layers
factored 4 layers
factored
BERT
Step
6 layers
Figure 5: GPT models trained on next token prediction tasks learn increasingly higher-order interactions during training. Left panel: Test loss of two-layer GPT-Neo model from [2 trained on TinyStories set and tested on clones this data set with a truncated maximum degree of many-body interactions between tokens. Right panel: We repeat the analysis of fig. 4 for the GPT-Nco model trained using next-token prediction, showing again clear signatures of sequential learning. The results are consistent with respect to different random seeds.
Limitations
The biggest limitation of our methodology is associated to the sampling procedure we used to generate clones. Although we used state-of-ther off-the-art ait-art methods to sample bidirectional models [2]] 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 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 [4]; 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 sequential way, going from lower to higher orders as training progresses. We showed this using 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 while training it on the original data set provides insights into the sequential learning patterns of the transformer. conducted numerical simulations to show that, when training a transformer on WikiText- 103 and TinyStories data sets, the test loss curves for clones with low-degree interactions (c.g., body)-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 many- body 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 an intriguing challenge.
specific order. Testing a transformer on these clones
Acknowledgements SG acknowledges co-funding from the European Union NextGenerationEU, in the context of the National Recovery and Resilience Plan, Investment PEI Project FAIR "Future Artificial Intelligence Research", and in the framework of the PRIN Project SELFMADE (code 2022E3WYTY CUP G53D230007880011,
10

---

References
[1] D. Saad and S.A. Solla. Exact Solution for On-Line Learning in Multilayer Neural Networks. Phys. Rev. Lett., 74421)44377-4444 1995.
[2] F. Farnia, J. Zhang, and D. Tse. A spectral approach to generalization and optimization in neural networks. In ICLR, 2018
[3] Guillermo Valle Pérez, Chico Q. Camargo, and Ard A. Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. In International Conference on Learning Representations. OpenReview.net, 2019.
[4] Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin L. Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. SGD on neural networks learns functions of increasing complexity. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 3491-3501, 2019.
[5] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron C. Courville. On the spectral bias of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proc. of ICML, volume 97 of Proceedings of Machine Learning Research, pages 5301-5310. PMLR, 2019.
[6] Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In Yoshua Bengio and Yann LeCun, editors, Proc. of ICLR, 2014.
[7] A. Saxe, J.L. McClelland, and S. Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 1 2019.
[8] M.S. Advani, A.M. Saxe, and H. Sompolinsky, High-dimensional dynamics of generalization error in neural networks. Neural Networks, 132:428- 446, 2020.
[9] S. Mei, A. Montanari, and P. Nguyen. mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115633))677665--7777) 2018.
[10] Benjamin Bowman and Guido Montúfar. Implicit bias of MSE gradient optimization in underparameterized neural networks. In Proc. of ICLR. OpenReview.net, 2022.
[11) Benjamin Bowman and Guido Montúfar. Spectral bias outside the training set for deep networks in the kernel regime. ArXiv preprint, abs/2206.0299 022927 2022.
[12] Etienne Boursier, Loucas Pillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances in Neural Information Processing Systems, 35:20105-20118, 2022.
[13] Maria Refinetti, Alessandro Ingrosso, and Sebastian Goldt. Neural networks trained with sgd learn distributions of increasing complexity. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
[14] Lorenzo Bardone and Sebastian Goldt. Sliding down the stairs: How correlated latent variables accelerate learning with neural networks. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41 International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 3024-3045. PMLR, 2024.
[15] Nora Belrose, Quintin Pope. Lucia Quirke, Alex Mallen, and Xiaoli Fern. Neural networks learn statistics of increasing complexity. arXiv preprint arXiv:2402. 04362, 2024.
[16] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pages 2397-2430. PMLR, 2023.

---

[17] Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Mapping of attention mechanisms a generalized potts model. Phys. Rev. Res., 6:023057, 2024.
[18] Jacob Devlin, Ming-Wei Weeg-Wei&& Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
[19] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019.
[20] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[21] Kartik Goyal, Chris Dyer, and Taylor Berg-Kirkpatren Exposing the implicit energy networks behind masked language models via netropolis-hastings 2022.
[22] Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english?, 2023.
[23] Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in trans- formers and their ability to learn sparse boolean functions, arXiv preprint arXiv:2211.12316. 2022.
[24] Emmanuel Abbe, Samy Bengio, Enric Boix-Adsera, Etai Littwin, and Joshua Susskind. Trans- formers learn through gradual rank increase. Advances in Neural Information Processing Systems, 36, 2024.
[25] Bhavya Vasudeva, Deqing Fu, Tianyi Zhou, Elliott Kau, Youqi Huang, and Vatsal Sharan. Sim- plicity bias of transformers to learn low sensitivity functions. arXiv preprint arXiv:2403.00695, 2024.
[26] Zeyu Liu, Yizhong Wang, Jungo Kasai, Hannaneh Hajishirzi, and Noah A Smith. Probing across time: What does roberta know and when? In Findings of the Association for Computational Linguistics:
EMNLP 2021, pages 820-842, 2021.
[27] Leshem Choshen, Guy Hacohen, Daphna Weinshall, and Omri Abend. The grammar-learnin trajectories of neural language models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pages 8281-8297, 2022.
[28) Francesco Cagnetta, Leonardo Petrini, Umberto M Tomasini, Alessandro Favero, and Matthieu Wyart. How deep neural networks learn compositional data: The random hierarchy model. Physical Review X, 14(3):031001, 2024.
[29] Francesco Cagnetta and Matthieu Wyart. Towards theory of how the structure of language is acquired by deep neural networks. arXiv preprint arXiv:2406.00044 2024.
[30] Jérôme Brunn Marc Mézard, Emanuele Moscato, and Luca Saglietti. How transformers learn structured data: insights from hierarchical filtering. arXiv:2408. 15138, 2024.
[31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30, 2017.
[32] Nicholas Bhattacharya, Neil Thomas, Roshan Rao, Justas Dauparas, Peter K. Koo, David Baker, Yun S. Song, and Sergey Ovchinnikov. Interpreting potts and transformer protein models through the lens of simplified attention. In Biocomputing 2022, pages 34 2020.
[33] Luciano Loris Viteritti, Riccardo Rende, and Federico Becca. Transformer variational wave functions for frustrated quantum spin systems. Phys. Rev. Lett. 130:236401,
[34] Luciano Loris Viteritti, Riccardo Rende, Alberto Parola, Sebastian Goldt, and Federico Becca. Transformer wave function for the shastry-sutherland model: emergence of a spin-liquid phase. arXiv preprint arXiv:2311.1688 2023.
12

---

[35] Riccardo Rende, Luciano Loris Viteritti, Lorenzo Bardone, Federico Becca, and Sebastian Goldt. A simple linear algebra identity optimize large-scale neural network quantum states. Communications Physics, 7(1), August 2024.
[36] Riccardo Rende, Sebastian Goldt, Federico Becca, and Luciano Loris Viteritti. Fine-tuning neural network quantum states. arXiv preprint arXiv:2403.07795, 2024.
[37] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. https://nsighheiivees...ss ess-e
I ammazzonws..om//aggggggnndddd.d
2018.
[38] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. J. Mach. Learn. Res., 22(1), 2021.
[39] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptoties of feature learning: How one gradient step improves the representation. Advances Neural Information Processing Systems, 35:37932-37946, 2022.
[40] Alex Damian, Eshaan Nichami, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2024.
[41] Emmanuel Abbé, Enric Boix Adserå, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-sadde dynamics. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 2552-2623. PMLR, 2023.
[42] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at time. arXiv:2305. /8270, 2023.
[43] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.0005. 2023.
[44] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics, 38(8):2102- 2110, 02 2022.
[45] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Yu Wang. Llion Jones, Tom Gibbs. Tamas Feher, Christoph Angerer, Martin Steinegger. Debsindhu Bhowmik, and Burkhard Rost. Prottrans: Toward understanding the language of life through self-supervised learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10):7122-712) 2022.
[46] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: -tate-of-the-art natural language processing, 2020.
[47] Peter Izsak, Moshe Berchansky, and Omer Levy. How train bert with academic budget, 2021.
[48] Barry C. Amold and S. James Press. Compatible conditional distributions. Journal of the American Statistical Association, 1989.
[49] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi Bertscore: Evaluating text generation with bert, 2020.
[50] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Van Liu. On layer normalization the transformer architecture, 2020.
13

---

A Experimental details
A.1 The data sets
TinyStories is a synthetic data set available on Huggingface [46] and generated by GPT 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-ar 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. (99, 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 to V 10000. We then extract sequences of tokens of length L = 128 [47] notice that, in the full 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 simultaneous masking approach.
(as implemented Huggingface) and setting the vocabulary size
WikiText We also used the WikiText- 103 data set 20 which another standard NLP data set. Similar to tinystories, we tokenized using the ByteL.evlBPEEToenniz,, setting the vocabulary size to |
500 and considering sequences of length L 64 tokens.
A.2 Obtaining the clones
We exploit the generative model described in section 2 to generate clones of the original TinyStories data set (see appendix IA. 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 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 E R and V E R|VX In our setup, the size of the vocabulary is | 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 E Rillxr and B E R²* In practice, we choose = 400. In this setting, we train two, four and six layers of factored attention via MLM. These models can then be sampled generate three distinct clones of the original TinyStories data set, with a different degree of interaction among tokens.
A.3 Sampling procedure
Auto-regressive architectures 37 are trained to learn the set of distributions {p(s{s<<))) and can be efficiently sampled. These models are associated 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()
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 and n different possible values. The conditional distributions are defined in terms of two matrices Aij = P(X- = x; LY = 3j) and B,j = P(Y = yソ X = æx), subject to the constraints Σ, Aij = 1 for every j and Σ, Bij 1 for every i. The conditionals A and B are compatible if and only if two vectors T; P(X x;)) and πi P(Y y₁), ). defining the marginals, exist such that T, B, = 7) Asj for every and j. The latter is a linear system, which does not always have 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 sampling scheme based the Metropolis Hastings (MH) algorithm, starting from the evidence that BERT can be used to reliably score a set of sequences 44). In this approach, for each given
input sequence (si...., sL), its score is defined as E(s $1) = - - h; 8;, with hy being the logits as defined in section [ Then, the MH sampler is run on the joint distribution defined as p( p(81,...,81, 8r.)
exp (- ((.., sr.)). Starting from an initial sequence (sf=1) ...., ..=,), at
each step local move is proposed, in which the i-th token is masked and replaced with s
14

---

Pmlm (s:)s The move is then accepted according to the standard MH criterion. Hence, in this framework, the conditional distributions obtained after training the transformer with ML. are used the transition kernel of the MH algorithm. In order to generate the clones of TtnyStories, 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 standard BERT-like transformer encoder [8] 19 with input-dependent attention weights (thus using queries and keys) and MIP after each self-attention layer. We use skip connec- tions 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 768 and size of the hidden layer of the MLPs equal 1536. Similar hyper-parameters are also used for the models trained in Ref. [22]. We choose SGD for the optimiser, setting the batch size 1024. We start with learning rate of 0.03, annealing with cosine decay scheduler. We trained the model for 10 hours on eight A100 GPUs.
For our experiments with next-token prediction, we trained 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 16 and embedding dimension equal 768. We used the AdamW optimiser, with the same hyperparameters used in Ref. (22). namely learning rate equal 5e 4, exponential decay B1 0.9 and B2 0.95. The weight decay was 0.1 and we used a batch size equal to 128.
GPT
B Analytical details
Solving for (2) in the case (1) 0
We aim to solve the dynamical equations presented in Eqs. 9 of the main text, specifically in the scenario where the first layer skipped, which corresponds to setting (t) 0. Under this condition, the dynamics of the second layer are described by:
1 h, - (CAA
To solve this differential equation, we first write C = U DU with D a diagonal matrix with diagonal elements D, and UTT = since the covariance matrix is symmetric. Thus we have
A + D = hy
where we have defined
=U and hy = Uhy In this new basis, the equations are decoupled. With the initial conditions 1 (1 = 0) = Ab (12) (0), the solution is:
h Dj
hy Dj
1 (t)
+
A (0)
e Djt
is easy to verify that, in the original basis, the stationary solution is A (1. + o0) = (C-1)) as discussed in the main text.
Loss lower bound
Consider the analytical setup presented section of the main text. The mean square error loss associated with the model in Eq. iis:
C(A0)4))) = 20 -
Σ
All Øj
Σ
A
A A
IA
j
1kl
15

---

where bj Bos- We want to show that, in the case of Gaussian distributed data, for any values of the parameters b, it is bounded from below by the value it has in A = 0. Taking the expectation, we get:
((A 1 = 1(162)
Σ
Σ
10
A3 17 by 15t E (mvoraa))
jki PKV
=5(AA²) Σ
by Aoj
Σ
A A00
Alik! A/t (Crawwv CRNCU Cw.Ckr)
kiv
2
) Σ by Aoy (T + = + 11'
Σ
T T3
+
by
(3 r 403
15
where ((A is function of A alone, and we defined CA For any possible choice of the vector 2 we have:
((A A 1 2 ((A = -0|4 = =014())=((4}}))
2 Indeed, EE- 10/ ;5 2 0 for all choices of A Moreover, Ejy by A0) matrix 1 are non-negative, therefore all its eigenvalues are non-negative.
r2 1 2 0, since all the
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and the introduction describe precisely what we present in the results (section 3 and 4)
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See Limitations section.
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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
17

---

Justification:
See section 3.1 and supplemental material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See appendix and captions of the Figures.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
18

---

Answer: [Yes]
Justification:
The code has been shared on zenodo.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7/nips.com wttpcccpppbccccceec guwdes/dde/eeee Cooeeeoodeesssseeeeeee /tte://bb///uudddd///eeeeeiessss ssiniiiiccypy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 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, as supplemental material.
7. 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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
19

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, see appendix A.4
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wwra.ccccaa/a//a//aaaaaa
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 't see any possible societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact 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.
20

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We don't develop any model which have risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
21

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
22


Reviews:
[{"review_id": "jA0ofIydf8", "reviewer": "Reviewer_Kh1E", "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.", "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 are properly stated in Section 5.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper attempts to validate the simplicity bias on neural networks trained on masked language modelling, which has not been extensively studied.", "weaknesses": "1.\tThe 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.\n\n2. 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.\n\n3.  The paper is hard to follow due to inconsistent and undefined notations.\n\n(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.\n\n(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.\n\n(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.\n\n(4) In Equation (2), $A$ and $V$ are not defined.\n\n(5) The meaning of the subscript $\\mu$ in Equation (3) seems to conflict with the meaning of the superscript $\\mu$ in Line 93.\n\n(6) In Line 134, $p$ is not defined. Moreover, it is not clear whether the superscript $p$ is an exponent, or simply an index.\n\n4.\tFigure 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the referee for their detailed questions and their feedback. Below we \nreply 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.\n\n> Support for our main claim in Figures 1 and 4: \n\nWe're glad the reviewer agrees that our results show that transformers learn\nhigher-order interactions between tokens in the later stages of training, point\n(ii). However, we think our results support also point (i) both in experiments\non real data and in the controlled setting with synthetic data. The left panel of\nFigure 1 shows that the test loss curves for the different clones overlap during\nthe initial stage of training. Specifically, even after the initial plateau,\nbetween 5000 and 6000 training steps, the loss curves on the various clones\ncontinue to overlap despite the quick improvement in test error.  The crossover\nbetween the two-body and the higher-order interactions is at step 15000, and\nonly at 22000 the models start differing significantly. These results show that,\nat the early stage of training, the model first learns low-order interactions\n_without learning high-order interactions_. Furthermore, point (i) is clearly\nsupported by the experiments on the synthetic dataset (Section 2) and from the\nanalytical computation (Section 3.1).\n\n> I encourage the authors show the test loss on more individual\n> instances as well as the average test loss across all instances to consolidate\n> this conclusion.\n\nWe have performed our experiments using various random seeds, as specified in\nthe figure captions, and found the results to be consistent across different\nseeds. We have updated the plots by adding the corresponding error bars; we show\nan example in the rebuttal figure 1, where we redrew Fig 1 with error bars. We\nwill update all figures in this way in the revised version.\n\n> Clarifying the nature of “factored attention”\n\nWe define factored attention, where the attention matrix depends only on the\nposition of the tokens, by writing the corresponding probability for the masked\ntoken $p_{\\mathrm{mlm}}$ in Eq (2). To  clarify this we will first define only\nthe factored self-attention layer, and then the equation for\n$p_{\\mathrm{mlm}}$. The trainable parameters in a factored self-attention layer\nare the input-independent attention weights $A$ (an $L\\times L$ matrix, with $L$\nthe length of the input sequence) and the value matrix $V$.\n\n> (2) In Line 93, $\\mu$ and $\\alpha$ are not defined. I guess $\\mu$ means the\n> $\\mu$-th sample in the training set, but I’m not sure what $\\alpha$ refers to.\n\nYou are right: $\\mu$ is an index that runs over training examples; $\\alpha$\nindexes positions along the sequence. We will clarify this in the revised\nmanuscript.\n\n> (3) According to Line 84 “each token is an element taking values in a discrete\n> vocabulary ”, the equation in Line 93 means that we are always measuring the\n> probability of generating the first token in the vocabulary. This seems\n> unreasonable.\n\nWe wrote the equation in line 93 for the first token to keep the equation\nreadable; this equation (and any other) naturally extends to predicting any\nother token.\n\n> (4) In Equation (2), $A$ and $V$ are not defined.\n\nThank you for pointing this out -- $A$ and $V$ are the attention and the value\nmatrix, respectively. We will add a comment.\n\n> (5) The meaning of the subscript $\\mu$ in Equation (3) seems to conflict with the\n> meaning of the superscript in Line 93.\n\nIn both equations, $\\mu$ is simply a \"dummy\" index that is summed over.\n\n> (6) In Line 134, $p$ is not defined. Moreover, it is not clear whether the\n> superscript is an exponent, or simply an index.\n\nIn this equation, $p$ is also a dummy index. We now see that this might clash\nwith the notation for probabilities, and use a different dummy index for this\nsum.\n\n> Regarding the design of Figure 1\n\nWe 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.   \n\n\n> Repeating the experiments on synthetic data with a \n> transformer with standard attention. \n\nThis is a great question, and inspired us to perform an additional analysis. We\ntrained a standard three-layer transformer encoder on clones of a synthetic\ndatasets characterized by interactions involving up to four bodies. Being a\nstandard transformer, this architecture also had several layers of Layer\nNormalization, which mitigate the occurrence of plateaus even on synthetic\ndatasets. However, the plot for this standard transformer in Fig\n1 of the rebuttal figures clearly shows sequential learning also in this\ncase.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YeqRvrGPr4", "reviewer": "Reviewer_fVhm", "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.", "questions": "It's a very clearly written and presented paper, I don't have further questions to ask.", "limitations": "n/a", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Novel Insight**. The paper provides novel insights into the learning dynamics of Transformers, highlighting a simplicity bias that was previously unexplored in the context of self-supervised learning with MLM. At the same time, most previous simplicity bias papers more focus on the bias at the **beginning** or **end** of training, and this papers provides this novel insight in studying the training dynamics and the arising simplicity bias explanation during training procedures. \n\n\n**Many-body interactions**. The introduction of a procedure to generate dataset clones that capture interactions up to specified orders is a significant methodological contribution, facilitating deeper analysis of model learning behaviors. The proposed method enables us to have a deeper understanding of token interaction and itself is of a high contribution.", "weaknesses": "**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.\n\n**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n- *Computational cost:* applying our method to train a many-body\n  approximation of a data set, or \"clone\", is _less_ expensive than training a\n  standard BERT-like model, since the architecture is much simplified. Likewise,\n  the sampling procedure that we use from Goyal et al. [ICLR '22,\n  arXiv:2106.02736] can sample the standard BERT models. Hence we are confident\n  that the method can be applied to any data set on which BERT was applied to.\n\n- *Generality:* The procedure, as presented in the manuscript, and in particular\n  the analytic derivations, are designed for BERT-style transformers, and\n  validated only on such models. Extending it to causal models and other NLP\n  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\n  topics in our future research.  To further support the empirical\n  evidence of sequential learning in BERT-like architectures trained with MLM on\n  NLP datasets, we replicated the analysis on the WikiText-2 dataset, observing\n  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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DAN8svwvv6", "reviewer": "Reviewer_ex9y", "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.", "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?\n\nCan you guarantee that the difference provably limited models actually differ by nothing but their limitations? What about random variation?", "limitations": "nothing obvious", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper has one contribution that I think is very valuable: it introduces a new family of neural architectures that have a guarantee on what order of feature interaction they are capable of representing. \n\nFurthermore, the paper is one of a few that use an approach of approximating a true distribution using an analytically characterized language; the most common way this is deployed is by training on ngram-generated text to guarantee that it is representable by only ngrams. This is a promising approach that sits at the midpoint between purely synthetic and purely natural data, and more work in machine learning should embrace it.\n\nThis paper provides another example of how models learn increasingly complex representations over the course of training. It is a good addition to the literature, although it is not the first such paper in language modeling or NLP.", "weaknesses": "*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.*\n\nThe paper should be updated to reflect related literature in language modeling or NLP. Here are a few other examples:\n- 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\n- there are well documented dependencies including the famous checkmate in one emergence example from https://arxiv.org/pdf/2206.04615\n- 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/\n- 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/\n- Section 4 explores increasingly complex ngram learning in https://arxiv.org/pdf/2402.04362\nin context learning approximates shorter ngrams before longer ngram distributions https://arxiv.org/abs/2402.11004\n\nThe 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. \n\nSome ways of demonstrating simplicity bias are:\n- 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\n- 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)\n- Correlating the accessibility of a rule and the degree to which the model ultimately captures a different rule https://arxiv.org/abs/2006.12433\n- analytic proofs of a preference towards simplicity https://arxiv.org/abs/1805.08522\nIt is not sufficient to simply demonstrate that the model happens to learn simple things and then more complex things later on.\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your careful reading of the manuscript, your\ndetailed 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.\n\n- *Related works* We thank the referee for the references. Indeed, the\n  connection with n-gram models is extremely important, since our clones\n  implement n-grams using a transformer architecture. We will explicitly mention\n  this in the revised version. On the other hand, we are aware that the NLP\n  community has made enormous steps forward in understanding learning dynamics,\n  with paradigms that go well beyond n-grams. In the revised version of the\n  manuscript we will explicitly mention ref\n  https://aclanthology.org/2022.acl-long.568.pdf and\n  https://aclanthology.org/2022.acl-long.568.pdf as examples of advanced\n  quantitative analysis in this direction. We will also mention\n  https://aclanthology.org/N19-1329/ as a technique which allow probing the\n  learning dynamics in language. Finally, we discuss the relation of our work to\n  the very recent paper https://arxiv.org/pdf/2402.04362. They show that deep\n  transformers first learn the token unigram and then the bigram\n  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\n  dimensionality. Moreover, we would like to highlight that none of these papers,\n  including arXiv:2402.04362, introduces an architecture which produces\n  distributions in which the simplicity bias can be _controlled_ by tuning\n  the interactions between tokens that are captured by the model, at an arbitrary order. This\n  analytically guaranteed control, as also recognized by the reviewer, is an\n  important and novel contribution of our work.\n\n- *The nature of simplicity bias* Thank you for sharing these references on\n  various ways in which learning simple rules might interfere with learning more\n  complex rules. These papers demonstrate \"simplicity biases\" by performing\n  experiments on appropriately fine-tuned training tasks, where simple rules\n  based on spurious correlations hinder the network from learning the\n  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\n  underlying rules of language rather well,  producing flawless chunks of\n  text. It is this process that we want to study through the lens of another\n  line of research. We hope that our paper convincingly shows that during a single training run, a model often\n  seems to be learning increasingly complex functions or distributions that\n  explain its training data. We discuss previous theoretical and experimental\n  work in the introduction, in particular in lines 18-26, many of which have\n  been published at NeurIPS/ICML/ICLR in the last couple of years (Refs\n  2,3,4,5,6,10,12,13). We will mention this separate line of work on simplicity\n  biases in the revised version of the manuscript.\n\nRegarding your *questions*:\n\n> Can you give a simple explanation of what a k-body interaction would look\n> like? is it essentially stating something about ngram models, but not\n> restricted to the most proximate tokens?\n\nA $k$-body interaction is any interaction between $k$ tokens anywhere along the\nsequence that can not be written as a product of lower order interactions. In\nthe energy-based models we use to construct the clones, each such interaction\ncorresponds to one term in the energy function involving $k$ tokens. Indeed, the tokens do not have to be proximate. N-gram\nmodels instead simply define the number of previous tokens on which a given\ntoken is conditionally dependent, i.e. how many conditions there are in $p(x_n |\nx_{n-1}, x_{n-2}, …)$.\n\n> Can you guarantee that the difference provably limited models actually differ\n> by nothing but their limitations? What about random variation?\n\nThe architecture of the \"limited models\" that we use to obtain the clones only\ndiffer by their \"limitations\", namely by the degree of the interactions between\ntokens that they cover. After training, the values of the parameters of these\nmodels could have some random variations in their values due to different\ninitial conditions; however, this will not change the degree of the interactions\nbetween tokens represented by these models.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HOoMUtRZgz", "reviewer": "Reviewer_NZGs", "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.​​​​​​​​​​​​​​​​\n\nThe key contributions of the paper are listed below:\n\n1) Introduction of a novel framework to study the learning dynamics of transformers:\n- They develop a method to create \"clones\" of natural language datasets that capture token interactions up to a specified order.\n- This is done by using transformers with factored attention layers and quadratic $x^2$ activation functions.\n\n2) 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.\n\n3) 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.\n\n4) 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.", "questions": "1) What does the \"last layer\" in line 163 mean?\n\nIt'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": "Please refer to my weaknesses section.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Overall, I enjoyed reading the paper very much, and I also love the hypothesis investigated in this work. The specific aspects of this work that I love are listed below.\n\n1) *The proposed hypothesis is of great significance*:\nBeyond the significance of this simplicity bias mentioned in the paper, I want to provide the perspectives from the cognitive science community. A series of language evolution studies, e.g. [1], have shown that simplicity bias plays a key role in shaping the evolution of natural languages. For syntax (i.e. generative functions of language data), humans also have such simplicity bias. It would be highly interesting and exciting (at least for me) to see that Transformer-based language models also gradually learn increasingly complicated functions for generating language.\n\n2) *The illustration is clear, and the structure is well-organised*:\nThe illustration of the method and the experiments is clear, and the structure of the paper is well-organised. Therefore, the whole idea is delivered very well. \n\n3) *The methodology is convincing*:\nThe study methodology, clones of datasets and factored attention, is convincing per se, and can quite well align with the hypothesis investigated in the paper. Though there is room for this work to be improved, it could be a pivot step towards understanding the learning dynamics of Transformer-based language models.\n\n\nReferences\n[1] Smith et al. (2013). Linguistic structure is an evolutionary trade-off between simplicity and expressivity", "weaknesses": "1) *Uncommonly used architecture*.\n\nAs 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.\n\nMoreover, 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.\n\n2) *Higher orders are not straightforward to sample.*\n\nAs 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.\n\n3) *The empirical evidence is not strong enough.*\n\nTo 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to thank you for your positive comments and your enthusiam for our work. We really appreciated the\nseries of language evolution studies on simplicity biases that you suggested,\nand we decided to include them in the revised version of our manuscript. \n\nBelow 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\nevidence. If so, we would appreciate it if you could reconsider your rating; if\nnot, we're looking forward to discussing further during the discussion period.\n\n- *Title*: We agree with your suggestion for the title, which would be more\n  specific.\n\n- *Activation function:* In our framework, we only use the $x^2$ activation\n  function in the architectures that are subsequently used to generate the\n  clones of a real NLP dataset, in which the maximum degree of interaction among\n  tokens can be controlled by tuning the number of layers. We use these clones\n  to demonstrate the simplicity bias of a standard transformer encoder\n  (BERT-like) with *standard activation function* trained under MLM on a real\n  NLP task. Thus, the main result of our paper is obtained for a standard\n  architecture, trained with a standard task on a standard dataset.\n\n- *Decoder-only transformers* are indeed popular language models. The procedure\n  we present in this manuscript, and in particular the analytic derivations, are\n  however designed for BERT-style transformers, and validated only on such\n  models. Extending it to causal models and other NLP tasks is for sure possible\n  and interesting, but non-trivial, and we plan to investigate this topic in our future research.\n\n- *Time complexity of sampling* The reviewer is right that the sampling of the\n  clones is not trivial; it is difficult to fully decorrelate the samples, and\n  the complexity of the method we use is not linear in the number of\n  samples. However, the modified transformer architectures that model the data\n  are computationally much more easy to handle than, say, estimating the\n  corresponding $n$-grams. The clone models require only a couple of layers of\n  factored self-attention, which is easier to compute than vanilla\n  self-attention layers, and they do not need fully-connected networks\n  in-between attention layers. This makes sampling of our clones, while\n  challenging, also way cheaper than sampling the distribution generated by an\n  ordinary transformer.\n\n- *Evidence for simplicity bias* We appreciate the need for additional experimental verification,\n  and have repeated our experiments on another data set, WikiText101. As we show\n  in our plot in the additional one-side PDF, we can see sequential learning of\n  increasingly higher-order interactions also on this data set. In our\n  manuscript, we further show sequential learning in a controlled setting with\n  synthetic data (cf. Section 2) and we see the same behaviour in our analytical\n  treatment of Sec. 3.1. We hope that this combination of experimental\n  verification on tinystories, WikiText101, synthetic data, combined with our\n  analytical results, is convincing evidence for our claims.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for taking the time to review of our manuscript. In our paper, we show that BERT-style transformers trained using masked language modeling learn increasingly complex interactions among input tokens. To conduct this analysis, we develop a procedure to generate clones of a given natural language data set, which rigorously capture the interactions between tokens up to a specified order. 

Prompted by the thoughtful questions in the replies, we conducted a number of additional experiments to solidify our points:
- On the left of Fig 1 of the attached PDF, we show sequential learning in the same setting as in Figure 1 of the main submission, but on a different data set: WikiText. 
- On the right of the same figure, we show sequential learning of a vanilla transformer with standard attention trained on our synthetic data set.
- We have also added an additional figure showing the statistical uncertainty in the test loss trajectories for Figure 1.

We hope that these results alleviate your concerns. If not, we look forward to discussing further during the discussion period.

Meta-review: 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.
Decision: Accept (poster)