告別3D高斯Splatting算法，帶神經(jīng)補(bǔ)償?shù)念l譜剪枝高斯場(chǎng)SUNDAE開(kāi)源了

作者：機(jī)器之心 2024-05-15 17:35:52

在這篇工作中，我們提出了一種新穎的具有神經(jīng)補(bǔ)償?shù)念l譜剪枝高斯場(chǎng) SUNDAE，通過(guò)引入圖信號(hào)處理，來(lái)建模高斯基元之間的關(guān)系，并混合不同基元的信息來(lái)補(bǔ)償剪枝造成的信息損失。

本論文作者包括帝國(guó)理工學(xué)院碩士生楊潤(rùn)一、北航二年級(jí)碩士生朱貞欣、北京理工大學(xué)二年級(jí)碩士生姜洲、北京理工大學(xué)四年級(jí)本科生葉柏均、中國(guó)科學(xué)院大學(xué)本科大三學(xué)生張逸飛、中國(guó)電信人工智能研究院多媒體認(rèn)知學(xué)習(xí)實(shí)驗(yàn)室（EVOL Lab）負(fù)責(zé)人趙健、清華大學(xué)智能產(chǎn)業(yè)研究院（AIR）助理教授趙昊等。

最近，3D Gaussian Splatting (3DGS) 作為一種新穎的 3D 表示方式，因其快速的渲染速度和高渲染質(zhì)量而受到關(guān)注。然而，這種方法也伴隨著高內(nèi)存消耗，例如，一個(gè)訓(xùn)練好的高斯場(chǎng)可能會(huì)使用超過(guò)三百萬(wàn)個(gè)高斯基元和超過(guò) 700 MB 的內(nèi)存。

近日，帝國(guó)理工學(xué)院、北航、北京理工大學(xué)、中國(guó)科學(xué)院大學(xué)、中國(guó)電信人工智能研究院多媒體認(rèn)知學(xué)習(xí)實(shí)驗(yàn)室（EVOL Lab）、清華大學(xué)智能產(chǎn)業(yè)研究院（AIR）等機(jī)構(gòu)的研究者聯(lián)合發(fā)表了一篇論文《SUNDAE: Spectrally Pruned Gaussian Fields with Neural Compensation》，我們認(rèn)為這種高內(nèi)存占用是由于沒(méi)有考慮基元之間的關(guān)系。在論文中，我們提出了一種名為 SUNDAE 的內(nèi)存高效的高斯場(chǎng)，采用頻譜修剪和神經(jīng)補(bǔ)償。

文章鏈接：https://arxiv.org/abs/2405.00676
項(xiàng)目主頁(yè)：https://runyiyang.github.io/projects/SUNDAE/

一方面，我們基于高斯基元的空間信息構(gòu)建了一個(gè)圖，用于模擬它們之間的關(guān)系，并設(shè)計(jì)了一個(gè)基于圖信號(hào)處理的降采樣模塊來(lái)剪枝，同時(shí)保留所需信號(hào)。另一方面，為了補(bǔ)償剪枝造成的質(zhì)量下降，我們利用了一個(gè)輕量級(jí)神經(jīng)網(wǎng)絡(luò)來(lái)混合渲染特征，有效地補(bǔ)償了質(zhì)量下降，同時(shí)在其權(quán)重中捕獲基元之間的關(guān)系。

我們通過(guò)大量的結(jié)果展示了 SUNDAE 的性能。例如，在 Mip-NeRF360 數(shù)據(jù)集上，SUNDAE 可以在使用 104 MB 內(nèi)存的情況下達(dá)到 26.80 PSNR 和 145 FPS，而標(biāo)準(zhǔn)的 3D Gaussian Splatting 算法在使用 523 MB 內(nèi)存的情況下達(dá)到 25.60 PSNR 和 160 FPS。

與此同時(shí)，自從開(kāi)源后，SUNDAE 受到國(guó)際上廣泛的關(guān)注，受到了知名 NeRF 社區(qū) MrNeRF，AI research 社區(qū)維護(hù)者 Ahsen Khaliq、以及多位相關(guān)領(lǐng)域研究人員轉(zhuǎn)發(fā)關(guān)注。

