Original: S AME P RE - TRAINING L OSS , B ETTER D OWNSTREAM : I MPLICIT B IAS M ATTERS FOR L ANGUAGE M ODELS
Revised: S AME P RE - TRAINING L OSS , B ETTER D OWNSTREAM : I MPLICIT B IAS M ATTERS FOR L ANGUAGE M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Language modeling on large-scale datasets leads to impressive performance gains on various downstream language tasks.
Revised: Language modeling on large-scale datasets leads to impressive performance gains on various downstream language tasks.

Original: The (validation) pre-training loss (or perplexity in autoregressive language modeling) is often used as the evaluation metric when developing language models since the pre-training loss tends to be well-correlated with downstream performance (which is itself difficult to evaluate comprehensively).
Revised: The (validation) pre-training loss (or perplexity in autoregressive language modeling) is often used as the evaluation metric when developing language models since the pre-training loss tends to be well-correlated with downstream performance (which is itself difﬁcult to evaluate comprehensively).

Original: Contrary to this conventional wisdom, this paper shows that 1) pre-training loss cannot fully explain downstream performance and 2) flatness of the model is well-correlated with downstream performance where pre-training loss is not.
Revised: Contrary to this conventional wisdom, this paper shows that 1) pre-training loss cannot fully explain downstream performance and 2) ﬂatness of the model is well-correlated with downstream performance where pre-training loss is not.

Original: On simplified datasets, we identify three ways to produce models with the same (statistically optimal) pre-training loss but different downstream performance: continue pre-training after convergence, increasing the model size, and changing the training algorithm.
Revised: On simpliﬁed datasets, we identify three ways to produce models with the same (statistically optimal) pre-training loss but different downstream performance: continue pre-training after convergence, increasing the model size, and changing the training algorithm.

Original: These experiments demonstrate the existence of implicit bias of pretraining algorithms/optimizers—among models with the same minimal pre-training loss, they implicitly prefer more transferable ones.
Revised: These experiments demonstrate the existence of implicit bias of pretraining algorithms/optimizers—among models with the same minimal pre-training loss, they implicitly prefer more transferable ones.

Original: Toward understanding this implicit bias, we prove that SGD with standard mini-batch noise implicitly prefers flatter minima in language models, and empirically observe a strong correlation between flatness and downstream performance among models with the same minimal pre-trainingloss.
Revised: Toward understanding this implicit bias, we prove that SGD with standard mini-batch noise implicitly prefers ﬂatter minima in language models, and empirically observe a strong correlation between ﬂatness and downstream performance among models with the same minimal pre-trainingloss.

Original: Wealsoproveinasyntheticlanguagesettingthatamongthemodels with the minimal pre-training loss, the flattest model transfers to downstream tasks.
Revised: Wealsoproveinasyntheticlanguagesettingthatamongthemodels with the minimal pre-training loss, the ﬂattest model transfers to downstream tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Large language models (LLMs) trained on internet-scale data have improved performance on a wide array of downstream tasks (Devlin et al., 2018;
Revised: Large language models (LLMs) trained on internet-scale data have improved performance on a wide array of downstream tasks (Devlin et al., 2018;

Original: Yang et al., 2019; Radford et al., 2019; Raffel et al., 2020; Brown et al., 2020).
Revised: Yang et al., 2019; Radford et al., 2019; Raffel et al., 2020; Brown et al., 2020).

Original: These models are trained with a language modeling pre-training loss to “fill in the blanks”—either predicting the next token/word (autoregressive language modeling loss, or perplexity) or masked tokens (masked language modeling (MLM) loss).
Revised: These models are trained with a language modeling pre-training loss to “ﬁll in the blanks”—either predicting the next token/word (autoregressive language modeling loss, or perplexity) or masked tokens (masked language modeling (MLM) loss).

Original: In common practice, the validation pre-training loss is used to monitor the training process (Brown et al., 2020; Zhang et al., 2022a) and compare different models since the pre-training loss is generally strongly correlated with downstream performance (Hernandez et al., 2021).
Revised: In common practice, the validation pre-training loss is used to monitor the training process (Brown et al., 2020; Zhang et al., 2022a) and compare different models since the pre-training loss is generally strongly correlated with downstream performance (Hernandez et al., 2021).

Original: Moreover, theoretical works on understanding LLMs also focus on how the pre-training loss affects downstream performance.
Revised: Moreover, theoretical works on understanding LLMs also focus on how the pre-training loss affects downstream performance.

Original: Saunshi et al.
Revised: Saunshi et al.

Original: Wei et al.
Revised: Wei et al.

Original: ; Xie et al.
Revised: ; Xie et al.

Original: (2021) show that good pre-training loss, or fitting the language modeling conditional probability well, is a main reason for downstream success of LLMs.
Revised: (2021) show that good pre-training loss, or ﬁtting the language modeling conditional probability well, is a main reason for downstream success of LLMs.

Original: Their analyses generally treat the language models as blackboxes and do not take into account how the models represents the conditional probability.
Revised: Their analyses generally treat the language models as blackboxes and do not take into account how the models represents the conditional probability.

Original: In this paper, we question the conventional wisdom on the correlation between the validation pre-training loss and downstream performance for language modeling.
Revised: In this paper, we question the conventional wisdom on the correlation between the validation pre-training loss and downstream performance for language modeling.

Original: Recent works have demonstrated that models with different architectures may have the same pre-training loss but different performance (Saunshi et al., 2022; Tay et al., 2021).
Revised: Recent works have demonstrated that models with different architectures may have the same pre-training loss but different performance (Saunshi et al., 2022; Tay et al., 2021).

Original: Due to the expressivity of modern neural nets, many parameter configurations even within the same architecture can still have the same pre-training loss.
Revised: Due to the expressivity of modern neural nets, many parameter conﬁgurations even within the same architecture can still have the same pre-training loss.

Original: A priori, it is unclear why all these configurations should have the same downstream performance.
Revised: A priori, it is unclear why all these conﬁgurations should have the same downstream performance.

Original: We find that different parameter configurations with the same pre-training loss can indeed have different downstream performance, especially when the pre-training loss reaches a near-optimal level.
Revised: We ﬁnd that different parameter conﬁgurations with the same pre-training loss can indeed have different downstream performance, especially when the pre-training loss reaches a near-optimal level.

Original: Concretely, using simplified text datasets, we find three situations that demonstrate such a phenomenon:
Revised: Concretely, using simpliﬁed text datasets, we ﬁnd three situations that demonstrate such a phenomenon:

Original: • Even after the pre-training loss converges, models at a later time step still tend to perform better.
Revised: • Even after the pre-training loss converges, models at a later time step still tend to perform better.

Original: • Models trained by standard algorithms have better performance than adversarially trained models with the same pre-training loss.
Revised: • Models trained by standard algorithms have better performance than adversarially trained models with the same pre-training loss.

Original: • Larger models tend to perform better downstream than smaller models even if they have the same pre-training loss.
Revised: • Larger models tend to perform better downstream than smaller models even if they have the same pre-training loss.

Original: These situations are most prominent in the saturation regime , where the models are close to the minimal possible pre-training loss (aka the entropy of the conditional probability, which can be estimated in our simplified datasets).
Revised: These situations are most prominent in the saturation regime , where the models are close to the minimal possible pre-training loss (aka the entropy of the conditional probability, which can be estimated in our simpliﬁed datasets).

Original: In the saturation regime, the pre-training loss of all models are almost the same, but the transferability to downstream tasks varies.
Revised: In the saturation regime, the pre-training loss of all models are almost the same, but the transferability to downstream tasks varies.

Original: Interestingly, this phenomenon also holds when linear probing on contextualized presentations is used for evaluating downstream performance instead of finetuning.
Revised: Interestingly, this phenomenon also holds when linear probing on contextualized presentations is used for evaluating downstream performance instead of ﬁnetuning.

Original: Thus, even though the predicted conditional probabilities of two models are the same (and correct), the contextualized representations can behave differently.
Revised: Thus, even though the predicted conditional probabilities of two models are the same (and correct), the contextualized representations can behave differently.

Original: In each of the first two cases above, we find two models with the same pre-training loss and the same architecture; but one has a better performance than the other.
Revised: In each of the ﬁrst two cases above, we ﬁnd two models with the same pre-training loss and the same architecture; but one has a better performance than the other.

Original: They only differ by the training algorithms that are used to produce them.
Revised: They only differ by the training algorithms that are used to produce them.

Original: Therefore, this suggests the training algorithms have an implicit bias toward one of these models—standard algorithms with more training steps biases towards parameter configurations that transfer better to downstream tasks.
Revised: Therefore, this suggests the training algorithms have an implicit bias toward one of these models—standard algorithms with more training steps biases towards parameter conﬁgurations that transfer better to downstream tasks.

Original: The third case has a more subtle but similar interpretation.
Revised: The third case has a more subtle but similar interpretation.

Original: There exists a hypothetical large model that represents the smaller model with worse performance (by padding zeros to the smaller model).
Revised: There exists a hypothetical large model that represents the smaller model with worse performance (by padding zeros to the smaller model).

Original: The training algorithm on the large architecture could have chosen it, but did not.
Revised: The training algorithm on the large architecture could have chosen it, but did not.

Original: This suggests the algorithm has an implicit bias against the hypothetical model (which has an equally good loss).
Revised: This suggests the algorithm has an implicit bias against the hypothetical model (which has an equally good loss).

Original: In supervised settings, optimizers are known to have an implicit bias toward selecting generalizable models among all models with small empirical loss.
Revised: In supervised settings, optimizers are known to have an implicit bias toward selecting generalizable models among all models with small empirical loss.

Original: E.g., see Damian et al.
Revised: E.g., see Damian et al.

Original: ; Li et al.
Revised: ; Li et al.

Original: (2021), which show that SGD implicitly biases toward flatter minima, and references therein.
Revised: (2021), which show that SGD implicitly biases toward ﬂatter minima, and references therein.

Original: However, the role of implicit bias in self-supervised learning has not been studied and is conceptually different.
Revised: However, the role of implicit bias in self-supervised learning has not been studied and is conceptually different.

Original: Unlike in supervised learning, the gap between empirical and population self-supervised losses is typically small, and thus implicit bias does not seem to contribute to bridging this gap.
Revised: Unlike in supervised learning, the gap between empirical and population self-supervised losses is typically small, and thus implicit bias does not seem to contribute to bridging this gap.

Original: Instead, the implicit bias selects local minima of the population self-supervised loss that transfer better to downstream tasks.
Revised: Instead, the implicit bias selects local minima of the population self-supervised loss that transfer better to downstream tasks.

Original: Why do the algorithms bias toward some type of models?
Revised: Why do the algorithms bias toward some type of models?

Original: In Section 3, we provide a first-cut theoretical analysis of the implicit bias in language modeling.
Revised: In Section 3, we provide a ﬁrst-cut theoretical analysis of the implicit bias in language modeling.

Original: Fortunately, despite the conceptual differences, mathematical tools from supervised settings can be straightforwardly adapted to language modeling settings.
Revised: Fortunately, despite the conceptual differences, mathematical tools from supervised settings can be straightforwardly adapted to language modeling settings.

Original: We prove that mini-batch SGD prefers flatter minima of population pre-training loss among all minima in the saturation regime.
Revised: We prove that mini-batch SGD prefers ﬂatter minima of population pre-training loss among all minima in the saturation regime.

Original: Interestingly, we obtain cleaner theoretical results for the standard mini-batch SGD, without the artificial label noise introduced in prior works (Damian et al., 2021; Li et al., 2021), partly because the mini-batch noise for LLMs does not vanish even at convergence.
Revised: Interestingly, we obtain cleaner theoretical results for the standard mini-batch SGD, without the artiﬁcial label noise introduced in prior works (Damian et al., 2021; Li et al., 2021), partly because the mini-batch noise for LLMs does not vanish even at convergence.

Original: We corroborate our theory with empirical evidence in Section 4.
Revised: We corroborate our theory with empirical evidence in Section 4.

Original: We show that for models with the same pre-training loss in the three situations above, flatness of the model (measured by the trace of
Revised: We show that for models with the same pre-training loss in the three situations above, ﬂatness of the model (measured by the trace of

Original: Hessian of the loss, as predicted by the theory) strongly correlates with the downstream performance.
Revised: Hessian of the loss, as predicted by the theory) strongly correlates with the downstream performance.

Original: Finally, to complement the theory and experiments above, we also rigorously formalize the connection between flatness and downstream performance in a simplified Dyck language setting in Section 5.
Revised: Finally, to complement the theory and experiments above, we also rigorously formalize the connection between ﬂatness and downstream performance in a simpliﬁed Dyck language setting in Section 5.

Original: In this setting, we prove that there are many models with good MLM pre-training loss; among them, the flattest model learns the most useful features for downstream tasks.
Revised: In this setting, we prove that there are many models with good MLM pre-training loss; among them, the ﬂattest model learns the most useful features for downstream tasks.

Original: Here, results from the supervised setting cannot be readily adapted since they are obtained (partially) via generalization bounds (Wei & Ma, 2019a;b), which do not apply to the language modeling setting where the implicit bias is not related to the gap between the empirical and population loss.
Revised: Here, results from the supervised setting cannot be readily adapted since they are obtained (partially) via generalization bounds (Wei & Ma, 2019a;b), which do not apply to the language modeling setting where the implicit bias is not related to the gap between the empirical and population loss.

Original: Proving the correlation between flatness and downstream performance in more general settings likely requires highly non-trivial and novel theoretical tools, and we hope to motivate future work on this topic.
Revised: Proving the correlation between ﬂatness and downstream performance in more general settings likely requires highly non-trivial and novel theoretical tools, and we hope to motivate future work on this topic.

Original:  T HE E XISTENCE OF I MPLICIT B IAS IN L ANGUAGE M ODELING
Revised:  T HE E XISTENCE OF I MPLICIT B IAS IN L ANGUAGE M ODELING

Original: In this section, we systematically investigate the relationship between pre-training loss and downstream performance with experiments.
Revised: In this section, we systematically investigate the relationship between pre-training loss and downstream performance with experiments.

Original: We find out that models with the same pre-training loss but different training procedures can have different downstream performance.
Revised: We ﬁnd out that models with the same pre-training loss but different training procedures can have different downstream performance.

Original: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Number of
Revised: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Number of

Original: [Figure]
Revised: [Figure]

Original: D o w n s t r e a m
Revised: D o w n s t r e a m

Original: A cc u r a c y %
Revised: A cc u r a c y %

Original: Linear Probe
Revised: Linear Probe

Original: P r e - t r a i n i n g
Revised: P r e - t r a i n i n g

Original: L o ss
Revised: L o ss

Original: PCFG, Task C, 41M
Revised: PCFG, Task C, 41M

