Can Elastic Latent Recurrence Remove the One-Loop Tax in Looped Language Models?

The one-loop tax is a starting-line deficit

A looped, or latent recurrent, language model runs a shared block of layers more than once per token. Because the block is reused, one set of weights can be served at several effective depths: more loops mean more inference compute and, in principle, more accuracy. The paper calls this property elastic depth, and treats the loop count as a quality-versus-test-time-compute knob that, unlike chain-of-thought effort settings, generates no extra tokens.

The problem is the endpoint. At one loop, earlier looped models fall well short of a non-recurrent model of the same size. The paper cites Huginn, trained with a mean recurrence of 32, scoring 34.9 against its fixed-depth baseline's 46.4 on ARC-Easy after 800B tokens, and notes Parcae sits below its Transformer baseline at one loop and above it at eight. The authors name this gap the one-loop tax. A model that pays it must spend several loops recovering lost ground before extra compute buys anything, so a separately trained non-recurrent model remains the better choice when latency matters.

The paper is explicit that a narrow loop range does not fix this by itself: in its ablations, sampling uniformly over one to five loops still leaves the one-loop endpoint behind the baseline. The question ELR addresses is therefore narrower than whether looping is a compute-efficient way to scale a model. It asks what it costs to make a non-recurrent baseline elastic without paying the tax.

Why one loop must execute each layer once

ELR splits a backbone into an input block, a shared recurrent block, and an output block. With k loops, the forward pass traverses L_in + k·L_rec + L_out layers, so each additional loop adds recurrent-block layers of compute and latency at a fixed parameter count. At k = 1, every unique layer executes exactly once, which is what makes the one-loop model cost the same to serve as a non-recurrent model of the same size.

This partitioning is a design choice, not a universal property of looped models. The paper contrasts it with Huginn, which initializes the recurrent state from a truncated Gaussian and projects the concatenation of the previous state and the input representation, and with Ouro, which has no separate input or output blocks and repeats the full stack without re-injecting the input. Those architectures do not guarantee that one loop equals one pass over unique layers, which is part of why the one-loop endpoint can lag.

Three choices that protect the endpoint

The first choice is the state update. Instead of replacing the state with the recurrent block's output, ELR takes a gated step along the proposed change, with a learned step size per token and channel. In ablations on a GPT-0.5B model, update rules with learned step sizes keep validation loss closer to the non-recurrent baseline across one to five loops than parameter-free rules such as residual addition or replacement.

The second choice is where loop counts are mixed. Sampling one loop count per optimizer step and varying it across steps, as in Huginn's locked-step sampling, does not preserve the one-loop endpoint even when one loop is drawn a third of the time. ELR instead mixes loop counts within each optimizer update, implemented by having a rotating subset of data-parallel ranks run one loop on every microbatch. The paper reports that the first two choices together decide whether the one-loop endpoint survives mixed-depth training.

The third choice is conditioning each pass on its loop index and budget. This does not protect the endpoint; it lowers loss at every additional loop by letting passes specialize. The paper separates these roles clearly, which matters for anyone deciding which ingredient to adopt first.

Mixed-depth training can start late

Training every step at mixed depth is expensive, so the paper tests introducing mixed depths only in the final quarter of pretraining. On GPT-0.5B, the delayed schedule with loop conditioning added at the switch reaches lower validation loss at every loop count than unconditioned mixed-depth training from scratch, 1.6147 versus 1.6273 at five loops, despite using mixed depths for only the last quarter of tokens. Conditioning is grafted on with zero-initialized heads, so the conditioned model starts as the trained one-loop checkpoint and deeper passes learn to depart from it.

The paper puts the added training cost of this late switch at about 1.16× the baseline's training compute. It also notes that it does not evaluate extrapolation beyond the trained depth, and that Huginn reported conditioning on the step index hurting extrapolation in its setting. The late-switch result is therefore about reaching lower loss within the trained loop range, not about generalizing to unseen depths.

What four loops buy at each scale

On GDP-Hybrid-4B, ELR's ten-task average at one loop is 65.13 against the non-recurrent baseline's 65.65, a 0.52-point deficit. On GDP-Hybrid-30B-A3B, ELR and ELR with route replay reach 66.08 and 65.95 against the baseline's 65.04, slightly above it. Individual tasks move in both directions at one loop: on the 4B model, ELR trails by 5.3 points on GSM8K and leads by 5.2 on MATH-500, so the aggregate parity hides task-level variation.

From one to four loops, the ten-task average rises by 7.25 points for the 4B model and 4.75 points for the 30B-A3B model. Most of that arrives early. Two and three loops capture 79% and 96% of the four-loop gain on the 4B model, and 87% and 91% on the 30B-A3B model. The authors attribute the diminishing returns to later loops appearing in fewer training examples: with the one-loop override, loops two, three, and four occur in 67.5%, 45%, and 22.5% of mixed-depth examples, giving a mean loop count of 2.35.

