标题党一下，顺便蹭一下 OpenAI Sora大模型的热点，主要也是回顾一下扩散模型的原理。
1. 简单理解扩散模型 简单理解，扩散模型如下图所示可以分成两部分，一个是 forward，另一个是 reverse 过程：
forward：这是加噪声的过程，表示为\(q(X_{0:T})\) \(t_0\) \(X_0\) \(t_T\) \(X_T\) \(X_t=X_0+ z_0+ z_1+...+ z_{t-1}\)
reverse：这很好理解，其实就是去噪过程,是\(q(X_{0:T})\) \(P_\theta(X_{0:T})\) \(X_T\) \(z_T\) \(X_0=X_t-z_t-z_{t-2}-...-z_1\) \(X_t\) \(z_t\) \(\overline z_{t-1}\) \(mse(z_t,\overline z_{t-1})\)
怎么理解这两个过程呢？一种简单的理解方法是我们可以假设世界上所有图像都是可以通过加密（就是 forward 过程）表示成隐变量，这些隐变量人眼看上去就是一堆噪声点。我们可以通过神经网络模型逐渐把这些噪声去掉，从而得到对应的原图（即 reverse 过程）。
2. 前向过程的数学表示
前向过程简单理解就是不断加噪声，加噪声的特点是越加越多：
前期加的噪声要少一点，这样是为了避免加太多噪声会导致模型不太好学习；
而当噪声量加的足够多后应该增加噪声的量，因为如果还是每次只加一点点，其实差别不大，而且这会导致前向过程太长，那么对应逆向过程也长，最终会增加计算量。所以噪声的量会有超参数\(\beta_t\) \(\beta_t\)
那我们可以很自然地知道，t 时刻的图像应该跟 t-1时刻的图像和噪声相关，所以有
\[X_t=\sqrt{\alpha_t}X_{t-1}+\sqrt{1-\alpha_t}z_1
其中\(\alpha_t=1-\beta_t\) \(z_1\) \(\beta_t\) \(\alpha_t\)
我们同样有\(X_{t-1}=\sqrt{\alpha_{t-1}}X_{t-2}+\sqrt{1-\alpha_{t-1}}z_2\)
\[\begin{align}
因为\(z_1,z_2\) \(\mathcal{N}(0,\sigma_{1}^{2})+\mathcal{N}(0,\sigma_{2}^{2})\sim\mathcal{N}(0,(\sigma_{1}^{2}+\sigma_{2}^{2}))\) \(\sqrt{(a_{t}(1-\alpha_{t-1})}^2+\sqrt{1-\alpha_{t}}^2=1-\alpha_t\alpha_{t-1}\) \(\tilde z_2\sim\mathcal{N}(0,1-\alpha_t\alpha_{t-1})\)
我们不断递归能够得到\(X_t\) \(X_0\)
\[\begin{align}
其中\(\overline{\alpha}_t=\alpha_t\alpha_{t-1}...\alpha_{1}\) \(\overline{z}_t\) \(\sigma=1-\overline{\alpha}_t\) \(z_t\) \(X_0\) \(\overline{\alpha}_t\) \(\overline z_t\) \(X_t\)
3. 逆向过程的数学表示 3.1 贝叶斯公式求解 扩散模型在应用的时候主要就是 reverse 过程，即给定一组随机噪声，通过逐步的还原得到想要的图像，可以表示为\(q(X_0|X_t)\) \(q(X_{t-1}|X_t)\)
\[q(X_{t-1}|X_t)=\frac{q(X_t|X_{t-1})q(X_{t-1})}{q(X_t)}
我们这里考虑扩散模型训练过程，我们默认是知道\(X_o\)
\[q(X_{t-1}|X_t,X_0)=\frac{q(X_t|X_{t-1},X_0)q(X_{t-1}|X_0)}{q(X_t|X_0)}
解释一下上面的公式：因为我们可以人为设置噪声分布，所以正向过程中每个时刻的数据也是知道的。例如，假设噪声\(z\) \(X_1=X_0+z\) \(q(X_1,X_0)\) \(q(X_{t-1},X_0),q(X_t,X_0)\) \(q(X_t|X_{t-1},X_0)\)
\(z\) \(a+\sqrt{b}z\)
3.2 高斯分布概率密度分布计算 下面公式中左边的概率分布其实就是右边三项概率分布的计算结果。
\[q(X_{t-1}|X_t,X_0)=\frac{q(X_t|X_{t-1},X_0)q(X_{t-1}|X_0)}{q(X_t|X_0)}
我们假设了噪声数据服从高斯分布\(\mathcal{N}(\mu,\sigma^2)\) \(exp{(-\frac{1}{2}\frac{(x-\mu)^2}{\sigma^2})}\)
\(q(X_t|X_0)\) \(exp(-\frac{1}{2}\frac{(X_t-\sqrt{\overline{a_t}}X_0)^2}{1-\overline{a_t}})\)
\[\propto\exp\left(-\,\frac{1}{2}\,(\frac{({X}_{t}-\sqrt{\alpha_{t}}{X}_{t-1})^{2}}{\beta_{t}}+\frac{({X}_{t-1}-\sqrt{\alpha}_{t-1}{X}_{0})^{2}}{1-\overline{{{\alpha}}}_{t-1}}-\frac{({X}_{t}-\sqrt{\overline{{{\alpha}}}_{t}}{X}_{0})^{2}}{1-\overline{{{\alpha}}}_{t}})\right)
其中上面公式中\(\beta_t=1-\alpha_t\) \(X_{t-1}\) \(X_{t-1}\)
\(exp{(-\frac{1}{2}\frac{(x-\mu)^2}{\sigma^2})}=exp(-\frac{1}{2}(\frac{1}{\sigma^2}x^2-\frac{2\mu}{\sigma^2}x+\frac{\mu^2}{\sigma^2}))\)
\[\tilde\sigma_t^2=\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\beta_t
方差等式中的\(\alpha,\beta\)
\[\tilde{\mu}_{t}({X}_{t},{X}_{0})\;=\frac{\sqrt{\alpha_{t}}({\bf1}-\bar{\alpha}_{t-1})}{{\bf1}-\bar{\alpha}_{t}}{X}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_{t}}{{\bf1}-\bar{\alpha}_{t}}{X}_{0}
均值跟\(X_t\) \(X_0\) \(X_t\) \(X_0\) \(X_0\) \(X_t=\sqrt{\overline{a}_t}X_0+\sqrt{1-\overline{a}_t}z_t\) \(X_0=\frac{1}{\sqrt{\overline{a}_t}}(X_t-\sqrt{1-\overline{a}_t}z_t)\) \(X_0\) \(X_t\)
\[\tilde{\mu}_{t}=\frac{1}{\sqrt{a_{t}}}(X_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{a}_{t}}}{z}_{t})
每个时刻的均值和方差的表达式就都有了。不过，每个时刻的方差是个定值，很容易求解，而均值却跟变量\(z_t\) \(z_t\) \(X_t\)
\[X_{t-1}=\tilde\mu_t+\tilde\sigma_t \epsilon
\(\epsilon\) \(z_t\) \(z_t=\text{diffusion_model}(x_t)\) \(X_{t-1}\)
那么要训练模型，我们肯定得有标签和损失函数啊。具体而言：
\(x_t\) \(z_t\) 标签其实就是 forward 过程中每个时刻产生的噪声数据\(\hat{z}_t\)
所以损失函数等于\(\text{loss}=mse(z_t, \hat{z}_t)\)
4. 代码实现 接下来我们结合代码来理解一下上述过程。
4.1 前向过程（加噪过程） 给定原始图像\(X_0\) \(\alpha_t=1-\beta_t\) \(X_t\)
\[\begin{align}
其中\(\overline{\alpha}_t=\alpha_t\alpha_{t-1}...\alpha_{1}\) \(\overline{z}_t\) \(\sigma=\sqrt{1-\overline{\alpha}_t}\)
下面是具体的代码实现，首先是与噪声相关超参数的设置和提前计算：
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from torchvision import transforms
# 定义线性beta时间表
def linear_beta_schedule(timesteps, start=0.0001, end=0.02):
# 在给定的时间步数内，线性地从 start 到 end 生成 beta 值
return torch.linspace(start, end, timesteps)
T = 300 # 总的时间步数
betas = linear_beta_schedule(timesteps=T) # β，迭代100个时刻
# 预计算不同的超参数（alpha和beta）
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, axis=0) # 累积乘积
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0) # 前一个累积乘积
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # alpha的平方根倒数
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) # alpha累积乘积的平方根
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) # 1-alpha累积乘积的平方根
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) # 计算后验分布q(x_{t-1}|x_t,x_0)的方差
接下来是具体的前向过程的计算，其中
get_index_from_list
forward_diffusion_sample
def get_index_from_list(vals, time_step, x_shape):
"""
返回传入的值列表vals（如β_t 或者α_t）中特定时刻t的值,同时考虑批量维度。
参数:
vals: 一个张量列表，包含了不同时间步的预计算值。
time_step: 一个包含时间步的张量，其值决定了要从vals中提取哪个时间步的值。
x_shape: 原始输入数据的形状，用于确保输出形状的一致性。
返回:
一个张量，其形状与原始输入数据x_shape相匹配，但是在每个批次中填充了特定时间步的vals值。
"""
batch_size = time_step.shape[0] # 获取批量大小
out = vals.gather(-1, time_step.cpu()) # 从vals中按照时间步收集对应的值
# 重新塑形为原始数据的形状，确保输出与输入在除批量外的维度上一致
return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(time_step.device)
# 前向扩散采样函数
def forward_diffusion_sample(x_0, time_step, device="cpu"):
"""
输入：一个图像和一个时间步
返回：图像对应时刻的噪声版本数据
"""
noise = torch.randn_like(x_0) # 生成和x_0形状相同的噪声
sqrt_alphas_cumprod_t = get_index_from_list(sqrt_alphas_cumprod, time_step, x_0.shape)
sqrt_one_minus_alphas_cumprod_t = get_index_from_list(sqrt_one_minus_alphas_cumprod, time_step, x_0.shape)
# 计算均值和方差
return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(
device
), noise.to(device)
image = Image.open('xiaoxin.jpg').convert('RGB')
img_tensor = transforms.ToTensor()(image)
for idx in range(T):
time_step = torch.Tensor([idx]).type(torch.int64)
img, noise = forward_diffusion_sample(img_tensor, time_step)
plt.imshow(transforms.ToPILImage()(img)) # 绘制加噪图像
4.2 训练
if __name__ == "__main__":
model = SimpleUnet()
T = 300
BATCH_SIZE = 128
epochs = 100
dataloader = load_transformed_dataset(batch_size=BATCH_SIZE)
device = "cuda" if torch.cuda.is_available() else "cpu"
logging.info(f"Using device: {device}")
model.to(device)
optimizer = Adam(model.parameters(), lr=0.001)
for epoch in range(epochs):
for batch_idx, (batch_data, _) in enumerate(dataloader):
optimizer.zero_grad()
# 对一个 batch 内的数据采样任意时刻的 time_step
t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()
x_noisy, noise = forward_diffusion_sample(batch_data, t, device) # 计算得到指定时刻的 加噪后的数据 和 对应的噪声数据
noise_pred = model(x_noisy, t) # 预测对应时刻的噪声
loss = F.mse_loss(noise, noise_pred) # 计算噪声预测的损失值
loss.backward()
optimizer.step()
这里我们忽略模型架构的具体细节，只需要知道每次模型的计算需要 噪声图像(
x_noisy
t
4.2 逆向过程（去噪采样过程）
给定某一时刻的数据\(X_t\) \(\mu\) \(\sigma\)
\[\tilde{\mu}_{t}=\frac{1}{\sqrt{a_{t}}}(X_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{a}_{t}}}{z}_{t})
\[\tilde\sigma_t^2=\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\beta_t
通过对\(\mathcal{N}(\tilde\mu_t,\tilde\sigma_t^2)\) \(X_{t-1}=\tilde\mu_t+\tilde\sigma_t\epsilon\) \(z_t\) \(\epsilon\)
实现代码如下：
@torch.no_grad()
def sample_timestep(model, x, t):
"""
使用模型预测图像中的噪声，并返回去噪后的图像。
如果不是最后一个时间步，则在此图像上应用噪声。
参数:
model - 预测去噪图像的模型
x - 当前带噪声的图像张量
t - 当前时间步的索引（整数或者整数型张量）
返回:
去噪后的图像张量，如果不是最后一步，返回添加了噪声的图像张量。
"""
# 从预设列表中获取当前时间步的beta值
betas_t = get_index_from_list(betas, t, x.shape)
# 获取当前时间步的累积乘积的平方根的补数
sqrt_one_minus_alphas_cumprod_t = get_index_from_list(sqrt_one_minus_alphas_cumprod, t, x.shape)
# 获取当前时间步的alpha值的平方根的倒数
sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)
# 调用模型来预测噪声并去噪（当前图像 - 噪声预测）
model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)
# 获取当前时间步的后验方差
posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)
if t == 0:
# 如Luis Pereira在YouTube评论中指出的，论文中的时间步t有偏移
return model_mean
else:
# 生成与x形状相同的随机噪声
noise = torch.randn_like(x)
# 返回模型均值加上根据后验方差缩放的噪声
return model_mean + torch.sqrt(posterior_variance_t) * noise
for i in reversed(range(0, T)):
t = torch.tensor([i], device='cpu', dtype=torch.long)
img = sample_timestep(model, img, t)
5. 总结 给定原始图像\(X_0\) \(\alpha_t=1-\beta_t\) \(X_t\)
\[\begin{align}
其中\(\overline{\alpha}_t=\alpha_t\alpha_{t-1}...\alpha_{1}\) \(\overline{z}_t\) \(\sigma=\sqrt{1-\overline{\alpha}_t}\)
给定某一时刻的数据\(X_t\) \(\mu\) \(\sigma\)
\[\tilde{\mu}_{t}=\frac{1}{\sqrt{a_{t}}}(X_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{a}_{t}}}{z}_{t})
\[\tilde\sigma_t^2=\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\beta_t
通过对\(\mathcal{N}(\tilde\mu_t,\tilde\sigma_t^2)\) \(X_{t-1}=\tilde\mu_t+\tilde\sigma_t\epsilon\) \(z_t\) \(\epsilon\)
参考
微信公众号：AutoML机器学习 MARSGGBO ♥ 如有意合作或学术讨论欢迎私戳联系~
