ELT: Elastic Looped Transformers for Visual Generation — part 1/3

Cleaned-text code points [0, 15567). Bibliography omitted here; source.md includes it.

\uselogo \reportnumber \teaser

# ELT: Elastic Looped Transformers for Visual Generation

###### Abstract

We introduce Elastic Looped Transformers (ELT), a highly parameter-efficient class of visual generative models based on a recurrent transformer architecture. While conventional generative models rely on deep stacks of unique transformer layers, our approach employs iterative, weight-shared transformer blocks to drastically reduce parameter counts while maintaining high synthesis quality. To effectively train these models for image and video generation, we propose the idea of Intra-Loop Self Distillation (ILSD), where student configurations (intermediate loops) are distilled from the teacher configuration (maximum training loops) to ensure consistency across the model’s depth in a single training step. Our framework yields a family of elastic models from a single training run, enabling Any-Time inference capability with dynamic trade-offs between computational cost and generation quality, with the same parameter count. ELT significantly shifts the efficiency frontier for visual synthesis. With 4 × \times reduction in parameter count under iso-inference-compute settings, ELT achieves a competitive FID of 2.0 on class-conditional ImageNet 256 × 256 256\times 256 and FVD of 72.8 on class-conditional UCF-101.

###### keywords

## 1 Introduction

Traditional techniques to increase compute capacity in deep learning models, such as stacking deeper layers or increasing network width, inevitably lead to a proportionally larger memory footprint. Recurrence offers a powerful alternative paradigm, enabling models to utilize large amounts of compute while maintaining a minimal memory footprint by leveraging the same set of parameters repeatedly. This architectural efficiency parallels the biological visual systems [ 37 , 38 ] , where recurrent processing, rather than strictly feedforward pathways, is essential for resolving complex visual inputs.

While looping of transformers was popularized by Universal Transformers [ 9 ] and has recently empowered language models with stronger reasoning capabilities [ 61 , 75 ] , its potential for high-fidelity visual generation remains largely untapped. From a practical standpoint, compared to the traditional models, Looped Transformers (a) are extremely parameter efficient and can perform significantly more compute (FLOPs) per parameter, (b) can have higher throughput by minimizing the “memory wall” bottleneck. By reusing a compact set of shared parameters across L L iterations, looped transformers achieve up to L × L\times higher arithmetic intensity (FLOPs per byte transferred from memory) compared to standard transformers of equivalent depth. This avoids the cost of repeated transfers between different units of the accelerator (GPUs/TPUs) required in large transformers, and (c) can exhibit robustness against overfitting in data-constrained settings.

Despite the parameter efficiency of looped architectures, their training remains challenging because intermediate representations often remain uninterpretable until the final loop (see Figure 2 ). We address this by introducing Elastic Looped Transformers (ELT), a class of generative models designed for progressive refinement. Unlike traditional recurrent transformers, ELT is optimized to provide meaningful, high-quality synthesis even at intermediate repeats. This enables Any-Time (elastic) inference capability - where a single model can scale its compute based on available resources without sacrificing much on generation quality. We present a pictorial representation of our proposed method in Figure 3 and generation results for Any-Time inference capability in 1 .

To achieve this functional flexibility across loops, we propose an Intra-Loop Self Distillation (ILSD) algorithm for training looped transformers. Rather than treating the loop as a fixed-depth process, our framework operates as a dual-path system: a teacher path executes the full loop count to provide a high-fidelity target, while a student path, defined strictly as a subset of the teacher’s trajectory, learns to produce comparable results with fewer iterations. Note that both the paths have the same parameter count. By framing the student process as a literal subset of the teacher’s forward propagation, we ensure there is no additional overhead during training. This approach forces the shared parameters to compress complex transformations into earlier loops. Consequently, the model does not just fix the output at the end. It learns an efficient, progressive refinement process that maintains the generation quality regardless of the exit point as motivated in Figure 2 . Our contributions and findings are summarized as follows:

State-of-the-Art Parameter Efficiency : Through the reuse of block of transformer layers across loops, ELT achieves a competitive FID of 2.0 on class-conditional ImageNet 256 × 256 256\times 256 and FVD of 72.8 on class-conditional UCF-101. This represents a 4 × \times reduction in parameters compared to baselines MaskGIT [ 7 ] (image generation) and MAGVIT [ 76 ] (video generation), while matching or improving their performance under iso-inference-compute settings.

Elastic/Any-Time Inference : By treating the looped block as an iterative refiner, our models enable Any-Time inference [ 83 ] , allowing to traverse the pareto frontier of quality versus compute at test-time without retraining. This allows to serve diverse deployment tiers: from latency-critical on-device generation (few loops) to high-fidelity cloud rendering (more loops).

Scalability : While model size remains a primary driver of quality, recursive looping provides a unique test-time compute lever that scales predictably across both Masked Generative Transformers [ 7 , 76 , 77 ] and Diffusion Transformers [ 53 ] .

## 2 Related Work

