Looped Transformers

A looped transformer runs the same weights more than once for each token. The output of a block, or of the whole stack, is fed back in and transformed again. The model behaves like a deeper network while storing only one copy of each block. The idea goes back to the Universal Transformer (2018). Recent work uses it to scale reasoning in latent space instead of in written tokens.

The extra depth is not free. Every loop still has to be computed, during training and for every generated token. On 29 October 2026, FAIR at Meta published How to Loop Your Transformer (Ekbote et al.), a controlled study of when that trade pays. It covers dense decoder-only models of 70M to 1B parameters trained from scratch, on up to 155B tokens with up to 32 loops. Its answer depends on which budget binds:

• Stored weights and data fixed: loop the whole stack. Training compute fixed: loop only the middle third. Gains flatten after a few loops, then reverse or diverge.
• Many loops: exact weight sharing becomes the constraint. Small per-loop LoRA adapters add 1–4% to the weights and help more the more loops there are.
• Decoding latency: ordinary loops run one after another for each token. Parallel Loop Transformers (PLT) let them run at once but lost most of the quality gain. The paper's fix is to keep the current-token embedding, normalize it and the incoming state separately, and gate them. With that fix, two parallel loops at 1B beat two sequential ones.

Depth without new weights. A standard transformer with $B$ blocks stores $B$ sets of weights and applies each once. Loop the stack $L$ times and the model is $LB$ blocks deep but still stores $B$. Compute per token grows with $L$; stored weights do not. Two schedules perform exactly the same evaluations in a different order. Full-Stack looping passes through all $B$ blocks, then starts again; Per-Layer looping repeats each block $L$ times before moving on. Middle looping keeps the first and last third unique and repeats only the centre. Looping can also be confined to attention or to the MLP inside each block.

Why decoding gets slow. Training sees the whole sequence, so every position runs in parallel within a loop. Generation produces one token at a time, and loop $\ell$ for that token needs loop $\ell-1$'s output for the same token. A two-loop model therefore makes two passes over its weights per token. At small batch sizes, decoding is limited by reading weights from memory, so two passes take about twice as long.

The PLT shift. Parallel Loop Transformers (Wu et al., 2025) feed loop $\ell$ the output of loop $\ell-1$ from the previous token, plus the current token's embedding. When a new token arrives, every one of those inputs already exists. The $L$ loop inputs can be stacked as a batch and pushed through the shared weights in one call. Total FLOPs are unchanged, but the sequential chain per token is one call long and the weights are read once.

The catch. The previous token's state is the model's prediction of what comes next. It is not the token that actually arrived. Without the current-token embedding, every shifted loop has to infer the token from that prediction, and the paper's runs without it diverged or did worse than no looping at all. Re-injecting the embedding is necessary but, as published, not sufficient. The embedding is added to a state whose magnitude is roughly 10 to 600 times larger, so the state drowns it out. Stage 3 measures that ratio.

Looping. For a reused function $F_\theta$ (one block, a group of blocks, or the whole stack),

If $F_\theta$ holds $P$ weights and costs $C$, the looped model stores about $P$ and costs about $LC$. Middle looping with $M$ repeated blocks out of $B$ has effective depth $B+(L-1)M$. The paper repeats the middle $M=B-2\lfloor B/3\rfloor$ blocks.

What a loop costs. Training FLOPs per token are about $6\times$ the weights each token passes through, plus the attention scores. The output head, with $128{,}256\times d$ weights and untied from the input embedding, runs once per token whatever $L$ is:

This expression reproduces all ten training-FLOP coordinates in the paper's Figure 1 to within 0.1%. The coefficient 7 is what matches their accounting; the paper does not state it. Because the head is so large, two loops cost 1.45× per token at 100M and 1.80× at 1B, not 2×. Training tokens are set per stored weight (tokens per parameter, TPP). A looped model therefore trains on the same data as its no-loop twin and spends more compute on it.

Loop-specific adapters. Exact sharing makes the last loop apply the same transformation as the first, although its input has already been transformed several times. The paper gives each loop a low-rank correction (Definition 1) and trains it from scratch together with the shared weights:

