PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Faisal Ladhak, Shantanu Acharya, Santiago Akle Serano, Simeng Sun, Boris Ginsburg, Krishna C Puvvada

arXiv:2610.12582v1cs.CL

TL;DR

Looped language models can offer adjustable inference compute, but typically pay a one-loop accuracy tax against same-size non-recurrent models. ELR combines a gated update, within-step depth mixing, and loop conditioning; its evaluations report near-baseline one-loop results and gains at four loops.

  • Problem

    Looped language models can vary inference compute by changing loop count, but existing models may fall far behind same-size non-recurrent baselines at one loop, limiting their appeal for latency-sensitive use.

    Source: §1 · pp. 1–2

  • Method

    ELR trains looped language models with a gated state update, loop counts mixed within each optimizer step, and conditioning on each pass's loop index and budget.

    Source: §1 · p. 2

  • Results

    At one loop, ELR is 0.5 points below its 4B baseline and 1.0 point above its 30B-A3B baseline; four loops add 7.25 and 4.75 points, respectively.

    Source: §8 · p. 15

  • Takeaways & Limitations

    A single checkpoint can trade inference compute for quality by changing loop count, while executing each unique layer once at one loop.

    Source: §1 · p. 2

    The evaluated serving setting fixes loop count per request; combining ELR with per-request or per-token adaptive depth remains future work.

    Source: §8 · p. 15

Abstract

from arXiv · show

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.

1. Introduction

ELR aims to remove the one-loop tax so a single checkpoint can trade inference compute for quality. Its training recipe preserves baseline-like accuracy at one loop while adding gains at higher loop counts.

1. Introduction

  • Existing looped models can lag non-recurrent baselines at one loop, forcing latency-sensitive deployments to use several loops before extra compute buys gains.The paper calls this gap the one-loop tax and asks how to make a model elastic without paying it.Source: §1 · pp. 1–2
  • ELR combines a gated state update, loop counts mixed within each optimizer step, and conditioning on loop index and budget to preserve the one-loop endpoint and lower loss with added loops.At one loop, each unique layer executes once and inference compute matches that of a same-size non-recurrent model.Source: §1 · p. 2, Figure 1 · p. 2
  • At one loop, ELR is 0.5 points below its 4B baseline and 1.0 point above its 30B-A3B baseline; at four loops it is 6.7 and 5.8 points above, respectively.These comparisons are on a ten-task average.Source: §1 · p. 3
  • Introducing mixed depths in the final quarter costs an estimated 16% more training FLOPs than the baseline, and zero-initialized conditioning makes this schedule outperform mixed-depth training from scratch.The conditioning heads start as the one-loop checkpoint at the switch.Source: §1 · pp. 2–3
  • Replaying first-loop expert routes on later loops cuts mixed-depth MoE update time by 11.7% with no measurable accuracy change.Source: §1 · p. 3
  • Post-trained Koi-4B spans scores of 18 to 29 on the Artificial Analysis Intelligence Index by changing loop count alone.Source: §1 · p. 3

2. Latent Recurrent Models

Latent recurrent models share a block of layers across repeated applications, separating effective depth from the number of unique layers. Elastic depth uses loop count as an inference-time compute control.

2. Latent Recurrent Models

  • The approach builds on depth-wise weight sharing and adaptive computation, and recent language models use it to separate effective depth from unique parameter count.The passage cites Dehghani et al. (2019), Graves (2016), and Banino et al. (2021) for earlier work, and Geiping et al. (2025), Zhu et al. (2025), and Bae et al. (2025b) for recent language models.Source: §2 · p. 3
  • A latent recurrent model composes input, recurrent, and output blocks; the recurrent block's shared parameters refine a state across repeated applications before the output block maps it to logits.Source: §2 · p. 3
  • Each additional loop traverses the recurrent layers again, adding inference compute and latency at fixed parameter count.Source: §2 · p. 3
  • Elastic depth means training over a range of loop counts and serving at any count within that range, making loop count an inference-time compute knob.Source: §2 · p. 4

3. Method

ELR trains one checkpoint to serve at multiple loop counts without sacrificing the one-loop endpoint. Its recipe combines a gated recurrent update, within-update depth mixing, loop conditioning, and cost-saving inference and training choices.

3.1. Elastic Latent Recurrence

  • At k=1, ELR executes each unique layer once; later passes update the recurrent state by a learned, token- and channel-specific gated step.The recurrent block is shared across loops, and its gate sets the step size along the proposed state revision.Source: §3.1 · p. 4

