来源: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 训练数据,也不需要对图像扩散模型进行修改,证明了预训练图像扩散模型作为先验的有效性。
引入:扩散模型
扩散模型是隐变量生成模型,它学习将样本从易处理的噪声分布逐渐转换为数据分布的过程,他由一个正向过程
和一个反向过程或生成模型
组成。正向过程
通过添加噪声从数据
中缓慢移除结构,反向过程或生成模型
则从噪声
开始缓慢添加结构。
前向过程
前向过程通常是一个高斯分布,它从先前在时间步 t 的噪声较小的隐含表示,转换为在时间步 t 1 的噪声较大的隐含表示。给定初始数据x,我们可以通过积分得到隐变量在t时间的边缘分布:
而如果利用这个结果,我们同时可以积分数据的密度以得到平滑后的数据分布。
为了使
在前向过程开始(
)的时候接近于数据分布,并且在前向传播的最后(
)接近于一个高斯分布,需要选择合适的
和
系数。
反向过程
反向过程,也即生成模型
则需要对随机噪声
一步步增加结构。状态转移函数通常可以参数化为
,其中
是从前向过程中推导出的后验分布,而
则是一个学习出的对优化去噪器的近似。
但在DDPM中,作者并没有直接学习
,而是训练了一个图片到图片的U-Net网络,
,用以预测隐变量中的噪声部分
:
这个噪声可以与一个预测损失函数关联:
而在此类似于VAE,生成模型的训练目标是变分下界(ELBO),把上述损失函数简化为:
其中
是一个时变的权重函数。因此,扩散模型的训练可以被视为学习一个隐变量模型,或者学习与数据的多个噪声版本相对应的一系列得分函数。文字生成图片的扩散模型学习
,这个网络的条件是文字的嵌入(embedding)。
方法
整体流程
DreamFusion的整体训练流程如下图所示:
对于一个描述文本“a DSLR photo of a peacock on a surfboard”,首先输入预训练好的、参数固定的文本特征提取器(图中Transformer),并以此为condition输入扩散模型(U-Net,同样参数固定)。每一步训练时,对Nerf渲染出的图像加入一个随机噪声进行扰动,输入U-Net后得到其估计出的噪声,并计算两者之差
,这个差值将会用于反向传播的梯度计算。
值得注意的是,Nerf模型渲染的图片将会选择一个随机的相机位置作为观察视角,并且使用从密度梯度计算的法线(normal)、以随机的光照方向对场景进行着色。这样可以保证在不同视角、光照情况下,渲染出的图像都符合文本描述的认知。
采样&损失函数
通常来说,扩散模型的采样方式是在观察到的数据上进行采样,也即采样像素。然而,在该任务中,作者们并不希望对像素进行采样;他们只希望学习出一个3D模型,使得该模型在任意视角下看上去像一张好的图片。这种模型通常可以被认为是一个可微图像参数化(differentiable image parameterization, DIP),也即用参数
和生成器
来获取图像:
。
在Nerf的背景下,参数
就是3D的Volume以及其他提供的参数,如观察视角、光照条件等等;而生成器
则是体积分渲染器。
如果我们按照DDPM等方案中的损失函数来面对该任务,那么对于任何一个生成的数据点
而言,我们的目标是学习到最优的3D Volume参数。需要注意的是,该任务中扩散模型的U-Net是固定的,所需要训练的仅仅是Nerf的3D Volume:
如果我们把这个目标的损失函数写成之前的形式,并且求其梯度,我们会得到这样一个形式:
我们可以发现,U-Net(也即扩散模型的主要模型)的雅各比矩阵计算起来的复杂度是巨大的。于是,作者决定将这一项省去,因此可以得到新的梯度:
这个损失函数和一般扩散模型一样,基于一个随机噪声在时间戳
处对
的扰动,使得参数的更新方向靠近一个有更高概率密度的区域。通过证明,作者发现这个梯度其实是一个加权概率密度蒸馏损失函数:
具体推导过程可参见论文。因此,作者命名他们使用的这个采样方案为SDS(Score Distillation Sampling)。这个方案不需要对扩散模型进行反向传播更新,而是只训练一个Nerf。下图中给出了该采样方案与扩散模型采样的对比:
这里,作者使用了简单的生成函数
作演示(也即生成一个对称的图片)。可以看到,DDPM类的扩散模型,是对于像素空间内的像素进行采样;而在右边,他们的方案可以直接对
使用SGD进行更新,在有生成函数的几何限制的情况下,获得可观的结果。
同时,作者也在附录中给出了这一损失函数反向传播的伪代码:
可以发现,该损失函数的反向传播实现也非常简练。
神经渲染
本文在渲染过程中使用的是和一般Nerf一致的设定,也即使用MLP输出3D位置的体密度以及从特定视角和光照下观察的RGB值。值得一提的是其着色(shading)部分。相比于传统的NeRF模型,他们的MLP仅仅得到表面处的RGB值,而这个RGB值随后将被一个可控的光照“点亮”。我们把这个表面上的RGB值表示为反照率(albedo)
,并把体密度表示为
,则他们的MLP得到的值分别为:
根据体密度,我们也可以计算某一位置
对应的法向量:
随后,我们假设有一个点光源,其坐标、颜色和环境光颜色分别为
,
和
,我们能够得到光照下该点的颜色:
训练过程
随机采样相机和光照
训练的每一步,相机的参数都是在球面坐标上随机采样实现的,包含了3个参数:仰角
, 方位角
, 与原点的距离
。除此之外他们还采样了一个up向量和一个聚焦点(look-at point),结合起来就可以获得相机的外参矩阵。对于内参矩阵,他们采样了一个乘数
,并联合图像的像素宽度得到其焦距
。实际操作中,他们渲染的图片大小为64x64。
对于光照而言,点光源位置是从一个以原点为中心附近的分布中采样的。
基于视角的文本嵌入
对于仰角较大的数据,文本中将会加入“overhead view”;对于小于60°的仰角,他们将会根据方位角加入文本“front view”,“side view”或是“back view”。
训练设置
硬件方面,作者使用了带有4块显卡的TPUv4的机器,进行每一个文本对应的3D场景进行训练。每一块卡渲染一个独立的视角,并且在每一张卡上的batchsize设置为1。
时间方面,训练过程需要迭代15000次,在上述硬件设置下共需要1.5小时。
效果展示
该方案可以产生许多质量高而且天马行空的模型,本文在这里只列举少数,更多结果可以在项目官网查看。
与其他方案的结果对比如下图所示:
可以看到,主观效果上,DreamFusion生成的模型看上去明显更加生动且接近描述。