Module 18a - Denoising Score Matching for Energy Based Models

This module is based on the work: How to Train Your Energy-Based Models by Yang Song and Diederik P. Kingma (2021).

Table of Contents

Theory for Energy-Based Models (EBM)

The density given by an EBM is:

where $E_\theta:\mathbb{R}^d \to \mathbb{R}$ and $Z_\theta=\int \exp(-E_\theta(x)) dx$.

Given samples $x_1,\dots, x_N$ in $\mathbb{R}^d$, we want to find the parameter $\theta$ maximizing the log-likelihood $\max_\theta \sum_{i=1}^N \log p_{\theta}(x_i)$. Since $Z_\theta$ is a function of $\theta$, evaluation and differentiation of $\log p_{\theta}(x)$ w.r.t. $\theta$ involves a typically intractable integral.

Maximum Likelihood Training with MCMC

We can estimate the gradient of the log-likelihood with MCMC approaches:

The first term is simple to compute (with automatic differentiation).

Maths: computing $\nabla_\theta \log Z_\theta$ We have:

Proof:

Thus, we can obtain an unbiased one-sample Monte Carlo estimate of the log-likelihood gradient by

with $\tilde{x}\sim p_\theta(x)$, i.e. a random sample from the distribution given by the EBM. Therefore, we need to draw random samples from the model. As explained during the course, this can be done using Langevin MCMC. First note that the gradient of the log-probability w.r.t. $x$ (which is the score) is easy to calculate:

Hence, in this case, Langevin MCMC is given by:

where $z_t\sim \mathcal{N}(0,I)$. When $\epsilon\to 0$ and $t\to \infty$, $x_t$ will be distributed as $p_\theta(x)$ (under some regularity conditions).

In this homework, we will consider an alternative learning procedure.

Score Matching

The score (which was used in Langevin MCMC above) is defined as

If $p(x)$ denote the (unknown) data distribution, the basic score matching objective minimizes:

The problem with this objective is that we cannot compute $\nabla_x \log p(x)$ as $p(x)$ is unknown. We can only compute (approximate) averages with respect to $p(x)$ with empirical averages. Fortunately, we can solve this issue as we have:

where $c$ is a constant (not depending on $\theta$).

Proof:

by integration by parts where for a vector valued function $v(x_1,x_2,x_3)$ $\nabla_x \cdot v = \frac{\partial v_1}{\partial x_1} + \frac{\partial v_2}{\partial x_2}+ \frac{\partial v_3}{\partial x_3}$. The statement follows.

Denoising Score Matching

There are several drawbacks about the score matching approach: computing the trace of the Hessian is expensive and scores will not be accurately estimated in low-density regions, see Generative Modeling by Estimating Gradients of the Data Distribution

Denoising score matching is an elegant and scalable solution. Consider the random variable $Y = X+\sigma Z$, where $X\sim p(x)$ and $Z\sim\mathcal{N}(0,I)$. We denote by $p^\sigma(y)$ the distribution of $Y$ so that we have:

Proof:

We denote by $\varphi$ the density of $\mathcal{N}(0,\sigma^2 I)$. We have $p^\sigma(y) = \int p(x) \varphi(y-x) dx$ so that using the fact that $\nabla_z \varphi(z) = -\frac{z}{\sigma^2} \varphi(z)$, we get

The denoising score matching objective is now

that we will minimize thanks to a gradient descent in the parameter $\theta$.

In practice, we use the following relation:

where $C$ does not depend on $\theta$ (made explicit below).

Proof: We have

Hence, in practice, we will minimize the (random) loss:

where the $z_i$ are iid Gaussian. As the dataset is too large, we will run SGD algorithm, i.e. make batches and use automatic differentiation to get the gradient w.r.t. $\theta$ over each batch.

Code for Energy Based Models

• Denoising Score Matching for Energy Based Models for a simple case Denoising_Score_Matching_Energy_Model.ipynb. Here is the corresponding solution: solution

Edit this page on Last modified: May 28, 2025. Website built with Franklin.jl and the Julia programming language.