PaperlayerView exact revision on arXiv

Source-linked AI summary

EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling

Theodoros Kouzelis, Ioannis Kakogeorgiou, Spyros Gidaris, Nikos Komodakis

arXiv:2502.09509v3cs.LG

TL;DR

Existing autoencoders produce latent representations that are not equivariant to scaling and rotation, complicating latent generative modeling and creating a trade-off with reconstruction quality. EQ-VAE regularizes pretrained autoencoders toward equivariance without architectural changes, improving downstream generative models while preserving reconstruction quality. Five epochs of SD-VAE fine-tuning yield a ×7 DiT-XL/2 speedup, and the method applies to both continuous and discrete autoencoders.

  • Problem

    Established autoencoders lack equivariance under scaling and rotation, unnecessarily complicating latent spaces, while stronger KL regularization can reduce reconstruction detail.

  • Method

    EQ-VAE modifies the autoencoder training objective by penalizing discrepancies between transformed-latent reconstructions and transformed input images, without architectural changes.

  • Results

    EQ-VAE improves downstream DiT, SiT, REPA, and MaskGIT performance while preserving reconstruction quality; five SD-VAE fine-tuning epochs yield a ×7 DiT-XL/2 speedup.

  • Takeaways & Limitations

    EQ-VAE is a plug-and-play enhancement for continuous and discrete autoencoders across diffusion and masked generative models.

  • Takeaways & Limitations

    Minimizing the equivariance loss alone can collapse the latent representation to a constant, making explicit regularization ineffective.

Abstract

from arXiv · show

Latent generative models have emerged as a leading approach for high-quality image synthesis. These models rely on an autoencoder to compress images into a latent space, followed by a generative model to learn the latent distribution. We identify that existing autoencoders lack equivariance to semantic-preserving transformations like scaling and rotation, resulting in complex latent spaces that hinder generative performance. To address this, we propose EQ-VAE, a simple regularization approach that enforces equivariance in the latent space, reducing its complexity without degrading reconstruction quality. By finetuning pre-trained autoencoders with EQ-VAE, we enhance the performance of several state-of-the-art generative models, including DiT, SiT, REPA and MaskGIT, achieving a 7 speedup on DiT-XL/2 with only five epochs of SD-VAE fine-tuning. EQ-VAE is compatible with both continuous and discrete autoencoders, thus offering a versatile enhancement for a wide range of latent generative models. Project page and code: https://eq-vae.github.io/.

1. Introduction

Latent generative models depend on autoencoders whose latent spaces can be unnecessarily complex because they lack equivariance to scaling and rotation. EQ-VAE regularizes pretrained autoencoders to improve latent structure and generative modeling without degrading reconstruction quality.

  • 1. Introduction: Latent generative models compress images with continuous or discrete autoencoders before learning their latent distribution with a generative model.Continuous examples include SD-VAE, while discrete examples include VQ-GAN.
  • 1. Introduction: Stronger KL regularization smooths latent spaces but reduces information capacity, degrading fine-grained reconstruction details.This trade-off can limit the overall performance of latent generative models.
  • 1. Introduction: Scaling and rotation are not equivariant in established autoencoders, introducing unnecessary latent-manifold complexity that forces generative models to learn avoidable nonlinear relationships.The paper identifies this as a key limitation of current autoencoders.
  • 1. Introduction: EQ-VAE penalizes discrepancies between reconstructions of transformed latent representations and corresponding transformed inputs, requiring no architectural changes or training from scratch.A few epochs of fine-tuning pretrained autoencoders are sufficient to encourage equivariant latent spaces.
  • 1. Introduction: EQ-VAE preserves reconstruction quality while reducing latent complexity and facilitating generative-model learning.The paper illustrates smoother latent representations and accelerated training in Figure 1.
  • 1. Introduction: ×7 speedup on DiT-XL/2 and ×4 speedup on REPA follow five-epoch SD-VAE fine-tuning with EQ-VAE.The method also improves DiT, SiT, REPA, and MaskGIT performance and supports continuous and discrete autoencoders.

2. Related work

Prior work structures autoencoder latent spaces through variational regularization and studies equivariance in computer vision. EQ-VAE is positioned as promoting equivariance rather than invariance for latent generative modeling.

  • 2. Related work: Latent diffusion models use autoencoders to produce compressed latent spaces, avoiding the computational inefficiency of direct pixel-space diffusion training.The related work discusses SD-VAE as a widely used KL-regularized autoencoder.
  • 2. Related work: Variational autoencoders regularize latent codes to provide meaningful structure, while EQ-VAE promotes equivariance rather than invariance under spatial transformations.This distinguishes the paper’s objective from the regularization framing described in prior work.
  • 2. Related work: The paper reports an extensive evaluation of equivariance regularization’s impact on latent generative modeling.This statement summarizes the work’s evaluation emphasis.
  • 2. Related work: Computer-vision research has developed group-equivariant convolutional networks to incorporate symmetries beyond the approximate translation equivariance of standard convolutions.Examples include roto-translation equivariance in two dimensions.

