사후평균 교정 플로우: 최소 MSE 사진 현실적 이미지 복원을 향하여

Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration

October 1, 2024

저자: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

초록

사실적인 이미지 복원 알고리즘은 일반적으로 왜곡 측정(예: PSNR, SSIM)과 지각적 품질 측정(예: FID, NIQE)을 통해 평가되며, 목표는 지각적 품질을 희생하지 않고 가능한 한 낮은 왜곡을 달성하는 것입니다. 이 목표를 달성하기 위해 현재 방법은 일반적으로 사후 분포에서 샘플링하거나 왜곡 손실(예: MSE)과 지각적 품질 손실(예: GAN)의 가중 합을 최적화하려고 노력합니다. 이 논문은 이전 연구와는 달리, 복원된 이미지의 분포가 원본 이미지의 분포와 동일한 경우에 MSE를 최소화하는 최적 추정자에 특히 관심을 가집니다. 최근의 이론적 결과는 이러한 추정자가 사후 평균 예측(MMSE 추정)을 원본 이미지의 분포로 최적으로 이동시킴으로써 구성될 수 있다는 것을 보여줍니다. 이 결과에 영감을 받아, 우리는 Posterior-Mean Rectified Flow (PMRF)를 소개합니다. 이는 이러한 최적 추정자를 근사하는 간단하면서 매우 효과적인 알고리즘입니다. 특히, PMRF는 먼저 사후 평균을 예측하고, 그 결과를 원하는 최적 이동 맵을 근사하는 rectified flow 모델을 사용하여 고품질 이미지로 이동시킵니다. 우리는 PMRF의 이론적 유효성을 조사하고, 다양한 이미지 복원 작업에서 이전 방법보다 일관되게 우수한 성능을 보여주는 것을 증명합니다.

English

Photo-realistic image restoration algorithms are typically evaluated by distortion measures (e.g., PSNR, SSIM) and by perceptual quality measures (e.g., FID, NIQE), where the desire is to attain the lowest possible distortion without compromising on perceptual quality. To achieve this goal, current methods typically attempt to sample from the posterior distribution, or to optimize a weighted sum of a distortion loss (e.g., MSE) and a perceptual quality loss (e.g., GAN). Unlike previous works, this paper is concerned specifically with the optimal estimator that minimizes the MSE under a constraint of perfect perceptual index, namely where the distribution of the reconstructed images is equal to that of the ground-truth ones. A recent theoretical result shows that such an estimator can be constructed by optimally transporting the posterior mean prediction (MMSE estimate) to the distribution of the ground-truth images. Inspired by this result, we introduce Posterior-Mean Rectified Flow (PMRF), a simple yet highly effective algorithm that approximates this optimal estimator. In particular, PMRF first predicts the posterior mean, and then transports the result to a high-quality image using a rectified flow model that approximates the desired optimal transport map. We investigate the theoretical utility of PMRF and demonstrate that it consistently outperforms previous methods on a variety of image restoration tasks.