Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models
Authors: Hong Liu
Sang Michael Xie
Zhiyuan Li
Tengyu Ma
Venue: Submitted to ICLR 2023 (2023)
Keywords: Language Modeling
Implicit Bias

Abstract: Language modeling on large-scale datasets leads to impressive performance gains on various downstream language tasks.  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). 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. 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.  These experiments demonstrate the existence of implicit bias of pre-training algorithms/optimizers---among models with the same minimal pre-training loss, they implicitly prefer more transferable ones. 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-training loss. We also prove in a synthetic language setting that among the models with the minimal pre-training loss, the flattest model transfers to downstream tasks.

Under review as conference paper at ICLR 2023
SAME PRE-TRAINING Loss, BETTER DOWNSTREAM: IMPLICIT BIAS MATTERS FOR LANGUAGE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Language modeling on large-scale datasets leads to impressive performance gains on various downstream language tasks. 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 be well-correlated with downstream performance (which is itself difficult to evaluate comprehen- sively). Contrary 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. On simplified datasets, we identify three ways to produce models with the same (sta- tistically optimal) pre-training loss but different downstream performance: continue pre-training after convergence, increasing the model size, and changing the training algorithm. These experiments demonstrate the existence of implicit bias of pre- training algoritbms//pptmiieess-omm amonn models with the same minimal pre-training loss, they implicitly prefer more transferable ones. Toward understanding this im- plicit bias, we prove that SGD with standard mini-batch noise implicitly prefers flatter minima in language models, and empirically observe a strong correlation be- tween flatness and downstream performance among models with the same minimal pre-training loss. We prove in asynthetic language setting that among the models with the minimal pre-training loss, the flattest model transfers to downstream tasks.
INTRODUCTION
Large language models (LLMs) trained on internet-scale data have improved performance wide array of downstream tasks (Devlin et al., 2018; Yang et al., 2019; Radford et al., 2019; Raffel et al., 2020; Brown al., 2020). These models trained with 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).
In common practice, the validation pre-training loss is used to monitor the training process (Brown 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). Moreover, theoretical works on understanding LLMs also focus on how the pre-training loss affects downstream performance. Saunshi et al. (2020); Wei et al. (2021); Xie et al. (2021) show that good pre- pre-training loss, or fitting the language modeling conditional probability well, is main reason for downstream success of LLMs. Their analyses generally treat the language models as blackboxes and do not take into account how the models represents the conditional probability.
In this paper, we question the conventional wisdom on the correlation between the validation pre-training loss and downstream performance for language modeling. 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). Due to the expressivity of modern neural nets, many parameter configurations even within the same architecture A priori, it is unclear why all these configurations should have the same downstream performance.
still have the same pre-training
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. Concretely, using simplified text datasets, find three situations that demonstrate such phenomenon:

---

Under review as conference paper at ICLR 2023
Even after the pre-training loss converges, models at a later time step still tend to perform better.
Models trained by standard algorithms have better performance than adversarially trained models with the same pre-training loss.
Larger models tend to perform better downstream than smaller models even if they have the same pre-training loss.
These situations are most prominent in the saturation regime, where the models are close the minimal possible pre-training loss (aka the entropy of the conditional probability, which can be estimated in our simplified datasets). In the saturation regime, the pre-training loss of all models are almost the same, but the transferability to downstream tasks varies. Interestingly, this phenomenon also holds when linear probing on contextualized presentations is used for evaluating downstream performance instead of finetuning. Thus, even though the predicted conditional probabilities of two models are the same (and correct), the contextualized representations can behave differently.
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. They only differ by the training algorithms that are used to produce them. Therefore, this suggests the training algorithms have an implicit bias toward one of these mooelss-standan algorithms with more training steps biases towards parameter configurations that transfer better to downstream tasks. The third case has more subtle but similar interpretation. There exists hypothetical large model that represents the smaller model with worse performance (by padding zeros to the smaller model). The training algorithm on the large architecture could have chosen it, but did not. This suggests the algorithm has an implicit bias against the hypothetical model (which has an equally good loss).
In supervised settings, optimizers are known to have an implicit bias toward selecting generalizable models among all models with small empirical loss. E.g., see see Damian al. (2021); Li al. (2021), which show that SGD implicitly biases toward flatter minima, and references therein. However, the role of implicit bias in sulfrsupe learning has not been studied and is conceptually different. Unlike in supervised learning, the gap between empirical and population self- supervsee losses is typically small, and thus implicit bias does not seem to contribute to bridging this gap. Instead, the implicit bias selects local minima of the population self-supervised loss that transfer better to downstream tasks.
Why do the algorithms bias toward some type of models? In Section 3, we provide a first-cut theoretical analysis of the implicit bias in language modeling. Fortunately, despite the conceptual differences, mathematical tools from supervised settings can be straightforwardly adapted to language modeling settings. We prove that mini-batch SGD prefers flatter minima of population pre-training loss among all minima in the saturation regime. 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; et al., 2021), partly because the mini-batch noise for L.Mss does not vanish even at convergence.
We corroborate our theory with empirical evidence in Section 4. 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 Hessian of the loss, as predicted by the theory) strongly correlates with the downstream performance.
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. 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. Here, results from the supervised setting cannot be readily adapted since they are obtained (partially) via generalization bounds (Wei & Ma, 201 2019a;b), which do not apply to the language modeling setting where the implicit bias not related to the gap between the empirical and population loss. 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.
2 THE EXISTENCE OF IMPLICIT BIAS IN LANGUAGE MODELING
In this section, we systematically investigate the relationship between pre pre-training loss and downstream performance with experiments. We find out that models with the same pre- pre-trainin loss but different training procedures can have different downstream performance.
2

---

