Harness Design Notes shared block unshared layers input injection
← All notes

Model architecture · the mechanism, the arithmetic, the reasoning

Looped Transformers

Take a block of layers and run it again, and again. A looped transformer gets its depth by reusing weights instead of adding them. That helps with reasoning, saves you memory, and does nothing for your latency bill.

An ordinary transformer is a stack of distinct layers. Layer 17 has its own weights, used once per token and never again. A looped transformer takes a short block of layers and applies it r times in a row, feeding the output of one pass back in as the input to the next. Depth goes up by the size of the block each time round. The parameter count stays where it is.

This is recurrence, but not the kind RNNs have. An RNN recurs across time, one token after another. A looped transformer recurs across depth, inside the forward pass for each token. Every position is still processed in parallel; each one just goes through the same weights more than once on the way up. Build one below. The parameter figures assume layers the width of Huginn-3.5B, the largest openly trained model of this kind, at about 395M parameters per layer.

Fig 1Build a looped modelinteractive
Presets
Same colour, same weights. Switch to the unrolled view to see what the hardware actually runs: a deep stack in which the shared block's layers recur. Parameter counts assume Huginn's layer shape (hidden size 5,280, gated MLP of 17,920) and a 65k-token vocabulary, with the concat adapter left out.

The anatomyFour parts, and why each one is there

The first looped transformer was the Universal Transformer (Dehghani et al., 2018): one layer, shared across all depth, applied as many times as each position needed. Recent designs keep that idea and add three things around it. In the 2025 recurrent-depth design of Geiping et al., the parts are:

  1. A prelude. A few ordinary, unshared layers that turn token embeddings into a latent representation e. Huginn uses two.
  2. The recurrent block. k layers applied r times to a state s. Huginn's block is four layers. The state starts as random noise rather than as e, which pushes the model to treat the loop as an iterative process instead of a fixed function.
  3. Input injection. Every iteration gets e again, either added to the state or concatenated with it and projected back down by an adapter. Without this, a deep loop slowly forgets what it was asked. Yang et al. found injection was what made looped models learn iterative algorithms well; Geiping et al. report that concatenation works better at scale and addition is fine for small models.
  4. A coda. A few more unshared layers that read the final state out to logits.

The training recipe matters as much as the wiring. If you always train at r = 32, you get a model that works at 32 and nowhere else. Huginn samples r fresh for every step from a heavy-tailed distribution with mean 32, mostly shallow with occasional very deep runs. To keep memory under control, it backpropagates through only the last 8 iterations. The result is a model whose depth you can choose at inference time. That dial is the whole reason to build it.

Adaptive depthNot every token needs the same number of loops

Once depth is a loop count rather than a fixed architecture, nothing forces every token to use the same count. A full stop doesn't need 32 passes; the digits of an arithmetic answer might. Something has to decide when each token is done, and four rules are in use. They differ in what they measure, when they decide, and whether the model has to be trained for them.

Fig 2Four stop rules: who gets how much depthinteractive
How to read it. Each column is one token of the sentence. Each row is one pass through the shared block, first pass at the top. A filled cell means that token went through the block on that pass; once a token stops, the cells below it stay empty. The number under each column is how many passes that token got. Pick a rule, read its stop condition, then move the slider.

An illustration, not a trained model. Each token has a made-up difficulty, and each rule's measured quantity (halting score, prediction change, router score) is simulated from it. The stop conditions themselves are the ones in the papers: Graves (2016) and Dehghani et al. (2018) for ACT, Geiping et al. (2025) for the KL exit, Bae et al. (2025) for the router.

On paper, every rule except the first saves compute: easy tokens leave early. Whether a GPU turns that into throughput is a separate question, and the honest answer is "not by default." There are two problems.

Serving adaptive depthKeeping the GPU busy when tokens stop at different times

Problem one: the batch moves in lockstep. A GPU serves many sequences at once by stacking their current tokens into one matrix and pushing it through the block together. If one token in the batch wants 16 passes and the rest want 4, the simple approach runs 16 batch-wide passes, and for 12 of them most rows compute nothing useful. The early exits saved no time at all.

