One Model to Loop Them All: Elastic Latent Recurrence for Language Models

AI summary

Overview

Research area: Natural Language Processing — efficient language model architectures and training recipes (latent recurrent / "looped" models, mixture-of-experts, inference-time compute scaling).

Technical level: Intermediate. The paper assumes familiarity with transformer blocks, attention, KV caching, and MoE routing, but its central argument is explained conceptually.

Scope: A single paper presenting a training recipe — Elastic Latent Recurrence (ELR) — that makes a looped language model match a non-recurrent model of the same size at one loop while improving with additional loops, validated on 0.5B, 4B, and 30B-A3B models plus a post-trained 4B model named Koi-4B.

What This Paper Is About

Latent recurrent (looped) language models reuse a shared block of layers repeatedly, so a single checkpoint can trade more inference compute for more accuracy simply by changing the loop count. The catch, which the authors call the "one-loop tax," is that at a single loop these models fall short of a non-recurrent model of the same size — Huginn, trained with a mean recurrence of 32, scores 34.9 against 46.4 on ARC-Easy even after 800B tokens, and Parcae is below its Transformer baseline at one loop. The paper's goal is to remove that tax so that one checkpoint is at least as good as the matched non-recurrent baseline at one loop and strictly better at additional loops.

Key Contributions

1. A recipe for elastic depth without a one-loop tax. A gated state update (adaptive delta with a learned per-token, per-channel step size) combined with mixing loop counts within each optimizer update, rather than across steps, preserves the one-loop endpoint; ablations show that uniform sampling, weighted sampling, and parameter-free update rules do not.

2. Elastic depth acquired late in pretraining. Training at one loop for the first 75% of tokens and switching to mixed depth only for the final quarter costs an estimated 16% more training FLOPs than the baseline (versus 62% for mixed-depth training throughout, in the 4B configuration); with zero-initialized loop conditioning grafted at the switch, this late schedule reaches lower loss than mixed-depth training from scratch.

3. Cross-loop route replay for looped MoE models. Reusing first-loop expert assignments on later loops cuts mixed-depth update time by 11.7% with no measurable change in accuracy.

4. Koi-4B. A post-trained elastic-depth model whose single checkpoint spans an estimated 18 to 29 on the Artificial Analysis Intelligence Index by loop count alone.

Main Findings

• One-loop parity on two scales. On GDP-Hybrid-4B, the ten-task average at one loop is 65.13 for ELR versus 65.65 for the matched non-recurrent baseline (0.52 points below). On GDP-Hybrid-30B-A3B, ELR reaches 66.08 and ELR with route replay 65.95, versus the baseline's 65.04.

• Each additional loop buys accuracy. From one to four loops, the ten-task average rises by 7.25 points on GDP-Hybrid-4B and 4.75 points on GDP-Hybrid-30B-A3B. Largest per-task gains are GSM8K (+11.8) and MATH-500 (+11.3) for the 4B model, and MMLU-Pro (+7.5) and GSM8K (+6.5) for the 30B-A3B model.

• Diminishing returns concentrate in the second loop. Two and three loops capture 79% and 96% of the one-to-four-loop gain on GDP-Hybrid-4B, and 87% and 91% on GDP-Hybrid-30B-A3B. The mixed-depth schedule gives p₁ = 0.325, p₂ = p₃ = p₄ = 0.225, and a mean loop count of 2.35, so later loops appear in fewer training examples.

• Adaptive delta is the only tested state update that matches the baseline at one loop. On GPT-0.5B it achieves the lowest loss at every loop count. A learned but input-independent "delta" rule ranks second, trailing by 0.006–0.007; the parameter-free "replace" rule has the highest one-loop loss and worsens from four to five loops, while "residual" stops improving after three loops.

• How loop counts are mixed matters more than the marginal distribution. At the same p₁ = 0.32, one-loop validation loss falls from 1.6871 (synchronized weighted sampling) to 1.6830 (15% one-loop rank override) to 1.6774 (independent-rank sampling), but independent-rank sampling costs about 36% more step time. Without overrides (f = 0), uniform sampling sits 0.0225 above the baseline at one loop; f = 0.15 is the smallest tested fraction below it.

