Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration
Guy Ohayon, Tomer Michaeli, Michael Elad
2024-10-02
Summary
This paper introduces Posterior-Mean Rectified Flow (PMRF), a new algorithm designed to improve the restoration of images, making them look more realistic while minimizing errors.
What's the problem?
Image restoration techniques often face a challenge: they need to reduce distortion (errors) in images while also ensuring that the restored images look good to the human eye. Traditional methods either focus too much on reducing errors or on making images visually appealing, which can lead to poor results in one area or the other.
What's the solution?
PMRF addresses this problem by using a two-step process. First, it predicts a basic version of the image that minimizes errors. Then, it refines this image using a method called rectified flow, which helps adjust the image to make it look more like a high-quality original. This approach allows PMRF to produce images that not only have low distortion but also maintain high perceptual quality, meaning they look realistic and appealing.
Why it matters?
This research is important because it provides a better way to restore images, which can be useful in many fields such as photography, film, and medical imaging. By improving how we fix and enhance images, PMRF can help create clearer and more realistic visuals, making it easier for people to analyze and enjoy visual content.
Abstract
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.