LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation
Authors: Ahmadreza Jeddi
Marco Ciccone
Babak Taati
Venue: ICLR 2026 Poster (2026)
Keywords: large language models
LLMs
reasoning
looped transformers
efficient inference
parameter sharing
Primary area: foundation or frontier models, including LLMs

Abstract: Looped Transformers have emerged as an efficient and powerful class of models for reasoning in the language domain. Recent studies show that these models achieve strong performance on algorithmic and reasoning tasks, suggesting that looped architectures possess an inductive bias toward latent reasoning. However, prior approaches fix the number of loop iterations during training and inference, leaving open the question of whether these models can flexibly adapt their computational depth under variable compute budgets. We introduce LoopFormer, a looped Transformer trained on variable-length trajectories to enable budget-conditioned reasoning. Our core contribution is a shortcut-consistency training scheme that aligns trajectories of different lengths, ensuring that shorter loops yield informative representations while longer loops continue to refine them. LoopFormer conditions each loop on the current time and step size, enabling representations to evolve consistently across trajectories of varying length rather than drifting or stagnating. Empirically, LoopFormer demonstrates robust performance on language modeling and reasoning benchmarks even under aggressive compute constraints, while scaling gracefully with additional budget. These results show that looped Transformers are inherently suited for adaptive language modeling, opening a path toward controllable and budget-aware large language models.

Published as a conference paper at ICLR 2026
LOOPFORMER: ELASTIC-DEPTH LOOPED TRANS- FORMERS FOR LATENT REASONING VIA SHORTCUT MODULATION
Ahmadreza Jeddi
Marco Ciccone
Babak T
University of Toronto
Vector Institute
University Health Network
Corresponding Author: ajeddi@cs toronto.edu
ABSTRACT
Looped Transformers have emerged as an efficient and powerful class of mod- els for reasoning in the language domain. Recent studies show that these mod- els achieve strong performance on algorithmic and reasoning tasks, suggesting that looped architectures possess inductive bias toward latent reasoning. How- ever, prior approaches fix the number of loop iterations during training and in- ference, leaving open the question of whether these models flexibly adapt their computational depth under variable compute budgets. We introduce Loop- Former, looped Transformer trained variable-length trajectories enable budget-conditioned reasoning. Our core contribution is a shorteut-consistency training scheme that aligns trajectories of different lengths, ensuring that shorter loops yield informative representations while longer loops continue to refine them. LoopFormer conditions each loop on the current time and step size, enabling rep- resentations to evolve consistently across trajectories of varying length rather than drifting or stagnating. Empirically, LoopFormer demonstrates robust performance language modeling and reasoning benchmarks pute constraints, while scaling gracefully with additional budget. These results show that looped Transformers inherently suited for adaptive language model- ing, opening a path toward controllable and budget-aware large language models. Project Page: https://loopformer: github.io
under aggressive com-
INTRODUCTION
Transformers with parameter sharing, often called looped or recurrent Transformers, have emerged as an efficient and capable alternative to deep non-shared stacks across vision and natural language Dehghani al. (208);; Lan al. (2019); Jaegle et al. (202I): }) Dutta al. (2021); Geiping et al.
(205)). In particular, looped Transformers in the language modeling setting have shown promising performance variety of algorithmic and reasoning tasks Geiping al. (225); Saunshi et al.
(2024) 2025); Gatmiry et al. (2024). These models appear to possess an inductive bias toward rea- soning: property sometimes referred latent reasoning, where they internalize reasoning skills akin to explicit chain-of-thought prompting in large language models. Moreover, studies indicate that such abilities scale gracefully with effective computational depth and number of loops. yielding improved results reasoning benchmarks Xu & Sato (2024); Saunshi et al. (2025); Geiping et al. (2025). However, existing approaches almost always train and evaluate with a fixed number of un- rolls. This raises fundamental question: do looped Transformers truly exploit their computational depth flexibly, and can they be trained operate effectively under variable compute budgets?
Despite their promise, current looped models remain tied to single trajectory length. Once trained, their representations collapse when evaluated at shorter or longer depths, since those settings are off-distribution [ae et al. (2024); Fan et al. (2024). In practice, this means looped models spend the same budget as non-looped iso FLOP baselines, forfeiting one of their key motivations: flexible compute. We instead consider budger-conene conditionn language modeling: inference, user specifies a compute budget M, and the model should produce high-quality representations without retraining. While early exiting, routing, and layer dropping have made non-looped Transformers more dynamic

---

Published as a conference paper at ICLR 2026
Schuster al. (2022); Fan et al. (2019): Elhoushi et al. (2024); Shazeer et al. (011); Raposo et al. (2024), little has been explored for looped models. Naively transplanting these techniques into looped architectures is fragile Bae et al. 2025); Geiping et al. (202): repeated passes of the shared block often converge to similar, stagmani states. We aim for elastic depth: a single looped model that performs well across user-chosen budgets without retraining late-step degeneracy. The core challenge is to train looped models whose internal trajectories remain stable across depths, so that shorter routes do not degrade and longer routes continue to refine rather than collapse.
We present LoopFormer, a shortcut-modulater looped language model that supports elastic depth, maintaining strong performance across a range of inference budgets. Inspired by diffusion models Frans et al. (2024); Lu & Song (2024) and neural ODEs Chen et al.] (2018). we cast iterative rep- resentation refnement as trajectory in representation space: token states evolve from an initial ho toward a target h1 over normalized unit-time horizon. Our key insight is to explicitly condition each loop step on the current time and a step size At (a "jump"), allowing coarser trajectories to approximate fine-grained ones with fewer steps. Training employs a shoricut-consistency objective that aligns the final representations of shorter routes with those of the full route via a stop-gradient target, effectively performing self-distillation within the loop. At inference, this conditioning yields elastic depth without retraining: the user selects a budget M S L (maximum loops) and a step schedule, and performance scales smoothly with compute, rather than collapsing shorter depths.
Departing from dynamic compute through routing or token halting Dehghani et al, (2018); Bae et al. 2025), we instead introduce the notion of loop trajectories, which enable effective and efficient language modeling, thereby giving the model the ability to internalize and refine reasoning and per- formance compute increases. Empirically, LoopFormer not only preserves the language modeling abilities reported in prior work, but also outperforms baseline looped models on language tasks and closes much of the gap with iso-FLOP non-looped models. This creates an opportunity to study loop trajectories as a new lens on reasoning in looped models. Moreover, unlike early-exit mechanisms that often collapse to degenerate states, LoopFormer maintains stable refinement: both perplexity and reasoning scale gracefully with the compute budget.
The key contributions of this paper are summarized as follows:
We formulate class of chorteut-modulast looped Transformers for language modeling that con- ditions each pass on internal time and step sizes.
We introduce a shortcut- consistt-- training protocol over families of step schedules that enables nompute-budgeted inference (elastic depth) without retraining.
We analyze representation dynamics across loop iterations using multiple geometric and information- -theoreei metrics (anisotropy, curvature, entropy, CKA), finding that adaptive early- exit looped models exhibit representational collapse, while our tain evolving, den-eeeeee states.
models main-
gains both perplexity and zero-shot rea- soning across diverse language benchmarks; ablations isolate the roles of time/step conditioning and offer practical guidance for choosing trajectories under fixed budgets.
We demonstrate consistent performance-per-comee
2 RELATED WORK
Recursive Looped Transformers. Parameter sharing provides an orthogonal route to efficiency and effective depth. The Universal Transformer demonstrated that repeatedly applying single block can match the representational capacity of deep non-shared stacks, while also introducing adaptive computation time Dehghani et al. (2018). ALBERT further showed that extensive cross- layer weight-tying yields substantial parameter savings during pretraining, without compromising downstream performance [Ln et al. (209)). Deep Equilibrium Models (DEQ) extend this paradigm by defining an implicit infinttlly deep, weight-tied transformation solved via fixed-point iteration and implicit differentiation Bai et al. al, (2019), with related equilibrium and recurrent architectures also explored in vision and multimodal contexts[Jaegh et al. (202) }; Yang et al. (2022). More recent work has studied looping as programmable computation Giannou et al. (2023). iterative kata-fitting or optimization solvers Yang et al. (2023), mechanisms for algorithmic length generalization Fan et al. (2024), and even the cornng-completeness of looped decoders for certain graph algorithms) De Luca & Fountoulakis 2024). Most relevant our work, Time-Modulated Looped Transformers (TMLT)
2