Original: Pre-training (
Revised: Pre-training (

Original: [Figure]
Revised: [Figure]

Original: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 Number of Steps / 

Revised: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 Number of Steps /

Original: D o w n s t r e a m
Revised: D o w n s t r e a m

Original: A cc u r a c y %
Revised: A cc u r a c y %

Original: Linear Probe
Revised: Linear Probe

Original: P r e - t r a i n i n g
Revised: P r e - t r a i n i n g

Original: L o ss
Revised: L o ss

Original: OPT, QNLI, 235M
Revised: OPT, QNLI, 235M

Original: Pre-training
Revised: Pre-training

Original: [Figure]
Revised: [Figure]

Original: 2. F ORMULATIONS
Revised: 2. F ORMULATIONS

Original: Masked language modeling.
Revised: Masked language modeling.

Original: Consider a vocabulary W = { 0 , 1 ,...,c } , where 0 is a special token for the mask.
Revised: Consider a vocabulary W = { 0 , 1 ,...,c } , where 0 is a special token for the mask.

Original: Let x =[ x 1 ,...,x T ] denote the input sequence of length T , and x − t =[ x 1 ,...,x t − 1 , 0 ,x t +1 ,...,x T ] denote the masked sentence, where t is sampled uniformly randomly and independently from [ T ] .
Revised: Let x =[ x 1 ,...,x T ] denote the input sequence of length T , and x − t =[ x 1 ,...,x t − 1 , 0 ,x t +1 ,...,x T ] denote the masked sentence, where t is sampled uniformly randomly and independently from [ T ] .

Original: The MLM conditional probability refers to the probability of x t given the rest of the sequence Pr ( x t | x − t ) .
Revised: The MLM conditional probability refers to the probability of x t given the rest of the sequence Pr ( x t | x − t ) .

Original: We use Pr ( ·| x − t ) to denote the c -dimensional probability vector
Revised: We use Pr ( ·| x − t ) to denote the c -dimensional probability vector

Original: Pr( ·| x − t ) :=
Revised: Pr( ·| x − t ) :=

Original: [Pr( x t =1 | x − t ) ,..., Pr( x t = c | x − t )] ∈ R c .
Revised: [Pr( x t =1 | x − t ) ,..., Pr( x t = c | x − t )] ∈ R c .

Original: In MLM pre-training, the model f θ ( · ) (parameterized by θ ) outputs the predicted MLM conditional probability vector f θ ( x − t )
Revised: In MLM pre-training, the model f θ ( · ) (parameterized by θ ) outputs the predicted MLM conditional probability vector f θ ( x − t )

Original: ∈ R c .
Revised: ∈ R c .

Original: The model is trained to predict the masked token x t given the rest of the sentence x − t with cross entropy loss, L ( θ )=
Revised: The model is trained to predict the masked token x t given the rest of the sentence x − t with cross entropy loss, L ( θ )=

Original: E x,t [ ℓ ( f θ ( x − t ) ,x t )]=
Revised: E x,t [ (cid:96) ( f θ ( x − t ) ,x t )]=

Original: E x,t [ − log([ f θ ( x − t )] x t )] .
Revised: E x,t [ − log([ f θ ( x − t )] x t )] .

Original: Downstream evaluation.
Revised: Downstream evaluation.

Original: The language model f θ is composed of a feature extractor h ψ , which outputs a sequence of contextual representations, and a linear classifier that outputs the conditional probability at every position.
Revised: The language model f θ is composed of a feature extractor h ψ , which outputs a sequence of contextual representations, and a linear classiﬁer that outputs the conditional probability at every position.

Original: On downstream tasks, we use a randomly initialized g ϕ on top of the pre-trained h ψ .
Revised: On downstream tasks, we use a randomly initialized g φ on top of the pre-trained h ψ .

Original: In fine-tuning, both g ϕ and h ψ are trained, while in linear probe, only g ϕ is updated.
Revised: In ﬁne-tuning, both g φ and h ψ are trained, while in linear probe, only g φ is updated.

Original: For fine-tuning, we use the contextual representations of the cls token.
Revised: For ﬁne-tuning, we use the contextual representations of the cls token.

Original: For linear probe, we concatenate the contextual representations of all the tokens together.
Revised: For linear probe, we concatenate the contextual representations of all the tokens together.

Original: Saturation regime.
Revised: Saturation regime.

Original: To study models with the same pre-training loss, we introduce the saturation regime in this paper, where the model output equals the true conditional probability, f θ ( x − t ) =
Revised: To study models with the same pre-training loss, we introduce the saturation regime in this paper, where the model output equals the true conditional probability, f θ ( x − t ) =

Original: Pr( ·| x − t ) .
Revised: Pr( ·| x − t ) .

Original: In the saturation regime, the MLM loss is equal to the entropy of the true conditional probability
Revised: In the saturation regime, the MLM loss is equal to the entropy of the true conditional probability

Original: L ( θ ) =
Revised: L ( θ ) =

Original: E x,t [ − log(Pr( x t | x − t ))]
Revised: E x,t [ − log(Pr( x t | x − t ))]

Original: = 1 T (cid:80)
Revised: = 1 T (cid:80)

Original: Tt =1 H ( x t | x − t ) , which is also the optimal pre-training loss.
Revised: Tt =1 H ( x t | x − t ) , which is also the optimal pre-training loss.

Original: Thus, all models in the saturation regime have the same, optimal pre-training loss, and we will show that they behave differently on downstream tasks.
Revised: Thus, all models in the saturation regime have the same, optimal pre-training loss, and we will show that they behave differently on downstream tasks.

Original: Our experiments use expressive enough architectures such that there are multiple parameter configurations in the saturation regime for our simplified datasets.
Revised: Our experiments use expressive enough architectures such that there are multiple parameter conﬁgurations in the saturation regime for our simpliﬁed datasets.

Original: For real large-scale data, it is currently computationally challenging to arrive at the saturation regime.
Revised: For real large-scale data, it is currently computationally challenging to arrive at the saturation regime.

Original: However, we hope that our experiments can provide insights for even larger models in the future and for other regimes where pre-training loss does not explain downstream performance.
Revised: However, we hope that our experiments can provide insights for even larger models in the future and for other regimes where pre-training loss does not explain downstream performance.

Original: 2. E XPERIMENTAL S ETUP
Revised: 2. E XPERIMENTAL S ETUP

Original: We design controlled experiments to study the correlation between pre-training loss and downstream performance.
Revised: We design controlled experiments to study the correlation between pre-training loss and downstream performance.

Original: In particular, we will find a set of models with almost the same pre-training loss.
Revised: In particular, we will ﬁnd a set of models with almost the same pre-training loss.

Original: We effectively use the same architecture family so that the main difference between the models only stems from training algorithms.
Revised: We effectively use the same architecture family so that the main difference between the models only stems from training algorithms.

Original: More details are provided in Section A.
Revised: More details are provided in Section A.

Original: Datasets.
Revised: Datasets.

Original: We introduce three generative models to produce simplified datasets, with which we can study various factors systematically.
Revised: We introduce three generative models to produce simpliﬁed datasets, with which we can study various factors systematically.

Original: With the knowledge of the true generative models that generate the data, we can compute the true conditional probability and scale up the models until they approach the saturation regime to ensure they have almost the same pre-training loss.
Revised: With the knowledge of the true generative models that generate the data, we can compute the true conditional probability and scale up the models until they approach the saturation regime to ensure they have almost the same pre-training loss.

Original: Moreover, we can generate unlimited amount of text for pre-training to avoid overfitting to the empirical pre-training loss.
Revised: Moreover, we can generate unlimited amount of text for pre-training to avoid overﬁtting to the empirical pre-training loss.

Original: 1) PCFG -generated dataset.
Revised: 1) PCFG -generated dataset.

Original: PCFG (Chomsky, 1956) generates sentences with probabilistic trees and is widely used to understand natural language (Johnson, 1998; Roark & Bacchiani, 2003;
Revised: PCFG (Chomsky, 1956) generates sentences with probabilistic trees and is widely used to understand natural language (Johnson, 1998; Roark & Bacchiani, 2003;

Original: Kim et al., 2019; Yang et al., 2021).
Revised: Kim et al., 2019; Yang et al., 2021).

Original: We randomly generate the production rules which satisfy the Chomsky Normal Form (Chomsky, 1956).
Revised: We randomly generate the production rules which satisfy the Chomsky Normal Form (Chomsky, 1956).

Original: The non-terminal symbols in the parse tree can be viewed as intrinsic
Revised: The non-terminal symbols in the parse tree can be viewed as intrinsic

Original: [Figure]
Revised: [Figure]

Original: / M
Revised: / M

Original: D o w n s t r e a m A cc u r a c
Revised: D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: PCFG, Task B
Revised: PCFG, Task B

Original: Fine-tuningLinear Probe 3 . 203 .. 243 . 263 .
Revised: Fine-tuningLinear Probe 3 . 203 .. 243 . 263 .

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre
Revised: Pre

Original: [Figure]
Revised: [Figure]

Original: training
Revised: training

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: D
Revised: D

Original: [Figure]
Revised: [Figure]

Original: t
Revised: t

Original: [Figure]
Revised: [Figure]

Original: cc u
Revised: cc u

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre
Revised: Pre

Original: [Figure]
Revised: [Figure]

Original: training (
Revised: training (

Original: [Figure]
Revised: [Figure]

Original: HMM
Revised: HMM

Original: [Figure]
Revised: [Figure]

Original: Task
Revised: Task

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: / M

Revised: / M

Original: D o w n s t r e a m A cc u r a c
Revised: D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: OPT, QNLI
Revised: OPT, QNLI

Original: Fine
Revised: Fine

Original: [Figure]
Revised: [Figure]

Original: Probe 1 .2 .2 .2 .2 .2 .
Revised: Probe 1 .2 .2 .2 .2 .2 .

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre
Revised: Pre

Original: [Figure]
Revised: [Figure]

Original: training (
Revised: training (

Original: [Figure]
Revised: [Figure]

Original: OPT
Revised: OPT

Original: [Figure]
Revised: [Figure]

Original: quantities associated with the sentence such as sentiment and syntax.
Revised: quantities associated with the sentence such as sentiment and syntax.

Original: Thus we design three downstream tasks A, B, and C to classify non-terminal symbols at different positions of the parse trees.
Revised: Thus we design three downstream tasks A, B, and C to classify non-terminal symbols at different positions of the parse trees.

Original: 2) HMM-generated dataset.
Revised: 2) HMM-generated dataset.

Original: HMM samples the hidden variables from the
Revised: HMM samples the hidden variables from the

Original: transition probabilities and the tokens from the emission probabilities.
Revised: transition probabilities and the tokens from the emission probabilities.

Original: (Wei et al., 2021; Xie et al., 2021) also analyze the properties of pre-trained language models with HMMs.
Revised: (Wei et al ., 2021; Xie et al., 2021) also analyze the properties of pre-trained language models with HMMs.

Original: We generate the transition and emission probabilities as random block-diagonal stochastic matrices.
Revised: We generate the transition and emission probabilities as random block-diagonal stochastic matrices.

Original: The downstream task is to classify the hidden variable in the sentence.
Revised: The downstream task is to classify the hidden variable in the sentence.

Original: We use task- k to refer to classifying the k -th hidden variable.
Revised: We use task- k to refer to classifying the k -th hidden variable.

Original: 3) OPT-generated dataset.
Revised: 3) OPT-generated dataset.

Original: We also introduce a more realistic pre-training dataset generated by the OPT models (Zhang et al., 2022a).
Revised: We also introduce a more realistic pre-training dataset generated by the OPT models (Zhang et al., 2022a).

Original: Starting from the bos token, we sample each token from the conditional probability output by the OPT model.
Revised: Starting from the bos token, we sample each token from the conditional probability output by the OPT model.

Original: For computational feasibility we only allow to generate the top-2000 most frequent tokens in the OPT vocabulary.
Revised: For computational feasibility we only allow to generate the top-2000 most frequent tokens in the OPT vocabulary.

Original: We use QNLI and SSTfrom GLUE (Wang et al., 2018) as downstream tasks.
Revised: We use QNLI and SSTfrom GLUE (Wang et al., 2018) as downstream tasks.

Original: Note that the true conditional probability can be computed efficiently for the three datasets given the knowledge of the generated models.
Revised: Note that the true conditional probability can be computed efﬁciently for the three datasets given the knowledge of the generated models.

Original: For PCFG and HMM-generated datasets, we can compute the true conditional probability with the inside algorithm (Lari & Young, 1990) and the Viterbi algorithm (Forney, 1973), respectively.
Revised: For PCFG and HMM-generated datasets, we can compute the true conditional probability with the inside algorithm (Lari & Young, 1990) and the Viterbi algorithm (Forney, 1973), respectively.

Original: For the OPT-generated dataset, we can calculate the MLM conditional probability from the joint probability, and the joint probability can be decomposed into the autoregressive conditional probability of the OPT model.
Revised: For the OPT-generated dataset, we can calculate the MLM conditional probability from the joint probability, and the joint probability can be decomposed into the autoregressive conditional probability of the OPT model.

Original: Models and algorithms.
Revised: Models and algorithms.

Original: For PCFG and OPT generated datasets, we use transformers (Vaswani et al., 2017) following the implementation of BERT (Devlin et al., 2018).
Revised: For PCFG and OPT generated datasets, we use transformers (Vaswani et al., 2017) following the implementation of BERT (Devlin et al., 2018).

Original: We use different model sizes ranging from 2M to 730M. For the HMM generated dataset, we use LSTM (Hochreiter & Schmidhuber, 1997) from 10M to 135M. In pre-training, all the models are pre-trained with AdamW following the protocol of Izsak et al.
Revised: We use different model sizes ranging from 2M to 730M. For the HMM generated dataset, we use LSTM (Hochreiter & Schmidhuber, 1997) from 10M to 135M. In pre-training, all the models are pre-trained with AdamW following the protocol of Izsak et al.

Original: We use batch size 4096, and train each model for more than 30K steps until the pre-training loss converges.
Revised: We use batch size 4096, and train each model for more than 30K steps until the pre-training loss converges.

Original: For comparison, we also consider other training algorithms.
Revised: For comparison, we also consider other training algorithms.

Original: The first adversarial algorithm is inspired by Liu et al.
Revised: The ﬁrst adversarial algorithm is inspired by Liu et al.

Original: ; Raghu et al.
Revised: ; Raghu et al.

Original: (2021), where the models are pre-trained with an additional meta-learning objective which messes up downstream performance.
Revised: (2021), where the models are pre-trained with an additional meta-learning objective which messes up downstream performance.

Original: The second algorithm is manually setting the weight of the model to represent a lookup table which memorizes all the masked sentences and the corresponding true conditional probability.
Revised: The second algorithm is manually setting the weight of the model to represent a lookup table which memorizes all the masked sentences and the corresponding true conditional probability.

Original: When the transformer is sufficiently large, the lookup table can be encoded as shown in Yun et al.
Revised: When the transformer is sufﬁciently large, the lookup table can be encoded as shown in Yun et al.

Original: 2. R ESULTS
Revised: 2. R ESULTS

Original: We compare the downstream performance of models with the same pre-training loss in the following situations: (2) training for different number of steps after the pre-training loss converges, (2) using different model sizes, and (3) training with normal training algorithms vs. adversarial training algorithms.
Revised: We compare the downstream performance of models with the same pre-training loss in the following situations: (1) training for different number of steps after the pre-training loss converges, (2) using different model sizes, and (3) training with normal training algorithms vs. adversarial training algorithms.

Original: In Figure 1, we plot the validation pre-training loss and the downstream performance of different models checkpoints along pre-training.
Revised: In Figure 1, we plot the validation pre-training loss and the downstream performance of different models checkpoints along pre-training.

Original: After the pre-training loss converges, although the pre-training loss does not improve, the downstream accuracy continues increasing.
Revised: After the pre-training loss converges, although the pre-training loss does not improve, the downstream accuracy continues increasing.

Original: Even with the same pre-training loss, larger models are better than smaller models.
Revised: Even with the same pre-training loss, larger models are better than smaller models.

Original: In Figure 2, we plot the pre-training loss and the downstream performance of models with different sizes.
Revised: In Figure 2, we plot the pre-training loss and the downstream performance of models with different sizes.

Original: As we increase the model size, the pre-training loss approaches the entropy of the true conditional probability, which is 3.196, 3.758, and 1.865 for PCFG, HMM, and OPT respectively.
Revised: As we increase the model size, the pre-training loss approaches the entropy of the true conditional probability, which is 3.196, 3.758, and 1.865 for PCFG, HMM, and OPT respectively.

Original: For PCFG and OPT-generated datasets, we use the vertical dashed line to indicate the place where the pre-training loss saturates as we scale up the model.
Revised: For PCFG and OPT-generated datasets, we use the vertical dashed line to indicate the place where the pre-training loss saturates as we scale up the model.

Original: For the much simpler HMM, the smallest 4M model can fit pre-training close to the entropy of the true conditional probability.
Revised: For the much simpler HMM, the smallest 4M model can ﬁt pre-training close to the entropy of the true conditional probability.

Original: With the same pre-training loss, scaling up the models improves linear probe performance by 6 . 9% , 4 . 5% , and 2 . 0% , on PCFG, HMM, and OPT generated data, respectively.
Revised: With the same pre-training loss, scaling up the models improves linear probe performance by 6 . 9% , 4 . 5% , and 2 . 0% , on PCFG, HMM, and OPT generated data, respectively.

Original: See Section A for results on other downstream tasks.
Revised: See Section A for results on other downstream tasks.

Original: [Table]
Revised: [Table]

Original: Naturally trained transformers are better than adversarially trained ones.
Revised: Naturally trained transformers are better than adversarially trained ones.

Original: In Table 1, we evaluate the 235M transformers on PCFG tasks A and B with different pre-training algorithms.
Revised: In Table 1, we evaluate the 235M transformers on PCFG tasks A and B with different pre-training algorithms.

Original: Although the adversarially trained transformer has almost the same pre-training loss as the normally trained 235M transformer, it is more than 6 % worse than the normal 235M model, and even worse than a normal 9M model on the downstream task
Revised: Although the adversarially trained transformer has almost the same pre-training loss as the normally trained 235M transformer, it is more than 6 % worse than the normal 235M model, and even worse than a normal 9M model on the downstream task

Original: B.
Revised: B.

Original: The lookup table has perfect pre-training loss, but it performs worse than all normally trained transformers in Figure 2(a) on task B. Note that this is different from the label-orthogonal training in Saunshi et al.
Revised: The lookup table has perfect pre-training loss, but it performs worse than all normally trained transformers in Figure 2(a) on task B. Note that this is different from the label-orthogonal training in Saunshi et al.

Original: They find out models with the same pre-training loss and different downstream performance by subtracting the mean of the representations, essentially changing the architecture, while our experiment compares models with the same architecture.
Revised: They ﬁnd out models with the same pre-training loss and different downstream performance by subtracting the mean of the representations, essentially changing the architecture, while our experiment compares models with the same architecture.

Original: The experiments above indicate that for models with the same architecture family and the same pretraining loss, the choice of training algorithms, model sizes, and the number of steps that the optimizer works can affect the downstream performance.
Revised: The experiments above indicate that for models with the same architecture family and the same pretraining loss, the choice of training algorithms, model sizes, and the number of steps that the optimizer works can affect the downstream performance.

Original: This indicates the existence of implicit bias of the training algorithms toward choosing more transferable parameter configurations among those with the same pre-training loss and architecture.
Revised: This indicates the existence of implicit bias of the training algorithms toward choosing more transferable parameter conﬁgurations among those with the same pre-training loss and architecture.

Original:  I MPLICIT B IAS L EADS TO F LAT S OLUTIONS IN L ANGUAGE M ODELING
Revised:  I MPLICIT B IAS L EADS TO F LAT S OLUTIONS IN L ANGUAGE M ODELING

Original: As discussed in the introduction, the difference between the role of implicit bias in supervised learning andlanguagemodelingisconceptual, becausethegapbetweenempiricalandpopulationself-supervised loss is small and thus implicit bias is not needed for bridging this gap.
Revised: As discussed in the introduction, the difference between the role of implicit bias in supervised learning andlanguagemodelingisconceptual, becausethegapbetweenempiricalandpopulationself-supervised loss is small and thus implicit bias is not needed for bridging this gap.

Original: Instead, the implicit bias benefits the performance on downstream tasks by picking networks that are more adaptable to those tasks.
Revised: Instead, the implicit bias beneﬁts the performance on downstream tasks by picking networks that are more adaptable to those tasks.

Original: Fortunately, the mathematical tools developed for supervised learning can be adapted to language modeling, which even allows cleaner results by removing some artificial assumptions like adding label noise.
Revised: Fortunately, the mathematical tools developed for supervised learning can be adapted to language modeling, which even allows cleaner results by removing some artiﬁcial assumptions like adding label noise.

Original: Concretely, we show that mini-batch SGD can find models in the flatter areas of pre-training loss landscape.
Revised: Concretely, we show that mini-batch SGD can ﬁnd models in the ﬂatter areas of pre-training loss landscape.

