Skip to content

ELT: Elastic Looped Transformers for Visual Generation

Conference: ECCV 2026
Paper: ECCV Official Page
Area: Image Generation
Keywords: looped transformers, parameter sharing, intra-loop self-distillation, elastic inference, visual generation

TL;DR

ELT repeatedly applies shared transformer layers within each generation step and distills the full-loop prediction into intermediate exits, allowing a single model to adjust its loop count to the available budget; at matched inference compute, it achieves ImageNet FID 2.0 with 111M parameters, compared with 446M for the corresponding MaskGIT-XL baseline.

Background & Motivation

Although MaskGIT, MAGVIT, and DiT operate on either discrete visual tokens or continuous latents, they all repeatedly invoke a generative network. Increasing the computational depth of that network usually means stacking more distinct transformer layers, which also increases parameter storage and weight movement. Recurrence offers another option: keep a small group of layers and repeatedly apply the same weights to obtain a deeper effective computation path. This reduces the number of unique parameters; it does not eliminate the repeated computation.

The difficulty is that a conventional looped model only needs its final output to be correct. Intermediate states can remain transitional representations that cannot be decoded directly. A model trained with 8 loops does not necessarily produce a slightly worse but usable image after 4 loops; an early exit can instead break generation. Weight sharing alone therefore yields a smaller model, but not automatically a model that works across compute budgets. The paper asks how a shared parameter set can support useful predictions at multiple exits without separately training a model for every depth.

The full-loop model and an early exit are not independent networks: the latter is already a prefix of the former's computation trajectory. ELT exploits this relationship by using the full-loop prediction as an online teacher and jointly supervising a randomly selected intermediate exit with ground truth and teacher predictions. Core Idea: train the recurrent block as an iterative refiner that remains useful with repeated application, turning parameter efficiency into usable elastic inference through intra-loop self-distillation.

Method

Overall Architecture

The input is the masked token sequence or noisy latent representation at the current generation step, together with conditions such as the class label. The output remains the token distribution or denoising prediction required by the original generation algorithm. ELT leaves the outer sampler intact and replaces the independent deep layer stack inside each network invocation with repeated applications of a shared block. Training reaches the maximum loop count while supervising an intermediate exit; inference selects an exit according to the compute budget.

The process comprises shared-block looping, a stochastic prefix exit, and curriculum-based intra-loop self-distillation. Solid arrows below indicate network computation and inference output; dashed arrows indicate training-time exit selection and supervision. The teacher is not a separate parameterized model.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Current generation input<br/>Masked or noisy representation"] --> B["Shared-Block Looping"]
    B -->|Reuse the same weights| B
    B -->|Choose loop count at inference| E["Shared Prediction Head<br/>Update sampling state"]
    B -.->|Select intermediate state in training| C["Stochastic Prefix Exit"]
    C -.->|Student prediction| D["Curriculum-Based<br/>Intra-Loop Self-Distillation"]
    B -.->|Full-loop teacher prediction| D
    Y["Ground-Truth Target"] -.-> D
    D -.->|Training gradients| B

Two kinds of iteration must be distinguished: the outer generation sampling steps and the shared-block loops within a single network invocation. Reducing the inner loop count lowers computation per step, but does not automatically reduce the outer sampling step count. The paper distinguishes this acceleration axis from few-step diffusion and treats their combination as a future possibility.

Key Designs

1. Shared-Block Looping: replace repeated weight storage with repeated computation

A shared block contains \(N\) distinct transformer layers and is repeated \(L\) times per generation step, giving an effective depth of \(D=N\times L\). Layers inside the block have different weights, whereas successive loops reuse the same block. An \(8\times4\) configuration therefore follows a computation path equivalent to 32 layer applications while storing only 8 distinct layers of weights. This differs from stacking 32 independent layers and does not mean executing only 8 layers. Input and output components also contribute parameters, so actual model sizes should be taken from the experimental tables.

This architecture separates physical parameter size from computational depth, but recurrence cannot replace capacity without limit. If the shared block is too small, too few parameters must implement every transformation repeatedly, restricting representational richness; Table 2 shows poor FID when a single layer is repeated 32 times. Conversely, a smaller weight set can reduce off-chip memory access, so equal FLOPs need not imply equal latency. Such benefits depend on model size, hardware caches, and implementation rather than following unconditionally from recurrence.

2. Stochastic Prefix Exit: extract a student from one full forward pass

Training fixes the maximum loop count \(L_{\max}\) and uniformly samples an intermediate integer exit satisfying \(L_{\min}\leq L_{\mathrm{int}}<L_{\max}\) at each iteration. The full-loop path is the teacher, and the selected intermediate state produces the student output through a shared prediction head. The student neither owns a separate parameter set nor restarts computation from the input. Different exits receive supervision over training, preventing the model from concentrating all useful generation behavior at one fixed endpoint.

