GRT Medium (medium-isoflop)
This is a Medium isoFLOP checkpoint of the Gated Recurrent Transformer (GRT), described in:
Gated Recurrent Transformers: Expressive Depth through Recurrent Modulation Amr Hegazy, Amr Alanwar, Mostafa Elhoushi arXiv:2608.15062 · Code
Architecture
GRT applies a shared transformer block recurrently with three key innovations:
- Prelude injection: the output of fixed prelude blocks is concatenated with the noised hidden state and re-projected at every recurrence step, grounding each iteration in the original input representation.
- Elementwise sigmoid gating: an MLP gate blends the shared block's output into the residual stream (initialized near-open at σ(+4) ≈ 0.98), allowing the model to learn which elements to overwrite as training proceeds.
- State and gate noise: Gaussian perturbation of both the hidden state (εx) and gate logits (εg) at each step discourages brittle exact-match patterns and prevents gate collapse.
The architecture is prelude → shared core × R → coda, where R (recurrence depth) is sampled uniformly during training and fixed at inference. This enables inference-time early exiting from a single checkpoint without auxiliary losses.
This checkpoint
| Value | |
|---|---|
| Regime | isoFLOP (matched inference FLOPs, fewer unique parameters) |
| Config | 2+20×4+2 |
| Parameters | ~358M |
| Block executions / token | 84 |
| Embedding dim | 1024 |
| Heads | 16 |
| Shared layers | 20 |
| Context length | 1024 |
| Training tokens | ~9.8B (CCCC Filtered) |
| Validation loss | 2.7632 |
Usage
import torch
from model import GPT, GPTConfig
# Load checkpoint
ckpt = torch.load("ckpt.pt", map_location="cpu")
config = GPTConfig.from_checkpoint_args(ckpt["model_args"])
model = GPT(config)
model.load_state_dict(ckpt["model"])
model.eval().cuda()
# Generate
import tiktoken
enc = tiktoken.get_encoding("gpt2")
prompt = enc.encode("The meaning of life is")
x = torch.tensor([prompt], dtype=torch.long, device="cuda")
y = model.generate(x, max_new_tokens=100, temperature=0.8, top_k=200, n=6)
print(enc.decode(y[0].tolist()))
Or use the provided sample.py:
git clone https://github.com/Amr-Hegazy1/gated-recurrent-transformer
cd gated-recurrent-transformer
# download ckpt.pt into logs/<config>/
python sample.py --out_dir=logs/<config> --recurrent_depth=6
Citation
@article{hegazy2026grt,
title = {Gated Recurrent Transformers: Expressive Depth through Recurrent
Modulation},
author = {Hegazy, Amr and Alanwar, Amr and Elhoushi, Mostafa},
journal = {arXiv preprint arXiv:2608.15062},
year = {2026},
}