告別3D高斯Splatting算法,帶神經補償的頻譜剪枝高斯場SUNDAE開源了
本論文作者包括帝國理工學院碩士生楊潤一、北航二年級碩士生朱貞欣、北京理工大學二年級碩士生姜洲、北京理工大學四年級本科生葉柏均、中國科學院大學本科大三學生張逸飛、中國電信人工智能研究院多媒體認知學習實驗室(EVOL Lab)負責人趙健、清華大學智能產業研究院(AIR)助理教授趙昊等。
最近,3D Gaussian Splatting (3DGS) 作為一種新穎的 3D 表示方式,因其快速的渲染速度和高渲染質量而受到關注。然而,這種方法也伴隨著高內存消耗,例如,一個訓練好的高斯場可能會使用超過三百萬個高斯基元和超過 700 MB 的內存。
近日,帝國理工學院、北航、北京理工大學、中國科學院大學、中國電信人工智能研究院多媒體認知學習實驗室(EVOL Lab)、清華大學智能產業研究院(AIR)等機構的研究者聯合發表了一篇論文《SUNDAE: Spectrally Pruned Gaussian Fields with Neural Compensation》,我們認為這種高內存占用是由于沒有考慮基元之間的關系。在論文中,我們提出了一種名為 SUNDAE 的內存高效的高斯場,采用頻譜修剪和神經補償。
- 文章鏈接:https://arxiv.org/abs/2405.00676
- 項目主頁:https://runyiyang.github.io/projects/SUNDAE/
一方面,我們基于高斯基元的空間信息構建了一個圖,用于模擬它們之間的關系,并設計了一個基于圖信號處理的降采樣模塊來剪枝,同時保留所需信號。另一方面,為了補償剪枝造成的質量下降,我們利用了一個輕量級神經網絡來混合渲染特征,有效地補償了質量下降,同時在其權重中捕獲基元之間的關系。
我們通過大量的結果展示了 SUNDAE 的性能。例如,在 Mip-NeRF360 數據集上,SUNDAE 可以在使用 104 MB 內存的情況下達到 26.80 PSNR 和 145 FPS,而標準的 3D Gaussian Splatting 算法在使用 523 MB 內存的情況下達到 25.60 PSNR 和 160 FPS。
與此同時,自從開源后,SUNDAE 受到國際上廣泛的關注,受到了知名 NeRF 社區 MrNeRF,AI research 社區維護者 Ahsen Khaliq、以及多位相關領域研究人員轉發關注。
一、帶神經補償的頻譜剪枝高斯場
1.1 基于頻譜圖的剪枝策略
3DGS 使用一組高斯基元來表示場景,由于這些基元在三維空間中的分布不規則,我們提出了基于圖的方法來捕獲基元之間的關系,而不是使用網格這樣的常規結構。
具體來說,我們采用圖信號處理理論來推導一個最優的采樣策略,該策略能夠基于圖信號保留特定頻譜的信息。通過控制頻譜帶寬,我們可以靈活地控制剪枝比例,建模高斯基元之間的關系。如圖 1 (c),我們可以控制剪枝 90% 的高斯基元而不降低渲染質量。
圖 1: (a) 3DGS 7k 迭代次數的結果;(b) 3DGS 30k 迭代次數的結果,使用了更多的高斯基元來表示三維場景,因此質量更高、速度更慢、存儲空間更大;(c) 剪枝了 90% 的高斯基元,在存儲空間上大幅減小,但是達到了相似的渲染效果。
我們使用高斯基元的中心來作為圖上的信號輸入,將高斯基元之間的距離作為圖的邊,圖的鄰接矩陣可以表示為
其中是高斯基元的中心點,
是一個閾值超參數,
是距離矩陣的方差。也就是說,如果兩個高斯基元之間的距離比一個閾值小,那么我們將其之間建立一條圖的邊。建立好圖的鄰接矩陣之后,我們可以根據 Haar-like 濾波器對圖上的信號進行處理,得到特定頻段的圖信號。最終根據想要的頻段信號進行剪枝,本文中我們使用了帶阻濾波器,保留表示物體細節的高頻信號和背景點的低頻信號。
1.2 神經補償機制
經過頻譜剪枝后,渲染質量因為刪去了過多的高斯基元不可避免會下降,為了解決這個問題,我們采用了一個神經網絡來補償這一質量損失,如圖 2 所示。
我們從 Gaussian Splatting 轉換到了 Feature Splatting,引入一個輕量級卷積神經網絡來輸出高斯基元映射到圖像上的 RGB 值,從而融合不同基元的信息。這使得補償網絡的權重在二維圖像空間中間接地來捕獲基元之間的關系。
圖 2: 左邊顯示的原版 3DGS,由于沒有捕捉到基元之間的關系,因此需要大量的存儲空間;中間展示了我們的頻譜剪枝策略,建模高斯基元之間的關系;右側顯示神經補償利用 2D 特征來改善渲染效果。
具體來說,我們不是像 3DGS 那樣直接渲染 RGB 圖像,而是通過用于 3D 高斯的可微分光柵化器獲得一個特征圖,該光柵化器將 3D 高斯基元的特征投影到二維特征圖上。
然后,我們利用一個輕量級神經網絡來建模基元之間的關系并補償頻譜剪枝后的質量下降。這個網絡由一個具有 skip-connection 的四層全卷積 U-Net 組成,它聚合來自不同基元的信息。使用平均池化進行下采樣,并使用雙線性插值進行圖像上采樣。該網絡以光柵化的特征圖為輸入,輸出 RGB 圖像。
SUNDAE 的總體框架如下圖 3 所示。
圖 3: (a) Pipeline: 對一個預訓練的 3D 高斯場,采用基于圖的剪枝策略對高斯基元進行降采樣,并使用卷積神經網絡來補償剪枝造成的損失。(b) 基于圖的剪枝:基于高斯基元之間空間關系的圖被用于剪枝。通過使用帶阻濾波器,這一過程便于從高頻組件中提取細節信息,同時捕捉低頻部分的一般特征,從而實現整個場景的全面而高效的表示。
1.3 連續剪枝策略
此外,我們還提出了一個連續剪枝的策略來降低峰值存儲,與訓練后剪枝不同,后者從一個完全密集的高斯場中剪除基元,連續剪枝涉及在整個訓練過程中的預定義間隔定期移除特定數量或比例的基元。這種方法旨在在訓練 3D 高斯場時持續控制基元的最大數量,從而降低訓練期間的峰值內存需求,并允許在 GPU 內存較低的 GPU 設備上進行訓練。
經驗表明,較低峰值內存的優勢以較弱的最終內存占用控制為代價。例如,如果我們每 2000 次迭代剪掉 20% 的基元,3D 高斯場的最終收斂狀態可能會偏離預期的 20% 減少。
此外,這種變化可能在不同場景中有所不同,增加了剪枝效果的可預測性和一致性的復雜性。因此,我們將連續剪枝策略視為必要時的替代方案。
二、實驗結果
2.1 定量結果
我們將 SUNDAE 與最先進的 3DGS 和 NeRF 算法進行對比,相比于 3DGS 來說,我們的模型只占用了 10% 的內存就可以達到相似的效果,并且使用 30% 或 50% 的內存便能超過原版 3DGS。并且在 FPS 上遠遠超過了 NeRF 相關的其他算法。
這是由于我們的模型能較好的捕捉高斯基元之間的關系,使用更少的高斯基元來高效地表征三維場景。
2.2 定性結果
定性結果中可以看到看到,我們們將 SUNDAE 在 1% 和 10% 采樣率的定性結果與 3DGS 和 InstantNGP 進行比較。
定性結果顯示,SUNDAE 能夠在只使用 10% 甚至 1% 的內存消耗下,達到類似的新視角合成質量。圖成功地構建了基元之間的關系,而神經補償頭部有效地維持了渲染質量。并且從圖 5 的第四行和最后一行可以看到,頻譜剪枝能夠移除靠近攝像機的漂浮物。
2.3 消融實驗
2.3.1 帶阻濾波器
帶阻濾波器的比率由一個參數表示。具體來說,在基于圖的剪枝過程中,我們采樣了若干基元,包括一定比例 () 的高通和剩余的 (1-) 低通。
結果表明,這個參數對渲染質量有顯著影響,50% 的比率提供了最佳的結果,而對低頻或高頻信號的不成比例強調會導致質量下降,因為通 50% 的比率保留了均衡的高頻細節和低頻背景所以效果更優。
2.3.2 補償網絡
如圖 6 和表 2 所示,我們定性和定量地展示了補償網絡的重要性。如表 2 所示,所有采樣率下,使用神經補償相比不使用都表現出了改善的性能。這一點通過圖 6 中展示的可視化結果得到了進一步支持,展示了該模塊在緩解頻譜剪枝造成的性能下降方面的補償能力。同時,也證明了基元之間的關系被很好地捕捉。
如表 3 所示,我們嘗試了不同大小的補償網絡,增加網絡大小并不一定能提高渲染質量,這與 ADOP 的發現一致,表明了相似的趨勢。我們采用 30MB 的 4 層 UNet 作為默認設置,以最佳平衡質量和內存。
2.3.3 更多的采樣點
如上表 1 所示,保留 50% 的基元在渲染質量上優于原始的 3DGS。我們還額外測試了保留 80% 和保留所有基元,以檢驗采樣率如何影響最終結果,如表 4 所示。
結果顯示,保留 80% 的基元提高了渲染質量,根據 LPIPS 顯示出改進,但在 PSNR 和 SSIM 上的視覺提升很小。保留所有基元(并訓練更多的周期)無法進一步提高質量,這也顯示了建模基元關系的重要性。如果沒有有效的關系建模,更多的基元會使模型難以收斂,且大量基元對場景表示產生負面影響。
此外,我們的目標是平衡渲染質量與存儲效率;然而,將存儲增加到 620MB 以保留 80% 的基元只帶來了輕微的質量提升,從而降低了存儲效率。
2.3.4 連續采樣策略
我們在 MipNeRF360 數據集中的 Bicycle 和 Counter 場景上測試了連續采樣策略,設定不同的剪枝間隔迭代次數和剪枝率。如表 5 所示,Points 是訓練后的基元數量,Ratio 是訓練后基元數量與原始 3DGS 的大致比率。
結果顯示,這種策略可以降低峰值內存,但難以控制最終內存(通過 Points 和 Ratio 反映)。因此,我們驗證了我們的訓練后剪枝策略,但仍然在我們的開源工具箱中提供連續剪枝策略作為一種替代方案。
2.3.5 高效性評估
關于訓練時間、CUDA 內存、渲染幀率和 ROM 存儲的詳細信息,請參見表 6。值得注意的是,「Ours-50%」版本在可接受的訓練時間內(1.41 小時)達到了最好的渲染質量,同時實現了實時渲染,并顯著降低了訓練期間的 CUDA 內存使用和 ROM 存儲。
三、結論
在這篇工作中,我們提出了一種新穎的具有神經補償的頻譜剪枝高斯場 SUNDAE,通過引入圖信號處理,來建模高斯基元之間的關系,并混合不同基元的信息來補償剪枝造成的信息損失。
我們使用高斯基元之間的空間信息構建圖來建模關系,并根據頻譜信息進行剪枝,去除冗雜的基元。一個輕量級神經網絡被用來補償剪枝后不可避免的渲染質量損失。
實驗結果表明,SUNDAE 在保持 3DGS 的效率的同時,顯著減小了內存,提升了效率并且保持了高保真的渲染質量。