DreamFusion 论文学习

扩散模型在 text-to-image 任务上有显著的效果，如 DALLE.E2。然而 text-to-3D 由于 3D 模型数据集不足，导致直接训练 text-to-3D 的扩散模型相当困难。DreamFusion¹ 跳过了这种训练框架，利用 text-to-image 与 NeRF 技术，达到 text-to-3D 的生成。

扩散模型算法²

训练

扩散模型在训练 noise predictor 时，会从目标分布中采样一张图片，接着在这张图片上叠加噪声。这里可以注意到，扩散模型会假设一个逐渐变小的数列 $({\bar{\alpha }}_1,{\bar{\alpha }}_2,...,{\bar{\alpha }}_T)$，随机叠加上对应的噪声。

可以看到随着${\bar{\alpha }}_T$ 越大，噪声加的就越厉害。而整个 noise predictor network 在做的事情就是将噪声图像和步骤3取样的$t$作为输入，去预测步骤4取样的噪声$ϵ$。

推理

产生图像的过程如下所示。步骤4简言之就是将噪声图像${\mathbf{x}}_t$渐渐地去噪，恢复到清晰的图像。这里暗藏玄机的部分在于，每个$t$循环都会在最后加上噪声$\mathbf{z}$。

影响模型生成的本质

不论是扩散模型、GAN 或 VAE 在这种图像生成模型，它们的本质都是去最小化图像生成模型产生的$P$分布，与搜集到的图像数据集产生的$P_{data}(x)$分布之间的某种 divergence。本质上就是去寻找$P_{\theta }(x)$与$P_{data}(x)$ 之间的最大似然。

式子的推导如下：

虽然我们无法得知$P_{\theta }(x)$，但我们可用生成模型输出高斯分布的方式去近似，我们先以 VAE 为例：

最大化 $\log P_{\theta }(x)$ 的过程经过一番推导，可以改用最大化 ELBO(evidence lower bound) 去表示：

KL散度（Kullback-Leibler divergence），也称为相对熵，是一种衡量两个概率分布 $A$ 和 $B$ 差异的指标。KL散度是不对称的，即$\text{KL}(A||B)\ne \text{KL}(B||A)$，
对于离散和连续的概率分布，分别为：
$$\begin{gathered} \text{KL}(A||B)=\sum _{x}A(x)\log (\frac{A(x)}{B(x)}), \\ \text{KL}(A||B)=\int A(x)\log (\frac{A(x)}{B(x)})dx, \\ \text{KL}(A||B)\ge 0 \end{gathered}$$

扩散模型³

如下所示，扩散模型是一个马尔可夫链，$q(x_t|x_{t-1})$ 可以看作为 encoder，$p(x_t|x_{t+1})$ 可以看作为 decoder。

接着，我们同样去试着最大化$\log p(x)$，可以得到 ELBO 如下：

ELBO 的部分就可如下展开：

把最后一项放大来看：

红框内在做的事，就是去试着让$q(x_{t-1}|x_t,x_0)$ 和 $P(x_{t-1}|x_t)$ 越接近越好，也就是下图的粉红色与绿色箭头的部分：

而我们可以通过马尔可夫性质与贝叶斯定理，将$q(x_{t-1}|x_t,x_0)$表示为：

$$q(x_t|x_{t-1},x_0)=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$$

由于我们知道整个编码的过程，其实就是在加高斯噪声，因此可以得到下面的关系式：

$$q(x_t|x_{t-1},x_0)=q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)\mathbf{I})$$

更改一下符号就可以表示为：

$$x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}ϵ\text{with}ϵ\sim \mathcal{N}(ϵ;0,\mathbf{I})$$

$t-1$ 这个时间点的图像同上逻辑可以得到：

$$x_{t-1}=\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}ϵ\text{with}ϵ\sim \mathcal{N}(ϵ;0,\mathbf{I})$$

透过递归的表示，可以得出$q(x_t|x_0)$:

而这也是我们在扩散模型中看到的噪声图像。接着我们继续试着找出 $q(x_{t-1}|x_t,x_0)$。

这里，我们引入新符号来表示 variance 部分，

$${\sigma }_q^2(t)=\frac{(1-{\alpha }_t)(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}$$

如上，由于我们可以控制 $q(x_{t-1}|x_t,x_0)$和$P(x_{t-1}|x_t)$的 variance，因此整个最小化 KL散度的过程就是在缩小$q(x_{t-1}|x_t,x_0)$和$P(x_{t-1}|x_t)$的均值。

详细推导如下：

其中$q(x_{t-1}|x_t,x_0)$的均值可以表示为：

$${\mu }_{q(x_t,x_0)}=\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})x_t+\sqrt{{\bar{\alpha }}_{t-1}}(1-{\alpha }_t)x_t}{1-{\bar{\alpha }}_t}$$

