Distribution Matching Variational AutoEncoder

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 Distribution-Matching VAE (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 a gFID of 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.

Most visual generative models follow a two-stage recipe: a first model compresses each image into a compact latent representation, and a second model learns to create new images by working in this latent space. How the latent space is organized strongly affects how well the second stage can learn — but until now its structure has been a black box, with no way to design or compare different choices. We propose Distribution-Matching VAE (DMVAE), a method that lets researchers explicitly specify the desired shape of this latent space and forces the first model to match it. This turns the latent space into a controllable design choice rather than something left to chance. Using DMVAE, we run the first systematic comparison of different latent space structures. Borrowing structures from vision models that have learned visual concepts on their own clearly works best, reaching strong image quality (gFID 3.22) on the ImageNet benchmark in only 64 training passes — far fewer than prior systems require. Deliberately shaping the latent space, rather than leaving it to chance, is a simple but powerful lever for better image generators.