Alternating Phase Langevin Sampling with Implicit Denoiser Priors for Phase Retrieval

Phase retrieval is the nonlinear inverse problem of recovering a true signal from its Fourier magnitude measurements. It arises in many applications such as astronomical imaging, X-Ray crystallography, microscopy, and more. The problem is highly ill-posed due to the phase-induced ambiguities and the large number of possible images that can fit to the given measurements. Thus, there's a rich history of enforcing structural priors to improve solutions including sparsity priors and deep-learning-based generative models. However, such priors are often limited in their representational capacity or generalizability to slightly different distributions. Recent advancements in using denoisers as regularizers for non-convex optimization algorithms have shown promising performance and generalization. We present a way of leveraging the prior implicitly learned by a denoiser to solve phase retrieval problems by incorporating it in a classical alternating minimization framework. Compared to performant denoising-based algorithms for phase retrieval, we showcase competitive performance with Fourier measurements on in-distribution images and notable improvement on out-of-distribution images.


INTRODUCTION
Phase retrieval is the reconstruction problem of obtaining lost phase information of a signal given only its Fourier magnitude. It arises in many physical measurement applications including X-ray crystallography, transmission electron microscopy, astronomical imaging, and more. More formally, the phase retrieval problem asks to recover a signal x * ∈ R n from phaseless measurements of the form where A ∈ C m×n is the measurement matrix and w is noise.
In the Fourier case, A is replaced with F, the Fourier transform, that operates on x * . Due to the phase ambiguity and noise present in the problem, phase retrieval is an ill-posed problem. There are many possible images that can fit to the given measurements, most of them lacking natural structure. Image priors combat this by constraining the space of possible solutions to images that do contain the structure that the object of interest is believed to obey. There is a rich history of image priors across computational imaging, and we highlight previous priors utilized in phase retrieval.
A common structural prior traditionally enforced in imaging applications is sparsity with respect to a wavelet basis. Several works have considered the use of sparsity priors in phase retrieval to improve reconstruction and yield more efficient algorithms [1]. However, sparsity priors are limiting as they inherently enforce a linear signal model, which may potentially lack the richness to model more complicated data of interest. On the learning side, there has been tremendous progress on priors in the form of generative models given by a deep neural network. For example, the works [2,3] have shown that generative networks can potentially better represent natural signals compared to sparsity priors and also demonstrate strong recovery performance. One drawback of this approach, however, is that generative models require an abundance of clean data to be trained. Moreover, such generators are often tailored specifically to a data distribution of interest, and struggle to reconstruct out-of-distribution images. These points ultimately limit their use in practical applications.
An interesting class of (implicit) priors are in the form of neural network-based denoisers, which are models that are trained remove additive noise from images. Denoisers implicitly learn natural image structure through many paired examples of noisy and clean images. Such models have shown impressive performance as regularizers in several linear inverse problems [4]. Additionally, these models have shown strong recovery performance as regularizers in phase retrieval and showcase the ability to recover images outside of their training distribution [5]. Recently, the work of [6] showed that it was possible to sample from the prior implicit in a denoiser, and demonstrated strong performance in using this prior to help solve linear inverse problems.
Our contributions: In this work, we propose an algorithm called Alternating Phase Langevin Sampling (APLS) to solve phase retrieval problems by combining classical alternating projection methods with modern denoiser priors. Our algorithm incorporates a posterior sampling method to access the implicit prior of a denoiser in an alternating minimiza-tion framework. Then, we empirically study our approach in the challenging Fourier measurement regime and showcase competitive results for in-distribution images and significantly outperform other approaches on out-of-distribution images. Our work showcases robustness and generalization performance for such priors, and highlights the potential for incorporating implicit denoiser priors more broadly for nonlinear inverse problems.

RELATED WORK
Phase retrieval methods can be classified as either convex or non-convex approaches. Convex techniques include semidefinite programming methods, such as PhaseLift [7], and methods based on linear programming, such as PhaseMax [8]. Non-convex methods for phase retrieval generally fall under the alternating projections framework, originating from the Gerchberg-Saxton [9] and Fienup [10] algorithms, or direct first-order methods, such as Wirtinger Flow [11]. We now detail some of the alternating minimization approaches that our work is inspired by.

Alternating Minimization
Alternating minimization provides an optimization framework for solving nonlinear inverse problems [12]. In particular, the nonlinear system b = |Ax * | can be reduced to a linear system if we had access to the true phase p of Ax * . However, p is not actually known and so the solution to the linear system is obtained by solving arg min Alternating minimization then alternates between updating the phase p and the solution x to iteratively minimize (2). Denoting the phase of z ∈ C by Ph(z) := z |z| , the phase p t and signal x t are updated via the following iteration for each t 1: The authors of [12] showcase geometric convergence of the algorithm with proper initialization and establish sample complexity guarantees for Gaussian measurements.

Denoisers for Regularization
Denoisers were first used as regularizers in solving inverse problems by [13]. This method is termed plug-and-play regularization because a denoiser can be "plugged in" as a regularizer when solving an inverse problem to produce more accurate recoveries [4].
The current state-of-the-art method for solving phase retrieval using a denoiser as a prior is prDeep [5]. They use the Regularization by Denoising framework (RED) [14] that explicitly defines the regularizer as where D(x) is the denoiser. For the denoiser, a powerful neural network-based denoiser DnCNN [15] is used.

Prior Implicit in a Denoiser
Why are denoisers such strong priors for inverse problems? When training to remove additive noise, denoisers must develop an understanding of what the important properties of structured, natural images are. This implicit prior embedded within denoisers was exposed by [6] through Empirical Bayes estimation and was used to achieve competitive performance in solving linear inverse problems. Specifically, it was shown in [6] that there is an interesting connection [16] between the score of the density p(y) induced by images contaminated with Gaussian noise y = x + z with variance σ 2 and MMSE denoisers: wherex(y) = E[x|y] is the MMSE estimate, which can be approximated by a denoiser trained on a large paired dataset of clean and noisy images. Based on this observation, [6] proposed an iterative algorithm that samples from a denoiser to take a step in the direction specified by the observation density gradient and gradually reduce the effective noise σ to eventually converge to a high-probability region of p(x). To achieve this, [6] employed a Langevin sampling algorithm derived from Langevin Dynamics to inject a decreasing amount of noise into the gradient ascent procedure.
Since the effective noise is decreasing, the denoiser must be able to adjust to the amount of noise as the stochasticity of the gradient changes. Most denoisers such as the one used in prDeep, are trained to only handle a specific noise value. Instead, the algorithm of [6] uses a bias-free CNN denoiser. By removing all the additive bias terms of the convolutional layers, [17] discovered that the resulting denoiser would become robust to the amount of additive noise.
Finally, this Langevin sampling algorithm was used to solve linear inverse problems such as deblurring, superresolution, and inpainting. Updating the gradient step to handle the conditional density p(x|Ax = Ax * ) allows the algorithm to sample the denoiser in a way that also fits to linear measurements Ax = Ax * .

APPROACH
We aim to incorporate an implicit denoiser prior into an efficient algorithm to solve phase retrieval problems. The de-  noiser prior has learned universal properties of natural images, and, thus, could potentially provide strong generalization guarantees to different datasets. One issue, however, is that the approach of [6] only allowed for sampling from the posterior conditioned on linear measurements p(x|Ax = Ax * ), rendering an extension to nonlinear inverse problems challenging. Inspired by alternating projection approaches, we propose an alternating minimization-sampling procedure that 1) fixes a candidate phase p t and then 2) updates the signal by sampling from the posterior x t ∼ p(x|Ax = p t b). We first describe the Langevin sampling algorithm for sampling from the posterior of the image conditioned on the linear part of the measurements.

Langevin Sampling
As noted in [6], the residual of the denoiser D(x) :=x(y)−y is proportional to the gradient of log observation density. This result can be utilized to perform the following update steps to approximately sample from the conditional density p(x|Ax = x c ), as outlined in Algorithm 1. The main update step is where hd x is the gradient step with step size h and γz is injected noise. The quantity γ controls the proportion of injected noise z ∼ N (0, I) and is given by where β ∈ (0, 1] is a hyperparameter. σ 2 = d 2 2 N (N is the size of the image) adjusts the amount of injected noise in accordance to the size of the gradient step. The gradient step, d x , is and uses D, the residual of the denoiser, to approximately sample from the conditional distribution p(x|Ax = x c ). See Algorithm 1 for the full details. Note that the quantity h t = h0t 1+h0(t−1) acts as an accelerator for the algorithm by increasing the step size and decreasing the amount of injected noise over time.
Algorithm 1: Posterior Sampling via the Implicit Denoiser Prior Data: We use the Alternating Minimization framework to decouple the non-linearity of the phase retrieval problem to incorporate a denoiser prior with Langevin Sampling. Specifically, we perform an alternating minimization-sampling algorithm that first fixes the phase and then approximately samples from the posterior conditioned on the phase-induced measurements by utilizing the implicit denoiser prior. That is, for initial points p 0 and x 0 , we compute for t 1, Additionally, we enforce the gradient of the Langevin Sampling to account for the phase p t by providing it the measurements x c = p t b. We outline more details in Algorithm 2.

EXPERIMENTAL RESULTS
The classical phase retrieval problem is to recover true signal x * from measurements b = |F(x * )| + w. Poisson shot noise is a common noise model that appears in many phase retrieval applications [18]. We approximate this as where α is a parameter that increases with the noise level. We use 4× over-sampled Fourier measurements by placing an n × n image in the center of a 2n × 2n square and taking its 2D Fourier transform. The location of the image in this square is the support which we assume is known a priori. As in [5], for the initial input x 0 we use the output of the Hybrid Input-Output (HIO) algorithm [10]. We take 50 random initializations x 1 , x 2 , . . . x 50 and run HIO for 50 iterations. We take the x i with the lowest residual b − |F(x i )| 2 and run HIO for another 1000 iterations with the initial iterate x i . The result is used as the input x 0 for APLS. Our algorithm is initialized with h 0 = 0.1, β = 0.0001, and T 2 = 500. Empirically, we achieved best results when the number of iterations to run the Langevin Sampling algorithm between phase updates was just T 1 = 1.
We compare to prDeep which is the state-of-the-art using the regularization by denoising framework. We use 4 DnCNN denoisers trained on the Berkeley Segmentation Dataset with noise standard deviations of 50, 25, 15, and 5 to implement the pipeline used in the original paper.
We evaluate the reconstructions of HIO, prDeep, and our algorithm on a set of 6 natural grayscale images (Figure 1) and a set 6 unnatural grayscale images (Figure 2). We calculate the Peak Signal to Noise Ratio, PSNR = 10 log 10 ( 255 2 mean((x−x * )) ), to assess the reconstruction quality between reconstructed imagex and original image x * . Table  1 displays the average PSNR across both test sets. We see that our method is comparable on natural images and we showcase considerable improvement for unnatural images over all noise levels.

CONCLUSION
We proposed the Alternating Phase Langevin Sampling (APLS) algorithm for phase retrieval, which incorporates an implicit denoiser prior in an alternating minimization framework. We demonstrated competitive performance with prDeep, a strong baseline incorporating a denoiser as a regularizer for phase retrieval. For natural images, our performance was comparable while for unnatural, out-ofdistribution images, we showcased superior performance. This work showcases the benefits of extracting the prior in a denoiser more directly. There are many interesting avenues of future work, including applying our approach to phase retrieval problems in more scientific domains, such as astronomical imaging, and theoretically analyzing the convergence and sample complexity of our method.