#

去噪扩散概率模型 (DDPM)

简而言之，我们从数据中获取图像并逐步添加噪点。然后，我们训练一个模型来预测每个步骤的噪声，并使用该模型生成图像。

以下定义和派生说明了其工作原理。详情请参阅论文。

转发进程

在 $T$ 时间步长内，转发过程会给数据 $x_{0} \sim q (x_{0})$ 增加噪音。

q (x_{t} ∣ x_{t - 1}) = N (x_{t}; 1 - β_{t} x_{t - 1}, β_{t} I) q (x_{1 : T} ∣ x_{0}) = t = 1 \prod T q (x_{t} ∣ x_{t - 1})

方差计划在哪里 $β_{1}, \dots, β_{T}$ 。

我们可以 $x_{t}$ 随时采样 $t$ ，

q (x_{t} ∣ x_{0}) = N (x_{t}; \overset{α_{t}}{ˉ} x_{0}, (1 - \overset{α_{t}}{ˉ}) I)

在哪里 $α_{t} = 1 - β_{t}$ 和 $\overset{α_{t}}{ˉ} = \prod_{s = 1}^{t} α_{s}$

反向处理

相反的过程会从四个 $T$ 时间步 $p (x_{T}) = N (x_{T}; 0, I)$ 长开始消除噪音。

p_{θ} (x_{t - 1} ∣ x_{t}) p_{θ} (x_{0 : T}) p_{θ} (x_{0}) = N (x_{t - 1}; μ_{θ} (x_{t}, t), Σ_{θ} (x_{t}, t)) = p_{θ} (x_{T}) t = 1 \prod T p_{θ} (x_{t - 1} ∣ x_{t}) = \int p_{θ} (x_{0 : T}) d x_{1 : T}

$θ$ 是我们训练的参数。

损失

我们根据负对数概率优化 ELBO（来自简森不等式）。

E [- lo g p_{θ} (x_{0})] \leq E_{q} [- lo g \frac{p _{θ} ( x _{0 : T} )}{q ( x _{1 : T} ∣ x _{0} )}] = L

损失可以改写如下。

L = E_{q} [- lo g \frac{p _{θ} ( x _{0 : T} )}{q ( x _{1 : T} ∣ x _{0} )}] = E_{q} [- lo g p (x_{T}) - t = 1 \sum T lo g \frac{p _{θ} ( x _{t - 1} ∣ x _{t} )}{q ( x _{t} ∣ x _{t - 1} )}] = E_{q} [- lo g \frac{p ( x _{T} )}{q ( x _{T} ∣ x _{0} )} - t = 2 \sum T lo g \frac{p _{θ} ( x _{t - 1} ∣ x _{t} )}{q ( x _{t - 1} ∣ x _{t} , x _{0} )} - l o g p_{θ} (x_{0} ∣ x_{1})] = E_{q} [D_{K L} (q (x_{T} ∣ x_{0}) ∥ p (x_{T})) + t = 2 \sum T D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t})) - l o g p_{θ} (x_{0} ∣ x_{1})]

$D_{K L} (q (x_{T} ∣ x_{0}) ∥ p (x_{T}))$ 是恒定的，因为我们保持 $β_{1}, \dots, β_{T}$ 不变。

计算 $L_{t - 1} = D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t}))$

后验的前向过程 $x_{0}$ 是，

q (x_{t - 1} ∣ x_{t}, x_{0}) \tilde{μ}_{t} (x_{t}, x_{0}) \tilde{β_{t}} = N (x_{t - 1}; \tilde{μ}_{t} (x_{t}, x_{0}), \tilde{β_{t}} I) = \frac{α ˉ _{t - 1} β _{t}}{1 - α _{t} ˉ} x_{0} + \frac{α _{t} ( 1 - α ˉ _{t - 1} )}{1 - α _{t} ˉ} x_{t} = \frac{1 - α ˉ _{t - 1}}{1 - α _{t} ˉ} β_{t}

