banner
Nagi-ovo

Nagi-ovo

Breezing homepage: [nagi.fun](nagi.fun)
github

拡散の直感と数学

本文主要以该视频的教学逻辑为主线,结合讲解内容进行整理和阐述,如有错误欢迎在评论区纠正!

直感部分#

理论支持#

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

这篇论文奠定了扩散模型的理论基础。作者提出了一种基于非平衡热力学的生成模型,通过逐步添加和去除噪声来实现数据生成。这为后续的扩散模型研究提供了重要的理论支持。

我们对图像应用大量噪声,然后用神经网络去噪。如果这个神经网络学得很好,那可以从完全随机的噪声开始最终得到我们训练数据中的图像。

image

正向扩散过程:#

迭代地对图像施加噪声,步骤足够多时图像会完全变成噪声,使用正态分布作为噪声源:

Screenshot 2024-12-09 at 21.03.59

反向扩散过程#

从纯粹的噪声到图像,涉及一个学习一步步去噪的神经网络。

为什么是逐渐去噪?作者在论文中提到 “一步直接完成去噪” 的结果很糟糕。

那么这个网络是什么样的?它又要预测什么?

算法改进#

Denoising Diffusion Probabilistic Models

这篇论文提出了去噪扩散概率模型(DDPM),显著提高了扩散模型的生成质量和效率。通过引入简单的去噪网络和优化的训练策略,DDPM 成为了扩散模型领域的一个重要里程碑。

作者讨论了神经网络可以预测的三种目标:

  1. 预测每个时间步的噪声均值 (predict the mean of the noise at each timestep)

    • 即预测条件分布 p(xt1xt)p(x_{t-1}|x_t) 的均值
    • 方差是 fixed 的,不可学习
  2. 预测原始图像 (predict x0x_0 directly)

    • 直接预测原始、未被污染的图像
    • 实验证明这种方式效果较差
  3. 预测添加的噪声 (predict the added noise)

    • 预测在正向过程中添加的噪声 εε
    • 和第一种方法(预测噪声均值)实际上是数学上等价的,只是参数化方式不同,它们可以通过简单的变换相互转换

论文最终选择了预测噪声(第三种方式)作为主要方法,因为这种方式训练更稳定且效果更好

这里每一步添加的噪声量是不固定的,通过一个 Linear Schedule 来控制噪声添加,防止训练过程不稳定。

大概长下面这样:
Screenshot 2024-12-09 at 21.34.56

可以看到最后最后几个时间步都接近完全噪声了,信息很少,此外整体来看信息摧毁得太快了,因此 OpenAI 使用了 Cosine Schedule 解决了这两个问题:

Screenshot 2024-12-09 at 21.39.35

模型架构#

U-Net#

随着这篇论文一起发表的是一个名叫 U-Net 的模型架构:

这个模型在中间有一个 Bottleneck(也就是参数量较小的层),用 Downsample-Block 和 Resnet-Block 来将输入图像投影到小分辨率,输出时用 Upsample-Block 将其投影回初始尺寸。

image

在某些分辨率下,作者加入了 Attention-Block 并在相同分辨率空间的层之间做 Skip-Connection。模型是被涉及为针对每一个时间步的,这是通过 Transformer 中的正弦位置编码嵌入来实现的,嵌入被投影到每个 Residual-Block 中。模型还能结合 Schedule,在不同时间步中去除不同量的噪声来提升生成效果,后面会详细讨论。

Bottleneck 和自编码器#

Bottleneck(瓶颈层)的概念最初是在无监督学习方法 "自编码器(Autoencoder)" 中提出并广泛使用的。作为自编码器架构中维度最低的隐藏层,它位于编码器和解码器之间,构成了网络中最窄的部分,强制网络学习数据的压缩表示,最小化重建误差并起到正则化作用:Lreconstruction=XDecoder(Encoder(X))2\mathcal{L}_{\text{reconstruction}} = \|X - \text{Decoder}(\text{Encoder}(X))\|^2

image

架构改进#

OpenAI 在他们的第二篇论文 Diffusion Models Beat GANs on Image Synthesis 中通过改进架构显著改善了整体效果:

  1. 增加网络深度(更多层),减少宽度(每层的通道数)
  2. 增加 Attention-Block 数量
  3. 并扩大每个 Attention Block 中的 heads 数量
  4. 引入 BigGAN 风格的 Residual Block 用于上采样和下采样
  5. 引入 Adaptive Group Normalization (AdaGN),通过条件信息(如时间步)来动态调整归一化的参数
  6. 用 Separate Classifier Guidance 帮助模型生成某类图片

