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

TL;DR

This paper introduces Posterior-Mean Rectified Flow (PMRF), an innovative method that optimally estimates images by transporting the posterior mean to ground-truth distribution, achieving minimal MSE while maintaining perceptual quality.

Contribution

The paper proposes PMRF, a novel algorithm combining posterior mean prediction with rectified flow transport to improve image restoration quality under perceptual constraints.

Findings

PMRF outperforms previous methods on multiple image restoration tasks.

Theoretical analysis confirms PMRF's effectiveness in minimizing MSE with perceptual constraints.

Empirical results demonstrate superior perceptual and distortion metrics.

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 commonly 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.