剪枝60%不損性能！上海AI Lab提出高斯剪枝新方法，入選CVPR 2025
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保留了濺射的α的完整性，使高斯的計算不會被跳過，能夠完全參與光柵化。然後，論文在光柵化框架中直接集成掩碼，修改兩行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