数学部分#

符号表#

  • XtX_t 代表在 tt 时间步的图像,即 X0X_0 是原始图像。可以简记 tt 越小噪声越少:

image

  • 噪声的最终图像是一个 isotropic(各方向相同) 的完全噪声,记作 XTX_T,在最开始的研究中 T=1000T=1000,后续工作会将其减小很多:

image

  • Forward Process:q(xtxt1)q(x_t|x_{t-1}),输入xt1x_{t-1}图像输出一张噪声更多的图像XtX_t

image

  • Backward Process:p(xt1xt)p(x_{t-1}|x_t),输入 xtx_t 图像用神经网络输出一个降噪的图像 xt1x_t-1

image

正向过程#

image

其中,

  • 1βtxt1\sqrt{1 - \beta_t} x_{t-1} 是分布的均值(mean);βt\beta_tnoise schedule 参数,范围在 0 到 1 之间,配合1βt\sqrt{1 - \beta_t}噪声完成缩放,随着时间步增加而减小,表示保留的原始信号部分

image

  • βtI\beta_t I 是分布的协方差矩阵(covariance matrix),II单位矩阵,表示协方差矩阵是对角的且各维度独立,随着时间步增加,添加的噪声量增大。

现在我们只需要将这个步骤不断迭代执行即可得到 1000 步后的结果,但其实这些可以一步完成。

重参数化技巧#

重参数化技巧在扩散模型和其他生成模型(如变分自编码器,VAE)中非常重要。它的核心思想是将随机变量的采样过程转化为一个确定性函数加上一个标准化的随机变量。这种转换使得模型可以通过梯度下降进行优化,因为它消除了采样过程中的随机性对梯度计算的影响

这里通过一个简单的例子来解释其意义

你要实现一个扔骰子可以有两种方式,

  • 第一种是随机性在函数内部:
# 1. 直接掷骰子(随机采样)
def roll_dice():
    return random.randint(1, 6)

result = roll_dice()
  • 第二种则是让随机性在函数外部,函数本身是确定性的:
# 2. 将随机性分离出来
random_number = random.random()  # 生成0到1之间的随机数

def transformed_dice(random_number):
    # 将0-1的随机数映射到1-6
    return math.floor(random_number * 6) + 1

result = transformed_dice(random_number)

概率论中我们学过:如果 XX 是一个随机变量,且 X𝒩(0,1)X ∼ 𝒩(0,1),那么有:aX+b𝒩(b,a2)aX + b ∼ 𝒩(b, a²)

因此对于正态分布 N(μ,σ2)\mathcal{N}(\mu, \sigma^2) 可以通过以下方式生成样本:

x=μ+σϵx = \mu + \sigma \cdot \epsilon

其中 ϵN(0,1)\epsilon \sim \mathcal{N}(0, 1) 是标准正态分布。

那同理,在正态分布中,

  • 不使用重参数化
# 直接从目标正态分布采样
x = np.random.normal(mu, sigma)
  • 使用重参数化
# 先从标准正态分布采样
epsilon = np.random.normal(0, 1)
# 然后通过确定性变换得到目标分布
x = mu + sigma * epsilon

对应到模型训练中涉及的梯度计算时,

不使用重参数化

def sample_direct(mu, sigma):
    return np.random.normal(mu, sigma)

# 这种情况下很难计算关于mu和sigma的梯度
# 因为随机采样操作阻断了梯度传播

使用重参数化

def sample_reparameterized(mu, sigma):
    epsilon = np.random.normal(0, 1)  # 梯度不需要通过这里传播
    return mu + sigma * epsilon        # 可以轻松计算mu和sigma的梯度

以 VAE(变分自编码器)为例:

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = Encoder()  # 输出mu和sigma
        self.decoder = Decoder()

    def reparameterize(self, mu, sigma):
        # 重参数化技巧
        epsilon = torch.randn_like(mu)  # 从标准正态分布采样
        z = mu + sigma * epsilon        # 确定性变换
        return z

    def forward(self, x):
        # 编码器输出mu和sigma
        mu, sigma = self.encoder(x)
        
        # 使用重参数化采样
        z = self.reparameterize(mu, sigma)
        
        # 解码器重建输入
        reconstruction = self.decoder(z)
        return reconstruction

