Последующее среднее исправленное потоковое распределение: к минимизации среднеквадратичной ошибки восстановления фотореалистичного изображения.

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

October 1, 2024

Авторы: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Аннотация

Алгоритмы восстановления фотореалистичных изображений обычно оцениваются с использованием мер искажения (например, PSNR, SSIM) и мер качества восприятия (например, FID, NIQE), где целью является достижение наименьшего искажения без ущерба для восприятия качества. Для достижения этой цели текущие методы обычно пытаются выбирать образцы из апостериорного распределения или оптимизировать взвешенную сумму потерь искажения (например, MSE) и потерь восприятия качества (например, GAN). В отличие от предыдущих работ, в данной статье рассматривается оптимальный оценщик, который минимизирует MSE при условии идеального восприятия, а именно когда распределение восстановленных изображений равно распределению исходных. Недавнее теоретическое открытие показывает, что такой оценщик можно построить, оптимально транспортируя апостериорное среднее предсказание (оценка MMSE) к распределению исходных изображений. Вдохновленные этим результатом, мы представляем Posterior-Mean Rectified Flow (PMRF) - простой, но очень эффективный алгоритм, приближающий этот оптимальный оценщик. В частности, PMRF сначала предсказывает апостериорное среднее, а затем транспортирует результат к высококачественному изображению с помощью модели исправленного потока, приближающей желаемое оптимальное транспортное отображение. Мы исследуем теоретическую полезность PMRF и демонстрируем, что он последовательно превосходит предыдущие методы на различных задачах восстановления изображений.

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.