Reproduce “SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training” on OpenResearch

About this paper

Mengmeng Wang, Dengyang Jiang, Liuzhuozheng Li, Yucheng Lin, Guojiang Shen, Xiangjie Kong, Yong Liu, Guang Dai and 1 more

SRA 2 introduces a lightweight framework that leverages pre-trained Variational Autoencoder (VAE) features to intrinsically guide and accelerate the training of diffusion transformers. The approach yields significant improvements in convergence speed and generative quality with minimal computational overhead, for instance, achieving a 7x acceleration for SiT-XL/2 on ImageNet while improving FID scores.

Key insights

• Pre-trained VAE features, typically used for latent space mapping, contain sufficient visual priors (texture, structure, semantics) to effectively guide diffusion transformer learning without additional online extraction costs.
• Aligning intermediate diffusion transformer representations with these VAE features, especially in earlier layers, significantly accelerates training convergence and enhances the fidelity and semantic coherence of generated images.
• The proposed intrinsic guidance mechanism avoids the computational burdens and external model dependencies common in other acceleration techniques, offering a more efficient and self-contained solution.

Summary by alphaXiv.