MaskGaussian: Adaptive 3D Gaussian Representation from Probabilistic Masks

Yifei Liu1,2, Zhihang Zhong1, Yifan Zhan1,3, Sheng Xu2, Xiao Sun1,
1 Shanghai Artificial Intelligence Laboratory, 2 Beihang University, 3 The University of Tokyo,
Paper Code

TL;DR: We introduce MaskGaussian, which probabilistically samples Gaussians and prunes them through gradients from both used and masked Gaussians. We embed masks into just 2 lines of CUDA forward code to make masked-out Gaussian's sampling probability trainable, preventing suboptimal pruning decisions from degenerated optimization state which loses scene details.

Ours
PSNR: 31.54
# GS (M): 0.51
FPS: 435.0
3DGS [Kerbl 2023]
PSNR: 31.54
# GS (M): 1.75
FPS: 195.1
Ours
PSNR: 27.34
# GS (M): 2.13
FPS: 227.2
Compact 3DGS [Lee 2024]
PSNR: 27.06
# GS (M): 2.46
FPS: 178.4

Left: Compared with 3DGS, MaskGaussian can reduce the number of Gaussian points without compromising the reconstruction quality. Right: MaskGaussian can use fewer Gaussians to better restore the "details" of the vines.



Abstract

While 3D Gaussian Splatting (3DGS) has demonstrated remarkable performance in novel view synthesis and real-time rendering, the high memory consumption due to the use of millions of Gaussians limits its practicality. To mitigate this issue, improvements have been made by pruning unnecessary Gaussians, either through a hand-crafted criterion or by using learned masks. However, these methods deterministically remove Gaussians based on a snapshot of the pruning moment, leading to sub-optimized reconstruction performance from a long-term perspective. To address this issue, we introduce MaskGaussian, which models Gaussians as probabilistic entities rather than permanently removing them, and utilize them according to their probability of existence. To achieve this, we propose a masked-rasterization technique that enables unused yet probabilistically existing Gaussians to receive gradients, allowing for dynamic assessment of their contribution to the evolving scene and adjustment of their probability of existence. Hence, the importance of Gaussians iteratively changes and the pruned Gaussians are selected diversely. Extensive experiments demonstrate the superiority of the proposed method in achieving better rendering quality with fewer Gaussians than previous pruning methods, pruning over 60% of Gaussians on average with only a 0.02 PSNR decline.



Method Overview

Overview of MaskGaussian: We illustrate our pipeline with five Gaussians, G1 through G5, where G2 and G5 are not sampled and masked. First, all Gaussians are splatted in the standard manner, and differentiable masks are sampled from their existence distributions. For each query pixel, a splat Gi has αi computed from normal attributes (center, scale, rotation). Splats with zero αi are filtered out, and the remaining splats and their masks are passed into the masked-rasterization. We apply the masks in two places: the transmittance evolution for Ti and the color rendering for ci. A masked splat Gi (e.g., i=2 in this figure) does not receive a gradient for αi, and thus does not update its normal attributes, but it receives a gradient for mask mi and updates its existence probability.



Results

Main results on Mip-Nerf360, Tanks & Temples, and Deep Blending.

Training Progress Comparison

Compared to Compact3DGS, MaskGaussian progressively prunes more Gaussians, leading to faster training and lower GPU memory requirements, while retainning higher reconstruction quality.


Peak GPU Memory Cost:

Also Supports Finetuning

Although not specifically designed for post-training, MaskGaussian can be used for finetuning an already existing Gaussian scene. You can use it as a drop-in replacement for LightGaussian — you can use it with exactly the same commands, no modifications needed. Try it out with the code here!


Can be seamlessly integrated into SOTA training pipelines

Our method can in theory prune Gaussians for any SOTA 3DGS-pipeline. For example, we add MaskGaussian on Taming-3DGS and show that it has better efficiency than directly controlling its targeting budget.