Under review as conference paper at ICLR 2023
PCFG, Task C, 41M
OPT. ONLJ 235M
Downdrrem Accuracy 2
Accuracy
57
Pre: fee-fraining Loos
Pre-training Unear Probe
Pre Pre-training
Pre tre-training
Downetream
Linear Probe
55
54
Number Steps/ 1000
Number of Steps/ 1000
(a) PCFG-Task C
(b) OPT-
Figure 1: Models a later time step performs better, even after the pre-training loss converges, (a) A model with 41M parameters pre-trained on the PCFG-generated dataset, and evaluated on task C. (b) A model with 235M parameters pre-traing pre-trained on the opT-ggenertted dataset, and evaluated on QNLI
2.1 FORMULATIONS
Masked language modeling. Consider a vocabulary [0.0.......).. where i special token for the mask. Let x = [TT denote the input sequence of length T, and z:
denote the masked sentence, where is sampled uniformly randomly and independently from [2]. The MLM conditional probability refers to the probability of Iz given the rest of the sequence Pr (Iz z=+). We use Pr x to denote the c-dimensional probability vector Pr(-|x-1) == [Pr (x, - ................................ Pr. (z, x-]]] E RR In MLM pre-training, the model fo(-) (parameterized by 8) outputs the predicted MLM conditional probability vector fo(x- E RR. The model is trained to predict the masked token 2; given the rest of the sentence =- with cross entropy loss, (()) = F.,,(<([o(2 ().;:)) = (oog (oooo(((((o (x)))))
Downstream evaluation. The language model is composed of a feature extractor hg, which outputs a sequence of contextual representations, and linear classifier that outputs the conditional probability every position. On downstream tasks, we use randomly initialized go on top of the pre-trained hgg In fine-tuning, both Io and ft: are trained, while in linear probe, only Ie is updated. For fine-tuning, we use the contextual representations of the token. For linear probe, we concatenate the contextual representations of all the tokens together.
regime. 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, fo(x = Prr(- ( In the saturation regime, the MLM loss equal to the entropy of the true conditional probability L(0) = E.. - log [=-))]] Thus, all models in the saturation regime have the same, optimal pre-training loss, and will show that they behave differently on downstream tasks. Our experiments use expressive enough architectures such that there are multiple parameter configurations in the saturation regime for our simplified datasets. For real large-scale data, it is currently computationally challenging to arrive at the saturation regime. 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.
Saturation
= = + H (x; ==))). which is also the optimal pre-training loss.
2.2 EXPERIMENTAL SETUP
We design controlled experiments study the correlation between pre-training loss and downstream performance. In particular, we will find a set of models with almost the same pre-training loss. We effectively use the same architecture family so that the main difference between the models only stems from training algorithms. More details are provided in Section A.
Datasets. We introduce three generative models to produce simplified datasets, with which we can study various factors systematically. 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. Moreover, we can generate unlimited amount of text for pre-training to avoid overfitting to the empirical pre- traininiing loss.
1) CCFG-generated dataset. PCFG (Chomsky, 1956) generates sentences with probabilistic trees and is widely used to understand natural language (Johnson, 1998; Roark & Bacchiani, 2003; Kim et al., 2019; Yang et al., 2021). We randomly generate the production rules which satisfy the Chomsky Nor- mal Form (Chomsky, 1956). The non-terminal symbols in the parse tree can be viewed as intrinsic
For simplicity, we only consider masking out one taken in each sentence.
3

---

Under review as conference paper at ICLR 2023
PCFG, Bakk0
1944 72ak-12
Treooeers bane
Freeariig
Frelarhi
Prbburtn
Uneer hope
We-tokkin
Prctrent
Wodl S
Model
(a) Task B
(b) HMM-Task--l 10
(c) OPT
Figure 2: Larger models perform better downstream than smaller models, even with almost the same pre-training loss. (a) Pre-train on the PCFG-generated dataset and evaluate on task B. (b) Pre-train on the HMM-generated dataset and evaluate on task-10. (c) Pre-train on the OPT-generated dataset and evaluate on QNLI. See Section 2 and Section A for details.
quantities associated with the sentence such as sentiment and syntax. Thus we design three down- stream tasks A, B, and C to classify non-terminal symbols at different positions of the parse trees.
2) HMM HMM-generated generated dataset. HMM samples the hidden variables from the transition probabilities and the tokens from the emission probabilities. (Wei et al., 2021; Xie et al., 2021) ) also analyze the properties of pre-trained language models with HMMs. We generate the transition and emission probabilities as random block -iik-dia stochastic matrices. The downstream task to classify the hidden variable in the sentence. We use task-k to refer to classifying the k-th hidden variable.
3) OPT--genertted dataset. We also introduce a more realistic pre-training dataset generated by the OPT models (Zhang et al., 2022a). Starting from the bos token, we sample each token from the conditional probability output by the OPT model. For computational feasibility we only allow to generate the top-2000 most frequent tokens in the OPT vocabulary. We use QNLI and SST-2 from GLUE (Wang et al., 201 8) as downstream tasks.
Note that the true conditional probability can be computed efficiently for the three datasets given the knowledge of the generated models. For PCFG and HMM--genratter datasets, we can compute the true conditional probability with the inside algorithm (Lari & Young, 1990) and the Viterbi algorithm (Forney, 1973), respectively. For the OPT-generated dataset, 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.
Models and algorithms. For PCFG and OPT generated datasets, we use transformers (Vaswani et al., 2017) following the implementation of BERT (Devlin et al., 2018). 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. (2021). We use batch size 4096, and train each model for more than 30K steps until the pre- pre-training loss converges. For comparison. we also consider other training algorithms. The first adversarial algorithm is inspired by Liu et al. (2020); Raghu et al. (2021), where the models are pre-trained with an additional meta- caatallinning objective which messes up downstream performance. The second algorithm is manually setting the weight of the model to represent lookup table which memorizes all the masked sentences and the corresponding true conditional probability. When the transformer is sufficiently large, the lookup table can be encoded shown in Yun al. (2019).
2.3 RESULTS
We compare the downstream performance of models with the same pre- pre-training loss in the following situations: (1) training for different number of steps after the pre-training loss converges, (2) using dif- ferent model sizes, and (3) training with normal training algorithms vs. adversarial training algorithms.
In Figure 1, we plot the validation pre-training loss and the downstream performance of different models checkpoints along pre- pre-rraning. After the pre-training loss converges, although the pre-training loss does not improve, the downstream accuracy continues increasing.
Even with the same pre-training loss, larger models are better than smaller models. In Figure 2, we plot the pre- gre-train loss and the downstream performance of models with different sizes. As we increase the model size, the pre-trate mee-training loss approaches the entropy of the true conditional probability, which is 3. 196,, 3. 7585, and 1.865 for PCFG, HMM, and OPT respectively. For PCFG and oeT-generated datasets, use the vertical dashed line indicate the place where the pre-training loss saturates as
4

---

Under review as conference paper at ICLR 2023
we scale up the model. For the much simpler HMM, the smallest 4M model can fit pre-training close to the entropy of the true conditional probability. With the same pre- vre-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. See Section A for results on other downstream tasks.
Naturally trained transformers are better than adver- sarially trained ones. In Table 1, we evaluate the 235M transformers on PCFG tasks A and B with different pre-training algorithms. Although the ad- versarially trained transformer has almost the same pre-training loss as the normally trained 235M trans-
Table 1: Different pre-training algorithms on PCFG.
Method
Pre-training
Task A %
Task B %
AdamW
3.204
89.9 10.3
49.2.10..8
Adversarial
3.206
83.1 10.6
4..
Lookup table 3.196
71.2
39.7
former, it more than 6% worse than the normal 235M model, and even worse than a normal 9M model on the downstream task B. 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. (2022). 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.
The experiments above indicate that for models with the same architecture family and the same pre- training loss, the choice of training algorithmss, model sizes, and the number of steps that the optimizer works can affect the downstream performance. 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.
3 IMPLICIT BIAS LEADS TO FLAT SOLUTIONS IN LANGUAGE MODELING
As discussed in the introduction, the difference between the role of implicit bias in supervised learning and language modeling is conceptual, because the gap between empirical and population self- superviid loss is small and thus implicit bias is not needed for bridging this gap. Instead, the implicit bias benefits the performance on downstream tasks by picking networks that are more adaptable to those tasks.
Fortunately, the mathematical tools developed for supervised learning can be adapted language modeling, which even allows cleaner results by removing some artificial assumptions like adding label noise. Concretely, we show that mini-batch SGD can find models in the flatter areas of pre-training loss landscape. The flatness is measured by the trace of Hessian of the pre-training loss Tr[V²2 L(O)]. See Figure 4 (Left) for an illustration of the implicit bias.
We analyze SGD on the population cross cross-entropy loss sampled data at every iteration, because, as argued, the difference between empirical and population pre-training loss is not our focus. 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). Let 1) be the learning rate and let og denote the parameter at step k. We drop the superscript 1) when there is no ambiguity. We will show that the implicit bias kicks when SGD reaches global minimizer- drives the iterate towards flatter global minimizers. For simplicity of demonstration, we analyze the process starting from global minimizer 0, i.e., we assume that 02 =0= 0 (for all 17). At each iteration k; we get fresh sample (a.,), where x is a sentence and & is the position of the masked token, and update the parameter 0 by 0x+1 02 V (fo, (x (x):) We assume the network is sufficiently expressive such that there are many fundamentally different global minimizers of the pre-training loss L. 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.
with freshly
A Assumption& 3.1. Assume that the loss L is a C3 -mooth function, and that the set of global minimizers, T. (d -dimensiila C² -u-summniffold of Rdd for some integer 1 < M S d, where for all 0e rran (()) =M.
A key observation for language model is that even if the model reaches the saturation regime, that is, the model reaches point on the manifold I' 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." Therefore, the dynamics of
This is in contrast with typical supervised setting where the empirical 0-1 loss and cross- cros--yntrppy loss can both achieve zero and consequently the mini-hatch noise vanishes. Such a difference enables us prove cleaner results (without the label noise) than the supervised setting (Damian al., 2021; Li al., 2021).
5

---

Under review as conference paper at ICLR 2023
PCFG, SGD, 235M
HMM, SGD, 67M
Downstream Accuracy %
Downstream Accuracy
55
54
Pre- Le--trainng Loss
Pre training Loss
92
Pre-training
Pra-training Hessian
Trace Hessian
50
48
Probe,
-3
2/g
Linear Task-10
43
8
3.00
1.4
33
2.5
Number Steps 1000
Number Steps 1000
Figure 3: The trace of Hessian correlates with downstream performance for model checkpoints with different number of steps after the pre- pre-training loss converges. Left: A model with 235M parameters pre-trained on the PCFG-generated dataset, and evaluated on task C. Right: A model with 67M parameters model pre-trained on the HMM-generated dataset, and evaluated 10.
SGD do not completely stop; instead, the iterate oscillates around the manifold r. turns out that this oscillation in turn encourages the parameter to move in certain direction along the manifold, determined by the covariance structure of the stochastic gradient. 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 pre-training loss.
Lemma 3.2 (Bartlett identity). For any 0 E r, (0) = L2)( where ()) the covariance of the stochastic gradient at that is, E(t) Er, Velogl (z ())x, (Volo[ff(==-llle) (L(O)TVL(C)
Though we give a proof of the lemma in Appendix for completeness, the formula holds for the MLE loss of any well-specified probabilistic models at global minimizer, and both the gradient covariance and the Hessian equals to the Fisher information matrix.
With Lemma 3.2, we can invoke Corollary 5.2 of Li at. (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.
ao( a0(1== - -
0(0)
where Vr= V is the Riemannian gradient on manifold r, or just the ordinary gradient projected back to the tangent space of O. In other words, the ODE (1) is essentially a projected gradient descent algorithm with loss function L()), the constraint set I, and infinitesimal learning rate. We show that SGD effectively minimizes the trace of the Hessian Tr(V with the constraint set similarly to ODE in (1).
Theorem 3.3. Suppose the loss function I, and the manifold of global minimizers I satisfy Assumption 3.1. For any K > 0 such that ODE (1) has solution { it holds that 0 k/p² converges in distribution to (K) 0.
Finally, we note that the above result can be extended to an arbitrary batch size B. The covariance of
stochastic gradient at O with batch size, denoted by (0), satisfies that -0)= (0). Therefore EB (0) h I()) 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& 1/ in equation (1) replaced by 1
4B
4 FLATTER MODELS HAVE BETTER DOWNSTREAM PERFORMANCE
In this section, we demonstrate with experiments that the flatness is well correlated with downstream performance the setting introduced in Section 2.
Evaluation of flatness. As in Theorem 3.3, we measure the flatness of different models in Section 2 by the trace Hessian of the pre-training loss (smaller trace of Hessian indicates flatter minima.) Note that when the model approaches the saturation regime, the trace of Hessian is approximately the second order derivative times the square the norm of the Jacobian, which is a high-dimensional matrix. For computational feasibility, we adopt a technique inspired by Wei et al. (2020) to unbiasedly estimate the trace of Hessian with random samples. Details are provided in Section B.

---

Under review as conference paper at ICLR 2023
Smoller architecture
Rabness, mppisi baa
Fatnees Implich blas
miNimum losses (alobal min
iglobal
Rotrest soluution within the smaller model architacturs
Optimization Prajectory iTheorum 3.3)
Flatbes: solution within the larger model archiecture
Figure 4: Left: The role of implicit bias. After the pre-training loss converges, the implicit bias drives the model toward flat minima, predicted by Theorem 3.3. Right: The interaction between model size and implicit bias. The implicit bias drives the model toward flat minima on both larger models and smaller models. The smaller model architecture can be viewed as a subset of the larger model architecture. Therefore, larger models can achieve flatter minima than smaller models.
PCPS
Lineal PiLie BIN
Unear Prte, Tark-3
Uneer Fribe
Track criesian
Tacs Hosian
Modal
(a) PCFG-Taske
0.92
(b) HMM->Task 10, 0.94
(c) OPT-QNLL, = 0.87
Figure 5: The trace of Hessian correlates with downstream performance for models with different sizes and almost the same pre- pre-training loss. On datasets generated by PCFG, HMM and OPT, we obeserve 0.92, 0.94, and 0.87 coefficient of determination with linear regression.
Results. In Figure 3, we compare the downstream accuracy and the trace of Hessian of different check- points obtained different times during pre-training On the PCFG and HMM datasets, the trace of Hes- sian demonstrates a clear decreasing trend after the validation pre-training loss converges, following the prediction of Theorem 3.3. Furthermore, as the trace of Hessian decreases, the downstream performance improves by 1.6% and 4.0% the PCFG and HMM datasets, respectively.
Table 235M transformer pre-trained with dif- ferent algorithms evaluated on PCFG Task C.
Method
Pre-training
Acc %
Tr(II)
AdamW
3.20
55.7.00,6
2.02.10.16 50.2.11.0
Adversarial
3.20
6.4511..1
We compare the trace of Hessian of the models pre-trained with adversarial algorithm and standard AdamW in Table 2. The trace of Hessian of the adversarially pre-trained model is times larger than the normally pre-trained model, corresponding to drop of 5.5% in downstream performance.
In Figure 5, we compare the downstream accuracy and the trace of Hessian of models with different sizes. On the dataset generated by a PCFG, the pre- pre-trainig loss 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.5 1.54, correlating with the increase of linear probe accuracy from 43.2% to 50.3%. On the OPT and HMM-generater datasets, we can also observe an increase linear probe accuracy with a sharp decrease in the trace of Hessian, as we increase the model size.
Interaction between implicit bias and model size. Intuitively, the implicit bias drives the model toward flat minima on both larger models and smaller models. The smaller transformer architecture a subset of the larger transformer architecture family (as justified in Section B). Thus the flattest minimum found within a larger transformer is flatter than the flattest minimum found within smaller transformer, and performs better downstream. (See Figure (Right).)
5 FLATNESS REGULARIZATION PROVABLY IDENTIFIES TRANSFERABLE MODELS ON SYNTHETIC LANGUAGE
Toward formally proving the connection between the flatness regularization (introduced by the stochas- tic gradient as argued in Section 3) and the downstream performance, consider setting with synthetic Dyck language. The simplicity of the data allows us to sharply analyze the internal working of single-layer transformer (with attention layer and MLP layer) for masked language modeling. We show that multiple parameter configurations can predict the conditional probability well, including one ideal model that leams the correct representations capturing the intrinsic structure of the sentence,
7

---

Under review as conference paper at ICLR 2023
Configuraiin1 (2)
y
Enoodings
Softrion Altottinn
Configumtoo 120
hppt
An example of input encodings with sentence length T- Right: Illustration of the two solutions. The softmax attention can sum the token encodings into Solution (1) contains two features transferable the downstream task. The neurons in solution (2) are sampled from Gaussian distribution, and not related to the downstream task. Both solutions can output the correct prediction for MLM pre-training, but solution (I) has much smaller trace of Hessian. and many "cheating" models that essentially memorize the conditional probability using random features. We will prove that the flattest model is the desired model that transfers to downstream tasks,
Figure 6: The synthetic language setting. Left:
Pre-training Distribution. Consider a variant of the Dyck language (Nivat, 1970) consisting of matching brackets. The vocabulary of the language has two brackets ( (add] ). Each sentence is composed of a sequence of tokens such that the total numbers of ('s and ''s are equal. To sample from the pre-training distribution P. we first draw sentence uniformly over all valid sentences with even length T. Then, we randomly select one position in [T] and replace the bracket with a mask token.
Downstream Task. 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 as the downstream task. Concretely, for any sequence z in {.) of length T, let g" (I:) count the number of mismatches in a::
g" (xx) 2 (x)) of) y''s in x - of ("s in
Thus, the sentence x a valid string in the language if and only (z)))= For MLM, the masked token can also be recovered from 9" (zr): 9 (a) 1 if the masked token is (. and g"(x) = -1 if the masked token is), ) To evaluate if model learns the structure, we consider downstream distribution of sentences which do not necessarily belong the the language. Each token is sampled from {
uniformly, randomly, and independently.
Encoding of the Inputs. With a slight abuse of notation, also denote by 2, the encoding of the t-th token. We encode the input as a one-hot vector in dimension = 2T, where the index of the nonzero element encodes the position and the sign encodes the bracket. Concretely, let e E Rd be the natural basis vector where the (-th entry is 1. Let x, e, if the 1-th token is ( and Te = -er otherwise. If the position / is mask, we set F, v, where v encodings with T 6 are provided in Figure 6 (Left). Note that the target function can be expressed as g" (z(x) (Ir, T 22,1. I:T) with this input encoding, where 1 is the all one vector in R² and [a] refers to the first T coordinates a.
(nif((+te=+) }). Examples of the
Models and Algorithms. Suppose K E ER are the query and key matrices, V E B is the value matrix and u E Th the parameter of the output layer. Let vj (Q, A single-layer transformer
composed of an attention layer and an MLP layer. (A
_ o O (Zj=] ax.j Væj), where the (z)== max (2,00 is the relu activation.
attention score 11. softmax (Qu,,, Kx.))...... (Q.,, We use the output of the first token, f (x) = (.))),..()))
We use the squared loss for both MLM and downstream adaptation. The loss function of MLM is L(bb,). In downstream adaptation, we have finite dataset (z y sampled i.i.d. from Pas. The training loss with n data is I (1b,u. and the population loss for the downstream task is L LP (v),u).
Main Intuitions. We are interested in two kinds of parameter configurations both with good pre- training loss: (1) learning the natural and transferable features 1y and (2) fitting the pre- pre-training task by memorizing the masked sentences. We construct the two solutions as follows. For solution (1), first note that the softmax attention layer can take the average of all the token encodings [xx.][= in a sentence.
Let us denote the sum by E Rd. = >T- Tt. Note that the first T coordinates in 2 are + indicating the bracket type and the last I coordinates indicate the position of the mask (See Figure 6 (Right)). On top of 2, two neurons can predict the masked token in MLM perfectly. Consider the two neurons Vi = [1:11::p1p V22 = 1 Or] Then g" (x) or (V= 2) oV 2), which is the transferable solution. For solution (2). we set the entries in V to i.i.d. samples from (0,T). If m is sufficiently large, we can find the coefficient u to express 9 (ar) with random Gaussian features, i.e. g" (x) = u (Vz).
8

---

Under review as conference paper at ICLR 2023
We observe that the trace of Hessian of configuration (1) is smaller than configuration (2), due to main difference between them the cancellation between activated neurons. In configuration (1), for every possible imput, only one of the neurons o (V 2) and (V2 2) is activated. In contrast, in configuration (2), many neurons can be activated at the same time. Among them, the output coefficient uy 's contain both positive and negative values, leading to cancellation between activated neurons. In Lemma F.1, we link the trace of the Hessian with the cancellation between neurons. Indeed, we show that the mimimum of trace of the Hessian can be achieved only if there is no such cancellation. Therefore solution (1) is also the minimizer of the trace of Hessian. The intuitions are formalized in Theorem 5.1.
Consider minimizing the trace of Hessian among all the solutions to the MLM pre-tra gre-training task: minimize,, Trr Tr ((),))) Tr (w)))), subject to L(b,u)=0..
Theorem 5.1. Suppose m 2 2 and T 2 6. The flattest solution à â are defined as the solution of the optimization problem above. û is the minimizer of downstream training loss on top of si. (1) =0
i E argmin, l|ull2 subject to P. (1b,u)= =0. 0. Then with probabilyy at least 1-2 I.
6 RELATED WORK
Language modeling and downstream adaptation. Large language modeling has revolutionized the NLP field. Starting from Devlin et al. (2018), a line of works improve the downstream performance on wide range of tasks with increasing model size and data amount (Yang al., 2019; Radford et al., 2019; Raffel et al., 2020). LLMs even exhibit unexpected emergent behaviors, such as in-context learning (Xie al., 2021; Min al., 2022), step-by-step reasoning (Wei et al., 2022), and zero-shot learning (Brown et al., 2020). Kaplan et al. (2020); Hernandez et al. (2021) 1) study the behavior of language models with increasing size, and find out that the pre-training loss is typically correlated with downstream performance evaluation metric for language models. A notable example the efficient transformer line of works, which benchmark the pre- pre-training loss given the same computation constraint (Dai et al., 2020; Wang et al., 2020; Choromanski et al., 2020; Liu et al., 2021).
model size increases. In practice, the pre-training loss is used as an
Understanding the success of language modeling. Empirical works on understanding MLM find out that the representations of language models encode rich semantic and syntactic information (Peters et al., 2018; Htut et al., 2019; Hewitt & Manning, 2019; Mamou et al., 2020). Theoretical works show that fitting the MLM conditional probability is a sufficient condition for good performance on downstream tasks. Zhang& & Hashimoto (2021) show MLM representations recover latent variables in graphical models. Wei et al. (2021) show linear probe on top of MLM models solves downstream tasks on datasets generated by HMMs. Recently, Saunshi et al. (2022) show that models with the same pre- pre-training loss but different architectures can have different downstream performance. Tay et al. (2021) find that narrow but deep transformer is better than a wide but shallow transformer with the same praitianning loss. Zhang et al. (2022h) demonstrate that Albert (Lan et al., 2019) generalizes better to OOD tasks than Bert on a synthetic reasoning task. These works indicate that the architecture is important factor for good downstream performance beyond pre-training loss. This paper discovers the role of implicit bias in language modeling, which happens with models in the same architecture.
Implicit bias in supervised learning. The training algorithm chooses solutions with certain properties, and usually leads to better generalization (Gunasckar et al., 2018: Soudry et al., 2018; Li et al., 2017: Ji & Telgarsky, 2018; Arora et al., 2019; Lyu & Li, 2019; Li et al., 2020; Woodworth et al., 2020; HaoChen et al., 2020). Recently, Blanc et al. (2019); Damian et al. (2021); Li et al. (2021) demonstrate label noise SGD biases the models toward flatter minima. However, the setting of implicit bias in supervised learning is different from language modeling. In language modeling, we have access to gigantic corpus, and cannot interpolate the pre-training dataset. Moreover, we care about the adaptability of the solution on downstream tasks instead of generalization in distribution.
CONCLUSION
We study the relationship between pre-training loss and downstream performance on language models. We discover that implicit bias matters beyond pre-training loss, and explore the mechanism of implicit bias in language modeling. Our experiments focus on simplified datasets due to constraint of computa- tional resources. We hope that the phenomenon can predict the implicit bias on more complex datasets as the community scale the models to even larger. Theoretically we provide cases where the flatness regularization can decide the performance on downstream performance. We wish this motivates future works the relationship between implicit bias and the internal working of the models.

