Unlocking Inverse Problems Using Deep Learning: Breaking Symmetries in Phase Retrieval
Abstract
In many physical systems, inputs related by intrinsic system symmetries generate the same output.
So when inverting such systems, an input is mapped to multiple symmetry-related outputs.
This causes fundamental difﬁculties for tackling these inverse problems by the emerging end-to-end deep learning approach.
Taking phase retrieval as an illustrative example, we show that careful symmetry breaking on the training data can help get rid of the difﬁculties and signiﬁcantly improve learning performance in real data experiments.
We also extract and highlight the underlying mathematical principle of the proposed solution, which is directly applicable to other inverse problems
Introduction
For many physical systems, we observe only the output and strive to infer the input.
The inference task is called inverse problem .
Formally, the underlying system is modeled by a forward mapping f , and solving the inverse problem amounts to identifying the inverse mapping f − 1 .
Let y denote the observed output.
Traditionally, inverse problems are phrased as regularized optimization problems: min x (cid:96) ( y , f ( x ))
+ λ Ω( x ) , where x represents the input to be estimated, (cid:96) ( y , f ( x )) ensures y ≈ f
( x ) , and Ω( x ) encodes structural information about x to make the problem well posed.
Deep learning has enabled the use of data-driven (cid:96) or Ω , or replacing mappings in iterative methods for solving the regularized formulation by data-adaptive ones.
The most radical is using neural networks directly to approximate f − 1 , also known as the end-to-end approach .
Several review articles [1–4] have covered these recent developments.
We note that most of the successes are about linear inverse problems, i.e., f is linear.
In this paper, we focus on the end-to-end approach applied to nonlinear inverse problems, and take phase retrieval (PR)—which is central to scientiﬁc imaging [5]—as an example.
Here, we focus on the 2D version, which is seen most frequently in applications.
Given
X ∈ C n 1 × n 2 , the forward mapping is
[Equation]
where F m 1 × n 1 consists of the ﬁrst n 1 columns of the Fourier matrix F m 1 × m 1 , often called partial Fourier matrix; similarly for F m 2 × n 2 .
The mapping is generically injective when m 1 ≥ 2 n 1 − 1 and m 2 ≥ 2 n 2 − 1 , up to three intrinsic symmetries: 1) 2D translation of the nonzero content of X ; 2) conjugate of 2D ﬂipping of X ; and 3) global phase transfer to X : X e iθ for any θ ∈ [ − π, π ) .
Any composition of these changes to X will leave the observation Y unaltered.
Fig.
1 illustrate the ﬁrst two symmetries, assuming X is a real-valued image.
by the Image
Symmetries can cause signiﬁcant difﬁculty for the end-to-end approach.
To see this, 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- function, allowing both positive and negative outputs, using the endto-end approach.
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 i (cid:55)→ x 2 i 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 capacity 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, similar problems can surface.
Solving PR is not free from the trouble.
Breaking symmetries for PR
The above difﬁculty has recently been articulated in [6,7].
Particularly, [6] has proposed a method to resolve the difﬁculty in Gaussian PR—a simpliﬁed version of PR—called symmetry breaking .
The proposed working principle is to identify a connected (in topological sense), smallest representative subset of the input space, so that all training samples are mapped to such subset to induce a smoother function to be approximated.
Take the above square root example again, the positive ray R + ⊂ R can be taken, as is it topologically connected, and any x ∈ R (except 0 ) can be represented by an element in R + by a sign ﬂipping, and it cannot be made smaller to remain representative.
Then the training data set (cid:8) x i , x 2 i (cid:9) will be processed so that any pair (cid:0) x i , x 2 i (cid:1) with negative x i will be changed to (cid:8) − x i , x 2 i (cid:9) , with the rest data points unchanged.
Obviously after this processing to the training set, the function determined will closely trace the upper branch of the square root function, which is much more smooth.
PR has three symmetries as discussed above.
Under the global phase transfer, equivalent data points form continuous curves that are easy to represent algebraically.
The conjugate 2D ﬂipping and nonzero content translation, however, induce irregular equivalent sets that are hard to represent.
Following [6] and prescribing a rule for symmetry breaking in the original X space seems hopeless.
Fortunately, the three symmetries can be equivalently represented in the complex phase e iθ space after the Fourier transform.
Let X denote the oversampled Fourier transform of X .
Now 1) for 2D translation, any allowable 2D translation t 1 , t 2 ∈ Z induces the change X ( k 1 , k 2 ) (cid:55)→ e i 2 π (cid:16) k 1 t 1 m 1 + k 2 t 2 m 2 (cid:17) X ( k 1 , k 2 ) ; 2) conjugate 2D ﬂipping induces the change X
(cid:55)→ X , i.e., change to the complex phase e i Θ (cid:55)→ e − i Θ ; and 3) global phase transfer induces the change X (cid:55)→ e iθ X .
The change due to 2) is a global sign ﬂipping in the angle space, and the equivalent set due to 3) is a line in the angle space.
But 1) is still relatively irregular whether represented in the angle or phase space.
Our strategy here is a combination of rigorous symmetry breaking for 2) and 3) in the complex phase space and heuristic symmetry breaking for 1) in the original space—our later real-data experiments conﬁrm that the combination is effective.
To break 1), we propose to simply center the nonzero content as a heuristic.
To break 2) and 3), we perform a geometric construction of a connected, smallest representative subset in the angle space and then represent it in the phase space to avoid the tricky 2 π periodicity issue in the angle space.
For brevity, we omit intuition behind the construction and directly present the results as follows.
Consider the following set in the phase domain
[Equation]
where S denotes the 1D complex circle and S + the upper half circle.
We can prove the following, stated in the equivalent vector space for convenience.
We write R ⊂ S m 1 m 2 to mean the equivalent of R in the vector notation.
Theorem 2.1.
Consider the conjugate ﬂipping and global phase transfer symmetries only.
The set R is a connected, smallest representative in the phase domain S m 1 m 2 with a negligible set
N = { 1 } × { ω ∈ S : Im ( ω ) = 0 } m 1 m 2 − 1 .
Proof.
See Appendix A.1.
To apply this, we work with end-to-end DNNs that directly predicts the m 1 × m 2 complex phases.
We ﬁrst center the nonzero content inside X
i ’s in the training set, and then take the oversampled
Fourier transform and perform the symmetry breaking as implied by Theorem 2.1 in the complex phase space.
For any phase matrix Ω , the symmetry breaking goes naturally as follows:
ﬁrst a global phase transfer is performed to make Ω(1 , 1) = 1 , and then a global angle (here we assume the angle has been transferred to the range of [ − π, π ) ) negation is performed, i.e., Ω (cid:55)→ − Ω if the second angle is negative.
Symmetry breaking for general inverse problems For general inverse problems, although the symmetries might be very different than here and the sample spaces could also be more complicated, the three properties, which concern only the geometric and topological aspects of the space, can be generalized as a basic mathematical principle for effective symmetry breaking.
Our symmetrybreaking solution for PR also suggests that for problems with multiple symmetries, one may need to look at a transformed space, or even mixture of spaces for different symmetries for efﬁcient representation and symmetry breaking.
Numerical Experiments
[Table]
In this section, we set up a preliminary experiment to verify our claim that effective symmetry breaking facilitates efﬁcient learning.
Particularly, we show that symmetry breaking substantially improves PR performance over alternative methods.
We conduct our experiments on the
Fashion-MNIST dataset [8].
We take their 60 , 000 training images and 10 , 000 test images to construct our training and test sets respectively.
Each example is a 28 × 28 grayscale image.
To simulate the typical black ground that causes the translation freedom in PR applications, we place all the images in a black background of 42 × 42 — most previous methods overlook this in their experiments, but practically the translation freedom, or what PR community call support estimation, is a major failing factor for most PR methods.
So n = 42 , and we take m = 96 here to ensure injectivity of the forward model 2 n − 1 = 83 is exceeded.
We create 4 variants of the dataset to test the impact of symmetries on learning—this is the ﬁrst time this kind of rigorous evaluation is performed; most previous methods use natural image datasets where the image contents are naturally centered and oriented, which does not match the scenarios in PR applications e.g., in coherent diffraction imaging.
We do this by modifying the images as described below, followed by the standard operation of taking Fourier magnitudes.
• No Symmetry : i.e., all images are placed in the center of the black background
i
e. padding 7 pixels on all side; samples shown in Fig.
3 (a)-left; • Flipping symmetry : all images are placed in the center of the black background and 50% of randomly selected training and test images are 2D ﬂipped; samples shown in Fig.(b)-left.
• Shift symmetry : all images placed in a larger dark background and randomly translated; samples shown in Fig.
3 (c)-left; • Shift and Flipping symmetries : random ﬂipping followed by random translation; samples shown in Fig.
3 (d)-left.
Results on randomly selected test images are presented in Fig.
We use U-Net [9] as our backbone neural network.
For each variant of the dataset, the left column is the groundtruth image, and the middle and right columns are results produced by U-Net- B , i.e., without symmetry breaking—this is exactly the method used in [10], one of the state-of-the-art methods based on the end-to-end approach, and U-Net- A , i.e., with symmetry breaking—our method, respectively. produces consistent recovery results without the symmetric
[Table]
First note that when no explicit symmetries are built into the dataset, U-Net- B , a representative end-to-end method for PR [10], gives good recovery. But it fails once we build in the essential symmetries.
The mode of failure is interesting, as the estimated images are almost always the superposition of the symmetric (translated or ﬂipped) copies of the groundtruth.
This is very similar to the failure mode of the classic iterative methods on PR.
Moreover, for images that are visually similar between the original and the ﬂipped copy g “handbag", “leggings", the reconstruction results are good with or without the ﬂipping symmetry, consistent with our intuition.
On the other hand, irrespective of the symmetries, U-Net- A consistently leads to good recovery.
Table 1 provides the average MSE adjusted to the symmetries (deﬁned in Appendix A.2) for the test set.
As noted above, absent symmetries, both U-Net- B and U-Net-
A work well and the average MSEs are the same.
However, once the dataset contains the essential symmetries, we see a substantial gap in the MSEs of the reconstructed images, which is consistent with the visual results.
For practical PR, mostly iterative methods are deployed.
However, these methods are known to fail when there is translation freedom in the image and the support (i.e., location of nonzero pixels) of the image content is not precisely known.
To see if our end-to-end approach makes progress on this, we compare it with a state-ofthe-art iterative method for PR recently proposed in [11] that has demonstrated good numerical stability and competitive performance, dubbed ALM.
Here we only experiment with the most realistic version of the dataset, i.e., with both shift and ﬂipping symmetries.
Results on randomly selected test images are presented in Fig.
4 and quantitatively results are presented in Table 2.
Visually, our method faithfully reconstructs the holistic content of the original images, whereas both U-Net-B and ALM fail miserably.
Quantitatively, our method leads the other two by a considerable gap in MSE.
These results show that symmetry breaking is signiﬁcant in unlocking the true potential of the endto-end approach for solving PR in particular, and nonlinear inverse problems with symmetries in general.
In a companion paper [12], we describe a less mathematical yet effective way of performing symmetry breaking for PR, that also seems to be generally applicable for other nonlinear inverse problems.
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 Deferred Proofs A.1 Proof of Theorem 2.
Proof.
It is clear that R is connected in S m [-
-]1 m 2 since R is path-connected set on S m 1 × m 2 with the inherited subspace Euclidean topology of C m 1 × m 2 .
Also, N is of Lebesgue measure 0 since it is a product of ﬁnite points.
Now we are going to prove R is a representative of S m 1 m 2 .
Let G be the set of all possible symmetry transfers composed of sequences of global phase transfer and global phase conjugation in PR.
For any given z = (e iθ 0 , e iθ 1 , · · · , e iθ m 1 m 2 − 1 ) , we need to ﬁnd a ω ∈ R such that there is a g ∈ G satisfying g ( ω )
=
z .
If Im(e i ( θ 1 − θ 0 ) ) > 0 , we take ω
=
(1 , e i ( θ 1 − θ 0 ) , e i ( θ 2 − θ 0 ) , · · · , e i ( θ m 1 m 2 − 1 − θ 0 ) ) then ω ∈ R and e iθ 0 ω =
z .
On the other hand, if Im(e i ( θ 1 − θ 0 ) ) < 0 , we can consider the conjugate format ω
= (1 , e i ( θ 1 − θ 0 ) , e i ( θ 2 − θ 0 ) , · · · , e i ( θ m 1 m 2 − 1 − θ 0 ) ) ∈ R and obviously a global phase negation followed by a global phase transfer e iθ 0 leads to z .
This proves that R is representative.
At last, we need to show the smallestness in the sense that with any point of R removed, we cannot recover it by other points in R .
That is, with arbitrary ˜ z ∈ R given, for all g ∈ G and all z ∈ R\{ ˜ z } , we have g ( z ) (cid:54) =
˜ z .
We ﬁrst claim that any g is equivalent to an optional global phase conjugation followed by a global phase transfer.
To see this, it is sufﬁcient to prove that the order of phase conjugation and phase transfer can be exchanged.
Let ψ denote a global phase transfer by e iψ and f phase conjugation. Now if ψ ◦ f = f ◦ ψ (cid:48) , or − ( ψ (cid:48) + θ )
= − θ + ψ + 2 kπ , we have ψ (cid:48)
= − ψ − 2 kπ . So one can keep exchanging conjugation and transfer so that all conjugations precede transfers.
The conjugations now can be equivalently written as an optional conjugation, and the transfers as a single transfer.
Now we can go back to the proof of smallestness.
Write ˜ z
= (e i ˜ θ 0 , e i ˜ θ 1 , · · · , e i ˜ θ m 1 m − 1 ) and z = (e iθ 0 , e iθ 1 , · · · , e iθ m 1 m 2 − 1 ) where
˜ θ
= θ 0 =
and Im (e i ˜ θ 1 ) , Im (e iθ 1 ) > 0 .
Suppose that there is a g ∈ G such that ˜
z
= g ( z ) .
we may assume g = f ◦ ψ or g = ψ where ψ is a phase transition with the total angles ψ and f is the conjugate ﬂipping.
If g = f ◦ ψ , ˜ z
= g ( z [-)implies for any 0 ≤ j ≤ M − 
-]) implies that [Equation] A
[Equation]
1) for some k j ∈ Z .
We can solve ψ = 2 πk 0 as j = 0 and this implies ˜ θ
j ≡ 2 π ( k j − k 0 ) − θ j mod 2 π ≡ − θ j for all j , especially, ˜ θ 1 = − θ 1 .
This contradict with the fact that Im (e i ˜ θ 1 ) ,
Im (e iθ 1 ) > 0 .
If g = ψ , we then have the relationship
[Equation]
Again, we can solve ψ = − 2 πk 0 as j =
and this indicates that ˜
z
= z which contradicts the assumption.
Hence, we prove the smallestness.
A.2 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 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.3) can be easily computed.
First we expand the square inside the objective and perform partial minimization with respect to θ , leading to
[Equation]
But Re (cid:10) A , B e iθ (cid:11) =
Re (cid:0) (cid:104) A , B (cid:105) e iθ (cid:1) ≤ (cid:12)(cid:12) (cid:104) A , B (cid:105) e iθ (cid:12)(cid:12) ≤ |(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.5) occurs either when η = 0 , which is (cid:107) A (cid:107) 2 F , or 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)
[Equation]
(A [Equation] 6) which is the smaller one.