Gated Recurrent Transformers: Expressive Depth through Recurrent Modulation
Amr Hegazy, Amr Alanwar, Mostafa Elhoushi
cs.CL, cs.LG
2026-08-15
A gated shared core between prelude and coda matches GPT-2 Small at isoFLOPs (3.14 vs 3.15). At large scale: 62% fewer params, +10% compiled latency.
Transformer depth and parameter count are glued together: another layer is another weight tensor. Sharing weights grows effective depth without growing storage, but the same map applied to the hidden state at step 1 and step 8 flattens the specialization that makes depth useful. Kaplan et al. already reported the split: recurrent models win at matched parameters and lose at matched compute. Later prelude-coda splits, Mixture-of-Recursions, and LoRA-untied recursive transformers all try to give a shared block a different job at each depth. This paper asks whether one shared core can behave like many layers, and whether that holds on both the isoFLOPs and isoPARAMS axes.
Gated Recurrent Transformer keeps a prelude-shared-coda layout, written npre + nrec × R + ncoda. The prelude runs once and freezes h^(pre). The shared block loops R times. The coda maps the last state to logits. Unique parameters do not grow with R. A 1+1×10+1 small model stores three unique blocks, loops ten times, and matches the forward FLOPs of 12-layer GPT-2 Small.
At each step the current state is concatenated with the prelude, projected by Wproj, perturbed with freshly sampled noise, and fed to the shared block to produce a proposal o^(r). An elementwise gate g^(r) comes from a two-layer MLP on the normalized state and prelude, plus gate noise: h^(r) = g ⊙ h^(r-1) + (1-g) ⊙ o. The gate bias starts at +4, so g≈0.98 at initialization and the residual stream passes through almost unchanged until the model learns where to overwrite. Conditioning on the current state, the prelude anchor, and noise gives the same weights a different input at every step. Training samples r uniformly in 1..R, so intermediate exits are trained without an auxiliary loss.
The recipe matches nanoGPT: sequence 1024, about 9.8B tokens, AdamW. GPT-2 Small/Medium/Large are trained from scratch, as are MoR, heavy-tail Poisson depth sampling, RRT, and Ouro.
IsoFLOPs. Small-scale GRT validation loss is 3.14 versus 3.15 dense, 35M versus 124M parameters. Across three seeds the worst GRT run (3.148) still beats the dense best (3.154). Medium is 2.89 versus 2.84 (127M vs 354M); large is 2.77 versus 2.71 (293M vs 774M). Dense still leads at the standard token budget for medium and large; GRT leads MoR and heavy-tail sampling in all nine scale-by-budget cells. IsoPARAMS, deeper recurrence: medium 2.76 versus 2.84, large 2.65 versus 2.71.
On nine zero-shot tasks at large scale, the isoFLOPs model averages 42.08 versus 42.05 dense. The isoPARAMS model averages 44.15, +2.10. Compiled generation for the large model costs +10% latency, stores 62% fewer parameters, and uses 59% less peak memory (639 MB vs 1570 MB). Uniform depth sampling makes early exit a byproduct; the abstract reports about 92% of full accuracy at half the recurrences. Ablations: recurrence alone is 0.107 nats worse; the elementwise gate is the largest single drop, 0.048 nats.
The bet is that a gate can buy functional diversity without unique layers. IsoFLOPs is the memory-tight training path that still matches dense compute. IsoPARAMS is the path that spends extra inference FLOPs for quality. One checkpoint becomes a compute-quality dial. The scale is GPT-2 Large on 10B tokens, not a modern LLM stack. Among recurrent-depth designs, gating plus a prelude anchor is simpler than LoRA untie or token routing, and it pulls away from MoR and Ouro on isoFLOPs.
The authors name three: R is fixed at inference, with no per-token halt; gate bias and noise may need retuning outside the GPT-2 family; the best sharing fraction shifts with scale and was not swept systematically. At medium and large scale, dense still wins validation loss at the standard token budget; the claimed catch-up appears after doubling tokens. A naive R-way KV cache eats the parameter win at large batch; averaging K/V across steps brings B=32 memory to 0.39× dense. The training mix is described only as diverse text.