Source-linked AI summary
The State-Prediction Separation Hypothesis
Giovanni Monea, Nathan Godey, Kianté Brantley, Yoav Artzi
TL;DR
Transformers entangle next-token prediction with storing reusable state, raising whether separating these roles improves language modeling. The paper introduces a two-stream Transformer and finds that separation improves validation and downstream performance while increasing data efficiency.
Problem
Transformers entangle immediate prediction with reusable state representation in one computation stream, motivating the question of whether separating these roles improves language modeling.
Method
The SPS Transformer separates prediction and state preparation through two interleaved computation streams, using non-persistent states for prediction and testing models across scales.
Results
2.3–3.1% improvement in zero-shot accuracy accompanies held-out-corpus NLL reductions of 0.09–0.11 across scales.
Takeaways & Limitations
SPS reaches standard Transformer quality with roughly half the training data, while retaining the same persistent KV-cache footprint and similar throughput.
Takeaways & Limitations
SPS adds a forward-pass slot per input position, roughly doubling per-step training compute, and whether separation can be achieved with lower overhead remains open.
Abstract
from arXiv · showhide
Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the \emph{state-prediction 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.
1 Introduction
Transformers reuse each hidden state for both immediate prediction and future state storage, creating competing computational demands. SPS separates these roles with two streams and improves training efficiency.
- Motivation: SPS hypothesizes that next-token prediction and state representation compete when forced through one computation stream.Separate routing is proposed to improve language modeling.
- Method: SPS inserts an additional computation step before each next-token prediction, creating state and prediction streams.The state step contributes key-value entries to the cache, while the prediction step emits the next token.
- Experiments: 53M to 1.678B parameters were evaluated in pretraining experiments spanning common research scales.The experiments compare SPS with standard Transformers under token-equivalent and compute-equivalent settings.
- Results: SPS reaches lower training loss than standard Transformers even when the baselines receive double the training tokens.The study also controls for compute and memory overheads to test whether separation itself drives the gains.
2 Prediction and State Preparation
A standard Transformer hidden state simultaneously serves the current prediction and prepares keys and values for all later predictions. These roles create distinct gradient pressures that are entangled in one representation.
- Dual roles: Each hidden state contributes to the immediate prediction of x_i+1 and to key-value entries consumed by later positions.Causal attention lets subsequent positions read the stored representations.
- Gradient decomposition: The language-modeling loss averages per-position cross-entropy terms, while shared parameters accumulate gradients from current and future losses.Causality restricts each position’s effects to its current and later losses.
- Gradient interference: The immediate prediction and future state-preparation gradients both backpropagate through the same hidden state.This shared path optimizes one activation set for competing objectives.
3 The State-Prediction Separation Transformer
SPS separates state storage from prediction by interleaving input and <predict> positions. Persistent input representations carry long-range state, while prediction representations emit tokens and are windowed.
- Architecture: SPS interleaves each input token with a learned <predict> token, producing separate input and prediction representation streams.The two tokens share a position encoding, but serve different computational roles.
- Prediction stream: Prediction x_i+1 is computed at ρ_i, so training loss is applied only at prediction positions.Input positions carry state forward rather than directly emitting the next-token loss.
- Gradient routing: The persistent input stream receives future-loss gradients, while prediction gradients are concentrated on the prediction stream.Prediction representations can still receive at most w −1 following state-preparation contributions.
- Attention and memory: Input entries persist in the KV cache, whereas <predict> entries are retained only within a sliding window of size w.This gives later queries access to long-range state and short-range token-choice representations.
- Inference: SPS has negligible additional inference cost because paired input and prediction steps are forwarded jointly.Its persistent cache matches standard Transformer cache size, with a bounded w-slot prediction buffer.
4 Experimental Setup
The experiments compare SPS with standard and resource-matched ablations across five model scales using FineWeb-Edu pretraining. Evaluation covers validation loss, out-of-distribution corpus loss, and zero-shot accuracy.
- Baselines: SPS is compared with STANDARD, 2X MEMORY, DELAYED STATE, and REVERSE SPS using a shared Transformer backbone.The ablations test doubled memory, delayed computation, and reversed stream persistence.
- Model scales: Five scales are evaluated with a 4,096-token context length and matched parameter counts and hyperparameters across variants.Table 1 summarizes the model configurations.
- Training data: Each model is trained for 20B tokens by default on FineWeb-Edu, with all runs seeing the same data order.The XL model is trained for 47B tokens until matching SPS validation loss.
- Efficiency evaluation: Figure 2 reports FineWeb-Edu validation NLL against tokens seen and GPU-hours, with learning-rate cooldown included in the top row.The figure is designed to compare both data and compute efficiency.
- Metrics: Evaluation includes held-out validation NLL, averaged corpus NLL over four out-of-distribution corpora, and averaged zero-shot accuracy over five benchmarks.These metrics measure pretraining performance and generalization.
5 Results
SPS improves validation loss, data efficiency, generalization, and downstream accuracy across scales while preserving standard-like inference costs. Controlled comparisons show the gains come from separating state and prediction roles, not merely adding computation or memory.
- Quality and efficiency: −0.042 to −0.068 validation NLL gap versus STANDARD from XS to XL shows SPS improves at every scale.The gap widens with scale.
- Quality and efficiency: Roughly half the training data reaches STANDARD’s quality, while matched GPU-hours eventually favor SPS at every scale.Even doubling STANDARD’s pre-decay budget from 18B to 36B tokens does not close the gap.
- Generalization: 0.09–0.11 lower Corpus NLL and 2.3–3.1% higher zero-shot accuracy show that SPS improvements transfer to held-out corpora and benchmarks.Corpus NLL averages four corpora, while task accuracy averages five zero-shot benchmarks.
- Quality and efficiency: 1.01 peak-memory ratio and throughput within 6–10% of STANDARD indicate that SPS preserves comparable inference efficiency.The throughput measurement uses a combined 1k-token prefill and 3k-token decode workload on H100.
- Controlled comparisons: SPS remains better than 2X MEMORY despite half the persistent KV cache, ruling out simple capacity as the explanation.Keeping prediction entries persistent re-couples prediction and state-carrying roles.
- Controlled comparisons: 0.019–0.021 lower validation NLL than DELAYED STATE shows that SPS’s role assignment matters beyond an extra computation step.DELAYED STATE improves over STANDARD but commits persistent state at the prediction slot.
- Window and gradient analysis: w=64 is the empirical best by a small margin, while REVERSE SPS degrades sharply at small windows.Persisting input entries tolerates a wider range of prediction-window sizes.
- Window and gradient analysis: 1.4–2.2× larger NLL degradation after persistent-state ablation shows SPS’s state carries more future-relevant information at inference.The comparison uses a forced persistent-cache window of ω=64.
6 Related Work
Prior work addresses the prediction–state tension through extra input-side computation or richer prediction targets. SPS instead structurally separates these roles during standard left-to-right training.
- Prediction–state tension: Transformer hidden states simultaneously encode current next-token predictions and persistent state for later predictions.This shared representation creates the tension SPS targets.
- Input-side computation: Pause-token methods add computation before answering, using extra forward passes to increase model capacity.These approaches resemble SPS’s added computation but apply it on the input side.
- Prediction targets: Multi-token-prediction methods enrich the future-prediction signal beyond the immediate next token.They address the prediction target rather than separating persistent state from prediction computation.
7 Discussion
SPS separates prediction from state preparation with two interleaved streams, improving language-modeling and transfer results across scales. The evidence supports the mechanism, while compute, scale, and theoretical limitations remain.
- Results: At every scale from XS to XL, SPS lowers FineWeb-Edu validation NLL, improves held-out-corpus NLL, and raises zero-shot accuracy.These gains occur at the same persistent KV-cache footprint as STANDARD and within a few percent of its throughput.
- Mechanism: SPS routes future-loss gradient onto the input stream and produces a persistent state more impactful for future states than alternatives.Restricted-state and gradient-flow analyses support separation as the source of the improvement.
- Scope: The evidence uses a single FineWeb-Edu pretraining corpus and scales only to 1.678B parameters.Transfer beyond the corpus and continuation beyond 1.6B remain suggested trends requiring further verification.
- Limitations: SPS adds a prediction-stream slot per input position, roughly doubling per-step training compute over STANDARD.Lower-overhead separation, sparse persistent state, and stream-specific parameters remain open directions.
A Full Transformer Notation
The notation defines an autoregressive Transformer as a distribution over the next token, then builds its representation through layered attention and feed-forward blocks. A final normalized, weight-tied representation produces the next-token distribution.
- Autoregressive model: An autoregressive Transformer maps each token sequence to a probability distribution over the next vocabulary token.The vocabulary is finite, and the model defines p over sequences up to length T.
- Representations: Each token is embedded into a d-dimensional vector before processing by L Transformer blocks.The blocks combine causal multi-head self-attention, feed-forward networks, and normalization layers.
- Attention: Each attention head uses projection matrices and rotary positional transforms within causally masked attention.Head outputs are concatenated and mixed through an output projection.
- Output distribution: The next-token distribution comes from a final RMSNorm followed by a weight-tied unembedding of the final hidden representation.This connects the last-layer representation to vocabulary probabilities.
B Full Main Results
Table 3 expands the main results into corpus-level language-modeling losses, benchmark-level zero-shot accuracies, and throughput comparisons. It exposes the individual measurements behind the aggregate metrics.
- Corpus evaluation: Per-corpus NLLs are reported for WikiText, C4, Pile-Books3, and GovReport.These values are averaged into the aggregate Corpus NLL metric.
- Downstream evaluation: Per-benchmark zero-shot accuracies are reported for ARC-Easy, HellaSwag, PIQA, SciQ, and LAMBADA.These values are averaged into Task Accuracy in the main text.
- Efficiency: The table also includes the prefill-throughput ratio for comparing inference efficiency.This complements the quality metrics with a throughput measure.
C Seed Variance and Statistical Tests
The authors test whether SPS’s advantage survives changes in data order and initialization, then apply formal tests to separate robust effects from noise. At the focused S, 10B setting, SPS remains statistically better than every baseline.
- Seed robustness: Three seeds vary both packed-sequence order and weight initialization for STANDARD, DELAYED STATE, 2X MEMORY, and SPS.The robustness check uses the headline run plus seed 0 and seed 1.
- Statistical tests: All three baselines reject the null that SPS does not improve validation NLL at p < 0.005 under one-sided Welch’s t-tests.The alternative tested is that SPS has lower validation NLL.
- Ablation equivalence: TOST p = 3.4×10−3 establishes DELAYED STATE and 2X MEMORY as statistically equivalent within ±0.01 NLL.Their equivalence margin is below the gaps to SPS (0.0153) and STANDARD (0.034).