三維高斯?jié)姙R（3D Gaussian Splatting）使得實時高質(zhì)量渲染成為可能，是當(dāng)前3D視覺領(lǐng)域最常用的算法之一。

問題在于，基于顯示高斯單元的表示方式，盡管可以高效濺射和光柵化，其密集化和優(yōu)化過程卻往往會生成冗余的高斯點，導(dǎo)致單個重建場景可能包含數(shù)百萬個高斯點。

這不僅降低了訓(xùn)練和渲染速度（本可能更快），還導(dǎo)致顯著的內(nèi)存消耗。

現(xiàn)在，來自上海AI Lab的研究團隊提出MaskGaussian，將掩碼融合進光柵化過程，首次為被使用和未被使用的高斯同時保留梯度，在剪枝高斯的同時，MaskGaussian極大限度地保持了重建質(zhì)量，提高了訓(xùn)練速度和減小內(nèi)存需求。

△與3DGS相比，MaskGaussian在不影響重建質(zhì)量的情況下減少高斯點數(shù)

該工作既支持從頭開始的訓(xùn)練，也支持對現(xiàn)有的高斯進行微調(diào)。

△MaskGaussian使用更少的高斯更優(yōu)秀的地還原出藤蔓的“細枝末節(jié)”

實驗結(jié)果顯示，MaskGaussian在Mip-NeRF360、Tanks & Temples和Deep Blending三個數(shù)據(jù)集上分別剪枝了62.4%、67.7%和75.3%的高斯點，且性能損失可以忽略不計。

無損修剪超過60%高斯點

對冗余高斯點進行剪枝，目前主要有兩種方法：

• 第一類方法基于手工設(shè)計的重要性評分，移除評分低于預(yù)設(shè)閾值的高斯點。這類方法通常需要掃描所有訓(xùn)練圖像以計算重要性評分，因此剪枝只能在訓(xùn)練期間執(zhí)行一次或兩次。
• 第二類方法使用可學(xué)習(xí)的掩碼，將其與高斯點的屬性相乘以接受梯度。盡管這種方法允許通過掩碼逐漸移除高斯點，但渲染的場景始終依賴于相同的高斯點子集: 如果一個高斯點未被移除，它會一直存在到當(dāng)前迭代； 一旦被移除, 它將被永久排除。這種確定性的掩碼生成未能考慮剪枝后場景的演化，可能導(dǎo)致一些當(dāng)前貢獻較小但在訓(xùn)練后期可能關(guān)鍵且難以恢復(fù)的高斯點被移除。這導(dǎo)致次優(yōu)的重建效果，尤其是細節(jié)或小物體的丟失。

研究團隊認為，當(dāng)前存在的掩碼剪枝策略表現(xiàn)受限的主要原因在于：掩碼與高斯的透明度等屬性相乘后，CUDA渲染器無法產(chǎn)生來自未使用高斯的梯度。

作為類比，在傳統(tǒng)的Transformer的掩碼剪枝中，未被使用的token會以零值參與計算并獲取梯度，使token不論是否被使用都能獲取更新，避免優(yōu)化過程朝著單一狀態(tài)發(fā)展；而在3DGS中，CUDA渲染器會直接跳過濺射到2D上透明度為0的高斯，不進行梯度計算，使得未被用的高斯無法得到更新。

針對這些問題，MaskGaussian旨在為每個高斯點學(xué)習(xí)一個掩碼分布。通過從該分布中采樣，可以生成一個二值掩碼，指示該高斯點的存在或不存在。然后，所有高斯點在不受掩碼干擾的情況下進行標準濺射，并和掩碼一起進入渲染計算。

掩碼為1和0的高斯都可以通過更改后的CUDA渲染器獲取正確梯度，計算其存在/不存在兩種情況下對當(dāng)前場景的貢獻影響。

如下圖所示，期望的顏色更接近掩碼為0的高斯點，而不是其后面的累積顏色，獲取的梯度便會使增加其存在概率，以便在后續(xù)迭代中被使用。注意該過程中并無人工設(shè)計，完全由梯度信息指引。

掩碼光柵化：前向過程

CUDA渲染器中的原始渲染公式如下：

其中，，，分別是像素點的顏色，第i個高斯的顏色，透明度和透光率。

