This is a Method paper. It introduces Latent Diffusion Models (LDMs), which train denoising diffusion models in the latent space of pretrained autoencoders rather than directly in pixel space. The key insight is that separating perceptual compression from generative learning enables high-resolution image synthesis at a fraction of the computational cost of pixel-based diffusion. The paper also introduces a cross-attention conditioning mechanism for flexible multi-modal generation.

Computational Cost of Pixel-Space Diffusion

Training diffusion models directly in pixel space is computationally expensive (150 to 1000 V100 GPU-days for leading models at the time) because the model must process high-dimensional RGB data at every denoising step. Much of this compute is spent modeling imperceptible high-frequency details. The authors observe that learning can be split into two stages: a perceptual compression stage that removes high-frequency detail, and a semantic compression stage where the generative model learns the conceptual composition. Prior two-stage approaches (VQGAN, DALL-E) relied on aggressive compression and autoregressive modeling in discrete latent spaces, trading off reconstruction quality for tractability.

Core Innovation: Diffusion in Latent Space

LDMs decompose image synthesis into two phases:

Phase 1: Perceptual Compression. A pretrained autoencoder (encoder $\mathcal{E}$, decoder $\mathcal{D}$) maps images $x \in \mathbb{R}^{H \times W \times 3}$ to a lower-dimensional latent representation $z = \mathcal{E}(x) \in \mathbb{R}^{h \times w \times c}$ with spatial downsampling factor $f = H/h$. The autoencoder is trained with a perceptual loss (matching deep features from a pretrained VGG network) and a patch-based adversarial objective, with either KL or VQ regularization on the latent space.

Phase 2: Latent Diffusion. A standard denoising diffusion model operates in this latent space. The training objective becomes:

$$L_{\text{LDM}} := \mathbb{E}_{\mathcal{E}(x), \epsilon \sim \mathcal{N}(0,1), t} \left[ \left| \epsilon - \epsilon_\theta(z_t, t) \right|_2^2 \right]$$

where $z_t$ is the noised latent at timestep $t$, and $\epsilon_\theta$ is a time-conditional UNet.

Cross-Attention Conditioning. To enable conditioning on text, semantic maps, or other modalities, the authors introduce cross-attention layers into the UNet. A domain-specific encoder $\tau_\theta$ maps conditioning input $y$ to an intermediate representation $\tau_\theta(y) \in \mathbb{R}^{M \times d_\tau}$, which interacts with the UNet features via:

$$Q = W_Q^{(i)} \cdot \varphi_i(z_t), \quad K = W_K^{(i)} \cdot \tau_\theta(y), \quad V = W_V^{(i)} \cdot \tau_\theta(y)$$

The conditional objective then becomes:

$$L_{\text{LDM}} := \mathbb{E}_{\mathcal{E}(x), y, \epsilon \sim \mathcal{N}(0,1), t} \left[ \left| \epsilon - \epsilon_\theta(z_t, t, \tau_\theta(y)) \right|_2^2 \right]$$

Both $\tau_\theta$ and $\epsilon_\theta$ are optimized jointly.

Experimental Setup and Results

The authors evaluate across multiple tasks and datasets:

Perceptual compression tradeoffs. Downsampling factors $f \in {1, 2, 4, 8, 16, 32}$ are compared on ImageNet class-conditional generation. LDM-1 (pixel-based) trains slowly; LDM-32 loses too much information. LDM-4 and LDM-8 achieve the best balance, with LDM-8 outperforming pixel-based diffusion by 38 FID points after 2M training steps on a single A100.

Unconditional image synthesis on CelebA-HQ 256, FFHQ 256, LSUN Churches/Bedrooms 256: LDM-4 achieves FID 5.11 on CelebA-HQ (state of the art at the time), outperforming LSGM, GANs, and other likelihood-based models. On LSUN-Bedrooms, LDM-4 achieves FID 2.95, close to ADM (1.90) with half the parameters and roughly 4x less training compute (see Appendix E.3.5).

Text-to-image synthesis on MS-COCO: A 1.45B parameter LDM-KL-8 model trained on LAION-400M achieves FID 12.63 with classifier-free guidance (a technique that amplifies the conditioning signal at the cost of diversity, by interpolating between conditional and unconditional predictions) at scale s=1.5, on par with GLIDE (FID 12.24, 6B params) and Make-A-Scene (FID 11.84, 4B params) with substantially fewer parameters.

Class-conditional ImageNet 256: LDM-4-G achieves FID 3.60, IS 247.67, outperforming ADM-G (FID 4.59) with fewer parameters and less compute.

Super-resolution: LDM-4 (big) achieves FID 2.4 on ImageNet 64-to-256 upscaling (validation split), outperforming SR3 in FID.

Inpainting on Places: LDM-4 (big, w/ ft) achieves FID 1.50, setting a new state of the art on image inpainting.

Key Findings and Limitations

• LDM-4 and LDM-8 offer the best tradeoff between perceptual compression and generation quality.
• The autoencoder only needs to be trained once and can be reused across different diffusion models and tasks.
• Cross-attention conditioning generalizes to text, semantic layouts, and bounding boxes without architecture changes.
• Convolutional sampling enables generation at resolutions higher than the training resolution (up to 1024x1024).
• Sequential sampling remains slower than GANs. The autoencoder reconstruction can become a bottleneck for tasks requiring pixel-level precision.

Reproducibility Details

Data

Algorithms

• Autoencoder regularization: KL-reg (KL penalty toward standard normal, weighted by ~$10^{-6}$) or VQ-reg (vector quantization layer on the latent space with a learned codebook)
• Diffusion: Standard DDPM denoising with reweighted objective
• Sampling: DDIM sampler with configurable steps (100 to 500 depending on task)
• Guidance: Classifier-free diffusion guidance with scale $s$ (1.5 for class-conditional and text-to-image quantitative evaluation; 10.0 for qualitative text-to-image samples)

Models

• Autoencoder: Based on VQGAN architecture with perceptual + adversarial loss
• UNet backbone: Time-conditional with cross-attention layers at multiple resolutions
• Text encoder: BERT-tokenizer with transformer $\tau_\theta$ for LAION text-to-image model
• LDM-4-G: 400M parameters, $f=4$ downsampling
• LDM-KL-8 (text): 1.45B parameters, $f=8$ downsampling, KL-regularized

Evaluation

Hardware

• Perceptual compression tradeoff experiments: single NVIDIA A100
• Inpainting model trained on eight V100
• Training at least 2.7x faster than pixel-based diffusion at equal parameters

Artifacts

Paper Information

Citation: Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2022). High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022. https://arxiv.org/abs/2112.10752

Publication: CVPR 2022

@inproceedings{rombach2022highresolution,
  title     = {High-Resolution Image Synthesis with Latent Diffusion Models},
  author    = {Rombach, Robin and Blattmann, Andreas and Lorenz, Dominik and Esser, Patrick and Ommer, Bj{\"o}rn},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages     = {10684--10695},
  year      = {2022}
}

Additional Resources:

• GitHub Repository
• Score-Based Generative Modeling with SDEs

This is a Method paper. It extends denoising diffusion probabilistic models (DDPMs) from continuous to discrete state-spaces by introducing structured Markov transition matrices for the corruption process. The paper unifies several corruption strategies, draws a formal connection between absorbing-state diffusion and masked language models, and demonstrates competitive results on both image and text generation.

Diffusion Beyond Continuous Spaces

Standard DDPMs operate in continuous state-spaces (e.g., pixel values treated as real numbers) and use Gaussian noise for corruption. Many important data types are inherently discrete: text (tokens from a vocabulary), quantized images (discrete pixel values), molecular structures, and segmentation maps. Prior work by Hoogeboom et al. extended binary diffusion to multinomial diffusion with uniform transition probabilities, but this limits the structure of the corruption process. D3PMs generalize this by allowing arbitrary transition matrices that encode domain-specific inductive biases.

Core Innovation: Structured Transition Matrices

D3PMs define a forward corruption process over discrete variables $\mathbf{x} \in {1, \ldots, K}^D$ using transition matrices $\mathbf{Q}_t \in \mathbb{R}^{K \times K}$:

$$q(\mathbf{x}_t | \mathbf{x}_{t-1}) = \text{Cat}(\mathbf{x}_t; \mathbf{p} = \mathbf{x}_{t-1} \mathbf{Q}_t)$$

where $\mathbf{x}_{t-1}$ is a one-hot row vector. The cumulative transition after $t$ steps is $\overline{\mathbf{Q}}_t = \mathbf{Q}_1 \mathbf{Q}_2 \cdots \mathbf{Q}_t$, giving:

$$q(\mathbf{x}_t | \mathbf{x}_0) = \text{Cat}(\mathbf{x}_t; \mathbf{p} = \mathbf{x}_0 \overline{\mathbf{Q}}_t)$$

The paper explores several transition matrix designs:

Uniform diffusion: $[\mathbf{Q}_t]_{ij} = (1 - \beta_t) \mathbf{1}_{i=j} + \beta_t / K$. Transitions with equal probability to any state. Stationary distribution is uniform.

Absorbing state: In absorbing-state diffusion, each non-mask token transitions to the mask state with probability $\beta_t$ per step, while tokens already at the mask state remain there:

$[\mathbf{Q}_t]_{ij} = (1-\beta_t)\mathbf{1}_{i=j\neq m} + \beta_t \mathbf{1}_{j=m} + \mathbf{1}_{i=j=m}$. Each token transitions to a designated absorbing state $m$ (e.g., [MASK] for text, gray pixel for images) with probability $\beta_t$. This establishes a direct connection to masked language models like BERT.

Discretized Gaussian: Transition probabilities decay as a function of the distance $|i-j|$ between states, mimicking Gaussian diffusion on ordinal data like pixel values.

Embedding-based nearest neighbor: For text, transitions are weighted by proximity in a pretrained word embedding space, so corruption preferentially swaps words with semantically similar ones.

Training objective. The reverse process $p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t)$ is parameterized by predicting $\tilde{p}_\theta(\tilde{\mathbf{x}}_0 | \mathbf{x}_t)$ and computing the posterior:

$$p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t) \propto \sum_{\tilde{\mathbf{x}}_0} q(\mathbf{x}_{t-1} | \mathbf{x}_t, \tilde{\mathbf{x}}_0) , \tilde{p}_\theta(\tilde{\mathbf{x}}_0 | \mathbf{x}_t)$$

The loss combines the variational lower bound (VLB) with an auxiliary cross-entropy loss $L_\lambda$:

$$L = L_{\text{VLB}} + \lambda , L_{\text{CE}}$$

where $L_{\text{CE}}$ is a reweighted cross-entropy loss on the $\mathbf{x}_0$ prediction that stabilizes training and improves sample quality. The VLB decomposes into per-timestep KL divergences between the true and predicted reverse transitions.

Experiments and Results

Image generation (CIFAR-10):

The best discrete D3PM variant is D3PM Gauss + logistic, which achieves FID 7.34 and NLL 3.44 bpd using the combined $L_\lambda$ loss with a truncated logistic parameterization. The truncated logistic parameterization replaces the standard softmax output with a discretized logistic distribution over pixel values, assigning probability mass to each discrete bin based on a continuous logistic CDF. This provides a smoother output distribution that better captures the ordinal structure of pixel intensities. This variant exceeds the continuous DDPM in log-likelihood (3.44 vs. 3.75 bpd) while approaching its sample quality (FID 7.34 vs. 3.17).

Text generation (text8, character-level, 1000 steps):

Among the D3PM variants and baselines evaluated, D3PM absorbing achieves the best bpc on text8 apart from Discrete Flow (Tran et al., 2019). On LM1B (sentencepiece vocabulary of 8192 tokens), D3PM absorbing achieves a perplexity of 76.9 at 1000 steps, compared to 137.9 for D3PM uniform and 43.6 for a comparable autoregressive transformer, demonstrating that discrete diffusion scales to large vocabularies.

Ablation findings:

• The auxiliary cross-entropy loss $L_\lambda$ is critical: for D3PM Gauss, it improves FID from 15.30 ($L_{\text{VLB}}$) to 8.34 ($L_\lambda$, $\lambda{=}0.001$). Adding the truncated logistic parameterization further improves FID to 7.34.
• Discretized Gaussian transitions outperform both uniform and absorbing-state transitions on CIFAR-10 across all metrics.
• For text, the absorbing-state (mask) model outperforms uniform and nearest-neighbor models. Nearest-neighbor diffusion provides only marginal improvement over uniform, a surprising negative result.
• The $\mathbf{x}_0$-parameterization ensures the learned reverse distribution has the correct sparsity pattern dictated by the transition matrix $\mathbf{Q}_t$.

Findings and Limitations

• The choice of transition matrix is an important design decision that encodes domain-specific inductive biases. Discretized Gaussian transitions work best for ordinal image data; absorbing-state transitions work best for text.
• D3PMs formally unify diffusion models and masked language models: absorbing-state diffusion with a [MASK] token is equivalent to a reweighted BERT-style training objective.
• The combined VLB + auxiliary loss ($L_\lambda$) achieves better density estimation (3.44 bpd) than continuous DDPMs (3.75 bpd) while producing competitive samples.
• Sample quality (best FID 7.34 for D3PM Gauss + logistic) still lags behind continuous-space DDPMs (FID 3.17) on CIFAR-10, though the gap narrows with structured transitions and the auxiliary loss.
• Scaling to very large numbers of categories $K$ requires special techniques (low-rank corruption or matrix exponentials) to manage the $O(K^2 T)$ memory cost of storing transition matrices.

Reproducibility Details

Data

Algorithms

• Noise schedules: Linear schedule for D3PM Gauss, cosine schedule for D3PM uniform, and a novel mutual information schedule for absorbing and nearest-neighbor models
• Reverse parameterization: $\mathbf{x}_0$-parameterization with posterior computation via Bayes’ rule
• Loss: $L_{\text{VLB}} + \lambda L_{\text{CE}}$ with $\lambda = 0.001$ for images and $\lambda = 0.01$ for text absorbing models
• Scaling: Low-rank corruption (absorbing, uniform) scales as $O(r^2 T)$; matrix exponentials for nearest-neighbor transitions

Models

• Image models: Modified U-Net architecture from Ho et al. (2020) adapted for categorical output via softmax over $K$ classes
• Text models: 12-layer T5-style transformer encoder with 70M parameters (12 heads, MLP dim 3072, QKV dim 768)
• Timesteps: $T = 1000$ for both images and text, though text models can be evaluated with fewer steps (e.g., 256 or 20)

Evaluation

Hardware

• All models trained for 1M steps with batch size 512 on TPUv2 or TPUv3
• Text models: 12-layer transformer encoder (T5 architecture), 70M parameters
• Image models: Modified U-Net architecture from Ho et al. (2020)

Artifacts

Paper Information

Citation: Austin, J., Johnson, D. D., Ho, J., Tarlow, D., & van den Berg, R. (2021). Structured Denoising Diffusion Models in Discrete State-Spaces. NeurIPS 2021. https://arxiv.org/abs/2107.03006

Publication: NeurIPS 2021

@inproceedings{austin2021structured,
  title     = {Structured Denoising Diffusion Models in Discrete State-Spaces},
  author    = {Austin, Jacob and Johnson, Daniel D. and Ho, Jonathan and Tarlow, Daniel and van den Berg, Rianne},
  booktitle = {Advances in Neural Information Processing Systems},
  volume    = {34},
  year      = {2021}
}

Additional Resources:

• Score-Based Generative Modeling with SDEs

This is a Method paper. It proposes consistency models, a new class of generative models designed for fast one-step (or few-step) generation. The models can be trained either by distilling pretrained diffusion models (consistency distillation) or as standalone generative models from scratch (consistency training). The paper provides theoretical analysis of both training modes and achieves FID 3.55 on CIFAR-10 for single-step non-adversarial generation (state of the art at the time of publication).

The Slow Sampling Problem in Diffusion

Diffusion models produce high-quality samples but require iterating through many denoising steps (often tens to hundreds), making generation slow compared to GANs or VAEs. Previous approaches to speed up sampling include faster ODE/SDE solvers (DDIM, DPM-Solver) and progressive distillation. These either still require multiple steps or depend on a complex multi-stage distillation pipeline. The goal is a model that can generate high-quality samples in a single forward pass while optionally allowing more steps for better quality.

Core Innovation: The Self-Consistency Property

The key idea builds on the Probability Flow (PF) ODE from the score-based SDE framework. The PF ODE describes a deterministic trajectory that converts noise into data, governed by the learned score function. For the VE-SDE parameterization used by EDM (Karras et al., 2022), this takes the form:

$$\frac{d\mathbf{x}_t}{dt} = -t , s_\phi(\mathbf{x}_t, t)$$

where $s_\phi$ is a pretrained score model, a consistency function $f(\mathbf{x}_t, t)$ maps any point on an ODE trajectory to the trajectory’s origin $\mathbf{x}_\epsilon$. The defining property is self-consistency:

$$f(\mathbf{x}_t, t) = f(\mathbf{x}_{t’}, t’) \quad \text{for all } t, t’ \in [\epsilon, T]$$

for any points $\mathbf{x}_t$ and $\mathbf{x}_{t’}$ on the same PF ODE trajectory.

Parameterization. The model enforces the boundary condition $f(\mathbf{x}_\epsilon, \epsilon) = \mathbf{x}_\epsilon$ using skip connections:

$$f_\theta(\mathbf{x}, t) = c_{\text{skip}}(t) , \mathbf{x} + c_{\text{out}}(t) , F_\theta(\mathbf{x}, t)$$

where $c_{\text{skip}}(\epsilon) = 1$ and $c_{\text{out}}(\epsilon) = 0$, ensuring the boundary condition is satisfied by construction.

Consistency Distillation (CD). Given a pretrained diffusion model, CD trains a consistency model by enforcing self-consistency between adjacent timesteps:

$$\mathcal{L}_{\text{CD}}^N(\theta, \theta^-; \phi) = \mathbb{E}\left[\lambda(t_n) , d!\left(f_\theta(\mathbf{x}_{t_{n+1}}, t_{n+1}), , f_{\theta^-}(\hat{\mathbf{x}}_{t_n}^\phi, t_n)\right)\right]$$

where $\hat{\mathbf{x}}_{t_n}^\phi$ is obtained by running one step of the ODE solver using the pretrained score model, $\theta^-$ is an exponential moving average (EMA) of $\theta$, and $d(\cdot, \cdot)$ is a distance metric. The use of a target network $\theta^-$ (updated via EMA) parallels techniques from deep Q-learning and momentum contrastive learning.

Consistency Training (CT). CT eliminates the need for a pretrained diffusion model. It replaces the ODE solver step with a score estimate derived from the denoising score matching identity:

$$\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t) = \mathbb{E}\left[\frac{\mathbf{x} - \mathbf{x}_t}{t^2} ,\middle|, \mathbf{x}_t\right]$$

Because this identity lets us estimate the score from noisy data alone (without a pretrained model), we can compute the ODE update directly from training samples. This allows training directly on data pairs $(\mathbf{x}, \mathbf{x} + t\mathbf{z})$ where $\mathbf{z} \sim \mathcal{N}(0, I)$.

Theoretical guarantee. If CD achieves zero loss, the consistency model error is bounded by $O((\Delta t)^p)$ where $\Delta t$ is the maximum timestep gap and $p$ is the order of the ODE solver.

Experiments and Benchmarks

Datasets: CIFAR-10 (32x32), ImageNet 64x64, LSUN Bedroom 256x256, LSUN Cat 256x256.

Architecture: All models use the NCSN++/EDM architecture. CD distills from pretrained EDM models.

Key results for consistency distillation (CD):

CD outperforms progressive distillation (PD) across all datasets and sampling steps, with the exception of single-step generation on Bedroom 256x256 where CD with $\ell_2$ slightly underperforms PD with $\ell_2$.

Key results for consistency training (CT):

CT outperforms existing single-step non-adversarial models (VAEs, normalizing flows), e.g., improving over DC-VAE’s FID of 17.90 on CIFAR-10. Samples from CT share structural similarity with EDM samples from the same initial noise, suggesting CT does not suffer from mode collapse.

Zero-shot editing: Consistency models support colorization, super-resolution, inpainting, stroke-guided generation, interpolation, and denoising at test time without task-specific training, by modifying the multi-step sampling algorithm.

Findings and Limitations

• Consistency distillation achieves state-of-the-art FID for one-step generation (3.55 on CIFAR-10, 6.20 on ImageNet 64x64).
• Multi-step sampling provides a smooth quality-compute tradeoff: more steps yield better FID.
• CT produces competitive results without any pretrained diffusion model, making consistency models a standalone generative model family.
• The LPIPS distance metric $d(\cdot, \cdot)$ generally outperforms $\ell_1$ and $\ell_2$ for training consistency models.
• At higher resolutions (LSUN 256x256), the gap between CD/CT and full EDM sampling widens.
• CT currently underperforms CD, suggesting room for improvement in the standalone training paradigm.

Reproducibility Details

Data

Algorithms

• ODE solver for CD: Euler and Heun (2nd order) solvers on the empirical PF ODE
• EMA for target network: Decay rate $\mu$ scheduled as a function of training step
• Schedule functions: $N$ (number of discretization steps) and $\mu$ (EMA rate) increase over training following specific schedules (see Appendix C of the paper)
• Distance metric: LPIPS performs best; $\ell_2$ and $\ell_1$ also evaluated

Models

