Almost Free State Prediction Separation
Not on the map. Why?
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 29 of its 32 references that are in the map. First posted on arXiv 3 Sep 2026.
Builds on
29 of 32 references.
References · 29
- 2020Scaling Laws for Neural Language Modelscited by 6,411
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2020Longformer: The Long-Document Transformercited by 3,362
- 2023Mistral 7Bcited by 3,065
- 2022Training Compute-Optimal Large Language Modelscited by 2,927
- 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,928
- 2022Fast Inference from Transformers via Speculative Decodingcited by 1,615
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2023Accelerating Large Language Model Decoding with Speculative Samplingcited by 975
- 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headscited by 821
17 more
- 2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 711
- 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 694
- 2024Phi-4 Technical Reportcited by 675
- 2024DataComp-LM: In search of the next generation of training sets for language modelscited by 381
- 2024Better & Faster Large Language Models via Multi-token Predictioncited by 344
- 2018Blockwise Parallel Decoding for Deep Autoregressive Modelscited by 318
- 2024Massive Activations in Large Language Modelscited by 277
- 2016Using the Output Embedding to Improve Language Modelscited by 266
- 2020Query-Key Normalization for Transformerscited by 266
- 2023Think before you speak: Training Language Models With Pause Tokenscited by 254
- 2023Language Modeling Is Compressioncited by 225
- 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 160
- 2023Microscaling Data Formats for Deep Learningcited by 132
- 2025Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-trainingcited by 62
- 2026FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalingcited by 58
- 2026The State-Prediction Separation Hypothesiscited by 2
- 2026Decode-Branch Transformers: Decoupling the Primary Prefill Path from Additional Decode Computationcited by 1
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2024Dynamic Vocabulary Pruning in Early-Exit LLMs
- 2026Dense Supervision Is Not Enough: The Readout Blind Spot in Looped Language Models
- 2026What to Preserve in Recursive Computation: A Local Predictive Sufficiency Principle
- 2025Pre-Training Curriculum for Multi-Token Prediction in Language Models
- 2026Compute Where it Counts: Self Optimizing Language Models
- 2023Learning to Skip for Language Modeling
- 2025One Jump Is All You Need: Short-Cutting Transformers for Early Exit Prediction with One Jump to Fit All Exit Levels
- 2024Dynamic layer selection in decoder-only transformers