PaperlayerView exact revision on arXiv

Source-linked AI summary

REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers

Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, Liang Zheng

arXiv:2504.10483v3cs.CVcs.LG

TL;DR

Latent diffusion models traditionally train VAEs and diffusion models separately, while naive end-to-end diffusion-loss tuning can degrade generation. REPA-E jointly tunes both with representation alignment, yielding faster training, improved latent structure, and state-of-the-art ImageNet generation results.

  • Problem

    The paper asks whether VAEs and latent diffusion models can be jointly tuned end-to-end, since naive diffusion-loss backpropagation is ineffective and can reduce generation performance.

  • Method

    REPA-E jointly trains VAE and diffusion-model features using representation-alignment loss, with diffusion-loss stop-gradient and VAE regularization.

  • Results

    Over 17× and 45× faster training than REPA and vanilla recipes, respectively, with ImageNet 256 × 256 FID of 1.12 with and 1.69 without classifier-free guidance.

  • Takeaways & Limitations

    End-to-end tuning improves latent-space structure across VAE architectures and improves downstream generation performance.

Abstract

from arXiv · show

In this paper we tackle a fundamental question: "Can we train latent diffusion models together with the variational auto-encoder (VAE) tokenizer in an end-to-end manner?" Traditional deep-learning wisdom dictates that end-to-end training is often preferable when possible. However, for latent diffusion transformers, it is observed that end-to-end training both VAE and diffusion-model using standard diffusion-loss is ineffective, even causing a degradation in final performance. We show that while diffusion loss is ineffective, end-to-end training can be unlocked through the representation-alignment (REPA) loss -- allowing both VAE and diffusion model to be jointly tuned during the training process. Despite its simplicity, the proposed training recipe (REPA-E) shows remarkable performance; speeding up diffusion model training by over 17x and 45x over REPA and vanilla training recipes, respectively. Interestingly, we observe that end-to-end tuning with REPA-E also improves the VAE itself; leading to improved latent space structure and downstream generation performance. In terms of final performance, our approach sets a new state-of-the-art; achieving FID of 1.12 and 1.69 with and without classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.

1. Introduction

The paper examines why standard two-stage latent diffusion training makes joint VAE–LDM optimization difficult, finding that naive diffusion-loss backpropagation harms generation. REPA-E instead uses representation alignment to jointly tune both components, accelerating training and improving latent structure and generation.

  • Motivation: Two-stage LDM training separately optimizes the VAE and diffusion model, leaving the VAE representation difficult to optimize for final generation.The paper frames this as the central challenge motivating joint training.
  • Results: End-to-end tuning adaptively improves latent-space structure across diverse VAE architectures and improves downstream generation performance.The paper reports smoother SD-VAE latents, more detailed IN-VAE latents, and better generation outcomes.
  • Motivation: Naive end-to-end diffusion-loss tuning is ineffective because it favors simpler latent structure that eases denoising but reduces generation performance.This behavior is reported directly from experiments.
  • Contribution: REPA-E uses representation-alignment loss to jointly tune the VAE and latent diffusion model during training.The recipe replaces diffusion-loss backpropagation to the VAE with REPA-based end-to-end tuning.
  • Results: Over 17× and 45× faster diffusion training is achieved versus REPA and vanilla training recipes, respectively.These speedups are reported as headline advantages of REPA-E.

2. Related Work

Prior work studies tokenizers, latent diffusion, and representation alignment largely as separate components. This paper positions REPA-E as joint end-to-end optimization of tokenizer and diffusion-model representations for faster convergence and improved generation.

  • Tokenizers and autoencoders: Existing tokenizers are commonly trained to minimize reconstruction error, which may not produce the optimal latent space for generation.The paper identifies latent-space quality for generation as distinct from reconstruction quality.
  • Latent diffusion models: Latent diffusion models typically use pretrained image tokenizers and train tokenizers and diffusion models separately.The separate training stages motivate joint optimization.
  • Representation alignment: Representation alignment has improved diffusion-model training speed and performance, but prior approaches apply REPA separately to LDMs or VAEs.The paper claims end-to-end application achieves better performance and speed than separate application.
  • End-to-end diffusion: LSGM jointly trains score-based generative models with a variational lower bound and entropy term, whereas REPA-E is reported to converge faster during end-to-end training.The comparison concerns convergence during joint training.

