Worklog: a looped transformer in Lean

Spent today on a small Lean 4 project: lean-recurrent-looped-transformer.

The Recurrent Looped Transformer (Zhang, Feng, Qin, this month) feeds the decoder's final hidden state back into the next token's input. That gives the state a path of $t \cdot L_D$ blocks after $t$ tokens, but it also means you cannot prefill a prompt in parallel. Their RLT-2 variant fixes the fed-back state inside a chunk of $B$ tokens and only updates it at the boundary, so a chunk can be computed at once. The paper says the chunked version computes the same thing as the token-by-token one, and that prefilling a prompt and then decoding gives the same result as processing everything sequentially. Both claims are a paragraph of prose. I wanted them as theorems.

What ended up working was not defining the algorithm at all. Define the answer instead: a sequence $y$ is the RLT output if every position equals the decoder's output on inputs merged with a register read from $y$ at the last chunk boundary. That is one line in Lean. Then two facts carry the rest: the register read at $t$ only looks before $t$, and the decoder is causal. Together they make the fixed point unique. Any way of filling in positions that goes in order and never straddles a chunk boundary reaches it, so chunk-wise, token-wise and prefill-then-decode are all equal by construction.

Numbers: one file, ~750 lines, import Std only, 54 theorems, no sorry. Vectors are Fin n → Rat the way Sasha Rush does it in his Lean transformer post; no exp, so softmax and sigmoid are rational stand-ins. The proofs never look at them anyway.

The part I would point someone to is fill_correct, the only lemma with any real content. Everything after it is folds over lists.

Later in the day I pulled the dependency graph out of the Lean environment, proof terms included: every declaration in the file and what it uses. The orange path is the argument above.