PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhengrong Yue, Taihang Hu, Mengting Chen, Haiyu Zhang, Zihao Pan, Tao Liu, Zikang Wang, Jinsong Lan, Xiaoyong Zhu, Bo Zheng, Yali Wang

arXiv:2605.07915v1cs.CV

TL;DR

Existing tokenizer objectives leave unclear which latent-manifold properties best support diffusion beyond reconstruction fidelity. PAE explicitly shapes those properties, matching RAE with up to 13× faster convergence and achieving a gFID of 1.03 on ImageNet 256×256.

  • Problem

    Existing tokenizer objectives emphasize reconstruction or inherited representations, while reconstruction quality alone is insufficient to explain generative performance.

  • Method

    PAE explicitly shapes latent manifolds by aligning refined VFM-derived priors with spatial structure, local continuity, and global semantics.

  • Results

    PAE reaches comparable quality to RAE with up to 13× fewer training epochs and achieves a state-of-the-art gFID of 1.03 on ImageNet 256×256.

  • Takeaways & Limitations

    Diffusion-friendly latent-manifold organization is more closely associated with downstream generation quality than reconstruction fidelity.

  • Takeaways & Limitations

    The empirical conclusions are validated mainly on ImageNet at 256×256 and do not fully test higher resolutions, diverse visual domains, or downstream tasks beyond class-conditional image synthesis.

Abstract

from arXiv · show

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.

1 Introduction

The introduction identifies a reconstruction–generation mismatch and argues that diffusion-friendly latent manifolds should be evaluated through their geometry. It motivates PAE, which explicitly shapes manifold properties and improves tokenizer and diffusion performance on ImageNet.

  • Motivation: Reconstruction quality alone can fail to predict generation quality, revealing a reconstruction–generation mismatch in vanilla VAE objectives.Vanilla VAE combines pixel-wise reconstruction loss with KL regularization, but improving reconstruction does not necessarily improve generation.
  • Latent Manifold Analysis: Controlled experiments investigate spatial structure coherence, local manifold continuity, and global manifold semantics as complementary properties of latent manifolds.The studies fix the latent channel budget and use eRank only as a supplementary diagnostic, attributing trends mainly to manifold geometry.
  • Latent Manifold Analysis: The three manifold properties strongly correlate with downstream gFID, making them effective indicators of a diffusion-friendly latent manifold.This conclusion comes from controlled pilot experiments designed to isolate differences in manifold geometry.
  • Proposed Method: PAE explicitly shapes the latent manifold through targeted regularizations for spatial structure, local continuity, and semantic organization.Spatial Structure Regularization aligns latents with corresponding VFM features, while Manifold Continuity Regularization perturbs latents and enforces perceivable consistency.
  • Experimental Results: 1.03 gFID is achieved by PAE with longer training, while it reaches RAE-comparable performance with up to 13× fewer training epochs.On ImageNet 256×256, PAE also reports an rFID of 0.26 under the stated experimental setting.

2 Related Work

Related work improves latent diffusion either by injecting representation priors into the generator or by constructing representation-rich latent spaces through autoencoders. The latter builds on latent diffusion’s reliance on a first-stage autoencoder, while earlier VAE designs primarily optimized reconstruction fidelity.

  • Representation Priors in Diffusion Generators: Representation-Guided DiT injects external representation priors into the diffusion generator to improve diffusion.Recent approaches reshape generator-side representations by aligning DiT features with vision foundation model representations or modeling high-level semantics before pixel-level synthesis.
  • Representation Autoencoders for Latent Diffusion: Representation-Native DiT constructs a representation-rich latent space through the autoencoder to improve downstream diffusion.Latent diffusion relies on a first-stage autoencoder to define the latent space used by downstream diffusion.
  • Representation Autoencoders for Latent Diffusion: Early VAE-based designs mainly optimize reconstruction fidelity, but reconstruction quality alone is an insufficient proxy for generative performance.

3 Method

