PaperlayerView exact revision on arXiv

Source-linked AI summary

ELT: Elastic Looped Transformers for Visual Generation

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

arXiv:2604.09168v3cs.CV

TL;DR

Visual generation needs high compute capacity without the parameter and memory costs of deep stacks of unique transformer layers. ELT uses recurrent weight-shared blocks trained with ILSD, yielding elastic inference across loop counts. It reports competitive image and video results with 4× fewer parameters under iso-inference-compute settings.

  • Problem

    Deep visual-generation models incur larger memory footprints as layers or width increase, while looped models struggle to make intermediate representations useful before the final loop.

  • Method

    ELT reuses transformer blocks across recurrent loops and applies ILSD between intermediate student configurations and the maximum-loop teacher configuration.

  • Results

    4× fewer parameters achieves FID of 2.0 on class-conditional ImageNet 256 × 256 and FVD of 72.8 on class-conditional UCF-101 under iso-inference-compute settings.

  • Takeaways & Limitations

    A single ILSD-trained model provides a family of performant configurations across variable iteration counts, enabling Any-Time inference without retraining.

  • Takeaways & Limitations

    Performance degrades when the number of unique layers is too small or inference loops significantly exceed the trained maximum loop count.

Abstract

from arXiv · show

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 \times 256$ and FVD of $72.8$ on class-conditional UCF-101.

1. Introduction

ELT uses recurrent, weight-shared transformer blocks to improve parameter efficiency while making intermediate loops useful for progressive visual refinement. ILSD enables Any-Time inference, and ELT reports competitive image and video quality with substantially fewer parameters.

  • Motivation: Recurrence reuses parameters to increase compute capacity without proportionally increasing memory footprint.Looped transformers can also improve arithmetic intensity and reduce memory-transfer costs relative to standard transformers of equivalent depth.
  • Motivation: ELT addresses intermediate representations that may remain uninterpretable until the final loop in conventional looped architectures.The model is designed for progressive refinement so intermediate repeats can produce meaningful synthesis.
  • Method: ILSD uses a full-loop teacher path and an intermediate-loop student path to train progressive refinement across loop depths.The student follows a subset of the teacher trajectory, rather than treating looping as a fixed-depth process.
  • Results: 4× fewer parameters yields FID of 2.0 on class-conditional ImageNet 256 × 256 and FVD of 72.8 on class-conditional UCF-101 under iso-inference-compute settings.The reported results match or improve baseline performance under the same inference-compute setting.
  • Implications: Any-Time inference lets one model traverse the quality-versus-compute frontier without retraining by varying its loop count.This supports deployment ranging from few-loop, latency-critical generation to more-loop, high-fidelity rendering.

2. Related Work

Related work establishes recursive shared-parameter architectures and elastic visual generation, while distinguishing ELT from fixed-point solvers, architectural skipping, and inter-step acceleration methods.

  • Recursive Architectures: Recursive architectures apply shared blocks repeatedly for parameter-efficient computation, with prior work spanning reasoning, adaptive depth, image understanding, and length generalization.Universal Transformers introduced single-layer iteration, while later studies extended looping across language and vision tasks.
  • Recursive Architectures: ELT explicitly optimizes unrolled intermediate states through ILSD rather than relying on a black-box solver to reach an analytical fixed point.This preserves Any-Time inference without requiring a strict equilibrium solution.
  • Parameter-Efficient Visual Generation: Prior parameter-efficient visual-generation methods use pruning, separable convolutions, lightweight diffusion variants, or step distillation.These approaches optimize architectures or sampling procedures rather than using ELT’s recurrent shared-block formulation.
  • Elastic Visual Generation: Elastic visual-generation methods vary computational budgets through block skipping, width reduction, uncertainty-based exits, or search-based sampling trajectories.ELT differs by obtaining elastic capability through recurrent refinement and ILSD rather than architectural skipping or external search.
  • Few-Step and Consistency Models: ELT reduces compute within each sampling step by varying loop count, whereas few-step and consistency methods reduce the number of sampling steps.The two axes are described as complementary and potentially combinable.

3. Preliminaries

