这是用户在 2024-6-14 16:08 为 https://zhuanlan.zhihu.com/p/702574549 保存的双语快照页面,由 沉浸式翻译 提供双语支持。了解如何保存?

论文学习笔记SDE:从EBM到Score-Matching,随机微分方程视角下的Diffusion大一统

40 人赞同了该文章
发布于 2024-06-12 16:29・IP 属地北京 ,编辑于 2024-06-12 16:33・IP 属地北京

1 写在前面

终于结束本学期核心专业课的期末考试,可以开始弄一些"fancy works"了。

前段时间主要在看Diffusion Model,其中纵观Diffusion理论工作的历史进程,基于SDE(随机微分方程Stochastic Differential Equations)框架对score-based模型NCSN和DDPM的统一是非常值得一读的work。但是笔者在看那篇讲SDE框架的Diffusion时,其中score-matching部分感到一头雾水,故而进行了一番“追根溯源”。

本文将从最原始版本的EBM(Energy-Based Model)讲起,到score-matching方法的引出,此后结合DDPM和NCSN(Generative Modeling by Estimating Gradients of the Data Distribution)来探究SDE框架下的生成模型(对应论文SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS)。

2 从Energy-Based Model讲起

2.1 Overview

基于能量的生成模型虽然相对于主流的GAN、VAEs、Diffusions较为小众,现在提及生成式模型的时候都鲜为提起(笔者也是在了解score-matching的时候才知道有这么个东西),不过,其联系物理热力学中布朗运动系统进行建模(前CPHO选手感到非常亲切:)),它避开直接对真实分布进行建模,而是选择通过神经网络拟合系统的“能量”(物理意义上的势能函数),再通过Langevin dynamics(朗之动力学)的采样手段,实现从系统势能到系统统计分布的转化。

参考论文:Implicit Generation and Modeling with Energy Based Models (neurips.cc)

2.2 EBM的物理背景

2.2.1 玻尔兹曼分布

在热力学中,由于粒子数量巨大没有办法实现对每一个例子的物理状态建模,我们对于一个热力学系统的刻画常常是通过系统粒子的统计分布实现的。

我们考虑一个 xRnx \in \mathbb{R^n} 的高维空间,这个空间中的势能函数可以被描述为 E(x)RE(x) \in \mathbb{R} ,从物理层面上,我们可以通过玻尔兹曼分布建模处于布朗运动下给定势能分布的系统粒子的统计分布:

pθ(x)=exp(Eθ(x))Z(θ)p_\theta(x) =\frac{exp(-E_\theta(x))}{Z(\theta)} \\其中 Z(θ)=exp(Eθ(x))dxZ(\theta) = \int exp(-E_\theta(x))dx ,仅仅做归一化的作用。

这个分布的物理意义可以解释为:对于这个给定能量分布的热力学系统,粒子做布朗运动,当达到稳定状态后,粒子的统计分布函数,即粒子出现在坐标 xx 处的概率。

在其他对于EBM解释的文章中,可能会直接说“使用玻尔兹曼分布建模真实样本的分布函数”,一个问题可能是玻尔兹曼分布真的能建模任意一个分布吗?

答案显然是可以的。事实上,对于建模任何一个概率密度函数,我们只需要满足:

(1) p(x)=1\int p(x) = 1 , (2) p(x)>=0p(x) >= 0,(3)概率密度函数的“任意性”

其中,(1)由归一化函数 Z(θ)Z(\theta) 保证,(2)由指数函数保证,对于建模分布的任意性(3),由神经网络可以拟合任何函数来满足。

2.2.2 Langevin dynamics

我们观察上面的分布函数,不难发现一个在实现时很直接的问题: Z(θ)Z(\theta) 怎么求?

要知道,我们使用数学方法对于这里积分的处理是颇为困难的。因而,我们需要一种新的方法,让我们能够实现避开 Z(θ)Z(\theta) 的积分计算,来实现仅通过采样 xx 、计算 Eθ(x)E_\theta(x) 来实现对于 pθ(x)p_\theta(x) 的求解。

在物理学中,我们可以通过Langevin dynamics的随机微分方程来实现对于一个热力系统的建模:

dx=xEθ(x)dt+2DdWdx = -\triangledown_xE_\theta(x)dt + \sqrt{2D}dW \\其中, DD 是扩散系数,用于控制布朗运动强度,W是布朗运动的随机扰动(服从高斯分布)。