3.2. Mixed-Depth Training with One-Loop Rank Overrides

  • Each microbatch samples a loop count, while rotating rank overrides force k=1 examples so one- and multi-loop gradients can share an optimizer update.The override keeps the marginal probability of k=1 approximately one third.Source: §3.2 · pp. 4–5

3.3. Loop-Conditioned Modulation

  • Each pass is conditioned on loop progress and total budget, and modulation of recurrent and output layers improves accuracy at every additional loop.Progress is i/k and budget is k/K, where K is the maximum training loop count.Source: §3.3 · p. 5
  • Zero-initialized conditioning heads initially preserve the one-loop checkpoint’s function at every loop count.The modulation is initially the identity, and later loops leave the state unchanged through zero exit factors.Source: §3.3 · p. 5

3.4. Reducing the Cost of Recurrence

  • Training at one loop for most tokens and continuing with mixed depths limits recurrence overhead while carrying over the architecture, optimizer state, and learning-rate schedule.The two phases differ only in the distribution over loop counts.Source: §3.4 · pp. 5–6
  • Cross-loop route replay reuses first-loop expert assignments and weights on later loops, while the selected experts still process updated representations.The method assumes first-loop routes remain useful at later loops.Source: §3.4 · p. 6
  • Last-loop KV caching overwrites one attention cache across loops, keeping attention KV memory at the one-loop level.Earlier-loop queries attend to earlier tokens’ final-loop keys and values and the current token’s keys and values for that loop.Source: §3.4 · p. 6

4. One-Loop Accuracy and Gains from Additional Loops

Across matched 4B dense and 30B-A3B MoE models, ELR retains near-baseline accuracy at one loop and improves ten-task averages as loops increase. The section also measures inference cost and tests cross-loop route replay in the MoE model.

4.1. Experimental Setup

  • Evaluation compares matched dense 4B and MoE 30B-A3B baselines on ten-task averages, testing each final checkpoint at every trained loop count.The models use delayed mixed-depth training after one-loop training for the first 75% of tokens.Source: §4 · p. 6, §4.1 · pp. 6–7

4.2. Main Results

  • At k=1, ELR is 0.52 points below baseline on 4B and 1.04 above on 30B-A3B; k=4 adds 7.25 and 4.75 points, respectively.The averages cover ten tasks; Figure 2 plots average accuracy against loop count.Source: §4.2 · p. 7, Figure 2 · p. 7
  • At 4B, two and three loops capture 79% and 96% of the four-loop gain; at 30B-A3B, they capture 87% and 91%.The reported pattern is diminishing returns as loop count increases.Source: §4.2 · pp. 7–8
  • At four loops, GDP-Hybrid-4B reaches 2.33× baseline TTFT and 2.06× decode latency; three loops capture 96% of its accuracy gain at 1.87× baseline TTFT.At one loop, ELR matches its non-recurrent baseline’s TTFT and per-token decode latency.Source: §4.2 · p. 8

4.3. Cross-Loop Route Replay

  • First-loop routes retain much, but not all, token-specific routing information; replay cuts mixed-depth update time by 11.7% while leaving the k=4 ten-task average unchanged within noise.Replay reduces router evaluations but retains expert dispatch and computation.Source: §4.3 · p. 8, Table 1 · p. 8

5. Koi-4B: A Post-Trained GDP-Hybrid ELR Model

Koi-4B’s benchmark score rises with inference depth, while last-loop KV caching keeps attention memory at its one-loop level with a small, benchmark-dependent accuracy cost. Its evaluation uses the authors’ harness and has no matched non-recurrent baseline.

5.1. Experimental Setup

  • Koi-4B uses a 7T-token pretraining run, a 50B-token mid-training stage, and supervised fine-tuning followed by reinforcement learning from verifiable rewards.Mixed-depth training is used after the switch from one-loop training; loop counts are sampled with a one-loop rank override in pretraining, mid-training, and supervised fine-tuning.Source: §5.1 · p. 9
  • Scores come from the authors’ harness and may differ from published figures; Koi-4B has no matched non-recurrent baseline and is compared with open models of similar size and larger.The evaluation reports Artificial Analysis Intelligence Index v4.1 and additional benchmarks outside that index.Source: §5.1 · p. 9

