扩散模型总结

扩散模型总结（附代码）

最近在回顾之前写的Understanding Diffusion Models时，发现有个问题，文章太细了，每一步都是数学推导，于是，这篇对其进行一次总结，使得对模型理解一目了然。

前向过程

$$q(x_t|x_{t-1})=\mathcal N(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I)$$

$x_0\sim q(x_0)$,$\beta_1,...,\beta_T$ is the variance schedule, $\alpha_t=1-\beta_t$

$$q(x_t|x_0)=\mathcal N(x_t;\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I)$$

$\bar\alpha_t=\prod_{s=1}^t\alpha_s$ .

$x_t\sim\mathcal q(x_t|x_0)$ 代码实现：

def q_xt_x0(self, x0: torch.Tensor, t:torch.tensor):
    mean = gather(self.alpha_bar, t) ** 0.5 * x0
    var = 1 - gather(self.alpha_bar, t)
    return mean, var
def q_sample(self, x0: torch.Tensor, t:torch.tensor, eps: Optional[torch.Tensor] = None):
    if eps is None:
        eps = torch.randn_like(x0)
    mean, var = self.q_xt_x0(x0, t)
    return mean + (var ** 0.5) * eps

反向过程

$$q(x_{t-1}|x_t,x_0)=\mathcal N(x_{t-1};\mu_q(x_t,x_0),\Sigma_q(t))$$

$\mu_q(x_t,x_0)=\frac1{\sqrt{\alpha_t}}x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}\sqrt{\alpha_t}}\epsilon_0.$

$$p_\theta(x_{t-1}|x_t)=\mathcal N(x_{t-1};\mu_\theta(x_t,t),\Sigma_\theta(x_t,t))$$

$\mu_\theta(x_t,t)=\frac1{\sqrt{\alpha_t}}x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}\sqrt{\alpha_t}}\hat\epsilon_\theta(x_t,t),\ \Sigma_q(t)=\Sigma_\theta(t)=\sigma_t^2I.$

损失

$$L_{simple}(\theta) = \mathbb E_{t,x_0,\epsilon}[\left\|\epsilon-\hat\epsilon_\theta(x_t,t)\right\|^2]\\ =\mathbb E_{t,x_0,\epsilon}[\left\|\epsilon-\hat\epsilon_\theta(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,t)\right\|^2]$$

代码实现：

def loss(self, x0: torch.Tensor, noise: Optional[torch.Tensor] = None):
    batch_size = x0.shape[0]
    t = torch.randint(0, self.n_steps, (batch_size,), device=x0.device, dtype=torch.long)
    if noise is None:
        noise = torch.randn_like(x_0)
    xt = q_sample(x0, t, eps=noise)
    eps_theta = self.eps_model(xt, t)
    return F.mse_loss(noise, eps_theta)

采样

$x_{t-1}\sim p_\theta(x_{t-1}|x_t)$ 代码实现：

def p_sample(self,xt: torch.Tensor, t:torch.tensor):
    eps_theta = self.eps_model(xt, t)
    alpha_bar = gather(self.alpha_bar, t)
    alpha = gather(self.alpha, t)
    eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5
    mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)
    var = gather(self.sigma2, t)
    eps = torch.randn(xt.shape, device=xt.device)
    return mean + (var ** 0.5) * eps

参考

[1] J. Ho, A. Jain, and P. Abbeel, “Denoising Diffusion Probabilistic Models.” arXiv, Dec. 16, 2020. doi: 10.48550/arXiv.2006.11239.

代码来源：https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/diffusion/ddpm/