其中这里的物理意义是:对于任意一个粒子,其运动取决于系统当前势能的梯度方向(受力)和一个布朗运动的随机扰动项决定;同时,对于一个给定势能分布的系统,对于任意一种初始分布,系统中的多个微观粒子在足够长时间后,这些粒子的分布特征将服从该系统平衡状态的统计分布,也就实现了微观和宏观规律的统一。

对于后面这种物理意义的解释可能在这里有些抽象,我们不妨先对该式子进行一番离散化与变形,也就得到我们所谓的Langevin dynamics采样方法:

xt=xt1λ2xEθ(xt1)+λε,εN(0,I)x_t = x_{t-1} - \frac{\lambda}{2}\triangledown_xE_\theta(x_{t-1}) + \sqrt\lambda\varepsilon, \varepsilon \sim N(0, I) \\ 可以证明,当 tt 足够大时, xtx_t 服从的分布将足够接近 pθ(x)p_\theta(x) ,对应论文中的

于是,我们就可以通过这样的采样方法来求解 pθ(x)p_\theta(x) ,以去除对于 Z(θ)Z(\theta) 的不可解部分。

这里,我们不妨写一个demo进行可视化来获取一些更加直观的理解,笔者参考了Langevin Monte Carlo: Sampling using Langevin Dynamics | Perceptron.blog这篇博客的代码

在这里,我们不妨假设此时真实分布就是一个标准高斯分布 p(x)=12πexp(x2)p(x) = \frac{1}{\sqrt{2\pi}}exp(-x^2) ,假设其 Z(θ)=1Z(\theta)=1 其对应的能量密度函数,则应该有 E(x)=x2xEθ(x)=2xE(x) = x^2, \triangledown_xE_\theta(x)=2x。我们系统的初始分布为一个均匀分布,我们通过多步的Langevin dynamics采样来观察系统统计分布的变化。

import matplotlib.pyplot as plt
import torch
import torch.distributions as D

eps = 0.01
steps = 500
n_samples = 10000

# 系统初始状态,X_0来自一个[-10, 10]之间的均匀分布
prior = D.Uniform(torch.Tensor([-10, -10]), torch.Tensor([10, 10])) 

x = prior.sample((n_samples, ))
print(x.shape)
colors = ['#8ECFC9','#FFBE7A','#FA7F6F','#82B0D2','#BEB8DC','#E7DAD2']

# 进行Langevin dynamics采样
for i in range(steps):
    if i % 100 == 0: # 每100次进行一次记录
        dim1, dim2 = torch.chunk(x, 2, dim=1)
        plt.scatter(dim1, dim2, alpha=0.5, c=colors[int(i / 100)])
        # plt.show()
        print(dim1.shape, dim2.shape)
    # Langevin dynamics采样的递归公式
    x = x - eps * x + torch.sqrt(2 * torch.ones_like(x) * eps) * torch.randn(size=x.shape)

plt.show()
print(x.mean().item(), x.std().item())
# result:
# -0.006517414934933186 1.0023117065429688 ~ N(0, 1)

结果如下图所示

不难发现随着采样不断进行,原本均匀分布在[-10, 10]之间的数据分布逐渐变为一个标准正态分布(对应绿色、橙色、红色、蓝色......),从物理意义上而言,可以认为是在指定系统能量的背景下,一个初始为均匀分布的系统粒子逐步达到该系统能量分布对应的平衡状态。

2.3 EBM & MLE Loss

基于上述物理背景,我们进行对EBM的介绍。

假设现在真实样本服从分布 pdata(x)p_{data}(x) ,在生成模型中,我们希望通过神经网络去拟合一个分布 pθ(x)p_\theta(x) 来逼近真实分布,在EBM中,我们不直接去拟合真实分布,而是将真实分布建模成玻尔兹曼分布的形式,使用神经网络去拟合能量函数 Eθ(x)E_\theta(x) pθ(x)=exp(Eθ(x))Z(θ)p_\theta(x) =\frac{exp(-E_\theta(x))}{Z(\theta)} \\由于在上文中提到 Z(θ)Z(\theta) 不可解,我们使用Langevin dynamics采样来代替对于 pθ(x)p_\theta(x) 的直接计算,即不断使用