• A one-loop-only model can beat the non-recurrent baseline. On GPT-0.5B it reaches 1.660 loss versus the baseline's 1.685, even though both execute each unique layer once.

• Delaying mixed depth is cheap. Mixed-depth training from scratch costs an estimated 1.79× baseline training compute on GPT-0.5B; starting after 25%, 50%, or 75% of tokens lowers this to 1.60×, 1.40×, and 1.20× while retaining 97%, 90%, and 84% of the five-loop improvement. At the 75% start, added compute falls from about 79% to 20% for a 0.009 increase in five-loop loss.

• Route replay trades a small loss increase for speed. Per-expert selection counts at loops one and four have cosine similarity 0.978 averaged over the ten recurrent MoE layers, yet individual tokens retain only 53% of their first-loop experts at loop four. Inference-only replay raises four-loop validation loss by 0.0055 and training with replay by 0.0038, while the ten-task average stays at 70.90 versus 70.83. Replay reduces router evaluations from 53 to 23 per four-loop forward pass, cutting update time by 11.7% over the mixed-depth phase and inference forward time by 1.2%.

• Inference cost at one loop is identical to the baseline. ELR at k = 1 matches the non-recurrent baseline's TTFT and per-token decode latency. At four loops it reaches 2.33× baseline TTFT and 2.06× decode latency on GDP-Hybrid-4B; three loops capture 96% of the accuracy gain at 1.87× TTFT.

• Koi-4B scales from 18 to 29 on the AA-Index by looping alone. Estimates rise from 18 at one loop to 26, 28, and 29 at two, three, and four loops. At one loop it is on par with Qwen3.5-4B; at two loops it 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 (about 6.5× as many total parameters) and trails only Qwen3.6-35B-A3B (nearly 9× as many total parameters). All scores come from the authors' own harness, not official Artificial Analysis results.

• Higher loop counts produce shorter outputs. Relative to one loop, output tokens per rollout fall by 21.8%, 29.5%, and 32.7% at two, three, and four loops. Because two-loop decode latency is about 1.36× but rollouts are 0.78× as long, decode time per rollout rises only about 6%; at four loops the figures are 2.06× and 0.67×, about 1.4× per rollout.

• Last-loop KV caching flattens memory. Per-loop caching at four loops needs 2.5× the one-loop KV memory (7.5 GiB versus 3.0 GiB per 128K-token request in BF16); last-loop caching stays at 1× at every loop count, with four-loop score reductions of 0–3% and an improvement on AA-Omniscience non-hallucination (+2.3). GDP layers keep per-loop state below 25 MiB per request at k = 4.

Methodology in Plain English

The authors start from a non-recurrent model as the reference point and ask what it costs to make that same model elastic. They split the backbone into three parts: an input block, a shared recurrent block holding roughly the middle half of the layers, and an output block. Because the blocks are disjoint, running one loop executes every unique layer exactly once — the model at k = 1 is structurally the same size and cost as the baseline.

Three design decisions do the work. First, instead of overwriting the hidden state with the recurrent block's output, each loop takes a gated step toward the proposed revision, with a learned step size per token and per channel (initialized so the first step is about 0.35). Second, loop counts are mixed inside a single optimizer update: a rotating subset of data-parallel ranks is forced to run one loop on every microbatch, so one-loop and multi-loop gradients are averaged together before the step. Third, each pass is told its loop index and total budget through a sinusoidal embedding fed into AdaLN/FiLM-style modulation of the layers, with all conditioning heads zero-initialized so the model starts out behaving exactly like the one-loop checkpoint.

To keep training affordable, the mixed-depth phase is postponed to the final quarter of pretraining. For MoE variants, the router is evaluated only on the first loop and its expert choices are reused afterwards. At inference, a single KV cache is overwritten each loop rather than kept per loop. The recipe is validated in stages: a 0.5B dense GPT model isolates each design choice, then 4B dense and 30B-A3B MoE hybrids are trained alongside matched non-recurrent baselines on the same data, and finally the full recipe including conditioning is carried through mid-training, SFT, and RLVR to produce Koi-4B.

Why This Matters

