Large-scale phase retrieval

High-throughput computational imaging requires efficient processing algorithms to retrieve multi-dimensional and multi-scale information. In computational phase imaging, phase retrieval (PR) is required to reconstruct both amplitude and phase in complex space from intensity-only measurements. The existing PR algorithms suffer from the tradeoff among low computational complexity, robustness to measurement noise and strong generalization on different modalities. In this work, we report an efficient large-scale phase retrieval technique termed as LPR. It extends the plug-and-play generalized-alternating-projection framework from real space to nonlinear complex space. The alternating projection solver and enhancing neural network are respectively derived to tackle the measurement formation and statistical prior regularization. This framework compensates the shortcomings of each operator, so as to realize high-fidelity phase retrieval with low computational complexity and strong generalization. We applied the technique for a series of computational phase imaging modalities including coherent diffraction imaging, coded diffraction pattern imaging, and Fourier ptychographic microscopy. Extensive simulations and experiments validate that the technique outperforms the existing PR algorithms with as much as 17dB enhancement on signal-to-noise ratio, and more than one order-of-magnitude increased running efficiency. Besides, we for the first time demonstrate ultra-large-scale phase retrieval at the 8K level (7680$\times$4320 pixels) in minute-level time.

Wide field of view and high resolution are both desirable for various imaging applications, such as medical imaging [1][2][3][4] and remote sensing 5 , providing multi-dimensional and multi-scale target information. As the recent development of computational imaging, large-scale detection has been widely employed in a variety of computational imaging modalities 3,4,6,7 . These computational imaging techniques largely extend the spatial-bandwidth product (SBP) of optical systems from million scale to billion scale. As an example, the SBP of the RUSH microscope platform 4 has reached as high as 1.7 × 10 8 . Such large amount of data poses great challenge for post software processing. Therefore, large-scale processing algorithms with low computational complexity and high fidelity are of great significance for those imaging and perception applications in various dimensions 8 .
In computational phase imaging, phase retrieval (PR) is required to reconstruct both amplitude and phase in complex space from intensity-only measurements. This problem originates from the limitation of low response speed of photodetectors that impedes direct acquisition of light wavefront. Mathematically, the underlying goal of PR is to estimate an unknown complexfield signal from the intensity-only measurements of its complex-valued transformation, which is described as where u is the underlying signal to be recovered (u ∈ C n×1 ), I contains the intensity-only measurements (I ∈ R m×1 ), A represents measurement matrix (A ∈ R n×n or C n×n ), and ω stands for measurement noise. Phase retrieval has been widely applied in plenty fields such as astronomy, crystallography, electron microscopy and optics 9 . It solves various nonlinear inverse problems in optical imaging, such as coherent diffraction imaging 10 (CDI), coded diffraction pattern imaging 11 (CDP), Fourier ptychographic microscopy 3 (FPM) and imaging through scattering medium 12 .
In the past few decades, different phase retrieval algorithms have been developed. Gerchberg and Saxton pioneered the earliest alternating projection (AP) algorithm in the 1970s 13 , which was then extended by Fienup et al. with several variants 14 . Due to its strong generalization ability, AP has been widely employed in multiple phase imaging models. Nevertheless, it is sensitive to measurement noise, suffering from poor noise robustness. Afterwards, researchers introduced optimization into PR, deriving a series of semi-definite programming (SDP) based algorithms 15,16 and Wirtinger flow (WF) based algorithms [17][18][19] . These techniques enhances robustness to measurement noise, but they require high computational complexity and high sampling rate, making them inapplicable for large-scale phase retrieval. Although the sparsity prior of natural images in transformed domains can be incorporated as an additional constraint to lower sampling rate [20][21][22] , it further increases computational complexity.
Recently, the booming deep learning (DL) technique has also been introduced for phase retrieval 23 . Following the large-scale training framework, the DL strategy outperforms the above traditional PR techniques with higher fidelity. However, it provides poor generalization that each suits only for specific models, such as holography 23 and FPM 24 . For different models and even different system parameters, the deep neural network requires to be retrained with new large-scale data sets. To sum, despite of different workflows, the above existing PR algorithms suffer from the tradeoff among low computational complexity, robustness to measurement noise and strong generalization, making them inapplicable for large-scale phase retrieval.
In this work, we report an efficient large-scale phase retrieval technique termed as LPR, as sketched in Fig. 1. It builds on the plug-and-play (PNP) 25 optimization framework, and extends the efficient generalized-alternating-projection (GAP) 8,26,27 strategy from real space to nonlinear complex space. The complex-field PNP-GAP scheme ensures strong generalization of LPR on various imaging modalities, and outperforms the conventional first-order PNP techniques (such as ISTA 28 and ADMM 25 ) with fewer auxiliary variables, lower computational complexity and faster convergence. As PNP-GAP decomposes reconstruction into separate sub-problems including measurement formation and statistical prior regularization 8,29 , we further introduce an alternating projection solver and an enhancing neural network respectively to solve the two sub-problems. These two solvers compensate the shortcomings of each other, allowing the optimization to bypass the poor generalization of deep learning and poor noise robustness of AP. As a result, LPR enables generalized large-scale phase retrieval with high fidelity and low computational complexity, making it a state-of-the-art method for various computational phase imaging applications.
We compared LPR with the existing PR algorithms on extensive simulation and experiment data of different imaging modalities. The results validate that compared to the AP based PR algorithms, LPR is robust to measurement noise with as much as 17dB enhancement on signle-to-noise ratio. Compared with the optimization based PR algorithms, the running time is significantly reduced by more than one order of magnitude. Finally, we for the first time demonstrated ultralarge-scale phase retrieval at the 8K level (7680×4320 pixels) in minute-level time, where most of the other PR algorithms failed due to unacceptable high computational complexity.

