Enabling Dynamic Computation in Looped LMs

Aayush Mishra; Arnau Padrés Masdemont; Victor Conchello Vendrell; Jordi Ros Giralt; Arash Behboodi; Fabio Valerio Massoli

cs.AIOct 6, 2026

Enabling Dynamic Computation in Looped LMs

Authors: Aayush Mishra, Arnau Padrés Masdemont, Victor Conchello Vendrell, Jordi Ros Giralt, Arash Behboodi, Fabio Valerio Massoli

Organizations: Qualcomm AI Research · Johns Hopkins University

Abstract

Looped LMs are parameter efficient and promise dynamic computation (saving memory and FLOPs on easy tokens). However, state-of-the-art open Looped LMs trained with this dynamic computation capability (Ouro models) do not realize it in practice as each loop iteration (depth) requires its own level of KV-cache, necessitating all loop computations. Moreover, Ouro's early-exit prior is enforced on each token equally, which results in static lower-depth like processing of all tokens regardless of difficulty. In this work, we propose a simple "best-available" KV caching strategy that works out-of-the-box, creating a new frontier in the performance vs depth space. Our approach enables up to 30% reduction in FLOPs and KV memory while retaining full-depth performance, showing the true flexibility of Looped LMs. Furthermore, training looped LMs with awareness about this KV caching strategy improves performance and efficiency. Finally, we apply a small but effective fix to the early-exit prior enforcement objective that makes tokens exit at truly heterogeneous depths based on effort. Our findings are validated on Ouro models as well as smaller looped LMs pre-trained from scratch.

Figures & tables

Figure 1: Ouro forces full-depth computation even after a token exits, as future tokens may need it. This makes early-exit practically useless. Ragged KV enables true early-exit by aliasing deeper blocks to the best-available cache. It avoids forced computation and saves memory+FLOPs during inference. Surprisingly, it works out-of-the-box for Ouro models ! On the right, we compare fixed-depth (1-4 loops) with Ragged KV at different exit thresholds on BigBenchHard (details in § E.1 ). Ragged KV unlocks significantly higher performance at low average decode depths.

Log-likelihood evaluation
Dataset	Metric	Fixed-Depth			Dense-Ragged		ETE-Ragged
	depth / thr.	2	3	4	.3	.5	.3	.5
	Avg. depth	2.00	3.00	4.00	2.39	3.05	2.63	3.21
MMLU	accuracy	60.4 $\pm$ 0.4	66.7 $\pm$ 0.4	67.5 $\pm$ 0.4	–	–	61.9 $\pm$ 0.4	66.2 $\pm$ 0.4
ARC-C	accuracy	54.9 $\pm$ 1.5	59.9 $\pm$ 1.4	59.8 $\pm$ 1.4	–	–	59.0 $\pm$ 1.4	60.0 $\pm$ 1.4
HellaSwag	accuracy	70.8 $\pm$ 0.5	73.1 $\pm$ 0.4	73.4 $\pm$ 0.4	–	–	72.8 $\pm$ 0.4	73.2 $\pm$ 0.4

Table 1: Performance of Ouro-1.4B in different Fixed-depth and Ragged inference modes. Values are percentages. Avg. depth denotes avg. decode depth for Dense-Ragged, and avg. overall depth for ETE-Ragged. Generation results report Qwen3.8-27B judged pass@

k

. Standard error across evaluation samples. Macro averages are unweighted and shown without an uncertainty estimate. Dense-Ragged columns for Log-likelihood evaluations are missing as it is equivalent to fixed@4.

Dataset	Metric	Fixed-Depth			Dense-Ragged			ETE-Ragged
	depth / thr.	2	3	4	.3	.45	.5	.3	.45	.5
	Avg. depth	2.00	3.00	4.00	2.24	2.45	2.83	2.30	2.60	3.00
AIME24	pass@1	40.4 $\pm$ 6.7	62.1 $\pm$ 7.0	65.8 $\pm$ 6.8	56.7 $\pm$ 7.3	57.9 $\pm$ 7.3	61.7 $\pm$ 7.2	55.0 $\pm$ 7.1	59.6 $\pm$ 7.2	61.7 $\pm$ 6.9
	pass@4	64.4 $\pm$ 7.7	79.9 $\pm$ 6.9	82.0 $\pm$ 6.9	76.0 $\pm$ 7.3	75.2 $\pm$ 7.4	78.7 $\pm$ 6.8	77.3 $\pm$ 6.4	76.4 $\pm$ 7.3	80.4 $\pm$ 6.9
