mlsanityData to 7 Oct 2026

Almost Free State Prediction Separation

arXiv 2609.03807 ↗First posted 3 Sep 2026cs.LG cs.AIcited by 0

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

  1. 2020Scaling Laws for Neural Language Modelscited by 6,411
  2. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
  3. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  4. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  5. 2020Longformer: The Long-Document Transformercited by 3,362
  6. 2023Mistral 7Bcited by 3,065
  7. 2022Training Compute-Optimal Large Language Modelscited by 2,927
  8. 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,928
  9. 2022Fast Inference from Transformers via Speculative Decodingcited by 1,615
  10. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
  11. 2023Accelerating Large Language Model Decoding with Speculative Samplingcited by 975
  12. 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headscited by 821
17 more
  1. 2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 711
  2. 2024MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategiescited by 694
  3. 2024Phi-4 Technical Reportcited by 675
  4. 2024DataComp-LM: In search of the next generation of training sets for language modelscited by 381
  5. 2024Better & Faster Large Language Models via Multi-token Predictioncited by 344
  6. 2018Blockwise Parallel Decoding for Deep Autoregressive Modelscited by 318
  7. 2024Massive Activations in Large Language Modelscited by 277
  8. 2016Using the Output Embedding to Improve Language Modelscited by 266
  9. 2020Query-Key Normalization for Transformerscited by 266
  10. 2023Think before you speak: Training Language Models With Pause Tokenscited by 254
  11. 2023Language Modeling Is Compressioncited by 225
  12. 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 160
  13. 2023Microscaling Data Formats for Deep Learningcited by 132
  14. 2025Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-trainingcited by 62
  15. 2026FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalingcited by 58
  16. 2026The State-Prediction Separation Hypothesiscited by 2
  17. 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

  1. 2024Dynamic Vocabulary Pruning in Early-Exit LLMs
  2. 2026Dense Supervision Is Not Enough: The Readout Blind Spot in Looped Language Models
  3. 2026What to Preserve in Recursive Computation: A Local Predictive Sufficiency Principle
  4. 2025Pre-Training Curriculum for Multi-Token Prediction in Language Models
  5. 2026Compute Where it Counts: Self Optimizing Language Models
  6. 2023Learning to Skip for Language Modeling
  7. 2025One Jump Is All You Need: Short-Cutting Transformers for Early Exit Prediction with One Jump to Fit All Exit Levels
  8. 2024Dynamic layer selection in decoder-only transformers