Looped Diffusion Transformer

2026-09-30

AI Summary

This study addresses the inefficiency of text-to-image models that rely on scaling parameters or sampling steps, as well as the feature degradation commonly induced by naive recurrent computation. To this end, we propose Looped-DiT, an architecture that increases computational depth by recurrently executing shared Transformer blocks. By integrating deep supervision with a self-modulating attention mechanism, the method stabilizes feature updates and enables efficient iterative refinement under a fixed parameter budget. Experimental results demonstrate that a model with merely 260 million parameters outperforms baselines 6.5 times larger, while reducing inference compute by 4.9 times. Furthermore, it significantly surpasses non-recurrent counterparts, exhibiting strong latent reasoning capabilities.

Abstract

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.

Problem

• text-to-image generation
• looped computation
• Diffusion Transformer
• deep supervision
• attention degradation

Innovation

• Looped Diffusion Transformer
• Shared Transformer Blocks
• Deep Supervision
• Self-Modulating Attention
• Latent Reasoning