• Architecture: NCSN++/EDM architecture from Karras et al. (2022)
• CD teacher: Pretrained EDM models
• Parameterization: Skip-connection formulation with $c_{\text{skip}}(t)$ and $c_{\text{out}}(t)$ from EDM

Evaluation

Hardware

• Training details follow EDM conventions
• CD and CT use the same batch sizes and learning rate schedules as EDM training

Artifacts

Paper Information

Citation: Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. (2023). Consistency Models. ICML 2023. https://arxiv.org/abs/2303.01469

Publication: ICML 2023

@inproceedings{song2023consistency,
  title     = {Consistency Models},
  author    = {Song, Yang and Dhariwal, Prafulla and Chen, Mark and Sutskever, Ilya},
  booktitle = {International Conference on Machine Learning},
  volume    = {202},
  year      = {2023},
  url       = {https://arxiv.org/abs/2303.01469}
}

Additional Resources:

• GitHub Repository
• Score-Based Generative Modeling with SDEs

This is primarily a Method paper. It proposes a unified framework that generalizes previous discrete score-based models (SMLD and DDPM) into continuous-time Stochastic Differential Equations (SDEs). The paper introduces algorithms for sampling (Predictor-Corrector) and likelihood computation (Probability Flow ODE), validated by setting new records on CIFAR-10 (FID 2.20, IS 9.89 at the time of publication). It also contains elements of Systematization by showing how existing methods are special cases of this broader framework.

What is the motivation?

Prior successful generative models, specifically Score Matching with Langevin Dynamics (SMLD) and Denoising Diffusion Probabilistic Models (DDPM), operate by sequentially corrupting data with slowly increasing noise and learning to reverse the process. Both methods treat the noise scales as a finite set of discrete steps. The authors aim to generalize this to a continuum of noise scales by modeling the diffusion process as a Stochastic Differential Equation (SDE). This continuous formulation enables:

• Flexible sampling: Use of general-purpose SDE solvers.
• Exact likelihood computation: Via connection to Neural ODEs.
• Controllable generation: Solving inverse problems (inpainting, colorization) without retraining.

What is the novelty here?

The core novelty is the SDE framework for score-based generative modeling:

• Continuous Generalization: Proving that SMLD and DDPM noise perturbations correspond to discretizations of Variance Exploding (VE) SDEs and Variance Preserving (VP) SDEs, respectively.
• Reverse-Time SDE: Leveraging Anderson’s result (Anderson, 1982: a result on time-reversal of diffusion processes showing that the reverse is also a diffusion, with the forward drift reversed and a correction term involving the score of the marginal density) that the reverse of a diffusion process is also a diffusion process, governed by the score (gradient of log density).
• Predictor-Corrector (PC) Samplers: A hybrid sampling strategy where a numerical SDE solver (Predictor) estimates the next step, and a score-based MCMC approach (Corrector) corrects the marginal distribution.
• Probability Flow ODE: Deriving a deterministic ODE that shares the same marginal densities as the SDE, enabling near-exact likelihood computation (accuracy is limited by both numerical ODE solver discretization and variance of the unbiased Hutchinson trace estimator) and latent space manipulation.
• Sub-VP SDE: A new SDE class proposed to improve likelihoods by bounding variance tighter than the VP SDE.

What experiments were performed?

The authors validated the framework on standard image benchmarks:

• Datasets: CIFAR-10 (32x32), CelebA (64x64), LSUN (Bedroom, Church), and CelebA-HQ (256x256 and 1024x1024).
• Ablation Studies: Comparing samplers (Ancestral vs. Reverse Diffusion vs. Probability Flow vs. PC) and SDE types (VE, VP, sub-VP).
• Architecture Search: Exploring improvements like FIR up/downsampling, rescaling skip connections, and increasing depth (leading to NCSN++ and DDPM++ architectures).
• Likelihood Evaluation: Computing Negative Log-Likelihood (NLL) in bits/dim using the Probability Flow ODE.
• Inverse Problems: Testing class-conditional generation, inpainting, and colorization using the conditional reverse-time SDE.

What outcomes/conclusions?

• Record Performance: The NCSN++ cont. (deep, VE) model achieved an Inception Score of 9.89 and FID of 2.20 on CIFAR-10 (as of ICLR 2021).
• High-Fidelity Generation: First score-based model to generate 1024x1024 images (CelebA-HQ).
• Competitive Likelihoods: The DDPM++ cont. (deep, sub-VP) model achieved 2.99 bits/dim on uniformly dequantized CIFAR-10, a record at the time.
• Sampling Efficiency: PC samplers consistently outperformed predictor-only methods (like standard ancestral sampling) for the same computational cost.
• Controllable Generation: Successful application to inpainting and colorization using a single unconditional model.
• Limitations: Sampling remains slower than GANs on the same datasets. The breadth of available samplers introduces many hyperparameters (SDE type, predictor, corrector, signal-to-noise ratio, number of steps) that require tuning.

Reproducibility Details

Data

• CIFAR-10: Used for main benchmarking (FID, Inception Score, NLL).
• CelebA-HQ: Used for high-resolution experiments at 256x256 and 1024x1024.
• LSUN: Bedroom and Church Outdoor categories (256x256) used for sampler comparison and controllable generation (inpainting, colorization).
• Preprocessing: CIFAR-10 images are 32x32; CelebA pre-processed to 64x64 following Song & Ermon (2020). Data is typically scaled to $[0, 1]$ or standardized depending on the specific SDE config.

Algorithms

Forward SDEs:

Here $dw$ denotes a Wiener process increment (a small, independent Gaussian noise burst at each timestep).

• VE SDE (Variance Exploding): $dx = \sqrt{\frac{d[\sigma^2(t)]}{dt}} dw$. Corresponds to SMLD. Used with $\sigma_{\min}=0.01$ and $\sigma_{\max}$ chosen via heuristics.
• VP SDE (Variance Preserving): $dx = -\frac{1}{2}\beta(t)x dt + \sqrt{\beta(t)} dw$. Corresponds to DDPM.
• Sub-VP SDE: $dx = -\frac{1}{2}\beta(t)x dt + \sqrt{\beta(t)(1 - e^{-2\int_0^t \beta(s)ds})} dw$. Bounded variance, good for likelihoods.

Reverse-Time SDE Solver (Predictor):

• Discretized via Reverse Diffusion Sampling, which matches the forward discretization.
• Euler-Maruyama solver used for continuously-trained models.

Corrector Algorithm:

• Langevin MCMC: Applies annealed Langevin dynamics: adds noise and takes a score-guided gradient step to correct the marginal distribution at each timestep.
• PC Sampling: Alternates between one step of the Predictor and one step of the Corrector.
• Signal-to-Noise Ratio ($r$): A hyperparameter for the corrector step size. Tuned values: $r \approx 0.16$ for VE SDEs on CIFAR-10.

Models

• NCSN++: Optimized architecture for VE SDEs. Key features:
  • 4 residual blocks per resolution.
  • BigGAN-type residual blocks.
  • Rescaling skip connections by $1/\sqrt{2}$.
  • FIR (Finite Impulse Response) up/downsampling.
  • “Residual” progressive architecture for input, no progressive growing for output.
• DDPM++: Optimized architecture for VP/sub-VP SDEs. Similar to NCSN++ but without FIR upsampling and no progressive growing.
• Deep Variants: “cont. (deep)” models double the depth (from 4 to 8 blocks per resolution) for the best reported results.
• Conditioning: Time $t$ is conditioned via random Fourier feature embeddings (scale 16) for continuous models.

Evaluation

Metrics:

• FID (Fréchet Inception Distance): Computed on 50k samples.
• Inception Score: Reported for CIFAR-10.
• NLL (Negative Log-Likelihood): Reported in bits/dim on uniformly dequantized data using the Probability Flow ODE.

Denoising: A single denoising step using Tweedie’s formula is applied at the end of sampling to remove residual noise, which significantly improves FID.

Hardware

Training:

• Batch size: 128 for CIFAR-10, 64 for LSUN, 8 for high-res CelebA-HQ.
• Iterations: Discrete-objective models trained for 1.3M iterations during architecture exploration. Continuous-objective models (cont.) trained for 0.95M iterations. High-res CelebA-HQ (1024x1024) trained for approximately 2.4M iterations.
• EMA: Exponential Moving Average rate of 0.999 used for VE models, 0.9999 for VP models.

Artifacts

All datasets used (CIFAR-10, CelebA-HQ, LSUN) are publicly available. Pretrained model checkpoints for CIFAR-10, CelebA-HQ, and FFHQ are provided in the repository. Specific hardware requirements (GPU type, training time) are not detailed in the paper.

Paper Information

Citation: Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2021). Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021. https://arxiv.org/abs/2011.13456

Publication: ICLR 2021

@inproceedings{song2021scorebased,
  title     = {Score-Based Generative Modeling through Stochastic Differential Equations},
  author    = {Song, Yang and Sohl-Dickstein, Jascha and Kingma, Diederik P and Kumar, Abhishek and Ermon, Stefano and Poole, Ben},
  booktitle = {International Conference on Learning Representations},
  year      = {2021},
  url       = {https://openreview.net/forum?id=PxTIG12RRHS}
}

Additional Resources:

• GitHub Repository
• Score Matching and Denoising Autoencoders

This is a Theory Paper.

Its primary contribution is a formal mathematical derivation connecting two previously distinct techniques: Score Matching (SM) and Denoising Autoencoders (DAE). It provides the “why” behind the empirical success of DAEs by grounding them in the probabilistic framework of energy-based models. It relies on proofs and equivalence relations (e.g., $J_{ESMq_{\sigma}} \sim J_{DSMq_{\sigma}}$).

What is the motivation?

The paper bridges a gap between two successful but disconnected approaches in unsupervised learning:

1. Denoising Autoencoders (DAE): Empirically successful for pre-training deep networks. They previously lacked a clear probabilistic interpretation.
2. Score Matching (SM): A theoretically sound method for estimating unnormalized density models that avoids the partition function problem but requires computing expensive second derivatives.

By connecting them, the authors aim to define a proper probabilistic model for DAEs (allowing sampling/ranking) and find a simpler way to apply score matching that avoids second derivatives.

What is the novelty here?

The core novelty is the Denoising Score Matching (DSM) framework and the proof of its equivalence to DAEs. Key contributions include:

• Equivalence Proof: Showing that training a DAE with Gaussian noise is equivalent to matching the score of a model against a non-parametric Parzen density estimator of the data.
• Denoising Score Matching ($J_{DSM}$): A new objective that learns a score function by trying to denoise corrupted samples. This avoids the explicit second derivatives required by standard Implicit Score Matching ($J_{ISM}$).
• Explicit Energy Function: Deriving the specific energy function $E(x;\theta)$ that corresponds to the standard sigmoid DAE architecture.
• Justification for Tied Weights: Providing a theoretical justification for tying encoder and decoder weights, which arises naturally from differentiating the energy function.

What experiments were performed?

The validation in this theoretical paper is purely mathematical and focuses on formal proofs:

• Derivation of Equivalence: The paper formally proves the chain of equivalences: $$J_{ISMq_{\sigma}} \sim J_{ESMq_{\sigma}} \sim J_{DSMq_{\sigma}} \sim J_{DAE\sigma}$$ where $q_{\sigma}$ is the Parzen density estimate.
• Appendix Proof: A detailed proof is provided to show that Explicit Score Matching ($J_{ESM}$) on the Parzen density is equivalent to the proposed Denoising Score Matching ($J_{DSM}$) objective.

What outcomes/conclusions?

• Theoretical Unification: DAE training is formally equivalent to Score Matching on a smoothed data distribution ($q_{\sigma}$).
• New Training Objective: The $J_{DSM}$ objective offers a computationally efficient way to perform score matching (no Hessian required) by using a denoising objective.
• Probabilistic Interpretation: DAEs can now be understood as Energy-Based Models (EBMs), allowing for operations like sampling (via Hybrid Monte Carlo) and likelihood ranking, which were previously ill-defined for standard autoencoders.
• Regularization Insight: The smoothing kernel width $\sigma$ in the Parzen estimator corresponds to the noise level in the DAE. This suggests that DAEs are learning a regularized version of the score, which may explain their robustness.
• Connection to Regularized Score Matching: The paper notes that Kingma and LeCun (2010) independently proposed a regularized score matching criterion $J_{ISMreg}$ derived by approximating $J_{ISMq_{\sigma}}$. The four $q_{\sigma}$-based objectives in this work (including the DAE objective) can be seen as approximation-free forms of regularized score matching, with the additional advantage that $J_{DSMq_{\sigma}}$ does not require second derivatives.

Key Concepts Explained

1. “Score” and “Score Matching”

What does “score” actually mean?

In this paper (and probabilistic modeling generally), the score is the gradient of the log-density with respect to the data vector $x$.

• Definition: $\psi(x) = \nabla_x \log p(x)$.
• Intuition: It is a vector field pointing in the direction of highest probability increase. Crucially, calculating the score avoids the intractable partition function $Z$, because $\nabla_x \log p(x) = \nabla_x \log \tilde{p}(x) - \nabla_x \log Z = \nabla_x \log \tilde{p}(x)$. The constant $Z$ vanishes upon differentiation.

What is Score Matching?

Score Matching is a training objective for unnormalized models. It minimizes the squared Euclidean distance between the model’s score $\psi(x;\theta)$ and the data’s true score $\nabla_x \log q(x)$.

2. The Parzen Density Estimator

What is it?

It is a non-parametric method for estimating a probability density function from finite data. It places a smooth kernel (here, a Gaussian) centered at every data point in the training set $D_n$.

• Formula: $q_{\sigma}(\tilde{x}) = \frac{1}{n} \sum_{t=1}^n \mathcal{N}(\tilde{x}; x^{(t)}, \sigma^2 I)$.

Why smooth the data?

1. To define the score: The empirical data distribution is a set of Dirac deltas (spikes). The gradient (score) of a Dirac delta is undefined. Smoothing creates a differentiable surface, allowing a valid target score $\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x})$ to be computed.

2. To model corruption: The Parzen estimator with Gaussian kernels mathematically models the process of taking a clean data point $x$ and adding Gaussian noise - the exact procedure used in Denoising Autoencoders.

3. Why avoiding second derivatives matters

Standard Implicit Score Matching (ISM) eliminates the need for the unknown data score, but introduces a new cost: it requires computing the trace of the Hessian (the sum of second partial derivatives) of the log-density.

• The Cost: For high-dimensional data (like images) and deep networks, computing second derivatives of the log-density is computationally expensive.
• This paper shows that Denoising Score Matching (DSM) allows you to bypass Hessian computation entirely. By using the Parzen target, the objective simplifies to matching a first-order vector, making it scalable to deep neural networks.

4. The equivalence chain - why each step?

The chain $J_{ISMq_{\sigma}} \sim J_{ESMq_{\sigma}} \sim J_{DSMq_{\sigma}} \sim J_{DAE\sigma}$ connects the concepts.

• $J_{ISMq_{\sigma}} \sim J_{ESMq_{\sigma}}$ (Implicit $\to$ Explicit): Why: Integration by parts. This is Hyvärinen’s original proof (2005): integration by parts moves the derivative from $\psi$ onto the data density $q$, producing a term involving $q$’s gradient (the score). The boundary term vanishes because $q_{\sigma}$ decays to zero at infinity (Hyvärinen’s 2005 regularity condition for Implicit Score Matching). The result allows replacing the unknown data score with a computable term involving only the model’s score and its Jacobian.

• $J_{ESMq_{\sigma}} \sim J_{DSMq_{\sigma}}$ (Explicit $\to$ Denoising): Why: The explicit score of the Parzen density is known. When $x$ is perturbed to $\tilde{x}$ by Gaussian noise $\epsilon \sim \mathcal{N}(0, \sigma^2 I)$, the gradient of the log-density pointing back to the mean is exactly $\frac{1}{\sigma^2}(x - \tilde{x})$. Minimizing the error against the true score becomes minimizing the error against this restoration vector.

• $J_{DSMq_{\sigma}} \sim J_{DAE\sigma}$ (Denoising $\to$ Autoencoder): Why: Algebraic substitution. If you define the model’s score $\psi(\tilde{x};\theta)$ to be proportional to the reconstruction error ($\propto x^r - \tilde{x}$), the score matching loss $J_{DSM}$ becomes proportional to the standard autoencoder squared loss $|x^r - x|^2$.

5. Energy-Based Models (EBMs) connection

What is an EBM?

An EBM defines a probability distribution via an energy function $E(x;\theta)$, where $p(x;\theta) \propto e^{-E(x;\theta)}$.

Why standard autoencoders lack probabilistic interpretation:

A standard autoencoder acts as a deterministic map $x \to x^r$, providing a reconstruction error. It lacks a normalization constant or a defined density function to support sampling or probability queries.

What does this enable?

By proving the equivalence, the DAE is formally defined as an EBM. This enables:

1. Sampling: Using MCMC methods (like Hybrid Monte Carlo) to generate new data from the DAE.
2. Ranking: Calculating the energy of inputs to determine which are more “likely” or “normal” (useful for anomaly detection).

6. The specific energy function form

The function is:

$$E(x; W, b, c) = - \frac{1}{\sigma^2} \left( \langle c, x \rangle - \frac{1}{2}|x|^2 + \sum_{j=1}^{d_h} \text{softplus}(\langle W_j, x \rangle + b_j) \right)$$

Why does it have that specific form?

It was derived via integration to ensure its derivative matches the DAE architecture. The authors worked backward from the DAE’s reconstruction function (sigmoid + linear) to find the scalar field that generates it.

Where does the quadratic term come from?

The score (negative energy gradient) needs to look like $\psi(x) \propto c - x + W^T\text{sigmoid}(Wx + b)$.

• The term $-x$ in the score arises because $\nabla_x(-\frac{1}{2}|x|^2) = -x$. Including $-\frac{1}{2}|x|^2$ inside the energy’s numerator produces this linear term after differentiation.

How does differentiating it recover the DAE reconstruction?

• $\nabla_x \sum_j \text{softplus}(\langle W_j, x \rangle + b_j) = W^T \sigma(Wx + b)$ (The encoder part).
• $\nabla_x \langle c, x \rangle = c$ (The bias).
• $\nabla_x (-\frac{1}{2}|x|^2) = -x$ (The input subtraction).
• Result: $-\nabla_x E \propto c + W^T h - x = x^r - x$.

7. “Tied weights” justification

What does it mean for weights to be “tied”?

The decoder matrix is the transpose of the encoder matrix ($W^T$).

Why is this theoretically justified?

Because the reconstruction function is interpreted as the gradient of an energy function. A vector field can only be the gradient of a scalar field if its Jacobian is symmetric.

• In the DAE energy derivative, the encoder contributes $W^T \sigma(Wx + b)$. If the decoder used a separate matrix $U$, the resulting vector field would not be a valid gradient of any scalar energy function (unless $U = W^T$).
• Therefore, for a DAE to correspond to a valid probabilistic Energy-Based Model, the weights must be tied.

The necessity of tied weights:

Within this parametrization, tied weights are a mathematical necessity: a separate decoder matrix $U \neq W^T$ would make the reconstruction function an invalid gradient of any scalar energy, breaking the EBM correspondence.

Reproducibility Details

Since this is a theoretical paper, the “reproducibility” lies in the mathematical formulations derived.

Data

• Input Data ($D_n$): The theory assumes a set of training examples $D_n = {x^{(1)}, …, x^{(n)}}$ drawn from an unknown true pdf $q(x)$.
• Parzen Density Estimate ($q_{\sigma}$): The theoretical targets are derived from a kernel-smoothed empirical distribution: $$q_{\sigma}(\tilde{x}) = \frac{1}{n} \sum_{t=1}^n q_{\sigma}(\tilde{x}|x^{(t)})$$ where the kernel is an isotropic Gaussian of variance $\sigma^2$.

Algorithms

1. Denoising Score Matching (DSM) Objective

The paper proposes this objective as a tractable alternative to standard score matching. It minimizes the distance between the model score and the gradient of the log-noise density:

$$J_{DSMq_{\sigma}}(\theta) = \mathbb{E}_{q_{\sigma}(x,\tilde{x})} \left[ \frac{1}{2} \left| \psi(\tilde{x};\theta) - \frac{\partial \log q_{\sigma}(\tilde{x}|x)}{\partial \tilde{x}} \right|^2 \right]$$

For Gaussian noise, the target score is simply $\frac{1}{\sigma^2}(x - \tilde{x})$.

2. Equivalence Chain

The central result connects four objectives:

$$J_{ISMq_{\sigma}} \sim J_{ESMq_{\sigma}} \sim J_{DSMq_{\sigma}} \sim J_{DAE\sigma}$$

This implies optimizing the DAE reconstruction error is minimizing a score matching objective.

Models

1. The Denoising Autoencoder (DAE)

• Corruption: Additive isotropic Gaussian noise $\tilde{x} = x + \epsilon, \epsilon \sim \mathcal{N}(0, \sigma^2 I)$.
• Encoder: $h = \text{sigmoid}(W\tilde{x} + b)$.
• Decoder: $x^r = W^T h + c$ (Tied weights $W$).
• Loss: Squared reconstruction error $|x^r - x|^2$. (The equivalence with DSM introduces a $\frac{1}{2\sigma^4}$ scaling factor.)

2. The Corresponding Energy Function

To make the DAE equivalent to Score Matching, the underlying Energy-Based Model $p(x;\theta) \propto e^{-E(x;\theta)}$ must have the following energy function:

$$E(x; W, b, c) = - \frac{1}{\sigma^2} \left( \langle c, x \rangle - \frac{1}{2}|x|^2 + \sum_{j=1}^{d_h} \text{softplus}(\langle W_j, x \rangle + b_j) \right)$$

Note the scaling by $1/\sigma^2$ and the quadratic term $|x|^2$.

Evaluation

• Metric: Theoretical Equivalence ($\sim$).
• Condition: The equivalence holds provided $\sigma > 0$ and the density $q_{\sigma}$ is differentiable and vanishes at infinity (Hyvärinen’s 2005 regularity condition for Implicit Score Matching).

Paper Information

Citation: Vincent, P. (2011). A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7), 1661-1674. https://doi.org/10.1162/NECO_a_00142

Publication: Neural Computation 2011

@article{vincentConnectionScoreMatching2011,
  title = {A {{Connection Between Score Matching}} and {{Denoising Autoencoders}}},
  author = {Vincent, Pascal},
  year = 2011,
  month = jul,
  journal = {Neural Computation},
  volume = {23},
  number = {7},
  pages = {1661--1674},
  doi = {10.1162/NECO_a_00142}
}

Additional Resources:

• Official PDF

This is primarily a Method paper, with a significant Theory component.

• Method: It proposes “Rectified Flow,” a novel generative framework that learns ordinary differential equations (ODEs) to transport distributions via straight paths. It introduces the “Reflow” algorithm to iteratively straighten these paths.
• Theory: It provides rigorous proofs connecting the method to Optimal Transport, showing that the rectification process yields a coupling with non-increasing convex transport costs and that recursive reflow reduces the curvature of trajectories.

What is the motivation?

The work addresses two main challenges in unsupervised learning: generative modeling (generating data from noise) and domain transfer (mapping between two observed distributions).

• Inefficiency of ODE/SDE Models: Continuous-time models (like Score-based Generative Models and DDPMs) require simulating diffusions over many steps, resulting in high computational costs during inference.
• Complexity of GANs: GANs provide fast (one-step) generation alongside challenges with training instability and mode collapse.
• Disconnection: Generative modeling and domain transfer are often treated as separate tasks requiring different techniques.

The authors aim to unify these tasks into a single “transport mapping” problem while bridging the gap between high-quality continuous models and fast one-step models.

What is the novelty here?

The core novelty is the Rectified Flow framework and the Reflow procedure.

• Straight-Line ODEs: Rectified Flow learns an ODE drift $v$ to follow the straight line connecting data pairs $(X_0, X_1)$, providing an alternative to diffusion models that rely on stochastic paths or specific forward processes. This is achieved via a simple least-squares optimization problem.
• Reflow (Iterative Straightening): The authors introduce a recursive training procedure where a new flow is trained on the data pairs $(Z_0, Z_1)$ generated by the previous flow. Theoretical analysis shows this reduces the “transport cost” and straightens the trajectories, allowing for accurate 1-step simulation (effectively converting the ODE into a one-step model).
• Unified Framework: The method uses the exact same algorithm for generation ($\pi_0$ is Gaussian) and domain transfer ($\pi_0$ is a source dataset), removing the need for adversarial losses or cycle-consistency constraints.

What experiments were performed?

The authors validated the method across image generation, translation, and domain adaptation tasks.

• Unconditioned Image Generation:
  • Dataset: CIFAR-10 ($32\times32$).
  • Baselines: Compared against GANs (StyleGAN2, TDPM), Diffusion/SDE Models (VP SDE, sub-VP SDE, VE SDE), ODE methods (VP ODE, sub-VP ODE, VE ODE), and distilled methods (DDIM Distillation).
  • High-Res: Validated on LSUN Bedroom/Church, CelebA-HQ, and AFHQ ($256\times256$).
• Image-to-Image Translation:
  • Datasets: AFHQ (Cat $\leftrightarrow$ Dog/Wild), MetFace $\leftrightarrow$ CelebA-HQ.
  • Setup: Transferring styles while preserving semantic identity (using a classifier-based feature mapping metric).
• Domain Adaptation:
  • Datasets: DomainNet, Office-Home.
  • Metric: Classification accuracy on the transferred testing data.

What outcomes/conclusions?

• Superior 1-Step Generation: On CIFAR-10 with a single Euler step (as of ICLR 2023), the distilled 2-Rectified Flow achieved an FID of 4.85, beating the best one-step U-Net model TDPM (FID 8.91, a truncated diffusion model using a GAN). The distilled 3-Rectified Flow reached a Recall of 0.51, beating the GAN baseline StyleGAN2+ADA (Recall 0.49).
• Straightening Effect: The “Reflow” procedure was empirically shown to reduce the “straightness” error and transport costs, validating the theoretical claims. “Straightness” is measured as $S(Z) = \mathbb{E}[\int_0^1 |\dot{Z}_t - (Z_1 - Z_0)|^2, dt]$ (zero means perfectly straight); “transport cost” is $\mathbb{E}[c(Z_1 - Z_0)]$ for a convex cost $c$, and Reflow reduces this for all convex costs.
• High-Quality Transfer: The model successfully performed image translation (e.g., Cat to Wild Animal) without paired data or cycle-consistency losses.
• Strong Full-Simulation Results: With RK45 adaptive ODE solving, 1-Rectified Flow achieves FID 2.58 and Recall 0.57 on CIFAR-10 (Table 1a), the best among ODE methods and comparable to fully simulated SDEs (VP SDE: FID 2.55).
• Fast Simulation: The method allows for extremely coarse time discretization (e.g., $N=1$) without significant quality loss after reflow, effectively solving the slow inference speed of standard ODE models.
• Domain Adaptation: On Office-Home, Rectified Flow achieves 69.2% accuracy, outperforming Deep CORAL (68.7%) and other baselines. On DomainNet, it achieves 41.4%, comparable to Deep CORAL (41.5%) and MLDG (41.2%).

Reproducibility Details

Data

The paper utilizes several standard computer vision benchmarks.

Algorithms

• Objective Function: The drift $v(Z_t, t)$ is trained by minimizing a least-squares regression objective: $$\min_{v} \int_{0}^{1} \mathbb{E}[|(X_1 - X_0) - v(X_t, t)|^2] dt$$ where $X_t = tX_1 + (1-t)X_0$ is the linear interpolation.

• Reflow Procedure: Iteratively updates the flow. Let $Z^k$ be the $k$-th rectified flow.

  1. Generate 4 million data pairs $(Z_0^k, Z_1^k)$ by simulating the current flow.
  2. Fine-tune the $i$-rectified flow model for 300,000 steps on these pairs to obtain the $(i+1)$-rectified flow.

• Distillation: For 1-step distillation ($k=1$), the L2 loss is replaced with LPIPS perceptual similarity, which empirically yields better image quality. For multi-step distillation, training samples $t$ from ${0, 1/k, \ldots, (k-1)/k}$ rather than the full $[0, 1]$ interval.

• ODE Solver:

  • Training: Analytical linear interpolation.
  • Inference: Euler method (constant step size $1/N$) or RK45 (adaptive).

Models

• Architecture:

  • Uses the DDPM++ U-Net architecture (from Song et al., 2020) across experiments. Implementation is modified from the open-source code of Song et al.

• Optimization:

  • Optimizer: Adam (CIFAR-10) or AdamW (Transfer/Adaptation).
  • Hyperparameters:
    • LR: $2 \times 10^{-4}$ (CIFAR), Grid search for transfer.
    • EMA: 0.999999 (CIFAR), 0.9999 (Transfer).
    • Batch Size: 4 (Transfer), 16 (Domain Adaptation).
    • Dropout: 0.15 (CIFAR), 0.1 (Transfer).

Evaluation

Hardware

The paper does not specify GPU models or training times. The DDPM++ U-Net architecture used in the experiments typically requires multi-GPU setups for training on high-resolution datasets.

Artifacts

Paper Information

Citation: Liu, X., Gong, C., & Liu, Q. (2023). Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. International Conference on Learning Representations (ICLR). https://openreview.net/forum?id=XVjTT1nw5z

Publication: ICLR 2023

@inproceedings{liuFlowStraightFast2023,
  title = {Flow {{Straight}} and {{Fast}}: {{Learning}} to {{Generate}} and {{Transfer Data}} with {{Rectified Flow}}},
  booktitle = {International Conference on Learning Representations},
  author = {Liu, Xingchao and Gong, Chengyue and Liu, Qiang},
  year = 2023,
  url = {https://openreview.net/forum?id=XVjTT1nw5z}
}

Additional Resources:

• Official Code Repository
• OpenReview Page

This is primarily a Method paper, as it introduces “Flow Matching” (FM), a novel simulation-free paradigm for training Continuous Normalizing Flows (CNFs) at scale. It is supported by a strong Theory basis, providing formal theorems that allow the intractable marginal vector field regression to be solved via a tractable conditional objective. It also touches on Systematization by showing that existing diffusion paths are specific instances of the proposed Gaussian probability path framework.

What is the motivation?

The paper aims to overcome the scaling limitations of Continuous Normalizing Flows (CNFs).

• Problem: Standard Maximum Likelihood training for CNFs requires expensive numerical ODE simulations during training, which scales poorly. Existing simulation-free methods often involve intractable integrals or result in biased gradients.
• Gap: Diffusion models scale well, yet they are restricted to specific, curved probability paths (e.g., VP, VE) that can result in slow sampling and long training times.
• Goal: To develop an efficient, simulation-free training method for CNFs that supports arbitrary probability paths, specifically allowing for straighter, more efficient trajectories like those from Optimal Transport.

What is the novelty here?

The core novelty is Flow Matching (FM) and specifically the Conditional Flow Matching (CFM) objective.

• Direct Vector Field Regression: The model regresses a target vector field $u_t$ that generates a desired probability path $p_t$.
• Conditional Flow Matching (CFM): The authors prove that regressing the vector field of conditional paths (e.g., $p_t(x|x_1)$ given a single data point) yields the same gradients as regressing the intractable marginal vector field. This bypasses the need to know the marginal score or vector field.
• Optimal Transport Paths: The framework enables the use of Optimal Transport (OT) displacement interpolation for probability paths. OT paths are straight lines with constant speed, leading to faster training and easier sampling.

Concurrent work note: Rectified Flow (Liu et al., 2023) and Stochastic Interpolants (Albergo & Vanden-Eijnden, 2023) were published concurrently at ICLR 2023 with structurally similar contributions under different names. All three independently propose simulation-free training of continuous flows via direct vector field regression; the differences lie in the specific interpolation schemes, theoretical framing, and experimental focus.

What experiments were performed?

• Domains: 2D Checkerboard data, CIFAR-10, and ImageNet at resolutions $32 \times 32$, $64 \times 64$, and $128 \times 128$.
• Task: Unconditional generative modeling (density estimation and sample quality) and conditional super-resolution ($64 \times 64 \to 256 \times 256$).
• Baselines: Compared against Diffusion-based methods on the same architecture (U-Net): DDPM, Score Matching (SM), and ScoreFlow.
• Ablations: Specifically compared FM with Diffusion paths vs. FM with Optimal Transport (OT) paths to isolate the benefit of the training objective vs. the path choice.

What outcomes/conclusions?

• Outperforms diffusion baselines: FM-OT consistently outperforms all diffusion-based methods (DDPM, Score Matching, ScoreFlow) in both Likelihood (NLL) and Sample Quality (FID) across CIFAR-10 and ImageNet, using the same U-Net architecture and training budget. Selected rows from Table 1 (NLL in bits per dimension, BPD; lower is better for all three metrics; “FM w/ OT” and “FM w/ Diffusion” refer to FM trained with OT paths and Diffusion paths respectively):

• Training stability: FM with diffusion paths (FM w/ Diffusion) is itself a more stable alternative to diffusion training than DDPM and Score Matching, as shown by training curves in the paper (Figure 5), even before switching to OT paths. The OT path then provides further gains.
• Sampling speed: The straight trajectories of OT paths allow accurate sampling with significantly fewer function evaluations (NFE) compared to diffusion paths.
• Generality: Diffusion is a specific instance of Gaussian probability paths within FM. OT paths are a better-optimized alternative available within the same framework.
• Downstream adoption: Flow matching has been adopted beyond image generation. DynamicFlow uses it as the generative backbone for simultaneously generating ligand molecules and transforming protein pockets, extending flow matching to structure-based drug design.

Reproducibility Details

Data

• Datasets: CIFAR-10, ImageNet ($32 \times 32$, $64 \times 64$, $128 \times 128$).
• Preprocessing:
  • Images are center-cropped and resized.
  • For $32 \times 32$ and $64 \times 64$, the preprocessing follows Chrabaszcz et al. (2017).
  • Data is transformed via $\varphi(y) = 2^7(y+1)$ mapping $[-1, 1]$ pixel values to $[0, 256]$ for BPD computation.

Algorithms

1. Conditional Flow Matching (CFM) Objective

The practical training objective used is the CFM loss, which bypasses intractable marginalization:

$$\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, q(x_1), p(x_0)} | v_t(\psi_t(x_0)) - u_t(\psi_t(x_0) | x_1) |^2$$

Where $t \sim \mathcal{U}[0,1]$, $x_1 \sim q(x_1)$ (data), and $x_0 \sim p(x_0)$ (noise).

2. Optimal Transport (OT) Probability Path

The authors recommend the OT path for efficiency.

• Mean/Std Schedule: $\mu_t(x) = t x_1$ and $\sigma_t(x) = 1 - (1 - \sigma_{min})t$.
• Conditional Flow Map: $\psi_t(x) = (1 - (1 - \sigma_{min})t)x + t x_1$.
• Target Vector Field: The closed-form regression target for OT is: $$u_t(x|x_1) = \frac{x_1 - (1 - \sigma_{min})x}{1 - (1 - \sigma_{min})t}$$

3. Sampling

Sampling is performed by solving the ODE $\frac{d}{dt}\phi_t(x) = v_t(\phi_t(x))$ from $t=0$ to $t=1$ using the learned vector field $v_t$.

• Solver: dopri5 (adaptive) is used for robust evaluation. Fixed-step solvers (Euler, Midpoint) are used for low-NFE efficiency tests.

Models

• Architecture: U-Net architecture from Dhariwal & Nichol (2021) is used for all image experiments.
• Toy Data: 5-layer MLP with 512 neurons.
• Hyperparameters:
  • Optimizer: Adam ($\beta_1=0.9, \beta_2=0.999$, weight decay=0.0).
  • Learning Rate: Polynomial decay or constant (see Table 3 in paper).
  • $\sigma_{min}$: Set to a small value (e.g., $1e-5$).

Evaluation

• Metrics:
  • NLL (BPD): Computed using the continuous change of variables formula, estimated via the Hutchinson trace estimator to bypass $O(d^3)$ divergence computation.
  • FID: Frechet Inception Distance for sample quality.
  • NFE: Number of Function Evaluations required by the solver.
• Likelihood Computation: Requires solving an augmented ODE to track the log-density change: $$\frac{d}{dt} \begin{bmatrix} \phi_t(x) \ f(t) \end{bmatrix} = \begin{bmatrix} v_t(\phi_t(x)) \ -\text{div}(v_t(\phi_t(x))) \end{bmatrix}$$

Hardware

• CIFAR-10: 2 GPUs.
• ImageNet-32: 4 GPUs.
• ImageNet-64: 16 GPUs.
• ImageNet-128: 32 GPUs.
• Precision: Full 32-bit for CIFAR/IM-32; 16-bit mixed precision for IM-64/128.

Artifacts

The paper does not release the original training code or model weights used in the experiments. The facebookresearch/flow_matching library was released later as a general-purpose PyTorch implementation of flow matching algorithms. Standard benchmark datasets (CIFAR-10, ImageNet) are publicly available.

Theoretical Notes: Why CFM Works

The paper relies on three key theorems to make training tractable.

Theorem 1 (Marginal Generation):

Marginalizing conditional vector fields $u_t(x|x_1)$ yields the correct marginal vector field $u_t(x)$ that generates the marginal probability path $p_t(x)$.

$$u_t(x) = \int u_t(x|x_1) \frac{p_t(x|x_1)q(x_1)}{p_t(x)} dx_1$$

Understanding the Proof:

To understand why this theorem holds, we have to look at the Continuity Equation, which is the fundamental partial differential equation (PDE) that links a probability density path $p_t$ to a vector field $u_t$.

A vector field $u_t$ is said to “generate” a probability path $p_t$ if and only if they satisfy the continuity equation:

$$\frac{\partial p_t(x)}{\partial t} + \nabla \cdot (p_t(x) u_t(x)) = 0$$

The proof of Theorem 1 relies on substituting the definitions of the marginal path and vector field into this equation to see if they balance out.

Step-by-Step Proof:

1. Start with the time derivative of the marginal path: We begin by differentiating the marginal probability path $p_t(x)$ with respect to time. By definition, the marginal path is the integral of the conditional paths over the data distribution: $$\frac{\partial p_t(x)}{\partial t} = \frac{\partial}{\partial t} \int p_t(x|x_1) q(x_1) dx_1$$

2. Swap derivative and integral: Assuming standard regularity conditions (Leibniz Rule), we can move the time derivative inside the integral: $$\frac{\partial p_t(x)}{\partial t} = \int \frac{\partial p_t(x|x_1)}{\partial t} q(x_1) dx_1$$

3. Apply the Conditional Continuity Equation: This is the critical step. We know that the conditional vector field $u_t(x|x_1)$ generates the conditional path $p_t(x|x_1)$. Therefore, for every single sample $x_1$, the pair satisfies the continuity equation: $$\frac{\partial p_t(x|x_1)}{\partial t} = -\nabla \cdot (p_t(x|x_1) u_t(x|x_1))$$

   Substituting this into our integral gives: $$\frac{\partial p_t(x)}{\partial t} = -\int \nabla \cdot (p_t(x|x_1) u_t(x|x_1)) q(x_1) dx_1$$

4. Pull the Divergence out: Since the divergence operator ($\nabla \cdot$) acts on $x$ and the integral is over $x_1$, we can pull the divergence operator outside the integral (by linearity): $$\frac{\partial p_t(x)}{\partial t} = -\nabla \cdot \left( \int p_t(x|x_1) u_t(x|x_1) q(x_1) dx_1 \right)$$

5. Match with the Marginal Vector Field Definition: Now, look at the term inside the parentheses. The paper defines the marginal vector field $u_t(x)$ specifically to make this term simpler. Rearranging the definition of $u_t(x)$ provided in the theorem: $$p_t(x) u_t(x) = \int p_t(x|x_1) u_t(x|x_1) q(x_1) dx_1$$

   Substitute $p_t(x) u_t(x)$ back into our equation from Step 4: $$\frac{\partial p_t(x)}{\partial t} = -\nabla \cdot (p_t(x) u_t(x))$$

Conclusion: We have just shown that $\frac{\partial p_t(x)}{\partial t} + \nabla \cdot (p_t(x) u_t(x)) = 0$. This is exactly the continuity equation. Because the marginal path and the aggregated marginal vector field satisfy this equation, the vector field is proven to generate the path.

Theorem 2 (Gradient Equivalence):

The intractable Flow Matching objective $\mathcal{L}_{FM}$ (which requires $u_t(x)$) has the same gradients as the tractable Conditional Flow Matching objective $\mathcal{L}_{CFM}$.

$$\nabla_\theta \mathcal{L}_{FM}(\theta) = \nabla_\theta \mathcal{L}_{CFM}(\theta)$$

This allows the model to learn the marginal vector field by only seeing conditional sample paths.

Understanding the Proof:

The reason Theorem 2 holds is that the “Conditional Flow Matching” (CFM) objective is essentially an unbiased estimator of the “Flow Matching” (FM) objective (up to a constant). When we average over all the conditional data points $x_1$, the “cross-term” in the loss function aligns perfectly with the marginal vector field.

1. Expand the Loss Functions

First, let’s look at the squared error in both objectives. Recall that $v_t$ is our neural network (parameterized by $\theta$), $u_t$ is the intractable marginal target, and $u_t(x|x_1)$ is the tractable conditional target.

Expanding the squared norms:

• FM Objective: $$\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t, p_t(x)} \left[ |v_t(x)|^2 - 2v_t(x) \cdot u_t(x) + |u_t(x)|^2 \right]$$

• CFM Objective: $$\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, q(x_1), p_t(x|x_1)} \left[ |v_t(x)|^2 - 2v_t(x) \cdot u_t(x|x_1) + |u_t(x|x_1)|^2 \right]$$

Key Insight: When we take the gradient $\nabla_\theta$, the last term in both equations disappears because the targets ($u_t$) are independent of the network weights $\theta$. We only need to show that the expectations of the first two terms match.

2. Matching the First Term ($|v_t(x)|^2$)

This part is straightforward. The expectation of $|v_t(x)|^2$ is the same in both cases because of how the marginal density $p_t(x)$ is defined.

• FM: averages over $p_t(x)$.
• CFM: averages over $p_t(x|x_1)q(x_1)$.

Since $p_t(x) = \int p_t(x|x_1) q(x_1) dx_1$ (by definition), averaging over the joint distribution is mathematically identical to averaging over the marginal $p_t(x)$.

3. Matching the Cross Term (The “Trick”)

This is the critical part of the proof. We need to show that the interaction between the network and the marginal field equals the interaction between the network and the conditional field.