论文将 $Σ_{θ} (x_{t}, t) = σ_{t}^{2} I$ 其中设置 $σ_{t}^{2}$ 为常量 $β_{t}$ 或 $\tilde{β_{t}}$ .

然后， $p_{θ} (x_{t - 1} ∣ x_{t}) = N (x_{t - 1}; μ_{θ} (x_{t}, t), σ_{t}^{2} I)$

对于给定的噪音， $ϵ \sim N (0, I)$ 使用 $q (x_{t} ∣ x_{0})$

x_{t} (x_{0}, ϵ) x_{0} = \overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ = \frac{1}{α _{t} ˉ} (x_{t} (x_{0}, ϵ) - 1 - \overset{α_{t}}{ˉ} ϵ)

这给了，

L_{t - 1} = D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t})) = E_{q} [\frac{1}{2 σ _{t}^{2}} ∥ ∥ \tilde{μ} (x_{t}, x_{0}) - μ_{θ} (x_{t}, t) ∥ ∥^{2}] = E_{x_{0}, ϵ} [\frac{1}{2 σ _{t}^{2}} ∥ ∥ \frac{1}{α _{t}} (x_{t} (x_{0}, ϵ) - \frac{β _{t}}{1 - α _{t} ˉ} ϵ) - μ_{θ} (x_{t} (x_{0}, ϵ), t) ∥ ∥^{2}]

使用模型重新参数化以预测噪声

μ_{θ} (x_{t}, t) = \tilde{μ} (x_{t}, \frac{1}{α _{t} ˉ} (x_{t} - 1 - \overset{α_{t}}{ˉ} ϵ_{θ} (x_{t}, t))) = \frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))

其中 $ϵ_{θ}$ 是预测 $ϵ$ 给定值的学习函数 $(x_{t}, t)$ 。

这给了，

L_{t - 1} = E_{x_{0}, ϵ} [\frac{β _{t} ^{2}}{2 σ _{t}^{2} α _{t} ( 1 - α _{t} ˉ )} ∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]

也就是说，我们正在训练预测噪音。

简化损失

$L_{simple} (θ) = E_{t, x_{0}, ϵ} [∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]$

这样可以最大限度地减少 $- l o g p_{θ} (x_{0} ∣ x_{1})$ 放 $t > 1$ 弃权重的时间 $t = 1$ 和 $L_{t - 1}$ 时间 $L_{t - 1}$ 。丢弃权重会 $\frac{β _{t} ^{2}}{2 σ _{t}^{2} α _{t} ( 1 - α _{t} ˉ )}$ 增加给出更高的权重 $t$ （噪声等级更高），从而提高样本质量。

该文件实现了损失计算和基本采样方法，我们在训练期间使用该方法生成图像。

这是提供 $ϵ_{θ} (x_{t}, t)$ 和训练代码的 UNet 模型。此文件可以从经过训练的模型生成样本和插值。

162from typing import Tuple, Optional
163
164import torch
165import torch.nn.functional as F
166import torch.utils.data
167from torch import nn
168
169from labml_nn.diffusion.ddpm.utils import gather

#

降噪扩散

172class DenoiseDiffusion:

#

eps_model 是 $ϵ_{θ} (x_{t}, t)$ 模特
n_steps 是 $t$
device 是用来放置常量的设备

177    def __init__(self, eps_model: nn.Module, n_steps: int, device: torch.device):

#

183        super().__init__()
184        self.eps_model = eps_model

#

创建 $β_{1}, \dots, β_{T}$ 线性增加的差异计划

187        self.beta = torch.linspace(0.0001, 0.02, n_steps).to(device)

#

$α_{t} = 1 - β_{t}$

190        self.alpha = 1. - self.beta

#

$\overset{α_{t}}{ˉ} = \prod_{s = 1}^{t} α_{s}$

192        self.alpha_bar = torch.cumprod(self.alpha, dim=0)

#

$T$

194        self.n_steps = n_steps