Masked generative and diffusion models iteratively refine visual representations across sampling steps. ELT replaces standard transformer processing with recurrent, weight-shared blocks, adding a loop-count compute lever within each step.

  • Masked Generative Models: MaskGIT generates image tokens simultaneously and iteratively refines them instead of decoding tokens sequentially in raster-scan order.Its sampling process progressively unmasks tokens according to model logits, randomness, and confidence.
  • Masked Generative Models: Masked generation starts from fully masked or partially masked tokens and repeats sampling over iterations k ∈ [1, K].The category label c conditions image or video generation, while the Sample function unmasks only selected high-confidence tokens.
  • Diffusion Models: Diffusion models iteratively reverse corruption from Gaussian noise toward data using a noise schedule and a model conditioned on timestep and class.The displayed update combines the current vector, model prediction, and standard-Normal noise with time-dependent scalars.
  • ELT Preliminaries: Figure 3 depicts shared-parameter recurrent training with ILSD and Any-Time inference through early exits using a shared MLM head.At inference, an intermediate block can predict the next state X_k+1 from the current sampling input X_k.
  • ELT Preliminaries: ELT implements the visual-generation model as recurrent, weight-shared transformer blocks that refine representations within each sampling step.This architecture aligns model computation with the recursive refinement used by masked and diffusion models.

4. Method

ELT decouples parameter count from computational depth by repeatedly applying shared transformer blocks, while ILSD trains intermediate loops to remain effective before the maximum loop count. The method jointly supervises stochastic intermediate-loop students and full-loop teachers, using task-appropriate ground-truth and distillation losses.

  • Looping Mechanism: A standard N-layer block looped L times achieves effective depth N × L while retaining only N unique parameterized layers.Looping reuses the same block parameters across successive applications, decoupling physical model size from computational depth.
  • Intra-Loop Self Distillation: ELT treats the looped transformer as an iterative refiner so intermediate configurations can produce meaningful representations before Lmax loops.The method targets the black-box intermediate trajectory problem in fixed-depth weight-tied transformers.
  • Intra-Loop Self Distillation: ILSD uses a full-loop teacher and an intermediate-loop student that shares the same parameters and follows a strict prefix of the teacher trajectory.The intermediate computation is already required within the maximum-loop forward pass, keeping training overhead minimal.
  • Training Objective: At each training iteration, ELT samples an intermediate loop count Lint between Lmin and Lmax and combines ground-truth and distillation supervision.The ground-truth-to-distillation weight λ is linearly decayed from 1 to 0 during training.
  • Loss Formulation: For masked discrete-token generation, ELT uses cross-entropy losses, while diffusion training uses sigmoid-weighted mean squared error for ground truth and distillation.The student uses Lint and the teacher uses Lmax in the diffusion formulation.

5. Experiments and Results

Experiments evaluate ELT for class-conditional image and video generation, showing competitive quality, parameter efficiency, elastic inference, and hardware efficiency. Across scaling studies, loop count and model capacity trade off against FID, while ILSD supports stable quality across inference depths.

  • Image Generation: FID 3.16 and 2.83 outperform the 32-layer DiT baseline’s FID 3.43 at 4× and 2× lower parameter counts, respectively.These results use 8N × 4L and 16N × 2L inference configurations.
  • Scaling and Elastic Inference: Adjusting inference loops traverses the FID–GFLOPs Pareto curve without retraining, enabling hardware-dependent quality–compute trade-offs.Increasing loops improves FID for fixed unique layers, but gains eventually diminish relative to scaling architecture size.
  • Image Generation: 3.5× peak throughput is achieved for model scale H as shared parameters reduce repeated HBM-to-SRAM transfers.Throughput ratios are measured on a Google Cloud TPU v6e with inference batch size 8.
  • Training Efficiency: 2× and 1.4× convergence speedups are reported for 16N × 2Lmax and 8N × 4Lmax ELT diffusion configurations at equal effective depth.Both configurations are compared with an N = 32 DiT baseline and effective depth D = 32.
  • Video Generation: FVD 72.8 versus 76 is achieved by a compact 76M ELT model against MAGVIT on UCF-101, reaching FVD 60.8 with more loops and sampling steps.The comparison is conducted under iso-inference-compute settings in a data-constrained regime.

6. Conclusion