One matrix then costs $d_{\text{in}}d_{\text{out}}+Lr(d_{\text{in}}+d_{\text{out}})$ weights, against $Ld_{\text{in}}d_{\text{out}}$ for fully separate loops. Loop Matrix lets loop $\ell$ mix all $L$ low-rank pairs, $\Delta W^{(\ell)}=\sum_j M_{\ell j}B_jA_j$, with a learned $M\in\mathbb R^{L\times L}$. Loop Softmax constrains each row of $M$ to a softmax. With rank 8 on all seven projections of the 300M model, the adapters add 1.1%, 2.1% and 4.3% to the weights at $L=2,4,8$. Doubling the depth with new layers would add 51.4%.

Parallel loops. With $\mathbf e_n$ the embedding of token $n$ and $\mathbf h^{(\ell)}_n$ the output of loop $\ell$, Cross-Loop Parallelism (Definition 2) is

with $\mathbf h^{(\ell)}_0=\mathbf 0$. The paper's unshifted control uses $\mathbf h^{(\ell-1)}_n$ instead. That is ordinary looping with the embedding added back at every loop, and it remains sequential.

Why the raw sum fails. If the state's RMS is $r$ times the embedding's and their directions are uncorrelated, the combined input points almost entirely along the state:

The paper measured $r$ between about 10 and 600, varying with loop and training budget. Its fix normalizes the two inputs separately and mixes them with a gate:

Here $W_{\text{down}}\in\mathbb R^{8\times2d}$ and $W_{\text{up}}\in\mathbb R^{d\times8}$. $W_{\text{up}}$ and $\mathbf b$ start at zero, so the gate starts at an exact 50/50 mix. A fixed $\mathbf g=\tfrac12$ works almost as well. Normalization is what matters; the gate choice is secondary.

Cache memory. Each loop attends over its own keys and values for earlier tokens, so separate caches multiply KV memory by $L$. Original PLT lets later loops reuse loop 1's cache and keep only a sliding window of $w$ positions of their own. That brings memory from $O(LND)$ to $O(ND+(L-1)wD)$. In this paper's experiments, separate caches gave the best quality. The normalized shifted PLT used in its 1B results keeps one cache per loop.

This executable Python example uses a toy shared block. Attention is replaced by a running mean over that loop's earlier inputs, but the block is still causal. The example trains PLT the way the paper's Algorithm 1 does, decodes it token by token as in Algorithm 2, and checks that the two give identical outputs. It then rebuilds Table 1's parameter count, the FLOP formula and the adapter overhead, and measures the cosine above on random vectors.

