ELT: Elastic Looped Transformers for Visual Generation

TL;DR

ELT (Elastic Looped Transformers) is a class of visual generative models that replaces a deep stack of unique transformer layers with a small block of weight-shared layers applied iteratively, decoupling parameter count from effective depth. The key training contribution is Intra-Loop Self Distillation (ILSD): a single forward pass through the maximum loop count $L$ acts as a teacher trajectory, and a stochastically-sampled intermediate loop $l_s \le L$ acts as a student that is a literal prefix of the teacher — so distillation adds essentially zero training overhead. The result is “Any-Time” inference: a single model traverses a quality-vs-compute Pareto curve at test time by varying the loop count, with no retraining. Headline numbers: 4× parameter reduction vs MaskGIT-XL at iso-inference- compute on class-conditional ImageNet (FID 2.0) and a 76M ELT beating the 306M MAGVIT-L on UCF-101 class-conditional video (FVD 72.8 vs 76).

Key claims

A weight-shared block of $L_b$ transformer layers applied $L$ times has the same effective depth $L_b \cdot L$ as a standard transformer with $L_b \cdot L$ unique layers, but only $L_b$ unique parameters — separating representational capacity from parameter count [§4].
Training a looped transformer for a fixed loop count $L$ alone produces uninterpretable intermediate-loop representations; quality at $l < L$ collapses [§4 / Fig. 8(b)].
Intra-Loop Self Distillation (ILSD): at every training step, sample $l_s \sim \mathrm{Uniform}(1, L-1)$ , compute losses for both the student exit at $l_s$ and the teacher exit at $L$ , and add a student-to-teacher distillation term (cross-entropy for masked generative models, sigmoid-weighted MSE for diffusion) with the teacher stop-gradded. Because the student trajectory is a strict prefix of the teacher’s, the only added forward cost is the intermediate-loop loss computation [§4, Algorithm 1].
A linearly-decayed mixing weight $\alpha$ from 1 → 0 over training anchors the student to ground truth early (while the teacher is still maturing) and shifts it to teacher mimicry later; the authors report no sensitivity to the decay rate as long as the transition is gradual [§4].
ELT-XL ( $L_b{=}6$ , $L{=}4$ ) matches MaskGIT-XL at FID 2.0 on class-conditional ImageNet-256 with 111M vs 446M parameters — a 4× parameter reduction at iso-inference-compute [Table 1].
On a DiT diffusion backbone with effective depth 32, ELT ( $L_b{=}16$ , $L{=}2$ , 1.1B params) reaches FID 2.83 vs 3.43 for a 2.1B 32-layer DiT baseline — same effective depth, half the parameters, lower FID [§5.2 / Table on DiT].
A minimum block size matters: ELT( $L_b{=}1$ , $L{=}32$ ) collapses to FID 10.30 despite 32 effective passes, so a single shared layer lacks expressiveness regardless of how many times it is iterated [§5.2].
Vanilla looped training (no ILSD) is competitive only when inference $L$ matches training $L$ ; performance degrades sharply at lower inference $L$ . ILSD recovers near-flat quality across $l = 1..L$ on both masked and diffusion frameworks [§5.4 / Fig. 8(a), 8(b)].
Beyond-training-depth extrapolation: on UCF-101 a model trained with $L{=}12$ reaches its peak FVD of 69.20 at $L{=}24$ — i.e. more loops at inference than at training improves quality [§5.4 / Fig. 7].
Parameter scaling is power-law in unique-layer count $L_b$ across model widths; gains are most pronounced when increasing both $L_b$ and width together [§5.2 / Fig. 5].
Throughput gains over the MaskGIT baseline: 2.9× (L), 3.3× (XL), 3.5× (H) at batch 8 on a single TPU slice — attributed to the shared-parameter footprint staying within accelerator memory capacity and avoiding repeated weight transfers across iterations. Throughput parity at the B scale because the baseline already fits in memory [§5.2 / throughput table].
ELT converges 2× and 1.6× faster than a 32-layer DiT baseline in iso-inference-compute settings at $(L_b{=}16, L{=}2)$ and $(L_b{=}8, L{=}4)$ respectively, on the same training token budget [§5.2 / Fig. 6].
On the data-constrained UCF-101 setup (13.7M training tokens) a 76M ELT outperforms a 306M MAGVIT-L baseline (FVD 72.8 vs 76) at matched inference compute, suggesting the weight-sharing prior regularizes effectively under data scarcity [Table 4].

Method

ELT defines a recurrent block of $L_b$ unique transformer layers parameterized by $\theta$ , and applies it $L$ times in sequence within a single sampling step. Effective depth is $L_b \cdot L$ ; unique-parameter depth is $L_b$ . This sits naturally inside both masked generative models (MaskGIT-style parallel decoding with discrete VQ tokens) and continuous-latent diffusion transformers (DiT-style flow / DDPM denoising), where the parameter sharing is across the loop dimension rather than across sampling steps.

ILSD is the training-side contribution. At each step the model is treated as a dual-path system over a single set of parameters: the teacher path runs $L$ loops and is supervised by the ground-truth loss; a student path sampled uniformly at $l_s \in \{1, \dots, L-1\}$ exits early and is supervised by both the ground-truth loss and a distillation loss against the (stop-gradded) teacher’s prediction. Crucially, because the student trajectory is a literal prefix of the teacher’s forward pass, the marginal training cost is the student’s loss computation — no separate teacher forward is needed, unlike standard online distillation. The joint loss is $\mathcal{L} = \mathcal{L}_{\text{gt}}^{T} + \alpha\,\mathcal{L}_{\text{gt}}^{S} + (1{-}\alpha)\,\mathcal{L}_{\text{ILSD}}$ with $\alpha$ linearly decayed from 1 to 0 over training. Loss form follows the framework: cross-entropy for masked generation, sigmoid- weighted MSE for diffusion.

At inference, varying $L$ traverses the quality–compute Pareto curve. The paper frames this as an intra-step compute lever, orthogonal to inter-step levers like consistency-model step reduction or progressive distillation [§2 / “Relation to Few-Step and Consistency Models”].

Results

ImageNet-256 class-conditional (MaskGIT-style discrete): ELT-XL ( $L_b{=}6$ , $L{=}4$ , 111M params) matches MaskGIT-XL FID 2.0 at 4× lower parameter count and the same 3.9k GFLOPs / 24 sampling steps [Table 1].
ImageNet-256 class-conditional (DiT-style continuous): ELT( $L_b{=}16$ , $L{=}2$ , 1.1B params, effective depth 32) reaches FID 2.83 vs 3.43 for a 2.1B 32-layer DiT baseline. The smaller ELT( $L_b{=}8$ , $L{=}4$ , 539M params) reaches FID 3.16, also beating the DiT-32 baseline at ~4× fewer parameters [§5.2 table on DiT].
UCF-101 class-conditional video (MAGVIT-style): 76M ELT( $L_b{=}6$ , $L{=}2$ ) reaches FVD 72.8 (vs 76 for 306M MAGVIT-L) at the same 4.3k GFLOPs and 12 sampling steps. Pushing to $L{=}4$ at the same 76M reaches FVD 60.8 at 13k GFLOPs and 24 sampling steps [Table 4].
Throughput: 2.9× / 3.3× / 3.5× speedup over MaskGIT baseline at scales L / XL / H, measured at batch 8 on a single TPU slice [§5.2].
Training convergence: 2× and 1.6× wall-clock speedup on ImageNet-256 vs a 32-layer DiT at matched effective depth and inference FLOPs [Fig. 6].
ILSD ablation: without ILSD, FID at $l < L$ degrades drastically (qualitative example in Fig. 8); with ILSD, FID is roughly flat from $l=1$ to $l=L$ , with the headline numbers achievable at exit points well below $L$ [Fig. 8(a)/(b)].
Extrapolation: UCF-101 model trained with $L{=}12$ achieves best FVD (69.20) at inference $L{=}24$ — ILSD enables modest extrapolation beyond training depth [§5.4].

Why it’s interesting

ELT is the first looped-transformer paper on the wiki to apply the idea to visual generation at scale (prior wiki references to recurrence have been in the LLM / reasoning context — see Universal Transformers, Mixture-of-Recursions, and the inference-time-scaling concept page). The Any-Time-inference framing makes loop count an explicit deployment-time lever, conceptually adjacent to Inference-Time Scaling‘s axis of “trade inference compute for capability” — but operating within a sampling step rather than across steps, parallel chains, or test-time weight updates. ILSD is essentially zero-overhead self-distillation because the student is a prefix of the teacher trajectory, which is a sharper version of the “internal self-distillation” framing explored in No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA) and Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis (Self-Flow) on the Diffusion training efficiency page — there the supervision is across noise levels and layers in a non-shared-weight network; here it is across loop counts in a shared-weight network. The data-constrained UCF-101 win (76M beating 306M) supports the authors’ argument that weight sharing acts as a regularizer when data is the bottleneck — a counterpoint to the “scale parameters proportionally to data” framing in On the Slow Death of Scaling (On the Slow Death of Scaling).