Source-linked AI summary

Full-bandwidth transformer

Xi Wang, Ziyang Cai, Zheng Zhan, Harry Dong, Ying Fan, Gustavo de Rosa, Tim Pearce, John Langford

arXiv:2608.08888v1cs.AI

TL;DR

Standard transformers have a narrow vertical feedback channel because top-layer states are discarded between decoding steps. This paper introduces latent feedback and finds that full-bandwidth transformers match baselines trained on twice as many tokens while often producing shorter reasoning traces at equal or better accuracy.

  • Problem

    Standard decoding leaves past representations depth-frozen, limiting how fully processed computation can re-enter the model across steps.

  • Method

    Full-bandwidth transformers fuse each previous top-layer hidden state with the sampled token embedding and train this recurrence through scheduled multi-pass parallel updates.

  • Results

    Full-bandwidth transformers match no-recurrence baselines trained on twice as many tokens in validation loss and multiple-choice accuracy, while improving free-form generation across math and coding tasks.

  • Takeaways & Limitations

    Latent feedback converts negligible extra inference compute into broader performance gains and often shorter reasoning traces at equal or better accuracy.

  • Takeaways & Limitations

    Experiments are limited to 1B-parameter models, and the feedback-pass schedule is heuristic.

Abstract

from arXiv · show

Autoregressive transformers compute along two axes: horizontally across generated tokens, and vertically through model depth. Dense attention gives each token broad horizontal access to the past, but the vertical feedback channel between decoding steps remains narrow: only the sampled token returns to the bottom of the stack, while the top-layer hidden state is discarded. We introduce the \emph{full-bandwidth transformer}, which widens this channel with \emph{latent feedback}: at each decoding step, the previous top-layer hidden state is fused with the sampled token embedding through a gated linear unit and fed back as the next input. Latent feedback lets non-verbalized computation re-enter the stack with a renewed depth budget, while preserving the standard transformer architecture, KV cache, and language-modeling objective. To train full-bandwidth transformers without losing parallel teacher forcing, we use a scheduled multi-pass objective that introduces latent feedback late in pretraining and mixes a small fraction of deeper feedback passes for stability. We train 1B-parameter full-bandwidth transformers up to 400B tokens and find that latent feedback improves validation loss, 5-shot language-model evaluation, math and coding generation, and instruction-tuned performance. With negligible per-token decoding overhead, full-bandwidth transformers match or approach standard transformers trained with roughly $1.5\times$ more tokens, and manage to produce shorter reasoning traces at equal or better accuracy.

1 Introduction

The paper introduces full-bandwidth transformers, which feed back the previous top-layer hidden state alongside the sampled token to widen inter-step communication. A scheduled multi-pass training regime enables this latent feedback while preserving efficient teacher forcing, yielding broad evaluation gains with negligible decoding overhead.

  • Motivation: Latent feedback widens the narrow autoregressive channel, allowing non-verbalized uncertainty, partial results, and plans to re-enter the stack with a renewed depth budget.Standard decoding passes back only the sampled token, while deeper past hidden states are otherwise unreachable.
  • Method: Latent feedback fuses the previous top-layer hidden state with the sampled token embedding through a gated linear unit and feeds the result back as the next input.The state uses the value pathway, while the token acts as the gate; the sampled token remains available for ordinary text generation.
  • Training: A multi-pass training regime introduces hidden-state feedback during training while retaining parallel teacher forcing efficiency.Each pass shifts the previous pass’s hidden states one position rightward, fuses them with token embeddings, and reruns the model.
  • Results: Negligible extra inference compute produces gains equivalent to substantially more training data, with recurrence-trained models matching no-recurrence baselines trained on twice the tokens in validation loss and multiple-choice accuracy.The comparison uses multiple forward passes for prefill.
  • Results: Latent feedback improves free-form generation over standard decoding with the same weights on GSM8K, Math500, HumanEval, and MBPP.The reported tasks span mathematics and code generation.

2 Background

