Code Paper

The State-Prediction Separation Hypothesis

Giovanni MoneaCornell University Nathan GodeyCornell University Kianté BrantleyHarvard University Yoav ArtziCornell University
Cornell University Harvard University
TL;DR

Every hidden state in a Transformer is asked to predict the next token and to store state for later predictions. We hypothesize that those two roles compete: they are different tasks that push the hidden state to store different information. We design a variant (SPS) that routes each onto its own token stream via an interleaved <predict> slot. Across five scales (53M–1.68B), SPS lowers validation NLL (negative log‑likelihood), improves out‑of‑distribution generalization, and raises zero‑shot accuracy by 2–3 points, all with the same persistent KV cache as a Standard Transformer. The advantage widens with scale: at 1.68B, SPS learns as much from 18B tokens as Standard does from 47B.

State‑Prediction Separation, in motion. Real tokens $x_i$ write persistent KV‑cache entries (filled slots); each inserted <predict> token $\rho_i$ holds an ephemeral entry (dashed) and emits the next‑token prediction $\hat x_{i+1}$ (orange). Each ephemeral entry is evicted once it leaves a short sliding window.

The tension SPS targets

In a standard autoregressive Transformer, the hidden state $h_i$ at position $i$ produces two things:

To see the two roles separately, picture each position as having its own copy of the parameters: write $\theta_i$ for the copy that produces $h_i$. By linearity, the gradient at that position splits into two terms, one for each role:

$$ \nabla_{\theta_i} \mathcal{L} \;=\; \underbrace{\tfrac{1}{T-1}\,\nabla_{\theta_i} \ell_i}_{\text{Prediction}} \;+\; \underbrace{\tfrac{1}{T-1}\,\sum_{j=i+1}^{T-1} \nabla_{\theta_i} \ell_j}_{\text{State}}. $$

Both signals are routed through the same $h_i$, with no architectural way to keep them apart. Empirical work (Wu et al., 2024) shows pretrained Transformers pre‑cache: part of $h_i$ carries information that is useful for future tokens but not for the current prediction, rather than information that serves both. This effect grows with scale.

Hypothesis (informal)

Next‑token prediction and state storage compete when forced through one shared hidden state. Routing them through separate streams yields better language modeling.

The SPS Transformer

We augment the input sequence with a learned <predict> token $\rho_i$ inserted after every real token:

$$ x = (x_1, \ldots, x_T) \;\;\longrightarrow\;\; (x_1, \rho_1,\; x_2, \rho_2,\; \ldots,\; x_T, \rho_T). $$

Loss is applied only at $\rho_i$ positions: each $\rho_i$ predicts $x_{i+1}$. The attention mask sets what persists: input entries $x_i$ are persistent in the KV cache, and <predict> entries $\rho_i$ are evicted once they leave a sliding window of size $w$. So input representations accumulate the full state‑preparation gradient $\sum_{j>i}\nabla_{\theta_{x_i}} \ell_j$, while $\rho_i$'s gradient is dominated by the immediate prediction term $\nabla_{\theta_{\rho_i}} \ell_i$.

At inference the overhead is minimal. The persistent KV cache holds only input tokens, exactly matching the Standard Transformer's footprint. A small $w$‑slot ring buffer keeps the most recent <predict> entries. Each generated token triggers one decode step that forwards $(x_i, \rho_i)$ jointly and reads next‑token logits from the <predict> hidden state. Because autoregressive decoding is memory‑bound, forwarding the extra token adds negligible latency, the same property that motivates speculative decoding.

SPS matches Standard on a fraction of the data

We pretrain five model sizes (XS 53M, S 131M, M 379M, L 831M, XL 1.68B) on FineWeb‑Edu for 20B tokens (18B before the final learning‑rate decay, which we call the “pre‑decay” budget), following the GPT‑2 recipe and using nanoGPT as the training scaffold. All variants share the same backbone, parameter count, and hyperparameters; they differ only in attention pattern. We also train Standard for 40B tokens (47B at XL) for a fair GPU‑hours comparison.

At 1.68B parameters, SPS matches a Standard Transformer trained on 47B tokens while using 2.6x fewer tokens.
At 1.68B parameters, State‑Prediction Separation matches a Standard Transformer trained on 47B tokens while using 2.6× fewer tokens (pre‑decay). At an 18B‑token pre‑decay budget it already reaches $\Delta\mathrm{NLL} = -0.071$ versus the Standard baseline.
Takeaway 1. SPS reaches Standard's quality on roughly half the training data, and the data‑efficiency ratio widens with scale. At the S scale (10B tokens), a 3‑seed sweep confirms the gap over Standard is significant at $p < 0.005$.