3. Method

EQ-VAE regularizes autoencoders so spatial transformations applied in latent space produce reconstructions aligned with transformed inputs. It uses implicit equivariance regularization, preserves reconstruction quality, and applies to continuous and discrete autoencoders.

  • 3.1. Preliminary: Continuous Autoencoders for Latent Generative Modeling: Autoencoders encode images into continuous or discrete latent spaces, which generative models subsequently learn to sample.Continuous models use an encoder E and decoder D; discrete models additionally quantize encoder features through a learned codebook.
  • 3.2. Lack of Equivarance under Spatial Tansformations: EQ-VAE preserves reconstruction quality under latent transformations and reduces RFID relative to baselines in the reported comparisons.Figure 2 compares SD-VAE and EQ-VAE at scaling factor s = 0.5, while Figure 3 evaluates scaling factors s = 0.75, 0.50, 0.25 and rotation angles including θ = π.
  • 3.2. Lack of Equivarance under Spatial Tansformations: Established autoencoders are not equivariant to scaling and rotation, so directly transforming latent codes can substantially degrade reconstruction quality.The evaluation compares D(τ ◦E(x)) against D(E(τ ◦x)) for transformed inputs and transformed latent representations.
  • 3.3. EQ-VAE: Regularization via equivariance constraints: Explicitly minimizing latent equivariance error alone can collapse representations to a constant, whereas EQ-VAE uses reconstruction and adversarial supervision to avoid mode collapse.The identity transform is sampled with probability pα so the objective returns to the standard reconstruction objective in that case.
  • 3.3. EQ-VAE: Regularization via equivariance constraints: EQ-VAE modifies the autoencoder objective to align reconstructions of transformed latent representations with correspondingly transformed input images.The approach encourages equivariance without changing the model architecture and can be applied by fine-tuning pre-trained autoencoders.
  • 3.3. EQ-VAE: Regularization via equivariance constraints: EQ-VAE samples anisotropic scaling and rotations, with transformations composed as τ = S(sx, sy) · R(θ).The method samples 0.25 < sx, sy < 1 and rotation angles π/2, π, and 3π/2; scaling equivariance is reported as more beneficial for generation than rotation equivariance.

4. Experiments

Experiments show that EQ-VAE improves generative performance and convergence across continuous and discrete autoencoders while preserving reconstruction quality. Fine-tuning pretrained autoencoders yields gains across DiT, SiT, REPA, and MaskGIT.

  • Continuous Autoencoders: EQ-VAE maintains reconstruction fidelity while reducing equivariance error and improving generative performance across continuous autoencoders.The evaluation covers SD-VAE, SD-VAE-16, SDXL-VAE, and SD3-VAE.
  • Discrete Autoencoders: 6.8 to 5.9 GFID improvement demonstrates that EQ-VAE also benefits discrete VQ-GAN autoencoders.The discrete-autoencoder experiment evaluates MaskGIT.
  • DiT & SiT: 14.5 GFID is achieved by DiT-XL/2 at 400K iterations with EQ-VAE, versus 19.5 without regularization.At 1.5M iterations, EQ-VAE reaches 8.8 GFID, outperforming the SD-VAE-FT-EMA model even at 7M iterations.
  • REPA: REPA reaches 5.9 GFID in 1M instead of 4M iterations with EQ-VAE, making convergence ×4 faster.EQ-VAE is applied in the first latent-modeling stage and is complementary to REPA.
  • MaskGIT: MaskGIT reaches 6.80 GFID in 130 epochs instead of 300, converging twice as fast with EQ-VAE.By epoch 300, it reaches 5.91 GFID and surpasses the cited prior results.

5. Conclusion

The paper argues that autoencoder latent structure is crucial to latent generative-model performance and introduces EQ-VAE to improve equivariance while preserving reconstruction capability.

  • EQ-VAE fine-tunes pretrained autoencoders with a modified objective that reduces equivariance error.The method requires only a few epochs of fine-tuning.
  • Established autoencoders produce latent representations that are not equivariant under simple spatial transformations.The paper highlights scaling and rotation as examples.
  • EQ-VAE significantly boosts latent generative-model performance while maintaining reconstruction capability.
  • The paper identifies the geometry of latent distributions and their relationship to generative-model performance as future research directions.