The Goal: Show $\mathbb{E}_{t, p_t(x)} [v_t(x) \cdot u_t(x)] = \mathbb{E}_{t, q(x_1), p_t(x|x_1)} [v_t(x) \cdot u_t(x|x_1)]$.

The Proof:

1. Start with the FM cross-term (marginal): $$\mathbb{E}_{t, p_t(x)} [v_t(x) \cdot u_t(x)]$$

2. Substitute the definition of the marginal vector field $u_t(x)$ derived in Theorem 1: $$u_t(x) = \int u_t(x|x_1) \frac{p_t(x|x_1) q(x_1)}{p_t(x)} dx_1$$

3. Plug this into the integral. The $p_t(x)$ terms cancel: $$\mathbb{E}_{t, p_t(x)} [v_t(x) \cdot u_t(x)] = \int_t \int_x p_t(x) v_t(x) \cdot \left[ \int_{x_1} u_t(x|x_1) \frac{p_t(x|x_1) q(x_1)}{p_t(x)} dx_1 \right] dx$$

4. This simplifies to: $$= \int_t \int_x \int_{x_1} v_t(x) \cdot u_t(x|x_1) p_t(x|x_1) q(x_1) dx_1 dx dt$$

5. This is exactly the definition of the expectation in the CFM objective: $$= \mathbb{E}_{t, q(x_1), p_t(x|x_1)} [v_t(x) \cdot u_t(x|x_1)]$$

Conclusion: Because the expectations of all terms involving $\theta$ are identical, the gradients must be identical.

Intuitively, this works like Denoising Score Matching or Stochastic Gradient Descent: even though each individual conditional vector field $u_t(x|x_1)$ points to a specific data point $x_1$ (which may differ from the true marginal direction), the average of all these pulls equals the true marginal vector field $u_t(x)$.

Theorem 3 (Gaussian Conditional VFs):

For any Gaussian probability path $p_t(x|x_1) = \mathcal{N}(x | \mu_t(x_1), \sigma_t(x_1)^2 I)$, the unique vector field generating it is available in closed form:

$$u_t(x|x_1) = \frac{\sigma’_t(x_1)}{\sigma_t(x_1)}(x - \mu_t(x_1)) + \mu’_t(x_1)$$

This theorem allows explicitly defining targets for both Diffusion (curved) and Optimal Transport (straight) paths.

Understanding the Proof:

The derivation of Theorem 3 comes from the direct relationship between a flow map $\psi_t$ and its generating vector field. Because we chose a specific, simple path (Gaussian), we can invert the flow map to find the vector field in closed form.

1. Define the Flow Map $\psi_t$

We start by defining the conditional probability path as a Gaussian:

$$p_t(x|x_1) = \mathcal{N}(x | \mu_t(x_1), \sigma_t(x_1)^2 I)$$

The simplest way to “push” a standard normal distribution (noise) $p_0 = \mathcal{N}(0, I)$ to this Gaussian is using an affine transformation (scaling and shifting). We define the flow map $\psi_t$ as:

$$\psi_t(x_0) = \sigma_t(x_1) x_0 + \mu_t(x_1)$$

This map takes a noise sample $x_0$ and transforms it into a sample $x$ at time $t$.

2. The Definition of a Generating Vector Field

By definition, a vector field $u_t$ generates a flow $\psi_t$ if the vector field describes the instantaneous velocity of the flow at any point. Mathematically:

$$u_t(\psi_t(x_0)) = \frac{d}{dt}\psi_t(x_0)$$

Let $x = \psi_t(x_0)$ be the position of the particle at time $t$. We want to find $u_t(x)$.

3. Invert the Flow Map

To find $u_t(x)$, we must express the equation in terms of $x$ rather than $x_0$. Since our flow map is a simple affine transformation (multiply and add), it is easily invertible (assuming $\sigma_t(x_1) \neq 0$):

$$x_0 = \frac{x - \mu_t(x_1)}{\sigma_t(x_1)}$$

We will call this inverse map $\psi_t^{-1}(x)$.

4. Differentiate the Flow Map

Now we calculate the left side of our definition equation (velocity): $\frac{d}{dt}\psi_t(x_0)$.

Taking the time derivative of $\psi_t(x_0) = \sigma_t(x_1) x_0 + \mu_t(x_1)$:

$$\frac{d}{dt}\psi_t(x_0) = \sigma’_t(x_1) x_0 + \mu’_t(x_1)$$

(Note: $\sigma’_t$ and $\mu’_t$ denote time derivatives).

5. Substitute and Solve

Now we combine everything. We know $u_t(\psi_t(x_0)) = \frac{d}{dt}\psi_t(x_0)$.

Substitute the result from Step 4 into this equation:

$$u_t(\psi_t(x_0)) = \sigma’_t(x_1) x_0 + \mu’_t(x_1)$$

This expresses the vector field in terms of the initial point $x_0$. We must express it in terms of the current point $x$. So, we plug in the inverse formula for $x_0$ derived in Step 3:

$$u_t(x|x_1) = \sigma’_t(x_1) \frac{x - \mu_t(x_1)}{\sigma_t(x_1)} + \mu’_t(x_1)$$

Rearranging terms gives the final closed form:

$$u_t(x|x_1) = \frac{\sigma’_t(x_1)}{\sigma_t(x_1)}(x - \mu_t(x_1)) + \mu’_t(x_1)$$

Why is this useful?

This formula means that as long as you can define a mean schedule $\mu_t(x_1)$ and a standard deviation schedule $\sigma_t(x_1)$ (which is easy to do for both Diffusion and Optimal Transport), you immediately get the exact vector field target $u_t(x|x_1)$ needed to train your neural network, bypassing complex ODE solving or score matching approximations.

Paper Information

Citation: Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. International Conference on Learning Representations (ICLR).

Publication: ICLR 2023

@inproceedings{lipmanFlowMatchingGenerative2023,
  title = {Flow Matching for Generative Modeling},
  author = {Lipman, Yaron and Chen, Ricky T. Q. and Ben-Hamu, Heli and Nickel, Maximilian and Le, Matt},
  booktitle = {International Conference on Learning Representations},
  year = {2023}
}

Additional Resources:

• ArXiv

This is primarily a Method paper, with significant Theory contributions.

The authors propose a specific algorithm (“InterFlow”) for constructing generative models based on continuous-time normalizing flows. The work is characterized by the derivation of a new training objective (a simple quadratic loss) that bypasses the computational bottlenecks of previous methods. It includes prominent baseline comparisons against continuous flow methods (FFJORD, OT-Flow) and diffusion models. The theoretical component establishes the validity of the interpolant density satisfying the continuity equation (a conservation law governing how probability mass flows) and bounds the Wasserstein-2 distance (a measure of transport cost between distributions, penalizing squared displacement) of the transport.

What is the motivation?

The primary motivation is to overcome the computational inefficiency of training Continuous Normalizing Flows (CNFs) using Maximum Likelihood Estimation (MLE). Standard CNF training requires backpropagating through numerical ODE solvers, which is costly and limits scalability.

Additionally, while score-based diffusion models (SDEs) have achieved high sample quality, they theoretically require infinite time integration and rely on specific noise schedules. The authors aim to establish a method that works strictly with Probability Flow ODEs on finite time intervals, retaining the flexibility to connect arbitrary densities without the complexity of SDEs or the cost of standard ODE adjoint methods.

What is the novelty here?

The core novelty is the Stochastic Interpolant framework:

• Explicit Interpolant Construction: The method defines a time-dependent interpolant $x_t = I_t(x_0, x_1)$ (e.g., trigonometric interpolation) that connects samples from the base density $\rho_0$ and target $\rho_1$.
• Simulation-Free Training: The velocity field $v_t(x)$ of the probability flow is learned by minimizing a simple quadratic objective: $G(\hat{v}) = \mathbb{E}[|\hat{v}_t(x_t)|^2 - 2\partial_t x_t \cdot \hat{v}_t(x_t)]$. Because $\partial_t I_t$ is known analytically from the interpolant definition, the expectation can be estimated by sampling $(x_0, x_1, t)$ directly. This avoids ODE integration during training (ODE integration is still required at inference).
• Decoupling Path and Optimization: The choice of path (interpolant) is separated from the optimization of the velocity field. MLE methods couple the path and objective.
• Connection to Score-Based Models: The authors show that for Gaussian base densities and trigonometric interpolants, the learned velocity field is explicitly related to the score function $\nabla \log \rho_t$, providing a theoretical bridge between CNFs and diffusion models.

What experiments were performed?

The authors performed validation across synthetic, tabular, and image domains:

• 2D Density Estimation: Benchmarked on “Checkerboard”, “8 Gaussians”, and anisotropic curved densities to visualize mode coverage and transport smoothness.
• High-Dimensional Tabular Data: Evaluated on standard benchmarks (POWER, GAS, HEPMASS, MINIBOONE, BSDS300) comparing Negative Log Likelihood (NLL) against FFJORD, OT-Flow, and others.
• Image Generation: Trained models on CIFAR-10 ($32 \times 32$), ImageNet ($32 \times 32$), and Oxford Flowers ($128 \times 128$) to test scalability.
• Ablations: Investigated optimizing the interpolant path itself (e.g., learning Fourier coefficients for the path) to approach optimal transport and minimize path length.

What outcomes/conclusions?

• Performance: The method matches or supersedes conventional ODE flows (like FFJORD) in terms of NLL while being significantly cheaper to train.
• Efficiency: The training cost per epoch is constant (simulation-free), whereas MLE-based ODE methods see growing costs as the dynamics become more complex.
• Scalability: The method successfully scales to $128 \times 128$ resolution on a single GPU, a resolution that prior ab-initio ODE flows had not demonstrated.
• Flexibility: The framework can connect any two arbitrary densities (e.g., connecting two different complex 2D distributions) without needing one to be Gaussian.
• Optimal Transport: For a fixed interpolant, minimizing $G(\hat{v})$ over the velocity field recovers the velocity for that specific path. Additionally optimizing over the interpolant family yields a solution to the Benamou-Brenier optimal transport problem.
• Limitations: The authors acknowledge that image FID scores trail dedicated diffusion models, noting that InterFlow was not optimized with standard training tricks such as exponential moving averages, truncation, or learning rate warm-ups. The framework’s sample quality could likely improve with these additions.

Reproducibility Details

Data

• Tabular Datasets: POWER (6D), GAS (8D), HEPMASS (21D), MINIBOONE (43D), BSDS300 (63D).
  • Training points range from ~30k (MINIBOONE) to ~1.6M (POWER).
• Image Datasets:
  • CIFAR-10 ($32 \times 32$, 50k training points).
  • ImageNet ($32 \times 32$, ~1.28M training points).
  • Oxford Flowers ($128 \times 128$, ~315k training points).
• Time Sampling: Time $t$ is sampled from a Beta distribution during training (reweighting) to focus learning near the target.

Algorithms

• Interpolant: The primary interpolant used is trigonometric: $I_t(x_0, x_1) = \cos(\frac{\pi t}{2})x_0 + \sin(\frac{\pi t}{2})x_1$.
  • Alternative linear interpolant: $I_t = a_t x_0 + b_t x_1$.
• Loss Function: $$G(\hat{v}) = \mathbb{E}_{t, x_0, x_1}[|\hat{v}_t(x_t)|^2 - 2\partial_t I_t(x_0, x_1) \cdot \hat{v}_t(x_t)]$$
  • The expectation is amenable to empirical estimation using batches of $x_0, x_1, t$.
• Sampling: Numerical integration using Dormand-Prince (Runge-Kutta 4/5).
• Optimization: SGD/Adam variants used for optimization.

Models

• Tabular Architectures:
  • Feed-forward networks with 4-5 hidden layers.
  • Hidden widths: 512 (POWER, GAS, HEPMASS, MINIBOONE) or 1024 (BSDS300).
  • Activation: ReLU (general) or ELU (BSDS300).
• Image Architectures:
  • U-Net based on the DDPM implementation.
  • Dimensions: 256 hidden dimension.
  • Sinusoidal time embeddings used.

Evaluation

• Metrics: Negative Log Likelihood (NLL) in nats (tabular) or bits per dim (images), Frechet Inception Distance (FID) for images.
• Baselines: FFJORD, Glow, Real NVP, OT-Flow, ScoreFlow, DDPM.

Tabular NLL (nats, lower is better; Table 2 Left):

Image Generation NLL and FID (Table 2 Right; NLL in bits per dim, lower is better):

Note: DDPM++ is from Song et al. (2021), the same work as ScoreSDE (it is the architecture optimized for VP/sub-VP SDEs). InterFlow matches ScoreSDE on CIFAR-10 NLL (2.99 bits per dim) while being simulation-free. FID is weaker than dedicated image models (10.27 vs 2.92 for ScoreSDE), reflecting the paper’s primary focus on tractable likelihood rather than sample quality.

Hardware

• Compute: All models were trained on a single NVIDIA A100 GPU.
• Training Time:
  • Tabular: $10^5$ steps.
  • Images: $1.5 \times 10^5$ to $6 \times 10^5$ steps.
  • Speedup: Demonstrated ~400x speedup compared to FFJORD on MiniBooNE dataset.

Artifacts

No official code release accompanies this paper. All tabular datasets (POWER, GAS, HEPMASS, MINIBOONE, BSDS300) are publicly available from prior work. CIFAR-10 and ImageNet are standard public benchmarks. Oxford Flowers 102 is also publicly available. Hyperparameters and architectures are fully specified in Tables 3 and 4 of the paper.

Paper Information

Citation: Albergo, M. S., & Vanden-Eijnden, E. (2023). Building Normalizing Flows with Stochastic Interpolants. The Eleventh International Conference on Learning Representations.

Publication: ICLR 2023

@inproceedings{albergoBuildingNormalizingFlows2022,
  title = {Building {{Normalizing Flows}} with {{Stochastic Interpolants}}},
  booktitle = {The {{Eleventh International Conference}} on {{Learning Representations}}},
  author = {Albergo, Michael Samuel and {Vanden-Eijnden}, Eric},
  year = 2023,
  url = {https://openreview.net/forum?id=li7qeBbCR1t}
}

Additional Resources:

• OpenReview
• arXiv

A natural question follows: “What if I use more samples? Won’t that make it better?”

Using more samples improves performance when paired with the correct objective function. Averaging the loss over $k$ samples yields minimal gains. Changing the objective itself is where the real gain comes from. This post explores the difference between a “multi-sample VAE” and the Importance Weighted Autoencoder (IWAE), a model that uses the same architecture as a VAE but is trained with a different objective that optimizes a tighter bound on the log-likelihood.

All ideas here are based on the fantastic paper: “Importance Weighted Autoencoders” by Burda, Grosse, and Salakhutdinov.

The Two Ways to Use $k$ Samples

Let’s say we have our encoder $q(h|x)$ and decoder $p(x,h)$. We decide to use $k=5$ samples instead of $k=1$. We have two main options for how to calculate our loss.

Option 1: The “Multi-Sample VAE” (The Naive Way)

This is the most straightforward idea. For each input $x$ in our batch:

1. Draw 5 samples ($h_1, …, h_5$) from $q(h|x)$.
2. Calculate the standard VAE $\mathcal{L}_1$ loss for each sample.
3. Average these 5 losses together.

This is an average of logs. As the IWAE paper shows experimentally, this approach gives you a more stable gradient, but the final performance (in terms of log-likelihood) is “only slightly” better. You’re paying a 5x computational cost for a marginal gain because you’re still optimizing the same “loose” $\mathcal{L}_1$ bound.

Option 2: The Importance Weighted Autoencoder (IWAE) (The Right Way)

The IWAE takes a different approach. For each input $x$:

1. Draw 5 samples ($h_1, …, h_5$) from $q(h|x)$.
2. Calculate an “importance weight” $w_i$ for each sample.
3. Average these 5 weights together.
4. Take the logarithm of that average.

This is a log of an average, and the difference matters.

The Math: Average-of-Logs vs. Log-of-Averages

Let’s make this concrete. The standard VAE $\mathcal{L}_1$ objective is:

$$ \mathcal{L}_1(x) = \mathbb{E} _{h\sim q(h|x)} \left[ \log \frac{p(x,h)}{q(h|x)} \right] $$

A multi-sample VAE simply gets a better estimate of this same value:

$$ \mathcal{L} _{\text{VAE}, k}(x) \approx \frac{1}{k} \sum _{i=1}^{k} \log w_i \quad \text{where} \quad w_i = \frac{p(x,h_i)}{q(h_i|x)} $$

The IWAE objective, $\mathcal{L}_k$, is fundamentally different:

$$ \mathcal{L} _k (x) = \mathbb{E} _{h_1..h_k \sim q(h|x)} \left[ \log \left( \frac{1}{k} \sum _{i=1}^{k} \frac{p(x,h_i)}{q(h_i|x)} \right) \right] $$

In practice, we estimate this with a single Monte Carlo sample (of $k$ latents):

$$ \mathcal{L} _k (x) \approx \log \left( \frac{1}{k} \sum _{i=1}^{k} w_i \right) $$

Because the logarithm is a concave function, Jensen’s Inequality tells us that the “log of an average” is always greater than or equal to the “average of logs.”

$$ \mathcal{L}_k(x) \ge \mathcal{L}_1(x) $$

This means the IWAE is optimizing a strictly tighter lower bound on the true log-likelihood of the data.

Why Does This “Log-of-Average” Matter?

This mathematical property provides two practical benefits.

1. Better Density Estimation

Because $\mathcal{L}_k$ is a tighter bound on the true $p(x)$, optimizing it pushes the model to learn a much better generative distribution. The paper shows that IWAEs achieve “significantly higher log-likelihoods” than VAEs.

2. Richer Latent Representations

This is the most interesting part. The standard VAE $\mathcal{L}_1$ objective “harshly penalizes” the model if its one sample $h$ is a poor explanation for $x$. This pressure forces the recognition network $q(h|x)$ to be “overly simplified” to avoid bad samples, which can lead to many latent dimensions becoming inactive (the paper reports the number of “active units” per model).

The IWAE objective is more flexible. It only needs one of the $k$ samples to be good. This “increased flexibility” allows the model to learn far more complex posterior distributions and “richer latent space representations.” The paper’s experiments confirm this, showing IWAEs learn to use many more “active units” in their latent space.

What This Looks Like in Code (PyTorch)

The implementation difference makes this crystal clear.

First, the “k-sample” trick: for a batch x of shape [B, D] and k=5 samples, we repeat x to get x_repeated of shape [B*k, D]. We do all our forward passes on this large tensor.

VAE (Multi-Sample, k > 1) Loss

Here, we can still use the analytical KL divergence, which is a big simplification.

# x_repeated has shape [B*k, 784]
# mu, logvar have shape [B*k, latent_dim]
# recon_x has shape [B*k, 784]

# recon_loss_all shape: [B*k]
recon_loss_all = F.binary_cross_entropy(recon_x, x_repeated, reduction='none').sum(dim=1)

# kl_loss_all shape: [B*k]
# We use the simple, analytical KL term!
kl_loss_all = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)

# total_loss_all shape: [B*k]
total_loss_all = recon_loss_all + kl_loss_all

# --- The Key Step ---
# Just average all B*k losses. This is the "average of logs".
loss = total_loss_all.mean()

IWAE (k > 1) Loss

Here, we must compute the exact log-probabilities of the specific samples we drew.

# Helper function to compute log-prob of a sample from a Gaussian
def log_prob_gaussian(sample, mu, logvar):
    const = -0.5 * sample.shape[-1] * torch.log(2 * torch.tensor(math.pi))
    log_det = -0.5 * torch.sum(logvar, dim=-1)
    log_exp = -0.5 * torch.sum((sample - mu)**2 / torch.exp(logvar), dim=-1)
    return const + log_det + log_exp

# --- Get the 3 log-prob components ---
# x_repeated, recon_x, z_samples, mu_repeated, logvar_repeated
# all have a first dimension of [B*k]

# 1. log p(x|h_i): Log-Reconstruction Probability
# log_p_x_given_h shape: [B*k]
log_p_x_given_h = -F.binary_cross_entropy(recon_x, x_repeated, reduction='none').sum(dim=1)

# 2. log p(h_i): Log-Prior Probability (under N(0, I))
# log_p_h shape: [B*k]
log_p_h = log_prob_gaussian(z_samples, 0.0, 0.0) # mu=0, logvar=0

# 3. log q(h_i|x): Log-Encoder Probability
# log_q_h_given_x shape: [B*k]
log_q_h_given_x = log_prob_gaussian(z_samples, mu_repeated, logvar_repeated)

# --- The Key Step ---
# Combine to get the log-importance-weight
# log_w shape: [B*k]
log_w = log_p_x_given_h + log_p_h - log_q_h_given_x

# Reshape to [B, k] to group samples by their original input
log_w_matrix = log_w.view(B, k) # B is original batch size

# --- Apply the IWAE Objective (Log-Sum-Exp Trick) ---
# This is the "log of the average"
# log( (1/k) * sum(exp(log_w)) ) = logsumexp(log_w) - log(k)
log_iwae_bound_per_x = torch.logsumexp(log_w_matrix, dim=1) - math.log(k)

# The objective is to MAXIMIZE this bound, so the loss is its negative
loss = -log_iwae_bound_per_x.mean()

The Critical Implementation Detail

Notice the key difference in the final step:

• VAE: loss = total_loss_all.mean() average of individual losses
• IWAE: loss = -torch.logsumexp(log_w_matrix, dim=1).mean() log of averaged weights

