PaperlayerView exact revision on arXiv

Source-linked AI summary

Geometric Autoencoder for Diffusion Models

Hangyu Liu, Jianyong Wang, Yutao Sun

arXiv:2603.10365v2cs.CV

TL;DR

Existing diffusion latent spaces are largely heuristic and must balance semantic discriminability, reconstruction fidelity, and compactness. GAE addresses this with VFM-derived semantic supervision, latent normalization without KL divergence, and dynamic noise sampling. On ImageNet 256 × 256, it reaches 1.31 gFID at 800 epochs and 1.82 at 80 epochs without Classifier-Free Guidance, while maintaining a favorable compression–semantic-depth–reconstruction balance.

  • Problem

    Existing latent-space designs for diffusion models lack principled guidance and struggle to jointly optimize semantic information, compactness, and reconstruction stability.

  • Method

    GAE analyzes alignment paradigms, derives a low-dimensional VFM semantic target, replaces KL divergence with latent normalization, and adds dynamic noise sampling.

  • Results

    1.31 gFID at 800 epochs and 1.82 at 80 epochs are achieved on ImageNet 256 × 256 without Classifier-Free Guidance.

  • Takeaways & Limitations

    GAE establishes a favorable equilibrium among compression, semantic depth, and robust reconstruction stability for latent diffusion modeling.

Abstract

from arXiv · show

Latent diffusion models have established a new state-of-the-art in high-resolution visual generation. Integrating Vision Foundation Model priors improves generative efficiency, yet existing latent designs remain largely heuristic. These approaches often struggle to unify semantic discriminability, reconstruction fidelity, and latent compactness. In this paper, we propose Geometric Autoencoder (GAE), a principled framework that systematically addresses these challenges. By analyzing various alignment paradigms, GAE constructs an optimized low-dimensional semantic supervision target from VFMs to provide guidance for the autoencoder. Furthermore, we leverage latent normalization that replaces the restrictive KL-divergence of standard VAEs, enabling a more stable latent manifold specifically optimized for diffusion learning. To ensure robust reconstruction under high-intensity noise, GAE incorporates a dynamic noise sampling mechanism. Empirically, GAE achieves compelling performance on the ImageNet-1K $256 \times 256$ benchmark, reaching a gFID of 1.82 at only 80 epochs and 1.31 at 800 epochs without Classifier-Free Guidance, significantly surpassing existing state-of-the-art methods. Beyond generative quality, GAE establishes a superior equilibrium between compression, semantic depth and robust reconstruction stability. These results validate our design considerations, offering a promising paradigm for latent diffusion modeling. Code and models are publicly available at https://github.com/sii-research/GAE.

1 Introduction

The paper identifies heuristic latent-space design as a central limitation in diffusion models and proposes GAE as a principled alternative combining semantic guidance, latent normalization, and robust reconstruction. GAE reports strong ImageNet generation, compression, and semantic-performance results.

  • Results: GAE establishes a superior Pareto curve between linear probing accuracy and latent dimension while accelerating convergence with 32 latent dimensions.The figure summarizes the relationship between semantic discriminability, compactness, and generation efficiency.
  • Motivation: Existing latent-space designs remain heuristic, with alignment strategies, dimensional constraints, and reconstruction stability limiting generative performance.The paper frames these issues as a lack of principled guidance for understanding latent-space properties.
  • Method: GAE analyzes alignment paradigms and constructs a low-dimensional semantic representation from VFMs to guide the autoencoder.This semantic target is presented as the most effective guidance for latent-space design among the analyzed approaches.
  • Method: GAE replaces the restrictive VAE KL-divergence objective with latent normalization and adds dynamic noise sampling for stable, noise-robust reconstruction.The framework targets a stable latent manifold optimized for diffusion learning under high-intensity noise.
  • Results: 1.31 gFID at 800 epochs and 1.82 at 80 epochs are achieved on ImageNet 256 × 256 without Classifier-Free Guidance.The paper reports that these results significantly surpass existing state-of-the-art methods.

2 Related Work

Related work integrates semantic representations into diffusion models through Vision Foundation Models, semantic supervision, alignment losses, and reconstruction-capable VFM integrations. These directions seek to combine perceptual understanding with generative reconstruction and improve learning efficiency.

  • Vision Foundation Models: Vision Foundation Models provide semantic priors through vision-language and visual self-supervised representation-learning paradigms.Examples include contrastive visual-language models and reconstructive or discriminative self-supervised methods.
  • Representation Alignment for Generation: Semantic representation integration has become a strategy for enhancing diffusion learning efficiency and accelerating convergence.The related-work landscape focuses on bridging reconstructive autoencoding and representation learning.
  • Representation Alignment for Generation: Prior approaches augment VAEs with semantic supervision, align VAE latents with VFM feature spaces, or integrate VFMs with reconstruction capability.These are described as three converging trajectories in the field.

3 Geometric Autoencoder