PAE explicitly shapes a compact latent manifold for diffusion using frozen VFM guidance, detail-aware modulation, and three prior-alignment objectives. Its training targets are refined to preserve semantic information while matching tokenizer bottleneck and spatial resolution.

  • Prior Alignment Objectives: PAE regularizes the latent space along instance-level spatial structure, local continuity, and global semantic organization.These three dimensions are implemented as SSR, MCR, and SCR beyond reconstruction.
  • Tokenizer Architecture: PAE extracts frozen VFM features, injects pixel detail with DAM, and projects them into a compact latent code for downstream diffusion.DAM keeps VFM features dominant through zero-initialized scale-and-shift fusion, gradually adding missing detail without uncontrolled feature mixing.
  • Spatial Structure Regularization (SSR): SSR aligns latent and target spatial Gram matrices to preserve instance-level topology through bottleneck compression.This alignment retains relative spatial relationships between latent tokens rather than relying on reconstruction alone.
  • Manifold Continuity Regularization (MCR): MCR enforces progressive consistency between neighboring perturbed latent levels, encouraging perceptually similar decodings while protecting anchor reconstruction quality.The cascaded design avoids forcing every perturbed latent to reconstruct the original image directly.
  • Semantic Consistency Regularization (SCR): SCR aligns pooled and patch-level compressed tokens with projected targets to preserve both concept-level organization and token-wise semantic directions.The two alignment levels address semantic distortion caused by bottleneck compression.
  • Refined VFM Priors: Refined VFM targets compress semantic features and spatially smooth them to better match the tokenizer bottleneck and latent resolution while preserving semantics and coarse spatial relations.Both refined targets remain fixed during tokenizer training.

4 Experiments

Experiments on ImageNet 256×256 show that PAE improves both short-horizon convergence and final generation quality, while ablations support jointly shaping latent structure, continuity, and semantics.

  • Convergence Speed and Final Performance: 1.03 gFID at 800 epochs establishes PAE (DINOv2) as a new state-of-the-art result under the same training setup.At 80 epochs, the same model achieves 1.27 gFID.
  • Convergence Speed and Final Performance: 1.27 guided gFID at 80 generator epochs lets PAE(DINOv2) outperform VTP (1.44) and GAE (1.48).PAE also surpasses RAE (DiTDH-XL) while using 675M versus 839M generator parameters and CFG instead of AutoGuidance.
  • Ablation Studies: Combining all three prior-alignment objectives gives the best overall generation performance, with each objective most strongly improving its intended dimension.The objectives target structure, continuity, and semantics jointly.
  • Generalization and Design Analysis: PAE remains effective across teacher encoders and stable under moderate changes in latent dimension and DAM depth.This generalization and design stability are summarized in Figure 8.
  • Ablation Studies: Refined VFM targets consistently improve SSC, GSQ, LPC, rFID, and gFID over raw targets, indicating cleaner, bottleneck-matched supervision.The improvement is described as modest relative to the larger gain from prior alignment.

5 Conclusion · A Extended Related Works · B Latent Manifold Geometry Metrics

