Flusso Raddrizzato della Media Posteriore: Verso il Ripristino Minimo dell'Immagine Fotorealistica con MSE Ridotto

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

October 1, 2024

Autori: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Abstract

Gli algoritmi di ripristino di immagini fotorealistiche vengono tipicamente valutati mediante misure di distorsione (ad esempio, PSNR, SSIM) e misure di qualità percettiva (ad esempio, FID, NIQE), dove l'obiettivo è ottenere la minima distorsione possibile senza compromettere la qualità percettiva. Per raggiungere questo obiettivo, i metodi attuali tentano tipicamente di campionare dalla distribuzione posteriore, o di ottimizzare una somma pesata di una perdita di distorsione (ad esempio, MSE) e una perdita di qualità percettiva (ad esempio, GAN). A differenza dei lavori precedenti, questo articolo si occupa specificamente dell'estimatore ottimale che minimizza l'MSE sotto un vincolo di indice percettivo perfetto, ovvero dove la distribuzione delle immagini ricostruite è uguale a quella delle immagini di riferimento. Un recente risultato teorico mostra che tale stimatore può essere costruito trasportando ottimamente la previsione media posteriore (stima MMSE) alla distribuzione delle immagini di riferimento. Ispirandoci a questo risultato, introduciamo il Flusso Raddrizzato della Media Posteriore (PMRF), un algoritmo semplice ma altamente efficace che approssima questo stimatore ottimale. In particolare, PMRF prevede prima la media posteriore, e poi trasporta il risultato a un'immagine di alta qualità utilizzando un modello di flusso raddrizzato che approssima la mappa di trasporto ottimale desiderata. Indaghiamo sull'utilità teorica di PMRF e dimostriamo che supera costantemente i metodi precedenti in una varietà di compiti di ripristino di immagini.

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.