AIME25	pass@1	31.7 $\pm$ 6.7	49.2 $\pm$ 7.7	47.9 $\pm$ 7.9	45.4 $\pm$ 7.2	45.4 $\pm$ 7.4	52.9 $\pm$ 7.2	46.2 $\pm$ 7.2	49.2 $\pm$ 7.3	52.1 $\pm$ 7.5
	pass@4	49.0 $\pm$ 8.5	65.6 $\pm$ 8.2	62.1 $\pm$ 8.5	67.7 $\pm$ 7.6	64.8 $\pm$ 7.9	70.8 $\pm$ 8.1	66.9 $\pm$ 8.0	68.7 $\pm$ 7.9	69.2 $\pm$ 8.0

Table 2: Performance of Ouro-2.6B Thinking. AIME24/25/26, GPQA-Diamond and OlympiadBench use the Qwen3.8-27B judge, and HumanEval+ uses executable EvalPlus tests. Values are pass@

k

percentages with standard error across evaluation samples. Macro averages include all six benchmarks and are unweighted.

Figure 2: (First row) Exit depth distribution of pre-trained Ouro models: predominantly depth 2/3 depending on the threshold

\tau

, illustrating the lack of true dynamic exits due to local prior enforcement. Noticeably, depth 1 exits are non-existent. (Second row) Jump rate fractions. A high majority of tokens do not jump around and exit at similar depths as their neighbors. Even when they do, they typically only jump one depth up or down. Large jumps (

\lvert e_{i}-e_{i-1}\rvert\geq 2

) rarely happen. These numbers undermine the claims of dynamic computation based on difficulty in Ouro models.

Figure 3: To enable early-exit, Ouro models are trained with a KL divergence term (w.r.t a uniform prior) over all tokens equally . This causes the same local pressure on all tokens regardless of difficulty , and makes the policy learn similar high entropy distributions. Our fix: enforce the KL to prior w.r.t the global expected exit distribution across tokens. On the right, we show that this simple fix creates confident (low-entropy), difficulty-aware exits per token while maintaining low KL with the prior (P: Poisson, G: Geometric, U: Uniform) overall. The exits are truly heterogeneous as evident by the large jump rates ( $\lvert e_{i}-e_{i-1}\rvert\geq 2$ ).

Figure 4: Exit depth and Jump rate fractions in local and global modes. global enables more natural

\tau

-insensitive exit distributions and dynamic jump rates. (prior: Poisson; for the other two, see Figure 10 ).

Figure 5: Progression of local and global training modes (prior: Poisson, see others in Figure 13 ). While both reach similar avg. exit depths (

\approx 3

), local behaves like fixed@3, and global produces heterogeneous exits. Ragged-aware (RA) training drastically improves ragged PPL of global models while improving avg. exit depth.

Appendix figures & tables31 assets

Supplementary material from the paper’s appendix.

Appendix

Metric	Measured	Comparator
Avg. decode depth	2.252	4.000 (full depth)
Active decode row-depths	18,373	32,640 (full depth)
Physical K/V rows	25,621	40,016 (dense)
Nested depth passes	823	574.2 (ideal ragged)

Appendix

Table 3: Resource accounting for a traced H100 run with batch size 32, threshold

0.3

, 1,812 prompt tokens, and 8,192 generated tokens. Resource rows compare against full-depth execution with

U=4

; the final row compares the current nested scheduler against an ideal work-conserving Ragged-KV schedule.

Figure 6: The final

\beta

-schedule used in pre-training experiments.

Figure 7: With a simple

\beta

-schedule, looped model pre-training collapses to a fixed-depth-1 like behavior. The model processes inputs in the first loop and then acts like identity in the next 3 loops.

Figure 8: Exit depth distribution of the three priors.

Figure 9: Priors From left to right: Uniform, Poisson, Geometric. Training progression and exit depth fractions (

\tau=0.5

) compared between local and global prior enforcement.

Figure 10: Priors From left to right: Uniform, Poisson, Geometric. Final exit depth fractions compared between local and global prior enforcement.

