grt-medium-isoflop

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:

1. 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.
2. 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.
3. 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

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},
}