AAAI 2023| 一即一切：轻松实现不同 NeRF 系列结构间的转换！

论文链接：https://arxiv.org/abs/2211.15977

开源代码：https://github.com/megvii-research/AAAI2023-PVD

一、导语

神经辐射场 (NeRF[1])方法已被证明是一种有效的，高质量的 3D 场景表示，并支持诸如编辑，检索，导航等下游任务。目前多种结构都在争夺 NeRF 的核心表示，包括多层感知器 (MLP)、稀疏张量(sparse Tensors)、低秩张量(low-rank Tensors)、哈希表(Hashtables)及其之间的组合。而这些表示目前都有一定程度上的 trade-off，例如基于哈希表的表示通常训练和渲染很快，但其不同分辨率的哈希混叠表示缺乏较为清晰的几何结构，可能会阻碍诸如基于空间关系感知的编辑操作。

为了解决上述问题，我们提出了渐进式的体素蒸馏方案 PVD，该方法可以实现 MLP、稀疏或低秩张量、哈希表等不同结构之间的两两互转。因此 PVD 可以根据不同的下游任务将手头已有模型调整为新的结构表示。PVD 是在不同层次表示上逐步进行的，这种从浅到深的方式使得结构间转换速度很快。比如使用 PVD 将基于哈希表的 Instant-NGP 模型转化为基于 MLP 的 NeRF 模型要比从头开始训练该 MLP 快 10 倍以上，且能实现相近甚至更高水平的图像合成质量。

二、背景

新视图合成(NVS)指在 3D 场景中的未知位置和视角上生成逼真的 2D 图像，其具有广泛的应用场景。NeRF[1]的出现极大地加速了这一问题及其相关问题的研究进程。NeRF 使用 MLP 从坐标空间构造出 3D 场景，MLP 的输入是一个 5D 坐标（空间位置+观察方向），输出是对应空间点的 density 和 color。隐式建模的优点是其表示有利于控制或改变场景的类纹理属性比如 NeRFW[2]以及时序属性如 DNeRF[3]，即比较容易地进行 3D 场景艺术设计和动态场景设计。但是基于 MLP 的 NeRF 需要对空间点进行密集采样，导致在训练和推理过程中需要对 MLP 进行多次查询，从而导致其运行速度很慢。

基于显式及混合形式的表达，则在一定程度上解决了完全基于 MLP 的 NeRF 运行较慢的问题。如 Plenoxels[4] 通过完全的显式表示，将场景直接放置在 3D 网格上，网格上的每个体素都存储其 density 和 color 相关信息。由于 Plenoxels 不带有任何 MLP 结构，只需要通过网格体素上存储的信息进行插值来得到空间中任意点的信息。这种方式大大缩减了 Plenoxels 的训练和推理时间。基于纯显性结构的表示有助于基于空间位置信息的场景几何编辑，包括多个场景的合并、特定位置的对象修复和操作等如 CCNeRF[5]。Plenoxels 一个较大的缺陷是模型参数量过大，其最终效果较为依赖网格的分辨率。

除了上述纯隐式和纯显式的表示之外，一些混合的表示方法也开始出现。比如最近受低秩张量近似算法的启发，TensoRF[6]将显式张量分解为低秩分量，从而显着减小了模型大小。INGP[7] 提出了多分辨率哈希编码 (MHE)，通过不同尺度的哈希表级联将给定坐标映射到一组特征，而后使用一个小的 MLP 将该特征解译为 density 和 color。低秩张量和 MHE 显着减少了模型内存占用并提高了推理速度。然而低秩张量与 MHE 的紧凑性都是以不直观的几何解释为代价，其后期引入的小 MLP 会损失一定的几何解释，且多个低秩张量的混合以及哈希表的散列机制导致了大量的空间混叠。

哪种结构的表示最好，没有固定的答案。不同下游任务所需的结构类型不同。因此本文尝试不同结构间的互转，在引入互转方法之前，需要考虑这样做的好处究竟有哪些：

（1）该类型的研究可以深入了解已经丰富且不断增长的 NeRF 架构群的建模能力和局限性。目前对这些结构的研究都是独立，这样的方式并不能充分发挥每一种结构的优势和规避其劣势。

