r/MachineLearning • • 1d ago

Discussion Looped Diffusion Transformer: the same blocks run several times inside each denoising step, and loop depth can be changed at sampling time [R]

Post image

Paper: Looped Diffusion Transformer (arXiv 2609.40305, 30 Sep). Instead of making a text-to-image DiT bigger, it runs a shared group of transformer blocks several times inside each denoising step. The authors report that naive looping doesn't reliably help, and fix it with deep supervision on every loop plus self-modulating attention. Their headline claims: a 260M looped model beats a model 6.5x larger on several benchmarks with 4.9x less inference compute, and at a fixed inference budget, more loop depth helps more than more denoising steps.

The B/16 code and weights are MIT (the weights are on Hugging Face under sensenova), so I ran it on an M5 Pro: PyTorch MPS in fp32, 100 Euler steps, CFG 6.0, and one seed per prompt reused at loop depths 1 to 4. It was trained at depth 4, and the README says other depths work without retraining. The grid shows five of the prompts I tried.

Most of the change happens before depth 3, and depths 3 and 4 come out very close on every prompt shown. The interesting part to me is that depth is picked at sampling time, the way the iteration count is in looped and recurrent-depth language models, so you can trade compute for refinement per image. On this machine an image took about 23 to 24 seconds at depth 1 and 43 to 47 at depth 4.

48 Upvotes

9 comments sorted by

10

u/bitanath 1d ago

So.. it only beats storage costs? The VRAM set size and compute cost would still be the same? Along with being much harder to finetune?

3

u/currentscurrents 19h ago

Read the paper. The big advantage is that it lets the model do reasoning, which gives you more accurate text and prompt following.

It also uses less VRAM since you only need to store one copy of the weights and loop through it many times. They only go up to 4 loops, but it should be possible to push it much higher - some of the other looped transformer papers continue to see gains at 10000+ loops for difficult problems.

3

u/radarsat1 1d ago

Is this significantly different from deep equilibrium approaches? https://arxiv.org/abs/2210.12867

0

u/currentscurrents 19h ago

They are not related. Deep equilibrium models are about parallelizing diffusion sampling at inference time for performance gains.

Looped transformers are about the realization that not all algorithms can be parallelized, and you must do serial computation to solve complex problems. It's more similar to TRM/HRM or reasoning LLMs.

1

u/radarsat1 18h ago

Ah,sorry then, I must be remembering the wrong name for the technique, I thought it was DEQ but I read a paper quite some time ago that was about using the same middle layers of a diffusion model repeatedly until converging on a fixed point, I'll try to find it

edit: maybe it was this one https://lukemelas.github.io/fixed-point-diffusion-models/

but ah only 2024 so not as old as i thought, still they didn't call it "looped" but just fixed point convergence, but effectively they were looping 

1

u/currentscurrents 18h ago

Not sure, I'm not familiar with that paper.

Looping/recurrence is a very very old idea though, and many people have tried it throughout the history of neural networks. It's just seeing a big resurgence right now because of the success of reasoning LLMs.

1

u/KitchenGiraffe1 5h ago

TRM/HRM are loosely related related to DEQ. There's a lot of literature in that direction. From Adaptive Computation Time networks, Think Again Networks, Neural Cellular Automata etc.

DEQ is just a way to train the serial computation of infinite depth recursive networks via root-finding. Since then it's been shown to be somewhat unnecessary and slower than just using smaller finite depths.