Original: The flatness is measured by the trace of Hessian of the pre-training loss
Revised: The ﬂatness is measured by the trace of Hessian of the pre-training loss

Original: Tr[ ∇ 2 L ( θ )] .
Revised: Tr[ ∇ 2 L ( θ )] .

Original: See
Revised: See

Original: Figure 4 (Left) for an illustration of the implicit bias.
Revised: Figure 4 (Left) for an illustration of the implicit bias.

Original: We analyze SGD on the population cross-entropy loss L ( θ )=
Revised: We analyze SGD on the population cross-entropy loss L ( θ )=

Original: E x,t [ − log([ f θ ( x − t )] x t )] with freshly sampled data at every iteration, because, as argued, the difference between empirical and population pre-training loss is not our focus.
Revised: E x,t [ − log([ f θ ( x − t )] x t )] with freshly sampled data at every iteration, because, as argued, the difference between empirical and population pre-training loss is not our focus.

Original: For simplicity, we present the results for batch size =1 , though they can be generalized to arbitrary batch size (see discussion below Theorem 3.3).
Revised: For simplicity, we present the results for batch size =1 , though they can be generalized to arbitrary batch size (see discussion below Theorem 3.3).

Original: Let η be the learning rate and let θ ηk denote the parameter at step k .
Revised: Let η be the learning rate and let θ ηk denote the parameter at step k .

Original: We drop the superscript η when there is no ambiguity.
Revised: We drop the superscript η when there is no ambiguity.

Original: We will show that the implicit bias kicks in when SGD reaches a global minimizer—it drives the iterate towards flatter global minimizers.
Revised: We will show that the implicit bias kicks in when SGD reaches a global minimizer—it drives the iterate towards ﬂatter global minimizers.

Original: For simplicity of demonstration, we analyze the process starting from a global minimizer θ , i.e. , we assume that θ η 0 = θ (for all η ).
Revised: For simplicity of demonstration, we analyze the process starting from a global minimizer θ , i.e. , we assume that θ η 0 = θ (for all η ).

Original: At each iteration k , we get a fresh sample ( x,t ) , where x is a sentence and t is the position of the masked token, and update the parameter θ by θ k +

Revised: At each iteration k , we get a fresh sample ( x,t ) , where x is a sentence and t is the position of the masked token, and update the parameter θ by θ k +

Original: =
Revised: =

Original: θ k − η ∇ θ ℓ ( f θ k ( x − t ) ,x t ) .
Revised: θ k − η ∇ θ (cid:96) ( f θ k ( x − t ) ,x t ) .

Original: We assume the network is sufficiently expressive such that there are many fundamentally different global minimizers of the pre-training loss L .
Revised: We assume the network is sufﬁciently expressive such that there are many fundamentally different global minimizers of the pre-training loss L .

Original: As a (non-trivial) regularity condition, following prior works (Fehrman et al., 2020; Li et al., 2021; Arora et al., 2022), we also assume that the minimizers of the loss function L are connected and form a smooth manifold.
Revised: As a (non-trivial) regularity condition, following prior works (Fehrman et al., 2020; Li et al., 2021; Arora et al., 2022), we also assume that the minimizers of the loss function L are connected and form a smooth manifold.

Original: Assumption 3.1.
Revised: Assumption 3.1.

Original: Assume that the loss L is a C 3 -smooth function, and that the set of global minimizers, Γ , is a ( d − M )
Revised: Assume that the loss L is a C 3 -smooth function, and that the set of global minimizers, Γ , is a ( d − M )

Original: -dimensional
Revised: -dimensional

Original: C 2 -submanifold of R d for some integer 1 ≤ M ≤ d , where for all θ ∈ Γ , rank (cid:0) ∇ 2 L ( θ ) (cid:1) =
Revised: C 2 -submanifold of R d for some integer 1 ≤ M ≤ d , where for all θ ∈ Γ , rank (cid:0) ∇ 2 L ( θ ) (cid:1) =

Original: M .
Revised: M .

Original: A key observation for language model is that even if the model reaches the saturation regime, that is, the model reaches a point on the manifold Γ of the minimizers, the optimization process still has non-vanishing gradient noise, because the cross-entropy loss is typically non-zero at the global minimizers and thus the stochastic gradient variance is also non-zero.
Revised: A key observation for language model is that even if the model reaches the saturation regime, that is, the model reaches a point on the manifold Γ of the minimizers, the optimization process still has non-vanishing gradient noise, because the cross-entropy loss is typically non-zero at the global minimizers and thus the stochastic gradient variance is also non-zero.

Original: Therefore, the dynamics of
Revised: Therefore, the dynamics of

Original: [Figure]
Revised: [Figure]

Original: SGD do not completely stop; instead, the iterate oscillates around the manifold Γ .
Revised: SGD do not completely stop; instead, the iterate oscillates around the manifold Γ .

Original: It turns out that this oscillation in turn encourages the parameter to move in a certain direction along the manifold, determined by the covariance structure of the stochastic gradient.
Revised: It turns out that this oscillation in turn encourages the parameter to move in a certain direction along the manifold, determined by the covariance structure of the stochastic gradient.

Original: The following lemma shows that the covariance of stochastic gradient for language models in the saturating regime has a favorable property, i.e. , it is equal to the Hessian of pre-training loss.
Revised: The following lemma shows that the covariance of stochastic gradient for language models in the saturating regime has a favorable property, i.e. , it is equal to the Hessian of pre-training loss.

Original: Lemma 3.2 (Bartlett identity) .
Revised: Lemma 3.2 (Bartlett identity) .

Original: For any θ ∈ Γ , Σ( θ ) = ∇ 2 L ( θ ) , where Σ( θ ) is the covariance of the stochastic gradient at θ , that is, Σ( θ ) =
Revised: For any θ ∈ Γ , Σ( θ ) = ∇ 2 L ( θ ) , where Σ( θ ) is the covariance of the stochastic gradient at θ , that is, Σ( θ ) =

Original: E t,x (cid:2) ∇ θ log[ f θ ( x − t )] x t ( ∇ θ log[ f θ ( x − t )] x t ) ⊤ (cid:3) − ∇ L ( θ ) ⊤ ∇ L ( θ ) .
Revised: E t,x (cid:2) ∇ θ log[ f θ ( x − t )] x t ( ∇ θ log[ f θ ( x − t )] x t ) (cid:62) (cid:3) − ∇ L ( θ ) (cid:62) ∇ L ( θ ) .

Original: Though we give a proof of the lemma in Appendix C for completeness, the formula holds for the MLE loss of any well-specified probabilistic models at a global minimizer, and both the gradient covariance and the Hessian equals to the Fisher information matrix.
Revised: Though we give a proof of the lemma in Appendix C for completeness, the formula holds for the MLE loss of any well-speciﬁed probabilistic models at a global minimizer, and both the gradient covariance and the Hessian equals to the Fisher information matrix.

Original: With Lemma 3.2, we can invoke Corollary 5.2 of Li et al.
Revised: With Lemma 3.2, we can invoke Corollary 5.2 of Li et al.

Original: (2021) to derive the following theorem which says that SGD will locally decrease the trace of Hessian along the solution of ordinary differential equation (1) defined below.
Revised: (2021) to derive the following theorem which says that SGD will locally decrease the trace of Hessian along the solution of ordinary differential equation (1) deﬁned below.

Original: [Equation]
Revised: [Equation]

Original: where ∇ Γ =
Revised: where ∇ Γ =

Original: P ⊥ Γ ∇ is the Riemannian gradient on manifold Γ , or just the ordinary gradient projected back to the tangent space of Γ at θ .
Revised: P ⊥ Γ ∇ is the Riemannian gradient on manifold Γ , or just the ordinary gradient projected back to the tangent space of Γ at θ .

Original: In other words, the ODE (1) is essentially a projected gradient descent algorithm with loss function Tr[ ∇ 2 L ( θ )] , the constraint set Γ , and infinitesimal learning rate.
Revised: In other words, the ODE (1) is essentially a projected gradient descent algorithm with loss function Tr[ ∇ 2 L ( θ )] , the constraint set Γ , and inﬁnitesimal learning rate.

Original: We show that SGD effectively minimizes the trace of the Hessian Tr[ ∇ 2 L ( θ )] with the constraint set Γ similarly to ODE in (1).
Revised: We show that SGD effectively minimizes the trace of the Hessian Tr[ ∇ 2 L ( θ )] with the constraint set Γ similarly to ODE in (1).

Original: Theorem 3.3.
Revised: Theorem 3.3.

Original: Suppose the loss function L and the manifold of global minimizers Γ satisfy Assumption 3.1.
Revised: Suppose the loss function L and the manifold of global minimizers Γ satisfy Assumption 3.1.

Original: For any K > 0 such that ODE (1) has a solution { ˆ θ ( t ) } Kt =0 , it holds that θ ηK/η 2 converges in distribution to ˆ θ ( K ) as η → 0 .
Revised: For any K > 0 such that ODE (1) has a solution { ˆ θ ( t ) } Kt =0 , it holds that θ ηK/η 2 converges in distribution to ˆ θ ( K ) as η → 0 .

Original: Finally, we note that the above result can be extended to an arbitrary batch size B .
Revised: Finally, we note that the above result can be extended to an arbitrary batch size B .

Original: The covariance of stochastic gradient at θ with batch size, denoted by Σ B ( θ ) , satisfies that Σ B ( θ )= 1 B Σ( θ ) .
Revised: The covariance of stochastic gradient at θ with batch size, denoted by Σ B ( θ ) , satisﬁes that Σ B ( θ )= 1 B Σ( θ ) .

Original: Therefore Σ B ( θ )= 1 B ∇ 2 L ( θ ) and we can again invoke Corollary 5.2 of Li et al. (2021) to derive the same result as in Theorem 3.3 but with the coefficient 14 in equation (1) replaced by 14 B .
Revised: Therefore Σ B ( θ )= 1 B ∇ 2 L ( θ ) and we can again invoke Corollary 5.2 of Li et al. (2021) to derive the same result as in Theorem 3.3 but with the coefﬁcient 14 in equation (1) replaced by 14 B .

Original:  F LATTER M ODELS H AVE B ETTER D OWNSTREAM P ERFORMANCE
Revised:  F LATTER M ODELS H AVE B ETTER D OWNSTREAM P ERFORMANCE

Original: In this section, we demonstrate with experiments that the flatness is well correlated with downstream performance in the setting introduced in Section 2.
Revised: In this section, we demonstrate with experiments that the ﬂatness is well correlated with downstream performance in the setting introduced in Section 2.

Original: Evaluation of flatness.
Revised: Evaluation of ﬂatness.

Original: As in Theorem 3.3, we measure the flatness of different models in Section 2 by the trace of Hessian of the pre-training loss (smaller trace of Hessian indicates flatter minima.)
Revised: As in Theorem 3.3, we measure the ﬂatness of different models in Section 2 by the trace of Hessian of the pre-training loss (smaller trace of Hessian indicates ﬂatter minima.)

Original: Note that when the model approaches the saturation regime, the trace of Hessian is approximately the second order derivative times the square of the norm of the Jacobian, which is a high-dimensional matrix.
Revised: Note that when the model approaches the saturation regime, the trace of Hessian is approximately the second order derivative times the square of the norm of the Jacobian, which is a high-dimensional matrix.

Original: For computational feasibility, we adopt a technique inspired by Wei et al.
Revised: For computational feasibility, we adopt a technique inspired by Wei et al.

Original: (2020) to unbiasedly estimate the trace of Hessian with random samples.
Revised: (2020) to unbiasedly estimate the trace of Hessian with random samples.

Original: Details are provided in Section B.
Revised: Details are provided in Section B.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original:  D o w n s t r e a m A cc u r a c
Revised: M D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: OPT
Revised: OPT

Original: Linear Probe, QNLI 1 . 901 .
Revised: Linear Probe, QNLI 1 . 901 .

Original: P r e - t r a i n i n g L o ss 26.
Revised: P r e - t r a i n i n g L o ss 26.

