#

去噪扩散概率模型 (DDPM)

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

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

转发流程

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

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}$

逆向流程

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

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}}{a}

本文将 $Σ_{θ} (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))

wh $ϵ_{θ}$ ere 是一个预测 $ϵ$ 给定的学习函数 $(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_{s} im pl e (θ) = E_{t, x_{0}, ϵ} [∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]$

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

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

这里是提供 $ϵ_{θ} (x_{t}, t)$ 和训练代码的 unET 模型。该文件可以根据训练后的模型生成样本和插值。

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

#

降噪扩散

173class DenoiseDiffusion:

#

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

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

#

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

#

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

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

#

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

191        self.alpha = 1. - self.beta

#

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

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

#

$T$

195        self.n_steps = n_steps

#

$σ^{2} = β$

197        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)

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

#

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

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

#

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

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

#

213        return mean, var

#

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

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

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

#

$ϵ \sim N (0, I)$

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

#

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

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

#

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

231        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))

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

#

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

247        eps_theta = self.eps_model(xt, t)

#

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

249        alpha_bar = gather(self.alpha_bar, t)

#

$α_{t}$

251        alpha = gather(self.alpha, t)

#

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

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

#

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

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

#

$σ^{2}$

258        var = gather(self.sigma2, t)

#

$ϵ \sim N (0, I)$

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

#

样本

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

#

简化损失

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

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

#

获取批次大小

274        batch_size = x0.shape[0]

#

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

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

#

$ϵ \sim N (0, I)$

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

#

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

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

#

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

285        eps_theta = self.eps_model(xt, t)

#

MSE 亏损

288        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})$