2D 扩散模型 + Nerf,实现文本生成 3D 模型

2022-11-07 15:05:13 浏览数 (1)

来源:Arxiv 论文作者:Ben Poole, Ajay Jain, Jonathan T. Barron, Ben Mildenhall 项目主页:https://dreamfusion3d.github.io/ 论文链接:https://arxiv.org/pdf/2209.14988.pdf 内容整理:王炅昊 本论文把近期发展火热的2D扩散模型和Nerf(神经辐射场)结合,提出DreamFusion,实现了从文字生成3D模型。

目录

  • 摘要
  • 引入:扩散模型
  • 方法

摘要

在数十亿图像-文本对上训练的扩散模型,在文字生成图像的任务上大获成功。但是,如果想要将这种方法应用于 3D 生成(synthesis),需要对大规模的 3D 数据集进行标注并且在其上面训练,除此之外,还需对 3D 数据去噪的有效架构,但目前这两者都不存在。在这项工作中,作者通过使用预训练的 2D 文本-图像的扩散模型,实现文本到 3D 合成。他们引入了基于概率密度蒸馏的损失函数,这也允许了2D扩散模型作为先验,用以优化参数图像生成器。在类似 DeepDream 的过程中使用这种损失函数,作者通过梯度下降优化随机初始化的 3D 模型(NeRF),使其从随机角度的 2D 渲染均能让损失函数值较低。

在该方法中,给定文本生成的 3D 模型可以从任意角度观察,通过任意照明重新点亮,或合成到任何 3D 环境中。值得注意的是,该方法不需要 3D 训练数据,也不需要对图像扩散模型进行修改,证明了预训练图像扩散模型作为先验的有效性。

引入:扩散模型

扩散模型是隐变量生成模型,它学习将样本从易处理的噪声分布逐渐转换为数据分布的过程,他由一个正向过程

q

和一个反向过程或生成模型

p

组成。正向过程

q

通过添加噪声从数据

mathbf{x}_t

中缓慢移除结构,反向过程或生成模型

p

则从噪声

mathbf{z}_t

开始缓慢添加结构。

前向过程

前向过程通常是一个高斯分布,它从先前在时间步 t 的噪声较小的隐含表示,转换为在时间步 t 1 的噪声较大的隐含表示。给定初始数据x,我们可以通过积分得到隐变量在t时间的边缘分布:

qleft(mathbf{z}_t mid mathbf{x}right)=mathcal{N}left(alpha_t mathbf{x}, sigma_t^2 mathbf{I}right)

而如果利用这个结果,我们同时可以积分数据的密度以得到平滑后的数据分布。

qleft(mathbf{z}_tright)=int qleft(mathbf{z}_t mid mathbf{x}right) q(mathbf{x}) d mathbf{x}

为了使

qleft(mathbf{z}_tright)

在前向过程开始(

sigma_0 approx 0

)的时候接近于数据分布,并且在前向传播的最后(

sigma_T approx 1

)接近于一个高斯分布,需要选择合适的

alpha_t

sigma_t

系数。

反向过程

反向过程,也即生成模型

p

则需要对随机噪声

pleft(mathbf{z}_Tright)=mathcal{N}(mathbf{0}, mathbf{I})

一步步增加结构。状态转移函数通常可以参数化为

p_phileft(mathbf{z}_{t-1} mid mathbf{z}_tright)=qleft(mathbf{z}_{t-1} mid mathbf{z}_t, mathbf{x}=hat{mathbf{x}}_phileft(mathbf{z}_t ; tright)right)

,其中

qleft(mathbf{z}_{t-1} mid mathbf{z}_t, mathbf{x}right)

是从前向过程中推导出的后验分布,而

hat{mathbf{x}}_phileft(mathbf{z}_t ; tright)

则是一个学习出的对优化去噪器的近似。

但在DDPM中,作者并没有直接学习

hat{mathbf{x}}_phileft(mathbf{z}_t ; tright)

,而是训练了一个图片到图片的U-Net网络,

epsilon_phileft(mathbf{z}_t ; tright)

,用以预测隐变量中的噪声部分

mathbf{z}_t

mathbb{E}left[mathbf{x} mid mathbf{z}_tright] approx hat{mathbf{x}}_phileft(mathbf{z}_t ; tright)=left(mathbf{z}_t-sigma_t epsilon_phileft(mathbf{z}_t ; tright)right) / alpha_t