Impact on research. The paper reframes the looped-model question. Rather than asking whether looping is a compute-efficient way to scale, it asks what it costs to make an existing model elastic — and shows the cost can be small in accuracy and roughly 16% extra training FLOPs in the 4B configuration. It also documents that a common practice (sampling one loop count per optimizer step) is what breaks the one-loop endpoint, and that per-depth-optimal fixed-depth models are not the right comparison for an elastic checkpoint.

Real-world applications (implications of the work, not claims tested in the paper):

• Serving a single checkpoint at several quality/latency tiers, much like a reasoning-effort setting but without generating extra tokens.
• Deploying on devices with a fixed memory budget, where a 4B model fits and a 30B-total MoE does not, and looping lets the model that fits spend compute to close the gap.
• Long-context serving, where last-loop KV caching holds attention KV memory at its one-loop size (3.0 GiB versus 7.5 GiB per 128K-token request at four loops in BF16).
• Reasoning workloads where extra latent depth substitutes for explicit chain-of-thought tokens, since higher loop counts produced 21.8–32.7% shorter rollouts at higher scores.

Industry relevance. The recipe is deliberately low-friction: every ELR model trains with its baseline's optimizer settings unchanged, delayed mixed depth avoids paying multi-loop cost across the whole run, and route replay plus last-loop caching cut real serving and training costs. A model family that ships as one checkpoint with an inference-time depth knob simplifies the model catalog and the capacity planning behind it.

Future Directions

• Post-training the MoE variant. The authors explicitly leave post-training of the 30B-A3B mixture-of-experts model to future work; only the dense 4B model is carried through mid-training, SFT, and RLVR.
• When first-loop routes stop being reusable. Route replay assumes first-loop expert assignments remain useful later; tokens retain only 53% of their first-loop experts at loop four, so the assumption's limits at larger loop counts or different MoE configurations are open.
• Whether last-loop KV caching holds beyond four loops. The scheme is applied at inference to a model trained with per-loop caches and costs 0–3% at four loops; its behavior at higher loop counts and longer contexts is not reported.
• Choosing the override fraction and switch point in general. The paper sets the one-loop override so p₁ is approximately one third and switches to mixed depth at 75% of tokens based on 0.5B ablations; how these choices transfer to other model sizes, loop ranges, and data budgets is not reported.

Target Audience

Machine learning researchers and engineers working on efficient language model architectures, inference-time compute scaling, and serving systems. It is most useful to readers who already understand transformer blocks, KV caches, and mixture-of-experts routing, and who want a concrete, ablation-backed recipe for building a checkpoint whose depth can be dialed at serving time. Practitioners focused on deployment cost and model-catalog simplification will find the training-cost and memory numbers directly actionable.

Authors’ abstract

A latent recurrent language model applies a shared block of layers repeatedly, so one checkpoint can trade inference compute for accuracy by changing the loop count. At one loop, however, such models typically fall well short of a non-recurrent model of the same size, a one-loop tax that makes a separate non-recurrent model the better choice whenever latency matters. We present Elastic Latent Recurrence (ELR), a recipe for training looped language models without this tax. At one loop, an ELR model costs the same to serve as a non-recurrent model of the same size and, trained on the same data, reaches similar accuracy. Each additional loop then buys more accuracy for more inference compute. Three choices make this work: a gated state update with a learned step size, loop counts mixed within each optimizer update, and conditioning each pass on its loop index and budget. The first two preserve one-loop accuracy, and the third lowers loss at every additional loop. Mixed depths can be introduced late in pretraining to limit the added training cost, and with zero-initialized conditioning added at the switch, this late schedule reaches lower loss than mixed-depth training from scratch. In MoE models, reusing first-loop expert routes on later loops cuts mixed-depth update time by 11.7% with no measurable change in accuracy. At inference, keeping only the last loop's attention KV cache holds attention KV memory at its one-loop level while retaining most of the gain from additional loops. On 4B dense and 30B-A3B MoE hybrid models, ELR at one loop is 0.52 points below and 1.04 points above its matched baseline on a ten-task average, and four loops add 7.25 and 4.75 points. Carried through post-training, the recipe yields Koi-4B, a single checkpoint that scores an estimated 18 on the Artificial Analysis Intelligence Index at one loop, matching Qwen3.5-4B, and 29 at four loops.