[論文レビュー] Bayesian Image Reconstruction using Deep Generative Models

Razvan Marinescu, Daniel Moyer+1 more|arXiv (Cornell University)|Dec 8, 2020

Advanced Image Processing Techniques参考文献 58被引用数 11

ひとこと要約

本稿では、事前学習済みのStyleGAN2生成器を事前分布として用い、潜在変数上のベイジアンMAP推定により画像再構成を実行する、ベイジアン再構成による生成モデル（BRGM）を提案する。微調整なしに、さまざまなデータセット（医療画像を含む）における超解像および画像補完において、推論時に変分推論を用いて複数の妥当な再構成をサンプリングすることで、競争力ある性能を達成する。

ABSTRACT

Machine learning models are commonly trained end-to-end and in a supervised setting, using paired (input, output) data. Examples include recent super-resolution methods that train on pairs of (low-resolution, high-resolution) images. However, these end-to-end approaches require re-training every time there is a distribution shift in the inputs (e.g., night images vs daylight) or relevant latent variables (e.g., camera blur or hand motion). In this work, we leverage state-of-the-art (SOTA) generative models (here StyleGAN2) for building powerful image priors, which enable application of Bayes' theorem for many downstream reconstruction tasks. Our method, Bayesian Reconstruction through Generative Models (BRGM), uses a single pre-trained generator model to solve different image restoration tasks, i.e., super-resolution and in-painting, by combining it with different forward corruption models. We keep the weights of the generator model fixed, and reconstruct the image by estimating the Bayesian maximum a-posteriori (MAP) estimate over the input latent vector that generated the reconstructed image. We further use variational inference to approximate the posterior distribution over the latent vectors, from which we sample multiple solutions. We demonstrate BRGM on three large and diverse datasets: (i) 60,000 images from the Flick Faces High Quality dataset (ii) 240,000 chest X-rays from MIMIC III and (iii) a combined collection of 5 brain MRI datasets with 7,329 scans. Across all three datasets and without any dataset-specific hyperparameter tuning, our simple approach yields performance competitive with current task-specific state-of-the-art methods on super-resolution and in-painting, while being more generalisable and without requiring any training. Our source code and pre-trained models are available online: https://razvanmarinescu.github.io/brgm/.

研究の動機と目的

• 画像再構成におけるエンドツーエンドの教師あり学習の限界、特に分布シフトや新しい汚損モデルに対して再学習が必要となる点を解決すること。
• タスク固有のファインチューニングを回避するために、1つの事前学習済み生成モデルを用いて複数の画像修復タスクに一般化可能であるようにすること。
• 単一の再構成ではなく、事後分布サンプリングにより妥当な解の分布を推定するベイジアンフレームワークを提供すること。
• 医療画像を含む多様で現実的なデータセットにおいて、データセット固有のハイパーパrameterチューニングなしに、本手法の有効性を示すこと。
• 大規模なペairedトレーニングデータや計算リソースへの依存を減らすために、強力な画像事前分布として事前学習済み生成モデルを活用すること。

提案手法

• 本手法は、推論中に固定された重みを持つ事前学習済みのStyleGAN2生成器を、自然画像の分布をモデル化する事前分布として用いる。
• 与えられた汚損画像Iに対して、潜在変数w*をベイジアンMAP推定により推定する。すなわち、w* = argmax_w p(w)p(I|f∘G(w)) となるw*を求める。
• 生成画像G(w)に前方汚損モデルfを適用して観測された入力Iをシミュレートする。fは任意の既知の汚損プロセス（例：ダウンサンプリング、マスキング、ぼかし）として定義可能である。
• 真の事後分布p(w|I)を近似するために変分推論を用い、潜在空間からの複数の妥当な再構成をサンプリング可能とする。
• 本手法は推論時のみに適用され、モデルのファインチューニングは一切不要であり、新しい汚損モデルや入力分布に対しても柔軟に適応可能である。
• 本手法はFlickFaces、MIMIC-IIIの胸部X線、および脳MRIスキャンの3つのデータセットで評価され、すべてのデータセットで一貫した性能を示した。

実験結果

リサーチクエスチョン

• RQ1 1つの事前学習済み生成モデルを用いて、再学習なしに複数の画像修復タスクを解決できるか？
• RQ2 深層生成事前分布を用いたベイジアンMAP推定は、エンドツーエンドの教師あり手法と比較して、再構成品質および多様性の面で優れているか？
• RQ3 事前学習済み生成器の潜在空間における変分推論は、意味のある妥当な再構成の分布を提供できるか？
• RQ4 本手法は、ドメインシフトを伴う医療画像を含む多様なデータ分布に一般化可能か？
• RQ5 ピクセル単位の損失関数に起因する平均化効果を避ける一方で、超解像および画像補完において最先端のモデルを上回る性能を示せるか？

主な発見

• FFHQの超解像ベンチマーク（16²解像度）において、BRGMはLPIPS 0.24、RMSE 25.66を達成し、PULSE（0.29/27.14）とESRGAN（0.35/29.32）を上回った。
• 64²解像度では、FFHQでLPIPS 0.36、RMSE 16.07を達成し、RMSEではSRFBN（0.23/9.40）を上回ったが、LPIPSではわずかに劣った。
• 画像補完タスクでは、FFHQでLPIPS 0.19、RMSE 24.28を達成し、SN-PatchGAN（0.24/30.75）を両方の指標で上回った。
• 人間評価では、16²解像度でBRGMが最良の再構成と選ばれる割合が42%に達し、PULSE（32%）とESRGAN（11%）を大きく上回った。
• MIMIC-IIIのX線データセットでは、16²解像度でLPIPS 0.18、RMSE 11.61を達成し、ESRGAN（0.32/14.67）とSRFBN（0.37/12.28）を上回った。
• 脳MRIスキャンでは、16²解像度でLPIPS 0.12、RMSE 12.42を達成し、ESRGAN（0.34/22.81）とSRFBN（0.33/12.57）を著しく上回った。