今天小編分享的科學經驗:剪枝60%不損性能!上海AI Lab提出高斯剪枝新方法,入選CVPR 2025,歡迎閱讀。
三維高斯潑濺(3D Gaussian Splatting)使得實時高質量渲染成為可能,是當前 3D 視覺領網域最常用的算法之一。
問題在于,基于顯示高斯單元的表示方式,盡管可以高效濺射和光栅化,其密集化和優化過程卻往往會生成冗餘的高斯點,導致單個重建場景可能包含數百萬個高斯點。
這不僅降低了訓練和渲染速度(本可能更快),還導致顯著的内存消耗。
現在,來自上海 AI Lab 的研究團隊提出MaskGaussian,将掩碼融合進光栅化過程,首次為被使用和未被使用的高斯同時保留梯度,在剪枝高斯的同時,MaskGaussian 極大限度地保持了重建質量,提高了訓練速度和減小内存需求。
△與 3DGS 相比,MaskGaussian 在不影響重建質量的情況下減少高斯點數
該工作既支持從頭開始的訓練,也支持對現有的高斯進行微調。
△MaskGaussian 使用更少的高斯更優秀的地還原出藤蔓的 " 細枝末節 "
實驗結果顯示,MaskGaussian 在Mip-NeRF360、Tanks & Temples和Deep Blending三個數據集上分别剪枝了62.4%、67.7%和75.3%的高斯點,且性能損失可以忽略不計。
對冗餘高斯點進行剪枝,目前主要有兩種方法:
第一類方法基于手工設計的重要性評分,移除評分低于預設阈值的高斯點。這類方法通常需要掃描所有訓練影像以計算重要性評分,因此剪枝只能在訓練期間執行一次或兩次。
第二類方法使用可學習的掩碼,将其與高斯點的屬性相乘以接受梯度。盡管這種方法允許通過掩碼逐漸移除高斯點,但渲染的場景始終依賴于相同的高斯點子集 : 如果一個高斯點未被移除,它會一直存在到當前迭代; 一旦被移除 , 它将被永久排除。這種确定性的掩碼生成未能考慮剪枝後場景的演化,可能導致一些當前貢獻較小但在訓練後期可能關鍵且難以恢復的高斯點被移除。這導致次優的重建效果,尤其是細節或小物體的丢失。
研究團隊認為,當前存在的掩碼剪枝策略表現受限的主要原因在于:掩碼與高斯的透明度等屬性相乘後,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 個高斯點的顏色和其後的顏色累計值。
梯度公式可以分為兩部分理解:
對顏色的權重:決定了該高斯點對最終顏色的影響程度,也間接決定了該點對梯度的貢獻大小。較大的權重意味着該點對渲染結果的重要性更高。
顏色優化方向:
中,表示損失函數對輸出顏色的優化方向,而
表示使用第 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
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法!
— 完 —
學術投稿請于工作日發郵件到:
标題注明【投稿】,告訴我們:
你是誰,從哪來,投稿内容
附上論文 / 項目主頁鏈接,以及聯系方式哦
我們會(盡量)及時回復你
點亮星标
科技前沿進展每日見