xt=xt1λ2xEθ(xt1)+λε,εN(0,I)x_t = x_{t-1} - \frac{\lambda}{2}\triangledown_xE_\theta(x_{t-1}) + \sqrt\lambda\varepsilon, \varepsilon \sim N(0, I) \\

得到一个解 xTx_T ,不妨记为 qθ(x)q_\theta(x) ,当 T+infT \rightarrow+inf 时,我们可以得到 pθ(x)qθ(x)p_\theta(x) \approx q_\theta(x) (在此不给出详细证明,论文中写为As shown by Welling and Teh [2011] )

在生成任务中,我们为了使得我们拟合的分布尽可能逼近 pdata(x)p_{data}(x) ,对于样本集合 x(i)Dx^{(i)} \sim D ,我们将Loss设计为 L(θ)=argmaxθExPdata[logpθ(x)]=argmaxθExPdata[Eθ(x)+logZ(θ)]L(\theta)= argmax_\theta \mathbb{E}_{x \sim P_{data}}[log p_\theta(x)] = argmax_\theta \mathbb{E}_{x \sim P_{data}} [-E_\theta(x) + logZ(\theta)]\\但是显然,我们引入Langevin dynamics采样的目的就是去除 Z(θ)Z(\theta) 的影响,因而,在Loss函数部分,我们仍然需要想办法剔除 Z(θ)Z(\theta) 项。幸运的是,在此前的理论工作中,已经给出了对于该问题的处理方法(see [Turner, 2005] for derivation),此情况下损失函数的梯度可以进行一个近似为文中的式(2)

在这里,使用损失函数的梯度近似代替计算损失函数本身,因而,对于Loss函数的部分也就变得可解了,我们就可以得到训练一个EBM的大致流程:

从伯努利分布or均匀分布中采样一个初始 x0x_0 ,通过K次朗之万动力学采样后得到 xkx_kΩ(·)\Omega(·) 代表截止梯度计算后得到 xix_i^- ,此后计算损失函数。

注意到这里对于损失函数的计算中还加入了一个L2正则项,是作者在训练过程中发现可以优化效果的trick

不过,从效果上来看,其实EBM不如DCGAN,虽然方法很novel,但是效果确实不好,可能这也是其比较小众的原因吧。

3 什么是"Score-Matching"?

3.1 Introduction

