Information theory explains generalization in predicting next tokens
Information-Theoretic Analysis of Next-Token Prediction under Markovian Data
Information TheoryMachine Learning
Summary
Predicting the next word or token in a sequence is harder when past data depends on earlier parts, like in stories or time series. The authors create a math framework that separates the effect of the learning algorithm from the natural time dependence in data. They find ways to estimate how well a model trained on some data will perform on new data, especially when the model sees long contexts. They also show that very long contexts don’t always improve predictions in practice. Their work helps clarify how memory length, model size, and data dependencies affect prediction accuracy.
What this means in practice
- •For machine learning engineers: Design next-token prediction systems that balance context length and model complexity to better predict performance on time-dependent data.
- •For natural language processing developers: Optimize language models by understanding when increasing input context no longer improves test predictions, reducing unnecessary computation.
Authors
Masoud Kavian, Abdellatif Zaidi, Milad Sefidgaran
Abstract
We develop an information-theoretic framework for generalization in next-token prediction under temporally dependent data. We consider independent trajectories generated by finite-memory Markov processes and distinguish algorithmic dependence, quantified by mutual information, from temporal dependence, characterized by mixing. For cross-entropy loss, we derive an expected generalization bound using the Donsker--Varadhan variational representation and a McDiarmid-type concentration inequality for Markov chains. A refinement captures the joint effect of context length and temporal mixing through the mixing properties of the history-state process. We then extend the bound through a rate--distortion formulation, replacing mutual information with the minimum information rate required to represent the learned model within a prescribed distortion in the generalization gap, yielding informative guarantees for deterministic algorithms over continuous hypothesis spaces. For margin-based prediction, we derive explicit bounds for linear and self-attention next-token predictors via noisy low-dimensional compression, revealing the roles of context length, model complexity, sample size, margin, and temporal mixing. Experiments on TinyStories and ETTh2 show that longer contexts can reduce both training and test losses, but typically reduce training loss more, enlarging the generalization gap. A complementary ETTh2 analysis identifies an effective predictive-memory scale near 24 hours, with no statistically supported improvement beyond this scale, offering a plausible explanation for test-performance saturation at larger contexts.