Looped models have an escape that ordinary transformers don't. Every pass runs the same weights, so a token on its 2nd pass and a token on its 11th can share one matrix multiply. When a sequence's token finishes, that sequence can start its next token in the freed slot right away, without waiting for the rest of the batch.

Fig 3Lockstep batching against depth-wise refillinteractive
How to read it. Each row is one sequence in a batch of 8 being decoded. Time runs left to right: each column is one batched pass through the block. Each teal bar is one token going through its passes, and its length is how many passes its stop rule gave it; alternating shades separate consecutive tokens. Red cells are slots where the GPU runs the block and the row has nothing to do.

A simulation, not a benchmark. Each token's pass count is drawn at random from 2 to 16, skewed short, the way an exit rule distributes depth. It ignores attention cost, which still differs by depth, and the missing keys and values in Fig 4.

This is continuous depth-wise batching, proposed with Relaxed Recursive Transformers (Bae et al., 2024). In their simulation with ideal exits, a recursive Gemma reached 2.66× the throughput of the original, and 1.88× a baseline that already used ordinary continuous batching. Their later Mixture-of-Recursions paper measured up to 2.06× throughput at 360M parameters. Only the matrix multiplies batch this neatly; attention still reads a different cache for each depth, which brings up the second problem.

Problem two: stopped tokens leave holes in the KV cache. In a transformer, a token at pass d attends to the keys and values every earlier token produced at pass d. If an earlier token stopped at pass 3, it never produced anything for passes 4 and up. A later token on pass 6 reaches for that entry and finds nothing. There are four published ways to deal with it.

Fig 4Four ways to handle the missing keys and valuesinteractive
How to read it. Each column is a token already in the context, and the rightmost is the new token being generated. Each row is a pass. A solid cell means that token has keys and values stored for that pass; earlier tokens stopped at different passes, so there are holes above them. The new token is now on pass d, and its attention needs row d from every earlier token: the outlined row. Pick a strategy to see what it actually reads.

Schematic. The exit depths are fixed and illustrative. "Stored" counts key/value entries across the 8 passes shown; "extra computed" counts entries that have to be produced for tokens that had already stopped.

None of the four is free. Copying fills every hole but brings back compute that early exit was supposed to save. Reading the deepest entry is free but stale. Per-depth caching is exact but shrinks what deep passes can see. Sharing the first pass is the cheapest and coarsest. Which one is acceptable is an empirical question about your model, so evaluate it on your own traffic rather than on the paper's benchmarks.

Two serving rules follow from all this. Decide depth at a granularity your kernels can actually exploit: fixed-capacity routing, or refill with depth-wise batching. And budget the KV handling as part of the cost of early exit, not as a detail to fix later.

What depth buysReasoning is serial; parameters are storage

The case for looping is an argument about what layers are for. Some of what a model does is recall: facts, vocabulary, associations, and those live in parameters. The rest is computation. Composing one lookup with another, carrying a digit, following a reference to its referent: each of these is a step that has to happen after the previous one. In a transformer, steps that depend on each other inside one forward pass have to sit at different depths. A model that is 12 layers deep cannot do 40 dependent steps in one go, however wide it is.

Looping lets you buy those steps without buying more storage. Saunshi et al. (ICLR 2025) put it in one sentence: many reasoning problems require a large depth but not necessarily many parameters. They show that a k-layer block looped L times nearly matches a non-looped kL-layer model on addition, multi-hop induction and several language reasoning benchmarks, and they prove a looped model can simulate T steps of chain-of-thought with T loops. The figure below makes the serial argument concrete with the simplest multi-step task there is.

Fig 5One loop, one hop: resolve the variableinteractive

The task. The prompt below is a list of definitions such as g = d, and one variable is defined as a number. The question asks for the value of one variable, marked with an orange ring. In the diagram, each circle is one variable, which is one line of the prompt, and its grey arrow points to what it is defined as. The amber square is the number, where the chain ends.

