RLT-pytorch 0.1.1

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)

To turn on Next-Latent Prediction (Teoh et al.):

model = RLT(
    num_tokens = 256,
    dim = 512,
    enc_depth = 4,
    dec_depth = 4,
    next_lat_loss = True
)

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

The recurrent block size can also vary across the sequence by passing recurrent_lengths - a sequence of block lengths that must sum to the sequence length

block_lengths = (1, 2, 5, 2, 1, 3)

tokens = torch.randint(0, 256, (2, sum(block_lengths)))

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

# during generation, the recurrent state only advances at the block boundaries

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

sampled = model.generate(prompt, recurrent_lengths = block_lengths)

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

@misc{teoh2025nextlatentpredictiontransformerslearn,
    title     = {Next-Latent Prediction Transformers Learn Compact World Models},
    author    = {Jayden Teoh and Manan Tomar and Kwangjun Ahn and Edward S. Hu and Tim Pearce and Pratyusha Sharma and Akshay Krishnamurthy and Riashat Islam and Alex Lamb and John Langford},
    year      = {2025},
    eprint    = {2511.05963},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url       = {https://arxiv.org/abs/2511.05963}
}