GAE uses dual pixel and frozen-semantic branches, RMSNorm-based latent normalization, dynamic noise sampling, and multi-objective reconstruction and semantic-preservation losses. The objective removes KL divergence while combining reconstruction fidelity, perceptual quality, adversarial training, and semantic alignment.

  • 3.1 Architecture: GAE uses a dual-branch architecture with pixel encoder-projector-decoder components and a frozen VFM semantic-teacher pipeline.The semantic branch includes a VFM followed by a downsampler, while the pixel branch processes the input for reconstruction.
  • 3.1 Architecture: RMSNorm maps projected pixel features onto a unit hypersphere, bounding latent values and preventing training collapse.The parameter-free normalization is intended to produce a numerically stable, well-distributed latent space.
  • 3.1 Architecture: Dynamic noise sampling perturbs the normalized latent mean with Gaussian noise at sampled scales before decoding the latent vector into a reconstruction.The process trains a continuous manifold across varying noise levels.
  • 3.2 Objective: The pixel reconstruction objective combines pixel-level L1, perceptual, and adversarial losses.These terms are weighted by λrec, λlpips, and λgan in the multi-objective reconstruction loss.
  • 3.2 Objective: Semantic Preservation aligns the pixel-branch latent mean with downsampled VFM features using an MSE loss.This alignment is designed to retain high-level semantic integrity and discriminability in the compressed latent space.
  • 3.2 Objective: GAE removes the standard VAE KL penalty and instead uses latent normalization with dynamic noise sampling to support bounded, robust representations for diffusion denoising.The total objective integrates reconstruction, perceptual, adversarial, and semantic-preservation terms.

4 Latent Alignment for Antoencoder

The section compares three semantic-alignment stages and identifies bottleneck-level Latent Alignment as the most effective way to preserve semantic integrity while retaining reconstruction and compactness.

  • Taxonomy of Alignment: Three paradigms reconcile VFM and autoencoder dimensionality at different stages: encoder features, expanded latents, or the compressed bottleneck.Pre Alignment supervises high-dimensional encoder features, Post Alignment expands compressed latents, and Latent Alignment projects VFM features into the compact latent space.
  • Latent Alignment: Latent Alignment uses a semantic downsampler to project high-dimensional VFM outputs into the compact latent space.This design addresses the mismatch between example latent dimensionality d = 32 and VFM dimensionality d = 1024 for ViT-L.
  • Pilot Study Setup: Pilot studies use ViT-L autoencoders, a frozen DINOv2-L/14 teacher, 60 training epochs, and Linear Probing alongside reconstruction evaluation.A pre-computed SVD matrix maps 1024-dimensional VFM features onto a 32-dimensional manifold for Latent Alignment.
  • Analysis and Findings: Latent Alignment preserves semantic discriminability better than Pre Alignment, which has slightly better reconstruction but loses substantial semantic information after compression.The findings indicate that supervising high-dimensional features before projection does not ensure semantic integrity survives bottleneck compression.
  • Training Regime of Semantic Teacher: A parametric downsampler replaces static SVD projection with learned nonlinear mapping to reconcile structural integrity and latent compactness.Its feature-autoencoder training uses cosine distillation between frozen VFM patch tokens and reconstructed counterparts, after which the decoder is discarded.
  • Downsampler Architecture: Patch Conv significantly outperforms Single Attention and Attention+Linear in Linear Probing, supporting a spatial-aware learned downsampler for discriminative teacher latents.The ablation evaluates three downsampler architectures: Single Attention, Attention+Linear, and Attention+Patch Conv.

5 Experiments

Experiments evaluate GAE across generation, reconstruction, semantic discriminability, noise robustness, scalability, and hyperparameter sensitivity on ImageNet. GAE generally combines compact latents with strong semantic, reconstruction, and generative performance.

  • Generative Performance: GAE achieves 1.31 gFID without CFG after 800 epochs and 1.13 with CFG, surpassing reported baselines on ImageNet 256 × 256.FAE reports 1.48 and 1.29, while RAE reports 1.51 and 1.13 for the corresponding settings.
  • Reconstruction Stability: Higher latent-noise training improves decoder tolerance to latent distribution shifts, while dynamic noise sampling balances reconstruction fidelity against robustness.For d = 32, Cσ = 0.1 leads early in rFID, whereas Cσ = 0.2 provides greater robustness as noise increases.
  • Training Efficiency and Convergence: 1.82 gFID at 80 epochs surpasses VA-VAE trained for 800 epochs, indicating rapid convergence for the diffusion model.The paper attributes this convergence to a semantically aligned latent space simplifying the subsequent learning objective.
  • Semantic-Reconstruction Trade-off: At 32 latent dimensions, GAE reaches 69.4% linear-probing accuracy, compared with 43.1% for VA-VAE at the same dimension.The comparison is presented as evidence that semantic alignment need not sacrifice visual fidelity.
  • Semantic and Reconstruction: At 64 dimensions, GAE achieves 78.3% linear-probing accuracy and 0.382 rFID, outperforming VTP-L on semantics and FAE on reconstruction.The corresponding baseline values are 73.9% for VTP-L and 0.660 for FAE.

6 Conclusion