Figure 11: Priors From left to right: Uniform, Poisson, Geometric. Exit depth fraction at different exit-thresholds for local prior enforcement, shows that the flipping-depth behavior exists across priors and is likely to continue with another flip between depths 2/3 if the prior enforcement is strengthened or applied for longer (see the dropping depth-3 fraction and rising depth-2 fraction for

\tau=0.1

and Poisson prior).

Panel of Figure 12 — Figure 12: global prior enforcement produces visibly heterogeneous hard exits. True Ragged-KV prefill at $\tau=0.5$ for the final 20B models (Poisson prior). Each column shows the first of five matched held-out 2,048-token sequences for one prior; the local and global rows within a column receive exactly the same tokens. Colored points are individual hard exit depths and the black curve is their 128-token rolling mean. A large jump is an adjacent pair with an absolute hard-depth change of at least two, $\lvert e_{i}-e_{i-1}\rvert\geq 2]$ : $3\!\to\!2$ is not a large jump, while $3\!\to\!1$ and $1\!\to\!4$ are. global traces have large-jump rates of 19.2%, 20.9%, and 13.2% for Poisson, Geometric, and Uniform priors, respectively, versus 0.2%, 2.5%, and 1.9% for the paired local traces. And we can see that the low depth tokens in global are actually placed meaningfully according to difficulty (commas, spaces, semicolons, simple words that follow obviously, etc).

Figure 13: Priors From left to right: Uniform, Poisson, Geometric. Training progression, avg. exit depth and fraction on each depth compared between local and global prior enforcement, along with ragged-aware alternate-step training branch-offs that greatly improve the ragged performance of global but not so much for local .

Figure 14: From left to right: Uniform, Poisson, and Geometric priors. Counterfactual and REINFORCE (RLK1) training provide no clear improvement over direct alternate-step Ragged-KV training despite their additional computation.

Benchmark	Examples	Shots	Metric
MMLU ( Hendrycks et al., 2021 )	14,042	5	Accuracy
ARC-Challenge ( Clark et al., 2018 )	1,172	25	Length-normalized accuracy
HellaSwag ( Zellers et al., 2019 )	10,042	10	Length-normalized accuracy
WinoGrande ( Sakaguchi et al., 2021 )	1,267	5	Accuracy

Appendix

Table 4: Non-thinking multiple-choice evaluation. All benchmarks use deterministic log-likelihood scoring without a chat template. Settings taken from the Ouro paper.

Benchmark	Examples	Shots	Max. tokens	Prompt
MMLU-Pro ( Wang et al., 2024 )	12,032	5	2,048	Answer generation
BBH ( Suzgun et al., 2022 )	6,511	3	1,024	Chain of thought
GSM8K ( Cobbe et al., 2021 )	1,319	3	1,024	Chain of thought
MATH500 ( Lightman et al., 2024 )	500	5	2,048	Answer generation

Appendix

Table 5: Non-thinking generated-answer evaluation. All benchmarks use 16 rollouts per example. Settings follow the Ouro paper.

Benchmark	Examples	Rollouts	Max. tokens	Scoring
AIME 2024	30	8	32,768	Qwen
AIME 2025	30	8	32,768	Qwen
AIME 2026	30	8	32,768	Qwen
GPQA-Diamond ( Rein et al., 2023 )	198	8	32,768	Qwen
OlympiadBench ( He et al., 2024 )	581	4	32,768	Qwen
HumanEval+ ( Chen et al., 2021 )	164	8	8,192	EvalPlus ( Liu et al., 2023 )

Appendix

Table 6: Thinking-model evaluation benchmarks. Every rollout is generated with thinking enabled. “Qwen” denotes Qwen3.8-27B correctness judgments; HumanEval+ is graded by executable tests.

Benchmarks	Link	License
MMLU	HuggingFace	MIT
ARC-Challenge	HuggingFace	CC BY-SA 4.0
HellaSwag	HuggingFace	MIT
Winogrande	HuggingFace	Apache 2.0
MMLU-Pro	HuggingFace	MIT
BBH	HuggingFace	Apache 2.0

Appendix

Table 7: The list of benchmarks.

Dataset	Link	License
FineWeb-Edu	HuggingFace	ODC-BY
AceReason-1.1-SFT	HuggingFace	CC BY 4.0
OpenThoughts3	HuggingFace	Apache 2.0

