Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fieldsを読んだメモ

　主な工夫は3つ

空間のWarping
オンライン蒸留
新しい正則化項

1. 空間のWarping

　遠いところをNeRFが扱える空間に持ってくるために、無限遠を一定球内に潰してくるようなWarpingを考える。

　 $f( \boldsymbol{x} )$ を $\mathbb{R} ^ 3 \to \mathbb{R} ^ 3$ の滑らかな座標変換関数とする。具体的には

$\mathrm{contract}( \boldsymbol{x} ) = \begin{cases} \boldsymbol{x} & \|\boldsymbol{x}\| \le 1 \\ \left( 2 - \frac{1}{\|\boldsymbol{x}\|} \right) \left( \frac{\boldsymbol{x}}{\|\boldsymbol{x}\|} \right) & \|\boldsymbol{x}\| \gt 1 \end{cases}$

とする。

　Mip-NeRFでは元の空間でのガウス分布を考えるので、ガウス分布が上記のcontract関数によってどのように変換されるかを考える。

　関数の線形近似により、ある点 $\boldsymbol{\mu}$ についてのヤコビアンを $\boldsymbol{J} _ f ( \boldsymbol{\mu} )$ としたとき

$f( \boldsymbol{x} ) \approx f( \boldsymbol{\mu} ) + \boldsymbol{J} _ f ( \boldsymbol{\mu} ) ( \boldsymbol{x} - \boldsymbol{\mu} )$

である。ここからExtended Kalman Filterと同様に考えると

$f( \boldsymbol{\mu}, \boldsymbol{\Sigma} ) = \left ( f( \boldsymbol{\mu} ), \boldsymbol{J} _ f ( \boldsymbol{\mu} ) \boldsymbol{\Sigma} \boldsymbol{J} _ f ( \boldsymbol{\mu} ) ^ \mathrm{T} \right)$

となる。この変形されたガウス分布についてMip-NeRFの処理を当てはめるようにする。

　そして、NeRFのレンダリング時のサンプリングにも気を使わないといけないらしい？ここがよくわからなかった。変換前のt-空間でいつも通りサンプリングしてからcontract関数で変換しても良さそうだし、変換後のs-空間で均一サンプリングするのでも良さそうなのに、(11)式みたいなものを持ち出してきて汎用的にしようとしているように感じる。

see [32] for a detailed motivation of this spacing

と言われても流石に厳しい（[32]の論文はこれ）。具体的にどこなんだ。ちょっと違う分野の論文っぽいし全部読むのしんどい……。

2. オンライン蒸留

　NeRFで使われるCoarse & Fineの2段階レンダリングという戦略において、Coarse側の色はほぼ最終的なピクセルの色に寄与しない。ならばCoarse側は密度を提案するだけのネットワークでも良いはず、という感じの考えだと思われる。なのでMip-NeRF360ではproposal MLPとNeRF MLPの2種類を使う。前者がCoarse、後者がFineに近い。

　ただ学習方法が変わる。proposalでは密度だけを出力するので、NeRF MLPをオンライン蒸留するような形で学習する。

　このオンライン蒸留において問題になるのは、ある光線について、proposal MLPで使用されるビンとNeRF MLPで用いられるビンが異なること。単純な比較が難しいため、損失としてはまず、少しでも被覆がある区間の和を取る関数boundを

$\mathrm{bound}(\hat{\boldsymbol{t}}, \hat{\boldsymbol{w}}, T) = \sum _ {j : T \cap \hat{T} _ j \neq \varnothing } \hat{w} _ j$

として考え、これを超えてしまうところで損失を計算する。

$\mathcal{L} _ {\mathrm{prop}} (\boldsymbol{t}, \boldsymbol{w}, \hat{\boldsymbol{t}}, \hat{\boldsymbol{w}}) = \sum _ i \frac{1}{w _ i} \max (0, w _ i - \mathrm{bound}(\hat{\boldsymbol{t}}, \hat{\boldsymbol{w}}, T) )$

　ここは動画を観るとわかりやすかった。

　図の通り、この勾配はproposal MLPの方にだけ流れる。そうでないと、学習序盤ではめちゃくちゃな提案をするproposalに本番NeRFの方が引っ張られてしまう。

3. 新しい正則化項

　NeRFをレンダリングして出来てしまうFloaterを減らしたい。Floaterというのは要するに浮いているべきでないものが浮いているということなので、そういう独立したところが減るように損失を設計したくなる。なので

$\mathcal{L} _ {\mathrm{dist}} (\boldsymbol{s}, \boldsymbol{w}) = \int \int _ {-\infty} ^ {\infty} \boldsymbol{w} _ s (u) \boldsymbol{w} _ s (v) \left| u - v \right| du dv$

というものを考えたくなる。重みが立っているところほど、そのu,v間の距離を縮めようとする力が働くようになる。とはいえこれは計算が大変なので、

$\mathcal{L} _ {\mathrm{dist}} (\boldsymbol{s}, \boldsymbol{w}) = \sum _ {i, j} w _ i w _ j \left| \frac{s _ i + s _ {i + 1}}{2} - \frac{s _ j + s _ {j + 1}}{2} \right| + \frac{1}{3} \sum _ i w _ i ^ 2 (s _ {i + 1} - s _ i)$

　第二項の1/3がどこから出てきたのかはよくわからない。