Implicit Bias of Next-Token Prediction
Christos Thrampoulidis
Next-token prediction (NTP), the go-to training paradigm in training large
language models, involves predicting the next token in a sequence. Departing
from traditional one-hot classification, in NTP, multiple tokens with varying
frequencies follow each given context. This work frames NTP training as
cross-entropy minimization over distinct contexts, each associated with a
sparse empirical probability vector across a finite vocabulary. It then
addresses the following question: do gradient-based optimizers exhibit a bias
towards solutions with specific structure as the NTP training loss reaches its
lower bound (entropy)? Specifically, for linear NTP models trained using
gradient descent (GD), we make the following contributions: Firstly, we
determine NTP-separability conditions on the data, under which GD can attain
its lower bound. We also demonstrate that these conditions hold under
overparameterization. Secondly, we establish that the parameters of GD
projected onto an appropriate data subspace converge to the unique solution of
a system of linear equations, which requires the logits' difference of
in-support tokens to be equal to the log-ratio of their respective
probabilities. Meanwhile, on the orthogonal subspace, the parameters diverge
and converge in the direction of the solution of a max-margin quadratic
program, minimizing the Euclidean norm of parameters satisfying the
\NTP-separability conditions. Akin to prior research on implicit bias of
one-hot classification, our work opens exciting avenues for future research
that can lead to better understanding optimization, generalization and
robustness principles of models trained with NTP.