扩散模型DDPM开源代码的剖析【对应公式与作者给的开源项目，diffusion model】

一、简介

论文地址：https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html
项目地址：https://github.com/hojonathanho/diffusion
公式推导参考这篇博客：https://blog.csdn.net/qq_45934285/article/details/129107994?spm=1001.2014.3001.5502
本文主要对扩散模型的关键公式给出原代码帮助理解和学习。有pytorch和TensorFlow版。
原作者给的代码不太好理解，给出了pytorch的好理解一些。

二、扩散过程：输入是x_0和时刻num_steps，输出是x_t

首先值得注意的是：x_0是一个二维数组，例如这里给的是一个10000行2列的数组，即每一行代表一个点。
这里取了s_curve的x轴和z轴的坐标，用点表示看起来就像一个s型

s_curve,_ = make_s_curve(10**4,noise=0.1)
s_curve = s_curve[:,[0,2]]/10.0
dataset = torch.Tensor(s_curve).float()

扩散过程其实就是一个不断加噪的过程，其不含参。可以给出最终公式。
$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0$ + ${\sqrt{1-\overline{\alpha }}}_t$ ${\overline{z}}_t$

在t不断变大的时候${\beta }_t$越来越大，${\alpha }_t=1-{\beta }_t$越来越小。即t增大的时候上面公式的前一项系数越来越小，后一项系数越来越大不断接近一个${\overline{z}}_t$的高斯分布。

代码来自diffusion_tf/diffusion_utils_2.py

  def q_sample(self, x_start, t, noise=None):"""Diffuse the data (t == 0 means diffused for 1 step)"""if noise is None:noise = tf.random_normal(shape=x_start.shape)assert noise.shape == x_start.shapereturn (self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)

pytorch：

#计算任意时刻的x采样值，基于x_0和重参数化
def q_x(x_0,t):"""可以基于x[0]得到任意时刻t的x[t]"""noise = torch.randn_like(x_0)# 创建了一个与 x_0 张量具有相同形状的名为 noise 的张量，并且该张量的值是从标准正态分布中随机采样得到的。alphas_t = alphas_bar_sqrt[t]alphas_1_m_t = one_minus_alphas_bar_sqrt[t]return (alphas_t * x_0 + alphas_1_m_t * noise)#在x[0]的基础上添加噪声

可见对于求$x_t$的公式最难理解的就是代码如何实现${\overline{z}}_t$在代码中是创建了一个与 x_0 张量具有**相同形状**的名为 noise 的张量，并且该张量的值是从标准正态分布中随机采样得到的。这个noise其元素的值是从均值为0、标准差为1的正态分布中随机采样得到的。这个张量可以被用于实现噪声注入，数据增强等操作，也可以被用于一些随机化的算法中。

值得一提的是原项目中用num_diffusion_timesteps=1000来表示t，假如num_steps=100，那么很多需要用到的参数都可以提前算出来。

三、逆扩散过程：输入x_t，不断采样最终输出x_0

最终公式是：

