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

Oct 4, 2026

LiFT: Loop Flow Transformers

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.

Mohammad Mahdi Derakhshani, Pedro M. P. Curvo, Gertjan J. Burghouts +2

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

Diffusion Transformer Recurrent Transformers Diffusion-Based Image Generation

Ask PDF

Oct 9, 2026

Learning Representations to Guide Pixel Diffusion Transformers

Visual representations are largely developed for understanding, yet what makes them useful for generation remains unclear. We study this question in pixel diffusion transformers and uncover two key insights. First, better representation quality does not necessarily translate into greater generative utility, as externally aligned semantics can be progressively lost during pixel synthesis. Second, architectural separation with specialized pixel decoders mainly redistributes semantic and spatial modeling, without necessarily yielding representations better suited to generation. These insights suggest that a useful generative representation must be both properly formed and actively exploited. Motivated by these insights, we introduce Representation Guidance (RG), which directly forms and exploits internal representations while retaining a plain Transformer architecture. A low-dimensional bottleneck forces generation-relevant information into a compact representation, which is then explicitly used to guide decoder generation by modifying the flow-matching target according to its induced prediction offset. RG learns clearer semantic and spatial structure while scaling efficiently in both parameters and computation. On ImageNet 256x256, RG-L/16 with only 459M parameters outperforms substantially larger pixel diffusion transformers, while RG-XL/16 achieves an FID of 1.55. RG further reaches an FID of 1.74 at 512x512 and extends to text-to-image generation, achieving 0.89 on GenEval and 87.1 on DPG-Bench.

Zihao Pan, Mengting Chen, Cheng Shi +6

Sun Yat-sen University · Alibaba Group · The University of Hong Kong +1

Diffusion Transformer Pixel-Space Diffusion Models Diffusion Model Guidance Representation Learning T2I Diffusion Models

Ask PDF

Sep 30, 2026

Looped Diffusion Transformer

Improving text-to-image models has traditionally relied on increasing model size or the number of denoising steps. In this work, we explore an alternative way to scale computation by repeatedly running shared Transformer blocks within each denoising step, effectively increasing computational depth while keeping the parameter count fixed. This looped computation enables iterative refinement of internal representations without explicit reasoning tokens. However, naive looping fails to consistently improve image quality. We trace this problem to weak supervision across intermediate loops and unregulated attention updates that progressively erode local information. To overcome these challenges, we propose Looped Diffusion Transformer (Looped-DiT), which combines deep supervision across intermediate loops with self-modulating attention to stabilize looped feature updates. Under matched-parameter and matched-compute settings, Looped-DiT consistently outperforms non-looped baselines. Notably, a 260M-parameter looped model can surpass a model 6.5x larger across multiple text-to-image benchmarks while requiring 4.9x lower inference compute. Beyond this performance gain, we find that looped computation can offer a more effective form of iterative computation for diffusion models, with increasing loop depth yielding larger gains than adding more denoising steps under a fixed inference budget. Furthermore, deeper loops can progressively correct mistakes made in earlier loops, exhibiting behaviors suggestive of latent reasoning. Together, these results show that looped computation offers a promising way to scale visual generation models.

Yong Xien Chng, Tianyi Chen, Wenwen Tong +7

SenseTime Research · LeapLab, Tsinghua University · Nanyang Tehnological Univesity

Diffusion Transformer Recurrent Transformers Diffusion Model Inference Acceleration T2I Diffusion Models

Ask PDF