透过先前的定义，可以将上式右边的 $x_0$ 替换：

可以得到均值为：

上式有换过符号，${\beta }_t$为$1-{\alpha }_t$。

这个过程，其实就是下图推理中红框内的部分。其中$ϵ$即为我们的 noise predictor network。

详细推导如下：

Score Function: Another Perspective of Diffusion Models⁴

Tweedie’s Formula 是 ${\hat{\theta }}^{TE}=x+{\sigma }^2\frac{d}{dx}\log p(x)$，可以看出，Tweedie’s Formula 估计的第一项就是 MLE(Maximum Likelihood Estimation)，而第二项可以看成在 MLE 基础上做的贝叶斯修正，这里要注意的是，Tweedie’s Formula 给的仅是对后验的期望。

通过 Tweedie’s Formula，我们可以将参数估计范化为：

$$\mathbb{E}[{\mu }_z|z]=z+{\Sigma }_z{\nabla }_z\log p(z)$$

其中$z\sim N(z;{\mu }_z,{\Sigma }_z)$。而扩散模型中，$q(x_t|x_0)$ 可以表示为：

$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$$

将上述两个式子结合，我们可以得到：

$$\mathbb{E}[{\mu }_{x_t}|x_t]=x_t+(1-{\bar{\alpha }}_t){\nabla }_{x_t}\log p(x_t)$$

而在 Tweedie’s Formula 中，均值的最佳估计就是我们实际 encoder 中的均值，也就是

$${\mu }_{x_t}=\sqrt{{\bar{\alpha }}_t}{x}_0$$

由此我们可以得到$x_0$：

$$\begin{gathered} \sqrt{{\bar{\alpha }}_t}{x}_0=x_t+(1-{\bar{\alpha }}_t)\nabla \log p(x_t) \\ \therefore x_0=\frac{x_t+(1-{\bar{\alpha }}_t)\nabla \log p(x_t)}{\sqrt{{\bar{\alpha }}_t}} \end{gathered}$$

将 $x_0$ 带入上面的 denoising transition mean（去噪转换均值）${\mu }_q(x_t,x_0)$ 中：

接着我们透过下式假设一个 denoising transition mean:

$${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}{x}_t+\frac{1-{\alpha }_t}{{\sqrt{\alpha }}_t}{s}_{\theta }(x_t,t)$$

因此根据扩散模型，最终对 denoising transition 的推导可以写为：

这里，我们可以对比一下之前的部分：

因此，整个优化工作就变成了训练一个神经网络$s_{\theta }$，使其接近$\nabla \log p(x_t)$，这个$\nabla \log p(x_t)$就叫做 score function。这种函数通常指的是概率分布的梯度，即概率密度函数相对于输入变量的导数。

利用 Tweedie’s Formula 以及初始图像 $x_0$，

我们可以发现 score function 其实和 source noise 相关。

Score function 的意义就是我们如何在数据空间中移动，区最大化 log probability。直觉上来说，将 source noise 加到图像上使图像变得模糊失真，在相反方向上进行调整，那么就可以对图像去噪，也因此最能增大后续的 log probability。

Score Distillation Sampling⁵

Differentiable Image Parameterization

DreamFusion¹ 的扩散过程并不是从 RGB 空间去采样图片，而是利用了 DIP。只要图片的参数化是可微分的，我们就可以通过它做反向传播。

简单来说，DIP 就是利用其他可微分的参数来表示一张图片。我们用$\theta$来表示 3D volume 的参数，那么采样到的图片就可以用$x=g(\theta )$来表示，$g$ 为一个 transformation function，将$\theta$参数化为图片$x$。

我们对采样像素并没兴趣，我们想要创建一个 3D 模型，从任意角度渲染的时候，图片看上去都很逼真。

Score-based Generative Modeling

我们可以将扩散模型理解为一个 score-based generative model。

首先介绍一下 EBM(Energy-based Model)：

$$p_{\theta }(x)=\frac{1}{Z_{\theta }}\exp (-f_{\theta }(x))$$

引用自物理力学，$f_{\theta }(x)$是灵活且可参数化的能量函数，因此可以用神经网络去建模。而$Z_{\theta }$称作 Normalizing Constant，确保$\int p_{\theta }(x)dx=1$。透过对 EBM 公式进行以下操作：

换句话说，我们用来建模 $f_{\theta }(x)$ 的神经网络，其实就是之前提到的 score function。这个神经网络可以通过最小化 Fisher Divergence 去优化：

$${\mathbb{E}}_{p(x)}=\left[{\Vert s_{\theta }(x)-\nabla \log (p(x))\Vert }_2^2\right]$$

