PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving the Diffusability of Autoencoders

Ivan Skorokhodov, Sharath Girish, Benran Hu, Willi Menapace, Yanyu Li, Rameen Abdal, Sergey Tulyakov, Aliaksandr Siarohin

arXiv:2502.14831v3cs.CVcs.AIcs.LG

TL;DR

Latent diffusion research has focused less on the interaction between autoencoders and diffusion backbones, particularly the diffusability of latent representations. The paper analyzes latent spectra, identifies excessive high-frequency components, and introduces scale equivariance in the decoder to align latent and RGB frequencies. The approach improves generation quality across image and video tasks, reducing FID by 19% and FVD by at least 44%.

  • Problem

    Prior autoencoder work emphasizes reconstruction quality and compression, leaving diffusability and autoencoder–diffusion interaction comparatively under-explored.

  • Method

    The paper analyzes autoencoder latent spectra and enforces decoder scale equivariance so downsampled latents correspond to downsampled RGB representations.

  • Results

    19% lower FID is reported for ImageNet-1K 256^2 image generation, and at least 44% lower FVD for Kinetics-700 17 × 256^2 video generation.

  • Takeaways & Limitations

    Aligning latent and RGB spaces across frequencies improves diffusability and LDM generation quality while maintaining reconstruction fidelity.

  • Takeaways & Limitations

    Fine-tuning CMS-AEI with reconstruction losses increased downstream FID by 14%, from 11.69 to 13.59.

Abstract

from arXiv · show

Latent diffusion models have emerged as the leading approach for generating high-quality images and videos, utilizing compressed latent representations to reduce the computational burden of the diffusion process. While recent advancements have primarily focused on scaling diffusion backbones and improving autoencoder reconstruction quality, the interaction between these components has received comparatively less attention. In this work, we perform a spectral analysis of modern autoencoders and identify inordinate high-frequency components in their latent spaces, which are especially pronounced in the autoencoders with a large bottleneck channel size. We hypothesize that this high-frequency component interferes with the coarse-to-fine nature of the diffusion synthesis process and hinders the generation quality. To mitigate the issue, we propose scale equivariance: a simple regularization strategy that aligns latent and RGB spaces across frequencies by enforcing scale equivariance in the decoder. It requires minimal code changes and only up to 20K autoencoder fine-tuning steps, yet significantly improves generation quality, reducing FID by 19% for image generation on ImageNet-1K $256^2$ and FVD by at least 44% for video generation on Kinetics-700 $17 \times 256^2$. The source code is available at https://github.com/snap-research/diffusability.

1. Introduction

Latent diffusion models depend not only on diffusion backbones and autoencoder reconstruction, but also on how well autoencoder latents support diffusion. The paper links poor diffusability to excessive latent high frequencies and proposes scale-equivariant decoder regularization to improve image and video generation.

  • Latent diffusion models combine an autoencoder with a diffusion backbone, while recent progress has emphasized scaling the backbone more than improving autoencoders.
  • Reconstruction quality, compression efficiency, and diffusability jointly affect LDM effectiveness, with poor diffusability requiring heavier diffusion backbones.
  • Diffusion denoising is coarse-to-fine, synthesizing low-frequency components before high-frequency details, but pretrained autoencoders can disrupt latent–RGB frequency correspondence.
  • Modern autoencoder latents contain prominent high-frequency components that diverge from RGB spectra and become stronger with larger channel sizes.
  • Scale equivariance aligns latent and RGB spaces across frequencies by enforcing decoder correspondence between downsampled latents and downsampled RGB representations.
  • 19% lower FID is reported for image generation on ImageNet-1K 256^2, while video generation on Kinetics-700 17 × 256^2 achieves at least 44% lower FVD.

2. Related Work