A hop is one substitution: take the variable you are holding, find its line, and replace it with the right-hand side. For a transformer that is one attention lookup. It is serial: you can't do hop 5 before you know where hop 4 landed. In this toy, one pass of the loop does one hop. The asked variable's own line gives the first hop for free, and the teal arc shows what the model currently believes it equals.

"Stop early once the answer stops changing." After each pass, compare the model's current answer with its answer after the previous pass. If they are the same, assume it's finished and stop, even if more loops were allowed. When the model reaches the number there is nothing left to substitute, so the next pass gives the same answer and it stops. It always spends one extra pass noticing. This is the KL rule from Fig 2 in its simplest form.

Pointer doubling is a smarter block. On every pass, every line substitutes once, all in parallel: g = d and d = l become g = l. Each line's reach doubles (1, 2, 4, 8 hops), so a 7-hop question takes 3 passes instead of 6. Attention can do this in one layer, since each line only has to look up one other line.

A toy algorithm in place of a trained block. The point it makes holds for real models: the number of passes needed is set by the problem, not by the parameter count, so the same weights solve longer problems when you let them loop longer.

Three things to try. First, allow fewer loops than the question needs: the model doesn't fail loudly. It answers with a variable halfway down the chain, which is what an under-deep model does on a problem that's too long. Second, raise the loops without touching anything else. The same weights now solve it. This is the length-generalisation result of Fan et al.: looped models trained with an adaptive number of steps extrapolate to longer inputs by looping longer at test time, where fixed-depth models stall. Third, tick "stop early" and allow plenty of loops. It stops one pass after the answer settles and hands back the unused budget.

The empirical record so far fits the argument:

Huginn-3.5B

Trained on 800B tokens. On GSM8K with chain-of-thought it scores 0% at one iteration and 42% at 32 (flexible match), and the authors report gains up to a compute load equivalent to a 50B-parameter model. Harder tasks keep improving with more iterations after easy ones have saturated.

Ouro 1.4B / 2.6B

Zhu et al. pre-trained looped models on 7.7T tokens and report that they match models up to 12B across a range of benchmarks. Their controlled experiments credit better manipulation of knowledge, not more knowledge stored.

Tiny Recursive Model

Jolicoeur-Martineau: a single 2-layer network recursed many times. 7M parameters, 45% on ARC-AGI-1 and 8% on ARC-AGI-2, ahead of most of the LLMs it was compared with on those puzzles.

The limit

Saunshi et al. find looped models lag on memorisation-heavy tasks such as closed-book QA, where non-looped models with more parameters win. Loops add thinking time. They add nothing to what the model knows.

Parameters are where a model keeps what it knows. Depth is how many steps it gets to think about it. Looping buys only the second.

The arithmeticIt stores like a 3.5B model and runs like a 53B one

This is where the free lunch ends. Recall from Speculative Decoding that generating a token is limited by memory bandwidth: the GPU has to stream every weight it uses out of HBM for each decode step. A looped model uses the same weights many times, so you might hope it reads them once and then works from on-chip memory.

It doesn't, because they don't fit. An H100 has 50 MB of L2 cache. One Huginn layer is about 0.8 GB in 16-bit. So every iteration streams the block's weights out of HBM again, exactly as a model with distinct layers would stream its own. Count the bytes moved and the FLOPs done for one token, and a looped model at r iterations matches an unshared model with the same effective depth, term for term. Looping at r = 32 gives Huginn 132 layer evaluations per token, and an unshared 132-layer model at that width has about 53B parameters. Geiping et al. put the same comparison at about 50B.