Original: Pre-training Trace of Hessian(
Revised: Pre-training Trace of Hessian

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Results.
Revised: Results.

Original: In Figure 3, we compare the downstream accuracy and the trace of Hessian of different checkpoints obtained at different times during pre-training.
Revised: In Figure 3, we compare the downstream accuracy and the trace of Hessian of different checkpoints obtained at different times during pre-training.

Original: On the PCFG and HMM datasets, the trace of Hessian demonstrates a clear decreasing trend after the validation pre-training loss converges, following the predictionofTheorem3.3.
Revised: On the PCFG and HMM datasets, the trace of Hessian demonstrates a clear decreasing trend after the validation pre-training loss converges, following the predictionofTheorem3.3.

Original: Furthermore, asthetraceofHessiandecreases, thedownstreamperformance improves by 1.6 % and 4.0 % on the PCFG and HMM datasets, respectively.
Revised: Furthermore, asthetraceofHessiandecreases, thedownstreamperformance improves by 1.6 % and 4.0 % on the PCFG and HMM datasets, respectively.

Original: We compare the trace of Hessian of the models pre-trained with adversarial algorithm and standard
Revised: We compare the trace of Hessian of the models pre-trained with adversarial algorithm and standard

Original: AdamW in Table 2.
Revised: AdamW in Table 2.

Original: The trace of Hessian of the adversarially pre-trained model is 3 times larger than the normally pre-trained model, corresponding to a drop of 5.5 % in downstream performance.
Revised: The trace of Hessian of the adversarially pre-trained model is 3 times larger than the normally pre-trained model, corresponding to a drop of 5.5 % in downstream performance.

Original: In Figure 5, we compare the downstream accuracy and the trace of Hessian of models with different sizes.
Revised: In Figure 5, we compare the downstream accuracy and the trace of Hessian of models with different sizes.

Original: On the dataset generated by a PCFG, the pre-training loss is almost the same for models larger than 9M. As we increase the model size, the trace of Hessian of the pre-training loss decreases from 2.68 to 1.54, correlating with the increase of linear probe accuracy from 43.2 % to 50.3 % .
Revised: On the dataset generated by a PCFG, the pre-training loss is almost the same for models larger than 9M. As we increase the model size, the trace of Hessian of the pre-training loss decreases from 2.68 to 1.54, correlating with the increase of linear probe accuracy from 43.2 % to 50.3 % .

Original: On the OPT and HMM-generated datasets, we can also observe an increase in linear probe accuracy with a sharp decrease in the trace of Hessian, as we increase the model size.
Revised: On the OPT and HMM-generated datasets, we can also observe an increase in linear probe accuracy with a sharp decrease in the trace of Hessian, as we increase the model size.

Original: Interaction between implicit bias and model size.
Revised: Interaction between implicit bias and model size.

Original: Intuitively, the implicit bias drives the model toward flat minima on both larger models and smaller models.
Revised: Intuitively, the implicit bias drives the model toward ﬂat minima on both larger models and smaller models.

Original: The smaller transformer architecture is a subset of the larger transformer architecture family (as justified in Section B).
Revised: The smaller transformer architecture is a subset of the larger transformer architecture family (as justiﬁed in Section B).

Original: Thus the flattest minimum found within a larger transformer is flatter than the flattest minimum found within a smaller transformer, and performs better downstream.
Revised: Thus the ﬂattest minimum found within a larger transformer is ﬂatter than the ﬂattest minimum found within a smaller transformer, and performs better downstream.

Original: (See Figure 4 (Right).)
Revised: (See Figure 4 (Right).)

Original:  F LATNESS R EGULARIZATION P ROVABLY I DENTIFIES T RANSFERABLE M ODELS ON S YNTHETIC L ANGUAGE
Revised:  F LATNESS R EGULARIZATION P ROVABLY I DENTIFIES T RANSFERABLE M ODELS ON S YNTHETIC L ANGUAGE

Original: Toward formally proving the connection between the flatness regularization (introduced by the stochastic gradient as argued in Section 3) and the downstream performance, we consider a setting with synthetic Dyck language.
Revised: Toward formally proving the connection between the ﬂatness regularization (introduced by the stochastic gradient as argued in Section 3) and the downstream performance, we consider a setting with synthetic Dyck language.

Original: The simplicity of the data allows us to sharply analyze the internal working of a single-layer transformer (with an attention layer and an MLP layer) for masked language modeling.
Revised: The simplicity of the data allows us to sharply analyze the internal working of a single-layer transformer (with an attention layer and an MLP layer) for masked language modeling.

Original: We show that multiple parameter configurations can predict the conditional probability well, including one ideal model that learns the correct representations capturing the intrinsic structure of the sentence, 𝑦 = − <<>-<> 1 1 - ± 1 - <-<><> ± 1 -1 1 -1 Encodings
Revised: We show that multiple parameter conﬁgurations can predict the conditional probability well, including one ideal model that learns the correct representations capturing the intrinsic structure of the sentence, 𝑦 = − <<>-<> 1 1 - ± 1 - <-<><> ± 1 -1 1 -1 Encodings

Original: Input 𝑋 𝑄 𝐾
Revised: Input 𝑋 𝑄 𝐾

Original: Softmax Attention 1- ± 1 1𝑧 1 1 1 1 1 

Revised: Softmax Attention 1- ± 1 1𝑧 1 1 1 1 1

Original: 𝑉 𝑉
Revised: 𝑉 𝑉

Original: Configuration (1)
Revised: Configuration (1)

Original: Configuration (2) -1 1 𝑢 𝑢 𝑦 = − 𝑥 !
Revised: Configuration (2) -1 1 𝑢 𝑢 𝑦 = − 𝑥 !

Original: 𝑥 " 𝑥 # 𝑥 $ 𝑥 % 𝑥 & and many “cheating” models that essentially memorize the conditional probability using random features.
Revised: 𝑥 " 𝑥 # 𝑥 $ 𝑥 % 𝑥 & and many “cheating” models that essentially memorize the conditional probability using random features.

Original: We will prove that the flattest model is the desired model that transfers to downstream tasks.
Revised: We will prove that the ﬂattest model is the desired model that transfers to downstream tasks.

Original: Pre-training Distribution.
Revised: Pre-training Distribution.

Original: Consider a variant of the Dyck language (Nivat, 1970) consisting of matching brackets.
Revised: Consider a variant of the Dyck language (Nivat, 1970) consisting of matching brackets.

Original: The vocabulary of the language has two brackets ⟨ and ⟩ .
Revised: The vocabulary of the language has two brackets (cid:104) and (cid:105) .

Original: Each sentence is composed of a sequence of tokens such that the total numbers of ⟨ ’s and ⟩ ’s are equal.
Revised: Each sentence is composed of a sequence of tokens such that the total numbers of (cid:104) ’s and (cid:105) ’s are equal.

Original: To sample from the pre-training distribution P , we first draw a sentence uniformly over all valid sentences with even length T .
Revised: To sample from the pre-training distribution P , we ﬁrst draw a sentence uniformly over all valid sentences with even length T .

Original: Then, we randomly select one position in [ T ] and replace the bracket with a mask token.
Revised: Then, we randomly select one position in [ T ] and replace the bracket with a mask token.

Original: Downstream Task.
Revised: Downstream Task.

Original: The most intrinsic property about the synthetic language is the difference in the number of left and right brackets in a prefix, and thus we use it as the downstream task.
Revised: The most intrinsic property about the synthetic language is the difference in the number of left and right brackets in a preﬁx, and thus we use it as the downstream task.

Original: Concretely, for any sequence x in {⟨ , ⟩} ∗ of length T , let g ∗ ( x ) count the number of mismatches in x :
Revised: Concretely, for any sequence x in {(cid:104) , (cid:105)} ∗ of length T , let g ∗ ( x ) count the number of mismatches in x :

Original: [Equation]
Revised: [Equation]

Revised: # of (cid:105) ’s in x − # of (cid:104) ’s in x (2)

Original: Thus, the sentence x is a valid string in the language if and only if g ∗ ( x )=0 .
Revised: Thus, the sentence x is a valid string in the language if and only if g ∗ ( x )=0 .

Original: For MLM, the masked token can also be recovered from g ∗ ( x ) : g ∗ ( x ) = 1 if the masked token is ⟨ , and g ∗ ( x ) = − 1 if the masked token is ⟩ .
Revised: For MLM, the masked token can also be recovered from g ∗ ( x ) : g ∗ ( x ) = 1 if the masked token is (cid:104) , and g ∗ ( x ) = − 1 if the masked token is (cid:105) .

Original: To evaluate if the model learns the structure, we consider a downstream distribution of sentences which do not necessarily belong to the the language.
Revised: To evaluate if the model learns the structure, we consider a downstream distribution of sentences which do not necessarily belong to the the language.

Original: Each token is sampled from {⟨ , ⟩} uniformly, randomly, and independently.
Revised: Each token is sampled from {(cid:104) , (cid:105)} uniformly, randomly, and independently.

Original: Encoding of the Inputs.
Revised: Encoding of the Inputs.

Original: With a slight abuse of notation, we also denote by x t the encoding of the t -th token.
Revised: With a slight abuse of notation, we also denote by x t the encoding of the t -th token.

Original: We encode the input as a one-hot vector in dimension d = 2 T , where the index of the nonzero element encodes the position and the sign encodes the bracket.
Revised: We encode the input as a one-hot vector in dimension d = 2 T , where the index of the nonzero element encodes the position and the sign encodes the bracket.

Original: Concretely, let e t ∈ R
Revised: Concretely, let e t ∈ R

Original: d be the natural basis vector where the t -th entry is 1.
Revised: d be the natural basis vector where the t -th entry is 1.

Original: Let x t =
Revised: Let x t =

Original: e t if the t -th token is ⟨ and x t = − e t otherwise.
Revised: e t if the t -th token is (cid:104) and x t = − e t otherwise.

Original: If the position t is a mask, we set x t to v , where v ∼ Unif ( {± e t + T } ) .
Revised: If the position t is a mask, we set x t to v , where v ∼ Unif ( {± e t + T } ) .

Original: Examples of the encodings with T =6 are provided in Figure 6 (Left).
Revised: Examples of the encodings with T =6 are provided in Figure 6 (Left).

Original: Note that the target function can be expressed as g ⋆ ( x )= 
Revised: Note that the target function can be expressed as g (cid:63) ( x )= −(cid:104) 1 T , [ (cid:80) Tt

Original: −⟨ 1 T , [ (cid:80) Tt =1 x t ] 1: T ⟩ with this input encoding, where 1 T is the all one vector in R T and [ a ] 1: T refers to the first T coordinates in a .
Revised: =1 x t ] 1: T (cid:105) with this input encoding, where 1 T is the all one vector in R T and [ a ] 1: T refers to the ﬁrst T coordinates in a .

Original: Models and Algorithms.
Revised: Models and Algorithms.

Original: Suppose Q,K ∈
Revised: Suppose Q,K ∈

Original: R
Revised: R

Original: k
Revised: k

Original: ×
Revised: ×

Original: d are the query and key matrices, V ∈ R m ×
Revised: d are the query and key matrices, V ∈ R m ×

Original: d is the value matrix and u ∈ R m is the parameter of the output layer.
Revised: d is the value matrix and u ∈ R m is the parameter of the output layer.

Original: Let ψ =( Q,K,V ) .
Revised: Let ψ =( Q,K,V ) .

Original: A single-layer transformer is composed of an attention layer and an MLP layer.
Revised: A single-layer transformer is composed of an attention layer and an MLP layer.

Original: [ Attn ψ,u ( x )] t = 1 m u ⊤ σ ( (cid:80) T j =1 a t,j V x j ) , where the attention score a t, 1: T = softmax ( ⟨ Qx t ,Kx t ⟩ , ···⟨ Qx t ,Kx T ⟩ ) .
Revised: [ Attn ψ,u ( x )] t = 1 m u (cid:62) σ ( (cid:80) T j =1 a t,j V x j ) , where the attention score a t, 1: T = softmax ( (cid:104) Qx t ,Kx t (cid:105) , ···(cid:104) Qx t ,Kx T (cid:105) ) .

Original: σ ( x )=max { x, 0 } is the relu activation.
Revised: σ ( x )=max { x, 0 } is the relu activation.

Original: We use the output of the first token, f ψ,u ( x )=[
Revised: We use the output of the ﬁrst token, f ψ,u ( x )=[

Original: Attn ψ,u ( x )] 1 .
Revised: Attn ψ,u ( x )] 1 .

Original: We use the squared loss for both MLM and downstream adaptation.
Revised: We use the squared loss for both MLM and downstream adaptation.

Original: The loss function of MLM is
Revised: The loss function of MLM is

Original: L ( ψ,u ) .
Revised: L ( ψ,u ) .

Original: In downstream adaptation, we have a finite dataset { x ( i ) } ni =1 sampled i.i.d.
Revised: In downstream adaptation, we have a ﬁnite dataset { x ( i ) } ni =1 sampled i.i.d.

Original: from P ds .
Revised: from P ds .

Original: The training loss with n data is (cid:98) L P ds ( ψ,u ) , and the population loss for the downstream task is L P ds ( ψ,u ) .
Revised: The training loss with n data is (cid:98) L P ds ( ψ,u ) , and the population loss for the downstream task is L P ds ( ψ,u ) .

Original: Main Intuitions.
Revised: Main Intuitions.

Original: We are interested in two kinds of parameter configurations both with good pretraining loss: (1) learning the natural and transferable features 1 T and (2) fitting the pre-training task by memorizing the masked sentences.
Revised: We are interested in two kinds of parameter conﬁgurations both with good pretraining loss: (1) learning the natural and transferable features 1 T and (2) ﬁtting the pre-training task by memorizing the masked sentences.

Original: We construct the two solutions as follows.
Revised: We construct the two solutions as follows.

Original: For solution (1), first note that the softmax attention layer can take the average of all the token encodings [ x t ] T t =1 in a sentence.
Revised: For solution (1), ﬁrst note that the softmax attention layer can take the average of all the token encodings [ x t ] T t =1 in a sentence.

Original: Let us denote the sum by z ∈ R d , z = (cid:80)
Revised: Let us denote the sum by z ∈ R d , z = (cid:80)

Original: Tt =1 x t .
Revised: Tt =1 x t .

Original: Note that the first T coordinates in z are ± 1 indicating the bracket type and the last T coordinates indicate the position of the mask (See Figure 6 (Right)).
Revised: Note that the ﬁrst T coordinates in z are ± 1 indicating the bracket type and the last T coordinates indicate the position of the mask (See Figure 6 (Right)).

Original: On top of z , two neurons can predict the masked token in MLM perfectly.
Revised: On top of z , two neurons can predict the masked token in MLM perfectly.

Original: Consider the two neurons
Revised: Consider the two neurons

Original: V 1 =[ 1 T ; 0 T ] , V 2 =[ − 1 T ; 0 T ] .
Revised: V 1 =[ 1 T ; 0 T ] , V 2 =[ − 1 T ; 0 T ] .

Original: Then g ∗ ( x )=
Revised: Then g ∗ ( x )=

Original: σ ( V ⊤ 2 z ) − σ ( V ⊤ 1 z ) , which is the transferable solution.
Revised: σ ( V (cid:62) 2 z ) − σ ( V (cid:62) 1 z ) , which is the transferable solution.

Original: For solution (2), we set the entries in V to i.i.d.
Revised: For solution (2), we set the entries in V to i.i.d.

Original: samples from N (0 ,T ) .
Revised: samples from N (0 ,T ) .

Original: If m is sufficiently large, we can find the coefficient u to express g ∗ ( x ) with random Gaussian features, i.e. g ∗ ( x )=
Revised: If m is sufﬁciently large, we can ﬁnd the coefﬁcient u to express g ∗ ( x ) with random Gaussian features, i.e. g ∗ ( x )=

Original: u ⊤ σ ( V z ) .
Revised: u (cid:62) σ ( V z ) .

Original: We observe that the trace of Hessian of configuration (1) is smaller than configuration (2), due to a main difference between them–the cancellation between activated neurons.
Revised: We observe that the trace of Hessian of conﬁguration (1) is smaller than conﬁguration (2), due to a main difference between them–the cancellation between activated neurons.

Original: In configuration (1), for every possible input, only one of the neurons σ ( V ⊤ 1 z ) and σ ( V ⊤ 2 z ) is activated.
Revised: In conﬁguration (1), for every possible input, only one of the neurons σ ( V (cid:62) 1 z ) and σ ( V (cid:62) 2 z ) is activated.

Original: In contrast, in configuration (2), many neurons can be activated at the same time.
Revised: In contrast, in conﬁguration (2), many neurons can be activated at the same time.

Original: Among them, the output coefficient u
Revised: Among them, the output coefﬁcient u

Original: i ’s contain both positive and negative values, leading to cancellation between activated neurons.
Revised: i ’s contain both positive and negative values, leading to cancellation between activated neurons.

Original: In Lemma D.1, we link the trace of the Hessian with the cancellation between neurons.
Revised: In Lemma F.1, we link the trace of the Hessian with the cancellation between neurons.

Original: Indeed, we show that the mimimum of trace of the Hessian can be achieved only if there is no such cancellation.
Revised: Indeed, we show that the mimimum of trace of the Hessian can be achieved only if there is no such cancellation.

Original: Therefore solution (1) is also the minimizer of the trace of Hessian.
Revised: Therefore solution (1) is also the minimizer of the trace of Hessian.

Original: The intuitions are formalized in Theorem 5.1.
Revised: The intuitions are formalized in Theorem 5.1.

Original: Consider minimizing the trace of Hessian among all the solutions to the MLM pre-training task: minimize ψ,
Revised: Consider minimizing the trace of Hessian among all the solutions to the MLM pre-training task: minimize ψ,

Original: u Tr[ ∇ 2 ψ L ( ψ,u )]
Revised: u Tr[ ∇ 2 ψ L ( ψ,u )]

Original: +Tr[ ∇ 2 u L ( ψ,u )] , subject to L ( ψ,u )=0 .
Revised: +Tr[ ∇ 2 u L ( ψ,u )] , subject to L ( ψ,u )=0 .

Original: Theorem 5.1.
Revised: Theorem 5.1.

Original: Suppose m ≥ 2 and T ≥ 6 .
Revised: Suppose m ≥ 2 and T ≥ 6 .

Original: The flattest solution ˆ ψ, ˆ u are defined as the solution of the optimization problem above.
Revised: The ﬂattest solution ˆ ψ, ˆ u are deﬁned as the solution of the optimization problem above.

Original: ˜ u is the minimizer of downstream training loss on top of ˆ ψ , ˜ u ∈ argmin u ∥ u ∥ 2 subject to (cid:98) L P ds ( ˆ ψ,u )=0 .
Revised: ˜ u is the minimizer of downstream training loss on top of ˆ ψ , ˜ u ∈ argmin u (cid:107) u (cid:107) 2 subject to (cid:98) L P ds ( ˆ ψ,u )=0 .

Original: Then with probability at least 1 − 2 − n , L P ds ( ˆ ψ, ˜ u )=0 .
Revised: Then with probability at least 1 − 2 − n , L P ds ( ˆ ψ, ˜ u )=0 .

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Language modeling and downstream adaptation.
Revised: Language modeling and downstream adaptation.

Original: Large language modeling has revolutionized the
Revised: Large language modeling has revolutionized the

Original: NLP field.
Revised: NLP ﬁeld.

Original: Starting from Devlin et al.
Revised: Starting from Devlin et al.

Original: (2018), a line of works improve the downstream performance on a wide range of tasks with increasing model size and data amount (Yang et al., 2019; Radford et al., 2019; Raffel et al., 2020).
Revised: (2018), a line of works improve the downstream performance on a wide range of tasks with increasing model size and data amount (Yang et al., 2019; Radford et al., 2019; Raffel et al., 2020).

Original: LLMs even exhibit unexpected emergent behaviors, such as in-context learning (Xie et al., 2021;
Revised: LLMs even exhibit unexpected emergent behaviors, such as in-context learning (Xie et al., 2021;

Original: Min et al., 2022), step-by-step reasoning (Wei et al., 2022), and zero-shot learning (Brown et al., 2020).
Revised: Min et al., 2022), step-by-step reasoning (Wei et al., 2022), and zero-shot learning (Brown et al., 2020).

Original: Kaplan et al.
Revised: Kaplan et al.

Original: ; Hernandez et al.
Revised: ; Hernandez et al.

Original: (2021) study the behavior of language models with increasing size, and find out that the pre-training loss is typically correlated with downstream performance as model size increases.
Revised: (2021) study the behavior of language models with increasing size, and ﬁnd out that the pre-training loss is typically correlated with downstream performance as model size increases.

Original: In practice, the pre-training loss is used as an evaluation metric for language models.
Revised: In practice, the pre-training loss is used as an evaluation metric for language models.

Original: A notable example is the efficient transformer line of works, which benchmark the pre-training loss given the same computation constraint (Dai et al., 2020;
Revised: A notable example is the efﬁcient transformer line of works, which benchmark the pre-training loss given the same computation constraint (Dai et al., 2020;

Original: Wang et al., 2020; Choromanski et al., 2020; Liu et al., 2021).
Revised: Wang et al., 2020; Choromanski et al., 2020; Liu et al., 2021).

Original: Understanding the success of language modeling.
Revised: Understanding the success of language modeling.

Original: Empirical works on understanding MLM find out that the representations of language models encode rich semantic and syntactic information (Peters et al., 2018;
Revised: Empirical works on understanding MLM ﬁnd out that the representations of language models encode rich semantic and syntactic information (Peters et al., 2018;

Original: Htut et al., 2019; Hewitt & Manning, 2019;
Revised: Htut et al., 2019; Hewitt & Manning, 2019;

Original: Mamou et al., 2020).
Revised: Mamou et al., 2020).

Original: Theoretical works show that fitting the MLM conditional probability is a sufficient condition for good performance on downstream tasks.
Revised: Theoretical works show that ﬁtting the MLM conditional probability is a sufﬁcient condition for good performance on downstream tasks.

Original: Zhang & Hashimoto (2021) show MLM representations recover latent variables in graphical models.
Revised: Zhang & Hashimoto (2021) show MLM representations recover latent variables in graphical models.

Original: Wei et al.
Revised: Wei et al.

Original: (2021) show linear probe on top of MLM models solves downstream tasks on datasets generated by HMMs.
Revised: (2021) show linear probe on top of MLM models solves downstream tasks on datasets generated by HMMs.

Original: Recently, Saunshi et al.
Revised: Recently, Saunshi et al.

Original: (2022) show that models with the same pre-training loss but different architectures can have different downstream performance.
Revised: (2022) show that models with the same pre-training loss but different architectures can have different downstream performance.

Original: Tay et al.
Revised: Tay et al.

Original: (2021) find out that a narrow but deep transformer is better than a wide but shallow transformer with the same pre-training loss.
Revised: ﬁnd out that a narrow but deep transformer is better than a wide but shallow transformer with the same pre-training loss.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2022b) demonstrate that Albert (Lan et al., 2019) generalizes better to OOD tasks than Bert on a synthetic reasoning task.
Revised: (2022b) demonstrate that Albert (Lan et al., 2019) generalizes better to OOD tasks than Bert on a synthetic reasoning task.

Original: These works indicate that the architecture is an important factor for good downstream performance beyond pre-training loss.
Revised: These works indicate that the architecture is an important factor for good downstream performance beyond pre-training loss.

Original: This paper discovers the role of implicit bias in language modeling, which happens with models in the same architecture.
Revised: This paper discovers the role of implicit bias in language modeling, which happens with models in the same architecture.

Original: Implicit bias in supervised learning.
Revised: Implicit bias in supervised learning.