ELT is a recurrent, parameter-efficient approach to visual generation that trains one model to perform well across variable iteration counts. It supports flexible inference and reports competitive image and video-generation performance with substantially fewer parameters.

  • ELT uses recurrent transformers to achieve strong visual-generation performance with 4× fewer parameters in iso-inference-compute settings.The reported results cover both image and video generation tasks.
  • ILSD trains a single model to remain performant across variable iteration counts, yielding a continuous family of models from one training run.Practitioners can balance image quality and GFLOPs dynamically during inference.
  • ELT can potentially allocate diffusion inference compute dynamically across denoising steps, spending more compute where it matters most.The paper also discusses elasticity for one-step generative models by varying loop counts at inference.
  • The ImageNet and video-generation figures provide qualitative comparisons of ELT configurations and vanilla looped transformers.The supplied figure passages identify comparisons and example configurations, but do not provide complete quantitative conclusions.

A. Implementation Details

The implementation covers ELT configurations for masked generative and diffusion models across image and video-generation settings. It also specifies optimizer, guidance, sampling, and architecture choices used in the experiments.

  • The experiments sweep classifier-free guidance scales for ELT and baselines, while UCF-101 class-conditional video generation uses no classifier-free guidance.Sampling temperature controls randomness when sampling from the categorical distribution.
  • The implementation uses bias=0.5 and scale=0.8 throughout the reported experiments.
  • Diffusion experiments use two model architectures with configurations and hyperparameters documented separately.The relevant architecture details are referenced in Table 7.

B. Additional Experiments

The additional experiments examine ELT’s training and inference efficiency. ELT reduces memory requirements and redundant transfers while preserving the effective depth used for comparison.

  • ELT trains intermediate loops in a single forward pass with negligible prediction-head cost, keeping training step time unchanged.The comparison profiles DiT training with an 8N × 4L looping configuration and effective depth D=32 for all methods.
  • ELT achieves 3.5x peak throughput speedup in iso-inference-compute settings by eliminating redundant HBM → VRAM memory transfers.This inference result is attributed to reduced memory transfers.

Algorithm 1 ELT Training with Intra-Loop Self Distillation (ILSD)

Algorithm 1 trains ELT with a shared recurrent transformer block by pairing a maximum-loop teacher trajectory with an intermediate-loop student state. The ILSD loss distills the teacher’s final representation into the sampled intermediate state.

  • The training algorithm defines a weight-shared transformer block, task-specific prediction head, inputs, targets, loop bounds, and curriculum weight.The curriculum weight λ decays from 1 to 0 over training.
  • Each training step samples an intermediate student loop count L_int between L_min and L_max.
  • A shared loop body records the intermediate feature when the sampled loop count is reached while continuing the forward trajectory.
  • The forward pass runs to L_max and returns both the maximum-loop feature F_max and intermediate feature F_int.
  • ILSD computes a distillation loss between F_int and a stop-gradient version of F_max, weighted by 1−λ.

Algorithm 2 ELT’s Any-Time Inference

ELT’s Any-Time inference dynamically scales computation by changing the loop budget while keeping the model parameters fixed. The iterative refinement analysis examines how updates vary across loops and sampling steps.

  • Any-Time inference: The loop budget L dynamically scales computation during a single sampling step without retraining the model.The weight-shared transformer block is recursively applied for the selected number of loops.
  • Algorithm components: g_Θ is the weight-shared composite transformer block, h_head is the task-specific prediction head, and x is the current-step input feature.L denotes the dynamic compute budget, or number of loops to execute.
  • Ablations: The ablations compare vanilla looping with ELT and separately evaluate the importance of the ILSD loss.These studies cover both video and image generation, including DiTs and masked generative models.
  • Iterative refinement: Initial denoising steps receive larger updates at higher loop counts than later denoising steps.The per-iteration update magnitudes are measured using the L2 norm across loops and sampling steps.

C. Limitations

ELT has two reported failure modes: insufficient unique-layer capacity and excessive inference looping beyond the trained regime. Deployment also requires hardware- and model-scale-specific calibration of the operating point.

  • Failure Cases: A configuration with too few unique layers can lack representational capacity regardless of the loop count.The paper gives the 1N × 32L configuration in Table 2 as an example.
  • Failure Cases: Quality can deteriorate when inference uses loop counts substantially beyond L_max because the shared block over-iterates beyond its trained convergence regime.This limitation concerns extrapolation beyond the maximum training loop count.
  • Deployment Considerations: Deployment requires selecting and calibrating the operating point (N, L) for the target hardware and deployment tier.The optimal loop count for a given quality target depends on model scale.
Loading 2604.09168v3…