Source-linked AI summary
Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models
Hong Liu, Sang Michael Xie, Zhiyuan Li, Tengyu Ma
TL;DR
Validation pre-training loss is often used to assess language models, yet its relationship with downstream performance leaves open why models at similar loss can transfer differently. The paper studies this implicit bias through simplified-data comparisons and theory, finding that flatness tracks transfer and that SGD favors flatter minima.
Problem
Although validation pre-training loss is commonly used to compare language models, different parameter configurations can share a loss and still have different downstream performance.
Source: §1 · p. 1Method
The paper compares matched-loss models across training time, model size, and training algorithms, and analyzes flatness theoretically and empirically.
Source: §1 · p. 2Results
At similar pre-training loss, later checkpoints, standard training, and larger models can perform better downstream, while SGD favors flatter minima.
Source: §1 · p. 2Takeaways & Limitations
In these settings, pre-training loss alone does not determine transferability, and flatter solutions are associated with better downstream performance.
Source: §1 · p. 2The experiments use simplified datasets and in-distribution downstream evaluation, while reaching saturation on real large-scale data is computationally challenging.
Source: §3.1 · p. 5, §7 · p. 12
Abstract
from arXiv · showhide
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.
1 Introduction
Validation pre-training loss is commonly used to compare language models, but equal loss—especially near the minimum—does not guarantee equal downstream performance. The paper links this variation to implicit bias toward flatter, more transferable solutions.
1 Introduction
- Validation pre-training loss is commonly used to monitor training and compare models because it generally correlates with downstream performance.Prior analyses often treat language models as black boxes and focus on fitting the conditional probability.Source: §1 · p. 1
- With the same pre-training loss, later checkpoints, standard rather than adversarial training, and larger models can each perform better downstream.These differences are most prominent near minimal pre-training loss.Source: §1 · p. 2
- The first two matched-loss comparisons use models with the same architecture, implicating training algorithms in their different downstream performance.For larger models, the authors also describe an equally good hypothetical solution that copies or zero-pads a smaller model, which training does not select.Source: §1 · p. 2
- SGD with standard mini-batch noise prefers flatter minima of population pre-training loss among minima in the saturation regime.The mini-batch noise does not vanish at convergence.Source: §1 · p. 2
- At equal pre-training loss, flatness measured by the Hessian trace strongly correlates with downstream performance across the three settings.The paper reports the same relationship for models evaluated by linear probing on contextualized representations.Source: §1 · p. 2
- In a simplified Dyck language setting, the flattest among models with good MLM pre-training loss learns the most useful downstream features.The authors note that supervised generalization-bound arguments do not directly apply to this language-modeling setting.Source: §1 · p. 2
2 Related Work
Related work connects language-model performance to scale, learned representations, architecture, and self-supervised objectives. This paper's distinct focus is how implicit bias can produce different transfer performance among same-architecture models with the same pre-training loss.
Language modeling and downstream adaptation
Understanding large language models
- Prior work finds that fitting language-model conditional probabilities can support downstream performance, while this paper reports additional architecture- and optimization-related effects.The paper also reports linear probing on language-model features outperforming linear probing on the true conditional probability.Source: §2 · p. 3
- Earlier studies show architecture-related downstream differences at similar pre-training loss; this paper examines implicit bias within the same architecture.Its evaluation focuses on fine-tuning and linear probing, unlike the cited in-context-learning comparison [62].Source: §2 · p. 3
Understanding self-supervised learning
Implicit bias in supervised learning
3 The Existence of Implicit Bias in Language Modeling
The experiments show that models with similar pre-training loss can differ in downstream performance, depending on model size, training duration, and algorithm. These differences indicate implicit bias toward more transferable configurations among models with the same architecture and loss.
3.1 Formulations
- Saturation means the model outputs the true conditional probabilities, giving all models the same optimal pre-training loss; downstream evaluation uses fine-tuning or linear probing.The MLM model predicts a masked token from the rest of the sentence.Source: §3.1 · p. 4
3.3 Results
- 6.9%, 4.5%, and 2.0%: at the same pre-training loss, scaling improved linear-probe performance on PCFG, HMM, and OPT, respectively.Figure 2 covers PCFG task B, HMM task-10, and OPT QNLI.Source: Figure 2 · p. 5, §3.3 · p. 6
3.1 Formulations
- Reaching saturation on real large-scale data is currently computationally challenging; the experiments instead use simplified datasets.Source: §3.1 · p. 5
3.2 Experimental Setup
- Controlled experiments use PCFG-, HMM-, and OPT-generated datasets with known conditional probabilities, focusing mainly on in-domain evaluation for PCFG and HMM.The study varies model size, training steps, and training algorithms.Source: §3.2 · pp. 5–6
3.3 Results
- After pre-training loss converged and stopped improving, downstream accuracy continued increasing across model checkpoints.The comparison tracks checkpoints along pre-training.Source: §3.3 · p. 6
- On PCFG task B, adversarially trained 235M transformers were more than 6% worse than normally trained 235M models and worse than normally trained 9M models.Table 1 compares different pre-training algorithms on PCFG.Source: Table 1 · p. 6, §3.3 · pp. 6–7
- A lookup table with perfect pre-training loss performed worse on task B than normally trained transformers; linear probes also favored standard-transformer representations over true conditional probabilities.The authors conclude that fitting conditional probabilities well does not fully explain language-model performance.Source: §3.3 · p. 7
- Across matched-loss models in the same architecture family, training algorithms, model sizes, and optimizer steps can affect downstream performance, indicating implicit bias toward more transferable configurations.Source: §3.3 · p. 7
4 Implicit Bias Leads to Flat Solutions in Language Modeling
The section argues that implicit bias in language modeling selects solutions more adaptable to downstream tasks, rather than closing a pre-training loss gap. It analyzes how mini-batch SGD favors flatter global minimizers, measured by Hessian trace.
4 Implicit Bias Leads to Flat Solutions in Language Modeling
- Because the population and empirical pre-training losses are close, the paper frames implicit bias as benefiting downstream adaptability rather than bridging that gap.The paper identifies selecting networks more adaptable to downstream tasks as the benefit.Source: §4 · p. 7
- The analysis uses population cross-entropy with fresh samples and starts at a global minimizer, assuming expressive models with connected, smooth minimizers.It presents batch size 1 for simplicity and states that results generalize to arbitrary batch size.Source: §4 · p. 7
- At global minimizers, cross-entropy typically remains nonzero, so gradient noise persists; its covariance equals the Hessian and steers movement along the minimizer manifold.For well-specified probabilistic models at a global minimizer, both the gradient covariance and Hessian equal the Fisher information matrix.Source: §4 · pp. 7–8
- Figure 3 compares Hessian trace and downstream performance across post-convergence checkpoints on PCFG/task C and HMM/task-10.The caption reports a correlation between trace and downstream performance for both panels.Source: Figure 3 · p. 8
- As η→0, SGD at step K/η² converges in distribution to an ODE trajectory that locally decreases Hessian trace along the global-minimizer manifold.The ODE is projected gradient descent on Hessian trace constrained to the manifold of global minimizers.Source: §4 · p. 8
- For arbitrary batch size B, the corresponding ODE coefficient is 1/(4B), replacing 1/4 in the batch-size-1 equation.The paper states the arbitrary-batch extension after presenting the batch-size-1 result.Source: §4 · p. 8
5 Flatter Models Have Better Downstream Performance
The section links lower Hessian-trace flatness with better downstream performance, including after pre-training loss convergence and across model sizes and algorithms.
5 Flatter Models Have Better Downstream Performance
- The experiments compare checkpoints over training time, model sizes, and standard versus adversarial training algorithms.For model-size comparisons, differently sized transformers are viewed as parameter configurations within a larger architecture with the same pre-training loss.Source: §5 · p. 8, Results · p. 9
- Flatness is measured by the trace of the pre-training-loss Hessian, with a smaller trace indicating a flatter minimum.The trace is estimated with random samples for computational feasibility.Source: §5 · p. 9
Results.
- As the Hessian trace decreased after validation pre-training loss converged, downstream performance improved by 1.6% on PCFG and 4.0% on HMM.The trace showed a clear decreasing trend on both datasets.Source: Results · p. 9
- Adversarially pre-trained models had twice the Hessian trace of standard AdamW models, corresponding to a 5.5% drop in downstream performance.Table 2 evaluates a 235M transformer on PCFG Task C.Source: Results · p. 9, Table 2 · p. 9
- On PCFG, among models larger than 9M with almost the same pre-training loss, increasing model size reduced Hessian trace from 19.8 to 12.6 as linear-probe accuracy rose from 40.4% to 50.5%.On OPT-generated data, increasing model size likewise coincided with higher linear-probe accuracy and sharply lower Hessian trace.Source: Results · p. 9
- The implicit bias favors flatter minima at both model sizes, while the larger architecture family can attain flatter minima than the smaller family.The smaller transformer architecture is a subset of the larger one.Source: Figure 4 · p. 9, Results · p. 10
6 Flatness Regularization Provably Identifies Transferable Models on Synthetic Language
In a synthetic Dyck-language setting, the paper analyzes models that achieve good masked-language-modeling loss but differ in whether they learn transferable structure or fit through random features. It argues that the flattest solution is the transferable one.
emphasis
- The synthetic analysis compares a solution that learns transferable features with “cheating” solutions that memorize the pre-training conditional probability using random features.The setting permits a sharp analysis of a single-layer transformer for masked language modeling.Source: §6 · p. 10
- The Dyck-language pre-training data are balanced-bracket sentences with one randomly masked position; the downstream task counts bracket mismatches in sequences that may be invalid.The downstream evaluation samples each bracket uniformly and independently.Source: §6 · p. 11
- The input encoding uses signed position-specific one-hot vectors, and attention can sum token encodings into a representation from which two neurons predict the masked token.The transferable construction uses two neurons to recover the masked token from the summed representation.Source: §6 · p. 11
- A random-Gaussian-feature solution can also predict the masked token correctly when its width is sufficiently large, despite not representing the downstream structure.Its output coefficients express the target function using the random features.Source: §6 · pp. 10–11
- The transferable two-neuron solution has a smaller Hessian trace because it avoids cancellation among simultaneously activated neurons, unlike the random-feature solution.The paper links minimum Hessian trace to the absence of such cancellation.Source: §6 · pp. 11–12
- The paper formalizes the comparison by minimizing Hessian trace among zero-pre-training-loss solutions and then adapting the resulting representation to the downstream task.The theorem setup assumes m ≥ 2 and T ≥ 6.Source: §6 · p. 12
7 Conclusion
The conclusion argues that implicit bias matters beyond pre-training loss and connects flatness regularization with downstream performance. It also identifies the study’s evaluation scope and directions for future work.
7 Conclusion
- The paper concludes that implicit bias matters beyond pre-training loss and explores its mechanism in language modeling.Its theory provides cases where flatness regularization can determine downstream performance.Source: §7 · p. 12
- Downstream evaluation focuses on in-distribution data to control for other factors; studying implicit bias on out-of-distribution tasks remains future work.Source: §7 · p. 12
- Experiments use simplified datasets because of computational-resource constraints, and the authors hope the phenomenon predicts implicit bias on more complex datasets as models scale.Source: §7 · p. 12
A.1 Generating Simplified Datasets
The appendix describes three generated datasets used for downstream tasks: PCFG, HMM, and OPT-generated text.
A.1 Generating Simplified Datasets
- The PCFG dataset uses vocabulary size 200, 50 non-terminal states, sentences up to length 32, and 2×10^7 generated sentences; downstream tasks classify non-terminals over spans of length 32, 16, or 8.Each task contains 0.1M examples.Source: §A.1 · p. 17
- The HMM dataset uses vocabulary size 200, 100 hidden states, length-16 sentences, and 1×10^7 generated sentences; tasks classify the sixth or tenth hidden variable.Each task contains 0.1M examples.Source: §A.1 · p. 17
- The OPT-generated dataset samples sentences from a 125M OPT model restricted to its top 2,000 tokens, with length capped at 24 and 2×10^8 sentences generated.The generated text examples are illustrated in Figure 8.Source: §A.1 · p. 17, Figure 8 · p. 17
A.2 Compute the True Conditional Probabilities
The appendix describes how the experiments compute true conditional probabilities, configure models, and evaluate downstream performance and pre-training loss.
A.2 Compute the True Conditional Probabilities
- True conditional probabilities are computed from the generative model’s joint probabilities, using the inside algorithm for PCFGs and the Viterbi algorithm for HMMs.The appendix also describes computing the true MLM conditional probability from the joint probability.Source: §A.2 · p. 17
A.3 Models
- Transformers are trained on PCFG- and OPT-generated datasets, with model dimensions and layer counts listed across the considered architectures.The setup uses learning rate 1e-3 and warmup proportion 0.06; LSTMs use the PyTorch implementation.Source: §A.3 · p. 18, Table 3 · p. 18
A.4 Algorithms
- Downstream evaluation uses a lookup-table baseline, an adversarial meta-learning algorithm, fine-tuning, and linear probing.The adversarial algorithm trains a downstream head on one subset, then updates representations to maximize validation loss on another; pre-training loss is evaluated using true conditional probabilities and 10^6 sentences.Source: §A.4 · p. 18
A.5 Results on Other Downstream Tasks.
Downstream performance improves with model size even when pre-training loss remains unchanged, including across several downstream tasks.
A.5 Results on Other Downstream Tasks.
- Downstream accuracy continues to rise as model size increases while pre-training loss remains unchanged.The pattern is also reported for PCFG Task A, OPT SST-2, and HMM Task-6 in the saturation regime.Source: Figure 9 · p. 19, §A.5 · p. 19
A.6 Evaluation of Pre-training Loss with KL Divergence.
The appendix reports that saturation remains visible under KL-divergence evaluation and details methods for estimating flatness and comparing models across sizes. It also describes functionality-preserving embeddings of smaller models into larger architectures.
A.6 Evaluation of Pre-training Loss with KL Divergence.
- Saturation remains observable during continued training and model scaling even when log KL divergence is used as the pre-training-loss measure.KL divergence reduces loss-evaluation variance when comparing models in the saturation regime.Source: §A.6 · p. 19, Figure 10 · p. 19
B.1 Unbiased Estimate of the Trace of Hessian.
- The trace of the Hessian is estimated unbiasedly by sampling contexts and then sampling predicted tokens for each context.The experiments use 10,000 contexts and 50 token samples per context.Source: §B.1 · p. 20
B.2 Details in Figure 3.
- SGD checkpoints are evaluated for downstream performance and Hessian trace on PCFG- and HMM-generated datasets.The appendix reports results across PCFG tasks A, B, and C and HMM tasks 6 and 10; Figure 11 describes the trace as correlating with downstream performance across checkpoints after convergence.Source: Figure 11 · p. 20, §B.2 · p. 20
B.3 Details in Figure 5.
- Comparing Hessian traces across transformer sizes requires embedding smaller models into larger ones, with the largest width set to a least common multiple of smaller widths.The evaluation uses linear probing under the Section 3 setting; Table 4 lists transformer shapes for Figure 5.Source: §B.3 · p. 20, Table 4 · p. 21
B.4 Embedding of a Smaller Transformer into a Larger Transformer.
- Smaller transformers can be embedded into larger ones without changing functionality by replicating weights to expand features and attention heads and using residual connections to add layers.The appendix develops this embedding first for MLPs and then for transformers with layer normalization and residual connections.Source: §B.4 · pp. 20–21, §B.4.1 · p. 21, §B.4.2 · p. 21
B.4.1 The Base Case with MLPs.
- The MLP construction embeds a network with doubled hidden dimensions while requiring its output to equal the original for every input.The appendix defines the original and enlarged multilayer MLPs and states the functionality-preservation condition.Source: §B.4.1 · p. 21
B.4.2 Real Transformers.
- For transformers, replicating weights preserves layer-normalization statistics and replicated representations, while zeroing added-block parameters lets residual connections add layers without changing functionality.The transformer setup and zero-parameter construction are described in the appendix.Source: §B.4.2 · pp. 21, 23
B.4.3 Viewing a Small Transformer as a Special Case of a Large Transformer
- The smaller architecture is a subset of the larger one, and the optimizer’s implicit bias can interact with architecture as it favors flatter minima in both.The appendix states that the flattest minima found in the larger architecture are flatter than those found in the smaller architecture.Source: §B.4.3 · p. 23
C Omitted Proofs in Section 4
The omitted proof invokes the true-conditional-probability assumption and the zero-gradient condition at a global minimizer.
C Omitted Proofs in Section 4
- The proof assumes parameters in Γ predict the true conditional probabilities and uses the zero gradient at a global minimizer.These steps are presented as part of the proof of Lemma 4.2.Source: C Omitted Proofs in Section 4 · p. 24
D.1 Omitted Proofs of Theorem 6.1
The proof connects pre-training Hessian flatness conditions to downstream transfer, showing that equality in a Hessian lower bound entails attention outputs in one direction and constant downstream sample complexity. It also establishes that random Gaussian features can fit the finite pre-training task.
D.1 Omitted Proofs of Theorem 6.1
- Equality in the Hessian lower bound requires attention outputs in one direction on all downstream inputs, a condition the proof translates into constant downstream sample complexity.Source: §D.1 · p. 25
- Inputs from D+ activate only I+ neurons and inputs from D− only I− neurons, preventing cross-label activation and cancellation.Source: §D.1 · p. 25
- When the downstream dataset contains both labels, the stated minimal-norm adaptation solution achieves zero population downstream loss.Source: §D.1 · p. 26
- At zero pre-training loss, the trace of the Hessian equals the squared norm of the gradient.Source: §D.1 · p. 26
D.2 The existence of random feature solutions D.6
- The appendix establishes that random Gaussian features can fit the pre-training task, using the finite number of possible pre-training inputs.Source: §D.2 · p. 29