Appendix

Table 8: The list of training datasets.

Model	Link	License
Ouro-1.4B	HuggingFace	Apache 2.0
Ouro-1.4B-Thinking	HuggingFace	Apache 2.0
Ouro-2.6B	HuggingFace	Apache 2.0
Ouro-2.6B-Thinking	HuggingFace	Apache 2.0
Nanbeige-4.2-3B	HuggingFace	Apache 2.0
Qwen-3.8-27B	HuggingFace	Apache 2.0

Appendix

Table 9: The list of models used.

Figure 15: Ouro model’s BBH performance frontier w.r.t the average decode depth in Dense prefill + Ragged decode mode. Fixed-depth performance for reference.

Figure 16: Ouro model’s BBH performance frontier w.r.t the average overall depth (prefill + decode) in Dense prefill + Ragged decode mode.

Figure 17: Ouro model’s BBH performance frontier w.r.t the average decode depth in End-to-end Ragged mode.

Figure 18: Ouro model’s BBH performance frontier w.r.t the average overall depth (prefill + decode) in End-to-end Ragged mode.

Figure 19: Exit depth fraction of all evaluation benchmark tokens processed in Ragged KV modes.

Figure 20: Jump rates in immediate token transitions for all evaluation benchmark tokens processed in Ragged KV modes.

Log-likelihood evaluation
Dataset	Metric	Fixed-Depth			Dense-Ragged		ETE-Ragged
	depth / thr.	2	3	4	.3	.5	.3	.5
	Avg. depth	2.00	3.00	4.00	2.39	3.05	2.63	3.21
MMLU	accuracy	60.4 $\pm$ 0.4	66.7 $\pm$ 0.4	67.5 $\pm$ 0.4	–	–	61.9 $\pm$ 0.4	66.2 $\pm$ 0.4
ARC-C	accuracy	54.9 $\pm$ 1.5	59.9 $\pm$ 1.4	59.8 $\pm$ 1.4	–	–	59.0 $\pm$ 1.4	60.0 $\pm$ 1.4
HellaSwag	accuracy	70.8 $\pm$ 0.5	73.1 $\pm$ 0.4	73.4 $\pm$ 0.4	–	–	72.8 $\pm$ 0.4	73.2 $\pm$ 0.4

Appendix

Table 10: Performance of Ouro-1.4B in different Fixed-depth and Ragged inference modes. Values are percentages. Avg. depth denotes avg. decode depth for Dense-Ragged, and avg. overall depth for ETE-Ragged. Generation results report Qwen3.8-27B judged pass@

k

. Standard error across evaluation samples. Macro averages are unweighted and shown without an uncertainty estimate.

Log-likelihood evaluation
Dataset	Metric	Fixed-Depth			Dense-Ragged
	depth / thr.	2	3	4	.3	.5
	Avg. depth	2.00	3.00	4.00	2.21	3.03
MMLU	accuracy	67.7 $\pm$ 0.4	73.4 $\pm$ 0.4	74.4 $\pm$ 0.4	–	–
ARC-C	accuracy	63.5 $\pm$ 1.4	65.4 $\pm$ 1.4	66.6 $\pm$ 1.4	–	–
HellaSwag	accuracy	76.8 $\pm$ 0.4	77.9 $\pm$ 0.4	78.4 $\pm$ 0.4	–	–

Appendix

Table 11: Performance of Ouro-2.6B under Fixed-Depth and Dense-Ragged inference. Values are percentages. Avg. depth denotes avg. decode depth for Dense-Ragged. Generation results report Qwen3.8-27B judged pass@

k

. Standard error across evaluation samples. Macro averages are unweighted and shown without an uncertainty estimate.

Figure 21: Prefix instability of ETE-Ragged in Ouro-2.6B non-thinking. Native ragged prefill collapses on BBH. Computing only the first two prompt tokens at fixed depth four restores P@1 and P@4 close to the fixed-depth frontier while retaining early exit for the remaining prompt and decode tokens. Point labels denote exit thresholds.

Dataset	Metric	Fixed-Depth			Dense-Ragged			ETE-Ragged
	depth / thr.	2	3	4	.3	.45	.5	.3	.45	.5
	Avg. depth	2.00	3.00	4.00	2.30	2.50	2.92	2.31	2.49	2.91
