Looped Transformers as Optimizers

Yulong Huang; Chen Jiang; Zhanpeng Zhou; Hongtao Zhang; Tianyu Li; Tianyu He; Xiangyu Zhang; Bojun Cheng

cs.LGSep 29, 2026

Looped Transformers as Optimizers

Authors: Yulong Huang, Chen Jiang, Zhanpeng Zhou, Hongtao Zhang, Tianyu Li, Tianyu He, Xiangyu Zhang, Bojun Cheng

Organizations: HKUST (GZ) · StepFun · SJTU · UCAS ·

Abstract

Looped Transformers provide a parameter-efficient approach to depth scaling by repeatedly applying shared Transformer blocks. Recent reasoning models have likewise highlighted the value of scaling test-time computation through longer computation trajectories. However, the principles for designing effective loop transitions remain poorly understood. We view the looped hidden state as a fast weight that is updated throughout the depth. We formulate loop transitions as local gradient-based updates, with recurrent blocks predicting implicit targets at each depth. Our framework derives loop transitions in closed form from a projection, a local objective and an optimizer update rule. Mapping representative loop transitions into this framework reveals mismatches between their transitions and projections. We first align the input maps of existing transitions. We then derive OperLoop, which combines explicit weight decay, adaptive step size and a delta objective. The aligned variants reduce training loss and improve average commonsense accuracy. OperLoop improves average generative performance over the compared looped and non-looped baselines under matched training FLOPs. These results support the framework's usefulness for loop design. We extend the analysis to additional loop models and outline a roadmap for future loop transition design.

Figures & tables

Figure 1: Optimization view of Looped Transformers. (a) A standard Transformer uses independently parameterized blocks. (b) A middle-looped Transformer repeatedly applies shared blocks; the Prelude and Coda may be identity maps in a fully looped model. (c) The recurrent state is treated as a fast weight. An input map $\mathbf{H}_{l}^{\rm in}$ reads out the state, and the shared blocks predict an implicit target. The local objective and optimizer update rule determine the state update. Setting $\mathbf{H}_{\rm in}=\mathbf{I}$ , $\lambda=\eta=1$ and $f^{\prime}(\mathbf{o}_{l})=-\operatorname{RecurrentBlocks}(\bm{o}_{l})$ recovers the vanilla loop in (b).

Models	Ouro ( Zhu et al., 2025 )	Huginn ( Geiping et al., 2025 )	Parcae ( Prairie et al., 2026 )	HyperLoop ( Zeitoun et al., 2026 )
Original loop transitions
Loop	$\begin{aligned} \bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})\end{aligned}$	$\begin{aligned} \bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})\end{aligned}$	$\begin{aligned} &\bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})\end{aligned}$	$\begin{aligned} &\mathbf{Y}_{l+1}=\mathbf{H}^{\mathrm{res}}_{l}\mathbf{Y}_{l}+\\ &(\mathbf{H}^{\mathrm{post}}_{l})^{\top}\operatorname{Blocks}_{\theta}(\bm{o}_{l})\end{aligned}$
Projection	$\bm{o}_{l}=\bm{y}_{l}$	$\bm{o}_{l}=\bm{y}_{l}+\mathbf{e}$	$\bm{o}_{l}=\bar{\mathbf{A}}\bm{y}_{l}+\bar{\mathbf{B}}\mathbf{e}$	$\bm{o}_{l}=\mathbf{H}^{\mathrm{pre}}_{l}\mathbf{Y}_{l}$
Optimization mapping to $\mathbf{W}_{i+1}=(1-\lambda\eta)\ \mathbf{W}_{i}-\eta\ (\bm{x}^{\top}f^{\prime})$
$\lambda$	$1$	$1$	$1$	$\mathbf{I}-\mathbf{H}^{\mathrm{res}}_{l}$
$\eta$	$1$	$1$	$1$	$1$

Table 1: Comparison of loop transitions under the optimizer update rule in Eq. 7 . The lower rows identify the decay coefficient (

\lambda

), step size (

\eta

), input map (

\bm{x}

), output error (

f^{\prime}

), and recurrent state (

\mathbf{W}

). Red marks mismatches between the update’s implied input map and the projection’s map. Aligning the transition with the input map used in the projection improves performance (§ 5.1 ).

Figure 2: Training-loss comparison for optimizer-aligned loop transitions. The left two panels show training-loss trajectories for Parcae Loop3 and HyperLoop Loop3, together with their input-map-aligned variants. The right two panels show the corresponding loss differences relative to Vanilla Loop3, defined as

