PaperlayerView exact revision on arXiv

Source-linked AI summary

Information-Theoretic Analysis of Next-Token Prediction under Markovian Data

Masoud Kavian, Abdellatif Zaidi, Milad Sefidgaran

arXiv:2609.34731v1stat.MLcs.ITcs.LGmath.IT

TL;DR

Next-token prediction uses context from temporally dependent data, raising the question of how context length and temporal dependence jointly affect generalization. The paper develops information-theoretic bounds and extensions for Markovian data. Its ETTh2 analysis identifies an effective predictive-memory scale near 24 hours, with no statistically supported improvement beyond that scale under the evaluation protocol.

  • Problem

    The paper asks how context length and temporal dependence jointly affect next-token predictor generalization, since longer contexts can improve prediction but also enlarge the train–test gap.

    Source: §I · p. 2

  • Method

    The paper derives information-theoretic generalization bounds for Markovian data, extends them using rate–distortion, and specializes them to margin-based linear and self-attention predictors.

    Source: §I · p. 3

  • Results

    The ETTh2 analysis identifies an effective predictive-memory scale near 24 hours, beyond which longer histories provide no statistically supported out-of-sample improvement under the evaluation protocol.

    Source: §I · p. 4

  • Takeaways & Limitations

    Longer memory can improve prediction when added context is informative, but it may simultaneously increase the statistical cost of generalization.

    Source: §I · p. 4

Abstract

from arXiv · show

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.

I. INTRODUCTION

The paper asks how context length and temporal dependence jointly shape next-token prediction generalization under Markovian data. It develops information-theoretic guarantees and studies the tradeoff between predictive benefit and generalization cost as memory grows.

III. BOUNDS ON THE GENERALIZATION GAP

  • The framework derives cross-entropy NTP bounds separating algorithmic dependence, measured by mutual information, from temporal dependence, characterized by mixing.A history-state refinement absorbs dependence from overlapping contexts and removes the explicit context-length factor from the leading constant.Source: §I · p. 3
  • The rate–distortion extension characterizes the minimum information rate needed to represent a learned model within a prescribed distortion in its generalization behavior.It remains meaningful for deterministic algorithms over continuous parameter spaces.Source: §I · p. 3

I. INTRODUCTION

  • Noisy low-dimensional compression yields margin-based bounds for linear and self-attention predictors, exposing the roles of context length, mixing, model norm, margin, dictionary size, and training-set size.Source: §I · p. 3
  • The ETTh2 analysis identifies an effective predictive-memory scale near 24 hours, beyond which longer histories provide no statistically supported out-of-sample improvement under the evaluation protocol.Source: §I · p. 4
  • Longer memory can improve prediction when added context is informative, but may also increase the statistical cost of generalization.The framework relates that cost to temporal mixing, context length, and predictor complexity and compressibility.Source: §I · p. 4

A. Data generation process

  • The study considers independent trajectories generated by non-stationary finite-memory Markov processes and asks how context length and temporal dependence jointly affect NTP generalization.A prediction uses preceding observations, so changing a token can affect several subsequent losses through overlapping contexts.Source: §I · p. 2, A. Data generation process · p. 4

A. Loss-less generalization bound

The section develops generalization bounds for next-token prediction under Markovian data, refining them with history-state mixing and rate–distortion. It then applies margin-based analysis to derive bounds for margin prediction and linear next-token predictors.

A. Loss-less generalization bound

  • Theorem 1 bounds the cross-entropy generalization gap under Assumptions 1–2, and Lemma 1 gives a history-state mixing refinement that is tighter under the stated condition.The proof sketch combines the Donsker–Varadhan inequality with Markov-chain concentration; changing a token can affect its own loss and up to ρ subsequent context windows.Source: A. Loss-less generalization bound · pp. 6–7

B. Lossy generalization bound

  • Theorem 2 replaces exact model information with a rate–distortion term, allowing prescribed distortion in the generalization gap.The refinement is intended to remain informative for deterministic algorithms over continuous hypothesis spaces and supports explicit scaling laws for LNTP and SANTP.Source: B. Lossy generalization bound · pp. 7–8