这个噪声可以与一个预测损失函数关联:

epsilon_phileft(mathbf{z}_t ; tright)=-sigma_t s_phileft(mathbf{z}_t ; tright)

而在此类似于VAE,生成模型的训练目标是变分下界(ELBO),把上述损失函数简化为:

mathcal{L}_{text {Diff }}(phi, mathbf{x})=mathbb{E}_{t sim mathcal{U}(0,1), epsilon sim mathcal{N}(mathbf{0}, mathbf{I})}left[w(t)left|epsilon_phileft(alpha_t mathbf{x} sigma_t epsilon ; tright)-epsilonright|_2^2right]

其中

w(t)

是一个时变的权重函数。因此,扩散模型的训练可以被视为学习一个隐变量模型,或者学习与数据的多个噪声版本相对应的一系列得分函数。文字生成图片的扩散模型学习

epsilon_phileft(mathbf{z}_t ; t, yright)

,这个网络的条件是文字的嵌入(embedding)。

方法

整体流程

DreamFusion的整体训练流程如下图所示:

对于一个描述文本“a DSLR photo of a peacock on a surfboard”,首先输入预训练好的、参数固定的文本特征提取器(图中Transformer),并以此为condition输入扩散模型(U-Net,同样参数固定)。每一步训练时,对Nerf渲染出的图像加入一个随机噪声进行扰动,输入U-Net后得到其估计出的噪声,并计算两者之差

hat{epsilon}_phileft(mathbf{z}_t mid y ; tright)-epsilon

,这个差值将会用于反向传播的梯度计算。

值得注意的是,Nerf模型渲染的图片将会选择一个随机的相机位置作为观察视角,并且使用从密度梯度计算的法线(normal)、以随机的光照方向对场景进行着色。这样可以保证在不同视角、光照情况下,渲染出的图像都符合文本描述的认知。

采样&损失函数

通常来说,扩散模型的采样方式是在观察到的数据上进行采样,也即采样像素。然而,在该任务中,作者们并不希望对像素进行采样;他们只希望学习出一个3D模型,使得该模型在任意视角下看上去像一张好的图片。这种模型通常可以被认为是一个可微图像参数化(differentiable image parameterization, DIP),也即用参数

theta

和生成器

g

来获取图像:

mathbf{x}=g(theta)

在Nerf的背景下,参数

theta

就是3D的Volume以及其他提供的参数,如观察视角、光照条件等等;而生成器

g

则是体积分渲染器。

如果我们按照DDPM等方案中的损失函数来面对该任务,那么对于任何一个生成的数据点

mathbf{x}=g(theta)

而言,我们的目标是学习到最优的3D Volume参数。需要注意的是,该任务中扩散模型的U-Net是固定的,所需要训练的仅仅是Nerf的3D Volume:

theta^*=arg min _theta mathcal{L}_{text {Diff }}(phi, mathbf{x}=g(theta))

如果我们把这个目标的损失函数写成之前的形式,并且求其梯度,我们会得到这样一个形式:

nabla_theta mathcal{L}_{text {Diff }}(phi, mathbf{x}=g(theta))=mathbb{E}_{t, epsilon}[w(t) underbrace{left(hat{epsilon}_phileft(mathbf{z}_t ; y, tright)-epsilonright)}_{text {Noise Residual }} underbrace{frac{partial hat{epsilon}_phileft(mathbf{z}_t ; y, tright)}{mathbf{z}_t}}_{text {U-Net Jacobian }} underbrace{frac{partial mathbf{x}}{partial theta}}_{text {Generator Jacobian }}]

我们可以发现,U-Net(也即扩散模型的主要模型)的雅各比矩阵计算起来的复杂度是巨大的。于是,作者决定将这一项省去,因此可以得到新的梯度:

nabla_theta mathcal{L}_{mathrm{SDS}}(phi, mathbf{x}=g(theta)) triangleq mathbb{E}_{t, epsilon}left[w(t)left(hat{epsilon}_phileft(mathbf{z}_t ; y, tright)-epsilonright) frac{partial mathbf{x}}{partial theta}right]

这个损失函数和一般扩散模型一样,基于一个随机噪声在时间戳

