告別3D高斯Splatting演算法,帶神經補償的頻譜剪枝高斯場SUNDAE開源了

机器之心發表於2024-05-15
圖片

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年,機器之心AIxiv專欄接收報導了2000多篇內容,覆蓋全球各大高校與企業的頂級實驗室,有效促進了學術交流與傳播。如果您有優秀的工作想要分享,歡迎投稿或者聯絡報導。投稿郵箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com


本論文作者包括帝國理工學院碩士生楊潤一、北航二年級碩士生朱貞欣、北京理工大學二年級碩士生薑洲、北京理工大學四年級本科生葉柏均、中國科學院大學本科大三學生張逸飛、中國電信人工智慧研究院多媒體認知學習實驗室(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 的效率的同時,顯著減小了記憶體,提升了效率並且保持了高保真的渲染質量。

相關文章