Dense-attention transformers preserve explicit representations of all past tokens through KV caching, giving new states direct access to the full history. However, information flows broadly across positions but remains depth-frozen vertically, motivating a wider feedback channel during sequential decoding.

  • KV cache: KV caching reuses previously computed keys and values, allowing each new hidden state to attend directly to every earlier token.Unlike RNNs and state-space models, transformers retain explicit representations of the full past rather than compressing history into a fixed-size recurrent state.
  • Bandwidths of a model’s horizonal axis vs. veritical axis: Dense attention provides effectively full-bandwidth horizontal access across positions, while vertical access across depth is restricted.A new token’s layer-ℓ state can read cached representations of earlier positions, but the depth-wise routing constraint limits information flow downward.
  • Bandwidths of a model’s horizonal axis vs. veritical axis: Past computation is depth-frozen: a representation produced at layer ℓ can be read only by layers above ℓ and cannot be routed back down.Thus, shallow layers of new tokens see only partially processed views of past positions even when deeper states are already cached.
  • Motivation: The same depth-wise dependency constraint enables parallel training across positions but offers little benefit during already sequential token generation.Training requires sequential computation across layers rather than tokens, whereas decoding opens the possibility of richer dependencies on past hidden states.

3 Widening the bandwidth with latent feedback decoding

Latent feedback widens decoding’s vertical channel by fusing each sampled token embedding with the previous top-layer hidden state before re-entering the transformer. The gated mechanism preserves standard architecture and KV-cache serving while improving computational accessibility and enabling more compact reasoning, with training approximated through parallel multi-pass updates.

  • Core mechanism: Latent feedback feeds the previous top-layer hidden state back into the next input alongside the sampled token embedding.Standard decoding is the special case where only the sampled token crosses between steps.
  • Core mechanism: The asymmetric gated fusion places the hidden state on the value pathway and uses the token embedding only as a multiplicative gate, blocking an additive shortcut to ordinary token inputs.This design prevents the model from suppressing the state pathway and recovering plain-token pretraining behavior.
  • Implementation: Under 1% per token, latent feedback adds inference cost independent of context length and model depth while preserving the input dimension, architecture, KV-cache layout, and serving stack.The extra computation is the fusion’s two D × D matrix multiplications.
  • Computational role: Latent feedback improves computational accessibility by letting every layer, including the lowest, read the full history rather than only a partially processed context view.The gain is computational rather than informational because the latent is a deterministic function of the token history.
  • Computational role: Latent feedback provides implicit draft space, allowing intermediate results to be updated through depth rather than only by extending the token sequence and predicting shorter reasoning rollouts.It does not add asymptotic decoding depth: T tokens still cost O(TL), and past states remain readable in the KV cache rather than being overwritten.
  • Training: Temporal parallelism approximates the sequential recurrence by shifting previous-pass hidden states rightward, fusing them with token embeddings, and rerunning the full stack in parallel across positions.Each additional pass performs a Jacobi-style update, retaining parallel teacher forcing while training on latent-feedback inputs.

4 Experiments

Experiments show that latent feedback improves validation and language-model evaluation, with most prefill gains arriving after the first recurrence step. It also improves math and coding generation, can match larger standard-training baselines, and often yields shorter reasoning traces without added per-token decoding cost.

  • Prefill feedback: Most improvement appears after the first fused prefill pass, while additional passes continue helping with diminishing returns.The first pass makes top-layer hidden states available at the input, acting as added effective depth for the prompt.
  • Prefill feedback: With two feedback passes, the 100B-token full-bandwidth transformer reaches the 200B-token standard baseline, while the 200B-token model reaches the 400B-token standard baseline.In this regime, fused prefilling corresponds to roughly 2× pretraining data efficiency.
  • Open-ended generation: SOFT improves over STANDARD on every evaluated open-ended generation task, while SOFT is strongest on math and FUSED is often strongest on coding.STANDARD and SOFT have identical prefill cost; FUSED doubles prefill cost but preserves the same per-token decoding cost.
  • Open-ended generation: 64.5 to 67.9 on GSM8K and 42.5 to 45.9 on HumanEval show that SOFT and FUSED improve over STANDARD and outperform the matched 200B-token standard baseline.On MBPP, FUSED reaches 41.2 versus 41.9 for the 1T-token standard baseline.
  • Reasoning traces: SOFT decoding often produces markedly shorter reasoning traces than STANDARD at equal or better accuracy on the base model.The shortening effect disappears after instruction tuning.

