Princeton Researcher Proposes Recurrent Looped Transformer to Carry Decoder State Across Tokens

In most decoder-only large language models, nothing computed at the last layer of token t feeds the first layer of token t+1; positions communicate only through attention over cached keys and values. RLT aims to close that loop by making the decoder recurrent from token to token.

RLT pairs a causal encoder with a recurrent decoder. The encoder processes tokens in parallel under a causal mask and produces representations e_t, from which key-value memory M≤t is projected. Memory groups can be shared across decoder layers (G = 1) or kept layer-specific (G = L_D). The decoder holds the recurrence. Its complete state is Ht = (st, CtD), where st is the final decoder output and CtD holds the retained SWA keys and values at every decoder layer. For each token, a gated merge combines e_t with the previous output s_{t-1}, then each decoder block runs causal SWA over decoder activations, cross-attention to encoder memory and a feed-forward network. The window W includes the current token, so at most W – 1 historical entries per layer are retained. The next-token distribution is read from st. Initialization happens once before BOS with a learned start state s* and an empty cache. The reference tied configuration uses 48 encoder and 48 decoder layers with compatible attention and feed-forward weights shared between them. Each token therefore executes 96 logical blocks, though decoder blocks add cross-attention, so per-block FLOPs are not equal. Zhang calls this parameter reuse, not activation copying.

The report lays out three design principles. The first is latent reasoning with unbounded temporal depth: after t processed tokens, the state path from s0 traverses t·LD decoder blocks, or 48t in the reference configuration. Per-token work stays fixed while the path's structural depth grows with the sequence. The report warns that gates and contraction may suppress long paths, and that structural depth is not a reasoning guarantee. The second is model-hardware co-design: encoder features and memory projections for known tokens use token-parallel kernels, while decoder transitions stay sequential within a sequence, though ready updates from independent sequences can share one batched kernel. The report states plainly that no exact parallel scan is assumed for the nonlinear decoder, no reduced-prefill speedup is claimed, and a standard parallel SWA decoder pass is not equivalent to the recurrence. Batching, kernel fusion and checkpointing are listed as implementation targets, not completed kernels. The third is model-RL algorithm co-design: pretraining, SFT, sampling and RL replay share one state transition. For RL, the sampler records each action's behavior log-probability under its actual sampling distribution, including temperature and truncation. The trainer rebuilds encoder memory, the recurrent output and every SWA cache from the sequence start under current parameters before scoring each action; old rollout states are never reused. Proposition 3.1 states that moving the prompt-response split leaves the conditional distribution unchanged for a fixed token history.

Pretraining is full-sequence next-token prediction with full backpropagation through time. SFT masks the loss to assistant targets but never masks state updates, so assistant losses backpropagate through user and tool tokens. Appendix B explains why partial detaching is risky: the state-to-state Jacobian has cross terms through decoder KV, so detaching only st leaves gradient paths through the cache; any truncated-BPTT scheme must name every detached tensor. For multi-turn serving, an exact prefix snapshot includes encoder cache and memory, the complete decoder state, position metadata, the window convention and model version. A fixed-weight snapshot can be reused because the state is independent of the serving split; weight updates invalidate old states, and editing a prefix forces recomputation from an earlier checkpoint. External tokens in multi-turn RL update the state but get no importance-ratio factors.

On prior work, encoder-derived memory follows YOCO, which caches KV once for a cross-decoder, and DeepSeek-V4.1-Flash, which projects decoder global KV from final encoder states; RLT keeps the memory but drops prompt-wide decoder skipping. Temporal feedback builds on Feedback Transformer and Recurrent Transformer; RLT instead feeds the previous final decoder output into the next decoder input and runs recurrence over the prompt too. Depth-wise reuse connects to Universal Transformers and recurrent-depth latent reasoning.