Recursive Architectures: The principle of recursively applying a shared block of parameters to enhance model efficiency is a well-established concept. This approach, often referred to as looping has been shaped by Universal Transformers [ 9 ] , which introduced the idea of iterating over a single transformer layer. Notably, there has been a resurgence of interest in this area. Saunshi et al. [61] demonstrated the power of looping for sophisticated reasoning tasks. Gatmiry et al. [19] showed that Looped Transformers can learn to implement multi-step gradient descent for in-context learning, providing a deeper understanding of their capabilities. Mixture-of-Recursions [ 3 ] further explored input-dependent dynamic and adaptive depth in looped transformers. Mixture-of-Recursions-VIT [ 45 ] extends it for image understanding. Fan et al. [16] utilized looping for length generalization. Geiping et al. Geiping et al. [21] demonstrated that scaling test-time compute via recurrent depth allows language models to perform complex latent reasoning.

Deep Equilibrium Models (DEQs) [ 4 , 54 , 22 , 47 , 1 , 17 ] , instead of unrolling a weight-tied layer for a fixed number of iterations, define the output as the fixed point of a non-linear transformation. Unlike DEQs that rely on black-box solvers for an analytical fixed point, our ELT framework explicitly optimizes unrolled intermediate states via Intra-Loop Self Distillation (ILSD), retaining the flexibility of Any-Time inference without requiring the network to reach a strict analytical fixed point.

Parameter-Efficient Visual Generation: Standard efficiency techniques [ 48 ] that work across deep learning models also help in speeding up visual generation models. MobileDiffusion [ 80 ] prunes redundant residual blocks and replaces standard layers with separable convolutions in UNet to get an optimized architecture ( ∼ \sim 400M) for on-device visual generation. EdgeFusion [ 6 ] employs BK-SDM, a lightweight Stable Diffusion variant, and refines the step distillation process of Latent Consistency Model. MaGNeTS [ 24 ] trains a family of nested transformers [ 41 , 42 ] with schedules of model sizes over the generation process, without increasing the parameter count.

Elastic Visual Generation: The paradigm of Any-Time or elastic generation focuses on decoupling model’s parameter count from its computational depth. E-DiT [ 70 ] introduces adaptive block skipping and MLP width reduction, allowing a single model to traverse varying computational budgets without retraining. In the context of visual reasoning, LoopViT [ 63 ] uses a weight-tied recursive architecture, employing a parameter-free dynamic exit mechanism to halt inference based on uncertainty of prediction. EvoSearch [ 26 ] proposes a search-based strategy that optimizes sampling trajectories at inference time. Unlike these methods that rely on architectural skipping or external search, our ELT framework equipped with Intra-Loop Self Distillation (ILSD) algorithm directly regularizes the recursive process, ensuring stability across the entire loop spectrum. Methods like ALIT [ 14 ] , FlexTok [ 2 ] , One-D-Piece [ 49 ] , ElasticTok [ 74 ] , CAT [ 62 ] , and DC-DiT [ 84 ] use tail dropping to allow elasticity in token sequence length.

Relation to Few-Step and Consistency Models : Recent approaches such as Consistency Models [ 67 ] and progressive distillation [ 58 ] address inference efficiency by reducing the number of sampling steps (inter-step acceleration). ELT is fundamentally orthogonal: it reduces compute within each sampling step by varying the loop count L L (intra-step acceleration). These two axes are complementary, we can combine ELT with few-step methods to achieve further efficiency gains. Moreover, unlike consistency models which require specialized training objectives and architectural constraints to enable variable-step inference, ELT’s elastic capability emerges naturally from ILSD applied to standard training objectives. We note that ELT is particularly compelling for one-step generative paradigms, where the loop count L L becomes the sole lever for controlling the compute-quality trade-off at inference time.

## 3 Preliminaries

Masked Generative Models : Masked Generative Image Transformer (MaskGIT) [ 7 ] introduced a novel approach to image generation that significantly differs from traditional autoregressive models. In autoregressive decoding, images are generated sequentially, one pixel/token at a time, following a raster scan order [ 15 , 40 , 71 , 44 ] . This sequential approach is computationally inefficient, as each token is conditioned only on the previously generated tokens, leading to a bottleneck in processing time. MaskGIT generates all tokens of an image simultaneously, while iteratively refining them. This method enables significant acceleration in the decoding process. The tokens are discrete and obtained using Vector Quantized (VQ) autoencoders, learned with self-reconstruction and photo-realism losses [ 76 ] . The iterative parallel decoding process is represented as: 𝐗 k ← Mask ∘ Sample ⁡ ( M ⁡ ( 𝐗 k − 1 , c ) , k ) \mathbf{X}_{k}\leftarrow\mathrm{Mask\circ Sample}({M}({\mathbf{X}}_{k-1},c),k) (1) where 𝐗 ∈ ℤ ≥ 0 N \mathbf{X}\in\mathbb{Z}^{N}_{\geq 0} , are the input tokens, N N is the number of tokens, k ∈ [ 1 , K ] k\in[1,K] denote the iteration number, with K K being the total number of iterations, 𝐗 0 \mathbf{X}_{0} is either completely masked for full generation, and partially masked for conditional generation tasks like frame prediction, c c is the category of image/video under generation. The Sample \mathrm{Sample} function utilizes logits predicted by the model M ( . ) {M(.)} , introduces certain randomness, and sorts them by confidence, unmasking only top-p tokens while masking the rest.