（2）这种转换的可能性将使设计人员从事先设计模型结构的负担中解放出来，因为现在他们可以简单地将训练有素的模型快速地转化为想要的结构，以满足后续不断变化的应用场景的需要。

（3）在 teacher 和 student 具有不同属性的情况下，二者可以实现有效互补，从而获得表达能力更强的模型。例如当使用哈希表示的 teacher 模型来获取张量表示的 student 模型时，可以受益于 teacher 更快的训练速度和更好的性能表现，来得到具有清晰几何结构且性能更优的 student 模型。

三、方法

如图所示，给定一个 teacher 和 student 后，我们将其统一为两部分的表示，即  和  。蒸馏过程分为三个 stage，stage1 只前传到  子网络。stage2 只前传到渲染前部分。stage3 则使用全部的信息进行训练。这样的方式有效缩短了反复查询 MLP 结构的巨大耗时，有效缩短转化时间。此外在过程中需要对 density 等信息进行值域约束。

我们的方法旨在实现神经辐射场不同架构之间的相互转换。由于此类架构的数量不断增加，我们无法尝试所有的转换,因此本文主要针对目前常见的几种结构间的互转，包括 NeRF 中的 MLP，Plenoxels 中的稀疏张量，以及两种混合表示：INGP 的哈希表和 TensoRF 的低秩张量（VM 分解）。我们首先以统一形式对不同结构的表示进行切分，然后基于该统一的视图来设计本文的蒸馏方案。

3.1 统一视图设计

在实验中发现如果只使用 ray 对应的 RGB 和 ray 上采样点对应的 density 和 color 等信息来进行蒸馏仍需要完全查询整个 MLP，是低效的。为此我们设计一个额外的中间特征来对齐不同结构之间的信息。我们观察到混合表示中的隐式和显式结构是自然分离的，分别对应于不同的学习目标，因此我们考虑将各类形式的模型均拆分成这种两部分表式形式。具体来说，给定一个模型  ，我们将它们表示为两个模块的级联，各个结构的划分如下表所示：

统一视角架构的两级划分方法

这里*可以是teacher也可以是student。对于混合表示,我们直接将显式部分视为  ，将隐式部分视为  。而对于纯隐式的MLP，我们根据网络的深度将其分为两部分。对于纯显式表示结构，我们通过让  为恒等映射将其表示为两部分。

3.2 Loss 设计

蒸馏最后一个阶段用到的 loss 如下所示：

其中的子项依次表示 volume-aligned loss, density loss, color loss, RGB loss和regularization loss。其中正则化项主要取决于 student 模型的形式，如对于 Plenoxels 和VM 分解，正则项为稀疏损失和 TV 损失。应该注意的是，我们只对 Plenoxels 进行density、color、RGB 和 regularization 损失。

3.3 density 值域约束

我们发现密度 density loss 项很难直接优化。我们将此问题归因于其特定的数值不稳定性。即density 可以反映空间中一点的透光率。当 density 大于或小于某个值时，其物理意义是基本一致的（即完全透明或完全不透明）。因此 density 的取值范围过大时，实际上只有一个区间的 density 值起着关键作用，而关键值之外的部分则会严重影响其优化。因此我们在实际计算density loss 时，只优化其值域上的一个区间段。

3.4 block-wsie 蒸馏

蒸馏分为三个 stage。stage1 只前传到  子网络，stage2 只前传到渲染前部分，stage3 则使用全部的信息进行训练。这样的方式有效缩短了反复查询 MLP 结构的巨大耗时，有效缩短转化时间。

四、实验

4.1 互转结果

Lego 场景上 Hash / VM分解 / MLP / 稀疏张量之间相互转换的定量和定性结果。我们对每个结构训练一个teacher 模型，然后用其蒸馏 4 种结构的 student 模型。图上数字表示 PSNR。“s-”表示 student.

NeRF-Synthetic 数据集上不同结构间的互转结果。

对于四种表示（Hash / VM-decomposition / MLP / sparse tensors），我们首先在 NeRF-Synthetic 数据集上的 8 个场景中从头开始训练每种结构的模型，总共获得了 32 个模型作为teacher。然后使用本文提出的 PVD 将这些 teacher 分别进行 4 种转换，获取 128 个 student 模型。转换完成后我们统计其平均指标(上表)和视觉效果(上图)。可以看出我们的方法对于不同结构间的转换是非常有效的。当一个模型转化为另一种形式时，student 性能与教师差异较小，这充分说明基于辐射场的不同结构表示间可以相互转换。进一步我们在 LLFF 和 TanksAndTemples 数据集上进行了类似的实验，并与其原始论文中的方法从头训练的结果进行对比，结果显示我们的方法仍具有可比较的性能表现。