\Delta\mathcal{L}=\mathcal{L}_{\mathrm{variant}}-\mathcal{L}_{\mathrm{Vanilla}}

. Negative values indicate lower loss than Vanilla Loop3.

Model	Param.	FLOPs	PPL $\downarrow$		Commonsense Acc. (%) $\uparrow$
Model	Param.	FLOPs	Lamb.	Wiki.	Lamb.	ARCc n	ARCe n	HellaS. n	WinoG.	PIQA n	Avg.	$\Delta$ Avg.
Baseline ${}_{\texttt{[18L]}}$	$1\times$	$1\times$	6.56	13.08	60.16	40.96	71.13	64.78	59.51	76.61	62.19	–
Parcae ${}_{\texttt{[4L-4L\texttimes 3-2L]}}$	$0.56\times$	$1\times$	7.71	13.98	57.44	38.57	65.91	62.00	56.12	74.81	59.14	–
Parcae (Aligned)	$0.56\times$	$1\times$	7.41	13.93	58.59	37.71	65.53	62.09	56.91	74.92	59.29	+0.15
HyperLoop ${}_{\texttt{[4L-4L\texttimes 3-2L]}}$	$0.56\times$	$1\times$	7.68	13.96	57.46	38.40	63.51	61.41	57.77	75.03	58.93	–
HyperLoop (Aligned)	$0.56\times$	$1\times$	7.35	13.93	58.63	35.15	66.67	62.83	58.56	74.97	59.47	+0.54

Table 2: Effect of input-map alignment on Parcae and HyperLoop in language modeling and commonsense evaluation. The looped variants differ only in their transitions. Parameters and training FLOPs are relative to the 18-layer baseline. The subscript

n

denotes normalized accuracy.

\Delta

Avg. is the improvement over the corresponding unaligned model, in percentage points.

		300B Training Tokens, 4k Sequence Length				500B Training Tokens, 4k Sequence Length
Category	Dataset	Base ${}_{\texttt{[18L]}}$	Loop ${}_{\texttt{[4L-4L\texttimes 3-2L]}}$			Base ${}_{\texttt{[32L]}}$	Loop ${}_{\texttt{[4L-8L\texttimes 3-4L]}}$
( Metrics )		5B MoE	Parcae	HyperLoop	OperLoop	10B MoE	Parcae	HyperLoop	OperLoop
	Parameters	$1\times$	$0.56\times$	$0.56\times$	$0.56\times$	$1\times$	$0.48\times$	$0.48\times$	$0.48\times$
	FLOPs	$1\times$	$1\times$	$1\times$	$1\times$	$1\times$	$1\times$	$1\times$	$1\times$
( PPL $\downarrow$ )	Wikitext	13.08	13.98	13.96	13.90	11.06	11.76	11.87	11.74
( PPL $\downarrow$ )	Lambada	6.56	7.71	7.68	7.21	4.82	5.20	5.32	5.14

Table 3: Main results on commonsense, math, code, and reasoning benchmarks. Parameters and training FLOPs are relative to the corresponding non-looped baseline. Bold and underlined values indicate the best and second-best results among looped models at each scale, respectively. The subscript

n

denotes normalized accuracy, and

*

denotes CoT evaluation. H.Eval and H.Eval + denote HumanEval and HumanEval+, respectively, both report pass@1 with greedy decoding.

Figure 3: Scaling loop from three (Loop3) to six (Loop6).

			PPL $\downarrow$		LL Acc $\uparrow$	EM Acc $\uparrow$
Model	Params.	FLOPs	Wiki.	Lamb.	Com. Avg.	GSM8K	GSM8K ∗	H.Eval	H.Eval +	MMLU	BBH ∗	EM Avg.	$\Delta$
Baseline (5B, 18L)	$1\times$	$1\times$	13.08	6.56	62.19	15.77	17.44	21.95	18.90	37.52	31.75	23.89	-
HyperLoop (Loop 3)	$0.56\times$	$1\times$	13.96	7.68	58.93	12.89	14.03	18.90	16.46	42.71	31.05	22.67	-1.22
OperLoop (Loop 3)	$0.56\times$	$1\times$	13.90	7.21	59.77	15.85	17.13	23.78	20.12	44.41	32.18	25.58	+1.69
HyperLoop (Loop 6)	$0.56\times$	$1.67\times$	13.35	6.51	61.39	22.44	27.14	24.39	20.73	48.25	35.54	29.75	+5.86
OperLoop (Loop 6)	$0.56\times$	$1.67\times$	13.16	6.39	61.72	24.87	31.08	27.44	21.95	50.15	35.89	31.90	+8.01

