Posterior-Mittelwert-geregelte Strömung: Auf dem Weg zur minimalen MSE fotorealistischen Bildwiederherstellung

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

October 1, 2024

Autoren: Guy Ohayon, Tomer Michaeli, Michael Elad

cs.AI

Zusammenfassung

Algorithmen zur photorealistischen Bildwiederherstellung werden in der Regel anhand von Verzerrungsmessungen (z. B. PSNR, SSIM) und wahrnehmungsbezogenen Qualitätsmessungen (z. B. FID, NIQE) bewertet, wobei das Ziel darin besteht, die geringstmögliche Verzerrung zu erreichen, ohne die wahrnehmungsbezogene Qualität zu beeinträchtigen. Um dieses Ziel zu erreichen, versuchen aktuelle Methoden in der Regel, aus der posterioren Verteilung zu sampeln oder eine gewichtete Summe aus einem Verzerrungsverlust (z. B. MSE) und einem wahrnehmungsbezogenen Qualitätsverlust (z. B. GAN) zu optimieren. Im Gegensatz zu früheren Arbeiten befasst sich dieser Artikel speziell mit dem optimalen Schätzer, der den MSE unter der Bedingung eines perfekten wahrnehmungsbezogenen Index minimiert, nämlich wenn die Verteilung der rekonstruierten Bilder derjenigen der Ground-Truth-Bilder entspricht. Ein kürzlich erzielter theoretischer Ergebnis zeigt, dass ein solcher Schätzer konstruiert werden kann, indem die posteriore Mittelwertvorhersage (MMSE-Schätzung) optimal zur Verteilung der Ground-Truth-Bilder transportiert wird. Inspiriert von diesem Ergebnis stellen wir Posterior-Mean Rectified Flow (PMRF) vor, einen einfachen, aber äußerst effektiven Algorithmus, der diesen optimalen Schätzer approximiert. Insbesondere sagt PMRF zunächst den posterioren Mittelwert vorher und transportiert dann das Ergebnis mithilfe eines rectified Flow-Modells zu einem hochwertigen Bild, das die gewünschte optimale Transportabbildung approximiert. Wir untersuchen den theoretischen Nutzen von PMRF und zeigen, dass er konsistent bessere Leistungen als frühere Methoden bei einer Vielzahl von Bildwiederherstellungsaufgaben erbringt.

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.