Related work establishes diffusion models and latent autoencoders as efficient frameworks for visual generation and compression. It also situates this paper among spectral-autoregression studies, autoencoder designs, compression methods, and concurrent equivariance approaches.

  • Diffusion models: Diffusion models generate content by progressively denoising an initial noise sample and have surpassed traditional GAN and VAE approaches.
  • Diffusion models: Latent diffusion models model high-resolution visual content in compact autoencoder-produced latent spaces through a two-stage framework.
  • Diffusion models: Prior work attributes diffusion success partly to implicit spectral autoregression, where progressive denoising generates natural images coarse-to-fine.
  • Image and video autoencoders: Image and video diffusion systems use spatial or spatio-temporal autoencoders with compression across spatial and temporal dimensions.
  • AEs for compression: Compression autoencoders target high-quality reconstruction with low bitrate, using downsampled latents and, for video, motion-aware or implicit temporal modeling.
  • Concurrent works: Concurrent methods apply shift or scale equivariance to autoencoders and diffusion models, motivated by spatial-transformation equivariance rather than diffusability.

3. Improving Diffusability

The analysis links latent spectral profiles to diffusability, finding excessive high-frequency components that grow with bottleneck size and are amplified by KL regularization. Scale equivariance regularization aligns latent spectra more closely with RGB and preserves content under high-frequency suppression.

  • Spectral analysis: DCT converts blockwise 2D signals from spatial values into frequency coefficients, whose normalized amplitudes form a zigzag-ordered frequency profile.The analysis treats channels independently and uses per-frame 2D DCT for videos or temporally extended latent codes.
  • Spectral analysis: FluxAE latents contain larger high-frequency components than RGB, and these components become more pronounced as bottleneck channel size increases.The study averages profiles across 256 samples, all channels, and all DCT blocks.
  • Spectral analysis: Higher KL regularization introduces more latent high frequencies through noise addition, creating a harmful side effect that reduces diffusability.KL regularization also aligns latents with the standard Gaussian prior used to initialize reverse diffusion.
  • Spectral analysis: Diffusion generates low frequencies before high frequencies, making flatter latent spectra less suitable because low frequencies provide less clean conditioning.Higher frequencies are also described as harder to model because they have higher dimensionality, emerge late, and accumulate errors more readily.
  • Scale equivariance regularization: High-frequency erasure quickly degrades vanilla FluxAE reconstructions, whereas high-frequency cutoff regularization improves low-frequency reliance, compression, and resilience to diffusion error accumulation.Scale-equivariant fine-tuning also preserves more content than the baseline as high-frequency components are progressively suppressed.
  • Scale equivariance regularization: Scale equivariance regularization enforces decoder consistency between downsampled latents and downsampled RGB representations, reducing latent high frequencies toward the RGB spectrum.Its objective adds a reconstruction term weighted by α; the main experiments use α = 0.25.

4. Experiments

Experiments evaluate scale-equivariant autoencoder regularization across image and video latent diffusion settings, including reconstruction, KL, spectral, and time-shifting ablations. The method improves downstream generation while generally preserving reconstruction quality, though its benefit varies across autoencoders.

  • Improving video autoencoders: 44% and 54% lower FVD10K are obtained for CV-AE and LTX-AE, respectively, with scale-equivariant regularization on Kinetics-700.Video autoencoders receive 20,000 fine-tuning steps, and the regularized variant improves performance in each case.
  • Reconstruction quality: Reconstruction quality remains similar across models on PSNR, SSIM, LPIPS, and FID despite improved downstream generation.The comparison uses 50,000 ImageNet and Kinetics samples for image and video autoencoders.
  • Ablations: Increasing KL strength can improve small-model LDM performance but worsens reconstruction and limits larger-model scaling, unlike scale equivariance.The best small-model score occurs at KL β = 0.1, while the corresponding larger model ranks among the worst.
  • Ablations: The regularized FluxAE-FT+SE performs best across tested time-shifting coefficients, separating its quality improvement from implicit time shifting.The comparison uses DiT-B/2 models and FDD5K results with 128 inference steps without classifier-free guidance.
  • Ablations: Increasing SE loss weight worsens total reconstruction quality, so 0.25 is selected to maintain base-autoencoder reconstruction while improving generation.The decoder must generalize across low- and high-frequency latents with fixed capacity.
  • Ablations: Scale equivariance reduces latent high-frequency components toward RGB spectra and preserves more content under progressive frequency removal.The ablation examines both spectral alignment and reconstruction under increasingly aggressive high-frequency suppression.

5. Conclusion

The paper argues that autoencoder latent spectra are as important as diffusion architectures for latent diffusion quality. It identifies excessive high-frequency content as harmful and proposes frequency-alignment regularization that preserves reconstruction while improving diffusion training.

  • Autoencoders are as critical as diffusion architectures for modern latent diffusion models.
  • Excessive high-frequency latent components create discrepancies with RGB spectra and lead to worse LDM synthesis quality.
  • The proposed regularization aligns latent and RGB spaces across frequencies while suppressing spurious high-frequency details without reducing reconstruction fidelity.

Impact Statement

The work improves autoencoder representations for latent diffusion and may support creative and educational generation applications while reducing computational overhead.

  • Improved autoencoder representations are intended to enhance generative performance in latent diffusion workflows.
  • The authors identify potential benefits for creativity, educational content creation, and computational overhead in generative workflows.

A. Limitations.

The study identifies scope, reconstruction, hyperparameter, and training-pipeline limitations that constrain how broadly its regularization results can be interpreted.

  • Evaluation scope: The evaluation covers four autoencoders across image and video domains, but not other diffusion parametrizations or architectures.The authors state that broader verification would make the study more complete.
  • Reconstruction trade-off: FluxAE reconstruction worsened under regularization, with FID increasing from 0.183 to 0.55, although adversarial losses might mitigate this effect.CogVideoX-AE showed improvement instead, so the reconstruction impact is not uniform across autoencoders.
  • Hyperparameters: Results show mild sensitivity to hyperparameters, including the SHF regularization weight and whether a small KL regularization is added.The authors disabled KL regularization in the final regularized setup for simplicity.
  • Reproducibility: None of the explored autoencoders released training pipelines, and fine-tuning DC-AE produced divergent reconstructions across many tested hyperparameter setups.This makes reproducing or extending the fine-tuning procedure non-trivial.
  • Future work: The authors leave these evaluation, reconstruction, hyperparameter, and training-pipeline limitations for future work.This is stated as the paper’s closing limitation statement.

B. Implementation Details

The implementation combines standard diffusion training components with autoencoder-specific spectral regularization experiments. The main regularization variants either remove or penalize high-frequency components, while scale-equivariant fine-tuning preserves decoder output layers and uses downsampling-based consistency.

  • DiT setup: The DiT models use self-conditioning, RoPE positional embeddings, and rectified-flow diffusion parametrization.These additions are presented as baseline-strengthening choices.
  • DiT setup: DiT models train for 400,000 steps with 10,000 warmup steps, cosine decay, AdamW, weight decay 0.01, and posterior sampling for VAE autoencoders.The learning rate decays from 0.0003 toward 0.00001.
  • DiT setup: DiT-S, B, L, and XL use progressively larger hidden dimensions, transformer depths, and attention-head counts.The listed hidden dimensions are 384, 768, 1024, and 1152, respectively.
  • Training infrastructure: Training uses gradient clipping and FSDP across one node with 8× A100 or H100 80GB GPUs, while CV-AE experiments use pre-extracted random 17-frame latents.The CV-AE procedure reduces dataset size but is applied consistently across the CogVideoX-AE family.
  • Autoencoder training: Autoencoder recipes individually tune learning rates and KL weights, use AdamW with mixed precision, maintain exponential moving averages, and disable KL regularization for scale-equivariant models.The EMA half-life is 5,000 steps, and other baseline hyperparameters remain unchanged when adding the regularizer.
  • Autoencoder training: The method freezes the decoder’s final normalization and output convolution layers because they were adversarially fine-tuned for texture reconstruction.The frozen parameters are described as negligible relative to the full model.
  • Spectral regularization variants: Explicit high-frequency chopping applies DCT masking to latents and RGB targets, then uses inverse DCT to return them to the spatial domain.The binary mask zeros selected frequencies, with N controlling the cutoff.

C.3. ImageNet 5122 experiments

