MaskGaussian团队 投稿
量子位 | 公众号 QbitAI
三维高斯泼溅（3D Gaussian Splatting）使得实时高质量渲染成为可能，是当前3D视觉领域最常用的算法之一。
问题在于，基于显示高斯单元的表示方式，尽管可以高效溅射和光栅化，其密集化和优化过程却往往会生成冗余的高斯点，导致单个重建场景可能包含数百万个高斯点。
这不仅降低了训练和渲染速度（本可能更快），还导致显著的内存消耗。
现在，来自上海AI Lab的研究团队提出MaskGaussian，将掩码融合进光栅化过程，首次为被使用和未被使用的高斯同时保留梯度，在剪枝高斯的同时，MaskGaussian极大限度地保持了重建质量，提高了训练速度和减小内存需求。
△与3DGS相比，MaskGaussian在不影响重建质量的情况下减少高斯点数
该工作既支持从头开始的训练，也支持对现有的高斯进行微调。
△MaskGaussian使用更少的高斯更优秀的地还原出藤蔓的“细枝末节”
实验结果显示，MaskGaussian在Mip-NeRF360、Tanks & Temples和Deep Blending三个数据集上分别剪枝了62.4%、67.7%和75.3%的高斯点，且性能损失可以忽略不计。
无损修剪超过60%高斯点
对冗余高斯点进行剪枝，目前主要有两种方法：
-
第一类方法基于手工设计的重要性评分，移除评分低于预设阈值的高斯点。这类方法通常需要扫描所有训练图像以计算重要性评分，因此剪枝只能在训练期间执行一次或两次。
-
第二类方法使用可学习的掩码，将其与高斯点的属性相乘以接受梯度。尽管这种方法允许通过掩码逐渐移除高斯点，但渲染的场景始终依赖于相同的高斯点子集: 如果一个高斯点未被移除，它会一直存在到当前迭代； 一旦被移除, 它将被永久排除。这种确定性的掩码生成未能考虑剪枝后场景的演化，可能导致一些当前贡献较小但在训练后期可能关键且难以恢复的高斯点被移除。这导致次优的重建效果，尤其是细节或小物体的丢失。
研究团队认为，当前存在的掩码剪枝策略表现受限的主要原因在于：掩码与高斯的透明度等属性相乘后，CUDA渲染器无法产生来自未使用高斯的梯度。
作为类比，在传统的Transformer的掩码剪枝中，未被使用的token会以零值参与计算并获取梯度，使token不论是否被使用都能获取更新，避免优化过程朝着单一状态发展；而在3DGS中，CUDA渲染器会直接跳过溅射到2D上透明度为0的高斯，不进行梯度计算，使得未被用的高斯无法得到更新。
针对这些问题，MaskGaussian旨在为每个高斯点学习一个掩码分布。通过从该分布中采样，可以生成一个二值掩码，指示该高斯点的存在或不存在。然后，所有高斯点在不受掩码干扰的情况下进行标准溅射，并和掩码一起进入渲染计算。
掩码为1和0的高斯都可以通过更改后的CUDA渲染器获取正确梯度，计算其存在/不存在两种情况下对当前场景的贡献影响。
如下图所示，期望的颜色更接近掩码为0的高斯点
，而不是其后面的累积颜色，获取的梯度便会使增加其存在概率，以便在后续迭代中被使用。注意该过程中并无人工设计，完全由梯度信息指引。
掩码光栅化：前向过程
CUDA渲染器中的原始渲染公式如下：
其中
，
，
，
分别是像素点的颜色，第i个高斯的颜色，透明度和透光率。
为了剪枝不重要的高斯点，MaskGaussian添加了掩码，该掩码可以与其他高斯点属性一起进行优化，以评估高斯点的贡献。
掩码的生成视为一个两类采样过程。具体来说，MaskGaussian为每个高斯点分配两个可学习的掩码分数，并应用Gumbel-Softmax从两个分数中采样一个可微分的类别，记为
。通过避免将掩码直接应用于高斯点属性，MaskGaussian保留了溅射的α的完整性，使高斯的计算不会被跳过，能够完全参与光栅化。
然后，论文在光栅化框架中直接集成掩码，修改两行CUDA代码，如下公式所示：
掩码被应用于颜色累积和透射率衰减过程。当
时，高斯点正常对颜色贡献并根据其消耗透射率；当
时，高斯点的颜色贡献被掩码，其透射率消耗被跳过。
这种公式确保了在处理被掩码高斯点缺失时，前向光栅化结果的正确性。需要注意的是，被掩码的高斯点仍然参与前向计算，并能够接收有意义的梯度。
掩码光栅化：反向过程
为了说明梯度公式，MaskGaussian定义
为从第i+1个高斯点到最后一个高斯点渲染的颜色，即:
其中
，
，
分别是第j个高斯点的颜色、散射密度和透射率。接着，掩码的梯度表示如下，证明详见论文附录：
其中
是总损失函数，
是指示高斯点存在或不存在的二值掩码，是像素x的最终输出颜色，和
分别为第i个高斯点的颜色和其后的颜色累计值。
梯度公式可以分为两部分理解：
-
对颜色的权重：
决定了该高斯点对最终颜色的影响程度，也间接决定了该点对梯度的贡献大小。较大的权重意味着该点对渲染结果的重要性更高。 -
颜色优化方向：
中，
表示损失函数对输出颜色的优化方向，而
表示使用第i个高斯点的颜色相较于背景颜色
的优势。例如，如果两者的点积为正，则说明使用该高斯点是有益的，掩码会接收到正梯度，从而增加该高斯点的存在概率，即便它当前未被采样和使用。
有趣的是，这个梯度公式已经包含了，即基于分数的剪枝方法所使用的重要性标准。
此外，这一公式还捕捉了所需颜色与被掩码高斯点颜色之间的关系，这是分数剪枝方法无法测量且容易忽略的内容。
与通过高斯点的不透明度和尺度与掩码相乘的方法相比，本文的方法并未将掩码梯度与这些属性绑定，从而避免了对小型高斯点的不利影响。此外，该方法还允许被掩码的高斯点接收梯度，以更新其掩码分布。
这一创新使得被掩码的高斯点，即使未被直接采样，在优化过程中依然能够对场景做出潜在贡献的调整。这解决了传统剪枝方法中未被采样点逐渐失效的“死亡螺旋”问题，从而有效提升了剪枝效果和渲染质量。
训练与剪枝
MaskGaussian使用平方损失约束高斯点的平均数量，其定义如下，并在实验中验证其优于
损失：
其中
是高斯点的总数，是第i个高斯点的掩码值（取值为 0 或 1）。最终的损失函数为：
其中，
是渲染的图像损失（通常是感知损失或像素级损失），
是平衡超参数，用于控制掩码约束的强度。
为了剪枝掉采样概率接近零的高斯点，论文对每个高斯点进行10次采样，并移除那些从未被采样的点。该剪枝过程在每次密集化步骤后以及每1000次迭代中执行一次。
通过广泛的实验结果和分析，研究团队证明了MaskGaussian的有效性。
在Mip-NeRF360、Tanks & Temples和Deep Blending三个数据集上，MaskGaussian分别剪枝了62.4%、67.7%和75.3%的高斯点，且性能损失可以忽略不计。
项目地址：https://maskgaussian.github.io/
代码链接：https://github.com/kaikai23/maskgaussian
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法！
— 完 —
学术投稿请于工作日发邮件到：
ai@qbitai.com
标题注明【投稿】，告诉我们：
你是谁，从哪来，投稿内容
附上论文/项目主页链接，以及联系方式哦
我们会（尽量）及时回复你
🌟 点亮星标 🌟
（文：量子位）