為了剪枝不重要的高斯點，MaskGaussian添加了掩碼，該掩碼可以與其他高斯點屬性一起進行優(yōu)化，以評估高斯點的貢獻。

掩碼的生成視為一個兩類采樣過程。具體來說，MaskGaussian為每個高斯點分配兩個可學(xué)習(xí)的掩碼分數(shù)，并應(yīng)用Gumbel-Softmax從兩個分數(shù)中采樣一個可微分的類別，記為。通過避免將掩碼直接應(yīng)用于高斯點屬性，MaskGaussian保留了濺射的α的完整性，使高斯的計算不會被跳過，能夠完全參與光柵化。

然后，論文在光柵化框架中直接集成掩碼，修改兩行CUDA代碼，如下公式所示：

掩碼被應(yīng)用于顏色累積和透射率衰減過程。當(dāng)時，高斯點正常對顏色貢獻并根據(jù)其消耗透射率；當(dāng)時，高斯點的顏色貢獻被掩碼，其透射率消耗被跳過。

這種公式確保了在處理被掩碼高斯點缺失時，前向光柵化結(jié)果的正確性。需要注意的是，被掩碼的高斯點仍然參與前向計算，并能夠接收有意義的梯度。

掩碼光柵化：反向過程

為了說明梯度公式，MaskGaussian定義為從第i+1個高斯點到最后一個高斯點渲染的顏色，即:

其中，，分別是第j個高斯點的顏色、散射密度和透射率。接著，掩碼的梯度表示如下，證明詳見論文附錄：

其中是總損失函數(shù)，是指示高斯點存在或不存在的二值掩碼，是像素x的最終輸出顏色，和分別為第i個高斯點的顏色和其后的顏色累計值。

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

• 對顏色的權(quán)重：決定了該高斯點對最終顏色的影響程度，也間接決定了該點對梯度的貢獻大小。較大的權(quán)重意味著該點對渲染結(jié)果的重要性更高。
• 顏色優(yōu)化方向：中，表示損失函數(shù)對輸出顏色的優(yōu)化方向，而表示使用第i個高斯點的顏色相較于背景顏色的優(yōu)勢。例如，如果兩者的點積為正，則說明使用該高斯點是有益的，掩碼會接收到正梯度，從而增加該高斯點的存在概率，即便它當(dāng)前未被采樣和使用。

有趣的是，這個梯度公式已經(jīng)包含了，即基于分數(shù)的剪枝方法所使用的重要性標準。

此外，這一公式還捕捉了所需顏色與被掩碼高斯點顏色之間的關(guān)系，這是分數(shù)剪枝方法無法測量且容易忽略的內(nèi)容。

與通過高斯點的不透明度和尺度與掩碼相乘的方法相比，本文的方法并未將掩碼梯度與這些屬性綁定，從而避免了對小型高斯點的不利影響。此外，該方法還允許被掩碼的高斯點接收梯度，以更新其掩碼分布。

這一創(chuàng)新使得被掩碼的高斯點，即使未被直接采樣，在優(yōu)化過程中依然能夠?qū)鼍白龀鰸撛谪暙I的調(diào)整。這解決了傳統(tǒng)剪枝方法中未被采樣點逐漸失效的“死亡螺旋”問題，從而有效提升了剪枝效果和渲染質(zhì)量。

訓(xùn)練與剪枝

MaskGaussian使用平方損失約束高斯點的平均數(shù)量，其定義如下，并在實驗中驗證其優(yōu)于損失：

其中是高斯點的總數(shù)，是第i個高斯點的掩碼值（取值為 0 或 1）。最終的損失函數(shù)為：

其中，   是渲染的圖像損失（通常是感知損失或像素級損失），是平衡超參數(shù)，用于控制掩碼約束的強度。

為了剪枝掉采樣概率接近零的高斯點，論文對每個高斯點進行10次采樣，并移除那些從未被采樣的點。該剪枝過程在每次密集化步驟后以及每1000次迭代中執(zhí)行一次。

通過廣泛的實驗結(jié)果和分析，研究團隊證明了MaskGaussian的有效性。

在Mip-NeRF360、Tanks & Temples和Deep Blending三個數(shù)據(jù)集上，MaskGaussian分別剪枝了62.4%、67.7%和75.3%的高斯點，且性能損失可以忽略不計。

項目地址：https://maskgaussian.github.io/
代碼鏈接：https://github.com/kaikai23/maskgaussian