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

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

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

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

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

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

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

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

無損修剪超過60%高斯點(diǎn)

對(duì)冗余高斯點(diǎn)進(jìn)行剪枝，目前主要有兩種方法：

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

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

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

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

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

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

掩碼光柵化：前向過程

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

其中，，，分別是像素點(diǎn)的顏色，第i個(gè)高斯的顏色，透明度和透光率。

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

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

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

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

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

掩碼光柵化：反向過程

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

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

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

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

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

有趣的是，這個(gè)梯度公式已經(jīng)包含了，即基于分?jǐn)?shù)的剪枝方法所使用的重要性標(biāo)準(zhǔn)。

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

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

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

訓(xùn)練與剪枝

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

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

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

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

通過廣泛的實(shí)驗(yàn)結(jié)果和分析，研究團(tuán)隊(duì)證明了MaskGaussian的有效性。

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

項(xiàng)目地址：https://maskgaussian.github.io/
代碼鏈接：https://github.com/kaikai23/maskgaussian