Wonder3D 论文学习

论文链接：https://arxiv.org/abs/2310.15008
代码链接：https://github.com/xxlong0/Wonder3D

解决了什么问题？

随着扩散模型的提出，3D 生成领域取得了长足进步。从单张图片重建出 3D 几何是计算机图形学和 3D 视觉的基础任务，在 VR、游戏、3D 内容生成、机器人领域有广泛的应用。由于该任务是不适定的，需要推理出可见和不可见区域的 3D 几何结构。该种能力的构建需要大量的 3D 世界的知识。

部分工作（如 DreamField, DreamFusion, Magic3D）需要通过 score distillation sampling 来蒸馏出 2D 图像扩散模型或视觉语言模型的知识，从文本或图像创建出 3D 模型。尽管效果不错，这些方法都面临着两个问题：效率和一致性。每个形状的优化通常需要几万次的迭代，这涉及到全图渲染和扩散模型的推理。因此，优化一个形状甚至会耗费数小时。此外，2D 先验模型在每次迭代时只考虑一个视角，使得每个视角都接近于输入图像，这会造成生成的 3D 形状不一致。

另一部分的工作则尝试通过网络推理来生成 3D 的几何结构如点云、meshes、神经场，避免耗时的优化。此类方法大多在 3D assets 上训练 3D 生成扩散模型。但是，由于现有的 3D 数据集规模有限，这些方法都缺乏足够的泛化性，只能生成特定类别的形状。

另一些方法如 SyncDreamer 和 MVDream 直接生成多视角 2D 图像。这些工作增强了图像生成的多视角一致性，能从生成的多视角 2D 图像中恢复出 3D 形状。本文方法也采用了多视角生成的机制，因为 2D 表征效率高、灵活性强。但是只依赖于彩色图像，生成形状的保真度就差一些，它们要么很难恢复几何细节，要么计算量惊人。

基于 score distillation sampling 的方法能够从 2D 的扩散先验信息中恢复出 3D 结构，但通常都非常耗时，生成的几何形状不一致。另一些工作则通过快速的网络推理来产生 3D 信息，但质量都比较差，缺乏几何细节信息。

提出了什么方法？

为了解决保真度、一致性、泛化性和效率的问题，本文提出了一个单视角 3D 重建方法，用一个跨域扩散模型来生成多视角、一致性的法线图和对应的彩色图像。核心思想是扩展 stable diffusion，对两个不同域（法线和色彩）的协同分布来建模。作者证明，可以通过一个 domain switcher 和跨域注意力机制来实现。Domain switcher 允许扩散模型生成法线图或彩色图，而跨域注意力机制则在这两个域之间传递信息，提升效率和质量。最后为了从生成视角中提取表面结构，作者提出了一个更加鲁棒的几何法线融合算法，能够重建出干净、高质量（如上图所示）的几何结构。

作者在 Google Scanned Object 数据集和各种不同风格的 2D 图像上做了实验。实验表明 Wonder3D 能够高效地恢复出高质量的几何结构。

1. Diffusion Models

扩散模型通过一个特殊的退化过程，逐步地恢复图像，它采用了一个前向马尔可夫链和反向马尔可夫链。给定一个来自于数据分布 $p (z)$ 的样本 $z_0$ ，去噪扩散模型的前向计算会输出一个序列的加了噪声的数据 $\{z_t | t\in(0,T)\}$ ，其中 $z_t = \alpha_t z_0 + \sigma_t \epsilon$ ， $\epsilon$ 是从分布 $\mathcal{N}(0,1)$ 中随机选择的噪声， $\alpha_t, \sigma_t$ 是 noise schedule 的固定序列。逐步地对目标图像进行前向计算，直到图像变成完全的高斯噪声。然后使用逆马尔可夫链，逐步地去除图像上的噪声，通过预测添加的随机噪声 $\epsilon$ ，从 $z_{t}$ 恢复出 $z_{t-1}$ 。

2. The Distribution of 3D assets

