The SPS Transformer: Solving the Secret Identity Crisis of Hidden States

TL;DR

A hidden state in a Transformer is a multitasker with a problematic workload: it must predict the very next word and prepare a "summary" (KV cache) for words that haven't been written yet. The State-Prediction Separation (SPS) Transformer solves this by giving these two jobs separate computation streams. The result? A model that learns twice as fast, hits higher accuracy (+2-3%), and keeps inference costs almost exactly the same.

The Problem: The "Dual-Role" Bottleneck

In standard architectures, every time a model processes a token, it generates a representation $h_i$. This $h_i$ is immediately used to calculate the loss for the next token ($j=i$). However, it is also saved into the KV cache to be read by every future token ($j>i$).

The authors argue this creates interference. The gradient for the immediate prediction often clashes with the gradient needed to build a robust long-term state. Imagine trying to write a sentence while simultaneously trying to memorize the entire plot of a book—doing both at once makes you worse at both.

Methodology: Structuring the Divorce

The SPS Transformer introduces a "State Stream" and a "Prediction Stream" by interleaving the input with a special <predict> token.

1. Input Stream (State): Processes the actual tokens. These are stored permanently in the KV cache.
2. Prediction Stream (Ephemeral): Processes the <predict> tokens. These are used to emit the next word but are discarded (evicted from the cache) shortly after.

By using an attention mask, the authors ensure that the "State" tokens accumulate gradients from all future positions, while the "Prediction" tokens focus almost entirely on the current loss.

Experimental Results: Faster Learning, Better Generalization

The researchers tested this across five scales (from 53M to 1.678B parameters) using the FineWeb-Edu dataset. The findings were consistent:

• Data Efficiency: SPS models reach the validation loss of a standard Transformer using half the tokens.
• Downstream Power: Average zero-shot accuracy improved by 2.3% to 3.1%.
• Inference Prowess: Unlike "Pause Tokens" which increase memory usage, SPS maintains the same persistent KV cache size as a standard Transformer.

Deep Insight: Why Does It Work?

The authors performed a "Gradient Ratio" analysis to prove their hypothesis. In standard models, the gradient signal for "future state" and "current prediction" is mixed. In SPS, the Input stream sustains significantly higher future-loss gradients, whereas the Prediction stream stays focused on the now.

When the authors tried to restrict the persistent memory of an SPS model, the performance collapsed much harder than in a standard model—proving that the SPS state is actually carrying more valuable, long-term information.

Critical Analysis & Conclusion

Takeaway

SPS proves that the "Next-Token Prediction" objective, while powerful, forces an architectural compromise. By structurally separating the "State" and "Prediction" roles, we can unlock latent performance in the Transformer backbone without increasing the permanent memory footprint.

Limitations

• Training Compute: While inference is cheap, training compute is roughly doubled because the context length effectively doubles.
• Parameter Sharing: In this study, both streams share the same parameters. Exploring distinct weights for each stream could yield even higher gains.

Future Outlook

As high-quality human data becomes scarce, the 2x data efficiency of SPS is a game-changer. This research opens the door to "asymmetric" Transformers—where the Prediction stream might be much shallower or narrower than the State stream, potentially slashing the training overhead while keeping the quality gains.