---

Published as a conference paper at ICLR 2026
Xu & Sato (2024) analyze the expressive power of looped Transformers in language modeling, showing that conditioning on timestep (loop index) improves scaling and perplexity. Our approach builds on this insight but extends it by conditioning on both normalized time and step size, and by training across families of trajectories rather than a single fixed schedule.
Dynamic compute. Dynamic compute allocation reduces inference cost by skipping reallo- cating computation where is not needed [Beniio et al. (2015); Huang et al. (2016); |aadda et al. (2016). Early exiting halts processing for "easy" inputs at intermediate layers. allowing deeper computation only for harder cases Elbaaad et al. (2019); |Schuster et al.] (2022); |Elhoushi et al. (2024). Other approaches include layer dropping and pruning strategies for BERT-style models Fan et al. 2019) Hou et al. (2020). Mfxture-of-Experts (MoE) increases capacity through sparse expert routing Shazeer al. (2017); Fedus al. (2022). while Mixture-of-Depiin
(MoD) Raposo et al. (2244) reframes adaptivity as token-wise routing across layers, enabling finegrained dynamic allo- cation of depth. Despite these advances, the adaptation of dynamic compute techniques to exploit the looping capability of recursive models has received little attention. A concurrent line of work, Mixture of Recursions Bae et al. (2025), extends routing to recursive stacks by varying the number of shared-block applications per token. In contrast, our approach does not route or halt computation at the token level; instead, we train looped models to be budget-conditione, ensuring that the same parameters operate robustly across user-specified loop budgets.
Latent reasoning. An increasing body of work investigates reasoning carried out within hid- den states rather than through explicit chain-of of-thuh prompting Goyal et al. (2023);
3 C Cheng & Van Durme (2024); Pfau et al. (2024); Kaissis et al. (2026); Chen et al. (2026); ; Zhu et al. (2025);
(2025); Wang et al. (2025); Zhang et al. (2025); Knupp et al. (2026); Liang & Pan (2026). Several approaches employ looped models to simulate multi-step reasoning or to ap-
proximate chain-of-thought dynamics directly HHo al. (2024); Saunshi et al. (2025); Wang al.
2025). In parallel, theoretical and empirical studies connect network depth and algorithmic gener- alization reasoning ability Merrill & Sabharwal] (2023;; Chen & Zou (202)) Ye al. (2024): Xu
& Sato (025)) More recently, analyses suggest that looped Transformers possess an inductive bias for reasoning that strengthens with increasing effective computational depth [Saunshi al.] (2024)
2025). Collectively, this line of work frames such abilities as lutent reasoning, wherein models perform iterative computations within their hidden-state space without explicit verbalization. Our work leverages this inductive bias of looped models but emphasizes trajectory robustness: ensuring that hidden-state computation remains informative even under shorter budgets, while continuing to refine effectively at greater depths.
Time / shorteut modulation and conditioning. Conditioning networks on continuous time or re- lated control variables has proven highly effective in generative diffusson modeling Ho et al. (2020); Rombach et al. (2022) Lipman et al. (002).. Diffusion Transformers (DiTyPeebless & (202) in- corporate timestep-conditeee conditioned adaptive normalization, while consistency-bert bassed training aligns lutions across discretizations Song et al. (203). Recent shortcut and one-step diffusion approaches further distill long trajectories into few steps by enforcing consistency between coarse and fine solvers Frans et al. 2024)): & Song (2024). These ideas are increasingly migrating into language modeling: for example, Mome-ladllaedd Looped Transformers (TMLT) Xu & Sato (2024) adapt DiT-style timestep conditioning to recursive language models, demonstrating improved scaling and perplexity. LoopFormer extends this trend by conditioning each loop not only on normalized time t but also on the step size At, and by training across families of sampled trajectories with shortcut- consistency objective. This design enables robust performance under arbitrary user-specified com- pute budgets, without relying per-token halting or routing mechanisms.
3 SHORTCUT-MODULATTI LOOPED MODELS FOR LATENT REASONING
3.1 NOTATION AND PROBLEM STATEMENT
We denote by X = (2), ........ FT) a sequence of T tokens drawn from vocabulary V. The token embeddings are obtained via E (X) E Txx Without loss of generality, positional embeddings can be added in one-shot manner applied through alternatives such RoPE; in this work,
3

---

Published as a conference paper at ICLR 2026
Query: What city the Eiffel Tower located in?
LM Head
Context tokens
Norm
token
Gate
FFN
M
Scale
V:
France
LoopFrrmer
RMS Norm
Block
t=0
[=]
next token
Parts
Gate
the
MLP
MLP
MHSA
Prance
Scale
Vi
RMS Norm
MLP
t=0
[=]
next token
Ae
Parts
M L
tily
the
France
A.
Inpui
Trajectory
(a) LoopFormer architecture
(b) Budget-conditioned trajectories
Figure 1: (a) LoopFormer conditions each loop normalized time and step size At, modulat- ing RMSNorm scales and gating the MHSA/FFN residuals. (b) During inference, shorter-budgee trajectories (M < L) approximate the full L-step route; more budget yields progressively refined next-token distributions while preserving utility at low budget.
adopt the former for simplicity. The initial hidden states are therefore
A (0)
(X) + Epos [1:2]
RTxd
range of looping mechanisms for looped Transformers has been studied in prior work Takase & Kiyono (2021): Saunshi et al. (2025); Bae et al. (2025 (024)) including cycle, middle-cycle, and relaxed-cycle variants. Since our focus on trajectories and their representation dynamics, we adopt the simplest cycle design, where stack of & Transformer blocks, denoted by (.).), is recursively applied on the hidden state. Following Saunshi al. (2025). we write (I @ L.) for a looped model
with & blocks repeated L times (approximate cost KL FLOPs), and (k&L @ 1) for non-looped Transformer comparable depth.
LoopFormer applies & (.) for M iterations (1 < M < I),, conditioning each loop on the pair (t,-1, A,). where A. t; 4-1 is the step size of the i-th iteration. We refer to AM (Ai...., AM) trajectory, constrained by MM
to < t < ... < tM = 1 are cumulative normalized timesteps, and
The maximum trajectory corresponds L uniform steps with A; J. for all i, At inferennce user specifies a compute budget M and provides step schedule A such that 2: =1
A,
3.2 ARCHITECTURE
We introduce LoopFrrmer, looped language modeling architecture designed to remain faithful to the standard Transformer while incorporating trajectory- based conditioning, Figure la provides an overview. In this section describe the forward pass and the components of the model.

---