Fisher Divergence（Fisher散度）是一种衡量两个概率分布差异的度量。对于随机变量$U$和$V$，概率密度函数分别为$f$和$g$，Fisher Divergence 定义为：$I(U||V)=I(f||g)=\int f(x){\Vert \nabla \log f(x)-\nabla \log g(x)\Vert }_2^{2}dx$。它基于 Fisher Information 的概念，Fisher Information 是关于概率分布参数的信息量度量。它的定义涉及到两个概率分布$P$和$Q$相对于某个参数$\theta$的 Fisher Information 矩阵$\mathcal{I}(\theta )$和它们的参数${\theta }_P$和${\theta }_Q$的差异。Fisher Information 矩阵是一个关于参数$\theta$的二阶导数的期望值构成的矩阵，可以表示为：$\mathcal{I}(\theta )=\mathbb{E}\left[-\frac{{\partial }^2}{\partial {\theta }^2}\log P(X|\theta )\right]$。

对图像空间中的$x$的 log likelihood 计算梯度，其实就是找到在数据空间中，可以增加 $p(x)$ likelihood 的方向。

Score function 的可视化如下图所示，水平面即$x$所在的数据空间，纵轴的尖峰（密度较高的区域）可以看作为模型最终收敛的集中可能模式。像这样透过 score function 表示 $p(x)$ 的分布并用 MCMC 去产生样本，就是所谓的 score-based generative modeling。

学习将一个分布表示为 score function，然后用这个函数通过 Markov Chain Monte Carlo(MCMC) 技术（如 Langevin dynamics）生成符合该分布的样本。这种方法在生成模型领域，特别是深度学习中，是一种新兴的研究方向。

马尔可夫链蒙特卡洛（MCMC）是一类用于从复杂概率分布中生成样本的算法。这些技术通过构建一个马尔可夫链来逐步探索状态空间，最终达到与目标分布相同的分布。

Langevin动力学是一种特定的MCMC技术，它通过模拟物理系统在热力学平衡状态下的行为来生成样本。这种动力学以Paul Langevin命名，他在研究布朗运动时引入了这一概念。在机器学习中，Langevin动力学可以用来模拟概率分布的梯度下降过程，从而生成样本。

SDS

DreamFusion¹ 定义扩散模型训练的损失如下，$w(t)$ 为加权函数，其余和下图的 Gradient Descent 一致：

$${\mathcal{L}}_{\text{Diff}}(\varphi ,x)={\mathbb{E}}_{t\sim \mathcal{U}(0,1),ϵ\sim \mathcal{N}(0,\mathbf{I})}\left[w(t){\Vert {ϵ}_{\varphi }({\alpha }_t\mathbf{x}+{\sigma }_tϵ;t)-ϵ\Vert }_2^2\right]$$

将上式中的图像$\mathbf{x}$专程 DIP $\mathbf{x}=g(\theta )$，我们的目的就变成最小化以下的损失函数：

$${\theta }^{\ast }=\underset{\theta }{\mathrm{argmin}}{\mathcal{L}}_{\text{Diff}}(\varphi ,\mathbf{x}=g(\theta ))$$

然而实验结果证实这样做的效果并不好，为何如此？我们首先对损失函数取梯度，可以得到：

$z(t)$ 是添加的噪声。

DreamFusion¹ 中提到，U-Net Jacobian 这个部分需要巨大的计算资源，并且在噪声小的情况下效果不佳，因此 DreamFusion¹ 直接省略了这个部分，并将剩下的部分定义为 SDS 的梯度：

$${\nabla }_{\theta }{\mathcal{L}}_{SDS}(\varphi ,\mathbf{x}=g(\theta ))\triangleq {\mathbb{E}}_{t,ϵ}\left[w(t)({\widehat{ϵ}}_{\varphi }({\mathbf{z}}_t;y,t)-ϵ)\frac{\partial \mathbf{x}}{\partial \theta }\right]$$

从 score-based generative model 的角度看，这其实就是去找到数据空间中的某个可能模式，也就是密度较高的区域。DreamFusion¹ 将这个方法称为 Score Distillation Sampling。

Algorithm⁶

扩散模型的部分，DreamFusion¹ 使用了 Imagen，而 3D 模型的生成则采用了 NeRF 的架构。首先随机初始化一个 NeRF，接着重复地随相机位置去渲染，再将渲染出来的图像和 Imagen 计算 SDS 损失，并做优化。整个流程如下图所示：

Neural Rendering of a 3D Model