A. Margin-Based Generalization Gap

  • The margin loss satisfies a bounded-difference condition, yielding an expected generalization bound and a history-state mixing refinement that removes the explicit factor ρ + 1.The refined bound is strictly tighter than the earlier margin bound whenever the stated condition holds.Source: A. Margin-Based Generalization Gap · p. 9

B. Linear Next Token Prediction

  • Theorem 3 upper-bounds the expected margin generalization gap for LNTP-SL using a lossy margin-based bound and an explicit compression scheme.The described scheme retains selected parameter vectors, projects them, discards atypically large projections, and adds bounded noise; the analysis controls distortion and information rate.Source: B. Linear Next Token Prediction · p. 10

C. Self-Attention Next Token Prediction

Longer contexts improve next-token prediction while widening the empirical train–test gap, and the theory bounds self-attention predictors while accounting for temporal mixing. The experiments also show that predictive gains can saturate as context grows.

C. Self-Attention Next Token Prediction

  • The SANTP-SL result bounds the expected generalization gap, with its proof sketch using separate compression of M3 and M4 to handle attention-weight perturbations.The analysis decomposes log-margin distortion into changes from replacing M4 and then M3.Source: C. Self-Attention Next Token Prediction · p. 11
  • History-state mixing yields strictly tighter bounds than scalar-mixing counterparts under the condition stated in the refinement.The refinement uses ˜τmin.Source: C. Self-Attention Next Token Prediction · p. 12

D. Effect of memory (mixing time)

  • Markovian dependence imposes a memory penalty captured by √τmin, inflating the generalization gap and requiring more data than an i.i.d. setting for comparable performance.The passage relates this penalty to redundancy from temporal coupling.Source: D. Effect of memory (mixing time) · p. 12

V. EXPERIMENTS

A. Language Model Pretraining on TinyStories

B. Next-Token Prediction on ETTh2

On ETTh2, longer contexts reduce training loss, but test loss improves only up to about ρ = 32 while the train–test gap widens overall. The accompanying Markov-chain tools characterize dependence through mixing, including for the history-state process.

B. Next-Token Prediction on ETTh2

  • The ETTh2 experiment quantizes the univariate OT sequence into 20 symbols and trains a causal self-attention predictor on preceding observations across context lengths 2 to 128.Each context-length model is trained from scratch with 15,000 updates and five random seeds.Source: B. Next-Token Prediction on ETTh2 · pp. 14–15
  • Test loss falls from approximately 0.82 at ρ = 2 to its low of approximately 0.73 around ρ = 32, then increases slightly as training loss remains substantially smaller.Source: B. Next-Token Prediction on ETTh2 · p. 14, Figure 2 · p. 15
  • The ETTh2 generalization gap widens overall from approximately 0.07 for the shortest contexts to around 0.17 for the largest, though it is not strictly monotone.Source: B. Next-Token Prediction on ETTh2 · p. 15
  • The authors describe the pattern as consistent with the analysis: out-of-sample gains from added history eventually saturate while the training–test discrepancy continues to widen.Source: B. Next-Token Prediction on ETTh2 · p. 15

A. On the Markovian random variables

  • The Markov-chain concentration results use mixing time, while the history-state process packages the current observation with its preceding ρ observations to characterize context dependence.Source: A. On the Markovian random variables · p. 17

B. Proof of Theorem 1

The proof combines a mutual-information step based on the Donsker–Varadhan representation with concentration for Markov chains. It controls how changing a trajectory affects the losses used in the generalization gap.

B. Proof of Theorem 1

  • The proof applies the Donsker–Varadhan variational representation and identifies the resulting information term with mutual information.Source: B. Proof of Theorem 1 · p. 18
  • For the concentration step, the proof invokes the Markov-chain McDiarmid inequality after establishing bounded differences for the relevant function.Source: B. Proof of Theorem 1 · pp. 18, 21
  • A changed token can affect at most ρ subsequent prediction losses because their contexts include that token.Source: B. Proof of Theorem 1 · pp. 19–21

C. Proof of Theorem 3

The proof of Theorem 3 uses lossy compression to bound the rate needed to represent a model within a prescribed distortion. The argument controls this rate through projected parameter vectors and retained-coordinate indicators.

