Learning Is Not A Race: Improving Retrieval in Language Models via Equal Learning (2025)

Abstract: Many applications that modern large language models (LLMs) are deployed on are retrieval tasks: the answer can be recovered from context and success is a matter of learning generalizable features from data. However, this is easier said than done. Overparametrized models trained on cross-entropy loss can overfit on noise. We argue that such overfitting is prone to happen when the model can identify mechanisms that rapidly drive down the loss of certain tokens early on in training. Fitting some tokens early reduce gradient signals in later iterations, as such, remaining tokens are more vulnerable to noise overfitting. We dub this phenomenon unequal learning and show that LLMs with longer contexts or larger embedding sizes are prone to this failure mode. In this work, we argue that learning training samples at an equal rate helps counter such biases. We highlight two mechanisms that promote equal learning: (i) loss functions that regularize uniform margins across training samples, (ii) small learning rates (e.g. by warming up) at the start of training. We demonstrate these approaches on various synthetic and natural language datasets.

## introduction if it looks like a duck, swims like a duck, and quacks like a duck, the maximum likelihood principle tells us we should expect -in all likelihood -a duck. this same principle underpins modern large language models (llms), which are largely transformer-based @xcite neural networks p θ trained to carry out next-token prediction by minimizing cross-entropy loss l(θ) = e x 1:t ∼ptrain t log p θ (x t |@xmath0. one key task that these models are often deployed on is retrieval from context. question answering @xcite , summarization @xcite , and sentiment analysis @xcite are amongst the many downstream language tasks whose success relies on the model's ability to retrieve and output the correct piece of information from context when given the appropriate prompt. when is a duck not a duck? in this work, we crystallize the insight that retrieval is not a corollary of doing language modeling well. llms that can achieve zero loss on the training objective do not necessarily learn generalizable features. accordingly, a well-trained transformer can fail to retrieve the correct information from context even when evaluated on samples drawn from the same distribution. let us briefly survey how such failure modes might arise. llms are overparametrized models and can learn multiple solutions from finite data. the solution (local optimum) that the model lands on depends on inductive biases arising from the complex interplay between factors such as model architecture, loss function, and optimization hyperparameters @xcite . these biases can cause the model to overfit, learning erroneous solutions that achieve zero loss on training samples but do not generalize at test time. recent work highlights one such bias: the ubiquitous cross-entropy loss. linear models trained on crossentropy learn a maximum-margin solution. to arrive at such a solution, the model provably exploits imperfect predictors ("shortcuts") to correctly classify some training samples -as long as enough noisy features exist to overfit the remaining samples @xcite . we show that llms are susceptible to the same pitfall. modern llm design increasingly favors long context lengths @xcite @xcite . however, when trained on cross-entropy, a transformer's ability to retrieve the correct content quickly erodes as context length increases. we observe that simply reducing batch size by squeezing samples into the same context window (thus keeping the overall number of samples per batch constant) can cause the model to overfit by increasing the amount of noisy features available. biases can arise from other design choices. for example, llms are commonly trained via next-token prediction, i.e. teacher-forcing. this process can itself induce bias @xcite . if earlier tokens in the response reveal information about the latter tokens in the sequence, a next-token predictor will inevitably leverage this correlation. this creates failure on "lookahead" tasks where the correct predictor requires planning later tokens ahead of earlier tokens. how, then, do we encourage our models to learn generalizable features for retrieval? in this work, we argue that an useful auxiliary objective is equal learning of the individual conditional distributions p train (x t |@xmath1. in the failure modes identified above, the model exploits correlations between tokens to quickly drive down the loss -that is, it learns to model certain parts of the sequence (or certain samples altogether) faster than others. this differential in learning is precisely how the model quickly learns a shortcut mechanism for part of the sequence. once those terms have been fitted, they no longer contribute to the batch gradient signal, allowing the model to more easily fit the remaining terms with noise. to combat this tendency, we argue that enforcing uniform learning is a useful inductive bias for learning generalizable features -in other words, all tokens should be learnt at an equal rate. how do we teach the model to learn all terms uniformly? we propose explicit regularization, by designing a family of loss functions that regularize for small margins. as the regularization term contradicts the cross-entropy term (which prefers maximizing margins), optimizing the overall objective requires the model to learn a fixed value for the margin and prevents it from fitting selective tokens too early on in training. furthermore, we recommend warming up the learning rate schedule. keeping learning rates small at the start of optimization dampens the model's ability to rapidly fit certain parts of the sequence. this allows all terms in the sequence to contribute to the gradient signal for more iterations, preventing the model from noise overfitting. we show that our approach can learn generalizable features on a number of diverse tasks that a standard cross-entropy-trained transformer fails on. we summarize our contributions as the following: feature learning and biases in transformers as the core architecture underpinning modern llms, understanding how transformers @xcite learn is a key research topic and much existing work has sought to precisely identify and describe the internal mechanisms that govern feature learning @xcite @xcite . more specifically, recent work has also studied biases that occur in transformers @xcite @xcite . our work broadly falls in this category, however, we specifically focus on improving feature learning for retrieval tasks where the correct mechanism can be learnt from context. we discuss a few closely-related papers below and how our approach differs from theirs: bachmann and nagarajan (2024) show that nexttoken prediction itself is a paradigm that can learn to learning features that do not generalize. more specifically, teacher-forcing can fail on "lookahead tasks", which are tasks that require the model to compute future tokens ahead of earlier ones in order. they demonstrate this insight on a constructed graph search problem, where the prefixes of the response are strongly correlated to the rest of response. as such, a model trained via teacher-forcing can achieve zero training loss simply by conditioning on earlier tokens, but will fail to generalize at test time. this work's chief contribution lies in identifying the problem; the authors provide a bespoke solution to the graph problem that relies on prior knowledge of the true solution. instead, we propose black-box solutions that do not require domain knowledge. focusing on shortcut learning, nagarajan et al. (2021) and puli et al. (2023) study a similar setting where the label is a deterministic function of the inputs, however, a shortcut mechanism exists that is strongly correlated to the label under the training distribution only. they show that empirical risk minimization (erm) with the crossentropy loss will prefer to learn the shortcut -provably so in the case of linear classifiers. both works attribute this phenomenon to the erm's bias towards learning a maximum-margin solution @xcite , and as a fix @xcite propose losses to control the margin for binary classification. our work extends their findings to the transformer architecture, showing that the same phenomenon can be observed in transformers under aggravating conditions such as long contexts. we relate the max-margin bias to unequal token learning and implement a set of solutions that operate over multiclass discrete sequences. lead to high generalization error. as such, the crossentropy loss can provide poor inductive biases. we first formalize the problem setting. we assume that there is a mechanism (function) that deterministically maps a context to the correct output tokens. formally, letting x, y denote the context and the output tokens, the conditional distribution of the output given context is a delta function: ## equal learning how can we promote the equal learning of terms in (1) during training? the idea is to constrain the model from learning specific terms too fast, especially during the early iterations of training. intuitively, we want to ensure that all tokens contribute to gradient signals for as long as possible during training, which makes noise overfitting harder since the model would need to overfit more tokens. our approach intervenes on the training loss by regularizing for this additional objective. ## experimental results in this section, we demonstrate the validity of our approach on two large-scale datasets: (i) the path-star graph task originally introduced in bachmann and nagarajan (2024), and (ii) two natural language texts on which we devise copy tasks susceptible to the same pitfalls introduced in section 3. experimental details can be found in appendix b. ## discussion as llms gain widespread traction and usage in society, understanding how and when they can fail is crucial to safe and robust performance. we believe that our work: (i) identifies an important pitfall that modern transformer-based llms are susceptible to -namely, overfitting due to unequal learning, (ii) provides empirical evidence for this problem, and (iii) proposes a working solution in the form of new loss functions.