VAE-REPA: Variational Autoencoder Representation Alignment for Efficient Diffusion Training

TL;DR

VAE-REPA introduces a lightweight intrinsic guidance mechanism for diffusion transformers by leveraging off-the-shelf pre-trained SD-VAE features. A small projection head aligns SiT’s intermediate latent representations with the rich priors encoded in VAE features via a smooth-L1 alignment loss, enabling faster training and better generation without external encoders or teacher models. The approach achieves notable gains on ImageNet 256×256, matching or surpassing state-of-the-art acceleration methods while adding only about 4% GFLOPs. It also generalizes to text-to-image tasks and remains compatible with other acceleration strategies, offering a practical path toward efficient diffusion training with minimal overhead.

Abstract

Denoising-based diffusion transformers, despite their strong generation performance, suffer from inefficient training convergence. Existing methods addressing this issue, such as REPA (relying on external representation encoders) or SRA (requiring dual-model setups), inevitably incur heavy computational overhead during training due to external dependencies. To tackle these challenges, this paper proposes \textbf{\namex}, a lightweight intrinsic guidance framework for efficient diffusion training. \name leverages off-the-shelf pre-trained Variational Autoencoder (VAE) features: their reconstruction property ensures inherent encoding of visual priors like rich texture details, structural patterns, and basic semantic information. Specifically, \name aligns the intermediate latent features of diffusion transformers with VAE features via a lightweight projection layer, supervised by a feature alignment loss. This design accelerates training without extra representation encoders or dual-model maintenance, resulting in a simple yet effective pipeline. Extensive experiments demonstrate that \name improves both generation quality and training convergence speed compared to vanilla diffusion transformers, matches or outperforms state-of-the-art acceleration methods, and incurs merely 4\% extra GFLOPs with zero additional cost for external guidance models.

Figures (4)

• Figure 1: We empirically visualize the feature information richness of SD-VAE sd-vae and SiT-XL/2 sit via PCA pca. Top: VAE features, extracted from original images by an SD-VAE encoder. Bottom: Latent features of SiT-XL/2 across different block layers and noise levels. We observe that SD-VAE features are significantly superior in delineating visual concepts compared to SiT’s latent representations, maintaining clearer details, structural integrity, and stronger semantic coherence. This motivates our use of VAE features for representation alignment.
• Figure 2: Comparison of typical SiT training paradigms.(a) Vanilla SiT Training: Images are encoded by a VAE, added with noise, and processed by the diffusion model for denoising. (b) SiT Training with External Representation Alignment (e.g., REPA repa): SiT training augmented with an external representative encoder and an MLP for alignment. (c) SiT Training with Dual-model Self-Alignment (e.g., SRA sra): SiT training leveraging a dual-model setup with an MLP for self-alignment, guided by a teacher diffusion model. (d) SiT Training with VAE Representation Alignment (ours): SiT utilizes VAE features as representation guidance and an MLP for alignment, efficiently combining VAE’s semantic richness with SiT’s denoising capability without introducing additional heavy models.
• Figure 2: FID comparison across training iterations for accelerated alignment methods. All experiments are conducted on ImageNet (256×256) with a batch size of 256 and without CFG.
• Figure 3: VAE-REPA Improves Visual Scaling.Top: Vanilla SiT-XL/2 and SiT-XL/2+VAE-REPA. Bottom: Vanilla REPA and REPA+VAE-REPA. Our method is verified to produce images with higher structural fidelity, finer details, and stronger semantic coherence at the same training steps compared with both vanilla SiT and vanilla REPA. Results for all methods are sampled using the same seed, noise, and class label, with a classifier-free guidance scale of 4.0 employed during sampling.