5.3. Last Loop KV-Cache Inference

  • 0–3% score reductions at k=4 with last-loop KV caching preserve most of the gain; AA-Omniscience non-hallucination improves.The cache keeps attention KV memory at its one-loop size at every loop count, without retraining.Source: Table 2 · p. 10, §5.3 · pp. 10–11

5.2. Inference-Time Depth Scaling

  • 18 at k=1 rises to 29 at k=4 on the Artificial Analysis Intelligence Index; at one loop Koi-4B 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 scores 3 points above Gemma4-26B-A4B.Source: §5.2 · p. 10
  • 21.8%, 29.5%, and 32.7% fewer output tokens per rollout accompany higher scores at k=2, 3, and 4, respectively.At two loops, decode time per rollout rises by about 6%; at four loops it rises by about 1.4×.Source: §5.2 · p. 10

6. Ablations

On GPT-0.5B, the ablations examine validation loss under alternative state updates, depth-training schedules, loop-sampling schemes, rank overrides, and loop conditioning. Together, they identify choices that preserve one-loop performance while supporting lower loss at deeper inference depths.

6.1. Adaptive Delta versus Alternative State Updates

  • Adaptive delta alone matches the baseline at k=1 and has the lowest loss at every loop; its input-dependent step size beats a fixed per-channel rate by 0.006–0.007.Parameter-free replace and residual updates perform worst; replace has the highest one-loop loss and rises from four to five loops, while residual stops improving after three.Source: §6.1 · p. 11

6.3. Synchronized versus Independent-Rank Loop Sampling

  • 1.6774 one-loop loss under independent-rank sampling is below 1.6830 with 15% overrides and 1.6871 with synchronized weighted sampling, but costs about 36% more step time.One-loop overrides recover part of the gain at synchronized step time, and the same loss ordering holds at every k≥2.Source: Table 3 · p. 12, §6.3 · p. 12

6.2. Fixed-Depth versus Mixed-Depth Training

  • A mixed-depth model improves monotonically from one to five loops, trading some per-depth quality for one checkpoint usable at all five depths instead of five specialists.Each fixed-depth model has its lowest loss at its training depth and degrades at other depths; the one-loop-only model reaches 1.660 versus the baseline’s 1.685.Source: §6.2 · p. 12

6.4. One-Loop Rank Overrides

  • f=0.15 (p1=0.32) is the smallest tested override fraction below the baseline; increasing f lowers one-loop loss but costs deeper-loop performance.The benefit is not explained by the higher one-loop marginal alone: synchronized weighted sampling at the same p1 remains above baseline.Source: §6.4 · pp. 12–13, Figure 5 · p. 13

6.5. Delayed Mixed-Depth Training

  • 1.20× baseline training compute at a 75% switch retains 84% of the five-loop improvement, with one-loop loss below baseline at every tested start point.Mixed-depth training from scratch costs 1.79× baseline compute; delaying it to 25%, 50%, or 75% of training lowers compute to 1.60×, 1.40×, and 1.20×, respectively.Source: Figure 5 · p. 13, §6.5 · p. 13

6.6. Loop-Conditioned Modulation

  • 1.6147 versus 1.6273 at five loops: zero-initialized loop conditioning on a delayed model yields lower validation-loss point estimates at every loop than mixed-depth training from scratch.One-loop loss changes from 1.6748 to 1.6739, while loss falls by 0.009 at k=2 and 0.022 at k=5; conditioning adds 31.6M parameters.Source: §6.6 · p. 13, Figure 6 · p. 14

7. Related Work

The related work covers elastic recurrent depth, methods for regulating recurrent dynamics, late acquisition of recurrence, and recurrent MoE serving. ELR is distinguished by its one-loop parity goal, within-update depth mixing, and delayed mixed-depth training without architectural change or optimizer restart.

Elastic recurrent depth

  • ELR uses zero-initialized conditioning grafted onto a trained one-loop checkpoint and mixes one-loop and deeper gradients within an update, unlike the cited depth-conditioning and shared-depth-sampling approaches.The paper does not evaluate extrapolation; it contrasts its approach with LoopFormer, Huginn, LoopMoE, and stochastic RINS.Source: §7 · p. 14

Regulating recurrent dynamics

  • ELR learns its adaptive-delta step size end to end without explicit stability constraints, while aiming for baseline parity at one loop and improvement at added loops.The cited alternatives regulate recurrent dynamics with hyper-connections, stability conditions, Jacobian regularization, or damped substeps.Source: §7 · p. 14