5 Related work

The paper situates full-bandwidth transformers among methods that add computation during sequential decoding, latent reasoning, recurrent-network training, and data-efficient pretraining. Its distinguishing choices include external latent reinjection without architecture changes and lower training and inference overhead than some latent-reasoning alternatives.

  • Alleviating the depth bottleneck at decoding time: Feedback Transformer mixes representations from each layer so future attention uses aggregated representations rather than same-layer key values.
  • Alleviating the depth bottleneck at decoding time: T 2MLR and Latent Recurrent Transformer inject or store hidden representations across positions, whereas this approach reinjects latents externally through input construction.The external reinjection introduces no architecture changes and uses fewer additional parameters.
  • Latent and continuous reasoning: Coconut and Soft Thinking also feed latent representations into context, but this work focuses on pretraining and augments rather than replaces discrete-token generation.Their multi-pass training doubles input length and KV-cache size, creating more training and inference overhead than the full-bandwidth transformer.
  • Parallel training of recurrent networks: Parallel recurrent-network methods commonly use linear architectures such as Mamba or Gated Deltanet, hybridized with transformer layers to compensate for missing representational capacity.
  • Data-efficient pre-training: The work also targets data-efficient pretraining, alongside representation objectives and loop transformers that use extra training FLOPs for repeated inference computation.Unlike loop transformers, it differs in where additional computation is paid at inference time.
  • Data-efficient pre-training: These methods reflect broader scaling beyond parameters and tokens as token-per-parameter ratios and high-quality unique data become binding constraints.

6 Limitation

The work is limited to 1B-parameter models and uses a heuristic schedule for feedback passes, leaving larger-scale validation and more rigorous recurrence-training ablations for future work.

  • Scale: The experiments cover only 1B-parameter models, so the approach was not verified at larger scales.The authors suggest deeper models might benefit more because their top-layer hidden states could contain richer information.
  • Training schedule: Feedback-pass scheduling is heuristic, motivating more rigorous ablations of recurrence-training length.

A Model architecture

The model is a decoder-only causal language model with tied embeddings and output head, 24 transformer layers, and a 1,536-dimensional hidden state. It combines SiLU GLU feed-forward blocks with gated grouped-query attention, RMS normalization, rotary positions, and mixed sliding-window/full attention over an 8,192-token context.

  • The decoder-only model ties a 100,352-token embedding to its output head and uses 24 transformer layers with a 1,536-dimensional hidden state.
  • Its feed-forward blocks are 6,656-dimensional SiLU GLUs, while attention uses 16 query heads, 8 shared key/value heads, and headwise gates.
  • Attention applies QK RMS normalization and rotary positions over an 8,192-token context, with most layers using a 2,048-token sliding window and every sixth layer using full attention.

B Comparison of LM eval performance with other models of similar scale

This section compares 0-shot language-model evaluation performance between EvoLM and other open-source models of similar scale.

  • B Comparison of LM eval performance with other models of similar scale: Table 2 reports 0-shot LM Eval performance for EvoLM alongside other open-source models.The comparison numbers for EvoLM and other open-sourced models are adopted from Qi et al. (2025).

C Full pseudo code for training · D vLLM compatibility