吃货视角的重参数化#

想象你在制作奶茶:

不使用重参数化

  • 直接制作一杯特定甜度的奶茶
  • 如果不好喝,你不知道是糖放多了还是放少了

使用重参数化

  1. 先准备一杯标准浓度的糖水(ϵ\epsilon
  2. 然后通过调整糖水的量(μ\mu)和稀释程度(σ\sigma)来达到目标甜度
  3. 如果不好喝,你可以清楚地知道是糖水量还是稀释程度需要调整(参数可优化)

image

总之,经过重参数化:

  • 梯度可以通过确定性变换传播
  • 参数可以通过梯度下降优化
  • 随机性被隔离,不影响梯度计算

正向数学推导#

xt1x_{t-1}xtx_t 的转移

  • 给定 xt1x_{t-1},我们希望生成 xtx_t
  • q(xtxt1)=N(xt;1βtxt1,βtI)q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I) 使用重新参数化技巧,
Σ=βtI,σ2=βtσ=βt\begin{align*} \because \Sigma &= \beta_t I, \sigma^2 = \beta_t \\ \therefore \sigma &= \sqrt{\beta_t} \end{align*}

可将 xtx_t 表示为 xt1x_{t-1} 的确定性变换加上噪声项:

xt=1βtxt1+βtϵ\begin{align*} x_t = \sqrt{1 - \beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon \end{align*}
  • 这里,1βtxt1\sqrt{1 - \beta_t} x_{t-1} 是均值部分,βtϵ\sqrt{\beta_t} \epsilon 是噪声部分。由于 ϵ\epsilon 是标准正态分布的样本,与模型参数无关,因此在反向传播时,梯度只需考虑 1βt\sqrt{1 - \beta_t}βt\sqrt{\beta_t} 对应的参数。这使得模型可以通过梯度下降进行有效优化。

我们用 αt\alpha_t 简化记法 & 记录其乘积的累计:

image

可得:

q(xtxt1)=αtxt1+1αtϵq(x_t \mid x_{t-1}) = \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t} \epsilon

计算两步转移:从 xt2x_{t-2}xtx_t

xt1=αt1xt2+1αt1ϵt1xt=αt(αt1xt2+1αt1ϵt1)+1αtϵtxt=αtαt1xt2+αt(1αt1)ϵt1+1αtϵt\begin{align*} x_{t-1} &= \sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_{t-1}} \epsilon_{t-1} \\ x_t &= \sqrt{\alpha_t} \left( \sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_{t-1}} \epsilon_{t-1} \right) + \sqrt{1 - \alpha_t} \epsilon_t \\ x_t &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{\alpha_t (1 - \alpha_{t-1})} \epsilon_{t-1} + \sqrt{1 - \alpha_t} \epsilon_t \end{align*}

因为 ϵt1\epsilon_{t-1}ϵt\epsilon_t 是独立的标准正态分布,合并噪声部分为一个新的噪声项 ϵN(0,I)\epsilon \sim \mathcal{N}(0, I)

xt=αtαt1xt2+1αtαt1ϵx_t = \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \epsilon

同理:

xt=αtαt1xt2+1αtαt1ϵxt=αtαt1αt2xt3+1αtαt1αt2ϵxt=αtαt1α1α0x0+1αtαt1α1α0ϵ通过归纳法,可以推出:xt=s=k+1tαsxk+1s=k+1tαsϵαˉt=s=1tαsk=0,xt=αˉtx0+1αˉtϵ(ϵN(0,I))\begin{align*} x_t &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \epsilon \\ x_t &= \sqrt{\alpha_t \alpha_{t-1} \alpha_{t-2}} x_{t-3} + \sqrt{1 - \alpha_t \alpha_{t-1} \alpha_{t-2}} \epsilon \\ x_t &= \sqrt{\alpha_t \alpha_{t-1} \cdots \alpha_1 \alpha_0} x_0 + \sqrt{1 - \alpha_t \alpha_{t-1} \cdots \alpha_1 \alpha_0} \epsilon \\ 通过&归纳法,可以推出:\\ x_t &= \sqrt{\prod_{s=k+1}^t \alpha_s} x_k + \sqrt{1 - \prod_{s=k+1}^t \alpha_s} \epsilon \\ \because \bar{\alpha}_t &= \prod_{s=1}^t \alpha_s \\ \therefore 当& k=0 时,x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon \quad (\epsilon \sim \mathcal{N}(0, I)) \end{align*}

