PaperlayerView exact revision on arXiv

Source-linked AI summary

Distribution Matching Variational AutoEncoder

Sen Ye, Jianning Pei, Mengde Xu, Shuyang Gu, Chunyu Wang, Liwei Wang, Han Hu

arXiv:2512.07778v2cs.CV

TL;DR

Existing approaches do not explicitly shape the aggregate latent distribution, leaving unclear which latent structures best support modeling. DMVAE aligns that distribution to an arbitrary reference and enables systematic prior comparisons. The paper reports that SSL-derived distributions balance reconstruction fidelity and diffusion tractability, with 3.22 gFID on ImageNet 256x256 after 64 epochs.

  • Problem

    Existing methods do not guarantee that the aggregate latent distribution is suitable for the second-stage generative model, leaving the relative benefits of different latent structures unclear.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    DMVAE aligns an autoencoder’s aggregate posterior to an arbitrary reference distribution using a distribution-matching constraint.

    Source: §1 · p. 2

  • Results

    3.22 gFID on ImageNet 256x256 after 64 epochs is reported, with DINO-derived distributions identified as optimal for diffusion modeling.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DINO’s semantic clustering is described as easier for diffusion models to learn while retaining sufficient richness for high-fidelity reconstruction.

    Source: §1 · p. 2

    Matching initially distant distributions remains challenging, often requires careful tuning, and can make the reference act more like a regularizer than a complete match.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Two-stage visual generation depends on latents that are both reconstructive and learnable, but existing approaches do not explicitly regulate their aggregate distribution. DMVAE aligns that distribution to a chosen reference, enabling systematic study of latent priors; the introduction reports strong results for self-supervised features.

1. Introduction

  • Existing methods provide no guarantee that the latent distribution suits the second-stage model, risking latents that are either difficult to learn or too rigid for detailed reconstruction.VAEs constrain individual samples rather than aggregate geometry, while fixed pre-trained encoders cannot be optimized for reconstruction.Source: §1 · pp. 1–2
  • DMVAE replaces the standard KL-divergence objective with distribution matching that aligns the encoder’s aggregate posterior to an arbitrary reference distribution.The framework uses a diffusion model to estimate the reference distribution’s score and makes latent structure controllable.Source: §1 · p. 2
  • The study evaluates Gaussian, text-embedding, supervised-feature, and self-supervised-feature reference distributions to compare their effects on diffusion learnability.The tested examples include SigLIP, ResNet, and DINO features.Source: §1 · p. 2
  • 3.22 gFID on ImageNet 256x256 after 64 epochs, and 1.82 gFID after 400 epochs, are reported for DMVAE with the selected SSL prior.The authors identify DINO-derived distributions as optimal for diffusion modeling, describing semantic clustering alongside sufficient richness for high-fidelity reconstruction.Source: §1 · p. 2

2. Preliminary

This section introduces image tokenization and reviews VAE regularization, flow-based distribution estimation, and approaches to shaping latent priors. It highlights the limits of per-sample regularization and fixed encoders as motivation for distribution-level control.

2. Preliminary

  • A tokenizer maps images x to latent codes z with encoder Eθ and reconstructs x with decoder Gω.Source: §2 · p. 2

2.1. Variational Autoencoders (VAEs)

  • Standard β-VAEs use a fixed, typically Gaussian prior and per-sample KL regularization; large β often causes collapse, so latent-diffusion models use very small β.One reported setting is 1 × 10−6.Source: §2.1 · p. 2

2.2. Flow Matching

  • Flow-based distribution estimation noises clean latents at sampled times and trains a time-conditioned velocity network; scores express distribution-matching gradients.Source: §2.2 · p. 2

2.3. Inject distribution prior into VAEs

  • VAVAE aligns encoder codes with pretrained DINO features, encouraging the codes to inherit those features’ semantic properties.Source: §2.3 · pp. 2–3
  • Diffusion-prior updates through the encoder can be unstable and cause posterior collapse, while AAEs face GAN instability and discriminator-capacity limits.AAEs match aggregate q(z) to a reference prior; the paper uses diffusion/flow models for stable matching to complex references.Source: §2.3 · p. 3

3. Distribution Matching VAE

The section argues that per-sample regularization leaves aggregate latent geometry uncontrolled, then presents score-based distribution matching and compares alternative objectives.

3. Distribution Matching VAE

  • Matching aggregate posterior q(z) to arbitrary pr(z) addresses the global geometry left unconstrained by per-sample VAE or VAVAE losses.Source: §3 · p. 3, §3.1 · p. 3

3.1. The DMVAE Objective

  • A frozen teacher represents the reference score while an online student tracks changing q(z), enabling score-based matching without directly evaluating high-dimensional divergences.Source: §3.1 · p. 3
  • Training pretrains and freezes the reference score model, then jointly updates VAE and student through separate reconstruction, student-score, and encoder-matching updates.Source: §3.1 · pp. 3–4

3.2. Variations of Distribution Matching Objective

  • Real-score maximization can drop modes for lacking coverage pressure, while end-to-end fake-score maximization can collapse latents to a few points.Source: §3.2 · p. 5
  • Direct score-difference minimization is unstable; loss-difference minimization aligns distributions but incurs costly Jacobian computation in high-dimensional image latents.Source: §3.2 · p. 5
  • DMVAE’s fake-minus-real score field pulls latents toward reference high-density modes while resisting single-mode collapse, aligning global structure without Jacobian overhead.Source: §3.2 · p. 5

4. Experiments

On ImageNet 256×256, the experiments use DMVAE to study how reference-distribution structure relates to diffusion modeling, then evaluate prior selection, design choices, and benchmark performance.

4. Experiments

  • The experiments test whether latent geometry affects diffusion-model efficiency and quality, then compare reference distributions, ablations, and benchmark performance on ImageNet 256×256.The reference-prior comparison uses a heuristic matching weight: 10 for data-derived priors and 1 for synthetic priors.Source: §4 · p. 6, §4.1 · p. 6

4.1. Selection of Reference Distribution

  • DINO features provide the best overall balance of reconstruction quality and generative modeling among the tested reference distributions.The study compares data-derived priors, including DINO, ResNet, text features, and diffusion noise, with data-independent priors such as Gaussian and GMM distributions.Source: §4.1 · p. 6, Table 1 · p. 6

4.2. Ablation Studies

  • λDM=10 balances quality; larger score networks improve matching and generation, CFG 3.0 slightly improves generation, extreme CFG destabilizes training, and annealing slightly improves low-noise matching.A higher λDM can degrade reconstruction by enforcing an overly strict constraint.Source: §4.2 · p. 7

4.3. Comparison with other methods

  • At 64 epochs, DMVAE reaches gFID 3.22 versus AlignTok's gFID 3.71, with rapid convergence in Figure 6; at 800 epochs, it reaches gFID 1.64.The paper also reports that DMVAE surpasses most prior work at 400 epochs.Source: §4.3 · p. 7, Figure 6 · p. 7

5. Related Work

Related work covers visual tokenizers and distribution matching. DMVAE shapes global latent structure through distribution matching rather than focusing on tokenizer architecture, data curation, or point-wise feature alignment.

Visual Tokenizers

  • Visual tokenizers map images into compact discrete or continuous latent spaces, with reconstruction quality and generative capability as a central trade-off.Related approaches include changing tokenizer architecture, curating data, and aligning latent spaces with pretrained vision features.Source: §5 · pp. 7–8
  • Some approaches use frozen vision-foundation encoders, but these can compromise fine-grained reconstruction details.Other cited efforts target high-compression tokenizers, masked-autoencoder tokenizers, or quality-tuning with curated data.Source: §5 · p. 8

Distribution Matching

  • Prior distribution-matching methods include VAEs and adversarial autoencoders; DMVAE adapts distribution-level regularization to align tokenizer latents with a predefined prior.The paper contrasts this global alignment with methods focused on point-wise feature alignment.Source: §5 · p. 8

6. Conclusion and Limitation

DMVAE enables distribution-level control of autoencoder latents, and the paper reports strong ImageNet generation results using DINO features. The discussion reviews related approaches and notes that matching distant distributions remains challenging.

6. Conclusion and Limitation

  • gFID 3.22 on ImageNet after 64 training epochs is reported for semantically rich self-supervised reference features such as DINO.The conclusion describes this approach as performing well in reconstruction fidelity and generative tractability.Source: §6 · p. 8
  • Matching distributions that are initially far apart remains challenging, often requires careful tuning, and limits precise control of the tokenizer output distribution.The authors describe the reference distribution as acting more like a regularizer than a complete match.Source: §6 · p. 9

