Free pause tokens cut 1B next-token loss 0.028 nats with near-zero inference cost

Free Pause Tokens

John Langford, Nathan Godey, Giovanni Monea, Yoav Artzi, Harry Dong, Ying Fan, Gustavo de Rosa, Zheng Zhan

cs.LG, cs.AI

2026-09-03

A KV-less prediction stream cuts 1B eval CE from 2.8957 to 2.8673. Phasing in at 42.5% costs 1.33× wall-clock; inference adds no position, cache, or step.

What problem this solves

A decoder Transformer stores one hidden state per position. That vector has to summarize the context for later positions and also be a good next-token predictor. The two jobs pull in different directions.

State-prediction separation (SPS) gives each job its own stream over a shared backbone. It lowers next-token loss, at a price: the prediction stream is a second pass, about 1.9× pretraining FLOPs. A flexible attention mask also refuses to run in stock FlashAttention, so wall-clock is worse than the FLOP ratio. An architecture that doubles training cost is not an upgrade if the old architecture can spend those FLOPs on extra tokens.

Method

The split is taken to the limit. The prediction stream writes no keys or values. It rides existing positions, which the authors call a free pause token. At every position it starts from one learned embedding, forms a query over the state stream's KV, and feeds the LM head. The loss sits only on the prediction stream. Inference adds no context length, no KV cache, and no extra decode steps. The only new parameter is that embedding.

Four cuts bring training cost down:

The testbed is a 1B decoder: 24 layers, width 1536, GQA 16/8, sliding window 2048 on most layers, sequence 8192. Data is Phi-4 derived. Global batch 524,288 tokens on 8×B200 with a Muon-family optimizer. The control is the same model with the prediction stream removed, matched on optimizer, data order, and batch.

Results

Iso-token, cooled at 100B:

schedule100B CEwall-clockvs full pause
control2.89571.00×-
full pause (w=0)2.86731.57×-
42.5% then pause2.86911.33×recovers 94%
75% then pause2.87561.14×recovers 71%

Full pause beats the control by 0.0284 nats, in the 2–3 centinat range. The phase switch has no loss spike. Most of the gain arrives within about 15B tokens after the switch.

Give the control the saved node-hours: the 75% schedule leads by 0.013 nats, 42.5% by 0.012, full pause by only 0.005. The cheapest schedule wins most at equal wall-clock. Under a stricter iso-FLOP (1.9×) reading, phased runs stay ahead and full pause goes slightly negative (+0.006).

Downstream, DCLM CORE moves from 0.327 (100B control) to 0.347 (full pause); climbmix bits-per-byte from 0.777 to 0.769. At equal compute, 100B pause matches a 150B control. Prefill costs the same as a vanilla Transformer. Decode fuses the two streams into one step; a B200 microbenchmark sits within about 1% of standard decode latency.

Why it matters

This is an iso-parameter, iso-token loss cut that also holds as an iso-compute (node-hour) win, and inference is nearly free. For a 1B run already in cooldown, turning on pause for the last 25% of tokens costs 1.14× wall-clock and recovers about 71% of the full-pause gain. The extra FLOPs sit inside a decode step that is usually depth-bound, the same spare parallel budget speculative decoding uses, spent on prediction quality rather than speed.

The gain is incremental. A few centinats are expensive to buy in pretraining and they inherit into every later stage. They are not a new architecture class, and they have been shown at 1B only.

Limitations

The authors flag a single scale (1B) and a single primary seed. Count FLOPs instead of wall-clock and full pause loses to a longer control. Shared-FFN runs used a slightly smaller global batch to fit a faster micro-batch; that control is already 0.011 nats worse, so the true iso-batch cost of sharing is smaller than the raw table, but it is still a quality-for-throughput trade. Per-task lm-eval scores sit inside 95% intervals; the downstream claim lives in pooled metrics. No combination with multi-token prediction. The learned pause embedding leans toward frequent continuations (comma, period, "the", "and"); the mechanistic story is thin.

Terms

Source

What people are saying

Related papers

All paper explainers