Table 4: Downstream evaluation at Loop3 and Loop6. Models are trained separately at each loop count. Com. Avg. denotes average commonsense accuracy. EM Avg. is the arithmetic mean of the six reported math, code, and reasoning scores, and

\Delta

is its difference from the baseline.

Variant	Delta Obj.	Step Size Schedule	Wiki. PPL $\downarrow$	Lamb. PPL $\downarrow$	LL Avg. $\uparrow$	EM Avg. $\uparrow$
OperLoop (Ours)	$\checkmark$	Causal	13.90	7.21	59.77	25.58
w/o Delta Obj.	$\times$	Causal	13.92	7.25	59.36	24.84

Table 5: Delta objective ablation at three loops with the causal step-size schedule fixed.

Variant	Delta Obj.	Step Size Schedule	Wiki. PPL $\downarrow$	Lamb. PPL $\downarrow$	LL Avg. $\uparrow$	EM Avg. $\uparrow$
OperLoop (Ours)	$\checkmark$	Causal	13.90	7.21	59.77	25.58
LR w Non-causal	$\checkmark$	Non-Causal	13.93	7.15	59.80	24.02
w/o LR	$\checkmark$	$\times$	13.95	7.48	59.38	24.40

Table 6: Step-size schedule ablation at three loops with the delta objective fixed.

		$\mathbf{W}_{i+1}=(1-\eta\lambda)\,\mathbf{W}_{i}-\eta\,(\bm{x}^{\top}f^{\prime}(\bm{o}_{l}))$
	Original Transition	$\bm{o}_{l}$	$\lambda$	$\eta$	$\bm{x}$	$\mathcal{L}$
Vanilla	$\bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})$	$\bm{y}_{l}$	$1$	$1$	$\mathbf{I}$	$-\langle\bm{t}_{l},\bm{o}_{l}\rangle$
Huginn ( Geiping et al., 2025 )	$\bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})$	$\bm{y}_{l}+\mathbf{e}$	$1$	$1$	$\mathbf{I}$	$-\langle\bm{t}_{l},\bm{o}_{l}\rangle$
Parcae ( Prairie et al., 2026 )	$\bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})$	${\color[rgb]{0.7266,0.1484,0.1328}\bar{\mathbf{A}}}\bm{y}_{l}+\bar{\mathbf{B}}\mathbf{e}$	$1$	$1$	${\color[rgb]{0.7266,0.1484,0.1328}\mathbf{I}}$	$-\langle\bm{t}_{l},\bm{o}_{l}\rangle$
Full-Bandwidth ∗ ( Wang et al., 2026 )	$\bm{y}_{l+1}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})$	${\color[rgb]{0.7266,0.1484,0.1328}\sigma(\mathbf{W}^{\rm g}\bm{e}_{l})\odot\mathbf{W}^{\rm u}}\bm{y}_{l}$	$1$	$1$	${\color[rgb]{0.7266,0.1484,0.1328}\mathbf{I}}$	$-\langle\bm{t}_{l},\bm{o}_{l}\rangle$
RecurrentGPT ( Hegazy et al., 2026 )	$\bm{y}_{l+1}=\bm{g}_{l}\odot\bm{y}_{l}+(1-\bm{g}_{l})\odot\operatorname{Blocks}_{\theta}(\bm{o}_{l})$	${\color[rgb]{0.7266,0.1484,0.1328}\mathbf{W}}\ [\bm{y}_{l}+\bm{\epsilon},\bm{e}]$	$1$	$1-\bm{g}_{l}$	${\color[rgb]{0.7266,0.1484,0.1328}\mathbf{I}}$	$-\langle\bm{t}_{l},\bm{o}_{l}\rangle$

Table 7: Extended comparison of loop transitions under the optimization framework. Recurrent states are treated as fast weights corresponding to

\mathbf{W}

. Vector and expanded states are denoted by

\bm{y}_{l}\in\mathbb{R}^{d}

and

\mathbf{Y}_{l}\in\mathbb{R}^{r\times d}

, respectively. Shared blocks predict the implicit target

\bm{t}_{l}=\operatorname{Blocks}_{\theta}(\bm{o}_{l})

. Red highlights input-map mismatches under the chosen optimizer update rule. Here,

