引

• 还记得我们曾经在lab1中用过一个$∇\log𝑝(𝑋_𝑡)$来进行ODE吗，本文将详细介绍与之有关的知识。
• 在上一讲中我们已经推导出了vector field的公式，并且给出了flow matching的训练算法，现在我们将学习如何处理diffusion model 并且如何使用score matching来训练它。

Conditional and Marginal Score Functions

• 什么是score function? 对于一个概率分布p(x)，它的score function定义为： $$s(x) = \nabla \log p(x)$$
• 这个函数的意义在于它告诉我们在x这个点上，对于x的log-likelihood的最快增长方向，也就是它的梯度。

Conditional Score Function

• 于是，我们定义Conditional Score Function为：

  $$s_t(x|z) = \nabla_x \log p_t(x|z)$$

• 我们可以以之前提到过的Gaussian conditional path为例来求解一下它的conditional score function。 已知：

  \[p_t(x|z) = \mathcal{N}(\alpha_t z, \beta_t^2 I_d)\]

  所以

  \[p_t(x|z) = \frac{1}{(2\pi)^{d/2} \beta_t^d} \exp\left(-\frac{1}{2\beta_t^2} \|x - \alpha_t z\|^2\right)\]
  \[\log p_t(x|z) = -\frac{d}{2} \log(2\pi) - d \log \beta_t - \frac{1}{2\beta_t^2} \|x - \alpha_t z\|^2\]
  \[s_t(x|z) = \nabla_x \log p_t(x|z) = -\frac{1}{\beta_t^2} (x - \alpha_t z)\]

• 我们会发现：哎嘿，这个score function恰好是x和z的一个线性组合，而我们之前求出的gaussian conditional path的vector field也是x和z的一个线性组合，我们完全可以将它们之间进行一个转化。 已知：(推导过程见上一篇)

  \[u_t^{\text{target}}(x|z) = \frac{\dot{\beta_t}}{\beta_t} x + \left( \dot{\alpha_t} - \frac{\dot{\beta_t}}{\beta_t} \alpha_t \right) z \]

令：

则有：

Marginal Score Function

• 同理定义一个

  $$s_t(x) = \nabla_x \log p_t(x)$$

• 我们可以先凭感觉定义一个Marginal Score Function:

  \[ s_t(x) = \nabla_x \log p_t(x) = \int \nabla_x \log p_t(x|z) \frac{p_t(x|z) \, p_{\text{data}}(z)}{p_t(x)} \, dz \]

  这个公式依然看起来很合理，对x求marginal score,我们就先固定x，然后对所有的z进行一个后验概率加权平均。

• 证明：

  $$ \begin{align*} \nabla \log p_t(x) &= \frac{\nabla p_t(x)}{p_t(x)} \\ &= \frac{\nabla \int p_t(x|z) p_{\text{data}}(z) \, dz}{p_t(x)} \\ &= \int \frac{\nabla p_t(x|z) p_{\text{data}}(z)}{p_t(x)} \, dz \\ &= \int \nabla \log p_t(x|z) \, \frac{p_t(x|z) p_{\text{data}}(z)}{p_t(x)} \, dz \end{align*} $$

  直接就推导出来了，奇妙，中间运用了两次$log$的微分法则。

• 然后我们可以将上面conditional score function在gaussian conditional path的例子中求出的结果带入$u_t^{\text{target}}(x)$的公式，得到：

  \[u_t^{\text{target}}(x) = a_t s_t(x) + b_t x\]

conditional and marginal denoiser

• 本段视频中老师没讲，纯纯个人理解，如果有错误欢迎指正。
• 在gaussian probility path的例子中，我们的$u_t^{\text{target}}(x|z)$和$s_t(x|z)$都是x和z的线性组合，而它们的marginal version则是在它们的基础上进行一个对z的后验概率加权积分。
• 于是note中提到它们都可以看成是 a linear reparameterization of the posterior mean $E[z|x]$, 也就是说我们给定一个噪声x，然后我们通过后验概率加权平均的方式得到一个对z的估计，然后就能求出score function 或者是 vector field了。这个东西有什么用呢，note中说一是形式简单，二是可能会数值/训练稳定好。
• conditional denoiser就是$D_t(x|z) = z$， marginal denoiser就是$D_t(x) = E[z|x] = \int z , \frac{p_t(x|z) p_{\text{data}}(z)}{p_t(x)} , dz \stackrel{(i)}{=} \frac{1}{\dot{\alpha}_t \beta_t - \alpha_t \dot{\beta}_t} \left( \beta_t u_t^{\text{target}}(x_t) - \dot{\beta}_t x_t \right)$，它们都是对z的一个估计。 其中(i)是通过带入$u_t^{\text{target}}(x|z)$与z的关系式就得到了。

The denoiser has a very intuitive interpre-tation: it is the expected value of clean data z given noisy data x.People often call such models denoising diffusion models as learning D_t and learning u_target are theoretically equivalent. 这里note中提到了一个很有意思的点，will the denoiser always output a “clean” data point? Why or why not, and what might this depend on? 这个让我想起了EAI课程中wh老师举过一个例子，在无人驾驶中，如果路中间有一个障碍物，那么我们记录左转和右转都是有效，写出一个loss函数是 (x-1)^2 + (x+1)^2, 但问题是这个loss函数最优解是x=0，完全撞上了障碍物，大错特错了。这里的道理是一样的，如果我们的data distribution有多个峰值，那么对于一个噪声数据点x来说，它的后验概率可能在多个峰值上都有比较大的值，那么它的denoiser输出的就是这些峰值的一个加权平均，可能并不是一个“clean”的数据点了。

Sampling with SDEs