Acquiring recurrence late

  • Unlike retrofitting a pretrained non-recurrent model, ELR continues a model recurrent by construction and changes only the loop-count distribution, with no architectural change or optimizer restart.The comparison is with Retrofitted Recurrence and Relaxed Recursive Transformers.Source: §7 · p. 14

Recurrent MoE and serving

  • ELR replays its own first-loop MoE routes on later loops and serves each request at a fixed loop count, unlike precomputed routing or batching requests at different depths.The cited work includes Read-ME and Continuous Depth Batching.Source: §7 · p. 15

8. Conclusion

The conclusion presents ELR as a single checkpoint spanning inference depths without a one-loop tax, and identifies open questions about depth gains, conditioning, adaptive serving, and MoE post-training.

8. Conclusion

  • At one loop, ELR is 0.5 points below its baseline at 4B and 1.0 point above at 30B-A3B; four loops add 7.25 and 4.75 points, respectively.The conclusion also reports Koi-4B at 18 on the Artificial Analysis Intelligence Index at one loop and 29 at four loops.Source: §8 · p. 15
  • Gains diminish after the second loop, intermediate-depth conditioning does not close the gap to fixed-depth specialists at the endpoints, and adaptive-depth serving and MoE post-training remain open.The experiments fix loop count per request; budget conditioning complicates adaptive per-request or per-token depth.Source: §8 · p. 15

A. Model Architecture and Training Details

  • Without loop-conditioned modulation, ELR’s adaptive-delta gate adds about 0.2% of baseline parameters for GPT-0.5B and GDP-Hybrid-4B; Koi-4B uses modulation and lacks a matched non-recurrent baseline.Appendix A.4 describes Koi-4B’s architecture and training stages.Source: A. Model Architecture and Training… · p. 18

A.1. Architectures

The section specifies the hybrid and dense model architectures, matched optimization settings, and mixed-depth training schedules, then details Koi-4B’s architecture and training stages.

A.1. Architectures

  • GDP-Hybrid models use a 14/24/14 input/recurrent/output partition, while GPT-0.5B uses 8/14/8; the MoE variant balances expert load with router bias and an auxiliary loss.GPT-0.5B alternates self-attention and SwiGLU MLP layers; the hybrid layout replaces Mamba-2 with GDP layers. All models use RMSNorm, bias-free linear layers, untied embeddings, BF16 training, and a 131,072-token vocabulary.Source: §A.1 · p. 18, Table 4 · p. 18

A.2. Optimization

  • Matched baselines share data, token budgets, and optimization settings with their ELR counterparts; the 30B-A3B run uses Muon for matrices and AdamW for other parameters.The learning rate follows warmup, a stable phase, and decay over the final 15% of training tokens.Source: Table 5 · p. 19, §A.2 · p. 19

A.3. Mixed-Depth Training

  • Mixed-depth updates assign a majority loop count per microbatch across data-parallel ranks, and delayed training resumes the one-loop checkpoint with its optimizer state and schedule.Gradients propagate through all executed loops without truncation.Source: §A.3 · p. 19

A.4. Koi-4B Architecture and Training

  • Koi-4B adds zero-initialized, loop-conditioned modulation to the GDP-Hybrid-4B backbone and switches from one-loop to mixed-depth training halfway through pretraining.It is pretrained for 7T tokens, then receives a 50B-token context-extension stage and four SFT epochs followed by 200 RLVR steps; mixed-depth settings are used in post-training.Source: §A.4 · pp. 19–20, Table 6 · p. 20

B.1. Evaluation Protocol

The section describes the benchmark protocols and reports where detailed scores and output-token accounting are presented. Its reported aggregate token counts decrease at higher loop counts, subject to a compute caveat.

B.1. Evaluation Protocol

  • Benchmark protocols use task-specific decoding and scoring, and the ELR checkpoints and non-recurrent baseline are evaluated under the same protocol.The benchmark list, few-shot counts, and metrics are reported in Table 7; GSM8K and MATH-500 use mathematical-equivalence scoring.Source: §B.1 · pp. 20–21, Table 7 · p. 21

B.2. GDP-Hybrid-4B / B.3. GDP-Hybrid-30B-A3B

  • Tables 8 and 9 report the per-benchmark scores underlying the GDP-Hybrid-4B and GDP-Hybrid-30B-A3B comparisons.The latter table includes dynamic routing and cross-loop route replay; both ELR models begin mixed-depth training from the same one-loop checkpoint.Source: §B.2 · p. 21, §B.3 · p. 21, Table 8 · p. 21

