RLT
An unofficial PyTorch implementation of the Recurrent Looped Transformer proposed by Yifan Zhang of Princeton, an encoder-decoder architecture that can share weights between encoder and decoder and vary its recurrent block size across a sequence.
RLT implements the Recurrent Looped Transformer as a PyTorch package, installable as rlt-pytorch. It supports weight tying between the encoder and decoder, an optional Next-Latent Prediction training objective, and recurrent block lengths that vary across a sequence rather than staying fixed, with the recurrent state advancing only at block boundaries during generation.
Usage is as a library: a model is constructed with token count, model dimension and encoder and decoder depth, then called on a batch of tokens to compute a loss for backpropagation, or used to generate continuations from a prompt respecting the same block-length structure. A training script for the enwik8 benchmark is included for testing.
It is a Python package licensed under MIT.