論文解説 StableNormal: Reducing Diffusion Variance for Stable and Sharp Normal

hugging face demo :https://huggingface.co/spaces/Stable-X/StableNormal
arxiv : https://arxiv.org/abs/2406.16864

ひとことまとめ

写真から高精度のnormal推定を可能にした

概要

以前の研究では、確率的な推論と、それらのアンサンブルに高いコストがかかっていた。そこで、アンサンブルを用いず、推論の分散を減らす方法を採用することで、安定かつ鮮明(Stable-and-Sharp)なnormalを生成する手法を提案する。

提案手法

DIffusion modelの分散

Diffusionベースのnormal推定は高い分散を持つことが知られている。

これは、diffusionのランダム性に依存するが、そのほとんどは

• モデルに入力する初期ノイズ

• 中間に挿入されるノイズ

の2つが原因である。そこで、分散を減らすため、2段階で推論する手法を提案する。まず、信頼できる初期値推定を行い(You-Only-Sample-Once Normal Initialization)、これをもとに微調整をdiffusionを用いて行う(Semantic-guided Normal Refinement)方法である。

You-Only-Sample-Once Normal Initialization

１ステップでnormalを推定する方法はGenPerceptで導入された。この方法ではガウスノイズを挿入せず、推論プロセスは決定的だが、その分出力が平滑化されてしまう問題がある。提案手法ではガウスノイズをもとにサンプリングすることで、安定性と鮮明さのバランスをとった推論を行う。しかし、ガウスノイズから直接画像を生成するのは、無数のガウスノイズから同じ画像を生成する必要があるため、一般に難しい。そこで、Shrinkage Regularizerを導入することで、これを解消する。

Shrinkage Regularizer

推論時の分散を減らすため、出力に対して何らかの制約が必要である。予測分布に対して直接エントロピーの制約を追加するのは難しい(分布推定に大量の生成が必要かつ分布が正規分布とは限らない)ため、ノイズから予測する場合と、ノイズの代わりに0で初期化した場合の2つをランダムに切り替えることで出力の分散を減らす。

ここで、$${p}$$は一様分布、$${\lambda}$$は0.4である。
式としては複雑に見えるが、要は60%はノイズからサンプリングした画像で推論を行い、40%は0で初期化された画像で推論を行うということである。

Semantic-guided Normal Refinement

YOSOで推定したnormalをベースに生成を行う場合、局所的(local)な情報のみをもとに画像を生成する傾向がある。しかし壁などにおいては大局的(global)な情報も使用しなければ正確な推定はできない。そこで事前学習されたモデル(DINO v2)の特徴量を意味的(大局的)な情報として追加することでこれを解消する。

モデルの構造はYOSOとほぼ同じだが、意味挿入ネットワーク(Semantic-injection Network)が追加されている。(図ではSGNと記載されている)

Semantic-injection Network

SGNは4層の$${3 \times 3}$$畳み込みで実装され、DINO特徴量をStable Diffusionの潜在変数と同じ解像度に変形する役割を持つ。DINO特徴量はそのままだと低解像度のため、FeatUp(特徴量の解像度を上げる(超解像)ネットワークを学習する手法)とbilinear補間で特徴量の空間解像度を高解像度にした。ノイズ付き潜在変数をU-Netに入力される前にDINOv2特徴量を加算する。学習時はSGNの最終層は0で初期化し、それ以外の層は正規分布で初期化する。

Heuristic Denoising Sampling

サンプリングにはDDIMを用い、最終出力を生成する。YOSOで予測したあと、10回のDDIMを経由して最終的な出力とした。

ここで、$${x_{t^+}}$$はYOSOで推測した初期normalで、$${t^+}$$は401を用いている。つまり、YOSOで$${t}$$=401の状態までもっていき、残りの400をSG-DRNを使いDDIMで計算している。ここで、$${t^+}$$に401を用いているのは、安定性と鮮明度のバランスがよかったためである。

実験

まず、提案手法の推定結果がどのくらいの分散になるかを計算した。

グラフにある通り、提案手法はアンサンブルの数によらず低い分散をもち、またアンサンブルの数にほとんど依存せず推論を行えることがわかる。

定量的評価では、すべてのデータセットで提案手法が1もしくは2番の精度を達成している。特にDIODE-indoorはMarigoldに大差をつけている。

定性的評価では、提案手法は他の手法と比べ、不鮮明な推定を行わず境界がはっきりした識別を行えていることがわかる。また反射光や透明な物体があっても問題なく推論できていることもわかる。

まとめ

• StableDiffusionを用いた画像から鮮明なNormalを推測する手法を提案

• 提案手法は2段階で推論を行い、1段階目でおおまかなNormalを、2段階目で詳細なNormalを推定する

• 他の手法よりも定量的・定性的に優れており、ボケが少なく鮮明なNormalを推定できていることを示した