LoopREG: Representation Supervision Enables Test-Time Recurrent Depth in Diffusion Transformers

Zipeng Wang; Anping Xiong; Jing Zhang; Xiaoya Zhu; Yi Jiang; Yanbing Xiao; Zige Zeng; Zezhou Chen

cs.CVOct 13, 2026

LoopREG: Representation Supervision Enables Test-Time Recurrent Depth in Diffusion Transformers

Authors: Zipeng Wang, Anping Xiong, Jing Zhang, Xiaoya Zhu, Yi Jiang, Yanbing Xiao, Zige Zeng, Zezhou Chen, +12 more

Organizations: Chongqing University of Posts and Telecommunications · China United Network Communications Group Co., Ltd. · University of Electronic Science and Technology of China · Dongbei University of Finance and Economics

Abstract

Reusing Transformer blocks offers a way to increase diffusion-model computation at test time, but naive recurrence can degrade generation beyond its training depth. We introduce LoopREG, a representation-supervised recurrent diffusion Transformer that repeatedly applies a shared block at the state supervised by a pretrained visual encoder. Alignment is applied after the final recurrent visit during training; inference extends the same learned operator to additional visits. On ImageNet 256 x 256, LoopREG reduces SiT-B/2 FID from the local REG baseline's 15.08 to 4.91 at 400K training steps, a 67.41% reduction. Within a checkpoint trained with two visits, one extra inference visit improves FID from 14.71 to 5.34 with unchanged weights. The benefit persists on SiT-L/2 and SiT-XL/2. Controlled interventions show that useful recurrence follows the location of representation supervision, while REPA alignment produces a smaller extrapolation benefit. These results make representation-supervised recurrence a practical route to improving diffusion generation through additional inference depth.

Figures & tables

Figure 1: LoopREG reduces REG’s FID by 67.41% on SiT-B/2. At 400K training steps, horizontal references show published REPA (24.40; Wu et al., 2025 ) and local REG (15.08). LoopREG (

K_{\mathrm{train}}=3

) reaches FID 4.91 at

K_{\mathrm{infer}}=5

(star). Figure 3 gives full sweeps. Local evaluations: EMA, 50K samples, 250 denoiser calls, CFG=1.

Figure 2: LoopREG: representation-supervised shared recurrence. SiT-B/2 is shown. (a) One clean image supplies both the VAE latent, noised after encoding, and frozen DINOv2 targets. Orange denotes CLS/semantic tokens; blue denotes image/patch tokens. Alignment reads the final recurrent state through a projector. (b) Shared weights are reused beyond the training horizon, illustrated by $K_{\mathrm{train}}=2$ and $K_{\mathrm{infer}}=3$ . Sampling starts from noise.

Model	Steps	$K_{\mathrm{train}}$	$K_{\mathrm{infer}}$	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	P $\uparrow$	R $\uparrow$
SiT-B/2 (local)	400K	1	1	35.25	6.65	42.14	0.530	0.650
REPA-B/2 *	400K	–	–	24.40	6.40	59.90	0.590	0.650
REG-B/2 (local)	400K	1	1	15.08	6.73	95.53	0.649	0.640
LoopSiT-B/2	400K	2	2	34.45	6.43	43.01	0.533	0.639
LoopSiT-B/2	400K	2	3	46.41	10.01	34.28	0.464	0.618
LoopREPA-B/2	400K	2	2	21.93	6.37	66.58	0.601	0.651

Table 1: Generation quality across model scales. Horizontal rules separate backbones; spacing separates method families. Shaded rows are baselines; an asterisk denotes published results from Wu et al. (2025) , Table 9 (SiT/REPA attributed there to Ma et al. (2024) ; Yu et al. (2025) ). Training steps are shown explicitly; the published REPA-XL/2 reference uses 4M steps. All other rows are local evaluations. Bold marks the best displayed metric within each multi-row checkpoint group; best depths are selected within the reported sweeps. P/R: precision/recall. Appendix A gives provenance, including the L/2 server summary.

