Deep S$^3$PR: Simultaneous Source Separation and Phase Retrieval Using Deep Generative Models

Introduction

For $L=1$ this problem is just standard phase retrieval (PR), a problem for which myriad solutions exist . However, when $L>1$ things become significantly more challenging. In particular, one is forced to disentangle the components of $y$ that came from $x_{i}$ from the components that came from $x_{j}$, with $i\neq j$. That is one must solve a source separation (SS) problem as well as multiple PR problems.

S3PR is intrinsic to a variety of different application domains where one measures the intensity of a field formed by multiple independent coherent sources. Under these conditions, the fields within a single source add but the intensities between distinct sources add. This situation appears in partially coherent phase imaging microscopy , correlation-based imaging through thin scattering media and around corners with highly separated or multi-spectral objects , transmission-matrix based imaging through thick scattering media with multiple independent sources , and even multiple source localization with mmWave 5G . A detailed description of S3PR’s role in correlation-based imaging through scattering media is provided in the supplement.

Despite the prevalence of the S3PR problem, no general solution to to S3PR yet exist. The lack of existing solutions is likely because S3PR is simply too non-convex and under-determined to be solved with conventional algorithms. While a cascaded solution, that is SS followed by PR, could work in principle, Section 5 demonstrates that in practice the components $|\mathbf{A}x_{1}|^{2},~{}|\mathbf{A}x_{2}|^{2},~{}\ldots,~{}|\mathbf{A}x_{L}|^{2}$ are too similar to reliably separate with traditional SS algorithms.

Our contribution

Using this constraint, we recover $x_{1},~{}x_{2},~{}\ldots,~{}x_{L}$, with our estimates $\hat{x}_{l}=G(\hat{z}_{l})$, from $y$ by solving the optimization problem

To our knowledge, this work represents the first general solution to S3PR. Accordingly, it opens up a variety of application domains. In Section 5 we apply this method to simulated Gaussian, coded-diffraction-pattern (CDP), and Fourier measurements and demonstrate the successful recovery of numbers ($\mathcal{C}=$ MNIST dataset) and articles of clothing ($\mathcal{C}=$ Fashion MNIST dataset ).

Limitations

Our results have a few limitations. First, our present reconstructions are low resolution and come from fairly restrictive classes — digits and articles of clothing. Second, while we provide extensive evidence that deep generative models can be used to solve S3PR, our current results are purely empirical. S3PR remains a highly under-determined and non-convex problem and deriving the conditions under which it can and cannot be solved remains an important but open problem.

Related work

While the S3PR problem is new, both PR and SS have been studied extensively and have a vast literature. Likewise, while not previously used for S3PR, deep generative models have recently been applied to a range of inverse problems. We now highlight a few of the most prominent of these works.

The optics community has studied the PR problem continuously since the 1970s . In the last decade, PR has caught the attention of the optimization and machine learning community as well . For a benchmark study of over a dozen popular PR algorithms, see PhasePack . No existing PR algorithm can handle S3PR — they are designed for a fundamentally different forward model that does not mix measurements.

recognized that if one imposed a prior on the reconstruction one could perform PR using significantly fewer measurements. This initial work has been followed up by numerous others that have imposed more and more elaborate priors on the reconstruction . Again, none of these algorithms can handle S3PR.

Multi-source phase retrieval

A handful of works have studied multi-source PR, defined as recovering $x_{1},~{}x_{2},~{}...,~{}x_{L}$ from

This model resembles Equation (1), but note that here the mixing occurs before the non-linearity. This difference makes the reconstruction problem significantly easier as it enables the application of a standard PR algorithm to recover $\sum_{l=1}^{L}\mathbf{A}_{l}x_{l}$ followed by a standard SS algorithm to recover $x_{1},~{}...,~{}x_{L}$. In contrast, the S3PR problem defined by Equation (1) does not lend itself to similar cascaded solutions as existing SS algorithms struggle to separate $|\mathbf{A}x_{1}|^{2}$ and $|\mathbf{A}x_{2}|^{2}$ when $x_{1}$ and $x_{2}$ are drawn from the same distribution/class; see the reconstructions in Section 5.

