Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models

TL;DR

This paper challenges the reliance on pre-training loss as a sole indicator of downstream performance in language models, highlighting the importance of model flatness and implicit bias of training algorithms in transferability.

Contribution

It demonstrates that models with identical pre-training loss can differ significantly in downstream performance due to implicit bias and flatness, supported by theoretical and empirical analysis.

Findings

Pre-training loss does not fully predict downstream performance.

Model flatness correlates strongly with transferability.

SGD implicitly favors flatter minima in language models.

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.