Separation helps at every scale, and for the right reason

SPS adds compute (a forward pass per input token) and an extra mechanism (the <predict> insertion). Two ablations isolate whether separation, specifically, is what helps.

The two baselines. 2× Memory keeps every <predict> entry persistent, doubling the cache; Delayed State persists the <predict> slot instead of the input. Flip to the Delayed State and 2× Memory tabs in the explorer below to see each pattern next to SPS.

3
x key ρ key (<predict>) in window evicted ℓ loss
Hover a row to inspect.

Standard versus SPS and the two baselines. A Standard Transformer uses the same hidden state for state and prediction. SPS interleaves a <predict> token $\rho_i$ after each input. Loss is taken at $\rho_i$ slots; only $x_i$ slots stay persistent in the KV cache, and $\rho_i$ slots are evicted after a $w$‑step window (we use $w{=}64$, tuned at the small scale). Each row of the matrix is a query token; reading across gives the keys it attends to, up to itself on the diagonal. The aligned strip below shows the same model as of the final step: the KV cache (solid = kept, hatched = evicted), the input tokens, and the loss $\ell_i$. Switch tabs to compare SPS against Standard and the two baselines; hover any row to see exactly which keys it reads.

FineWeb-Edu validation NLL vs tokens seen (top) and GPU-hours (bottom) for five scales.
SPS reaches lower loss at every scale, whether we match tokens or compute. FineWeb‑Edu validation NLL against training tokens (top row, token parity) and against GPU‑hours (bottom row, compute parity). SPS stays below Standard in both views, and even doubling Standard’s pre‑decay token budget from 18B to 36B does not close the gap.

At token parity, SPS simply wins. Compared at the same number of training tokens (the top row of the figure), SPS reaches lower validation NLL at every scale. Three results hold consistently across XS–XL:

  1. Validation NLL. At a matched token budget, SPS lowers validation NLL by 0.042 at XS, growing to 0.068 at XL.
  2. Generalization. Held‑out NLL on four out‑of‑distribution corpora drops by 0.09–0.15, and zero‑shot accuracy on five standard benchmarks improves by 2.3–3.1 points.
  3. Inference efficiency. Same persistent KV footprint as Standard (1.01×), throughput within 5–10% of Standard.

And it holds at compute parity too. Each SPS training step costs roughly 2× a Standard step, since every input token is paired with an inserted <predict> token, so a token‑matched comparison alone could look unfair to Standard. The bottom row of the figure settles it by plotting loss against GPU‑hours instead of tokens: after a short initial phase, SPS reaches any given validation loss in fewer total GPU‑hours than Standard, and training Standard on twice its pre‑decay budget (36B, and 47B at XL) still does not catch up. The doubled per‑step compute more than pays for itself: SPS is the cheaper model to train.

Hover a bar for the exact value.
Each group is a model scale; within it, one bar per method shows the metric's absolute value, with Standard included as the grey bar. The y-axis is zoomed to the data range so the gaps are legible, and its label marks whether lower or higher is better for the selected metric. Hover any bar for its exact value and margin over Standard.

It is not just “more memory.”

2× Memory uses the same per‑step compute as SPS but keeps every <predict> entry persistent. So <predict> entries now serve both as a prediction site and as a state carrier for later queries (exactly the conflation SPS removes), at the cost of doubling the persistent KV cache. SPS, with half the persistent state, beats it on validation NLL at every scale. The advantage is not capacity‑based.

Takeaway 2. SPS outperforms a model with twice its persistent state. Keeping <predict> entries persistent forces them to play both roles: the conflation SPS removes.

It is not just “extra compute before each prediction.”

Delayed State matches SPS's per‑step compute and Standard's persistent‑cache size, but commits the persistent state at the $\rho_i$ slot, one step after the input. The two streams aren't separated: $\rho_i$ still serves as both prediction and state. Delayed State does improve over Standard (the extra compute helps), but SPS beats it by 0.019–0.021 in validation NLL and 0.05–0.10 in corpus NLL at every scale.

Takeaway 3. Even given an extra computation step before writing the persistent state, Delayed State underperforms SPS. Separation matters more than extra computation.