A. Additional Ablations

The ablations compare explicit and implicit equivariance regularization and examine the regularization-strength hyperparameter. Explicit regularization lowers equivariance errors but triggers mode collapse, whereas implicit regularization improves generative performance and is robust to pα choices.

  • The objective is designed to reduce encoder equivariance error while avoiding mode collapse and preserving reconstruction performance.
  • Each objective is evaluated by fine-tuning SD-VAE and training DiT-B/2 on the resulting latent distribution.
  • Explicit regularization drastically lowers equivariance errors but triggers mode collapse, while implicit regularization significantly enhances generative performance.
  • The regularization-strength hyperparameter pα controls the strength of the regularization.
  • The method is relatively robust to pα choices, with experiments using pα = 0.7 and pα = 0.3.
  • Latent-representation complexity is quantified using intrinsic dimensionality computed with the TwoNN estimator from distances to each point’s two nearest neighbors.

C. Details on Evaluation Metrics

The evaluation uses established generative and reconstruction metrics on ImageNet-based data. Generative quality is assessed with distributional, spatial, classification-based, realism, and coverage measures, while reconstruction quality uses signal, structural, and perceptual similarity metrics.

  • Generative evaluation follows ADM’s setup and uses the same reference batches with the official implementation on NVIDIA A100 GPUs.
  • FID measures feature-distribution distance between two image datasets using Inception-v3 features under a multivariate Gaussian assumption.
  • sFID computes FID from intermediate spatial Inception-v3 features to capture generated-image spatial distributions.
  • IS measures KL-divergence between the original label distribution and softmax-normalized Inception-v3 logits.
  • Precision measures the fraction of realistic generated images, while Recall measures the fraction of training data covered by generated data.
  • Reconstruction is evaluated on ImageNet’s 50K-image validation set using PSNR, SSIM, and LPIPS.

C.3. Equivariance Error

Equivariance error quantifies how effectively EQ-VAE constrains autoencoder latent representations to equivary under scale and rotation transformations. The evaluation uses ImageNet validation samples and specified rotation and scale sets.

  • Equivariance error measures latent equivariance under scale and rotation transformations.
  • The evaluation uses N = 50K ImageNet validation samples and defines T as the set of considered transformations.
  • Rotation transformations are Tr = {π/2, π, 3π/2}, while scale transformations are Ts = {0.25, 0.50, 0.75}.

D. Detailed Benchmarks

The detailed benchmarks report additional performance, training-iteration, and reconstruction metrics across DiT, SiT, REPA, MaskGIT, and autoencoder evaluations. Separate tables cover generative-model results and comparisons of autoencoders with and without EQ-VAE regularization.

  • Table 9 details performance for DiT-XL/2 and SiT-XL/2 models.
  • Table 10 evaluates REPA models using SiT-XL/2, including SD-VAE-FT-EMA and EQ-VAE configurations.The listed results are reported without classifier-free guidance, with CFG = 1.0.
  • Table 11 presents detailed evaluations of MaskGIT models trained with VQ-GAN and EQ-VAE.The experiments use an open-source PyTorch reproduction, and results are reported with CFG = 3.0.
  • Table 12 compares reconstruction evaluation metrics for autoencoders with and without EQ-VAE regularization.

E. Specifications of Autoencoder Models

The appendix specifies the autoencoders used in the experiments and situates several latent generative models by their architectures and latent-space modeling approaches. Autoencoder specifications include training dataset, latent channels, and compression rate.

  • Table 13 records each autoencoder’s original training dataset, latent channels c, and compression rate f.
  • MaskGIT uses a bidirectional transformer decoder to iteratively predict masked visual tokens produced by VQ-GAN.
  • LDM models image distributions in compressed latent spaces produced by KL- or VQ-regularized autoencoders.
  • DiT uses a pure transformer backbone for diffusion-model training and incorporates AdaIN-zero modules.

G. Additional Qualitative Results

Additional qualitative results visualize reconstruction behavior under latent transformations and compare latent representations across VAE models. EQ-VAE preserves reconstruction quality under scaling and rotation transformations and produces smoother representations in the shown comparisons.

  • Figure 6 compares reconstructed images when transformations are applied to inputs versus directly to latent representations.The comparison covers scaling and rotation for SD-VAE and EQ-VAE.
  • EQ-VAE preserves reconstruction quality under latent transformations, whereas SD-VAE shows significant degradation.
  • EQ-VAE produces smoother latent representations for both SD-VAE and SDXL-VAE.
  • Figure 8 shows uncurated 256 × 256 samples from DiT-XL/2 with EQ-VAE using classifier-free guidance scale 4.0.
Loading 2502.09509v3…