Original: The training algorithm chooses solutions with certain properties, and usually leads to better generalization (Gunasekar et al., 2018;
Revised: The training algorithm chooses solutions with certain properties, and usually leads to better generalization (Gunasekar et al., 2018;

Original: Soudry et al., 2018; Li et al., 2017; Ji & Telgarsky, 2018; Arora et al., 2019;
Revised: Soudry et al., 2018; Li et al., 2017; Ji & Telgarsky, 2018; Arora et al., 2019;

Original: Lyu & Li, 2019; Li et al., 2020; Woodworth et al., 2020; HaoChen et al., 2020).
Revised: Lyu & Li, 2019; Li et al., 2020; Woodworth et al., 2020; HaoChen et al., 2020).

Original: Recently, Blanc et al.
Revised: Recently, Blanc et al.

Original: ; Damian et al.
Revised: ; Damian et al.

Original: ; Li et al.
Revised: ; Li et al.

Original: (2021) demonstrate label noise SGD biases the models toward flatter minima.
Revised: (2021) demonstrate label noise SGD biases the models toward ﬂatter minima.

Original: However, the setting of implicit bias in supervised learning is different from language modeling.
Revised: However, the setting of implicit bias in supervised learning is different from language modeling.

Original: In language modeling, we have access to gigantic corpus, and cannot interpolate the pre-training dataset.
Revised: In language modeling, we have access to gigantic corpus, and cannot interpolate the pre-training dataset.

Original: Moreover, we care about the adaptability of the solution on downstream tasks instead of generalization in distribution.
Revised: Moreover, we care about the adaptability of the solution on downstream tasks instead of generalization in distribution.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We study the relationship between pre-training loss and downstream performance on language models.
Revised: We study the relationship between pre-training loss and downstream performance on language models.

Original: We discover that implicit bias matters beyond pre-training loss, and explore the mechanism of implicit bias in language modeling.
Revised: We discover that implicit bias matters beyond pre-training loss, and explore the mechanism of implicit bias in language modeling.

Original: Our experiments focus on simplified datasets due to constraint of computational resources.
Revised: Our experiments focus on simpliﬁed datasets due to constraint of computational resources.

Original: We hope that the phenomenon can predict the implicit bias on more complex datasets as the community scale the models to even larger.
Revised: We hope that the phenomenon can predict the implicit bias on more complex datasets as the community scale the models to even larger.

Original: Theoretically we provide cases where the flatness regularization can decide the performance on downstream performance.
Revised: Theoretically we provide cases where the ﬂatness regularization can decide the performance on downstream performance.

Original: We wish this motivates future works on the relationship between implicit bias and the internal working of the models.
Revised: We wish this motivates future works on the relationship between implicit bias and the internal working of the models.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: To ensure reproducibility, we describe the implementation details of the algorithms and the construction of the datasets in Section A and Section
Revised: To ensure reproducibility, we describe the implementation details of the algorithms and the construction of the datasets in Section A and Section

Original: B. The code of the experiments is provided in the supplementary material.
Revised: B. The code of the experiments is provided in the supplementary material.

Original: We provide the proof in Section C and Section D.
Revised: We provide the proof in Section C and Section F.

Original: A D ETAILS IN S ECTIONA. G ENERATING S IMPLIFIED D ATASETS
Revised: A D ETAILS IN S ECTIONA. G ENERATING S IMPLIFIED D ATASETS

Original: PCFG-generated dataset.
Revised: PCFG-generated dataset.

Original: We consider a PCFG with vocabulary size 200.
Revised: We consider a PCFG with vocabulary size 200.

Original: The state space is S , and | S | =50 .
Revised: The state space is S , and | S | =50 .

Original: All the production rules have two symbols on the right side.
Revised: All the production rules have two symbols on the right side.

Original: The sentence length is limited to 32, which means the depth of the parse tree is limited to 6.
Revised: The sentence length is limited to 32, which means the depth of the parse tree is limited to 6.

Original: We generate a total of 2 × 10 7 sentences, which is 3 .
Revised: We generate a total of 2 × 10 7 sentences, which is 3 .

Original: 4 × 10 8 tokens.
Revised: 4 × 10 8 tokens.

Original: The downstream tasks are classifying the non-terminal symbols in the parse tree of the PCFG (50-way classification).
Revised: The downstream tasks are classifying the non-terminal symbols in the parse tree of the PCFG (50-way classiﬁcation).

Original: The label is defined as y =argmax s ∈ S Pr( s | x 1 ,x 2 ,...,x 1+ L ) .
Revised: The label is deﬁned as y =argmax s ∈ S Pr( s | x 1 ,x 2 ,...,x 1+ L ) .

Original: Tasks A, B and C are defined on the symbols corresponding to span length L =32 , 16 and 8 , respectively.
Revised: Tasks A, B and C are deﬁned on the symbols corresponding to span length L =32 , 16 and 8 , respectively.

Original: Each of the downstream task contains 0.1M examples.
Revised: Each of the downstream task contains 0.1M examples.

Original: Examples of the generated trees are provided in
Revised: Examples of the generated trees are provided in

Original: Figure 7.
Revised: Figure 7.

Original: HMM-generated dataset.
Revised: HMM-generated dataset.

Original: We consider an HMM with vocabulary size 200 and state space size 100.
Revised: We consider an HMM with vocabulary size 200 and state space size 100.

Original: The sentence length is restricted to 16.
Revised: The sentence length is restricted to 16.

Original: We generate a total of 1 × 10 7 sentences, which is 1 . 6 × 10tokens.
Revised: We generate a total of 1 × 10 7 sentences, which is 1 . 6 × 10tokens.

Original: The downstream task is to classify the latent variable in the HMM generative model.
Revised: The downstream task is to classify the latent variable in the HMM generative model.

Original: We consider task-6 and task-10, which classify the 6-th and 10-th hidden variables respectively.
Revised: We consider task-6 and task-10, which classify the 6-th and 10-th hidden variables respectively.

Original: Each of the downstream task contains 0.1M examples.
Revised: Each of the downstream task contains 0.1M examples.

Original: OPT-generated dataset.
Revised: OPT-generated dataset.

Original: We use the 125M OPT model to generate the training dataset.
Revised: We use the 125M OPT model to generate the training dataset.

Original: To simplify the dataset, we further process the logit of OPT to select only from the top-2000 tokens in the vocabulary.
Revised: To simplify the dataset, we further process the logit of OPT to select only from the top-2000 tokens in the vocabulary.

Original: Starting from the bos token, we sample every token of the sentence from the predicted autoregressive
Revised: Starting from the bos token, we sample every token of the sentence from the predicted autoregressive

Original: LM probability.
Revised: LM probability.

Original: The sentence length is restricted to 24.
Revised: The sentence length is restricted to 24.

Original: We generate a total of 2 × 10 8 sentences, which is 3 .
Revised: We generate a total of 2 × 10 8 sentences, which is 3 .

Original: 2 × 10 9 tokens.
Revised: 2 × 10 9 tokens.

Original: Examples of the generated text are provided in Figure 8.[Figure]
Revised: Examples of the generated text are provided in Figure 8.

Original:  THE T RUE C ONDITIONAL P ROBABILITIES
Revised: A. C OMPUTE THE T RUE C ONDITIONAL P ROBABILITIES

Original: We can compute the true MLM conditional probability Pr ( x t | x − t ) from the joint probability
Revised: We can compute the true MLM conditional probability Pr ( x t | x − t ) from the joint probability

Original: Pr( x t ,x − t ) ,
Revised: Pr( x t ,x − t ) with one mask per token,

Original: [Equation]
Revised: [Equation]

Original: Since we already know the generative model, we can compute the joint probability efficiently.
Revised: Since we already know the generative model, we can compute the joint probability efﬁciently.

Original: For
Revised: For

Original: PCFG, we can compute the joint probability with the inside algorithm, which decomposes the joint probability into lower layers in the parse tree.
Revised: PCFG, we can compute the joint probability with the inside algorithm, which decomposes the joint probability into lower layers in the parse tree.

Original: For HMM, we can compute the joint probability with the
Revised: For HMM, we can compute the joint probability with the

Original: Viterbi algorithm.
Revised: Viterbi algorithm.

Original: For OPT, we have Pr( x 1 ,...,x T )=
Revised: For OPT, we have Pr( x 1 ,...,x T )=

Original: Pr( x 1 ) (cid:81) T − 1 t =1 Pr
Revised: Pr( x 1 ) (cid:81) T − 1 t =1 Pr

Original: x t +1 | x 1 ,...,x t ) .
Revised: x t +1 | x 1 ,...,x t ) .

Original: A. M ODELS
Revised: A. M ODELS

Original: We use transformers on PCFG and OPT-generated datasets.
Revised: We use transformers on PCFG and OPT-generated datasets.

Original: We use learning rate 1e-3 and warmup proportion 0.06.
Revised: We use learning rate 1e-3 and warmup proportion 0.06.

Original: All the models are trained based on the implementation of Izsak et al.
Revised: All the models are trained based on the implementation of Izsak et al.

Original: We list the sizes of the transformers in Table 3.
Revised: We list the sizes of the transformers in Table 3.

Original: d-model is the size of the hidden layers.
Revised: d-model is the size of the hidden layers.

Original: d-inter is the size of the
Revised: d-inter is the size of the

Original: [Table]
Revised: [Table]

Original: intermediate layers in MLP.
Revised: intermediate layers in MLP.

Original: n-head is the number of heads per layer.
Revised: n-head is the number of heads per layer.

Original: # layers is the number of layers.
Revised: # layers is the number of layers.

Original: For LSTMs, we use the implementation of PyTorch.
Revised: For LSTMs, we use the implementation of PyTorch.

Original: We consider d-hidden in [128,256,512,768,1024], and # layers in [4,6,8,12,16].
Revised: We consider d-hidden in [128,256,512,768,1024], and # layers in [4,6,8,12,16].

Original: A. A LGORITHMS
Revised: A. A LGORITHMS

Original: The lookup table.
Revised: The lookup table.

Original: To evaluate the downstream performance of the lookup table, we first create the lookup table with the data of the downstream task.
Revised: To evaluate the downstream performance of the lookup table, we ﬁrst create the lookup table with the data of the downstream task.

Original: With the method mentioned in Section A.2, we can generate the true conditional probability of each token and use it as the contextual embeddings.
Revised: With the method mentioned in Section A.2, we can generate the true conditional probability of each token and use it as the contextual embeddings.

Original: The adversarial algorithm.
Revised: The adversarial algorithm.

Original: The adversarial algorithm we use to mess up the downstream performance is maximizing a meta-learning objective in pre-training.
Revised: The adversarial algorithm we use to mess up the downstream performance is maximizing a meta-learning objective in pre-training.

Original: Suppose the linear head of the downstream task is g ϕ and the feature representation is h ψ .
Revised: Suppose the linear head of the downstream task is g φ and the feature representation is h ψ .

Original: The meta-learning algorithm first trains the head g ϕ to minimize the training loss of the downstream task, and then update h ψ to maximize the validation loss on the downstream tasks.
Revised: The meta-learning algorithm ﬁrst trains the head g φ to minimize the training loss of the downstream task, and then update h ψ to maximize the validation loss on the downstream tasks.

Original: Concretely, we randomly sample two disjoint subsets D 1 and D 2 from the downstream training dataset D .
Revised: Concretely, we randomly sample two disjoint subsets D 1 and D 2 from the downstream training dataset D .

Original: We train g ϕ to minimize the loss of downstream tasks on D 1 , ˆ ϕ ( ψ ) ∈ argmin 1 | D 1 | (cid:80) ( 
Revised: We train g φ to minimize the loss of downstream tasks on D 1 , ˆ φ ( ψ ) ∈ argmin 1 | D 1 | (cid:80) ( x,y ) ∈

Original: x,y ) ∈ D 1 ℓ ( g ϕ ( h ψ ( x )) ,y ) .
Revised: D 1 (cid:96) ( g φ ( h ψ ( x )) ,y ) .

Original: Then we train h ψ to maximize the validation loss on
Revised: Then we train h ψ to maximize the validation loss on

Original: D 2 during pre-training, minimizeψ L ( ψ ) − λ
Revised: D 2 during pre-training, minimize

Original: [Equation]
Revised: [Table]

Original: | D 2 | (cid:80) ( x,y ) ∈ D 2 ℓ ( g ˆ ϕ ( ψ ) ( h ψ ( x )) ,y ) .
Revised: h ψ ( x )) ,y ) .

Original: The optimization can be efficiently carried out with closed form solution of ϕ as shown in Liu et al.
Revised: The optimization can be efﬁciently carried out with closed form solution of φ as shown in Liu et al.

Original: Fine-tuning.
Revised: Fine-tuning.

Original: Following the standard protocol of Devlin et al.
Revised: Following the standard protocol of Devlin et al.

Original: (2018), we use the contextual embeddings of the CLS token for fine-tuning.
Revised: (2018), we use the contextual embeddings of the CLS token for ﬁne-tuning.

Original: We use AdamW with learning rate 1e-4.
Revised: We use AdamW with learning rate 1e-4.

Original: We perform 200 warmup steps and train on the downstream tasks for 10 epochs.
Revised: We perform 200 warmup steps and train on the downstream tasks for 10 epochs.

Original: Linear probe.
Revised: Linear probe.

Original: Since the CLS token is not trained in pre-training, we concatenate embeddings of all the tokens in the sentence as the representations.
Revised: Since the CLS token is not trained in pre-training, we concatenate embeddings of all the tokens in the sentence as the representations.

Original: We use AdamW with learning rate 1e-3 to train the linear head.
Revised: We use AdamW with learning rate 1e-3 to train the linear head.

Original: We train on the downstream tasks for 100 epochs.
Revised: We train on the downstream tasks for 100 epochs.

Original: Note that to make the capacity of the linear probe itself controlled, we adopt a random Gaussian projection to dimension 512 on the concatenation of the embeddings.
Revised: Note that to make the capacity of the linear probe itself controlled, we adopt a random Gaussian projection to dimension 512 on the concatenation of the embeddings.

Original: We report the standard deviation of linear probe and fine-tuning from 5 random seeds.
Revised: We report the standard deviation of linear probe and ﬁne-tuning from 5 random seeds.

Original: Evaluation of pre-training loss.
Revised: Evaluation of pre-training loss.

Original: Since we have access calculate the true conditional probability, we can calculate the cross entropy loss as the sum of the entropy of the true conditional probability and the
Revised: Since we have access calculate the true conditional probability, we can calculate the cross entropy loss as the sum of the entropy of the true conditional probability and the

Original: KL divergence between the predicted and true conditional probabilities.
Revised: KL divergence between the predicted and true conditional probabilities.

Original: This is more accurate than evaluating on the validation datasets in the standard ways.
Revised: This is more accurate than evaluating on the validation datasets in the standard ways.

Original: We report the number of pre-training loss with 10 6 sentences, and calculate the standard deviation on 5 subsets, each of which has size 2 × 10 5 .
Revised: We report the number of pre-training loss with 10 6 sentences, and calculate the standard deviation on 5 subsets, each of which has size 2 × 10 5 .

Original: A. R ESULTS ON O THER D OWNSTREAM T ASKS.
Revised: A. R ESULTS ON O THER D OWNSTREAM T ASKS

Original: We also provide results on other downstream tasks in this subsection.
Revised: We also provide results on other downstream tasks in this subsection.

Original: On PCFG Task A, OPT SSTand the Task-6 of HMM, we can also observe the increase in downstream performance as we scale up the models in the saturation regime.
Revised: On PCFG Task A, OPT SSTand the Task-6 of HMM, we can also observe the increase in downstream performance as we scale up the models in the saturation regime.

Original: 2 5 9 12 41 80 235 436 730 Model Size / M

Revised: 2 5 9 12 41 80 235 436 730 Model Size / M

Original: D o w n s t r e a m A cc u r a c
Revised: D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: PCFG, Task A
Revised: PCFG, Task A

Original: Fine-tuningLinear Probe 3 . 203 .. 243 .. 283 .. 323 ..
Revised: Fine-tuningLinear Probe 3 . 203 .. 243 .. 283 .. 323 ..

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre-training (a) PCFG → Task A      730 Model Size / M

Revised: Pre-training (a) PCFG → Task A      730 Model Size / M

Original: D o w n s t r e a m A cc u r a c
Revised: D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: OPT, SST
Revised: OPT, SST

Original: Fine-tuningLinear Probe
Revised: Fine-tuningLinear Probe

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre-training (b) OPT → SST135 Model Size / M
Revised: Pre-training (b) OPT → SST135 Model Size / M

Original: D o w n s t r e a m A cc u r a c
Revised: D o w n s t r e a m A cc u r a c

Original: y %
Revised: y %

Original: HMM, TaskFine-tuningLinear Probe 3 .3 .3 .3 .3 .3 .
Revised: HMM, TaskFine-tuningLinear Probe 3 .3 .3 .3 .3 .3 .

Original: P r e - t r a i n i n g L o ss
Revised: P r e - t r a i n i n g L o ss

Original: Pre-training (c) HMM →
Revised: Pre-training (c) HMM →

Original: Task-
Revised: Task-

Original: B D ETAILS IN S ECTION
Revised: B D ETAILS IN S ECTION

Original: UnbiasedestimateofthetraceofHessian.
Revised: Unbiased estimate of the trace of Hessian.

Revised: Evaluating the trace of Hessian requires the norm of the Jacobian ∇ θ log[ f θ ( x − t )] .

Original: Sincetheoutputdimension c andtheumberofparameters are all very large, computing the Jacobian ∇ θ  f θ ( x − t ) will be very inefficient.
Revised: Since the output dimension c and the number of parameters are all very large, computing the Jacobian ∇ θ log[ f θ ( x − t )] will be very inefﬁcient.

