Deep Learning Initialized Phase Retrieval
Abstract
Phase retrieval (PR) consists of estimating 2D or 3D objects from their Fourier magnitudes and takes a central place in scientiﬁc imaging.
At present, most iterative methods for PR work well only when the initialization is close enough to the solution and fail otherwise.
But there has been no general way of obtaining desired initialization.
In this paper, we show that a carefully designed deep learning pipeline can consistently generate reliable initialization, so that the subsequent iterative methods can solve the PR problem and produce high-quality solutions.
Technically, PR is an inverse problem containing three forward symmetries, and naive deployment of end-to-end deep learning for PR yields poor initialization.
We explain why the symmetries cause the learning difﬁculty and propose a novel strategy that substantially improves the estimation.
Overall, the proposed method solves PR in regimes not accessible by the previous methods, and our work synergizes deep learning and iterative methods to solve a difﬁcult scientiﬁc inverse problem.
Introduction
Given the oversampled Fourier magnitudes Y = |F ( X ) | 2 ∈ R m 1 × m 2 of X ∈ C n 1 × n 2 ( m i ≥ n i for i = 1 , 2 ) is it possible to recover X ?
This is the phase retrieval (PR) problem central to scientiﬁc imaging and numerous other ﬁelds [1, 2].
In imaging, the complex phases of F ( X ) are missing because practical detectors cannot record complex phases.
As a nonlinear inverse problem, PR is tricky to solve: (1) for any X , global phase shift X (cid:55)→ e iθ X for all θ ∈ ( − π, π ] , overall spatial shift of the nonzero content of X and 2D conjugate ﬂipping of X all leave the measurement Y unchanged (see Fig.
So the best one can hope for is recovery up to these intrinsic symmetries; and (2) the mapping X (cid:55)→
|F ( X ) | 2 is generically injective up to the intrinsic symmetries when m i ≥ 2 n i − 1 [4] for i = 1 , 2 .
So in this paper, we always assume m i ≥ 2 n i − 1 for i = 1 , 2 and solving PR is up to the intrinsic symmetries.
In practice, PR is often solved by iterative methods such as HIO [5],
RAAR [6] and difference map [7]; see [8] for an updated review of these methods.
However, they work well when support (i.e., locations of nonzero elements) of X is provided with reasonable accuracy.
Otherwise, even the most sophisticated iterative methods fail to work in practice.
The failure is mostly due to stagnation caused by the intrinsic symmetries [9].
Deep learning has brought about new prospects of solving difﬁcult inverse problems, of which PR is an instance.
One can phrase PR as a regularized optimization problem: min X (cid:96) ( Y , |F
( X )
| 2 ) + λ Ω( X ) , and then replace (cid:96) , or Ω , or mappings inside speciﬁc numerical methods for PR using data-driven neural network modules.
This approach is taken in, e.g., [10,11], where HIO is needed to produce good initialization—for simpler inverse problems, such special initialization is not required [12].
More radical is the end-to-end approach , where a neural network is trained to directly approximate the inverse mapping or its proxies.
[13–16] have taken this approach and shown promising results.
Here, we take a critical view of the initial successes.
Difﬁculty of learning with symmetries When solving nonlinear inverse problems with symmetries, the end-to-end approach may face the difﬁculty of approximating highly oscillatory functions.
This issue has recently been elucidated in [17].
We summarize the main argument using the learning square root example: suppose we randomly sample real values x i ’s and form a training set (cid:8) x i , x 2 i (cid:9) and try to learn the square-root function using the end-to-end approach, allowing both positive and negative outputs.
Now if we think of the function determined by the training set, which the neural network is trying to approximate, it is highly oscillatory (see Fig.
2): the sign symmetry in the forward mapping x (cid:55)→ x 2 dictates that in the training set, there are frequent cases where x 2 i and x 2 j are close but x i and x j have different signs and are far apart.
Although in theory neural networks with adequate capacities are universal function approximators, in practice they will struggle to learn such irregular functions.
For general inverse problems, so long as the forward symmetries can relate remote inputs to the same output, such as all the three symmetries in PR, similar problems can surface.
images of simulated crystal structures in Bragg
CDI applications [18,19] (only real parts of complex-valued images are shown), which do not have any natural orientation or centering.
Biases in practical image datasets Strangely, the difﬁculty is only brieﬂy touched in [16] among all recent work applying the end-to-end approach to PR.
When we examine the training and test data that previous works use, it becomes clear that the issue is probably covered by intrinsic dataset biases.
Previous experiments typically use images from standard computer vision datasets such as
MNIST, ImageNet, CelebA(faces), where the image contents tend to be centralized and naturally oriented (see Fig.
3 (a)–(b)).
This helps break the shift and ﬂipping symmetries naturally, as these images are relatively close to each other compared to when mixed up with some of their symmetric copies.
1 Our analysis has been conﬁrmed in a preliminary version of this work, where we showed that augmenting natural datasets to account for symmetries fails a state-of-the-art method that performs well without the augmentation [3].
In short, the fundamental difﬁculty has been concealed by biased data which do not reﬂect the essential properties of data in PR applications: nano-scale crystals (see Fig.
3 (c)), astronomical objects [20], where there is no natural orientation or centering of the image contents.
Toward Practical PR
In this section, we sketch a two-stage method for solving PR, which consists of an initialization stage based on the end-to-end approach, and a reﬁnement stage based on a robust iterative method.
2.1 Passive symmetry breaking
Let T .
=
{ X i , Y i = |F
( X ) | 2 } be the training dataset we set up for implementing the end-to-end approach, and g W be a chosen neural network parametrized by weights W .
Then the naive end-toend approach will take the form min W (cid:80) i (cid:96) ( X i , g W ( Y i )) , where (cid:96) is the loss function.
Due to the symmetries, Y i , Y j that are close may correspond to X i and X j that are centered at very different locations or ﬂipped and hence far apart.
This forces g W to approximate a rapidly changing function, i.
The difﬁculty occurs because we require g W ( Y i ) to match X i as possible, where the latter induces mixed symmetries.
How about we getting rid of X i ’s?
A natural alternative formulation is
[Equation]
as whether g W ( Y i ) outputs X i or any of its symmetric copies, |F ◦ g W ( Y i )
| 2 ≈ Y i ∀ i .
Why it might work?
We toss away the difﬁcult X i ’s, but we also supply less information to the learning model.
Now g W has much more freedom, and it can still generate distant outputs for nearby inputs.
Why is there hope?
We draw our inspiration from the growing pile of evidence that neural networks optimized with ﬁrst-order stochastic methods tend to learn simple functions over complicated ones, known as implicit regularization [21].
For our problem, g W is simple when all the symmetries are broken and complicated when there are symmetries.
So if implicit regularization occurs, symmetries are naturally broken.
Precursors Eq.
(2.1) can be considered as an autoencoder objective with a known decoder.
It is also similar to the cycle consistency idea [22–24] used in deep learning for several computer vision/graphics tasks.
A uniﬁed theme is to approximate identity maps.
But our motivation here is for implicit symmetry breaking taking advantage of implicit regularization, which differ from all other works.
The same formulation and its equivalent form in the autocorrelation form has been independently proposed in [16].
They have not articulated the learning difﬁculty caused by all symmetries but the shift.
Our previous work [17] proposes proprocessing steps for the training set to break symmetries for Gaussian PR—simpliﬁed version of PR where the shift and ﬂipping symmetries are erased.
The method can be generalized for PR as we show in a companion paper [25], but here we follow a different routine that seems simpler.
Jacobian regularization When the underlying distribution is not densely sampled, the learning difﬁculty is (counterintuitively) ameliorated, as the function determined by the data points becomes less oscillatory.
This can be seen from Fig.
2 again and generalized to PR.
Although complete symmetry breaking probably still corresponds to the “simplest” function, partial symmetry breaking leads to functions that are not worse off.
This could lead to competitive local solutions to Eq.
To promote the simplest solution and hence complete symmetry breaking, we can optionally add in certain regularization terms to enhance regularity of the solution.
Empirically, we ﬁnd the Jacobian regularization [26] that encourages slowly varying functions helps, leading to the regularized objective
[Equation]
To save computation, we adopt the approximation scheme proposed in [26] for the Jacobian term.
2.2 Reﬁnement
The end-to-end learning performs prediction based on regression, the quality of which is often only mediocre and depends on how densely the underlying data distribution can be sampled.
For many scientiﬁc inverse problems such as PR, high quality solution for each instance is always sought. In this work we propose to perform subsequent reﬁnement using iterative methods.
For PR, we observe our ﬁrst stage produces results that often lead to reasonable support estimation, which sufﬁces as initialization for popular iterative methods to solve PR.
Here, we choose a recently proposed iterative method based on second-order augmented Lagrangian method, dubbed ALM [27], that exhibits fast convergence and good stability.
Experiments
Data We conduct our experiments on the FashionMNIST dataset [28], which is used by several previous works on PR.
We take their 60 , 000 training images and 10 , 000 test images to construct our training and test sets, respectively.
The images are 28 × 28 .
To simulate the typical black ground that causes the translation freedom in PR applications, we place all the images in a black background of 36 × 36 .
So n = 36 , and we take m = 72 here to ensure injectivity of the forward model 2 n − 1 = 71 is exceeded.
We further modify both train and test images with a random ﬂipping and random translation operation to test the impact of symmetries on learning; samples are shown in Fig. 4-(a).
We use U-Net [29] (DNN) as our backbone DNN.
We term our method (i.e., the two-stage method described in Section 2) Passive Reﬁned .
We compare it with HIO [5], ALM [27], Naive DNN [30], Passive DNN (DNN trained using in Eq.
(2.2)), Naive Reﬁned (Naive DNN reﬁnement with ALM), prDeep [10].
HIO and ALM are representive iterative methods, and Naive DNN and prDeep are two of the state-of-the-art methods based on deep learning: Naive DNN implements the end-to-end approach, and prDeep is a hybrid approach that integrates deep learning and an iterative method, initialized with HIO.
Reconstruction results Results on randomly selected test images are presented in Fig.
Due to symmetries, both the iterative methods (ALM, HIO) and Naive DNN fail to recover the original images.
The reconstruction from the passive DNN captures the object’s shape, although the image details are not completely recovered.
Thus, there is a need for additional reﬁnement.
We observe that the shape information from the ﬁrst stage can be used to derive relatively accurate support of the object, which is ideal for initializing reﬁnement stage which is based on classic iterative methods.
After reﬁnement, the ﬁnal image is a reasonably accurate reconstruction of the original image. [Table]
Table 1 provides the average MSE adjusted to the symmetries (deﬁned in Appendix A.1) for the test set.
As noted, in the presence of symmetries, we see a huge gap in MSE of the reconstructed images between Passive Reﬁned and other competing methods, consistent with the visual results.
The high-quality solution by our Passive Reﬁned method depends crucially on good support estimation enabled by our
Passive DNN stage. However, Naive DNN and the classic HIO—which is used to initialize prDeep [10], are unable to produce good support estimation to enable further reﬁnement.
Acknowledgments and Disclosure of Funding
KT and VK are supported by NSF HDR 1934721. The authors acknowledge the Minnesota Supercomputing Institute (MSI) at the University of Minnesota for providing resources that contributed to the research results reported within this paper.
A Appendix A.1 Mean Square Error (MSE)
Our reconstructed image is in C m × m , where our original image is in C n × n .
To account for the three symmetries when taking MSE measure, we take the following steps: we take the original image, and scan through the larger reconstructed image to account for the translation symmetry.
At each scan position, we calculate an adjusted MSE between the current patch B ∈ C n × n and the original image A .
A λ > 0 and a global phase factor e iθ (to account for the global phase) are introduced when calculating the MSE, i.e.,
[Equation]
The smallest adjusted MSE is recorded over all scan positions.
Then, the original image A is 2D ﬂipped and conjugated and the same scanning process is repeated to calculate another smallest MSE, to account for the ﬂipping symmetry.
The smaller of the smallest MSE values is ﬁnally taken.
Below, we show that the optimal value in Eq.
(A.1) can be easily computed.
First we expand the square inside the objective and perform partial minimization with respect to θ , leading to
[Equation]
(A.
[Equation]
But Re (cid:10) A , B e iθ (cid:11) = Re (cid:0) (cid:104)
A
[Equation]
B
[Equation]
e
[Equation]
≤
[Equation]
e
[Equation]
≤ |(cid:104) A , B (cid:105)| and the upper bound is achievable when θ = − ∠ (cid:104) A , B (cid:105) . So the optimization problem now becomes
[Equation]
The minimum of Eq. (A.3) occurs either when η = 0 , which is (cid:107) A (cid:107) 2 F , or by unconstrained optimality condition when 2 η (cid:107) B (cid:107) 2 F = 2 |(cid:104) A , B (cid:105)| = ⇒ η = |(cid:104) A , B (cid:105)| / (cid:107) B (cid:107) 2 F , leading to the function value (cid:107) A (cid:107) 2 F − | (cid:104) A , B (cid:105) | (cid:107) B (cid:107) 2 F , (A.4) which is the smaller one.