PaperlayerView exact revision on arXiv

Source-linked AI summary

Looped Diffusion Transformer

Yong Xien Chng, Tianyi Chen, Wenwen Tong, Haiwen Diao, Zhongang Cai, Lei Yang, Ziwei Liu, Lewei Lu, Dahua Lin, Gao Huang

arXiv:2609.40305v1cs.CVcs.LG

TL;DR

Deeper or wider Transformers and longer chain-of-thought sequences each raise the cost of scaling computation. Looped-DiT repeatedly applies shared Transformer blocks with intermediate-loop supervision and self-modulating attention; a 260M-parameter model outperforms non-looped models with roughly 6.5× more parameters across multiple benchmarks while requiring 4.9× lower inference compute.

  • Problem

    Deeper or wider Transformers require proportionally more parameters, while chain-of-thought reasoning increases inference-time computation through longer token sequences.

    Source: §1 · p. 2

  • Method

    Looped-DiT repeatedly applies shared Transformer blocks within each denoising step, combining intermediate-loop deep supervision with self-modulating attention.

    Source: §1 · p. 3

  • Results

    A 260M-parameter Looped-DiT outperforms non-looped models with roughly 6.5× more parameters across multiple text-to-image benchmarks while requiring 4.9× lower inference compute.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Under matched inference compute, increasing loop depth yields greater gains than adding denoising steps, and deeper loops progressively correct earlier errors.

    Source: §1 · p. 3

    Experiments focus on MiniT2I-based MMDiT models at approximately 260M parameters and 512 × 512 resolution; scaling to larger models, latent-space architectures, and higher resolutions remains unclear.

    Source: §A.4 · p. 21

Abstract

from arXiv · show

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.

1 Introduction

Deeper or wider Transformers increase parameter costs, while chain-of-thought increases inference computation through longer token sequences, motivating more efficient ways to scale text-to-image models. Looped-DiT applies shared Transformer blocks repeatedly within denoising steps and reports gains in benchmark performance and inference-time scaling.

1 Introduction

  • Looped-DiT runs shared middle Transformer blocks repeatedly within each denoising step and combines intermediate-loop deep supervision with self-modulating attention.Source: §1 · pp. 2–3
  • Naive looping can remain below the non-looped baseline, saturate, and decline beyond training depth, alongside progressively less linearly decodable spatial information.A ridge-regression probe's R2 drops from 0.865 after the 1st loop to 0.562 after 8th loops.Source: §1 · pp. 2–3, Figure 3 · p. 3
  • A 260M-parameter Looped-DiT outperforms non-looped models with roughly 6.5× more parameters across multiple text-to-image benchmarks while requiring 4.9× lower inference compute.Source: §1 · p. 3
  • Under matched inference compute, increasing loop depth yields greater gains than allocating additional computation to more denoising steps.Source: §1 · p. 3
  • The authors report that deeper loops progressively correct earlier errors and resolve interdependent visual constraints without explicit textual reasoning traces.The motivation is that coherent image generation requires inferring visual content and identifying and correcting inconsistencies as generation evolves.Source: §1 · pp. 2–3

2 Looped Multimodal Diffusion Transformer

Looped-DiT repeatedly applies a shared Transformer stage within each denoising step, increasing computational depth without increasing parameter count. Deep Supervision trains intermediate-loop predictions, while Self-Modulating Attention regulates repeated attention updates.

2 Looped Multimodal Diffusion Transformer

  • A pre-loop stage A, shared looped stage B, and post-loop stage C process image and text representations; B runs N times within each denoising step.Because B's parameters are shared, increasing N adds computational depth without increasing parameter count, independently of the number of denoising steps.Source: §2 · pp. 3–4

2.1 Deep Supervision

  • Deep Supervision applies the flow-matching objective to predictions at every loop depth, giving earlier loop states direct learning signals.The predictions reuse the shared post-loop stage and are decoded only during training, adding no parameters or inference overhead.Source: §2.1 · p. 4

2.2 Self-Modulating Attention

  • Self-Modulating Attention regulates repeated updates using either a learned, token-dependent Gated Attention scalar or Exclusive Self Attention's state-dependent projection removing the token's own value-direction component.These mechanisms modulate head outputs before concatenation and output projection and are applied only within the looped stage.Source: §2.2 · pp. 4–5

3 Experiments

The experiments compare Looped-DiT with state-of-the-art models and matched parameter- and compute-budget alternatives, then assess loop-depth allocation and component ablations. They report benchmark gains, more effective inference-time compute allocation, and improvements from deep supervision and self-modulating attention.

3.1 Main Results

  • Looped-DiT B/16 achieves a 2.5-point higher avg. score than next-best InternVL-U with 6.5× fewer parameters.It achieves the best results on DPG-Bench, PRISM, T2I-CoReBench, SpatialGenEval, and TIIF-Short.Source: §3.1 · p. 5
  • Looped-DiT B/16 traces the Pareto frontier under both inference-compute and latency metrics, with loop depth and denoising steps controlling the inference budget.Average Score is the mean across six benchmarks.Source: §3.1 · p. 6, Figure 4 · p. 6
  • Successive loops introduce missing content, satisfy spatial constraints, and correct rendering errors or remove extraneous objects in the shown examples.The paper reports consistent performance gains as loop count increases from 1 to 4.Source: §3.1 · p. 6, Figure 5 · p. 6

