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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Ohayon, Guy, Michaeli, Tomer, Elad, Michael
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915136785088512
author Ohayon, Guy
Michaeli, Tomer
Elad, Michael
author_facet Ohayon, Guy
Michaeli, Tomer
Elad, Michael
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_00418
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration
Ohayon, Guy
Michaeli, Tomer
Elad, Michael
Image and Video Processing
Artificial Intelligence
Computer Vision and Pattern Recognition
Signal Processing
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.
title Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration
topic Image and Video Processing
Artificial Intelligence
Computer Vision and Pattern Recognition
Signal Processing
url https://arxiv.org/abs/2410.00418