Flujo Rectificado de la Media Posterior: Hacia la Restauración de Imágenes Foto-Realistas con Mínimo Error Cuadrático Medio

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

October 1, 2024

Autores: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Resumen

Los algoritmos de restauración de imágenes fotorealistas suelen evaluarse mediante medidas de distorsión (por ejemplo, PSNR, SSIM) y medidas de calidad perceptual (por ejemplo, FID, NIQE), donde el objetivo es lograr la menor distorsión posible sin comprometer la calidad perceptual. Para alcanzar este objetivo, los métodos actuales suelen intentar muestrear de la distribución posterior, u optimizar una suma ponderada de una pérdida de distorsión (por ejemplo, ECM) y una pérdida de calidad perceptual (por ejemplo, GAN). A diferencia de trabajos anteriores, este artículo se centra específicamente en el estimador óptimo que minimiza el ECM bajo una restricción de índice perceptual perfecto, es decir, donde la distribución de las imágenes reconstruidas es igual a la de las imágenes de referencia. Un resultado teórico reciente muestra que dicho estimador se puede construir transportando de manera óptima la predicción media posterior (estimación MMSE) a la distribución de las imágenes de referencia. Inspirados por este resultado, presentamos Posterior-Mean Rectified Flow (PMRF), un algoritmo simple pero altamente efectivo que aproxima este estimador óptimo. En particular, PMRF primero predice la media posterior, y luego transporta el resultado a una imagen de alta calidad utilizando un modelo de flujo rectificado que aproxima el mapa de transporte óptimo deseado. Investigamos la utilidad teórica de PMRF y demostramos que supera consistentemente a métodos anteriores en una variedad de tareas de restauración de imágenes.

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.