$$q\left(X_{t-1}\mid X_t{X}_0\right)=N\left(X_{t-1};\frac{1}{\sqrt{{\alpha }_t}}(X_t-\frac{{\beta }_t}{\sqrt{\left(1-{\bar{\alpha }}_t\right)}}Z),\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\right),Z\sim N(0,I)$$
在论文中方差设置为一个常数${\beta }_t$或${\tilde{\beta }}_t$其中：
$${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$$
因此可训练的参数只存在与其均值之中。

就是这个公式，方差变为${\beta }_t$，其中${ϵ}_{\theta }$是模型model

def p_sample(model,x,t,betas,one_minus_alphas_bar_sqrt):"""从x[T]采样t时刻的重构值"""t = torch.tensor([t])coeff = betas[t] / one_minus_alphas_bar_sqrt[t]eps_theta = model(x,t)mean = (1/(1-betas[t]).sqrt())*(x-(coeff*eps_theta))z = torch.randn_like(x)sigma_t = betas[t].sqrt()sample = mean + sigma_t * zreturn (sample)

代码来自diffusion_tf/diffusion_utils_2.py

  def p_sample(self, denoise_fn, *, x, t, noise_fn, clip_denoised=True, return_pred_xstart: bool):"""Sample from the model"""model_mean, _, model_log_variance, pred_xstart = self.p_mean_variance(denoise_fn, x=x, t=t, clip_denoised=clip_denoised, return_pred_xstart=True)noise = noise_fn(shape=x.shape, dtype=x.dtype)assert noise.shape == x.shape# no noise when t == 0nonzero_mask = tf.reshape(1 - tf.cast(tf.equal(t, 0), tf.float32), [x.shape[0]] + [1] * (len(x.shape) - 1))sample = model_mean + nonzero_mask * tf.exp(0.5 * model_log_variance) * noiseassert sample.shape == pred_xstart.shapereturn (sample, pred_xstart) if return_pred_xstart else sample

循环恢复。
可见初始的x_t完全是一个随机噪声。torch.randn(shape)
cur_x可以看做是一个当前的采样，是一个二维数组，就是上面说的10000行2列。
然后x_seq可以看做是一个三维数组，即元素为cur_x的一个数组。
i是时刻，从n_steps的反向开始。

def p_sample_loop(model,shape,n_steps,betas,one_minus_alphas_bar_sqrt):"""从x[T]恢复x[T-1]、x[T-2]|...x[0]"""cur_x = torch.randn(shape)x_seq = [cur_x]for i in reversed(range(n_steps)):cur_x = p_sample(model,cur_x,i,betas,one_minus_alphas_bar_sqrt)x_seq.append(cur_x)return x_seq

代码来自diffusion_tf/diffusion_utils_2.py

  def p_sample_loop(self, denoise_fn, *, shape, noise_fn=tf.random_normal):"""Generate samples"""assert isinstance(shape, (tuple, list))i_0 = tf.constant(self.num_timesteps - 1, dtype=tf.int32)img_0 = noise_fn(shape=shape, dtype=tf.float32)_, img_final = tf.while_loop(cond=lambda i_, _: tf.greater_equal(i_, 0),body=lambda i_, img_: [i_ - 1,self.p_sample(denoise_fn=denoise_fn, x=img_, t=tf.fill([shape[0]], i_), noise_fn=noise_fn, return_pred_xstart=False)],loop_vars=[i_0, img_0],shape_invariants=[i_0.shape, img_0.shape],back_prop=False)assert img_final.shape == shapereturn img_final

那么最终得到的x_seq就是最终从噪声恢复出的x_T到x_0序列。

恢复图像过程中利用了：

1. ${\beta }_t$数组，其可以引申出很多参数(带$\alpha$的)
2. num_steps即时刻。
3. 一个model

参数只有model，那么我们来看model到底是什么。

四、具体参考算法流程图

五、模型model和损失函数（最重要！）

1、先看损失函数

其中x_0就是一个batch的数据，batch_size取x_0的行数，就是batch_size
值得注意的是这里t是一个列向量，那么a也是一个列向量，其与x_0相乘的时候用了，广播机制（第二列完全复制第一列）然后与x_0对应位置上的元素相乘。

def diffusion_loss_fn(model,x_0,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,n_steps):"""对任意时刻t进行采样计算loss"""batch_size = x_0.shape[0]    #对一个batchsize样本生成随机的时刻t 覆盖到更多不同的tt = torch.randint(0,n_steps,size=(batch_size//2,))t = torch.cat([t,n_steps-1-t],dim=0)# [batchsize]行向量t = t.unsqueeze(-1)#压缩最后的时刻1[batchsize，1]列向量  #x0的系数a = alphas_bar_sqrt[t]    #eps的系数aml = one_minus_alphas_bar_sqrt[t]    #生成随机噪音epse = torch.randn_like(x_0)     #构造模型的输入x = x_0*a+e*aml    #送入模型，得到t时刻的随机噪声预测值output = model(x,t.squeeze(-1))   #这里t又变为一维向量[batchsize] #与真实噪声一起计算误差，求平均值return (e - output).square().mean()

其调用过程见：
重要的：batch_x是dataset的一个batch_size行的数据。例如dataset是10000行，batch_size=128那么batch_x是dataset的128行。需要注意的是，如果 dataset 的大小不能被 batch_size 整除，那么最后一个批次的大小可能会小于 batch_size。即一个epoch的batch数量为dataset.size()/batch_size向上取整。

dataloader = torch.utils.data.DataLoader(dataset,batch_size=batch_size,shuffle=True)
for idx,batch_x in enumerate(dataloader):loss = diffusion_loss_fn(model,batch_x,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,num_steps)

dataloader这行代码是使用 PyTorch 中的 DataLoader 函数来创建一个数据加载器对象。这个对象可以用来迭代访问输入数据，以便将其输入到神经网络中进行训练或评估。

dataset：一个包含输入数据和对应标签的数据集对象。这个对象需要实现 getitem 和 len 方法，以便能够通过索引或长度访问数据集中的数据。
batch_size：一个整数值，表示每个批次的数据量。
shuffle：一个布尔值，表示是否需要在每个 epoch 开始前随机打乱数据集中的数据。

这个函数返回的是一个数据加载器对象，可以通过迭代器的方式来访问其中的数据。例如，可以使用 for 循环来遍历整个数据集，每次迭代会返回一个批次的数据和对应的标签。这样就可以方便地将数据输入到神经网络中进行训练或评估。

for idx,batch_x in enumerate(dataloader):

这行代码是使用 enumerate 函数和 DataLoader 对象来遍历整个数据集，并按照 batch_size 的大小分成若干批次进行迭代。

具体来说，enumerate 函数用于同时返回迭代对象中的元素以及它们的索引。在这里，dataloader 对象就是要被迭代的对象，而 idx 变量则表示当前迭代的批次索引。batch_x 变量则是当前批次中的输入数据，它是一个包含 batch_size 个样本的张量对象。

需要注意的是，在使用 enumerate 函数时，可以通过设置 start 参数来指定起始索引的值，默认为 0。例如，如果设置 start=1，那么第一个迭代的索引就是 1 而不是 0

2、model（看解释）

这里给出了一个比较简单的model，论文中使用的是Unet结构的模型。具体可以参考：https://nn.labml.ai/diffusion/ddpm/unet.html
具体训练过程可以参考：https://nn.labml.ai/diffusion/ddpm/experiment.html

import torch
import torch.nn as nnclass MLPDiffusion(nn.Module):def __init__(self,n_steps,num_units=128):super(MLPDiffusion,self).__init__()self.linears = nn.ModuleList([nn.Linear(2,num_units),nn.ReLU(),nn.Linear(num_units,num_units),nn.ReLU(),nn.Linear(num_units,num_units),nn.ReLU(),nn.Linear(num_units,2),])self.step_embeddings = nn.ModuleList([nn.Embedding(n_steps,num_units),nn.Embedding(n_steps,num_units),nn.Embedding(n_steps,num_units),])
#     第6步的modeldef forward(self,x,t):
#         x = x_0for idx,embedding_layer in enumerate(self.step_embeddings):t_embedding = embedding_layer(t)x = self.linears[2*idx](x)x += t_embeddingx = self.linears[2*idx+1](x)x = self.linears[-1](x)return x

首先：

batch_size = 128
dataloader = torch.utils.data.DataLoader(dataset,batch_size=batch_size,shuffle=True)
num_epoch = 4000 # 4000个散点图
plt.rc('text',color='blue')
model = MLPDiffusion(num_steps)#输出维度是2，输入是x和step
optimizer = torch.optim.Adam(model.parameters(),lr=1e-3)

然后：训练这个model

for t in range(num_epoch):for idx,batch_x in enumerate(dataloader):loss = diffusion_loss_fn(model,batch_x,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,num_steps)optimizer.zero_grad()loss.backward()torch.nn.utils.clip_grad_norm_(model.parameters(),1.)optimizer.step()

该模型的参数包括多个线性层和Embedding层的权重和偏置，在代码中用nn.ModuleList保存。在优化器更新时，通过model.parameters()方法获取所有可更新的参数，包括线性层和Embedding层的权重和偏置。因此，优化器更新的参数就是这些权重和偏置。

全连接层可以对输入数据进行线性变换，并输出与变换后数据维度相同的结果。例如，对于输入向量 $x$，全连接层可以通过将其与权重矩阵 $W$ 相乘，加上偏置向量 $b$，并应用某种非线性激活函数 $\sigma$，来计算输出向量 $y$：

$$y=\sigma (Wx+b)$$
更新的是这个W和b

Embedding 的参数也是模型中需要被训练的权重矩阵，它的维度是 (vocab_size, embedding_dim)，其中 vocab_size 是词汇表大小，embedding_dim 是词嵌入的维度。Embedding 的参数在训练过程中被不断地更新，以最小化模型在训练集上的损失函数。每个词都被编码成一个 embedding_dim 维的向量，这些向量将作为输入传递到后续的网络层中。Embedding 层通过学习将每个词映射到一个低维空间中的向量表示，从而使得相似的词在这个向量空间中也更加接近，提高了模型的泛化能力。

更新的是(vocab_size, embedding_dim)

optimizer.zero_grad()的作用是将模型参数的梯度清零。在每次反向传播计算梯度之前，需要先清除之前的梯度，否则之前计算的梯度会累加到当前计算中，导致梯度计算出错。因此，在每次训练之前，需要调用zero_grad()来清除梯度。

loss.backward()是PyTorch中用于计算梯度的函数。在神经网络训练中，首先需要将输入数据通过前向传递计算出预测值，然后根据预测值和真实标签计算出损失函数值。接下来，需要计算损失函数对于模型参数的梯度，以便使用优化算法对参数进行更新。==loss.backward()函数的作用就是计算损失函数对于所有需要梯度的参数的导数。==在调用该函数之前，需要将所有需要梯度计算的参数设置requires_grad=True。该函数计算得到的梯度会保存在参数的.grad属性中。通常，计算完梯度后需要对其进行裁剪，以避免梯度爆炸问题，然后使用优化器对模型参数进行更新。

torch.nn.utils.clip_grad_norm_()是PyTorch中用于梯度裁剪的函数。梯度裁剪是一种防止梯度爆炸的技术，它限制了梯度的最大范数，从而避免梯度值过大。函数的第一个参数是一个可迭代对象，通常是模型的参数列表，第二个参数是梯度的最大范数。在该代码中，model.parameters()返回模型的参数列表，1.指定了梯度的最大范数为1。因此，该函数的作用是将模型参数的梯度限制在最大范数为1的范围内。梯度裁剪通常在计算梯度后和参数更新之前进行，以保证梯度不会过大。

optimizer.step()是PyTorch中用于更新模型参数的函数。在每次梯度计算之后，需要使用优化器来更新模型的参数。优化器根据梯度和学习率等参数来计算参数更新的值，并将计算得到的值应用到模型参数中。调用optimizer.step()函数可以实现该操作。该函数会更新优化器内部的参数状态，以便下一次迭代使用。

在Python中，super() 函数是用来调用父类(超类)的一个方法。在面向对象编程中，通常使用 super() 函数来初始化父类的构造方法。在这里，super(MLPDiffusion,self).init() 调用了父类 nn.Module 的构造函数，相当于显式调用 nn.Module 的构造函数，初始化了当前类 MLPDiffusion 的基类 nn.Module。这样做是为了确保子类继承父类中的属性和方法，并且能够正确地使用它们。

for idx,embedding_layer in enumerate(self.step_embeddings):

这行代码是在遍历self.step_embeddings列表中的每个embedding层，同时记录下每个embedding层在列表中的索引idx。

在MLPDiffusion模型的forward函数中，self.step_embeddings是一个由三个Embedding层组成的ModuleList，每个Embedding层的作用是将当前时间步t转化为一个num_units维的向量。在模型的forward函数中，需要遍历这三个Embedding层，并将其输出的向量与模型的输入x相加。这样做的目的是为了将当前时间步t的信息融入到模型的计算中。因此，遍历每个Embedding层，并将其输出的向量与模型的输入x相加，是非常必要的。

x = self.linears[2*idx](x)

这行代码的作用是对模型的输入x进行线性变换，即将x通过一个全连接层（线性层）进行变换。其中，2idx表示要使用==第idx个Embedding层的输出，==而self.linears[2idx]则表示要使用与之对应的全连接层。

需要注意的是，由于self.linears是一个由多个层组成的ModuleList，而不是一个单独的层，因此需要通过下标来获取其中的某个层。2*idx表示要使用第idx个Embedding层的输出，因为每个Embedding层输出的维度都是num_units，所以输入到全连接层的维度也是num_units。这里采用的激活函数是ReLU。

第idx个Embedding层的输出是一个维度为（batch_size，num_units）的张量，表示对输入的时间步t进行了嵌入（embedding）后的结果。其中，batch_size表示输入数据的批次大小，num_units是预定义的嵌入向量的维度，这里为128。Embedding层将整数编码转换为密集向量，这些密集向量在整个模型的训练过程中逐渐学习得到，类似于单词的词向量。在本模型中，使用三个Embedding层分别嵌入了当前时间步t，前一时间步t-1和后一时间步t+1，这些嵌入向量会在全连接层中与输入向量x进行加权求和，以产生输出。

x = self.linears[-1](x)

这行代码对应的是在模型的最后一层加上一个全连接层，输出维度为2。该全连接层对前面所有层的输出进行线性变换，并输出一个维度为2的向量作为最终的预测值。

ReLU它是一个激活函数，用于在神经网络的前向传播过程中对输入进行非线性变换。ReLU函数的形式为$f(x)=\max (0,x)$，可以将输入的负值部分清零，保留正值部分。在反向传播中，ReLU层的导数可以被有效地计算，因此可以通过反向传播算法对模型的其他参数进行更新。

六、损失函数的推导

我们根据负对数似然优化 ELBO（来自简森不等式）。
$$\begin{gathered} \mathbb{E}[-\log p_{\theta }(x_0)]\le {\mathbb{E}}_q[-\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}] \\ =L \end{gathered}$$
损失可以改写如下。
$$\begin{gathered} L={\mathbb{E}}_q[-\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}] \\ ={\mathbb{E}}_q[-\log p(x_T)-\sum _{t=1}^T\log \frac{p_{\theta }(x_{t-1}|x_t)}{q(x_t|x_{t-1})}] \\ ={\mathbb{E}}_q[-\log \frac{p(x_T)}{q(x_T|x_0)}-\sum _{t=2}^T\log \frac{p_{\theta }(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}-\log p_{\theta }(x_0|x_1)] \\ ={\mathbb{E}}_q[D_{KL}(q(x_T|x_0)\Vert p(x_T))+\sum _{t=2}^T{D}_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))-\log p_{\theta }(x_0|x_1)] \end{gathered}$$
$D_{KL}(q(x_T|x_0)\Vert p(x_T))$ 是常数，因为我们保持 ${\beta }_1,\dots ,{\beta }_T$ 不变。

