💬 NLP

T^2MLR: Transformer with Temporal Middle-Layer Recurrence

The paper introduces T^2MLR, a Transformer architecture that enhances latent reasoning by fusing cached middle-layer representations from previous tokens into earlier layers of the current token, a targeted approach that outperforms standard baselines and full-layer recurrence while allowing for efficient retrofitting of existing pretrained models.

Original authors: Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora

Published 2026-07-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to solve a tricky puzzle. In the world of artificial intelligence, the most popular robots today are built on a design called the "Transformer." Think of a Transformer like a very fast, very eager reader who processes a story one word at a time. Every time it reads a new word, it has to forget almost everything about its deep, complex thoughts from the previous word and start fresh, compressing all that rich thinking into a tiny, simple note before moving to the next word. This is like trying to solve a math problem while constantly erasing your scratch paper and only writing down the final answer before moving to the next step. It works okay for simple tasks, but when the robot needs to hold a long chain of logic in its head—like tracking a secret code or solving a multi-step riddle—it often gets lost because it can't keep its "thoughts" alive between steps.

Scientists have been trying to fix this by letting the robot keep a "hidden notebook" of its thoughts, but most attempts either make the robot incredibly slow (because it has to re-read its whole notebook every time) or force it to write its thoughts in a weird, invisible language that's hard to train. The big question is: Can we give the robot a way to keep its deep, abstract thinking alive from one word to the next without slowing it down or making it impossible to teach? This paper introduces a clever new trick called T2MLR (Transformers with Temporal Middle-Layer Recurrence) that suggests the answer is yes, but only if we change where the robot keeps its thoughts.

The researchers propose a new way for these AI models to "remember" their thinking process. Instead of forcing the robot to write its deep thoughts down in a tiny note at the very beginning or the very end of its processing, they let the robot pass a "secret message" from the middle of its brain directly to the middle of its brain for the next word. Imagine a relay race where, instead of passing the baton only at the start or finish line, the runners pass a special note to each other right in the middle of the track. This allows the complex, abstract reasoning happening in the "middle layers" of the AI to persist and evolve as the robot reads, without having to stop and re-calculate everything from scratch.

The paper finds that this "middle-layer" approach works surprisingly well. When they tested it on tasks that require tracking states (like remembering a sequence of moves in a game) or solving multi-step math problems, the new T2MLR model consistently beat standard models of the same size. Interestingly, they discovered that you don't need to make the whole robot remember everything; just letting a small chunk of the middle layers (about 20% of the network) pass these notes back and forth was often better than making the entire robot loop over its own thoughts. Even more exciting for the real world, they showed that you don't have to build a new robot from scratch to use this. They took an existing, pre-trained robot (a 1.7 billion parameter model) and simply "grafted" this new middle-layer memory system onto it. After just a little bit of extra training on math problems, this upgraded robot got significantly better at solving them, jumping from a 35.8% success rate to nearly 40% on one test, and from 12.8% to 18% on another.

The authors suggest that this works because the "middle" of a Transformer is where the most abstract reasoning happens, and by letting those specific thoughts persist over time, the model can do much better "latent reasoning" (thinking without writing everything down). However, they are careful to note that this comes with a trade-off: while the robot doesn't get much slower when it's actually answering questions (it only adds about 8% extra work per word), it does take longer to train because the computer has to do some extra math to figure out how these memory notes should work. Despite this training cost, the results suggest that effective reasoning in AI doesn't require looping over every single layer of the brain; instead, a targeted, middle-layer memory system might be the key to unlocking smarter, more persistent thinking.

Technical Summary: T2MLR: Transformer with Temporal Middle-Layer Recurrence

Problem Statement

The authors identify a fundamental limitation in standard Transformer architectures regarding reasoning capabilities. While Transformers excel at in-context retrieval, their autoregressive decoding mechanism creates an information bottleneck. At each decoding step, rich, high-dimensional latent representations are repeatedly compressed into a sparse token space (one-hot vectors) before being fed as input to the next step. This process prevents intermediate abstract reasoning states from persisting across time, forcing the model to reconstruct necessary information indirectly via attention or through highly compressed input embeddings.

Existing approaches to latent reasoning attempt to mitigate this by either performing reasoning entirely in continuous space (often at the final layer) or by propagating uncertainty through linear combinations of token embeddings. However, these methods typically feed recurrent signals back through the input embedding, bypassing the middle layers where abstract reasoning is known to occur. Alternatively, "looped" Transformer architectures increase effective depth by iterating over all layers, but this incurs inference costs that scale linearly with the number of loops, making them impractical for deployment.

Methodology

The paper introduces T2MLR (Transformers with Temporal Middle-Layer Recurrence), a novel architecture designed to allow abstract intermediate representations to persist and evolve across decoding steps without increasing inference-time depth.

Architectural Design

T2MLR modifies a standard $L$ -layer decoder-only Transformer by introducing a gated recurrent pathway that connects a deeper layer from the previous token to an earlier layer of the current token.