完整推导流程如下:

q(xtxt1)=N(xt;1βtxt1,βtI)=1βtxt1+βtϵ=αtxt1+1αtϵq(xtxt2)=αtαt1xt2+1αtαt1ϵq(xtxt3)=αtαt1αt2xt3+1αtαt1αt2ϵq(xtx0)=αtαt1α1α0x0+1αtαt1α1α0ϵ=αˉtx0+1αˉtϵ(ϵN(0,I))=N(xt;αˉtx0,(1αˉt)I)\begin{align} q(x_t \mid x_{t-1}) &= \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I) \\ &= \sqrt{1 - \beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon \\ &= \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t} \epsilon \\ q(x_t \mid x_{t-2}) &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \epsilon \\ q(x_t \mid x_{t-3}) &= \sqrt{\alpha_t \alpha_{t-1} \alpha_{t-2}} x_{t-3} + \sqrt{1 - \alpha_t \alpha_{t-1} \alpha_{t-2}} \epsilon \\ q(x_t \mid x_0) &= \sqrt{\alpha_t \alpha_{t-1} \cdots \alpha_1 \alpha_0} x_0 + \sqrt{1 - \alpha_t \alpha_{t-1} \cdots \alpha_1 \alpha_0} \epsilon \\ &= \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon \quad (\epsilon \sim \mathcal{N}(0, I))\\ &= \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) I) \end{align}

逆向数学推导#

由于方差是固定的而不需要去学习 (见 1.3),所以我们只需要神经网络去预测均值:

image

我们的最终目标是预测两个时间步中的噪声,现在从损失函数开始分析:

log(pθ(x0))-log(p_\theta(x_0))

但是这个负对数似然中,x0x_0的概率依赖前面的所有其它时间步。我们可以学习一个逼近这些条件概率的模型作为解决方案,这里就需要用到 Varational Lower Bound (变分下界) 来得到一个更好计算的公式。

变分下界#

image

假设我们有一个无法计算的函数f(x)f(x),在我们的场景中是负对数似然,我们可以找一个始终满足 g(x)f(x)g(x) \leq f(x) 的可计算函数 g(x)g(x) ,那么优化 g(x)g(x) 也可以让 f(x)f(x) 增加:
Screenshot 2024-12-12 at 20.06.47

我们这里通过减去 KL 散度来保证这点,KL 散度是衡量两个分布相似度的指标,始终非负

DKL(pq)=xp(x)logp(x)q(x)dxD_{KL}(p \| q) = \int_x p(x) \log \frac{p(x)}{q(x)} \, dx

减去一个始终非负的项可以保证结果始终小于原始函数,这里用 “+” 是因为我们想要最小化损失,所以加上后始终保证原始的负对数似然大:

log(pθ(x0))log(pθ(x0))+DKL(q(x1:Tx0)pθ(x1:Tx0))-\log(p_\theta(x_0)) \leq -\log(p_\theta(x_0)) + D_{KL}(q(x_{1:T} \mid x_0) \| p_\theta(x_{1:T} \mid x_0))

这种形式下,因为负对数似然还在,下界仍然是不可计算的,我们需要得到一个更好的表达。首先将 KL 散度重写为两个项的对数比:

log(pθ(x0))log(pθ(x0))+DKL(q(x1:Tx0)pθ(x1:Tx0))=log(pθ(x0))+log(q(x1:Tx0)pθ(x1:Tx0))\begin{align*} -\log(p_\theta(x_0)) &\leq -\log(p_\theta(x_0)) + D_{KL}(q(x_{1:T} \mid x_0) \| p_\theta(x_{1:T} \mid x_0)) \\ &=-\log(p_\theta(x_0)) + \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta(x_{1:T} \mid x_0)} \right) \\ \end{align*}

再对其分母应用贝叶斯法则:

pθ(x1:Tx0)=pθ(x0x1:T)pθ(x1:T)pθ(x0)p_\theta(x_{1:T} \mid x_0)= \frac{p_\theta(x_0 \mid x_{1:T}) p_\theta(x_{1:T})}{p_\theta(x_0)}

Note

贝叶斯法则: p(AB)=p(BA)p(A)p(B)p(A \mid B) = \frac{p(B \mid A) p(A)}{p(B)}

上式的分子部分 pθ(x0x1:T)pθ(x1:T)p_\theta(x_0 \mid x_{1:T}) p_\theta(x_{1:T}) 实际上是联合概率 pθ(x0,x1:T)p_\theta(x_0, x_{1:T}),因为:

pθ(x0,x1:T)=pθ(x0x1:T)pθ(x1:T)p_\theta(x_0, x_{1:T}) = p_\theta(x_0 \mid x_{1:T}) p_\theta(x_{1:T})

通常, pθ(x0:T)p_\theta(x_{0:T}) 表示 x0x_0 和所有中间步骤 x1:Tx_{1:T} 的联合概率,即:

pθ(x0:T)=pθ(x0,x1:T)p_\theta(x_{0:T}) = p_\theta(x_0, x_{1:T})

Note

pθ(x0:T)p_\theta(x_{0:T}) 表示从时间步 0 到 TT 的所有状态 x0,x1,,xTx_0, x_1, \ldots, x_T 的联合概率分布。

pθ(x0:T)=p(xT)t=1Tpθ(xt1xt) p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)

代入有:

log(q(x1:Tx0)pθ(x1:Tx0))=log(q(x1:Tx0)pθ(x0:T)pθ(x0))将分母pθ(x0:T)pθ(x0)转化为乘法形式:1pθ(x0:T)pθ(x0)=pθ(x0)pθ(x0:T)=log(q(x1:Tx0)pθ(x0:T))+log(pθ(x0))\begin{align*} \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta(x_{1:T} \mid x_0)} \right) = \log \left( \frac{q(x_{1:T} \mid x_0)}{\frac{p_\theta(x_{0:T})}{p_\theta(x_0)}} \right) \\ 将分母\frac{p_\theta(x_{0:T})}{p_\theta(x_0)}转化为乘法形式:\frac{1}{\frac{p_\theta(x_{0:T})}{p_\theta(x_0)}} = \frac{p_\theta(x_0)}{p_\theta(x_{0:T})} \\ = \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta (x_{0:T})} \right) &+ \log(p_\theta(x_0)) \end{align*}

即按照下图中的流程得到最后形式:

Screenshot 2024-12-12 at 22.58.23

豁然开朗,烦人的两项消掉了:

log(pθ(x0))log(pθ(x0))+log(q(x1:Tx0)pθ(x0:T))+log(pθ(x0))=log(q(x1:Tx0)pθ(x0:T))\begin{align*} -\log(p_\theta(x_0)) &\leq -\log(p_\theta(x_0)) + \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta (x_{0:T})} \right) + \log(p_\theta(x_0)) \\ &= \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta (x_{0:T})} \right) \end{align*}

这样就得到了可以最小化的下限,并且式中的内容都是已知的:

  • 分子是正向过程的联合概率分布:q(x1:Tx0)=t=1Tq(xtxt1)q(x_{1:T} \mid x_0)=\prod_{t=1}^T q(x_t \mid x_{t-1})
  • 分母是逆向过程的联合概率分布:pθ(x0:T)=p(xT)t=1Tpθ(xt1xt)p_\theta (x_{0:T})=p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)

为了让其有解析解,还需要几个额外的重组步骤:

log(q(x1:Tx0)pθ(x0:T))=log(t=1Tq(xtxt1)p(xT)t=1Tpθ(xt1xt))=log(1p(xT)t=1Tq(xtxt1)t=1Tpθ(xt1xt))=log(1p(xT))+log(t=1Tq(xtxt1)t=1Tpθ(xt1xt))=log(p(xT))+log(t=1Tq(xtxt1)t=1Tpθ(xt1xt))=log(p(xT))+t=1Tlog(q(xtxt1)pθ(xt1xt))=log(p(xT))+t=2Tlog(q(xtxt1)pθ(xt1xt))+log(q(x1x0)pθ(x0x1))\begin{align} \log \left( \frac{q(x_{1:T} \mid x_0)}{p_\theta(x_{0:T})} \right) &= \log \left( \frac{\prod_{t=1}^T q(x_t \mid x_{t-1})}{p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)} \right) \\ &= \log \left( \frac{1}{p(x_T)} \cdot \frac{\prod_{t=1}^T q(x_t \mid x_{t-1})}{\prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)} \right)\\ &= \log \left( \frac{1}{p(x_T)} \right) + \log \left( \frac{\prod_{t=1}^T q(x_t \mid x_{t-1})}{\prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)} \right) \\ &= -\log(p(x_T)) + \log \left( \frac{\prod_{t=1}^T q(x_t \mid x_{t-1})}{\prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)} \right) \\ &=-\log(p(x_T)) + \sum_{t=1}^T \log \left( \frac{q(x_t \mid x_{t-1})}{p_\theta(x_{t-1} \mid x_t)} \right) \\ &=- \log(p(x_T)) + \sum_{t=2}^T \log \left( \frac{q(x_t \mid x_{t-1})}{p_\theta(x_{t-1} \mid x_t)} \right) + \log \left( \frac{q(x_1 \mid x_0)}{p_\theta(x_0 \mid x_1)} \right) \end{align}