This seemingly small change implements the fundamental mathematical difference between optimizing an “average of logs” versus a “log of averages.”

When to Use Each Approach

The Computational Trade-off

Both approaches scale linearly with $k$. If you use $k=5$ samples, you’re doing roughly 5x the computation. The question is whether you get 5x the benefit.

For multi-sample VAEs, the answer is usually “no”. You get more stable gradients but only marginal performance improvements.

For IWAEs, the answer is often “yes”. You get meaningfully better log-likelihoods and richer latent representations that can be worth the computational cost.

Conclusion

The next time you use more samples with your VAE, switch to the IWAE objective to get the full benefit of the computational cost of $k > 1$.

The mathematical insight is simple but powerful: Jensen’s Inequality tells us that the “log of an average” is always greater than or equal to the “average of logs.” By optimizing this tighter bound, IWAEs achieve better density estimation and learn richer latent representations than standard VAEs.

The implementation requires computing exact log-probabilities to evaluate the specific samples. The result is a fundamentally more powerful model using the exact same architecture.

Want to dive deeper? Check out the original IWAE paper for experimental results and theoretical analysis, or explore my VAE tutorial for hands-on implementation details.

This is a Method paper that introduces the Importance Weighted Autoencoder (IWAE), a generative model that shares the same architecture as the Variational Autoencoder (VAE) but uses a different, tighter objective function. The key innovation is using importance weighting to derive a strictly tighter log-likelihood lower bound than the standard VAE objective.

What is the motivation?

The standard VAE has several limitations that motivated this work:

• Strong assumptions: VAEs typically assume the posterior distribution is simple (e.g., approximately factorial) and that its parameters can be easily approximated from observations.
• Simplified representations: The VAE objective can force models to learn overly simplified representations that underutilize the network’s full modeling capacity.
• Harsh penalization: The VAE objective harshly penalizes approximate posterior samples that are poor explanations for the data, which can be overly restrictive.
• Inactive units: VAEs tend to learn latent spaces with effective dimensions far below their capacity, where many latent units are ignored (a phenomenon later termed posterior collapse, where the approximate posterior collapses to the prior and conveys no information). The authors wanted to investigate whether a new objective could address this issue.

What is the novelty here?

The core novelty is the IWAE objective function, denoted as $\mathcal{L}_{k}$.

• VAE ($\mathcal{L}_{1}$ Bound): The standard VAE maximizes $\mathcal{L}(x)=\mathbb{E}_{q(h|x)}[\log\frac{p(x,h)}{q(h|x)}]$. This is equivalent to the new bound when $k=1$.

• IWAE ($\mathcal{L}_{k}$ Bound): The IWAE maximizes a tighter bound that uses $k$ samples drawn from the recognition model $q(h|x)$:

$$\mathcal{L}_{k}(x)=\mathbb{E}_{h_{1},…,h_{k}\sim q(h|x)}\left[\log\frac{1}{k}\sum_{i=1}^{k}\frac{p(x,h_{i})}{q(h_{i}|x)}\right]$$

• Tighter Bound: The authors prove that this bound is always tighter than or equal to the VAE bound ($\mathcal{L}_{k+1} \geq \mathcal{L}_{k}$) and that as $k$ approaches infinity, $\mathcal{L}_{k}$ approaches the true log-likelihood $\log p(x)$.

• Increased Flexibility: Using multiple samples gives the IWAE additional flexibility to learn generative models whose posterior distributions are complex and violate the VAE’s simplifying assumptions.

Key Concept: Averaging Inside vs. Outside the Log

A crucial distinction exists between how VAE and IWAE utilize $k$ samples. Understanding this difference explains why increasing $k$ in IWAE improves the bound. In VAE, it reduces variance.

VAE (Average of Logs):

For a VAE, using $k$ samples approximates:

$$\mathbb{E}\left[ \frac{1}{k} \sum_{i=1}^k \log w_i \right] \approx \text{ELBO}$$

where $w_i = p(x, h_i) / q(h_i | x)$. Increasing $k$ here only reduces the variance of the gradient estimator; the model still targets the same ELBO bound, so performance gains saturate quickly.

IWAE (Log of Average):

IWAE performs the averaging inside the logarithm:

$$\mathbb{E}\left[ \log \left( \frac{1}{k} \sum_{i=1}^k w_i \right) \right] = \mathcal{L}_k$$

By Jensen’s Inequality ($\log(\mathbb{E}[X]) \geq \mathbb{E}[\log(X)]$ for concave functions), this bound is mathematically guaranteed to be at least as tight as the VAE bound. Each increase in $k$ defines a new, strictly tighter lower bound on the log-likelihood.

Why This Matters for Gradients:

In IWAE, the gradient weights are normalized importance weights $\tilde{w}_i = w_i / \sum_j w_j$. This means “bad” samples (those with low $w_i$) contribute very little to the gradient update since they vanish from the weighted sum. VAE uses unweighted samples, so a single sample with extremely low probability produces a massive negative log value that can dominate the loss and harshly penalize the model. IWAE’s formulation allows the model to focus learning on the samples that explain the data well.

What experiments were performed?

The authors compared VAE and IWAE on density estimation tasks using the MNIST and Omniglot datasets. They evaluated two main network architectures: one with a single stochastic layer and another with two stochastic layers. The models were trained with varying numbers of importance samples ($k \in {1, 5, 50}$) to observe the effect on performance and latent space utilization. The primary metrics for evaluation were the test log-likelihood (estimated using 5000 samples) and the number of “active” latent units, which quantifies the richness of the learned representations.

What outcomes/conclusions?

• Better Performance: IWAE achieved higher log-likelihoods than VAEs across all configurations. On MNIST with two stochastic layers and $k=50$, IWAE reached $-82.90$ nats compared to $-84.78$ for VAE. On Omniglot, the best IWAE achieved $-103.38$ nats versus $-106.30$ for VAE. IWAE performance improved consistently with increasing $k$, while VAE performance benefited only slightly from using more samples ($k>1$).

• Richer Representations: In all experiments with $k>1$, IWAE learned more active latent dimensions than VAE, suggesting richer latent representations.

• Objective Drives Representation: The authors found that latent dimension inactivation is driven by the objective function. They demonstrated this through an “objective swap” experiment:

This experiment provides evidence that the objective function itself influences latent utilization:

• VAE → IWAE: A converged VAE model, when fine-tuned with the IWAE objective ($k=50$), gained 3 active units (19 → 22) and improved test NLL from 86.76 to 84.88.
• IWAE → VAE: A converged IWAE model fine-tuned with the VAE objective lost 2 active units (25 → 23) and worsened test NLL from 84.78 to 86.02.

These results strongly suggest that inactivation of latent dimensions is driven by the objective function rather than by optimization dynamics, initialization, or architecture. The authors note that optimization also plays a role, as the swap results do not exactly match training from scratch.

• Comparison to Other Models: On MNIST, the best IWAE ($-82.90$ nats) outperformed deep belief networks ($-84.55$ nats) and deep autoregressive networks ($-84.13$ nats), though DRAW ($-80.97$ nats), which exploits spatial structure, achieved better results. On Omniglot, the best IWAE ($-103.38$ nats) fell slightly behind RBMs trained with persistent contrastive divergence ($-100.46$ nats).

• Conclusion: IWAEs learn richer latent representations and achieve better generative performance than VAEs with equivalent architectures and training time.

Reproducibility Details

Data

• MNIST: $28 \times 28$ binarized handwritten digits (60,000 training / 10,000 test).
• Omniglot: $28 \times 28$ binarized handwritten characters from various alphabets (24,345 training / 8,070 test).
• Binarization: Dynamic sampling where binary values are sampled with expectations equal to the real pixel intensities (following Salakhutdinov & Murray, 2008).
• Fixed Binarization: Results on a fixed binarization of MNIST (Larochelle, 2011) confirm that IWAE outperforms VAE across preprocessing methods. It exhibits notably more overfitting compared to dynamic sampling.

Models

Two main network architectures were tested:

1. One stochastic layer (50 units) with two deterministic layers (200 units each).
2. Two stochastic layers (100 and 50 units). Between x and h1 were two deterministic layers with 200 units each. Between h1 and h2 were two deterministic layers with 100 units each.

• Activations: tanh for deterministic layers; exp applied to variance predictions to ensure positivity.
• Distributions: Gaussian latent layers with diagonal covariance; Bernoulli observation layer.
• Initialization: Glorot & Bengio (2010) heuristic.

Algorithms

• Optimizer: Adam ($\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-4}$).
• Batch Size: 20.
• Learning Rate Schedule: Annealed rate of $0.001 \cdot 10^{-i/7}$ for $3^i$ epochs (where $i=0…7$), totaling 3,280 passes over the data.
• Variance Control: A common concern with importance sampling is high variance. The authors prove that the Mean Absolute Deviation of their estimator is bounded by $2 + 2\delta$, where $\delta$ is the gap between the bound and true log-likelihood. As the bound tightens, variance remains controlled.
• Computational trick: In the basic IWAE implementation, both forward and backward passes must be done independently for each of the $k$ samples, so the cost scales linearly with $k$. However, the authors describe an optional optimization: stochastically approximate the gradient sum by sampling a single $\epsilon_i$ proportional to its normalized weight $\tilde{w}_i$, then computing only that one backward pass. This reduces the cost to $k$ forward passes and one backward pass. Since the backward pass costs roughly twice the forward pass, this yields approximately a 3x speedup for large $k$ at the cost of increased gradient variance.

Relationship to Reweighted Wake-Sleep (RWS): Both IWAE and Reweighted Wake-Sleep (Bornschein & Bengio, 2015) use importance-weighted samples and have closely related generative model updates. The key difference is that IWAE derives a single unified lower bound $\mathcal{L}_k$ and uses the reparameterization trick to train the recognition network jointly. RWS instead uses separate wake and sleep phases for the recognition network, which are not derived from $\mathcal{L}_k$.

Evaluation

1. Test Log-Likelihood: Primary measure of generative performance, estimated as the mean of $\mathcal{L}_{5000}$ (5000 samples) on the test set.
2. Active Units: To quantify latent space richness, the authors measured “active” latent dimensions. A unit $u$ was defined as active if its activity statistic $A_{u}=\text{Cov}_{x}(\mathbb{E}_{u\sim q(u|x)}[u])$ exceeded $10^{-2}$. The $10^{-2}$ threshold is justified by a bimodal distribution of the log activity statistic, showing clear separation between active and inactive units.

Hardware

• Hardware: GPU-based implementation using mini-batch replication to parallelize the $k$ samples. Specific GPU type and training times are not reported.

Artifacts

Paper Information

Citation: Burda, Y., Grosse, R., & Salakhutdinov, R. (2016). Importance Weighted Autoencoders. International Conference on Learning Representations (ICLR) 2016. https://arxiv.org/abs/1509.00519

Publication: ICLR 2016

@inproceedings{burda2016importance,
  title={Importance Weighted Autoencoders},
  author={Yuri Burda and Roger Grosse and Ruslan Salakhutdinov},
  booktitle={International Conference on Learning Representations},
  year={2016},
  url={https://arxiv.org/abs/1509.00519}
}

Additional Resources:

• ArXiv

This is a Method paper that introduces a generative mechanism (the VAE) and an optimization technique (the reparameterization trick), with formal theoretical derivation. The method, called the Auto-Encoding VB (AEVB) algorithm, leads to what we now know as the variational auto-encoder (VAE) when neural networks are used as the recognition model.

What is the motivation?

The authors address two central intractabilities in directed probabilistic models with continuous latent variables:

1. Intractable Posteriors: In models with continuous latent variables (like those with non-linear hidden layers), the true posterior $p_{\theta}(z|x)$ cannot be calculated analytically, preventing the use of standard EM algorithms.

2. Large Datasets: Sampling-based solutions like Monte Carlo EM (MCEM) require expensive sampling loops per datapoint. This makes them too slow for large datasets where batch optimization is too costly and efficient minibatch updates are required.

What is the novelty here?

The Reparameterization Trick (SGVB Estimator)

The core innovation is the Stochastic Gradient Variational Bayes (SGVB) estimator. The authors solve the high variance of standard gradient estimation by “reparameterizing” the random variable $\tilde{z}$.

They express $z$ as a deterministic function of the input $x$ and an auxiliary noise variable $\epsilon$:

$$\tilde{z} = g_{\phi}(\epsilon, x) \quad \text{with} \quad \epsilon \sim p(\epsilon)$$

• Mechanism: For a Gaussian posterior, $z = \mu + \sigma \odot \epsilon$ where $\epsilon \sim \mathcal{N}(0, I)$.
• Impact: This makes the Monte Carlo estimate differentiable with respect to the variational parameters $\phi$, allowing the variational lower bound to be optimized via standard stochastic gradient ascent (like SGD or Adagrad).

The AEVB Algorithm (The VAE)

The Auto-Encoding VB (AEVB) algorithm amortizes inference by learning a global recognition model (encoder) $q_{\phi}(z|x)$ jointly with the generative model (decoder) $p_{\theta}(x|z)$.

Objective Function: Maximize the variational lower bound $\mathcal{L}(\theta, \phi; x^{(i)})$:

$$\mathcal{L} \simeq -D_{KL}(q_\phi(z|x^{(i)}) | p_\theta(z)) + \frac{1}{L} \sum_{l=1}^L \log p_\theta(x^{(i)}|z^{(i,l)})$$

• First Term (Regularizer): Forces the approximate posterior to match the prior (integrated analytically for Gaussians).
• Second Term (Reconstruction Error): The expected negative reconstruction error (estimated via sampling).

This mirrors the standard auto-encoder objective, adding a variational regularizer.

What experiments were performed?

The method was benchmarked against the Wake-Sleep algorithm and Monte Carlo EM (MCEM) using the MNIST (digits) and Frey Face (continuous faces) datasets.

What outcomes/conclusions?

• Efficiency: AEVB converged faster and reached a better lower bound than Wake-Sleep (Figure 2). It scaled efficiently to the full MNIST dataset. MCEM’s per-datapoint sampling cost made it impractical at full dataset scale, so comparisons were limited to small subsets (Figure 3).

• Regularization: The KL-divergence term provided a regularizing effect, preventing overfitting while increasing latent dimensions ($N_z$).

• Manifold Learning: The model successfully learned smooth 2D latent manifolds (visualized in Appendix A), grouping similar digits/faces together.

Reproducibility Details

Data

Evaluation Data: For the marginal likelihood comparison (Figure 3), the paper used MNIST with $N_{\text{train}} = 100$ and $N_{\text{train}} = 5000$ to compare data efficiency (marginal log-likelihood vs. training samples seen) across algorithms. A smaller network (100 hidden units, 3 latent variables) was used for this comparison because the marginal likelihood estimator only works reliably in low-dimensional latent spaces.

Algorithms

• Algorithm: Stochastic gradient ascent with Adagrad (global stepsizes chosen from ${0.01, 0.02, 0.1}$).
• Regularization: The objective included a weight decay term corresponding to a prior $p(\theta)=\mathcal{N}(0,I)$.
• Minibatches: Size $M=100$ with $L=1$ sample per datapoint.
• Initialization: Parameters sampled from $\mathcal{N}(0, 0.01)$.

Models

The original VAE used simple Multi-Layered Perceptrons (MLPs):

• Symmetry: The encoder and decoder were symmetric, having an equal number of hidden units.
• Hidden Units: 500 units for MNIST, 200 for Frey Face (to prevent overfitting on the smaller dataset).
• Activations: Tanh activation functions for the hidden layers.
• Latent Space: Experimented with $N_z$ ranging from 2 to 200.
• Outputs:
  • MNIST: Bernoulli MLP (sigmoid output).
  • Frey Face: Gaussian MLP, with means constrained to $(0,1)$ via sigmoid.
• Encoder Architecture: For the Gaussian encoder, the mean $\mu$ and log-variance $\log(\sigma^2)$ are linear outputs from the shared hidden layer (they share the hidden layer weights and have separate output weights).
• Log-Variance: The encoder predicted $\log(\sigma^2)$ for numerical stability.

Evaluation

The paper distinguishes between two metrics:

• Variational Lower Bound: Used as the training objective (what the model optimizes).
• Marginal Likelihood: Used for final evaluation (Figure 3). The true marginal likelihood $p_\theta(x)$ was estimated using an Importance Sampling estimator constructed from samples drawn via Hybrid Monte Carlo (HMC), as detailed in Appendix D. This estimator uses: $p_{\theta}(x^{(i)}) \simeq (\frac{1}{L}\sum \frac{q(z)}{p(z)p(x|z)})^{-1}$.

This distinction is critical: the training metric (lower bound) differs from the evaluation metric (estimated marginal likelihood).

Hardware

• Hardware: Trained on a standard Intel Xeon CPU (approx. 40 GFLOPS); no GPUs were used.
• Training Time: Approximately 20-40 minutes per million training samples.

Key Implementation Details from Appendices

• Appendix A: Visualizations of 2D latent manifolds learned for MNIST and Frey Face datasets.
• Appendix B: Closed-form solution for the KL divergence of two Gaussians, essential for implementing the efficient version of the estimator (Equation 10).
• Appendix C: Exact MLP equations, including the use of tanh hidden layers and specific output layers for Bernoulli vs. Gaussian data. Includes specifications for Bernoulli MLPs (binary data) and Gaussian MLPs (real-valued data).
• Appendix D: Marginal likelihood estimation protocol using Hybrid Monte Carlo (HMC) and importance sampling for evaluation (Figure 3).

Paper Information

Citation: Diederik P. Kingma and Max Welling. “Auto-Encoding Variational Bayes.” arXiv:1312.6114 [stat.ML], 2013. https://doi.org/10.48550/arXiv.1312.6114

Publication: ICLR 2014 (arXiv preprint December 2013)

@misc{kingma2022autoencodingvariationalbayes,
      title={Auto-Encoding Variational Bayes},
      author={Diederik P Kingma and Max Welling},
      year={2013},
      eprint={1312.6114},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/1312.6114},
}

Additional Resources:

• Wikipedia: Variational Autoencoder - General overview
• OpenReview - Original peer review with author responses
• Modern VAE in PyTorch - Implementation tutorial on this site

This is a method paper that introduces a training approach for Variational Autoencoders (VAEs) to address fundamental limitations in their generative quality through improved prior learning.

What is the motivation?

The work is motivated by a critical limitation in Variational Autoencoders known as the “prior hole” problem, where the prior distribution p(z) fails to match the aggregate approximate posterior q(z). This mismatch leads to areas in the latent space with high density under the prior that don’t map to realistic data samples, resulting in poor generative quality compared to GANs and other generative models.

The figure above illustrates this mismatch. The blue dots represent where a trained encoder actually places data in the latent space (the aggregate posterior $q(z)$), which often forms complex, non-Gaussian shapes. The red dashed contours show the standard Gaussian prior $p(z) = \mathcal{N}(0, I)$, which assumes data is centered at the origin. When generating new samples, we draw from this prior, making it likely to sample from the empty “hole” where the decoder has never seen training data, producing unrealistic outputs.

A natural question arises: the prior $p(z)$ is used for sampling at inference time, so why does learning a better prior also improve likelihood (NLL)? The answer lies in the VAE objective. VAEs maximize the Evidence Lower Bound (ELBO):

$$ \log p(x) \geq \mathcal{L}_{\text{ELBO}}(x) = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{Reconstruction}} - \underbrace{\text{KL}(q(z|x) \parallel p(z))}_{\text{Regularization}} $$

The KL divergence term penalizes the mismatch between each data point’s approximate posterior $q(z|x)$ and the prior $p(z)$. When the prior is a simple Gaussian but the aggregate posterior forms a complex shape (as in the figure above), this KL term remains unnecessarily high for every data point.

By replacing the simple prior with a learned $p_{\text{NCP}}(z)$ that matches the aggregate posterior, the KL penalty decreases, tightening the ELBO and improving NLL. The learned prior thus provides a unified solution: better likelihood during training (tighter bound) and better sampling at inference (no “holes”).

The OpenReview discussion contains a significant theoretical debate regarding the paper’s core premise. Reviewers argued that the “prior hole” problem is actually a failure of the posterior to match the prior, or a failure of the encoder. The authors defended their approach by noting that even with a perfect posterior, a simple Normal prior might fail because the decoder lacks capacity to map a simple distribution to complex data without dropping modes. This justifies fixing the prior by making it learned and complex.

What is the novelty here?

The authors propose an energy-based model (EBM) prior that is trained using Noise Contrastive Estimation (NCE), which they term a Noise Contrastive Prior (NCP). The key innovations are:

• Two-Stage Training Process: First, a standard VAE is trained with a simple base prior. Then, the VAE weights are frozen and a binary classifier learns to distinguish between samples from the aggregate posterior q(z) and the base prior p(z).
• Reweighting Strategy: The core idea is to reweight a base prior distribution p(z) with a learned reweighting factor r(z) to make the resulting prior $p_{\text{NCP}}(z)$ better match the aggregate posterior q(z).
• NCE for EBM Training: The method frames EBM training as a binary classification task to avoid computationally expensive MCMC sampling.
• Scalability to Hierarchical Models: For hierarchical VAEs with multiple latent groups, the NCP approach can be applied independently and in parallel to each group’s conditional prior.

What experiments were performed?

The method was evaluated on several standard image generation benchmarks:

• MNIST (dynamically binarized): Likelihood evaluation on a controlled, small-latent-space task
• CIFAR-10: Standard computer vision benchmark for generative modeling
• CelebA 64x64: Applied to both standard VAE architectures and more advanced VAEs with GMM priors (RAE model)
• CelebA HQ 256x256: High-resolution face generation task

The hierarchical NVAE models used 30 latent groups for CIFAR-10 and CelebA-64, 20 groups for CelebA-HQ-256, and 10 groups of $4 \times 4$ latent variables for MNIST (deliberately small to enable reliable partition function estimation). The experiments compared FID scores, likelihood metrics, and qualitative sample quality between baseline VAEs and NCP-enhanced versions, with particular focus on hierarchical VAEs (NVAE).

What outcomes/conclusions?

The proposed NCP method demonstrated improvements in generative quality across evaluated datasets, with modest gains on standard VAEs and particularly large gains on hierarchical models like NVAE:

• CelebA-64: NCP improved FID scores from 48.12 to 41.28 for standard VAEs, and from 40.95 to 39.00 for RAE models with GMM priors.
• Hierarchical Models (NVAE): The impact was particularly pronounced on hierarchical VAEs:
  • CIFAR-10: FID improved from 51.71 to 24.08
  • CelebA-64: FID improved from 13.48 to 5.25, making it competitive with GANs
  • CelebA HQ 256x256: FID reduced from 40.26 to 24.79
• Likelihood Performance: On MNIST, NCP-VAE achieved 78.10 nats NLL vs. baseline NVAE’s 78.67 nats

On CIFAR-10 and CelebA-HQ-256, the concurrent VAEBM method (which forms an EBM on the data space rather than the latent space) outperforms NCP-VAE. However, the authors argue the two approaches are complementary: NCP-VAE targets the latent space while VAEBM operates in data space, and combining them could yield further gains. NCP-VAE also has the advantage of applicability to discrete data (e.g., binarized MNIST) and simpler setup since it only requires training binary classifiers rather than MCMC-based training and sampling.

The key conclusions are that two-stage training with noise contrastive estimation provides an effective framework for learning expressive energy-based priors that addresses the prior hole problem while scaling efficiently to hierarchical models.

Reproducibility Details

Algorithms

The Reweighting Mechanism

The core innovation is defining the NCP prior as $p_{\text{NCP}}(z) \propto p(z)r(z)$. The reweighting factor $r(z)$ is derived from the binary classifier $D(z)$ using the likelihood ratio trick:

$$ r(z) \approx \frac{D(z)}{1 - D(z)} $$

Here, $D(z)$ is the sigmoid output of the trained discriminator, representing the probability that sample $z$ came from the aggregate posterior $q(z)$ (“real”). For an optimal discriminator $D^*(z)$, this ratio exactly equals $\frac{q(z)}{p(z)}$, allowing the model to approximate the density ratio without explicit density estimation.

Hierarchical Architecture Strategy

For hierarchical models (like NVAE), the method trains $K$ binary classifiers in parallel (one for each latent group). Crucially, to ensure efficiency, the classifiers reuse the context feature $c(z_{<k})$ extracted by the frozen VAE’s prior network. This architectural choice provides significant computational savings.

Test-Time Sampling (Inference)

Since $p_{\text{NCP}}(z)$ is an energy-based model, direct sampling is impossible. The paper employs two methods to generate samples:

• Sampling-Importance-Resampling (SIR): Used for most results. It draws $M$ samples (e.g., $M=5000$) from the base prior $p(z)$ and resamples them based on weights $w^{(m)} = r(z^{(m)})$.
• Langevin Dynamics (LD): An iterative refinement method using the gradient of the energy function $E(z) = -\log r(z) - \log p(z)$.

Models

Decoder Architecture

For RGB datasets (CIFAR-10, CelebA), the output likelihood must be changed from Discretized Logistic (standard NVAE) to a Normal distribution. The authors note this change alone led to “significant improvements in the base model performance.” Using the standard NVAE decoder will result in a weaker baseline than reported.

Discriminator Architecture

The binary classifier uses a ResNet-style architecture with Squeeze-and-Excitation (SE) blocks:

• Activation: Swish
• Normalization: Batch Normalization
• Optimization: Adam with Cosine Annealing (learning rate: $10^{-3} \to 10^{-7}$)

The SE blocks help the model focus on channel-wise feature recalibration, which is important for distinguishing subtle differences between prior and aggregate posterior in high-dimensional latent spaces.

Hardware

The main paper is vague on training time, but the OpenReview rebuttal explicitly lists hardware costs:

• Hardware: NVIDIA Tesla V100 (32GB) GPUs
• Per-Discriminator Training: ~13 hours for 100 epochs
• Parallelization: Because latent groups are independent, all discriminators can train in parallel
• Total Cost (CelebA-64): ~8.1 GPU-days
• Comparison: The authors argue this is efficient compared to VDVAE, which requires ~560 GPU-days

Evaluation

Inference Speed vs. Quality Trade-off

Reviewers flagged that SIR sampling can be prohibitively slow. The authors clarified the exact trade-off:

The quality gain from 500 to 5,000 samples is modest (FID difference of 1.51) while inference time increases roughly 8x, suggesting $M=500$ may be a practical default.

Hyperparameters

• FID Calculation: 50,000 samples
• SIR Proposals: 5,000 samples (paper default) or 500 (practical)
• MNIST: Dynamically binarized version used for likelihood evaluation
• Optimizers: The study largely adopts hyperparameters from baseline papers (e.g., Lawson et al. for MNIST, Ghosh et al. for RAE)

Debugging Benchmark: 25-Gaussians

The supplement provides a toy experiment ideal for verifying a new implementation before running on expensive image datasets:

• Setup: Synthetic dataset of 25 2D-Gaussians arranged on a grid
• Target NLL: ~-0.954 nats (NCP) vs. ~-2.753 nats (Vanilla VAE)
• Success Criterion: Samples should avoid low-density regions between grid points. A standard VAE will generate samples in these “prior holes,” while a working NCP implementation should cleanly remove these artifacts.

Implementation Warnings

• SIR Failure Mode: If the learned prior $p_{\text{NCP}}$ deviates too far from the base prior, SIR sampling collapses (low effective sample size). The paper shows a strong correlation between the NCE classification loss and the effective sample size (Fig. 5(b)), indicating that SIR reliability depends on how well the base prior matches the aggregate posterior.
• Temperature Scaling: The qualitative images in the paper use reduced temperature for improved visual sharpness (Section 5.3). The FID tables do not specify a temperature, so results may or may not use $T=1.0$.

Data

The method was evaluated on several standard image generation benchmarks:

• MNIST (dynamically binarized): Likelihood evaluation on a controlled, small-latent-space task
• CIFAR-10: Standard computer vision benchmark for generative modeling (32x32 RGB images)
• CelebA 64x64: Face generation task with moderate resolution
• CelebA HQ 256x256: High-resolution face generation task

All datasets use standard train/test splits from the computer vision literature.

Additional Metrics

Beyond FID and NLL, the paper uses:

• Effective Sample Size (ESS): Validates reliability of the SIR sampling procedure
• Maximum Mean Discrepancy (MMD): Measures distance between aggregate posterior and NCP prior distributions

Paper Information

Citation: Aneja, J., Schwing, A. G., Kautz, J., & Vahdat, A. (2021). A contrastive learning approach for training variational autoencoder priors. Advances in Neural Information Processing Systems, 34, 29604-29616. https://proceedings.neurips.cc/paper/2021/hash/0496604c1d80f66fbeb963c12e570a26-Abstract.html

Publication: NeurIPS 2021

@inproceedings{aneja2021contrastive,
  title={A Contrastive Learning Approach for Training Variational Autoencoder Priors},
  author={Aneja, Jyoti and Schwing, Alexander G and Kautz, Jan and Vahdat, Arash},
  booktitle={Advances in Neural Information Processing Systems},
  volume={34},
  pages={29604--29616},
  year={2021}
}

Additional Resources:

• OpenReview Discussion
• Code Repository (Google Drive; link may become inaccessible over time)

A Variational Autoencoder (VAE) is a type of generative model, meaning its primary purpose is to learn the underlying structure of a dataset so it can generate new, similar data.

Whether the data is images, raw audio clips, or 2D graphs of drug-like molecules, a VAE aims to capture the essential features that define the data distribution. Once trained, it should be able to create entirely new samples that resemble the training data without simply copying specific examples.

Introduced by Kingma and Welling in 2013 (Auto-Encoding Variational Bayes, Paper), VAEs are used for:

• Generation: Creating new data (images, music, text).
• Dimensionality Reduction: Compressing data into a much smaller, meaningful representation (a “latent space”).
• Imputation: Intelligently filling in missing data (e.g., denoising images).

Importantly, they aim to provide a structured and continuous latent space, which allows for smooth interpolation between data points and meaningful manipulations of generated samples (think: optimization).

TL;DR: The Complete PyTorch Implementation

For those who just want the code, here is a complete, modern VAE implementation in PyTorch. It features softplus standard deviation parameterization for numerical stability and a custom training step that handles the ELBO loss correctly.

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class VAEOutput:
    z: torch.Tensor
    mu: torch.Tensor
    std: torch.Tensor
    x_recon: torch.Tensor
    loss: torch.Tensor
    loss_recon: torch.Tensor
    loss_kl: torch.Tensor

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=512, latent_dim=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_std = nn.Linear(hidden_dim, latent_dim)

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, input_dim)
        )

    def encode(self, x):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        # Softplus + epsilon for stable std deviation
        std = F.softplus(self.fc_std(h)) + 1e-6
        return mu, std

    def reparameterize(self, mu, std):
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x, kl_weight=1.0):
        mu, std = self.encode(x)
        z = self.reparameterize(mu, std)
        x_recon = self.decode(z)

        # 1. Reconstruction Loss (Binary Cross Entropy for MNIST)
        # Sum over features, mean over batch
        recon_loss = F.binary_cross_entropy_with_logits(x_recon, x, reduction='none').sum(dim=1).mean()

        # 2. KL Divergence
        # Analytic KL for Normal distributions
        kl_loss = -0.5 * torch.sum(1 + torch.log(std**2) - mu**2 - std**2, dim=1).mean()

        # 3. Total Loss (ELBO)
        loss = recon_loss + (kl_weight * kl_loss)

        return VAEOutput(z, mu, std, x_recon, loss, recon_loss, kl_loss)

# --- Training Loop Example ---
def train_step(model, batch, optimizer, kl_weight=1.0):
    model.train()
    optimizer.zero_grad()

    # Forward pass
    output = model(batch, kl_weight)

    # Backward pass
    output.loss.backward()

    # Gradient clipping (recommended)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()
    return output.loss.item()

The Core Idea: Learning to Generate

The VAE is built on a key assumption: our complex, high-dimensional data (like a $28 \times 28$ pixel image, $\mathbf{x}$) is actually generated by some simpler, low-dimensional, unobserved variable (a “latent” variable, $\mathbf{z}$).

A Physical Metaphor: Water Molecules and Phase Diagrams

Consider a glass of water. At the microscopic level, you have more than $10^{24}$ $\text{H}_2\text{O}$ molecules bouncing around in an incredibly high-dimensional space. Each molecule has position, velocity, and interactions with its neighbors, computationally intractable to track directly. Yet we can describe the macroscopic behavior of all these molecules using just two simple variables: temperature and pressure. These two dimensions create a “phase diagram” that tells us whether our water will be ice, liquid, or vapor. The temperature and pressure are “latent variables” that capture the essential physics governing this complex molecular dance.

A VAE makes the same assumption: complex data (like images) emerges from simpler underlying factors. A handwritten digit might be generated by latent factors like “pen thickness,” “writing angle,” “digit style,” and “size,” a much simpler description than tracking all 784 pixel values independently.

The VAE learns two functions: one that maps from complex data ($\mathbf{x}$) to these descriptive factors ($\mathbf{z}$), and another that maps from these factors back to the data. It accomplishes this with two main components, typically implemented as neural networks:

1. The Encoder (Recognition Model)

This network takes a complex data point $\mathbf{x}$ (an image) and determines the “knob settings” $\mathbf{z}$ that could explain or generate it. This allows us to compress or understand the data.

$$q_{\phi}(\mathbf{z} | \mathbf{x})$$

It’s like examining a container of molecules and summarizing their complex arrangement into key parameters like temperature and pressure.

Crucially, the encoder outputs the parameters of a probability distribution (a simple Gaussian) that describes $\mathbf{z}$.

For each input $\mathbf{x}$, the encoder network outputs:

• A vector of means, $\mathbf{\mu}$
• A vector of standard deviations, $\mathbf{\sigma}$

These parameters define our approximation $q_{\phi}(\mathbf{z} | \mathbf{x}) = \mathcal{N}(\mathbf{z} \mid \mathbf{\mu}, \mathbf{\sigma}^2\mathbf{I})$. We then sample from this distribution to get the $\mathbf{z}$ that we feed to the decoder. This probabilistic step is what forces the latent space to be continuous and structured. It forces similar inputs to map to nearby regions in latent space, enabling smooth interpolation and generation.

2. The Decoder (Generative Model)

This network learns the “generative process.” It takes a simple latent vector $\mathbf{z}$ and reconstructs the complex data $\mathbf{x}$. This allows us to generate new data by feeding it a random $\mathbf{z}$ and observing what image $\mathbf{x}$ it produces.

$$p_{\theta}(\mathbf{x} | \mathbf{z})$$

The decoder reverses the encoder: it takes the simple latent representation and “paints” the full, complex image from it. It’s like taking temperature and pressure values and producing a detailed arrangement of water molecules consistent with those conditions. The goal is to reproduce the exact input as closely as possible.

After training, we have two networks that can be used for a variety of purposes:

• Generation: If the latent space is well-structured, we can sample random $\mathbf{z}$ vectors from a simple distribution (like a standard normal) and feed them into the Decoder to generate new images. This is particularly useful for searching for data points with desired properties, like in drug discovery, where we might want to generate molecules with specific characteristics.
• Compression: The Encoder can compress complex data into a low-dimensional latent space, which can be useful for visualization or as a feature extractor for other tasks.

The “Variational” Problem

Calculating the true distribution of latent variables $p_{\theta}(\mathbf{z}|\mathbf{x})$ (the posterior) is mathematically intractable.

This intractability arises from Bayes’ theorem:

$$p_{\theta}(\mathbf{z} | \mathbf{x}) = \frac{p_{\theta}(\mathbf{x} | \mathbf{z}) p_{\theta}(\mathbf{z})}{p_{\theta}(\mathbf{x})}$$

Breaking down each component:

• $p_{\theta}(\mathbf{x} | \mathbf{z})$ is our decoder, which is straightforward to compute given our likelihood model.
• $p_{\theta}(\mathbf{z})$ is our prior over latent variables, typically a simple distribution like a standard normal, making it easy to compute.
• $p_{\theta}(\mathbf{x})$ is the marginal likelihood of the data. And here lies the problem. It requires integrating over all possible latent variables that could have generated $\mathbf{x}$: $$p_{\theta}(\mathbf{x}) = \int p_{\theta}(\mathbf{x} | \mathbf{z}) p_{\theta}(\mathbf{z}) d\mathbf{z}$$ It is the normalization factor that ensures the posterior is a valid probability distribution (i.e., sums to 1 over all $\mathbf{z}$).

This integral is intractable because it involves integrating over a high-dimensional latent space with a complex likelihood function. No closed-form solution exists, and numerical integration is computationally prohibitive.

This is where the “variational” approach provides the solution. We approximate the true posterior by learning an encoder, $q_{\phi}(\mathbf{z} | \mathbf{x})$, that serves as a variational approximation to this intractable true distribution. The VAE’s training process optimizes this approximation to be as accurate as possible, pushing this learned distribution closer to the true posterior.

The VAE Objective: A Balancing Act

To get these two networks (parameterized by $\theta$ and $\phi$) to work together, we train them jointly with a special loss function. This objective has two parts that balance two different goals:

1. Reconstruction Loss

$$E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})]$$

This term asks: “How well can we reconstruct our original image?” It forces the VAE to be good at its job. The process goes:

1. Take an input point $\mathbf{x}$.
2. Use the Encoder to get its latent representation $\mathbf{z} \sim q_{\phi}(\mathbf{z} | \mathbf{x})$.
3. Use the Decoder to generate a new image $\mathbf{x}’$ from $\mathbf{z}$, $\mathbf{x}’ \sim p_{\theta}(\mathbf{x} | \mathbf{z})$.
4. Compare $\mathbf{x}$ and $\mathbf{x}’$.

The reconstruction loss measures the difference between the original and the reconstructed image.

• For continuous inputs (like general images), this is often Mean Squared Error (MSE).
• For inputs in $[0, 1]$ (like MNIST pixel intensities, which after ToTensor() are continuous values in $[0, 1]$), we use Binary Cross-Entropy (BCE). We treat each pixel as an independent Bernoulli variable whose target is its intensity. The decoder outputs the logits for each pixel, and the BCE-with-logits loss (e.g., F.binary_cross_entropy_with_logits) is the numerically stable way to compute the negative log-likelihood.
• More generally, you can output parameters of a desired output distribution. What if you wanted a mixture of Gaussians? The decoder could output the means, variances, and mixture weights, and you could compute the negative log-likelihood accordingly.

This loss pushes the encoder to produce useful $\mathbf{z}$ vectors and pushes the decoder to learn how to interpret them accurately.

2. The KL Divergence (The Regularizer)

$$D_{KL}(q_{\phi}(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

On its own, the reconstruction loss might “cheat.” The encoder could learn to map every image to a different, specific point in the latent space, essentially “memorizing” the data. While this minimizes reconstruction error, it creates a meaningless latent space that fails at generation.

The KL divergence term fixes this. It’s a regularizer that forces the latent space to be organized and smooth.

We force the encoder’s output, $q_{\phi}(\mathbf{z} | \mathbf{x})$, to be close to a simple, predefined prior distribution, $p_{\theta}(\mathbf{z})$. This prior is almost always a standard normal distribution because it is mathematically convenient, easy to sample from, and encourages a well-behaved latent space.

This regularization term acts as a penalty, measuring how much the encoder’s output distribution diverges from the simple prior. By minimizing this KL divergence, we encourage the model to:

• Avoid overfitting by preventing the encoder from memorizing specific locations for each input
• Create meaningful clusters where similar inputs map to nearby regions in the latent space
• Maintain continuity so that points close together in latent space (like different variations of the digit “7”) decode into visually similar outputs

This smooth, structured latent space is what enables generation: we can sample random points from our prior distribution and decode them into realistic new data.

Ultimately, the optimizer finds a balance between these two objectives: reconstructing the data well while keeping the latent space organized and regularized.

The Reparameterization Trick: Making it All Trainable

We have a problem. The training process requires sampling:

1. Encoder produces $\mathbf{\mu}$ and $\mathbf{\sigma}$.
2. We sample $\mathbf{z} \sim \mathcal{N}(\mathbf{\mu}, \mathbf{\sigma}^2\mathbf{I})$.
3. Decoder uses $\mathbf{z}$ to reconstruct $\mathbf{x}’$.
4. We calculate the loss.

The “sampling” step is a random, non-differentiable operation. We can’t backpropagate the reconstruction loss from the decoder through this random node to update the encoder’s weights.

The reparameterization trick makes the sampling process differentiable. We generate $\mathbf{z}$ deterministically by sampling a random noise vector and transforming it:

1. Sample a random noise vector $\mathbf{\epsilon}$ from a simple, fixed distribution (e.g., the standard normal $\mathcal{N}(\mathbf{0}, \mathbf{I})$).
2. Compute $\mathbf{z}$ as: $\mathbf{z} = \mathbf{\mu} + \mathbf{\sigma} \odot \mathbf{\epsilon}$

This simple change moves the randomness “outside” the network. The gradient can now flow deterministically from $\mathbf{z}$ back through the $\mathbf{\mu}$ and $\mathbf{\sigma}$ nodes to the encoder network. This is the key engineering insight that allows us to train the entire model end-to-end with standard backpropagation.

Where Does This Objective Come From? (The Math)

This two-part loss function is derived directly from the goal of maximizing the marginal likelihood of the data, $\log p_{\theta}(\mathbf{x})$.

For a single data point $\mathbf{x}^{(i)}$, we can write: $$\log p_{\theta}(\mathbf{x}^{(i)}) = D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}^{(i)}) || p_{\theta}(\mathbf{z} | \mathbf{x}^{(i)})) + \mathcal{L}(\theta, \phi; \mathbf{x}^{(i)})$$

• The first term is the KL divergence between our encoder’s approximation and the (intractable) true posterior. This is non-negative, and unfortunately we cannot compute it.
• The second term, $\mathcal{L}$, is the Variational Lower Bound (also known as the Evidence Lower Bound, or ELBO). Since the KL term is $\ge 0$, we know that $\log p_{\theta}(\mathbf{x}^{(i)}) \ge \mathcal{L}$.

By maximizing this lower bound $\mathcal{L}$, we push up the “floor” on the true likelihood of our data. This is a problem we can solve.

When we expand this $\mathcal{L}$ term, we get our famous two-part objective:

$$\mathcal{L}(\theta, \phi; \mathbf{x}^{(i)}) = E_{q_{\phi}(\mathbf{z} | \mathbf{x}^{(i)})}[\log p_{\theta}(\mathbf{x}^{(i)} | \mathbf{z})] - D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}^{(i)}) || p_{\theta}(\mathbf{z}))$$