3.2 Necessity of Looping

  • Under matched settings, Looped-DiT B/32 improves average score by 4.0 points over the parameter-matched baseline and scores 59.1 versus 58.1 for the training-compute-matched deeper baseline.It also outperforms both inference-compute-matched baselines despite using fewer parameters.Source: §3.2 · p. 7, Table 2 · p. 7
  • From 25- to 50-step settings, allocating compute to loop depth consistently yields higher performance than allocating it to additional denoising steps.The comparisons include a same-checkpoint single-pass baseline and a model trained without looping; single-pass settings match the four-loop inference FLOPs.Source: §3.2 · p. 7, Figure 6 · p. 7
  • Looping yields larger gains on subtasks with relational, procedural, or spatial constraints, while CoT helps more with inferring implied content from examples and rules.Combining looping and CoT achieves the highest overall score.Source: §3.2 · p. 8, Table 3 · p. 8

3.3 Ablation Study

  • Final + Mean supervision and XSA together yield the highest average; XSA gains are +1.6 for the looped model, −0.1 for the parameter-matched baseline, and +0.7 for the compute-matched baseline.Both Gated Attention and XSA improve over unmodulated attention, with XSA giving the larger gain; the paper relates this to smaller attention updates and better token-position decodability.Source: §3.3 · pp. 8, 10, Table 4 · p. 8, Table 5 · p. 10
  • Deep supervision improves performance across 1–8 inference loops, reduces early-exit penalties, and maintains performance beyond the four-loop training depth; adaptive looping beats fixed-loop inference at matched mean loop counts.Adaptive looping varies per-image loop counts without retraining.Source: §3.3 · p. 9, Figure 7 · p. 9, Figure 8 · p. 9

4 Related Works

The work connects looped Transformers, text-to-image diffusion, and efficient inference. It studies shared-block loop depth as a parameter-efficient way to increase computational depth and reallocate inference compute.

4 Related Works

  • Under a fixed inference budget, increasing loop depth yields larger gains than adding more denoising steps.This positions looped computation as a potentially more effective form of iterative computation for diffusion models.Source: §4 · p. 11
  • The study applies looped computation to text-to-image generation, focusing on parameter efficiency and compositional and spatial reasoning.It builds on MMDiT and examines how inference compute can be allocated between loop depth and denoising steps.Source: §4 · p. 11
  • Unlike few-step methods that reduce denoising steps, this work investigates internalizing iterative computation through loop depth.The passage places this approach alongside few-step and one-step methods as a different way to improve inference efficiency.Source: §4 · p. 11

5 Conclusions

Looped-DiT repeatedly applies shared Transformer blocks within each denoising step, with intermediate supervision and self-modulating attention. The paper reports that looping can improve text-to-image models efficiently and that deeper loops show behaviors suggestive of latent visual reasoning.

5 Conclusions

  • Looped-DiT increases computational depth by repeatedly applying shared Transformer blocks within each denoising step without increasing parameter count.The method combines intermediate-loop Deep Supervision with Self-Modulating Attention to regulate updates as representations evolve.Source: §5 · p. 11
  • Naive looping is unreliable and progressively degrades spatial information, motivating intermediate supervision and adaptive attention regulation.Controlled comparisons indicate that the resulting gains are not simply explained by scaling model depth or width, using the same inference compute for more denoising steps, or adding explicit textual reasoning.Source: §5 · p. 11
  • Successive loops exhibit behaviors suggestive of latent visual reasoning as hidden-state refinement progressively resolves constraints and corrects earlier errors.Source: §5 · p. 11
  • The findings present looping as a promising parameter-efficient way to improve future text-to-image models.Source: §5 · p. 11

A.1 Implementation Details

The implementation uses a pixel-space MMDiT backbone with shared middle blocks repeated within denoising steps. Training adds intermediate-loop supervision, while inference varies loop depth without changing parameters.

