[논문 리뷰] Bayesian Image Reconstruction using Deep Generative Models

Razvan Marinescu, Daniel Moyer외 1명|arXiv (Cornell University)|2020. 12. 08.

Advanced Image Processing Techniques참고 문헌 58인용 수 11

한 줄 요약

이 논문은 사전에 훈련된 StyleGAN2 생성기의 사전분포를 활용하여 잠재 벡터에 대한 베이지안 최대사후추정(MAP) 추정을 통해 이미지 복원을 수행하는 Bayesian Reconstruction through Generative Models(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/.

연구 동기 및 목표

• 이미지 복원에서 종단간 지도 학습의 한계를 해결하기 위해, 특히 분포 이탈이나 새로운 손상 모델에서 재학습이 필요로 하는 문제를 해결하고자 한다.
• 특정 작업에 맞게 미세조정하지 않고도 하나의 사전에 훈련된 생성 모델을 통해 다양한 이미지 복원 작업에 일반화할 수 있도록 하고자 한다.
• 단일 복원 결과가 아니라 타당한 솔루션의 분포를 사후 분포 샘플링을 통해 추정하는 베이지안 프레임워크를 제공하고자 한다.
• 의료 영상 포함 다양한 실제 데이터셋에서의 효과성을 입증함으로써, 데이터셋 전용 초모수 조정 없이도 성능을 유지하고자 한다.
• 대규모 쌍체 훈련 데이터와 계산 자원 의존도를 줄이기 위해 사전에 훈련된 생성 모델을 강력한 이미지 사전분포로 활용하고자 한다.

제안 방법

• 이 방법은 사전에 훈련된 StyleGAN2 생성기를 자연 이미지의 분포를 모델링하기 위한 사전분포로 사용하며, 추론 기간 동안 생성기 가중치는 고정된다.
• 주어진 손상된 입력 이미지 I에 대해, 사후분포 p(w|I)를 최대화하는 잠재 벡터 w*를 베이지안 MAP 추정을 통해 추정한다: w* = argmax_w p(w)p(I|f∘G(w)).
• 전방 손상 모델 f는 생성된 이미지 G(w)에 적용되어 관측된 입력 I를 시뮬레이션하며, f는 알려진된 모든 손상 과정(예: 다운샘플링, 마스킹, 블러링)이 될 수 있다.
• 진짜 사후분포 p(w|I)를 근사하기 위해 변분 추론을 사용하여 잠재 공간에서 다수의 타당한 복원 결과를 샘플링할 수 있다.
• 이 방법은 모델의 미세조정 없이 추론 시점에 적용되므로, 새로운 손상 모델이나 입력 분포에 쉽게 적응할 수 있다.
• 이 접근법은 FlickFaces, MIMIC-III 체스트 X-ray, 뇌 MRI 스캔의 세 가지 데이터셋에서 평가되었으며, 모든 데이터셋에서 일관된 성능을 보였다.

실험 결과

연구 질문

• RQ1 재학습 없이도 하나의 사전에 훈련된 생성 모델로 다양한 이미지 복원 작업을 해결할 수 있는가?
• RQ2 깊이 있는 생성 사전분포를 활용한 베이지안 MAP 추정은 종단간 지도 학습 방법에 비해 복원 품질과 다양성 측면에서 어떻게 비교되는가?
• RQ3 사전에 훈련된 생성기의 잠재 공간에서 변분 추론을 통해 의미 있는 타당한 복원 결과의 분포를 제공할 수 있는가?
• RQ4 의료 영상에서 도메인 이탈이 발생하는 경우를 포함한 다양한 데이터 분포에 대해 이 방법은 얼마나 일반화 가능한가?
• RQ5 픽셀 단위 손실 방법에서 흔히 발생하는 평균화 효과를 피하면서도 초해상도 및 완성 작업에서 최신 기술보다 뛰어난 성능을 내는가?

주요 결과

• 16² 해상도에서 FFHQ 초해상도 벤치마크에서 BRGM는 LPIPS 0.24와 RMSE 25.66을 기록하여 PULSE(0.29/27.14)와 ESRGAN(0.35/29.32)를 모두 앞섰다.
• 64² 해상도에서 FFHQ에서 BRGM는 LPIPS 0.36와 RMSE 16.07을 기록했으며, RMSE 측면에선 SRFBN(0.23/9.40)을 앞섰지만 LPIPS 측면에선 약간 뒤진 성과를 보였다.
• 완성 작업에서 BRGM는 FFHQ에서 LPIPS 0.19와 RMSE 24.28을 기록하여 SN-PatchGAN(0.24/30.75)을 두 지표에서 모두 앞섰다.
• 사람 평가에서 BRGM는 16² 해상도에서 최고의 복원 결과로 42%의 표를 얻었으며, PULSE(32%)와 ESRGAN(11%)에 비해 뚜렷하게 뛰어났다.
• MIMIC-III X-ray 데이터셋에서 BRGM는 16² 해상도에서 LPIPS 0.18과 RMSE 11.61을 기록하여 ESRGAN(0.32/14.67)과 SRFBN(0.37/12.28)을 모두 앞섰다.
• 뇌 MRI 스캔에서 BRGM는 16² 해상도에서 LPIPS 0.12와 RMSE 12.42를 기록하여 ESRGAN(0.34/22.81)과 SRFBN(0.33/12.57)을 크게 앞서는 성과를 보였다.