[Paper Review] Maximum Likelihood Training of Score-Based Diffusion Models

Yang Song, Conor Durkan+2 more|arXiv (Cornell University)|Jan 22, 2021

Advanced Neuroimaging Techniques and Applications54 references46 citations

TL;DR

The paper derives a likelihood-weighted objective that upper-bounds the negative log-likelihood of score-based diffusion models, enabling approximate maximum likelihood training that improves model likelihood across datasets and SDEs, achieving competitive bits/dim on CIFAR-10 and ImageNet32x32 without augmentation.

ABSTRACT

Score-based diffusion models synthesize samples by reversing a stochastic process that diffuses data to noise, and are trained by minimizing a weighted combination of score matching losses. The log-likelihood of score-based diffusion models can be tractably computed through a connection to continuous normalizing flows, but log-likelihood is not directly optimized by the weighted combination of score matching losses. We show that for a specific weighting scheme, the objective upper bounds the negative log-likelihood, thus enabling approximate maximum likelihood training of score-based diffusion models. We empirically observe that maximum likelihood training consistently improves the likelihood of score-based diffusion models across multiple datasets, stochastic processes, and model architectures. Our best models achieve negative log-likelihoods of 2.83 and 3.76 bits/dim on CIFAR-10 and ImageNet 32x32 without any data augmentation, on a par with state-of-the-art autoregressive models on these tasks.

Motivation & Objective

• Motivate and enable higher-likelihood training for score-based diffusion models (SBDMs) by connecting their weighted score-matching losses to an upper bound on negative log-likelihood.
• Introduce likelihood weighting (lambda(t) = g(t)^2) that tightens the bound and improves model likelihood.
• Provide theoretical guarantees linking SDE/ODE formulations to likelihood, and propose variance reduction and variational techniques to practical training.
• Demonstrate improved log-likelihoods across datasets, SDE families (VE/VP/subVP) and model architectures, with competitive CIFAR-10 and ImageNet-32x32 results.

Proposed method

• Formulate score-based diffusion models with forward SDEs and reverse-time SDEs, and define time-dependent scores to be learned by a neural network.
• Show that with likelihood weighting lambda(t) = g(t)^2 the weighted score-matching objective upper-bounds the KL divergence to the model p_theta^SDE, linking to negative log-likelihood optimization.
• Prove equivalence under certain conditions between p_theta^SDE and p_theta^ODE (CNF) when the score matches the true time-dependent score.
• Provide practical bounds for individual datapoints via L^SM_theta(x) and L^DSM_theta(x) to enable efficient estimation and training.
• Address numerical stability by using a small epsilon > 0 and discuss bias corrections via Jensen’s inequality.
• Introduce variance reduction via importance sampling to mitigate high variance from likelihood weighting and combine with variational dequantization to improve likelihoods.

Experimental results

Research questions

• RQ1 Can likelihood weighting provide a principled upper bound to the negative log-likelihood of score-based diffusion models?
• RQ2 Does optimizing the likelihood-weighted score-matching objective improve the actual log-likelihood of p_theta^SDE and p_theta^ODE across datasets and SDE types?
• RQ3 Under what conditions is the likelihood bound tight or equivalent to the true data score dynamics?
• RQ4 How can variance be controlled when using likelihood weighting in training, and can importance sampling and variational dequantization further improve likelihoods?

Key findings

• Likelihood weighting (lambda(t) = g(t)^2) upper bounds the KL divergence from data to p_theta^SDE, enabling approximate maximum likelihood training.
• Training with likelihood weighting consistently improves the likelihood of p_theta^SDE across multiple datasets, SDE types (VE/VP/subVP), and model architectures.
• In experiments, the approach yields competitive log-likelihoods (2.83 bits/dim on CIFAR-10 and 3.76 bits/dim on ImageNet-32x32) without data augmentation, on par with state-of-the-art autoregressive models.
• Theoretical results connect SDE-based SBDMs to CNFs via probability flow ODEs, and show equivalence to CNF likelihoods under exact score matching.
• Variance inherent in likelihood weighting is effectively reduced with importance sampling, improving training stability and performance.
• The method enables higher-quality likelihood estimates without sacrificing too much from standard Fréchet Inception Distance metrics.