A.1 Implementation Details

  • The backbone is MiniT2I, a pixel-space denoiser based on MMDiT, retaining its image and text pathways and joint attention.The B/32 and B/16 variants contain 260.2M and 258.1M parameters, respectively.Source: §A.1 · p. 16
  • The 17-block model divides its blocks into pre-loop, looped, and post-loop groups, repeating the shared middle five blocks four times during training.This gives 32 effective block applications per denoising step while retaining 17 unique parameterized blocks; inference loop depth can vary without changing model parameters.Source: §A.1 · p. 16, Table 6 · p. 16
  • Training uses pixel-space flow matching with noise scale 2.0 and drops the text condition with probability 0.1.Source: §A.1 · pp. 16–17
  • Deep Supervision optimizes predictions at loop depths 1, 2, and 3 against the same flow-matching target as the final prediction.Final+mean weighting (1/3, 1/3, 1/3, 1) performs best among the two tested schemes and is adopted for B/16; intermediate predictions add no inference-time computation or parameters.Source: §A.1 · p. 17
  • B/32 is pretrained for 250K steps and fine-tuned for 40K, while B/16 is pretrained for 500K and fine-tuned for 80K.The models use AdamW with global batch size 1024 and peak learning rate 4 × 10−4; hyperparameters are shared across scales and stages unless noted.Source: §A.1 · p. 17, Table 7 · p. 17
  • Training data follows MiniT2I, with CC12M pretraining and fine-tuning on a mixture of BLIP3o-60K, DALL-E 3, and ShareGPT-4o-Image.Images are resized, center-cropped to 512 × 512, and normalized to [−1, 1], without random flipping or additional image augmentation.Source: §A.1 · p. 17
  • Evaluation covers six complementary text-to-image benchmarks for B/16 and four reasoning-focused benchmarks for B/32 analyses and ablations.Unless specified otherwise, evaluation uses EMA checkpoints, Euler sampling with 100 denoising steps, classifier-free guidance scale 6.0, and loop depth 4.Source: §A.1 · pp. 17–18
  • Deep Supervision raises training compute to 1,246 GFLOPs per sample per step versus 809 and 805 GFLOPs for Deeper and Wider baselines, respectively.The overhead is training-only; at inference, intermediate predictions are not evaluated, and Self-Modulating Attention adds negligible compute. The cost table measures inference GFLOPs per denoising forward pass.Source: §A.1 · p. 18, Table 8 · p. 18

A.2 Additional Related Work

The related work situates Looped-DiT among diffusion Transformers, looped computation, and methods for reducing sampling steps. It distinguishes internal hidden-state refinement from approaches that reduce external denoising steps.

A.2 Additional Related Work

  • Diffusion Transformer research spans scalable backbones, efficient training, modality-specific joint attention, linear attention, and compressed latent representations.The passage contrasts those architectural and scaling efforts with this work's use of repeatedly applied shared blocks to increase computational depth.Source: §A.2 · pp. 18–19
  • Few-step and one-step methods reduce external model evaluations, whereas this work investigates internalizing iterative computation through loop depth.Source: §A.2 · p. 19
  • Looped Transformer research covers parameter sharing, iterative algorithm execution, reasoning, latent computation, adaptive recursion, and memory or multi-resolution extensions.Source: §A.2 · p. 19
  • Elastic Looped Transformers study variable loop depth for class-conditional image and video generation, using intra-loop self-distillation to address early-exit degradation.Looped-DiT instead directly supervises intermediate predictions with the final prediction's training objective and studies open-ended text-to-image generation.Source: §A.2 · p. 19

A.3 Additional Details on Self-Modulating Attention

Self-Modulating Attention (SMA) changes attention-head updates within the looped stage, without changing attention weights. Gated Attention and Exclusive Self Attention (XSA) implement this modulation through distinct, representation-dependent mechanisms.

A.3 Additional Details on Self-Modulating Attention

  • SMA transforms each attention head’s output after scaled dot-product attention and before concatenation and output projection, changing residual updates without changing attention weights.It is applied only in the looped stage; pre-loop and post-loop stages use standard attention.Source: §A.3 · pp. 19–20
  • Gated Attention uses a learned, token-dependent scalar gate for each head to control the magnitude of its contribution.Gate parameters are modality-specific, with separate parameters for image and text tokens.Source: §A.3 · p. 20
  • XSA is parameter-free and projects the attention output orthogonally to the token’s own value direction, removing its direct self-value contribution.Unlike simply scaling the output by 1 −αii,h, the projection also removes components of other value vectors along that direction; it is non-expansive at the individual-head output before projection.Source: §A.3 · p. 20
  • Gated Attention and XSA both modulate attention outputs before they enter the residual stream, and recompute their modulation from current representations at every loop.Gated Attention uses learned scalar gates, whereas XSA uses a parameter-free, state-dependent projection.Source: §A.3 · p. 20

A.4 Limitations and Future Work

The study identifies limits in the settings explored, the design configurations tested, the training cost of Deep Supervision, and the evidence for latent visual reasoning.

A.4 Limitations and Future Work

  • Experiments are limited to MiniT2I-based MMDiT models at approximately 260M parameters and 512 × 512 resolution, leaving scaling to larger models, latent-space architectures, and higher resolutions unclear.The authors identify evaluating those settings as future work.Source: §A.4 · p. 21
  • Resource constraints limited the study to a fixed [6] stage partition and training loop depth of four rather than an exhaustive design-space exploration.Future work could examine how loop placement, training depth, and the fraction of shared blocks interact with scale and inference budget.Source: §A.4 · p. 21
  • Deep Supervision increases training computation because intermediate loop states are additionally decoded during optimization.This overhead is absent at inference; reducing its training cost or using more efficient objectives remains a direction for improving overall efficiency.Source: §A.4 · p. 21
  • Evidence for latent visual reasoning is primarily behavioral and representational, not a complete mechanistic account of how the behaviors emerge.The authors suggest more direct analysis of information flow and computation across loop iterations.Source: §A.4 · p. 21

Related papers

Loading 2609.40305v1…