Original: We approximate the Jacobian with random samples as follows.
Revised: Instead, we can estimate the trace of Hessian unbiasedly with random samples as follows.

Original: Suppose f  ( x − t ) is the predicted logits of the conditional  probability, thetrace of Hessian  can be expressed as
Revised: Suppose f θ ( x − t ) is the predicted probability of the conditional probability. In the saturation regime, as f θ ( x − t ) approaches the true conditional probability, the Hessian of the pre-training loss w.r.t. the parameters can be expressed as

Original: [Equation]
Revised: [Equation]

Original: ∇ θ f θ ( x − t )]
Revised: Therefore we have

Original: ˆ y ∥ 22 ] .
Revised: [Equation]

Original: To approximate this expectation, we can  sample  i.i.d.
Revised: To approximate this expectation, we can ﬁrst sample t,x − t from the language, then draw i.i.d.

Original: samples ˆ y fromsoftmax ( f  ( x − t ))  and use the average as the unbiased estimate.
Revised: samples x t from ( f θ ( x − t )) , and use the average as the unbiased estimate.

Original: For all experiments, we use number of samples 50.
Revised: For all experiments, we samplex − t and sample 50 x t for each x − t .

Original: Details in Figure 3.
Revised: Details in Figure 3.

Original: To verify Theorem 3.3 that SGD biases the model towards flatter minima, we conduct MLM on PCFG and HMM-generated datasets with SGD.
Revised: To verify Theorem 3.3 that SGD biases the model towards ﬂatter minima, we conduct MLM on PCFG and HMM-generated datasets with SGD.

Original: We set the proportion of warmup stage to 12 % total number of steps, and fix the learning rate to 1e-3 after the warmup.
Revised: We set the proportion of warmup stage to 12 % total number of steps, and ﬁx the learning rate to 1e-3 after the warmup.

Original: We evaluate the downstream performance and the trace of Hessian of different checkpoints along pre-training.
Revised: We evaluate the downstream performance and the trace of Hessian of different checkpoints along pre-training.

Original: The standard deviation of trace of Hessian is calculated based on 5 times of sampling 50 examples as mentioned above.
Revised: The standard deviation of trace of Hessian is calculated based on 5 times of sampling 50 examples as mentioned above.

Original: Apart from the PCFG task C and HMM task-10, we also provide results on PCFG tasks A, B and HMM task-6 in Figure 10.
Revised: Apart from the PCFG task C and HMM task-10, we also provide results on PCFG tasks A, B and HMM task-6 in Figure 10.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] n s t r e a
Revised: B. E MBEDDING OF A S MALLER T RANSFORMER INTO A L ARGER T RANSFORMER .

Original: [Figure] A
Revised: In this subsection, we show that a smaller transformer can be embed into a larger transformer without changing the functionality.

Original: [Figure] u r a c y % [Figure]
Revised: We enable the embedding by considering two techniques (1) adding additional layers using residual connections without changing the functionality and (2) increasing feature dimension / adding more attention heads without change the functionality by duplicating the weights.

Original: P r [Figure] t r a i n i n g L o ss [Figure]
Revised: B.1. T HE B ASE C ASE WITH MLP S .

Original: Viewing a small model as a special case of large model.

Original: We show that a small neural network is a special case of a large neural network by considering two cases (1) adding additional layers without changing the functionality and (2) adding more channels/attention heads without change the functionality.
Revised: To gain some insights of how to increase the feature dimension without changing the functionality, we start with vanilla MLPs without layer-norm and residual connections He et al. Consider a multi-layer MLP f W,a ( x ) . The weight matrices are W = [ W 0 ,...,W L − 1 ] .The

Original: For (1), since the transformer contains residual connections, we can add new layers on top of a smaller model and fill in zeros to the added parameters.
Revised: dimensionality is W l ∈ R d l +1 × d l . The representations are deﬁned recursively, h l +1 ( x )

Revised: = σ ( W l h l ( x )) .

Revised: We denote the input by h 0 ( x )=

Revised: x and the ﬁnal output is deﬁned as f W,a ( x )= a (cid:62) h L ( x ) .

Revised: The activation σ here is relu or leaky relu.

Revised: We aim to embed f W,a ( x ) into f ˜ W, ˜ a ( x ) , where ˜ W =[ ˜

Revised: W 0 ,..., ˜ W L − 1 ] and

Revised: ˜ W l ∈ R 2 d l +1 × 2 d l .

Revised: We need to make sure f W,a ( x )=

Revised: f ˜ W, ˜ a ( x ) for all x .

Original: In this way, we can obtain a larger model with the same functionality as the small model.
Revised: For this case, we can set ˜ W l =1 /(cid:20) [Equation] (cid:21) for l ∈ [ L − 1] , ˜ W 0 = 1 √(cid:20)

Revised: WW (cid:21) , and ˜ a = 1 √ 2 (cid:20) aa (cid:21) .

Revised: We can verify that ˜ h l ( x )= 1 √(cid:20)

Revised: h l ( x ) h l ( x ) (cid:21) inductively.

Revised: Therefore we have

Revised: [Equation]

Revised: Next we turn to transformers with residual connections and layer norm.

Revised: We ﬁrst use the same strategy as the MLP case to add additional feature dimension and attention heads by replicating the weights, and then show how to add new layers using residual connections.

Revised: At a high level, replicating the weight maintains the mean and the variance calculated by the layernorm.

Revised: Therefore the representations inside the transformer also get replicated, without changing the values in each of the replicated groups.

Revised: Setup.

Revised: A transformer is composed of an input embedding W E , L blocks of self-attention, and an output layer.

Revised: Transformers also contain layer-norm and residual connections.

Revised: Suppose the input x = [ x 1 ,...,x T ] , where x i ∈ R d .

Revised: Each block of the self-attention contains of an attention layer and an MLP layer, both equipped with residual connections and layer-norm.

Revised: Let us denote by [ h 0 ( x )]

Revised: i =LN( W E x i ) the input embeddings.

Revised: Suppose the hidden size is d h , i.e. [ h l ( x )] i ∈ R d h .

Revised: The attention layer is deﬁned as [ v l ( x )]

Revised: i = LN

Revised: ([ h l ( x )] i +[ Attn l ( h l ( x ))] i ) , and the MLP layer is deﬁned as [ h l +1 ( x )] i = LN

Revised: ([ v l ( x )] i + U l σ ( W l [ v l ( x )]

Revised: i + b l )) .

Revised: The activation σ is GeLU.

Original: Concretely, suppose a small model is g θ , an additional block is  h ϕ , and the large model can be expressed as f θ,ϕ ( · )
Revised: The ﬁnal output is [ f ( x )] i = W (cid:62) E [ h L ( x )] i . Note that W E is both the input embedding and the weight of the output layer.

Revised: They are tied in training.

Revised: The layer norm is on the feature dimension.

Revised: [LN( x i )]

Revised: j = γ j ∗ ˆ x ij + β j .

Revised: ˆ x

Revised: i is the normalized version of x i with zero mean and unit variance.

Revised: γ and β are trainable.

Revised: The multi-head attention consists of n h self-attention heads.

Revised: The deﬁnition of the multi-head attention is Attn l ( h l ( x )) =

Revised: [( A l 1 h l ( x ) V l 1 ) (cid:62) ,..., ( A ln h h l ( x ) V ln h ) (cid:62) ] (cid:62) O l .

Revised: The output matrix O l ∈ R d h × d h .

Revised: The attention heads composes of the attention score times the feature matrix times the value matrix.

Revised: The attention score A lk ∈ R d (cid:48) × d (cid:48) is computed with softmax dot product.

Revised: For each k ∈ [ n h ] , A lk = softmax ( h l ( x ) Q lk K (cid:62) lk h l ( x ) (cid:62) ) .

Revised: Following the implementation of Devlin et al.

Revised: , the dimension of the attention head is always d (cid:48)

Revised: = 64 , thus d h = 64 n h .

Revised: The dimension of the intermediate layer in the MLP is set to 4 d h , which means U l ∈ R d h × 4 d h and W l ∈ R 4 d h × d h .

Revised: We aim to embed the smaller transformer f

Revised: ( x ) into ˜ f ( x ) , where ˜ d h =2 d h , ˜ n h =2 n h , and ˜ L =

Revised: L + L (cid:48) .

Revised: Increasing feature dimension with replication of the parameters.

Revised: Although the transformers have layer-norm and residual connections, we can still modify the strategy in the base case with MLPs slightly to increase the width of the model and the number of attention heads without changing the functionality.

Revised: Consider the following weight replication method.

Revised: For l ∈ [0 ,...,L − 1] ,

Revised: [Equation]

Revised: We observe that the intermediate layers of the transformers are also replicated for the ﬁrst L blocks, i.e. ˜ h l ( x )

Original: =
Revised: =

Revised: (cid:20) h l ( x ) h l ( x ) (cid:21) and ˜ v l ( x )=

Revised: (cid:20) v l ( x ) v l ( x ) (cid:21) for l ∈ [1 ,...,L ] .

Revised: First note that [ h 0 ( x )]

Revised: i =LN( W E x i ) .

Revised: Since replicating the features will not change the mean and the variance, we have ˜ h 0 ( x )=

Original:  h ϕ{+
+} ( g θ ( · ))
Revised: (cid:20) h {+
+} ( x ) h 0 ( x ) (cid:21) .

Revised: Then we can show that replicating the features will not change the attention scores as well.

Original: + g θ ( · )  . 
Revised: This makes ˜ v l ( x ) = (cid:20) v l ( x ) v l ( x ) (cid:21) . Finally note that we can apply the base case of the MLP to reason about the MLP layer, and show ˜ h l +1 ( x ) =

Revised: (cid:20) h l

Original: Then setting the parameters ϕ to all 0 can make f θ,ϕ{++
+} ( · )=
Revised: {++
+} ( x ) h l +1 ( x ) (cid:21) . Therefore we have shown ˜ h l ( x ) = (cid:20) h l ( x ) h l ( x ) (cid:21) inductively.

Original: g θ ( · )  . 
Revised: Adding additional layers using residual connections. We have demonstrated that ˜ h l ( x )= (cid:20) h l ( x ) h l ( x ) (cid:21) for l ∈ [0 ,...,L ] . Now let’s consider the added L (cid:48) blocks on top of the small model.

Original: For (2), since the transformers contain layer norm, we can duplicate the parameters channelwise toduplicate the feature dimension.
Revised: Since the transformer contains residual connections, we can add new blocks on top of a small model and ﬁll in zeros to the added parameters. We will show that in this way, ˜ h l ( x )=˜ h L ( x ) , for any l ∈ [ L,...,L + L (cid:48) ] .

Revised: This will indicate that ˜ h

Revised: L + L (cid:48)

Original: Since we keep the mean and variance across channels unchanged, the functionality will also be preserved.
Revised: ( x )= ˜ h L ( x )= (cid:20) h l ( x ) h l ( x ) (cid:21) .

Revised: Recall that ˜ W E =(cid:20) W E

Revised: W E (cid:21) .

Revised: This indicate that

Revised: [Equation]

Revised: [ f ( x )]

Revised: i =[ f

Revised: ( x )] i , which means we can add new layers on top of a small transformer without changing the functionality.

Revised: Now we show that U l = 0 and O l = 0 for l ∈ [ L,...

Revised: L + L (cid:48) − 1] will make ˜ h l ( x ) =

Revised: ˜

Revised: h L ( x ) , for any l ∈ [ L,...,L + L (cid:48) ] .

Revised: This holds because [ v l ( x )]

Revised: i = LN

Revised: ([ h l ( x )] i + [ Attn l ( h l ( x ))] i ) and [ h l +1 ( x )] i =

Revised: LN([ v l ( x )]

Original: Similarly, we can also add more heads with this duplication technique.
Revised: i + U l σ ( W l [ v l ( x )] i + b l )) . If O l = 0 , we have v l ( x ) = h l ( x ) from the ﬁrst equation.

Original: Therefore, the small models can be enlarged into a larger models with functionality preserved.
Revised: If U l =0 , we have h l +1 ( x )= v l ( x ) from the second equation. Therefore we have ˜ h l ( x )=

Revised: ˜

Original:  The smaller transformer architecture can  be viewed as a subset of the larger transformer architecture.
Revised: h L ( x ) , for any l ∈ [ L,...,L + L (cid:48) ] . B.1. V IEWING A S MALL T RANSFORMER AS A S PECIAL C ASE OF A L ARGE T RANSFORMER As demonstrated above, smaller transformers can be embedded into larger transformers with functionality preserved. The smaller transformer architecture can therefore be viewed as a subset of the larger transformer architecture.

Original: In this sense, a set of transformers with different sizes and the same pre-training loss found in Section 2 can be viewed as a set of transformers with the same size.
Revised: In this sense, a set of transformers with different sizes and the same pre-training loss found in Section 2 can be viewed as a set of transformers with the same size after the embedding.

Original: Note that the training algorithm only finds out the natural larger models, instead of the larger models which are augmented from the smaller models.
Revised: Note that the training algorithm only ﬁnds out the natural larger models, instead of the larger models which are embedded from the smaller models.

Original: This indicates that the implicit bias of the optimizer can interact with the model architecture.
Revised: This indicates that the implicit bias of the optimizer can interact with the model architecture.

Original: The implicit bias drives the model toward flat minima on both larger models and smaller models.
Revised: The implicit bias drives the model toward ﬂat minima on both larger models and smaller models.

Original: The smaller transformer architecture is a subset of the larger transformer architecture, thus the flattest minima found with a larger transformer is flatter than the minima found with a smaller transformer.
Revised: The smaller transformer architecture is a subset of the larger transformer architecture, thus the ﬂattest minima found with a larger transformer is ﬂatter than the minima found with a smaller transformer.

Original: (See Figure 6).
Revised: (See Figure 6).

Original: C O MITTED P ROOFS IN S ECTION
Revised: C O MITTED P ROOFS IN S ECTION

Original: Proof of Lemma 3.2.
Revised: Proof of Lemma 3.2.

Original: We first recall loss
Revised: We ﬁrst recall loss

Original: [Equation]
Revised: [Equation]

Original: Note that conditioned on any x − t ,t , it holds that
Revised: Note that conditioned on any x − t ,t , it holds that

Original: [Equation]
Revised: [Equation]

Original: where in the last step, we use the assumption that θ ∈ Γ
Revised: where in the last step, we use the assumption that θ ∈ Γ

Original: [Equation]
Revised: [Equation]

Original: that is, for all x,t
Revised: that is

Original: [Equation]
Revised: [Equation]

Original:  f θ ( x − t )=
Revised: for all x,t , f θ ( x − t )=

Original: Pr( ·| x − t )
Revised: Pr( ·| x − t )

Original: [Equation]
Revised: [Equation]

Original: which implies the following
Revised: which implies the following

Original: [Equation]
Revised: [Equation]

Original: Since θ is a global minimizer of L
Revised: Since θ is a global minimizer of L

Original: [Equation]
Revised: [Equation]

Original: we have that ∇ L ( θ )=
Revised: we have that ∇ L ( θ )=

Original: E t,x ∇ θ log[ [Equation] x − t )] x t =0 .
Revised: E t,x ∇ θ log[ f θ ( x − t )] x t =0 .

Original: Therefore, we have that
Revised: Therefore, we have that

Original: [Equation]
Revised: [Equation]

Original: which completes the proof
Revised: which completes the proof

Original: [Equation]
Revised: [Equation]

Revised: D P RACTICAL I MPLICATIONS

Revised: Pre-training Algorithms.

Revised: While we focus on the saturation regime in the paper as a controlled way to compare models with the same pre-training loss, the overall takeaway is that the implicit bias of the training algorithms matters for downstream performance (no matter whether we are in the saturation regime or not).

Revised: Moreover, understanding the implicit biases needed for downstream performance may also lead to better training methods (instead of better evaluation methods) that might encourage the correct biases more strongly.

Revised: Therefore, a practical direction is to design better pre-training algorithms with more favorable biases which can lead to better downstream performance than AdamW and SGD.

Revised: Better Metrics for Language Models.

Revised: In common practice, the validation pre-training loss is used to monitor the training process (Brown et al., 2020; Zhang et al., 2022a) and compare different models (Hernandez et al., 2021).

Revised: However, Saunshi et al.

Revised: ; Tay et al.

Revised: (2021) show that pretraining loss is not necessarily correlated with downstream performance when comparing different architectures.

Revised: We further show that pre-training loss may not always be a reliable indicator even for the same architecture.

Revised: While downstream tasks could be used as a proxy metric for evaluation, the main issue is that large language models are trained to be general / multi-purpose models where the space of downstream tasks is large and unknown during the time of pre-training.

Revised: Thus from a fundamental standpoint, it is beneﬁcial to design a more reliable indicator that is agnostic to downstream tasks.

Revised: Explicit regularization.

Revised: We show that implicit bias, especially the implicit bias of ﬂatness matters for downstream performance in language modeling.

Revised: Leveraging the implicit bias to design better explicit regularization in language modeling is also an important direction.

Revised: Bahri et al.