Per layer, hidden size d, gated MLP width f L = 4d² + 3df      (≈ 395M at Huginn's shape)
Stored versus run, P + k + C distinct layers, E = P + k·r + C evaluated params stored   = (P + k + C)·L + V·d
bytes read per decode step = 2·(E·L + V·d) + batch · context · 4d · E
KV cache stored  = batch · context · 4d · (P + C + k·rcached)
step time       ≈ max(bytes / bandwidth, FLOPs / peak)

So the gain is in capacity, not speed. The weights take 7 GB instead of 105 GB, which is the difference between one GPU and a tensor-parallel group, and the freed memory can go to KV cache and larger batches. Price it yourself:

Fig 6Serving a looped model on one H100interactive
batch

A roofline model, not a benchmark. H100 SXM: 80 GB, 3.35 TB/s, 989 dense BF16 TFLOPS, all at 100% efficiency, so real numbers will be slower in proportion. BF16 weights and cache, full multi-head attention, no weight caching across iterations. The two KV-sharing options follow Geiping et al. (a cache budget that reuses slots every few iterations) and MoR (reuse the first recursion's keys and values); both papers report small quality costs, which this model doesn't price in.

Look at the latency chart first: the looped line and the same-depth line are the same line. At batch 1 and 4k context, Huginn at 32 loops takes about 35 ms per token. Without the loop it would take about 2 ms. The extra 33 ms is the price of the extra thinking, and the loop gives no discount on it. What changes is memory: the 53B-equivalent no longer fits on the card, while the looped model leaves most of the card free.

Two exceptions. A model small enough to keep its block on-chip really does avoid re-reading weights. TRM's 7M parameters take 14 MB in 16-bit and would fit in L2 several times over, and SRAM-heavy accelerators shift the line further. And KV sharing saves memory but not time: each layer evaluation still attends over the full context, so the reads remain.

The comparison that flatters looping isn't with a deeper model. It's with chain-of-thought. Reasoning in text costs one full forward pass per thought token, and every thought sits in the KV cache for the rest of the request. Reasoning in loops costs extra depth per token, with nothing added to the context. Saunshi et al.'s result that T loops can simulate T steps of CoT is what makes that trade available, though whether the trade is good depends on how many tokens your workload's reasoning would have taken.

The risksWhere it bites in practice

Latency scales with the dial

Every loop is another pass over the block's weights. If users choose the depth, or the model halts adaptively, your latency distribution now has a tail set by the hardest prompts. Cap r per request and track it like any other SLO input.

The KV cache grows with every iteration

Naively, each iteration's attention layers add their own keys and values: Huginn at 32 loops caches 132 layers' worth per token, as much as the 53B model would. Sharing schemes cut this sharply, but they're an approximation you have to evaluate on your own traffic.

Ragged depth fights batching

Per-token halting saves FLOPs only if your kernels can compact the batch as tokens exit. Otherwise the batch runs at the depth of its deepest token. Fixed-capacity routing (MoR) keeps the batch shape known in advance, at the cost of a capacity limit the router must respect.

Training is costlier and touchier

Backpropagating through r iterations multiplies activation memory, which is why Huginn truncates to the last 8. Deep weight-tied recurrences are also prone to instability; Geiping et al. describe failed early runs in which the model learned to ignore the loop.

Extrapolation has a ceiling

Looping past the depths seen in training helps only if the model learned a process that converges. Gains on easy tasks saturate after a handful of iterations, and a model trained at one fixed depth may degrade rather than improve when you loop it further.

It cannot learn facts it has no room for

A 3.5B looped model has the storage of a 3.5B model. On knowledge-bound workloads like closed-book QA, long-tail entities and API trivia, a larger flat model wins, and no loop count changes that.

When it paysA capacity tool with a thinking dial

Put together, a looped transformer is worth considering when memory capacity, not bandwidth or FLOPs, is what you're short of: one GPU instead of eight, an edge device, or a deployment where KV cache headroom sets your batch size. It fits best when the workload is reasoning-heavy and uneven in difficulty, so a per-request depth dial earns its keep, and when the alternative is a long chain-of-thought whose tokens you would otherwise pay for in time and cache.

It is the wrong tool when you're serving a busy endpoint at high batch, where you're compute-bound and a loop costs the same FLOPs as a bigger model; when the workload is knowledge-bound; or when the requirement is simply lower latency. Then a smaller model, or speculative decoding on the one you have, buys more.

The honest summary of the architecture is one line: it decouples what a model knows from how long it thinks, and you pay for the thinking in full every time.