RLT-pytorch

RLT (Recurrent Looped Transformer)

Unofficial implementation of the Recurrent Looped Transformer proposed by Yifan Zhang of Princeton.

Will also do some exploration of the Recurrent Transformer proposed by Costin-Andrei Oncescu et al. of Harvard, if I have any remaining time

Install

$ pip install rlt-pytorch

Usage

import torch
from RLT import RLT

model = RLT(
    num_tokens = 256,
    dim = 512,
    enc_depth = 4,
    dec_depth = 4,
    dec_sliding_window_size = 16,
    tbptt_step_size = 16 # optional truncated bptt
)

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, max_len = 128) # (2, 96)

Test

Train on enwik8

$ uv run train_enwik8.py

Citations

@techreport{zhang2026recurrentlooped,
    title  = {Recurrent Looped Transformer},
    author = {Zhang, Yifan},
    year   = {2026},
    month  = {Sep},
    url    = {https://github.com/yifanzhang-pro/recurrent-looped-tranformer}
}

@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{kimiteam2026attentionresiduals,
    title   = {Attention Residuals},
    author  = {Kimi Team and Guangyu Chen and Yu Zhang and Jianlin Su and Weixin Xu and Siyuan Pan and Yaoyu Wang and Yucheng Wang and Guanduo Chen and Bohong Yin and Yutian Chen and Junjie Yan and Ming Wei and Y. Zhang and Fanqing Meng and Chao Hong and Xiaotong Xie and Shaowei Liu and Enzhe Lu and Yunpeng Tai and Yanru Chen and Xin Men and Haiqing Guo and Y. Charles and Haoyu Lu and Lin Sui and Jinguo Zhu and Zaida Zhou and Weiran He and Weixiao Huang and Xinran Xu and Yuzhi Wang and Guokun Lai and Yulun Du and Yuxin Wu and Zhilin Yang and Xinyu Zhou},
    year    = {2026},
    eprint  = {2603.15031},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL},
    url     = {https://arxiv.org/abs/2603.15031},
}

Release files for RLT-pytorch 0.0.15

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for RLT-pytorch 0.0.15
File	Size	Uploaded
rlt_pytorch-0.0.15.tar.gz	10.9 kB	Sep 16, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for RLT-pytorch 0.0.15
File	Interpreter	ABI	Platform	Reset
rlt_pytorch-0.0.15-py3-none-any.whl 10.2 kB Sep 16, 2026	Python 3	none	any	Details

Total release size: 21.1 kB

Release files / rlt_pytorch-0.0.15.tar.gz

Download URL	rlt_pytorch-0.0.15.tar.gz
Size	10.9 kB
Tags	Source
SHA-256 checksum How to use checksums	`ca5d821c1595ab5e9ff0dcdea98b8f46fd95d2e6a126227d066a98f23ffcf4c3`
BLAKE2b-256 checksum How to use checksums	`6d633c7ef0becc9c704a01aa88bf5e18873bbb17f61c59ba5aea051645ad79b0`
Upload date	Sep 16, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / rlt_pytorch-0.0.15-py3-none-any.whl

Download URL	rlt_pytorch-0.0.15-py3-none-any.whl
Size	10.2 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`2edf29f54d41c86f11a9836390e5f58aea68a1b9ff7b7491594959b0ae6805f4`
BLAKE2b-256 checksum How to use checksums	`0a9c193bc32f13febfe76cf61ae88534f602e0b46dbd891ea5746813e21203f2`
Upload date	Sep 16, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`