The conclusion presents GAE as a principled response to heuristic latent-space design. It reports a balance among semantic information, compression, reconstruction stability, and efficient generation.

  • Conclusion: GAE addresses heuristic latent-space design by jointly considering semantic alignment, latent distribution, and reconstruction stability.The framework is intended to bridge perceptual understanding with generative fidelity.
  • Conclusion: 1.31 gFID after 800 epochs and 1.82 after 80 epochs are achieved without Classifier-Free Guidance on ImageNet 256 × 256.The conclusion also characterizes GAE as balancing compression, semantic depth, and robust reconstruction stability.

A Qualitative Results

The qualitative-results figure shows samples generated by GAE with a 32-dimensional latent space after 800 training epochs. The samples use CFG to enhance visual fidelity.

  • Qualitative Results: Qualitative samples are generated with GAE using d = 32 after 800 epochs of training.The figure description identifies the latent dimensionality and training duration for the displayed samples.
  • Qualitative Results: A CFG scale of w = 3.3 is used to enhance visual fidelity in the displayed samples.This setting applies to the qualitative generation shown in the figure.

C Reconstruction Quality

The reconstruction analysis compares GAE with alignment-heavy and foundation-model architectures using reconstruction and semantic metrics. GAE shows strong reconstruction while semantic priors substantially affect generation quality.

  • Reconstruction Comparison: At 64 latent dimensions, GAE achieves 29.085 PSNR and 0.834 SSIM, compared with RAE’s 18.86 PSNR and 0.62 SSIM.These results are reported as a quantitative reconstruction comparison.
  • Architecture Comparison: DINOv2-L reaches 69.2% linear-probing accuracy and 2.36 gFID, whereas MAE-L reaches 29.8% and 4.20 despite better reconstruction metrics.MAE-L reports rFID 0.355 and PSNR 28.712, while DINOv2-L has stronger semantic priors with slightly lower reconstruction performance.

E Ablation on CFG Parameters

The section defines CFG weight and interval as guidance controls, while documenting the patch-convolution downsampler and associated training and evaluation configurations.

  • CFG Parameters: CFG weight w controls guidance scale to balance sample fidelity and diversity.The CFG interval defines the temporal range during sampling where guidance is applied.
  • Architecture: The patch-wise module first projects local spatial blocks through convolution, then repartitions and flattens features into latent patches.This two-step procedure is designed to optimize computational efficiency.
  • Training Configuration: GAE uses a ViT-L backbone with downsampling factor f = 16 and DINOv2-L-based semantic supervision.The framework supports 32-dimensional and 64-dimensional latent spaces.
  • Training Configuration: The semantic preservation loss weight is λsp = 1.0, latent standard deviation is Cσ = 0.2, and GAE training lasts 200 epochs with batch size 1024.These settings are specified for the main GAE training configuration.
  • Diffusion Optimization: Diffusion training uses a constant learning rate of 2.0 × 10^-4, batch size 1024, EMA weight 0.9999, and protocol-dependent QK-Norm usage.QK-Norm is used for 800-epoch runs but omitted for 80-epoch benchmarks.
  • Evaluation: Evaluation uses an SDE-based sampler without CFG and an ODE-based sampler with CFG, with 250 sampling steps and class-uniform sampling.The protocol follows established procedures for sampling and evaluation.

H Detailed Configuration of Ablation Study

The ablation study uses a reduced training setup with fixed latent dimensions, noise scale, optimizer settings, and evaluation protocol to measure relative performance gains.

  • Training Configuration: Ablation variants use a ViT-L backbone, latent dimension d = 32, 100 GAE training epochs, and latent standard deviation Cσ = 0.1.The reduced duration is intended to accelerate the experimental cycle.
  • Diffusion Optimization: Ablation diffusion models are trained for 80 epochs with QK-Norm disabled, while other optimization hyperparameters match the primary setup.The 80-epoch duration was found sufficient to reflect relative performance gains.
  • Sampling and Evaluation: Ablation evaluation primarily reports unguided gFID and rFID using an ODE-based sampler with 250 steps and timeshift fixed at 0.7.Class labels are sampled randomly following the VAVAE evaluation standard.

I Linear Probing Evaluation Protocol

The linear probing protocol evaluates latent semantic discriminability using the latents alone, with different readouts for compressed and high-dimensional representations.

  • Evaluation Metric: Semantic discriminability is evaluated with ImageNet-1K linear probing Top-1 accuracy.The metric is used to quantify the semantic quality of the learned latent space.
  • Protocol Design: Unlike prior DINO-style probing, the protocol uses learned latents as stand-alone representations without concatenating pooled latents with DINO CLS tokens.This design measures the inherent semantic quality of the learned latents.
  • Readout Configurations: Flatten maps compressed latent tensors such as d = 32 or 64 directly to the 1,000 ImageNet classes.This readout captures spatial and semantic information within the bottleneck.
  • Readout Configurations: GAP applies Global Average Pooling to high-dimensional representations such as 1,024 dimensions before classification.Pooling avoids the excessive parameter count associated with flattening high-dimensional tensors.
  • Probe Optimization: Linear probes are trained for 100 epochs with batch size 2,048 using SGD at a constant learning rate of 0.008.No weight decay or complex learning-rate scheduling is applied.
Loading 2603.10365v2…