What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion

2026-05-08

AI Summary

Existing latent diffusion models suffer from encoders that lack explicit optimization for generation-friendly latent manifolds, limiting sample quality. This work proposes the Prior-Aligned AutoEncoder (PAE), which systematically identifies and formalizes three key properties of diffusion-friendly latent manifolds—spatial consistency, local continuity, and global semantic coherence—and incorporates them as explicit training objectives. By integrating prior knowledge distilled from a Variational Flow Model (VFM) with perturbation-based regularization, PAE directly optimizes the latent structure, departing from conventional paradigms that rely solely on reconstruction fidelity or pretrained representations. On ImageNet at 256×256 resolution, PAE achieves a new state-of-the-art gFID of 1.03 and demonstrates a 13× faster training convergence under identical settings, significantly advancing both generation quality and efficiency.

Abstract

Tokenizers are a crucial component of latent diffusion models, as they define the latent space in which diffusion models operate. However, existing tokenizers are primarily designed to improve reconstruction fidelity or inherit pretrained representations, leaving unclear what kind of latent space is truly friendly for generative modeling. In this paper, we study this question from the perspective of latent manifold organization. By constructing controlled tokenizer variants, we identify three key properties of a diffusion-friendly latent manifold: coherent spatial structure, local manifold continuity, and global manifold semantics. We find that these properties are more consistent with downstream generation quality than reconstruction fidelity. Motivated by this finding, we propose the Prior-Aligned AutoEncoder (PAE), which explicitly shapes the latent manifold instead of leaving diffusion-friendly manifold to emerge indirectly from reconstruction or inheritance. Specifically, PAE leverages refined priors derived from VFMs and perturbation-based regularization to turn spatial structure, local continuity, and global semantics into explicit training objectives. On ImageNet 256x256, PAE improves both training efficiency and generation quality over existing tokenizers, reaching performance comparable to RAE with up to 13x faster convergence under the same training setup and achieving a new state-of-the-art gFID of 1.03. These results highlight the importance of organizing the latent manifold for latent diffusion models.

Problem

• latent diffusion
• latent manifold
• tokenizer
• generative modeling
• diffusion-friendly

Innovation

• latent manifold
• diffusion-friendly
• Prior-Aligned AutoEncoder
• VFM priors
• perturbation-based regularization