ELT: Elastic Looped Transformers for Visual Generation

Sahil Goyal; Swayam Agrawal; Gautham Govind Anil; Prateek Jain; Sujoy Paul; Aditya Kusupati

Figure 1: Class-conditional Image Generation on ImageNet

\bm{256\times 256}

. Comparing our proposed Elastic Looped Transformers (left) against vanilla looped transformers (right), which simply repeat transformer layers during a forward pass of diffusion. As shown, standard looping only produces a coherent image when the inference loop count exactly matches its training setup (

L=8

), with quality severely degrading at all other values. In contrast, our approach maintains high-fidelity generation across several evaluated compute budgets.

Figure 2: Latent Trajectories of Standard vs. Elastic Looped Transformers .

{X_{k}}^{\text{int}}

&

{X_{k}}^{\text{max}}

represent output of intermediate (

L_{\text{int}}

) & final loops (

L_{\text{max}}

) respectively for

k^{th}

generation sampling step. Unlike standard recurrent models (left) where only the final iteration

{X_{k}}^{\text{max}}

reaches the solution space, our ILSD training (right) guides intermediate states

{X_{k}}^{\text{int}}

also toward the target space. This transformation shifts the model from a fixed-depth architecture to an Any-Time inference framework, supporting flexible computational budgets through early exits within a sampling step.

Figure 3: Overview of ELT framework . Training (left) utilizes shared parameters

\Theta

with recurrent loops and Intra-Loop Self Distillation (ILSD) to improve intermediate representations. This enables Any-Time inference (right), allowing the model to exit early and predict

X_{k+1}

from any intermediate block via a shared MLM head. Note that

X_{k}

represents the input for

k^{th}

sampling iteration and output images in ELT Any-Time Inference (right) assume

k=K

(last sampling step). Refer Section 3 for details of annotations.

Model	AR	FID $\downarrow$	IS $\uparrow$	# params	# steps	# Gflops
BigGAN-deep Δ [ 5 ]		7.0	171.4	160M	1	-
StyleGAN-XL Δg [ 60 ]		2.3	265.1	166M	1	-
Improved DDPM Δ [ 52 ]		12.3	-	280M	250	$>$ 150k
ADM + Upsample Δg [ 13 ]		3.9	215.8	554M	250	371k
LDM-4 Δg∗ [ 56 ]		3.6	247.7	400M	250	51.5k
DiT-XL/2 Δg∗ [ 53 ]		2.3	278.2	675M	250	59.5k

Table 1: Class-conditional Image Generation on ImageNet

256\times 256

. “# steps” refers to the number of neural network runs. Δ denotes values taken from prior publications. ∗ indicates usage of extra training data. g denotes use of classifier-free guidance [ 28 ] . Note that

L

in (

N\times L

) notation is inference loop count per sampling step.

Figure 4: Pareto front of FID vs. Inference GFLOPs . The black curve denotes the fit (FID =

1922.5\cdot G^{-0.95}+1.48

) over pareto-optimal configurations, representing the best achievable FID for a given computational budget. Points are labeled as

N\text{ layers}\times L\text{ loops}

. Results demonstrate that ELT scales as effectively as the baseline while remaining significantly more parameter-efficient. Scaling both the model dimension (

d

) and the number of loops (

L

) follows a predictable efficiency frontier, where larger models with fewer loops often compete with smaller models with more loops at specific target GFLOPs. Trends across faded points show scaling along number of loops

L

in inference, from a single run trained with a certain

L_{max}

loops. Refer Table 6 for the exact baseline configuration, including the specific number of layers for each comparison point.

Figure 5: Scaling with Parameters . We show the best-achievable FID as a function of model parameters (log scale). Each series corresponds to iso-inference-compute configuration for a specific model width (

d

). Points are labeled as

N\ \text{layers}\times L\ \text{loops}

. Results show that increasing parameter count via model width yields superior FID. Proposed ELT mechanism shows that even parameter-constrained models can achieve similar performance through recursive depth. Each ELT point in the figure is best inference configuration (

N\times L

) chosen from its corresponding training run (

N\times L_{\text{max}}

). Refer Table 6 for the exact baseline configuration, including the specific number of layers for each comparison point.

Table 7

Figure 6: Faster Training Convergence with ELT . While using

2\times

less parameters, ELT improves the training efficiency in diffusion framework.

Figure 9

Figure 8: Impact of ILSD across generation processes. ILSD improves performance for all

L

in inference especially when

L\neq L_{max}

. For w/o ILSD,

L_{max}

here refers to the number of loops the model is trained with. Results are on ImageNet

256\times 256

.

Figure 9: Class-conditional Image Generation on ImageNet

\bm{256\times 256}

. Comparing our proposed Elastic Looped Transformers (left) against vanilla looped transformers (right), which simply repeat transformer layers during a forward pass of diffusion. As shown, standard looping only produces a coherent image when the inference loop count exactly matches its training setup (

L=8

), with quality severely degrading at all other values. In contrast, our proposed ELT enhances vanilla looping with Self Distillation, maintaining a high-fidelity generation across several evaluated compute budgets.

Figure 10: Class-conditional Image Generation on ImageNet

\bm{256\times 256}

. Comparing our proposed Elastic Looped Transformers (left) against vanilla looped transformers (right), which simply repeat transformer layers during a forward pass of diffusion. As shown, standard looping only produces a coherent image when the inference loop count exactly matches its training setup (

L=8

), with quality severely degrading at all other values. In contrast, our proposed ELT enhances vanilla looping with Self Distillation, maintaining a high-fidelity generation across several evaluated compute budgets.

Figure 11: Class-conditional Image Generation on ImageNet

\bm{256\times 256}

. Qualitative results for ELT in diffusion framework,

16N\times 2L

inference configuration & trained with

L_{\text{max}}=2

.

Figure 12: Class-conditional Image Generation on ImageNet

\bm{256\times 256}

. Qualitative results for ELT in masked generative framework. Model scale: ELT-G,

8N\times 3L

inference config & trained with

L_{\text{max}}=4

(FID=1.9).

	Masked Generative Models		Diffusion Models
Config	ImageNet	UCF-101	ImageNet
Model Input
Latent shape	(1, 16, 16)	(4, 16, 16)	(4, 32, 32)
Codebook Size	1024	1024	—
Embed / Latent dim	1024	1024	4096
Training / Optimization

Table 5: Implementation Details. Hyperparameters and Configurations for Masked Generative Models (on ImageNet and UCF-101) and Diffusion Models (on ImageNet).

Model Scale	d_model	MLP dim	Heads	Layers (in Baseline)
S	384	1536	6	12
B	768	3072	12	12
L	1024	4096	16	24
XL	1152	4608	16	28
H	1280	5120	16	32
g	1408	6144	16	40

Table 6: Transformer Model Scaling . Architecture details for various model sizes.

Model Scale	d_model	MLP dim	Heads	Layers (in Baseline)
Small	2048	8192	16	16
Large	2048	8192	16	32

Table 7: Architecture Details for ELT in Diffusion framework.

Method	Step Time (ms)	Peak HBM (gb)
Baseline (32 layers)	130	7.7
Just Looping	130	5.1
ELT (w/ ILSD)	130	5.1
ELT opt (w/ ILSD)	118	5.8

ELT: Elastic Looped Transformers for Visual Generation

Abstract

Figures & tables

Appendix

Image Generation

LiFT: Loop Flow Transformers

UniGen-AR: Unifying Visual Generation with Auto-Regressive Modeling

One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers

Abstract

Image Generation

Explore similar work

LiFT: Loop Flow Transformers

UniGen-AR: Unifying Visual Generation with Auto-Regressive Modeling

One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers