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 blocks after 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 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 is the RLT output if every position equals the decoder's output on inputs merged with a register read from at the last chunk boundary. That is one line in Lean. Then two facts carry the rest: the register read at only looks before , 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.
