Distribution Matching Variational AutoEncoder

TL;DR

DMVAE introduces an explicit distribution-matching constraint to align the encoder's latent space with arbitrary reference distributions, improving image modeling and synthesis quality beyond traditional Gaussian priors.

Contribution

It generalizes VAEs by enabling explicit alignment with various reference distributions, facilitating systematic investigation of optimal latent distributions for image modeling.

Findings

SSL-derived latent distributions improve reconstruction and efficiency

Achieves gFID of 3.2 on ImageNet with 64 epochs

Explicit distribution matching enhances latent space quality

Abstract

Most visual generative models compress images into a latent space before applying diffusion or autoregressive modelling. Yet, existing approaches such as VAEs and foundation model aligned encoders implicitly constrain the latent space without explicitly shaping its distribution, making it unclear which types of distributions are optimal for modeling. We introduce \textbf{Distribution-Matching VAE} (\textbf{DMVAE}), which explicitly aligns the encoder's latent distribution with an arbitrary reference distribution via a distribution matching constraint. This generalizes beyond the Gaussian prior of conventional VAEs, enabling alignment with distributions derived from self-supervised features, diffusion noise, or other prior distributions. With DMVAE, we can systematically investigate which latent distributions are more conducive to modeling, and we find that SSL-derived distributions provide an excellent balance between reconstruction fidelity and modeling efficiency, reaching gFID equals 3.2 on ImageNet with only 64 training epochs. Our results suggest that choosing a suitable latent distribution structure (achieved via distribution-level alignment), rather than relying on fixed priors, is key to bridging the gap between easy-to-model latents and high-fidelity image synthesis. Code is avaliable at https://github.com/sen-ye/dmvae.