t

处对

mathbf{x}

的扰动,使得参数的更新方向靠近一个有更高概率密度的区域。通过证明,作者发现这个梯度其实是一个加权概率密度蒸馏损失函数:

nabla_theta mathcal{L}_{mathrm{SDS}}(phi, mathbf{x}=g(theta))=nabla_theta mathbb{E}_tleft[sigma_t / alpha_t w(t) mathrm{KL}left(qleft(mathbf{z}_t mid g(theta) ; y, tright) | p_phileft(mathbf{z}_t ; y, tright)right)right]

具体推导过程可参见论文。因此,作者命名他们使用的这个采样方案为SDS(Score Distillation Sampling)。这个方案不需要对扩散模型进行反向传播更新,而是只训练一个Nerf。下图中给出了该采样方案与扩散模型采样的对比:

这里,作者使用了简单的生成函数

mathbf{x}=(operatorname{flip}(theta), theta)

作演示(也即生成一个对称的图片)。可以看到,DDPM类的扩散模型,是对于像素空间内的像素进行采样;而在右边,他们的方案可以直接对

theta

使用SGD进行更新,在有生成函数的几何限制的情况下,获得可观的结果。

同时,作者也在附录中给出了这一损失函数反向传播的伪代码:

可以发现,该损失函数的反向传播实现也非常简练。

神经渲染

本文在渲染过程中使用的是和一般Nerf一致的设定,也即使用MLP输出3D位置的体密度以及从特定视角和光照下观察的RGB值。值得一提的是其着色(shading)部分。相比于传统的NeRF模型,他们的MLP仅仅得到表面处的RGB值,而这个RGB值随后将被一个可控的光照“点亮”。我们把这个表面上的RGB值表示为反照率(albedo)

rho

,并把体密度表示为

tau

,则他们的MLP得到的值分别为:

(tau, boldsymbol{rho})=operatorname{MLP}(boldsymbol{mu} ; theta)

根据体密度,我们也可以计算某一位置

boldsymbol{mu}

对应的法向量:

boldsymbol{n}=-nabla_{boldsymbol{mu}} tau /left|nabla_{boldsymbol{mu}} tauright|

随后,我们假设有一个点光源,其坐标、颜色和环境光颜色分别为

ell

ell_rho

ell_a

,我们能够得到光照下该点的颜色:

mathbf{c}=boldsymbol{rho} circleft(boldsymbol{ell}_rho circ max (0, boldsymbol{n} cdot(boldsymbol{ell}-boldsymbol{mu}) /|boldsymbol{ell}-boldsymbol{mu}|) boldsymbol{ell}_aright)
训练过程
随机采样相机和光照

训练的每一步,相机的参数都是在球面坐标上随机采样实现的,包含了3个参数:仰角

phi_{text {cam }} inleft[-10^{circ}, 90^{circ}right]

, 方位角

theta_{text {cam }} inleft[0^{circ}, 360^{circ}right]

, 与原点的距离

[1,1.5]

。除此之外他们还采样了一个up向量和一个聚焦点(look-at point),结合起来就可以获得相机的外参矩阵。对于内参矩阵,他们采样了一个乘数

lambda_{text {focal }} in mathcal{U}(0.7,1.35)

,并联合图像的像素宽度得到其焦距

lambda_{text {focal }} w

。实际操作中,他们渲染的图片大小为64x64。

对于光照而言,点光源位置是从一个以原点为中心附近的分布中采样的。

基于视角的文本嵌入

对于仰角较大的数据,文本中将会加入“overhead view”;对于小于60°的仰角,他们将会根据方位角加入文本“front view”,“side view”或是“back view”。

训练设置

硬件方面,作者使用了带有4块显卡的TPUv4的机器,进行每一个文本对应的3D场景进行训练。每一块卡渲染一个独立的视角,并且在每一张卡上的batchsize设置为1。

时间方面,训练过程需要迭代15000次,在上述硬件设置下共需要1.5小时。

效果展示

该方案可以产生许多质量高而且天马行空的模型,本文在这里只列举少数,更多结果可以在项目官网查看。

与其他方案的结果对比如下图所示:

可以看到,主观效果上,DreamFusion生成的模型看上去明显更加生动且接近描述。

0 人点赞