The State-Prediction Separation Hypothesis
AuthorsGiovanni Monea, Nathan Godey, Kianté Brantley, Yoav Artzi
Resources
This paper argues that transformers should split remembering from predicting, and that doing so can make language models train more efficiently and work better downstream.
Key results
largest SPS model size reported across the five scales
smallest SPS model size reported across the five scales
standard Transformer validation NLL at XS
SPS validation NLL at XS
standard Transformer validation NLL at XL
SPS validation NLL at XL
What the paper found
The State-Prediction Separation Hypothesis argues that a Transformer’s hidden state is overloaded: it must both predict the next token and store reusable state for later tokens. To test this, the authors at Cornell and Harvard introduce the State-Prediction Separation Transformer, which interleaves each input token with a dedicated <predict> slot so prediction gradients and state-preparation gradients flow through different streams. Across five model scales from 53M to 1.678B parameters, trained on FineWeb-Edu, SPS consistently beats a standard Transformer and two controls, 2X Memory and Delayed State. On FineWeb-Edu validation NLL, SPS improves from 3.107 to 3.065 at XS, from 2.858 to 2.810 at S, and from 2.458 to 2.390 at XL, while generalization also improves: corpus NLL drops by 0.09–0.15 and average zero-shot accuracy rises by 2.3–3.1 percentage points over ARC-Easy, HellaSwag, PIQA, SciQ, and LAMBADA. The key result is efficiency: SPS reaches Standard quality with roughly half the training data, matching Standard’s persistent KV-cache footprint and staying within about 6–10% of its throughput, with peak memory held at 1.01× Standard. The analysis shows the gain is structural, not just extra compute or memory, because SPS routes more future-loss gradient to the input stream and makes the persistent state 1.4–2.2× more important at inference than Delayed State.
Original abstract
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.
Read the original paperMore in Foundation Models
Browse all 47 papers →How Much Is an AI Token Worth? Scaling Laws for Wild AI-Generated Web Text
Jenna Russell, Ben Glickenhaus, Katherine Thai, John Wieting, Mohit Iyyer, Max Spero, Bradley Emi
AI-generated web text can help language models at first, but beyond a tipping point it degrades performance on human writing, making data filtering and separate evaluation increasingly important.
TabFM: A Zero-Shot Foundation Model for Tabular Data
Weihao Kong, Erez Louidor Ilan, Shuxin Nie, Taman Narayan, Rajat Sen, Yichen Zhou, Deqing Fu, Samet Oymak, Abhimanyu Das
TabFM is a large synthetic-data-trained model that aims to make accurate tabular predictions instantly, without retraining for each new dataset.
When Do Biological Reasoning Models Use Their Biological Inputs?
Ada Fang, Nikitha Thoduguli, Lukas Fesser, Hanlin Zhang, Sham M. Kakade, Marinka Zitnik
The study finds that many biological reasoning systems appear to succeed without meaningfully using the biological inputs they were designed to reason over.