以前的工作使用了 3D 表征如点云和神经辐射场，作者认为，3D assets 的分布（记做 $p_{a}(\mathbf{z})$ ）可以建模成对应的 2D 多视角法线图和彩色图像之间的协同分布。特别地，给定一组相机 $\{\pi_1, \pi_2, ..., \pi_K\}$ 以及一个条件输入图像 $y$ ，

$p_a(\mathbf{z})=p_{nc}(n^{1:K}, x^{1:K}|y)$

其中 $p_{nc}$ 是图像 $y$ 的法线图 $n^{1:K}$ 和彩色图 $x^{1:K}$ 的分布。在后续讨论中，作者省去了记号 $y$ 。因此，我们的目标是学习一个模型 $f$ ，合成不同相机姿态的多个法线图和彩色图，记做：

$(n^{1:K}, x^{1:K}) = f(y, \pi_{1:K})$

通过 2D 表征，我们的方法用 2D 扩散模型（基于几十亿张图像训练得到）构建，它具有很强的 zero-shot 泛化能力。另一方面，法线图描述了物体表面的起伏波动，编码了丰富的细节几何信息。我们可以从 2D 法线图提取高质量的 3D 几何信息。最后，我们可以将扩散机制的这个跨域协同分布表示成一个马尔可夫链：

$p(n_T^{1:K}, x_T^{1:K}) \prod_t{p_{\theta}(n_{t-1}^{1:K}, x_{t-1}^{1:K} | n_{t}^{1:K}, x_{t}^{1:K})}$

其中 $p(n_T^{1:K}, x_T^{1:K})$ 是高斯噪声。核心问题是描述出分布 $p_\theta$ ，这样我们就能从这个马尔可夫链中采样，生成出法线图和彩色图像。

3. 方法

本文提出了一个多视角跨域的扩散机制，在两个不同的域上运行，生成多视角一致的法线图和彩色图。整体架构如下图所示。首先，采用一个多视角扩散模型，生成多视角法线图和彩色图，使用多视角注意力要求不同的视角具有一致性。其次，domain switcher 允许扩散模型在多个域中运行，无需重新训练现有的 Stable Diffusion。因此，我们可以利用基础模型的泛化能力，它们一般是在大规模数据上训练得到。作者提出了一个跨域注意力，在法线域和彩色图像域之间传递信息，确保几何和视觉信息的一致性。最终，该几何法线融合方法能从 2D 法线图和彩色图像中恢复高质量的几何信息和外观信息。

3.1 Consistent Multi-view Generation

2D 扩散模型单独地生成每个图像，生成的图像在多个视角的视觉和几何上并非一致的。为了增强不同视角的一致性，像 SyncDreamer 和 MVDream 一样，作者利用了注意力机制来促进不同视角的信息传递，隐式地编码多视角图像之间的关系，如下图所示。

将原来的自注意力层扩展至全局范围，允许注意力层连接到其它视角的图像。不同视角的 keys 和 values 彼此连接，促进信息的传递。通过在注意力层共享不同视角的信息，扩散模型能够感知到多视角的关系，能够生成一致的多视角彩色图像和法线图。

3.2 Cross-Domain Diffusion

本文模型基于一个预训练的 2D stable diffusion 模型构建。但是，当前的 2D 扩散模型都是针对单个域设计的，那么挑战就是如何有效地将 stable diffusion 扩展至多个域。

Naive Solutions

为了实现此目的，作者探索了多个可能的设计。直接方案就是在 UNet 模块的输出增加四个通道，表示其它的域。因此，扩散模型能同时输出法线图和彩色图像。但是，我们注意到该设计收敛速度慢，泛化性差。这是因为通道扩张可能影响 stable diffusion 的预训练权重，造成灾难性的模型遗忘。

回到第一个等式，可以将协同分布拆分为两个条件分布：

$q_a(\mathbf{z}) = q_n(n^{1:K})\cdot q_c(x^{1:K} | n^{1:K})$

我们可以一开始训练一个扩散模型来生成法线图，然后基于生成的法线图，训练另一个扩散模型来生成彩色图像。但是，这个双阶段框架会增加计算成本，也会造成表现退化的问题。

Domain switcher