The section describes ImageNet-1K 512^2 experiments and regularization-strength ablations for FluxAE, alongside additional from-scratch evaluations for image and video autoencoders.

  • ImageNet-1K 512^2: The ImageNet-1K 512^2 generation experiment trains DiT-L/2 for 400,000 steps on FluxAE latents.The setup uses class-conditional generation without guidance.
  • Ablations: The regularization-strength ablation trains FluxAE for 10,000 steps with varying strength α.The corresponding results are reported in Table 9.
  • Experimental materials: The appendix includes a DCT zigzag indexing order for frequency selection and a table describing the ImageNet-1K 512^2 generation results.The supplied passages identify these artifacts but do not report their numerical outcomes.
  • From-scratch experiments: Additional from-scratch experiments train FluxAE on ImageNet 256^2 and CogVideoX-AE on Kinetics-700 17 × 256^2, with DiT models trained on their resulting latents.The experiments compare autoencoders trained with and without the proposed regularization.

C.6. Compute cost analysis

The compute analysis tests whether scale-equivariant gains arise from extra autoencoder computation or simply from longer fine-tuning. The reported comparisons favor the regularized approach under the tested setup.

  • Additional analyses: The appendix includes from-scratch result tables and latent-spectrum figures for WanAE and LTX-AE under varying KL regularization or channel size.The supplied passages identify these materials without reporting their numerical results.
  • Additional compute: FluxAE requires 447 GFLOPs per encoder-decoder pass, while regularization adds 77.6 or 19.4 GFLOPs through reduced-resolution decoder passes.These correspond to approximately +17% or +4% extra compute, respectively.
  • Additional compute: The regularization reuses the encoder pass and samples ×2 or ×4 reduced resolution for its additional decoder computation.The resulting overhead is approximately one-quarter or one-sixteenth of decoder compute.
  • Longer fine-tuning comparison: With 20K baseline fine-tuning steps, DiT-B/2 reached FID5K 33.99 and FDD5K 642.7, compared with 25.87 and 551.27 for regularized FluxAE fine-tuned for 10K steps.The comparison is used to assess longer baseline fine-tuning versus regularization.

D.2. Samples visualizations

The visualizations show uncurated class-conditional image and video samples generated with different autoencoders, fine-tuning variants, and diffusion backbones.

  • ImageNet samples compare FluxAE, FluxAE + FT, and FluxAE + SE using DiT-XL/2 at 256 × 256.Inference used 256 steps and guidance scale 3.0 for random classes.
  • Additional ImageNet visualizations use DiT-XL/2 trained for 1M steps with FluxAE + SE at 256 × 256.
  • An ImageNet visualization uses DiT-XL/2 trained for 400K steps with FluxAE + SE.
  • A separate ImageNet comparison uses DiT-B/1 with CMS-AEI, CMS-AEI + FT, and CMS-AEI + SE.The supplied caption specifies class-conditional ImageNet samples at 256 × 256 and guidance scale 1.5.
  • Video visualizations compare CogVideoX-AE, CogVideoX-AE + FT, and CogVideoX-AE + SE on Kinetics at 17 × 256 × 256.They use DiT-XL/2 or DiT-B/2, with 256 inference steps and guidance scale 3.0.
  • Another video visualization compares LTX-AE, LTX-AE + FT, and LTX-AE + SE using DiT-B/1.

E. Failed experiments

The paper reports several regularization experiments that did not improve over the selected scale-equivariance approach or the vanilla baseline, often with practical drawbacks.

  • The explored alternatives were presented as unsuccessful directions retained to encourage future exploration.
  • MiniLDM-regulzation training: Training a small LDM jointly with the autoencoder did not outperform the vanilla baseline despite extensive hyperparameter search.This experiment was motivated by optimizing the autoencoder and LDM together to make the latent space more diffusable.
  • Lipszhitz regularization: R1 regularization performed almost on par with scale equivariance but caused much slower training and engineering difficulties.The slowdown came from second-order differentiation and poor compatibility with FSDP; scale equivariance performed slightly better and was simpler conceptually.
  • Temporal scale equivariance: Temporal scale-equivariance regularization did not lead to improved results.The authors hypothesize that temporal high-frequency components are more noticeable to humans than low-frequency components.
Loading 2502.14831v3…