Results
We applied LPR and the existing PR algorithms on both simulation and experiment data of three computational phase imaging modalities including CDI, CDP and FPM, to investigate respective pros and cons. The competing algorithms for comparison includes the alternating projection technique (AP) 13,14 , the SDP based techniques (PhaseMax (PMAX) 30  AmpFlow (TAF), Reweighted AmpFlow (RAF)), Coordinate Descent (CD) 35 , KACzmarz (KAC) 36 and the deep learning based prDeep technique 22 . All the algorithm parameters were tuned based on the Phasepack 37 for respective best performance. The convergence is determined when the intensity difference of reconstructed image between two successive iterations is smaller than a preset threshold. We employed the peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM) 38 to quantify reconstruction quality. All the calculation was tested on a desktop PC with an Intel i7-9700 CPU, 16G RAM and an Nvidia GTX 1660s GPU.
Coherent diffraction imaging. CDI is a representative non-interferometric phase imaging technique, and has been widely applied in physics, chemistry and biology due to its simple setup 9 .
It illuminates a target using coherent plane wave, and records the intensity of far-field diffraction pattern. By oversampling the diffracted light field and applying phase retrieval, both the target's amplitude and phase information can be reconstructed. Mathematically, the measurement formation of CDI is where u denotes the target information, and F represents the Fourier transformation that approximates the far-field diffraction.
Following the above formation model, we employed a high-resolution image (1356×2040 pixels) from the DIV2K 39 dataset as the latent signal to synthesize CDI measurements. Due to the uniqueness guarantee of solution, CDI requires at least 4 times oversampling in the Fourier domain 40 . Correspondingly, we padded zeros around the image matrix to generate a 2712×4080 image. We implemented Fourier transform to the image and retained only its intensity as measurements. Additionally, to investigate the techniques' robustness to measurement noise, we further added different levels of white Gaussian noise (WGN) to the measurements. Table 1 presents the quantitative reconstruction evaluation of different techniques. The results show that the CD and KAC methods failed with no convergence. This is because these techniques require higher sampling ratio. The PLIFT and PLAMP methods do not work as well, because they require matrix lifting and involve higher dimensional matrix that is out of memory in large-scale reconstruction. The other methods except for prDeep obtain little improvement compared to the AP algorithm. Specifically, the WF, AF and PMAX methods even degrade due to limited sampling ratio and noise corruption. The reconstruction of prDeep is better than the conventional algorithms, but with only 2dB enhancement on PSNR, and almost no SSIM improvement compared to AP. In contrast, LPR produces significant enhancement on reconstruction quality, with as much as 6dB and 0.29 improvement on PSNR and SSIM, respectively. Due to limited space, the detailed visual comparison of different techniques is presented in Fig. S1 (supplementary information), which coincides with the above quantitative results. We further compared these algorithms on experiment CDI data 41  Coded diffraction pattern imaging. CDP 11 is a coded version of CDI, which introduces wavefront modulation to increase observation diversity. The strategy of multiple modulations and acquisitions enables to effectively bypass the oversampling limitation of the conventional CDI. Generally, the target light field is additionally modulated by a spatial light modulator (SLM), and the measurements after far-field Fraunhofer diffraction can be modeled as where d represents the modulation pattern, and denotes the Hadamard product.
We simulated CDP measurements with five and single phase modulations, respectively. The modulation patterns d are subject to Gaussian distribution 11 . We employed the same image as CDI to be the ground-truth signal, and added various levels of WGN to the measurements. Table   2 where F −1 is inverse Fourier transform, P denotes system's pupil function, and S represents the wave function of incident light.
Following the formation model, we first implemented a simulation comparison with the following setup parameters: the wavelength is 625nm, the numerical aperture (NA) of objective lens is 0.08, the height from the light source to the target is 84.8mm, and the distance between adjacent light sources is 4mm. The pixel size of camera is 3.4µm. Two microscopy images of blood cells 42 (2048×2048 pixels) were employed as the latent high-resolution (HR) amplitude and phase, respectively. The size of captured low-resolution images (LR) was one fourth of the HR images. Ultra-large-scale phase retrieval. In ultra-large-scale imaging applications such as 4K (4096×2160 pixels) or 8K (7680×4320 pixels), most reconstruction algorithms are not applicable due to either highly large memory requirement or extremely long running time. Nevertheless, the reported LPR technique still works well in such applications. As a demonstration, we implemented a simulation of 8K-level CDP (5 modulations), using an 8K outer space color image as ground truth (released by NASA using the Hubble Telescope). Its spatial resolution is 7680×4320 (each color channel) with in total 33.1 million pixels. Figure 6 presents the reconstruction results of AP and LPR, with the input SNR being 5dB. The colse-ups show that the result of AP is drowned out by measurement noise, leading to dimness and loss of target details. In comparison, LPR outperforms a lot with strong robustness. Both their running time lie in the minute level. Another set of 8K reconstruction results is shown in Fig. S8 (supplementary information).

Methods
Following optimization theory, the phase retrieval task can be modeled aŝ where u denotes the target complex filed to be recovered, f (u) is a data-fidelity regularizer that ensures consistency between the reconstructed result and measurements, and g(u) is a regularizer that imposes certain statistical prior knowledge. Conventionally, Eq. (5) is solved following the first-order proximal gradient methods, such as ISTA and ADMM that are time-consuming to calculate gradients in large-scale nonlinear tasks 29 . In this work, instead, we employ the efficient generalized-alternating-projection (GAP) strategy 29 to transform Eq. (5) with fewer variables to where v is an auxiliary variable balancing the data fidelity term and prior regularization, A denotes measurement matrix, and I represents measurement. The difference between the conventional ADMM and GAP optimization is the constraint on the measurement 29 . ADMM minimizes I − |Au| 2 , while GAP imposes the constraint I = |Au| 2 .
To tackle the large-scale phase retrieval task, we extend the efficient plug-and-play (PNP) optimization framework 25 from real space to nonlinear complex space. Fundamentally, PNP decomposes optimization into two separate sub-problems including measurement formation and prior regularization, so as to incorporating inverse recovery solvers together with various image enhancing solvers to improve reconstruction accuracy, providing high flexibility for different applications.
Mathematically, Eq. (6) is decomposed into the following two sub-problems, to alternatively update the two variables u and v.
• Updating u: given v (k) , u (k+1) is updated via a Euclidean projection of v (k) on the manifold where P R is phase retrieval solver. Considering its great generalization ability on various imaging modalities and low computational complexity, we employ the AP method as the P R solver. It alternates between the target and observation planes allowing to incorporate any information available for the variables, providing low sampling rate requirement.
When the intensity difference of reconstructed image between two successive iterations is smaller than a given threshold, the iteration stops with convergence. Since both the two solvers P R and EN are highly efficient and flexible, the entire reconstruction maintains low computational complexity and great generalization. The demo code has been released at bianlab.github.io.

Conclusion and Discussion
In this work, we engaged to tackle the large-scale phase retrieval problem, and reported a generalized LPR optimization technique with low computational complexity and strong robustness.
It extends the efficient PNP-GAP framework from real space to nonlinear complex space, and incorporates the alternating projection solver and enhancing neural network. As validated by ex-