Phase retrieval and blind demodulation

In various imaging applications, one records measurements of a signal $x_{1}$ that is illuminated by an unknown signal $x_{2}$. In this context, the measurement model becomes

where $\circ$ denotes the Hadamard (elementwise) product. While solutions to this problem exist , the fact that mixing occurs before the non-linearity makes the problem fundamentally different from S3PR.

Partially coherent phase retrieval and source recovery

When one images a fixed object $x$ with multiple unknown Fourier plane (Köhler) illumination sources, one is presented with a constrained S3PR problem

where $e_{1}$, $e_{2}$, … $e_{L}$ are unknown 1-hot vectors representing the location of an illumination source and $\mathbf{F}$ is a 2-D discrete Fourier transform matrix which models their propagation to the object plane. By capturing multiple measurements of this form at different stand-offs (which corresponds to selecting a particular measurement matrix $\mathbf{A}$) one can reframe this problem as deconvolution and solve it using conventional algorithms .

2 Source separation

Traditional SS algorithms assume that the sources to be separated are statistically independent and that there are as many or more observations as there are unknowns . When this is not the case, the problem is known as under-determined SS and is much more challenging.

Under-determined SS, which features prominently in reflection removal and hyperspectral imaging, has been accomplished by imposing priors on the reconstructed signals, for instance that they are sparse or group sparse in some basis or have other structural properties that can be exploited . Under-determined SS can also be accomplished with convolutional neural networks .

3 Inverse problems with deep generative models

Deep generative models have been used to solve a variety of under-determined inverse problems.

The idea of solving the compressive sensing problem by recovering the latent vectors of a deep generative model was first proposed in . This work has been followed up by a number of papers which have sought to improve the speed and generalizability of the method.

Phase retrieval

PR with generative models was introduced in and . This method has since been accelerated in and applied to various PR problems in .

We note in passing that optimizing the latent variables of untrained networks, following ideas proposed in , can also be used to impose priors in order to help solve the PR problem .

Other inverse problems

For the sake of completeness, we note that optimizing the latent variables of a deep generative models has also been used to perform blind demodulation , blind deconvolution , matrix decomposition , and source separation . We also note that optimizing the latent variables of untrained neural networks has been used for a variety of image decomposition tasks, including source separation .

Naive S3PR solutions

Here we describe two naive solution with which one might approach S3PR. As we will demonstrate in Section 5, the first approach is largely ineffective. The second does not generalize.

one could first use an under-determined SS algorithm to estimate $|\mathbf{A}x_{1}|^{2},~{}|\mathbf{A}x_{2}|^{2},~{}\ldots,~{}|\mathbf{A}x_{L}|^{2}$ and then use a PR algorithm ($L$ times) to recover $x_{1},~{}x_{2},~{}\ldots,~{}x_{L}$ from these estimates. That is, one could solve potentially solve S3PR by performing under-determined SS followed by PR.

The aforementioned SS step requires one to separate measurements $|\mathbf{A}x_{1}|^{2},~{}|\mathbf{A}x_{2}|^{2},~{}\ldots,~{}|\mathbf{A}x_{L}|^{2}$ from their sum. Representative measurements and their sum are presented in Figure 1. Separating such signals is a highly under-determined SS problem for which standard SS algorithms, like non-negative matrix factorization, do not apply. Instead, one needs to use prior information, e.g. sparsity in some basis/dictionary, about the measurements to unmix them.

As a best attempt at dictionary-based under-determined SS, for each dataset and measurement matrix pair we learn a unique 500 element dictionary, $\mathbf{D}$ (for a total of 6 dictionaries). Each dictionary was formed using the K-SVD algorithm and was optimized to form one-hot representations of $60\,000$ training measurements, which were specific to the dataset and the measurement matrix.

With a dictionary in hand, we then perform SS by finding an $L$-hot representation of $y$ in this dictionary using orthogonal matching pursuit (OMP) . That is, we use OMP to approximately solve

