Denoising Diffusion Models¶
Study forward noising and reverse sampling on a one-dimensional Gaussian mixture. Because the exact score is available, we can compare numerical sampling with the target density before training a neural approximation by denoising score matching.
Run this tour¶
Run the cells in order with a Python 3 kernel. The first cell locates the companion data and toolbox and installs missing dependencies when needed. All worked examples include their implementation directly in this notebook. Random seeds make comparisons reproducible; you can change them to explore other samples.
The default experiment uses 2,000 particles and a small neural network so that the complete calculation remains practical on a CPU. Increase these sizes for a more accurate density estimate.
# Locate the companion toolbox locally, or fetch it for a standalone/Colab copy.
from pathlib import Path
import importlib.util
import os
import subprocess
import sys
working = Path.cwd()
candidates = [working, working / "python", working.parent / "python"]
python_dir = next((p for p in candidates if (p / "nt_toolbox").is_dir()), None)
if python_dir is None:
checkout = working / "numerical-tours-support"
if not checkout.exists():
subprocess.run(
[
"git",
"clone",
"--depth",
"1",
"--branch",
"master",
"https://github.com/gpeyre/numerical-tours.git",
str(checkout),
],
check=True,
)
python_dir = checkout / "python"
os.chdir(python_dir)
if str(python_dir) not in sys.path:
sys.path.insert(0, str(python_dir))
requirements = python_dir / "requirements.txt"
if any(
importlib.util.find_spec(name) is None
for name in [
"numpy",
"scipy",
"matplotlib",
"skimage",
"sklearn",
"pywt",
"ipywidgets",
"cvxpy",
"skfmm",
"autograd",
"progressbar",
"celer",
]
):
subprocess.run(
[sys.executable, "-m", "pip", "install", "-r", str(requirements)], check=True
)
if (
importlib.util.find_spec("torch") is None
or importlib.util.find_spec("torchvision") is None
):
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"-r",
str(python_dir / "requirements-torch.txt"),
],
check=True,
)
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(0)
plt.rcParams.update(
{
"figure.figsize": (8, 4),
"figure.dpi": 100,
"axes.spines.top": False,
"axes.spines.right": False,
"font.size": 11,
"image.cmap": "gray",
}
)
%matplotlib inline
This numerical tours explains denoising diffusion model on a simple 1-D example. The density to sample from is a mixture of Gaussian, so that the diffused densities can be computed in closed form.
To learn more about diffiusion model, you can have a look at Valentin de Bortoli courses's slides. I have written a short mathematical memo here.
import torch
torch.manual_seed(0)
torch.set_num_threads(1)
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
Mixture of Gaussians¶
In order to be able to have closed form formula for the flows simulated bellow, we consider 1-D mixtures of Gaussians $\sum_i a_i \mathcal{N}(\mu_i,\sigma_i^2)$ parameterized by $(a,\mu,\sigma)$. We define their density functions and a sampling operator.
def gauss(x, m, s):
return (
1
/ (torch.sqrt(2 * torch.tensor(torch.pi)) * s)
* torch.exp(-((x - m) ** 2) / (2 * s**2))
)
def gaus_mixture(x, mu, sigma, a):
y = x * 0
for i in range(len(mu)):
y = y + a[i] * gauss(x, mu[i], sigma[i])
return y
def sample_mixture(P, mu, sigma, a):
a1 = np.array(a / torch.sum(a))
I = np.random.choice(np.arange(0, len(a)), size=(P, 1), p=a1)
return (mu[I] + torch.randn((P, 1)) * sigma[I]).flatten()
Define also helpers for the display of histograms and trajectories.
fs = (6, 3) # size of figures for display
xmax = 3
nbins = 200
def display_histogram(X0):
P = len(X0)
h0, b = torch.histogram(X0.flatten(), bins=nbins, range=[-xmax, xmax])
plt.bar(
torch.linspace(-xmax, xmax, nbins),
h0 / P * nbins / (2 * xmax),
width=2 * xmax / (nbins - 1),
align="center",
color="red",
)
def display_trajectories(
Y, m=1000, alpha=0.3, linewidth=0.2
): # m=number of trajectories to display
disp_ind = torch.round(torch.linspace(0, P - 1, m)).type(torch.int32)
I = torch.argsort(Y[:, -1])
Y = Y[I, :]
for i in range(m):
k = disp_ind[i]
s = k / (P - 1)
plt.plot(
torch.linspace(0, T, N),
Y[k, :],
color=[s.item(), 0, 1 - s.item()],
alpha=alpha,
linewidth=linewidth,
)
Define a measure $\rho_0(x)$ defined as a mixture and display.
mu = torch.tensor([-0.7, 0, 0.5]) * 3 # mean
sigma = torch.tensor([0.02, 0.05, 0.03]) * 9 # std
a = torch.tensor([2, 3, 2]) # weight, should sum to 1
a = a / torch.sum(a)
def rho0(x):
return gaus_mixture(x, mu, sigma, a)
Evaluate the density on a grid and display it.
x = torch.linspace(-xmax, xmax, 1000)
plt.figure(figsize=fs)
plt.fill_between(x, rho0(x));
Langevin sampling¶
To sample from a density $\rho_0(x)$, Langevin diffusion performs a gradient descent on the negative log-density, which reads, for a step size $\tau>0$, $$ X_{k+1} := X_k + \tau \eta_0(X_k) + \sqrt{2\tau}W_k $$ where the $W_{k} \sim \mathcal{N}(0,\text{Id})$ are i.d.d. Gaussian white noise, and where the ``non-smoothed'' (we will introduced its smoothed version later) score reads $$ \eta_0(x) := \nabla \log \rho_0(x) = \frac{\nabla \rho_0(x)}{\rho_0(x)}. $$
When $\tau \rightarrow 0$, this corresponds to the following Langevin SDE $$ d X_t = \eta_0(X_t) d t + \sqrt{2\tau} d W_t $$ where $W_t$ is a Wiener process.
One can show that the law of $X_T$ converges in law toward $\rho_0$.
We first compute and display this score.
def eta0(x):
x1 = x.detach().clone()
x1.requires_grad = True
L = torch.sum(rho0(x1))
L.backward()
return x1.grad / rho0(x)
plt.figure(figsize=fs)
plt.plot(x, rho0(x).detach().numpy(), label="$\\rho_0$")
plt.plot(x, eta0(x).detach().numpy() / 40, label="$\\eta_0$")
plt.legend();
Sample $Z_0$ from the initial density at time $t=0$, which is uniform on $[-w/2,w/2]$.
P = 2000 # number of particles
w = 6 # support of the indicator
Z0 = w * torch.rand((P,)) - w / 2
Run the recursion of Langevin.
T = 10 # final time
N = 2500 # number of steps
tau = T / N # step size
Z = torch.zeros((P, N))
Z[:, 0] = Z0
for i in range(N - 1):
Z[:, i + 1] = (
Z[:, i]
+ tau * eta0(Z[:, i])
+ torch.sqrt(2 * torch.tensor(tau)) * torch.randn((P,))
)
Display evolving empirical densities by showing the histograms of the samples.
disp_ind = (torch.tensor([0, 0.001, 0.005, 0.01, 0.99]) * N).type(torch.int32)
for i in range(len(disp_ind)):
k = disp_ind[i]
plt.subplot(len(disp_ind), 1, i + 1)
display_histogram(Z[:, k])
plt.fill_between(x, rho0(x).detach().numpy(), color=[0, 0, 1, 0.25])
Display sample path $t \mapsto X_t$.
plt.figure(figsize=fs)
display_trajectories(Z, m=100)
Forward noising process¶
The two major issues with Langevin is that it converges slowly for singular $\eta_0$ (in particular with highly non concave log-density) and that it needs $\eta_0$ to be known in closed form or estimated. For applications to generative model, this is unacceptable, because $\rho_0$ is only known approximately from sample. It is very hard to estimate $\eta_0$ from sample, but as we will see, it is possible to instead estimated smoothed version of it. Diffusion model leverage this by replacing Langevin sampling by another sampling method that rely on such smoothed score.
This sampling method is achieved by inversing a forward ``noising process''. We consider the following noising model $$ d X_t = -X_t + \sqrt{2} dW_t $$ where $W_t$ is a Wiener process. From the knowledge of $X_0$, the law $\rho_t$ of $X_t$ can be computed exactly as the law of $$ X_t \sim e^{-t} X_0 + \sqrt{1-e^{-2t}} Z $$ where $Z \sim \mathcal{N}(0,1)$.
In the case of a mixture of Gaussian considered here, $$ X_0 \sim \rho_0 \triangleq \sum_i a_i \mathcal{N}(\mu_i,\sigma_i^2) $$ the law of $X_t$ can thus be computed exactly in closed form as $$ X_t \sim \rho_t \triangleq \sum_i a_i \mathcal{N}(e^{-t} \mu_i, e^{-2t} \sigma_i^2 + 1-e^{-2t}) $$
def rho(x, t):
sigma_t = torch.sqrt(torch.exp(-2 * t) * sigma**2 + 1 - torch.exp(-2 * t))
return gaus_mixture(x, mu * torch.exp(-t), sigma_t, a)
Display evolution with $t$ of the density $\rho_t(x)$
plt.figure(figsize=fs)
t_list = torch.tensor([0, 0.1, 0.5, 2, 100]) / 10
for i in range(len(t_list)):
s = i / len(t_list)
t = t_list[i]
plt.fill_between(
x, rho(x, t), color=[s, 0, 1 - s, 0.25], label=f"t={t:.1f}", linewidth=0
)
plt.legend();
We first generate $P$ sample $X_0$ from the mixture of Gaussian at time $t=0$, and compare the associated empirical histogram with the density.
X0 = sample_mixture(P, mu, sigma, a)
plt.figure(figsize=fs)
display_histogram(X0)
plt.plot(x, rho0(x).detach().numpy())
The forward noising process is computed numerically approximately with using an Euler-Maruyama scheme $$ X_{k+1} = X_k - \tau X_k + \sqrt{2\tau} W_k, $$ where $W_k$ are i.i.d. white noises.
Set the final time $T$, which should be large enough so that the final distribution of $X_T$ is approximately $\mathcal{N}(0,1)$.
T = 2 # final time
tau = T / N # step size
Now run the recursion.
X = torch.zeros((P, N))
X[:, 0] = X0
for i in range(N - 1):
X[:, i + 1] = (
X[:, i] - tau * X[:, i] + torch.sqrt(2 * torch.tensor(tau)) * torch.randn((P,))
)
# display evolving density
t_list = torch.tensor([0, 0.01, 0.05, 0.2, T * 0.99])
for i in range(len(t_list)):
s = i / len(t_list)
t = t_list[i]
k = torch.round(t / tau).type(torch.int32)
plt.subplot(len(t_list), 1, i + 1)
plt.plot(x, rho(x, t), color="blue")
display_histogram(X[:, k])
Reverse diffusion¶
The score is the vector field $$ \eta_t(x) := \nabla \log(\rho(x,t)) = \frac{\nabla \rho(x,t)}{\rho(x,t)}. $$ Informally, it points toward regions of high densities.
Since here $\rho(x,t)$ can be computed in closed form, the score can itself be computed in closed form. Of course, in real situation, it is impossible, and we will see in the next section how to approximate it numerically from sampled of $\rho(x,t)$.
def eta(x, t):
x1 = x.detach().clone()
x1.requires_grad = True
L = torch.sum(rho(x1, t))
L.backward()
return x1.grad / rho(x, t)
for i in range(len(t_list)):
s = i / len(t_list)
t = t_list[i]
plt.subplot(len(t_list), 1, i + 1)
plt.plot(x, rho(x, t).detach().numpy())
plt.plot(x, eta(x, t).detach().numpy() / 20)
The actual sampling (the generative process) is now done by reverting in time this process, i.e., for a large enough $T \gg 0$, one seeks to approximate $X_{T-t}$.
The main result of diffusion model is that the law of $X_{T-t}$ is exactly equal to the law of $Y_t$ satisfying the following Langevin SDE, initialized with $Y_0=X_T$ $$ \mathrm{d} Y_t = [ Y_t + (1+\alpha) \nabla \log(\rho_{T-t})(Y_t) ] \mathrm{d} t + \sqrt{2 \alpha} \mathrm{d} W_t $$ where $\alpha \geq 0$ is a parameter controlling the randomness introduced during the backward flow. This can be discretized using an Euler-Maruyama scheme, starting from $Y_0=X_{T/\tau}$ $$ Y_{k+1} = Y_k + \tau Y_t + (1+\alpha) \nabla \log(\rho_{T-t})(Y_t) + \sqrt{2\tau\alpha} W_k. $$
The case $\alpha=1$ is the one used in diffusion model. Setting different value of $\alpha$ is equivalent to changing the time of the diffusion. The case $\alpha=0$ corresponds to a singular limit which is only random at $Y_0$ but after follow an advection equation $$ \frac{\mathrm{d} Y_t}{\mathrm{d} t} = Y_t + \nabla \log(\rho_{T-t})(Y_t) $$
Sample $Y_0$ from the stationary distribution $\mathcal{N}(0,1)$. Since we are in 1-D, instead of using random sample, we use deterministic "optimal" (in quantization error aka Optimal Transport) sample location, given by the uniform sampling of the inverse cumulative distribution (quantile) of the Gaussian.
import scipy as sp
Y0 = torch.tensor(sp.stats.norm.ppf(np.arange(0, P) / P + 1 / (2 * P)))
plt.figure(figsize=fs)
display_histogram(Y0)
plt.plot(x, rho(x, torch.tensor(10)));
Implement the backward recursion.
alpha = 0 # pure advection
alpha = 0.01 # original diffusion is alpha=1
Y = torch.zeros((P, N))
Y[:, 0] = Y0
t = torch.tensor(T)
for i in range(N - 1):
Y[:, i + 1] = (
Y[:, i]
+ tau * (Y[:, i] + (1 + alpha) * eta(Y[:, i], t))
+ torch.sqrt(2 * torch.tensor(tau * alpha)) * torch.randn((P,))
)
t = t - tau
Display the evolution of the density of the law of $Y_{T-t}$. It is close to $\rho_t(x)$, and one can show that it matches exactly when $T \to +\infty$ and $\tau \to 0$.
# display evolving density
t_list = T - torch.tensor([0.01, 0.01, 0.05, 0.2, 0.99 * T])
for i in range(len(t_list)):
s = i / len(t_list)
t = t_list[i]
k = torch.round(t / tau).type(torch.int32)
plt.subplot(len(t_list), 1, i + 1)
plt.plot(x, rho(x, T - t), color="blue")
display_histogram(Y[:, k])
Displays sample of trajectories $t \mapsto Y_t$.
display_trajectories(Y, m=100, alpha=0.5, linewidth=1)
plt.axis([0.4 * T, T, -3, 3]);
Denoising Score Matching¶
In order to be able to implement the reverse diffusion, one needs to compute the score $\nabla \log(\rho_t)$, where $\rho_t$ is the density of the distribution of $X_t$. The idea is to approximate this score using a function computed from samples of $X_t$.
$$ \nabla \log \rho_t(z) \approx \phi_\theta(z,t) $$ $$ \underset{\theta}{\min} \quad \int_t \int_x \int_z \| \frac{e^{-t} x - z}{1-e^{-2t}} - \phi_\theta(z,t) \|^2 \, \text{d} \mathbb{P}(z|x) \text{d} \rho_0(x) \lambda_t \text{d} t $$ where $\lambda_t$ is a density of sampled time.
To generate the pairs $(X_0,X_t)$ we first sample $X_0$ and then sample $\mathbb{P}(X_t|X_0)$ which we know in closed form.
N = 2000 # number of training samples
X0 = sample_mixture(N, mu, sigma, a)
For the sampling density $\lambda_t$ of time, we use a uniform distribution of time within $[10^{-2},T]$.
Tmin = 1e-1
Tmax = 0.5
t_list = torch.linspace(Tmin, Tmax, N)
To sample from $\mathbb{P}(X_t|X_0)$, recall that given $X_0$, then $X_t$ can be sampled as $$ X_t \sim e^{-t} X_0 + \sqrt{1-e^{-2t}} Z $$ where $Z \sim \mathcal{N}(0,1)$.
W = torch.randn((N,))
Xt = torch.exp(-t_list) * X0 + torch.sqrt(1 - torch.exp(-2 * t_list)) * W
Generate for $i \in \{1,\ldots,N\}$ the input data points $U_i := (z_i \sim X_{t_i},t_i) \in \mathbb{R}^2$ and the output data to predict $V_i := \frac{e^{-t_i} x_i - z_i}{1-e^{-2t_i}}$ for $x_i \sim X_0$.
def tMap(t):
return (
((t - Tmin) / (Tmax - Tmin) - 1 / 2) * 2 * xmax
) # map t to the same range as x to ease training
U = torch.vstack((Xt, tMap(t_list))).T
V = (torch.exp(-t_list) * X0 - Xt) / (1 - torch.exp(-2 * t_list))
Uncomment the following code if you want to use ground truth $\eta_t$ values for training.
if 0:
rho_v = torch.func.vmap(rho) # parrallel version of rho
def eta_v(x, t):
x1 = x.detach().clone()
x1.requires_grad = True
L = torch.sum(rho_v(x1, t))
L.backward()
return x1.grad / rho_v(x, t)
Data_out = eta_v(Xt, t_list)
For $\phi_\theta(z,t)$, we use a Multilayers Perceptron (MLP) layers neural network with input dimension $2$ (for $(z,t)$) and the output should approximate $\frac{e^{-t} x - z}{1-e^{-2t}}$ and has dimension $1$. The MLP has a varying number of layers and dimensions, encoded in d_hidden.
def weights_init(m):
if isinstance(m, nn.Linear):
torch.nn.init.xavier_normal_(m.weight)
if m.bias is not None:
torch.nn.init.zeros_(m.bias)
class MLP(nn.Module):
def __init__(self, d_in, d_hidden, d_out, dropout_probability=0.1):
super(MLP, self).__init__()
layers = []
layers.append(nn.Linear(d_in, d_hidden[0]))
layers.append(nn.Sigmoid())
layers.append(nn.Dropout(dropout_probability))
for i in range(1, len(d_hidden)):
layers.append(nn.Linear(d_hidden[i - 1], d_hidden[i]))
layers.append(nn.Sigmoid())
layers.append(nn.Dropout(dropout_probability))
layers.append(nn.Linear(d_hidden[-1], d_out))
self.mlp = nn.Sequential(*layers)
self.apply(weights_init)
def forward(self, x):
return self.mlp(x)
Create the neural network with 3 hidden layers and initialize its weights.
d_in = 2
d_hidden = [64, 32]
d_out = 1
model = MLP(d_in, d_hidden, d_out, dropout_probability=0)
Train the network by minimzing $$ \text{Loss}(\theta) := \frac{1}{N} \sum_{i=1}^N \| \phi_\theta(U_i)-V_i \|^2 $$ with respect to the neural network weights and biases $\theta$ using the Adam optimizer.
optimizer = torch.optim.Adam(model.parameters(), lr=0.001 * 10)
num_epochs = 200
Loss_list = []
for epoch in range(num_epochs):
optimizer.zero_grad()
loss = torch.sum((model(U).flatten() - V.flatten()) ** 2) / N
loss.backward()
optimizer.step()
Loss_list.append(loss.item())
plt.plot(Loss_list);
Evaluate $g_\theta(z,t)$ on a grid and compare it with $\eta_t(z)$.
nT = 100
nX = 200
Tlist = torch.linspace(Tmin, Tmax, nT)
Xlist = torch.linspace(-xmax, xmax, nX)
A = torch.zeros((nX, nT))
B = torch.zeros((nX, nT))
for i in range(nT):
A[:, i] = eta(Xlist, Tlist[i]) # ground trust to predict
Z1 = torch.vstack((Xlist, Xlist * 0 + tMap(Tlist[i]))).T
B[:, i] = model(Z1).flatten() # approximation with neural network
plt.subplot(1, 2, 1)
plt.imshow(A, extent=[Tmin, Tmax, -xmax, xmax], aspect="auto")
plt.title("$\\eta_t(z)$")
plt.ylabel("z")
plt.xlabel("t")
plt.subplot(1, 2, 2)
plt.imshow(B.detach().numpy(), extent=[Tmin, Tmax, -xmax, xmax], aspect="auto")
plt.title("$\\phi_{\\theta}(z,t)$")
plt.xlabel("t");
References and further reading¶
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. 2020, NeurIPS. Learning a generative reverse process by denoising.
Yang Song et al.. Score-Based Generative Modeling through Stochastic Differential Equations. 2021, ICLR. Reverse-time SDEs and probability-flow ODEs.
Pascal Vincent. A Connection Between Score Matching and Denoising Autoencoders. 2011, Neural Computation 23(7), 1661–1674. The denoising score-matching objective used in diffusion training.
Aapo Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. 2005, JMLR 6, 695–709. Estimating log-density gradients without a normalization constant.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. 2016, MIT Press. Neural networks, automatic differentiation, and optimization.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 2015, ICLR. Adaptive stochastic gradient updates for neural network training.