Diffusion Models : Diffusion models [ 29 , 66 ] generate data by learning to reverse a process that gradually corrupts a signal 𝐗 0 \mathbf{X}_{0} into Gaussian noise 𝐗 T \mathbf{X}_{T} through a predefined noise schedule. At its core, diffusion denoising/sampling proceeds iteratively as follows: 𝐗 t − 1 ← α 1 ​ ( t ) ⋅ 𝐗 t + α 2 ​ ( t ) ⋅ M ⁡ ( 𝐗 t , t , c ) + α 3 ​ ( t ) ⋅ 𝐳 \displaystyle\mathbf{X}_{t-1}\leftarrow\alpha_{1}(t)\cdot\mathbf{X}_{t}+\alpha_{2}(t)\cdot M(\mathbf{X}_{t},t,c)+\alpha_{3}(t)\cdot\mathbf{z} where 𝐗 t \mathbf{X}_{t} is the (partially) denoised vector at time t t , α 1 ​ ( t ) , α 2 ​ ( t ) , α 3 ​ ( t ) \alpha_{1}(t),\alpha_{2}(t),\alpha_{3}(t) are time-dependent scalars defined by the noise schedule and 𝐳 \mathbf{z} is a sample drawn from standard Normal distribution. The model M M takes as input 𝐗 t \mathbf{X}_{t} along with time t t and class label c c . From the perspective of the forward diffusion process, the generative task can be thought of as predicting the noise added to a latent representation 𝐗 t \mathbf{X}_{t} at timestep t t , conditioned on a class label c c . While traditionally architectures based on U-Nets [ 57 ] have been used, the Diffusion Transformer (DiT) [ 53 ] architecture shifts away from this design by treating image latents as sequences of tokens, and using transformer blocks for processing these tokens. Similar to MaskGIT, we explore replacing the typical DiT transformers with Elastic Looped Transformers for denoising at time t t .

In summary, both masked generative models and diffusion models involve recursive refinement over multiple sampling steps, sharing parameters across sampling steps. Unlike standard transformers, our ELT framework aligns the architecture of model M M with the progressive refinement process by implementing M M as a recurrent, weight-shared transformer blocks. This allows the model to perform recursive refinement within each sampling step, providing a test-time compute lever to trade-off inference speed and generation quality.

## 4 Method

Looping Mechanism : Let the number of transformer layers to be looped be N N and number of loops per sampling step be L L . The total effective depth for a single sampling or denoising step of the generation process is then N × L N\times L . Let f θ i ​ ( 𝐱 ) f_{\theta_{i}}(\mathbf{x}) denote a single transformer layer with parameters θ i \theta_{i} . We define a composite block g Θ ​ ( x ) g_{\Theta}(x) consisting of N N unique transformer layers with parameters Θ = { θ 1 , θ 2 , … , θ N } \Theta=\{\theta_{1},\theta_{2},\dots,\theta_{N}\} as follows: g Θ ( 𝐱 ) = f θ N ( f θ N − 1 ( ⋯ f θ 1 ( 𝐱 ) ) ) \displaystyle g_{\Theta}(\mathbf{x})=f_{\theta_{N}}(f_{\theta_{N-1}}(\cdots f_{\theta_{1}}(\mathbf{x})))

In a standard transformer model with total depth 𝒟 = N × L \mathcal{D}=N\times L , the effective transformation F 𝒟 ​ ( x ) F_{\mathcal{D}}(x) would require 𝒟 \mathcal{D} sets of unique parameters. In looping , we reuse the same block of parameters Θ \Theta for L L successive applications, resulting in only N N unique layers of parameters. The effective transformation for a N × L N\times L configuration is given by: F ( N , L ) ​ ( 𝐱 ) = g Θ ( g Θ ( ⋯ g Θ ( 𝐱 ) ) ) ⏟ L ​ loops ≡ g Θ L ​ ( 𝐱 ) \displaystyle F_{(N,L)}(\mathbf{x})=\underbrace{g_{\Theta}\left(g_{\Theta}(\cdots g_{\Theta}(\mathbf{x}))\right)}_{L\text{ loops}}\equiv g_{\Theta}^{L}(\mathbf{x})

This looping architecture decouples physical model size from computational depth (see Figure 3 for a visual overview). The parameter count ( Θ \Theta ) is constrained by the number of unique blocks ( N N ), while representational capacity and depth ( 𝒟 \mathcal{D} ) scale with loop count ( L L ). This setup provides two primary advantages: extreme parameter efficiency and high throughput. See Section 5.2 for details.