Fluxo Retificado da Média Posterior: Rumo à Restauração de Imagens Foto-realistas com Mínimo Erro Quadrático Médio

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

October 1, 2024

Autores: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Resumo

Algoritmos de restauração de imagens fotorrealistas são tipicamente avaliados por medidas de distorção (por exemplo, PSNR, SSIM) e por medidas de qualidade perceptual (por exemplo, FID, NIQE), onde o objetivo é alcançar a menor distorção possível sem comprometer a qualidade perceptual. Para atingir esse objetivo, os métodos atuais normalmente tentam amostrar da distribuição posterior, ou otimizar a soma ponderada de uma perda de distorção (por exemplo, MSE) e uma perda de qualidade perceptual (por exemplo, GAN). Ao contrário de trabalhos anteriores, este artigo trata especificamente do estimador ótimo que minimiza o MSE sob uma restrição de índice perceptual perfeito, ou seja, onde a distribuição das imagens reconstruídas é igual à das imagens verdadeiras. Um resultado teórico recente mostra que tal estimador pode ser construído transportando de forma ótima a previsão da média posterior (estimativa MMSE) para a distribuição das imagens verdadeiras. Inspirados por esse resultado, introduzimos o Fluxo Retificado da Média Posterior (PMRF), um algoritmo simples, porém altamente eficaz, que aproxima esse estimador ótimo. Em particular, o PMRF primeiro prevê a média posterior e, em seguida, transporta o resultado para uma imagem de alta qualidade usando um modelo de fluxo retificado que aproxima o mapa de transporte ótimo desejado. Investigamos a utilidade teórica do PMRF e demonstramos que ele supera consistentemente métodos anteriores em uma variedade de tarefas de restauração de imagens.

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.