Source-linked AI summary
Almost Free State Prediction Separation
John Langford, Nathan Godey, Giovanni Monea, Yoav Artzi, Harry Dong, Ying Fan, Gustavo de Rosa, Zheng Zhan
TL;DR
Standard transformers must use one representation for both context summarization and prediction, while state–prediction separation improves loss at roughly 1.9× pretraining FLOPs. This paper introduces a free pause stream that adds prediction computation without a new sequence position, reducing training overhead while preserving gains at equal tokens and compute.
Problem
Standard transformers use one position representation for both context summarization and next-token prediction, while state–prediction separation improves loss at roughly 1.9× pretraining FLOPs.
Method
A weight-shared prediction stream queries the state stream’s keys and values without writing its own, and training cost is reduced through shared FFNs, zero prediction self-window, and phased training.
Results
2.8673 versus 2.8957 yields a −0.0284 nats iso-token gain, while a 42.5% phased switch reaches 2.8691 at 1.33× wall-clock overhead.
Takeaways & Limitations
At equal node-hours, every free-pause variant is ahead; under stricter approximately 1.9× iso-FLOP accounting, phased runs retain a −0.005 to −0.009 advantage.
Takeaways & Limitations
Experiments are limited to a single 1B scale and a single primary seed.
Abstract
from arXiv · showhide
A free pause token gives a language model extra compute to form each next-token prediction (as a pause, or thinking, token does) but carries that compute in a parallel prediction stream over a weight-shared backbone rather than as an extra token in the sequence. It improves next-token prediction by 2-3 centinats in practice on a 1B parameter model. Because the pause rides an existing position instead of adding one, it is free to use: at inference it adds no context length, no KV cache, and essentially no latency with the growth in inference flops typically irrelevant as it is not the active bottleneck on throughput. The only primary cost is in training, where additional training compute versus an optimized pretraining pipeline is reduced to as low as x1.14 while preserving most of the benefits. The result is an isoflop, isoparameter, and isotoken improvement over standard next token trained transformers.
1 Introduction
The paper separates summarizing context from predicting the next token using a weight-shared prediction stream that adds no sequence position. This free pause improves loss while reducing the training-cost burden of prior state–prediction separation.
- A standard decoder state must both summarize prior tokens and predict the next token, creating competing representational demands.
- Prior state–prediction separation improves next-token loss but costs roughly 1.9× the pretraining FLOPs.
- The free pause gives prediction its own weight-shared stream, initialized by one learned embedding and querying the state stream’s keys and values without writing any.
- Because the prediction stream adds no keys, values, or sequence position, the pause is free at inference and adds only the learned pause embedding as a parameter.
- Against a matched control, tail phasing lowers next-token cross-entropy throughout training and remains below the control’s compute-for-loss frontier at equal wall-clock compute.
2 Method
The method separates context summarization from next-token prediction using two weight-shared streams, while reducing the added training cost through fused passes, zero prediction self-window, shared FFNs, and phased training.
- Two-stream architecture: The state stream produces persistent keys and values, while the prediction stream uses a shared learned vector and queries the state stream to produce next-token predictions.The prediction stream writes no keys or values of its own, so it adds no sequence position.
- Training efficiency: The two-pass training split enables FlashAttention-friendly state and prediction passes instead of an inefficient interleaved mask.The state pass runs first, producing cached keys and values for prediction cross-attention.
- Shared gated FFN: One FFN evaluation per position replaces two by pooling the streams with scalar gates and routing the shared output back into each stream.The gates are negligible single-output linears initialized at 0.5.
- Training efficiency: Setting the prediction attention window to w=0 removes a second merged attention call, while incurring only a millinats performance cost.This design choice avoids prediction-written keys and values.
- Phasing: Training can start as standard next-token prediction and switch to the separated architecture for only the final 1−f fraction of the schedule.The backbone optimizer, schedule, dataloader, and training state continue across the switch; only predict_embedding is initialized fresh.
3 Related work
The free pause extends state–prediction separation and pause-token computation without adding sequence positions, while differing from speculative decoding in its purpose.
- State–prediction separation: The free pause is the w=0 limit of state–prediction separation: its prediction stream forms queries over state keys and values but writes nothing itself.Earlier state–prediction separation retained prediction keys and values within a sliding window.
- Pause and thinking tokens: Unlike a conventional pause token, the free pause uses a parallel stream, leaving context length, KV-cache size, and decode-step count unchanged.A conventional pause token occupies a new sequence position and enlarges all three.
- Speculative decoding: Speculative decoding uses extra queries for speculative future continuations, whereas the free pause predicts the current position and keeps that prediction.The free pause improves prediction quality rather than generation speed.
4 Model and optimization
Experiments use a 1B decoder transformer and a matched strong baseline, with identical optimization, data order, schedule, and global batch except for the prediction stream.
- Model: The model is a 1B decoder transformer with 24 layers, hidden size 1536, grouped-query attention, sliding-window attention, periodic full-attention layers, and sequence length 8192.It also uses QK-normalization, partial rotary embeddings on full-attention layers, and tied input/output embeddings.
- Optimization: Optimization uses a Muon-family optimizer with a warmup–stable–cooldown schedule, global batch 524,288 tokens, bf16 activations, and mxfp8 matrix multiplications.The schedule ends with a final-25% linear cooldown.
- Baseline: The strong baseline is the identical model with the prediction stream removed, matched on optimizer, data order, schedule, and global batch.The only architectural difference is the added second pass and its resulting wall-clock throughput.
5 Measurement
The prediction pass costs about 1.9× FLOPs but only about 1.57× wall-clock time because it skips key-value projections and runs at higher utilization.
- 1.9× FLOPs is the prediction pass cost, rather than 2×, because it skips key-value projections.
- 1.57× wall-clock overhead is lower than the FLOP overhead because the compute-dense pass reaches about 20% higher utilization.
6 Results
The free pause improves next-token loss and remains advantageous at equal tokens or compute, while shared FFNs and phasing reduce its training overhead and inference remains nearly unchanged.
- Loss results: 2.8673 versus 2.8957 yields a −0.0284 nats iso-token gain for the free pause over the strong control.The comparison uses matched 100B-token runs.
- Shared gated FFN: The shared gated FFN reduces full-pause overhead from 1.57× to 1.35× wall-clock and gives up approximately 0.005–0.010 nats against the two-pass form.It evaluates the position-wise FFN once per position rather than once per stream.
- Phasing: Switching at 42.5% reaches 2.8691, within noise of the full pause’s 2.8673, while reducing overhead from 1.57× to 1.33× wall-clock.The phased run tracks the control until switching, then diverges below it without a visible loss spike.
- Memory and throughput: 14 GB/GPU less peak memory enables a faster micro-batch, raising throughput to 0.74× control versus 0.64× for the two-pass form.The shared form uses 91 GB rather than 105 GB at sequence length 8192 and micro-batch 4.
- Iso-compute analysis: At equal node-hours, every free-pause variant leads the control, with advantages ranging from −0.013 for the shortest phase to −0.005 for the full pause.Under stricter approximately 1.9× iso-FLOP accounting, phased runs remain positive at −0.005 to −0.009 while the full pause becomes +0.006.
- Inference: Inference remains close to vanilla decoding: a fused two-stream decode is within approximately 1% of standard decode latency.The prediction reads the state stream’s newly appended keys and values while both streams advance together.
7 Conclusion
Within the tested scope, the method offers a modest isotoken, isoparameter, and compute advantage, while its learned pause embedding has a concentrated set of unusually large coordinates.
- 7 Conclusion: A single 1B scale and primary seed limit the experiments, although the authors report the results as beyond the noise level.The conclusion explicitly bounds the evidence to one model scale and seed.
- 7 Conclusion: The learned pause embedding has a near-Gaussian bulk, but its top-10 dimensions contain 11% of the energy and several channels reach approximately 4σ.The measured embedding uses RMS 0.88, with outlying coordinates around ±3–4.
- 7 Conclusion: The approach retains essentially all of the gain at 1.33× wall-clock training cost and reaches 1.09× when some gain is traded away, while inference stays free.These costs are contrasted with the roughly 1.9× pretraining-FLOP cost of the original separation.
A Engineering
The serving implementation shares the state stream’s paged KV cache with the prediction stream, avoiding an extra cache while executing both streams together.
- A Engineering: The state stream writes paged KV while the prediction stream reads the same layer’s KV read-only through cross-layer KV-sharing, so no extra cache is needed.All weights are shared, and both streams run as one 2T-row batch so each layer’s weights load once.
B Cross-entropy summary
The cross-entropy summary reports lower loss for pause variants than the listed control baselines, with the best value among these entries from pause w=64 at 100B.
- B Cross-entropy summary: 2.8626 is the lowest listed cross-entropy, achieved by pause w=64 at 100B.The same table lists pause w=0 at 2.8673 and phased 42.5% at 2.8691.
- B Cross-entropy summary: 2.8957 is the listed cross-entropy for the 100B control baseline, compared with 2.8742 for the extended 150B control.These are the two control entries in the summary.
- B Cross-entropy summary: 2.8756 is the listed cross-entropy for the phased 75% variant.This value is reported alongside the other pause and control variants.