Towards Looped Models Done Right, Part II: Rethinking at Fixed Points

Benhao Huang; Chufan Shi; Junlin Chen; Shicheng Wen; Zhengzhong Liu; Eric Xing; Xuezhe Ma

cs.LGOct 5, 2026

Towards Looped Models Done Right, Part II: Rethinking at Fixed Points

Authors: Benhao Huang, Chufan Shi, Junlin Chen, Shicheng Wen, Zhengzhong Liu, Eric Xing, Xuezhe Ma

Organizations: Institute of Foundation Models · CMU · USC

Abstract

Every recurrence of a looped language model adds cost in training, decoding, prefill, and reinforcement learning (RL). The closer recurrent states get to fixed points, the less the path to them matters. This enables truncated backpropagation in training; terminal key-value (KV) sharing for decoding with almost no loss in accuracy; a distilled student that prefills up to 1.79x faster; and RL updates that compute gradients from saved rollout states, 2x faster than backpropagating through the replayed trajectory. We therefore improve the two components of training that shape these fixed points: the depth prior and input injection. Fixed-depth training breaks KV sharing, and Huginn's broad depth prior supports sharing but dilutes supervision at the target depth more than sharing requires; we learn the prior from prediction feedback, with an entropy term that keeps it broad. Existing injection schemes let the state's component along the input amplify or cancel the injection; we remove this component with orthogonal injection. From 100M to 1.6B parameters, the learned prior and orthogonal injection lower perplexity at every scale relative to Huginn's prior and existing injection schemes, respectively. At 1.6B, the learned prior with a 3x smaller KV cache matches the downstream average of fixed-depth training with the full cache.

Figures & tables

Figure 2 : Heterogeneous token convergence in Huginn. (a) Relative hidden-state changes $\rho_{r,t}$ , with convergence threshold $\tau=2\%$ . (b) Later tokens can approach approximate fixed points before earlier ones. (c) Position versus convergence depth; the purple line shows medians over $128$ -position intervals. Details in Appendix .

Figure 3 : The endpoint stands in for the recurrent trajectory. Recurrence runs upward; ellipses mark omitted positions and depths. (a) In training, each update reads prefix states from the previous depth; darker states are closer to convergence (as in Figure 2 ). Prefix tokens converge at different depths (Section 2 ), so the sampled depth

R

decides whether reads inside the backpropagation window see converged context; gradients stop at sg. (b) Decoding reads only the terminal prefix KV banks

C^{\star}_{<t}

and discards the others, with nearly unchanged downstream accuracy. (c) During prefill, a distilled student replaces the recurrence and supplies approximate terminal banks

\widehat{C}_{<t}

Figure 4 : The learned prior narrows the trade-off between fixed-depth training and PLN sampling. Native WikiText-103 test PPL across test-time depths at matched training budgets. Fixed-depth training is best at training-matched depth

R=5

but degrades sharply at others. PLN sampling performs well across all depths but falls behind at

R=5

. The learned prior lies between them, narrowing the gap at

R=5

while staying stable across depths.

Panel of Table 2 — Table 2: The learned prior far surpasses Untied 4 at the same parameter count and KV cache, and approaches Untied 12 with $3\times$ fewer of both. The looped models ( $1{+}2{\times}R{+}1$ ) are evaluated at $R=5$ , a logical depth of $12$ blocks, with terminal KV sharing, which caches $4$ KV banks, one per physical block, instead of $12$ . Learned priors start from the Fixed PLN-5 prior. FLOPs are in units of $10^{19}$ . Generation: GSM8K accuracy, DROP F1, and MBPP+ pass@1. Untied 4 and Untied 12 stack four and twelve distinct blocks and decode with their own full KV caches; Untied 4 matches the looped models’ parameters and cache size, and Untied 12 their logical depth, with $3\times$ their non-embedding parameters and KV cache. Fixed-depth training breaks KV sharing: it loses over half its GSM8K accuracy at M and L relative to the full cache (Appendix Table 8 ). See Appendix for evaluation details.

Panel of Table 3 — Table 3: OrthoInj has the highest downstream average and the lowest validation PPL at every scale. All models train with PLN-5 and are evaluated at $R=5$ with terminal KV sharing, as in Table 2 . Recipes are in Appendix B.4 .

Figure 5 : OrthoInj ends training with the lowest loss at every scale. Training CE of the Table 3 runs over the final

20\%

of training, as a centered rolling median over

5\%

of each run; the gray band marks the learning-rate decay (final

10\%

of updates). Within each scale, all runs see the same sampled depths.

Panel of Table 4 — Table 4: Reusing rollout states cuts scoring and backward time. The L learned-prior model of Table 2 is fine-tuned with Dr. GRPO at $R=6$ on GSM8K or MBPP+; the GSM8K runs are also evaluated on MATH500, which RL does not train on. Pass@ $k$ uses eight samples per question. Time: median scoring and backward time per update, excluding rollout and the optimizer step. See Appendix for training and evaluation details.

Panel of Table 5 — Table 5: Distilled prefill trades downstream accuracy for lower prefill latency. Each path prefills the prompt, then the learned-prior teacher decodes at $R=5$ from the path’s four terminal KV banks; Untied 4 prefills and decodes on its own. The distilled student is initialized from the teacher and trained with the hidden and KL objective of equation ( 4 ) on one quarter of its teacher’s pretraining tokens; the teacher rows prefill at $R=5$ , $3$ , $2$ , or $1$ . Latency: mean $\pm$ s.d. of six $8$ K prefills at batch size $8$ . See Appendix for evaluation details.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6 : Freezing a converged prefix reproduces joint iteration (Lemma 1 ). Relative gap between the current-token endpoint of joint iteration and of iteration with the prefix frozen after convergence or after five recurrences. Dots are prefixes that pass the convergence checks (

206

304

for Fixed

R=5

191

304

for PLN-5); markers and bars give medians and interquartile ranges. The dotted line marks the

