ELT: Elastic Looped Transformers for Visual Generation — part 2/3
Cleaned-text code points [15567, 31388). Bibliography omitted here; source.md includes it.
Intra-Loop Self Distillation (ILSD) : In a standard weight-tied (looped) transformer, the model is typically optimized only for its final output after fixed L max L_{\text{max}} iterations i.e. the loop count for which it is trained. However, this creates a “black box” internal trajectory where intermediate loops ( L < L max L<L_{\text{max}} ) may produce suboptimal representations until the final projection layer. By treating the looped block g Θ g_{\Theta} as an iterative refiner, we can motivate a training objective that ensures the model remains useful at multiple depths. This not only encourages the model to learn a more robust, incremental transformation but also allows for elastic inference, where the model can be exited early with minimal performance drop. Towards this end, we propose Intra-Loop Self Distillation (ILSD).
From a distillation perspective, ILSD leverages the fact that a model with more loops ( L max L_{\text{max}} ) naturally possesses a more mature and refined representational space than its shallower version ( L int L_{\text{int}} ), even though they have the same unique parameters. By treating the full-depth model as an internal teacher, we provide a high-fidelity, low-variance signal for the shallower student to follow (see Figure 3 ). This forces the shared parameters Θ \Theta to compress complex transformations into fewer steps, effectively distilling the knowledge of the deep model into the early stages of the computation.
Stochastic Student Sampling ( S 3 S^{3} ) : We train with a fixed number of loops, L max L_{\text{max}} , for a block of N N layers with unique parameters. During each training step, we treat the model as a dual-path system sharing a single set of parameters Θ \Theta . We define a teacher path that executes the full L max L_{\text{max}} loops and a stochastic student path that exits early at an intermediate loop L int L_{\text{int}} . The intermediate path receives supervision from the ground truth labels as well as the online teacher with maximum training loop count L max L_{\text{max}} .
We denote the training loss for the output of a configuration N × L N\times L as ℒ Θ ( F ( N , L ) ( 𝐱 ) , 𝐲 ) \mathcal{L}_{\Theta}(F_{(N,L)}(\mathbf{x}),\mathbf{y}) where 𝐲 \mathbf{y} denote the ground-truth for a certain input masked / denoise level. At every training iteration, we randomly sample an intermediate loop count L int L_{\text{int}} from uniform distribution such that L min ≤ L int < L max L_{\text{min}}\leq L_{\text{int}}<L_{\text{max}} . Note that L min L_{\text{min}} is just used for constraining the student sampling distribution. The effective joint loss, ℒ Θ ILSD \mathcal{L}^{\text{ILSD}}_{\Theta} , is computed as:
ℒ Θ ILSD = \displaystyle\mathcal{L}^{\text{ILSD}}_{\Theta}= ℒ GT ( F ( N , L max ) ( 𝐱 ) , 𝐲 ) \displaystyle\mathcal{L}^{\text{GT}}(F_{(N,L_{\text{max}})}(\mathbf{x}),\mathbf{y}) (1) Ground-truth for teacher + \displaystyle+ λ ℒ GT ( F ( N , L int ) ( 𝐱 ) , 𝐲 ) \displaystyle\lambda\mathcal{L}^{\text{GT}}(F_{(N,L_{\text{int}})}(\mathbf{x}),\mathbf{y}) (2) Ground-truth for student + \displaystyle+ ( 1 − λ ) ℒ dist ( F ( N , L int ) ( 𝐱 ) , sg ( F ( N , L max ) ( 𝐱 ) ) ) \displaystyle(1-\lambda)\mathcal{L}^{\text{dist}}(F_{(N,L_{\text{int}})}(\mathbf{x}),\text{sg}(F_{(N,L_{\text{max}})}(\mathbf{x}))) (3) Intra-Loop Self Distillation with L int ∼ 𝒰 ( L min , L max ) \displaystyle L_{\text{int}}\sim\mathcal{U}(L_{\text{min}},L_{\text{max}}) Stochastic Student Sampling where sg is stop-grad for teacher ( L max L_{\text{max}} ) in ILSD, λ \lambda is hyperparameter controlling the weight between the ground truth and distillation losses. We introduce a curriculum for λ \lambda and linearly decay it from 1 1 to 0 0 as training progresses. This initially anchors the student to reliable ground-truth labels while the teacher is still untrained and gradually shift to mimicking the teacher’s predictions once they have matured. We found this linear schedule to be effective across all our experiments and did not observe sensitivity to the decay rate, as long as the transition is gradual.
Loss Formulation : The exact forms of ℒ GT \mathcal{L}^{\text{GT}} and ℒ dist \mathcal{L}^{\text{dist}} depend on the algorithm used for training. For masked generative models with discrete tokens, we use the cross-entropy loss for both ground-truth and distillation: ℒ GT \displaystyle\mathcal{L}^{\text{GT}} = − ∑ i ∈ M a s k log P ( N , L int ) ( y i ∣ 𝐱 m a s k ) \displaystyle=-\sum_{i\in Mask}\log P_{(N,L_{\text{int}})}(y_{i}\mid\mathbf{x}_{mask}) ℒ dist \displaystyle\mathcal{L}^{\text{dist}} = − ∑ i ∈ M a s k ∑ v ∈ 𝒱 P ( N , L max ) ( v ∣ 𝐱 m a s k ) log P ( N , L int ) ( v ∣ 𝐱 m a s k ) \displaystyle=-\sum_{i\in Mask}\sum_{v\in\mathcal{V}}P_{(N,L_{\text{max}})}(v\mid\mathbf{x}_{mask})\log P_{(N,L_{\text{int}})}(v\mid\mathbf{x}_{mask}) where 𝐱 m a s k \mathbf{x}_{mask} is the masked input, y = { y i } i ∈ Mask y=\{y_{i}\}_{i\in\text{Mask}} represents the ground-truth tokens for the masked positions, and 𝒱 \mathcal{V} is the full vocabulary of the tokenizer. For diffusion, we use the sigmoid-weighted Mean Squared Error (MSE) for both ground-truth and distillation losses. Let 𝐱 t \mathbf{x}_{t} be the noised version of ground truth latent 𝐱 0 \mathbf{x}_{0} , the ground-truth loss is: ℒ GT = w ( t ) ‖ F ( N , L ) ( 𝐱 t ) − 𝐱 0 ‖ 2 2 \displaystyle\mathcal{L}^{\text{GT}}=w(t)\norm{F_{(N, L)}(\mathbf{x}_t) - \mathbf{x}_0}_{2}^{2} where L L is L int L_{\text{int}} for student and L max L_{\text{max}} for teacher. w ( t ) w(t) is a time-dependent sigmoid weighting term [ 34 ] . The distillation loss is: ℒ dist = w ( t ) ‖ F ( N , L max ) ( 𝐱 t ) − F ( N , L int ) ( 𝐱 t ) ‖ 2 2 \displaystyle\mathcal{L}^{\text{dist}}=w(t)\norm{F_{(N, L_{\text{max}})}(\mathbf{x}_t) - F_{(N, L_{\text{int}})}(\mathbf{x}_t)}_{2}^{2}
Note that gradients from both computational paths, corresponding to L int L_{\text{int}} and L max L_{\text{max}} , update the single, shared set of block parameters Θ \Theta . This joint optimization provides a significantly richer training signal; the shared block g Θ g_{\Theta} is forced to learn a transformation that is not only effective at L int L_{\text{int}} loops but remains incrementally useful for the subsequent iterations up to L max L_{\text{max}} loops. This constraint prevents the model from learning a shortcut that might minimize loss at a specific depth but fail when composed further. Consequently, the shared block generalizes better to lower depths, leading to higher performance even with fewer loops. It is interesting to note that unlike traditional distillation, where we have to forward propagate through both the student and teacher models separately, in the proposed way of Intra-Loop Self Distillation (ILSD), the training overhead is minimal, as the computation of F ( N , L int ) ( 𝐱 ) F_{(N,L_{\text{int}})}(\mathbf{x}) is a strictly required intermediate step for computing F ( N , L max ) ( 𝐱 ) F_{(N,L_{\text{max}})}(\mathbf{x}) i.e. the student trajectory ( L int L_{\text{int}} ) is a strict prefix of the teacher trajectory ( L max L_{\text{max}} ). Refer Algorithm 1 and Algorithm 2 for details of our training and inference algorithms respectively.
## 5 Experiments and Results
We conduct extensive experiments using masked generative transformers and diffusion transformers to demonstrate the generality and efficacy of our approach on class-conditional image generation. We also experiment with class-conditional video generation using masked generative transformers. We first detail our experimental setup and then present the results.
### 5.1 Experimental Setup
Datasets : We experiment on ImageNet 256 × 256 256\times 256 [ 10 ] for image generation and UCF-101 [ 68 ] for class-conditional video generation.
Implementation Details : (i) Masked Generative Transformers : We use pretrained tokenizers from MaskGIT [ 7 ] (images) and MAGVIT [ 76 ] (videos) with a codebook size of 1024 tokens. Image models are trained at 256 × 256 256\times 256 resolution, compressed to 16 × 16 16\times 16 discrete tokens with an embedding dimension of 1024. Video models are trained for 16 × 128 × 128 16\times 128\times 128 sequences, compressed to 4 × 16 × 16 4\times 16\times 16 tokens. Following MaskGIT, we adopt the BERT [ 12 ] architecture as the transformer backbone and perform experiments at several model scales to understand the scaling behaviors of ELT. We train all models for 270 epochs unless otherwise specified. We employ a cosine schedule for unmasking tokens during inference. For image generation, we use classifier-free guidance by dropping class condition labels for 10 % 10\% of the training batches. (ii) Diffusion Transformers: We use a pretrained Stable Diffusion v1.4 VAE [ 56 ] model to map 256 × 256 256\times 256 images into a continuous 32 × 32 × 4 32\times 32\times 4 latent space (8× spatial downsampling). We train a DDPM-style diffusion model which operates on these latents using a DiT architecture. We employ a shifted cosine noise schedule and sigmoid-weighted MSE loss for training [ 34 ] . Models are trained for 500K steps with a batch size of 512 512 using Adam. Unless mentioned otherwise, sampling uses 512-step DDPM with classifier-free guidance scale 3.0 3.0 . See Appendix for more details.
Evaluation Metrics and Efficiency : To evaluate the quality of synthesized content, we employ Fréchet Inception Distance (FID) [ 27 , 13 ] and Inception Score (IS) [ 59 ] for image generation tasks, and Fréchet Video Distance (FVD) [ 69 ] for video generation. Beyond generative quality, we evaluate model efficiency using inference-time GFLOPs and throughput (samples generated per second). In the proposed N × L N\times L design space, for a fixed model scale and block size N N , both computational complexity and latency scale linearly with the number of loops L L . This relationship allows us to precisely navigate the trade-off between generation quality and inference speed by modulating the loop count.
### 5.2 Image Generation
Comparison with Baselines : We present the results for 256 × 256 256\times 256 image generation on ImageNet-1k in Table 1 . Despite using 4 × \times less parameters, ELT-XL achieves same FID of 2.0 as MaskGIT-XL, which is the base setup for ELT. As shown in recent literature, using a superior tokenizer [ 79 , 77 , 73 ] or optimized training & inference configurations [ 50 , 51 ] can further boost ELT’s performance.
We also present comparisons of ELT using diffusion models in Section 5.2 . Notably, ELT outperforms the baseline model with 32 layers (FID 3.43) using iso-inference-compute 8 N × 4 L 8N\times 4L (FID 3.16) and 16 N × 2 L 16N\times 2L (FID 2.83) settings, achieving 4 × \textbf{4}\times and 2 × \textbf{2}\times reduction in parameter count respectively. The 1 N × 32 L 1N\times 32L configuration (FID 10.30) reveals that a single unique transformer layer, despite 32 effective passes, lacks the representational capacity for high-fidelity generation, highlighting that a minimum block size N N is necessary to provide sufficient architectural expressiveness within each iteration. While vanilla looping (without ILSD) gives competitive performance when running inference with same loops as training ( L = L m a x L=L_{max} ), performance degrades drastically for lower number of loops which is mitigated by ILSD as show in Figure 8(b) .
Qualitative Results : 1 , Figure 9 , and Figure 10 compare ELT against vanilla looped transformers within a diffusion framework. It is clear that ELT unlocks Any-Time inference capability providing dynamic trade-off between generation quality and inference speed. Figure 11 and Figure 12 visualize the generation results of ELT in diffusion and masked generative framework respectively.
Scaling Inference GFLOPs and Pareto Front : Figure 4 illustrates the trade-off between generation quality (FID) and inference compute (GFLOPs). The pareto front (black curve) reveals that while increasing the loop count ( L L ) consistently improves FID (faded points) for a fixed unique layers ( N N ), the gains eventually diminish, where transitioning to the next architecture scale becomes more performant than over-looping smaller models. Crucially, ELT allows for Any-Time inference, we can traverse the pareto curve at test-time by simply adjusting L L to meet specific hardware constraints without retraining.
Scaling Parameters : We further investigate the relationship between model capacity and performance in Figure 5 . By plotting the best FID achieved at each parameter budget, we observe a clear power-law scaling trend across all model widths. While increasing the number of unique layers ( N N ) reduces FID, the gains are most pronounced when accompanied by an increase in model width ( d d ). Specifically, the d = 1536 d=1536 configuration ( G G ) achieves the lowest overall FID with only 8 8 unique layers, while the full G G model has 48 48 unique layers. However, for on-device budgets, smaller looped variants (L and XL) remain highly efficient, providing a flexible scale-to-performance ratio that is critical for resource-constrained visual generation.
Faster Training Convergence : Our method demonstrates significantly faster convergence than standard DiT architectures [ 53 ] . As illustrated in Figure 6 , ELT-based diffusion models achieve 2 × 2\times and 1.4 × 1.4\times speedups when using configurations of 16 N × 2 L max 16N\times 2L_{\text{max}} and 8 N × 4 L max 8N\times 4L_{\text{max}} , respectively, in iso-inference-compute settings with N = 32 N=32 DiT baseline. Note that effective depth 𝒟 \mathcal{D} is same for both ELT and DiT baseline (32).
ELT has high Throughput : ELT utilizes a compact set of shared parameters and maintain its major weight footprint closer to the accelerator computation unit. This avoids the cost of repeated memory transfers typically required in standard models. We evaluate the efficiency of our proposed architecture by measuring the throughput ratio relative to the baseline across various model scales with an inference batch size of 8, on 1 TPU slice. We choose the best performing ELT inference configuration from each model scale (refer Figure 5 ) and report their throughput gains in Section 5.2 . Our method delivers significant performance gains across model scales L, XL and H. Note that these gains are contingent on the model size: as long as the shared parameters remain within memory capacity, the architecture eliminates redundant memory movement across iterations. For model scale B, the baseline (MaskGIT [ 7 ] ) is small enough to fit entirely within device memory capacity, eliminating the transfer penalty. Generally speaking, ELT offers extreme parameter efficiency, which would reduce significant memory transfers even for bigger models which need to be sharded along multiple devices.
### 5.3 Video Generation
We use the MAGVIT [ 76 ] framework to train parallel decoding based video generation models. We summarize the results for class-conditional video generation on UCF-101 in Table 4 . Our compact 76M ELT model outperforms MAGVIT baseline (FVD 72.8 vs 76) on data-constrained settings of UCF-101 ( ∼ \sim 13.7M training tokens) in iso-inference-compute settings. Scaling the compute further with number of loops and sampling steps gives a boost in performance, reaching FVD of 60.8. This suggests that looped transformers can exhibit robustness against overfitting in data-constrained regimes like UCF-101, effectively regularizing the learning process while maintaining the expressive capacity for high-quality generation.
### 5.4 Intra-Loop Self Distillation (ILSD) drives Elasticity