1import math, random
2
3rng = random.Random(0)
4D = 4
5W = [[rng.uniform(-0.8, 0.8) for _ in range(D)] for _ in range(D)]
6add = lambda a, b: [x + y for x, y in zip(a, b)]
7
8def block(x, history):
9    """A toy shared transformer for one position: mix with this loop's earlier inputs (attention), then transform."""
10    ctx = [sum(col) / len(history) for col in zip(*history)]
11    return [math.tanh(sum(w * v for w, v in zip(row, add(x, ctx)))) for row in W]
12
13def run_loop(inputs):                       # training: every position at once, causally
14    return [block(inputs[n], inputs[: n + 1]) for n in range(len(inputs))]
15
16def plt_train(E, L):                        # Algorithm 1: B(l) = E + Shift(H(l-1))
17    H = run_loop(E)
18    for _ in range(1, L):
19        H = run_loop([add(e, h) for e, h in zip(E, [[0.0] * D] + H[:-1])])
20    return H
21
22def plt_decode(E, L):                       # Algorithm 2: one token at a time
23    caches, prev, out, calls = [[] for _ in range(L)], [[0.0] * D] * L, [], 0
24    for e in E:
25        inputs = [e] + [add(e, prev[l - 1]) for l in range(1, L)]   # all ready before any loop runs
26        for l in range(L):
27            caches[l].append(inputs[l])     # each loop keeps its own cache
28        prev = [block(inputs[l], caches[l]) for l in range(L)]      # independent: one batched call
29        calls += 1
30        out.append(prev[-1])
31    return out, calls
32
33def looped_decode(E, L):                    # ordinary looping: loop l waits for loop l-1, same token
34    caches, out, calls = [[] for _ in range(L)], [], 0
35    for e in E:
36        x = e
37        for l in range(L):
38            caches[l].append(x)
39            x = block(x, caches[l])
40            calls += 1
41        out.append(x)
42    return out, calls
43
44E = [[rng.gauss(0, 1) for _ in range(D)] for _ in range(6)]
45decoded, plt_calls = plt_decode(E, 3)
46assert all(abs(a - b) < 1e-12 for u, v in zip(plt_train(E, 3), decoded) for a, b in zip(u, v))
47_, seq_calls = looped_decode(E, 3)
48print(f"6 tokens, 3 loops: {plt_calls} batched calls with PLT, {seq_calls} sequential calls without")
49
50# Table 1 shapes reproduce the reported non-embedding counts (blocks + untied output head).
51def shapes(blocks, d, mlp, q, kv, vocab=128_256, hd=128):
52    block = 2 * d * q * hd + 2 * d * kv * hd + 3 * d * mlp
53    return block, vocab * d, blocks * block + vocab * d
54block, head, total = shapes(16, 2048, 5632, 16, 2)
55assert total == 967_311_360                 # "1B": 967.311M in Table 1
56
57# Training FLOPs per token: 6 x weights per evaluation, 7 * ctx * d_attn per attention evaluation.
58def flops(blocks, d, mlp, q, kv, loops, tpp, ctx=8192):
59    block, head, total = shapes(blocks, d, mlp, q, kv)
60    per_token = 6 * head + loops * blocks * (6 * block + 7 * ctx * q * 128)
61    return per_token * tpp * total
62two = flops(16, 2048, 5632, 16, 2, loops=2, tpp=80)
63print(f"1B, two loops, TPP 80: {two:.3e} training FLOPs (paper: 1.07e21)")
64
65# Loop-specific LoRA, rank 8 on all seven projections, against doubling the depth (300M).
66blocks, d, mlp, q, kv = 10, 1280, 3584, 10, 1
67block, head, total = shapes(blocks, d, mlp, q, kv)
68per_loop = blocks * 8 * ((d + q * 128) + 2 * (d + kv * 128) + (q * 128 + d) + 3 * (d + mlp))
69print("LoRA adds", [f"{L * per_loop / total:.1%}" for L in (2, 4, 8)], f"vs {blocks * block / total:.1%} for 2x depth")
70
71# Raw re-injection: if the state is r times larger, b = e + h barely points at the current token.
72rms = lambda v: math.sqrt(sum(x * x for x in v) / len(v))
73cos = lambda a, b: sum(x * y for x, y in zip(a, b)) / math.sqrt(sum(x * x for x in a) * sum(y * y for y in b))
74e = [rng.gauss(0, 1) for _ in range(4096)]
75for r in (10, 100, 600):
76    h = [r * rng.gauss(0, 1) for _ in range(4096)]
77    raw = add(e, h)
78    mixed = [0.5 * x / rms(e) + 0.5 * y / rms(h) for x, y in zip(e, h)]   # RMSNorm both, gate 1/2
79    assert abs(cos(raw, e) - 1 / math.sqrt(1 + r * r)) < 0.05 and abs(cos(mixed, e) - 0.707) < 0.05
80    print(f"r = {r:3}: cos(e + h, e) = {cos(raw, e):.4f}; normalized 50/50: {cos(mixed, e):.3f}")

The equality check is the point of PLT. Decoding builds every loop's input from states cached at the previous token, so all three block calls for a token are independent. On real hardware they become one batched call over the shared weights. Training is unchanged: the loops still run in order, each over the whole sequence. In a real model, block is the full transformer stack, each loop's cache holds keys and values, and the RMSNorm-plus-gate combination replaces add(e, prev[l - 1]).

Which part should loop? At a fixed model size and data budget, more loops lower the loss only up to a point. At 100M the best loop count is between 4 and 16, and lower when the training budget is small. Beyond it some designs diverge, and at 32 loops Full-Stack looping jumps from 1.077 to 1.21 bits per byte. Looping attention and the MLP separately degrades first. Ranked per loop, Full-Stack looping is the most consistently strong design at 300M and 500M. Ranked per training FLOP, Middle-Stack looping gives the better return when compute is short, because two thirds of the network still runs once.

When should loops specialize? Only when there are several of them. At 300M with two loops the adapters change nothing measurable; at eight loops Loop Matrix lowers bits per byte by 0.0063. The paper makes Loop Matrix its default. Adapters do not fix decoding latency, and they are a parameter-efficient gain, not a compute-efficient one: a model with genuinely new layers is better at the same FLOPs, if you can store it.

Can parallel loops keep the gain? Not as originally published. The paper rebuilds PLT one piece at a time at 200M. The shift alone, KV sharing and the gated sliding window recover only a fraction of the gain: they sit close to the no-loop baseline, and the raw shift gets worse with more loops. The unshifted control, meanwhile, beats Full-Stack looping. Separate RMS normalization recovers most of the gap. At 1B with two loops the normalized shifted model beats sequential Full-Stack looping at every budget tested. At four loops it is only level with Full-Stack looping.

What does it buy overall? In this over-trained regime, the efficiency claims are relative to a scaling curve fitted to the no-loop runs. At 1B, a no-loop model would need 1.3–3.3× as many training tokens to match two loops. Charging the second loop, two-loop models trained on half the tokens of a no-loop model still match losses that the no-loop curve needs 6–37% more training FLOPs to reach, across 100M–1B. Multiple-choice accuracy at 1B rises by about 0.5–1.7 points, with no consistent change on MMLU.

What the study does not show. Every model is dense, at most 1B parameters, trained from scratch and evaluated mostly by likelihood. Most settings use one seed. The paper does not test mixture-of-experts or post-trained models, free-form generation, or parallel loops beyond four. It reports training FLOPs, not measured decoding latency. The batching argument for PLT is structural, and the speed-up you see will depend on whether decoding is limited by memory bandwidth.

You must ship a model with at most 1B stored weights, decode at batch size 1 on a memory-bound GPU, and train within $1.1\times10^{21}$ FLOPs. Should you loop, how many times, and with which design?

Loop twice. Using the paper's configurations, a no-loop 1B model at 160 tokens per weight costs $1.19\times10^{21}$ FLOPs, which is over budget. Two loops at 80 tokens per weight cost $1.07\times10^{21}$. By the paper's fitted curve, a no-loop model would need $1.13\times10^{21}$ FLOPs to match ordinary two-loop looping and $1.29\times10^{21}$ to match the normalized shifted PLT, so both beat anything the no-loop model can reach within budget. At those budgets, Figure 8 gives 0.8776 bits per byte for the normalized shifted model against 0.8835 for sequential Full-Stack looping. Ordinary looping would read the weights twice per generated token, roughly doubling latency at batch size 1. The normalized shifted PLT reads them once, as a batch of two. It keeps a separate KV cache per loop: $16\ \text{layers}\times2\times2\ \text{heads}\times128\times2\ \text{bytes}=16$ KiB per token per loop, or 256 MiB for both loops at 8,192 tokens. That is negligible next to the 2.5 GB of BF16 weights. Do not go to four loops: the paper finds less return per FLOP and a wider gap to the sequential control.

Source and scope

Sources:

• Ekbote, Pagnoni, Sachan, Bhosale, Lewis, Fazel-Zarandi, Rajput & Gonugondla, How to Loop Your Transformer: Loop Scaling, Specialization, and Parallel Inference (FAIR at Meta, arXiv:2610.40325, 29 October 2026).
• Wu et al., Parallel Loop Transformer for Efficient Test-Time Computation Scaling: the original PLT, with KV sharing and gated sliding-window attention.
• Bae et al., Relaxed Recursive Transformers: per-loop LoRA for a pretrained model converted into a looped one.
• Dehghani et al., Universal Transformers; Saunshi et al., Reasoning with Latent Thoughts; Geiping et al., Scaling up Test-Time Compute with Latent Reasoning.

Every measured value on this page was read from the vector paths of the paper's figures, so it carries the precision of the plot rather than of a table. The paper publishes one table (Table 1). Parameter counts, the FLOP formula, the cosine illustration, the KV-cache sizes, the code and the scenario are LiminalML derivations. We have not trained these models.