Figure 3: Training horizon shifts the useful inference horizon. Dotted lines mark training horizons; stars and labels mark the best observed FID. The dashed horizontal line shows the local REG baseline. All three training horizons produce large improvements over this baseline, with optima moving to later recurrent depths. The logarithmic vertical axis preserves the severe under-roll failures.

Backbone	Blocks/call	Added blocks (%)	FID ( $K_{\mathrm{infer}}$ : $2\to 3$ )	FID reduction (%)
SiT-B/2	13 $\to$ 14	7.69	14.71 $\to$ 5.34	63.67
SiT-L/2	25 $\to$ 26	4.00	4.23 $\to$ 3.21	24.13
SiT-XL/2	29 $\to$ 30	3.45	3.38 $\to$ 2.84	16.10

Table 2: Return from one extra shared-block application. All models train with

K_{\mathrm{train}}=2

for 400K steps. Added block executions are relative to matched

K_{\mathrm{infer}}=2

; FID reductions compare the same checkpoint at

K_{\mathrm{infer}}=2

and

3

Figure 4: Extra recurrent computation improves the optimization-step–FID curve. Left: fixed

K_{\mathrm{infer}}=2,3,4

curves from 50K steps along one XL training run, with published REG reference points. Right: 400K results;

K_{\mathrm{infer}}=2,3

share one checkpoint, while REG is a published reference. Lines connect evaluated checkpoints; 210K/220K include

K_{\mathrm{infer}}=3

only.

Figure 5: Useful recurrence follows representation supervision in the tested interventions. Left: fixed supervision at depth four, varied loop block. Right: supervision and recurrence moved together. Both panels reuse the d4/B4 anchor. Each model trains at

K_{\mathrm{train}}=2

Figure 6: Useful extra depth improves representation geometry. Cells show measured values; darker shading indicates a better rank within each row , and outlines mark that row’s optimum. Blue depth labels mark $K_{\mathrm{train}}$ . Patch CKA, spatial correlation, and FID prefer the same extra depth for each LoopREG checkpoint. Native projector cosine instead peaks at the trained depth. Representation measurements use the final N=1,024 probe (CKA subset 256), averaged over five diffusion times; FID uses the separate 50K generation evaluation.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Transient optima and horizon-dependent latent progress. A contractive trajectory can pass through a task-preferred state before reaching its fixed point. Under a fixed total-progress assumption, a larger training horizon corresponds to smaller recurrent steps and a later optimal inference horizon.

Condition	$K_{x}$	$K_{c}$	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	Calls
$K_{\mathrm{train}}=2$ , generation seed 2; 50,000 images per condition
A	2	2	14.709	6.702	97.933	0.6473	0.6326	250
B	3	2	7.533	6.236	142.720	0.7455	0.5599	500
C	2	3	9.362	6.097	132.561	0.7021	0.5974	500
D	3	3	5.339	5.783	185.375	0.7934	0.5143	250
$K_{\mathrm{train}}=2$ , generation seed 3; 50,000 images per condition

Appendix

Table 3: All paired 50K output-depth evaluations. Bold marks the best metric within each checkpoint and generation seed. Calls counts denoiser evaluations per image: A/D use 250 and B/C use 500, all with 250 state updates.

Fixed supervision at depth four
Sup. depth	Loop block	FID $K_{\mathrm{infer}}=2$	FID $K_{\mathrm{infer}}=3$	FID $K_{\mathrm{infer}}=4$	Gain (%)	R at $K_{\mathrm{infer}}=3$
4	2	13.92	32.56	58.85	-133.93	0.676
4	4	14.71	5.34	6.20	63.67	0.508
4	6	14.53	31.11	69.58	-114.13	0.611
4	8	14.18	29.97	59.55	-111.33	0.618
Supervision and recurrence moved together

Appendix