Because the student trajectory is a strict prefix of the teacher trajectory, the method reuses an already computed hidden state instead of training a separate small deployment model. The paper emphasizes minimal extra computation: this means no additional student-backbone forward pass, not that the intermediate prediction head, loss, and gradient handling are entirely free. Random exit sampling is a training procedure. At inference, the user chooses the loop budget; the paper does not provide a controller that automatically decides when to exit based on image difficulty.

3. Curriculum-Based Intra-Loop Self-Distillation: trust ground truth first, then the mature teacher

Intra-Loop Self Distillation (ILSD) retains the ground-truth loss at the full-loop output and supervises the student with both ground truth and teacher predictions. The teacher prediction is stop-gradient when used as a distillation target, preventing that target from changing merely to accommodate the student. The teacher path still learns through its own ground-truth loss, and both paths ultimately update the same shared parameters. Only the distillation-target gradient route is stopped; a separate teacher model is not frozen.

The student's ground-truth term has weight \(\lambda\), while its teacher-distillation term has weight \(1-\lambda\); training linearly reduces \(\lambda\) from 1 to 0. Early in training, the teacher is unreliable, so ground truth stabilizes the intermediate exit. Later, the student learns the teacher's more mature distribution or denoising prediction. This encourages early loops to produce useful transformations while preserving pressure for later loops to refine them, rather than requiring identical predictions at every loop.

A Worked Example

Consider \(N=4\) and \(L_{\max}=8\), and suppose one training iteration samples \(L_{\mathrm{int}}=3\). This is a mechanism illustration, not an additional experimental result. The current noisy image passes through the shared block 3 times, producing a student state after the equivalent of 12 layer applications. Computation continues with the same 4-layer weight set until loop 8, producing the teacher state after the equivalent of 32 layer applications.

The shared prediction head reads both exits. The full exit is supervised by the true denoising target, while the intermediate exit uses both that target and the stop-gradient teacher prediction. Joint back-propagation updates the shared weights. A different exit may be sampled in the next training iteration, progressively covering different compute budgets.

At deployment, choosing 4 loops executes the equivalent of 16 layer applications inside every outer sampling step before returning the prediction to the original sampler; choosing 8 loops executes 32 layer applications. Both use the same parameter set. The former does not skip half the outer sampling steps, nor does it guarantee unchanged quality for every image. Any-time inference here is an empirical capability over evaluated budgets, not a theoretical guarantee that every positive integer loop count works reliably.

Loss & Training

Discrete masked generation uses cross-entropy at masked positions: the ground-truth term uses the original token labels, and the distillation term uses the full-loop teacher's soft distribution over the entire vocabulary. The student can therefore learn the teacher's relative assessment of alternative tokens rather than only the correct label. Images use the MaskGIT tokenizer, mapping \(256\times256\) images to \(16\times16\) tokens. Videos use the MAGVIT tokenizer, mapping \(16\times128\times128\) sequences to \(4\times16\times16\) tokens. Both codebooks contain 1024 entries.

Continuous diffusion uses time-dependent sigmoid-weighted mean squared error. The ground-truth term compares the prediction with the clean latent, while the distillation term compares student and teacher denoising predictions. The cached loss equations contain PDF extraction corruption; this note explains the mechanism from the surrounding prose rather than presenting guessed typesetting as the authors' exact formulas. Reproduction should consult the original paper. The diffusion model uses a Stable Diffusion v1.4 VAE to map images into a \(32\times32\times4\) latent space, together with a shifted cosine noise schedule.

Masked models train for 270 epochs by default and use a cosine unmasking schedule at inference. Image training drops class conditions in 10% of batches to support classifier-free guidance. Diffusion models use Adam with batch size 512 for 500K steps, with default sampling based on 512-step DDPM and guidance scale 3.0. Loop budgets and sampling budgets must be recorded separately; training loop counts are not the generation step counts listed in the tables.

Key Experimental Results

Main Results

The following selection from original Table 1 covers class-conditional ImageNet \(256\times256\) generation and retains direct comparisons within the same framework. Lower FID and higher IS are better. GFLOPs and sampling steps follow the original table.

Model FID IS Parameters Sampling Steps GFLOPs
MaskGIT-L 2.1 270.1 303M 24 3.7k
ELT-L, 8 layers ร— 3 loops 2.2 254.3 101M 24 3.7k
ELT-L, 12 layers ร— 2 loops 2.1 281.8 152M 24 3.7k
MaskGIT-XL 2.0 294.8 446M 24 3.9k
ELT-XL, 7 layers ร— 4 loops 2.0 266.1 111M 24 3.9k

ELT-XL matches FID with approximately one quarter of the parameters, but IS falls from 294.8 to 266.1, so the result is not parity on every metric. The unchanged 3.9k GFLOPs also shows that the central result is parameter efficiency, not a fourfold FLOPs reduction.

Ablation Study

Original Table 2 compares unique layer counts and loop counts in the diffusion framework. Every configuration below has effective depth 32, exposing capacity allocation at fixed computational depth. This is not a single-factor ablation that removes ILSD.

