AHPA: Adaptive Hierarchical Prior Alignment for Diffusion Transformers

TL;DR

AHPA introduces a hierarchical, adaptive alignment method for diffusion transformers that dynamically adjusts supervision granularity during training, improving convergence and output quality without extra inference costs.

Contribution

The paper proposes AHPA, a novel framework that adaptively aligns hierarchical VAE features at different levels to match the denoising process's changing needs.

Findings

AHPA improves training convergence of diffusion transformers.

AHPA enhances generation quality compared to baseline methods.

AHPA requires no additional inference cost or external supervision during training.

Abstract

Representation alignment has recently emerged as an effective paradigm for accelerating Diffusion Transformer training. Despite their success, existing alignment methods typically impose a fixed supervision target or a fixed alignment granularity throughout the entire denoising trajectory, whether the guidance is provided by external vision encoders, internal self-representations, or VAE-derived features. We argue that such timestep-agnostic alignment is suboptimal because the useful granularity of representation supervision changes systematically with the signal-to-noise ratio. In high-noise regimes, diffusion models benefit more from coarse semantic and layout-level anchoring, whereas in low-noise regimes, the training signal should emphasize spatially detailed and structurally faithful refinement. This non-stationary alignment behavior creates a representational mismatch for static single-level supervisors. To address this issue, we propose Adaptive Hierarchical Prior Alignment (AHPA), a lightweight alignment framework that exploits the hierarchical representations naturally embedded in the frozen VAE encoder. Instead of using only a single compressed latent as the alignment target, AHPA extracts multi-level VAE features that provide complementary priors ranging from local geometry and spatial topology to coarse semantic layout. A timestep-conditioned Dynamic Router adaptively selects and weights these hierarchical priors along the denoising trajectory, thereby synchronizing the alignment granularity with the model's evolving training needs. Extensive experiments show that AHPA improves convergence and generation quality over baselines and incurs no additional inference cost while avoiding external encoder supervision during training.