C. Proof of Theorem 3

  • The proof upper-bounds the rate–distortion term by constructing a compressed model that satisfies the required distortion condition.Source: C. Proof of Theorem 3 · p. 24
  • The proof represents prediction scores as inner products between model-dependent and context-dependent vectors before applying a random projection.Source: C. Proof of Theorem 3 · p. 25
  • It bounds the compression information using data processing and the entropy cost of indicators identifying retained parameter vectors.The retained-set cost is controlled by a binary-entropy estimate.Source: C. Proof of Theorem 3 · pp. 27–28

D. Proof of Theorem 4

Theorem 4 is proved using the same lossy approach as Theorem 3. The proof uses an extension of Theorem 2 to upper-bound the rate–distortion term for an attention model parameterized by two matrices.

D. Proof of Theorem 4

  • The proof of Theorem 4 follows the same lossy approach as the proof of Theorem 3.The parameters introduced here are used throughout the remainder of the proof.Source: D. Proof of Theorem 4 · pp. 29–30
  • An extension of Theorem 2 is used to derive an upper bound on the rate–distortion term RD(·).The attention mechanism depends on the product of the transposed Key matrix and the Query matrix.Source: D. Proof of Theorem 4 · p. 30
  • The proof introduces M3 and M4 to define the attention problem’s parameterization.The supplied passage identifies M3 with the transposed Key matrix and defines M4 from the Query matrix.Source: D. Proof of Theorem 4 · p. 30

2 M1. Accordingly, we consider the equivalent parameterization

The proof treats the attention predictor through a two-matrix parameterization, bounds its distortion and mutual-information terms, and applies the lossy margin-generalization bound. The stated distortion scale is ϵ = O(1/(NT)).

2 M1. Accordingly, we consider the equivalent parameterization

  • The attention model is parameterized as W = (M3, M4), with M4 introduced from the Key–Query matrix product.The proof defines row and column vectors of these matrices for its bounds.Source: §2 · pp. 30–31
  • The proof bounds distortion terms and then derives an upper bound on the mutual-information term before combining them.The distortion analysis uses reconstructed matrices and invokes the data-processing inequality in later steps.Source: §2 · pp. 38–39
  • Combining the rate bound with distortion ϵ = O(1/(NT)) and the lossy margin-generalization bound yields the theorem’s final bound.The supplied passage does not specify the resulting bound’s expression.Source: §2 · p. 40

E. Proof of Lemma 1

The proof of Lemma 1 bounds an MGF term using the paper’s McDiarmid-type inequality for dependent random variables and a uniform upper bound on the loss variables.

E. Proof of Lemma 1

  • The proof bounds the MGF term on the right-hand side of (120) to complete Lemma 1.It follows the same steps as (20)–(23) and obtains an analogous result to (24).Source: E. Proof of Lemma 1 · p. 40
  • The proof invokes τmin as defined in (10).The supplied passages do not state further details about its role here.Source: E. Proof of Lemma 1 · p. 40
  • The McDiarmid-type inequality in Lemma 3 is applied using nonnegative Yt values and their uniform upper bound from Lemma 4.The passage states that these facts establish the subsequent inequality.Source: E. Proof of Lemma 1 · p. 40

F. Proof of Lemma 3

The proof of Lemma 3 constructs a Marton coupling for the loss sequence in blocks sized by the mixing time. Maximal coupling at block boundaries yields disagreement probability at most δ, while common continuation preserves agreement after states coincide.

F. Proof of Lemma 3

  • The coupling partitions the trajectory into consecutive blocks of length L = ˜τ(δ), with a possibly shorter final block.The two histories agree through the preceding block and may differ at the current block; the Markov property permits coupling the state at its end and then continuing the Markov process.Source: F. Proof of Lemma 3 · p. 41
  • At each block boundary, maximal coupling of the L-step transitions bounds endpoint disagreement by δ; once states coincide, a common continuation keeps them equal.Intervening paths are sampled from their conditional Markov bridge distributions, preserving each trajectory’s law.Source: F. Proof of Lemma 3 · p. 41
  • The proof uses the complete context in the state preceding a block to connect state agreement with identical losses in that block.A common continuation makes every loss in the block identical when those preceding states coincide.Source: F. Proof of Lemma 3 · p. 41
  • The constructed coupling yields mixing-matrix entries Γii = 1, Γi,i+1 ≤ 1, and Γi,i+k ≤ δk−1 for k ≥ 2.Source: F. Proof of Lemma 3 · p. 41