where $\|\alpha\|_{0}$ is the cardinality of $\alpha$. The result allows us to form the estimates

where $\alpha_{l}$ is the $l^{th}$ non-zero element of $\alpha$ and $D_{l}$ is the column of $\mathbf{D}$ (i.e., the element of the dictionary) associated with $\alpha_{l}$.

Following SS, we reconstruct $x_{1},~{}x_{2},~{}\ldots,~{}x_{L}$ by solving $L$ phase retrieval problems by applying gradient descent to the loss

We also experimented with the widely used Gerchberg–Saxton algorithm but found that it ran far slower and offered worse performance.

As we will demonstrate in Section 5, this approach is largely ineffective at solving S3PR.

2 Discriminative neural networks

One could potentially use discriminative deep neural networks to perform the source separation step within the above algorithm. One could even use a discriminative neural network to directly learn a mapping from mixed measurements, $y$, to images, $x_{1},~{}x_{2},\ldots,x_{L}$. However, with either approach such a network would become specific to the forward model, the number of signals in the mixture, the distribution of the signals in the mixture, and potentially the signal-to-noise ratio (SNR) – obtaining the results in our paper would have required training 36 separate networks.

In contrast, our proposed method generalizes across forward models, SNRs, and the number of signals in the mixture – we trained only two networks.

S3PR using deep generative models

A better way to perform S3PR is to leverage deep generative models as priors.

We recover images $x_{1},~{}x_{2},~{}\ldots,~{}x_{L}$, with our estimates $\hat{x}_{l}=G(\hat{z}_{l})$, from $y$ by solving the optimization problem

using an alternating descent algorithm. That is, we iteratively compute the loss (6) and take a gradient step (with momentum) with respect to $z_{1}$, then compute the loss and take a gradient step (with momentum) with respect to $z_{2}$, etc. In practice, we found alternating descent (using the ADAM optimizer ) provided a near monotonic reduction of the loss and ran in two minutes on an Nvidia Titan RTX GPU.

For time/resource-sensitive applications, one could also use an alternating projection algorithm or more advanced methods like AMP or ADMM , which can provide significantly faster convergence.

Generative model

We train one such DC-GAN network to produce MNIST digits and another to produce Fashion MNIST articles of clothing. Each network was based off the Pytorch implementation of DC-GAN from and was trained using the code’s default parameters. The networks were trained using the training portion of their respective datasets and tested, as described in the next section, on a subset of the testing portion.

Experimental results

We now apply the two proposed methods to Gaussian, coded diffraction pattern (CDP), and Fourier measurements of images of numbers and articles of clothing. We denote the under-determined SS followed by PR approach described in Section 3.1 with “USS + PR”. We call the deep generative model based approach from Section 4 “Deep S3PR”.

We apply the algorithms to Gaussian, CDP, and Fourier measurements to generate Figures 1, 2, and 3, respectively. To generate each of the quantitative results presented in Tables 1, 2, and 3 we apply the algorithms to 10 sets of images and compute the average normalized mean squared error (NMSE). To account for the labeling ambiguity, that is the solution $\hat{x}_{1}=A$, $\hat{x}_{2}=B$ is equivalent to the solution $\hat{x}_{1}=B$, $\hat{x}_{2}=A$, we report the loss associated with the ordering of the solutions that produces the minimum error. There is a sign ambiguity, that is $\hat{x}_{1}=A$ is an equivalent solution to $\hat{x}_{1}=-A$, that we similarly account for. Likewise, for Fourier measurements, we account for the flip ambiguities of the solutions by searching, over all flips left-right and up-down, for the one that minimizes the error.

Algorithm settings