The paper proposes PAE to improve latent diffusion by explicitly shaping latent manifolds, emphasizing instance-level structure, local continuity, and global semantics beyond reconstruction quality. Related work is organized around representation, spectral, and structural approaches, while appendix metrics operationalize complementary latent-geometry properties for empirical diagnosis.

  • 5 Conclusion: PAE explicitly shapes tokenizer-induced latent manifolds because reconstruction quality alone does not explain tokenizer effectiveness.The framework targets latent diffusion through prior-aligned autoencoders.
  • 5 Conclusion: Stronger generation is associated with latent spaces preserving instance-level structure, local continuity, and global semantics.These properties are presented as more informative than reconstruction quality alone.
  • A Extended Related Works: Visual-tokenizer research spans reconstruction-oriented autoencoders, representation-based tokenizers, unified understanding-generation tokenizers, and diffusion-specific regularization.The literature is organized by how methods improve the diffusability or learnability of induced latent spaces.
  • A Extended Related Works: The broader tokenizer landscape evolved from compression and reconstruction methods toward improved quantization efficiency, codebook usage, scalability, and adaptive token allocation.Examples include AEs, VAEs, VQ-VAEs, VQGAN, BSQ, IBQ, DC-AE, MAE-Tok, CAT, ElasticTok, and FlexTok.
  • A Extended Related Works: Representation-centric methods leverage pretrained visual foundation models as tokenizers, encoders, teachers, alignment targets, or supervision sources.The passage lists methods including RAE, Scale-RAE, FlatDINO, LV-RAE, DINO-SAE, FAE, SVG, VFMTok, RepTok, and VFM-VAE.
  • A Extended Related Works: Spectral and structure-centric methods improve diffusability through spectral bias, lower-frequency organization, local smoothness, reduced redundancy, local correlation regularization, and latent masked reconstruction.These approaches indicate that diffusion-friendly latents depend on information organization across features, not only semantics or reconstruction quality.
  • A Extended Related Works: The manifold-centered perspective unifies semantic priors with lower-frequency organization, local smoothness, and reduced redundancy as complementary aspects of latent organization.Representation-centric methods emphasize semantic priors, whereas spectral and structure-centric methods emphasize organization and regularity.
  • B Latent Manifold Geometry Metrics: The appendix formalizes three complementary latent-geometry properties as empirical diagnostics for diffusion-friendly spaces beyond reconstruction quality alone.The metrics operationalize intuitively relevant properties and are not claimed to be formally derived complexity measures for diffusion training.

B.1 Metric Definitions … C.4 Detailed Configuration of Ablation Study

The paper defines three primary latent-geometry metrics—SSC, LPC, and GSQ—plus supplementary eRank, then details their diffusion-learning relevance and PAE implementation. It also specifies refined VFM-derived priors, experimental configurations, and the sampling protocol used for evaluation.

  • B.1 Metric Definitions: SSC measures whether latent-token grouping preserves object-aware spatial structure, with higher values indicating better alignment.It compares spectral clustering of latent-token affinities against projected COCO panoptic labels using geometric-mean normalized mutual information.
  • B.1 Metric Definitions: LPC measures decoded perceptual stability under local latent perturbations, where smaller values indicate less perceptual change.The multi-scale diagnostic uses relative perturbation scales R = {0.1, 0.5, 1.0, 2.0}, with larger weights assigned to smaller scales.
  • B.1 Metric Definitions: GSQ measures semantic neighborhood quality through nearest-neighbor class purity, with larger values indicating more class-consistent latent neighborhoods.It is designed for retrieval-style or representation-rich tokenizers whose class manifolds may be locally pure without being globally unimodal.
  • B.1 Metric Definitions: eRank is a supplementary diagnostic of latent-channel utilization rather than a primary geometry objective.A larger eRank indicates more even channel usage, while a smaller eRank indicates concentration in a few dominant directions.
  • B.2 Why These Metrics Matter for Diffusion Learning: SSC, LPC, and GSQ provide complementary empirical views of spatial organization, local continuity, and semantic neighborhood quality relevant to diffusion learning.Their usefulness is supported primarily by controlled comparisons and correlation analysis with downstream generation quality, rather than a strict causal or theorem-level relationship.
  • C.1 Main Experiment Configurations: The implementation uses standard LightningDiT-XL generator configurations on ImageNet 256 × 256, with 80-epoch and 800-epoch setups.For MCR, normalized latent perturbations use maximum angular deviations of 42.5° and 85° for small and large perturbations.
  • C.2 Refining VFM Prior: PAE constructs fixed semantic and structural priors through separate VFM refinement paths before final tokenizer training.The semantic path compresses frozen VFM features to d = 32, while the structural path uses AnyUp upsampling, low-pass normalization, and resizing; Gram alignment preserves coarse spatial relations.
  • C.3 Sampling and Evaluation Protocol: Generative evaluation uses SDE sampling without CFG, ODE sampling with CFG, 250 sampling steps, and class-uniform ImageNet sampling.These settings follow the evaluation protocol used in prior representation-native autoencoder works.