一、帶神經(jīng)補(bǔ)償?shù)念l譜剪枝高斯場(chǎng)

1.1 基于頻譜圖的剪枝策略

3DGS 使用一組高斯基元來(lái)表示場(chǎng)景，由于這些基元在三維空間中的分布不規(guī)則，我們提出了基于圖的方法來(lái)捕獲基元之間的關(guān)系，而不是使用網(wǎng)格這樣的常規(guī)結(jié)構(gòu)。

具體來(lái)說(shuō)，我們采用圖信號(hào)處理理論來(lái)推導(dǎo)一個(gè)最優(yōu)的采樣策略，該策略能夠基于圖信號(hào)保留特定頻譜的信息。通過(guò)控制頻譜帶寬，我們可以靈活地控制剪枝比例，建模高斯基元之間的關(guān)系。如圖 1 (c)，我們可以控制剪枝 90% 的高斯基元而不降低渲染質(zhì)量。

圖 1: (a) 3DGS 7k 迭代次數(shù)的結(jié)果；(b) 3DGS 30k 迭代次數(shù)的結(jié)果，使用了更多的高斯基元來(lái)表示三維場(chǎng)景，因此質(zhì)量更高、速度更慢、存儲(chǔ)空間更大；(c) 剪枝了 90% 的高斯基元，在存儲(chǔ)空間上大幅減小，但是達(dá)到了相似的渲染效果。

我們使用高斯基元的中心來(lái)作為圖上的信號(hào)輸入，將高斯基元之間的距離作為圖的邊，圖的鄰接矩陣可以表示為

其中是高斯基元的中心點(diǎn)，是一個(gè)閾值超參數(shù)，是距離矩陣的方差。也就是說(shuō)，如果兩個(gè)高斯基元之間的距離比一個(gè)閾值小，那么我們將其之間建立一條圖的邊。建立好圖的鄰接矩陣之后，我們可以根據(jù) Haar-like 濾波器對(duì)圖上的信號(hào)進(jìn)行處理，得到特定頻段的圖信號(hào)。最終根據(jù)想要的頻段信號(hào)進(jìn)行剪枝，本文中我們使用了帶阻濾波器，保留表示物體細(xì)節(jié)的高頻信號(hào)和背景點(diǎn)的低頻信號(hào)。

1.2 神經(jīng)補(bǔ)償機(jī)制

經(jīng)過(guò)頻譜剪枝后，渲染質(zhì)量因?yàn)閯h去了過(guò)多的高斯基元不可避免會(huì)下降，為了解決這個(gè)問(wèn)題，我們采用了一個(gè)神經(jīng)網(wǎng)絡(luò)來(lái)補(bǔ)償這一質(zhì)量損失，如圖 2 所示。

我們從 Gaussian Splatting 轉(zhuǎn)換到了 Feature Splatting，引入一個(gè)輕量級(jí)卷積神經(jīng)網(wǎng)絡(luò)來(lái)輸出高斯基元映射到圖像上的 RGB 值，從而融合不同基元的信息。這使得補(bǔ)償網(wǎng)絡(luò)的權(quán)重在二維圖像空間中間接地來(lái)捕獲基元之間的關(guān)系。

圖 2: 左邊顯示的原版 3DGS，由于沒(méi)有捕捉到基元之間的關(guān)系，因此需要大量的存儲(chǔ)空間；中間展示了我們的頻譜剪枝策略，建模高斯基元之間的關(guān)系；右側(cè)顯示神經(jīng)補(bǔ)償利用 2D 特征來(lái)改善渲染效果。

具體來(lái)說(shuō)，我們不是像 3DGS 那樣直接渲染 RGB 圖像，而是通過(guò)用于 3D 高斯的可微分光柵化器獲得一個(gè)特征圖，該光柵化器將 3D 高斯基元的特征投影到二維特征圖上。

然后，我們利用一個(gè)輕量級(jí)神經(jīng)網(wǎng)絡(luò)來(lái)建模基元之間的關(guān)系并補(bǔ)償頻譜剪枝后的質(zhì)量下降。這個(gè)網(wǎng)絡(luò)由一個(gè)具有 skip-connection 的四層全卷積 U-Net 組成，它聚合來(lái)自不同基元的信息。使用平均池化進(jìn)行下采樣，并使用雙線性插值進(jìn)行圖像上采樣。該網(wǎng)絡(luò)以光柵化的特征圖為輸入，輸出 RGB 圖像。