Both the sequential USS + PR algorithm and Deep S3PR method are implemented in Pytorch. Their respective optimization problems are solved with the ADAM optimizer with a learning rate of $0.02$. ADAM’s momentum decay terms $b_{1}$ and $b_{2}$ are set to their default values of $0.9$ and $0.999$ respectively. For USS + PR, we minimize the under-determined SS loss (3) by running the ADAM optimizer for $2\,000$ iterations. The PR loss (5) is similarly minimized by running the ADAM optimizer for $2\,000$ iterations. For Deep S3PR we minimize the loss (6) by running the ADAM optimizer for $2\,000$ iterations. For both USS + PR and Deep S3PR we perform 5 restarts. For USS + PR, at each restart the estimates $\alpha$ and $x_{1},~{}x_{2},~{}\ldots,~{}x_{L}$ are initialized with random i.i.d. Gaussian vectors with mean zero and unit variance. Similarly, for Deep S3PR at each restart the latent vectors $z_{1},~{}z_{2},~{}\ldots,~{}z_{L}$ are initialized with random i.i.d. Gaussian vectors with mean zero and unit variance. For both, we use the result with the smallest residual error, $\|y-\sum_{l}|\mathbf{A}\hat{x}_{l}|\|^{2}$, as our final solution.

1 Gaussian measurements

We first test the proposed S3PR methods on complex-valued Gaussian measurement matrices. The elements of our measurement matrices are drawn from an i.i.d. circular Gaussian distribution (the real and imaginary parts of each element are drawn from i.i.d. distributions).

Figure 1 demonstrates that Deep S3PR is very effective with Gaussian measurement matrices. Even with mixtures of four images, Deep S3PR produces near perfect reconstructions of MNIST digits and recognizable, though imperfect, reconstructions of Fashion MNIST articles of clothing as well. In contrast, the sequential solution to S3PR produces significant errors with even just two images. The quantitative results, presented in Table 1, mirror these findings.

2 Coded diffraction pattern measurements

We next test our methods on simulated CDP measurements, which were first proposed in . The CDP measurement matrix can be written as

where $\mathbf{F}$ represents the two dimensional Fourier transform and $\mathbf{D}_{1},~{}\mathbf{D}_{2},~{}\mathbf{D}_{3},$ and $\mathbf{D}_{4}$ are diagonal matrices whose diagonal entries are drawn uniformly from the unit circle in the complex plane.

Figure 2 and Table 2 demonstrate that Deep S3PR is largely effective with CDP measurement matrices as well. Meanwhile, as it did with Gaussian measurements, the sequential solution to S3PR makes errors starting with just two images.

3 Fourier measurements

Finally, we test the proposed methods with Fourier measurements, which is arguably their most important use case. Fourier measurements form the basis of most coherence diffraction imaging systems as well as various correlation-based imaging systems . See the supplement for more information.

When dealing with Fourier measurements, we minimize the loss (6) by minimizing

where $\star$ denotes autocorrelation. The equivalence of (6) and (12) follows from Parseval’s theorem and the relationship $\mathbf{F}(x\star x)=|\mathbf{F}x|^{2}$. We found this formulation of the problem offered faster convergence than directly minimizing (6).

Results

Fourier measurements prove to be significantly more challenging than Gaussian or CDP measurements. Figure 3 demonstrates that with three MNIST images, Deep S3PR starts to make errors. While the technique can often reconstruct the general shape of Fashion MNIST images, there are artifacts in most reconstructions. Table 3 shows these errors are reflected in the average NMSE as well. With Fourier measurements the sequential USS + PR algorithm fails with even two images.

Discussion

This work introduces and formalizes the simultaneous source separation and phase retrieval (S3PR) problem. It then demonstrates how S3PR can be solved by optimizing the latent variables of a deep generative model. We effectively apply the proposed Deep S3PR technique to various mixtures of Gaussian, CDP, and Fourier measurements.

By leveraging the powerful image priors encoded in a pretrained deep generative model, Deep S3PR is able to solve a problem that stymies classical, dictionary-based algorithms. Moreover, because the generative model is problem agnostic, the proposed method generalizes across forward models, mixtures sizes, noise levels and more. This stands in stark contrast to a discriminative neural network approach, which would need to be retrained each time one of these parameters changes.

Finally, Deep S3PR represents a major step forward in addressing a long-standing challenge in computational optics and we look forward to experimentally validating its benefits with physical experiments; such as imaging extended objects through scattering media or around corners . Deep S3PR stands to enable major strides in these and other applications.