#

$σ^{2} = β$

196        self.sigma2 = self.beta

#

获取 $q (x_{t} ∣ x_{0})$ 分发

q (x_{t} ∣ x_{0}) = N (x_{t}; \overset{α_{t}}{ˉ} x_{0}, (1 - \overset{α_{t}}{ˉ}) I)

198    def q_xt_x0(self, x0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

#

收集 $α_{t}$ 和计算 $\overset{α_{t}}{ˉ} x_{0}$

208        mean = gather(self.alpha_bar, t) ** 0.5 * x0

#

$(1 - \overset{α_{t}}{ˉ}) I$

210        var = 1 - gather(self.alpha_bar, t)

#

212        return mean, var

#

样本来自 $q (x_{t} ∣ x_{0})$

q (x_{t} ∣ x_{0}) = N (x_{t}; \overset{α_{t}}{ˉ} x_{0}, (1 - \overset{α_{t}}{ˉ}) I)

214    def q_sample(self, x0: torch.Tensor, t: torch.Tensor, eps: Optional[torch.Tensor] = None):

#

$ϵ \sim N (0, I)$

224        if eps is None:
225            eps = torch.randn_like(x0)

#

得到 $q (x_{t} ∣ x_{0})$

228        mean, var = self.q_xt_x0(x0, t)

#

样本来自 $q (x_{t} ∣ x_{0})$

230        return mean + (var ** 0.5) * eps

#

样本来自 $p_{θ} (x_{t - 1} ∣ x_{t})$

p_{θ} (x_{t - 1} ∣ x_{t}) μ_{θ} (x_{t}, t) = N (x_{t - 1}; μ_{θ} (x_{t}, t), σ_{t}^{2} I) = \frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))

232    def p_sample(self, xt: torch.Tensor, t: torch.Tensor):

#

$ϵ_{θ} (x_{t}, t)$

246        eps_theta = self.eps_model(xt, t)

#

收集 $\overset{α_{t}}{ˉ}$

248        alpha_bar = gather(self.alpha_bar, t)

#

$α_{t}$

250        alpha = gather(self.alpha, t)

#

$\frac{β}{1 - α _{t} ˉ}$

252        eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5

#

$\frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))$

255        mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)

#

$σ^{2}$

257        var = gather(self.sigma2, t)

#

$ϵ \sim N (0, I)$

260        eps = torch.randn(xt.shape, device=xt.device)

#

样本

262        return mean + (var ** .5) * eps

#

简化损失

$L_{simple} (θ) = E_{t, x_{0}, ϵ} [∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]$

264    def loss(self, x0: torch.Tensor, noise: Optional[torch.Tensor] = None):

#

获取批次大小

273        batch_size = x0.shape[0]

#

批次中 $t$ 每个样本的随机获取

275        t = torch.randint(0, self.n_steps, (batch_size,), device=x0.device, dtype=torch.long)

#

$ϵ \sim N (0, I)$

278        if noise is None:
279            noise = torch.randn_like(x0)

#

$x_{t}$ 的样本 $q (x_{t} ∣ x_{0})$

282        xt = self.q_sample(x0, t, eps=noise)

#

得到 $ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t)$

284        eps_theta = self.eps_model(xt, t)

#

MSE 亏损

287        return F.mse_loss(noise, eps_theta)

去噪扩散概率模型 (DDPM)

转发进程

反向处理

损失

计算Lt−1​=DKL​(q(xt−1​∣xt​,x0​)∥pθ​(xt−1​∣xt​))

简化损失

降噪扩散

获取q(xt​∣x0​)分发

样本来自q(xt​∣x0​)

样本来自pθ​(xt−1​∣xt​)

简化损失

计算 $L_{t - 1} = D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t}))$

获取 $q (x_{t} ∣ x_{0})$ 分发

样本来自 $q (x_{t} ∣ x_{0})$

样本来自 $p_{θ} (x_{t - 1} ∣ x_{t})$