AIME24	pass@1	24.2 $\pm$ 5.4	47.9 $\pm$ 7.3	50.4 $\pm$ 7.1	45.8 $\pm$ 7.1	47.9 $\pm$ 7.0	46.2 $\pm$ 6.9	40.8 $\pm$ 6.7	45.8 $\pm$ 7.0	46.2 $\pm$ 6.8
	pass@4	48.5 $\pm$ 7.6	67.3 $\pm$ 7.9	71.7 $\pm$ 7.7	67.1 $\pm$ 8.2	69.7 $\pm$ 7.7	68.6 $\pm$ 7.9	64.9 $\pm$ 7.8	68.0 $\pm$ 7.9	69.5 $\pm$ 7.6
AIME25	pass@1	24.2 $\pm$ 7.0	36.2 $\pm$ 7.2	37.9 $\pm$ 7.0	33.3 $\pm$ 7.3	36.7 $\pm$ 7.2	40.8 $\pm$ 7.4	30.0 $\pm$ 7.2	35.0 $\pm$ 6.9	41.2 $\pm$ 7.2
	pass@4	34.0 $\pm$ 8.1	55.0 $\pm$ 8.5	58.1 $\pm$ 8.5	51.0 $\pm$ 7.9	55.9 $\pm$ 8.2	59.2 $\pm$ 8.5	43.3 $\pm$ 8.4	58.3 $\pm$ 7.6	59.7 $\pm$ 8.6

Appendix

Table 12: Performance of Ouro-1.4B Thinking. AIME24/25/26, GPQA-Diamond and OlympiadBench use the Qwen3.8-27B judge, and HumanEval+ uses executable EvalPlus tests. Values are pass@

k

percentages with standard error across evaluation samples. Macro averages include all six benchmarks and are unweighted.

Dataset	Metric	Fixed-Depth			Dense-Ragged			ETE-Ragged
	depth / thr.	2	3	4	.3	.45	.5	.3	.45	.5
	Avg. depth	2.00	3.00	4.00	2.24	2.45	2.83	2.30	2.60	3.00
AIME24	pass@1	40.4 $\pm$ 6.7	62.1 $\pm$ 7.0	65.8 $\pm$ 6.8	56.7 $\pm$ 7.3	57.9 $\pm$ 7.3	61.7 $\pm$ 7.2	55.0 $\pm$ 7.1	59.6 $\pm$ 7.2	61.7 $\pm$ 6.9
	pass@4	64.4 $\pm$ 7.7	79.9 $\pm$ 6.9	82.0 $\pm$ 6.9	76.0 $\pm$ 7.3	75.2 $\pm$ 7.4	78.7 $\pm$ 6.8	77.3 $\pm$ 6.4	76.4 $\pm$ 7.3	80.4 $\pm$ 6.9
AIME25	pass@1	31.7 $\pm$ 6.7	49.2 $\pm$ 7.7	47.9 $\pm$ 7.9	45.4 $\pm$ 7.2	45.4 $\pm$ 7.4	52.9 $\pm$ 7.2	46.2 $\pm$ 7.2	49.2 $\pm$ 7.3	52.1 $\pm$ 7.5
	pass@4	49.0 $\pm$ 8.5	65.6 $\pm$ 8.2	62.1 $\pm$ 8.5	67.7 $\pm$ 7.6	64.8 $\pm$ 7.9	70.8 $\pm$ 8.1	66.9 $\pm$ 8.0	68.7 $\pm$ 7.9	69.2 $\pm$ 8.0

Appendix

Table 13: Performance of Ouro-2.6B Thinking. AIME24/25/26, GPQA-Diamond and OlympiadBench use the Qwen3.8-27B judge, and HumanEval+ uses executable EvalPlus tests. Values are pass@

k

percentages with standard error across evaluation samples. Macro averages include all six benchmarks and are unweighted.

Log-likelihood evaluation
Dataset	Metric	Fixed-Depth		ETE-Ragged
	depth / schedule	3	4	cos. $1\!\rightarrow\!.5$ / $.5$	$.7/.5$
	Avg. depth (P / D / O)			3.71 / 3.05 / 3.61	3.66 / 3.05 / 3.57
