剪枝60％不損效能！上海AILab提出高斯剪枝新方法，入選CVPR2025

三維高斯潑濺（3D Gaussian Splatting）使得即時高質量渲染成為可能，是當前3D視覺領域最常用的演算法之一。

問題在於，基於顯示高斯單元的表示方式，儘管可以高效濺射和光柵化，其密集化和最佳化過程卻往往會生成冗餘的高斯點，導致單個重建場景可能包含數百萬個高斯點。

這不僅降低了訓練和渲染速度（本可能更快），還導致顯著的記憶體消耗。

現在，來自上海AI Lab的研究團隊提出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個高斯點的顏色和其後的顏色累計值。

梯度公式可以分為兩部分理解：