Flux redressé de la moyenne postérieure : vers une restauration d'images photoréalistes avec erreur quadratique moyenne minimale

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

October 1, 2024

Auteurs: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Résumé

Les algorithmes de restauration d'images photoréalistes sont généralement évalués par des mesures de distorsion (par exemple, PSNR, SSIM) et par des mesures de qualité perceptuelle (par exemple, FID, NIQE), où l'objectif est d'atteindre la plus faible distorsion possible sans compromettre la qualité perceptuelle. Pour atteindre cet objectif, les méthodes actuelles tentent généralement d'échantillonner à partir de la distribution postérieure, ou d'optimiser une somme pondérée d'une perte de distorsion (par exemple, MSE) et d'une perte de qualité perceptuelle (par exemple, GAN). Contrairement aux travaux précédents, cet article s'intéresse spécifiquement à l'estimateur optimal qui minimise le MSE sous une contrainte d'indice perceptuel parfait, c'est-à-dire lorsque la distribution des images reconstruites est égale à celle des images de référence. Un résultat théorique récent montre qu'un tel estimateur peut être construit en transportant de manière optimale la prédiction moyenne postérieure (estimation MMSE) vers la distribution des images de référence. Inspiré par ce résultat, nous introduisons le Flux Rectifié de la Moyenne Postérieure (PMRF), un algorithme simple mais très efficace qui approxime cet estimateur optimal. En particulier, PMRF prédit d'abord la moyenne postérieure, puis transporte le résultat vers une image de haute qualité en utilisant un modèle de flux rectifié qui approxime la carte de transport optimale souhaitée. Nous étudions l'utilité théorique de PMRF et démontrons qu'il surpasse systématiquement les méthodes précédentes dans diverses tâches de restauration d'images.

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.