C.4.1 Pilot Studies · C.4.2 Ablation on regularization strategy · D More Ablation and Discussion

The pilot studies use controlled, single-factor interventions to test how reconstruction, spatial structure, local continuity, and global semantics relate to generation quality. Regularization ablations then compare prior alignment with KL and auxiliary diffusion alternatives under matched tokenizer and generator settings.

  • C.4.1 Pilot Studies: Pilot groups vary one latent-manifold factor while keeping the tokenizer scaffold, optimization, downstream generator, and evaluation protocol fixed within each group.The groups target reconstruction-versus-generation mismatch, spatial structure, local continuity, and global semantics.
  • C.4.1 Pilot Studies: Group 1 sweeps bottleneck dimension d ∈{32, 48, 64, 96, 128} using only L1 and LPIPS losses to compare reconstruction fidelity with generation quality.No prior-alignment or manifold-regularization terms are used in this group.
  • C.4.1 Pilot Studies: Group 2 activates only SSR and sweeps λSSR ∈{0, 0.05, 0.1, 0.2, 0.5}, while fixing λMCR=0 and disabling SCR, to isolate spatial structure.The baseline latent shape is 16×16×32 with L1+LPIPS reconstruction losses.
  • C.4.1 Pilot Studies: Group 3 activates only MCR and sweeps λMCR ∈{0, 0.05, 0.15, 0.3, 0.5}, while fixing λSSR=0 and disabling SCR, to isolate local continuity.This intervention uses local perturbation regularization without additional spatial-structure or semantic-alignment supervision.
  • C.4.1 Pilot Studies: Group 4 activates only SCR and sweeps λSCR ∈{0, 0.1, 0.3, 0.6, 1.0}, applying it only to the globally pooled token to isolate global semantics.λSSR=0 and λMCR=0 throughout this group.
  • C.4.2 Ablation on regularization strategy: Baseline 1 removes the full prior-alignment loss Lp, meaning SSR, MCR, and SCR are all absent during tokenizer training.It is the reference model for measuring prior alignment’s contribution.
  • C.4.2 Ablation on regularization strategy: Baseline 2 adds KL regularization LKL = DKL(q(z | x) ∥N(0, I)) with weight 10−6 to test generic distributional regularization without geometry-targeted prior alignment.It starts from Baseline 1.
  • C.4.2 Ablation on regularization strategy: Baseline 3 adds a 2-layer DiT diffusion-regularizer branch during tokenizer training, predicting standard diffusion targets on noisy latent codes without using the branch at inference.The auxiliary diffusion loss tests direct diffusion-oriented encouragement of the latent space.

D.1 Pearson correlation analysis for Manifold Metrics

Across diverse tokenizer families, diffusion-friendly manifold metrics generally align with downstream SiT-XL generation quality, with LPC showing the clearest monotonic relationship. These correlations support broader relevance while indicating that metrics should be treated as complementary rather than individually decisive.

  • Cross-tokenizer validation: Cross-tokenizer evaluations show that manifold metrics relate to SiT-XL gFID across VAE-based, vector-quantized, masked-token, and representation autoencoders.The analysis covers diverse existing autoencoders and tokenizers rather than only the proposed tokenizer design.
  • Correlation convention: After coordinate normalization, positive slopes or Pearson coefficients consistently indicate that better latent geometry aligns with better generation quality.Normalization uses sign flips for gFID and LPC and a log transform for GSQ to reduce scale skew.
  • Observed trends: LPC exhibits the clearest monotonic relation among the four metrics, suggesting that local path continuity is strongly associated with diffusion quality across tokenizers.The observed trends are directionally consistent with the design motivation across different tokenizer families.
  • Generalization: The correlations extend beyond the main experimental setting because metrics are computed independently while generation quality is measured after training diffusion transformers on each latent space.This supports interpreting the manifold properties as broader influences on diffusion performance rather than artifacts of one tokenizer or training recipe.
  • Limitations: No single metric fully determines generation quality, so the metrics are best viewed as complementary descriptors whose correlations can vary across tokenizer families.Correlation strength may also be affected by latent dimensionality, representation scale, and encoder-decoder inductive bias.