Table 4: Site/depth interventions. Bold marks the lowest FID within each row. Gain is the relative FID reduction from

K_{\mathrm{infer}}=2

K_{\mathrm{infer}}=3

; negative values indicate degradation. The last column reports recall at the extrapolated depth. The d4/B4 anchor is repeated across groups. At d8/B8, FID improves while sFID worsens.

Figure 8: Fixed and random training horizons. The full observed depth range exposes under-roll failures, finite optima, and recall differences. The dotted recall line is the local REG baseline. Random

\{3,4,5\}

and

\{1,2,3,4,5\}

are still improving in FID at

K_{\mathrm{infer}}=8

, the evaluated boundary.

$K_{\mathrm{train}}$ / Policy	$K_{\mathrm{infer}}$	Blocks	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	P $\uparrow$	R $\uparrow$
2 / fixed	1	12	151.9461	50.3472	6.8845	0.1480	0.1539
2 / fixed	2	13	14.1520	6.6645	99.3560	0.6541	0.6383
2 / periodic	2	13	14.1520	6.6645	99.3560	0.6541	0.6383
2 / periodic	4	15	12.2748	8.2114	137.2766	0.6315	0.5338
2 / periodic	6	17	61.1025	19.6300	39.9223	0.2887	0.4964
2 / periodic	8	19	147.5555	52.2028	10.3376	0.0813	0.2113

Appendix

Table 5: Unshared controls at 400K. Every evaluated depth is included, grouped by training horizon and inference policy. The fixed and periodic pair share their matched-depth anchor. Bold marks metric optima within each group.

Model	$K_{\mathrm{train}}$	NMSE $K_{\mathrm{infer}}$	FID $K_{\mathrm{infer}}$	NMSE train	NMSE FID-best
LoopSiT-x2	2	2	2	0.46074	–
LoopREG-x2	2	2	3	0.45706	0.45898
LoopREG-x3	3	3	5	0.45684	0.46031

Appendix

Table 6: The separate N=256 teacher-forced velocity probe, averaged over five noise times. The FID-best horizon differs from the NMSE-best horizon for both LoopREG checkpoints. “–” means that the two horizons coincide.

Forecast	MAE, 200K	MAE, 400K	Depth hits, 400K	Selected $(K_{\mathrm{train}},K_{\mathrm{infer}})$
Component exponential	0.638	3.262	1/3	(3,5)
Total-FID exponential	0.652	3.299	1/3	(3,5)
Latest observation (150K)	2.237	5.258	1/3	(3,5)

Appendix

Table 7: Early-curve forecasts on 5K-sample FID. MAE averages 18 candidates. Depth hits count correct inference-depth selections across three trained models.

Method	Direction hits	Depth hits	Change MAE	Mean selection loss
Quadratic feature extrapolation	3/8	0/3	8550.500	8.744
Quadratic, half step	3/8	0/3	82100.946	8.744
Linear feature extrapolation	3/8	0/3	242.900	8.744
Zero change	0/8	0/3	3.560	8.744
$K_{\mathrm{train}}+1$	–	1/3	–	0.405
Observed calibration anchors	–	3/3	–	0.000

Appendix

Table 8: Local feature extrapolation on 4,096-sample FID. Direction hits compare signs over eight edges. Selection loss is measured FID at the selected depth minus the best measured FID; depths are chosen by accumulated predicted changes, with smaller-depth tie-breaking.

Steps	Tested $K_{\mathrm{infer}}$	FID at 2	Best $K_{\mathrm{infer}}$	Best FID	Gain (%)	Recall
50K	2,3,4,5	25.1962	4	11.8154	53.11	0.3407
100K	2,3,4,5	8.3360	3	4.7552	42.96	0.4682
200K	2,3,4,5	4.6383	3	3.4942	24.67	0.5010
220K	3	–	–	3.3675	–	0.5180
400K	1,2,3,4,5	3.3823	3	2.8376	16.10	0.5490

Appendix