10^{-4}

agreement tolerance.

A: Architecture and attention geometry
	S	M	L
Blocks (prelude/core/coda)	1 / 2 / 1
Width	1,536	3,072	6,144
FFN	4,352	8,704	17,408
Q heads	24	48	96
KV heads	6	12	24

Appendix

Table 6: Model configurations at three scales. Blocks count prelude, recurrent-core, and coda blocks of the looped models; Untied 4 and Untied 12 stack

4

and

12

blocks. Width and FFN give hidden and feedforward dimensions, and Q and KV give query and key–value head counts. Non-emb. and Total give Untied 4 parameter counts without and with embeddings; the looped rows add Huginn’s width

\times

width input adapter. The main text rounds the scales to

100

400

M, and

1.6

B. Parameter and token counts are in billions; the peak learning rate is tuned per run (Appendix B.1 ).

Injection recipe	Injected input $\tilde{e}$	Input to recurrent core
DEQ-QKV	$e$	$W_{\rm QKV}^{(j)}h^{r,(j)}+W_{e}^{(j)}\tilde{e}$
Huginn-Linear [ 2 ]	$e$	$W_{\rm lin}[h^{r};\tilde{e};1]$
Parcae-Decay [ 15 ]	$\operatorname{RMSNorm}(e)$	$\Lambda h^{r}+\Delta\odot W\tilde{e}$
Parcae w/o norm	$e$	$\Lambda h^{r}+\Delta\odot W\tilde{e}$
OrthoInj (ours)	$e$	$Q_{q}\Lambda h^{r}+q$

Appendix

Table 7: Input-injection recipes.

e=P_{\theta}(x)

is the prelude output,

\tilde{e}

is the injected input, and RMSNorm uses learned gains. DEQ-QKV injects at internal sites (superscript

(j)

indexes a recurrent layer); the others inject at the core entrance.

\Lambda=\operatorname{diag}(\exp(-\Delta\odot\alpha))

with

\Delta=\operatorname{softplus}(\Delta_{\rm raw})

and

\alpha=\exp(\alpha_{\rm raw})

, and OrthoInj applies

Q_{q}

of equation ( 8 ) with

q=\Delta\odot W\tilde{e}

and

\epsilon=10^{-6}

Panel of Table 8 — Table 8: Table 2 without KV reuse. Same checkpoints as Table 2 , evaluated at $R=5$ . Val. and Wiki are native PPL over full sequences, choice tasks are scored without a cache, and generation keeps all twelve visit-specific KV banks. Table 2 instead scores Val. and Wiki only on the tokens after a $4{,}096$ -token prompt that continue a document begun inside it (Appendix ), so only the task columns compare directly across the two tables. Untied 12 stacks twelve distinct blocks ( $0.31$ B, $1.25$ B and $4.98$ B non-embedding parameters at S, M and L, three times the looped models) and has twelve KV banks; it is scored the same way and not bolded.

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models
Sep 30, 2026Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin +5Recurrent Model Loop
RND NLP, DAIMLD, Russian Federation
Retrofitting Recurrent Depth into a Pretrained Language Model: Installation, Extrapolation, Transfer, and Retention at Two Parameter Budgets
Jul 31, 2026Mark ShapiroLarge Language Model Pretraining Recurrent Model
RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory
Sep 3, 2026Yuxiang Wang, Kunyu Feng, Yingda Shen +3Efficient Latent Reasoning Recursive Models
The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2

Sep 30, 2026, cs.CLJ/K move · Enter open · S save

Closing the Loop: Practical Training Recipes for Looped Language Models

Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin, Inessa Fedorova+4

RND NLP, DAIMLD, Russian Federation

Recurrent Model Loop Multi-Stage Training Dense Baseline Recipe

Open paper AskPDF arXiv

Looped language models increase effective depth by repeatedly applying a shared block of layers, but existing large-scale recipes require multi-stage training over trillions of tokens, while the benefits of recurrence remain difficult to separate from differences in data and training. In this work, we establish practical training recipes for looped language models, with three main results. (1) We develop a compute-efficient from-scratch pipeline that reduces the training budget from 7.7T tokens in Ouro to 310B tokens while retaining strong reasoning performance. Pretraining followed by high-quality mid-training, together with learning-rate warmup and stronger exit-gate regularization, enables stable recurrent training without prior multi-stage schedules. (2) Under controlled comparisons, our 1.4B LoopLM outperforms a parameter-matched dense model trained on the same data and token budget on all 12 evaluated benchmarks, including +14 points on GSM8K, +10 on MATH, and +22 on DROP. At matched inference compute, it approaches a 3.9B dense model on mathematical reasoning and reading comprehension while using only 36% as many parameters. (3) We introduce a minimal recipe for converting pretrained dense models into looped ones: a single learned input-mixing scalar and a smoothed exit loss, with no step-specific parameters. Applied to Qwen3-1.7B-Base, Looped Qwen improves over an identically continued dense baseline on every evaluated benchmark across two data regimes, with statistically clear gains on GSM8K, MATH, and MMLU-Pro on the curated mixture. Together, these results make looped language models substantially cheaper to train from scratch and practical to introduce into existing pretrained checkpoints, while isolating the gains due to recurrence itself.

Abstract

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models

Retrofitting Recurrent Depth into a Pretrained Language Model: Installation, Extrapolation, Transfer, and Retention at Two Parameter Budgets

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory