Full-bandwidth transformer
Xi Wang, Ziyang Cai, Zheng Zhan, Harry Dong, Ying Fan, Gustavo de Rosa, Tim Pearce, John Langford
cs.AI
2026-08-10
Microsoft fuses the previous top-layer hidden state back into each decoding step at under 1% overhead; a 1B full-bandwidth transformer matches a standard model trained on 1.5x the tokens and approaches 5x-token baselines on math.
Autoregressive transformers compute along two axes: horizontally across tokens and vertically through depth. The horizontal axis is effectively full-bandwidth, since attention lets each token read the whole past. The vertical axis is narrow. At each step the entire top-layer hidden state, a D-dimensional vector, is compressed into a single discrete token and fed back, while the top-layer state itself is discarded. Deeper intermediate states persist in the KV cache but are depth-frozen: a shallow layer of a new token only sees a partially processed view of the past, and the already-computed deep states can never return to the bottom of the stack for further work. The model must either verbalize intermediate state (chain of thought) or recompute it at every position. As unique high-quality data runs low, the question becomes whether more learning signal can be squeezed from each token without adding data.
The core idea is latent feedback. At each decoding step, the previous top-layer hidden state is fused with the current token embedding through a gated linear unit and fed back to the bottom of the stack:
et (x) h{t-1} = WU h{t-1} sigma(WG et)
The design is deliberately asymmetric: the hidden state rides the value pathway and the token acts only as a multiplicative gate. A symmetric additive fusion would open a shortcut where the model suppresses the state pathway, recovers the plain token input, and reaches ordinary pretraining loss while leaving the wide channel unused. The gate makes reading the state mandatory. The added cost is two D-by-D matrix multiplies, under 1% per token; the architecture, KV cache, and serving stack are untouched, and it is vLLM-compatible.
The hard part is training. A pretrained checkpoint has never seen hidden states at its input, so feedback cannot simply be switched on at inference. The recurrence is also sequential across positions, so training on it directly forfeits the parallel teacher forcing that makes transformers efficient. The authors approximate it with multiple passes (temporal parallelism): each pass shifts the previous pass's states one position right, fuses them, and re-runs the full stack in parallel, paying sequentiality across a few passes rather than across the sequence. The schedule is progressive: most of training uses the ordinary single pass, and feedback passes enter only late in training, at a typical mix of 75% one-pass, 22% two-pass, 3% three-pass batches. A counterintuitive finding is that a model trained with only one- and two-pass batches diverges once feedback depth exceeds what it saw, while adding just 3% three-pass batches turns the learned map into a contraction toward a fixed point that stays stable out to k=1000 passes. A few stabilization tricks (prefix mixin, depth scaling, weight tying, jitter noise at sigma=0.02) round it out.
1B-parameter models on the Phi-4 data mix, trained up to 400B tokens, with baselines up to 1T.
| Metric | Standard decode | Latent feedback | Reference |
| Math500 Pass@1 (200B) | 0.27 | 0.37 | beats 1T baseline |
| HumanEval Pass@3 (200B) | 0.31 | 0.34 | |
| GSM8K Pass@1 (200B, instr-tuned) | 64.5 | 67.9 (FUSED) | |
| GSM8K Pass@1 (400B, instr-tuned) | 67.9 | 71.8 (FUSED) | 1T std 70.13 |
| MATH-500 (400B, FUSED) | 46.0 | 48.4 | 1T std 47.4 |
| MBPP Pass@3 (200B, FUSED) | 38.4 | 41.2 | 1T std 41.9 |
On validation loss and 5-shot LM Eval, a 100B full-bandwidth model with two fused prefill passes matches a 200B standard model, and 200B matches 400B, roughly 2x data efficiency. On base models, latent feedback often produces shorter reasoning traces at equal or better accuracy (the median Math500 token count drops), exactly what the widened channel predicts: intermediate computation that would otherwise be verbalized rides the hidden state. A controlled probe shows that one feedback step lifts a layer-0 linear probe from near-chance to 99.6% (completion tracking) and 100% (delayed memory).
It drops into an existing stack with no architecture or serving change and two lines of decoding loop, which is the easiest path to adoption. It also opens a new scaling axis: spend inference compute, or an auxiliary training objective, to extract more signal from the same tokens, which matters as unique data runs out. Even when feedback is not used at decoding time, training with the feedback objective improves the model's representations, so it is also a training-efficiency win.
Results are only at 1B scale; whether gains grow or shrink for larger models is unknown. It needs a dedicated training recipe and cannot be switched on for a stock checkpoint. The most practical benefit, shorter reasoning, disappears after instruction tuning because SFT data is produced in a verbose token-by-token style that re-imposes verbalization; preserving conciseness with on-policy post-training is left to future work. The headline 1.5x is the baseline; the best numbers usually require FUSED (doubled prefill), which is not always free.