Acknowledgments and Disclosure of Funding

C.M. was supported by an appointment to the Intelligence Community Postdoctoral Research Fellowship Program at Stanford University administered by Oak Ridge Institute for Science and Education (ORISE) through an interagency agreement between the U.S. Department of Energy and the Office of the Director of National Intelligence (ODN). G.W. was supported by an NSF CAREER Award (IIS 1553333), a Sloan Fellowship, and a PECASE by the ARL.

References

S3PR in computational optics

Phase retrieval (PR) is a fundamental part of many computational optics/imaging systems. For instance, it is used in microscopy to exceed the diffraction limit , X-ray crystallography to perform imaging without a lens , and astronomical imaging to see through atmospheric aberrations .

Implicit in many of these systems is an assumption that there is a single coherent illumination source. (A source is said to be coherent if its constituent fields maintain a constant phase offset, and thus produce construct and destructive interference.) When this assumption is broken, and there are instead multiple independent coherent sources, the reconstruction problem becomes S3PR rather than PR. Accordingly, our inability to effectively solve S3PR places practical limitations on the performance of these systems.

We next describe the physics and mathematics behind a real-world imaging systems where S3PR limits performance. Similar limits play out in many other imaging applications.

Speckle correlation imaging through scattering media was first introduced in 2012 and has since been extended in and many other works. It is based off of ideas first developed for astronomical imaging nearly 50 years ago .

This family of techniques allows one to non-invasively image simple objects ( reconstructed phantoms consisting of letters, numbers, and smiley faces) through thin scattering media, like a layer of soft tissue. The technique can reconstruct millimeter-scale features and requires only a temporally coherent, spatially incoherent light source and a camera . The physical setup associated with speckle correlation imaging is illustrated in Figure 6(a).

Speckle correlation imaging through scattering media is based on a physical phenomenon known as the angular memory-effect . It states that if an object $O$ illuminates a scattering media with a temporally coherent light, each point of the scattering media will produce a spatially invariant interference pattern $S$ (known as “speckle”), so long as, from the perspective of the scattering media, the angle subtended by the object is small.

Because of the memory-effect–induced spatial invariance of the speckle, the measurement model associated with Figure 6(a) is convolutional and is described by

where $S$ is the unknown speckle pattern, $O$ is the object of interest, and $*$ denotes convolution.

Because the autocorrelation function of a speckle pattern is Dirac-like , ignoring DC terms we have that

With this estimate of the autocorrelation of the hidden object $O$ in hand, we can reconstruct the object by using phase retrieval algorithms and the relationship

where $\mathcal{F}(\cdot)$ denotes the Fourier transform operator.

2 Speckle correlation imaging over an extended field-of-view

Now consider two small, spatially separated objects, $O_{1}$ and $O_{2}$, as illustrated in Figure 6(b). Each object itself subtends a small angle, and thus experiences spatially invariant speckle. However, they are far apart from one another and so each experiences a different speckle pattern. Thus, the measurement formation model is described by

where $S_{1}$ and $S_{2}$ denote two independent speckle realizations.

Because the two speckle realizations are uncorrelated, again ignoring DC terms, $S_{1}\star S_{2}\approx 0$. Thus, the autocorrelation of $I$ becomes

Taking the Fourier transform of the result, we arrive at

In this way, speckle correlation imaging over an extended field of view naturally leads to an S3PR reconstruction problem. Until now, correlation-based imaging through scattering media systems have avoided or ignored S3PR by either (1) restricting themselves to imaging only a single small object, (2) using invasive illumination sources that illuminate only one portion of a hidden object at a time , or (3) capturing thousands of training images and throwing a discriminative neural network at the problem . Our work makes strides towards removing these limitations.

Additional Results

Figures 4, 5, and 6 and Tables 4, 5, and 6 provide additional experimental results, captured at an SNR of 15 rather than 50. The resulting Gaussian and coded diffraction pattern (CDP) reconstructions are similar to the higher SNR reconstructions from the main text, while the low SNR Fourier results exhibit a few more artifacts.