菜单

扩散模型

相关源文件

本页面记录了代码库中去噪扩散概率模型(DDPM)的实现。扩散模型是一类生成模型,它通过逐步去噪一个用随机噪声初始化的变量来学习生成数据。该实现基于Ho等人发表的论文《Denoising Diffusion Probabilistic Models》。

系统概览

扩散模型的实现包括两个主要过程:

  1. 前向过程:在T个时间步长内逐渐向数据添加噪声
  2. 逆向过程:通过在每一步预测并去除噪声来学习去噪数据

扩散模型架构

来源

核心组件

DenoiseDiffusion 类

DenoiseDiffusion 类是扩散过程的核心实现,负责处理噪声调度、采样和损失计算。

来源

主要方法包括

  • q_xt_x0:计算分布 q(x_t|x_0) 的参数
  • q_sample:通过向 x_0 添加噪声从 q(x_t|x_0) 中采样
  • p_sample:从 p_θ(x_{t-1}|x_t) 中采样,进行一步去噪
  • loss:计算用于训练的简化损失

使用 U-Net 进行噪声预测

该系统使用 U-Net 架构来预测噪声。U-Net 模型具有:

  1. 收缩路径(下采样)
  2. 扩展路径(上采样)
  3. 相应层之间的跳跃连接
  4. 时间步嵌入以调节模型

来源

数学基础

前向过程

前向过程在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}^T q(x_t | x_{t-1})

其中 β_1, ..., β_T 是方差调度。这可以简化为直接采样 x_t

q(x_t|x_0) = N(x_t; √(α̅_t) x_0, (1-α̅_t) I)

其中 α_t = 1 - β_t 和 α̅_t = ∏_{s=1}^t α_s

在代码中,这通过 q_sample 方法实现,该方法向输入添加噪声。

来源

逆向过程

逆向过程从 x_T ~ N(0, I) 开始去除噪声

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_t^2 I)
p_θ(x_{0:T}) = p_θ(x_T) ∏_{t=1}^T p_θ(x_{t-1} | x_t)

其中均值 μ_θ(x_t, t) 由以下参数化:

μ_θ(x_t, t) = (1/√α_t) * (x_t - (β_t/√(1-α̅_t)) * ε_θ(x_t, t))

在代码中,p_sample 方法通过预测噪声并使用它来确定先验均值来实现这一点。

来源

损失函数

简化的损失函数为:

L_simple(θ) = E_{t,x_0,ε} [||ε - ε_θ(√(α̅_t) x_0 + √(1-α̅_t)ε, t)||²]

这训练模型以预测添加到图像中的噪声。

该实现为批次中的每张图像选择一个随机时间步长,并计算真实噪声与预测噪声之间的均方误差(MSE)。

来源

实现细节

U-Net 架构组件

U-Net 架构由几个专门组件组成:

组件描述
TimeEmbedding(时间嵌入)创建时间步 t 的正弦嵌入
ResidualBlock(残差块)包含归一化、激活、卷积和跳跃连接的块
AttentionBlock(注意力块)实现多头自注意力机制以捕捉全局结构
DownBlock(下采样块)结合 ResidualBlock 和可选的 AttentionBlock 用于下采样路径
UpBlock(上采样块)结合 ResidualBlock 和可选的 AttentionBlock 用于上采样路径
MiddleBlock(中间块)最低分辨率处的特殊块,带有注意力机制
Upsample(上采样)将空间分辨率提高 2 倍
Downsample(下采样)将空间分辨率降低 2 倍

来源

U-Net 前向传播

U-Net 前向传播过程:

来源

训练和采样

训练过程

训练过程包括:

  1. 采样一批真实图像 x₀
  2. 采样随机时间步长 t
  3. 添加噪声以创建 x_t
  4. 使用模型预测噪声
  5. 计算真实噪声与预测噪声之间的 MSE 损失
  6. 更新模型参数

来源

采样过程

要生成样本,从随机噪声开始应用逆向过程:

  1. 采样 x_T ~ N(0, I)
  2. 对于 t = T, T-1, ..., 1
    • 预测噪声 ε_θ(x_t, t)
    • 采样 x_{t-1} ~ p_θ(x_{t-1}|x_t)
  3. 返回 x_0 作为生成的图像

来源

图像插值

该模型还可以对两张图像进行插值。

来源

配置和集成

该系统包含用于训练和评估的配置类。

来源

结论

本文档涵盖了代码库中去噪扩散概率模型(DDPM)的实现。该系统提供了一个完整的管道,用于训练扩散模型、生成样本和执行图像插值。该实现严格遵循原始 DDPM 论文中描述的数学框架,并针对高效训练和推理进行了定制。

来源