G. Proof of Lemma 4

The proof of Lemma 4 reduces two steps to prior definitions and a boundedness assumption.

G. Proof of Lemma 4

  • Equation (124) follows from the definitions in (3) and (4).The passage gives the derivation source but not the equation’s further role.Source: G. Proof of Lemma 4 · p. 42
  • Assumption (1) bounds each coordinate g_w,j by B through the norm bound on g_w.The stated implication is |g_w,j| ≤ ∥g_w∥ ≤ B for every j in [d].Source: G. Proof of Lemma 4 · p. 42

H. TinyStories: Experimental and Implementation Details

The TinyStories experiment uses one tokenized stream and shared optimization settings to compare models trained with different context lengths.

H. TinyStories: Experimental and Implementation Details

  • Eight percent of TinyStories stories are selected, byte-level BPE-tokenized with an 8,000-token vocabulary, and concatenated into one stream.The first 10% of the stream is held out; the paper calls this test loss although the implementation calls the subset validation.Source: H. TinyStories: Experimental and… · p. 43
  • Sliding windows of ρ + 1 consecutive tokens create each context-and-target example, with only window length changing across contexts.The same tokenized corpus is used for every ρ.Source: H. TinyStories: Experimental and… · p. 43
  • Independent six-layer GPT-style decoder models are trained from scratch at contexts 2, 4, 8, 16, 32, 64, and 128.Each model has approximately 25 million trainable parameters; total gradient updates and optimization hyperparameters are shared across context lengths.Source: H. TinyStories: Experimental and… · p. 43
  • Models are trained for 10,000 AdamW updates, and the reported training and held-out losses are measured at the final optimization step.Losses are periodically evaluated using at most 100 mini-batches; figures show means across independent realizations with one-standard-deviation bands.Source: H. TinyStories: Experimental and… · pp. 43–44

I. ETTh2: Experimental and Implementation Details

The ETTh2 experiments compare next-token prediction across contexts and separately estimate a predictive-memory scale using multivariate temporal prediction. The context experiment’s test loss bottoms around ρ = 32, while the auxiliary analysis finds a 24-hour minimum and no statistically supported gain beyond it.

1) Context-Length Generalization Experiment:

  • The discrete ETTh2 experiment quantizes the univariate OT series into 20 bins and trains a two-layer causal self-attention predictor at seven context lengths.Contiguous blocks are assigned to train or test, with a context-width buffer at adjacent boundaries; models receive 15,000 updates and results average five seeds.Source: 1) Context-Length Generalization… · p. 44

2) Empirical Estimation of the Effective Predictive Memory:

  • Discrete ETTh2 test cross-entropy reaches its lowest observed value around ρ = 32, while training loss remains substantially smaller; the auxiliary mean NLL minimum is 24 hours, with no statistically supported gain beyond it.Some moderately larger contexts, including 34 hours, remain near-optimal, indicating a plateau rather than a sharp memory cutoff.Source: 1) Context-Length Generalization… · p. 44, 2) Empirical Estimation of the… · pp. 46–47
  • The auxiliary analysis operationally estimates useful history length, rather than testing an exact finite-order Markov representation or zero information in later lags.A Transformer predicts the current multivariate observation from preceding observations using learned conditional means and variances.Source: 2) Empirical Estimation of the… · p. 45
  • The estimated 24-hour scale is specific to the considered model and evaluation protocol and is not an exact finite Markov order or a sharp boundary on temporal dependence.The auxiliary analysis is offered as a plausible empirical interpretation of the main experiment’s context-length behavior.Source: 2) Empirical Estimation of the… · pp. 45, 47
  • The auxiliary search compares out-of-sample NLL across history lengths using four expanding-window folds and three seeds, with paired moving-block bootstrap tests and Holm correction.The paired evaluation uses identical target timestamps; the bootstrap uses 24-hour blocks and controls family-wise error at α = 0.05.Source: 2) Empirical Estimation of the… · pp. 45–46
  • The coarse search spans 1–168 hours, and a two-hour-resolution search from 20–36 hours leaves the minimum at 24 hours.The refined comparison also finds a broad near-optimal plateau, including a 34-hour context.Source: 2) Empirical Estimation of the… · pp. 45–46

Related papers

Loading 2609.34731v1…