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

Computer Vision and Pattern Recognition

Summary

The gist is being written…

Authors

Zipeng Wang, Anping Xiong, Jing Zhang, Xiaoya Zhu, Yi Jiang, Yanbing Xiao, Zige Zeng, Zezhou Chen, Xiang Liu, Huan Hu, Ping Chen, Kai Wang, Zhaoxiang Liu, Shiguo Lian, Ye LV, Fang Zhao, Huanlin Gao, Zhengrui Huang, Bo Shen, Yuren You

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.