import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
import einops
from tqdm import tqdmWhat to diffuse? The basics!
diffusion models
My own notes on how diffusion models work, and why they are so cool.
 |
|---|
| Diagram from Denoising Diffusion Probabilistic Models |
This is a direct quote from the paper, that introduced diffusion models, but how do you make one and what is the intuition behind these diffusion models? That is what I want to explain in this post. It consist out of the following parts;
- What can we do with diffusion models and why do they exists?
- What are the underlying assumptions we make about this diffusion process?
- How do we convert these underlying assumptions to a loss function that we can use to train a model.
- How do we go from loss function to a trainable model and train it?
 |
|---|
At its core, diffusion models are generative models used to generate new samples of a distribution. The idea behind generative modelling is to derive for a given dataset \(\mathcal{D}\) of observations \(x^{(1)}, \cdots x^{(N)}\) to learn or derive the underlying distribution from which the data is generated. Thus derive \(x \sim p_{data}\). Which enables us to sample from this distribution and thus generate more samples, which is the goal of diffusion models.
The sample space for images is \(x \in \mathbb{R}^{H\times W}\), is incredibly big, while most of this space is occupied mostly be noisy samples. As I attempted to visualise in the picture above.
As mentioned before, we want to sample previously unseen examples, thus we need a mapping that produces a sample inside the space \(x \in \mathbb{R}^{H\times W}\). A Generative Adversarial Network (GAN), is a technique that produces such a mapping (the generator). This generator takes as input noise and outputs an image
\[ g_\theta(x_{noise}) = x_{img} \]
GANs however have two major problems;
- Mode collapse, this meant that the generator only learned to predict one specific sample. Thus it collapses the whole space \(\{x | p_{true}(x) \geq 0.9 \}\) into one sample.
- Unstable training, due to the adversarial nature of a GAN, you have to train a generator and discriminator together. When the complexity of the discriminator and the generator are not balanced, the two models will just stop learning from each other and nothing sensible happens.
Diffusion models solve these problems by modeling the underlying data distribution in a sequential way. 
In this way the training objective becomes more stable, since the whole problem is split into many smaller problems. And it doesn’t require the training of a discriminator, since the it just learns to denoise samples.
A different way to look at diffusion models that I found very useful is that of the perspective of score models, that basically says that diffusion models try to a score function \(s(x)\), which is the gradient of the log likelihood and points, for each point \(x \in \mathbb{R}^{H\times W}\)to the direction that makes your sample more likely.
\[ s(x) \approx \nabla_x \log{p(x)} \Big(= \epsilon_\theta(x, t) \Big) \]
This part explain the math behind diffusion models and how to get a loss function we will employ to train a model. Most of the math comes from the blogpost What are Diffusion Models? and from the paper link. I have rewritten the equations and added the schematics in way that it makes more sense to me, and I hope it makes it helps you as well.
As just mentioned, diffusion models make the assumption that samples are sequentially constructed from noise. This means that there is a sequence of hidden latent variables that starts with \(x_T \sim \mathcal{N}(0, \sigma^2)\) where gradually information is added (or de-noised), until there is a clean sample without noise \(x_0\).
Before I can explain how to generate data, I need to start with the forward process, the process where data is destroyed.
This is defined by the approximate posterior.
\[ q(x_{1:T}|x_0) = \prod_{t=1}^T q(x_t|x_{t-1}) \]
With the transition probabilities.
\[ q(x_t|x_{t-1}) = \mathcal{N}(x_t|\sqrt{1-\beta_t} x_{t-1}, \beta_tI) \]
Where the sequence \(\beta_0, \cdots, \beta_{T}\) is the variance schedule that controls how fast information is destroyed in the process. Using the reparameterization trick we can expess \(x_t\)
\[ x_{t} = \sqrt{1-\beta_{t-1}} \cdot x_{t-1} + \sqrt{\beta_{t-1}} \cdot \epsilon_{t-1} \]
Were
\[ \epsilon \sim \mathcal{N}(0, I) \]
The reason why this setup is so convenient is that we can use this property to derive a tractable form for \(x_t\)
\[ \begin{align*} x_{t} &= \sqrt{1-\beta_{t-1}} \cdot x_{t-1} + \sqrt{\beta_{t-1}} \cdot \epsilon_{t-1} \\ x_{t} &= \sqrt{1-\beta_{t-1}} \cdot (\sqrt{1-\beta_{t-2}} \cdot x_{t-2} + \sqrt{\beta_{t-2}} \cdot \epsilon_{t-2}) + \sqrt{\beta_{t-1}} \cdot \epsilon_{t-1} \\ x_{t} &= \sqrt{1-\beta_{t-1}} \sqrt{1-\beta_{t-2}} \cdot x_{t-2} + \sqrt{1-\beta_{t-1}}\sqrt{\beta_{t-2}} \cdot \epsilon_{t-2} + \sqrt{\beta_{t-1}} \cdot \epsilon_{t-1} \\ \vdots \\ x_{t} &= \sqrt{\bar{\alpha}_t} x_0 + \sqrt{(1-\bar{\alpha}_t)}\cdot\epsilon_{combined} \end{align*} \]
Gaussian distribution to derive a tractable form;
\[ q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I) \]
Where we use the identities
\[ \alpha_t = 1-\beta_t \\ \bar{\alpha}=\prod_{s=1}^t\alpha_t \]
This simply means that we can sample \(x_t\) for arbitrary \(t\) using the reparameterization trick.
Before we are finished with the forward process, I want to define one more probability, the forward process posteriors;
\[ q(x_{t-1}|x_t ,x_0) = \mathcal{N}(x_{t-1};\hat{\mu}_t(x_t, x_0), \hat{\beta}_t) \]
Where
\[ \hat{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar\alpha_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{{1-\bar\alpha}_t}x_t \]
And
\[ \hat{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar\alpha_t}\beta_t \]
The reason why this guy is important will be clear at the end of the blogpost. Now we know, how we can derive, given a sample \(x \in \mathcal{D}\), a latent \(x_t\) for an arbitrary \(t\). |
So, now we have defined the forward process, lets look at what we actually want to learn, namely the reverse process. This process is defined by
\[ p_\theta(x_0) = \int p_\theta(x_0, x_1, \cdots, x_{T-1}) d{x_{1:T}} \]
\[ p_\theta(x_0, \cdots, x_{T-1}) = p_\theta(x_0)\prod_{t=1}^Tp_\theta(x_{t-1}|x_t) \]
In this reverse process the transitions are assumed to be
\[ \begin{align*} p_{\theta}(x_{t-1}|x_t) &= \mathcal{N}\Big(x_{t-1};\mu_\theta(x_t, t), \hat{\beta}_t \mathbf{I}\Big ) \end{align*} \]
Where \(\mu_\theta(x_t, t)\) is the function distribution we want to learn! This function we can rewrite to;
\[ \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \Big(x_t - \frac{\beta_t}{\sqrt{1-\hat\alpha_t}}\epsilon_\theta(x_t, t) \Big) \]
When we have this function \(p_\theta(x_{t-1}|x_t)\), we can use the resample trick to sequentially reconstruct samples from noise.
So we just defined the diffusion process as a Markov chain, now we need to come up with a loss which we can use to optimise the parameters of our model \(\theta\).
Starting with the joint distribution can be done by minimising the log-likelihood for the model;
\[ \begin{align*} & -\log \big( p_{\theta}(x) \big) = -\log \big( \int p_\theta(x_{0:T})dx_{1:T} \big) \\ &=-\log \big( \int p_\theta(x_{0:T})\cdot \frac{q(x_{1:T}|x_0)}{q(x_{1:T}|x_0)} dx_{1:T} \big) \\ &= -\log \big(\mathbb{E}_q\Big[ \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} \Big] \big)\leq -\mathbb{E}_q\Big[\log \big( \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)} \Big) \Big] \\ &=-\mathbb{E}_q\Big[-\log p(x_T) - \sum_{t\geq 1} \log \{\frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})}\} \Big] \end{align*} \]
The problem with this is that we are now comparing \(p_\theta(x_{t-1}|x_t)\) with \(q(x_t|x_{t-1})\). It would be prefered if we could compare for each timestep \(t\), the approximate reverse forward process with the posterior forward process, \(q(x_{t-1}|x_t, x_0)\) (which we just derived). How this works can be seen if we realise that we can rewrite;
\[ q(x_t|x_{t-1}) = \frac{q(x_t|x_0)}{q(x_{t-1}|x_0)} \cdot q(x_{t-1}|x_t, x_0)\\ \]
To obtain
\[ \begin{align*} &=-\mathbb{E}_q\Big[-\log p(x_T) - \sum_{t\geq 1} \log \{\frac{p_\theta(x_{t-1} |x_t)}{q(x_t|x_{t-1})}\} \Big] \\ &=-\mathbb{E}_q\Big[-\log p(x_T) - \sum_{t>1} \log \{\frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})}\} -\log \big\{ \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\big\}\Big] \\ &=-\mathbb{E}_q\Big[-\log p(x_T) - \sum_{t>1} \log \Bigg\{\frac{p_\theta(x_{t-1}|x_t}{q(x_{t-1}|x_t, x_0)}\cdot\frac{q(x_{t-1}|x_0)} {q(x_t|x_0)}\Bigg\} -\log \big\{ \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\big\}\Big] \\ \end{align*} \]
Rewritting
\[ -\sum_{t>1}\log\Bigg\{\frac{q(x_{t-1}|x_0)} {q(x_t|x_0)}\Bigg\} = \log q(x_1|x_0) - \log q(x_T|x_0) \]
And substitution gives us
\[ \begin{align*} &=-\mathbb{E}_q\Big[-\log p(x_T) - \sum_{t>1} \log \Bigg\{\frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)}\cdot\frac{q(x_{t-1}|x_0)} {q(x_t|x_0)}\Bigg\} -\log \big\{ \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\big\}\Big] \\ &=-\mathbb{E}_q\Big[ - \log \Big\{\frac{p(x_T)}{q(x_T|x_0)}\Big\} - \sum_{t>1} \log \Bigg\{\frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)}\Bigg\} - \log\big\{p_\theta(x_0|x_1)\big\}\Big] \end{align*} \]
This means that if we want parameters \(\theta\) that minimises the log-likelihood, we need to make sure that these parameters minimises a lower bound which consist out of the following terms;
\[ \begin{align*} \mathbb{E}_q\Big[&D_{KL}\big[q(x_1|x_0)||p(x_1|x_0)\big] \\ &+ \sum_{t>1}D_{KL}\big[q(x_{t-1}|x_t, x_0) || p_\theta(x_{t-1}|x_t) \big] \\ &-\log p_\theta(x_0|x_1) \Big] \end{align*} \]
Which we can decompose into;
\[ \begin{align*} &L_T = D_{KL}\big[q(x_T|x_0)||p(x_T)\big] \\ &L_{t-1} = \sum_{t>1}D_{KL}\big[q(x_{t-1}|x_t, x_0) || p_\theta(x_{t-1}|x_t) \big]\\ &L_0 = - \log p_\theta(x_0|x_1) \end{align*} \]
This basically says that we can optimise the lower bound by only looking at \(L_{t-1}\), because \(L_T\) is constant during training and can be ignored.
Because both \(p_\theta\) and \(q\) are both Gaussians, the divergence becomes;
\[ L_{t-1} = \mathbb{E}_q\Big[\frac{1}{2\sigma^2_t}||\hat\mu(x_t, x_0) - \mu_\theta(x_t, t)||^2\Big] \]
This means that the most straightforward way to implement the reverse diffusion process is to simply predict \(\mu(x_t, x_0)\). This is the first way to implement a diffusion model.
The second way is rewriting to;
\[ \mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \Big(x_t - \frac{\beta_t}{\sqrt{1-\hat\alpha_t}}\epsilon_\theta(x_t, t) \Big) \]
The authors of the paper found this more convenient. Which leads to an even simpler version of the lower bound \[ L_{simple}(\theta) = \mathbb{E}\Big[||\epsilon - \epsilon_\theta(\sqrt{\hat\alpha_t}x_0 +\sqrt{1-\hat\alpha_t}\epsilon, t)||^2\Big] \]
In this way we don’t directly predict, given a variable \(x_t\) what was the previous \(x_{t-1}\), but we predict, given \(x_t\), what the noise \(\epsilon_{\theta}(x_t, t)\) at that given timestep. Intuitively, we predict which direction we have to move to make our sample less noisy. |Probabilistic Graphical Model Diagram
|Probabilistic Graphical Model Diagram|
UNET diagram from U-Net: Convolutional Networks for Biomedical Image Segmentation
So now we know what we want, namely \(\epsilon_\theta(x, t)\), a model that takes as input a distorted sample \(x_t\) and outputs the noise \(\epsilon\) that was added for timestep \(t\) and we know that we want to apply a mse loss on this model. So for this we use an UNET, just like the authors of the original diffusion models, and we use a convolution with a shift scale operation to condition this on the timestep.
This UNET consist out of a few building blocks which I will explain from small to big, starting with a basic convolutional block;
class ConvBlock(torch.nn.Module):
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: int = 1,
padding: int = 1,
bias: bool = False,
nrm: bool = True,
act: bool = True,
) -> None:
"""A basic convolutional block, convolution, activation and normalization. Uses Groupnorm.
"""
super().__init__()
self.conv = torch.nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias)
self.nrm_fn = torch.nn.GroupNorm(num_groups=8, num_channels=out_channels) if nrm else torch.nn.Identity()
self.act_fn = torch.nn.LeakyReLU(inplace=True) if act else torch.nn.Identity()
def forward(self, x) -> torch.Tensor:
return self.act_fn(self.nrm_fn(self.conv(x)))However, we want a model that can be conditioned on time \(t\), we use a shift-scale operation on the output of a normal convolution layer.
class TimeConvBlock(ConvBlock):
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
t_scale, t_shift = t.chunk(2, dim=1)
x = self.nrm_fn(self.conv(x))
return self.act_fn(x * (1 + t_scale[..., None, None]) + t_shift[..., None, None])These TimeConvBlocks will be combined to make a TimeResidualBlock.
class TimeResidualBlock(torch.nn.Module):
"""Basic residual block with time conditioning.
Time conditioning is done by employing a scale/shift operator on the projected features before the activation.
Contains 3 convolutional blocks, the last one without activation.
The residual connection is added to the last output.
"""
def __init__(self, channels: int) -> None:
super().__init__()
self.conv_block_1 = TimeConvBlock(in_channels=channels, out_channels=channels)
self.conv_block_2 = TimeConvBlock(in_channels=channels, out_channels=channels)
self.conv_block_3 = TimeConvBlock(in_channels=channels, out_channels=channels, act=False)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
r = x.clone()
x = self.conv_block_1(x=x,t=t)
x = self.conv_block_2(x=x,t=t)
return self.conv_block_3(x=x,t=t) + xBoth an an EncoderBlock and a DecoderBlock consist out of one TimeResidualBlock and a Resample Block. In the EncoderBlock the resample operation is done first, while in the DecoderBlock the resample operation is done last, such that the TimeResidualBlock operates on features maps of low resolution.
class ResidualUnetBlock(torch.nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
upsample: bool,
time_channels: int) -> None:
super().__init__()
self.upsample = upsample
if self.upsample:
#This means it is an decoder block, so we do the first the convolution block, then the resample block
self.residual_block = TimeResidualBlock(channels=in_channels)
self.resample_block = ConvUpsample(in_channels=in_channels, out_channels=out_channels)
else:
#This means that it is an encoder block so first resample, then convolutional block
self.resample_block = ConvDownSample(in_channels=in_channels, out_channels=out_channels)
self.residual_block = TimeResidualBlock(channels=out_channels)
self.timestep_block = torch.nn.Sequential(
torch.nn.Linear(time_channels, out_channels*2 if not upsample else in_channels*2),
torch.nn.SiLU(),
)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
h = x.clone()
t = self.timestep_block(t)
if not self.upsample:
x = self.resample_block(x)
x = self.residual_block(x=x, t=t)
if self.upsample:
x = self.resample_block(x)
return x, hCombining these blocks gives our model
class UNET(torch.nn.Module):
def __init__(self, in_channels: int, channels: list[int]) -> None:
super().__init__()
self.in_channels = in_channels
self.channels = channels
self.init_cnv = ConvBlock(in_channels=in_channels, out_channels=channels[0])
time_channels = channels[0] * 4
self.time_enc = torch.nn.Sequential(
SinusoidalPositionEmbeddings(time_channels=channels[0]*4),
torch.nn.Linear(channels[0]*4, channels[0]*4),
torch.nn.SiLU(),
)
#Encoder/decoder
self.encoder = torch.nn.ModuleList([])
self.decoder = torch.nn.ModuleList([])
for in_channels, out_channels in zip(channels[:-1], channels[1:]):
self.encoder.append(ResidualUnetBlock(in_channels, out_channels, time_channels=time_channels, upsample=False))
self.decoder.append(ResidualUnetBlock(out_channels*2, in_channels,time_channels=time_channels, upsample=True))
self.attn_block = Attention2D(in_channels=channels[-1], out_channels=channels[-1]*2, head_channels=16, n_heads=4)
self.attn_norm = torch.nn.GroupNorm(num_groups=1, num_channels=channels[-1]*2)
self.final_cnv = ConvBlock(in_channels=channels[0]*2, out_channels=self.in_channels, nrm=False, act=False, )
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
x = self.init_cnv(x)
t = self.time_enc(t)
skip_connection = []
for layer in self.encoder:
x, h = layer(x=x, t=t)
skip_connection.append(h)
x = self.attn_block(x)
x = self.attn_norm(x)
for layer in self.decoder[::-1]:
x, _ = layer(x=x, t=t)
x = torch.cat((x, skip_connection.pop(-1)), dim=1)
return self.final_cnv(x)batch_size = 64
eval_steps = 500
learning_rate = 0.0001
beta_start = 0.0001
beta_end = 0.015
diffusion_timesteps = 1000
n_trn_steps = 5_000transforms = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.0,), (1.0,))])
dataset = torchvision.datasets.MNIST(root=".", transform=transforms, train=True, download=True)
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=batch_size,
drop_last=True,
shuffle=True)def get_diffusion_dict_np(
beta_start: float, beta_end: float, diffusion_timesteps: int
) -> dict[str, np.ndarray]:
betas = np.linspace(beta_start, beta_end, diffusion_timesteps)
alphas = 1.0 - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.pad(alphas_cumprod[:-1], (1, 0), constant_values=1.0)
sqrt_recip_alphas = np.sqrt(1.0 / alphas)
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
return {
"betas": betas,
"alphas": alphas,
"alphas_cumprod": alphas_cumprod,
"alphas_cumprod_prev": alphas_cumprod_prev,
"sqrt_recip_alphas": sqrt_recip_alphas,
"sqrt_alphas_cumprod": sqrt_alphas_cumprod,
"sqrt_one_minus_alphas_cumprod": sqrt_one_minus_alphas_cumprod,
"posterior_variance": posterior_variance,
}def q_sample(
gaussian_dict_np: dict[str, np.ndarray],
x_start: np.ndarray,
timesteps: np.ndarray,
noise: np.ndarray,
) -> np.ndarray:
"""Forward process, this we add noise to x_start=x0.
The amount of noise we add is defined by the variance schedule defined in the gaussian_dict_np.
Args:
x_start (torch.Tensor): Start
t (int): Diffusion timestep, this value determines the scale of the noise we will add.
noise (torch.Tensor, optional): The noise to add, this is the amount of noise the network should predict.
Returns:
torch.Tensor: x_t
"""
sqrt_alphas_cumprod_t = gaussian_dict_np["sqrt_alphas_cumprod"][timesteps]
sqrt_one_minus_alphas_cumprod_t = gaussian_dict_np["sqrt_one_minus_alphas_cumprod"][
timesteps
]
return sqrt_alphas_cumprod_t.reshape(-1, 1, 1, 1) * x_start + sqrt_one_minus_alphas_cumprod_t.reshape(-1, 1, 1, 1) * noiseclass ConvBlock(torch.nn.Module):
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: int = 1,
padding: int = 1,
bias: bool = False,
nrm: bool = True,
act: bool = True,
) -> None:
"""A basic convolutional block, convolution, activation and normalization. Uses Groupnorm.
"""
super().__init__()
self.conv = torch.nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias)
self.nrm_fn = torch.nn.GroupNorm(num_groups=8, num_channels=out_channels) if nrm else torch.nn.Identity()
self.act_fn = torch.nn.LeakyReLU(inplace=True) if act else torch.nn.Identity()
def forward(self, x) -> torch.Tensor:
return self.act_fn(self.nrm_fn(self.conv(x)))class TimeConvBlock(ConvBlock):
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
t_scale, t_shift = t.chunk(2, dim=1)
x = self.nrm_fn(self.conv(x))
return self.act_fn(x * (1 + t_scale[..., None, None]) + t_shift[..., None, None])class ConvDownSample(torch.nn.Module):
def __init__(self, in_channels: int, out_channels: int) -> None:
super().__init__()
self.conv = torch.nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1, )
self.nrm_fn = torch.nn.GroupNorm(num_groups=8, num_channels=out_channels)
self.act_fn = torch.nn.LeakyReLU(inplace=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.act_fn(self.nrm_fn(self.conv(x)))class ConvUpsample(torch.nn.Module):
def __init__(self, in_channels: int, out_channels:int, ) -> None:
super().__init__()
self.conv = torch.nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, )
self.nrm_fn = torch.nn.GroupNorm(num_groups=8, num_channels=out_channels)
self.act_fn = torch.nn.LeakyReLU(inplace=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = torch.nn.functional.interpolate(x, scale_factor=2, )
return self.act_fn(self.nrm_fn(self.conv(x)))class TimeResidualBlock(torch.nn.Module):
"""Basic residual block with time conditioning.
Time conditioning is done by employing a scale/shift operator on the projected features before the activation.
Contains 3 convolutional blocks, the last one without activation.
The residual connection is added to the last output.
"""
def __init__(self, channels: int) -> None:
super().__init__()
self.conv_block_1 = TimeConvBlock(in_channels=channels, out_channels=channels)
self.conv_block_2 = TimeConvBlock(in_channels=channels, out_channels=channels)
self.conv_block_3 = TimeConvBlock(in_channels=channels, out_channels=channels, act=False)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
r = x.clone()
x = self.conv_block_1(x=x,t=t)
x = self.conv_block_2(x=x,t=t)
return self.conv_block_3(x=x,t=t) + xclass ResidualUnetBlock(torch.nn.Module):
def __init__(self,
in_channels: int,
out_channels: int,
upsample: bool,
time_channels: int) -> None:
super().__init__()
self.upsample = upsample
if self.upsample:
#This means it is an decoder block, so we do the first the convolution block, then the resample block
self.residual_block = TimeResidualBlock(channels=in_channels)
self.resample_block = ConvUpsample(in_channels=in_channels, out_channels=out_channels)
else:
#This means that it is an encoder block so first resample, then convolutional block
self.resample_block = ConvDownSample(in_channels=in_channels, out_channels=out_channels)
self.residual_block = TimeResidualBlock(channels=out_channels)
self.timestep_block = torch.nn.Sequential(
torch.nn.Linear(time_channels, out_channels*2 if not upsample else in_channels*2),
torch.nn.SiLU(),
)
def forward(self, x: torch.Tensor, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
h = x.clone()
t = self.timestep_block(t)
if not self.upsample:
x = self.resample_block(x)
x = self.residual_block(x=x, t=t)
if self.upsample:
x = self.resample_block(x)
return x, hclass Attention2D(torch.nn.Module):
def __init__(self,
in_channels: int,
out_channels: int,
head_channels: int, n_heads: int, ) -> None:
"""Attention block that computes the self attention for each feature.
First projects the features to Q, K, V embeddings.
Reshapes these blocks by collapsing the H, W axis.
Then computes a similarity score.
Then computes the output values.
"""
super().__init__()
self.embedding_dim = head_channels * n_heads
self.q_scale = head_channels ** -0.5
self.n_heads = n_heads
self.out_channels = out_channels
self.conv_q = torch.nn.Conv2d(in_channels=in_channels, out_channels=self.embedding_dim, kernel_size=1, stride=1, padding=0, bias=False)
self.conv_k = torch.nn.Conv2d(in_channels=in_channels, out_channels=self.embedding_dim, kernel_size=1, stride=1, padding=0, bias=False)
self.conv_v = torch.nn.Conv2d(in_channels=in_channels, out_channels=self.embedding_dim, kernel_size=1, stride=1, padding=0, bias=False)
self.conv_out = torch.nn.Conv2d(in_channels=self.embedding_dim, out_channels=self.out_channels, kernel_size=1, stride=1, padding=0, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
N, C, H, W = x.shape
q = self.conv_q(x)
k = self.conv_k(x)
v = self.conv_v(x)
# re-arrange the q, k, v from (N, C, H, W) to (N, n_heads, head_channels, HxW)
q = einops.rearrange(q, pattern="N (a b) H W -> N a b (H W)", a=self.n_heads) * self.q_scale
k = einops.rearrange(k, pattern="N (a b) H W -> N a b (H W)", a=self.n_heads)
v = einops.rearrange(v, pattern="N (a b) H W -> N a b (H W)", a=self.n_heads)
similarity_matrix = torch.einsum("N a b i, N a b j -> N a i j", q, k)
similarity_matrix = torch.nn.functional.softmax(similarity_matrix, dim=1)
output = torch.einsum("N a i j, N a d j -> N a i d", similarity_matrix, v)
output = einops.rearrange(output, "N a (H W) d->N (a d) H W", H=H, W=W)
return self.conv_out(output)class SinusoidalPositionEmbeddings(torch.nn.Module):
def __init__(self, time_channels: int) -> None:
super().__init__()
self.time_channels = time_channels
self.embeddings = torch.log(torch.Tensor([10_000])) / (self.time_channels // 2 - 1)
def forward(self, t: torch.Tensor) -> torch.Tensor:
embeddings = torch.exp(torch.arange(self.time_channels // 2, device=t.device) * -self.embeddings.to(t.device))
embeddings = t[:, None] * embeddings[None, :]
return torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)class UNET(torch.nn.Module):
def __init__(self,
in_channels: int,
channels: list[int]) -> None:
super().__init__()
self.in_channels = in_channels
self.channels = channels
self.init_cnv = ConvBlock(in_channels=in_channels, out_channels=channels[0])
time_channels = channels[0] * 4
self.time_enc = torch.nn.Sequential(
SinusoidalPositionEmbeddings(time_channels=channels[0]*4),
torch.nn.Linear(channels[0]*4, channels[0]*4),
torch.nn.SiLU(),
)
#Encoder/decoder
self.encoder = torch.nn.ModuleList([])
self.decoder = torch.nn.ModuleList([])
for in_channels, out_channels in zip(channels[:-1], channels[1:]):
self.encoder.append(ResidualUnetBlock(in_channels, out_channels, time_channels=time_channels, upsample=False))
self.decoder.append(ResidualUnetBlock(out_channels*2, in_channels,time_channels=time_channels, upsample=True))
self.attn_block = Attention2D(in_channels=channels[-1], out_channels=channels[-1]*2, head_channels=16, n_heads=4)
self.attn_norm = torch.nn.GroupNorm(num_groups=1, num_channels=channels[-1]*2)
self.final_cnv = ConvBlock(in_channels=channels[0]*2, out_channels=self.in_channels, nrm=False, act=False, )
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
x = self.init_cnv(x)
t = self.time_enc(t)
skip_connection = []
for layer in self.encoder:
x, h = layer(x=x, t=t)
skip_connection.append(h)
x = self.attn_block(x)
x = self.attn_norm(x)
for layer in self.decoder[::-1]:
x, _ = layer(x=x, t=t)
x = torch.cat((x, skip_connection.pop(-1)), dim=1)
return self.final_cnv(x)def batch_distort(x_trn: torch.Tensor, batch_size: int, t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if t == None:
t = torch.randint(0, diffusion_timesteps-1, (batch_size,))
noise = torch.randn_like(x_trn)
xt = q_sample(gaussian_dict_np=diffusion_dict_np, x_start = x_trn.numpy(), timesteps=t.numpy(), noise=noise.numpy())
return torch.from_numpy(xt).to(torch.float32), noise, tdef plot_q_sample(x: torch.Tensor, xt: torch.Tensor, ) -> None:
fig, axs = plt.subplots(1, 2)
axs[0].imshow(xt, cmap="binary")
axs[1].imshow(x, cmap="binary")
axs[0].set_title(r"$x_t$")
axs[1].set_title(r"$x_0$")
_ = [ax.axis("off") for ax in axs.flatten()]Okey, now we have everything defined we can finally train our model.
unet = UNET(in_channels=1, channels=[16, 32, 64])
optimizer = torch.optim.Adam(params=unet.parameters(), lr=learning_rate, )
diffusion_dict_np = get_diffusion_dict_np(beta_start=beta_start, beta_end=beta_end, diffusion_timesteps=diffusion_timesteps)iter_loader = iter(loader)
for step in range(n_trn_steps):
try:
x_trn, _ = next(iter_loader)
except StopIteration:
iter_loader = iter(loader)
x_trn, _ = next(iter_loader)
xt, noise, t = batch_distort(x_trn=x_trn, batch_size=batch_size, t=None)
predicted_noise = unet(x=xt, t=t)
loss = torch.nn.functional.mse_loss(predicted_noise, noise)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % eval_steps == 0:
print(f"{step}/{n_trn_steps} loss: {loss.item():.4f}")0/5000 loss: 0.9969
500/5000 loss: 0.0757
1000/5000 loss: 0.0661
1500/5000 loss: 0.0395
2000/5000 loss: 0.0466
2500/5000 loss: 0.0342
3000/5000 loss: 0.0432
3500/5000 loss: 0.0368
4000/5000 loss: 0.0354
4500/5000 loss: 0.0391Okey, now we have our model, we can sample some digits!
with torch.no_grad():
xt = torch.randn_like(xt)
for timestep in tqdm(reversed(range(0, diffusion_timesteps-1))):
t = torch.ones_like(xt[:, 0, 0, 0]) * timestep
et = unet(x=xt.float(), t=t.float())
model_mean = diffusion_dict_np["sqrt_recip_alphas"][timestep] * (xt - diffusion_dict_np["betas"][timestep] * et.detach().numpy() /diffusion_dict_np["sqrt_one_minus_alphas_cumprod"][timestep])
if timestep != 0:
xt = model_mean + np.sqrt(
diffusion_dict_np["posterior_variance"][timestep]
) * torch.randn_like(model_mean)
else:
xt = model_mean
0it [00:00, ?it/s]/tmp/ipykernel_35124/127156701.py:6: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)
model_mean = diffusion_dict_np["sqrt_recip_alphas"][timestep] * (xt - diffusion_dict_np["betas"][timestep] * et.detach().numpy() /diffusion_dict_np["sqrt_one_minus_alphas_cumprod"][timestep])
3it [00:00, 28.05it/s]
7it [00:00, 29.94it/s]
11it [00:00, 30.51it/s]
15it [00:00, 30.71it/s]
19it [00:00, 30.87it/s]
23it [00:00, 30.94it/s]
27it [00:00, 31.01it/s]
31it [00:01, 31.04it/s]
35it [00:01, 31.06it/s]
39it [00:01, 31.11it/s]
43it [00:01, 31.13it/s]
47it [00:01, 31.06it/s]
51it [00:01, 31.05it/s]
55it [00:01, 30.91it/s]
59it [00:01, 30.45it/s]
63it [00:02, 30.58it/s]
67it [00:02, 30.75it/s]
71it [00:02, 30.81it/s]
75it [00:02, 30.90it/s]
79it [00:02, 30.91it/s]
83it [00:02, 30.99it/s]
87it [00:02, 31.03it/s]
91it [00:02, 31.04it/s]
95it [00:03, 30.39it/s]
99it [00:03, 30.52it/s]
103it [00:03, 30.75it/s]
107it [00:03, 30.76it/s]
111it [00:03, 30.85it/s]
115it [00:03, 30.93it/s]
119it [00:03, 30.95it/s]
123it [00:03, 30.51it/s]
127it [00:04, 30.63it/s]
131it [00:04, 30.87it/s]
135it [00:04, 30.87it/s]
139it [00:04, 30.98it/s]
143it [00:04, 30.93it/s]
147it [00:04, 31.01it/s]
151it [00:04, 31.04it/s]
155it [00:05, 31.05it/s]
159it [00:05, 30.16it/s]
163it [00:05, 29.95it/s]
167it [00:05, 30.32it/s]
171it [00:05, 30.52it/s]
175it [00:05, 30.66it/s]
179it [00:05, 30.78it/s]
183it [00:05, 30.86it/s]
187it [00:06, 30.60it/s]
191it [00:06, 29.76it/s]
194it [00:06, 29.63it/s]
198it [00:06, 29.96it/s]
202it [00:06, 30.30it/s]
206it [00:06, 30.53it/s]
210it [00:06, 30.25it/s]
214it [00:06, 30.43it/s]
218it [00:07, 30.58it/s]
222it [00:07, 30.74it/s]
226it [00:07, 30.76it/s]
230it [00:07, 30.87it/s]
234it [00:07, 30.86it/s]
238it [00:07, 30.92it/s]
242it [00:07, 30.94it/s]
246it [00:08, 30.87it/s]
250it [00:08, 30.96it/s]
254it [00:08, 30.63it/s]
258it [00:08, 30.44it/s]
262it [00:08, 30.58it/s]
266it [00:08, 30.63it/s]
270it [00:08, 30.70it/s]
274it [00:08, 30.77it/s]
278it [00:09, 30.81it/s]
282it [00:09, 30.84it/s]
286it [00:09, 30.85it/s]
290it [00:09, 30.53it/s]
294it [00:09, 30.50it/s]
298it [00:09, 30.63it/s]
302it [00:09, 30.73it/s]
306it [00:09, 30.86it/s]
310it [00:10, 30.90it/s]
314it [00:10, 30.89it/s]
318it [00:10, 30.94it/s]
322it [00:10, 30.42it/s]
326it [00:10, 30.62it/s]
330it [00:10, 30.79it/s]
334it [00:10, 30.86it/s]
338it [00:11, 30.95it/s]
342it [00:11, 30.99it/s]
346it [00:11, 30.97it/s]
350it [00:11, 31.09it/s]
354it [00:11, 31.07it/s]
358it [00:11, 30.55it/s]
362it [00:11, 30.76it/s]
366it [00:11, 30.56it/s]
370it [00:12, 30.39it/s]
374it [00:12, 30.61it/s]
378it [00:12, 30.79it/s]
382it [00:12, 30.91it/s]
386it [00:12, 30.62it/s]
390it [00:12, 30.66it/s]
394it [00:12, 30.81it/s]
398it [00:12, 31.00it/s]
402it [00:13, 30.99it/s]
406it [00:13, 31.06it/s]
410it [00:13, 31.05it/s]
414it [00:13, 31.14it/s]
418it [00:13, 31.12it/s]
422it [00:13, 31.19it/s]
426it [00:13, 31.09it/s]
430it [00:13, 31.14it/s]
434it [00:14, 31.13it/s]
438it [00:14, 31.03it/s]
442it [00:14, 31.04it/s]
446it [00:14, 31.07it/s]
450it [00:14, 30.87it/s]
454it [00:14, 30.44it/s]
458it [00:14, 30.63it/s]
462it [00:15, 30.73it/s]
466it [00:15, 30.86it/s]
470it [00:15, 30.77it/s]
474it [00:15, 30.61it/s]
478it [00:15, 30.73it/s]
482it [00:15, 30.27it/s]
486it [00:15, 30.56it/s]
490it [00:15, 30.67it/s]
494it [00:16, 30.80it/s]
498it [00:16, 30.87it/s]
502it [00:16, 30.97it/s]
506it [00:16, 30.97it/s]
510it [00:16, 30.98it/s]
514it [00:16, 31.08it/s]
518it [00:16, 30.50it/s]
522it [00:16, 30.57it/s]
526it [00:17, 30.75it/s]
530it [00:17, 30.87it/s]
534it [00:17, 30.98it/s]
538it [00:17, 30.97it/s]
542it [00:17, 30.53it/s]
546it [00:17, 29.72it/s]
550it [00:17, 30.16it/s]
554it [00:18, 30.43it/s]
558it [00:18, 30.59it/s]
562it [00:18, 30.70it/s]
566it [00:18, 30.84it/s]
570it [00:18, 30.90it/s]
574it [00:18, 30.89it/s]
578it [00:18, 30.91it/s]
582it [00:18, 30.52it/s]
586it [00:19, 29.68it/s]
589it [00:19, 29.48it/s]
593it [00:19, 29.93it/s]
597it [00:19, 30.27it/s]
601it [00:19, 30.47it/s]
605it [00:19, 30.62it/s]
609it [00:19, 30.75it/s]
613it [00:19, 30.79it/s]
617it [00:20, 30.86it/s]
621it [00:20, 30.87it/s]
625it [00:20, 30.95it/s]
629it [00:20, 30.91it/s]
633it [00:20, 31.03it/s]
637it [00:20, 31.00it/s]
641it [00:20, 31.05it/s]
645it [00:20, 31.11it/s]
649it [00:21, 30.69it/s]
653it [00:21, 29.75it/s]
657it [00:21, 30.22it/s]
661it [00:21, 30.53it/s]
665it [00:21, 29.88it/s]
669it [00:21, 30.30it/s]
673it [00:21, 30.64it/s]
677it [00:22, 30.84it/s]
681it [00:22, 30.93it/s]
685it [00:22, 31.01it/s]
689it [00:22, 31.01it/s]
693it [00:22, 31.13it/s]
697it [00:22, 31.12it/s]
701it [00:22, 31.16it/s]
705it [00:22, 31.15it/s]
709it [00:23, 31.14it/s]
713it [00:23, 30.91it/s]
717it [00:23, 29.25it/s]
721it [00:23, 29.88it/s]
725it [00:23, 30.26it/s]
729it [00:23, 30.60it/s]
733it [00:23, 30.78it/s]
737it [00:23, 30.93it/s]
741it [00:24, 30.99it/s]
745it [00:24, 30.97it/s]
749it [00:24, 31.02it/s]
753it [00:24, 31.01it/s]
757it [00:24, 31.12it/s]
761it [00:24, 31.09it/s]
765it [00:24, 31.11it/s]
769it [00:25, 31.12it/s]
773it [00:25, 31.12it/s]
777it [00:25, 31.07it/s]
781it [00:25, 30.54it/s]
785it [00:25, 30.75it/s]
789it [00:25, 30.47it/s]
793it [00:25, 30.32it/s]
797it [00:25, 30.62it/s]
801it [00:26, 30.83it/s]
805it [00:26, 30.87it/s]
809it [00:26, 30.96it/s]
813it [00:26, 30.94it/s]
817it [00:26, 31.06it/s]
821it [00:26, 30.69it/s]
825it [00:26, 30.82it/s]
829it [00:26, 30.89it/s]
833it [00:27, 30.98it/s]
837it [00:27, 30.99it/s]
841it [00:27, 31.05it/s]
845it [00:27, 30.59it/s]
849it [00:27, 30.63it/s]
853it [00:27, 30.69it/s]
857it [00:27, 30.18it/s]
861it [00:28, 30.51it/s]
865it [00:28, 30.59it/s]
869it [00:28, 30.77it/s]
873it [00:28, 30.94it/s]
877it [00:28, 31.00it/s]
881it [00:28, 31.04it/s]
885it [00:28, 31.07it/s]
889it [00:28, 30.44it/s]
893it [00:29, 30.66it/s]
897it [00:29, 30.80it/s]
901it [00:29, 30.86it/s]
905it [00:29, 31.07it/s]
909it [00:29, 30.75it/s]
913it [00:29, 30.61it/s]
917it [00:29, 30.15it/s]
921it [00:29, 30.47it/s]
925it [00:30, 30.24it/s]
929it [00:30, 29.65it/s]
933it [00:30, 30.04it/s]
937it [00:30, 30.40it/s]
941it [00:30, 29.99it/s]
945it [00:30, 30.37it/s]
949it [00:30, 30.56it/s]
953it [00:31, 30.70it/s]
957it [00:31, 30.83it/s]
961it [00:31, 30.88it/s]
965it [00:31, 30.94it/s]
969it [00:31, 31.00it/s]
973it [00:31, 30.89it/s]
977it [00:31, 30.50it/s]
981it [00:31, 30.70it/s]
985it [00:32, 30.83it/s]
989it [00:32, 30.81it/s]
993it [00:32, 30.90it/s]
997it [00:32, 30.92it/s]
999it [00:32, 30.72it/s]fig, axs = plt.subplots(1, 5, figsize=(10, 5))
axs[0].imshow(xt[0, 0].numpy(), cmap="binary")
axs[1].imshow(xt[1, 0].numpy(), cmap="binary")
axs[2].imshow(xt[2, 0].numpy(), cmap="binary")
axs[3].imshow(xt[3, 0].numpy(), cmap="binary")
axs[4].imshow(xt[4, 0].numpy(), cmap="binary")
This amazing blogpost helped me understand (most of) the math behind diffusion models and inspired me to make my own blogpost combining both the math and python code.
@article{weng2021diffusion,
title = "What are diffusion models?",
author = "Weng, Lilian",
journal = "lilianweng.github.io",
year = "2021",
month = "Jul",
url = "https://lilianweng.github.io/posts/2021-07-11-diffusion-models/"
}This blogpost explained the reparameterization trick.
@article{weng2018VAE,
title = "From Autoencoder to Beta-VAE",
author = "Weng, Lilian",
journal = "lilianweng.github.io",
year = "2018",
url = "https://lilianweng.github.io/posts/2018-08-12-vae/"
}This work and the blogpost inspired most of the schematics
@inproceedings{
song2021scorebased,
title={Score-Based Generative Modeling through Stochastic Differential Equations},
author={Yang Song and Jascha Sohl-Dickstein and Diederik P Kingma and Abhishek Kumar and Stefano Ermon and Ben Poole},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=PxTIG12RRHS}
}The pytorch code I wrote for this blog was written from scratch to help me explain the algorithm, but it was inspired by the code from this work
@article{DBLP:journals/corr/abs-2102-09672,
author = {Alex Nichol and
Prafulla Dhariwal},
title = {Improved Denoising Diffusion Probabilistic Models},
journal = {CoRR},
volume = {abs/2102.09672},
year = {2021},
url = {https://arxiv.org/abs/2102.09672},
eprinttype = {arXiv},
eprint = {2102.09672},
timestamp = {Wed, 24 Feb 2021 15:42:45 +0100},
biburl = {https://dblp.org/rec/journals/corr/abs-2102-09672.bib},
bibsource = {dblp computer science bibliography, https://dblp.org}
}And ofcourse, to the work on diffusion models;
@article{DBLP:journals/corr/abs-2006-11239,
author = {Jonathan Ho and
Ajay Jain and
Pieter Abbeel},
title = {Denoising Diffusion Probabilistic Models},
journal = {CoRR},
volume = {abs/2006.11239},
year = {2020},
url = {https://arxiv.org/abs/2006.11239},
eprinttype = {arXiv},
eprint = {2006.11239},
timestamp = {Tue, 23 Jun 2020 17:57:22 +0200},
biburl = {https://dblp.org/rec/journals/corr/abs-2006-11239.bib},
bibsource = {dblp computer science bibliography, https://dblp.org}
}