SUNDAE 的總體框架如下圖 3 所示。

圖 3: (a) Pipeline: 對(duì)一個(gè)預(yù)訓(xùn)練的 3D 高斯場(chǎng)，采用基于圖的剪枝策略對(duì)高斯基元進(jìn)行降采樣，并使用卷積神經(jīng)網(wǎng)絡(luò)來(lái)補(bǔ)償剪枝造成的損失。(b) 基于圖的剪枝：基于高斯基元之間空間關(guān)系的圖被用于剪枝。通過(guò)使用帶阻濾波器，這一過(guò)程便于從高頻組件中提取細(xì)節(jié)信息，同時(shí)捕捉低頻部分的一般特征，從而實(shí)現(xiàn)整個(gè)場(chǎng)景的全面而高效的表示。

1.3 連續(xù)剪枝策略

此外，我們還提出了一個(gè)連續(xù)剪枝的策略來(lái)降低峰值存儲(chǔ)，與訓(xùn)練后剪枝不同，后者從一個(gè)完全密集的高斯場(chǎng)中剪除基元，連續(xù)剪枝涉及在整個(gè)訓(xùn)練過(guò)程中的預(yù)定義間隔定期移除特定數(shù)量或比例的基元。這種方法旨在在訓(xùn)練 3D 高斯場(chǎng)時(shí)持續(xù)控制基元的最大數(shù)量，從而降低訓(xùn)練期間的峰值內(nèi)存需求，并允許在 GPU 內(nèi)存較低的 GPU 設(shè)備上進(jìn)行訓(xùn)練。

經(jīng)驗(yàn)表明，較低峰值內(nèi)存的優(yōu)勢(shì)以較弱的最終內(nèi)存占用控制為代價(jià)。例如，如果我們每 2000 次迭代剪掉 20% 的基元，3D 高斯場(chǎng)的最終收斂狀態(tài)可能會(huì)偏離預(yù)期的 20% 減少。

此外，這種變化可能在不同場(chǎng)景中有所不同，增加了剪枝效果的可預(yù)測(cè)性和一致性的復(fù)雜性。因此，我們將連續(xù)剪枝策略視為必要時(shí)的替代方案。

二、實(shí)驗(yàn)結(jié)果

2.1 定量結(jié)果

我們將 SUNDAE 與最先進(jìn)的 3DGS 和 NeRF 算法進(jìn)行對(duì)比，相比于 3DGS 來(lái)說(shuō)，我們的模型只占用了 10% 的內(nèi)存就可以達(dá)到相似的效果，并且使用 30% 或 50% 的內(nèi)存便能超過(guò)原版 3DGS。并且在 FPS 上遠(yuǎn)遠(yuǎn)超過(guò)了 NeRF 相關(guān)的其他算法。

這是由于我們的模型能較好的捕捉高斯基元之間的關(guān)系，使用更少的高斯基元來(lái)高效地表征三維場(chǎng)景。

2.2 定性結(jié)果

定性結(jié)果中可以看到看到，我們們將 SUNDAE 在 1% 和 10% 采樣率的定性結(jié)果與 3DGS 和 InstantNGP 進(jìn)行比較。

定性結(jié)果顯示，SUNDAE 能夠在只使用 10% 甚至 1% 的內(nèi)存消耗下，達(dá)到類(lèi)似的新視角合成質(zhì)量。圖成功地構(gòu)建了基元之間的關(guān)系，而神經(jīng)補(bǔ)償頭部有效地維持了渲染質(zhì)量。并且從圖 5 的第四行和最后一行可以看到，頻譜剪枝能夠移除靠近攝像機(jī)的漂浮物。

2.3 消融實(shí)驗(yàn)

2.3.1 帶阻濾波器

帶阻濾波器的比率由一個(gè)參數(shù)表示。具體來(lái)說(shuō)，在基于圖的剪枝過(guò)程中，我們采樣了若干基元，包括一定比例 () 的高通和剩余的 (1-) 低通。