Inference cost follows the same shape. At one loop, ELR matches its non-recurrent baseline's time-to-first-token and per-token decode latency. At four loops on the 4B model, those rise to 2.33× and 2.06× baseline. Three loops therefore capture 96% of the one-to-four accuracy gain at 1.87× baseline time-to-first-token, which is the paper's clearest quality-per-cost comparison.

Reusing first-loop expert routes saves training time

In the MoE hybrid, routing drifts across loops per token even though aggregate expert selection is stable: per-expert selection counts at loops one and four have a cosine similarity of 0.978 averaged over the ten recurrent MoE layers, yet individual tokens retain only 53% of their first-loop experts at loop four. The paper measures the effect of that drift by feeding each recurrent MoE layer its loop-four input and computing the routed-expert output twice, once with the router's current choice and once with the token's loop-one routes. The two outputs have a mean cosine similarity of 0.767, against 0.227 for a control that shifts first-loop routes across positions.

Reusing first-loop routes on later loops therefore keeps much of the token-specific routing information. Training with replay raises four-loop validation loss by 0.0038 relative to dynamically routed ELR, while leaving the ten-task average at four loops unchanged within noise, 70.90 versus 70.83. The gain is in training time: update time falls by 11.7% over the mixed-depth training phase and by 19.6% at fixed four loops, because replay removes later-loop auxiliary-loss calculations and backward work through repeated router evaluations. Inference forward time falls by only 1.2% at 8K tokens, since expert dispatch and computation remain.

Keeping KV memory at its one-loop level

Standard inference keeps a separate attention KV cache for each loop, so memory grows with depth. In Koi-4B, three of six attention layers sit in the recurrent block, so at four loops per-loop caching needs 2.5× the one-loop KV memory, 7.5 GiB against 3.0 GiB per 128K-token request in BF16. Last-loop KV inference instead shares one cache across loops: each loop overwrites the cache, and later tokens read earlier tokens' final-loop keys and values. The scheme needs no retraining and holds attention KV memory at its one-loop size at every loop count. GDP layers keep per-loop state that does not grow with sequence length and stays below 25 MiB per request at four loops.

The accuracy cost is small but not zero. At four loops on Koi-4B, last-loop caching changes benchmark scores by between −3.0 and +2.3 points relative to per-loop caching, with the largest drop on AA-LCR and a gain on AA-Omniscience non-hallucination. Output length also falls, to 0.80× to 1.01× the per-loop count depending on benchmark. The paper's claim is that most of the gain from additional loops survives, not that the two caching schemes are equivalent.

Koi-4B trades latent depth for output tokens

Carried through pretraining, mid-training, SFT, and RLVR, the recipe yields Koi-4B, a single checkpoint evaluated at one to four loops on the Artificial Analysis Intelligence Index v4.1 suite using the authors' own harness. Its estimated index score is 18 at one loop, matching Qwen3.5-4B, then 26, 28, and 29 at two, three, and four loops. The second loop contributes 8 of the 11 points gained from one to four loops. At two loops the model is above Qwen3.5-9B and Nemotron-3.5-Lightning and level with Gemma4-26B-A4B; at four loops it is 3 points above Gemma4-26B-A4B, which has about 6.5× as many total parameters, and trails only Qwen3.6-35B-A3B, which has nearly 9× as many.

These comparisons are not matched-baseline comparisons. Koi-4B has no non-recurrent baseline trained on the same data, so its one-loop score is set against open models of similar size, and all scores come from the authors' harness rather than from Artificial Analysis, so they may differ from published figures. The matched-baseline evidence for the one-loop claim comes from the 4B and 30B-A3B ablations, not from Koi-4B.

The same checkpoint also produces shorter responses at higher loop counts while scoring higher. Relative to one loop, output tokens per rollout fall by 21.8%, 29.5%, and 32.7% at two, three, and four loops, as a geometric mean over benchmarks. The authors read this as the model substituting latent depth for explicit chain-of-thought tokens. The shorter generations partly offset the added per-token cost: at two loops, per-token decode latency is about 1.36× that of one loop, but rollouts are 0.78× as long, so decode time per rollout rises by only about 6%. At four loops the figures are 2.06× and 0.67×, or about 1.4× per rollout.

Where the recipe stops paying off

The paper's own accounting points to the second loop as the decision threshold. It supplies 8 of the 11 index points Koi-4B gains from one to four loops, and two and three loops capture 79% and 96% of the four-loop gain on the 4B ablation. Going to four loops costs 2.33× baseline time-to-first-token and 2.06× per-token decode latency on that model, while three loops already capture 96% of the accuracy gain at 1.87× time-to-first-token. For a latency-sensitive deployment, the evidence favors two or three loops over four.

The authors leave several questions open. They have not tested whether reweighting the loop-count distribution recovers the diminishing returns, though they note later loops appear in fewer training examples. Conditioning closes most of the gap to fixed-depth specialists at intermediate depths but not at the endpoints, and they have not isolated why. Their serving setting fixes the loop count per request, and combining ELR with per-request or per-token adaptive depth is left as future work, as is post-training the MoE variant. The one-loop parity claim is also established against matched baselines only at 4B dense and 30B-A3B MoE scale, so it should not be assumed to hold at other sizes without evidence.

Sources used in this answer