根据贝叶斯法则重写 sum 项的分子:q(xtxt1)=q(xt1xt)q(xt)q(xt1)q(x_t \mid x_{t-1})=\frac{q(x_{t-1}\mid x_t)q(x_t)}{q(x_{t-1})}

但这又回到了前面,这些项都是需要估计全部样本导致 high variance,如给出下图所示的 xtx_t,你很难确定上一个状态是什么样的:

image

改进思路则为通过直接条件化于原始数据 x0x_0

q(xt1xt,x0)q(xtx0)q(xt1x0)\Longrightarrow \frac{q(x_{t-1} \mid x_t, x_0) q(x_t \mid x_0)}{q(x_{t-1} \mid x_0)}

这样同时给出无噪声的图像,侯选的 xt1x_{t-1} 就少了,方差会减小:

image

代入回原式:

=log(p(xT))+t=2Tlog(q(xt1xt,x0)q(xtx0)pθ(xt1xt)q(xt1x0))+log(q(x1x0)pθ(x0x1))=log(p(xT))+t=2Tlog(q(xt1xt,x0)pθ(xt1xt))+t=2Tlog(q(xtx0)q(xt1x0))+log(q(x1x0)pθ(x0x1))\begin{align} &= - \log(p(x_T)) + \sum_{t=2}^T \log \left( \frac{q(x_{t-1} \mid x_t, x_0) q(x_t \mid x_0)}{p_\theta(x_{t-1} \mid x_t) q(x_{t-1} \mid x_0)} \right) + \log \left( \frac{q(x_1 \mid x_0)}{p_\theta(x_0 \mid x_1)} \right) \\ &= - \log(p(x_T)) + \sum_{t=2}^T \log \left( \frac{q(x_{t-1} \mid x_t, x_0)}{p_\theta(x_{t-1} \mid x_t)} \right) + \sum_{t=2}^T \log \left( \frac{q(x_t \mid x_0)}{q(x_{t-1} \mid x_0)} \right) + \log \left( \frac{q(x_1 \mid x_0)}{p_\theta(x_0 \mid x_1)} \right) \end{align}

展开第二个 sum 项,可以发现大部分项都被化简掉了:

image

=log(p(xT))+t=2Tlog(q(xt1xt,x0)pθ(xt1xt))+log(q(xTx0)q(x1x0))+log(q(x1x0)pθ(x0x1))\begin{align} &= - \log(p(x_T)) + \sum_{t=2}^T \log \left( \frac{q(x_{t-1} \mid x_t, x_0)}{p_\theta(x_{t-1} \mid x_t)} \right) + \log \left( \frac{q(x_T \mid x_0)}{q(x_{1} \mid x_0)} \right) + \log \left( \frac{q(x_1 \mid x_0)}{p_\theta(x_0 \mid x_1)} \right) \end{align}

对最后两项应用 log rules 可以化简一些项:

log(q(xTx0)q(x1x0))+log(q(x1x0)pθ(x0x1))=[logq(xTx0)logq(x1x0)]+[logq(x1x0)logpθ(x0x1)]=logq(xTx0)logpθ(x0x1)\begin{align*} \log \left( \frac{q(x_T \mid x_0)}{q(x_{1} \mid x_0)} \right) + \log \left( \frac{q(x_1 \mid x_0)}{p_\theta(x_0 \mid x_1)} \right)&=\left[ \log q(x_T \mid x_0) - \log q(x_{1} \mid x_0) \right] + \left[ \log q(x_1 \mid x_0) - \log p_\theta(x_0 \mid x_1) \right] \\ &=\log q(x_T \mid x_0) - \log p_\theta(x_0 \mid x_1) \end{align*}

再将化简后的第一个项移到前面,合并成一个对数得到最终解析形式:

=log(p(xT))+t=2Tlog(q(xt1xt,x0)pθ(xt1xt))+logq(xTx0)logpθ(x0x1)=log(q(xTx0)p(xT))+t=2Tlog(q(xt1xt,x0)pθ(xt1xt))log(pθ(x0x1))=DKL(q(xTx0)p(xT))+t=2TDKL(q(xt1xt,x0)pθ(xt1xt))log(pθ(x0x1))=t=2TDKL(q(xt1xt,x0)pθ(xt1xt))log(pθ(x0x1))\begin{align} &= - \log(p(x_T)) + \sum_{t=2}^T \log \left( \frac{q(x_{t-1} \mid x_t, x_0)}{p_\theta(x_{t-1} \mid x_t)} \right) + \log q(x_T \mid x_0)- \log p_\theta(x_0 \mid x_1) \\ &= \log(\frac{q(x_T\mid x_0)}{p(x_T)}) + \sum_{t=2}^T \log \left( \frac{q(x_{t-1} \mid x_t, x_0)}{p_\theta(x_{t-1} \mid x_t)} \right) - \log(p_\theta(x_0 \mid x_1)) \\ &= D_{KL}(q(x_T | x_0) \| p(x_T)) + \sum_{t=2}^T D_{KL}(q(x_{t-1} | x_t, x_0) \| p_\theta(x_{t-1} | x_t)) - \log(p_\theta(x_0 | x_1)) \\ &= \sum_{t=2}^T D_{KL}(q(x_{t-1} | x_t, x_0) \| p_\theta(x_{t-1} | x_t)) - \log(p_\theta(x_0 | x_1)) \end{align}

这个形式的第一个项是可以忽略的,因为 qq 没有可学习参数,只是加噪声的正向过程,会收敛为正态分布,而 p(xT)p(x_T) 只是从高斯分布中随机采样的噪声,因此可以确定该项 KL 散度会很小。

剩余两项的推导结果如下(过程省略,详见 Lilian's Blog)
Screenshot 2024-12-13 at 14.43.24

β\beta 是固定的,那么就关注 μ\mu 的形式:

μ~t(xt,x0)=αˉt(1αˉt1)1αˉtxt+αˉt1βt1αˉtx0\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} x_0

正向过程生成的闭合形式 xt=αˉtx0+1αˉtϵx_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon 可以重写为 x0x_0 形式:

x0=1αˉt(xt1αˉtϵ)x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon \right)

将上述 x0x_0 的表达式代入预测均值公式 μ~t\tilde{\mu}_t

μ~t=αˉt(1αˉt1)1αˉtxt+αˉt1βt1αˉt1αˉt(xt1αˉtϵ)\tilde{\mu}_t = \frac{\sqrt{\bar{\alpha}_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} \cdot \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon \right)

现在 μ\mu 不再依赖 x0x_0。继续化简,首先展开第二项:

αˉt1βt1αˉt1αˉt(xt1αˉtϵ)=αˉt1βtαˉt(1αˉt)xtαˉt1βt1αˉtαˉt(1αˉt)ϵ\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} \cdot \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon \right) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{\sqrt{\bar{\alpha}_t} (1 - \bar{\alpha}_t)} x_t - \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t \sqrt{1 - \bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t} (1 - \bar{\alpha}_t)} \epsilon

xtx_t 项进行合并:

μ~t=(αˉt(1αˉt1)1αˉt+αˉt1βtαˉt(1αˉt))xtαˉt1βt1αˉtαˉt(1αˉt)ϵ\tilde{\mu}_t = \left( \frac{\sqrt{\bar{\alpha}_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} + \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{\sqrt{\bar{\alpha}_t} (1 - \bar{\alpha}_t)} \right) x_t - \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t \sqrt{1 - \bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t} (1 - \bar{\alpha}_t)} \epsilon

xtx_t 的系数进行进一步合并和化简,最终得到:

μ~t=1αˉt(xtβt1αˉtϵ)\tilde{\mu}_t = \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon \right)

这表示我们基本上只是减去 xtx_t 生成的随机缩放噪声,这就是神经网络要预测的东西。

代入后的损失函数 LtL_t 定义为一个均方误差:

Lt=12σt21αˉt(xtβt1αˉtϵ)μθ(xt,t)2=12σt21αˉt(xtβt1αˉtϵ)1αˉt(xtβt1αˉtϵθ(xt,t))2=12σt2βtαˉt(1αˉt)(ϵϵθ(xt,t))2=βt22σt2αˉt(1αˉt)ϵϵθ(xt,t)2\begin{align*} L_t &= \frac{1}{2\sigma_t^2} \left\| \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon \right) - \mu_\theta(x_t, t) \right\|^2 \\ &= \frac{1}{2\sigma_t^2} \left\| \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon \right) - \frac{1}{\sqrt{\bar{\alpha}_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) \right\|^2 \\ &= \frac{1}{2\sigma_t^2} \left\| \frac{\beta_t}{\sqrt{\bar{\alpha}_t (1-\bar{\alpha}_t)}} (\epsilon - \epsilon_\theta(x_t, t)) \right\|^2 \\ &= \frac{\beta_t^2}{2\sigma_t^2 \bar{\alpha}_t (1-\bar{\alpha}_t)} \|\epsilon - \epsilon_\theta(x_t, t)\|^2 \end{align*}

最后的形式就是时间步 tt 的实际噪声和神经网络预测噪声之间的均方误差。研究人员发现忽略前面的缩放项会得到更好的采样质量并且更容易实现。

βt22σt2αt(1α^t)ϵϵθ(xt,t)2ϵϵθ(xt,t)2\frac{\beta_t^2}{2 \sigma_t^2 \alpha_t (1 - \hat{\alpha}_t)} \left\| \epsilon - \epsilon_\theta(x_t, t) \right\|^2 \longrightarrow \left\| \epsilon - \epsilon_\theta(x_t, t) \right\|^2

回到原始公式

N(xt1;1αt(xtβt1αˉtϵθ(xt,t)),βt)\mathcal{N}\left(x_{t-1}; \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t)\right), \beta_t\right)

作者决定在最后一步的采样中,不再添加额外的随机噪声使生成过程更加稳定:

Screenshot 2024-12-13 at 15.51.11

最后的形式为:

Lsimple=Et,x0,ϵ[ϵϵθ(αˉtx0+1αˉtϵ,t)2]    Et,x0,ϵ[ϵϵθ(xt,t)2]\begin{align} L_{\text{simple}} &= \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta \left( \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}, t \right) \right\|^2 \right] \\ &\implies \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta \left( \mathbf{x}_t, t \right) \right\|^2 \right] \end{align}
  • Et,x0,ϵ\mathbb{E}_{t, \mathbf{x}_0,\boldsymbol{\epsilon}} 表示对时间步 tt、原始数据 x0\mathbf{x}_0 和噪声 ϵ\boldsymbol{\epsilon} 取期望
  • ϵ\boldsymbol{\epsilon} 是实际添加的随机噪声
  • ϵθ\boldsymbol{\epsilon}_\theta 是神经网络预测的噪声
  • αˉtx0+1αˉtϵ\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon} 是前向过程的闭式解,表示在时间步 tt 的噪声数据,因此可以简化为:
    • xt\mathbf{x}_t 直接表示时间步 tt 的噪声数据,即 αˉtx0+1αˉtϵ\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}
    • 整个损失函数本质上是在衡量预测噪声与实际噪声之间的均方误差

其中,时间步 tt 通常是从均匀分布中采样的(即 tUniform(1,T)t∼Uniform(1,T),T 是总的时间步数)。这种选择确保了在训练过程中,每个时间步都有相同的概率被选择,从而使模型在所有时间步上都能有效地学习去噪过程。

训练#

image

首先我们从数据集中采样一些图像,然后采样 tt 和来自正态分布的噪声,然后通过梯度下降优化目标

采样#

首先从正态分布中采样 xtx_t,然后用前面展示过的公式通过重参数化来采样 xt1x_{t-1}

image

注意这里 t=1t=1 时是不增加噪声的,根据公式

x0=1αt(x1βt1αˉtϵθ(x1,1))x_0 = \frac{1}{\sqrt{\alpha_t}} \left( x_1 - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_1, 1) \right)

t=1t = 1 时,公式用于从 x1x_1 恢复到 x0x_0,这是去噪过程的最后一步。此时,我们希望尽可能准确地重建原始图像。在最后一步不添加噪声(即没有 βtϵ\sqrt{\beta_t} \epsilon 项),可以避免在生成最终图像时引入不必要的随机性,从而保持图像的清晰度和细节。

代码实现#

推荐 知乎 Sunrise 的 MLP 简化实现
后面有时间的话考虑做一下 Stable Diffusion 的代码手撕,挖个坑...

参考资料#

読み込み中...
文章は、創作者によって署名され、ブロックチェーンに安全に保存されています。