NeRF 采用体积渲染，由相机位置向各个像素点方向发射射线，射线经过的 3D 位置 $\mu$ 会得到不透明度 $\tau$ 和 RGB 颜色 $c$。最终像素的值即为该射线一路从相机过来所经过的所有 $\mu$ 之加权总和。

$$\mathbf{C}=\sum _i{w}_i{\mathbf{c}}_i,w_i={\alpha }_i\prod _{j<i}(1-{\alpha }_j),{\alpha }_i=1-\exp (-{\tau }_i\Vert {\mu }_i-{\mu }_{i+1}\Vert )$$

DreamFusion¹ 中，针对原始 NeRF 有以下几项改进：

Shading

传统的 NeRF 针对每个 3D 位置计算不透明度$\tau$和 RGB 颜色 $\mathbf{c}$，而在 DreamFusion¹ 中则是计算 3D 模型的每个表面如何着色，也就是着色器（Shader）的概念。该 Shader MLP 会以 3D 位置 $\mu$ 和相机位置参数 $\theta$ 作为输入，输出不透明度 $\tau$ 与物体表面的 RGB 反射率 $\rho$（可视为材质的颜色）：

$$(\tau ,\rho )=\text{MLP}(\mu ;\theta )$$

有了$\tau$和$\rho$，我们需要物体表面的法向量$\mathbf{n}$，去计算 3D 位置在该物体表面最终的颜色。$\mathbf{n}$ 可以透过 3D 位置$\mu$上的$\tau$的反向梯度计算出来：

$$\mathbf{n}=-{\nabla }_{\mu }\tau /\Vert {\nabla }_{\mu }\tau \Vert$$

接着假设点光源位置$l$与颜色$l_p$，采用漫反射（Diffuse Reflection）模型，可对 3D 位置 $\mu$ 进行着色：

$$\mathbf{c}=\rho \circ (l_p\circ \max (0,\mathbf{n}\cdot (l-\mu )/\Vert l-\mu \Vert )+l_a)$$

Scene Structure

DreamFusion¹ 将整个 NeRF 的范围限制在一个球体内，并且该球体的背景是通过另一组 MLP 参数生成的。

Geometry Regularizers

DreamFusion¹ 针对射线经过的 3D 位置所计算出的不透明度，会有额外的正则惩罚，这么做是为了避免在 empty space 中塞入错误的噪声。

Text-to-3D Synthesis

文本到 3D 的生成可以拆解为四个步骤：

• 随机采样一个相机位姿和光照；
• 从该相机视角和光线阴影，渲染 NeRF 的一张图像；
• 计算 SDS 损失关于 NeRF 参数的梯度；
• 使用优化器更新 NeRF 的参数。

下面针对这四个步骤进行解释：

1. 随机采样一个相机和光照

在每次训练步骤，相机的极坐标参数会从下面的范围随机选择：

• 仰角 $\varphi \in [-10°,90°]$
• 水平角 $\theta \in [0°,360°]$
• 与原点的距离 $r\in [1,1.5]$

另外，焦距$\lambda$与点光源位置$l$从以下范围得到：

• 焦距$\lambda \in U(0.7,1.35)$
• $l$从以相机位置为均值的分布中抽样

2. 渲染

NeRF渲染的解析度为 $64\times 64$，搭配预训练的 Text-to-Image Model。

3. Diffusion loss with view-dependent conditioning

在文字提示上，DreamFusion¹ 会提供针对相机参数的额外提示，比如仰角$>60°$时，会加上 “overhead view” 的字眼。另外，DreamFusion¹ 采用的是 T5-XXL Text Embedding。

4. 优化

针对训练环境，在一台装有四块 TPUv4 芯片的机器上进行 3D 场景的优化。每张芯片渲染一个单独的视图，每个设备的 batch size 为1 评估扩散 U-Net 模型。进行了1.5万次迭代优化，这个过程需要大概 1.5 小时。计算时间平均分配在渲染 NeRF 和评估扩散模型上。使用 Distributed Shampoo 优化器来优化参数。

---

1. https://arxiv.org/pdf/2209.14988.pdf ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎

2. https://medium.com/@tonytsai225/3d%E7%94%9F%E6%88%90-dreamfusion-part-1-fb5027805389 ↩︎

3. https://arxiv.org/pdf/2208.11970.pdf ↩︎

4. https://medium.com/@tonytsai225/3d%E7%94%9F%E6%88%90-dreamfusion-part-2-score-based-generative-model-d432650b2011 ↩︎

5. https://medium.com/@tonytsai225/3d%E7%94%9F%E6%88%90-dreamfusion-part-3-score-distillation-sampling-fdfc766ecf74 ↩︎

6. https://medium.com/@tonytsai225/3d%E7%94%9F%E6%88%90-dreamfusion-part-4-algorithm-2ed445393430 ↩︎