recurrent-transformer-pytorch 0.0.17

Recurrent Transformer (wip)

Explorations into the Recurrent Transformer proposed by Costin-Andrei Oncescu et al. of Harvard University

Usage

import torch
from recurrent_transformer import RecurrentTransformer

model = RecurrentTransformer(
    num_tokens = 256,
    dim = 512,
    depth = 6,
    recurrent = True
)

tokens = torch.randint(0, 256, (2, 1024))

# forward for loss

loss = model(tokens, return_loss = True)
loss.backward()

# generate

prompt = torch.randint(0, 256, (2, 32))

sampled = model.generate(prompt, 128) # (2, 128)

Train

Train on enwik8, validating at the trained length and extrapolated length

$ uv run train_enwik8.py

Citations

@misc{oncescu2026recurrenttransformergreatereffective,
    title   = {The Recurrent Transformer: Greater Effective Depth and Efficient Decoding},
    author  = {Costin-Andrei Oncescu and Depen Morwani and Samy Jelassi and Alexandru Meterez and Mujin Kwun and Sham Kakade},
    year    = {2026},
    eprint  = {2604.21215},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2604.21215},
}

@misc{bahdanau2014neural,
    title   = {Neural Machine Translation by Jointly Learning to Align and Translate},
    author  = {Dzmitry Bahdanau and Kyunghyun Cho and Yoshua Bengio},
    year    = {2014},
    eprint  = {1409.0473},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL},
    url     = {https://arxiv.org/abs/1409.0473}
}

@Article{AlphaFold2021,
    author  = {Jumper, John and Evans, Richard and Pritzel, Alexander and Green, Tim and Figurnov, Michael and Ronneberger, Olaf and Tunyasuvunakool, Kathryn and Bates, Russ and {\v{Z}}{\'\i}dek, Augustin and Potapenko, Anna and Bridgland, Alex and Meyer, Clemens and Kohl, Simon A A and Ballard, Andrew J and Cowie, Andrew and Romera-Paredes, Bernardino and Nikolov, Stanislav and Jain, Rishub and Adler, Jonas and Back, Trevor and Petersen, Stig and Reiman, David and Clancy, Ellen and Zielinski, Michal and Steinegger, Martin and Pacholska, Michalina and Berghammer, Tamas and Bodenstein, Sebastian and Silver, David and Vinyals, Oriol and Senior, Andrew W and Kavukcuoglu, Koray and Kohli, Pushmeet and Hassabis, Demis},
    journal = {Nature},
    title   = {Highly accurate protein structure prediction with {AlphaFold}},
    year    = {2021},
    doi     = {10.1038/s41586-021-03819-2},
    note    = {(Accelerated article preview)},
}

@inproceedings{ren2026rethinking,
    title   = {Rethinking Expressivity and Degradation-Awareness in Attention for All-in-One Blind Image Restoration},
    author  = {Bin Ren and Runyi Yang and Qi Ma and Xu Zheng and Mengyuan Liu and Danda Pani Paudel and Luc Van Gool and Rita Cucchiara and Nicu Sebe},
    booktitle = {The Fourteenth International Conference on Learning Representations},
    year    = {2026},
    url     = {https://openreview.net/forum?id=IBzmQVia88}
}

@inproceedings{press2022train,
    title   = {Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation},
    author  = {Ofir Press and Noah A. Smith and Mike Lewis},
    booktitle = {International Conference on Learning Representations},
    year    = {2022},
    url     = {https://openreview.net/forum?id=R8sQPpGCv0}
}