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

Official response

We thank the reviewer for their positive feedback on our paper, noting that it is “clear, interesting, and essential.” We address the reviewer’s concerns and questions below. ## Relation to SDEdit Please let us clarify two key distinctions between our PMRF method and SDEdit: (1) SDEdit uses a diffusion model which is trained to transport **white noise** to clean images. In contrast, we use a flow model trained to transport the **posterior mean predictions** (MMSE estimates) to clean images. (2) When relying on a diffusion model that goes from white noise to clean images (as in SDEdit), adding a bit of noise to the image is like going midway into the diffusion process. In our case, it is not the same as going midway into the flow process from MMSE to images: the noise we add is supposed to solve a technical problem, as we explain in Sections 3 and 6. ## Scope of experiments on the face image domain While blind face image restoration remains a highly challenging task, we agree that evaluating PMRF on more diverse datasets (e.g., ImageNet) would better showcase its versatility. Therefore, we have added additional experiments in Section 5.2, where we trained PMRF and the other methods on the ImageNet dataset to solve various tasks: denoising, super-resolution, and colorization. Notably, we observed a consistent trend across both the ImageNet and face restoration tasks: PMRF outperforms all prior frameworks on the distortion-perception plane. We have revised the main text to incorporate these additional experiments (Section 5.2 and Appendix C). We kindly invite you to review our updated results. Thank you for this valuable suggestion! ## No quantitative results in the main body We respectfully clarify that we do provide quantitative results in the main text, for example Table 1, Figure 3, and Figure 4. Figures 3 and 4 illustrate quantitative results on the perception-distortion plane through plots rather than tables, as this visualization offers an intuitive view of the perception-distortion tradeoff. This layout aims to provide readers with a clear understanding of our method’s performance compared to previous methodologies.

Official response

We sincerely thank the reviewer for describing our work as “comprehensive, well-written, and convincing.” We address the noted weaknesses and questions below: ## Evaluation on general image datasets While blind face image restoration remains a highly challenging task, we agree that evaluating PMRF on more diverse datasets (e.g., ImageNet) would better showcase its versatility. Therefore, we have added additional experiments in Section 5.2, where we trained PMRF and the other methods on the ImageNet dataset to solve various tasks: denoising, super-resolution, and colorization. Notably, we observed a consistent trend across both the ImageNet and face restoration tasks: PMRF outperforms all prior frameworks on the distortion-perception plane. We have revised the main text to incorporate these additional experiments (Section 5.2 and Appendix C). We kindly invite you to review our updated results. Thank you for this valuable suggestion! ## Ablation studies on key components We agree that ablation studies on components such as the rectified flow model, posterior mean prediction, and noise levels would offer valuable insights. Due to resource constraints and the extensive set of experiments already conducted, we prioritized core evaluations.

Official response

We thank the reviewer for highlighting strengths in our work, including our contributions to "enhancing the accuracy of image restoration, providing theoretical foundations, and conducting extensive experiments across multiple image restoration tasks." We address the reviewer’s points in detail below: ## Discussing FlowIE in the main text Thank you for the great suggestion to highlight the differences between PMRF (our method) and FlowIE in the main text. As we mention in the appendix (L1031), FlowIE is a conditional method that uses a ControlNet (similarly to DiffBIR), which distinguishes it from our PMRF approach. Specifically, FlowIE falls under the category of methods that attempt to sample from the posterior distribution, which differs fundamentally from PMRF’s optimal transport-based approach. We revised the main paper to emphasize this distinction (see L360 in the revised paper). Unfortunately, FlowIE does not currently provide a public checkpoint, so a direct comparison was not possible. ## Clarification on artifacts Regarding the artifacts noted in Figures 1 and 4, we carefully examined the reconstructed images and could not identify specific areas with the reflective spots mentioned (forehead, inner eye corners, and cheeks). If possible, could you share a highlighted example of these artifacts for further clarification? As with most generative models, minor artifacts are not uncommon and may appear in certain reconstructions. ## Discussing limitations Kindly note that we discuss several limitations in Section 6. For instance, our method requires tuning a hyper-parameter $\sigma_s$, which is the standard deviation of the noise added to the posterior mean predictor in the initial PMRF stage. In contrast, posterior sampling methods generally do not require this hyper-parameter adjustment. ## Could PMRF be adopted for more general image restoration? Thank you for pointing out this interesting question. We added new ImageNet experiments to the revised paper (see our new Section 5.2). Importantly, we see that PMRF still dominates the other baseline frameworks on the distortion-perception plane, achieving better distortion with either better or on-par perceptual quality. Namely, these experiments demonstrate the applicability of PMRF to general-content images. While our current experiments are limited to the natural image domain, we believe that PMRF can also be adopted for other image modalities, such as medical imaging. We hope to see future research exploring PMRF’s potential in such domains.

Official response

We would like to thank the reviewer for recognizing our work as "well-motivated, well-written, well-supported, and easy-to-understand." Below, we address the reviewer's questions and points for improvement: ## Posterior mean predictor model performance Kindly note that we report the performance of the baseline models (posterior mean predictors) in Figure 4 in the main text (the red-star ``posterior-mean'' in the controlled experiments) and in Tables 7-10 in the appendix (blind image restoration task). As the reviewer anticipates, these models achieve the best MSE (equivalently, PSNR) in all experiments. However, due to the perception-distortion tradeoff, they exhibit lower perceptual quality. Moreover, we would like to clarify that these models are trained to minimize the MSE (L2) loss, as our goal in the first stage of PMRF is to predict the posterior mean $\mathbb{E}[X|Y]$, which minimizes the MSE. While training a model with the L1 loss (Mean Absolute Error) often achieves similarly high PSNR, we selected the L2 loss to align with our theoretical framework and objective. ## Identity preservation metric Please note that we report the identity preservation metric for the blind face image restoration task in Table 1, denoted as "Deg" - this refers to the embedding angle of ArcFace. Our model achieves the second-best score in identity preservation, demonstrating competitive performance in maintaining face identity compared to previous methods. ## Inference Speed, Memory Consumption, and FLOPs. We report below the inference speed, GFLOPs and memory consumption of the models from all experiments: ### Blind image restoration 1. Inference speed/forward pass: 33 milliseconds for the vector field (HDiT model), 38 milliseconds for the posterior mean predictor (SwinIR model). The results are averaged over 1000 forward passes. 2. GFLOPs/forward pass: 100.67 for the vector field (HDiT model), 86.8 for the posterior mean predictor (SwinIR model). 3. Memory consumption: 612MB for the vector field, 60MB for the posterior mean predictor. ### Face image restoration experiments in Section 5.2 1. Inference speed/forward pass: 20 milliseconds for the vector field (HDiT model), 33 milliseconds for the posterior mean predictor (SwinIR model). The results are averaged over 1000 forward passes. 2. GFLOPs/forward pass: 44.83 for the vector field (HDiT model), 10.4 for the posterior mean predictor (SwinIR model). 3. Memory consumption: 463MB for the vector field, 18MB for the posterior mean predictor. ### ImageNet restoration experiments in Section 5.2 (newly added for the rebuttal - see the revised paper pdf) 1. Inference speed/forward pass: 19 milliseconds for the vector field (HDiT model), 19 milliseconds for the posterior mean predictor (HDiT model). The results are averaged over 1000 forward passes. 2. GFLOPs/forward pass: 11.21 for the vector field (HDiT model), 11.21 for the posterior mean predictor (HDiT model). 3. Memory consumption: 463MB for the vector field, 463MB for the posterior mean predictor. We added these details to the appendix (see Table 12 in the paper revision). Thanks for the suggestion! ## Relation to Hierarchical Conditional Flow (HCF) As we understand, HCF is a conditional normalizing flow model trained to generate samples from the posterior distribution. While HCF uses an L1 regularization term to center its outputs around the posterior mean (as noted by the authors of HCF, after Equation 7 in their paper), our approach differs fundamentally by focusing on directly transporting the posterior mean to the target distribution, rather than sampling around it. This superficial difference leads to the compelling theoretical and practical results of our method.