MMLU	accuracy	66.7 $\pm$ 0.4	67.5 $\pm$ 0.4	66.3 $\pm$ 0.4	–
ARC-C	accuracy	59.9 $\pm$ 1.4	59.8 $\pm$ 1.4	60.1 $\pm$ 1.4	–
HellaSwag	accuracy	73.1 $\pm$ 0.4	73.4 $\pm$ 0.4	73.2 $\pm$ 0.4	–

Appendix

Table 14: ETE-Ragged with different prefill schedules. Performance on Ouro-1.4B. Cosine decreases the prefill threshold from

1

.5

and uses

.5

during decode;

.7/.5

uses fixed prefill and decode thresholds. Generation results use the Qwen3.8-27B judge.

Log-likelihood evaluation
Dataset	Metric	Fixed-Depth		ETE-Ragged
	depth / schedule	3	4	cos. $1\!\rightarrow\!.5$ / $.5$	$.7/.5$
	Avg. depth (P / D / O)			3.63 / 3.03 / 3.51	4.00 / 3.45 / 3.88
MMLU	accuracy	73.4 $\pm$ 0.4	74.4 $\pm$ 0.4	72.0 $\pm$ 0.4	–
ARC-C	accuracy	65.4 $\pm$ 1.4	66.6 $\pm$ 1.4	65.4 $\pm$ 1.4	–
HellaSwag	accuracy	77.9 $\pm$ 0.4	78.4 $\pm$ 0.4	78.1 $\pm$ 0.4	–

Appendix

Table 15: ETE-Ragged with different prefill schedules. Performance on Ouro-2.6B. Cosine decreases the prefill threshold from

1

.5

and uses

.5

during decode;

.7/.5

uses fixed prefill and decode thresholds. Generation results use the Qwen3.8-27B judge.

Dataset	Metric	Fixed-Depth		ETE-Ragged
	depth / schedule	3	4	cos. $1\!\rightarrow\!.5$ / $.5$	$.7/.5$
	Avg. depth (P / D / O)			3.47 / 2.92 / 2.93	3.03 / 2.91 / 2.91
AIME25	pass@1	36.2 $\pm$ 7.2	37.9 $\pm$ 7.0	35.8 $\pm$ 7.6	38.3 $\pm$ 7.7
	pass@4	55.0 $\pm$ 8.5	58.1 $\pm$ 8.5	56.7 $\pm$ 9.2	60.0 $\pm$ 9.1
GPQA-Diamond	pass@1	38.3 $\pm$ 2.3	38.4 $\pm$ 2.3	39.6 $\pm$ 2.5	37.1 $\pm$ 2.5
	pass@4	67.9 $\pm$ 2.7	67.5 $\pm$ 2.7	70.7 $\pm$ 3.2	66.2 $\pm$ 3.4

Appendix

Table 16: ETE-Ragged with different prefill schedules. Performance on Ouro-1.4B-thinking. Cosine decreases the prefill threshold from

1

.5

and uses

.5

during decode;

.7/.5

uses fixed prefill and decode thresholds. Generation results use the Qwen3.8-27B judge.

Inference	P@1	P@4
Fixed@1	2.6	4.4
Fixed@2	26.7	43.7
Fixed@3	57.0	76.3
Fixed@4	65.6	81.5

Appendix

Table 17: Naive latest-only caches fail on Ouro-1.4B. BBH exact-answer results (%) for the non-thinking model. Ragged-Decode-Shared-KV uses depth-4 K/V for every prompt position, then appends each decode token’s exit-depth K/V. Ragged-Collapse-KV instead initializes the cache with every prompt token’s exit-depth K/V. Results use the raw strict answer-marker parser rather than the Qwen judge used in the headline generation tables. Every point has 135 samples and four generated rows. Fixed-depth controls are sampled (temperature 1.0, top-p .7); the cache baselines are greedy, so their P@1 and P@4 coincide.

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
Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching
Aug 10, 2026Kristian Schwethelm, Daniel Rueckert, Georgios KaissisEfficient Inference
Chair for AI in Healthcare and Medicine, Technical University of Munich (TUM), Germany · Department of Computing, Imperial College London, UK · Munich Center for Machine Learning (MCML), Germany +1
FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates
Sep 24, 2026Wanqi Yang, Shiwei LiuBlock Sparse Flash Attention Transformer Architectures
ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center

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

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates