prNet: Data-Driven Phase Retrieval via Stochastic Refinement

TL;DR

This work tackles the ill-posed phase retrieval problem by introducing prNet, a Langevin-dynamics–based framework that samples from the posterior $p(oldsymbol{x}|oldsymbol{y})$ to balance distortion with perceptual realism. It advances three architectures—prNet-Small, prNet-Large, and prNet-Large-Adversarial—each refining classical HIO-based estimates with learned denoisers and data-consistency steps, and itFurther leverages warm-start initialization and test-time augmentation. Empirical results on Fourier phase retrieval show state-of-the-art fidelity and perceptual quality, with robustness to initialization and noise and added improvements from adversarial aggregation. The approach offers a principled, scalable stochastic solver for nonlinear inverse problems and provides code and models for broader adoption in computational imaging.

Abstract

Phase retrieval is an ill-posed inverse problem in which classical and deep learning-based methods struggle to jointly achieve measurement fidelity and perceptual realism. We propose a novel framework for phase retrieval that leverages Langevin dynamics to enable efficient posterior sampling, yielding reconstructions that explicitly balance distortion and perceptual quality. Unlike conventional approaches that prioritize pixel-wise accuracy, our methods navigate the perception-distortion tradeoff through a principled combination of stochastic sampling, learned denoising, and model-based updates. The framework comprises three variants of increasing complexity, integrating theoretically grounded Langevin inference, adaptive noise schedule learning, parallel reconstruction sampling, and warm-start initialization from classical solvers. Extensive experiments demonstrate that our methods achieve state-of-the-art performance across multiple benchmarks, both in terms of fidelity and perceptual quality. The source code and trained models are available at https://github.com/METU-SPACE-Lab/prNet-for-Phase-Retrieval

Figures (9)

• Figure 1: The overall pipeline of prNet-Large-Adversarial. Our approach begins with the Initialization Stage, where $m$ random initializations are refined using HIO for $s$ steps. The top $k$ candidates with the lowest residuals are selected and further refined with $n$ additional HIO iterations to produce initial estimates $\mathbf{x}_0'$. In the Main Loop, each estimate is stochastically perturbed with Gaussian noise and iteratively refined using a combination of classical HIO updates and a learned denoiser $D_\theta$. This process is repeated for $T$ iterations. In the Final Stage, the refined outputs are passed through a learned denoiser $D_\phi$, trained adversarially via a critic model to enhance realism and perceptual quality. Compared to prNet-Large, which uses simple averaging in the Final Stage, prNet-Large-Adversarial incorporates a learned denoiser for aggregation. Additionally, prNet-Large-Adversarial refines multiple reconstructions in parallel, while prNet-Small operates on a single initialization throughout.
• Figure 2: Architecture of the UNet denoiser with timestep input. For prNet-Small, the input and output are single images, as illustrated. In the main loop denoiser of prNet-Large, the network processes multiple input images and produces multiple output images.
• Figure 3: Test time augmentation (TTA): We execute the full pipeline on both the original initialization outputs and their flipped versions, then average the results to produce the final output.
• Figure 4: Test time augmentation using dihedral group $D_4$ (TTA $D_4$).
• Figure 5: The histograms of PSNR (left column) and SSIM (right column) for the reconstructions produced by various methods across 236 test images and 5 Monte Carlo runs for the $\alpha=3$ scenario. Vertical dashed lines indicate the mean PSNR and SSIM values.