後驗均值修正流：朝向最小均方誤差的照片逼真圖像恢復

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.