Table 9: Representative XL training checkpoints. Best depths are minima over the listed sets. The 220K checkpoint supplies a single-depth observation. All metrics at the final 400K checkpoint appear in Appendix H .

Figure 9: Representation geometry and native alignment along recurrent depth. FID, patch-mean CKA/spatial correlation, and native projector cosine use aligned depth axes. Geometry is averaged over five noise times from the final 1,024-image probe (CKA subset: 256). These are teacher-forced hidden trajectories; FID comes from separate full-generation evaluations. Vertical scales differ across panels.

$K_{\mathrm{infer}}$	Blocks	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	P $\uparrow$	R $\uparrow$
LoopSiT-B/2, B4, $K_{\mathrm{train}}=2$
1	12	52.0813	11.3411	28.9425	0.3676	0.6608
2	13	34.4524	6.4313	43.0118	0.5334	0.6387
3	14	46.4096	10.0098	34.2806	0.4637	0.6182
4	15	78.8033	19.1662	20.7103	0.3153	0.5667
5	16	118.9361	31.6541	12.4613	0.1959	0.4317

Appendix

Table 19

$K_{\mathrm{infer}}$	Blocks	FID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	P $\uparrow$	R $\uparrow$
LoopREG-L/2, d=8, B8, $K_{\mathrm{train}}=2$
1	24	59.4139	12.6041	26.0546	0.3214	0.7664
2	25	4.2286	5.0924	170.3622	0.7516	0.6420
3	26	3.2082	4.4353	262.6577	0.8593	0.5201
4	27	5.4971	4.5898	291.8665	0.8827	0.4438
5	28	7.0317	4.9216	267.8294	0.8450	0.4206

Appendix

Table 20

Explore similar work

LiFT: Loop Flow Transformers
Oct 4, 2026Mohammad Mahdi Derakhshani, Pedro M. P. Curvo, Gertjan J. Burghouts +2Diffusion Transformer Recurrent Transformers
VISLab · University of Amsterdam · TNO, Intelligent Imaging +1
Learning Representations to Guide Pixel Diffusion Transformers
Oct 9, 2026Zihao Pan, Mengting Chen, Cheng Shi +6Diffusion Transformer Pixel-Space Diffusion Models
Sun Yat-sen University · Alibaba Group · The University of Hong Kong +1
Looped Diffusion Transformer
Sep 30, 2026Yong Xien Chng, Tianyi Chen, Wenwen Tong +7Diffusion Transformer Recurrent Transformers
SenseTime Research · LeapLab, Tsinghua University · Nanyang Tehnological Univesity

Oct 4, 2026J/K move · Enter open · S save

LiFT: Loop Flow Transformers

Mohammad Mahdi Derakhshani, Pedro M. P. Curvo, Gertjan J. Burghouts, Jan-Willem van de Meent+1

VISLab · University of Amsterdam · TNO, Intelligent Imaging+1

Diffusion Transformer Recurrent Transformers Diffusion-Based Image Generation

Open paper AskPDF arXiv

We introduce Loop Flow Transformers (LiFT), a family of looped generative models that scales computation by repeatedly applying a shared Diffusion Transformer (DiT) core, with only light changes to the standard architecture. Rather than asking every recurrent step for the final prediction, LiFT trains each step with a single regression target: a point on a straight path from the model's initial estimate to the flow-matching target. Because we index these targets by a continuous depth coordinate, a trained model can loop far beyond its training depth with no retraining, early exits, or other modifications. In our experiments, these longer rollouts improve generation, so inference computation can grow without adding parameters. On ImageNet at 256x256, LiFT-L/2 achieves an FID 3.34 points lower than our dense DiT-XL/2 baseline while using approximately 60% fewer parameters, 32% fewer training FLOPs, and 52% fewer inference FLOPs.

Abstract

Explore similar work

LiFT: Loop Flow Transformers

Learning Representations to Guide Pixel Diffusion Transformers

Looped Diffusion Transformer