为了克服这些问题，作者通过 domain switcher 设计了一个跨域的扩散机制，记做 $s$ 。Switcher $s$ 是一个一维的向量，标记不同的域，将该 switcher 输进扩散模型作为额外的输入使用。因此，第二个等式就可以扩展成：

$n^{1:K}, x^{1:K} = f(y, \pi_{1:K}, s_n), f(y, \pi_{1:K}, s_c)$

Switcher $s$ 首先用 positional encoding 做编码，再和 time embedding 拼接到一起。然后注入 stable diffusion 的 UNet。实验表明，这个微小的变动不会改变预训练先验，而收敛速度会更快、鲁棒性和泛化性会更强。

Cross-domain Attention

有了 switcher 后，扩散模型能生成两个不同的域。但是要注意到，对于单个视角而言，彩色图像和法线图并不是几何一致的。为了解决这个问题，需要确保法线图和彩色图的一致性，于是作者加入了一个跨域注意力机制来促进信息在两个域之间传递。该机制的目的是确保生成的图像能在几何和外观上对齐。

该跨域注意力层与原版的自注意力层结构相同，位于 UNet 每个 transformer 模块的 cross-attention 层之前，如上图所示。在跨域注意力层，将法线域和彩色图像域的 keys 和 values 结合起来，经过注意力操作计算。该设计确保了生成的彩色图像和法线图紧密是关联的，增强二者的几何一致性。

3.3 Textured Mesh Extraction

为了从 2D 法线图和彩色图中直接提取 3D 几何结构，作者优化了一个 neural implicit signed distance field(SDF)，融合生成的 2D 数据。不像其它的表征（如 meshes），SDF 更加简洁，判别度更强。

SDF 是一个函数，它为空间中的每个点返回一个值，表示该点到最近表面的距离。如果点在表面内部，返回负值；如果点在表面外部，返回正值。SDF函数通常定义为： $f(x)=\text{距离}-\epsilon$ 。其中 $\epsilon$ 是一个小的正数，用于确保函数在表面上的值为零。利用神经网络来学习 SDF 函数的参数。这意味着神经网络通过训练，学习如何从输入数据（如图像或点云）中推断出距离场。

但是，无法使用现有的基于 SDF 的重建方法（如 NeuS）。这些方法都是为真实拍摄的图像设计的，需要密集的输入视角。然而，本文方法生成的视角相对稀疏，生成的法线图和彩色图像可能会有一些错误预测的像素。在几何优化时，这些错误可能会累加，造成几何结构的畸变或不完整。于是作者提出了一个几何优化机制。

Optimization Objectives

有了法线图 $G_{0:N}$ 和彩色图 $H_{0:N}$ 后，首先利用分割模型从法线图和彩色图中分割出物体的 masks $M_{0:N}$ 。随机选取一个 batch 的像素点和它们在世界空间 $P=\{g_k, h_k, m_k, \mathbf{v}_k\}$ 对应的射线，其中 $g_k$ 是第 $k$ 个采样像素的法线值， $h_k$ 是第 $k$ 个像素的颜色值， $m_k\in \{0,1\}$ 是第 $k$ 个像素的 mask 值， $\mathbf{v}_k$ 是第 $k$ 个射线的方向。

整体的优化函数如下，

$\mathcal{L}=\mathcal{L}_{normal}+\mathcal{L}_{rgb}+\mathcal{L}_{mask}+\mathcal{R}_{eik}+\mathcal{L}_{sparse}+\mathcal{L}_{smooth}$

其中 $\mathcal{L}_{normal}$ 表示法线损失， $\mathcal{L}_{rgb}$ 是 MSE 损失，计算渲染颜色 $\hat{h}_k$ 和生成颜色 $h_k$ 的损失。 $\mathcal{L}_{mask}$ 是一个二元交叉熵损失，计算渲染 mask $\hat{m}_k$ 和生成 mask $m_k$ 的损失。 $\mathcal{R}_{eik}$ 表示 eikonal 正则项，促使 SDF 的梯度的级数为单位长度。 $\mathcal{L}_{sparse}$ 是稀疏正则项。 $\mathcal{L}_{smooth}$ 是一个 3D 平滑正则项，迫使 SDF 梯度在 3D 空间平滑。