LLFF 和 TanksAndTemples 数据集上从头开始训练一个模型与通过本文方法蒸馏得到的模型间的效果对比。

4.2 加速收敛和性能提升

左侧是 1.5 小时内通过 PVD 蒸馏出的 NeRF 模型结果（INGP 作为 teacher）。右侧是 NeRF 25 小时从头开始训练的结果。PVD 有效提高了合成质量并减少了训练时间。

从头训练 NeRF 与使用本文方法蒸馏得到 NeRF 的耗时和性能对比。

从上图和表可以看到，当使用一个性能较高的结构作为 teacher 时，我们的方法可以显著减少 student 模型训练时间并提高student模型性能。以往相同结构间的蒸馏方式(从大参数量的 NeRF 模型蒸馏到小参数量的 NeRF 模型)的过程异常低效，因为它需要在训练和蒸馏过程中不断查询大参数量 NeRF 模型。而我们异构形式之间的蒸馏可以高效查询 teacher，能够实现更快速的蒸馏。

4.3 消融实验

在 NeRF-Synthetic 数据上的消融实验(VM 分解蒸馏为 MLP).

有无 density 值域限制的蒸馏结果对比。

消融实验证明了我们方法中每个组件的必要性。

4.4 模型受限

Synthetic-NeRF 数据集上相互转换后的 PSNR 差距。PSNRstu 表示通过蒸馏得到的 student 的 PSNR。PSNRself 表示从头开始训练得到的 PSNR。PSNRtea 表示 teacher 的 PSNR。

可以看到 student 的性能主要受限于两方面：一是 teacher 的建模能力，二是 student 的建模能力。当 teacher 能力更强时往往有助于提升 student 的最终表现(如 INGP 作为 teacher，MLP 作为 student)；当二者能力相当时，二者的最终表现相近(如同种结构间的蒸馏)；而当student能力更强时，student的性能表现会受限于 teacher 的能力(如 MLP 为teacher，INGP 为 student)，此时可以进一步 finetune student 来提升其性能。

五、总结

在这项工作中，我们提出了一种系统的蒸馏方法 PVD，使得不同 NeRF 架构（包括 MLP、稀疏张量、低秩张量和哈希表）之间进行转换，同时保持较高的合成质量。PVD 成功的核心是统一结构视角的表示，损失函数的设计、渐进式蒸馏方案以及密度值的特殊处理等。通过打破不同结构间的隔阂，PVD 允许下游任务在已有模型上进行简单处理。实验证明了我们的方法在合成数据集和真实数据集上的有效性。

参考文献

[1] Mildenhall B, Srinivasan P P, Tancik M, et al. Nerf: Representing scenes as neural radiance fields for view synthesis[J]. Communications of the ACM, 2021, 65(1): 99-106.

[2] Martin-Brualla R, Radwan N, Sajjadi M S M, et al. Nerf in the wild: Neural radiance fields for unconstrained photo collections[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 7210-7219.

[3] Pumarola A, Corona E, Pons-Moll G, et al. D-nerf: Neural radiance fields for dynamic scenes[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 10318-10327.

[4] Fridovich-Keil S, Yu A, Tancik M, et al. Plenoxels: Radiance fields without neural networks[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 5501-5510.

[5] Tang J, Chen X, Wang J, et al. Compressible-composable nerf via rank-residual decomposition[J]. arXiv preprint arXiv:2205.14870, 2022.

[6] Chen A, Xu Z, Geiger A, et al. Tensorf: Tensorial radiance fields[C]//Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XXXII. Cham: Springer Nature Switzerland, 2022: 333-350.

[7] Müller T, Evans A, Schied C, et al. Instant neural graphics primitives with a multiresolution hash encoding[J]. ACM Transactions on Graphics (ToG), 2022, 41(4): 1-15.

论文链接：

https://arxiv.org/abs/2211.15977

技术信仰，价值务实

旷厂小伙伴们还将带来怎样的技术创新？

让我们一起期待吧！