其实关于Score-Matching的概念早在2005年一篇JMLR就提出了(hyvarinen05a.pdf (jmlr.org)

我们仍然把之前EBM使用玻尔兹曼分布对于真实分布的建模公式拿过来看:

pθ(x)=exp(Eθ(x))Z(θ)p_\theta(x) =\frac{exp(-E_\theta(x))}{Z(\theta)} \\我们提到,其实我们常希望通过样本集合的 x(i)Dx^{(i)} \sim D 计算极大似然函数 logpθ(x)logp_\theta(x) 来优化神经网络,但是显然这里对于极大似然的计算会牵扯到我们希望避开的 Z(θ)Z(\theta) : L(θ)=logpθ(x)=Eθ(x)+logZ(θ)L(\theta) = logp_\theta(x) =-E_\theta(x)+logZ(\theta) \\\\ 在EBM中我们使用的 θ\triangledown_\theta 求梯度近似来避开Z(θ)Z(\theta) 的计算,在这篇论文中,还提出了一种基于score的方法,或者说基于对 x\triangledown_x 求梯度的方法。观察上式,我们不难发现我们所困扰的 Z(θ)Z(\theta) 不受变量 xx 的影响,因此,我们不妨先对 xx 求一次梯度,于是得到score函数的定义: s(x;θ)=xlogpθ(x)=xEθ(x)s(x;\theta) = \triangledown_x logp_\theta(x)=- \triangledown_x E_\theta(x) \\

不难看出,其实score函数说白了就是极大似然的梯度,我们将对于极大似然的优化问题转换为对其梯度的优化。

相应的,对于我们的真实分布"ground truth" ,我们可以得到其得分函数的定义:

sreal(x)=xlogpreal(x)=xEreal(x),RnRns_{real}(x) = \triangledown_x logp_{real}(x)=- \triangledown_x E_{real}(x) , \mathbb{R}^n \rightarrow \mathbb{R}^n \\ 接下来,所谓"score-matching"的过程,就是优化参数 θ\theta ,使得两个得分函数尽量接近的过程: J(θ)=12xRnpreal(x)||sreal(x)s(x;θ)||2dx=Expreal[||sreal(x)s(x;θ)||2]J(\theta) = \frac{1}2 \int_{x \in \mathbb{R}^n} p_{real}(x)|| s_{real}(x) - s(x;\theta) ||^2 dx = \mathbb{E}_{x \sim p_{real}}[ ||s_{real}(x) - s(x;\theta) ||^2] \\在这里,可能会出现两点疑问:得分完全匹配是否能够等价于两个分布匹配呢?我们希望通过 pθ(x)p_\theta(x) 去近似 preal(x)p_{real}(x) ,但是我们只有后者的样本,没有分布函数,不知道 sreal(x)s_{real}(x) 是什么的情况下上面这个式子要怎么用呢?

3.2 Score-Matching to MLE

在文章的Appendix部分,我们可以得到证明(论文中的 q(ξ;θ)q(\xi;\theta) 就是本文中的 exp(Eθ(x))exp(-E_\theta(x))

两个得分函数相等,两边分别取积分会得到同一位置的极大似然值相差一个常量 cc ,即 logp(x;θ)=logp(x)+clogp(x;\theta) = logp(x) + c ,然而,我们不妨对两边求积分 inf+infdx\int_{-inf}^{+inf}dx ,可知 c=0c=0 必然成立。

因而,得分匹配可以完全等价于其概率密度函数等价,即两个分布一致。

3.3 关于 sreals_{real}

能够通过得分匹配来替代含 Z(θ)Z(\theta) 的极大似然函数计算固然非常美好,但是观察得分匹配损失函数,不难发现我们其实并不知道ground truth的 sreal(x)s_{real}(x) ,对于这个问题,我们试图对原式子进行一些假设与变形。

首先展开这个L2正则项,考虑到我们的优化参数是 θ\theta ,我们可以抛去 L(θ)L(\theta) 中与 θ\theta 无关的项,基于此认识我们进行变形:

J(θ)=12xRnpreal(x)||sreal(x)s(x;θ)||2dx12xRnpreal(x)[s(x;θ)22sreal(x)s(x;θ)]dxJ(\theta) = \frac{1}2 \int_{x \in \mathbb{R}^n} p_{real}(x)|| s_{real}(x) - s(x;\theta) ||^2 dx \simeq \frac{1}2 \int_{x \in \mathbb{R}^n} p_{real}(x)[ s(x;\theta)^2 - 2 s_{real}(x) s(x;\theta)] dx\\对于后半部分是仍然与 sreal(x)s_{real}(x) 相关项,可以证明后半部分等价于 i=1kxis(xi;θ)\sum_{i=1}^k\triangledown_{x_i} s(x_i;\theta) ,具体证明过程可以看这篇回答(有没有谁通俗的讲一下Denoising score matching? - 知乎 (zhihu.com)),最终,我们可以将其损失变形为只与 s(x;θ)s(x;\theta) 相关的形式:

J(θ)=xpreal(x)i=1n[xis(xi;θ)+12s(xi;θ)2]dx=Expreal[i=1n[xis(xi;θ)+12s(xi;θ)2]]J(\theta) = \int_x p_{real}(x)\sum_{i=1}^n[\triangledown_{x_i} s(x_i;\theta) + \frac12 s(x_i;\theta)^2]dx = \mathbb{E}_{x \sim p_{real}}[\sum_{i=1}^n[\triangledown_{x_i} s(x_i;\theta) + \frac12 s(x_i;\theta)^2]]

在SDE论文中,我们将会看到这样形式的Score-Matching损失,用于代替极大似然损失进行概率分布相似度的度量。

4 SDE视角下的DDPM与NCSN

4.1 Overview

以上讨论了这么多关于score函数的内容,这和我们的SDE框架下的Diffusion有什么关系呢?

一言以蔽之,当我们使用SDE框架来描述扩散过程时,我们需要使用score函数来描述我们的反向过程和损失函数。大体可见下图:

此前的生成模型

4.2 SDE视角下的Score-based Generative Model

4.2.1 Diffusion Process with SDE

我们使用SDE(随机微分方程)的标准形式来刻画一个扩散过程,在论文中表示为:
dx=f(x,t)dt+g(t)dwdx = f(x, t)dt + g(t)dw \\其中 dwdw 是一个布朗运动过程(也就是我们在EBM中、DDPM中提到的高斯噪声); f(x,t),RnRnf(x, t), \mathbb{R}^n \rightarrow \mathbb{R}^n 被称为" x(t)x(t) 的漂移系数",在DDPM中,它表现为加噪过程对 xtx_t 的放缩; g(t),RnRng(t), \mathbb{R}^n \rightarrow \mathbb{R}^n 是加噪过程中对于噪声尺度放缩的系数

显然这个公式中是微分的连续形式,但是我们的扩散场景事实上是从 x0x_0xTx_T 的逐步加噪,显然,我们需要将该式进行离散化处理。

xt+Δtxt=f(x,t)Δt+g(t)Δtε,εN(0,I)x_{t+\Delta t} - x_t = f(x, t)\Delta t + g(t) \sqrt {\Delta t} \varepsilon, \varepsilon \sim N(0, I)\\

将上式重写成概率的形式,即为:
p(xt+Δt|xt)N(xt+f(x,t)Δt,g2(t)Δt)p(x_{t+\Delta t} | x_t) \sim N(x_t + f(x, t)\Delta t, g^2(t) \Delta t) \\

4.2.2 Reverse Process with SDE

考虑扩散过程的逆过程,即试图从 xtx_t 重建 xt1x_{t-1} ,我们仍然写为一个概率的形式 p(xt|xt+Δt)p(x_t | x_{t+\Delta t}) ,进行推导变形:

p(xt|xt+Δt)=p(xt,xt+Δt)p(xt+Δt)=p(xt+Δt|xt)p(xt)p(xt+Δt)=p(xt+Δt|xt)exp[logp(xt)logp(xt+Δt)]p(x_t | x_{t+\Delta t}) = \frac{p(x_t, x_{t+\Delta t})}{p(x_{t+\Delta t})} = \frac{p(x_{t+\Delta t} | x_t) p(x_t) }{p(x_{t+\Delta t})} = p(x_{t+\Delta t} | x_t) exp[logp(x_t) - log p(x_{t+\Delta t})]\\此时,我们考虑 Δt0\Delta t\rightarrow 0 ,在 x=xt,t=tx = x_t, t= t 处对 logp(xt+Δt)logp(x_{t + \Delta t}) 进行泰勒展开:

logp(xt+Δt)logp(xt)+(xt+Δtxt)xlogp(xt)+Δtφφtlogp(xt)logp(x_{t + \Delta t}) \approx logp(x_t) + (x_{t + \Delta t} - x_t)\triangledown_xlogp(x_t) + \Delta t \frac{\varphi}{\varphi t} logp(x_t) \\对此结果代回原式:
p(xt|xt+Δt)=p(xt+Δt|xt)exp[(xt+Δtxt)xlogp(xt)Δtφφtlogp(xt)]p(x_t | x_{t+\Delta t}) = p(x_{t+\Delta t} | x_t) exp[ -(x_{t + \Delta t} - x_t)\triangledown_xlogp(x_t)- \Delta t \frac{\varphi}{\varphi t} logp(x_t)] \\由于我们知道 p(xt+Δt|xt)N(xt+f(x,t)Δt,g2(t)Δt)p(x_{t+\Delta t} | x_t) \sim N(x_t + f(x, t)\Delta t, g^2(t) \Delta t) ,同时考虑到 Δt0\Delta t \rightarrow 0 ,一阶泰勒展开后忽视高阶小量 Δt2\Delta t^2 ,故可以得到:

p(xt|xt+Δt)exp[||xt+Δtxtf(x,t)Δt||22g2(t)Δt(xt+Δtxt)xlogp(xt)Δtφφtlogp(xt)]=exp[||xtxt+Δt+[f(xt+Δt,t+Δt)g2(t+Δt)xlogp(xt+Δt)]Δt||22g2(t+Δt)Δt]p(x_t | x_{t+\Delta t}) \simeq exp[-\frac{|| x_{t + \Delta t} - x_t - f(x, t)\Delta t ||^2 }{2 g^2(t) \Delta t} -(x_{t + \Delta t} - x_t)\triangledown_xlogp(x_t)- \Delta t \frac{\varphi}{\varphi t} logp(x_t)] \\ = exp [ - \frac{|| x_t - x_{t + \Delta t} + [f(x_{t + \Delta t}, t + \Delta t) - g^2(t + \Delta t) \triangledown_x logp(x_{t + \Delta t})] \Delta t ||^2}{2 g^2(t + \Delta t) \Delta t}] \\在这里,我们将其还原为概率形式与离散形式,则为:
p(xt|xt+Δt)N(xt+Δt[f(xt+Δt,t+Δt)g2(t+Δt)xlogp(xt+Δt)]Δt,g2(t+Δt)ΔtI)xtxt+Δt=[f(xt+Δt,t+Δt)g2(t+Δt)xlogp(xt+Δt)]Δt+g(t+Δt)Δtε,εN(0,I)p(x_t | x_{t+\Delta t}) \sim N(x_{t + \Delta t} - [f(x_{t + \Delta t}, t + \Delta t) - g^2(t + \Delta t) \triangledown_x logp(x_{t + \Delta t})] \Delta t, g^2(t + \Delta t) \Delta tI) \\ x_t - x_{t+\Delta t} = - [f(x_{t + \Delta t}, t + \Delta t) - g^2(t + \Delta t) \triangledown_x logp(x_{t + \Delta t})] \Delta t + g(t + \Delta t) \sqrt{\Delta t}\varepsilon, \varepsilon \sim N(0, I) \\ 考虑 Δt0\Delta t \rightarrow 0 ,重写为微分方程形式,即为:

dxt=[f(x,t)g2(t)xlogp(xt)]dt+g(t)dwdx_t = [f(x, t) - g^2(t) \triangledown_x logp(x_{t})]dt + g(t)dw \\

4.2.3 Score-Matching Loss

通过上述微分方程不难看出,如果我们希望实现从 xtx_{t}xt1x_{t-1} 的重建过程,已知 f(x,t)f(x, t)g(t)g(t) ,我们的核心需要是求解当前的 s(t)=xlogp(xt)s(t) = \triangledown_x logp(x_{t}) ,显然这是t时刻对应分布的得分函数,因此我们可以使用一个神经网络去拟合此时的得分 sθ(xt,t)s_\theta(x_t,t) ,通过得分匹配来实现优化(在此前的推导中,我们得知在这里我们并不需要真实得分"ground truth",只需要其样本就可以实现优化),让我们拟合的分布足够逼近真正的得分函数:
Jt(θ)=xtp(xt)||sθ(x,t)xlogp(xt)||2dxtJ_t(\theta) = \int_{x_t} p(x_t)|| s_\theta(x, t) - \triangledown_x logp(x_{t}) ||^2 dx_t\\ 虽然事实上,我们所拥有的是来自0时刻的样本 x0x_0 ,虽然扩散过程是完全确定的,我们可以通过 x0x_0 连续推导出任意时刻的分布 xtx_t ,我们在这里不妨还是归到 x0x_0 的形式,以便于将其写为期望。p(xt)=p(xt|x0)p(x0)dx=Ex0preal[p(xt|x0)]p(xt|x0)=p(xt|xt1)p(xt1|xt2)...p(x1|x0)dx1...dxtp(x_t) = \int p(x_t | x_0) p(x_0) dx = \mathbb{E}_{x_0 \sim p_{real}}[p(x_t | x_0)] \\ p(x_t | x_0) = \int p(x_t | x_{t-1})p(x_{t-1} | x_{t-2})...p(x_1 | x_0) dx_1...dx_t \\ 由此,我们就可以将损失函数写为

L(θ)=EtJt(θ)=Etxtp(xt)||sθ(x,t)xlogp(xt)||2dxt=Et{Ext|x0||sθ(x,t)xlogp(xt)||2}=Et{Ext|x0Ex0preal[||sθ(x,t)xlogp(xt|x0)||2]}L(\theta) = \mathbb{E}_t J_t(\theta) = \mathbb{E}_t \int_{x_t} p(x_t)|| s_\theta(x, t) - \triangledown_x logp(x_{t}) ||^2 dx_t\\ = \mathbb{E}_t \{ \mathbb{E}_{x_t | x_0} || s_\theta(x, t) - \triangledown_x logp(x_{t}) ||^2\} \\ = \mathbb{E}_t \{ \mathbb{E}_{x_t | x_0} \mathbb{E}_{x_0 \sim p_{real}} [|| s_\theta(x, t) - \triangledown_x logp(x_{t}|x_0) || ^2 ] \} \\

也就和论文中公式对应:

至此,我们完成了对于结合SDE与score函数的视角来刻画扩散模型的扩散过程、反向过程与损失函数。

4.3 SDE视角下的NCSN与DDPM

4.3.1 Overview of NCSN (Noise Conditional Score Networks)

这个模型来自本文作者宋飏博士在2019NIPS上的一篇工作:[1907.05600] Generative Modeling by Estimating Gradients of the Data Distribution (arxiv.org)

在SDE中,被归纳为"SMLD"(SCORE MATCHING WITH LANGEVIN DYNAMICS),在本文中,给出一种较为naive的解释。在本文开头介绍的EBM中,我们试图拟合一个"能量函数"(其实就是得分)来据此通过Langevin dynamics从一个先验分布中生成样本。

而在NCSN中,仍然试图去拟合真正的得分函数 sθ(x,σ)s_\theta(x, \sigma) ,不同的是我们将生成拆成多时刻进行,每一时刻之间的递推通过Langevin dynamics采样来实现生成,就像是将多个EBM连接起来一样(如DDPM就像多个VAE的连接)。

其原文的算法流程图如下:

注意,这篇论文和SDE这篇的Notation并不一样,
具体表现为:算法流程图中刻画了从一个先验采样逐步去噪的过程,本质上是扩散过程的reverse process。因而,图中的 x0x_0 对应的是一个从先验分布的采样。然而,对于DDPM代表的扩散模型而言,一般特征空间先验是用加噪后的 xTx_T 表示,而 x0x_0 代表一个真实样本。
在SDE的background部分,其实是在描述正向过程,是和扩散模型的一般习惯保持一致的,因此本文后续将和SDE保持Notation一致性。

其中,刻画了一个噪声序列 {σi}i=1N\{ \sigma_i \}_{i=1}^{N} ,其中 σ1<σ2<...<σN\sigma_1 < \sigma_2< ...< \sigma_N ,对任意噪声水平有:

pσi(xi)=preal(x0)pσi(xi|x)dxpσi(xi|x0)N(x0,σi2I)p_{\sigma_i}({x_i}) = \int p_{real}(x_0) p_{\sigma_i}({x_i}|x) dx \\ p_{\sigma_i}(x_i|x_0) \sim N(x_0, \sigma_i^2I) \\ 其中,为了将样本空间能够和特征空间先验建立一个过程的映射,我们应该有 σ1\sigma_1 足够小,使得 pσ1pdatap_{\sigma_1} \approx p_{data} ,有 σN\sigma_N 足够大,使得 pσNN(0,σN2I)p_{\sigma_N} \approx N(0, \sigma_N^2I) ,这样我们就能实现逐步加噪与去噪,我们将上式子从概率形式重写写为方程形式:

xt=x0+σtε,εN(0,I)x_t = x_0 + \sigma_t \varepsilon, \varepsilon \sim N(0, I) \\由于上式子对任意t成立,我们可以将其转化为一个递归形式,注意到由于 xt+1x_{t+1}xtx_t 的噪声项是相互独立的,因而我们可以将其统一(类似我们在DDPM中从递推式变为通项一样):

xt+1=xt+σt+12σt2ε,εN(0,I)x_{t+1} = x_t + \sqrt{\sigma_{t+1}^2 - \sigma_{t}^2}\varepsilon, \varepsilon \sim N(0, I) \\ 这里,就可以体现所谓逐步加噪的特性,与DDPM所不同的是,NCSN的操作并不会进行对 x0x_0 本身的放缩,而是通过不断调整噪声水平实现加噪。相应的,其去噪的反向过程也可以认为是通过逐步缩小噪声水平,将一张噪声还原为原先图像。

4.3.2 SDE 框架

对于NCSN而言,我们将其中离散化变量先转化为连续量 {σi}σ(t),ziz(t),{xi}x(t)\{\sigma_i\} \rightarrow \sigma(t), z_i \rightarrow z(t), \{x_i\} \rightarrow x(t) ,变形得到:

x(t+Δt)=x(t)+σ(t+Δt)2σ(t)2εdx=dσ(t)2dtdwx(t+\Delta t) = x(t) + \sqrt{\sigma(t+\Delta t)^2 - \sigma(t)^2} \varepsilon \\ dx = \sqrt{\frac{ d\sigma(t)^2}{dt}}dw \\ 因而,我们可以说其实NCSN的扩散过程就是 f(x,t)=0,g(t)=dσ(t)2dtf(x, t) = 0, g(t) = \sqrt{\frac{ d\sigma(t)^2}{dt}} 的SDE,在文中被称为VE-SDE

对于DDPM而言,我们回顾其递归形式 xt+1=1βt+1xt+βt+1εx_{t+1} = \sqrt{1-\beta_{t+1}} x_t + \sqrt{\beta_{t+1}}\varepsilon (介绍可见论文学习笔记DDPM: 结合flows、VAE看Diffusion - 知乎 (zhihu.com)),类似上述进行转化:

x(t+Δt)=1β(t+Δt)x(t)+β(t+Δt)εdx=12β(t)xdt+β(t)dwx(t+\Delta t)= \sqrt{1-\beta(t+\Delta t)} x(t) + \sqrt{\beta(t+\Delta t)}\varepsilon \\ dx = -\frac12 \beta(t)xdt + \sqrt{\beta(t)}dw \\ 从这里,我们可以得到DDPM的扩散过程可以被描述为 f(x,t)=12β(t)x,g(t)=β(t)f(x, t) = -\frac12 \beta(t)x, g(t) = \sqrt{\beta(t)} 的SDE,在文中被成为VQ-SDE。

4.3.3 从SDE构造Diffusion Model

建立SDE和Diffusion Model的关联的一大益处在于,我们可以从SDE的角度开始设计新的Diffusion了。仅需要给定一个扩散过程的SDE,其反向过程、Loss函数也就可以随之确定,并可以使用score-based方法进行优化。当其与此前的Diffusion Model有所关联时,仅需要在backbone上较小的改动就可以开始训练。

在本文中,作者就提出一种sub-VQ-SDE,其形式如下:

也取得了较为不错的效果:

5 结语

作为Diffusion领域里程碑式的工作,建立SDE与Diffusion Model的关联后,使得我们可以用SDE理论的已知结论去优化Diffusion的设计,提供更多的设计灵活性和理论支撑。

6 写在后面

课内的活可算闲下来些了,半周时间读了四篇论文,感觉收获颇丰。感觉确实必须要进行一些理论层面的工作才能对深度学习这块有够深的理解,对于生成模型在专栏中仅剩下GAN没有写了。不过本周应该是搞不出来了,毕竟还有OO博客 & OS挑战性任务以及一些工程层面上的任务。理论的推导告一段落,该开始工程的实践了!

Reference

Implicit Generation and Modeling with Energy Based Models (neurips.cc)

hyvarinen05a.pdf (jmlr.org)

[1907.05600] Generative Modeling by Estimating Gradients of the Data Distribution (arxiv.org)

[2011.13456] Score-Based Generative Modeling through Stochastic Differential Equations (arxiv.org)

Generative Modeling by Estimating Gradients of the Data Distribution | Yang Song (yang-song.net)

生成扩散模型漫谈(五):一般框架之SDE篇 - 科学空间|Scientific Spaces

朗之万方程 - 维基百科,自由的百科全书 (wikipedia.org)

Langevin Monte Carlo: Sampling using Langevin Dynamics | Perceptron.blog

【概率方法】朗之万动力学 Langevin Dynamics-CSDN博客

【基础知识】DDPM中的“Langevin动力学”的概念_langevin dynamics-CSDN博客

有没有谁通俗的讲一下Denoising score matching? - 知乎 (zhihu.com)

[生成模型新方向]: score-based generative models-CSDN博客

扩散模型NCSN,使用分数匹配的生成模型,Generative Modeling by Estimating Gradients of the Data Distribution-CSDN博客

【理论推导】随机微分方程(SDE)视角下的Diffusion Model与Score-based Model_score based diffusion model-CSDN博客

发布于 2024-06-12 16:29・IP 属地北京 ,编辑于 2024-06-12 16:33・IP 属地北京
「真诚赞赏,手留余香」
还没有人赞赏,快来当第一个赞赏的人吧!
欢迎参与讨论

还没有评论,发表第一个评论吧

文章被以下专栏收录

想来知乎工作?请发送邮件到 jobs@zhihu.com