• 我们之前展示过如何利用vector field来进行ODE sampling了，那么对于diffusion model来说，我们就需要用到SDE了。
• SDE Extension Trick, 仿照flow matching的ODE版本，我们定义： $$ \begin{aligned} X_0 &\sim p_{\text{init}}, \\ dX_t &= \left( u_t^{\text{target}}(X_t) + \frac{\sigma_t^2}{2} \nabla \log p_t(X_t) \right) dt + \sigma_t \, dW_t, \\ \Rightarrow \quad X_t &\sim p_t \quad (0 \le t \le 1). \end{aligned} $$ 你可能觉得很奇怪，我们原本的SDE想法是一个确定的vector field加上一个noise项，为什么这里要加上一个score function呢？这一点其实很好理解，我们一般是先确定了一个probability path, 然后计算vector field来进行ODE sampling的，但是对于SDE来说，加上了一个noise，会导致中间的$p_t$发生变化，所以我们需要加上这个score function来修正noise带来的影响。

加上这个修正项，可以保证每个时刻的p_t总概率分布是正确的，只不过每一个x采样的路径可能会受到noise影响而歪歪扭扭。这里是对于任何一个σ_t都成立的，但是由于训练中的误差，采样时的误差，我们可以挑选一个最佳的σ_t来得到最好的采样效果。

• 更进一步的，可以代入marginal score function：（Only for gaussian probability path‼️）

  $$X_0 \sim p_{\text{init}}, \quad dX_t = \left( (a_t + \frac{\sigma_t^2}{2}) \nabla \log p_t(X_t) + b_t X_t \right) dt + \sigma_t \, dW_t$$
  $$\Rightarrow \quad X_t \sim p_t \quad (0 \le t \le 1)$$

• 如何证明这个SDE Extension Trick呢？只要证明它满足Fokker-Planck equation就可以了。

首先定义 $\Delta$ ：

for scalar field $w_t: \mathbb{R}^d \to \mathbb{R}$ 其实就是对一阶导数算div。

证明：

符合Fokker-Planck equation的格式，于是新的$u_t$就是 $u_t^{\text{target}} + \frac{\sigma_t^2}{2} \nabla \log p_t$。

Langevin Dynamics

• 小表格中提到了Langevin Dynamics, 其实就是一个$p_t = p$的特殊情况，也就是说我们在采样过程中保持$p_t$不变，此时我们令$u_t^{\text{target}} = 0$ (带入Fokker-Planck equation你会发现是成立的)，于是得到了Langevin Dynamics的SDE： $$ dX_t = \frac{\sigma^2}{2} \nabla \log p(X_t) dt + \sigma dW_t $$
• 这个SDE的意义在于它可以让我们采样到一个固定的分布p，而且它具有收敛性，即便$p_{init}$离p十万八千里，只要有足够的时间，最后一定会趋向于p.
• 开头提到的我们曾经在lab1中就使用了这个公式从一个随机噪声分布收敛到一个高斯分布！用的就是这个公式了。

Optional: GLASS Flows, Stochastic evolution with ODEs

• The remarkable property of SDE sampling (compared to ODEs) is that the evolution becomes stochastic, i.e. the initial point X0 does not fully determine Xt for t > 0. Perhaps surprisingly, it is also possible to get the same stochastic transitions purely via ODEs via a simple sampling trick called GLASS Flows [20]. This allows to exploit the stochastic nature of SDEs (e.g. via search algorithms) while keeping the efficiency of ODEs.
• SDE能够让采样更多样性，可以在采样中跳过一些训练误差带来的不良影响，但是效率更低；GLASS Flows则是一个采样技巧，可以让我们在保持ODE效率的同时，也能得到SDE的采样多样性。

Score Matching

• 首先定义loss函数： $$\begin{align} \mathcal{L}_{\text{SM}}(\theta) &= \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, x \sim p_t(\cdot|z)} \left\| s_t^\theta(x) - \nabla \log p_t(x) \right\|^2 \\ \mathcal{L}_{\text{CSM}}(\theta) &= \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, x \sim p_t(\cdot|z)} \left\| s_t^\theta(x) - \nabla \log p_t(x|z) \right\|^2 \end{align}$$ 问题是这个$\nabla \log p_t(x)$是一个关于z的积分，算个🥚，所以第一个loss直接放弃，和之前flow matching类似的我们可以证明：

where $C$ is independent of parameters $\theta$. Therefore, their gradients coincide:

In particular, for the minimizer $\theta^$, it will hold that $s_t^{\theta^} = \nabla \log p_t$.

• note中又以gaussian为例子详细描述了Denoising Diffusion Models 以及最后的训练算法，这里最重要的insight就是为什么要叫denoising, 通过带入公式我们可以发现，如果$\beta_t$比较小，那么很容易数值爆炸(也就是要想训练出来的话必须加很多的噪声)。我们可以先提取出$\beta_t$, 舍掉它，然后你会发现优化项里是$\beta_t s_t^\theta(x) + \epsilon$, 于是可以令

  $$-\beta_t s_t^\theta(x) = \epsilon_t^\theta(x) \quad \Rightarrow \quad \mathcal{L}_{\text{DDPM}}(\theta) = \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, \epsilon \sim \mathcal{N}(0, I_d)} \left\| \epsilon_t^\theta(\alpha_t z + \beta_t \epsilon) - \epsilon \right\|^2$$

  这实际上就是用$\theta$ 来在固定z，固定t，固定x的情况下去预测原始噪声$\epsilon$，所以叫denoising了。

• note的最后也是给出了训练的伪代码，当然loss函数我们既可以用score版本也可以用noise版本。其实，在使用gaussian probability path的情况下，这些量之间都是线性关系，都可以相互转化，我们完全可以只搞一个最简单的训练器，比如(noise或者D(z|x)),然后转化得到vector field和score function, 最后应用SDE采样就好了。