D.2 Few-Step Sampling Results · D.3 Why does PAE achieve a better fidelity–learnability balance?

PAE reaches strong few-step generation rapidly, matching FAE’s 250-step gFID in 15 steps and achieving 1.05 gFID at 45 steps. Its fidelity–learnability balance reflects strong, balanced latent geometry across diffusion-relevant dimensions rather than a single dominant factor.

  • D.2 Few-Step Sampling Results: Under identical LightningDiT-XL/1 training and classifier-free guidance settings, PAE is compared fairly against FAE for few-step sampling.Both models use 800 training epochs; PAE uses the DINOv2 tokenizer.
  • D.2 Few-Step Sampling Results: 16.7× fewer inference steps: PAE matches FAE’s 250-step gFID using only 15 sampling steps.PAE quickly approaches its full-sampling performance as inference steps increase.
  • D.2 Few-Step Sampling Results: 1.05 gFID at 45 sampling steps is PAE’s reported few-step result, while PAE also achieves substantially higher IS than FAE in the few-step regime.Figure 13 reports consistently higher IS for PAE than FAE across few-step sampling.
  • D.3 Why does PAE achieve a better fidelity–learnability balance?: Existing tokenizers exhibit a reconstruction-fidelity versus downstream-learnability trade-off, whereas PAE combines strong reconstruction with the best learnability.PAE’s gain therefore does not come from sacrificing either reconstruction fidelity or learnability.
  • D.3 Why does PAE achieve a better fidelity–learnability balance?: PAE is simultaneously strong in spatial structure (SSC), local continuity (LPC), and global semantics (GSQ), while maintaining high latent utilization measured by eRank.The results suggest that PAE succeeds through balanced latent geometry rather than one dominant property.
  • D.3 Why does PAE achieve a better fidelity–learnability balance?: DINO-based PAE is most balanced across SSC, LPC, and GSQ, consistent with its strongest downstream generation, while SigLIP-based PAE has stronger semantics but weaker spatial structure and moderate continuity.Different VFM backbones therefore induce different geometry profiles under the same tokenizer design; SigLIP remains competitive but does not match DINO-based PAE.

D.4 Ablation on perturbation design for MCR

The ablation isolates perturbation design while holding the tokenizer backbone and other objectives fixed, showing that cascaded perturbations best support latent continuity and generation quality. Small and large perturbations provide complementary benefits, with cascaded perturbation combining both levels progressively.

  • Experimental setup: All variants share the tokenizer backbone, reconstruction, SSR, and SCR objectives, differing only in perturbation design to test explicit continuity regularization.Perturbations are applied along a random normalized direction in the RMS-normalized, sphere-like latent space.
  • Perturbation designs: 42.5° maximum angular deviation defines Small Perturb, while 85° defines Large Perturb.Cascaded Perturb uses both perturbation levels with the progressive consistency objective in Eq. (4).
  • Ablation findings: Generic perturbation consistency improves over removing MCR, confirming that local latent regularization is useful.The comparison indicates that MCR benefits are not explained solely by generic robustness.
  • Ablation findings: Small perturbations improve LPC and gFID but have limited effect, whereas the proposed cascaded design achieves the best continuity and generation quality.The cascaded design combines small and large perturbations through progressive consistency.

D.5 Cross-Encoder Quantitative Results. … E.5 Tokenizer Latent Interpolation