\bm{e}

is the Prelude output,

\bm{\epsilon}

is a noise vector, and

{\ast}

denotes token shifts along the context. For training-free Recirculation, we compare only its transition. The Vanilla category includes Universal Transformers ( Dehghani et al., 2019 ) , MoEUT ( Csordás et al., 2024 ) , Recursive Transformer ( Bae et al., 2025 ) and Ouro ( Zhu et al., 2025 ) .

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Setting	Value
Model family	Architecture	Decoder-only Transformer
Numerics	Parameter / router / LM-head precision	BF16 / FP32 / FP32
Representation	Hidden size / vocabulary size	$1024$ / $128{,}896$
Attention	Heads / head dimension / GQA groups	$16$ / $128$ / $4$
Context	Maximum sequence length	$4096$
Attention pattern	Repeated every four layers	Sliding–Sliding–Sliding–Full

Appendix

Table 8: Architecture configuration used in the experiments. The sliding–sliding–sliding–full pattern repeats every four layers.

Setting	Prelude $i$	Shared $K$	Loops $R$	Coda $j$	Effective Depth	Distinct Blocks
18-layer comparison (Loop3)	4	4	3	2	18	10
32-layer comparison (Loop3)	4	8	3	4	32	16
Compute scaling (Loop6)	4	4	6	2	30	10

Appendix

Table 9: The number of distinct blocks is

i+K+j

, and the effective depth is

i+RK+j

Category	Setting	Value
Objective	Training objective / loss	Causal LM / vocabulary-parallel cross-entropy
Batching	Global batch / micro batch	$1024$ / $8$ sequences
Batching	Nominal tokens per update	$4{,}194{,}304$
Precision	Training parameter precision	BF16
Optimizer	Muon momentum / Nesterov	$0.95$ / enabled
Optimizer	Newton–Schulz variant / iterations	Polar Express / $6$

Appendix

Table 10: Training configuration. Nominal token counts are computed from the sequence length, global batch size, and number of updates.

Fast Weights

See all 8 papers

DeepLoop: Depth Scaling for Looped Transformers
Jul 15, 2026Shuzhen Li, Yifan Zhang, Jiacheng Guo +2Residual Learning Recurrent Transformers
Princeton University · University of California, Los Angeles
Loop- $μ$ P: Residual Scaling and Hyperparameter Transfer in Looped Transformers
Jun 16, 2026Shaowen Wang, Bingrui Li, Tansheng Zhu +4Residual Learning Hyperparameter Transfer
Tsinghua University · ByteDance Seed · M-A-P
FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates
Sep 24, 2026Wanqi Yang, Shiwei LiuEfficient Transformer Inference Recurrent Transformers
ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center

Jul 15, 2026J/K move · Enter open · S save

DeepLoop: Depth Scaling for Looped Transformers

Shuzhen Li, Yifan Zhang, Jiacheng Guo, Quanquan Gu+1

Princeton University · University of California, Los Angeles

Residual Learning Recurrent Transformers

Open paper AskPDF arXiv

Looped Transformers scale sequential computation by applying a compact stack of physical blocks for multiple rounds, increasing unrolled depth without increasing stored parameters. This reuse changes the residual-scaling problem: in an untied Transformer, each residual branch receives and applies its own parameter update, whereas in a looped Transformer one shared update aggregates gradients from repeated visits and is read back by those same visits in the next linearized forward pass. We formalize this tied-depth effect through a first-order perturbation bound controlled by a visit-alignment coefficient

κ_R

. The bound recovers the DeepNorm exponent when visits decorrelate, but in the conservative aligned regime it requires the exponent to increase from

1/4

1/2

as loop count grows at fixed physical depth. The resulting method, DeepLoop, keeps the Post-LN DeepNorm architecture and sets

α=(2N)^{1/2}

and

β=(8N)^{-1/2}

for unrolled depth

N

. On GPT-style looped language models at GPT-2 small and GPT-2 medium scale, DeepLoop is neutral when no physical block is revisited and improves validation loss and downstream accuracy once recurrent depth is activated. These results show that stable recurrent depth requires residual scaling rules that account for parameter visits, not only nominal layer count.

Abstract

Fast Weights

Explore similar work

DeepLoop: Depth Scaling for Looped Transformers

Loop-μμμP: Residual Scaling and Hyperparameter Transfer in Looped Transformers

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

Loop- $μ$ P: Residual Scaling and Hyperparameter Transfer in Looped Transformers