Gated Recurrent Transformers (GRT)

Official implementation of Gated Recurrent Transformers: Expressive Depth
through Recurrent Modulation.

GRT is a recurrent-depth transformer where fixed-depth prelude and coda
blocks bracket a single shared core iterated R times. A lightweight
elementwise update gate, conditioned on the hidden state, the fixed prelude
output, and resampled noise, modulates each recurrence so the model can
specialize the input to the same few layers instead of relying on many
unique layers for functional diversity. Under an isoFLOP constraint, a
3-layer GRT matches the accuracy of a 12-layer GPT-2 Small baseline at
similar training and inference FLOPs; under an isoParam constraint, deeper
recurrence reaches a 2.76 validation loss versus 2.84 for a non-recurrent
counterpart at matched parameters and data.

The code is a trimmed fork of nanoGPT.
model.py and train.py cover everything reported in the paper; the config
files under config/ reproduce the individual runs.

Layout

model.py            GRT (prelude -> shared core -> coda)
train.py            training loop, DDP-ready
sample.py           generation from a checkpoint
configurator.py     config-file / CLI override mechanism (from nanoGPT)
utils.py            run provenance logging (config, command, git state)
config/             one file per run reported in the paper
data/cccc_filtered/ dataset preparation
scripts/            lm-evaluation-harness integration
assets/             figures used in this README

Setup

pip install -r requirements.txt
python data/cccc_filtered/prepare.py   # writes data/cccc_filtered/{train,val}.bin

Preparation downloads ~58 GB into the HuggingFace cache and takes a while.

Training

All commands are run from the repository root.

torchrun --standalone --nproc_per_node=2 train.py config/grt_medium_isoflop.py

Any config value can be overridden from the command line:

torchrun --standalone --nproc_per_node=2 train.py config/grt_medium_isoflop.py \
    --max_iters=2000 --out_dir=logs/quick --wandb_log=False

Every run writes ckpt.pt plus the resolved config, launch command, and git
state to out_dir. Set --init_from=resume to continue from a checkpoint in
the same directory.

Configurations

Notation is p+n×R+c: p prelude blocks, n shared blocks applied R times,
c coda blocks. Every run uses common-pile/cccc_filtered, sequence length 1024, ~491,520
tokens per optimizer step, and 20,000 steps (~9.8B tokens).

Config	Layout	Unique blocks	Regime
`gpt2_small.py`	12L	12	dense baseline
`gpt2_medium.py`	24L	24	dense baseline
`gpt2_large.py`	36L	36	dense baseline
`grt_small_isoflop.py`	1+1×10+1	3	isoFLOP vs. 12L
`grt_medium_isoflop.py`	2+5×4+2	9	isoFLOP vs. 24L
`grt_large_isoflop.py`	1+5×6+5	11	isoFLOP vs. 36L
`grt_medium_isoparam.py`	2+20×4+2	24	isoParam vs. 24L
`grt_large_isoparam.py`	3+30×6+3	36	isoParam vs. 36L

Architecture flags

Flag	Default	Effect
`n_layers_prelude`, `n_layer`, `n_layers_coda`	1, 1, 1	prelude, shared core, coda depth
`recurrence_min`, `recurrence_max`	2, 10	depth sampling range; `recurrence_max` is `R`
`depth_sampling`	`"uniform"`	`"uniform"` or `"lognormal"` (heavy-tail baseline)
`prelude_injection`	`True`	concatenate the prelude anchor with the state each step
`gating`	`True`	blend the block output in with the elementwise gate
`gate_init_bias`	`4.0`	gate ≈ 0.98 at initialisation
`gate_temperature`	`1.0`	sigmoid temperature
`noise_std`	`0.1`	state and gate noise
`noise_first_step_only`	`False`	inject noise only at step 0

Sampling

python sample.py --out_dir=logs/grt-medium-isoflop --recurrent_depth=4 \
    --start="The recurrent block"

--recurrent_depth is the compute dial: any value up to the trained R can be used, and quality increases with depth.

Downstream evaluation

Requires pip install lm-eval.

python scripts/lm_eval_grt.py \
    --checkpoint logs/grt-large-isoflop/ckpt.pt \
    --tasks hellaswag,arc_easy,arc_challenge,piqa,lambada,winogrande,openbookqa,boolq \
    --device cuda --dtype bfloat16 --recurrences 6 \
    --output_path eval_results/large-isoflop.json

Citation

@misc{hegazy2026gatedrecurrenttransformersexpressive,
      title={Gated Recurrent Transformers: Expressive Depth through Recurrent Modulation}, 
      author={Amr Hegazy and Amr Alanwar and Mostafa Elhoushi},
      year={2026},
      eprint={2608.15062},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2608.15062}, 
}

gated-recurrent-transformer

Specs

README