PAE generalizes across frozen teachers and improves reconstruction fidelity while explicitly strengthening spatial structure, global semantics, and local continuity in its latent manifold. Qualitative visualizations show smoother semantic transitions and better preservation of fine details, supporting diffusion-friendly representations.

  • D.5 Cross-Encoder Quantitative Results.: PAE consistently improves gFID and IS across DINOv2, SigLIP2, DINOv3, and MAE, with DINOv2 and DINOv3 achieving the strongest overall performance.The gains show that prior alignment is not tied to a particular frozen representation backbone.
  • E.1 Additional Reconstruction Visualizations: PAE achieves superior reconstruction fidelity to SD-VAE and RAE, especially preserving thin structural lines and complex textual information.These high-frequency details are often blurred or lost by reconstruction-oriented baselines.
  • E.3 Global Semantic Organization: Semantic Consistency Regularization produces tighter class-wise latent clustering by aligning representations with compressed VFM semantic priors.Improved Global Semantic Quality facilitates conditional generative modeling and faster convergence.
  • E.4 DiT Latent Interpolation: PAE maintains semantic coherence and image quality during linear and spherical linear interpolation between noise vectors in a trained DiT latent space.The visualization indicates a robust, continuous feature space supporting smooth semantic transitions.
  • E.5 Tokenizer Latent Interpolation: Manifold Continuity Regularization yields smooth interpolations across semantic categories by enforcing a locally Lipschitz-continuous manifold.This local smoothness minimizes the LPC and makes small latent perturbations correspond to gradual perceptual changes.

E.6 More Generation results

The paper provides additional PAE visualization results using classifier-free guidance with w = 3.3 in Figs. 21–35.

  • Additional PAE visualizations are presented in Figs. 21–35 with CFG (w = 3.3).These figures extend the reported generation visualizations for PAE.

F Limitations and Future Works · G Broader Impacts · Input PAE RAE SVG SD-VAE

The study’s limitations concern validation scope, fixed-resolution settings, and reliance on explicit VFM-derived supervision, motivating broader and more adaptive evaluation. PAE is presented as a principled tokenizer-design framework, with additional visualizations illustrating improved reconstruction, manifold organization, and interpolation behavior.

  • F Limitations and Future Works: The main experiments use only ImageNet at 256×256, leaving generalization to higher resolutions, diverse visual domains, and downstream tasks beyond class-conditional synthesis untested.The authors describe this as a standard but relatively controlled setting.
  • F Limitations and Future Works: The framework focuses on fixed-resolution latent diffusion and does not yet address variable spatial scales, dynamic token allocation, or resolution-adaptive generation.These settings may interact more strongly with compression ratio and spatial capacity.
  • F Limitations and Future Works: PAE’s effectiveness currently depends on refined VFM-derived supervision and carefully designed regularization terms, leaving more natural alternatives unresolved.Open alternatives include stronger tokenizer pretraining, larger-scale data, and unified self-supervised objectives without handcrafted alignment losses.
  • F Limitations and Future Works: Future work will evaluate the perspective with larger-scale training, higher and dynamic resolutions, broader generation settings, and stronger tokenizer pretraining.The goal is to induce diffusion-friendly manifold organization more directly and robustly.
  • G Broader Impacts: PAE reframes latent-diffusion tokenizer design around explicit latent-manifold organization, improving generation quality and training efficiency beyond reconstruction-oriented objectives alone.This is presented as a principled framework for tokenizer design.
  • Input PAE RAE SVG SD-VAE: Additional reconstruction comparisons show that PAE consistently preserves finer visual details than representative tokenizer baselines.The same finding is reported in Figures 14 and 15.
  • Input PAE RAE SVG SD-VAE: Prior alignment produces clearer patch-wise spatial structure, more compact class-consistent neighborhoods, and smoother local and downstream DiT-space interpolations.These visualizations respectively support improved spatial relations, semantic organization, local manifold continuity, and a stable latent space for diffusion.
Loading 2605.07915v1…