---

Under review conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
To ensure reproducibility, we describe the implementation details of the algorithms and the construction of the datasets in Section and Section B. The code of the experiments is provided in the supplementary material. We provide the proof in Section C and Section F.
REFERENCES
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems, pp. 7411-7422, 2019.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri. Stefanie Jegelka, Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 Proceedings of Machine Learning Research, pp. 948-1024. PMLR, 17-23 mlr www.rrppreseeeee. www.pppesss...... wttp://vvvll//// wtr.////arra22a2222222 html.
Jul 2022. URL wttps::////o....... http::////ooccedeeeeeeee
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. arXiv preprint arXiv:21 2021.
Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. International Conference on Learning Representations (ICLR), 2020.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbee like process. arXiv preprint 1rXiv:1904.000000 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing 33: 13:1877711011, 2020.
Noam Chomsky. Three models for the description of language. IRE Transactions on information theory, 2(3): 1956.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 204Y:49099..4994 2020.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe Le. Humnel-transformer Filtering out sequential redundancy for efficient language processing. Advances in neural information processing systems, 33:4271 1284271-4282,,1 2020.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461, 2021.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 110.v:18800.480000 2018.
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions, Journal of Machine Learning Research, 21:136, 2020.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2001000142222000
G David Forney. The viterbi algorithm. Proceedings of the IEFF, 61(3):268-278 1973.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, pp. 9461-9471, 2018.
Jeff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.0868/ 2020.
10