计算 $L_{t-1}=D_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))$

以 $x_0$ 为条件的前向过程后验是，
$$\begin{gathered} q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_t\mathbf{I}) \\ {\tilde{\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 \\ {\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t \end{gathered}$$
该论文设置 ${\Sigma }_{\theta }(x_t,t)={\sigma }_t^2\mathbf{I}$ 其中 ${\sigma }_t^2$ 设置为常量 ${\beta }_t$ 或 ${\tilde{\beta }}_t$。
然后,
$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t^2\mathbf{I})$$
对于给定噪声 $ϵ\sim \mathcal{N}(0,\mathbf{I})$ 使用 $q(x_t|x_0)$
$$\begin{gathered} x_t(x_0,ϵ)=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ \\ {x}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t(x_0,ϵ)-\sqrt{1-{\bar{\alpha }}_t}ϵ) \end{gathered}$$
这给出了：
$$\begin{gathered} L_{t-1}=D_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t)) \\ ={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}\Vert \tilde{\mu }(x_t,x_0)-{\mu }_{\theta }(x_t,t){\Vert }^2] \\ ={\mathbb{E}}_{x_0,ϵ}[\frac{1}{2{\sigma }_t^2}\Vert \frac{1}{\sqrt{{\alpha }_t}}(x_t(x_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)-{\mu }_{\theta }(x_t(x_0,ϵ),t){\Vert }^2] \end{gathered}$$
使用模型重新参数化以预测噪声
$$\begin{gathered} {\mu }_{\theta }(x_t,t)=\tilde{\mu }(x_t,\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t)\left)\right) \\ =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)) \end{gathered}$$
其中 ${ϵ}_{\theta }$ 是一个学习函数，它在给定 $(x_t,t)$ 的情况下预测 $ϵ$。

这给出了：
$$L_{t-1}={\mathbb{E}}_{x_0,ϵ}[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2]$$
也就是说，我们正在训练预测噪声。
Simplified loss
$$L_{\text{simple}}(\theta )={\mathbb{E}}_{t,x_0,ϵ}[\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2]$$
这最小化 $-\log p_{\theta }(x_0|x_1)$ when $t=1$ 和 $L_{t-1}$ for $t>1$ 丢弃 在 $L_{t-1}$ 中加权。丢弃权重 $\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}$ 增加赋予较高 $t$（具有较高噪声水平）的权重，从而提高样本质量。

如果觉得有帮助的话，记得点赞收藏支持一下哦。如果有问题请在评论区一起交流讨论。
Reference：
https://www.bilibili.com/video/BV1b541197HX/?spm_id_from=333.999.0.0
https://nn.labml.ai/diffusion/ddpm/index.html