扩散模型DDPM开源代码的剖析【对应公式与作者给的开源项目,diffusion model】
- 一、简介
- 二、扩散过程:输入是x_0和时刻num_steps,输出是x_t
- 三、逆扩散过程:输入x_t,不断采样最终输出x_0
- 四、具体参考算法流程图
- 五、模型model和损失函数(最重要!)
- 1、先看损失函数
- 2、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 = α ‾ t x 0 x_t=\sqrt {\overline{\alpha}_ {t}}x_ {0} xt=αtx0 + 1 − α ‾ t \sqrt {1-\overline {\alpha} }_ {t} 1−αt z ‾ t \overline {z}_t zt
在t不断变大的时候 β t \beta_t βt越来越大, α t = 1 − β t \alpha_t=1-\beta_t αt=1−βt越来越小。即t增大的时候上面公式的前一项系数越来越小,后一项系数越来越大不断接近一个 z ‾ t \overline {z}_t zt的高斯分布。
代码来自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.shape
return (
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 x_t xt的公式最难理解的就是代码如何实现 z ‾ t \overline {z}_t zt在代码中是创建了一个与 x_0 张量具有**相同形状**的名为 noise 的张量,并且该张量的值是从标准正态分布中随机采样得到的。这个noise其元素的值是从均值为0、标准差为1的正态分布中随机采样得到的。这个张量可以被用于实现噪声注入,数据增强等操作,也可以被用于一些随机化的算法中。
值得一提的是原项目中用num_diffusion_timesteps=1000来表示t,假如num_steps=100,那么很多需要用到的参数都可以提前算出来。
三、逆扩散过程:输入x_t,不断采样最终输出x_0
最终公式是:
q
(
X
t
−
1
∣
X
t
X
0
)
=
N
(
X
t
−
1
;
1
α
t
(
X
t
−
β
t
(
1
−
α
ˉ
t
)
Z
)
,
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
)
,
Z
∼
N
(
0
,
I
)
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)
q(Xt−1∣XtX0)=N(Xt−1;αt1(Xt−(1−αˉt)βtZ),1−αˉt1−αˉt−1βt),Z∼N(0,I)
在论文中方差设置为一个常数
β
t
\beta _t
βt或
β
~
t
\tilde{\beta }_t
β~t其中:
β
~
t
=
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
\tilde{\beta }_t=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \beta_{t}
β~t=1−αˉt1−αˉt−1βt
因此可训练的参数只存在与其均值之中。
就是这个公式,方差变为
β
t
\beta_t
βt,其中
ϵ
θ
\epsilon_\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 * z
return (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 == 0
nonzero_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) * noise
assert sample.shape == pred_xstart.shape
return (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 == shape
return img_final
那么最终得到的x_seq就是最终从噪声恢复出的x_T到x_0序列。
恢复图像过程中利用了:
- β t \beta_t βt数组,其可以引申出很多参数(带 α \alpha α的)
- num_steps即时刻。
- 一个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 覆盖到更多不同的t
t = 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]
#生成随机噪音eps
e = 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 nn
class 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步的model
def forward(self,x,t):
# x = x_0
for idx,embedding_layer in enumerate(self.step_embeddings):
t_embedding = embedding_layer(t)
x = self.linears[2*idx](x)
x += t_embedding
x = 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 x x,全连接层可以通过将其与权重矩阵 W W W 相乘,加上偏置向量 b b b,并应用某种非线性激活函数 σ \sigma σ,来计算输出向量 y y y:
y
=
σ
(
W
x
+
b
)
y = \sigma(Wx + b)
y=σ(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 ) f(x) = \max(0, x) f(x)=max(0,x),可以将输入的负值部分清零,保留正值部分。在反向传播中,ReLU层的导数可以被有效地计算,因此可以通过反向传播算法对模型的其他参数进行更新。
六、损失函数的推导
我们根据负对数似然优化 ELBO(来自简森不等式)。
E
[
−
log
p
θ
(
x
0
)
]
≤
E
q
[
−
log
p
θ
(
x
0
:
T
)
q
(
x
1
:
T
∣
x
0
)
]
=
L
\mathbb{E}[-\log \textcolor{lightgreen}{p_\theta}(x_0)] \le \mathbb{E}_q [ -\log \frac{\textcolor{lightgreen}{p_\theta}(x_{0:T})}{q(x_{1:T}|x_0)} ] \\ =L
E[−logpθ(x0)]≤Eq[−logq(x1:T∣x0)pθ(x0:T)]=L
损失可以改写如下。
L
=
E
q
[
−
log
p
θ
(
x
0
:
T
)
q
(
x
1
:
T
∣
x
0
)
]
=
E
q
[
−
log
p
(
x
T
)
−
∑
t
=
1
T
log
p
θ
(
x
t
−
1
∣
x
t
)
q
(
x
t
∣
x
t
−
1
)
]
=
E
q
[
−
log
p
(
x
T
)
q
(
x
T
∣
x
0
)
−
∑
t
=
2
T
log
p
θ
(
x
t
−
1
∣
x
t
)
q
(
x
t
−
1
∣
x
t
,
x
0
)
−
log
p
θ
(
x
0
∣
x
1
)
]
=
E
q
[
D
K
L
(
q
(
x
T
∣
x
0
)
∥
p
(
x
T
)
)
+
∑
t
=
2
T
D
K
L
(
q
(
x
t
−
1
∣
x
t
,
x
0
)
∥
p
θ
(
x
t
−
1
∣
x
t
)
)
−
log
p
θ
(
x
0
∣
x
1
)
]
L = \mathbb{E}_q [ -\log \frac{\textcolor{lightgreen}{p_\theta}(x_{0:T})}{q(x_{1:T}|x_0)} ] \\ = \mathbb{E}_q [ -\log p(x_T) - \sum_{t=1}^T \log \frac{\textcolor{lightgreen}{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{\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)} -\log \textcolor{lightgreen}{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 \textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)) -\log \textcolor{lightgreen}{p_\theta}(x_0|x_1)]
L=Eq[−logq(x1:T∣x0)pθ(x0:T)]=Eq[−logp(xT)−t=1∑Tlogq(xt∣xt−1)pθ(xt−1∣xt)]=Eq[−logq(xT∣x0)p(xT)−t=2∑Tlogq(xt−1∣xt,x0)pθ(xt−1∣xt)−logpθ(x0∣x1)]=Eq[DKL(q(xT∣x0)∥p(xT))+t=2∑TDKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))−logpθ(x0∣x1)]
D
K
L
(
q
(
x
T
∣
x
0
)
∥
p
(
x
T
)
)
D_{KL}(q(x_T|x_0) \Vert p(x_T))
DKL(q(xT∣x0)∥p(xT)) 是常数,因为我们保持
β
1
,
…
,
β
T
\beta_1, \dots, \beta_T
β1,…,βT 不变。
计算 L t − 1 = D K L ( q ( x t − 1 ∣ x t , x 0 ) ∥ p θ ( x t − 1 ∣ x t ) ) L_{t-1} = D_{KL}(q(x_{t-1}|x_t,x_0) \Vert \textcolor{lightgreen} {p_\theta}(x_{t-1}| x_t)) Lt−1=DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))
以
x
0
x_0
x0 为条件的前向过程后验是,
q
(
x
t
−
1
∣
x
t
,
x
0
)
=
N
(
x
t
−
1
;
μ
~
t
(
x
t
,
x
0
)
,
β
~
t
I
)
μ
~
t
(
x
t
,
x
0
)
=
α
ˉ
t
−
1
β
t
1
−
α
ˉ
t
x
0
+
α
t
(
1
−
α
ˉ
t
−
1
)
1
−
α
ˉ
t
x
t
β
~
t
=
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
q(x_{t-1}|x_t, x_0) = \mathcal{N} \Big(x_{t-1}; \tilde\mu_t(x_t, x_0), \tilde\beta_t \mathbf{I} \Big) \\ \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
q(xt−1∣xt,x0)=N(xt−1;μ~t(xt,x0),β~tI)μ~t(xt,x0)=1−αˉtαˉt−1βtx0+1−αˉtαt(1−αˉt−1)xtβ~t=1−αˉt1−αˉt−1βt
该论文设置
Σ
θ
(
x
t
,
t
)
=
σ
t
2
I
\textcolor{lightgreen}{\Sigma_\theta}(x_t, t) = \sigma_t^2 \mathbf{I}
Σθ(xt,t)=σt2I 其中
σ
t
2
\sigma_t^2
σt2 设置为常量
β
t
\beta_t
βt 或
β
~
t
\tilde\beta_t
β~t。
然后,
p
θ
(
x
t
−
1
∣
x
t
)
=
N
(
x
t
−
1
;
μ
θ
(
x
t
,
t
)
,
σ
t
2
I
)
\textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) = \mathcal{N}\big(x_{t-1}; \textcolor{lightgreen}{\mu_\theta}(x_t, t), \sigma_t^2 \mathbf{I} \big)
pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)
对于给定噪声
ϵ
∼
N
(
0
,
I
)
\epsilon\sim\mathcal{N}(\mathbf{0}, \mathbf{I})
ϵ∼N(0,I) 使用
q
(
x
t
∣
x
0
)
q(x_t|x_0)
q(xt∣x0)
x
t
(
x
0
,
ϵ
)
=
α
ˉ
t
x
0
+
1
−
α
ˉ
t
ϵ
x
0
=
1
α
ˉ
t
(
x
t
(
x
0
,
ϵ
)
−
1
−
α
ˉ
t
ϵ
)
x_t(x_0, \epsilon) = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon \\ x_0 = \frac{1}{\sqrt{\bar\alpha_t}} \Big(x_t(x_0, \epsilon) - \sqrt{1-\bar\alpha_t}\epsilon\Big)
xt(x0,ϵ)=αˉtx0+1−αˉtϵx0=αˉt1(xt(x0,ϵ)−1−αˉtϵ)
这给出了:
L
t
−
1
=
D
K
L
(
q
(
x
t
−
1
∣
x
t
,
x
0
)
∥
p
θ
(
x
t
−
1
∣
x
t
)
)
=
E
q
[
1
2
σ
t
2
∥
μ
~
(
x
t
,
x
0
)
−
μ
θ
(
x
t
,
t
)
∥
2
]
=
E
x
0
,
ϵ
[
1
2
σ
t
2
∥
1
α
t
(
x
t
(
x
0
,
ϵ
)
−
β
t
1
−
α
ˉ
t
ϵ
)
−
μ
θ
(
x
t
(
x
0
,
ϵ
)
,
t
)
∥
2
]
L_{t-1} = D_{KL}(q(x_{t-1}|x_t,x_0) \Vert \textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)) \\ = \mathbb{E}_q \Bigg[ \frac{1}{2\sigma_t^2} \Big \Vert \tilde\mu(x_t, x_0) - \textcolor{lightgreen}{\mu_\theta}(x_t, t) \Big \Vert^2 \Bigg] \\ = \mathbb{E}_{x_0, \epsilon} \Bigg[ \frac{1}{2\sigma_t^2} \bigg\Vert \frac{1}{\sqrt{\alpha_t}} \Big( x_t(x_0, \epsilon) - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}} \epsilon \Big) - \textcolor{lightgreen}{\mu_\theta}(x_t(x_0, \epsilon), t) \bigg\Vert^2 \Bigg] \\
Lt−1=DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))=Eq[2σt21
μ~(xt,x0)−μθ(xt,t)
2]=Ex0,ϵ[2σt21
αt1(xt(x0,ϵ)−1−αˉtβtϵ)−μθ(xt(x0,ϵ),t)
2]
使用模型重新参数化以预测噪声
μ
θ
(
x
t
,
t
)
=
μ
~
(
x
t
,
1
α
ˉ
t
(
x
t
−
1
−
α
ˉ
t
ϵ
θ
(
x
t
,
t
)
)
)
=
1
α
t
(
x
t
−
β
t
1
−
α
ˉ
t
ϵ
θ
(
x
t
,
t
)
)
\textcolor{lightgreen}{\mu_\theta}(x_t, t) = \tilde\mu \bigg(x_t, \frac{1}{\sqrt{\bar\alpha_t}} \Big(x_t - \sqrt{1-\bar\alpha_t}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big) \bigg) \\ = \frac{1}{\sqrt{\alpha_t}} \Big(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)
μθ(xt,t)=μ~(xt,αˉt1(xt−1−αˉtϵθ(xt,t)))=αt1(xt−1−αˉtβtϵθ(xt,t))
其中
ϵ
θ
\epsilon_\theta
ϵθ 是一个学习函数,它在给定
(
x
t
,
t
)
(x_t, t)
(xt,t) 的情况下预测
ϵ
\epsilon
ϵ。
这给出了:
L
t
−
1
=
E
x
0
,
ϵ
[
β
t
2
2
σ
t
2
α
t
(
1
−
α
ˉ
t
)
∥
ϵ
−
ϵ
θ
(
α
ˉ
t
x
0
+
1
−
α
ˉ
t
ϵ
,
t
)
∥
2
]
L_{t-1} = \mathbb{E}_{x_0, \epsilon} \Bigg[ \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar\alpha_t)} \Big\Vert \epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t) \Big\Vert^2 \Bigg]
Lt−1=Ex0,ϵ[2σt2αt(1−αˉt)βt2
ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)
2]
也就是说,我们正在训练预测噪声。
Simplified loss
L
simple
(
θ
)
=
E
t
,
x
0
,
ϵ
[
∥
ϵ
−
ϵ
θ
(
α
ˉ
t
x
0
+
1
−
α
ˉ
t
ϵ
,
t
)
∥
2
]
L_{\text{simple}}(\theta) = \mathbb{E}_{t,x_0, \epsilon} \Bigg[ \bigg\Vert \epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t) \bigg\Vert^2 \Bigg]
Lsimple(θ)=Et,x0,ϵ[
ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)
2]
这最小化
−
log
p
θ
(
x
0
∣
x
1
)
-\log \textcolor{lightgreen}{p_\theta}(x_0|x_1)
−logpθ(x0∣x1) when
t
=
1
t=1
t=1 和
L
t
−
1
L_{t-1}
Lt−1 for
t
>
1
t\gt1
t>1 丢弃 在
L
t
−
1
L_{t-1}
Lt−1 中加权。丢弃权重
β
t
2
2
σ
t
2
α
t
(
1
−
α
ˉ
t
)
\frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar\alpha_t)}
2σt2αt(1−αˉt)βt2 增加赋予较高
t
t
t(具有较高噪声水平)的权重,从而提高样本质量。
如果觉得有帮助的话,记得点赞收藏支持一下哦。如果有问题请在评论区一起交流讨论。
Reference:
https://www.bilibili.com/video/BV1b541197HX/?spm_id_from=333.999.0.0
https://nn.labml.ai/diffusion/ddpm/index.html