Why separation helps: the gradients

We probe how each architecture allocates the prediction‑vs‑state gradient at training time. Again give each position its own copy of the parameters, writing $\theta_p$ for the copy that produces the hidden state at position $p$. In SPS and Delayed State each step has two such copies: the input slot $\theta_{x_i}$ and the predict slot $\theta_{\rho_i}$. For each source position $p$ and offset $k$, we back‑propagate the single step‑$k$‑ahead loss in isolation and form the ratio of the position‑$p$ gradient from $\ell_{i+k}$ to the one from the current‑step loss $\ell_i$:

$$ r(p,\,k) \;=\; \frac{\bigl\lVert \nabla_{\theta_p}\, \ell_{i+k} \bigr\rVert_2} {\bigl\lVert \nabla_{\theta_p}\, \ell_i \bigr\rVert_2}. $$
Per-offset future/present gradient ratio for Standard, SPS, and Delayed State at five scales, log axis.
Where future‑loss gradient lands during training. Per‑offset gradient ratio $r(p,\,k)$ for Standard (single stream), SPS, and Delayed State at every scale. For SPS and Delayed State, solid curves are the input stream ($p{=}x_i$), dotted curves the prediction stream ($p{=}\rho_i$). SPS keeps the input stream persistent, and it sustains future‑loss gradient at every offset; Delayed State persists the prediction stream instead, and that stream stays low while the future‑loss gradient collects on its non‑persistent input stream, which falls below SPS past the $w{=}64$ window.

SPS's input stream carries strictly more future‑loss gradient than Standard at every offset, while its prediction stream carries strictly less. The two roles end up on different tokens, and because the input stream is the persistent one, that gradient lands on exactly the state that survives to serve later queries. Delayed State inverts which stream persists: there the prediction stream $\rho_i$ carries the persistent state, and it is that stream that stays uniformly low. The gradient collects instead on the input stream, which stays respectable (comparable to or above Standard). But in Delayed State the input is the non‑persistent, windowed stream: its slot is evicted after the $w{=}64$ window, so past offset $k{=}64$ it falls below SPS and can carry gradient forward only indirectly. In short, Delayed State piles the future‑loss signal onto the stream it discards and starves the one it keeps.

Carrying future‑loss gradient at training time is necessary, but not sufficient: we need to check that SPS's persistent state is actually used at inference. For a trained model $M$, let $M_\omega$ be the same model run with its persistent cache forced to a sliding window of size $\omega$ (distinct from the prediction window $w$), and let $\ell_i(M)$ be the loss at document‑relative position $i$. The degradation from that restriction is:

$$ \Delta\ell_{i}(M_{\omega}) \;=\; \ell_{i}(M_{\omega}) - \ell_{i}(M). $$

We average each $\Delta\ell_i$ over 8,000 documents and plot it against the query position $i$.

Per-position loss degradation when each method's persistent state is restricted to a sliding window of size 64, plotted against query position.
How much the persistent state is actually used at inference. Per‑position loss degradation $\Delta\ell_i(M_\omega)$ when each method's persistent state is restricted to a sliding window of size $\omega{=}64$, plotted against document‑relative query position. SPS's curve sits uniformly above Delayed State's, with a late‑position gap that narrows from $\sim$2.0× at XS to $\sim$1.3× at XL.
Takeaway 4. SPS better separates current‑step‑loss and future‑loss gradients during training and produces a persistent state that is more important at inference than Delayed State's.

What's next

SPS's per‑step training compute is roughly double Standard's because of the inserted <predict> token. An obvious follow‑up is whether the same separation can be obtained more cheaply: a shallower or narrower prediction stream, or a sparser persistent state. The two streams currently share all parameters; giving each distinct attention/FFN parameters now that the roles are decoupled is a natural next step.

Two further compute‑limited gaps remain in the evidence. We pretrain on a single corpus (FineWeb‑Edu); the consistent gains on out‑of‑distribution corpora and zero‑shot tasks suggest the trend transfers, but we have not tested alternative mixtures. And our largest scale is 1.68B. The SPS‑vs‑Standard gap widens monotonically across XS–XL, suggesting the trend extends to still‑larger models, but verifying that directly would require substantially more compute than we had available.

This matters in a regime where high‑quality human‑generated text is projected to run out. Learning more from each token directly extends the runway for pretraining, and in our scaling trend the advantage grows with model size.