Technical Summary: Full-Bandwidth Transformer

Problem Statement

Autoregressive transformers compute along two axes: horizontally across generated tokens and vertically through model depth. While dense attention provides broad horizontal access to past tokens, the vertical feedback channel between decoding steps remains narrow. In standard decoding, the model compresses its entire $D$-dimensional top-layer hidden state into a single discrete token (carrying at most $\log_2 |V|$ bits) to serve as the input for the next step. Consequently, non-verbalized computation (intermediate activations, uncertainty, partial plans) is "depth-frozen" in the KV cache; it remains accessible to layers above its production point but cannot return to the bottom of the stack for further processing. This forces the model to either verbalize intermediate states token-by-token or recompute them from scratch, limiting the effective depth of reasoning per token.

Methodology: Latent Feedback Decoding

The authors propose the Full-Bandwidth Transformer, which widens the inter-step feedback channel by feeding the previous top-layer hidden state back into the input alongside the sampled token embedding.

1. Latent Feedback Mechanism

At each decoding step $t$, the input $u_t$ is formed by fusing the sampled token embedding $e_t$ with the previous top-layer hidden state $h^L_{t-1}$:
$$u_t = e_t \otimes h^L_{t-1}$$
The fusion operator $\otimes$ is a Gated Linear Unit (GLU):
$$e_t \otimes h^L_{t-1} = (W_U h^L_{t-1}) \odot \sigma(W_G e_t)$$

• Asymmetry: The hidden state occupies the value pathway, while the token acts as a multiplicative gate. This design prevents the model from learning to ignore the state pathway (a potential shortcut in symmetric fusion), forcing it to utilize the full hidden state.
• Architecture: This adds only two $D \times D$ matrix multiplications per token. The standard transformer architecture, KV cache layout, and serving stack remain untouched.
• Effect: This allows non-verbalized states to re-enter the bottom of the stack with a renewed depth budget. Every layer, including the shallowest, can now access the full history as processed by the entire stack, rather than just the layers beneath it.

2. Parallel Training Strategy

Training this recurrence directly would forfeit parallel teacher forcing, as the input at position $t$ depends on the output at $t-1$. The authors introduce a multi-pass objective to resolve this:

• Temporal Parallelism: Instead of unrolling sequentially over tokens, the model runs $K$ parallel forward passes over the sequence.
  • Pass 1: Standard forward pass with plain token embeddings.
  • Pass $k$ ($k > 1$): The hidden states from Pass $k-1$ are shifted right by one position, fused with the token embeddings, and fed back into the stack in parallel.
• Loss Function: The objective combines the standard Next-Token Prediction (NTP) loss from the first pass with auxiliary NTP losses from subsequent passes. Gradients from later passes backpropagate into earlier latent states, encouraging them to be reusable inputs.
• Scheduling: To ensure stability and manage compute costs, the authors use a progressive schedule:
  • Training begins with standard single-pass objectives.
  • Latent feedback is introduced mid-training (e.g., 2-pass batches).
  • A small fraction (e.g., 3%) of batches with higher passes (e.g., 3-pass) is mixed in. Empirically, this small fraction stabilizes the learned feedback map, making it a contraction toward a fixed point that remains stable even when extrapolated to long horizons (hundreds of passes) during inference.
• Prefix Mix-in: To handle the distribution mismatch between training (fully fused inputs) and inference (plain prompt + fused generation), the training process randomizes the prefix length where fused inputs begin, matching the prompt-then-generate structure of inference.

3. Stability Techniques

To prevent divergence during long-horizon inference, the authors employ:

• Depth Scaling: Ensuring the norm of the top-layer state remains $O(1)$ rather than growing with depth.
• Weight Tying: Tying the embedding layer and readout layer weights to ensure compatibility between plain token and fused state inputs.
• Noise Regularization: Adding small jitter noise to the carried hidden state during training to improve robustness.

Key Results

The authors trained 1B-parameter models up to 400B tokens and evaluated them against standard transformers trained on 100B to 1T tokens.

• Data Efficiency: Full-bandwidth transformers match or approach the performance of standard transformers trained with roughly 1.5× to 2× more tokens. For instance, a 200B-token full-bandwidth model approaches the performance of a 400B-token standard baseline.
• Performance Gains:
  • Validation Loss: Improved across all scales.
  • Reasoning & Coding: On free-form generation tasks (GSM8K, Math500, HumanEval, MBPP), latent feedback decoding (SOFT and FUSED regimes) consistently outperforms standard decoding of the same weights.
    • SOFT (Latent Feedback): Best for math tasks, suggesting hidden state carries reasoning better than verbalized tokens.
    • FUSED (Multi-pass Prefill): Best for coding tasks, suggesting refined prompt representations are crucial.
  • Instruction Tuning: Gains persist after long-context extension and instruction tuning.
• Conciseness: On base models (pre-instruction tuning), latent feedback produces markedly shorter reasoning traces at equal or better accuracy. The model performs computation "silently" in the hidden state rather than narrating it token-by-token. This effect diminishes after instruction tuning, likely because the tuning data imitates verbose standard reasoning.
• State Accessibility: Probing experiments confirm that recurrent prefilling exposes global state information to shallow layers (Layer 0) much more effectively than standard prefilling, validating the "full-bandwidth" claim.

Significance and Claims

The paper claims that the Full-Bandwidth Transformer offers a new scaling axis for Large Language Models: allocating more computation per token rather than just more tokens or parameters.

• Revisiting Scaling: As high-quality unique data becomes a constraint, the authors argue for extracting more learning signal from each token. Latent feedback achieves this by allowing the model to iterate on its internal state without the bottleneck of verbalization.
• Inference Efficiency: The approach provides performance gains with negligible per-token decoding overhead (only two matrix multiplications). The additional compute is primarily incurred during optional multi-pass prefilling or training, not during the sequential generation of each token.
• Architectural Minimalism: The method preserves the standard transformer architecture and KV cache, requiring only a change in the input construction and a specific training schedule.
• Stability: The work demonstrates that recurrent feedback loops can be trained effectively in transformers using parallel multi-pass objectives, overcoming the instability often associated with training recurrent networks.

The authors conclude that full-bandwidth transformers represent a shift from relying solely on data scale to leveraging recurrent, feedback-based computation to improve representation quality and reasoning efficiency.