C. Koi-4B Evaluation and Detailed Results

  • Table 10 reports benchmark scores for one post-trained Koi-4B checkpoint evaluated at loop counts one through four.CritPt is omitted because models of this size score near zero on it and it contributes zero to the index at every loop count; Figure 3 rounds index estimates to the nearest integer.Source: C. Koi-4B Evaluation and Detailed… · p. 22, Table 10 · p. 22

C.1. Output Token Accounting and Aggregation

  • Output-token counts include capped completions and exclude judges and the τ3 user simulator; fewer output tokens alone do not establish lower inference compute or latency.GDPval token averages include timed-out tasks even though those tasks receive no Elo judgment, so the token and score averages cover different task sets.Source: §C.1 · p. 23
  • 35.1%: total output tokens across the evaluated benchmarks fall from one to four loops; excluding GDPval, the geometric-mean ratio at four loops is 0.7006.The aggregation gives equal weight to the eight evaluated AA-Index benchmarks, regardless of output length, and counts AA-Omniscience once because its two index components share generations.Source: §C.1 · p. 23

C.2. Last-Loop KV-Cache Results

The section reports last-loop KV-cache comparisons and timing measurements, while noting uncertainty in single-run score comparisons and workload-dependent latency results.

C.2. Last-Loop KV-Cache Results

  • Seven of 27 per-question score comparisons reach |z| ≥ 2, but the statistics capture variation across questions within a single run, not across runs.At this uncorrected threshold, about one such result is expected by chance; |z| < 2 indicates no detected difference, not equivalence.Source: §C.2 · pp. 23–24
  • Last-loop KV produces fewer output tokens on all ten benchmarks at two loops and all but MATH-500 at three and four loops, largely by reducing output-cap truncations.At four loops, the Terminal-Bench truncation rate falls from 0.6% to 0.2%; capped-request token savings account for 96–98% of its reduction despite 3–7% more mean agent turns.Source: §C.2 · p. 24

D.1. Measurement Setup

  • The primary timing setup uses one H100, batch size one, 1,024 prompt tokens, and 128 generated tokens; reported intervals reflect timing variability only.Times are medians over seven repetitions after two warm-up batches, normalized to the non-recurrent baseline.Source: §D.1 · p. 24

D.2. Sensitivity to Workload

  • Across tested prompt lengths and batch sizes, one-loop ELR latency stays near baseline; Figure 7 plots ten-task accuracy against TTFT and per-token decode latency.Decode-latency differences range from −0.26% to +2.16%, and TTFT differences from −1.54% to +1.53%, depending on execution mode.Source: §D.2 · p. 25, Figure 7 · p. 25

E.1. Evaluation Protocol

The section evaluates route replay and input re-injection, including validation sampling, routing comparisons, and timing measurements. Replay reduces mixed-depth update time, while input re-injection preserves the reported ablation ordering.

E.1. Evaluation Protocol

  • The route-replay evaluation compares dynamic routing, inference-only replay, and a separately trained replay checkpoint on 256 validation documents.Each document contributes one 2,048-token window; results on 64 separate documents with 8,192-token windows follow the same pattern.Source: §E.1 · p. 25

E.2. Routing and Output Similarity

  • Aggregate expert-use counts across loops one and four are highly similar, but per-token expert retention is only 53%; replay outputs remain closer to fresh routing than to a shifted-route control.The layer-output comparison holds each layer’s fourth-loop input fixed and does not propagate comparison outputs to later layers.Source: §E.2 · p. 25, Table 13 · p. 26

E.3. Timing Measurements

  • 11.0–12.7%: cross-loop route replay reduces mixed-depth training update time across the measured training intervals.At four loops, inference forward time falls by 1.21–2.96% across sequence lengths of 128 to 8,192 tokens, with the smallest reduction at 8,192 tokens.Source: §E.3 · pp. 25–26

F. Ablation Results / F.1. Input Re-injection

  • With input re-injection, the fixed-versus-mixed-depth, loop-sampling, one-loop-override, and delayed-training ablations retain the same ordering as before.The paper expects the state-update comparison to transfer as well; Table 14 reports validation loss.Source: §F.1 · p. 26, Table 14 · p. 26
Loading 2610.12582v1…