Distribution Matching VAE (DMVAE): Exploring the Best Latent Space for Visual Generation
The Tokenizer’s Dilemma
Mainstream generative models follow a “compress first, generate later” paradigm. An ideal latent space must satisfy two conflicting criteria:
- Ease of modeling: The latent distribution \(q(z)\) should be simple enough for a downstream diffusion or autoregressive model to learn.
- Reconstruction fidelity: The latent code \(z\) must retain enough information to reconstruct the original image perfectly.
Existing methods often struggle to strike a good balance:
- Standard VAEs: These use KL divergence to pull the posterior for each individual sample toward a standard Gaussian \(\mathcal{N}(0, I)\). In practice, a small KL weight is typically used to preserve reconstruction quality, leaving a latent-space manifold that is difficult to model.
- RAE: Directly using a frozen model such as DINO as the encoder makes the latent distribution easy to model, but prevents reconstruction of the original image, ruling out downstream tasks such as image editing.
- VFM-based Methods (e.g., VAVAE, AlignTok): These use pointwise losses, such as MSE, to bring latent features closer to those of a pretrained representation encoder. However, pointwise alignment does not guarantee that the aggregated posterior \(q(z) = \int q(z\mid x)p(x)dx\) forms a manifold that is easy to model.
Illustration of VAE, RAE, VFM-based methods, and DMVAE
Our Method: Distribution Matching VAE
We propose DMVAE, whose central idea is to constrain \(q(z)\) directly at the distribution level, aligning it with an arbitrary reference distribution \(p_r(z)\). Because both \(q(z)\) and \(p_r(z)\) are high-dimensional, implicit distributions, directly computing their KL divergence is infeasible. Inspired by Distribution Matching Distillation (DMD), we characterize the two distributions using their score functions, using this “general-purpose distribution estimator” to address the problem of measuring the distance between distributions.
Core Training Procedure
Our training involves three models:
- Teacher (Real Score Model): We pretrain a diffusion model on the reference distribution \(p_r(z)\), such as the distribution of DINOv2 features, to learn its score function \(\nabla_z \log p_r(z)\), then freeze it.
- Student (Fake Score Model): A diffusion model trained jointly with the VAE to track the score function of the encoder’s output distribution \(q(z)\) as it evolves.
- VAE Encoder: The encoder learns through both reconstruction and distribution matching objectives.
DMD Gradient: The distribution matching gradient is
\[\nabla_z \mathcal{L}_{DM} \propto s_{fake}(z) - s_{real}(z)\]The DMD gradient pushes latent codes toward high-density regions of the reference distribution, guided by \(s_{real}\), while using \(s_{fake}\) to prevent the entire distribution from concentrating at a mode of the reference distribution. Compared with GAN-based or pointwise constraints, this approach is more stable and preserves the global structure of the distribution’s manifold.
Overview of the DMVAE training procedure
Choosing the Distribution Matching Objective
There are several possible choices for the distribution matching loss function:
1. Maximizing the Real Score Function
FlowPrior adopts this idea, seeking to maximize the likelihood of the encoder’s outputs under the teacher model. In our implementation, we use v-prediction and can place stop-gradient operators at different points in the loss. We compare three variants, shown below:
Distribution matching objective: maximizing the real score function
As the figure shows, maximum likelihood drives samples toward high-probability regions, causing mode dropping and failing to preserve the global structure of the reference distribution.
2. Maximizing the Fake Score Function
This approach trains the VAE and diffusion model simultaneously, similar to the unsuccessful end-to-end training variant discussed in REPA-E. It also has three variants, shown below. Training with this approach leads to severe mode collapse, with the latent space collapsing around one or a few points.
Distribution matching objective: maximizing the fake score function
3. Minimizing the L2 Distance Between the Student and Teacher Score Functions
Here too, there are three implementation variants, as shown in the figure. This approach also fails to achieve distribution matching.
Distribution matching objective: minimizing the L2 distance between the student and teacher score functions
4. Minimizing the Difference Between the Teacher and Student Flow Matching Losses
As shown on the right in the figure below, this approach can effectively match distributions. However, computing the loss requires the Jacobians of both score models, which is computationally expensive for high-dimensional image generation.
5. The Method Used in This Work
As shown in the center of the figure below, our method successfully matches the initial distribution to the target distribution. It also avoids computing Jacobians, overcoming issues encountered by the four approaches above, including mode dropping and poor computational efficiency.
Distribution matching objectives: DMD and minimizing the difference between teacher and student flow matching losses
What Kind of Distribution Is Best for Generation?
DMVAE reframes tokenizer design as a problem of “reference distribution selection.” This allows us to systematically study the properties of different distributions and compare which can achieve the perfect balance between reconstruction and modeling.
We compare two broad categories of distributions:
- Data-Derived: Self-supervised features (DINOv2), supervised features (ResNet), text features (SigLIP-text), and noise at intermediate diffusion steps.
- Data-Independent: Standard Gaussian distributions, Gaussian mixture models (GMMs), and randomly sampled features.
Our experimental finding: feature distributions from self-supervised learning (SSL), such as DINO, are the best choice.
Through t-SNE visualizations, we find that:
- Gaussian distributions are too homogeneous and lack semantic structure, requiring downstream generative models to spend substantial compute reconstructing semantic clusters.
- DINO distributions naturally exhibit strong semantic clustering, with moderate within-class variance and clear boundaries between classes.
- DMVAE successfully “distills” this favorable manifold structure into the VAE’s latent space, reducing confusion for the generative model and making the distribution easier to model.
t-SNE visualizations of different distributions
We find that using a small CFG value in the DMD objective effectively improves generation quality. We visualize distributions under different CFG settings using t-SNE and find that CFG strengthens semantic clustering in the distribution.
Distributions under different CFG settings
Experimental Results
We evaluate on ImageNet at \(256 \times 256\) resolution. To isolate the effects of our method, we do not use CFG at evaluation time.
- Generation Quality: After only 64 epochs of training, DMVAE achieves 3.22 gFID. Training for 400 epochs further reduces gFID to 1.82.
- Convergence: Compared with methods such as AlignTok and VAVAE, DMVAE converges faster and achieves better results.
Reconstruction and generation results on ImageNet256 without guidance
Comparison of DMVAE’s convergence speed with baselines
Conclusion
DMVAE breaks with the paradigm that VAEs must use a Gaussian prior. Through distribution matching, we can align the latent-space distribution with an arbitrary reference distribution.
Our study shows that aligning a VAE’s latent distribution with the distribution of a visual representation encoder at the distribution level offers a direct path to high-fidelity, efficient visual generation.