Revised: (2021) show explicit ﬂatness regularization with SAM (Foret et al., 2020) can boost downstream performance when applying to downstream tasks themselves and the intermediate stages between pre-training and ﬁne-tuning, but they did not study this on pre-training, partly because SAM is not efﬁcient enough for pre-training (SAM requires back prop for 2 times per step, and more steps to reach the same level of pre-

Revised: [Table]

Revised: E R ESULTS OF 15 % M ASK R ATE

Revised: As we mentioned in Section 2 and Section A.2, the true conditional probability of MLM Pr( x t | x − t ) can be computed from the joint probability

Revised: [Equation]

Revised: and the joint probability can be computed efﬁciently with knowledge of the generative algorithms.

Revised: The number of computation of the joint probability is the size of W .

Revised: However, when we have multiple masks, the time complexity of computing the true conditional probability can be extremely large.

Revised: For example, if we have two masks in one sentence,

Revised: [Equation]

Revised: Now to compute the denominator, we will need to compute the joint probability for | W | 2 times.

Revised: In general, the time complexity of computing the true MLM conditional probability is exponential in the number of masks.

Revised: Still, we can pre-train the model with 15% mask rate and evaluate the loss with one mask per sentence.

Revised: The results on PCFG is provided in Figure 11.

Revised: Although in pre-training we use 15% mask rate, the validation loss evaluated with one mask per sentence does not change much compared with training with one mask per sentence.

Revised: The downstream performance does not change signiﬁcantly with 15% mask rate either.

Revised: The conclusion of Section 2 and Section 4 still holds.

Revised: [Figure]

Revised: D

Revised: [Figure]

Revised: n

Revised: [Figure]

Revised: t r

Original: D O MITTED P ROOFS IN S ECTIOND. O MITTED P ROOFS OF T HEOREM 5. 
Revised: [Figure] cc [Figure] F O MITTED P ROOFS IN S ECTIONF. O MITTED P ROOFS OF T HEOREM 5.

Original: Recall that the loss function of MLM is L ( ψ, u ) =
Revised: Recall that the loss function of MLM is L ( ψ, u ) =

Original: E x ∼ P [( f ψ,u ( x ) − g ∗ ( x )) 2 ] .
Revised: E x ∼ P [( f ψ,u ( x ) − g ∗ ( x )) 2 ] .

Original: In downstream adaptation, we have access to a finite dataset { x ( i ) } ni =1 sampled i.i.d.
Revised: In downstream adaptation, we have access to a ﬁnite dataset { x ( i ) } ni =1 sampled i.i.d.

Original: from P ds .
Revised: from P ds .

Original: The training loss is (cid:98) L P ds ( ψ,u ) = 1 n (cid:80) ni =

Revised: The training loss is (cid:98) L P ds ( ψ,u ) = 1 n (cid:80) ni =

Original: [( f ψ,u ( x ( i ) ) − g ∗ ( x ( i ) )) 2 ] , and the population loss for the downstream task is
Revised: [( f ψ,u ( x ( i ) ) − g ∗ ( x ( i ) )) 2 ] , and the population loss for the downstream task is

Original: L P ds ( ψ,u )=
Revised: L P ds ( ψ,u )=

Original: E x ∼ P ds
Revised: E x ∼ P ds

Original: [( f ψ,u ( x ) − g ∗ ( x )) 2 ] .
Revised: [( f ψ,u ( x ) − g ∗ ( x )) 2 ] .

Original: Proof of Theorem 5.1.
Revised: Proof of Theorem 5.1.

Original: We first calculate the trace of Hessian of the pre-training loss and then derive a lower bound for it in Lemma D.1.
Revised: We ﬁrst calculate the trace of Hessian of the pre-training loss and then derive a lower bound for it in Lemma F.1.

Original: We then show that the lower bound can be achieved only if the output of the attention are in one direction for all the downstream input in Lemma D.3.
Revised: We then show that the lower bound can be achieved only if the output of the attention are in one direction for all the downstream input in Lemma F.3.

Original: This translates to constant sample complexity for the downstream task.
Revised: This translates to constant sample complexity for the downstream task.

Original: Lemma D.1.
Revised: Lemma F.1.

Original: Denote by h Q,K ( x ) = (cid:80)
Revised: Denote by h Q,K ( x ) = (cid:80)

Original: Tj =1 a j x j the output of the attention head.
Revised: Tj =1 a j x j the output of the attention head.

Original: In the setting of
Revised: In the setting of

Original: Theorem 5.1, I + = { i ∈ [ m ] | u i > 0 } and I − =
Revised: Theorem 5.1, I + = { i ∈ [ m ] | u i > 0 } and I − =

Original: { i ∈ [ m ] | u i < 0 } .
Revised: { i ∈ [ m ] | u i < 0 } .

Original: The trace of Hessian can be lower bounded,
Revised: The trace of Hessian can be lower bounded,

Original: [Equation]
Revised: [Equation]

Original: where the lower bound is achieved if and only if the following conditions are satisfied,
Revised: where the lower bound is achieved if and only if the following conditions are satisﬁed,

Original: [Equation]
Revised: [Equation]

Original: Denote by D + = { x | y = 1 } and D − = { x | y = − 1 } .
Revised: Denote by D + = { x | y = 1 } and D − = { x | y = − 1 } .

Original: I x is the set of index of neurons which is activated on x , I x
Revised: I x is the set of index of neurons which is activated on x , I x

Original: = { i ∈ [ T ] | V ⊤ i x > 0 } .
Revised: = { i ∈ [ T ] | V (cid:62) i x > 0 } .

Original: By the condition in equation 5, the attention is taking the average of [ x t ]
Revised: By the condition in equation 5, the attention is taking the average of [ x t ]

Original: Tt =1 .
Revised: Tt =1 .

Original: We can map D + and D − to the feature space, H + = { h Q,K ( x ) | y
Revised: We can map D + and D − to the feature space, H + = { h Q,K ( x ) | y

Original: = 1 } and
Revised: = 1 } and

Original: H − = { h Q,K ( x ) | y = − 1 } .
Revised: H − = { h Q,K ( x ) | y = − 1 } .

Original: We first show that one neuron cannot be activated on inputs from both D + and D − , and all non zero neuron has to be activated on some input.
Revised: We ﬁrst show that one neuron cannot be activated on inputs from both D + and D − , and all non zero neuron has to be activated on some input.

Original: Also note that a neuron cannot be activated on no input, unless the weight is 0 .
Revised: Also note that a neuron cannot be activated on no input, unless the weight is 0 .

Original: Fact D.2.
Revised: Fact F.2.

Original: (1) ∀ x ∈ D
Revised: (1) ∀ x ∈ D

Original: + , I x ⊆
Revised: + , I x ⊆

Original: I + .
Revised: I + .

Original: Similarly we have ∀ x ∈ D − , I x ⊆
Revised: Similarly we have ∀ x ∈ D − , I x ⊆

Original: I − .
Revised: I − .

Original: (2) Suppose V i ̸ =0 , then there exists h ∈ H + ∪ H − , V ⊤ i h> 0 .
Revised: (2) Suppose V i (cid:54) =0 , then there exists h ∈ H + ∪ H − , V (cid:62) i h> 0 .

Original: Proof of Fact D.2.
Revised: Proof of Fact F.2.

Original: (1) Otherwise, suppose j ∈
Revised: (1) Otherwise, suppose j ∈

Original: I x
Revised: I x

Original: ∩ I − , since y = 1 m (cid:80) mi =1 u i σ ( V ⊤ i h Q,K ( x )) > 0 , there has to be j ′{+(cid:
 )+} ∈
Revised: ∩ I − , since y = 1 m (cid:80) mi =1 u i σ ( V (cid:62) i h Q,K ( x )) > 0 , there has to be j {+(cid:
 )+} ∈

Original: I x ∩ I + , which contradicts the condition in equation 4.
Revised: I x ∩ I + , which contradicts the condition in equation 4.

Original: (2) Suppose v ⊤ h ≤ 0 for all h ∈ H + ∪ H − .
Revised: (2) Suppose v (cid:62) h ≤ 0 for all h ∈ H + ∪ H − .

Original: Then we have v ⊤ h =0 for all h , since v ⊤ h< 0 indicates v ⊤ ( − h ) > 0 , and − h belongs to the support of P due to the symmetry of the distribution.
Revised: Then we have v (cid:62) h =0 for all h , since v (cid:62) h< 0 indicates v (cid:62) ( − h ) > 0 , and − h belongs to the support of P due to the symmetry of the distribution.

Original: However, in Lemma D.5, we show that the matrix stacking all input together has full row rank, thus v has to be 0 , leading to a contradiction.
Revised: However, in Lemma F.5, we show that the matrix stacking all input together has full row rank, thus v has to be 0 , leading to a contradiction.

Original: We have the following lemma characterizing the solutions achieving all the qualities in Lemma D.1.
Revised: We have the following lemma characterizing the solutions achieving all the qualities in Lemma F.1.

Original: Intuitively, all the neurons can be divided into two sets, and each input can only activate neurons in one of the sets, leading to no cancellation between activated neurons.
Revised: Intuitively, all the neurons can be divided into two sets, and each input can only activate neurons in one of the sets, leading to no cancellation between activated neurons.

Original: This holds because of equation 4 and the properties of the input distribution.
Revised: This holds because of equation 4 and the properties of the input distribution.

Original: Lemma D.3.
Revised: Lemma F.3.

Original: Suppose Q,K,V satisfy the equality in Lemma D.1.
Revised: Suppose Q,K,V satisfy the equality in Lemma F.1.

Original: For all i ∈ I − , on downstream data x , if g ∗ ( x )=1 , we have V ⊤ i h Q,K ( x )=
Revised: For all i ∈ I − , on downstream data x , if g ∗ ( x ) = 1 , we have V (cid:62) i h Q,K ( x ) =

Original: c i > 0 .
Revised: c i > 0 .

Original: c i is a constant which holds for every x if g ∗ ( x )=1 .
Revised: c i is a constant which holds for every x if g ∗ ( x ) = 1 .

Original: If g ∗ ( x )= − 1 , we have V ⊤ i h Q,K ( x )=0 .
Revised: If g ∗ ( x )= − 1 , we have V (cid:62) i h Q,K ( x )=0 .

Original: For all i ∈ I + , on downstream data x , if g ∗ ( x )= − 1 , we have V ⊤ i h Q,K ( x )=
Revised: For all i ∈ I + , on downstream data x , if g ∗ ( x )= − 1 , we have V (cid:62) i h Q,K ( x )=

Original: c i > 0 .
Revised: c i > 0 .

Original: c i is a constant which holds for every x if g ∗ ( x )=1 .
Revised: c i is a constant which holds for every x if g ∗ ( x )=1 .

Original: If g ∗ ( x )=1 , we have V ⊤ i h Q,K ( x )=0 .
Revised: If g ∗ ( x )=1 , we have V (cid:62) i h Q,K ( x )=0 .

Original: Now let us consider the downstream task.
Revised: Now let us consider the downstream task.

Original: It suffices to consider the constant vector c .
Revised: It sufﬁces to consider the constant vector c .

Original: If samples satisfying g ∗ ( x ) = 

Revised: If samples satisfying g ∗ ( x ) =

Original: and g ∗ ( x ) = − 1 both show up in the downstream dataset, the minimal norm solution ˜ u is ˜ u
Revised: and g ∗ ( x ) = − 1 both show up in the downstream dataset, the minimal norm solution ˜ u is ˜ u

Original: I − =
Revised: I − =

Original: mc I − ∥
Revised: mc I − (cid:107) c I − (cid:107) 2 2 , ˜ u I + = − mc I + (cid:107)

Original: c
Revised: c

Original: I− ∥ 2 2 , ˜ u I + = − mc I + ∥ c I + ∥ 2 2 , and ˜ u [ m ] \ ( I + ∪ I − ) =0 .
Revised: I + (cid:107) 2 2 , and ˜ u [ m ] \ ( I + ∪ I − ) =0 .

Original: Then we can verify that
Revised: Then we can verify that

Original: [Equation]
Revised: [Equation]

Original: = [Equation] [ g
Revised: g

Original: [Equation]
Revised: [Equation]

Original: ( x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: =1] −

Original: [Equation]
Revised: [Equation]

Original: [ g
Revised: [ g

Original: [Equation]
Revised: [Equation]

Original: ( x
Revised: ( x

Original: [Equation]
Revised: [Equation]

Revised: −

Revised: [Equation]

Original: = g ∗ ( x )
Revised: = g ∗ ( x )

Original: [Equation]
Revised: [Equation]

Original: Therefore, L P ds ( ˆ ψ, ˜ u )=
Revised: Therefore, L P ds ( ˆ ψ, ˜ u )=

Original: E x ∼ P ds
Revised: E x ∼ P ds

Original: [( f ˆ ψ, ˜ u ( x ) − g ∗ ( x )) 2 ]=0 , which completes the proof.
Revised: [( f ˆ ψ, ˜ u ( x ) − g ∗ ( x )) 2 ]=0 , which completes the proof.

Original: We first show that when the pre-training loss equals 0 , the trace of Hessian equals the square of the norm of the gradient.
Revised: We ﬁrst show that when the pre-training loss equals 0 , the trace of Hessian equals the square of the norm of the gradient.

Original: Lemma D.4.
Revised: Lemma F.4.

Original: For any parameters θ , if the pre-training loss L ( θ )=
Revised: For any parameters θ , if the pre-training loss L ( θ )=

Original: E x [( f θ ( x ) − y ) 2 ]=0 , the trace of
Revised: E x [( f θ ( x ) − y ) 2 ]=0 , the trace of

Original: Hessian equals the square of the norm of the gradient,
Revised: Hessian equals the square of the norm of the gradient,

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma D.4.
Revised: Proof of Lemma F.4.

Original: We can express the Hessian as follows.
Revised: We can express the Hessian as follows.

Original: [Equation]
Revised: [Equation]

Original: Since L ( θ )=
Revised: Since L ( θ )=

Original: E x [( f θ ( x ) − y ) 2 ]=0 , we have with probability 1 , ℓ ′ ( f θ ( x ) ,y ){+)=
+}
Revised: E x [( f θ ( x ) − y ) 2 ]=0 , we have with probability 1 , (cid:96) (cid:48) ( f θ ( x ) ,y {+)=
+}

Revised: and (cid:96) (cid:48)(cid:48)

Original: =0 and ℓ ′′ ( f θ ( x ) ,y )
Revised: ( f θ ( x ) ,y )

Original: =2 is a constant.
Revised: =2 is a constant.

Original: Proof of Lemma D.1.
Revised: Proof of Lemma F.1.

Original: [Equation] [Equation] 
Revised: [Equation] [Equation] [Equation]

Original: The equality in step 6 is achieved if and only if the gradient of Q and K is 0 .
Revised: The equality in step 6 is achieved if and only if the gradient of Q and K is 0 .

Original: Equation (7) is from
Revised: Equation (7) is from

Original: AM-GM [Equation] and the equality is achieved iff
Revised: AM-GM, and the equality is achieved iff

Original: [Equation]
Revised: [Equation]

Original: The equality in step 8 is achieved iff on all input, there is no cancellation between activated neurons,
Revised: The equality in step 8 is achieved iff on all input, there is no cancellation between activated neurons,

Original: [Equation]
Revised: [Equation]

Original: Since the attention score a j satisfies a j > 0 and (cid:80) Tj =1 a j =1 , and all embeddings x t in one masked sentence are orthogonal to each other with norm 1 , we have ∥ h Q,K ( x ) ∥ 2 ≥ 1 √ T .
Revised: Since the attention score a j satisﬁes a j > 0 and (cid:80) Tj =1 a j =1 , and all embeddings x t in one masked sentence are orthogonal to each other with norm 1 , we have (cid:107) h Q,K ( x ) (cid:107) 2 ≥ 1 √ T .

Original: The equality is achieved iff a j = 1 T for all x and all j ∈ [ T ] .
Revised: The equality is achieved iff a j = 1 T for all x and all j ∈ [ T ] .

Original: Proof of Lemma D.3 .
Revised: Proof of Lemma F.3 .

Original: Suppose V i is a neuron with i ∈ I − .
Revised: Suppose V i is a neuron with i ∈ I − .

Original: Then there exists h ∈ H − , V ⊤ i h > 0 .
Revised: Then there exists h ∈ H − , V (cid:62) i h > 0 .

Original: Without loss of generality, suppose the masked position in h is 1 , i.e. h 1 = 0 , h 2 = 1 .
Revised: Without loss of generality, suppose the masked position in h is 1 , i.e. h 1 = 0 , h 2 = 1 .

Original: Now let us consider the components in V i corresponding to the input positions and the mask positions separately.
Revised: Now let us consider the components in V i corresponding to the input positions and the mask positions separately.

Original: [Equation]
Revised: [Equation]

Original: We claim that V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 is either 0 or c .
Revised: We claim that V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 is either 0 or c .

Original: To prove this, consider ˜ h , which is only different from h on the mask, h − ˜ h
Revised: To prove this, consider ˜ h , which is only different from h on the mask, h − ˜ h

Original: =2 e 2 .
Revised: =2 e 2 .

Original: Also consider − h and − ˜ h .
Revised: Also consider − h and − ˜ h .

Original: Due to the symmetry of the distribution, − h and − ˜ h are in H + .
Revised: Due to the symmetry of the distribution, − h and − ˜ h are in H + .

Original: By Fact D.2, V ⊤ i ( − h ) ≤ 0 and V ⊤ i ( − ˜ h ) ≤ 0 .
Revised: By Fact F.2, V (cid:62) i ( − h ) ≤ 0 and V (cid:62) i ( − ˜ h ) ≤ 0 .

Original: V ⊤ i ( − h ) cannot be 0 , because this will leads to V ⊤ i h
Revised: V (cid:62) i ( − h ) cannot be 0 , because this will leads to V (cid:62) i h

Original: =0 .
Revised: =0 .

Original: Case 1.
Revised: Case 1.

Original: V ⊤ i ( − h ) < 0 and V ⊤ i ( − ˜ h ) < 0 .
Revised: V (cid:62) i ( − h ) < 0 and V (cid:62) i ( − ˜ h ) < 0 .

Original: We have V ⊤ i h
Revised: We have V (cid:62) i h

Original: = V ⊤ i ˜ h > 0 , indicating that V ⊤ i ( h − ˜ h )
Revised: = V (cid:62) i ˜ h > 0 , indicating that V (cid:62) i ( h − ˜ h )

Original: = 2 V ( p ) i 1 =0 .
Revised: = 2 V ( p ) i 1 =0 .

Original: Now we show that V ( c ) i 2: T = c 1 .
Revised: Now we show that V ( c ) i 2: T = c 1 .

Original: Due to the condition of equation 3, we know that for any h ∈ H − masked on the first token, either V ⊤ i h
Revised: Due to the condition of equation 3, we know that for any h ∈ H − masked on the ﬁrst token, either V (cid:62) i h

Original: =0 or they equal to the same positive value c for all h .
Revised: =0 or they equal to the same positive value c for all h .

Original: We claim that V ⊤ i 
Revised: We claim that V (cid:62) i h

Original: h = c for any h ∈ H − .
Revised: = c for any h ∈ H − .

Original: Otherwise there exist H − , 1 and H − , 0 , H − , 0 ∩ H − , 1 = ∅ and
Revised: Otherwise there exist H − , 1 and H − , 0 , H − , 0 ∩ H − , 1 = ∅ and

Original: H − , 0 ∪ H − , 1 = H − ∩{ h | h 2 = ± 1 } .
Revised: H − , 0 ∪ H − , 1 = H − ∩{ h | h 2 = ± 1 } .

Original: V ⊤ i h
Revised: V (cid:62) i h

Original: = c for any h ∈ H − , 1 and V ⊤ i h =0 for any h ∈ H − , 0 .
Revised: = c for any h ∈ H − , 1 and V (cid:62) i h =0 for any h ∈ H − , 0 .

Original: This cannot happen for T ≥ 6 .
Revised: This cannot happen for T ≥ 6 .

Original: By Lemma [-D.
-] we know that the matrix stacking all such h 2: T together has full row rank, thus V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 =0 .
Revised: By Lemma [-D.
-]F.5 we know that the matrix stacking all such h 2: T together has full row rank, thus V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i [Equation]

Original: Case 2.
Revised: Case 2.

Original: V ⊤ i ( − h ) < 0 and V ⊤ i ( − ˜ h )
Revised: V (cid:62) i ( − h ) < 0 and V (cid:62) i ( − ˜ h )

Original: = 0 , which indicates V ⊤ i
Revised: = 0 , which indicates V (cid:62) i

Original: h
Revised: h

Original: = V ( p ) i 1 .
Revised: = V ( p ) i 1 .

Original: Consider another h ′ ∈
Revised: Consider another h (cid:48)

Original:  H − which is not equal to h and h ′ 2 =1 .
Revised: ∈ H − which is not equal to h and h (cid:48) 2 =1 .

Original: Similarly we can find ˜ h ′ , h ′ − ˜ h 
Revised: Similarly we can ﬁnd ˜ h (cid:48) , h (cid:48) − ˜ h (cid:48)

Original: ′ =2 e 2 .
Revised: =2 e 2 .

Original: Still we have V ⊤ i ( − h ′ ) <
Revised: Still we have V (cid:62) i ( − h (cid:48) ) <

Original: and V ⊤ i ( − ˜ h ′ )
Revised: and V (cid:62) i ( − ˜ h (cid:48) )

Original: = 0 , this tells us V ⊤ i h ′ =
Revised: = 0 , this tells us V (cid:62) i h (cid:48) =

Original: V ⊤ i h , due to the condition of equation 3.
Revised: V (cid:62) i h , due to the condition of equation 3.

Original: Applying this to different h ′ s, we have that V ⊤
Revised: Applying this to different h (cid:48) s, we have that V (cid:62)

Original: i h equals the same positive value for all h ∈ H − and the masked position is 0 .
Revised: i h equals the same positive value for all h ∈ H − and the masked position is 0 .

Original: By Lemma [-D.
-] we know that the matrix stacking all such h 2: T together has full row rank, thus
Revised: By Lemma [-D.
-]F.5 we know that the matrix stacking all such h 2: T together has full row rank, thus

Original: V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 = c .
Revised: V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 = c .

Original: We have proved that V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 is either 0 or c .
Revised: We have proved that V ( c ) i 2: T = c 1 for some c> 0 and V ( p ) i 1 is either 0 or c .

Original: We continue to show that
Revised: We continue to show that

Original: V ( c ) i = c 1 for the same c> 0 and either V ( p ) i is 0 or its coordinates is ± c .
Revised: V ( c ) i = c 1 for the same c> 0 and either V ( p ) i is 0 or its coordinates is ± c .

Original: For case 1, consider h ′′ ∈ H − whose masked position is 2 , h ′′ 4 = 1 .
Revised: For case 1, consider h (cid:48)(cid:48) ∈ H − whose masked position is 2 , h (cid:48)(cid:48) 4 =1 .

Original: Also suppose that h ′′ 1 = 1 By{+(cid:48)(cid:48) 
+}
Revised: Also suppose that h {+(cid:48)(cid:48) 
+}

Original:  Fact D.2, we know that V ⊤ i ( − h ′′ ) ≤ 0 and V ⊤ i ( − ˜ h ′′ ) ≤ 0 , this implies that − V ( c ) i 1 ≤ V ( p ) i
Revised: =1 By Fact F.2, we know that V (cid:62) i ( − h (cid:48)(cid:48) ) ≤ 0 and V (cid:62) i ( − ˜ h (cid:48)(cid:48) ) ≤ 0 , this implies that − V ( c ) i 1 ≤ V ( p ) i

Original: 2 ≤ V ( c ) i 1 .
Revised: 2 ≤ V ( c ) i 1 .

Original: Applying the same argument above, we know that either V ( p ) i 

Revised: Applying the same argument above, we know that either V ( p ) i

Original: = 0 or V ( p ) i 2 = ± V ( c )
Revised: =0 or V ( p ) i 2 = ± V ( c )

Original: i 1 , otherwise both V ⊤ i h ′′ > 0 and V ⊤ i ˜ h ′′ > 0 hold, and V ⊤ i
Revised: i 1 , otherwise both V (cid:62) i h (cid:48)(cid:48) >and V (cid:62) i ˜ h (cid:48)(cid:48) > 0 hold, and V (cid:62)

Original:  h ′′ ̸ = V ⊤ i ˜ h ′′ , contradicting condition in equation 3.
Revised: i h (cid:48)(cid:48) (cid:54) = V (cid:62) i ˜ h (cid:48)(cid:48) , contradicting condition in equation 3.

Original: If V ( p ) i 2 = V ( c ) i 1 , from equation 3 we know V ⊤ i h ′′ =
Revised: If V ( p ) i 2 = V ( c ) i 1 , from equation 3 we know V (cid:62) i h (cid:48)(cid:48)

Original:  V ⊤ i h , which indicates V ( p ) i [-
-]
Revised: = V (cid:62) i h , which indicates V ( p ) i [-
-]2 =

Original: = V ( c ) i 1 =
Revised: V ( c ) i 1 =

Original: c 2 .
Revised: c 2 .

Original: In this case we can find another h ′′′ whose masked position is 2 , h ′′′ 4 [-=
-] but h ′′′ 1 = − 1 .
Revised: In this case we can ﬁnd another h (cid:48)(cid:48)(cid:48) whose masked position is 2 , h (cid:48)(cid:48)(cid:48) 4 [-=
-]= 1 but h (cid:48)(cid:48)(cid:48) 1 = − 1 .

Original: Then V ⊤ i 
Revised: Then V (cid:62) i h (cid:48)(cid:48)(cid:48) (cid:54) =

Original: h ′′′ ̸ = V ⊤ i h , contradicting equation 3.
Revised: V (cid:62) i h , contradicting equation 3.

Original: Thus V ( p ) i 2 ̸ = V ( c ) i 1 .
Revised: Thus V ( p ) i 2 (cid:54) = V ( c ) i 1 .

Original: Similarly V ( p ) i 2 ̸ = − V ( c ) i 1 .
Revised: Similarly V ( p ) i 2 (cid:54) = − V ( c ) i 1 .

Original: The only possible situation is V ( p ) i 2 = 0 .
Revised: The only possible situation is V ( p ) i 2 =0 .

Original: Applying the argument in this paragraph to other masked position, we have V ( p ) i = 0 .
Revised: Applying the argument in this paragraph to other masked position, we have V ( p ) i = 0 .

Original: Since H − is invariant under permutation, V ( c ) i 

Revised: Since H − is invariant under permutation, V ( c ) i

Original: = c , and V ( c ) i =
Revised: = c , and V ( c ) i =

Original: c 1 .
Revised: c 1 .

Original: For case 2, exactly the same argument as the above paragraph with the same h ′′ and h ′′′ shows that the coordinates of V ( p ) i is ± c .
Revised: For case 2, exactly the same argument as the above paragraph with the same h (cid:48)(cid:48) and h (cid:48)(cid:48)(cid:48) shows that the coordinates of V ( p ) i is ± c .

Original: Therefore, we have shown that for any i ∈ I − , V ( c ) i =
Revised: Therefore, we have shown that for any i ∈ I − , V ( c ) i =

Original: c 1 for same c> 0 .
Revised: c 1 for same c> 0 .

Original: The symmetry of distribution immediately tells us for any i ∈
Revised: The symmetry of distribution immediately tells us for any i ∈

Original: I + , V ( c ) i =
Revised: I + , V ( c ) i =

Original: c 1 for c< 0 .
Revised: c 1 for c< 0 .

Original: On the downstream distribution P ∗ , since there is no masked token, only V ( c ) i is working.
Revised: On the downstream distribution P ∗ , since there is no masked token, only V ( c ) i is working.

Original: Since
Revised: Since

Original: V ( c ) i =
Revised: V ( c ) i =

Original: c 1 always holds, we complete the proof.
Revised: c 1 always holds, we complete the proof.

Original: Lemma D.5.
Revised: Lemma F.5.

Original: Suppose M ∈ R(2 k +1) ×
Revised: Suppose M ∈ R

Original:  (2 k +1) is a matrix composed of ± 1 .
Revised: 2 k +1) × (2 k +1) is a matrix composed of ± 1 .

Original: The first row of M is [1 ,...,(cid:124) (cid:123)(cid:122) (cid:125) k+11’s , − 1 ,..., −(cid:124) (cid:123)(cid:122) (cid:125) k-1’s ] .
Revised: The ﬁrst row of M is [1 ,...,(cid:124) (cid:123)(cid:122) (cid:125) k+11’s , − 1 ,..., −(cid:124) (cid:123)(cid:122) (cid:125) k-1’s ] .

Original: Define a permutation ρ (1)=2 , ρ (2)=3 ... ρ (2 k +1)=1 .
Revised: Deﬁne a permutation ρ (1)=2 , ρ (2)=3 ... ρ (2 k +1)=1 .

Original: For all i ≥ 2 , M i,ρ ( j ) =
Revised: For all i ≥ 2 , M i,ρ ( j ) =

Original: M i − 1 ,j .
Revised: M i − 1 ,j .

Original: Then the rank of M is 2 k +1 .
Revised: Then the rank of M is 2 k +1 .

Original: Proof of Lemma D.5.
Revised: Proof of Lemma F.5.

Original: Note that M i + M ρ k +

Revised: Note that M i + M ρ k +

Original: ( i ) =
Revised: ( i ) =

Original: 2 e i for all i ∈
Revised: 2 e i for all i ∈

Original: [2 k +1] , which means we can express the orthonormal basis as linear combination of the rows in M .
Revised: [2 k +1] , which means we can express the orthonormal basis as linear combination of the rows in M .

Original: Therefore, the rank of M is 2 k +1 .
Revised: Therefore, the rank of M is 2 k +1 .

Original: D. T HE EXISTENCE OF RANDOM FEATURE SOLUTIONS D. 
Revised: F. T HE EXISTENCE OF RANDOM FEATURE SOLUTIONS F.

Original: Theorem D.6.
Revised: Theorem F.6.

Original: Suppose m ≥ ˜ O
Revised: Suppose m ≥ ˜ O

Original: (2 T T 3 ϵ − 2 ) , and V i ∼N (0 ,TI 2 T ) for all i ∈ [ m ] .
Revised: (2 T T 3 (cid:15) − 2 ) , and V i ∼N (0 ,TI 2 T ) for all i ∈ [ m ] .

Original: With probability at least 1 − δ over V , there exists ψ ′ ,u ′ , satisfying L ( ψ ′ ,u ′ ) ≤ ϵ and ∥ u ′ ∥
Revised: With probability at least 1 − δ over V , there exists ψ (cid:48) ,u (cid:48) , satisfying L ( ψ (cid:48) ,u (cid:48) ) ≤ (cid:15) and (cid:107) u (cid:48)

Original:  22 ≤ O ( T 2 δ − 1 ) .
Revised: (cid:107) 22 ≤ O ( T 2 δ − 1 ) .

Original: Proof of Theorem D.6.
Revised: Proof of Theorem F.6.

Original: Since the number of possible input in pre-training is finite, we can invoke
Revised: Since the number of possible input in pre-training is ﬁnite, we can invoke

Original: Lemma 9 in Bai & Lee (2020) to show that random Gaussian features can fit the pre-training task.
Revised: Lemma 9 in Bai & Lee (2020) to show that random Gaussian features can ﬁt the pre-training task.

Original: Lemma D.7 (Lemma 9 in Bai & Lee (2020)) .
Revised: Lemma F.7 (Lemma 9 in Bai & Lee (2020)) .

Original: Suppose ∥ h ∥ 2 = (cid:113) 1 T , v ∼N (0 ,TI 2 T ) .
Revised: Suppose (cid:107) h (cid:107) 2 = (cid:113) 1 T , v ∼N (0 ,TI 2 T ) .

Original: There exists a random variable a ( v ) such that
Revised: There exists a random variable a ( v ) such that

Original: [Equation]
Revised: [Equation]

Original: and a satisfies E v [ a 2 ]= O ( T 2 ) .
Revised: and a satisﬁes E v [ a 2 ]= O ( T 2 ) .

Original: Consider Q = K = 0 .
Revised: Consider Q = K = 0 .

Original: In this case we have h Q,K ( x ) = 1 T (cid:80)
Revised: In this case we have h Q,K ( x ) = 1 T (cid:80)

Original: Tj =1 x j .
Revised: Tj =1 x j .

Original: Also note that for the pretraining task, y
Revised: Also note that for the pretraining task, y

Original: = −√
Revised: = −√

Original: T 1 ⊤ h Q,K ( x ) for all x .
Revised: T 1 (cid:62) h Q,K ( x ) for all x .

Original: Since x j are norm 1 orthogonal to each other, we have ∥ h Q,K ( x ) ∥ 2 = 1 T for all x .
Revised: Since x j are norm 1 orthogonal to each other, we have (cid:107) h Q,K ( x ) (cid:107) 2 = 1 T for all x .

Original: Now we can show that V i ∼ N (0 ,TI 2 T ) , i ∈ [ m ] independently is the random feature solution which can solve the pre-training task.
Revised: Now we can show that V i ∼ N (0 ,TI 2 T ) , i ∈ [ m ] independently is the random feature solution which can solve the pre-training task.

Original: Suppose g ( h ) =
Revised: Suppose g ( h ) =

Original: 1 m (cid:80)
Revised: 1 m (cid:80)

Original: mr =1 σ ( V ⊤ r h ) a ( V r ) , and g ( R ) ( h ) :
Revised: mr =1 σ ( V (cid:62) r h ) a ( V r ) , and g ( R ) ( h ) :

Original: = 1 m (cid:80)
Revised: = 1 m (cid:80)

Original: mr =1 σ ( V ⊤ r h ) a ( V r ) I ( ∥ V r ∥ 2 ≤ √
Revised: mr =1 σ ( V (cid:62) r h ) a ( V r ) I ( (cid:107) V r (cid:107) 2 ≤ √

Original: TR ) .
Revised: TR ) .

Original: R is large enough such that Pr (cid:16) sup r ∈ [ m ] ∥ V r ∥ 2 ≥√ TR (cid:17) ≥ 1 − δ/ 2 .
Revised: R is large enough such that Pr (cid:16) sup r ∈ [ m ] (cid:107) V r (cid:107) 2 ≥√ TR (cid:17) ≥ 1 − δ/ 2 .

Original: We have g ( h ) =
Revised: We have g ( h ) =

Original: g ( R ) ( h ) on this event.
Revised: g ( R ) ( h ) on this event.

Original: Let g ∗ ( R ) ( h ) be the truncated version of g ∗ ( h )
Revised: Let g ∗ ( R ) ( h ) be the truncated version of g ∗ ( h )

Original: = −√ T 1 ⊤ h , g ∗ ( h ) ( R ) =
Revised: = −√ T 1 (cid:62) h , g ∗ ( h ) ( R ) =

Original: E v [ σ ( v ⊤ h ) a ( v ) I ( ∥ v ∥ 2 ≤√ TR )] .
Revised: E v [ σ ( v (cid:62) h ) a ( v ) I ( (cid:107) v (cid:107) 2 ≤√ TR )] .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: By Chebyshev and a union bound, we have
Revised: By Chebyshev and a union bound, we have

Original: Pr

Original: [Equation]
Revised: [Equation]

Original: For t = ϵ 2 , we have m ≥ n ( h ) R 2 T 2 ϵ − 2 . 
Revised: For t = (cid:15) 2 , we have m ≥ n ( h ) R 2 T 2 (cid:15) − 2 . |

Original: [Equation]
Revised: [Equation]

Revised: ≤

Revised: [Equation]

Original: Choosing R =
Revised: Choosing R =

Original: ˜ O
Revised: ˜ O

Original: ( (cid:112) ( T )) will make Pr (cid:16)∥ v ∥ > √ TR (cid:17) 14 ≤ cϵT .
Revised: ( (cid:112) ( T )) will make Pr (cid:16)

Original: [-Also note that n ( h )=( T/ 2 − 1) (cid:0) T T/ 
 +1 (cid:1) .-]
Revised: [-Also note that n ( h )=( T/ 2 − 1) (cid:0) T T/ 
 +1 (cid:1) .-][Table]

Original: Thus m ≥ ˜ O
Revised: Thus m ≥ ˜ O

Original: (2 T T 3 ϵ − 2 ) suffices.
Revised: (2 T T 3 (cid:15) − 2 ) sufﬁces.