从数学原理到代码实现：扩散模型的变分下界推导
从数学原理到代码实现：扩散模型的变分下界推导
【免费下载链接】guided-diffusion 
项目地址: https://gitcode.com/gh_mirrors/gu/guided-diffusion
引言：为什么变分下界是扩散模型的核心？
你是否曾疑惑：为什么扩散模型（Diffusion Model）能生成如此高质量的图像？其背后的数学原理是什么？本文将从数学推导到代码实现，全面解析扩散模型中的变分下界（Variational Lower Bound, VLB），揭示其在模型训练中的关键作用。读完本文，你将能够：
- 理解扩散过程的数学建模
- 掌握变分下界的完整推导过程
- 学会如何在代码中实现变分下界计算
- 分析变分下界与模型性能的关系
扩散模型基础：前向与反向过程
前向扩散过程（Forward Diffusion Process）
扩散模型的核心思想是通过一个马尔可夫链（Markov Chain）将数据逐步转换为噪声，然后学习反向过程从噪声重建数据。前向扩散过程定义为：
$$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I)$$
其中 $\beta_t$ 是一个随时间递增的小常数。通过迭代应用上述分布，我们可以得到 $q(x_t | x_0)$ 的闭式解：
$$q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\alpha_t}x_0, (1-\alpha_t)I)$$
其中 $\alpha_t = 1 - \beta_t$，$\bar{\alpha}t = \prod{s=1}^t \alpha_s$。
在 guided_diffusion/gaussian_diffusion.py 中，这一过程通过 q_sample 函数实现：
def q_sample(self, x_start, t, noise=None):
"""
Diffuse the data for a given number of diffusion steps.
In other words, sample from q(x_t | x_0).
"""
if noise is None:
noise = th.randn_like(x_start)
assert noise.shape == x_start.shape
return (
_extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+ _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
* noise
)
反向扩散过程（Reverse Diffusion Process）
反向扩散过程是前向过程的逆过程，定义为：
$$p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$$
我们的目标是学习参数 $\theta$，使得反向过程能够从噪声 $x_T$ 重建出真实数据 $x_0$。
变分下界（VLB）的数学推导
证据下界（ELBO）
根据贝叶斯公式，我们有：
$$\log p_\theta(x_0) = \log \int p_\theta(x_{0:T}) dx_{1:T} \geq \mathbb{E}{q(x{1:T}|x_0)}[\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}]$$
右侧即为证据下界（Evidence Lower Bound, ELBO），也称为变分下界（VLB）：
$$L_{VLB} = \mathbb{E}{q(x{1:T}|x_0)}[\log p_\theta(x_0) + \sum_{t=1}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})}]$$
VLB的分解
通过数学变换，VLB可以分解为：
$$L_{VLB} = -D_{KL}(q(x_T|x_0)||p_\theta(x_T)) + \sum_{t=1}^T \mathbb{E}{q(x_t|x_0)}[D_{KL}(q(x_{t-1}|x_t, x_0)||p_\theta(x_{t-1}|x_t))] - \log p\theta(x_0|x_1)$$
其中：
- 第一项是先验匹配项
- 中间项是KL散度的累加
- 最后一项是重建似然
在实际应用中，通常忽略常数项，将损失函数简化为：
$$L_{VLB} \approx \mathbb{E}_{t,x_0,\epsilon}[\text{KL}(q(x_{t-1}|x_t, x_0)||p_\theta(x_{t-1}|x_t))]$$
变分下界的代码实现
KL散度计算
在 guided_diffusion/losses.py 中，normal_kl 函数实现了两个高斯分布之间的KL散度计算：
def normal_kl(mean1, logvar1, mean2, logvar2):
"""
Compute the KL divergence between two gaussians.
Shapes are automatically broadcasted.
"""
tensor = None
for obj in (mean1, logvar1, mean2, logvar2):
if isinstance(obj, th.Tensor):
tensor = obj
break
assert tensor is not None, "at least one argument must be a Tensor"
logvar1, logvar2 = [
x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
for x in (logvar1, logvar2)
]
return 0.5 * (
-1.0
+ logvar2
- logvar1
+ th.exp(logvar1 - logvar2)
+ ((mean1 - mean2) ** 2) * th.exp(-logvar2)
)
变分下界项计算
在 guided_diffusion/gaussian_diffusion.py 中，_vb_terms_bpd 函数计算了变分下界的各项：
def _vb_terms_bpd(
self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
):
"""
Get a term for the variational lower-bound.
The resulting units are bits per dimension.
"""
true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
x_start=x_start, x_t=x_t, t=t
)
out = self.p_mean_variance(
model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
)
kl = normal_kl(
true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
)
kl = mean_flat(kl) / np.log(2.0) # Convert to bits
decoder_nll = -discretized_gaussian_log_likelihood(
x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
)
decoder_nll = mean_flat(decoder_nll) / np.log(2.0) # Convert to bits
# At the first timestep return the decoder NLL,
# otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
output = th.where((t == 0), decoder_nll, kl)
return {"output": output, "pred_xstart": out["pred_xstart"]}
训练损失计算
在 training_losses 函数中，根据不同的损失类型（如KL、MSE等），使用变分下界项计算损失：
def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
"""
Compute training losses for a single timestep.
"""
if model_kwargs is None:
model_kwargs = {}
if noise is None:
noise = th.randn_like(x_start)
x_t = self.q_sample(x_start, t, noise=noise)
terms = {}
if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
terms["loss"] = self._vb_terms_bpd(
model=model,
x_start=x_start,
x_t=x_t,
t=t,
clip_denoised=False,
model_kwargs=model_kwargs,
)["output"]
if self.loss_type == LossType.RESCALED_KL:
terms["loss"] *= self.num_timesteps
# ... (其他损失类型的实现)
变分下界的优化策略
损失类型选择
guided_diffusion 支持多种损失类型，通过 LossType 枚举定义：
class LossType(enum.Enum):
MSE = enum.auto() # 使用原始MSE损失（学习方差时使用KL）
RESCALED_MSE = enum.auto() # 使用原始MSE损失（学习方差时使用RESCALED_KL）
KL = enum.auto() # 使用变分下界
RESCALED_KL = enum.auto() # 类似KL，但重新缩放以估计完整的VLB
def is_vb(self):
return self == LossType.KL or self == LossType.RESCALED_KL
其中，KL 和 RESCALED_KL 直接使用变分下界作为损失函数。
β调度策略
β值的选择对扩散过程至关重要。guided_diffusion 提供了多种β调度策略，如线性调度和余弦调度：
def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
if schedule_name == "linear":
scale = 1000 / num_diffusion_timesteps
beta_start = scale * 0.0001
beta_end = scale * 0.02
return np.linspace(
beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
)
elif schedule_name == "cosine":
return betas_for_alpha_bar(
num_diffusion_timesteps,
lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
)
else:
raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
不同的β调度策略会影响变分下界的形状，进而影响模型的学习效果。
总结与展望
本文详细介绍了扩散模型中变分下界的数学推导和代码实现。我们从扩散过程的基本原理出发，推导出变分下界的数学表达式，并深入分析了其在 guided_diffusion 库中的实现细节。通过理解变分下界，我们可以更好地把握扩散模型的训练过程和优化方向。
未来，随着扩散模型的不断发展，变分下界的优化方法也将不断创新。例如，如何设计更有效的β调度策略，如何在保持生成质量的同时降低计算复杂度，都是值得探索的方向。
希望本文能帮助你深入理解扩散模型的数学基础，为你的研究或工程实践提供参考。如果你有任何问题或建议，欢迎在评论区留言讨论。
参考资料
- Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models.
- Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., & Ganguli, D. (2015). Deep unsupervised learning using nonequilibrium thermodynamics.
- Dhariwal, P., & Nichol, A. (2021). Diffusion models beat GANs on image synthesis.
【免费下载链接】guided-diffusion 项目地址: https://gitcode.com/gh_mirrors/gu/guided-diffusion
北京人形旗下天工造物具身智能开源社区，聚焦具身天工与慧思开物两大平台
更多推荐
19
0- 0
已为社区贡献6条内容
所有评论(0)