Note: This post was written with the help of AI.
Loops within each denoising step
I have been reading a few papers on Looped Transformers for generation and wanted to compare them here. Diffusion and flow models start from noise and generate an image step by step. These papers repeat the same Transformer blocks K times within each network call, increasing computation depth without adding parameters. The three 2026 papers outperform larger dense models in their respective experiments. Another Google paper carries hidden states across denoising steps, which makes a useful comparison. I also look at their connections to recurrent depth in language models and to DEQ.
6 + 5×K + 6
T × blocks per call
Independent of K
The outer dial tracks sampler time t, from noise to image; each tick is one network call. The inner dial tracks repetitions of the shared blocks within that call. The [6, 5, 6] split comes from Looped-DiT. The image is a schematic, not a model output.
Sampling steps determine how finely we follow the flow; loops determine how accurately we estimate the flow at each step. The two computation budgets can be adjusted separately.
Adapted from the LiFT introduction: "Steps follow the flow more closely; loops estimate it more accurately."
Basic structure: Prelude → Core × K → Coda
The three papers use nearly the same structure. Split a DiT into three parts: the prelude P runs once and encodes the noisy input into a hidden state; the core R shares weights across K repetitions; and the coda D reads out a prediction, either velocity v or a clean image x̂0. All of this happens within one denoising call, before the sampler advances to the next t.
Parameter count and executed depth can then be counted separately: there are LP + LR + LD unique blocks, but LP + K·LR + LD block executions. Looped Transformers already used this idea in language models, including Universal Transformers and Geiping et al.'s recurrent depth; here it becomes the inner loop of a diffusion model. The names prelude, core and coda come from Geiping et al.'s Huginn. LiFT explicitly adopts this decomposition (see §8).
| Prelude | Core × K | Coda | Additional details | |
|---|---|---|---|---|
| ● ELT | embedding | N layers repeated L times, e.g. 8N × 4L | Shared head | Tested with both MaskGIT and DiT |
| ● Looped-DiT | 6 MMDiT blocks | 5 blocks × 4 loops | 6 blocks | 17 unique blocks, 32 executions; XSA attention in the loop |
| ● LiFT | 4 DiT blocks | R blocks × K loops | Prediction head only | RMSNorm then h₀ reinjection each loop; the core also receives a depth-coordinate s embedding |
What happens with final-loop supervision only
The simplest approach supervises only the last loop. All three papers start here and observe a similar failure: the model learns a fixed computation of exactly Ktrain loops, rather than an improvement operator that can be applied repeatedly. Fewer loops produce incomplete images; extra loops can damage an already usable output.
A vanilla looped DiT trained with 8 loops on ImageNet produces coherent images only when inference also uses exactly 8 loops. At 2, 4, 6 or 10 loops, quality degrades severely. Intermediate states are not near the solution space.
In MMDiT, performance saturates and then declines when inference uses more loops than training. A ridge-regression probe predicts image tokens' 2D positions from their hidden states: R² drops from 0.865 after the first loop to 0.562 after the eighth. Repeated attention updates wash out local information.
Existing intermediate supervision makes outputs usable at multiple depths, but performance soon deteriorates beyond training depth. ImageNet FID rises as soon as that depth is exceeded; video performance peaks at 1.5 times training depth. In a separate LoopDiT experiment, looping did not beat a dense model at equal compute.
How to supervise intermediate loops
The structures are similar; I think the main difference is how the loss treats intermediate loops. The diagram places each readout uk in prediction space, with arrows showing its supervision target. Select a scheme to compare them.
ELT: Elastic Looped Transformers
ELT uses the deepest path as a teacher for shallower exits. The same model can then choose its loop count according to the inference budget or exit early.
Method: Intra-Loop Self Distillation
- Each training step samples a student exit Lint ~ U(Lmin, Lmax). The student path is a prefix of the teacher path, so the added training cost is small.
- The student learns from both ground truth and the teacher's stop-gradient output. The weight λ decreases linearly from 1 to 0: ground truth dominates early, and the teacher dominates later.
- Discrete tokens use cross-entropy; diffusion uses sigmoid-weighted MSE. The method applies to both generators.
- The aim is Any-Time inference: use the same weights and choose the loop count for the device's budget.
Parameter sharing and throughput
- If the shared parameters fit on chip, L repetitions can increase arithmetic intensity by up to L times, avoiding repeated HBM–SRAM transfers.
- On TPU v6e, throughput is 2.9×, 3.3× and 3.5× the baseline at L, XL and H scales. The B-scale baseline already fits on chip, so there is no speedup (1.0×).
- The DiT variants also converge faster: 16N×2L by 2× and 8N×4L by 1.4×, at the same inference compute as a 32-layer DiT.
With half the parameters (1.1B versus 2.1B), 16N×2L beats the 32-layer dense model; 8N×4L also wins with roughly one quarter. Repeating a single layer 32 times lacks capacity (FID 10.30), so each loop still needs enough blocks. Data: ELT Table 2.
Looped-DiT: Looped Diffusion Transformer
Text-to-image generation sometimes requires inferring unstated content, satisfying several constraints and correcting earlier errors. Looped-DiT attempts these computations through loops in hidden space, without generating a textual chain of thought.
Deep Supervision
- Hidden states from loops 1, 2 and 3 also pass through the shared 6-block post-loop segment, using a flow-matching loss against the same clean image x0.
- Two weighting schemes are compared: exponential (1/8, 1/4, 1/2, 1) and final+mean (1/3, 1/3, 1/3, 1). The latter performs better.
- Extra computation is needed only during training, with no added inference parameters or compute. Training FLOPs are about 1.54× those of a baseline matched for inference compute.
Self-Modulating Attention
- Softmax controls which tokens supply information, but not how much is written back. Repeated updates with the same blocks can erase local token information.
- Gated Attention multiplies each head by a sigmoid gate conditioned on the current token.
- XSA (Exclusive Self Attention) projects out the output component along the token's own value direction: z = (I − v̂v̂ᵀ) Σ α v. It adds no parameters, performs best, and is used only in the looped segment.
- XSA adds +1.6 points to the looped model, versus −0.1 to a non-looped baseline with equal parameters and +0.7 to one with equal compute. It mainly addresses a problem introduced by looping.
The non-looped MiniT2I-B/16 with the same basic structure scores 66.4; looping raises this to 71.5. InternVL-U (1.7B) scores 69.0. The authors report 6.5× fewer parameters, 4.9× less inference compute and 3.2× lower latency. Data: Looped-DiT Table 1.
| Model | Matched condition | Parameters | Effective depth | Inference GFLOPs | Mean score |
|---|---|---|---|---|---|
| MiniT2I B/32 | Parameters | 260M | 17 | 146 | 55.2 |
| Deeper MiniT2I | Inference compute | 473M | 32 | 267 | 57.7 |
| Wider MiniT2I | Inference compute | 489M | 17 | 268 | 58.6 |
| Deeper + Deep Supervision | Training compute | 473M | 32 | 267 | 58.1 |
| Looped-DiT B/32 | — | 260M | 32 | 267 | 59.1 |
The looped model leads under all three matching conditions and has the fewest parameters. Data: Looped-DiT Table 2.
More loops or more steps at the same budget
- At the same inference budget, adding loops improves scores more than adding denoising steps, consistently over the 25–50-step range.
- There are two comparisons: disabling loops in the same checkpoint, and a model trained without loops using the same architecture.
- A lightweight gating network chooses the number of loops per image. At the same average loop count, adaptive exits outperform a fixed count, and the budget can be adjusted without retraining.
Loops and textual CoT are complementary
- Loops help more on tasks requiring simultaneous relational, spatial and procedural constraints. Rewriting prompts with CoT adds little on these tasks.
- CoT helps more on tasks requiring inference of unstated content from examples and rules.
- Combining both gives the best score. Qualitatively, later loops add missing objects, move objects to satisfy spatial constraints, correct text and remove extra objects.
LiFT: Loop Flow Transformers
LiFT gives each loop a different target. Draw a straight line from the model's initial estimate to ground truth; loop k predicts the corresponding point on that line. Depth is represented by a continuous coordinate. In the experiments, models trained with 2 loops still benefit from increasing inference to 16 loops.
Method: trajectory supervision
ūs = (1 − s)·b + s·u★ // s ∈ [0, 1]
L = 1/K Σk ‖uk − ūsk‖² + λ‖u₀ − u★‖²
- Training samples Ktrain − 1 coordinates per example and sorts them, with the endpoint fixed at s = 1. The final loop therefore always learns the standard flow-matching target.
- Inference uses a uniform grid sk = k / Kinf. The loop count determines how finely the path is discretized, much as integration steps discretize generation time.
- In the ablation, fixed-grid training gives FID 75.26 and 86.95 at 16 and 32 loops, while random-coordinate training gives 67.71 and 68.38.
Extrapolation, equal-compute comparisons and limits
- Extrapolation: the earlier two papers mainly make outputs useful within training depth; LiFT makes extra loops beyond training depth continue to help.
- Equal compute: comparisons use the same TFLOPs as dense DiT, rather than only the same parameter count.
- Limits: at B/2 scale, dense models are stronger at every budget. At training depth, every LiFT loses to the dense model of the same scale. The advantage comes from reallocating inference compute to loops.
- The paper also cites an iso-depth scaling result: one shared recurrence contributes roughly as much as half an independent block.
At 3 loops, L/2 R10 already beats dense XL/2 (14.64 versus 15.54), using slightly less inference compute (11.43 versus 11.86 TFLOPs) and about 60% fewer parameters. Both curves flatten around 8–16 loops and rise slightly at 32. Data: LiFT Table 5.
Darker cells indicate lower FID. The black outline marks training (50 steps, 2 loops). Dense L/2 has comparable training compute to this LiFT checkpoint and about 41% more parameters. One efficient setting is 25 steps × 4 loops: FID 13.16 at 7.40 TFLOPs. Dense L/2 reaches 16.94 at 50 steps and 8.07 TFLOPs. Data: LiFT Tables 6 and 7.
Diffusion as a Training Curriculum for Timestep-Free Iterative Reasoning
The first three papers loop within one denoising call and restart at the next step. This paper carries the hidden state to the next step and removes timestep conditioning. At inference, gradual denoising is unnecessary; diffusion mainly provides the training curriculum.
Method
- The denoiser is itself recurrent: Sudoku uses 1 Transformer block with row, column and box masks, repeated J = 8 times; mazes use a 3×3 convolution block repeated 8 times.
- Each step has two LayerNorm outputs: one is projected into a denoising prediction ŷ, and the other becomes the hidden state h for the next step. There is no direct loss on h.
- The network receives neither a timestep nor an iteration index, so the same update can be applied for any number of steps.
- Training samples trajectory length K ~ U[20, 160]. Each step's input is re-noised from ground truth, and backpropagation is truncated to 4 steps. The hidden state h is the only channel that carries earlier information forward without corruption.
Findings
- Accuracy improves monotonically up to 10,000 steps, far beyond the 160 training steps and 4-step backpropagation window.
- Keeping inference noise at its maximum gives the best result (99.90%). Removing noise drops accuracy to 49.39%. Noise acts as exploration: a single trajectory can escape an incorrect solution and eventually settle.
- Removing h slows convergence and lowers the ceiling. Adding or removing timestep conditioning makes little difference.
- An ordered annealing schedule during training matters: fixed noise or independently sampled noise levels at each step substantially reduce accuracy.
The upper group shows means over three seeds, with standard-deviation error bars (Appendix Table 2). The lower group comes from Section 3. The two annealing values (99.54 and 99.56) come from different experiments.
Relation to within-step loops
- Here diffusion provides a curriculum and recurrence provides computation. The first three papers make diffusion models recurrent; this paper uses diffusion to train a recurrent solver.
- It uses a different handoff between states. LiFT explicitly reinitializes the loop state at every integration step, while RIN and this paper carry hidden states across steps. The text-to-image and video papers discussed here use the first approach.
- Transfer to generation needs care. These reasoning tasks have unique solutions. On multi-solution Sudoku, the authors find that 57.93% of trajectories move to another solution after finding one. Image generation also admits many solutions, so keeping maximum noise at inference may not transfer well.
Background: Huginn and DEQ
There are two sources for these architectures. Language-model recurrent depth unrolls shared blocks a finite number of times, samples depth during training and backpropagates through only the final loops. DEQ (Deep Equilibrium Models) instead solves the fixed point z* = f(z*, x), corresponding to a shared-weight loop run to convergence at infinite depth. Both study how to separate computation depth from parameter count.
Grouping the papers by how they unroll recurrence helps identify which can benefit from more loops at inference.
Huginn's Remark 3.1 reports a relevant negative result: conditioning the core Ri on an absolute step index breaks path independence and prevents extrapolation. LiFT uses a continuous depth coordinate normalized to [0, 1], while Drozdova et al. remove timestep inputs altogether. Both avoid this problem.
Huginn: Recurrent Depth
- The layout is (2, 4, 2): 2 prelude layers, 4 core layers and 2 coda layers, giving 132 executed layers at r = 32. Each loop concatenates prelude output e with the current state and passes them through an adapter. The initial state is random noise.
- Training samples r from a log-normal Poisson distribution with mean 32, and backpropagates through only the last 8 loops. Activation memory does not depend on r.
- 3.5B parameters and 800B training tokens. Inference is evaluated up to r = 64, with GSM8K CoT at 34.80 versus OLMo-7B's 6.07. The authors describe the compute as equivalent to a 50B-parameter model.
- Inference can exit adaptively per token when adjacent-loop KL falls below 5×10⁻⁴, and can share the KV cache.
LiFT and Looped-DiT's prelude–core–coda layout, along with random-depth training for inference extrapolation, trace back to this work.
LoopVL: Recurrent Visual Intelligence
- Two 16-layer stacks: a fast L and a slow H. Module-Loop repeats L with fixed high-level conditioning; Model-Loop repeats the entire L→H cycle. The default [L×3→H]×2 executes 128 layers using 32 unique layers.
- About 1B parameters, trained from scratch: 75B text-pretraining tokens, 49B multimodal tokens, 16.8B SFT tokens, then GRPO.
- At matched tokens (0.14T), MMStar is 63.47 versus 55.33 for the non-looped baseline of the same size. Attention Gini rises from 0.405 to 0.858 at the boundary between Model-Loop cycles, which the authors call a Visual Aha Moment.
- It does not extrapolate: increasing cycles from the training count of 2 to 4 drops MMStar to 24.00.
This is visual understanding rather than generation. Training uses a fixed loop count without random depth; extra loops make it worse, providing a comparison with random-depth methods such as Huginn and LiFT.
DEQ: Deep Equilibrium Models
- Solve z* = fθ(z*; x) directly, equivalent to an infinitely deep shared-weight network with x injected at every layer. Broyden's method finds the root.
- Backpropagation uses the implicit function theorem. Activation memory is comparable to one layer and independent of depth; the abstract reports savings of up to 88%.
- On WikiText-103, DEQ-Transformer has perplexity 24.2 at 2.7GB memory, versus Transformer-XL's 24.3 at 8.5GB. The cost is 2.82× slower training and 1.76× slower inference.
This is the limiting case of recurrence. If each loop approaches the same fixed point, more loops solve it more accurately.
DEQ-Flow: fixed points in optical flow
- Rewrite RAFT's iterative ConvGRU update as a fixed-point problem, solved with Anderson or Broyden until residual falls below 10⁻³. Phantom gradients make training nearly 2× faster.
- Sparse fixed-point correction losses on selected intermediate solutions address DEQ training instability.
- Each frame is initialized from the previous frame's fixed point. Inference is about 1.6× faster on Sintel, with 4–6× lower training memory.
- KITTI-15 (C+T) F1-all: 16.3 for DEQ-RAFT-B versus 17.4 for RAFT at the same size.
One early example of reusing iterative states across video frames, with a structure directly relevant to frame-by-frame autoregressive world models.
FPDM: Fixed Point Diffusion Models
- Starting from DiT, keep two explicit blocks at the ends and replace the middle with a timestep-conditioned implicit fixed-point layer that receives input injection.
- Training first iterates n ~ U[0, N] times without gradients, then m ~ U[1, 12] times with gradients.
- Sampling distributes the iteration budget across timesteps, with each timestep warm-started from the previous fixed point.
- 85M parameters versus 674M for DiT-XL/2; training memory is 10.2GB versus 25.2GB.
A DEQ version of within-step loops. The outer loop is sampling, the inner loop solves a fixed point, and the solution is reused across outer steps. This lies between the designs of LiFT and Drozdova et al.
GET: a DEQ student for one-call generation
- Generative Equilibrium Transformer: an injection Transformer processes noise and class, an equilibrium Transformer solves a fixed point, and a decoder produces an image in one network call.
- Offline distillation uses one million noise–image pairs generated by EDM. The main experiment unrolls 6 iterations with direct backpropagation; gradient checkpointing can also give O(1) activation memory.
- One network call on CIFAR-10: unconditional FID 6.91 versus PD-EDM's 8.34, and class-conditional FID 6.25. GET-Mini (19.2M) has FID 10.72 versus ViT's 11.49 at 85.2M.
ELT's discussion suggests adjusting one-call model quality through loop count. GET already did this in 2023 using fixed-point iterations.
| Total block calls | FPDM | DiT-XL/2 |
|---|---|---|
| 140 | 33.9 | 110.0 |
| 280 | 22.4 | 35.2 |
| 560 | 19.6 | 16.5 |
At tight budgets, the smaller fixed-point model leads by a large margin; with more compute, the larger dense model overtakes it. At 560 calls with DDPM, FPDM still leads: 26.1 versus 37.9. This resembles the capacity limits seen in LiFT B/2 and ELT with 1 layer repeated 32 times. Data: FPDM Table 1.
Reusing iterative states in the outer loop
- DEQ-Flow: reuse optical-flow fixed points across video frames.
- FPDM: reuse fixed points across diffusion timesteps.
- Huginn: continuous CoT initializes the next token from the previous token's final state.
- Drozdova et al.: carry hidden states across denoising steps and re-noise samples at each step.
ELT, Looped-DiT and LiFT instead restart recurrence at every denoising step. These four reuse strategies offer possible designs for frame-by-frame video world models.
Comparing the four papers
| ELT | Looped-DiT | LiFT | Diffusion as Curriculum | |
|---|---|---|---|---|
| Task | Class-conditional images and UCF-101 video | Open-ended text-to-image | Class-conditional images | Sudoku and mazes with unique solutions |
| Generator | MaskGIT-style masked generation; DDPM-style DiT | Pixel-space MMDiT, flow matching | Latent-space DiT, flow matching | Continuous VP diffusion; recurrent Transformer or convolution block |
| Intermediate targets | Ground truth plus deepest loop (stop-gradient) | Same x₀ at every loop | Different points ūs on a straight path | No separate supervision within 8 inner loops; supervise ŷ at each step, with no loss on h |
| Architecture changes | Minimal | XSA or gated attention in the core | h₀ reinjection and depth-coordinate embedding | Hidden states across steps; no timestep conditioning |
| Main capability | Any-Time exits and throughput | Parameter efficiency, implicit visual reasoning and adaptive loops | Extrapolation beyond training depth; gains at equal compute | Anytime solving; noise-driven exploration within one trajectory |
| Beyond training depth | Modest video extrapolation (4 → 6) | Stable up to 8 loops with deep supervision | Continued improvement from 2 → 16 loops | Training up to 160 steps; improvement continues to 10,000 |
| Representative result | 4× fewer parameters at matched FID 2.0 | 0.26B scores 71.5 versus 69.0 at 1.7B | FID lower by 3.34 with 60% fewer parameters | Sudoku-Extreme 99.90% |
The intermediate targets compare as follows. Looped-DiT asks every loop to predict the final answer, making early loops useful through a simple objective. ELT asks shallow exits to match deeper ones, whose outputs are more reachable than ground truth and have lower variance. LiFT assigns each loop a fraction of the correction, giving loops different roles. Its extrapolation follows this design: more loops discretize the same path more finely.
The fourth paper focuses on what is passed between steps. The first three pass noisy samples; it also passes a clean hidden state with no direct supervision. It shows the largest extrapolation factor in the table (160 → 10,000 steps), but has only been validated on unique-solution reasoning tasks.
Looped-DiT explicitly distinguishes itself from ELT in related work: ELT uses self-distillation, while Looped-DiT directly supervises the same target at every exit without a teacher configuration. LiFT groups ELT and DeepFlow as methods that train intermediate predictions toward the complete output, which it aims to improve upon.
Timeline of related work
Recurrent depth has a long history in language models. In generation, iteration within a denoising call and carrying states across calls should be considered separately.
Universal Transformers
Dehghani et al. (Google) repeatedly apply a Transformer layer, separating parameters from depth. ALBERT later shared parameters across layers.
DEQ · Deep Equilibrium Models
Bai, Kolter and Koltun solve the fixed point of a shared-weight network and backpropagate with the implicit function theorem. Memory is independent of depth: the limit of recurrence run to convergence.
DEQ-Flow
Rewrite RAFT's iterative optical flow as fixed-point solving, and reuse fixed points across video frames.
RIN · Recurrent Interface Networks
Jabri, Fleet and Chen (Google) carry and reuse latent states across denoising steps. LiFT distinguishes this across-step recurrence from within-step loops. Drozdova et al.'s 2026 Google paper also uses the across-step approach.
Theory and reasoning in Looped Transformers
Giannou et al. show that looped Transformers can serve as programmable computers; Saunshi et al. connect looping with latent reasoning; Mixture-of-Recursions chooses recursion depth adaptively per token.
GET · FPDM
DEQ enters generation: GET (2023) uses a fixed point for a one-call distilled student; FPDM (2024) replaces DiT's middle layers with a timestep-conditioned fixed-point layer.
Huginn · Recurrent Depth
Geiping et al. use prelude–core–coda, random training depth and truncated backpropagation. Their 3.5B model continues to improve up to 64 inference loops. LiFT and Looped-DiT later adopt this layout.
Related designs in generation
DeepFlow uses intermediate-layer supervision and velocity alignment in flow models; DiT-Air studies layer sharing in DiT.
ELT · Google DeepMind
The first systematic study of looping for image and video generation, introducing ILSD and Any-Time inference for both masked generation and diffusion.
LoopDiT (technical blog) · XSA
Wenhao Chai's LoopDiT trains a recurrent DiT-B/4 with MeanFlow. It does not beat dense models at equal compute, and extra steps help more than extra loops. Shuangfei Zhai's Exclusive Self Attention is later used in Looped-DiT's core.
Diffusion as a Training Curriculum · Google
A recurrent denoiser with hidden states across steps and no timestep conditioning becomes an anytime solver for Sudoku and mazes. Diffusion provides the training curriculum; inference need not gradually denoise.
LoopVL
A 1B vision-language model trained from scratch with loops substantially beats a non-looped baseline at matched tokens. Training uses a fixed loop count; extra loops reduce performance.
Looped-DiT
Loops for open-ended text-to-image. The paper diagnoses washed-out spatial information, addresses it with deep supervision and self-modulating attention, and provides qualitative evidence of implicit visual reasoning.
LiFT
Trajectory targets at continuous depth coordinates replace final-answer targets at every loop. Extra inference loops continue to help, beating larger dense DiTs at equal compute.
Open questions
Does it hold at larger scales?
Looped-DiT uses 260M parameters at 512 resolution. LiFT uses ImageNet 256 without CFG, and loses to dense models at B/2 scale. In ELT, width remains the main scaling dimension; loops add inference computation. Billion-parameter latent text-to-image models have not yet been studied in this setting.
Video experiments are still limited
Only ELT among the three tests video, on the data-limited UCF-101 setting, where the authors suggest looping also reduces overfitting. LiFT lists budget-aware video generation as future work. Allocating loops by frame or difficulty may matter more for long videos.
How should loops be allocated across steps?
Current models use the same loop count at every denoising step. ELT suggests spending compute where it is most needed; Looped-DiT's gate already chooses loops per image. LiFT finds that the best (T, K) depends on budget: 16 loops × 10 steps is both worse and more expensive than 4 loops × 25 steps.
Can within-step loops and across-step memory be combined?
These text-to-image and video models reinitialize recurrence at each denoising step. Drozdova et al. show that an unsupervised hidden state carried across steps can accumulate substantial computation with backpropagation truncated to just 4 steps. Video generation naturally has states between frames and may be a suitable place to try this. DEQ-Flow's across-frame fixed points and FPDM's across-timestep reuse already save substantial compute.
What should inference noise do?
For unique-solution reasoning, maximum noise works best, acting as search within one trajectory. Generation needs diverse but stable samples, while trajectories in multi-solution tasks can drift between solutions. How to jointly design noise schedules and loop counts remains unanswered.
Combining loops with one-call generation
ELT notes that for one-call models such as consistency or drifting models, loop count would be the only adjustable quality parameter. Looped-DiT and LiFT position looping as another axis alongside few-step distillation; their combined effect still needs experiments.