Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration (2410.00418v3)

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.

• The paper introduces PMRF, which combines posterior mean prediction with rectified flow to optimally minimize MSE under perfect perceptual quality constraints.
• The paper demonstrates significant improvements in FID, KID, PSNR, and SSIM across tasks like denoising, super-resolution, and blind face restoration.
• The paper provides theoretical insights proving that its optimal estimator approach yields MSE equal to or lower than that of traditional posterior sampling methods.

Posterior-Mean Rectified Flow for Photo-Realistic Image Restoration

The paper presents a novel approach, Posterior-Mean Rectified Flow (PMRF), designed to enhance the performance of photo-realistic image restoration tasks by minimizing Mean Squared Error (MSE) under a perfect perceptual quality constraint. Unlike existing methods that primarily focus on posterior sampling or weighted sums of distortion and perceptual quality losses, PMRF aims for a theoretically optimal estimator that directly addresses this tradeoff.

Key Contributions and Methodology

PMRF is conceptualized by leveraging a theoretical result which suggests that the optimal estimator in terms of MSE can be constructed by transporting the posterior mean prediction to align with the distribution of ground-truth images. Inspired by this, PMRF involves a two-stage process:

1. Posterior Mean Prediction: A model is trained to predict the posterior mean by minimizing the MSE between the predicted and ground-truth images. This is instrumental in achieving the Minimum MSE (MMSE).
2. Optimal Transport via Rectified Flow: A rectified flow model is then employed to transform the posterior mean prediction into a high-quality image. This model, trained to approximate the desired optimal transport map, leverages straight-line paths interpolated between the posterior predictions and the high-quality image distribution.

Results and Comparative Analysis

PMRF demonstrates a consistent performance advantage across various image restoration tasks including denoising, super-resolution, inpainting, colorization, and notably, blind face image restoration. It sets new benchmarks on challenging datasets such as CelebA-Test, where it surpasses state-of-the-art methods in key metrics: FID, KID, PSNR, and SSIM, indicating high perceptual quality and low distortion.

Theoretical Implications

The paper underscores an important theoretical insight by showing that PMRF approximates the optimal estimator which minimizes MSE under a perfect perceptual constraint. The authors prove that PMRF's MSE is either equal to or less than that achieved by posterior sampling methods—a significant enhancement given its propensity to achieve perfect perceptual index theoretically.

Practical Applications

On a practical front, PMRF is poised to significantly impact applications requiring high fidelity and visually appealing image reconstructions, such as medical imaging and mobile photography. Its robust performance across multiple distortion and perceptual quality measures makes it a versatile tool in the arsenal of image restoration techniques.

Future Directions

The research opens avenues for further exploration in applying PMRF to a broader set of image domains, including non-face images or larger datasets. Additionally, integrating PMRF with architectures capable of handling more complex dependencies and larger-scale datasets holds potential for advancements in image restoration capabilities.

In conclusion, PMRF represents a significant contribution to the field by not only advancing the theoretical understanding of image restoration under the distortion-perception tradeoff but also by demonstrating empirical success across challenging tasks, paving the way for future innovations in photo-realistic image restoration methodologies.