3. REPA-E: Unlocking VAE for Joint Training

REPA-E motivates joint VAE–LDM training by showing that diffusion loss damages latent structure and that alignment is linked to generation quality but bottlenecked by VAE features. Its recipe uses REPA for joint updates, normalization, stop-gradient diffusion loss, and VAE regularization.

  • 3.1. Motivating End-to-End Training with REPA: Higher representation-alignment scores correlate with improved generation performance, making alignment a proxy for the generation objective.The paper uses this relationship to motivate REPA-based end-to-end tuning.
  • 3.1. Motivating End-to-End Training with REPA: Naive diffusion-loss backpropagation simplifies VAE latents, making denoising easier but reducing generation performance.Lower spatial variance in the VAE latent contributes to this simpler structure.
  • 3.1. Motivating End-to-End Training with REPA: Vanilla REPA alignment saturates around 0.4 because VAE features bottleneck the achievable score, while REPA gradients to the VAE address this limitation.The figure passage reports saturation near 0.42; the prose describes the VAE-feature bottleneck.
  • 3.2. End-to-End Training with REPA: REPA-E jointly trains VAE and LDM features with representation-alignment loss instead of backpropagating diffusion loss into the VAE.The alignment objective is applied to diffusion hidden states and pretrained perceptual features.
  • 3.2. End-to-End Training with REPA: Batch normalization provides differentiable latent normalization using exponential moving-average statistics without recomputing dataset statistics after every update.This normalization is inserted between the VAE and latent diffusion model.
  • 3.2. End-to-End Training with REPA: A stop-gradient restricts diffusion-loss updates to the LDM, while VAE regularization preserves the original VAE’s reconstruction performance.The regularizers include reconstruction, GAN, and KL-divergence losses.
  • 3.2. End-to-End Training with REPA: REPA-E produces higher-quality and more structurally meaningful images than vanilla REPA at matched training stages, including 400K steps.The visual comparison uses identical seeds, noise, labels, and classifier-free guidance scale 4.0.

4. Experiments

Experiments show that REPA-E accelerates latent-diffusion training, improves generation quality across settings, and adaptively improves VAE latent spaces for downstream generation.

  • Impact on Training Performance and Speed: REPA-E improves generation FID from 19.40 to 12.83 at 20 epochs, 11.10 to 7.17 at 40 epochs, and 7.90 to 4.07 at 80 epochs.These comparisons are against REPA on ImageNet 256 × 256 without classifier-free guidance.
  • Impact on Training Performance and Speed: At 80 epochs, REPA-E surpasses baselines trained for substantially longer, reaching FID 4.07 within 400K steps versus REPA’s 5.9 after 4M steps.The cited comparison includes FasterDiT, MaskDiT, DiT, SiT, and REPA.
  • Generalization and Scalability of REPA-E: REPA-E improves gFID over REPA across model scales, with gains of 29.6% for SiT-B, 32.3% for SiT-L, and 34.0% for SiT-XL.The increasing gains indicate improved scaling behavior across the evaluated SiT model sizes.
  • Generalization and Scalability of REPA-E: REPA-E consistently improves gFID across perceptual encoders, VAE architectures, and alignment depths.Examples include gFID reductions from 24.1 to 16.3 with DINOv2-B, from 22.7 to 12.7 with IN-VAE, and from 23.0 to 16.4 at the sixth layer.
  • Generalization and Scalability of REPA-E: REPA-E can train both VAE and LDM from scratch, achieving FID 4.34 within 400K iterations versus REPA’s 5.90 after 4M iterations.Pretrained VAE initialization provides only a slight additional performance benefit.
  • Impact of End-to-End Tuning on VAE: End-to-end tuning improves latent spaces across VAE architectures and enables tuned VAEs to serve as drop-in replacements with improved downstream generation.REPA-E reduces noise in SD-VAE latents and adds detail to over-smoothed VA-VAE latents.

