Posterior-Meangewogen Stroom: Naar Minimale MSE Foto-Realistische Afbeeldingsherstel

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

October 1, 2024

Auteurs: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Samenvatting

Algoritmes voor het herstellen van fotorealistische afbeeldingen worden doorgaans geëvalueerd aan de hand van vervormingsmetingen (bijv. PSNR, SSIM) en perceptuele kwaliteitsmetingen (bijv. FID, NIQE), waarbij het streven is om de laagst mogelijke vervorming te bereiken zonder in te leveren op de perceptuele kwaliteit. Om dit doel te bereiken, proberen huidige methoden typisch te samplen uit de posterior distributie, of een gewogen som te optimaliseren van een vervormingsverlies (bijv. MSE) en een perceptuele kwaliteitsverlies (bijv. GAN). In tegenstelling tot eerdere werken, richt dit artikel zich specifiek op de optimale schatter die de MSE minimaliseert onder een beperking van een perfect perceptueel index, namelijk waar de distributie van de gereconstrueerde afbeeldingen gelijk is aan die van de werkelijke afbeeldingen. Een recent theoretisch resultaat toont aan dat zo'n schatter geconstrueerd kan worden door de posterior mean voorspelling (MMSE schatting) optimaal te transporteren naar de distributie van de werkelijke afbeeldingen. Geïnspireerd door dit resultaat introduceren we Posterior-Mean Rectified Flow (PMRF), een eenvoudig maar zeer effectief algoritme dat deze optimale schatter benadert. In het bijzonder voorspelt PMRF eerst de posterior mean, en transporteert vervolgens het resultaat naar een hoogwaardige afbeelding met behulp van een gerechtvaardigd stroommodel dat de gewenste optimale transportmap benadert. We onderzoeken het theoretische nut van PMRF en tonen aan dat het consequent beter presteert dan eerdere methoden bij verschillende taken voor het herstellen van afbeeldingen.

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.