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

2024-10-01

arXiv.org

AI Summary

How can pixel-level distortion (e.g., MSE) be minimized in image inpainting while strictly preserving perceptual naturalness—i.e., ensuring the reconstructed distribution matches the true data distribution? This paper proposes PMRF, the first framework that unifies Bayesian posterior mean estimation with optimal transport theory to achieve provably MSE-optimal reconstruction under strict perceptual fidelity constraints. PMRF jointly models the posterior mean and a rectified flow-based transport map, eliminating the need for GAN-based adversarial training or iterative sampling. Evaluated across multiple image restoration tasks, PMRF consistently surpasses state-of-the-art methods in both reconstruction fidelity (PSNR) and perceptual quality (FID), simultaneously achieving high pixel accuracy and realistic synthesis. By bridging Bayesian inference and optimal transport, PMRF overcomes fundamental limitations of conventional weighted-loss heuristics and posterior sampling strategies—both theoretically and empirically.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

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.

Problem

• Image Inpainting
• Distortion Minimization
• Visual Naturalness

Innovation

• PMRF
• Image Restoration
• Lower Distortion