---

Under review as conference paper at ICLR 2023
leamnin Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Danny Hemnandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer. 2021.
arXiv preprini
John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pp.
4129-4138, 2019.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-17888 1997.
Phu Mon Htut, Jason Phang, Shikha Bordia, and Samuel R Bowman. Do attention heads in bert track syntactic dependencies? arXiv preprint arXiv: 1rXiv: 2019.
Peter Izsak, Moshe Berchansky, and Omer Levy. How to train bert with an academic budget. arXiv preprint arXiv:2104.07705, 2021.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint 188.i 02032, 2018.
Mark Johnson. Pcfg models of linguistic tree representations. Computational Linguistics, 24(4): 613-632, 1998.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, arXiv preprint arXiv:2001. 003661, 2020.
Yoon Kim, Chris Dyer, and Alexander M Rush. Compound probabilistic context-free grammars for grammar induction. Xiv preprint arXiv: 10225, 2019.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint 1991:19909.1192.. 2019.
Karim Lari and Steve Young. The estimation of stochastic context- -fee grammars using the inside- outside algorithm. Computer speech & language, 4(1):35-56 1990.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterizee matrix sensing and neural networks with quadratic activations. arXiv preprint arXiv:1712.09203. pp. 2 47,47,20 2017.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. arXiv preprint 1xXi::0022.09899 2020.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss? mathe- matical framework. arXiv preprint w 2021.
Hanxiao Liu, Zihang Dai, David So, and Quoe V Le. Pay attention to mlps. Advances in Neural Information Processing Systems, 34:9204-9215 2021.
Hong Liu, Jeff Z HaoChen, Colin Wei, and Tengyu Ma. Meta-learning transferable representations with single target domain. arXiv preprint arXiv:201 1.01418, 2020.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 05890, 2019.
Jonathan Mamou, Hang Le, Miguel Del Rio, Cory Stephenson, Hanlin Tang, Yoon Kim, and Yeon Chung. Emergence of separable manifolds in deep language representations. arXiv preprint arXiv:2006.0 01095, 2020.
11

---

Under review as conference paper at ICLR 2023
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role demonstrations: What makes in-context learning work? arXiv preprint arX1v:2202.128.77 2022.
Maurice Nivat. On some families of languages related to the dyck language. In Proceedings the second anmual ACM symposium on Theory of computing, pp. 22.221 24.22211,5 1970.
Matthew E Peters, Mark Neumann, Luke Zettlemoyer, and Wen-tau Yih. Dissecting contextual word embeddings: Architecture and representation. arXiv preprint arXiv: 1808. 0098,9 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188);9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
Aniruddh Raghu, Jonathan Lorraine, Simon Kornblith, Matthew McDermott, and David K Duvenaud. Meta-learning to improve mre-training. Advances in Neural Information Processing Systems, 34: 23231-232444 2021.
Brian Roark and Michiel Bacchiani. Supervised and unsupervised pefg adaptation to novel domains. In Proceedings of the 2003 Human Language Technology Conference of the North American Chapter the the Association for Computational Linguistics, pp. 205-212, 2003.
Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. mathematical exploration of why language models help solve downstream tasks. arXiv preprint arXiv:2010.03664: 2020.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. arXiv preprint arXiv:2202. 14037, 2022.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018.
Yi Tay, Mostafa Dehghani, Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. arXiv preprint arXiv:2109. 10686, 2021.
Jingfenn,,Jnneng Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: I 776600366,,001
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018.
Sinong Wang, Belinda Lii, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04777 04766. 2020.
Colin Wei and Tengyu Ma. Data-dependent sample complexity of deep neural networks via lipschitz augmentation. In Advances in Neural Information Processing Systems, pp. 9722-9733, 2019a.
Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification via all-layer margin. arXiv preprini arXiv: /9X:v:/990..0 191004441 2019b.
Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. arXiv preprint arXiv:2002. /2915, 2020.
Colin Wei, Sang Michael Xie, and Tengyu u Ma. Why do pretrained language modelshelp in downstream tasks? an analysis of head and prompt tuning. arXiv preprint arXiv:2/06.09226, 2021.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le. and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11900 2022.
12

---

Under review as conference paper at ICLR 2023
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparameetriee models. arXiv preprint arXiv:2002. 0220000..0927.... 2020.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
Songlin Yang. Yanpeng Zhao, and Kewei Tu. Neural bi-lexicalized pefg induction. arXiv preprint arXiv:2105.1502/ 2021.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Russ Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.1007, 2019.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. arXiv preprint 27Xiv:2205.01666 2022a.
Tianyi Zhang and Tatsunori Hashimoto. On the inductive bias of masked language modeling: From statistical to syntactic dependencies. arXiv preprint arXiv:2104.05662 2021.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint arXiv:2206. 04301, 2022b.
13

---

