# Almost Free State Prediction Separation

arXiv: https://arxiv.org/abs/2609.03807 · HTML: https://mlsanity.com/paper/2609.03807/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 32 references that are in the map.
- First posted on arXiv 3 Sep 2026.

## Builds on (most cited first, 29 of 29)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Accelerating Large Language Model Decoding with Speculative Sampling](https://mlsanity.com/paper/2302.01318/) (2023)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Phi-4 Technical Report](https://mlsanity.com/paper/2412.08905/) (2024)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [Blockwise Parallel Decoding for Deep Autoregressive Models](https://mlsanity.com/paper/1811.03115/) (2018)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Think before you speak: Training Language Models With Pause Tokens](https://mlsanity.com/paper/2310.02226/) (2023)
- [Language Modeling Is Compression](https://mlsanity.com/paper/2309.10668/) (2023)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Microscaling Data Formats for Deep Learning](https://mlsanity.com/paper/2310.10537/) (2023)
- [Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training](https://mlsanity.com/paper/2504.13161/) (2025)
- [FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling](https://mlsanity.com/paper/2603.05451/) (2026)
- [The State-Prediction Separation Hypothesis](https://mlsanity.com/paper/2607.01218/) (2026)
- [Decode-Branch Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation](https://mlsanity.com/paper/2608.12385/) (2026)

## Built on by (most cited first, 1 of 1)

- [The Surprising Effectiveness of Shared Memory in Looped Transformers](https://mlsanity.com/paper/2610.02383/) (2026)

## Abstract (arXiv)

State--prediction separation (SPS) relieves a language model's hidden state of two competing burdens---summarizing the context and predicting the next token---by splitting the forward pass into a state stream and a prediction stream. The separation works, but it is expensive: the prediction stream is a second pass over the whole backbone, costing ~1.9× the pretraining FLOPs, and even more in terms of wall-clock time when using a flexible attention mask. This paper makes state--prediction separation almost free. We take the separation to its limit with a free pause token: a prediction stream that writes no keys or values at all and so rides the sequence's existing positions. It improves next-token prediction of a standard Transformer by 2-3 centinats in practice on a 1B parameter model, and because it adds no position it costs nothing at inference---no added context length, no KV cache, no decode steps, and essentially no latency, with the growth in inference flops typically irrelevant as it is not the active bottleneck on throughput. The cost is therefore entirely in training where we use four mechanisms to drive it down: a two-pass split that keeps FlashAttention kernels viable, the w=0 prediction window, a shared gated FFN that evaluates one FFN per position rather than one per stream, and phasing the separation onto the tail of the run. Together these bring the overhead versus an optimized pretraining pipeline to 1.33× wall-clock while recovering ~94% of the gain compared to SPS, and to as low as 1.09× along a graceful quality/compute tradeoff. Furthermore, the FFN optimization reduces the raw flops required at inference time. The result is an isoflop, isoparameter, and isotoken improvement over standard next token trained transformers.