Published as a conference paper at ICLR 2026
LoopFormer is looped decoder-only Transformer in which single shared stack is applied iter- atively; the key novelty is how each loop is conditioned. At iteration i, the model conditions on the pair (t,-1, A,). where fę-1 [0. 1] is the cumulative normalized time and A; E (0, 1] is the step size. Both scalars are encoded with sine-cosine frequency embeddings and projected by small MLPs obtain et and ea, which are summed to form & 6 ea. This signal modulates the LoopFormer Block: an MLP maps e; to scaling (71, 722 for the two RMSNorm layers and to gating (a); a2) applied immediately before the residual connections of MHSA and FFN. As a result, each loop is explicitly aware of its location and granularity in the trajectory, enabling consistent behavior across coarse- and fine- nhe-graindd schedules.
Relation to prior work. For the design of our architecture, we follow the overall approach of looped models such as ALBERT and recent looped decoders et al. (2019); (Sausshi al.) (2025).
and focus our novelty on trajectory conditioning. Our modulation mechanism is inspired by Diffu- sion Transformers (DiT) Peebles & Xie (2023) and shortcut/one-step diffusion Frans et al. (024),
which condition blocks timestep via adaptive normalization (adaLN), regressing scale/shift and gates from a time embedding. In the looped language model setting. TMLT Xu & Sato (2024) adopts this DiT-style idea by using the loop index time (with RMSNorm). LoopFormer extends this conditioning by using both normalized time / and step size Al, and by training over families of trajectories with shorteut-consistennyy objective, yielding trajectory-consistent representations suitable for budget-condditoned inference.
3.3 TRAINING & INFERENCE WITH LOOPFORMER
A key goal of LoopFormer is to train models that perform well even with fewer than L loops, thereby enabling elastic depth at inference time. Unlike prior adaptive-compute approaches that rely on per-token halting or early exits, our framework uses a user-defined compute budget, similar in spirit to diffusion models. Empirically, we find that naive early exiting in looped architectures leads to stagnant representations in later iterations, under-utilizing depth. To address this, Loop- Former leverages shortcut conditioning together with consistency loss that encourages trajectories of different lengths to converge toward the full trajectory of length L. The overall objective is:
f + Ls + A Leonst
where LL and Ls denote next-token prediction losses for the longest and sampled shortcut trajecto- ries, respectively, and Leons is ttop-gradient consistency loss aligning per-token logits of shorter trajectories to those of the longest trajectory. We set a X = 0.1 in all experiments. Algorithm formalizes the training procedure.
Shorteut trajectory sampling. Given a LoopFormer (A: @ L) and the maximum trajectory AL, during training, we additionally sample shortcut trajectory As with budget 1 simplicity, at each batch we first sample a shortcut length ~ [III..... L 1}. then uniformly draw the discrete step schedule As over [0, 1] such that > A, 1. This ensures exposure to both long and short trajectories during training.
L. For
Inference with elastic depth. At inference time, LoopFormer be deployed flexibly any budget M < L. A user specifies M and Av. and the model produces outputs that scale smoothly with compute, without retraining. |Figure 1b 1b|provides conceptual illustration of budget-conditioned refinement: for any chosen budget, shorter routes are trained to approximate the t=1 endpoint, with quality improving as steps increase. Algorithm 2 outlines the inference procedure.
4 EXPERIMENTS
We evaluate latent reasoning, scalability, efficiency, and representation dynamics of our shortcut- modulated looped language models. Following Tay et al. (Suns) et (2025), we compare 24-layer, ~ -BB-arrmeterr non-looped Transformer with use a GPT-style decoder Radford et al. (2019) with NanoGPT configurations. Training is performed on deduplicated subset of The Pile Gao et al. (2020) for 25B tokens in accordance with Chinchilla scaling Hoffmann et al. (2022). See See[Appendi: A for details. Unless otherwise specified, all reported results use uniform step sizes at inference: for compute budget M, A 1/M M for all i.
looped variants. All models
5

---

Published as a conference paper at ICLR 2026
Algorithm LoopFormer Training
Algorithm 2 LoopFormer Inference
while not converged do
Given input X
Sample batch (X, Y
Choose budget M < L Sample schedule AM
Construct max trajectory AL
Sample shortcut length S 1 -
Initialize h()) & Brok (X) + Epss for to M do
Sample trajectory As of length 5
h(*)) <- (b L)
() t P& (1) As)
h - Pa(h ;; As) return LMHead (MM)) -1])
CL < CE(LMHead ).
Ls < ( (,
Ccons < stopgrad(h(4) h (1)
Update using Vo + As + A2Coons)
Table 1: Perplexity and zero-shot reasoning for @3 looped models under three inference budgets
(24x, 12x, 6x). At x x we also report fixed-depth baselines; at 12x/6x we compare against Base (12 s 1) and (6 & 1). While depth-elastic, LoopFormer narrows the perplexity gap Base and is competitive on reasoning, outperforming other looped variants, especially at higher budgets.
Perplesits
Langunge Tasks (Accuracy))
Parsmn / FLOPs
Pile
Fine Weh-EEh OpenWabTest
COPA
HS
OBQA
PIQA
Race
SaQ
ARC
SIQA
WG
Avg Ace
Budget: 24x
Bas (21
24h, 24x
9.49
20.7
20008
61
35.04
41.96
27.6
66
29
78.1
33.43
38.18
50.4
45.27
Beae-Lar cup (5
3x M4
10.9
53
34.53
90.46
34.68
63.32
28.8
63.7
31.71
38.43
498
47.83
())
10:38
2967
68.8
36.95
Lorp (3-0
Bosee Loop- Bowe-Loop-EB-Come EB-pp-----Cns
/ 24x
11.56
25
24.41
30.54
32.19
26.8
62.13
28.25
64.9
31.45
37.4
$2.88
63,25
TMLT BE (3:00
3x/
10.7
24.07
23.17
65
51,03
35.14
25.4
63.11
25.8
68.8
31.74
37.41
50.85
44.03
Locplormer Ours (3
0 24x
10.28
22.87
21.98
66
32.3
33.27
26.8
63.33
30.81
68
32.71
37.99
51.94
44.81
Budget:
Buse (12x1)
12x/ 12x
9:88
22.24
21.41
67
32.72
37.78
26.2
64.69
29,86
09.1
32.29
30.38
51.3
44.93
Naived Lmp-RR (3-g)
12x
11:66
26.74
26.93
368
43.99
mop--FF-Cmmn (3
2804
TMLT-BE (3:00
/ 12x
12.18
28.28
27.11
60
20,88
27.73
26.4
62.02
28.8
61.4
30.50
36,69
51.54
41.5
Looplormee- Ours (3
Sx / 12x
11:12
25.02
24:21
68
31
30.35
25.4
63.06
28.25
66.3
51.78
50.77
$3.43
43.73
Budget:
Buse (6x3)
66. 6x
11.13
25.28
24.28
64
30.45
33.45
25.6
62.51
2804
67.5
31,07
36.18
48.54
42.73
Naive Loop BG (3
3x. 6x
12:61
29.36
28.38
63
29.64
21.28
24.6
61.48
26.41
62.3
29.91
35.41
50,6
60,97
Loop- (3.92)
3x.
15.79
37.84
26.41
50.01
34.9
50.27
LocpFormer Ouzs (3
3x./64
14.3
33.45
32.46
63
28.09
26.14
26.6
60.1
25.74
55.6
29.29
35.26
50.2
40.36
4.1 LATENT REASONING AND PERPLEXITY
We train looped models under different parameter and compute budgets. Using the notation (k: & L) for a k-block & unrolled L times, we consider E (1,2,3) and train with maximum loops I. E {8, 12, 24). Fixed-loop models are trained and evaluated with the same L, whereas depth-elastic models support inference at any M < L. In this section we use kL as a proxy for FLOPs when comparing baselines, ignoring embeeding/unemmedding costs; Table summarizes the results for k- == 3.
Evaluation metrics and benchmarks. Following/Saunshi ctt al. (2025), we report perplexity and downstream zero-shot accuracy. Perplexity is measured on Fine Edu Penedo et al. (2024), OpenWebText Gooaaslan & Cohen (019), and The Pile. For
Saunshi et al. (2025): Geiping et al. (2025); Bae

---

Published as a conference paper at ICLR 2026
block iters)
blocks iters)
Iters)
Leephener
hanecary Al
Perplesity
14-
13-
12-
10
12x
23x
24c
64
18x
21x
24x
6x
12x
15x
18x
23x
24k
FLOPs
FLOPs
FLOPs
Figure 2: Scaling with layers and loops. Perplexity on The Pile across (k, I.) and inference budgets
M S L; larger k lowers perplexity at fixed M, and additional loops further reduce it.
block
k blocks iters)
k= blocks iters)
Average Aucuresc 196
42-
42-
LeepTomerr
THESEE
THETEE
Duse Cenp-ECCE
36
01
12c
15c
24c
1.5
20x
7.0
17x
1hs
10c
21x
FLOPs
F1.OPs
FLDPs
Figure 3: Scaling with layers and loops. Average zero-shot reasoning accuracy on 10 tasks for the same settings; more loops improve reasoning at fixed k. Across budgets, LoopFormer scales smoothly without collapse and consistently outperforms other looped baselines.
latent reasoning, we report zero-shot accuracy on ten established benchmarks spanning a range of reasoning difficulty: COPA Roemmele et al. (201T), HellaSwag (HS) Zellers et al. (019)) LAM- BADA (LB) Paperno et al. (2016), OpenBook QA (OBQA) Mihaylov et al. (2018). PIQA Bisk et al.
22000, RACE Lai et al. (2017), Social IQA (SIQA) Sap et al. (2019). ARC Clack et al. (2018), SciQ Welbl et al. (201) 7), and WinoGrande (WG) Sakaguchi et al. (202I). Where available, we use normalized accuracy; for ARC we report the average over Easy and Challenge.
Baselines. We compare LoopFormer against two groups of baselines.
Fixed-depth models; (a) Base: non-looped Transformer; (b) Base-Loop: standard looped model as in Saunshi et al. (0225); (c) TMLT: looped model with timestep conditioning Xu & Sato (0024).
Depth-elastic models: (i) Base-Loop-EE: naive early exiting applied to the basic looped model!" (ii) Base-Loop- Loop-EE-Cons: (i) augmented with our consistency loss during training; (iii) TMLT-EE:
early exiting and consistency training applied to TMLT to enable depth elasticity.
Findings.
Table highlights three trends:
Consistent with Mohtashami et al. (2023); Saunshi et al. (2025). ) looped models trail non-looped baselines on perplexity, reflecting the role of parameters memorization. However, LoopFormer closes much of this gap, surpassing even fixed-depth looped variants.
In zero-shot reasoning, looped models benefit from iterative refinement and can approach iso- FLOP non-looped baselines; LoopFormer is the most competitive among looped models.
Under reduced budgets, LoopFormer maintains high utility: at 12x it remains close to Base (12@1) on both perplexity and reasoning, indicating that budget-conditioned trajectories preserve informative representations rather than collapsing.
We next examine the effect of the number of layers k and the number of loops L,, and compare against depth-elastic alternatives under more compute configurations.
Akin to Recurrent Depth Geiping al. (025) without sandwich normalizations or randomly initialized recurrent state. We trained both variants; the simpler one performs better.
7

---

Published as a conference paper at ICLR 2026
0.90
Ateence Contee Stntionts
2.54
Curatta
Logfornet
0.84
fowe Lon
TML7
0.53
Necasalaed
Nartaalood
(a) Curvature
(b) Anisotropy
(c) Prompt Entropy
Figure 4: Representation metrics over normalized depth. Panels show (a) curvature, (b) anisotropy.
and (c) normalized prompt entropy. Early-exit baselines remain flat, indicating minimal change with additional loop steps, whereas LoopFormer exhibits sustained evolution that rises through mid- depths and tapers near the end, suggesting useful depth elastic dynamics.
11
Lrop Lanp-Bar
414
(a) LoopFormer
(h) Base-Loop-EF-Cons
(c) TMLT-EE
Figure 5: CKA similarity across loop steps. Each heatmap reports cross-step CKA within a model family. Depth-elastic baselines show high CKA (indicating little change across loops), suggesting stagnation. LoopFormer exhibits progressive drift, especially toward the later steps.
4.2 NUMBER OF LAYERS VS. NUMBER OF LOOPS
We study how perplexity and reasoning change the number of Transformer layers per block (k) varies. We train & E {1, 3) with maximum loops L E (8, 12, 24). then evaluate across multiple inference budgets (M < L). Figure 2 reports perplexity and Figure reports average zero-shot reasoning accuracy across benchmarks.
As the plots show, both perplexity and reasoning improve with larger and more iterations. Loop- Former preserves these trends under budgets (M < L): shorter trajectories remain informative and added loops yield smooth gains without collapse. Our consistency augmeette training also improves the scaling of Base-Loop and TMLT in the depth-elastic regime.
4.3 ANALYZING REPRESENTATION COLLAPSE IN LOOPED TRANSFORMERS
The role of depth in Transformers has been widely examined through scaling laws and theory Ka- plan al. (2020): Csordás al.) (2025), with recurring observation that very deep stacks can exhibit representation degeneration (or collapxe). where hidden states change little across layers Ethayarajh (2019), 9): Dong al. (2021) Godey et al. (2024). Several works attribute this to an induc- live bias of self-attention toward uniformity, including rank decay in attention maps with depth [ong] ctt al. (2021). Related studies anisotropy Razzhigaev et al. (2023); Godey et al. (2024) report pro- nounced effects in language modeling, and these effects could be amplified in looped architectures where a single block is repeatedly applied Dong et al. (202T).
We analyze token dynamics along the computation depth using four complementary metrics: (i) anisotropy Godey et al. (2024) within layer, measured as average pairwise cosine similarity among tokens in the prompt; higher values indicate more aligned (less diverse) directions. (ii) Curva-
8

