AHPA: Adaptive Hierarchical Prior Alignment for Diffusion Transformers

2026-05-04

AI Summary

Existing diffusion Transformers employ fixed-granularity supervision signals for representation alignment, which fail to adapt to the varying semantic requirements across noise levels during denoising, leading to representational mismatch. To address this, this work proposes a lightweight, adaptive hierarchical prior alignment framework that introduces, for the first time, a timestep-aware dynamic routing mechanism. By leveraging multi-level features from a frozen VAE encoder as complementary priors, the method adaptively selects and weights alignment granularities along the denoising trajectory. Notably, it requires no external encoder, achieves granularity-adaptive alignment during training, and incurs no additional computational overhead at inference. The approach significantly accelerates convergence and enhances generation quality, outperforming current alignment baselines.

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.

Problem

• representation alignment
• diffusion transformers
• denoising trajectory
• hierarchical priors
• timestep-agnostic alignment

Innovation

• Adaptive Hierarchical Prior Alignment
• Diffusion Transformers
• Representation Alignment
• Dynamic Router
• Hierarchical VAE Features