Recurrence Boundaries: The architecture is parameterized by two layer indices, $\ell_{start}$ and $\ell_{end}$ , defining a "middle" block of layers.
Recurrent Cache ( $R_t$ ): A constant-size cache vector is maintained for each decoding step $t$ . It stores the representation after layer $\ell_{end}$ from the previous step.
Fusion Module ( $\Phi$ ): Before processing the current token at layer $\ell_{start}$ , the model fuses the current token's representation ( $h^{(\ell_{start}-1)}_t$ ) with the recurrent cache ( $R_{t-1}$ ) from the previous token. The fusion is governed by a learnable gated mechanism:
$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$
This design allows the model to learn when to rely on the current token's stream versus the temporally accumulated latent state.
Cache Update: After passing through layer $\ell_{end}$ , the recurrent cache is updated using the new representation and the previous cache via a residual connection and RMSNorm.

Training Strategy: Temporal Parallelism

A key challenge for latent reasoning architectures is that standard sequence-parallel training (token-parallel) breaks when dependencies exist across time steps. To address this, the authors propose an approximated temporal-parallel training scheme using Jacobi fixed-point iterations:

Instead of sequential token processing, the model performs $d_{forward}$ iterations of refinement in parallel.
It starts with an initial cache derived from a standard forward pass, then iteratively fuses this cache with the current token's representation, refines the middle block, and updates the cache.
Gradients are backpropagated through these iterations with a truncated depth ( $d_{backward}$ ), analogous to Truncated Backpropagation Through Time (TBPTT).
This approach allows the model to be trained with dense teacher forcing and sequence parallelism, a capability the authors claim is unique among continuous chain-of-thought variants.

Key Contributions

Architectural Innovation: T2MLR introduces a targeted recurrent pathway specifically in the middle layers of the Transformer, rather than at the input embedding or final output. This preserves the standard autoregressive interface while enabling latent state persistence.
Scalable Training: The authors demonstrate a method to train these recurrent dependencies using approximate temporal parallelism, enabling pretraining and fine-tuning on large datasets without sacrificing scalability.
Retrofitting Capability: The paper shows that T2MLR does not require training from scratch. The recurrent pathway can be inserted into an existing pretrained model (e.g., SmolLM2-1.7B) and fine-tuned to yield significant reasoning improvements.
Efficiency: Unlike looped Transformers, T2MLR adds negligible inference overhead (at most ~8% per token) because the recurrence is handled via a constant-size cache and a single forward pass during inference.

Experimental Results

The authors evaluate T2MLR across synthetic benchmarks, natural language pretraining, and downstream reasoning tasks.

Synthetic Reasoning (S5-Retrieval): T2MLR significantly outperforms both standard Transformers and LSTMs on a task requiring simultaneous state tracking and in-context retrieval. While Transformers fail at long sequence lengths due to depth constraints and LSTMs fail at retrieval due to fixed state compression, T2MLR maintains high accuracy.
Pretraining (FineWeb-Edu): On 135M parameter models, T2MLR variants consistently match or exceed parameter-matched Transformer baselines on zero-shot NLP benchmarks (ARC, HellaSwag, etc.). Notably, middle-layer recurrence (looping over ~20% of layers, e.g., layers 13–18) often outperforms full-layer recurrence.
Downstream Reasoning: On tasks like GSM-Aug (math), HotPotQA (multi-hop QA), and Variable Assignment, T2MLR shows consistent gains over baselines. The middle-layer recurrence variants again outperform full-layer recurrence, suggesting that targeted recurrence is more effective than global looping.
Scaling: The performance gains persist and grow when scaling to 361M and 1B parameters and extending pretraining to 50B tokens.
Retrofitting: Applying T2MLR to a pretrained SmolLM2-1.7B-Instruct model and fine-tuning on math data improved GSM8K accuracy from 35.8% to 39.9% and MATH500 from 12.8% to 18.0%, demonstrating practical applicability without full pretraining.

Significance and Claims

The paper posits that effective latent reasoning in Transformers does not require looping over all layers (which increases inference cost) nor confining recurrence to the input embedding (which bypasses abstract computation). Instead, the authors claim that targeted middle-layer recurrence is sufficient to unlock significant reasoning capabilities.

The significance of T2MLR lies in its ability to:

Decouple Reasoning Depth from Inference Cost: It introduces a temporal latent pathway that persists across time steps without the linear inference penalty associated with looped architectures.
Lower the Barrier to Adoption: By demonstrating that the architecture can be retrofitted onto existing pretrained models, the authors argue that advanced reasoning capabilities can be added to current models at a fraction of the cost of pretraining from scratch.
Validate Inductive Biases: The results suggest that the "sweet spot" for recurrence in Transformers is the middle layers, where abstract computation occurs, rather than the input or output interfaces.

The authors acknowledge that T2MLR incurs a training-time overhead (2–4x wall-clock time depending on the recurrent block size) due to the Jacobi approximation steps. They frame this as a trade-off: T2MLR sacrifices training compute to gain a low-overhead, high-performance inference architecture for reasoning tasks.