A. Extended Related Work

  • Prior tokenizer work targets reconstruction, feature alignment, equivariance, architecture, compression, or data curation, while leaving the aggregate latent distribution unregulated.The related-work discussion contrasts these approaches with directly shaping the entire latent distribution toward a desired prior.Source: A. Extended Related Work · p. 12
  • Earlier distribution-matching approaches include Gaussian-prior VAEs, expressive priors, adversarial autoencoders, and score matching used for diffusion distillation.The paper applies score matching to align visual-tokenizer representations with a predefined prior.Source: A. Extended Related Work · p. 12

B. Implementation Details

The implementation uses a staged pipeline, first pretraining an image tokenizer and then training a teacher to model its latent representations.

B. Implementation Details

B.1. Tokenizer Pretraining

  • Tokenizer pretraining uses frozen DINO-v2-large features projected to 32-dimensional latents, with a convolution-based decoder and reconstruction-plus-adversarial objective.The 256 × 256 input yields a 16 × 16 feature map; the tokenizer is trained for 8 epochs, and the adversarial weight is 0.5.Source: §B.1 · pp. 12–13

B.2. Reference Distribution (Teacher) Pretraining

  • The reference-distribution teacher is trained on tokenizer latents with LightningDiT-XL and QKNorm using flow-matching loss.Training runs for 400 epochs (500K steps).Source: §B.2 · p. 13

B.3. Distribution Matching VAE Training

DMVAE training jointly optimizes the VAE and a student model to align the evolving latent space with a teacher-defined reference distribution, followed by decoder fine-tuning.

B.3. Distribution Matching VAE Training

  • The VAE minimizes reconstruction and distribution-matching losses while a student models its evolving latent distribution through a score-based loss.The student is optimized every step, while the VAE updates once every five iterations; joint training runs for 350K steps.Source: §B.3 · p. 13

B.4. Decoder Fine-tuning

  • Decoder fine-tuning freezes the encoder and projector and optimizes only the decoder with the reconstruction objective for 250K steps.The fine-tuning uses batch size 256 and learning rate 1 × 10−4.Source: §B.4 · p. 13

B.5. Diffusion Model Training

After DMVAE training, a new high-fidelity generative model is trained from scratch in the learned latent space.

B.5. Diffusion Model Training

  • A new LightningDiT-XL model is trained from scratch on DMVAE latents using flow-matching loss for 800 epochs (1M steps).Training uses batch size 1024 and learning rate 2 × 10−4.Source: §B.5 · p. 13

C. Visualization of DMD Gradient Direction

DMD update directions trade off coarse guidance at high noise against precise but potentially uninformative guidance at low noise. Matching distant distributions also creates training instability, motivating multiple stabilization strategies.

C. Visualization of DMD Gradient Direction

  • Large noise scales can yield coarse, inaccurate DMD update directions, while small scales give precision but may fail to guide distributions that are far apart.Averaging across timesteps motivates dynamically adjusting the noise schedule during training.Source: C. Visualization of DMD Gradient… · pp. 13–14

D. Stabilizing the Distant Distribution Matching.

  • High noise can smooth away fine distribution differences and cause vanishing gradients, whereas low noise with disjoint supports can produce explosive, high-variance gradients and mode collapse.Source: D. Stabilizing the Distant… · p. 14
  • DMVAE matches an autoencoder’s entire aggregate posterior to an unpaired reference, making potentially large initial distribution discrepancies a source of instability.Source: D. Stabilizing the Distant… · p. 14
  • Training stabilization uses pretrained score-model weights and encoder initialization, alternating VAE and score-model training, and reducing latent spaces and priors to low dimension.The paper gives d = 32 as an example and says dimensionality reduction mitigates sparsity and computational challenges.Source: D. Stabilizing the Distant… · p. 14

E. Convergence Speed Visualization

A visual comparison reports that DMVAE generates higher-quality images than VAVAE at equivalent training steps. The comparison is conditional and uses no classifier-free guidance.

E. Convergence Speed Visualization

Related papers

Loading 2512.07778v2…