---

Published as a conference paper at ICLR 2026
1.-
1.-
0.-
0.8
0.-
0.-
0.4
0.
Timee
0.4
2000
(a) Perplexity (3 & 8: M
(b) Accuracy (3 M
(c) Perplexity (2 & 12; M 6)
Figure 6: Performance across length-M trajectories. (a.b) enumerate all M schedules for Loop-
Former (3 & 8), reporting perplexity and average zero-shot accuracy: (c) shows perplexity for M- schedules of (2 @ 12). Even at fixed budget, trajectory choice matters: spreads are large, and top schedules allocate coarser steps early and finer steps late.
ture Hosseini & Fedorenko 2023). local geometric measure of how rapidly token representations
change direction across neighboring positions. (iii) Prompt entropy (Skean et al. (2025), matrix- based estimate of how spread-out token embeddings are across feature dimensions; higher entropy suggests greater diversity and lower redundancy. (iv) CKA similarity Kornblith et al. (2019) across loop steps, quantifying representational similarity between different iterations of the shared block. While prior work studies these metrics in various architectures and their correlation with downstream performance Garrido et al. (2023); (Sken et al. (2025), here we use them primarily as diagnostics to assess whether looped models use additional computation (iterative depth) stagnate.
Figure 4 summarizes curvature, anisotropy, and prompt entropy across loop steps; Figure 5 reports cross-step CKA. Early-exit looped baselines show flat trajectories on all three metrics and high CKA, indicating stagnation. In contrast, LoopFormer evolves: early steps have lower curvature, weaker angular alignment, and lower entropy; steps progress, curvature and entropy rise and alignment increases, then all three taper near the final step as the model prepares for unembedding.
Overall, these patterns suggest that LoopFormer maintains useful representational dynamics across loops, with shorter trajectories remaining informative and longer trajectories providing additional refinement, rather than converging prematurely to a static state. We view this as evidence that shortcut conditioning and consistency training help avert collapse in depth-elastic looped models.
4.4 How TO CHOOSE TRAJECTORIES UNDER A FIXED BUDGET
A practical question for our depth-elastic models is: given budget M < L. how should we choose A A Beyond average gains, we ask which trajectories work best and which parts of time contribute most. The representation analyses in Figure 4 and Figuure 5 suggest pattern: early steps produce more similar states, activity rises mid/ then tapers near the end.
To probe this directly,
toy yet exhaustive study. For LoopFormer (3 8) with budget M 4, we enumerate all trajectories with step sizes summing to 1 (aligned to training granularity) and measure perplexity and average zero-shot accuracy (Figure 6a Figure 6b). Despite identical compute, performance varies across schedules (spread 1.44 perplexity and ~1. accuracy points).
Repeating for (2 12) with M 6, the perplexity spread grows to nearly (Figure 6ch.
Findings. (1) Even under uniform training over budgets and times, some schedules outperform others by wide margins.
(2) The best schedules for perplexity and for downstream reasoning are close but not identical; both favor allocating larger steps early and finer steps late.
DISCUSSION AND FUTURE DIRECTIONS
We introduced LoopFormer a looped Transformer that conditions each iteration on normalized time I and step size At, trained across trajectories with chonteut-consistency loss. This framing induces

---

Published as a conference paper at ICLR 2026
loop trajectories in hidden space: under fixed budget, shorter routes yield useful intermediate states, and additional steps refine them toward a shared t=1 endpoint, realizing latent reasoning while supporting dedget-conditionee (elastic) inference without retraining. Empirically, Loop- Former delivers strong performance-ee--oopprr on perplexity and consistently improves down- stream zero-shot reasoning, while avoiding the representational stagnation seen in naive early-exit baselines, Limitations include global (sequence-level) rather than instance/loken-adppty budget- ing, added training overhead from multi-trajectory consistency, and correlational (not causal) repre- sentation analyses. Promising directions include theory/diagnosties of the representation space.
schedule policies and deeper
ACKNOWLEDGEMENTS
We thank Prof. Colin Raffel (University of Toronto) and Prof. Ali Etemad (Queen's University) for valuable feedback throughout this project. We also thank Negin Baghbanzadeh and Dev Shah for their help during the early stages of our work on looped vision models.
Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, the Digital Research Alliance of Canada, and companies sponsoring the Vector Institute)
REFERENCES
Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, and Tal Schuster. Re- laxed recursive transformers: Effective parameter sharing with layer-wise lora. arXiv preprint arXiv:2410.20672 2024.
Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, al. Mixture-of-recursionns Learning dynamic recur- sive depths for adaptive token-level computation. arXiv preprint arXiv:2507.1053 2025.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in neural information processing systems, 32, 2019.
Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv: 1.1.00697, 2015.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelli- gence, volume 34, pp.
2020.
Guanxu Chen, Dongrui Liu, and Jing Shao. Loop as bridge: Can looped transformers truly link representation space and natural language outputs? arXiv preprint arXiv:2601.10242, 2026.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Xingwu Chen and Difan Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404.0160) 2024.
Jeffrey Cheng and Benjamin Van Durme. Compressed chain of thought: Efficient reasoning through dense representations. arXiv preprint 24X1::24/2.13/7/1 2024,
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1rXiv:1803.05457 2018.
Róbert Csordás, Christopher D Manning, and Christopher Potts. Do language models use their depth efficiently? arXiv preprint arXiv:2505.13804 2025.
Artur Back De Luca and Kimon Fountoulakis. Simulation of graph algorithms with looped trans- formers. arXiv preprint 10X 2024.
10

---

Published as a conference paper at ICLR 2026
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819. 2018.
Yihe Dong, Jean- Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. International conference on machine learning, pp. 2793-2803. PMLR, 2021.
Subhabrata Dutta, Tanya Gautam, Soumen Chakrabarti, and Tanmoy Chakraborty. Redesigning the transformer architecture with insights from multi-particle dynamical systems. Advances in Neural Information Processing Systems, 34:5531-5544, 2021.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. preprint arXiv:1910.10000 2019.
transformer. arXiv
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, et al. Layerskip: Enabling early exit inference and self-speculative decoding. arXiv preprint arXiv:2404. 16710, 2024.
Kawin Ethayarajh. How contextual are contextualized word representations? comparing the geom- etry of BERT, ELMo, and GPT-2 embeddings. arXiv preprint arXiv: 2019.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1556, 2019.
Ying Fan, Yilun Du, Kannan Ramchandran, and Kangwook Lee. Looped transformers for length generalization. arXiv preprint 24Xv:2449./5647 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120): 2022.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv preprint arXhv:2410.12557 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster. Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
Quentin Garrido, Randall Balestriero, Laurent Najman, and Yann Lecun. Rankme: Assessing the downstream performance of pretrained self-supervised representations by their rank. In Interna- tional conference on machine learning, pp. 10929-10974, PMLR, 2023.
Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? arXiv preprint arXiv:2410.08220 2024.
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh. Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and 'om Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint wwxiv:2502.05/7/ 2025.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. International Conference on Machine Learning. pp. -1398-11442, PMLR, 2023.
Nathan Godey, Éric de Clergerie, and Benoît Sagot. Anisotropy inherent to self-attention in transformers. arXiv preprint arXiv:2401.12143, 2024.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus http:///s///////ooo00oottttt. TpenWebTextCozpuu 2019.
io/
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint arXiv:2310.0222 02.00020222.... 2023.
11

---

Published as a conference paper at ICLR 2026
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06769 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilisttc models. Advances in neural information processing systems, 33:6840-6851, 2020.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optiman large language models. arXiv preprint arXiv:2203. 15556, 2022.
Eghbal Hosseini and Evelina Fedorenko. Large language models implicitly learn to straighten neural sentence trajectories construct predictive representation of natural language. Advances in Neural Information Processing Systems, 36:43918 43-330,, 2023.
Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. DynaBERT: Dynamic BERT with adaptive width and depth. Advances in Neural Information Processing Systems, 33: 2020.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference computer vision, pp. 646-661. Springer, 2016.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pp. 4651 1664. PMLR, 2021.
Alexia Jolicoour-Martinen. Less is more: Recursive reasoning with tiny networks. arXiv preprint trXhv:2510.048711 2025.
Georgios Kaissis, David Mildenberger, Juan Felipe Gomez, Martin J. Menten, and Eleni Triantafil- lou. Step-resolved data attribution for looped transformers, 2026. URL https://arxiv:
org/abs/2602.1005
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836 2020.
Jonas Knupp, Jan Hendrik Metzen, Jeremias Bohn, Georg Groh, and Kristian Kersting. Depth- recurrent attention mixtures: Giving latent reasoning the attention it deserves. arXiv preprint 2rXiv:2601.21582. 2026.
Simon Kornblihh, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference machine learning, pp. 3519- 3529. PMIR, 2019.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. RACE: Large-scale reading comprehension dataset from examinations. arXiv preprint arXiv:17044 04683, 2017.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. ALBERT: lite BERT for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019.
Jia Liang and Liangming Pan. Do latent-cot models think step-by-step? mechanistic study sequential reasoning tasks. arXiv preprint 2xXiv26602.00449 2026.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous- time consistency models. arXiv preprint arXiv:2410.1108: 2024.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923, 2023.
12

---

Published as a conference paper at ICLR 2026
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02885 2018.
Amirkeivan Mohtashami, Matteo Pagliardini, and Martin Jaggi. CoTFormer: More tokens with attention make up for less depth. In Workshop on Advancing Neural Network Training: Compu- tational Efficiency, Scalability, and Resource Optimization (WANT@ NeurlPS 2023), 2023.
Priyadarshini Panda, Abhronil Sengupta, and Kaushik Roy. Conditional deep learning for energy- efficient and enhanced pattern recognition. In 2016 design, automation & test in europe confer- ence & exhibition (DATE), pp. 475 IEBE, 2016.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoe Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. arXiv preprint ar0i:: 06031, 2016.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195-4205, 2023.
Guilherme Penedo, Hynek Kydlíčck, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances Neural Information Processing Systems, 37:308 1-30849, 2024.
Jacob Pfau, William Merrill, and Samuel R Bowman. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint 2rX1v:2404.15758 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. -f Dynamically allocating compute transformer-basee lan- guage models. arXiv preprint 24xi::2200.022258, 2024.
Anton Razzhigaev, Matvey Mikhalchuk, Elizaveta Goncharova, Ivan Oseledets, Denis Dimitrov, and Andrey Kuznetsov. The shape of learning: Anisotropy and intrinsic dimensions in transformer- based models. arXiv preprint arXiv:2311.05928 2023.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In AAAI spring symposium: logical formaliza- tions of commonsense reasoning, pp. 90-95, 2011.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pp. 10684-10695, 2022.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. arXiv preprint arXiv:1904.0 09728, 2019.
Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank Jakkam Reddi, and Sanjiv Kumar. On the inductive bias of stacking towards improving reasoning. Advances in Neural Information Processing Systems, 37:7 1437-71464,, 2024.
Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank Reddi. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint 27X1: 2025.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems,
2022.
13

---

Published as a conference paper at ICLR 2026
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outragcously large neural networks: The cransely-gated caaturr-of-expertr layer. arXiv preprint arXiv:1701.065300 2017.
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv. Layer by layer: Uncovering hidden representations in language models. arXiv preprint arXiv:2502.02011 2025.
Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Sho Takase and Shun Kiyono. Lessons parameter sharing across layers in transformers, arXiv preprint arXiv:2104.06022 2021.
Yi Tay, Mostafa Dehghani. Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Siamak Shakeri, Dara Bahri, Tal Schuster, et al. U12: Unifying language learning paradigms. arXiv preprint arXiv:2205.055
2022.
Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song. and Yasin Abbasi Yadkori. Hierarchical reasoning model. arXiv preprint arXiv:2506.21334 2025.
Johannes Welbl, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. arXiv preprint arXiv:1777.00 01777007000 2017.
Kevin Xu and Issei Sato. expressive power of looped transformers: Theoretical analysis and enhancement via timestep encoding. arXiv preprint arXiv:2410.01405,, 2024.
Kevin Xu and Issei Sato. To cot to loop? formal comparison between chain-of chain-of-thougl! offthhught and looped transformers. arXiv preprint arXiv:2505. 19245, 2025.
Jiewen Yang, Xingbo Dong, Liujun Liu, Chao Zhang, Jiajun Shen, and Dahai Yu. Recurring the transformer for video action recognition. In Proceedings of the IEEE/CVF Conference Com- puter Vision and Pattern Recognition, pp. 14063-14073, 2022.
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. arXiv preprint arXiv:2311.12424. 2023.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process. arXiv preprint arXhv:2407.20311 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma-
chine really finish your sentence? arXiv preprint arXiv:1905.078300 2019.
Alex L Zhang, Tim Kraska, and Omar Khattab. Recursive language models. arXiv preprint arXiv:25/2.24601, 2025.
Rui-Jie Zhu, Zixuan Wang, Kai Hua, Tianyu Zhang, Ziniu Li, Haoran Que, Boyi Wei, Zixin Wen, Fan Yin, He Xing, al. Scaling latent reasoning via looped language models. arXiv preprint arXiv:2510.25441, 2025.
14

---

Published as a conference paper at ICLR 2026
APPENDIX
A TRAINING AND IMPLEMENTATION DETAILS
Hardware and framework. All models trained on 4XH000 (80 GB) GPUs using the open- source NanoGPT training stack as a reference implementation.
Data and tokens. Unless otherwise specified, we run each experiment for 50,000 optimizer steps with global batch size of 48 sequences and block size 1024 (context length). This corresponds to approximately 25B training tokens in total.
Optimization. We AdamW with weight decay 2 10-1, cosine learning-rate decay (per NanoGPT), peak learning rate 1r = 6 x 10-4, minimum learning rate min.lr = x 10 and 4,000 warmup steps, which we found important for stability. Similar observations in Geip- ing et al. (2025), we occasionally observe training instabilities at depth; warmup and cosine decay mitigate these in practice. Unless noted, other optimizer and training defaults follow NanoGPT.
Model hy perparrameeess.. Following Saunshi et al. (2025), Theads = 32 for all L)) configurations. The soowwrd dimension dæ = 5120 with a standard two-layer GELU MLP. All normalizations RMSNorm. We learned positional embeddings added to token embeddings (NanoGPT default). We also tested RoPE[SSu et al.) (2024) and observed slightly better small-scale performance, but we use learned positions for simplicity and efficiency. We tie input and output embeddings (weight tying).
hidden size d 2048 and
LoopFormer conditioning. We use two embedding modules, one for normalized time t E [o, 1] and one for step size At E 1). Each maps scalar input to d-dimensional conditioning vector via (i) fixed sinusoidal Fourier features width D, = 256 and max period 10,000, followed by (ii) a 2-layer MLP with hidden size d and SiLU nonlinearity:
((T) = MLP HILP((co((TO))) ([cos(Tw)),sm))
cos(rwp, ,2); sin(Twp,,1))) E Rd
where wi = exp( k-1 log 10,000) for 1, .,..., D2/2. Given batchwise scalars and Al, we compute o(l) and es s(Al) and sum them obtain the per-iteration conditioning signal e e + + Rdd
D
Conditioning is applied inside each LoopFormer block via Adal.N-style modulator: small MLP takes and outputs 4d parameters, which we split into (amsa Amlp: Ymea, 7mlp). We use RMSNorm (with no learned affinity) before MHSA and FFN, and apply multiplicative scaling and residual gating as
x +- 2 + Cmma C MHSA (z) @ (1 %mssa)
x + x + amlp @ FFN FAFFFFNMMSSSE (FFM(NNNNom((1) @ (1 + 7mmp))
broadcast over the sequence length. The modulator SiLU followed by a linear layer with output size 4d, and is zero-initialized (weights and bias), ensuring the initial behavior matches the unmodulated backbone and that conditioning is learned stably.
TMLT baseline. For Modulladd Looped Transformers, we follow the authors' setup and condition each iteration on the loop index, implementing the timestep modulation as described in their paper.
B ADDITIONAL EXPERIMENTS AND ABLATIONS
B. COMPUTATIONAL OVERHEAD OF LOOPFORMER
Algorithm I uses dual-trajectonn objective: for every batch compute (i) the full L-loop trajec- tory for the main LM loss, and (ii) a sampled shorter M-loop trajectory for shorteut-consisteneyy This design adds overhead during training, however, during the inference, the computational
15

---

Published as a conference paper at ICLR 2026
Table 2: Perplexity and zero-shot reasoning for K=3 models. We compare a 3-layer vanilla Trans- former against LoopFormer evaluated at three inference budgets: (3 @ 2), (3 @ 4), and (3 & 8). LoopFormer matches vanilla performance at low budgets, and consistently improves with additional loops, substantially outperforming the vanilla baseline as compute increases, while retaining depth elasticity.
Perplesit
Tisks
Params/ FLOPs
Edi
OpenWelTeed
PIQA
SdQ
Base (Rxx 11
Bu/ 3x
12:95
30.23
29:10
59
24.7
26.47
26.4
59.85
2708
64.1
300
35.57
$1.45
40.93
LingFonncr Ours (3
36/5%
143
33.45
32.46
63
25,9
26.14
26.6
001
2574
586
29.39
35.26
50.
40.36
Lopplooner- Oues (3.9
1s/11c
11:12
15.02
34:21
611M
25
16.
11.7
53.47
43.31
LorFFoneer (3.9)))
Flops and inference times LoopFormer scale similar to that of a vanilla transformer. the fol- lowing we report the training overhead.
For clarity, we decompose the per-batch FLOPs into a loop-independent overhead plus loop- dependent term. Let Cio denote the FLOPs of the embedding / unembedding (non-loop compute), and let C₁ denote the FLOPs one loop through the shared K-block stack. Then running & loops costs
C(2) Cio &C1.
LoopFormer computes a full L-loop trajectory plus one sampled shorter trajectory of length M, where M Unif(I,..... 13. Thus the expected per-batch FLOPs are
C(L) + +MMIC(M) = (Cio + LC1) + (Cio + E(M]
Since M is uniform on II....,,I I 1), we have B[M] = L.2. Plugging in,
C(L) EM 20 + (L+ + 4 C₁ 2Cio + 4
Relative to the fixed-loop/vanilla training cost C(L) = Cio + LC1, the multiplicative overhead is
CLL) EM((M( 20, + 8/ Ci C(L)
Co + LC1
which, for the regimes considered in this work, is approximately 1.5x the FLOPs of fixed-loop training.
We confirm this empirically. Across settings, LoopFormer requires about 1.5x the training FLOPs of fixed-loop or vanilla transformer baseline trained under the same token budget and number of iterations. In wall-clock time, this corresponds to an approximate 1.3x slowdown in our setup (measured on 4xH100 GPUs with identical batch size, optimizer, and data). It is worth mentioning that other depth-elaste baselines studied in this work (e.g. TMLT-EE Base-Loop-EEs also have the same overhead as our LoopFormer.
Overall, LoopFormer pays a modest training overhead enable elastic-depth inference: a single Shared-parameter model that remains robust under aggressive truncation and continues to refine representations as more loops are allocated. The inference on the other has similar computational FLOPs to that of vanilla or fixed-loop baselines.
B.2 COMPARISON WITH VANILLA TRANSFORMERS UNDER A SIMILAR PARAMETER BUDGET
Here we compare LoopFormer to a vanilla Transformer under similar parameter budget. We fix the shared stack to K = 3 layers, and evaluate LoopFormer at different inference budgets, namely (3 @ 2). (3 @ 4), and (3 @ 8). Table reports perplexity and downstream language task accuracy for LoopFormer and a 3 layer vanilla Transformer with a similar number of parameters. We observe that at small compute (few loops), LoopFormer is par with the vanilla baseline, while increasing the loop budget leads to consistent improvements, allowing our depth-elaste model to substantially outperform the vanilla Transformer.
16

---

Published as a conference paper at ICLR 2026
Table Perplexity and zero-shot reasoning under FLOPs-matched training and inference. We com- pare (3 @ 8) looped models as well Base (24 @ 1). To match total training FLOPs, LoopFormer trained for 34k iterations (seeing ~8B fewer Pile tokens), while baselines are trained for 50k iter- ations. Despite the reduced token budget, LoopFormer matches TMLT and remains competitive on reasoning across budgets, while retaining depth-elastic behavior.
Perplesity
Languagy Tasks (Aruraиy))
Parans// FLOPs
File
Fine FebeWebEEde
OpenWebllest
COPA
HBS
LB
OBQA
PIOA
Race
SeQ
ARC
SIQA
WG
Avg Ace
34s
20.08
35.04
33.43
50.4
45.27
Base Loop (3 6s 8)
3x. 24k
10.91
24.53
24.53
61
30,46
34.68
27
63.22
28.5
63.7
31.71
38.43
49.8
42.88
TMLT (3 8)
3x/ 24x
10.38
22.87
21,99
63
32.34
39.06
27.8
63.11
29.67
09.8
31,67
36.95
$1.54
44.69
LoopPommer Ouns (3.2) 5)
31.96
31.68
B.3 COMPARISON UNDER FLOPS-MATCHED TRAINING
As discussed in lnluubeection B. elastic-depth training incurs approximately 1.5x the training FLOPs of vanilla or fixed-loop Transformers due to the additional sampled shortcut trajectory. To evaluate LoopFormer under FLOPP-matched training regime, we train LoopFormer model for 34k iter- ations (all other hyperparameters unchanged), resulting in roughly the same total training FLOPs as the baselines trained for 50k iterations. This compute-matcher LoopFormer therefore sees about 8B fewer Pile tokens during training. Table reports perplexity and zero-shot reasoning when both training and inference FLOPs are matched. Even under this reduced training budget, LoopFormer remains on par with TMLT and other baselines, while additionally providing depth-elaste inference.
B.4 PYTORCII PSEUDOCODE FOR LOOPFORMER
Algorithm presents self-contained PyTorch-style pseudocode snippet for the core LoopFormer shared stack. The module consists of LoopF o rme womerreeloo rBlock, which applies AdaLN-style modula- tion to the attention and MLP residual branches using the loop-conditioning vector e (derived from (1, At)), and SharedBlow ooc that stacks K such blocks to form the shared Transformer 4x.
Within each loop iteration, Sha redBl ock applies all blocks sequentially the token states x; this entire K-block stack is then reused across loops and repeated for A or L iterations during training or inference. This snippet clarifies how LoopFormer handles multiple blocks per loop and how the same shared stack supports variable compute budgets.
17

---

Published as a conference paper at ICLR 2026
Algorithm 3: PyTorch implementation of LoopFormer shared block
Inputs:
torch. Tensor [B, T, token states.
torch. Tensor [B, DJ loop-conditinniag vector.
config model config with n_enbd, etc.
depth int (K) number of distinct Trransorrmer blocks
class CocpEormerBlock (nn Modull))
def _iit inini_nssell config): super in_itt init_()
config.n_emiod
self.In_l nn. RMSNorm (d, deeeentwwse_afffneeeee self.atth Causalsel (config)
self 1n_ nn RMSNorm
seli.mip MLP (config)
Conditioning -> (gate_msa, gate_mlp, scale_msa, scale_mlp) self self..daal_moculatioo nn. Sequential nn. in.SiLU()
nn. (d, 4 d, bias=True),
AdaLN-2ero init: start from identity updates mn.init. zenitt (t..itt.orro__coff...dooo
[1]. nn. an.init. htt.::rrroo......caaaooaa.aaa... (1).cu]llaaioo
def (crwarddse1f, x:
6:
torch. Tensor:
gate_msa, gate_mip, scale_msa, scale_mip -
self.ada aaa_alll_mmouuii i (c) . chunk (4, dim=l)
Multi-heas attention branch x x gate msa, unsqueeze (1) self.attm(
self.in_1(s + (1.0 + scale_msa unw.uueeeus (1))
Feed-forward branch
x x gate_mlp........ unsqueeze (1) self.mip(
self.In_2(s (1.0 + scale_mlp mlp. unsqueeze (1))
return
class Shared81 ock (nn. Module)
def Liit (self, depth: int, config):
super . init_ ()
= nn ModuleeListt
[LoopFormerBlockKCC) (config) for in range (depth) ]
def forward (self, x: torch. Tensor, a: torch. Tensor) -> torch, Tensor: for blk in self.blocks: x blk (x, c)
return
18


Reviews:
[{"review_id": "bbq93B1zK1", "reviewer": "Reviewer_aD1e", "summary": "This paper presents a new model composed of k blocks each trained to be repeated up to L times. The architecture adds in each block a side information about the scheduling to visit [0, 1] with M steps, and modulate a scale and gate layer before/after the multi-head attention and FFN blocks.\n\nThe training procedure uses a loss that combines the cross-entropy of the full-compute (L iterations), the cross-entropy of a lesser compute (M iterations), and a L2 between the pre-readout activations of both.\n\nExperimental validation shows that this approach is not as good as a vanilla architecture for a given compute budget, but has a lesser parameter count.", "questions": "It is unclear in Fig 1 and in the loss definition of 3.3 how the method deals with multiple blocks. It looks as if there was only one there.\n\nRegarding the training cost and the overall performance, can you clarify in what regime you see the usefulness of this model?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The proposed architecture is reasonably simple and well motivated. The overall direction, transformer-like that can dynamically modulate compute, is important. The results showing the nice monotonicity of perplexity or accuracy vs. FLOPs (Fig 2) is great.", "weaknesses": "The experimental results are underwhelming, and maybe I missed the point, but it is unclear to me how this model improves wrt a vanilla transformer. The results are presented in such a way that the training budget is not equalized (if I am correct? I do not understand the first sentence of 4.1), the inference flops are, and the parameter counts are not, although this is generally not the limiting factor.\n\nTraining budgets should have been equalized (e.g. pick the nb of training iterations per model), and results should be summarized with scatter plots on e.g. an average accuracy vs. wall-clock time or flops frame.\n\nThe base model should be added as a point on Fig 2 (a) and (b)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u1wgxeZOjs", "reviewer": "Reviewer_exw9", "summary": "This paper tackles the limitation of looped Transformers, whose performance typically degrades when inference loops don't match the fixed number used in training. They propose LoopFormer which novelly condition each loop on both normalized time and step size to adapt to different computational depth. The model is trained on \"shortcut-consistency\" loss that forces shorter trajectories to match the representation of the full trajectory. They empirically show that the LoopFormer perform more robustly at smaller compute budgets than existing baselines.", "questions": "1. What is the practical training overhead (e.g., in FLOPs or wall-clock time) of the dual forward pass required for the consistency loss, compared to a standard Base-Loop or TMLT baseline?\n2. Given that Figure 5 shows significant performance variance based on the chosen step schedule Δ_M, have you explored methods for learning an optimal, budget-aware scheduling policy?\n3. The model is trained to interpolate (M <= L). Have you tested its ability to extrapolate to M > L loops? Does performance continue to improve, or does it diverge given the t=1 training target?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper tackles a clear, practical, and important problem. Enabling flexible, elastic compute in parameter-efficient models like looped Transformers is a highly valuable research direction.\n2. The paper is well written and easy to follow.\n3. The experiments are thorough and convincing. In addition to strong task performance, the authors provide a compelling explanation for why LoopFormer works by analyzing metrics like curvature, anisotropy, and CKA similarity. They demonstrate that baselines suffer from representational stagnation (flat metrics, high CKA), while LoopFormer's representations continue to evolve and refine with each loop. This analysis significantly strengthens the paper beyond just reporting better numbers.", "weaknesses": "1. The training procedure (Algorithm 1) requires two forward passes per batch (for the full and short trajectories) to compute the consistency loss. This appears to roughly double the training cost compared to a standard looped model. The paper mentions this as a limitation but does not quantify it. A brief analysis of the training FLOPs/time overhead vs. a Base-Loop or TMLT baseline would be valuable for assessing the practical trade-offs.\n2. The paper heavily emphasizes \"latent reasoning,\" using terms like \"thought trajectories.\" However, the benchmarks (COPA, PIQA, HellaSwag) are standard zero-shot LM evaluation tasks, not complex, multi-step algorithmic reasoning tasks. The model shows improved general performance (including perplexity), and the representation analysis shows refinement, not necessarily reasoning in a formal sense. So the terming could be a bit misleading to readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j2942wDN1I", "reviewer": "Reviewer_xXws", "summary": "This paper presents LoopFormer, which allows reasoning under variable compute budgets. It extends prior looped Transformers by introducing time- and step-size modulation, where each iteration receives sinusoidal embeddings of layer index and step size to dynamically modulate RMSNorm and residual scaling.\nA shortcut-consistency loss aligns representations across different loop lengths, enabling stable performance even with fewer inference steps.", "questions": "NA", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well written and easy to follow, with clear motivation and setups\n* The motivation is clearly presented, and the transition from fixed-depth looped Transformers to elastic-depth design feels natural.\n* Experiments are reasonably comprehensive, evaluating both variable loop lengths and the effect of the proposed shortcut-consistency loss.\n* The main claims are well supported", "weaknesses": "* The degree of novelty is not bad but moderate. While the proposed elastic-depth formulation and shortcut-consistency loss are well designed, they extend existing time-modulated looped Transformer frameworks rather than introducing a fundamentally new paradigm.\n* the paper does not provide theoretical intuition or analysis explaining why combining t and $\\Delta t$ through sinusoidal modulation is a good choice here", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4ElYTzhkRD", "reviewer": "Reviewer_pdxW", "summary": "This paper introduces LoopFormer, a looped Transformer architecture that can flexibly adapt its computational depth based on a user-specified budget, a feature termed \"elastic depth\". The model is designed to perform robustly across a range of loop iterations at inference time without needing to be retrained, addressing a key limitation of prior looped models.\n\nTraditional looped Transformers are trained and evaluated with a fixed number of loop iterations. This rigidity means they cannot adapt to variable compute budgets; their internal representations tend to collapse or stagnate when evaluated at depths different from their training configuration, leading to degraded performance.\n\nThe paper introduces a shortcut-consistency training scheme that enables compute-budgeted inference (elastic depth) without retraining the model. \nThe work demonstrates that naive early-exiting in looped models leads to representational collapse, where hidden states stagnate across iterations. In contrast, LoopFormer's representations continue to evolve, showing that it uses additional depth effectively for refinement.", "questions": "Your results show that performance is sensitive to the choice of the inference-time schedule. Could you provide a practical heuristic for selecting a high-performing schedule without an exhaustive search? Have you considered methods for learning an optimal, input-dependent schedule to automate this process? A response here would clarify if this is a minor tuning step or a significant practical hurdle.\n\nThe shortcut-consistency training algorithm appears to double the training compute by requiring two forward passes per batch . Could you provide an analysis of this trade-off? Specifically, how does LoopFormer compare against a baseline (e.g., TMLT) that is trained for twice as many steps using the same total compute budget? This would clarify whether the inference-time flexibility justifies the increased training cost.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper introduces the novel concept of \"elastic depth\" for looped Transformers, creatively adapting ideas from diffusion models to frame iterative refinement as a continuous-time \"thought trajectory\" . This unique problem formulation and synthesis of ideas results in a highly original approach to budget-conditioned reasoning.\n\nThe work is of high quality, featuring a rigorous experimental setup with strong, appropriate baselines and a comprehensive evaluation on both perplexity and a wide range of reasoning tasks . The claims are further substantiated by an in-depth representational analysis that uses multiple metrics to convincingly demonstrate that LoopFormer avoids the representational collapse that plagues naive adaptive methods.\n\nThe paper's contribution is significant because it provides a practical solution to the important problem of adaptive computation, making efficient looped architectures more versatile and deployable.", "weaknesses": "The shortcut-consistency training algorithm requires two forward passes per batch (one for the full trajectory and one for a shortcut), which effectively doubles the training compute compared to baseline looped models . The paper acknowledges this \"added training overhead\" but does not analyze the trade-off. A key missing experiment is a comparison against a baseline model trained for twice as many steps, which would clarify if the inference-time flexibility is worth the significant increase in training cost.\n\nThe model's performance is highly sensitive to the choice of the inference-time step schedule, yet this schedule must be selected manually by the user1111. This introduces a required hyperparameter tuning step that complicates deployment and may prevent users from achieving the model's optimal performance for a given budget. The work would be stronger if it included a method to learn or predict an optimal schedule automatically.\n\nThe experiments are conducted on ~1B parameter models, and it is unclear if the training dynamics and performance benefits will scale to much larger, state-of-the-art foundation models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)