• Term 1: The expected log-likelihood of reconstructing $\mathbf{x}^{(i)}$ from $\mathbf{z}$. Maximizing this is the same as minimizing the Reconstruction Loss.
• Term 2: The negative KL divergence between our encoder and the simple prior. Maximizing this is the same as minimizing the KL Divergence Loss.

Thus, the VAE’s objective balances these two critical goals: faithfully reconstructing the data while maintaining a simple, regularized latent structure that is useful for generation.

From ELBO to Practical Loss

Remember, our goal is to maximize the ELBO:

$$\mathcal{L}(\theta, \phi; \mathbf{x}) = E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})] - D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

Since deep learning libraries are built to minimize a loss function, we simply flip the sign and minimize the negative ELBO ($-\mathcal{L}$).

This gives us our final, practical loss function:

$$\text{Loss} = -\mathcal{L} = -E_{q_{\phi}(\mathbf{z} | \mathbf{x})}[\log p_{\theta}(\mathbf{x} | \mathbf{z})] + D_{KL}(q_\phi(\mathbf{z} | \mathbf{x}) || p_{\theta}(\mathbf{z}))$$

This is the function you actually implement. Minimizing this loss achieves both of our goals:

1. It minimizes the Reconstruction Loss (which is the same as maximizing the log-likelihood).
2. It minimizes the KL Divergence, forcing the encoder to match the prior.

Modern PyTorch VAE Implementation

Now that we understand the VAE architecture and objective, let’s implement a modern VAE in PyTorch. I’ll focus primarily on the model and loss function here, though the full code is available on GitHub.

My VAE implementation uses an output dataclass and a VAE class extending nn.Module.

"""Variational Autoencoder (VAE) model implementation."""

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


def get_activation(activation: str) -> nn.Module:
    """Get activation function by name."""
    activation_lower = activation.lower()
    ACTIVATION_MAP = {
        "relu": nn.ReLU(),
        "tanh": nn.Tanh(),
        "sigmoid": nn.Sigmoid(),
        "leaky_relu": nn.LeakyReLU(),
        "elu": nn.ELU(),
        "gelu": nn.GELU(),
    }
    if activation_lower not in ACTIVATION_MAP:
        supported = ", ".join(ACTIVATION_MAP.keys())
        raise ValueError(
            f"Unsupported activation '{activation}'. Supported: {supported}"
        )
    return ACTIVATION_MAP[activation_lower]


@dataclass
class VAEConfig:
    """VAE model configuration specifying architecture and behavior."""

    hidden_dim: int
    latent_dim: int

    input_shape: tuple[int, int, int] = (1, 28, 28)  # Default: MNIST
    activation: str = "tanh"  # Default: tanh, what was used in the original VAE paper
    use_softplus_std: bool = False  # Whether to use softplus for std parameterization
    n_samples: int = 1  # Number of latent samples per input during training


@dataclass
class VAEOutput:
    """VAE forward pass output containing all relevant tensors and optional losses."""

    x_logits: torch.Tensor
    z: torch.Tensor
    mu: torch.Tensor
    std: torch.Tensor

    x_recon: torch.Tensor | None = None
    loss: torch.Tensor | None = None
    loss_recon: torch.Tensor | None = None
    loss_kl: torch.Tensor | None = None


class VAE(nn.Module):
    """Variational Autoencoder with support for deterministic and probabilistic reconstruction."""

    DEFAULT_EPS = 1e-8

    def __init__(self, config: VAEConfig) -> None:
        """Initialize VAE with given configuration.

        Args:
            config: VAE configuration specifying architecture and behavior
        """
        super().__init__()
        self.config = config

        # Build encoder: input -> hidden -> latent parameters (mu, sigma)
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(
                int(torch.prod(torch.tensor(config.input_shape))), config.hidden_dim
            ),
            get_activation(config.activation),
            nn.Linear(config.hidden_dim, config.latent_dim * 2),
        )

        # Build decoder: latent -> hidden -> reconstructed input
        self.decoder = nn.Sequential(
            nn.Linear(config.latent_dim, config.hidden_dim),
            get_activation(config.activation),
            nn.Linear(
                config.hidden_dim, int(torch.prod(torch.tensor(config.input_shape)))
            ),
            nn.Unflatten(1, config.input_shape),
        )

    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Encode input to latent distribution parameters."""
        encoder_output = self.encoder(x)
        mu, sigma = torch.chunk(encoder_output, 2, dim=-1)
        return mu, sigma

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """Decode latent representation to reconstruction logits"""
        return self.decoder(z)

    def reparameterize(self, mu: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
        """Apply reparameterization trick for differentiable sampling."""
        epsilon = torch.randn_like(std)
        return mu + std * epsilon

    def forward(
        self,
        x: torch.Tensor,
        compute_loss: bool = True,
        reconstruct: bool = False,
        eps: float = DEFAULT_EPS,
    ) -> VAEOutput:
        """Forward pass through the VAE.

        Args:
            x: Input tensor of shape (batch_size, *input_shape)
            compute_loss: Whether to compute VAE loss components
            reconstruct: Whether to return reconstructions or distributions
            eps: Small epsilon value for numerical stability

        Returns:
            VAEOutput containing all relevant tensors and optionally computed losses
        """
        # Prepare input for multiple sampling if needed
        x_expanded = self._expand_for_sampling(x) if self.config.n_samples > 1 else x

        # Encode and sample from latent space
        mu, sigma = self.encode(x)
        std = self._sigma_to_std(sigma, eps=eps)
        mu_expanded, std_expanded = self._expand_latent_params(mu, std)
        z = self.reparameterize(mu_expanded, std_expanded)

        # Decode latent samples
        x_logits = self.decode(z)

        # Create output object
        output = VAEOutput(
            x_logits=x_logits,
            z=z,
            mu=mu,
            std=std,
            x_recon=torch.sigmoid(x_logits) if reconstruct else None,
        )

        # Compute losses if requested
        if compute_loss:
            loss, loss_recon, loss_kl = self._compute_loss(
                x_expanded, x_logits, mu, sigma, std
            )
            output.loss = loss
            output.loss_recon = loss_recon
            output.loss_kl = loss_kl

        return output

    # ==================== Helper Methods ====================

    def _sigma_to_std(
        self, sigma: torch.Tensor, eps: float = DEFAULT_EPS
    ) -> torch.Tensor:
        """Convert sigma parameter to standard deviation."""
        if self.config.use_softplus_std:
            return F.softplus(sigma) + eps
        else:
            return torch.exp(0.5 * sigma)  # sigma represents log-variance

    def _expand_for_sampling(self, x: torch.Tensor) -> torch.Tensor:
        """Expand input tensor for multiple sampling."""
        shape_dims = [1] * len(self.config.input_shape)
        x_expanded = x.unsqueeze(1).repeat(1, self.config.n_samples, *shape_dims)
        return x_expanded.view(-1, *self.config.input_shape)

    def _expand_latent_params(
        self, mu: torch.Tensor, std: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Expand latent parameters for multiple sampling."""
        if self.config.n_samples == 1:
            return mu, std

        mu_expanded = (
            mu.unsqueeze(1)
            .repeat(1, self.config.n_samples, 1)
            .view(-1, self.config.latent_dim)
        )
        std_expanded = (
            std.unsqueeze(1)
            .repeat(1, self.config.n_samples, 1)
            .view(-1, self.config.latent_dim)
        )

        return mu_expanded, std_expanded

    # ==================== Loss Computation ====================

    def _compute_loss(
        self,
        x: torch.Tensor,
        x_logits: torch.Tensor,
        mu: torch.Tensor,
        sigma: torch.Tensor,
        std: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute VAE loss components for deterministic reconstruction."""
        loss_recon = self._compute_reconstruction_loss(x, x_logits)
        loss_kl = self._compute_kl_loss(mu, sigma, std)
        return loss_recon + loss_kl, loss_recon, loss_kl

    def _compute_reconstruction_loss(
        self, x: torch.Tensor, x_logits: torch.Tensor
    ) -> torch.Tensor:
        """Compute reconstruction loss using binary cross-entropy."""
        return F.binary_cross_entropy_with_logits(
            x_logits, x, reduction="sum"
        ) / x.size(0)

    def _compute_kl_loss(
        self,
        mu: torch.Tensor,
        sigma: torch.Tensor,
        std: torch.Tensor,
        eps: float = DEFAULT_EPS,
    ) -> torch.Tensor:
        """Compute KL divergence between latent distribution and standard normal prior."""
        # Analytical KL: KL(N(μ,σ²) || N(0,1)) = 0.5 * Σ(μ² + σ² - 1 - log(σ²))
        if self.config.use_softplus_std:
            # sigma is just the raw output, need to use std directly: σ
            kl_per_sample = 0.5 * torch.sum(
                mu.pow(2) + std.pow(2) - 1 - torch.log(std.pow(2) + eps), dim=1
            )
        else:
            # sigma represents log-variance parameterization: log(σ²)
            kl_per_sample = 0.5 * torch.sum(mu.pow(2) + sigma.exp() - 1 - sigma, dim=1)

        return kl_per_sample.mean()

Loss Scaling

Both components of the VAE loss should be summed over data dimensions and averaged over the batch size. A common mistake is using the reduction="mean" option in PyTorch loss functions, which averages over all elements in the tensor.

• The KL Divergence (loss_kl) is a penalization term. Each dimension of the latent space has the potential to add complexity and deviate from the prior. As you increase the latent dimensionality, you typically see the KL loss increase in magnitude. That’s the cost of having a more expressive latent space.
• The Reconstruction Loss (loss_recon) measures how well the model reconstructs the input data, and it should scale with input dimensionality (this can bias the model toward better reconstruction for higher-dimensional data).

In the case of MNIST, if we used reduction="mean" for BCE, it would be averaged over all $784 \times \text{batch size}$ pixels, making it tiny compared to the KL loss. The KL term would dominate, and the model would learn to ignore the input, potentially leading to posterior collapse.

While modern optimizers can handle a variety of scenarios and you can still learn effective models with imperfect scaling, the original VAE paper used the scaling described above, and I recommend following that convention.

Mitigating Posterior Collapse: KL Annealing/Warmup

One common issue in training VAEs, especially with powerful decoders (like RNNs or deep CNNs), is posterior collapse. This happens when the KL term dominates the loss early in training. The model quickly learns to just output the prior distribution ($q(z|x) \approx p(z)$) to drive the KL loss to zero, effectively ignoring the latent code $z$. The decoder then becomes a powerful autoregressive model that ignores the latent input.

To prevent this, we often use KL Annealing (or Warmup). We introduce a weight $\beta$ for the KL term that starts at 0 and slowly increases to 1 over the first $N$ steps or epochs.

$$ \mathcal{L} = \mathcal{L}_{recon} + \beta \cdot D_{KL} $$

This allows the model to focus purely on reconstruction first (using the full latent capacity), and then slowly adds the regularization pressure.

# Simple Linear Annealing Scheduler
def get_kl_weight(step, total_steps, max_val=1.0):
    val = (step / total_steps) * max_val
    return min(max(val, 0.0), max_val)

# In your training loop:
for epoch in range(epochs):
    beta = get_kl_weight(epoch, warmup_epochs)
    loss = recon_loss + beta * kl_loss

Parameterizing Standard Deviation

    def _sigma_to_std(self, sigma: torch.Tensor, eps: float = DEFAULT_EPS) -> torch.Tensor:
        """Convert sigma parameter to standard deviation."""
        if self.config.bound_std is not None:
            return torch.sigmoid(sigma) * self.config.bound_std + eps
        elif self.config.use_softplus_std:
            return F.softplus(sigma) + eps
        else:
            return torch.exp(0.5 * sigma)  # sigma represents log-variance

Parameterizing the mean of the latent distribution is straightforward since $\mu \in \mathbb{R}$. However, the standard deviation $\sigma$ must be strictly positive (as must the variance $\sigma^2$). This type of constrained optimization is challenging for neural networks.

Log-Variance One common approach is to have the network output the log-variance ($\log \sigma^2$). This is what the original VAE paper did. The idea is to allow the network to output any real number and treat that value as the log-variance, $s = \log \sigma^2$. We can then compute the standard deviation as $\sigma = \exp(0.5 s)$, which is always positive.

The KL divergence formula simplifies nicely with this parameterization:

$$ \text{KL}( \mathcal{N}(\mu, \sigma^2) || \mathcal{N}(0, 1) ) = \frac{1}{2} \sum_{i=1}^d (\mu_i^2 + \sigma_i^2 - 1 - \log \sigma_i^2) $$

0.5 * torch.sum(mu.pow(2) + s.exp() - 1 - s, dim=1)

Softplus Standard Deviation An alternative is to have the network output $\sigma$ directly. This must be handled with care to ensure positivity. Strictly positive activations like softplus are required. Activations like ReLU can output zero, leading to numerical instability (during training) and deterministic behavior (during sampling). Additionally, adding a small epsilon value ensures numerical stability by preventing $\sigma$ from being exactly zero.

The KL divergence formula becomes slightly more complex:

0.5 * torch.sum(
    mu.pow(2) + std.pow(2) - 1 - torch.log(std.pow(2) + eps), dim=1
)

Bounded Standard Deviation Another option is to bound the standard deviation to a maximum value using a sigmoid transformation (or similar). This replaces mapping to $(0, \infty)$ with mapping to $(0, \text{bound})$. This helps prevent extremely high variance values that might destabilize training, while limiting the expressiveness of the latent distribution. Like with softplus, adding a small epsilon ensures numerical stability by preventing $\sigma$ from being exactly zero or approaching it too closely.

Gradient Behavior All parameterizations can work well in practice and have different gradient behaviors. Think of $g(s)$ as a transformation function from the network output to the proper domain of $\sigma$ (or $\sigma^2$); in the log-variance case, $g(s) = \exp(s)$, while in the softplus case, $g(s) = \text{softplus}(s) + \epsilon$.

The gradient of the loss with respect to these outputs can be written using the chain rule:

$$ \frac{\partial \mathcal{L}}{\partial s} = \frac{\partial \mathcal{L}}{\partial \sigma} \cdot \frac{\partial g(s)}{\partial s} $$

where $\frac{\partial g(s)}{\partial s}$ is the derivative of the transformation function.

We need to guard against two pathological cases:

• $\frac{\partial g(s)}{\partial s} \rightarrow 0$: This leads to vanishing gradients, making it hard for the network to learn.
• $\frac{\partial g(s)}{\partial s} \rightarrow \infty$: This leads to exploding gradients, causing instability during training and potentially divergence.

The log-variance parameterization, with its exponential transformation that is its own derivative, exhibits both issues at extreme values. If $s \rightarrow -\infty$, then $\sigma \rightarrow 0$ and the gradient vanishes. If $s \rightarrow \infty$, then $\sigma \rightarrow \infty$ and the gradient explodes. Since the interval $(0, 1)$ is mapped to $(-\infty, 0)$ in log-space, it’s much more difficult for the network to drive $\sigma$ to small values. In practice, exploding gradients at high values have been more problematic in my experience. Gradient clipping, learning rate scheduling, and clamping the log-variance output to a maximum value can help mitigate this.

What about softplus? The derivative of softplus is the sigmoid function, which smoothly maps $(-\infty, \infty)$ to $(0, 1)$. Gradients are always bounded by unity, preventing explosion (barring explosion from other parts of the network). However, as $s \rightarrow -\infty$, the gradient approaches zero, leading to vanishing gradients. Adding a small epsilon helps mitigate this, ensuring that $\sigma$ never gets too close to zero. Nonetheless, learning can still slow down.

For bounded standard deviation, the derivative of the sigmoid function is also bounded, preventing exploding gradients. (The gradient of sigmoid is defined in terms of itself: $\text{sig}’(x) = \text{sig}(x)(1 - \text{sig}(x))$; its maximum value is $0.25$ at $x=0$.)

Experiments

2D MNIST VAE with Different Std. Dev. Parameterizations

First, let’s run an experiment that is close to what was done in the original VAE paper. We’ll use MNIST as our dataset, a simple feedforward architecture with tanh activations, and the log-variance parameterization for the latent distribution.

Some of the differences from the original paper include:

• Using a hidden size of 512 (the original used 500)
• Using the AdamW optimizer (the original used vanilla Adagrad)
• Applying similar weight decay, doing so quite differently due to the optimizer change
• Focusing primarily on 2D latent spaces (for now)

This results in a network with 807,700 parameters. I train each model for 150 epochs at most and highlight the best based on the reconstruction loss on the test set. Just for fun, I sweep across different standard deviation parameterizations and learning rate warmup strategies.

From this summary table, all three parameterizations work well. The differences in final loss values are quite small. This could be due to the simplicity of the dataset and model architecture, further amplified by forcing the network to compress images into a very low-dimensional latent space (2D).

Since the softplus parameterization with learning rate warmup achieved the best reconstruction loss, let’s visualize some of its training dynamics and results more closely.

Loss Dynamics

To understand the VAE’s behavior, we must look at the ELBO and its two components: the Reconstruction Loss and the KL Divergence.

These plots reveal a clear narrative:

1. Rapid Initial Learning: Performance skyrockets in the first ~15 epochs.
2. Overfitting: The Reconstruction Loss (middle) flatlines for the test set while continuing to improve for training, a classic sign of memorization.
3. The Balancing Act: The KL Divergence (bottom) initially rises (“The Cost of Learning”) as the model stretches the latent space to encode digits, then saturates.
4. Equilibrium: The total ELBO (top) improves slowly, driven by the model finding the optimal trade-off between reconstruction and regularization. Notice that Test and Train KL tracks closely: a sign of good regularization!

Visualizing the VAE Trade-Off: BCE vs. KL

While the line plots visualize progress over time, they miss the evolving relationship between our two competing objectives.

A VAE is fundamentally a multi-objective optimization problem. We want to:

1. Minimize Reconstruction Loss (BCE)
2. Minimize KL Divergence

Combining them as the ELBO is common and effective, though it can mask some of the underlying dynamics.

These two goals are in direct conflict. To get perfect reconstruction (BCE = 0), the encoder would need to “memorize” each input, mapping it to a unique, precise point in latent space. This would cause the KL divergence to skyrocket, as these specific, “pointy” distributions are nothing like our smooth N(0, 1) prior.

Conversely, to get perfect KL divergence (KL = 0), the encoder must always output N(0, 1), regardless of the input. This perfectly matches the prior. Since the latent code $\mathbf{z}$ now contains zero information about the input $\mathbf{x}$, the decoder can only learn to output the “average” image, resulting in terrible reconstruction.

The training process is a search for the best compromise.

This plot shows the test set’s BCE (y-axis) vs. KL Divergence (x-axis) at every evaluation step. The color gradient from cool (blue) to warm (red) represents the training progress from Epoch 0 to 150.

Here’s how to interpret this training path:

1. The Start (Green Diamond, ~Epoch 0): The model starts at the top-left.

   • High BCE (Reconstruction): The decoder is random and hasn’t learned to reconstruct anything. Reconstruction is terrible.
   • Low KL Divergence: The encoder is also random. Its output distributions $q_{\phi}(\mathbf{z} | \mathbf{x})$ are a random mess. On average, this “mess” is coincidentally close to the “mess” of the prior $p_{\theta}(\mathbf{z})$, so the KL penalty is low. The model isn’t encoding any useful information yet, so it’s not paying a high price for it.

2. Phase 1: The Initial Plunge (Blue Path): The path moves almost straight down.

   • BCE Plummets: The model’s first and easiest task is to learn to reconstruct something. The optimizer finds massive, easy gains by making the decoder output “blurry digits” to replace the initial noise.
   • KL Stays Low: The model achieves this huge reconstruction win without needing to learn a very complex latent space. It’s the “low-hanging fruit” of training.

3. Phase 2: The Trade-Off (The “Elbow”): The path stops dropping vertically and starts moving to the right and down.

   • “Spending” KL to “Buy” Reconstruction: This is the true VAE trade-off in action. The easy wins are gone. To make the reconstructions sharper and more accurate (lowering BCE further), the model must now learn a more complex, informative latent representation.
   • It “stretches” the latent distributions $q_{\phi}(\mathbf{z} | \mathbf{x})$ to encode more details about each specific digit. This “stretching” moves it further from the simple N(0, 1) prior, and the KL divergence (the “cost”) goes up.

4. The End Game (Red Path & Star): The path settles in the bottom-right corner.

   • Finding the “Elbow”: The model finds an equilibrium. It has pushed the KL divergence as high as it’s “worth” for the reconstruction gains it gets. Trying to get even better reconstruction (moving further down) would cost an enormous, disproportionate amount in KL divergence (moving far to the right), and the total loss would increase.
   • Best Recon (Orange Star): The best reconstruction model (Epoch 118) is found right at this “elbow,” representing the best-found balance point on the trade-off frontier.

This single plot visualizes the entire training dynamic as a journey along the Pareto frontier: the set of optimal solutions where you can’t improve one objective (BCE) without worsening the other (KL).

Generative Performance

Let’s take a look at how well this model can decode samples.

Reconstruction Performance

Immediately, we see a couple of key points:

• Reconstructions are quite blurry compared to the originals. This is expected given the low capacity of the model and the extreme compression into a 2D latent space. General structure is typically preserved, while fine details are lost.
• The network struggles with 4s and 9s, often mixing them up or producing ambiguous shapes. This is a common failure mode in MNIST models due to the similarity of these digits.

Sampling from the Prior

If we sample from the prior N(0, 1) and decode those samples, we get a variety of digit-like images. From this, we get a pretty rich representation of digits. Almost all digits appear to be featured in this random sampling. Again, we see the standard blurriness.

Sweeping the Latent Space

We can select two points at random (here, two zeros), embed them into our latent space and then walk across that latent space to interpolate between two data points. Here, we see a walk that takes us from a zero that is askew to one that is more upright.

Finally, we can also sweep each latent dimension independently to see how they affect the generated images.

1. Sweeping z_1 (top row), we see a 5 become an 8 and then a 9. The slant shifts from left to right as we sweep the dimension.
2. Sweeping z_2 (bottom row), we see a 4 become a 9 and then an 8. Then it becomes a 3, a 2, some nonsense, and a 6.

So clearly each latent dimension is encoding some high-level features of the digits, and we can manipulate those features by moving in latent space.

Inspecting the Latent Space

What does the actual latent space look like?

Even without class information, the network organizes the latent space to encode digit structure effectively. It also becomes immediately apparent why 4s and 9s are so confused by the model. That region is a dense mixture of the two.

We can also look at the marginal distributions of each latent dimension to see how well they match the prior N(0, 1). Here, z_1 is closer to the prior than z_2. z_2 exhibits a bimodal marginal distribution, indicating that the encoder is using this dimension to separate two distinct clusters of data.

We also might want to understand how the log-variance of the latent distributions behaves.

For the most part, we see similar concentration. Some digits are more concentrated than others, though in general the difference is slight.

Beyond 2D: Higher-Dimensional Latent Spaces

What happens as we increase the latent dimensionality? We must do dimensionality reduction to visualize latent spaces, giving us an approximate sense of how the latent space is organized.

As we double the dimensionality, we see a dominant trend at first:

• The reconstruction loss goes down
• The KL loss goes up
• The KL loss per dimension goes down

Something odd happens when we jump from 16 to 32 latent dimensions: some of our latent dimensions become degenerate and stop encoding useful information. This could be an indication we need to choose our hyperparameters a little more cautiously. Perhaps we need a different architecture. Or maybe there is an intrinsic limit to the dimensionality needed for this dataset past which it’s not really helpful to keep scaling the latent dimension.

Training Dynamics

The training dynamics show the battle between reconstruction and KL divergence for different latent dimensionalities. As we increase the latent dimensionality, the oscillation in the KL divergence becomes more pronounced. Particularly chaotic is the $D=16$ case, which struggles to find a stable equilibrium. By the time we expand to $D=32$, the KL penalty seems to overpower the ability to encode information in the latent space, leading to many inactive dimensions. The drop in KL complexity has staircase-like steps without clearly gaining reconstruction ability.

Reconstruction and Generation

As we increase the latent dimensionality, the reconstruction quality improves significantly.

As we increase the dimensionality, we see the increase in quality we’d expect given the reduction in BCE reconstruction loss. In the jump to 4D, we’re able to better resolve the differences between 4s and 9s. Images become much sharper by the time we hit 16 dimensions. The differences between 16 and 32 dimensions, however, are marginal.

Sampling quality also improves with latent dimensionality. Images are sharper as we increase the dimensionality. However, the space seems to get sparser as we increase to the largest dimensionalities, which makes sense given the size and nature of our dataset.

Latent Space Visualizations

The challenge with visualizing higher-dimensional latent spaces is that we must reduce their dimensionality to 2D. PCA struggles to capture the variance of higher dimensionalities. The 4D and 8D plots suggest increasingly better separation of the numeric classes. However, the 16D and 32D plots only show 10-20% of the variance and give a misleading image of overlap.

Conclusion

In this tutorial, we’ve journeyed from the core theory of Variational Autoencoders to a practical, modern PyTorch implementation and a series of experiments on the MNIST dataset. Our findings highlight several key takeaways for practitioners:

1. The VAE is a Balancing Act: The fundamental tension between reconstruction fidelity and latent space regularization is the core of the VAE. Our visualization of the BCE vs. KL loss trade-off clearly showed training as a search for an optimal point on this Pareto frontier, where improving one objective necessarily means sacrificing the other.

2. Latent Dimensionality is a Critical Hyperparameter: Increasing the latent dimension consistently improved reconstruction quality with diminishing returns. As we saw in the jump from 16 to 32 dimensions, too much capacity can lead to “inactive” dimensions, where the KL penalty overpowers the model’s ability to encode useful information. This demonstrates that choosing the right latent size is crucial for both performance and efficiency.

3. VAEs Learn Meaningful Unsupervised Representations: Without any labels, our VAE successfully organized the latent space, clustering similar digits and enabling smooth interpolations. This underscores the power of VAEs for unsupervised learning, dimensionality reduction, and discovering the underlying structure in complex data.

4. Implementation Details Matter: While different standard deviation parameterizations yielded similar results on this simple problem, understanding their gradient behaviors is key for tackling more complex datasets where training stability can be a major challenge. Proper loss scaling is similarly crucial to prevent one term from dominating the other and leading to issues like posterior collapse.

While the classic VAE produces characteristically blurry reconstructions, it remains a foundational generative model. The principles we’ve explored here (the ELBO, the reparameterization trick, and the trade-off between reconstruction and regularization) are central to many more advanced generative models used today.

Questions or feedback? Feel free to reach out. I’d love to hear about your experiences with VAE experiments!

Modern generative AI is dominated by diffusion models and autoregressive transformers. The adversarial training dynamics and objective functions introduced by Generative Adversarial Networks (GANs) still inform how the field thinks about loss-function design and training stability today. Before diving into GANs, let’s establish what we’re trying to accomplish with generative models.

The core goal: Create a system that can generate new, realistic data that appears to come from the same distribution as our training data.

Think of having a model that can create images, text, or audio that are difficult to distinguish from human-created content. This is what generative modeling aims to achieve.

The Mathematical Foundation

Generative models aim to estimate the probability distribution of real data. If we have parameters $\theta$, we want to find the optimal $\theta^*$ that maximizes the likelihood of observing our real samples:

$$ \theta^* = \arg\max_\theta \prod_{i=1}^{n} p_\theta(x_i) $$

This is equivalent to minimizing the distance between our estimated distribution and the true data distribution. A common distance measure is the Kullback-Leibler Divergence. Maximizing log-likelihood equals minimizing KL divergence.

Two Approaches to Generative Modeling

Explicit Distribution Models

These models define an explicit probability distribution and refine it through training.

Example: Variational Auto-Encoders (VAEs) require:

• An explicitly assumed prior distribution
• A likelihood distribution
• A “variational approximation” to evaluate performance

Implicit Distribution Models

These models learn to generate data by indirectly sampling from a learned distribution. GANs exemplify this implicit approach, learning distributions through adversarial competition.

The GAN Architecture: A Game of Deception

Generative Adversarial Networks get their name from three key components:

• Generative: They create new data
• Adversarial: Two networks compete against each other
• Networks: Built using neural networks

The core innovation is the adversarial setup: two neural networks compete against each other, driving mutual improvement.

The Generator: The Forger

Role: Create convincing fake data from random noise

The generator network $G$ learns a mapping function: $$z \rightarrow G(z) \approx x_{\text{real}}$$

Where:

• $z$ is a random latent vector (the “noise”)
• $G(z)$ is the generated sample
• The goal is making $G(z)$ indistinguishable from real data

Key insight: The latent space $z$ is continuous, meaning small changes in $z$ produce smooth, meaningful changes in the generated output.

The Discriminator: The Detective

Role: Distinguish between real and generated samples

The discriminator network $D$ outputs a probability: $$D(x) = P(\text{x is real})$$

• $D(x) \approx 1$ for real samples
• $D(x) \approx 0$ for fake samples

It functions as an “authenticity detector” that progressively improves.

The Adversarial Competition

This adversarial dynamic drives the training process. The generator and discriminator have directly opposing objectives:

This creates a dynamic where both networks continuously improve:

• Generator creates better fakes to fool the discriminator
• Discriminator becomes better at detecting fakes
• The cycle continues until equilibrium

Learning Through Metaphors

Relatable analogies often clarify complex concepts. Here are two metaphors that capture different aspects of how GANs work.

The Art Forger vs. Critic

Generator = Art Forger
Discriminator = Art Critic

A criminal forger tries to create fake masterpieces, while an art critic must identify authentic works. Each interaction teaches both parties:

• The forger learns what makes art look authentic
• The critic develops a keener eye for detecting fakes
• Eventually, the forger becomes so skilled that even experts can’t tell the difference

This captures the adversarial nature and continuous improvement aspect of GANs.

The Counterfeiter vs. Bank Teller

Generator = Counterfeiter
Discriminator = Bank Teller

Day 1: Criminal brings a crayon drawing of a dollar bill. Even a new teller spots this fake.

Day 100: The counterfeiter has learned better techniques. The teller has developed expertise in security features.

Day 1000: The fake money is so convincing that detecting it requires advanced equipment.

This illustrates the progressive improvement and escalating sophistication in both networks.

The Mathematical Foundation

Now let’s examine the mathematical framework that makes GANs work. The core of GAN training is solving a minimax optimization problem.

The Minimax Objective

$$ \min_{G} \max_{D} V(D, G) = \mathbb{E}_{x \sim p_{\text{data}}(x)}[\log D(x)] + \mathbb{E}_{z \sim p_z(z)}[\log(1 - D(G(z)))] $$

Breaking this down:

• $\mathbb{E}_{x \sim p_{\text{data}}(x)}[\log D(x)]$: The expected log-probability for real data.
  • Discriminator’s Goal: Maximize this term to correctly classify real samples.
• $\mathbb{E}_{z \sim p_z(z)}[\log(1 - D(G(z)))]$: The expected log-probability for fake data being correctly identified as fake.
  • Discriminator’s Goal: Maximize this term.
  • Generator’s Goal: Minimize this term to fool the discriminator.

Why “Minimax”?

• Discriminator ($D$): Tries to maximize the objective → Better at distinguishing real from fake.
• Generator ($G$): Tries to minimize the objective → Better at fooling the discriminator.

A Practical Challenge: Vanishing Gradients

The minimax objective presents a practical problem early in training. When the generator is poor, the discriminator can easily distinguish real from fake samples with high confidence ($D(G(z)) \approx 0$). This causes $\log(1 - D(G(z)))$ to saturate and results in vanishing gradients for the generator, which effectively stalls learning.

The Solution: Practitioners typically train the generator to maximize $\log(D(G(z)))$ to provide stronger gradients early in training. This non-saturating heuristic prevents the learning process from stalling.

The Training Process

The beauty of GANs lies in their alternating optimization:

1. Fix $G$, train $D$: Make the discriminator optimal for the current generator
2. Fix $D$, train $G$: Improve the generator against the current discriminator
3. Repeat: Continue until reaching Nash equilibrium

Theoretical Goal: Nash Equilibrium

At convergence, the discriminator outputs $D(x) = 0.5$ for all samples, meaning it can’t distinguish between real and fake data. This indicates that $p_{\text{generator}} = p_{\text{data}}$. Our generator has learned the true data distribution.

The Evolution of Objective Functions

The objective function is the mathematical heart of any GAN. It defines how we measure the “distance” between our generated distribution and the real data distribution. This choice profoundly impacts:

• Training stability: Some objectives lead to more stable convergence
• Sample quality: Different losses emphasize different aspects of realism
• Mode collapse: The tendency to generate limited variety
• Computational efficiency: Some objectives are faster to compute

The original GAN uses Jensen-Shannon Divergence (JSD), but researchers have discovered many alternatives that address specific limitations. Let’s explore this evolution.

The Original GAN: Jensen-Shannon Divergence

The foundational GAN minimizes the Jensen-Shannon Divergence:

$$ \text{JSD}(P, Q) = \frac{1}{2} \text{KL}(P | M) + \frac{1}{2} \text{KL}(Q | M) $$

Where $M = \frac{1}{2}(P + Q)$ is the average distribution, and $\text{KL}$ is the Kullback-Leibler Divergence.

Strengths: Solid theoretical foundation, introduced adversarial training
Limitations: Can suffer from vanishing gradients and mode collapse

Wasserstein GAN (WGAN)

The Wasserstein GAN replaced Jensen-Shannon divergence with the Earth-Mover (Wasserstein) distance, which gives meaningful gradients even when the real and generated distributions do not overlap.

Understanding Earth-Mover Distance

The Wasserstein distance, also known as Earth-Mover distance, has an intuitive interpretation:

Imagine two probability distributions as piles of dirt. The Earth-Mover distance measures the minimum cost to transform one pile into the other, where cost = mass x distance moved.

Mathematically:

$$ W_p(\mu, \nu) = \left( \inf_{\gamma \in \Gamma(\mu, \nu)} \int_{M xM} d(x, y)^p , d\gamma(x, y) \right)^{1/p} $$

Why Earth-Mover Distance Matters

WGAN Implementation

Since we can’t compute Wasserstein distance directly, WGAN uses the Kantorovich-Rubinstein duality:

1. Train a critic function $f$ to approximate the Wasserstein distance
2. Constrain the critic to be 1-Lipschitz (using weight clipping)
3. Optimize the generator to minimize this distance

Key WGAN Benefits

Meaningful loss function: Loss correlates with sample quality
Improved stability: Less prone to mode collapse
Theoretical guarantees: Solid mathematical foundation
Better convergence: Works even when distributions don’t overlap

Improved WGAN: Solving the Weight Clipping Problem

Improved WGAN (WGAN-GP) addresses a critical flaw in the original WGAN: weight clipping.

The Problem with Weight Clipping

Original WGAN clips weights to maintain the 1-Lipschitz constraint:

# Problematic approach
for param in critic.parameters():
    param.data.clamp_(-0.01, 0.01)

Issues with clipping:

• Forces critic to use extremely simple functions
• Pushes weights toward extreme values ($\pm c$)
• Can lead to poor gradient flow
• Capacity limitations hurt performance

The Gradient Penalty Solution

WGAN-GP introduces a gradient penalty term to constrain the critic:

$$ L = E_{\tilde{x} \sim P_g}[D(\tilde{x})] - E_{x \sim P_r}[D(x)] + \lambda E_{\hat{x}}[(||\nabla_{\hat{x}} D(\hat{x})||_2 - 1)^2] $$

Where $\hat{x}$ are points sampled uniformly along straight lines between real and generated data points.

Advantages:

• No capacity limitations
• Better gradient flow
• More stable training
• Works across different architectures

LSGAN: The Power of Least Squares

Least Squares GAN takes a different approach. It replaces the logarithmic loss with L2 (least squares) loss.

Motivation: Beyond Binary Classification

Traditional GANs use log loss, which focuses primarily on correct classification:

• Real sample correctly classified → minimal penalty
• Fake sample correctly classified → minimal penalty
• Distance from decision boundary ignored

L2 Loss: Distance Matters

LSGAN uses L2 loss, which penalizes proportionally to distance:

$$ \min_D V_{LSGAN}(D) = \frac{1}{2}E_{x \sim p_{data}(x)}[(D(x) - b)^2] + \frac{1}{2}E_{z \sim p_z(z)}[(D(G(z)) - a)^2] $$

$$ \min_G V_{LSGAN}(G) = \frac{1}{2}E_{z \sim p_z(z)}[(D(G(z)) - c)^2] $$

Where typically: $a = 0$ (fake label), $b = c = 1$ (real label)

Benefits of L2 Loss

Key insight: LSGAN minimizes the Pearson χ² divergence, providing smoother optimization landscape than JSD.

Relaxed Wasserstein GAN (RWGAN)

Relaxed WGAN bridges the gap between WGAN and WGAN-GP, proposing a general framework for designing GAN objectives.

Key Innovations

Asymmetric weight clamping: RWGAN introduces an asymmetric approach that provides better balance.

Relaxed Wasserstein divergences: A generalized framework that extends the Wasserstein distance, enabling systematic design of new GAN variants while maintaining theoretical guarantees.

Benefits

• Better convergence properties than standard WGAN
• Framework for designing new loss functions and GAN architectures
• Competitive performance with other Wasserstein-based methods

Key insight: RWGAN parameterized with KL divergence shows excellent performance while maintaining the theoretical foundations that make Wasserstein GANs attractive.

Statistical Distance Approaches

Several GAN variants focus on minimizing specific statistical distances between distributions.

McGAN: Mean and Covariance Matching

McGAN belongs to the Integral Probability Metric (IPM) family, using statistical moments as the distance measure.

Approach: Match first and second-order statistics:

• Mean matching: Align distribution centers
• Covariance matching: Align distribution shapes

Moment-matching objectives like this are conceptually related to settings where aligning statistical moments matters, such as matching a generated distribution to a target physical distribution (e.g., molecular conformations). McGAN itself, however, was introduced and demonstrated as an IPM method for image generation.

Limitation: Relies on weight clipping like original WGAN.

GMMN: Maximum Mean Discrepancy

Generative Moment Matching Networks eliminates the discriminator entirely, directly minimizing Maximum Mean Discrepancy (MMD).

MMD Intuition: Compare distributions by their means in a high-dimensional feature space:

$$ \text{MMD}^2(X, Y) = ||E[\phi(x)] - E[\phi(y)]||^2 $$

Benefits:

• Simple, discriminator-free training
• Theoretical guarantees
• Can incorporate autoencoders for better MMD estimation

Drawbacks:

• Computationally expensive
• Often weaker empirical results

MMD GAN: Learning Better Kernels

MMD GAN improves GMMN by learning optimal kernels adversarially to improve upon fixed Gaussian kernels.

Innovation: Combine GAN adversarial training with MMD objective for the best of both worlds.

Different Distance Metrics

Cramer GAN: Addressing Sample Bias

Cramer GAN identifies a critical issue with WGAN: biased sample gradients.

The Problem: WGAN’s Wasserstein distance lacks three important properties:

1. Sum invariance (satisfied)
2. Scale sensitivity (satisfied)
3. Unbiased sample gradients (not satisfied)

The Solution: Use the Cramer distance, which satisfies all three properties:

$$ d_C^2(\mu, \nu) = \int ||E_{X \sim \mu}[X - x] - E_{Y \sim \nu}[Y - x]||^2 d\pi(x) $$

Benefit: More reliable gradients lead to better training dynamics.

Fisher GAN: Chi-Square Distance

Fisher GAN uses a data-dependent constraint on the critic’s second-order moments (variance).

Key Innovation: The constraint naturally bounds the critic without manual techniques:

• No weight clipping needed
• No gradient penalties required
• Constraint emerges from the objective itself

Distance: Approximates the Chi-square distance as critic capacity increases:

$$ \chi^2(P, Q) = \int \frac{(P(x) - Q(x))^2}{Q(x)} dx $$

The Fisher GAN essentially measures the Mahalanobis distance, which accounts for correlated variables relative to the distribution’s centroid. This ensures the generator and critic remain bounded, and as the critic’s capacity increases, it estimates the Chi-square distance.

Benefits:

• Efficient computation
• Training stability
• Unconstrained critic capacity

Beyond Traditional GANs: Alternative Approaches

The following variants explore fundamentally different architectures and training paradigms.

EBGAN: Energy-Based Discrimination

Energy-Based GAN replaces the discriminator with an autoencoder.

Key insight: Use reconstruction error as the discrimination signal:

• Good data → Low reconstruction error
• Poor data → High reconstruction error

Architecture:

1. Train autoencoder on real data
2. Generator creates samples
3. Poor generated samples have high reconstruction loss
4. This loss drives generator improvement

Benefits:

• Fast and stable training
• Robust to hyperparameter changes
• No need to balance discriminator/generator

BEGAN: Boundary Equilibrium

BEGAN combines EBGAN’s autoencoder approach with WGAN-style loss functions.

Innovation: Dynamic equilibrium parameter $k_t$ that balances:

• Real data reconstruction quality
• Generated data reconstruction quality

Equilibrium equation:

$$ L_D = L(x) - k_t L(G(z)) $$

$$ k_{t+1} = k_t + \lambda(\gamma L(x) - L(G(z))) $$

MAGAN: Adaptive Margins

MAGAN improves EBGAN by making the margin in the hinge loss adaptive over time.

Concept: Start with a large margin, gradually reduce it as training progresses:

• Early training: Focus on major differences
• Later training: Fine-tune subtle details

Result: Better sample quality and training stability.

Summary: The Evolution of GAN Objectives

The evolution of GAN objective functions reflects the field’s progression toward more stable and theoretically grounded training procedures. Each variant addresses specific limitations in earlier approaches.

Complete Reference Table

Practical Recommendations

For practitioners, the choice depends on specific requirements and engineering tradeoffs:

• WGAN-GP: Best balance of stability and performance for most applications. However, tuning the gradient penalty $\lambda$ can be sensitive in practice.
• LSGAN: Simpler implementation with good empirical results.
• EBGAN: Fast experimentation and prototyping.
• Original GAN: Educational purposes and understanding fundamentals.

Real-World Impact: In my work training VLMs on terabyte-scale multimodal data and forecasting chaotic physical systems, these foundational dynamics still matter. Most generation today runs on diffusion models or autoregressive transformers, but the loss-design and training-stability lessons that came out of GAN research carry over. The choice of objective function shapes generation quality, training stability, and compute cost.

Acknowledgments: This post was inspired by the excellent survey “How Generative Adversarial Networks and Their Variants Work: An Overview of GAN”.