論文題目:One is All: Bridging the Gap Between Neural Radiance Fields Architectures with Progressive Volume Distillation
論文連結:https://arxiv.org/abs/2211.15977
導語
神經輻射場 (NeRF) 方法已被證明是有效的,可用于 3D 場景高品質的表示,并支援諸如編輯,檢索,導航等下遊任務。目前多種結構都在争奪 NeRF 的核心表示,包括多層感覺器 (MLP)、稀疏張量(sparse Tensors)、低秩張量(low-rank Tensors)、哈希表(Hashtables)及其之間的組合。而這些表示目前都有一定程度上的 trade-off,例如基于哈希表的表示通常訓練和渲染很快,但由于不同分辨率的哈希混疊表示,其缺乏較為清晰的幾何結構,可能會阻礙諸如基于空間關系感覺的編輯操作。
為了解決上述問題,在本文中我們嘗試進行不同結構間的轉化。我們提出了漸進式的體素蒸餾方案 PVD,該方法可以實作 MLP、稀疏或低秩張量、哈希表等不同結構之間的兩兩互轉。是以 PVD 可以根據不同的下遊任務将手頭已有模型調整為新的結構表示。PVD 是在不同層次表示上逐漸進行的,這種從淺到深的方式使得結構間轉換速度很快。比如使用 PVD将基于哈希表的 Instant-NGP 模型轉化為基于 MLP 的 NeRF 模型要比從頭開始訓練該 MLP 快 10 倍以上,且能實作相近甚至更高水準的圖像合成品質。
背景
新視圖合成(NVS)指在 3D 場景中的未知位置和視角上生成逼真的 2D 圖像,其具有廣泛的應用場景。NeRF[1]的出現極大地加速了這一問題及其相關問題的研究程序。NeRF 使用 MLP 從坐标空間構造出 3D 場景,MLP 的輸入是一個 5D 坐标(空間位置+觀察方向),輸出是對應空間點的 density 和 color。隐式模組化的優點是其表示有利于控制或改變場景的類紋理屬性比如 NeRFW[2]以及時序屬性如 DNeRF[3],即比較容易地進行 3D 場景藝術設計和動态場景設計。但是基于 MLP 的 NeRF 需要對空間點進行密集采樣,導緻在訓練和推理過程中需要對 MLP 進行多次查詢,進而導緻其運作速度很慢。
基于顯式及其混合形式的表達,則在一定程度上解決了完全基于 MLP 的 NeRF 運作較慢的問題。如 Plenoxels[4] 通過完全的顯式表示,将場景直接放置在 3D 網格上,網格上的每個體素都存儲其 density 和 color 相關資訊。由于 Plenoxels 不帶有任何 MLP 結構,其隻需要通過網格體素上存儲的資訊進行插來值得到空間中任意點的資訊。這種方式大大縮減了 Plenoxels 的訓練和推理時間。基于純顯性結構的表示有助于基于空間位置資訊的場景幾何編輯,包括多個場景的合并、特定位置的對象修複和操作等如 CCNeRF[5]。Plenoxels 一個較大的缺陷是模型參數量過大,其最終效果較為依賴網格的分辨率。
除了上述純隐式和純顯式的表示之外,一些混合的表示方法也開始出現。比如最近受低秩張量近似算法的啟發,TensoRF[6] 将顯式張量分解為低秩分量,進而顯着減小了模型大小。INGP[7] 提出了多分辨率哈希編碼 (MHE),通過不同尺度的哈希表級聯将給定坐标映射到一組特征,而後使用一個小的 MLP 将該特征解譯為 density 和 color。MHE 顯着減少了模型記憶體占用并提高了推理速度。然而低秩張量與 MHE 的緊湊性都是以不直覺的幾何解釋為代價,其後期引入的小 MLP 會損失一定的幾何解釋,且多個低秩張量的混合以及哈希表的散列機制導緻了大量的空間混疊。
哪種結構的表示最好,沒有固定的答案。不同下遊任務所需的結構類型不同。是以本文嘗試不同結構間的互轉,在引入互轉方法之前,需要考慮這樣做的好處究竟有哪些:
(1)該類型的研究可以深入了解已經豐富且不斷增長的 NeRF 架構群的模組化能力和局限性。目前對這些結構的研究都是獨立,這樣的方式并不能充分發揮每一種結構的優勢和規避其劣勢。
(2)這種轉換的可能性将使設計人員從事先設計模型結構的負擔中解放出來,因為現在他們可以簡單地将訓練有素的模型快速地轉化為想要的結構,以滿足後續不斷變化的應用場景的需要。
(3)在 teacher 和 student 具有不同屬性的情況下,二者可以實作有效互補,進而獲得表達能力更強的模型。例如當使用哈希表示的 teacher 模型來擷取張量表示的 student 模型時,可以受益于 teacher 更快的訓練速度和更好的性能表現,來得到具有清晰幾何結構且性能更優的 student 模型。
方法
如圖所示,給定一個 teacher 和 student 後,我們将其統一為兩部分的表示,即和 。蒸餾過程分為三個 stage,stage1 隻前傳到 子網絡。stage2 隻前傳到渲染前部分。stage3 則使用全部的資訊進行訓練。這樣的方式有效縮短了反複查詢 MLP 結構的巨大耗時,有效縮短轉化時間。此外在過程中需要對 density 等資訊進行值域限制。
我們的方法旨在實作神經輻射場不同架構之間的互相轉換。由于此類架構的數量不斷增加,是以我們不會嘗試一一實作這些轉換,是以本文主要針對目前常見的幾種結構間的互轉,包括 NeRF 中的 MLP,Plenoxels 中的稀疏張量,以及兩種混合表示:INGP 的哈希表和 TensoRF 的低秩張量(VM 分解)。我們首先以統一形式對不同結構的表示進行切分,然後基于該統一的視圖來設計本文的蒸餾方案。
3.1 統一視圖設計
在實驗中發現如果隻使用 ray 對應的 RGB 和 ray 上采樣點對應的 density 和 color 等資訊來進行蒸餾仍需要完全查詢整個 MLP,這是低效的。為此我們設計一個額外的中間特征來對齊不同結構之間的資訊。我們觀察到混合表示中的隐式和顯式結構是自然分離的,分别對應于不同的學習目标,是以我們考慮将各類形式的模型均拆分成這種兩部分表式形式。具體來說,給定一個模型,我們将它們表示為兩個子產品的級聯,各個結構的劃分如下表所示:
(統一視角架構的兩級劃分方法)
這裡*可以是 teacher 也可以是 student。對于混合表示,我們直接将顯式部分視為,将隐式部分視為。而對于純隐式的 MLP,我們根據網絡的深度将其分為兩部分。對于純顯式表示結構,我們通過讓為恒等映射将其表示為兩部分。
3.2 Loss 設計
蒸餾最後一個階段用到的 loss 如下所示:
其中的子項依次表示 volume-aligned loss, density loss, color loss, RGB loss 和 regularization loss。其中正則化項主要取決于 student 模型的形式,如對于 Plenoxels 和 VM 分解,正則項為稀疏損失和 TV 損失。應該注意的是,我們隻對 Plenoxels 進行 density、color、RGB 和 regularization 損失。
3.2 density 值域限制
我們發現密度 density loss 項很難直接優化。我們将此問題歸因于其特定的數值不穩定性。即 density 可以反映空間中一點的透光率。當 density 大于或小于某個值時,其實體意義是基本一緻的(即完全透明或完全不透明)。是以 density 的取值範圍過大時,實際上隻有一個區間的 density 值起着關鍵作用,而關鍵值之外的部分則會嚴重影響其優化。是以我們在實際計算 density loss 時,隻優化其值域上的一個區間段。
3.2 block-wsie蒸餾
蒸餾分為三個 stage。stage1 隻前傳到子網絡,stage2 隻前傳到渲染前部分,stage3 則使用全部的資訊進行訓練。這樣的方式有效縮短了反複查詢 MLP 結構的巨大耗時,有效縮短轉化時間。
實驗
4.1 互轉結果
Lego 場景上 Hash / VM分解 / MLP / 稀疏張量之間互相轉換的定量和定性結果。我們對每個結構訓練一個 teacher 模型,然後用其蒸餾 4 種結構的 student 模型。圖上數字表示 PSNR。“s-”表示 student.
NeRF-Synthetic 資料集上不同結構間的互轉結果。
對于四種表示(Hash / VM-decomposition / MLP / sparse tensors),我們首先在 NeRF-Synthetic 資料集上的 8 個場景中從頭開始訓練每種結構的模型,總共獲得了 32 個模型作為 teacher。然後使用本文提出的 PVD 将這些 teacher 分别進行 4 種轉換,擷取 128 個 student 模型。轉換完成後我們統計其平均名額(上表)和視覺效果(上圖)。可以看出我們的方法對于不同結構間的轉換是非常有效的。當一個模型轉化為另一種形式時,student 性能與教師差異較小,這充分說明基于輻射場的不同結構表示間可以互相轉換。進一步我們在 LLFF 和 TanksAndTemples 資料集上進行了類似的實驗,并與其原始論文中的方法從頭訓練的結果進行對比,結果顯示我們的方法仍具有可比較的性能表現。
LLFF 和 TanksAndTemples 資料集上從頭開始訓練一個模型與通過本文方法蒸餾得到的模型間的效果對比。
4.2 加速收斂和性能提升
左側是 1.5 小時内通過 PVD 蒸餾出的NeRF模型結果(INGP 作為 teacher)。右側是 NeRF 25 小時從頭開始訓練的結果。 PVD 有效提高了合成品質并減少了訓練時間。
從頭訓練 NeRF 與使用本文方法蒸餾得到 NeRF 的耗時和性能對比。
從上圖和表可以看到,當使用一個性能較高的結構作為 teacher 時,我們的方法可以顯著減少 student 模型訓練時間并提高 student 模型性能。以往相同結構間的蒸餾方式(從大參數量的 NeRF 模型蒸餾到小參數量的 NeRF 模型)的過程異常低效,因為它需要在訓練和蒸餾過程中不斷查詢大參數量 NeRF 模型。而我們異構形式之間的蒸餾可以高效查詢 teacher,能夠實作更快速的蒸餾。
4.3 消融實驗
在 NeRF-Synthetic 資料上的消融實驗(VM 分解蒸餾為 MLP).
有無 density 值域限制的蒸餾結果對比。
消融實驗證明了我們方法中每個元件的必要性。
4.4 模型受限
Synthetic-NeRF 資料集上互相轉換後的 PSNR 差距。 PSNRstu 表示通過蒸餾得到的 student 的 PSNR。 PSNRself 表示從頭開始訓練得到的 PSNR。 PSNRtea 表示 teacher 的 PSNR。
可以看到 student 的性能主要受限于兩方面:一是 teacher 的模組化能力,二是student的模組化能力。當 teacher 能力更強時往往有助于提升 student 的最終表現(如 INGP 作為 teacher,MLP作為 student);當二者能力相當時,二者的最終表現相近(如同種結構間的蒸餾);而當student能力更強時,student 的性能表現會受限于 teacher 的能力(如 MLP 為 teacher,INGP 為 student),此時可以進一步 finetune student 來提升其性能。
總結
在這項工作中,我們提出了一種系統的蒸餾方法 PVD,使得不同 NeRF 架構(包括 MLP、稀疏張量、低秩張量和哈希表)之間進行轉換,同時保持較高的合成品質。PVD 成功的核心是統一結構視角的表示,損失函數的設計、漸進式蒸餾方案以及密度值的特殊處理等。通過打破不同結構間的隔閡,PVD 允許下遊任務在已有模型上進行簡單處理。實驗證明了我們的方法在合成資料集和真實資料集上的有效性。
參考文獻
[1] Mildenhall B, Srinivasan P P, Tancik M, et al. Nerf: Representing scenes as neural radiance fields for view synthesis[J]. Communications of the ACM, 2021, 65(1): 99-106.
[2] Martin-Brualla R, Radwan N, Sajjadi M S M, et al. Nerf in the wild: Neural radiance fields for unconstrained photo collections[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 7210-7219.
[3] Pumarola A, Corona E, Pons-Moll G, et al. D-nerf: Neural radiance fields for dynamic scenes[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 10318-10327.
[4] Fridovich-Keil S, Yu A, Tancik M, et al. Plenoxels: Radiance fields without neural networks[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 5501-5510.
[5] Tang J, Chen X, Wang J, et al. Compressible-composable nerf via rank-residual decomposition[J]. arXiv preprint arXiv:2205.14870, 2022.
[6] Chen A, Xu Z, Geiger A, et al. Tensorf: Tensorial radiance fields[C]//Computer Vision–ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XXXII. Cham: Springer Nature Switzerland, 2022: 333-350.
[7] Müller T, Evans A, Schied C, et al. Instant neural graphics primitives with a multiresolution hash encoding[J]. ACM Transactions on Graphics (ToG), 2022, 41(4): 1-15.