Optimizing Intermediate Representations of Generative Models for Phase Retrieval
Abstract
Fourier phase retrieval is the problem of reconstructing images from magnitudeonly measurements.
It is relevant in many areas of science, e.g., in X-ray crystallography, astronomy, microscopy, array imaging and optics.
When training data isavailable, generative models can be used to constrain the solution set.
However, not all possible solutions are within the range of the generator.
Instead, they arerepresented with some error.
To reduce this representation error in the contextof phase retrieval, we first leverage a novel variation of intermediate layer optimization (ILO) to extend the range of the generator while still producing imagesconsistent with the training data.
Second, we introduce new initialization schemesthat further improve the quality of the reconstruction.
With extensive experiments,we can show the benefits of our modified ILO and the new initialization schemes
Introduction and Related Work
Recently, generative models have shown impressive performance to regularize inverse problems.
However, the solutions is required to lay in the range of the generator.
To decrease the representationerror and thereby boost the expressiveness of these generative models, the idea of optimizingintermediate representations has been applied to various linear inverse problems [1; 2].
In our paper, we show that this idea can be extended to solve the more difficult non-linear inverseproblem of phase retrieval.
In particular, we tackle the challenging problem of non-oversampled
Fourier phase retrieval.
For brevity, we define the problem for one dimension.
Fourier Phase Retrieval: In many imaging applications in physics, we are only able to measurethe magnitude of the Fourier transform of an image, e.g., in X-ray crystallography [3], astronomy [4],microscopy [5], array imaging [6] or optics [7].
The magnitude measurements are given as
[Equation]
[Equation]
) where F is the discrete Fourier transform and x ∈ R n is the original image (here represented as avector).
As the problem is highly underdetermined we require additional information about x , which,in this paper, will be learned from example images to constrain the solution space.
We assume thatthe Fourier measurements are not oversampled, i.e., we consider the case, where the image x hasnot been zero-padded which corresponds to a practical more relevant setup.
A related problem is the
Gaussian phase retrieval problem, where F in Equation 1 is replaced by a (not necessarily square)matrix A that has normally-distributed entries.
Note, that the problem of phase retrieval are more difficult than linear inverse problems due to thenon-linearity (the absolute value) in the forward model.
In the following, we give a short overview ofthe various approaches distinguishing between optimization-based and learning-based methods.
1.1 Methods without Learning for Phase Retrieval
One of the first (Fourier) phase retrieval algorithms is the Gerchberg-Saxton (GS) method [8]that is based on alternating projections.
Fienup [9] later introduced the hybrid-input-output (HIO)algorithm which improved the Gerchberg-Saxton algorithm.
The Gaussian phase retrieval problemwas approached by Candes et al.
[10] who used methods based on Wirtinger derivatives to minimizea least-squares loss.
Wang et al.
[11] considered a similar idea but used a different loss functionand so-called truncated generalized gradient iterations.
Holographic phase retrieval is a relatedproblem, which aims to reconstruct images from magnitude measurements where a known referenceis assumed to be added onto the image.
Lawrence et al.
[12] approached this problem using untrainedneural-network priors.
1.2 Learning-based Methods for Phase Retrieval
Recently, learning-based methods for solving phase retrieval problems gained a lot of momentum.
They can be classified into two groups: supervised methods and unsupervised methods.
Supervised methods: These methods directly learn a mapping that reconstructs images from themeasurements.
Learning feed-forward networks for solving Fourier phase retrieval was done by
Nishizaki et al.
This approach was later extended to a neural network cascade by Uelwer et al.[14] , who also used conditional generative adversarial networks to solve various phase retrievalproblems [15].
Supervised learning of reference images for holographic phase retrieval was done by
Hyder et al.
[16] and Uelwer et al.
[17] gave further insights.
Unsupervised methods: Unsupervised methods are agnostic with respect to the measurementoperator, because they are trained on a dataset of images without the corresponding measurements.
Generative priors for compressive phase retrieval have been analyzed by Hand et al.
[18] and Liu et al.[19].
Alternating updates for compressive phase retrieval with generative models was proposed by
Hyder et al.
Manekar et al.
[21] used a passive loss formulation to tackle the non-oversampled
Fourier phase retrieval problem.
In the context of these related works, the method that we propose in this paper can be seen as anunsupervised learning-based approach that can optionally be extended with a supervised component,as we detail in Section 2.2.
1.3 Optimizing Representations of Generators for Inverse Problems
Generative models have been used to solve linear inverse problems, e.g., compressed sensing [22],image inpainting [23].
To decrease the representation error and thereby boost the expressiveness ofthese generative models, the idea of optimizing intermediate representations was successfully appliedin the context of compressed sensing [1], image inpainting and super-resolution [2].
The formerpaper called this approach generator surgery, the latter intermediate layer optimization (ILO).
1.4 Our Contributions
The contributions of this paper are the following:
We reformulate and extend the idea of optimizing intermediate representations in generative
models for solving different phase retrieval problems.
In particular, we propose a subsequentoptimization that further improves the reconstruction.
We introduce new learned and non-learned initialization schemes for generator-based ap-
proaches to solve inverse problems, reducing the need for multiple runs with random restarts. [Equation]
[Figure] In extensive experiments, we consider the underdetermined Fourier phase retrieval and showthat our method provides state-of-the-art results.4.
By analyzing various combinations of methods and generator networks, we show theinfluence of the choice of the generator.
In further ablation studies, we show the importanceof each component of our method.
Phase Retrieval with Generative Models
The basic idea of applying a trained generative model G : R k
→ R n to the phase retrieval problem isto plug G ( z ) into the least-squares data fitting term and to optimize over the latent variable z , i.e.,
[Equation]
This method has been used for Gaussian phase retrieval by Hand et al.
While this approachyields good results, the reconstruction quality is limited by the range of the generator network G .
This issue is, for example, also discussed in the work of Asim et al.
In the following, we explainhow the range of G can be extended by optimizing intermediate representations of G instead of onlysolving Equation (2) with respect to the latent variable z .
Notation: For the exposition of the method we use the following notation: The generator network
G = G k ◦ · · · ◦ G 1 can be written as the concatenation of k layers G 1 , . . .
, G k .
The subnetworkconsisting of layers i through j is denoted by G ji = G j ◦ · · · ◦ G i .
Note, that G = G k
=
G i 1 ◦ G ki +1for 1 ≤ i ≤ k .
The output of layer G i is written as z i , i.e., z i
=
G i
( z i − 1 ) .
The input of G is z 0 .
Furthermore, we define the ℓ 1 -ball with radius r around z as B r ( z ) = (cid:8) x (cid:12)(cid:12)
∥ x − z ∥ 1 ≤ r (cid:9) .
2.1 Phase Retrieval with Intermediate Layer Optimization (PRILO)
The key idea of ILO is to vary the intermediate representations learned by a sub-network of thegenerator G while ensuring an overall consistency of the solution.
The latter is particularly challengingfor non-linear problems like phase retrieval.
As we will show in the experiments, additional stepsand special initialization schemes are essential to obtain useful results.
All of these steps will beexplained next.
The generator network captures the prior knowledge and restricts the solutions of the (underdetermined) phase retrieval problem.
In this work, we use two important ideas:
Range extension: Hidden representations z i for i > 0 are allowed to vary outside the range of G i 1 .
This will reduce the representation error of the generator.
Image consistency: Solutions must be realistic, non-degenerate and similar to the the trainingdataset.
This is achieved by back-projection ensuring that z i is not too far away from theimage of some z 0 .
Note that one has to strive for a trade-off between range extension and image consistency, whererange extension facilitates the optimization, while image consistency regularizes the solution.
Concretely, we repeatedly split the generator G into sub-networks G i 1 and G ki +1 and iterativelyoptimize the intermediate representations
z i resulting from G i 1 (instead of simply optimizing theinitial latent variable z 0 ).
In this way, we go outside the range of the sub-network G i 1 to reduce themean-squared magnitude error.
To make the resulting image more consistent we back-project z icloser to the range of G i 1 .
We will show that z 0 from the back-projection is not only a good candidatefor the overall optimization, but also serves as a good initialization for further optimization leading toa significant refinement of the overall solution.
To get started, we initially optimize the input variable z 0 (minimize Equation (2)).
This providesstarting points for all intermediate representations z i for i > 0 .
The overall procedure then iterativelyapplies the following three steps to different intermediate layers:
A. Forward optimization: vary the intermediate representation z
i to minimize the magnitudeerror while staying in the ball with radius r i around the current value of z i , i.e.,
[Equation]
[Equation]
Note, that this step possibly leaves the range of G i 1 (range extension).
By introducingball-constraints on the optimized variable we avoid overfitting the measurements.
B. Back projection: find an intermediate representation G i 1 ( z 0 ) that is close to the optimal z ∗ ifrom step A, i.e.,
[Equation]
where 0 ∈ R l denotes the vector of all zeros.
By doing so we ensure that there exists a latentz 0 that yields the reconstruction.
Thereby, we regularize the solution from step A to obtainimage consistency.
While these two steps have been shown to improve the overall performance of linear inverse problems, we found that for phase retrieval the following refinement step significantly improves thereconstructions.
C. Refinement: Our intuition here is that the back-projected latent ¯ z 0 serves as a good initialization for further optimization.
Starting from the ¯ z 0 found in step B, we again optimize themeasurement error
[Equation]
Note, that as we only optimize the latent variable of G , we do not have to apply theback-projection again for this refinement step.
These steps can be repeated for multiple intermediate representations.
After that we obtain the finalreconstruction as x ∗ =
G ( z ∗ 0 ) .
Steps A, B and C are visualized in Figure 1.
The three optimization problems stated in Equations (3)-(5) are solved using projected gradientdescent.
This means that after each gradient descent step the iterate is projected back onto theappropriate ℓ 1 -ball.
In our implementation this projection is implemented using the method describedby Duchi et al.
2.2 Initialization Schemes
Due to the difficulty of the phase retrieval problem, many methods are sensitive to the initialization.
Candes et al.
[10] proposed a spectral initialization technique that is often used for compressive
Gaussian phase retrieval.
Since it requires one to solve a large eigenvalue problem, this methodcan be quite costly.
A common approach to improve reconstruction performance is to use randomrestarts and to select the solution with the lowest magnitude error after the optimization.
However,this requires running the method multiple times.
Instead, we propose two fast initialization schemesthat use the generative model and the available magnitude information before the optimization.
Magnitude-Informed Initialization (MII): Deep generative models can generate a lot of imagesat relatively low cost: instead of optimizing Equation (2) with random restarts, we sample a set ofstarting points Z
= { z (0)0 , . . .
, z ( p ) 0 } and select the one with the lowest magnitude mean-squarederror (MSE),
[Equation]
before the optimization.
The empirical reason for this choice is that the MSE between the unknowntarget image and the initial reconstructions G ( z ) for
z ∈
Z strongly correlates with the MSE of theirmagnitudes (correlation ρ = 0 . 91 ).
This initialization scheme is applicable to other reconstructionalgorithms that search in the latent space of a generative model as well.
Learned Initialization (LI): The generator is trained on a dataset of images that are characteristic for the problem.
We use the generator to train an encoder network E θ that maps magnitudemeasurements y to latent representations z 0 .
Once the encoder is trained, we can use it to predict aninitialization for the optimization discussed in Section 2.1.
A naive way to train the encoder network is to create input/output pairs by starting with random latentvectors z 0 and combining them with the magnitudes
|F G ( z 0 ) | of the corresponding image.
Thedisadvantage is that the original training images are only implicitly used in this approach (becausethose were used to train the generator).
To leverage the generator and also the training images, weestimate the weights θ of the encoder E θ , such that encoded magnitudes E θ ( y ) generate an image
G ( E θ ( y ))
that is close to the original image x .
This idea originates from GAN inversion [26].
More precisely, we minimize a combination of three loss functions to train the encoder E θ ,
[Equation]
where L MSE is the image MSE, L perc the LPIPS loss [27], and L adv is a Wasserstein adversarialloss [28] with gradient penalty [29] (using the discriminator network D ϕ ).
In our experiments, we setλ perc = 5 · 10 − 5 and λ adv = 0 .
Note, that the generator network is fixed and we only optimize theencoder weights θ
and the discriminator weights ϕ to solve a learning objective.
Additionally, we also found it helpful to apply a small normally-distributed perturbation with mean0 and standard deviation σ = 0 .
05 to the predicted latent representation.
We do so because usinga fixed initialization leads to a deterministic optimization trajectory when running the optimizationmultiple times.
For better exploration of the optimization landscape we also apply gradient noise during optimization.
In combination with the projection onto the feasible set the update reads as
[Equation]
(8) where f is the objective function corresponding to the current step, u ( k ) ∼ N (0 , σ 2 k I ) with σ 2 k =η (1+ k ) γ and P is the projection onto B r ( x (0) ) .
In our experiments we set η = 0 .
02 and γ = 0 .
Thisnoise decay schedule was proposed by Welling and Teh [30] and is also discussed by Neelakantanet al.
[31] in the context of neural network training.
One drawback of the learned initialization is that it requires one to retrain the encoder network whenthe measurement matrix changes.
In the following, we only evaluate this approach for the Fourierphase retrieval problem. [Table] [Table]
Experimental Evaluation
We evaluate our method on the following datasets: MNIST [32], EMNIST [33], FMNIST [34] [Table] and
CelebA [35].
The first three datasets consist of 28 × 28 grayscale images, whereas the latter dataset is a collection of 200 , 000 color images that we cropped and rescaled to a resolution of 64 × 64 .
Similar to Hand et al.
[18], we use a fully-connected variational autoencoder (VAE) for the MNISTlike datasets and a DCGAN [36] for the CelebA dataset.
Both architectures were also used by Boraet al.
[22] for compressed sensing.
Going beyond existing works, we are interested in improvingthe performance on the CelebA dataset even further, thus we considered deeper, more expressivegenerators, like the Progressive GAN [37] and the StyleGAN [38].
Since the optimization of theinitial latent space for deeper generators is more difficult, we expect significantly better results byadaptively optimizing the successive layers of these models.
HIO [9]
[Figure]
3.1 Phase Retrieval with Fourier Measurements
For the problem of Fourier phase retrieval, we compare our method first with the HIO algorithm [9]and the GS algorithm [8] (both having no learning component) and then with the following supervisedlearning methods: an end-to-end (E2E) learned multi-layer-perceptron [15], a residual network [13],a cascaded multilayer-perceptron [14] and a conditional GAN approach [15].
Additionally, we apply
DPR which was originally only tested on Gaussian phase retrieval [18].
Table 1 and Table 2 show the mean peak-signal-to-noise-ratio (PSNR) and the mean structuralsimilarity index measure (SSIM) [39].
Each reported number was calculated on the reconstructionsof 1024 test samples.
We allow four random restarts and select the generated sample resulting in thelowest measurement error.
As one can see, more expressive models result in better reconstructions.
Furthermore, we note that our proposed initialization schemes, i.e., the LI and MII, lead to improvedperformance.
Figure 2 shows typical reconstructions from the CelebA dataset.
Surprisingly, our unsupervised PRILO-MII approach often quantitatively outperforms the supervisedcompetitors (MNIST and EMNIST).
Only for FMNIST we get slightly worse results which weattribute to the performance of the underlying generator of the VAE.
Notably, we outperform DPRthat uses the same generator network, which shows that the modifications explained in this paper arebeneficial.
Summarizing, among classical and unsupervised methods our new approach performsbest, often even better than the supervised methods.
The CelebA dataset is more challenging and the influence of the choice of the generator networksis substantial.
We consider the DCGAN [36], the Progressive GAN [37] and the StyleGAN [38]as base models for our approach (second column in Table 2).
Furthermore, the new initializationschemes have a strong impact.
For fairness, we also combined the existing DPR approach [18] withthe various generator models and the initialization schemes.
Thus we are able to study the influenceof the different components.
Among the generators, StyleGAN performs best.
MII improves theresults, while LI achieves even better reconstructions.
For all previously mentioned combinations,our new method PRILO is better than DPR. [Figure] [Table]
3.2 Ablation Studies
In our experiments, we observed that the performance can be improved by re-running the optimizationprocedure with a different initialization and selecting the result with the lowest magnitude errorwhich is a common practice in image reconstruction.
However, this approach is quite costly.
Ourinitialization schemes described in Section 2.2 can be used to achieve a similar effect without theneed of re-running the whole optimization procedure multiple times.
In Figure 3, we compare thesingle randomly initialized latent variable, 5000 initialization using the MII approach and slightlyperturbed learned initialization (LI).
Again, we consider a test set of 1024 samples.
Although, we arestill using restarts in combination with our initialization, we observe that already for a single run ourinitialization outperforms the results of 5 runs.
In order to assess the impact of each step on the reconstruction performance, we perform an ablationstudy.
We compare our complete PRILO-MII and PRILO-LI models with different modified versionsof the model: for one, we consider omitting every optimization, i.e., we only compare with theinitialization.
Next, we omit the back-projection step (B) and the refinement step (C).
We alsocompare with the variant of our model which we only omit step C. Finally, we analyze the impact ofthe ball-constraints.
Table 3 shows that each of the components is important to reach the results.
Conclusion
Generative models play an important role in solving inverse problem.
In the context of phase retrieval,we observe that optimizing intermediate representations is essential to obtain excellent reconstructions.
Our method PRILO, used in combination with our new initialization schemes, produces betterreconstruction results for Fourier phase retrieval than existing (supervised and unsupervised) methods.
Notably, in some cases our unsupervised variant PRILO-MII even outperforms existing supervisedmethods.
We also show that the initialization schemes we introduced can easily be adapted to differentmethods for inverse problems that are based on generative models, e.g., DPR.
Our ablation studyjustifies that each of our used components is essential to achieve the reported results.
Limitations: During our experiments, we observed that some hyperparameter tuning is necessaryto achieve good results: the selected intermediate layer highly influences the results and also the radiifor the constraints play an important role.
Furthermore, our approach is in some cases not able toreconstruct the finer details and is still (to some extend) limited by the generative model. Appendix [Figure]
A.1 CelebA Reconstructions using the Progressive
[Figure]
[Figure]
[Figure] A.3 Ablation Experiments: PSNR Throughout the Optimization [Figure]
P S N
[Figure]
A
BC
A
BC
A
BC
PRILO
[Figure]
LI PRILO
[Figure]
w/o
[Figure]
PRILO
[Figure] LI w/o [Figure] C