Guy Ohayon Tomer Michaeli Michael Elad

Abstract

Photo-realistic image restoration algorithms are typically evaluated bydistortion measures (e.g., PSNR, SSIM) and by perceptual quality measures(e.g., FID, NIQE), where the desire is to attain the lowest possible distortionwithout compromising on perceptual quality. To achieve this goal, currentmethods typically attempt to sample from the posterior distribution, or tooptimize a weighted sum of a distortion loss (e.g., MSE) and a perceptualquality loss (e.g., GAN). Unlike previous works, this paper is concernedspecifically with the optimal estimator that minimizes the MSE under aconstraint of perfect perceptual index, namely where the distribution of thereconstructed images is equal to that of the ground-truth ones. A recenttheoretical result shows that such an estimator can be constructed by optimallytransporting the posterior mean prediction (MMSE estimate) to the distributionof the ground-truth images. Inspired by this result, we introducePosterior-Mean Rectified Flow (PMRF), a simple yet highly effective algorithmthat approximates this optimal estimator. In particular, PMRF first predictsthe posterior mean, and then transports the result to a high-quality imageusing a rectified flow model that approximates the desired optimal transportmap. We investigate the theoretical utility of PMRF and demonstrate that itconsistently outperforms previous methods on a variety of image restorationtasks.