Research map / Image generation / Autoregressive and latent generation
EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EQ-VAE, a fine-tuning regularizer that makes autoencoder latents equivariant to scaling and rotation, yielding faster convergence and better FID for latent diffusion and masked generative models.
Problem
Latent generative models rely on autoencoders whose latent spaces are not equivariant to semantic-preserving spatial transformations (scaling, rotation), making the latent distribution complex and slowing or limiting generative training.
Approach
Fine-tune pre-trained continuous or discrete autoencoders with an implicit equivariance loss: decode transformed latents and match them to transformed input images, with the identity transform sampled with probability 0.5 to preserve reconstruction. This avoids the mode collapse of direct latent-equivariance losses.
Key results
- EQ-VAE fine-tuning of SD-VAE for 5 epochs reduces DiT-B/2 gFID from 43.5 to 34.1 while rFID stays near 0.82 vs 0.90.
- DiT-XL/2 with EQ-VAE reaches gFID 14.5 at 400K iterations vs 19.5 baseline, a claimed roughly 7x speedup.
- REPA (SiT-XL/2) with EQ-VAE reaches gFID 5.9 at 1M iterations vs 4M for standard REPA, about 4x faster; MaskGIT with EQ-VAE reaches gFID 5.91 at 300 epochs and 6.80 at 130 epochs versus 300 epochs baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add EQ-VAE-regularized latents | DiT-B/2 · ImageNet 256×256 | gFID: 43.5 -> 34.1 | Helped |
| Add EQ-VAE-regularized latents | DiT-XL/2 · ImageNet 256×256, 400K iterations | gFID: 19.5 -> 14.5 | Helped |
| Add EQ-VAE-regularized latents | SiT-XL/2+REPA · ImageNet 256×256 | gFID: 5.9 -> 5.9 | No clear effect |
| Add EQ-VAE regularization | SD-VAE latents used by DiT-B · ImageNet 256×256 | gFID: 43.8 -> 34.1 | Helped |
| Add EQ-VAE regularization | VQ-GAN latents used by MaskGIT · ImageNet 256×256 | gFID: 6.8 to 5.9 | Helped |
| Use rotation transformations only (ablation) | EQ-VAE with SD-VAE latents and DiT-B/2 · ImageNet 256×256 | gFID: 43.5 -> 41.2 | Helped |
| Use isotropic scaling only (ablation) | EQ-VAE with SD-VAE latents and DiT-B/2 · ImageNet 256×256 | gFID: 43.5 -> 35.8 | Helped |
| Combine rotations and isotropic scaling (ablation) | EQ-VAE with SD-VAE latents and DiT-B/2 · ImageNet 256×256 | gFID: 35.8 -> 34.1 | Helped |
| Replace isotropic with anisotropic scaling (ablation) | EQ-VAE with SD-VAE latents and DiT-B/2 · ImageNet 256×256 | rFID: 0.82 -> 0.92 | Hurt |
| Replace implicit regularization with explicit equivariance loss (ablation) | SD-VAE · DiT-B/2 on SD-VAE latents | gFID: 43.5 -> 141.3 | Hurt |
| Add stop-gradient to explicit equivariance loss (ablation) | SD-VAE · DiT-B/2 on SD-VAE latents | gFID: 43.5 -> 134.7 | Hurt |
| Use 5 extra fine-tuning epochs with the standard objective instead of EQ-VAE (ablation) | SD-VAE · DiT-B/2 on SD-VAE latents | gFID: 43.5 -> 43.5 | No clear effect |
Surprising findings
- Explicit equivariance loss alone collapses the latent space, raising gFID to 141.3 and rFID to 117.9, so the implicit reconstruction-based objective is necessary.
Limitations
- Anisotropic scaling improves generation but degrades reconstruction quality, creating a trade-off.
- Intrinsic dimension estimates via TwoNN are noisy and tend to overestimate ID; the analysis is relative only.
- Main gains are measured on ImageNet 256x256 with class-conditional generators; generality to text-to-image is not tested.
Experimental setup
Datasets: OpenImages ImageNet LAION-Aesthetics
Benchmarks: ImageNet 256x256
Metrics: gFID rFID sFID IS precision recall PSNR SSIM LPIPS equivariance error
Baselines: SD-VAE SDXL-VAE SD3-VAE VQ-GAN DiT SiT REPA MaskGIT
Builds on: latent diffusion models VQ-GAN REPA DiT
Closest papers by meaning
- MAE: Mutual Posterior-Divergence Regularization for Variational AutoEncoders (ICLR 2019 · similarity 0.90)
- From Variational to Deterministic Autoencoders (ICLR 2020 · similarity 0.90)
- Perceptual Generative Autoencoders (ICML 2020 · similarity 0.89)
- Generating Diverse High-Fidelity Images with VQ-VAE-2 (NeurIPS 2019 · similarity 0.89)
- CV-VAE: A Compatible Video VAE for Latent Generative Video Models (NeurIPS 2024 · similarity 0.88)
- SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer (CVPR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.