Configuration FID Parameters Effective Depth
DiT, 32 independent layers 3.43 2.1B 32
ELT, 1 layer ร— 32 loops 10.30 69M 32
ELT, 4 layers ร— 8 loops 3.96 271M 32
ELT, 8 layers ร— 4 loops 3.16 539M 32
ELT, 16 layers ร— 2 loops 2.83 1.1B 32

Excessive sharing constrains capacity: repeating a single layer does not match the quality of a deep network. The 8-layer and 16-layer configurations outperform the 32-layer DiT while using fewer parameters, indicating that sufficient shared-block capacity matters more than maximizing the loop count.

Original Table 4 reports class-conditional UCF-101 video results without guidance. The last row increases both loop count and sampling steps, so it is not an iso-compute comparison.

Model FVD IS Parameters Sampling Steps GFLOPs
MAGVIT-L 76 ยฑ 2 89.27 ยฑ 0.15 306M 12 Approximately 4.3k
ELT, 6 layers ร— 4 loops 72.8 ยฑ 2.5 88.27 ยฑ 0.33 76M 12 Approximately 4.3k
ELT, 6 layers ร— 6 loops 60.8 ยฑ 2.7 87.88 ยฑ 0.39 76M 24 Approximately 13k

Key Findings

  • Figure 8 compares training with and without ILSD: conventional looped models deteriorate away from their training depth, whereas ILSD improves the usability of multiple exits. Text extracted from the plots is interleaved, so no additional exact values are inferred from ambiguous curves here.
  • Figure 7 and Section 5.4 report that a video model trained with a maximum of 4 loops reaches FVD 69.20 at 6 inference loops. This belongs to the depth-extrapolation analysis and must not be conflated with the 60.8 configuration in Table 4.
  • Table 3 measures throughput ratios on TPU v6e with 1 ร— 1 topology and batch size 8: 3.3 for XL, 3.5 for H, and 1.0 for B. Benefits clearly depend on model scale and memory-access conditions.
  • The 2.0-fold and 1.4-fold improvements in Figure 6 concern training-step convergence to comparable FID, not necessarily proportional end-to-end wall-clock training speedups.

Highlights & Insights

  • Shared weights learn both early usability and continued refinement. Final-only supervision allows intermediate states to drift, while stochastic exit supervision brings decodability forward. This explains elasticity more directly than reducing the number of distinct layers alone.
  • Distillation reuses the existing trajectory. The student reads a teacher-prefix state instead of running another backbone. The idea may transfer to other iterative refinement networks, provided shallow and deep exits predict in compatible spaces.
  • Capacity, per-step depth, and sampling steps are separate controls. They affect parameter storage, computation per invocation, and generation trajectory length, respectively. ELT introduces an independently adjustable inner-depth axis; its efficiency should not be attributed entirely to few-step sampling.

Limitations & Future Work

  • Limited task coverage. Experiments focus on class-conditional ImageNet images and short UCF-101 videos, without establishing equivalent gains for text-to-image generation, long videos, or higher-resolution deployment.
  • Loop extrapolation is not universally guaranteed. The authors explicitly call for further study across datasets and scales, and undersized shared blocks already expose a capacity bottleneck.
  • Hardware evidence has a specific scope. Throughput measurements use a particular TPU configuration and batch size. Fewer parameters do not imply proportional reductions in activations, back-propagation storage, or latency on every GPU.
  • Adaptive scheduling remains future work. Allocating different loop budgets across denoising steps and combining ELT with one-step or few-step generation are promising, but the paper does not systematically validate these combinations or a general automatic exit policy.
  • Compared with Universal Transformers and conventional looped networks: weight sharing itself is established. ELT focuses on making multiple recurrent exits usable for visual generation through ILSD, rather than merely showing that recurrent computation can work.
  • Compared with MaskGIT, MAGVIT, and DiT: ELT retains their discrete decoding or diffusion paradigms and changes parameter organization and supervision inside each network invocation, making it a compatible architectural modification.
  • Compared with E-DiT and MaGNeTS: the former adjusts budgets through block skipping and width reduction, while the latter uses nested generative models. ELT fixes the shared weight set and varies effective depth through repetition.
  • Compared with consistency models and progressive distillation: these primarily reduce outer sampling steps, whereas ELT adjusts inner computational depth. The approaches are conceptually complementary, but that complementarity is not an experimentally established combined gain.

Rating

  • Novelty: 4/5. Recurrence and sharing have established foundations, but prefix self-distillation and curriculum supervision form a clear method for elastic visual generation.
  • Experimental Thoroughness: 4/5. Results cover discrete image generation, continuous diffusion, and video, with capacity and throughput analyses; open-domain tasks and cross-hardware validation remain limited.
  • Writing Quality: 4/5. The mechanism and experiments are clearly organized, but claims of no extra overhead and any-time behavior require implementation and evaluation context.
  • Value: 4/5. The method provides a practical compute-budget control for parameter-constrained generative models, with deployment value dependent on hardware and quality requirements.