5. Conclusion

The conclusion presents REPA-E as a way to unlock end-to-end VAE and diffusion-model training after naive diffusion-loss backpropagation fails.

  • 5. Conclusion: REPA-E unlocks joint VAE and diffusion-model tuning through representation-alignment loss after direct diffusion-loss backpropagation degrades performance.The recipe also improves latent-space structure and accelerates diffusion training.
  • 5. Conclusion: REPA-E speeds diffusion-model training by over 17× versus REPA and 45× versus vanilla training.The paper reports these speedups alongside improved latent-space structure and generation performance.
  • 5. Conclusion: REPA-E achieves generation FID 1.12 with classifier-free guidance and 1.69 without it on ImageNet 256 × 256.The conclusion describes these results as state-of-the-art.

Transformers

The diffusion-loss ablation examines how naive end-to-end tuning changes latent-space structure and compares it with standard training and REPA-E.

  • Transformers: Table 10 compares total variation and mean per-channel latent variance across standard LDM training, naive diffusion-loss tuning, and REPA-E.All experiments initialize the VAE with SD-VAE.

A. Impact of Diffusion Loss on Latent Space

Naive diffusion-loss backpropagation simplifies the latent space for denoising but harms image-generation performance, whereas REPA-E improves both outcomes.

  • A. Impact of Diffusion Loss on Latent Space: Direct diffusion-loss backpropagation reduces spatial variance, making denoising easier by hacking the latent space while reducing image-generation performance.The analysis uses spatial variance and total variation to quantify latent structure.
  • A. Impact of Diffusion Loss on Latent Space: REPA-E improves generation performance while also improving the underlying VAE’s latent-space structure.The paper links this result to observations in Figures 3 and 6.

B. Additional Analysis

Additional analyses compare REPA-E with alternative end-to-end objectives and test its performance across resolutions, tasks, and VAE reconstruction quality. These evaluations consistently report stronger generation performance, faster convergence, and competitive reconstruction quality.

  • Scaling REPA-E to Higher Latent Resolution: REPA-E significantly improves generation quality over vanilla REPA on ImageNet-512 with 64×64 latents at 100K steps without classifier-free guidance.The experiment uses SD-VAE as tokenizer and SiT-L as the diffusion model.
  • MSCOCO Text-to-Image Generation with E2E-VAE: End-to-end tuned VAEs improve MSCOCO text-to-image generation performance when used with MMDiT and REPA.Models are trained for 100K steps and evaluated with classifier-free guidance using αcfg = 2.0 and EMA weights.
  • Comparison of End-to-End Training Objectives: REPA-E outperforms naive diffusion-loss training and the LSGM entropy-regularized objective across metrics after 400K SiT-XL training steps.The comparison uses consistent settings and reports faster convergence and stronger generation quality for REPA-E.
  • VAE Reconstruction Evaluation on ImageNet-256: REPA-E maintains competitive VAE reconstruction quality across all reported metrics on ImageNet-256.The method primarily improves generative capability while preserving reconstruction quality.

C. Remarks on FID Evaluation

The ImageNet evaluation protocol generally uses 50,000 randomly sampled images, while the results in Table 9 use class-balanced sampling because it produced slightly better FID.

  • FID Evaluation: Table 9 uses class-balanced ImageNet sampling, generating 50 images per class across 1,000 categories, because it yielded slightly better FID.The standard protocol instead generates 50,000 images by randomly sampling class labels.
Loading 2504.10483v3…