1 Introduction↩︎

Attention-based architectures, including the Transformer [1] and earlier recurrent designs [2], 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):

Technically, we separate the state and prediction functionalities by inserting an additional computation time step before predicting the next token (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.¹ 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. 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 \(\theta\). 8 details the full architecture. The input is a sequence \(x = (x_1,\ldots,x_T)\). At sequence position \(i\) with input token \(x_i\), the model computes a per-layer hidden state \(h_i^{(l)} \in \mathbb{R}^d\), where \(l=1,\dots,L\), \(h_i^{(0)}\) is the token embedding, and \(h_i^{(L)}\) the final representation. Each position \(i\) also contributes entries to the key-value (KV) cache. Through causal attention, all the past keys and values (i.e., from positions \(k < i\)) contribute to the final representation \(h_i^{(L)}\), which is then used to compute the next-token distribution. Each hidden state \(h_i^{(l)}\) plays two roles: it is part of the computation of the immediate prediction for \(x_{i+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 backpropagation. The language modeling training loss is \(\mathcal{L} = \frac{1}{T-1}\sum_{i=1}^{T-1} \ell_i\), with the per-position cross-entropy loss \(\ell_i = -\log p(x_{i+1}\mid x_{\le i})\). The parameters \(\theta\) are used repeatedly in each position \(i=1,\dots,T\) in the Transformer. We isolate the gradients for each position \(i\) by denoting \(\nabla_{\theta_i}\mathcal{L}\), and can similarly denote \(\nabla_{\theta_i}\ell_j\) to denote the gradients for position \(i\) from the loss at position \(j\).² The gradients of the loss \(\mathcal{L}\) by \(\theta\) are the linear sum of all per-position gradients: \[\nabla_\theta\mathcal{L} \;=\; \sum_{i=1}^{T-1} \nabla_{\theta_i}\mathcal{L} \;=\; \frac{1}{T-1}\sum_{i=1}^{T-1}\sum_{j=1}^{T-1}\nabla_{\theta_i}\ell_j \;=\; \frac{1}{T-1}\sum_{i=1}^{T-1}\sum_{j=i}^{T-1}\nabla_{\theta_i}\ell_j \;\;.\label{eq:causal-grad}\tag{1}\] The last term follows pruning every non-causal \(\ell_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 losses back-propagated only through the KV cache (\(j>i\)) decomposes the gradient by source: \[\label{eq:aggregate-grad} \nabla_{\theta}\mathcal{L} \;=\; \sum_{i=1}^{T-1}\;[\; \underbrace{\vphantom{\sum_{i=j+1}^{T-1}}\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}}\;]\;\;.\tag{2}\]

Time step \(i\) contributes gradients for the prediction of \(x_{i+1}\) — the prediction task — and for the preparation of keys and values that help all later positions \(j>i\) make better predictions — the state representation task. Both components flow (i.e., back-propagate) through the same hidden state \(h_i^{(l)}\), 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=(x_1,\ldots,x_T)\), we form an augmented sequence by interleaving dummy tokens \(\rho_i\), all set to a new <predict> token: \[\label{eq:interleaved} x \rightarrow (x_1,\rho_1,\; x_2,\rho_2,\; \ldots,\; x_T,\rho_T)\;\;.\tag{3}\] The two tokens \(x_i\) and \(\rho_i\) at index \(i\) share the same position encoding. The model now maintains two interleaved streams of representations: an input stream \(\{x_i\}_{i=1}^T\) that we use to carry the state forward, and a prediction stream \(\{\rho_i\}_{i=1}^T\) 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 \(x_{i+1}\) is done at the position of \(\rho_i\), so the loss is applied only at \(\rho_i\) positions: \[\label{eq:sps-loss} \mathcal{L} = \frac{1}{T-1}\sum_{i=1}^{T-1} -\log p\bigl(x_{i+1} \mid x_1,\rho_1,\ldots,x_i,\rho_i\bigr)\;\;.\tag{4}\]

In a standard Transformer, the two streams are tied together: \(\rho_i\) does not exist, and the same representations \(h_i^{(l)}\) at each position must simultaneously pack the information to emit the prediction for \(x_{i+1}\) and produce the keys and values read by every later position. The two gradient components of 2 (the prediction term \(\nabla_{\theta_i}\ell_i\) and the state-preparation term \(\sum_{j>i}\nabla_{\theta_i}\ell_j\)) are routed through one and the same \(h_i^{(\ell)}\), with no architectural separation. 1 illustrates the State-Prediction Separation Transformer (SPS), and compares it to the standard architecture.

We can now make the informal hypothesis from 1 precise in this two-stream notation:

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 \(x_i\) or \(\rho_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: \[\label{eq:sps-mask} \mathcal{A}_{\mathrm{SPS}}(i,q) = \underbrace{\{x_k : k \leq i\}}_{\text{All causal inputs}} \;\cup\; \underbrace{ \begin{cases} \{\rho_k : i-w \leq k < i\} & \text{if } q=x_i\\ \{\rho_k : i-w \leq k \leq i\} & \text{if } q=\rho_i \end{cases}}_{\text{Recent \texttt{<predict>} entries}}\;\;.\tag{5}\] 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 2 to different streams. Input representations \(h_{x_i}^{(l)}\) are visible to every later query, so they accumulate the full state-preparation gradient \(\sum_{j>i}\nabla_{h_{x_i}^{(l)}}\ell_j\). <predict> representations \(h_{\rho_i}^{(l)}\) are visible only within a window. Their gradient is dominated by the immediate prediction term \(\nabla_{\theta_{\rho_i}}\ell_i\), with a contribution limited to at most \(w-1\) following state-preparation losses. 1 contrasts a standard Transformer with our SPS Transformer.

3.0.0.1 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.

3.0.0.2 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 [3], [4]. 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 \((x_i, \rho_i)\) jointly and reads next-token logits from the <predict> hidden state.

4 Experimental Setup↩︎

4.0.0.1 Baselines

We compare SPSto a standard Transformer (Standard) and to two ablations. The first, 2x Memory, 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, Delayed State, 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 SPSbut populated by input entries rather than <predict> entries, and to all causal <predict> entries.

Delayed Statekeeps the persistent KV cache size roughly equivalent to Standard’s. However, unlike SPS, no separation between roles is enforced. The <predict> stream carries both prediction and state. Delayed Statetherefore differs from SPSonly in whether the two streams are separated.

r0.55 4pt

All variants share the same backbone, a pre-normalized Transformer blocks with RMSNorm [5], SwiGLU feed-forward networks [6] of intermediate size \(3d\), rotary positional embeddings [7], no biases on linear layers, a weight-tied unembedding, and a context length of \(4{,}096\) tokens. We evaluate five scales, summarized in 1. For S, M, L, and XL, we follow the GPT-2 [8] recipe, and we add XS as a smaller scale. SPS, 2x Memory, Delayed State, and Reverse SPSuse the same backbone, parameter count, and hyperparameters as Standardat every scale, differing only in attention pattern. Unless otherwise specified, all sliding-window variants use \(w{=}64\) at every scale.

4.0.0.2 Data

We pretrain on FineWeb-Edu [9], 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 \(20\)B tokens by default. This budget meets or exceeds the Chinchilla compute-optimal ratio of \({\approx}20\) tokens per parameter [10] at every scale, except for XL (which is under-trained due to its higher cost of training).³ To obtain a fair GPU-hours comparison, we additionally train Standard for \(40\)B tokens (except for XL, that is trained for 47B tokens, until it matches SPS validation loss). All runs see the same data in the same order.

4.0.0.3 Training

We base our training code on nanoGPT [11], including its standard hyperparameters.⁴ The global batch size is \(96\) sequences of length \(4{,}096\), i.e.,\(\approx\!400{,}000\) tokens per gradient update. All models are trained in bfloat16 mixed precision. We use a learning-rate schedule [12] 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 open-source Triton implementation of FlashAttention [13] to support the sliding-window and <predict>-token attention patterns of SPS, Delayed State, and 2x Memory, so that all variants train at comparable throughput; we found alternatives such as FlexAttention [14] 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).

4.0.0.4 Evaluation

We report three families of metrics. (a) Validation loss. Held-out NLL on FineWeb-Edu, our pretraining distribution. (b) Generalization. Corpus NLL averaged over four out-of-distribution corpora (WikiText [15], C4 [16], Pile-Books3 [17], GovReport [18]), and zero-shot accuracy averaged over five standard benchmarks (ARC-Easy [19], HellaSwag [20], PIQA [21], SciQ [22], LAMBADA [23]), evaluated as standard practice via the LM Evaluation Harness [24]. (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 Standard 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↩︎

6pt

5.0.0.1 Performance and Efficiency

2 summarizes performance and efficiency findings. SPSattains the lowest validation NLL on FineWeb-Edu at every scale, with the gap over Standardwidening from \(-0.042\) at XS to \(-0.068\) at XL. 2 shows that at matched training tokens SPSreaches lower validation NLL than Standardthroughout training (top), and that at matched GPU-hours SPSeventually overtakes Standardat every scale (bottom). Even doubling Standard’s pre-decay training budget from \(18\)B to \(36\)B tokens does not close the gap. SPSthus reaches Standard’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 [25]. The improvement carries over to held-out generalization. Corpus NLL on four out-of-distribution corpora drops by \(0.09\)–\(0.11\) across scales, and zero-shot accuracy on five standard benchmarks improves by \(2.3\)–\(3.1\%\). 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 Standard’s (peak memory ratio \(1.01\)), and end-to-end throughput is within \(6\)–\(10\%\) of Standardat all scales. While each result above is from a single training run, we verify with a \(3\)-seed sweep at S, \(10\)B that SPS’s gap over Standard, Delayed State, and 2x Memoryis significant at \(p<0.005\) (one-sided Welch’s \(t\)-test; 10).

5.0.0.2 State and Prediction Role Separation

2x Memorykeeps every <predict> entry persistent at the same per-step compute as SPS, which discards <predict> entries beyond the window. SPSis 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 SPSseparates.

A second hypothesis is that SPSbenefits only from giving the model an additional Transformer step before the persistent state is committed. Delayed Statetests this directly. It has the same per-step compute as SPSand the same persistent-cache size as Standardand SPS, but commits the persistent state at the <predict> slot, one step after the input, and so does not separate the two roles. Delayed Statedoes improve over Standard, confirming that the extra computation step carries some benefit, but at every scale SPSremains consistently better, by \(0.019\)–\(0.021\) in validation NLL and by \(0.06\)–\(0.07\) in corpus NLL (2).

5.0.0.3 Which Stream Should Persist, and at What Window?

We additionally run a Reverse 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 ([fig:window95ablation]). Both SPSand Delayed Stateshow nearly constant performance across all \(w \in \{0, 16, 64, 256\}\), with \(w{=}64\) as the empirical best by a small margin; we therefore fix \(w{=}64\) at every scale. Reverse SPS matches SPSat 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.

5.0.0.4 Analysis: Why Does Separation Help?

Following 2 , we probe how each architecture allocates gradients between the prediction role (the immediate loss \(\ell_i\)) and the state representation role (future losses \(\ell_j\), \(j > i\)) at training time, and what this implies at inference. Recall that we denote as \(\nabla_i \ell_j\) the gradients from the loss term \(\ell_j\) due to the use of the parameters in position \(i\).

In SPSand Delayed Stateeach step \(i\) occupies two positions in the interleaved sequence: an input slot \(x_i\) where the input token is read, and a predict slot \(\rho_i\) where the cross-entropy \(\ell_i\) is computed (the only position where a prediction happens). Both positions use the parameters \(\theta\). Therefore, we can separate the gradients to \(\nabla_{x_i}\ell_j\) and \(\nabla_{\rho_i}\ell_j\). Because \(x_i\)’s hidden states feed into \(\rho_i\), both slots receive non-zero gradient from \(\ell_i\) even though the loss is realized only at \(\rho_i\). This is in contrast to Standard, where each step occupies a single position with gradients \(\nabla_i \ell _j\).

For every source position \(p\) at step \(i\) (so \(p \in \{x_i, \rho_i\}\) in SPSand Delayed State, \(p = x_i\) in Standard), we isolate \(\nabla_p \ell_{i+k}\), the gradients from the step-\(k\)-ahead loss, and compute the ratio⁵ \[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}\;\;,\] 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 \le 512\), separately on input positions (\(p{=}x_i\)) and prediction positions (\(p{=}\rho_i\)) for SPS and Delayed State. 4 (a) shows a clean dichotomy. SPS’s input stream carries more future-loss gradient than Standard at every offset, and its prediction stream carries strictly less, meaning the two roles are routed to different tokens. Delayed State 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 \(\omega\) (distinct from the prediction window \(w\)) and measuring the resulting NLL degradation. For a trained model \(M \in \{\mathrm{\small SPS},\,\mathrm{\small Delayed State}\}\), we define \(M_\omega\) as \(M\) used with a forced sliding persistent-cache of size \(\omega\), \(\ell_i(M)\) as the loss of the vanilla model \(M\) at position \(i\), and \(\ell_i(M_\omega)\) is the loss of the altered \(M_\omega\) at the same position \(i\). We measure \[\Delta\ell_i(M_\omega) \;=\; \ell_i(M_\omega) - \ell_i(M),\] for document positions \(i \in \{1,\dots,2048\}\). We set a small \(\omega=64\) (as opposed to the vanilla \(4{,}096\)) and average each \(\Delta\ell_i\) across \(8{,}000\) documents. A larger \(\Delta\ell_i\) means more of \(M\)’s long-range prediction depends on persistent keys outside the \(\omega\)-window. 4 (b) shows that ablating SPS’s out-of-window persistent state hurts NLL \(1.4\)–\(2.2\times\) more than ablating Delayed State’s across scales, although with the full persistent cache SPS performs better than Delayed State. Therefore, the future-loss gradient SPS routes onto the input stream actually translates into a persistent state the model relies on at inference.

6 Related Work↩︎

6.0.0.1 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. [27] 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 [28], [29]. The early two-stream attention of XLNet [30] also distinguishes prediction from content, but in service of permutation language modeling rather than to relieve pre-caching under standard left-to-right training. SPSis 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.

6.0.0.2 Adding Compute on the Input Side

One approach to relieve this tension is adding computation at input positions. [31] 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 [32]. [33] 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 Delayed Stateand 2x Memorybaselines isolate this distinction. Both retain the inserted-token mechanism and the extra compute, but do not enable the separation, and both underperform SPSat every scale.

6.0.0.3 Enriching the Future-Prediction Signal

A complementary line of work intervenes on the prediction target. [34] 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 [35]–[40] and belief-state objectives [41], [42] address this by adding auxiliary losses at non-current positions to strengthen the future-prediction signal itself. SPSpursues 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, SPSrelieves 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, SPSlowers 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 its throughput. Our experiments show that separation is key to the observed improvement. Our gradient-flow and restricted-state analyses confirm the mechanism: SPSroutes future-loss gradient onto the input stream and produces a persistent state significantly more impactful for future states than alternatives. SPSshows 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 [25]. 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.678\)B parameters; the SPS-vs-StandardNLL 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.