The training appendix provides one-step pseudocode with k passes and documents a vLLM implementation that preserves latent feedback during CUDA-graph decoding. The implementation stores each request’s latest hidden state, fuses it with the next token embedding, and recycles the resulting state within the same model.

  • C Full pseudo code for training: The training procedure is specified as one step containing k passes.
  • C Full pseudo code for training: The full training code includes a normalization layer and regularization noise.
  • D vLLM compatibility: vLLM retains each request’s latest trunk hidden state and copies it into a persistent, fixed-address model buffer before the next decode step.
  • D vLLM compatibility: This buffer arrangement allows CUDA graphs to capture the GLU cross gate inside the forward pass.
  • D vLLM compatibility: A patched GPUModelRunner._model_forward stores detached hidden states in a request-ID-keyed dictionary and maps packed rows using query_start_loc.
  • D vLLM compatibility: The forward function fuses the saved state with the next token embedding through the learned GLU cross gate and recycles the resulting hidden state.
  • D vLLM compatibility: Unlike EAGLE and MTP, the model feeds its own state back into the same model to define the actual next-token distribution.

E Extended extrapolation results

Extended extrapolation experiments show that feedback remains stable up to 1,000 passes, far beyond the 3 passes used during training.

  • E Extended extrapolation results: 1,000 feedback passes remain stable in the extended extrapolation experiment.The experiment extends the number of feedback passes to 1,000.
  • E Extended extrapolation results: The stable extrapolation extends far beyond the 3 feedback passes used during training.The figure compares the extrapolation range with the 3-pass training setting.
  • E Extended extrapolation results: The extended experiment tests feedback-pass extrapolation beyond the training-time setting.Figure 10 extends the number of feedback passes to assess stability outside the training range.

F Explanation on state tracking tasks

The section evaluates state tracking with matched synthetic tasks that vary distractor length while preserving targets, testing relational computation, delayed memory, and multi-register latest-write retrieval. Linear probes and controlled recurrent-prefilling experiments measure where this information is represented across model depth.

  • Synthetic state-tracking setup: Matched examples append 0, 8, 32, 128, or 256 semantically null scratch updates, varying sequence length without changing the target.At the final colon, the study records the layer-0 input and every Transformer block’s output.
  • Completion tracking: Completion tracking requires predicting DONE when required and completed counts match, and MORE otherwise, using balanced assignments across numeral fields and labels.Matched pairs share distractors and scratch context; only the counter relation changes.
  • Delayed memory: Delayed memory requires retaining an initially assigned binary state through identical, label-independent scratch operations before predicting ZERO or ONE.This isolates persistent transport of a specified state, unlike completion tracking’s relational computation across fields.
  • Multi-register latest-write tracking: Multi-register latest-write tracking queries one register after independent updates, requiring preservation of every register’s latest value and binding the query to the correct state component.Matched inputs share the complete update history and differ only in the queried register.
  • Probe construction: L2-regularized linear probes use four-fold grouped cross-validation, holding out numeral-pair groups for completion and complete scratch-context groups for memory.The enlarged experiment contains 1,600 completion examples from 80 groups and 1,280 memory examples from 128 groups.
  • Controlled experiments: Register-count and recurrent-suffix controls vary maintained variables and recurrently prefilled suffix length while probing representations across layers.Register counts are m ∈{1, 2, 4, 8}, while recurrent suffixes contain k ∈{1, 2, 4} final tokens.

G Model outpout

The 200b model correctly simplifies the tested arithmetic and algebraic expressions under both standard and soft decoding, including 137/500 = 0.274 and an algebraic result of 5r^5.

  • 200b model, soft: Under soft decoding, the model gives 137/500 = 0.274 and identifies the last nonzero decimal digit.The passage explicitly presents the quotient as 0.274, though the digit itself is truncated from the supplied text.
  • 200b model, standard decoding: Under standard decoding, the model simplifies (10r^3)(4r^6)/(8r^4) to 5r^5.The solution combines exponent rules, simplifies coefficients, and then reduces the fraction.
  • 200b model, soft: Under soft decoding, the model independently simplifies the same algebraic expression to 5r^5.It first obtains the numerator 40r^9 and then divides by the denominator.
Loading 2608.08888v1…