后验均值修正流：朝向最小均方误差照片逼真图像恢复

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

October 1, 2024

作者: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

摘要

通常，逼真图像恢复算法的评估通过失真度量（例如 PSNR、SSIM）和感知质量度量（例如 FID、NIQE）进行，其目标是在不影响感知质量的前提下实现尽可能低的失真。为实现这一目标，当前方法通常尝试从后验分布中采样，或优化失真损失（例如 MSE）和感知质量损失（例如 GAN）的加权和。与以往研究不同，本文专注于在完美感知指数约束下最小化 MSE 的最优估计器，即重建图像的分布等于地面真实图像的分布。最近的理论结果表明，通过将后验均值预测（MMSE 估计）最优地传输到地面真实图像的分布，可以构建这样的估计器。受此结果启发，我们引入后验均值矫正流（PMRF），这是一种简单而高效的算法，可近似实现这一最优估计器。具体而言，PMRF 首先预测后验均值，然后使用近似所需最优传输映射的矫正流模型将结果传输到高质量图像。我们研究了 PMRF 的理论效用，并证明在各种图像恢复任务中，它始终优于以往方法。

English

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.