The State-Prediction Separation Hypothesis Giovanni Monea†
Nathan Godey†
Kianté Brantley⋄
Yoav Artzi†
Cornell University ⋄ Harvard University [email protected], {ng554, yoavartzi}@cornell.edu [email protected]
Abstract Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the stateprediction separation hypothesis: disentangling the two roles yields better language modeling performance. We design a Transformer variant that uses two computation streams to separate the two functions, and conduct pretraining experiments across various scales. Our experiments show that state-prediction separation consistently offers better data and compute efficiencies, improving validation loss and outperforming standard Transformers by 2–3 percentage points on average on downstream tasks. We also conduct extensive empirical analysis that rules out potential confounders and demonstrates the fundamental difference in the gradients our design entails.
Standard Transformer 𝑥ො 2
𝑥ො 3
𝑥ො 4
𝑥ො 5
State-Prediction Separation Transformer
𝑥ො 6
𝑥ො 2
𝑥ො 3
𝑥ො 4
𝑥ො 5
𝑥ො 6 Ln
Ln
…
… L1
L1
x1
x3
x2
x4
x5
x1
KV cached activations
ρ1
x2
ρ2
x3
ρ3
x4
ρ4
x5
ρ5
KV ephemeral activations
2.61
Pre-Decay NLL
arXiv:2607.01218v1 [cs.CL] 1 Jul 2026
†
2.58 2.55 2.52 10B
20B
S TANDARD (47B tokens)
30B
Training Tokens
40B
S TATE -P REDICTION S EPARATION (18B tokens)
Figure 1: Standard versus State-Prediction Separation Transformer. Top: The standard Transformer uses the same hidden states for both memory and prediction. Our variant separates these roles: input token xi time steps form a persistent state, while prediction token steps ρi produce next-token predictions. Bottom: At 1.6B parameters, State-Prediction Separation matches the validation loss of a standard Transformer trained on 47B tokens while using 2.6× fewer tokens (pre-decay). At an 18B-token pre-decay budget, it already achieves ∆NLL = −0.071 versus standard. Preprint.
1
Introduction
Attention-based architectures, including the Transformer [Vaswani et al., 2017] and earlier recurrent designs [Bahdanau et al., 2015], have dual use for the activations computed at each time step: they are used to predict the output of that time step (i.e., token in language models) and are attended to by subsequent steps. The first role is focused on prediction; the second on capturing state information to be reused later on. Generally, these two functions are entangled in the same representation and computation stream (i.e., forward path). In this paper, we propose and study the following hypothesis in large language models (LLMs): Hypothesis: State-Prediction Separation (SPS; informal) The next-token prediction computation and state representation compete when forced through the same computation. Routing them through separate streams yields better language modeling. Technically, we separate the state and prediction functionalities by inserting an additional computation time step before predicting the next token (Figure 1). Time steps then appear in pairs: first, the token previously generated is processed, but no new token is emitted. An additional time step follows, which emits the next token to generate. The key-value (KV) entries from the first of the two steps are added to the KV cache, while the entries of the latter of the two are discarded.1 This design distinguishes between two streams: a state stream and a prediction stream. We conduct extensive experiments pretraining a set of LLMs at common research scales, from 53M to 1.678B parameters. The main result is that SPS significantly improves pretraining performance over standard Transformers in both token-equivalent and compute-equivalent settings. Figure 1 illustrates one of the key results: non-separating baselines cannot match the training loss of an SPS Transformer even with double the number of training tokens. We also show that state-prediction separation outperforms several variants controlling for SPS’s compute and memory overheads, proving that separation is the key component in the improvements we report. Our code is available at https://github.com/lil-lab/sps.
2
Prediction and State Preparation
We consider a standard autoregressive Transformer with vocabulary V , depth L, and parameters θ. Appendix A details the full architecture. The input is a sequence x = (x1 , . . . , xT ). At sequence (l) position i with input token xi , the model computes a per-layer hidden state hi ∈ Rd , where (0) (L) l = 1, . . . , L, hi is the token embedding, and hi the final representation. Each position i also contributes entries to the key-value (KV) cache. Through causal attention, all the past keys and values (L) (i.e., from positions k < i) contribute to the final representation hi , which is then used to compute (l) the next-token distribution. Each hidden state hi plays two roles: it is part of the computation of the immediate prediction for xi+1 , and it produces KV entries read by every later position. This double role is reflected in the optimization gradients, as they are computed through backpropagaPT −1 1 tion. The language modeling training loss is L = T −1 i=1 ℓi , with the per-position cross-entropy loss ℓi = − log p(xi+1 | x≤i ). The parameters θ are used repeatedly in each position i = 1, . . . , T in the Transformer. We isolate the gradients for each position i by denoting ∇θi L, and can similarly denote ∇θi ℓj to denote the gradients for position i from the loss at position j.2 The gradients of the loss L by θ are the linear sum of all per-position gradients: ∇θ L =
T −1 X i=1
∇θ i L =
T −1 T −1 T −1 T −1 1 XX 1 XX ∇θi ℓj = ∇θi ℓj . T − 1 i=1 j=1 T − 1 i=1 j=i
(1)
The last term follows pruning every non-causal ℓj , because of the causality in attention, each step’s parameters affect only the current and future losses. Separating the step’s own loss (j = i) from the 1 In practice, as we describe in Section 3, we retain prediction KV activations within a sliding window.
2 This follows how the forward pass creates a rolled-out computation graph with the parameters θ used repeatedly.
2
losses back-propagated only through the KV cache (j > i) decomposes the gradient by source: ∇θ L =
T −1 X i=1
1 1 ∇θi ℓi + T −1 [ T −1
|
{z
Prediction
}
|
T −1 X j=i+1
∇θi ℓj ] .
{z
State
(2)
}
Time step i contributes gradients for the prediction of xi+1 — the prediction task — and for the preparation of keys and values that help all later positions j > i make better predictions — the state (l) representation task. Both components flow (i.e., back-propagate) through the same hidden state hi , which is therefore optimized to conflate the two roles in a single set of activations.
3
The State-Prediction Separation Transformer
We separate the two roles by augmenting the standard Transformer with an additional learned token, <predict>, inserted after every input token. Given an input sequence x = (x1 , . . . , xT ), we form an augmented sequence by interleaving dummy tokens ρi , all set to a new <predict> token: x → (x1 , ρ1 , x2 , ρ2 , . . . , xT , ρT ) .
(3)
The two tokens xi and ρi at index i share the same position encoding. The model now maintains two interleaved streams of representations: an input stream {xi }Ti=1 that we use to carry the state forward, and a prediction stream {ρi }Ti=1 that we use to emit next-token predictions. Beyond a sliding window of size w, only key-value elements from the input stream positions are available in the KV cache, to be attended by later positions. The sliding window allows to attend to the specific token-choice representations for a short horizon (i.e., for local coherence). The prediction xi+1 is done at the position of ρi , so the loss is applied only at ρi positions: L=
T −1 1 X − log p xi+1 | x1 , ρ1 , . . . , xi , ρi . T − 1 i=1
(4)
In a standard Transformer, the two streams are tied together: ρi does not exist, and the same (l) representations hi at each position must simultaneously pack the information to emit the prediction for xi+1 and produce the keys and values read by every later position. P The two gradient components of Equation 2 (the prediction term ∇θi ℓi and the state-preparation term j>i ∇θi ℓj ) are routed through (ℓ)
one and the same hi , with no architectural separation. Figure 1 illustrates the State-Prediction Separation Transformer (SPS), and compares it to the standard architecture. We can now make the informal hypothesis from Section 1 precise in this two-stream notation: Hypothesis: State-Prediction Separation (formal) P The prediction gradients ∇θρi ℓi and the state-preparation gradients j>i ∇θxi ℓj are both used to optimize the computation of the same representations hi , thereby competing with each other in the standard Transformer. Separating the hidden representations to hρi and hxi and routing the gradients appropriately separates the two functions, and yields lower next-token loss at matched parameter count.
At training time, we realize this separation through an attention mask. In SPS, input entries are persistent, while <predict> entries are evicted once they leave a sliding window of size w. A query q at step i (i.e., either xi or ρi ) attends to all causal input entries and only recent <predict> entries. The only difference between input and prediction positions is that prediction positions attend to their corresponding input position: {ρk : i − w ≤ k < i} if q = xi . (5) ASPS (i, q) = {xk : k ≤ i} ∪ {ρk : i − w ≤ k ≤ i} if q = ρi | {z } {z } | All causal inputs Recent <predict> entries
The persistent KV cache of SPS contains only input entries; <predict> entries are read by at most w later queries before being discarded. This routes the two gradient components of Equation 2 to 3
(l)
different streams. Input representations hxi are visible to every later query, so they accumulate the P (l) full state-preparation gradient j>i ∇h(l) ℓj . <predict> representations hρi are visible only within xi a window. Their gradient is dominated by the immediate prediction term ∇θρi ℓi , with a contribution limited to at most w − 1 following state-preparation losses. Figure 1 contrasts a standard Transformer with our SPS Transformer. Training Efficiency Our method increases compute in order to separate the state and prediction streams, which makes training more expensive due to the doubled context length. We efficiently simulate the sliding window through attention masking, and apply the same mechanism to prevent attention from crossing document boundaries. Inference Efficiency The additional cost is negligible at inference. Forwarding one or a few tokens simultaneously incurs essentially the same latency. This is a well-known property that motivates speculative decoding [Chen et al., 2023, Leviathan et al., 2023]. Concretely, SPS’s persistent KV cache contains only input tokens, matching the size of a standard Transformer cache, with a bounded w-slot ring buffer holding the most recent <predict> entries. Each generated token triggers a single decode step that forwards the pair (xi , ρi ) jointly and reads next-token logits from the <predict> hidden state.
4
Experimental Setup
Baselines We compare SPS to a standard Transformer (S TANDARD) and to two ablations. The first, 2 X M EMORY, retains both input and <predict> entries in the KV cache throughout the sequence. The model gains computational capacity from its doubled context length, but its persistent memory footprint also doubles, and <predict> entries still serve both prediction and state-preparation. The second, D ELAYED S TATE, inserts a <predict> token after every input, giving the model an extra computation step before each prediction, but commits the persistent state at the <predict> slot, one step after the input. Compared to SPS, this variant delays state preparation as well, and performs both prediction and state preparation together at the <predict> slot. A query q at step i attends to the w most recent input entries, using the same fixed-size ephemeral window as SPS but populated by input entries rather than <predict> entries, and to all causal <predict> entries. D ELAYED S TATE keeps the persistent KV cache size roughly equivalent to S TANDARD’s. However, unlike SPS, no separation between roles is enforced. The <predict> stream carries both prediction and state. D ELAYED S TATE therefore differs from SPS only in whether the two streams are separated. All variants share the same backbone, a XS S M L XL pre-normalized Transformer blocks with Layers (L) 8 12 24 36 48 RMSNorm [Zhang and Sennrich, 2019], Hidden size (d) 512 768 1024 1280 1600 SwiGLU feed-forward networks [Shazeer, Heads (H) 8 12 16 20 25 2020] of intermediate size 3d, rotary poFFN size 1536 2304 3072 3840 4800 sitional embeddings [Su et al., 2021], no Parameters 53M 131M 379M 831M 1.678B biases on linear layers, a weight-tied unembedding, and a context length of 4,096 Table 1: Model configurations across the five scales. tokens. We evaluate five scales, summarized in Table 1. For S, M, L, and XL, we follow the GPT-2 [Radford et al., 2019] recipe, and we add XS as a smaller scale. SPS, 2 X M EMORY, D ELAYED S TATE, and R EVERSE SPS use the same backbone, parameter count, and hyperparameters as S TANDARD at every scale, differing only in attention pattern. Unless otherwise specified, all sliding-window variants use w=64 at every scale. Data We pretrain on FineWeb-Edu [Penedo et al., 2024], a high-quality educational subset of FineWeb, and tokenize with the GPT-2 tokenizer. Sequences are packed across document boundaries, with an end-of-sequence token (<eos>) inserted between consecutive documents to delimit them. Attention is masked so that queries within a document cannot attend to keys from any other document, and we exclude <eos> positions from the next-token loss so that the model is never trained to predict the start of an unrelated document. Each model is trained for 20B tokens by default. This budget meets or exceeds the Chinchilla compute-optimal ratio of ≈20 tokens per parameter [Hoffmann et al., 2022] at every scale, except for XL (which is under-trained due to its higher cost of training).3 To obtain a fair GPU-hours comparison, we additionally train S TANDARD for 40B tokens (except for 3 Chinchilla-optimal ≈ 1.1B, 2.6B, 7.6B, 16.6B, and 33.6B tokens for XS, S, M, L, and XL.
4
XS
Validation NLL
3.26
S
3.03
M
2.84
L
2.74
3.21
2.97
2.77
2.67
2.60
3.16
2.92
2.71
2.60
2.52
3.11
2.86
2.65
2.53
2.45
3.06
10B
15B
20B
2.80
10B
15B
20B
2.58
10B
15B
Training Tokens
20B
2.46
10B
15B
20B
2.38
3.26
3.03
2.84
2.73
2.67
3.24
3.00
2.81
2.69
2.63
3.21
2.97
2.77
2.66
2.59
3.19
2.94
2.74
2.62
2.56
3.16
5
10
15
20
S TANDARD
25
2.91
10
20
30
40
50
2.71
25
50
75 100 125
2.59
GPU Hours
S TANDARD (M ORE T OKENS ) 2 X M EMORY
D ELAYED S TATE
50 100 150 200 250
SPS
XL
2.67
2.52
10B
100
15B
200
300
20B
400
Tokens to Match Pre-Decay Standard
Figure 2: SPS trains faster and reaches lower loss at every scale. FineWeb-Edu validation NLL vs. tokens seen (top) and GPU-hours (bottom). The top row includes LR cool down. XL, that is trained for 47B tokens, until it matches SPS validation loss). All runs see the same data in the same order. Training We base our training code on nanoGPT [Karpathy, 2022], including its standard hyperparameters.4 The global batch size is 96 sequences of length 4,096, i.e., ≈ 400,000 tokens per gradient update. All models are trained in bfloat16 mixed precision. We use a learning-rate schedule [Hägele et al., 2024] consisting of a brief linear warmup, a constant phase at the peak learning rate, and a linear decay covering the final 10% of training tokens. For faster training, we adapt the opensource Triton implementation of FlashAttention [Dao et al., 2022] to support the sliding-window and <predict>-token attention patterns of SPS, D ELAYED S TATE, and 2 X M EMORY, so that all variants train at comparable throughput; we found alternatives such as FlexAttention [Dong et al., 2024] to be either memory-inefficient or substantially slower in our setting. All XS/S/M/L runs use 2 NVIDIA H100 80 GB GPUs, while all XL runs use 2 NVIDIA B200 GPUs, with data-parallel distributed training. All runs in the main results use a single seed (i.e., the data ordering and weight-initialization seed are matched across methods). Evaluation We report three families of metrics. (a) Validation loss. Held-out NLL on FineWebEdu, our pretraining distribution. (b) Generalization. Corpus NLL averaged over four out-ofdistribution corpora (WikiText [Merity et al., 2017], C4 [Raffel et al., 2020], Pile-Books3 [Gao et al., 2020], GovReport [Huang et al., 2021]), and zero-shot accuracy averaged over five standard benchmarks (ARC-Easy [Clark et al., 2018], HellaSwag [Zellers et al., 2019], PIQA [Bisk et al., 2020], SciQ [Welbl et al., 2017], LAMBADA [Paperno et al., 2016]), evaluated as standard practice via the LM Evaluation Harness [Gao et al., 2024]. (c) Inference efficiency. End-to-end throughput (tokens/s) and peak GPU memory measured on a single NVIDIA H100 for a batch of 16 sequences with a prefill of 1024 tokens followed by 3072 decode steps, reported as ratios relative to S TANDARD at the same scale. Every method runs through the same generation loop, KV-cache layout, and Triton attention path, with each method using the fused kernel matched to its own attention pattern.
5
Results
Performance and Efficiency Table 2 summarizes performance and efficiency findings. SPS attains the lowest validation NLL on FineWeb-Edu at every scale, with the gap over S TANDARD widening from −0.042 at XS to −0.068 at XL. Figure 2 shows that at matched training tokens SPS reaches 4 AdamW with β =0.9, β =0.95, weight decay 0.1, gradient clipping at 1.0, and peak learning rate 6×10−4 . 1 2
5
VALIDATION L OSS Size
Method
G ENERALIZATION
I NFERENCE E FFICIENCY
FineWeb-Edu (↓)
Corpus NLL (↓)
Throughput (×, ↑)
Peak Memory (×, ↓)
XS
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
Task Accuracy (%, ↑)
3.107 3.073 (−0.034) 3.086 (−0.021) 3.065 (−0.042)
4.335 4.311 (−0.024) 4.312 (−0.023) 4.243 (−0.092)
44.9 47.3 (+2.4) 46.9 (+2.0) 47.3 (+2.5)
1.00 0.93 (−0.07) 0.94 (−0.06) 0.94 (−0.06)
1.00 1.81 (+81%) 1.01 (+1%) 1.01 (+1%)
S
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
2.858 2.829 (−0.030) 2.829 (−0.029) 2.810 (−0.048)
4.018 3.981 (−0.036) 3.992 (−0.026) 3.913 (−0.105)
49.5 51.1 (+1.5) 51.8 (+2.2) 51.8 (+2.3)
1.00 0.93 (−0.07) 0.94 (−0.06) 0.94 (−0.06)
1.00 1.81 (+81%) 1.01 (+1%) 1.01 (+1%)
M
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
2.648 2.610 (−0.038) 2.611 (−0.037) 2.591 (−0.058)
3.732 3.687 (−0.045) 3.701 (−0.031) 3.627 (−0.106)
55.8 57.3 (+1.4) 57.5 (+1.7) 58.7 (+2.9)
1.00 0.93 (−0.07) 0.95 (−0.05) 0.95 (−0.05)
1.00 1.81 (+81%) 1.01 (+1%) 1.01 (+1%)
L
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
2.537 2.495 (−0.042) 2.491 (−0.045) 2.470 (−0.067)
3.598 3.542 (−0.055) 3.532 (−0.066) 3.484 (−0.113)
60.1 62.1 (+1.9) 61.3 (+1.2) 62.6 (+2.5)
1.00 0.83 (−0.17) 0.90 (−0.10) 0.90 (−0.10)
1.00 1.78 (+78%) 1.01 (+1%) 1.01 (+1%)
XL
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
2.458 2.411 (−0.047) 2.410 (−0.048) 2.390 (−0.068)
3.487 3.423 (−0.064) 3.433 (−0.055) 3.338 (−0.149)
63.2 64.3 (+1.0) 64.8 (+1.6) 66.3 (+3.1)
1.00 0.64 (−0.36) 0.94 (−0.06) 0.94 (−0.06)
1.00 1.75 (+75%) 1.01 (+1%) 1.01 (+1%)
Table 2: SPS outperforms all baselines on quality while remaining comparable to S TANDARD in memory and throughput. Main results across XS–XL. Bold marks the best per column within each size; SPS rows are shaded. Task accuracy averages 5 zero-shot benchmarks; Corpus NLL averages 4 corpora. Throughput is end-to-end tokens/sec for a combined prefill 1k + decode 3k workload relative to S TANDARD on H100; Peak Memory is the ratio of peak GPU memory used during decode.
lower validation NLL than S TANDARD throughout training (top), and that at matched GPU-hours SPS eventually overtakes S TANDARD at every scale (bottom). Even doubling S TANDARD’s pre-decay training budget from 18B to 36B tokens does not close the gap. SPS thus reaches S TANDARD’s quality on roughly half the training data, with the data-efficiency ratio widening as scale grows. This is an important property as high-quality human-generated text approaches projected exhaustion [Villalobos et al., 2024]. The improvement carries over to held-out generalization. Corpus NLL on four outof-distribution corpora drops by 0.09–0.11 across scales, and zero-shot accuracy on five standard benchmarks improves by 2.3–3.1%. Figure 3 shows this trend directly. Crucially, this quality gain comes at minimal increase in inference cost. SPS’s persistent KV cache is the same size as S TANDARD’s (peak memory ratio 1.01), and end-to-end throughput is within 6–10% of S TANDARD at all scales. While each result above is from a single training run, we verify with a 3-seed sweep at S, 10B that SPS’s gap over S TANDARD, D ELAYED S TATE, and 2 X M EMORY is significant at p < 0.005 (one-sided Welch’s t-test; Appendix C). Takeaway 1: SPS outperforms S TANDARD in validation loss, generalization, and learning speed at every scale, matching S TANDARD’s persistent memory footprint and inference throughput while reaching S TANDARD’s quality on roughly half the training data. State and Prediction Role Separation 2 X M EMORY keeps every <predict> entry persistent at the same per-step compute as SPS, which discards <predict> entries beyond the window. SPS is consistently better in validation NLL across XS–XL, even though it has half the persistent KV cache. This shows that the gain is not capacity-based. Keeping <predict> entries persistent forces each one to serve both as a prediction site and as a state carrier for later queries, re-coupling the two streams that SPS separates. 6
Validation NLL
Accuracy Gain (pp)
2.91 3 2 1 0 XS 53M
S 131M
M 379M
L 831M
2.85 2.82 2.79
XL 1.7B
Model Scale
D ELAYED S TATE 2 X M EMORY
2.88
0
16
64
256
Temporary Window Size
SPS S TANDARD
D ELAYED S TATE SPS
Figure 3: SPS gives consistent downstream accuracy gains, with the largest observed gain at the largest scale. SPS improves average accuracy at every scale, with gains of roughly 2–3 percentage points and the largest gain observed at XL.
R EVERSE SPS S TANDARD
Figure 4: SPS works best at small but nonzero w, and outperforms R EVERSE SPS at every window. Final FineWeb-Edu validation NLL vs. <predict>-window size for SPS, D E LAYED S TATE, and Reverse SPS at S, after 20B training tokens.
A second hypothesis is that SPS benefits only from giving the model an additional Transformer step before the persistent state is committed. D ELAYED S TATE tests this directly. It has the same per-step compute as SPS and the same persistent-cache size as S TANDARD and SPS, but commits the persistent state at the <predict> slot, one step after the input, and so does not separate the two roles. D ELAYED S TATE does improve over S TANDARD, confirming that the extra computation step carries some benefit, but at every scale SPS remains consistently better, by 0.019–0.021 in validation NLL and by 0.06–0.07 in corpus NLL (Table 2). Takeaway 2: The specific structure of state-prediction separation in SPS matters more than extra computations or even simply doubling the memory. Which Stream Should Persist, and at What Window? We additionally run a R EVERSE SPS variant that swaps the two roles. <predict> entries from the persistent state and input entries are windowed and used to emit the next-token prediction. This is the mirror image of SPS, which predicts from <predict> and stores state from input. This isolates whether the specific role assignment matters. We jointly ablate the <predict>-window size w and the choice of persistent stream at the S scale, where a sweep is cheap to run, and reuse the resulting w across all scales in the main experiments (Figure 4). Both SPS and D ELAYED S TATE show nearly constant performance across all w ∈ {0, 16, 64, 256}, with w=64 as the empirical best by a small margin; we therefore fix w=64 at every scale. R EVERSE SPS matches SPS at moderate w but degrades sharply at small w, where windowing the input stream cuts off recent input visibility. SPS’s ordering is the more robust default: persisting inputs tolerate a wider range of w before quality drops. Takeaway 3: SPS is robust to window size, while R EVERSE SPS collapses at small w. Persisting the input stream is the more forgiving design. Analysis: Why Does Separation Help? Following Equation 2, we probe how each architecture allocates gradients between the prediction role (the immediate loss ℓi ) and the state representation role (future losses ℓj , j > i) at training time, and what this implies at inference. Recall that we denote as ∇i ℓj the gradients from the loss term ℓj due to the use of the parameters in position i.
In SPS and D ELAYED S TATE each step i occupies two positions in the interleaved sequence: an input slot xi where the input token is read, and a predict slot ρi where the cross-entropy ℓi is computed (the only position where a prediction happens). Both positions use the parameters θ. Therefore, we can separate the gradients to ∇xi ℓj and ∇ρi ℓj . Because xi ’s hidden states feed into ρi , both slots receive non-zero gradient from ℓi even though the loss is realized only at ρi . This is in contrast to S TANDARD, where each step occupies a single position with gradients ∇i ℓj . 7
Future / Present Ratio
XS
S
M
L
XL
101 100 10−1 0
128
256
384
512 0
128
256
384
512 0
128
256
384
512 0
Relative Future Offset
128
256
Method S TANDARD
SPS
384
512 0
128
256
384
512
Gradient Stream D ELAYED S TATE
Input Stream
Prediction Stream
(a) Per-offset gradient ratio r(p, k) for S TANDARD, SPS, and D ELAYED S TATE. For SPS and D ELAYED S TATE, solid curves are the input stream (p=xi ) and dotted curves the prediction stream (p=ρi ). Bold curves are Savitzky–Golay–smothed [Savitzky and Golay, 1964] trends. The faint curves underneath are the raw per-offset means. SPS’s input stream consistently sustains more future-loss gradient; D ELAYED S TATE’s prediction stream stays low, and its input stream collapses past the ephemeral window k=64.
XS
S
M
L
XL
∆`
0.6 0.4 0.2 0.0 64
512 1024 1536 2048 64
512 1024 1536 2048 64
512 1024 1536 2048 64
Document-relative query position
SPS
512 1024 1536 2048 64
512 1024 1536 2048
D ELAYED S TATE
(b) Per-position NLL degradation when each method’s persistent state is restricted to a window of size w=64, plotted against the document-relative query position t at each scale. Bold curves are Savitzky–Golay– smoothed [Savitzky and Golay, 1964] trends. The faint curves underneath are the raw per query position means. SPS’s curve sits uniformly above D ELAYED S TATE’s, with the late-position gap ranging from ∼ 2.0× at XS to ∼ 1.3× at XL. SPS’s persistent keys carry more future-relevant information and evicting them is more hurtful.
Figure 5: Analysis of SPS. (a) Where future-loss gradient lands during training. (b) How much the persistent state is actually used at inference. For every source position p at step i (so p ∈ {xi , ρi } in SPS and D ELAYED S TATE, p = xi in S TANDARD), we isolate ∇p ℓi+k , the gradients from the step-k-ahead loss, and compute the ratio5 r(p, k) =
∇θp ℓi+k 2 ∇θp ℓi 2
,
(6)
the magnitude of position-p gradients coming from the loss k steps ahead relative to the current time step gradients. We average r over 8,000 examples (1,000 documents, 8 source positions each) for k ≤ 512, separately on input positions (p=xi ) and prediction positions (p=ρi ) for SPS and D ELAYED S TATE. Figure 5a shows a clean dichotomy. SPS’s input stream carries more future-loss gradient than S TANDARD at every offset, and its prediction stream carries strictly less, meaning the two roles are routed to different tokens. D ELAYED S TATE reduces this separation. Its prediction stream stays uniformly low, and its input stream falls below SPS past the <predict> window k=64, beyond which gradient can only flow indirectly. A single stream does not allow the model to effectively learn to predict and to represent state. Carrying future-loss gradient is necessary but not sufficient, the persistent state must also be important at inference. We test this by restricting each variant’s persistent state to a sliding window of size ω (distinct from the prediction window w) and measuring the resulting NLL degradation. For a trained model M ∈ {SPS, D ELAYED S TATE}, we define Mω as M used with a forced sliding persistent-cache of size ω, ℓi (M ) as the loss of the vanilla model M at position i, and ℓi (Mω ) is the 5We exclude the language-model head from θ
p , since it carries only the present-loss prediction role we are trying to isolate.
8
loss of the altered Mω at the same position i. We measure ∆ℓi (Mω ) = ℓi (Mω ) − ℓi (M ),
for document positions i ∈ {1, . . . , 2048}. We set a small ω = 64 (as opposed to the vanilla 4,096) and average each ∆ℓi across 8,000 documents. A larger ∆ℓi means more of M ’s long-range prediction depends on persistent keys outside the ω-window. Figure 5b shows that ablating SPS’s out-of-window persistent state hurts NLL 1.4–2.2× more than ablating D ELAYED S TATE’s across scales, although with the full persistent cache SPS performs better than D ELAYED S TATE. Therefore, the future-loss gradient SPS routes onto the input stream actually translates into a persistent state the model relies on at inference. Takeaway 4: SPS better separates present-loss and future-loss gradients and produces a persistent state more important at inference compared to D ELAYED S TATE.
6
Related Work
Tension Between Present and Future Predictions Each hidden state in a Transformer is asked to do two jobs at once: encode the next-token prediction at its own position, and prepare the persistent state that later positions will read from. Wu et al. [2024] study this tension precisely by contrasting two hypotheses: breadcrumbs (the keys and values useful for the current prediction also serve future ones) and pre-caching (some computation in early positions is targeted at later predictions and would be wasted on the current one). They find pre-caching in pretrained Pythia, increasing with scale, consistent with mechanistic evidence that earlier-position representations already encode upcoming-token information [Elhage et al., 2021, Pal et al., 2023]. The early two-stream attention of XLNet [Yang et al., 2019] also distinguishes prediction from content, but in service of permutation language modeling rather than to relieve pre-caching under standard left-to-right training. SPS is a direct architectural response to this tension. Rather than asking one stream to serve both jobs, it inserts a dedicated <predict> slot at every position to carry the next-token prediction, freeing the input stream to specialize as persistent state. If the two-jobs view is right, separating the roles should help. Our experiments confirm this at every scale. Adding Compute on the Input Side One approach to relieve this tension is adding computation at input positions. Goyal et al. [2024] append <pause> tokens to the prompt so the model gets extra forward passes before answering, motivated by the fact that Transformer expressivity is bounded by context length [Merrill and Sabharwal, 2024]. Pfau et al. [2024] use the same insertion as filler tokens during training. These methods share SPS’s mechanism of adding extra tokens, but use it to add capacity rather than to separate the two roles. Our D ELAYED S TATE and 2 X M EMORY baselines isolate this distinction. Both retain the inserted-token mechanism and the extra compute, but do not enable the separation, and both underperform SPS at every scale. Enriching the Future-Prediction Signal A complementary line of work intervenes on the prediction target. Bachmann and Nagarajan [2024] show that teacher-forced next-token prediction can silently fail on planning tasks where one step is hard and the rest are easy, motivating training signals that reach beyond the immediate next token. Multi-token-prediction methods [Stern et al., 2018, Monea et al., 2023, Gloeckle et al., 2024, DeepSeek-AI, 2024, Ahn et al., 2025, Gerontopoulos et al., 2025] and belief-state objectives [Hu et al., 2025, Teoh et al., 2026] address this by adding auxiliary losses at non-current positions to strengthen the future-prediction signal itself. SPS pursues a different goal. We do not enrich what is predicted, but separate where prediction and state-preparation take place. By routing the next-token loss onto a dedicated <predict> slot, SPS relieves the present-future tension structurally, while remaining compatible with these approaches.
7
Discussion
We introduce the SPS hypothesis that posits that the two tasks each hidden state must perform, predicting the next token and preparing state for later predictions, interfere when forced through one representation, and that separating them structurally should help. We study the hypothesis with the SPS Transformer, which realizes this separation via two interleaved streams, using non-persistent states for prediction. The experiments are decisive: at every scale from XS to XL, SPS lowers FineWeb-Edu validation NLL, improves held-out-corpus NLL, and raises zero-shot accuracy, at the same persistent KV-cache footprint as the standard Transformer and within a few percent of 9
its throughput. Our experiments show that separation is key to the observed improvement. Our gradient-flow and restricted-state analyses confirm the mechanism: SPS routes future-loss gradient onto the input stream and produces a persistent state significantly more impactful for future states than alternatives. SPS shows dramatic data efficiency gains, which increase monotonically across XS–XL, suggesting it would only grow with more compute. This matters in a regime where high-quality data is finite and approaching projected exhaustion [Villalobos et al., 2024]. Learning more from each token directly extends the runway for pretraining. Compute constrains the scope of our evidence in two ways. First, we pretrain on a single corpus (FineWeb-Edu); the consistent gains on out-of-distribution corpora and zero-shot benchmarks suggest the trends transfer beyond it, but we could not tested alternative pretraining mixtures. Second, our largest scale is 1.678B parameters; the SPS-vs-S TANDARD NLL gap monotonically widens across XS–XL, suggesting the trend should continue past 1.6B, but this requires further verification. Our argument that mixing prediction and state-preparation in one hidden state is suboptimal rests on controlled ablations and gradient/state analyses; a formal characterization of when and how much this conflation hurts, as a function of capacity, depth, or data, would tighten the case and is left for future work. SPS’s prediction stream adds a forward-pass slot per input position, roughly doubling per-step training compute over the standard Transformer; whether the same separation can be obtained at lower overhead, via shallower or narrower computation on the prediction stream, or a sparser persistent state, is an open and practically valuable question. The two streams currently also share all parameters. We leave for future work whether further separating them (e.g., via distinct attention/FFN parameters per stream) could yield further gains now that the roles are decoupled.
Acknowledgments This research was partially supported by a gift to the LinkedIn–Cornell Bowers Strategic Partnership; AI-MI and NSF Award 2433348; the National Science Foundation NSF under award OAC-2311521; and NASA under award No. 20-OSTFL20-0053. NG is supported by an Empire AI Postdoctoral Fellowship. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation or of NASA. We thank the members of the Cornell LIL Lab for helpful discussions. KB acknowledges this work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
References Kwangjun Ahn, Alex Lamb, and John Langford. Efficient joint prediction of multiple future tokens, 2025. URL https://arxiv.org/abs/2503.21801. Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 2296–2318. PMLR, 21–27 Jul 2024. URL https://proceedings.mlr.press/v235/bachmann24a.html. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv.org/abs/1409.0473. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432–7439. AAAI Press, 2020. doi: 10.1609/AAAI.V34I05.6239. URL https://doi.org/10.1609/aaai.v34i05.6239. Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling, 2023. URL https://arxiv.org/abs/2302.01318. 10
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL https://arxiv.org/abs/1803.05457. Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Re. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=H4DqfPSibmx. DeepSeek-AI. Deepseek-v3 technical report, 2024. URL https://arxiv.org/abs/2412.19437. Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels, 2024. URL https://arxiv. org/abs/2412.05496. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://transformer-circuits.pub/2021/framework/index.html. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL https://zenodo.org/records/12608602. Anastasios Gerontopoulos, Spyros Gidaris, and Nikos Komodakis. Multi-token prediction needs registers. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=WDdBhcwzGe. Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction, 2024. URL https://arxiv. org/abs/2404.19737. Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/ forum?id=ph04CRkPdC. Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: //openreview.net/forum?id=Y13gSfTjGr. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack William Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=iBBcRUlOAPR. Edward S. Hu, Kwangjun Ahn, Qinghua Liu, Haoran Xu, Manan Tomar, Ada Langford, Dinesh Jayaraman, Alex Lamb, and John Langford. The belief state transformer. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/ forum?id=ThRMTCgpvo. 11
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1419–1436, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.naacl-main.112. URL https://aclanthology.org/2021.naacl-main.112/. Andrej Karpathy. NanoGPT. https://github.com/karpathy/nanoGPT, 2022. Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. URL https: //openreview.net/forum?id=Byj72udxe. William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview.net/forum?id=NjNGlPh8Wh. Giovanni Monea, Armand Joulin, and Edouard Grave. PaSS: Parallel speculative sampling. arXiv preprint arXiv:2311.13581, 2023. Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state. In Proceedings of the 27th Conference on Computational Natural Language Learning (CoNLL), page 548–560. Association for Computational Linguistics, 2023. doi: 10.18653/v1/2023.conll-1.37. URL http://dx.doi.org/10.18653/ v1/2023.conll-1.37. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Katrin Erk and Noah A. Smith, editors, Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1525–1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/P16-1144. URL https://aclanthology.org/P16-1144/. Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https://openreview.net/forum?id= n6SCkn2QaG. Jacob Pfau, William Merrill, and Samuel R. Bowman. Let’s think dot by dot: Hidden computation in transformer language models. In First Conference on Language Modeling, 2024. URL https: //openreview.net/forum?id=NikbrdtYvG. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAI technical report, 2019. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(1), January 2020. ISSN 1532-4435. Abraham. Savitzky and M. J. E. Golay. Smoothing and differentiation of data by simplified least squares procedures. Analytical Chemistry, 36(8):1627–1639, 1964. doi: 10.1021/ac60214a047. URL https://doi.org/10.1021/ac60214a047. Noam Shazeer. Glu variants improve transformer, 2020. URL https://arxiv.org/abs/2002. 05202. 12
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips.cc/paper_files/paper/2018/ file/c4127b9194fe8562c64dc0f5bf2c93bc-Paper.pdf. Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2021. Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Pratyusha Sharma, Riashat Islam, Alex Lamb, and John Langford. Next-latent prediction transformers learn compact world models. In Bridging Planning and Reasoning in Natural Language with Foundational Models, 2026. URL https://openreview.net/forum?id=Lh4ayjJIAW. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf. Pablo Villalobos, Anson Ho, Jaime Sevilla, Tamay Besiroglu, Lennart Heim, and Marius Hobbhahn. Will we run out of data? limits of llm scaling based on human-generated data, 2024. Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Leon Derczynski, Wei Xu, Alan Ritter, and Tim Baldwin, editors, Proceedings of the 3rd Workshop on Noisy User-generated Text, pages 94–106, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10.18653/v1/W17-4413. URL https://aclanthology.org/W17-4413/. Wilson Wu, John Xavier Morris, and Lionel Levine. Do language models plan ahead for future tokens? In First Conference on Language Modeling, 2024. URL https://openreview.net/ forum?id=BaOAvPUyBO. Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ dc6a7e655d7e5840e66733e9ee67cc69-Paper.pdf. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Màrquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1472. URL https://aclanthology.org/P19-1472/. Biao Zhang and Rico Sennrich. Root mean square layer normalization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ 1e8a19426224ca89e83cef47f1e7f53b-Paper.pdf.
A
Full Transformer Notation
Let V denote a finite vocabulary. An autoregressive Transformer defines a distribution p : V ≤T → ∆(V ) mapping a sequence of tokens to a distribution over the next token. The model consists of L Transformer layers, with H attention heads each, with dimensions dh = d/H. (0) For any given sequence position i, each token xi is mapped to an embedding hi = Exi ∈ Rd , where E ∈ R|V |×d is a learned embedding matrix (so Exi denotes its xi -th row). 13
The embeddings are processed by L blocks combining causal multi-head self-attention (MHA) with a position-wise feed-forward network (FFN) and normalization layers (Norm), producing the per-layer (l) hidden states hi ∈ Rd . For l = 1, . . . , L: (l) (l−1) (l) h̃i = hi + MHA(l) NormMHA (h(l−1) ) i , (l) (l) (l) (l) hi = h̃i + FFN(l) NormFFN (h̃i ) . (l)
(l)
(l−1)
) denote the normalized input to layer l. Each head η ∈ {1, . . . , H} is Let h̄i := NormMHA (hi (l,η) (l,η) (l,η) parameterized by projection matrices WQ , WK , WV ∈ Rdh ×d and computes (l,η)
qi
(l,η) (l) h̄i ,
= WQ
(l,η)
ki
(l,η) (l) h̄i ,
= WK
(l,η)
vi
(l,η) (l) h̄i .
= WV
(7)
A rotary positional transform Ri is applied to queries and keys, and the per-head output is obtained by causally masked attention, then concatenated across heads and mixed by an output projection (l) WO ∈ Rd×d : ! (l,η) (l,η) X (Ri qi )⊤ (Rj kj ) (l,η) (l,η) √ oi = softmaxj vj , (8) d h j≤i (l) (l,1) (l,H) MHA(l) (h̄(l) )i = WO oi ; . . . ; oi . (9) The next-token distribution is obtained by applying a final RMSNorm and a weight-tied unembedding (L) to the final representation hi : (L) p(· | x≤i ) = softmax E RMSNormf (hi ) . (10)
B
Full Main Results
Table 3 expands Table 2 with the per-corpus NLLs (WikiText, C4, Pile-Books3, GovReport) and per-benchmark zero-shot accuracies (ARC-Easy, HellaSwag, PIQA, SciQ, LAMBADA) that are averaged into Corpus NLL and Task Accuracy in the main text, along with the prefill-throughput ratio.
C
Seed Variance and Statistical Tests
Re-training each variant at every scale across multiple seeds is computationally prohibitive at pretraining cost, so we run a focused seed-robustness check at the S, 10B setting. Each of the four variants (S TANDARD, D ELAYED S TATE, 2 X M EMORY, SPS) is re-trained with three seeds: the headline run plus seed 0 and seed 1. The seeds vary both the training-data ordering (the order in which packed sequences are streamed by the data loader) and the weight-initialization random seed; all other hyperparameters are held fixed at their main-table values. Figure 6 shows the mean and 95% confidence interval of the final FineWeb-Edu validation NLL across the three seeds; the confidence interval is computed from the Student-t distribution with the multiplier t0.025, 2 ≈ 4.30 for n = 3. We test whether SPS improves over each baseline by a one-sided Welch’s t-test against the alternative “SPS has lower validation NLL”. All three baselines reject the null at p < 0.005 even with the small-n Student-t penalty: Comparison
gap (NLL)
t
p
SPS vs. S TANDARD SPS vs. D ELAYED S TATE SPS vs. 2 X M EMORY
−0.0495 −0.0153 −0.0161
−26.6 −6.78 −7.88
4.4×10−4 1.7×10−3 2.2×10−3
Table 2 also suggests that D ELAYED S TATE and 2 X M EMORY end up at indistinguishable validation NLL despite very different memory footprints. We confirm this with a two one-sided test (TOST) for equivalence within ±0.01 NLL: TOST p = 3.4×10−3 , so the two are statistically equivalent at this scale within a margin well below the gap to either SPS (0.0153) or S TANDARD (0.034). 14
NLL Generalization (↓)
Task Generalization (%, ↑)
Size
Method
WT
C4
Books3
GR
ARC-E
HS
PIQA
SciQ
LAMB
S
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
3.466 3.417 3.423 3.368
4.340 4.302 4.304 4.274
4.908 4.896 4.928 4.735
3.357 3.310 3.312 3.273
50.3 49.5 51.1 50.3
34.8 36.2 36.5 37.4
64.0 65.0 64.3 65.4
71.5 73.4 76.1 74.1
27.0 31.3 30.9 31.8
M
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
3.182 3.141 3.150 3.101
4.058 4.026 4.006 3.974
4.603 4.552 4.618 4.443
3.087 3.030 3.032 2.988
56.6 55.9 57.0 59.5
42.5 44.9 44.7 45.8
67.7 68.4 68.7 69.0
77.3 79.7 79.9 80.4
35.0 37.4 37.4 39.0
L
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
3.063 3.006 2.996 2.953
3.913 3.856 3.853 3.830
4.468 4.421 4.396 4.313
2.946 2.886 2.881 2.841
61.9 62.9 61.0 62.8
47.5 50.6 50.8 52.2
69.3 70.5 71.4 70.7
81.7 84.6 81.9 83.1
40.4 41.8 41.6 44.3
XL
S TANDARD 2 X M EMORY D ELAYED S TATE SPS
2.954 2.895 2.893 2.831
3.812 3.749 3.752 3.718
4.336 4.265 4.303 4.061
2.846 2.783 2.783 2.740
64.5 64.0 64.7 66.3
52.6 55.1 55.4 56.3
71.6 71.6 72.0 71.9
84.2 84.1 85.9 87.5
43.2 46.6 46.1 49.5
Validation NLL
Table 3: Per-corpus and per-benchmark expansion of the main results in Table 2. Per-corpus held-out NLL on four out-of-distribution corpora (WikiText, C4, Pile-Books3, GovReport) and per-benchmark zero-shot accuracy on five standard tasks (ARC-Easy, HellaSwag, PIQA, SciQ, LAMBADA), averaged into Corpus NLL and Task Accuracy in the main text. Bold marks the best per column within each size; SPS rows are shaded.
2.92 2.91 2.90 2.89 2.88 2.87
N S TA
DA
RD
DE
Y LA
S ED
TE TA
2
E XM
MO
RY
SP
S
Figure 6: Seed-level robustness at S, 10B. Final FineWeb-Edu validation NLL across n=3 seeds per method (the headline run plus seed 0 and seed 1). Bars are 95% confidence intervals (Student-t, t0.025,2 ≈ 4.30).
15