8 Full Transformer Notation↩︎

Let \(V\) denote a finite vocabulary. An autoregressive Transformer defines a distribution \(p \colon V^{\leq T} \to \Delta(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 \(d_h = d/H\). For any given sequence position \(i\), each token \(x_i\) is mapped to an embedding \(h_i^{(0)} = E_{x_i} \in \mathbb{R}^d\), where \(E \in \mathbb{R}^{|V| \times d}\) is a learned embedding matrix (so \(E_{x_i}\) denotes its \(x_i\)-th row).

The embeddings are processed by \(L\) blocks combining causal multi-head self-attention (\(\mathrm{MHA}\)) with a position-wise feed-forward network (\(\mathrm{FFN}\)) and normalization layers (\(\mathrm{Norm}\)), producing the per-layer hidden states \(h_i^{(l)} \in \mathbb{R}^d\). For \(l = 1, \ldots, L\):\[\begin{align} \tilde{h}_i^{(l)} &= h_i^{(l-1)} + \mathrm{MHA}^{(l)}\!\bigl( \mathrm{Norm}^{(l)}_{\mathrm{MHA}}(h^{(l-1)})\bigr)_i,\\ h_i^{(l)} &= \tilde{h}_i^{(l)} + \mathrm{FFN}^{(l)}\!\bigl( \mathrm{Norm}^{(l)}_{\mathrm{FFN}}(\tilde{h}_i^{(l)})\bigr). \end{align}\] Let \(\bar{h}_i^{(l)} \mathrel{\vcenter{:}}= \mathrm{Norm}^{(l)}_{\mathrm{MHA}}(h_i^{(l-1)})\) denote the normalized input to layer \(l\). Each head \(\eta \in \{1, \ldots, H\}\) is parameterized by projection matrices \(W_Q^{(l,\eta)}, W_K^{(l,\eta)}, W_V^{(l,\eta)} \in \mathbb{R}^{d_h \times d}\) and computes \[\label{eq:qkv} q_i^{(l,\eta)} = W_Q^{(l,\eta)} \bar{h}_i^{(l)}, \qquad k_i^{(l,\eta)} = W_K^{(l,\eta)} \bar{h}_i^{(l)}, \qquad v_i^{(l,\eta)} = W_V^{(l,\eta)} \bar{h}_i^{(l)}.\tag{6}\] A rotary positional transform \(\mathcal{R}_i\) 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 \(W_O^{(l)} \in \mathbb{R}^{d \times d}\): \[\begin{align} o_i^{(l,\eta)} &= \sum_{j \leq i} \mathrm{softmax}_{j}\!\left( \frac{(\mathcal{R}_i q_i^{(l,\eta)})^{\!\top} (\mathcal{R}_j k_j^{(l,\eta)})}{\sqrt{d_h}} \right) v_j^{(l,\eta)}, \tag{7}\\ \mathrm{MHA}^{(l)}(\bar{h}^{(l)})_i &= W_O^{(l)} \bigl[\, o_i^{(l,1)};\, \ldots;\, o_i^{(l,H)} \,\bigr]. \tag{8} \end{align}\]

The next-token distribution is obtained by applying a final RMSNorm and a weight-tied unembedding to the final representation \(h_i^{(L)}\): \[\label{eq:unembed} p(\cdot \mid x_{\leq i}) = \mathrm{softmax}\!\bigl( E\, \mathrm{RMSNorm}_{\text{f}}(h_i^{(L)}) \bigr).\tag{9}\]

9 Full Main Results↩︎

3 expands 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.

6pt

10 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, \(10\)B setting. Each of the four variants (Standard, Delayed State, 2x Memory, 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. 5 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 \(t_{0.025,\,2}\approx 4.30\) for \(n=3\).

We test whether SPSimproves over each baseline by a one-sided Welch’s \(t\)-test against the alternative “SPShas lower validation NLL”. All three baselines reject the null at \(p<0.005\) even with the small-\(n\) Student-\(t\) penalty:

2 also suggests that Delayed Stateand 2x Memoryend up at indistinguishable validation NLL despite very different memory footprints. We confirm this with a two one-sided test (TOST) for equivalence within \(\pm 0.01\) NLL: TOST \(p=3.4{\times}10^{-3}\), so the two are statistically equivalent at this scale within a margin well below the gap to either SPS(\(0.0153\)) or Standard(\(0.034\)).

References↩︎

---

1. In practice, as we describe in 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 \(\theta\) used repeatedly.↩︎

3. Chinchilla-optimal \(\approx\) \(1.1\)B, \(2.6\)B, \(7.6\)B, \(16.6\)B, and \(33.6\)B tokens for XS, S, M, L, and XL.↩︎

4. AdamW with \(\beta_1{=}0.9\), \(\beta_2{=}0.95\), weight decay \(0.1\), gradient clipping at \(1.0\), and peak learning rate \(6\!\times\!10^{-4}\).↩︎

5. We exclude the language-model head from \(\theta_p\), since it carries only the present-loss prediction role we are trying to isolate.↩︎