The State-Prediction Separation Hypothesis

2026-07-01

AI Summary

This work addresses a fundamental limitation in conventional Transformers, where state retention and next-token prediction are tightly coupled within a single computational stream, constraining modeling efficiency. The authors propose the state-prediction disentanglement hypothesis and introduce a novel dual-stream Transformer architecture that explicitly decouples state maintenance from prediction. Through large-scale pretraining, gradient analysis, and systematic ablation studies, they provide the first empirical validation that such architectural disentanglement substantially enhances performance. Across multiple model scales, the proposed method consistently achieves lower validation loss and improved data and computational efficiency, outperforming standard Transformers by an average of 2–3 percentage points on downstream tasks.

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.

Problem

• state-prediction separation
• Transformer
• language modeling
• computation stream
• next token prediction

Innovation

• state-prediction separation
• dual-stream Transformer
• language modeling efficiency
• gradient analysis
• compute efficiency