T^2MLR: Transformer with Temporal Middle-Layer Recurrence
Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora
cs.CL, cs.AI
2026-07-17
T²MLR injects the last token's middle-layer state into an earlier layer. Looping 20% of layers at 135M: 44.14 vs 42.83; a retrofitted 1.7B hits 18.0 MATH500 at ~8% decode cost.
Autoregressive Transformers squash a rich hidden state back to a token at every step, then use that token as the only input to the next forward pass. Middle layers do most of the abstract work, yet the next token's shallow layers cannot see the previous token's already-computed mid-depth state. They must rebuild it through depth, or hope the compressed embedding still carries it. Coconut-style methods put recurrence at the last layer or the embedding. Looped Transformers buy depth by running the same block several times per token, and decode cost scales with the loop count.
Princeton Language and Intelligence wants a third option: keep abstract middle-layer states alive across time, keep the standard autoregressive interface, and keep per-token decode close to a vanilla Transformer.
T²MLR adds a constant-size recurrent cache R on a decoder-only backbone. Two layer indices mark a middle block, ℓstart to ℓend. At token t, a gated fusion module Φ mixes R{t-1} into the pre-ℓstart state; the middle block then runs as usual; after ℓend the cache becomes RMSNorm(ht^{ℓend} + R{t-1}). Φ uses two learnable scalar gates plus per-dimension sigmoids, with the scalars initializable at zero for early stability.
Dense teacher forcing cannot see R for every token in parallel, so training approximates the cache with Jacobi fixed-point iteration: a cache-free pass seeds R from ℓend, then dforward=16 middle-block refinements; backward depth is dbackward=4, like truncated BPTT. At decode time the extra work is the fusion module. Measured generation overhead stays under about 8% and shrinks as attention dominates on longer sequences and larger models.
S5-Retrieval asks for both S5 group state tracking and in-context lookup. A 4-layer LSTM and a 4-layer Transformer collapse on exact match. A T²MLR of the same depth is near-perfect in-distribution and keeps non-trivial token accuracy past the training length, after 150k steps against 400k for the baselines.
At 135M parameters, 10B FineWeb-Edu tokens, and matched size (baseline hidden 576 raised to 584, so the Transformer is slightly larger):
| Setup | Zero-shot average |
| Transformer | 42.83 |
| T²MLR full D=30 | 43.36 |
| T²MLR (13,18) D=6 | 44.14 |
| Pause-token ×2 | 43.31 |
| Full-looped ×2 | 42.99 |
| Middle-looped ×3 | 42.68 |
Looping about 20% of the middle layers beats looping all of them. A fixed-width ablation at D=6 and D=14 puts the same block on early, middle, or late layers; middle wins both times.
On multi-hop and grade-school math finetunes, middle-layer variants again beat full-layer recurrence. Relative gains grow on reasoning as width scales: at 361M, HotpotQA-Easy 24.43→28.28 (+15.8%) and GSM-Aug NL 31.08→34.12; at 1B, HotpotQA 23.28→26.52 (+14.0%). With 50B pretraining tokens, the 361M zero-shot average moves 52.83→54.78.
The pathway can be grafted on. Inserting T²MLR (5,28) into SmolLM2-1.7B-Instruct and finetuning one epoch on OpenMathReasoning lifts GSM8K 35.78→39.88 and MATH500 12.80→18.00.
If the goal is better reasoning without looped depth or longer written chains, this is an adoption path: KV cache stays, the decode API stays, and the paper's measured stepwise overhead is at most about 8%. Middle-layer recurrence beating full-layer recurrence is the design lesson. Where the loop sits matters more than whether a loop exists.
Read the compute claim narrowly. Jacobi training is about 2-4× slower in wall-clock. Matching that budget by training the 135M Transformer for 2.24 epochs yields a 45.30 zero-shot average, above T²MLR's 44.14. The authors frame the main tables as parameter, data, and inference-compute matched, not training-time matched.
Training cost is the limitation they put first. They point to fewer refinement steps, or exact recurrent states in on-policy RL, as future work. Scale is still modest: from-scratch runs to 1B, retrofit to 1.7B, no multi-seed error bars. S5-Retrieval numbers live in a figure. There is no head-to-head with Coconut-style embedding recurrence, because those methods lack scalable dense teacher forcing. Several downstream sets are hardened (assignment depth 5, ProsQA 60 nodes), so they are not comparable to the original papers' scores.