Looped language models stop paying the one-loop tax: one checkpoint matches its non-recurrent baseline at one loop and adds 7.25 points at four
Synopsis
Looped language models used to trail a non-recurrent model of the same size at one loop; a single checkpoint now matches a non-recurrent baseline trained on the same data at one loop and gains 7.25 points on a ten-task average at four loops.
Interpretation
A looped language model can keep one checkpoint and trade inference compute for accuracy by changing the loop count, without falling behind a non-recurrent model of the same size at one loop. Earlier elastic-depth models fell well short of a non-recurrent model of equal depth at one loop: Huginn scored 34.9 against 46.4 on ARC-Easy, and Parcae sat below its Transformer baseline at one loop and above it at eight, a gap the authors call the one-loop tax. At two scales, a 4B dense hybrid and a 30B-A3B mixture-of-experts hybrid, ELR was compared with non-recurrent baselines matched in parameters, data, and training horizon: ten-task averages were 65.13 against 65.65 and 66.08 against 65.04 at one loop, with four loops adding 7.25 and 4.75 points.
What keeps the one-loop endpoint alive is how the recurrent state is updated and where loop counts are mixed inside an optimizer update, not a narrower loop range. Sampling one loop count per optimizer step and varying it across steps fails to preserve the one-loop endpoint even when one loop is drawn a third of the time, and uniform sampling over one to five loops still leaves the one-loop endpoint behind the baseline. In 0.5B ablations, synchronized weighted sampling gave a one-loop loss of 1.6871, a 15% one-loop rank override lowered it to 1.6830, independent-rank sampling reached 1.6774, and the non-recurrent baseline was 1.6852.
Elastic depth can be acquired late in pretraining, and conditioning heads initialized to zero at the switch make that late schedule beat mixed-depth training from scratch. Mixed-depth training from scratch costs an estimated 1.62 times the baseline's training compute, while switching only in the final quarter costs about 1.16 times; in the 0.5B ablation, the 75% switch lowers the estimated added training compute from about 79% to about 20%. On the 0.5B model, switching after 25%, 50%, and 75% of tokens retained 97%, 90%, and 84% of the five-loop improvement over the baseline; at the 75% switch the five-loop loss rose by 0.009, and with conditioning the five-loop validation loss was 1.6147 against 1.6273 for mixed-depth training from scratch.
A recurrent mixture-of-experts model can reuse first-loop routing on later loops and keep only the last loop's attention KV cache, cutting training time and attention memory. Re-executing each router on every loop repeats work, and standard inference keeps a separate attention KV cache per loop, so memory grows with depth. Cross-loop route replay cut mixed-depth update time by 11.7% while raising four-loop validation loss by 0.0038; on Koi-4B, last-loop KV caching held attention KV memory at 3.0 GiB against 7.5 GiB per 128K-token request, with four-loop scores 0–3% below per-loop caching.
Perspective
The recipe is validated on a 4B dense hybrid and a 30B-A3B mixture-of-experts hybrid, and Koi-4B carries it through pretraining, a long-context mid-training stage, supervised fine-tuning, and reinforcement learning from verifiable rewards, all scored in the authors' own harness, so figures may differ from published ones. Koi-4B has no matched non-recurrent baseline, so its one-loop accuracy is compared with open models of similar size. The authors propose combining ELR with per-request or per-token adaptive depth and post-training the mixture-of-experts variant as next steps.
Gains diminish after the second loop, which the authors say is consistent with later loops appearing in fewer training examples, but they have not tested whether reweighting the loop-count distribution recovers them. Conditioning closes most of the gap to fixed-depth specialists at intermediate depths but not at the endpoints, and the authors have not isolated why. The serving setting fixes the loop count per request, and per-request or per-token adaptive depth would complicate budget conditioning. The last-loop KV statistics capture variation across questions within a single run, not across runs, and GDPval is excluded from that comparison because timeout failures leave the two cache policies covering different task subsets.
