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

Authors: Guy Ohayon, Tomer Michaeli, Michael Elad
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Image restoration, blind face image restoration, inverse problems, perception-distortion tradeoff, image processing, computer vision, machine learning, generative models, optimal transport, rectified flow, flow matching

Abstract:
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 commonly 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.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730193992349

thread_type:
official_review

<thread_object>
creator: reviewer_rjjp

<content>
summary:
The paper introduces posterior-mean rectified flow framework, which approximates optimal estimator under perceptual constraint. The method includes predicting posterior mean and subsequent distribution transport using rectified flow. The experiments on face image restoration tasks validate its efficacy.

strengths:
1.	The motivation is clearly stated. The main idea that considering image restoration problem using MSE optimization under perceptual constraint is interesting and essential.
2.	The experiments on face image restoration demonstrate superior performance over previous methods.

weaknesses:
1. The method seems like such a pipeline: Supervised restoration model trained using MSE loss + SDEdit method (with added $\sigma_s$ noise). Can authors provide further interpretations?
2. The experiments focus on specific the face image domain. The extension to other image domains (such as natural images) is unclear and unexplored.
3. The structure of the paper should be reorganized. The main body does not contain any quantitative results.

questions:
I suggest the authors add experiments on other image domains rather than face domain since the main topic is image restoration problem.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730367068883

thread_type:
official_review

<thread_object>
creator: reviewer_amj7

<content>
summary:
The work introduces Posterior-Mean Rectified Flow (PMRF) for image restoration, addressing the challenge of minimizing Mean Squared Error (MSE) while maintaining high perceptual quality. PMRF operates in two stages: first predicting the posterior mean, then using a rectified flow model to transport the result to match the ground-truth image distribution. The authors provide theoretical foundations showing that PMRF can achieve better MSE than posterior sampling methods while maintaining perfect perceptual quality under certain conditions. Through extensive experiments on various image restoration tasks (including blind face restoration, denoising, super-resolution, inpainting, and colorization), they demonstrate that PMRF consistently outperforms or matches state-of-the-art methods in both distortion metrics and perceptual quality measures.

strengths:
+ The paper provides comprehensive experiments and solid theoretical proof, supporting the effectiveness of the proposed method.
+ It is well-written, with clear and informative figures and tables that enhance the reader's understanding. 
+ The motivation behind the approach is convincing, and the experimental results substantiate the motivation to a significant extent, demonstrating the practical impact of the proposed framework.

weaknesses:
- While the method shows strong results in blind face image restoration, its evaluation is limited to this domain. Given the theoretical generality of the proposed framework, it would be valuable to see experiments on more diverse and general image datasets to assess the broader applicability of PMRF. Extending their experimental evaluation to include general image restoration would demonstrate the versatility of their method.
- Including ablation studies on the rectified flow model, posterior mean prediction, noise levels, and other architectural choices would provide deeper insights into the importance of each component and help the reader understand which parts of the framework drive the performance improvements.

questions:
If I understand correctly, the idea proposed by the authors should also be effective on general images, beyond the specific tasks tested. It would be interesting to know if the authors have considered conducting experiments or making attempts on general image datasets to further validate the versatility of their method. Of course, the paper is still completely acceptable without these additional experiments. However, including them would significantly enhance the impact of the work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645284141

thread_type:
official_review

<thread_object>
creator: reviewer_pkin

<content>
summary:
This paper presents a novel approach to photo-realistic image restoration (PIR) that aims to minimize mean squared error (MSE) while maintaining perceptual quality. The authors introduce the Posterior-Mean Rectified Flow (PMRF) algorithm, which optimally transports the posterior mean prediction to the distribution of ground-truth images. This method contrasts with traditional approaches that often prioritize perceptual quality at the expense of distortion. By leveraging a rectified flow model, PMRF effectively approximates the optimal estimator that minimizes MSE under a perfect perceptual index constraint. Experimental results demonstrate that PMRF outperforms existing methods across several restoration tasks.

strengths:
+ PMRF utilizes an optimal transport approach to align the posterior mean with the ground-truth distribution, enhancing the accuracy of image restoration.
+ The paper provides theoretical foundations that demonstrate PMRF’s ability to achieve lower MSE compared to traditional posterior sampling methods.
+ The authors conduct extensive experiments across multiple image restoration tasks, showcasing PMRF’s effectiveness and robustness compared to state-of-the-art techniques.

weaknesses:
- The recent FlowIE work also employs rectified flow for image restoration. I noticed that this paper references FlowIE in the supplementary material. It would be beneficial to include a discussion of this related work and highlight the differences in the main paper.
- In Figures 1 and 4, there are some unnaturally abrupt reflective spots or areas on the forehead, inner eye corners, and cheeks. Are these artifacts caused by the Posterior-Mean Rectified Flow? Please explain the reason behind these effects.

questions:
There are noticeable reflective spots or areas on the forehead, inner eye corners, and cheeks. Are these artifacts a result of the PMRF ? Please clarify the cause of these effects. Additionally, are there any limitations to the proposed method? Could this design be adapted for more general image restoration tasks beyond photo-realistic images?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688010112

thread_type:
official_review

<thread_object>
creator: reviewer_zdeq

<content>
summary:
This paper starts from the distortion-perception tradeoff in image restoration. It first generates the posterior mean and then transports the result to a high quality image following the ground-truth image distoration. The first stage is done by a traditional end-to-end model SwinIR trained with L1 loss, while the second stage is implemented with a rectified flow model by approximating the desired optimal transport map.

strengths:
1, It minimizes the MSE under a perfect perceptual index constraint and achieves a smaller MSE (higher fidelity) than posterior sampling methods (e.g., diffusion-based methods directly restoraing the image from noises).

2, With such a two-stage model, it could be easy-to-train (training a L1 loss-based model is stable) and fast-to-inference (both of the two stages are few-step models).

3, This paper is well-motived, well-written, well-supported and easy-to-understand.

weaknesses:
1, What’s the performance of the baseline model trained the L1 loss? Will it achieve the best PSNR?

2, In addition to distortion metrics, could you please the model’s ability for human face identity perservation?  

3, What are the inference speed, memory consumption and FLOPs?

4， The core idea of generating a L1 loss result and then transport it to visually-pleasing image is somewhat similar to the idea of [1] (using the normalizing flow to generate high-quality images conditional the low-quality input). What are their relations?

[1] Hierarchical Conditional Flow: A Unified Framework for Image Super-Resolution and Image Rescaling

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1732222861613

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of paper revision

comment:
We sincerely thank the reviewers for their efforts to review our paper and provide valuable feedback. In response to the reviews, we have added new experiments on the ImageNet dataset (discussed in Section 5.2). To facilitate your review of the revised manuscript, we summarize the changes below:

1. **Section 5.1.1**: We now clarify the difference between PMRF (our method) and FlowIE in L360.
2. **Section 5.2**: We have slightly revised this section to include details about the new ImageNet experiments.
3. **Figure 4**: The primary quantitative results of the new ImageNet experiments are now presented alongside the results of the face image restoration experiments, where we also include visual examples. Additionally, **Figure 5 in the appendix** extends these results by evaluating the algorithms with varying numbers of flow inference steps (as we did before for the face image restoration experiments).
4. **Appendix C**: Comprehensive details of the ImageNet experiments have been added.
5. **Table 12 in the appendix**: We updated this table to reflect the training hyper-parameters used in the new experiments, and included additional information such as memory consumption and GFLOPs per forward pass.
6. **New Figures 14–16 in the appendix**: These figures include additional visual examples of the ImageNet experiments.

We truly hope this summary of changes will streamline your review process and facilitate your evaluation of the updates.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The reviewers found the work with good theoretical background and validation. Author rebuttal addressed most of their concerns. Reviewers rJJP and PKin were also satisfied with the author responses but did not change the original score. Given that all reviewers were positive and convinced after the rebuttal, the paper is recommended for acceptance.