Eikonal 正则项是一种在数学和物理中用于描述波动现象的数学模型中常见的项。它与波动方程中的相位速度有关，尤其是在描述光波或其他类型的波的传播时。Eikonal 方程是一个非线性偏微分方程，通常用于描述波动现象中的相位波前。它的形式可以简化为 $|\nabla \phi|^2=n^2$ ，其中 $\phi$ 是相位函数， $\nabla$ 是梯度算子， $n$ 是介质的折射率。在波动方程中，正则项是指那些描述波动传播的项。Eikonal 正则项特别关注波前的形状和传播速度。

Geometry-aware Normal Loss

由于 SDF 表征的本质不同，我们可以通过计算 SDF 的二阶导数获得优化后的 SDF 的法线值 $\hat{g}$ 。最大化 SDF $\hat{g}$ 的法线和生成法线 $g$ 的相似度，提供 3D 几何监督信号。为了克服不同视角下法线的细微错误，作者提出了一个几何法线损失：

$\mathcal{L}_{normal}=\frac{1}{\sum{w_k}}\sum{w_k \cdot e_k}, \quad e_k=(1-\cos{(\hat{g_k}, g_k)})$

其中 $e_k$ 是第 $k$ 个射线的 SDF 的法线 $\hat{g}_k$ 和生成法线 $g_k$ 之间的损失， $\cos(\cdot, \cdot)$ 表示余弦函数， $w_k$ 是几何权重：

$w_k=\left\{ \begin{aligned} 0 &, & \cos(\mathbf{v}_k, \mathbf{g}_k)>\epsilon, \\ \exp{(|\cos(\mathbf{v}_k, \mathbf{g}_k)|)} &, & \cos(\mathbf{v}_k, \mathbf{g}_k)\leq \epsilon \end{aligned} \right.$
其中 $\exp$ 表示指数函数， $|\cdot|$ 表示绝对值函数， $\epsilon$ 是接近于 $0$ 的负数阈值，我们计算生成法线 $g_k$ 和第 $k$ 个射线的观测方向 $\mathbf{v}_k$ 的角度余弦值。

在计算机视觉中，射线的观察方向可以用于估计场景的深度和结构。例如，通过分析从不同视点发出的射线的交点，可以重建三维场景。

该方法的设计逻辑就是法线的朝向角，设定为面向外的，而观测方向是面向里的。这就确保了法线向量和观测射线的角度不低于 $90^\circ$ 。如果偏离了这个评价标准，说明生成的法线不够准确。

而且，值得注意的是，从多个不同观测点都能看到优化后形状的 3D 点。因此，它会受到对应这些视角的多条法线影响。但是，如果多条法线无法展现完美的一致性，几何监督信号就会变得模楞两可，产生不准确的几何结构。为了解决这个问题，作者引入了一个加权机制，而非同等地看待这些法线。如果一条法线与观测射线有更大的夹角，则权重就更大。这增强了几何监督过程的准确率。

法线是一个向量，它在几何体的表面上的某一点与该表面垂直。在二维空间中，法线是一个垂直于曲线的直线；在三维空间中，法线是一个垂直于表面的平面。法线通常用向量表示。如果一个曲面由函数 $z = f (x, y)$ 定义，则该曲面上点 $(x, y, f (x, y))$ 的法线可以通过曲面的梯度向量 $\nabla{f}$ 来确定。

Outlier-dropping Losses

除了法线损失，mask 损失和颜色损失也用于优化几何结构和外观特征。但是，在彩色图像和 masks 中不可避免地会出现一些错误的点，在优化的过程中会逐步累加，产生异常的表面和空洞。于是作者提出了一个简单而有效的策略，叫做 outlier-dropping loss。以颜色损失计算为例，除了简单地将每一步的所有采样射线的颜色求和，作者首先以倒序的顺序将这些损失排序，舍弃损失最大的那部分。这是因为错误预测和其它视角之间缺乏足够的一致性。加入了该策略后，优化后的几何结构就剔除了错误的几何结构和畸变区域。