結(jié)果表明，這個(gè)參數(shù)對(duì)渲染質(zhì)量有顯著影響，50% 的比率提供了最佳的結(jié)果，而對(duì)低頻或高頻信號(hào)的不成比例強(qiáng)調(diào)會(huì)導(dǎo)致質(zhì)量下降，因?yàn)橥?50% 的比率保留了均衡的高頻細(xì)節(jié)和低頻背景所以效果更優(yōu)。

2.3.2 補(bǔ)償網(wǎng)絡(luò)

如圖 6 和表 2 所示，我們定性和定量地展示了補(bǔ)償網(wǎng)絡(luò)的重要性。如表 2 所示，所有采樣率下，使用神經(jīng)補(bǔ)償相比不使用都表現(xiàn)出了改善的性能。這一點(diǎn)通過(guò)圖 6 中展示的可視化結(jié)果得到了進(jìn)一步支持，展示了該模塊在緩解頻譜剪枝造成的性能下降方面的補(bǔ)償能力。同時(shí)，也證明了基元之間的關(guān)系被很好地捕捉。

如表 3 所示，我們嘗試了不同大小的補(bǔ)償網(wǎng)絡(luò)，增加網(wǎng)絡(luò)大小并不一定能提高渲染質(zhì)量，這與 ADOP 的發(fā)現(xiàn)一致，表明了相似的趨勢(shì)。我們采用 30MB 的 4 層 UNet 作為默認(rèn)設(shè)置，以最佳平衡質(zhì)量和內(nèi)存。

2.3.3 更多的采樣點(diǎn)

如上表 1 所示，保留 50% 的基元在渲染質(zhì)量上優(yōu)于原始的 3DGS。我們還額外測(cè)試了保留 80% 和保留所有基元，以檢驗(yàn)采樣率如何影響最終結(jié)果，如表 4 所示。

結(jié)果顯示，保留 80% 的基元提高了渲染質(zhì)量，根據(jù) LPIPS 顯示出改進(jìn)，但在 PSNR 和 SSIM 上的視覺(jué)提升很小。保留所有基元（并訓(xùn)練更多的周期）無(wú)法進(jìn)一步提高質(zhì)量，這也顯示了建模基元關(guān)系的重要性。如果沒(méi)有有效的關(guān)系建模，更多的基元會(huì)使模型難以收斂，且大量基元對(duì)場(chǎng)景表示產(chǎn)生負(fù)面影響。

此外，我們的目標(biāo)是平衡渲染質(zhì)量與存儲(chǔ)效率；然而，將存儲(chǔ)增加到 620MB 以保留 80% 的基元只帶來(lái)了輕微的質(zhì)量提升，從而降低了存儲(chǔ)效率。

2.3.4 連續(xù)采樣策略

我們?cè)?MipNeRF360 數(shù)據(jù)集中的 Bicycle 和 Counter 場(chǎng)景上測(cè)試了連續(xù)采樣策略，設(shè)定不同的剪枝間隔迭代次數(shù)和剪枝率。如表 5 所示，Points 是訓(xùn)練后的基元數(shù)量，Ratio 是訓(xùn)練后基元數(shù)量與原始 3DGS 的大致比率。

結(jié)果顯示，這種策略可以降低峰值內(nèi)存，但難以控制最終內(nèi)存（通過(guò) Points 和 Ratio 反映）。因此，我們驗(yàn)證了我們的訓(xùn)練后剪枝策略，但仍然在我們的開(kāi)源工具箱中提供連續(xù)剪枝策略作為一種替代方案。

2.3.5 高效性評(píng)估

關(guān)于訓(xùn)練時(shí)間、CUDA 內(nèi)存、渲染幀率和 ROM 存儲(chǔ)的詳細(xì)信息，請(qǐng)參見(jiàn)表 6。值得注意的是，「Ours-50%」版本在可接受的訓(xùn)練時(shí)間內(nèi)（1.41 小時(shí)）達(dá)到了最好的渲染質(zhì)量，同時(shí)實(shí)現(xiàn)了實(shí)時(shí)渲染，并顯著降低了訓(xùn)練期間的 CUDA 內(nèi)存使用和 ROM 存儲(chǔ)。