SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training

TL;DR

SRA 2 introduces a lightweight, VAE-based intrinsic guidance method for diffusion transformers, significantly improving training efficiency and generation quality without external models or high computational costs.

Contribution

The paper proposes SRA 2, a novel VAE feature alignment framework that accelerates diffusion training and enhances output quality without external encoders or dual-model setups.

Findings

Faster training convergence compared to vanilla diffusion models

Improved generation quality over baseline methods

Minimal additional computational overhead (4% GFLOPs)

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 SRA 2, a lightweight intrinsic guidance framework for efficient diffusion training. SRA 2 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, SRA 2 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 SRA 2 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.