Under review conference paper at ICLR 2023
A
DETAILS IN SECTION 2
A.1 GENERATING SIMPLIFIED DATASETS
PCFG-gene CCGG-genertted dataset. We consider PCFG with vocabulary size 200. The state space is S, and |S| 50. All the production rules have two symbols the right side. The sentence length is limited to 32, which means the depth of the parse tree is limited to We generate a total 2 x 10 sentences, which is 3.4 x 10a tokens. The downstream tasks classifying the non non-terminal symbols in the parse tree of the PCFG (50-way classification). The label is defined as y argmax, PPr(s ( 22....1.. +L). Tasks A, B and are defined on the symbols corresponding span length L= 32, 16 and respectively. Each of the downstream task contains 0. M examples. Examples of the generated trees are provided in Figure
A2.5
22.
M.7
anding officiate
obextred
vetnequeet appitiection
ertetnet maccoeded
Figure 7: An example the genearated PCFG sentence.
HMM-generated dataset. We consider an HMM with vocabulary size 200 and state space size 100. The sentence length restricted tokens. The downstream task is to classify the latent variable in the HMM generative model. We consider task-6 and 10, which classify the 6-th and 10-th hidden variables respectively. Each of the downstream task contains O. 1M examples.
We generate a total of x 10" sentences, which is 1.6 10"
OPT-generated dataset. We use the 125M OPT model to generate the training dataset. To simplify the dataset, further process the logit of OPT to select only from the top-2000 tokens in the vocabulary. Starting from the bos token, we sample every token the sentence from the predicted autoregressive LM probability. The sentence length is restricted to 24. We generate total of 2 10 sentences, which is x.2x 10 tokens. Examples of the generated text are provided in Figure 8.

I really either, why you feel wouldn't great you 
I want the person in the photo tell there are children under years old in the
Figure 8: An example of the genearated OPT sentence.
A.2 COMPUTE THE TRUE CONDITIONAL. PROBABILITIES
We can compute the true MLM conditional probability Pr (x from the joint probability P with one mask per token,
-r
Pr( EcEw Pr
Since we already know the generative model, we can compute the joint probability efficiently. For PCFG, we can compute the joint probability with the inside algorithm, which decomposes the joint probability into lower layers in the parse tree. For HMM, we compute the joint probability with the Viterbi algorithm. For OPT, have Pr(xI,..,.TT) Pr(m 1 (31,...,.+..
A.3 MODELS
We transformers PCFG and OPT OPT-generated datasets. We use learning rate 1e-3 and warmup proportion 0.06. All the models are trained based on the implementation Izsak et al. (2021). We list the sizes of the transformers in Table 3. d-model is the size of the hidden layers. d-inter is the size of the
14

---

Under review conference paper at ICLR 2023
Table 3: Shape of the transformers.
d-model
n-nead
#layers
d-inter
1
256
128
2
4
512
192
3
5
768
256
6
1024
512
8
768
12
3072
1024
16
16
4096
1280
20
20
5120
1536
24
24
6144
intermediate layers in MLP. n-head is the number of heads per layer. # is the number of layers. For LSTMs, use the implementation of PyTorch. We consider d-hidden 1128,256,512,7666 1024), and #layers in 12,
A.4 ALGORITHMS
The lookup table. To evaluate the downstream performance of the lookup table, we first create lookup table with the data of the downstream task. With the method mentioned in Section A.2, generate the true conditional probability each token and use it as the contextual embeddings.
The ad verssaiia algorithm. The adversarial algorithm use to mess up the downstream performance maximizing a meta-learning objective in pre-training. Suppose the linear head of the downstream task is Is and the feature representation is hgg The meta-learning algorithm first trains the head Je to minimize the training loss of the downstream task, and then update he to maximize the validation loss the downstream tasks. Concretely, randomly sample two disjoint subsets Dj and D2 from the downstream training dataset D. We train De to minimize the loss of downstream tasks on Di, ((v)) argmin /., [( (2)()))))))) Then we train hes to maximize the validation loss on D2 during pre-training, minimize, A be efficiently carried out with closed form solution of as shown in Liu et al. (2020).
E( ( (( (z)()))) The optimization can
Fine-tuning, Following the standard protocol of Devlin et al. (2018), we use the contextual embeddings of the CLS token for fine-tuning We use AdamW with learning rate le-4. We perform 200 warmup steps and train on the downstream tasks for 10 epochs.
Linear probe. Since the CLS token is not trained in pre-training, concatenate embeddings of all the tokens the sentence as the representations. We use AdamW with learning rate le-3 to train the linear head. We train the downstream tasks for 100 epochs. Note that to make the capacity of the linear probe itself controlled, we adopt random Gaussian projection to dimension 512 on the concatenation of the embeddings.
We report the standard deviation linear probe and fine-tuning from 5 random seeds.
Evaluation of pre-training loss. Since have access calculate the true conditional probability, can calculate the cross entropy loss as the sum of the entropy of the true conditional probability and the KI. divergence between the predicted and true conditional probabilities. This is more accurate than evaluating on the validation datasets in the standard ways. We report the number of pre- pre-training loss with 10 sentences, and calculate the standard deviation subsets, each of which has size 2 10a.
A.5 RESULTS ON OTHER DOWNSTREAM TASKS.
We also provide results other downstream tasks in this subsection. On PCFG Task A, OPT SST-2 and the T of HMM, we can also observe the increase in downstream performance as we scale up the models in the saturation regime.
15

---

Under review as conference paper at ICLR 2023
PCRG TA
Dewichoom Accuray
Proet la
So. powiturr Lncs
Inn-larerg Unear Prose
Unear Prete
Leeer
22n
Pr-tmiiiigg
Nodel Sbe/ M
Model Size/
Mods Soe/M
(a) PCFG-Task A
(b) OPT-SST-2
(c) HMM-Task--
Figure 9: The downstream accuracy continues to rise as we increase the model size, althought the pre-training loss remains unchanged.
B DETAILS IN SECTION
Unbiased estimate of the trace of Hessian. Evaluating the trace of Hessian requires the norm of the Jacobian Ve fog]fo((=_))))) Since the output dimension and the number of parameters are all very large, computing the Jacobian fologl.fo (x] will be very inefficient. Instead, we can estimate the trace of Hessian unbiasedly with random samples as follows. Suppose fo (x=,) is the predicted probability of the conditional probability. In the saturation regime, as fo(x approaches the true conditional probability, the Hessian of the pre-training loss w.r.t. the parameters can be expressed as
L(0)= Br Velog]fo (2 (=
Therefore have
E 1t.= Volog] [o
To approximate this expectation, we can first sample f from the language, then draw i.i.d. samples Iy from (fo(x-e)), and use the average as the unbiased estimate. For all experiments, we sample 10000 v- and sample 50 for each ----
Details in Figure 3. To verify Theorem 3.3 that SGD biases the model towards flatter minima, we conduct MLM on PCFG and HLMM-genratte datasets with SGD. We set the proportion of warmup stage to 12% total number of steps, and fix the learning rate to 1c-3 after the warmup. We evaluate the downstream performance and the trace of Hessian of different checkpoints along pre-training. The standard deviation of trace of Hessian is calculated based on 5 times sampling 50 examples as mentioned above. Apart from the PCFG task and HMM task- 10, also provide results on PCFG tasks A, B and HMM task-6 in Figure 10.
Dowrstrom Accurary %
Fre Kiiigg Lss
Freetra Tre-iiiingg loss
Preipiing ng
Peeiirinn Tast of Pookan
Bots Roadn
Ureat Prbbo Tolk
Uner Prabe, Tosk
Namner of Sleno/ 1000
Marber of Mtgoo 1383
Number of Sigo 3300
(a) PCFG, SGD, Task A, 235M
(b) PCFG, SGD, Task B, 235M
(c) HMM, SGD, Task-6,67M
Figure 10: The trace of Hessian correlates with downstream performance for model checkpoints with different number of steps after the pre-training loss converges.
B. EMBEDDING OF A SMALLER TRANSFORMER INTO A LARGER TRANSFORMER.
In this subsection, we show that a smaller transformer can be embed into larger transformer without changing the functionality. 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.
16

---

Under review as conference paper at ICLR 2023
B.1.1 THE BASE CASE WITH MLPs.
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. (2016). Consider a multi-layer MLP (a). The weight matrices are W = [Wo....; W1.-1]. T.hh The dimensionality is Wi E RR xd; The representations are defined recursively, h441 = a (x)). We denote the input by hg (a=) x and the final output defined as fw,+(2) =a (x)). The activation a here is relu or leaky relu. We aim to embed Sw,a (1r) into Jw, (x)), where W WW.. and Wi ER2d 1x2d; We need to make sure fw, (z)= -f fw.a (z) for all æ.
WL-1]
W Wi W
W
[a] v2
For this case, we can set Wi 1/22
for for v2 W
and ā= 1
We can
Wi
(x) verify that h=(x)= //2 (1)
1
inductively. Therefore we have
fw.a (z) =ā h hx())
V2 la T "a" _a (x)
1
V2 (x)
1
2 Sw,a(()) fw.a(x)
= ff,
B.1.2 REAL TRANSFORMERS.
Next we turn to transformers with residual connections and layer norm. We first use the same strategy as the MI. case to add additional feature dimension and attention heads by replicating the weights, and then show how to add new layers using residual connections. At a high level, replicating the weight maintains the mean and the variance calculated by the layernorm. Therefore the representations inside the transformer also get replicated, without changing the values in each of the replicated groups.
Setup. A transformer is composed of an input embedding We. L blocks of self-attention, and an output layer. Transformers also contain layer-norm and residual connections. Suppose the input [a] , Drr, where 2; Rd.. Each block of the self-attention contains of an attention layer and an MLP layer, both equipped with residual connections and layer-norm. Let us denote by the input embeddings. Suppose the hidden size is d,, i.e. ( The attention layer is defined as [=]], = [[1 (z)]: ) + [Attn; (h(())))+ and the MLP layer is defined as [//+1 (x) LN((NN[(L ([r((()))) ( Uio (W] (2)]; + b1)). The activation o is GeLU. The final output is
[f(=)]= WF ( [x)]; Note that We is both the input embedding and the weight of the output layer.
They are tied in training.
The layer norm is on the feature dimension. [LN(2;)], = + By. L. is the normalized version of F, with zero mean and unit variance. y and are trainable.
The multi- head attention consists of n., self-attention heads. The definition of the multi-head attention Attny (hy (x)) = (z) T..., (Ain,, by (in. T O₁. The output matrix O₂ E Rd xdn, The attention heads composes of the attention score times the feature matrix times the value matrix. The attention score AiR E IRd' xd' is computed with softmax dot product. For each & E [Ma) Au- softmax ( (z) hT)
Following the implementation of Devlin al. (2018), the dimension of the attention head is always d' = 64, thus d, = 64ng. The dimension of the intermediate layer in the MLP is set to 4dh, which means Uj E dh Ads and W E R4da xds.
We aim to embed the smaller transformer f(x) into f(x), where d, = =2d,, nn = 2ng, and Ż- L+L+'.
Increasing feature dimension with replication of the parameters. 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
17

---

Under review as conference paper at ICLR 2023
functionality. Consider the following weight replication method. For 1 E O.....L- 11,
WE
F
-
n
b
bi 61
3
5t [t]
2 WF
WW Wi
U
U1
0
Oi
Wi Wi
U1
Wi]
U1
Ui
01
01
1 Qi = 21Quu Qutl, Kis = 2 [KiN, Kit), Vix= = Vike Vix] for k (u,.
Qu - 1 Quck- ) Kin = Ki(k Ki Ki(k-na)l+ Vi= 2 VV(A Vi(=-aa) for & [M +1 +1...,. 2mm)
We observe that the intermediate layers of the transformers are also replicated for the first L blocks, Lee
(x) and ūz(x))
(x) for (1,..., First note that [oo (2)) (W k...). Since
ha(x)= = (x)
uz (a)
ho (z)
replicating the features will not change the mean and the variance, we have (o))))
ho (x) Then
we can show that replicating the features will not change the attention scores as well. This makes
[x(()))
Finally note that we can apply the base case of the MLP to reason about the MLP
v((2))
Th441 (x) (141
(1)( (1((r) inductively.
layer, and show (x))
Therefore we have shown (x))))
(h(())
Adding additional layers using residual connections. We have demonstrated that h1(x))
(1)
for lE [0, Now let's consider the added L' blocks on top of the small model. Since the transformer contains residual connections, we can add new blocks on top of a small model and fill in zeros the
added parameters. We will show that in this way, hy = h (x). for any I [L,...,L++ L'J. This will
[h(x))
WF
indicate that h (z)==
Recall that We -
This indicate that
hi(z))
We
[f(x));= [:];::Wv]
W I W
(x)
(x)
2 --()))++((((0)
=[{())],
which means can add new layers top of small transformer without changing the functionality.
Now show that Uj 0 and O 0 for / [L.... + -' will make h((x) hz (z), for any I e LL....,- L + L'). This holds because [u)( = (N4( (h( and [h(+1 (x))e = LN([ur(=)]:
If +x)):
O, we have (z)) hz( from the first equation. If for any
U1 =0, we have h4 (xr) = (1)) from the second equation. Therefore have
1E [.].
B.1.3 VIEWING A SMALL TRANSFORMER AS SPECIAL CASE OF A LARGE TRANSFORMER
As demonstrated above, smaller transformers can be embedded into larger transformers with func- tionality preserved. The smaller transformer architecture can therefore be viewed as a subset of the larger transformer architecture. In this sense, a set of transformers with different sizes and the same pre-training loss found in Section 2 can be viewed as set of transformers with the same size after the embedding. Note that the training algorithm only finds out the natural larger models, instead of the larger models which are embedded from the smaller models. This indicates that the implicit bias of the optimizer can interact with the model architecture. The implicit bias drives the model toward flat minima both larger models and smaller models. The smaller transformer architecture is subset of the larger transformer architecture, thus the flattest minima found with larger transformer is flatter than the minima found with smaller transformer. (See Figure 6).
18

---

Under review as conference paper at ICLR 2023
C
OMITTED PROOFS IN SECTION
Proof of Lemma 3.2. We first recall loss
L(0) E.,, [-lo [x Br, B (x
Note that conditioned on any x x-,,,, e,t, it holds that
Ex,
Volog (fo [))e
Volfo(s- (Volool
+ +Bzilt
[f (x)
-0+ Bx,
[Volog (z-)))) (Volog(fo(<=-))l))
where in the last step, use the assumption that 1 I', that is, for all xt,, fo PP(- xx=,), which implies the following
[-
VāLfo(r-
Bay (t,m
Vilf(o= -V Σ [o 1). -V21=0
[fo(o=
Since 0 is a global minimizer of L, we have that Volog...elll =0. Therefore, we have that
=(0)=Er,,) Vabogfffee.Dlee//lllllttttt-l
E Volog (ell
= -E,, Bag a (o (ollgg//oo
-I,,-_ V Em Enng log ())
=(),(
which completes the proof.
19

---

Under review conference paper at ICLR 2023
D PRACTICAL IMPLICATIONS
Pre-training Algorithms. While we focus on the saturation regime in the paper as a controlled way 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). 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. Therefore, a practical direction to design better pre-training algorithms with more favorable biases which can lead to better downstream performance than Adam W and SGD.
common practice, the validation pre- pre- vre-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). However, Saunshi et al. (2022); Tay et al. (2021) show that pre- training loss is not necessarily correlated with downstream performance when comparing different architectures. We further show that pre-training loss may not always be a reliable indicator even for the same architecture. While downstream tasks could be used issue is that large language models are trained to be general multi-purpose models where the space of downstream tasks large and unknown during the time of pre-training. Thus from fundamental standpoint, it is beneficial to design more reliable indicator that is agnostic to downstream tasks.
Better Metrics for Language Models.
proxy metric for evaluation, the main
Explicit regularization. We show that implicit bias, especially the implicit bias of flatness matters for downstream performance in language modeling. Leveraging the implicit bias to design better explicit regularization in language modeling is also an important direction. Bahri et al. (2021) show explicit flatness regularization with SAM (Foret et al., 2020) boost downstream performance when applying to downstream tasks themselves and the intermediate stages between pre-training and fine-tuning, but they did not study this on pre-training, partly because SAM is not efficient enough for pre- vre-training (SAM requires back prop for 2 times per step, and more steps to reach the same level of pre-training loss (Foret al., 2020)).
20

---

Under review as conference paper at ICLR 2023
E RESULTS OF 15% MASK RATE
As we mentioned in Section and Section A.2, the true conditional probability of MLM Pr(x//11) can be computed from the joint probability
Pr(2;=c|(-d)= Pr(xt=
and the joint probability be computed efficiently with knowledge of the generative algorithms. The number of computation of the joint probability is the size of W.
However, when we have multiple masks, the time complexity of computing the true conditional probability can be extremely large. For example, if we have two masks in one sentence,
Pr(a, C, =d
Pr(x= -,00 = x Ex.'ew - C,, =
Now to compute the denominator, we will need to compute the joint probability for [W times. In general, the time complexity of computing the true MLM conditional probability is exponential in the number of masks.
Still, we can pre-train the model with 15% mask rate and evaluate the loss with one mask sentence. The results on PCFG is provided in Figure 11. Although in pre-training we use 15% mask rate, the validation loss evaluated with mask per sentence does not change much compared with training with one mask per sentence. The downstream performance does not change significantly with 15% mask rate either. The conclusion of Section 2 and Section 4 still holds.
PCFG, Task
PCRG
by % 16 1o Pre-traitt miingg Lowe
Frobe Task
Downstmam
Uncar prace: fectunes
fretusing
Preltsiinng Cen. Tiree of Hestlan
Predraaing
Nodel
Model Size
Figure 11: The result of pre-training with 15% mask rate and evaluating with mask per sentence on PCFG. Left: Larger models perform better downstream than smaller models, even with almost the same pre- pre-training loss. Right: The trace of Hessian correlates with downstream performance for models with different sizes and almost the same pre-training loss.
21

---

Under review as conference paper at ICLR 2023
F OMITTED PROOFS IN SECTION 5
F.1 OMITTED PROOFS OF THEOREM 5.1
Recall that the loss function of MLM is L(v, u) = Ex-rl(fe, 9" (z))))] In downstream adaptation, we have access to a finite dataset {x(*) sampled i.i.d. from Pa. The training loss is (1, > 12-i-1l00a. of (x(1)) 9" (x())))) and the population loss for the downstream task is L (H, -
Proof of Theorem 5.1. We first calculate the trace of Hessian of the pre-training loss and then derive lower bound for it in Lemma F.1. 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. This translates constant sample complexity for the downstream task.
Lemma F.1. Denote by hQ, (x)) - d.zz the output of the attention head. In the setting of Theorem 5. I+ fie [m] u, >0} 0} and I_ = {ie [m] u, < 0}. The trace Hessian can be lower
bounded,
/
Tr [
T
where the lower bound is achieved if and only if the following conditions are satisfied,
V Vie {x1 V, "e.hhhha (x)=0),
(3)
Yz, 1...
V2V,
(4)
1 (r,jE(T], a= T
(5)
Denote by D+ (x y 1) and D_ fx y= Ig is the set of index of neurons which is activated on 1, = ii [T] V. x 0}. By the condition in equation 5, the attention is taking the average of [x=]]__
We can map D+ and D_ to the feature space, H = {hQ x (x) y 1 and We first show that one neuron cannot be activated on inputs from both D.,
II = (hQ x (->))y=--l and D_, and all non zero neuron has to be activated on some input. Also note that a neuron cannot be activated on no input, unless the weight is O.
Fact F.2. (1) D+, I, I.. Similarly we have Vx LLC. (2) Suppose F, then there exists he - V, h.0.
Proof of Fact F.2. (1) Otherwise, suppose jEnnee-E,,,LL since y- = (V, l1 (x)) >0, there has to be j' E I NT, which contradicts the condition in equation 4. (2) Suppose v. Thh < 0 for all E H4UHH Then we have v T h= for all h, since v' < indicates v (-h) >0,0 0, and -h belongs to the support of P due to the symmetry of the distribution. However, in Lemma F.5, we show that the matrix stacking all input together has full row rank, thus v has be O, leading to a contradiction.
We have the following lemma characterizing the solutions achieving all the qualities in Lemma F.I. Intuitively, all the neurons can be divided into two sets, and each input can only activate neurons in one of the sets, leading the properties of the input distribution.
cancellation between activated neurons. This holds because of equation 4 and
Lemma F.3. Suppose Q,K.V satisfy the equality in Lemma F.I. For all i I_, on downstream data x, if (x) 1, we have V, hQ.K (x) c; C; is constant which holds for every x if g" (z) =1 (x) = we have V hQ. (x)==) -.
For all I+. on downstream data x, g" (x) -1, we have V. (T)))(())=)) C; is constant which holds for every if (x)) -0)=1. g (x) = we have V, ho =0.
22

---

Under review as conference paper at ICLR 2023
Now let us consider the downstream task. It suffices to consider the constant vector If samples satisfying 9 (a) and 9" (z) = -- both show up in the downstream dataset, the minimal norm solution is
mcz
+ and [mm[(44u1_-- Then can verify that
1 wr,
|00 2
f (z)
mc
(V, (z)x
Σ
mc;
2 (V,
hq.x(T))
l
c1
iET_
- mca
[g (z) (x)==1]
mea
ci
[g (x)=
| 11
VE7
c
2
=I[g"
"gg (x)
(x)
Therefore, L (1) Pa
[[/3 (2) - (z)) =0, which completes the proof.
We first show that when the pre-training loss equals 0, the trace of Hessian equals the square of the norm of the gradient.
Lemma FA. For any parameters the pre-training loss L(0) - Ex [ 00, the trace of =)
Hessian equals the square of the norm of the gradient,
-
Proof of Lemma F.4. We can express the Hessian as follows.
Since ([)=))))((())))))) we have with probability 1, (, 0 8" (fo(x)u))= 2 constant.
Proof of Lemma F.I.
Σ (f
(By Lemma F.4)
0eIQ.K,Vm)
2 E[ V. fo. (1
(6)
E
[12]
E
(V
h (e.
2 E
[Ma,x(a)]22
(8)
2E [1l (z) (2)f.
2V T
(9)
The equality in step 6 achieved if and only if the gradient of Q and K is 0. Equation (7) is from AM-GM, and the equality is achieved iff
V (x) =////////o. (x)
(
The equality in step 8 is achieved iff on all input, there no cancellation between activated neurons.
(e,iE4,,''EL, V VVvVTx5'
23

---

Under review as conference paper at ICLR 2023
Since the attention score aj satisfies o, 0 and Σ)= = and all embeddings xy in one masked sentence are orthogonal to each other with norm 1, we have [|)o.x(*)llz 2 - The equality is achieved iff a, = + for all a and all a je
Proof of Lemma F3. Suppose V; is a neuron with i I.. Then there exists H-, V; > O. Without loss of generality, suppose the masked position in h is 1, i.e. h1 0, h2 1. Now let us consider the components in V; corresponding to the input positions and the mask positions separately. vi(0) [.,,, Vi.r] and v,(0) = [V..T4.
We claim that V,(00 2:7" = c1 for some c 0 and v,(") is either 0 or c. To prove this, consider h. which is only different from & on the mask, h- 20a. Also consider -h and -h. Due to the symmetry of the distribution, -h and - h are in H.. By Fact F.2, V (-h) < and V. (-h) 50. V,T (-h) cannot be O, because this will leads to V; h-0.
= v,Th> 0, indicating that V,T (h-h) 2V/(2) 0. Now we show that V.(0) 2:T cc. Due to the condition of equation 3, we know that for any E H_ masked on the first token, either V. h= or they equal to the same positive value c for all h. We claim that V. = for any E H. Otherwise there exist H-,1 and H H H H_ = and H H- - = H (h +1). Vi h e for any ài E H_ and V. h 1 for any h H-.o- This cannot happen for 7 6. By Lemma F.5 we know that the matrix stacking all such h2.r together has
Case 1. V,, (-h) < O and V (-h) <.0 O. We have =
full row rank, thus v/(0
2:7 c1 for some e> and V (0)) =0.
Case 2. V. (-h) < 0 and V.T (-h)- =0, which indicates V.T /2 V.(P), Consider another h E H which is not equal to h and h/2 = Similarly we can find h', h' 2 Still we have (-N)<0 and V (-h') = 0, this tells us V. h/ = V, T h, due to the condition of equation 3. Applying this to different M's we have that V.T h equals the same positive value for all h H_ and the masked position is 0. By Lemma F.5 we know that the matrix stacking all such h2.r together has full row rank, thus ((()) 2.7 = cl for some c>> 0 and (P) =c
We have proved that v/(0) 2:T el for some <>> and V((P) is either 0 or c. We continue to show that vi(0) cl for the same e>> 0 and either v.(2) is 0 or its coordinates is tc.
For case consider h² H_ whose masked iosition 2, A²" = Also suppose that h" = By Fact F.2.
we know that (-Th") < and V.T (-A") 00, this implies that - ((()) S <V(P) 2 <v((e) the same argument above, we know that either v.(2) 2 or v. ((()) 2 - vV( 1. otherwise both V. N² >0
1- Applying
and V. A > 0 hold, and V,T A² * vVTh"" contradicting condition in equation 3. If ((()) from equation 3 we know V. VV;"V,Th which indicates v.(0) v-(e) =.. In this case
2 (((0)
find
another RM whose masked position is 2, h²" 1 but h"" = -1 Then V,T A" HV; h, contradicting equation 3. Thus (((P) * (((), 1- Similarly yー(r)
2 , v/()) 1. The only possible situation is y(()) = 0.
Applying the argument in this paragraph to other masked position, we have V.(P) =0. 0. Since H_ invariant under permutation, Vi(c) = and Vg(e) =cl.
For case 2, exactly the same argument as the above paragraph with the same h² and h shows that the coordinates of (p) is tc.
Therefore, we have shown that for any SnyieL. v/(0) =cl for same e> 0. The symmetry of distribution immediately tells us for any i E In,, v.() cl for c 0.
On the downstream distribution P., since there is no masked token, only v((0) is working. Since (e()) cl always holds, complete the proof.
24

---

Under review as conference paper at ICLR 2023
Lemma F5, Suppose M ((+++))((++)) is matrix composed =1. The first row of M is
[1... -. Define permutation p(1) 22 p(2) = p 11 For all 222, Mi,p(j) =
ktll's
k-l's Then the rank of is 2k
Mi-1..
MM+MMpi Proof of Lemma F.5. Note that M, + pp (2) 2e; for all e [2k++1], which means we can express the orthonormal basis as linear combination of the rows in M. Therefore, the rank of M is 2 O
25

---

Under review conference paper at ICLR 2023
F.2 THE EXISTENCE RANDOM FEATURE SOLUTIONS F.6
Theorem F... Suppose in 2 0 2 and V ~ (1,, for all iE [m]. With probability at least 1 over V, there exists vú ud, satisfying L( su) ) Se and || O(T")-1)
Proof of Theorem F... Since the number of possible input in pre-training is finite, we can invoke Lemma Bai Lee (2020) to show that random Gaussian features can the pre-training task.
Lemma F.7 (Lemma 9 in Bai & Lee (2020)). Suppose | V e~N(0,TI2r) There exists random variable (v) such that
Blo(u ho]]=
and a satisfies Ey,, =
Consider Q = K 0. In this case have hQ.k (z) EE training task, y VT T (x)) for all x. Since 2jj are norm orthogonal to each other, we have [lhq. (2) || 1
zy. Also note that for the pre-
T for all x. Now we can show that V; ~ (N,,TTTzr) E [m] independently the
random feature solution which can solve the pre-training task.
Suppose g(h) = Σ (V,TT)))(( h))(V,) and 9(R) (h) == 1 E oVVTTh)a(VV)) (()VVl) S
m
VTR). R is large enough such that Pr (supre(m) (V 2TR 2 1 - 1-8/2. We have g(h)
(RR) (h) on this event. Let 9iR) (h) be the truncated version of gh(h) - h. g (h) We have
(
Evl(((h) - (x)) _ m
R²T2
m
By Chebyshev and a union bound, we have
m( mf²
Pr -gx h
Fort= 5. we have m m2m(h))
=E.
<E.,((t)3))EE((00 (1 à
SCTPr(10|>>V)
Choosing R=Ō(((TT) will make Pr (101|>>/T)) S (loo)
Thus m suffices.
26


Reviews:
[{"review_id": "tmiuCLjKwP", "reviewer": "Reviewer_KN6A", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "the submission conducted experiments to study the impact of configurations on the transferability of a pre-trained model through the lens of Hessian matrices, and the studied configurations include the number of training iterations, the training objectives, and the size of a model. Unsurprisingly, the conclusion is that more training iterations, standard training objectives, and larger models would lead to better performance on the downstream tasks. ", "summary_of_the_review": "I don't recommend accepting the submission since the empirical results are well-known and the theorems don't provide theoretical justifications to the observations.", "strength_and_weaknesses": "----Strength----\n\n1. the submission found a correlated trend between the decay of the trace of the Hessian matrix evaluated on the training data and the increase of performance on the downstream tasks.\n\n2. given the above observation, it becomes straightforward to infer that a pre-trained model with parameters at a flat minimum would lead to decent performance on the downstream tasks.\n\n----Weaknesses----\n\n\n1. the trace of a Hessian matrix in the formulation from this paper is very easy to obtain and one doesn't even need to materialise the matrix. The issue is that the Hessian matrix is evaluated around the global minimum of a model being trained on the pre-training data, and, without any knowledge of the pre-training data, it becomes very difficult to compute the Hessian matrix. The other issue is that it does require many samples to obtain a reasonable estimate of the trace.\n\n2. the trace of the Hessian matrix is only an overview of the magnitude of all eigenvalues. A smaller trace could mean that the minimum might be flatter than other minima with the same loss, but it could also mean that the minimum might be steep in only a few directions, and completely non-informative in others. The latter case wouldn't necessarily give us a strong transferable model, and yet the trace of the Hessian matrix wouldn't be able to tell.\n\n3. maybe I am missing the point and I am happy to be corrected, but the arguments made in this submission or the observations obtained are already well-known empirically, especially in the case of zero-shot and few-shot learning with very large pre-trained transformer models. In addition, the theorems mentioned in the submission don't seem necessary, or, in any case, support the arguments or the observations.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "eYYcLlXeZB", "reviewer": "Reviewer_6MPN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper addresses the issue of pre-training loss cannot fully explain downstream performance, they instead claim that the flatness of the model is well-correlated with downstream performance. The author showed that at the same level of pretraining loss, large models have better downstream performance, because of the flatness of large models. The authors conducted experiments with PCFG/HMM/OPT-generated data as downstream tasks. The authors also provided theories to support their claims.\n", "summary_of_the_review": "The paper is good in terms of its qualities, but there are some problems that need to be answered to further prove the correctness of their claims. I would tend to accept this paper if they can give good answers to my questions. For now, I will give a borderline score.", "strength_and_weaknesses": "Strength:\n- The paper is well-written and easy to follow.\n- The paper is supported by many experiments to verify the claims that the authors proposed.\n- Besides empirical results, the author also provided theoretical points of view to support their claim.\n\nWeaknesses:\n- The claim that flatness can decide the downstream performance is not well-supported. There are many factors that can affect the downstream performance, for example, the model size. One possible explanation is that the flatness could be a consequence of scaling up the model size, and the good performance is brought by the large model size, too. If that is the case, then flatness is not the reason for good downstream performance. The authors need to have more evidence to prove that flatness is the main reason that leads to good performance, otherwise, they may be both the consequence of another factor (like model size). \n- When computing Hessian, I assume that you are using the loss of pretraining datasets, not downstream datasets. Then how can the flatness on the pretraining task reflect the situation on downstream tasks? This claim is not very intuitive. It may need more explanations.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OE53br0FL4b", "reviewer": "Reviewer_8JLN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Practitioners tend to gauge the downstream performance of large language models by comparing their pretraining losses. However, this paper shows that when the pretraining loss is near convergence, downstream performance can vary depending on the training methods, despite near identical pretraining loss. While there isn’t necessarily a correlation between the pretraining loss and downstream performance, the “flatness” of the solution, as characterized by the trace of the Hessian, does correlate with downstream performance. This observation is then formalized and proved on a synthetic language task.", "summary_of_the_review": "Interesting empirical observations and theoretical result. However, more empirical results on real data, especially on how the discovered insight can enable better downstream performance, would strength this work given the rather empirical motivation.", "strength_and_weaknesses": "This work seeks to better understand an important question in representation learning: how does pretraining performance correlate with downstream performance?\n\nThis question is adequately answered in controlled settings where we observe the non-correlation between pretraining loss and downstream performance and the correlation between solution flatness and downstream performance. The investigation is both well-motivated and nicely executed. In addition, a theoretical result is provided in support of the empirical observations.\n\nIt is, however, unclear what the practical implications of this work are. First of all, current large language models are not in the saturation regime, and it is hard to estimate when they will be as datasets grow with model size in tandem. Second, while this paper points that pretraining loss is not a reliable indicator of downstream performance, a simple remedy is to evaluate on downstream tasks during pretraining and compare models accordingly, which is likely already done in practice. Finally, this paper does not demonstrate if the insight gleaned in this work can lead to additional “flatness regularization” that induces better downstream performance on real datasets. It is understandable that large-scale experiments are expensive and are not expected, but given the rather empirical motivation of the paper, some validation on real data seems desirable.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Nieybt9GPY", "reviewer": "Reviewer_mwuY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper questions the common sense that correlation between pre-training validation loss such as masked language modeling loss and generalization performance of downstream tasks after fine-tuning. It empirically shows that the models with the same pre-training loss can achieve varying test accuracy on downstream tasks. Instead, the models converged to flat minima obtain better test accuracy.", "summary_of_the_review": "As previously mentioned, the authors do not properly tackle the open question --- flat minima and generalization which might still seem to be controversial topic. Moreover, they need to tone down their claim since their theoretical analysis is limited to very specific controlled experiment which is far from realistic scenario. Thus, I am inclined to reject.", "strength_and_weaknesses": "\n## Strength\n- The paper explores the interesting and practical research question --- correlation between pre-training loss and generalization performance of downstream tasks.\n\n\n- In the controlled experiments, the authors empirically show the correlation between flatness and generalization. Even though language model achieves optimal pre-training loss, we can improve generalization performance of downstream task as we continue pre-training, which results in flatter minima.\n\n- Authors provide some theoretical analysis that connects flatter solution to the pre-training loss and generalization performance of downstream tasks in the limited settings. \n--- \nPlease address my following concerns. If they are properly addressed, I am willing to raise my score.\n## Weakness\n\n- My expertise not learning theory, but in my humble opinion, what the authors claims is too bold. They claim that they \"prove that SGD with standard mini-batch noise implicitly prefers flatter minim in language models\", but it is limited to the special case ---  Dyck language experiment which is far from practical scenario.  It would be better to tone down the claim.\n\n- I am a bit confused about the correlation between flatness and generalization. For example, Dinh et al., [1] show that we can build equivalent models but with sharper minima, which might contradict to  what the authors claim, specifically theorem 5.1. However, the authors did not properly tackle this issue even in the related work section. \n\n\n- It is not clear how the size of models and flat minima is related to each other. The authors claim that smaller transformer architecture is a subset of the larger transformer architecture family, but it does not make sense. If we increase the width of transformer or increase the vocabulary size, the smaller model cannot be a special case of larger model. It would be better to explain why larger model converges to flatter minima than smaller one.\n\n- In the experiments, it is unclear why the authors compare vanilla sgd training with adversarial training. I think adversarial weight perturbation [2] is more relevant baseline since it explicitly regularize the flatness of weight loss landscape.\n\n- Minor typo: In page 4, (2) training for different number of steps -> (1) training for different number of steps\n\n\n## Questions\n\n- Regarding the flat minima, what happens if we explicitly enforce the language model to converge flatter minima? For example, we can explicitly regularize the trace of Hessian to be small or use SAM [3] optimizer to enforce such regularization.\n\n- Why is the cross-entropy loss is typically non-zero at the global minimizers?\n\n- Why do you assume that we only mask a single token in a sentence for masked language model? In practice, we usually mask 15% of tokens  of a sentence (e.g. BERT [4], RoBERTA [5]).\n\n## References\n[1] Dinh, Laurent, et al. \"Sharp minima can generalize for deep nets.\" International Conference on Machine Learning. PMLR, 2017.\n\n[2] Wu, Dongxian, Shu-Tao Xia, and Yisen Wang. \"Adversarial weight perturbation helps robust generalization.\" Advances in Neural Information Processing Systems 33 (2020): 2958-2969.\n\n[3] Foret, Pierre, et al. \"Sharpness-aware Minimization for Efficiently Improving Generalization.\" International Conference on Learning Representations. 2021.\n\n[4] Devlin, Jacob, et al. \"BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.\" Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). 2019.\n\n\n[5] Liu, Yinhan, et al. \"Roberta: A robustly optimized bert pretraining approach.\" arXiv preprint arXiv:1907.11692 (2019).\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject