I2I-PR: Deep Iterative Refinement for Phase Retrieval using Image-to-Image Diffusion Models

TL;DR

I2I-PR tackles phase retrieval by combining a robust initialization strategy with a learned image-to-image diffusion refinement under a physics-informed framework. By starting from multiple classical initial estimates and using InDI-based denoising, followed by data-consistency steps and geometric self-ensemble, the method achieves higher reconstruction quality and training efficiency than traditional solvers and recent diffusion-based approaches. The work also provides uncertainty quantification via ensembles and demonstrates the benefits of aggregation for perceptual and distortion metrics, with strong generalization to unseen data. Overall, the approach offers a practical, robust, and scalable solution for phase retrieval across diverse applications, with code and models publicly available.

Abstract

Phase retrieval aims to recover a signal from intensity-only measurements, a fundamental problem in many fields such as imaging, holography, optical computing, crystallography, and microscopy. Although there are several well-known phase retrieval algorithms, including classical alternating projection-based solvers, the reconstruction performance often remains sensitive to initialization and measurement noise. Recently, diffusion models have gained traction in various image reconstruction tasks, yielding significant theoretical insights and practical advances. In this work, we introduce a deep iterative refinement framework that redefines the role of diffusion models in phase retrieval. Instead of generating images from random noise, our method starts with multiple physically consistent initial estimates and iteratively refines them through a learned image-to-image diffusion process. This enables data-driven phase retrieval that is both interpretable and robust, leveraging the strengths of classical solvers while mitigating their weaknesses. Furthermore, we propose an enhanced initialization strategy that integrates classical algorithms with a novel acceleration mechanism to obtain reliable initial estimates. During inference, we adopt a geometric self-ensemble strategy based on input flipping, together with output aggregation to further improve the final reconstruction quality. Comprehensive experiments demonstrate that our approach achieves substantial gains in both training efficiency and reconstruction quality, consistently outperforming classical and recent state-of-the-art methods. These results highlight the potential of diffusion-driven refinement as an effective and general framework for robust phase retrieval across diverse applications. The source code and trained models are available at https://github.com/METU-SPACE-Lab/I2I-PR-for-Phase-Retrieval

Figures (10)

• Figure 1: The overall pipeline of I2I-PR. The method begins with an initialization stage combining HIO and AER to generate multiple candidate solutions. These are refined through an iterative image-to-image denoising process that incorporates data consistency with classical HIO updates.
• Figure 2: An example of the defined gradual degradation process used during training. The timestep increases from left to right. The rightmost image ($t=7$) corresponds to the output of the initialization procedure, and the leftmost image ($t=0$) corresponds to the clean image.
• Figure 3: Architecture of the used denoiser with multiple input images and a timestep input, producing a single denoised image. The network follows a UNet structure. Following the conventions defined in diffusers, this architecture first utilizes a downsampling path composed of DownBlock2D layers that progressively reduce spatial resolution while increasing the number of channels and enhancing feature extraction. AttnDownBlock2D layers incorporate attention to focus on relevant spatial information. At the bottleneck, the UNetMidBlock2D refines the lowest resolution features before upsampling begins. The upsampling path mirrors the downsampling, with UpBlock2D and AttnUpBlock2D layers restoring spatial resolution with attention-based refinement. Skip connections between corresponding downsampling and upsampling blocks help preserve spatial information across resolutions. The final denoised image is obtained by subtracting the predicted noise (residual) from the mean of the input images.
• Figure 4: Geometric interpretation of the acceleration mechanism used during the ER phase in the initialization procedure. The mechanism employs momentum-like updates between successive projections to enhance convergence speed and stability.
• Figure 5: Augmentation of the initialization outputs with an equivariant transform.