简介

分数蒸馏采样（Score Distillation Sampling, SDS）是当今3D生成的一大流派，这个方法利用了预训练的2D扩散模型就能生成3D物体，而无需真实的3D数据。

原理推导

扩散模型的训练方式是：

$$\mathcal{L}_{\textrm{Diff}}(x;\phi)=\mathbb{E}_{t,\epsilon}\left[w(t)\|\epsilon_\phi(x_t, t)-\epsilon\|_2^2\right] \tag{1}$$

对于一个训练好的扩散模型，当给定一张图片$x$，这张图片在这个扩散模型训练集的分布内，那么这个损失函数将会较小；反之则较大。 SDS正是利用这个原理：如果一个3D模型的每一个角度渲染出来的图像都在预训练扩散模型的分布内（看起来比较真实），那么这个3D模型整体也就看起来比较真实。 在SDS算法中，优化一个NeRF，这个NeRF不断从不同的角度渲染图像，然后将这个图像交给预训练扩散模型计算公式(1)中的损失，即： 假设这个NeRF由$\theta$参数化，从角度$v$渲染图像$x=x(\theta,v)$，则使用SDS优化$\theta$的方法是：

$$$$

但是，公式(2)是不高效的，求解梯度得到：

$$=\mathbb{E}_{t,\epsilon,v}\left[w(t)\nabla_\theta\| \epsilon_\phi(x_t(\theta,v))-\epsilon \|^2_2\right]\\ =2\mathbb{E}_{t,\epsilon,v}\left[w(t)\frac{\partial\epsilon_\phi(x_t(\theta,v))}{\partial x_t(\theta,v)}\frac{\partial x_t(\theta,v)}{\partial\theta}\left(\epsilon_\phi(x_t(\theta,v))-\epsilon\right)\right]\tag{3}$$

求解(*)需要计算扩散模型UNet的Jacobian（反向传播），效率很低。所以在SDS算法中，不直接使用公式(2)进行优化。 在SDS算法中，发现直接将梯度中的(*)省去也能有很好的效果，即：

$$=\nabla_\theta\mathbb{E}_{t,\epsilon,v}\left[w(t)\| \epsilon_\phi(x_t(\theta,v))-\epsilon \|^2_2\right]\\ =2\mathbb{E}_{t,\epsilon,v}\left[w(t)\frac{\partial x_t(\theta,v)}{\partial\theta}\left(\epsilon_\phi(x_t(\theta,v))-\epsilon\right)\right]\tag{4}$$

这个梯度相当于最小化以下KL散度：

$$\nabla_\theta\mathcal{L}_{\mathrm{SDS}}=\nabla_\theta\mathbb{E}_{t,\epsilon,v}\left[ \sigma_t/\alpha_tw(t)\mathrm{KL}\left( q(x_t(\theta,v)|x_0(\theta,v))\|p_\phi(x_t(\theta,v)) \right) \right]\tag{5}$$

证明： $(5)=\mathbb{E}_{t,\epsilon,v}\left[ \sigma_t/\alpha_tw(t)\nabla_\theta\mathrm{KL}\left( q(x_t(\theta,v)|x_0(\theta,v))\|p_\phi(x_t(\theta,v)) \right) \right]$ 其中，$\nabla_\theta\mathrm{KL}\left( q(x_t(\theta,v)|x_0(\theta,v))\|p_\phi(x_t(\theta,v)) \right) =\mathbb{E}_{q}\left[\nabla_\theta \log q(x_t(\theta,v)|x_0(\theta,v)) - \nabla_\theta\log p_\phi(x_t(\theta,v)) \right]$ 第一项：$\nabla_\theta \log q(x_t(\theta,v)|x_0(\theta,v)) =\nabla_\theta\log\left( -\frac{1}{\sqrt{2\pi}\sigma_t}e^{-\frac{(x_t(\theta,v)-\alpha_t x_0(\theta,v))^2}{2\sigma_t^2}} \right) =\nabla_\theta\log\left( -\frac{1}{\sqrt{2\pi}\sigma_t}e^{-\frac{\epsilon^2\sigma_t^2}{2\sigma_t^2}} \right)=0$ 第二项：$-\nabla_\theta\log p_\phi(x_t(\theta,v))\approx \frac{\alpha_t}{\sigma_t}\epsilon_\phi(x_t(\theta,v),t)\frac{\partial x_t(\theta,v)}{\partial\theta}$ 所以，$(5)=\mathbb{E}_{t,\epsilon,v}\left[ w(t)\epsilon_\phi(x_t(\theta,v),t)\frac{\partial x_t(\theta,v)}{\partial\theta} \right]$ 对于一个0均值的高斯变量$\epsilon$，在期望的意义下为0，可以继续将(5)改写为： $(5)=\mathbb{E}_{t,\epsilon,v}\left[ w(t)(\epsilon_\phi(x_t(\theta,v),t)-\epsilon)\frac{\partial x_t(\theta,v))}{\partial\theta} \right]$ 证明完毕！

SDS分析

公式(5)说明了SDS的优化目标是最小化NeRF从各个角度采样的图片加噪后的分布和扩散模型先验分布的KL散度，但是对于一个NeRF来说，它只表示一个场景，而扩散模型是大量图片形成的分布。SDS强行减少两个分布的KL散度，会导致NeRF渲染图的分布更加地平滑。所以SDS就会导致NeRF渲染图尽量和先验分布的所有图的距离都一样大，会导致多样性低、模糊的现象。 从优化方向（梯度）的角度来说，(5)式的是让NeRF渲染图朝着score，即最大化似然的方向优化，这就变成了mode seeking的方法，不论是什么样的初始化，最终会收敛到同一个似然最大的地方（注意这是一个高斯分布），多样性降低。