Technical Summary: Latent Recurrent Transformer (LRT)

Problem Statement

Autoregressive transformers are the standard for language modeling, yet they rely on fixed-depth feedforward computations for each generated token. While recurrence can theoretically increase computational capacity for iterative refinement, existing approaches introduce significant inference costs. Methods that add recurrence in depth (repeatedly applying blocks to the same token) or in time (inserting pause or thinking tokens) increase inference latency through extra decoding steps or block applications.

The authors observe that autoregressive decoding already computes high-level recurrent signals "for free": the hidden states of the previous token. Specifically, upper-layer states trained for next-token prediction contain useful latent context. The core research question is whether these already-computed high-level representations from the previous token can be reused as recurrent memory for the next token without adding extra decoding steps.

Methodology

Architecture: Latent Recurrent Transformer (LRT)

LRT is a lightweight augmentation of standard autoregressive transformers that introduces a cross-layer, cross-token latent pathway.

• Mechanism: At token position $t$, LRT reuses a source-layer hidden state from the previous position, $m_{t-1} = h^{\ell_{src}}_{t-1}$, as recurrent memory. Here, $\ell_{src}$ is a selected source layer (typically an upper-middle layer).
• Injection: The memory $m_{t-1}$ is injected into the transformer layers of the current token $t$ via two complementary mechanisms:
  1. KV Projection: Projects the memory into the attention key-value spaces ($\tilde{k}_{rec}, \tilde{v}_{rec}$) and combines them with local keys/values using input-dependent gates. This provides a direct route for the previous token's representation into the attention mechanism.
  2. Residual Injection: Adds the memory directly to the residual stream ($\bar{x}^{(\ell)}_t = \alpha_\ell x^{(\ell)}_t + \gamma_\ell m_{t-1}$), exposing the entire block to the recurrent signal.
• Inference Efficiency: Crucially, LRT preserves the standard decoder-only backbone, attention mechanism, and KV-cache interface. It does not insert memory tokens, pause tokens, or depth loops. Each generated token still utilizes exactly one normal transformer forward pass, with the recurrent state simply being updated and passed to the next step. The default variant adds only 0.3% parameters.

Training Strategy: Interleaved Parallel Training

A major challenge is pretraining this recurrence at scale. Exact training would require sequentially unrolling the transformer over the full sequence (Backpropagation Through Time), destroying the parallelism essential for efficient transformer pretraining. Chunked training (processing contiguous blocks sequentially) preserves some parallelism but only propagates memory across chunk boundaries, offering a coarse approximation.

The authors introduce Interleaved Parallel Training to approximate token-level recurrence while retaining parallel computation:

1. Initialization Pass: A single full-sequence forward pass builds a shared buffer containing keys, values, and source-layer recurrent states for all positions.
2. Refinement Passes: The sequence positions are partitioned into $S$ disjoint subsets (e.g., strided partitions). These subsets are refined in parallel.
3. Write-Back: Updated states from a subset are written back to the shared buffer. Subsequent subsets can consume the refined memory from earlier subsets.
4. Objective: The loss is the average of the initialization loss and the refinement losses for each subset.

This strategy ensures every token receives a recurrent-memory-aware supervision step. In terms of compute, it costs approximately 2× a standard transformer update (one initialization + one effective refinement), but it mimics token-level recurrence more closely than chunked training.

Key Contributions

1. LRT Architecture: A lightweight recurrent extension that reuses high-level source-layer hidden states from the previous token as memory, enabling cross-layer recurrence without extra decoding steps or architectural changes to the KV-cache interface.
2. Interleaved Parallel Training: A parallel training approximation that refines disjoint token subsets to mimic token-level recurrence under a fixed training budget, avoiding the sequential unrolling required by standard RNN-style training.
3. Empirical Validation: Demonstration that LRT improves the compute-quality trade-off across various model sizes and tokens-per-parameter budgets, with the default shared-projection variant adding minimal parameter overhead (0.3%).

Results

Experiments were conducted on nanochat-style GPT backbones (20L and 24L layers) pretrained on FineWeb-Edu.

• Scaling Behavior: Under baseline-equivalent training compute (accounting for the 2× cost of interleaved training), LRT shifts scaling curves toward lower Bits Per Byte (BPB) and higher CORE (a metric for in-context learning across 22 tasks).
• Performance Gains:
  • On the 24L backbone at baseline-equivalent compute 80, LRT-shared achieved 0.695 BPB compared to the baseline's 0.699.
  • In-context learning (CORE) also improved; on the 20L backbone at compute 80, LRT-shared reached 0.274 vs. baseline 0.271.
• Ablation Findings:
  • Source Layer: An upper-middle layer (e.g., layer 12 in a 20L model) performs better than the final layer, suggesting useful memory should be high-level but not overly specialized for logits.
  • Temporal Source: The immediately preceding state ($m_{t-1}$) is the strongest single memory source that preserves the one-forward-pass constraint.
  • Injection: KV Projection and Residual Injection are complementary; using both yields the best results.
  • Training: Interleaved parallel training significantly outperforms chunked training, which provides a coarser approximation of recurrence.

Significance and Claims

The paper claims that LRT successfully shifts the compute–quality trade-off for autoregressive transformers. By reusing "free" high-level signals from the previous token and training them via a parallel approximation, LRT enables iterative refinement and better context utilization without the inference latency penalties of depth-recurrent or pause-token methods.

The authors emphasize that LRT is a lightweight augmentation that preserves the standard transformer interface. The gains are attributed to the recurrent memory mechanism itself rather than simply adding extra parameters or gates, as generic architectural additions (like gated attention or layer scaling) on the same strong baseline yielded negligible improvements. The work suggests that useful recurrent memory should be high-level but not the final logit-specialized layer, and that token-level recurrence can be effectively approximated in parallel during pretraining.