Phase Retrieval via Incremental Truncated Wirtinger Flow

Introduction

Despite the hardness, impressive theoretical and empirical performance guarantees have been obtained recently by making some assumptions about the model, such as assuming the sensing vectors to be i.i.d. samples from, say $\mathcal{N}(0,I)$ or $\mathcal{CN}(0,I)$. The first work in this direction was , which employed the squared loss function and attempted to solve

Follow-up works such as , , made progress on the computational complexity front. All of these works developed algorithms to solve the problem (2) directly without performing the lifting step, using a first-order optimization method for iterative refinement, after an appropriate initialization. The work also provided stability guarantees if the measurements are corrupted with noise.

However, each iteration in these algorithms requires one pass through the entire data. This can be highly undesirable when the dimensions of the problem are large, since a single update can require a large amount of time, in part due to the communication delays introduced when the entire data does not fit in the available memory. Large dimensions naturally arise in the phase retrieval problem since the object of interest usually represents an image, so $n$ is the product of the image dimensions.

In this paper, we build on the idea of Truncated Wirtinger Flow (TWF) from and modify it to obtain the Incremental Truncated Wirtinger Flow (ITWF). By incremental, we mean that each iteration of the algorithm only accesses one randomly chosen data point, i.e. one sensing vector and the corresponding measurement. Thus, each iteration of ITWF is cheaper than that of TWF by a factor $m$, similar to what happens, e.g., by going from full gradient descent to stochastic gradient descent (SGD) in the case of standard empirical risk minimization problems. Unfortunately, this benefit is not obtained readily for the problem at hand, since the truncation performed at each iteration in makes use of a threshold that is a function of all the sensing vectors and measurements. Excluding the truncation leads to a severe hit in the performance as observed in . Furthermore, similar to SGD, sampling one data point instead of all data points introduces variance in the descent direction. The main contribution of this paper is the design of the incremental method ITWF that matches the excellent performance of TWF in terms of the statistical complexity, computational complexity and robustness to noisy measurements despite being an incremental method. In fact, our numerical experiments demonstrate that ITWF far surpasses TWF on the computational complexity front.

Table 1 provides a comparison of algorithms. Though our original intention was to simply develop an incremental version of TWF to allow efficient handling of large data, we find that ITWF provides a remarkable speedup as compared to TWF, even at $n\approx 1000$. To provide the reader an idea about the speed-up, we mention here some observations from the numerical experiments in Section 4. In Example 2, after an initialization stage requiring 10 passes through the data, the second stage of TWF requires at least 120 further passes to get a high accuracy solution. In contrast, the second stage of ITWF recovers a high accuracy solution in less than 15 passes. Example 3 presents an even more compelling case. Here, after 50 passes through the data for initialization, the solution returned by TWF after 50 further passes can instead be obtained by making 3 passes using ITWF.

There has also been a lot of recent work on different formulations of the phase retrieval problem. In the sparse phase retrieval problem, it is additionally assumed that the vector $x$ has only a few non-zero entries, and algorithms are sought for which the sample complexity and computational complexity have optimal dependence on the number of non-zero entries. The interested reader is referred to works such as , , . Phase retrieval under the assumption of structured sensing vectors has also been a topic of interest, e.g. coded diffraction patterns , STFT . We also present numerical experiments in this paper demonstrating the performance of ITWF when the sensing vectors are structured. Recent developments on different tractable formulations of the phase retrieval problem have been compiled in the survey article . Finally, while we focus on devising an incremental update for the iterative refinement stage of TWF, the initialization stage can also be made incremental in a straightforward manner by using incremental algorithms for PCA, such as the algorithm from .

Main Idea

In the following section, we first describe the algorithm TWF from , and then describe the proposed algorithm ITWF.

Since we can only hope to recover the solution upto a global phase, we define $\text{dist}(\bm{z},\bm{x})$ to be $\min_{\varphi\in[0,2\pi)}\|e^{-j\varphi}\bm{z}-\bm{x}\|,$ and we implicitly assume in the remainder of the paper that $\bm{z}$ is $e^{-j\varphi(\bm{z})}\bm{z}$, where $\varphi(\bm{z})$ is the argument of the above minimization problem. Then, we use $\bm{h}$ to denote $\bm{z}-\bm{x}.$ Thus, $\|\bm{h}\|=\text{dist}(\bm{z},\bm{x}).$

(1) Truncated Spectral Initialization Set $\bm{z}^{(0)}$ to be the principal eigenvector (appropriately scaled) of

with $\alpha_{y}$ set to, say 3. The point $\bm{z}^{(0)}$ satisfies

for any $\delta>0$, as long as $m/n$ exceeds a sufficiently large constant. For intuition and proof of this fact, we refer the reader to .

(2) Truncated Wirtinger Flow For $t=0,1,\dots,T-1$, perform the update

and the events $\mathcal{E}^{i}_{1,t}$ and $\mathcal{E}^{i}_{2,t}$ are defined as follows,

The idea is to perform an update similar to gradient descent. The truncation results in dropping those indices for which the numerator and denominator magnitudes are very different from their expected values respectively, since such terms can exert a large atypical influence causing the full gradient to point in an undesirable direction.

2 Incremental Truncated Wirtinger Flow

The straightforward way of turning the above algorithm into an incremental one would be to replace (4) by:

where $i_{t}$ would be chosen uniformly at random from $\{1,2,\dots,m\}$. However, as can be seen from (5), checking if $\mathcal{E}^{i}_{2,t}$ has occurred for any $i$ requires a full pass through the entire data. Thus, (6) is as costly as (4). To address this difficulty, we replace $\mathcal{E}^{i}_{2,t}$ by the event $\mathcal{E}^{i}_{3}$, which is defined as

where $\alpha_{x}$ is a constant to be chosen. This is amenable for an incremental update, which we call Incremental Truncated Wirtinger Flow (ITWF):

Note in fact that not only does $\mathcal{E}^{i}_{3}$ not depend on the entire data, but it is non-adaptive, i.e. it does not depend on $\bm{z}^{(t)}$. So, $\mathcal{E}^{i}_{3}$ simply discards measurements in which $|\bm{a}_{i}^{T}\bm{x}|$ deviates a lot from $\|\bm{x}\|$, similar to the rule used during the initialization stage.

The main implication of (5) used in is that the magnitude of the term $\bm{a}_{i}^{T}\bm{h}$ can be controlled, which turns out to be crucial in obtaining the optimal sample complexity $m=O(n)$. Since ITWF cannot employ this truncation rule, we are not able to control the magnitude of $\bm{a}_{i}^{T}\bm{h}$. However, replacing (5) by $\mathcal{E}^{i}_{3}$ does not result in any deterioration in the sample complexity and computational complexity.

Furthermore, we would like to point out an empirical observation that we did not observe a significant difference in the numerical experiments even if we only employed truncation based on $\mathcal{E}_{1,t}^{i}$, i.e. by excluding $\mathcal{E}^{i}_{3}$! In fact, we also observed that TWF also performs similarly whether $\mathcal{E}_{2,t}^{i}$ is included or not. This suggests that while the theoretical analysis of TWF and ITWF require these additional truncation events, it might be possible to prove the convergence results via a different line of analysis that does not require introducing these events. However, we observed that truncation based on $\mathcal{E}_{1,t}^{i}$ is indeed crucial in the numerical experiments.

Main Results and Discussion

We focus on the real valued case for simplicity of exposition, and for concreteness, we fix the constants $\alpha_{z}^{\text{lb}}$, $\alpha_{z}^{\text{ub}}$ and $\alpha_{x}$ to be $0.3$, $5$ and $5$ respectively. The following two theorems are counterparts of the two main results in . The first result, Theorem 1, focuses on the noiseless case (1) and provides a linear convergence guarantee.

Under noiseless measurements (1) with $\{\bm{a}_{i}\}_{i=1}^{m}$ $\sim\mathcal{N}(0,I)$ independent, there exist universal constants $C,c_{0},c_{1},c_{2}>0$ and $0<\rho,\nu<1$, such that with probability at least $1-Cm\exp(-c_{1}n)$ and $\mu=c_{2}/n$, the iterates in Algorithm 1 satisfy {IEEEeqnarray}l E_I^t[dist^2(z^(t),x)] ≤ν(1 - ρn)^t∥x∥^2,\IEEEeqnarraynumspace if $m\geq c_{0}n$.

Thus, the (mean squared error) MSE is reduced by a factor $\left(1-\rho/n\right)^{m}$ after one pass through the data.

Note that the expectation is only with respect to the algorithm randomness, not with respect to the data randomness $\{\bm{a}_{i}\}_{i=1}^{m}$. This is crucial since the data could be provided as it is, thus necessitating the need for a convergence guarantee that holds with high probability with respect to the data randomness. Since the convergence is linear, a simple application of Markov’s inequality already provides a strong convergence guarantee in which the high probability also refers to the algorithm randomness.

Also note that linear convergence is achieved for our setup even with an incremental method since the effect of the variance of the stochastic gradient can be controlled without having to choose a step-size that decreases with iteration, while in general, incremental methods suffer from slow convergence unless variance-reducing modifications are used. The reason for this is explained in the next subsection.

Instead of (1), if the measurements are noisy such that

where $\eta_{i}$ denotes the noise term (need not be stochastic), then we have the following theorem which provides a stability guarantee.

Under noisy measurements (8) with $\{\bm{a}_{i}\}_{i=1}^{m}$ $\sim\mathcal{N}(0,I)$ independent and the noise satisfying $\|\bm{\eta}\|_{\infty}\leq\epsilon_{\eta}\|\bm{x}\|^{2}$ for some small constant $\epsilon_{\eta}>0$, there exist universal constants $C,c_{0},c_{1},c_{2}>0$ and $0<\rho,\nu<1$, such that with probability at least $1-Cm\exp(-c_{1}n)$ and $\mu=c_{2}/n$, the iterates in Algorithm 1 satisfy {IEEEeqnarray}l E_I^t[dist^2(z^(t),x)] ≲∥η∥2m∥x∥2 + (1 - ρn)^t∥x∥^2,\IEEEeqnarraynumspace if $m\geq c_{0}n$.

As described in , this result can be applied to the case when the measurements are obtained independently according to a Poisson noise model:

and the solution satisfies $\|\bm{x}\|^{2}\geq\log^{3}m$ (required to ensure that the condition $\|\bm{\eta}\|_{\infty}\leq\epsilon_{\eta}\|\bm{x}\|^{2}$ holds with high probability).

Consider the noiseless case. The reason why we can expect linear convergence from ITWF can be understood by the following thought experiment. Say, after the $t^{\text{th}}$ iteration of ITWF, we had the luxury of obtaining a new measurement by sampling $\bm{a}_{t}$ independently from the population distribution $\mathcal{N}(0,I)$, instead of being restricted to sample from the empirical distribution $\{\bm{a}_{i}\}_{i=1}^{m}$.

Let $\bm{h}$ denote $\bm{z}^{(t)}-\bm{x}$, and $\mathcal{E}_{t}=\mathcal{E}_{1,t}\cap\mathcal{E}_{3}=\left\{\alpha_{z}^{\text{lb}}\leq\frac{|\bm{a}_{t}^{*}\bm{z}^{(t)}|}{\|\bm{z}^{(t)}\|}\leq\alpha_{z}^{\text{ub}}\right\}\cap\left\{\frac{|\bm{a}_{t}^{T}\bm{x}|}{\|\bm{x}\|}\leq\alpha_{x}\right\}$ denote the overall truncation event. Then the expected distance to the optimal solution after performing the ITWF update (7) using the new independent measurement, is {IEEEeqnarray}rCl E_t[dist^2(z^(t+1),x)] & = ∥h∥^2 - 4μE_t[ |a_t^Th|^2(1 + atTxatTz(t))1_E_t] + 4μ^2E_t[ ∥a_t∥^2|a_t^Th|^2(1 + atTxatTz(t))^2 1_E_t]

Since $\bm{a}_{t}\sim\mathcal{N}(0,I)$ independent of $\bm{h},\bm{z},\bm{x}$, it is not difficult to show that if $\|\bm{h}\|$ is sufficiently small compared to $\|\bm{x}\|$, {IEEEeqnarray*}l E_t[ |a_t^Th|^2(1 + atTxatTz)1_E_t] ≥(1 - c_1)∥h∥^2 , E_t[ ∥a_t∥^2|a_t^Th|^2(1 + atTxatTz)^2 1_E_t ]≲n ∥h∥^2, for an appropriate constant $c_{1}$. As can be observed from the latter inequality, the variance of the stochastic gradient is proportional to $\|\bm{h}\|^{2}$. As a result, the variance is not only bounded, but reduces as we get closer to the solution, which allows us to choose a non-diminishing step-size $\mu=\Theta\left(\frac{1}{n}\right)$, resulting in linear convergence:

In the above thought experiment, we had the luxury of an infinite amount of measurements at our disposal. Back in reality, we have the restriction of using a finite set of $m$ measurements. Since we need to keep sampling from this set to simulate the stochastic gradient descent from our thought experiment, the proof of Theorem 1 mainly involves showing that we can draw similar conclusions despite the finiteness of the data set and the resulting dependence between the sensing vectors and the path traversed by the algorithm.

Numerical Experiments

To provide evidence for the performance of ITWF and comparisons to existing algorithms, we provide numerous examples in this section.

We compare the sample complexity of TWF and ITWF empirically when the sensing vectors $\bm{a}_{i}$ are i.i.d. $\mathcal{N}(0,I)$. The dimension $n$ is chosen to be 1000, and the number of measurements $m$ is varied from $2n$ to $6n$, and success is declared if the relative root mean squared error $\frac{\text{dist}(\bm{z},\bm{x})}{\|\bm{x}\|}$ is less than $10^{-5}$ within 1000 passes through the data. The initialization uses 50 truncated power iterations. It can be seen from Figure 1, which is obtained by averaging over 100 Monte Carlo trials at each value of $m$, that the empirical success rate of ITWF is comparable with that of TWF, even slightly better.

2 Example 2: Computational Complexity (Stage II)

We consider the same setup as the previous example, with $m$ fixed to be 8000. Figure 2 shows the relative root mean squared error $\frac{\text{dist}(\bm{z},\bm{x})}{\|\bm{x}\|}$ of three algorithms, each run with the best step size, as a function of the number of passes through the data. All algorithms were initialized using 10 power iterations. The black line corresponds to the without-replacement variant of ITWF, in which all data points are visited exactly once in every block of $m$ iterations, each time in an independently chosen random order. The performance of the without-replacement variant is marginally better than the with-replacement variant, as has also been observed in numerous other contexts. However, the main point of the example is to show that ITWF offers substantial improvements in computational complexity over TWF. Owing to the better performance of the without-replacement variant, we adopt this sampling method for the remainder of the numerical experiments.

3 Example 3: Structured Sensing Vectors

To demonstrate the performance of the algorithm on a real signal when the assumption of random Gaussian sensing vectors do not hold, we consider another example from . An image of Stanford main quad of size 320 $\times$ 1280 pixels is used, and measurements are obtained via a set of $L$ coded diffraction patterns as

where $\bm{F}$ is the DFT matrix and $\bm{D}^{(l)}$ is a diagonal matrix (representing a mask) containing independent entries, each uniformly distributed over $\{+1,-1,j,-j\}$. The number of measurements is $nL$ with $L=12$, and the notation $|\cdot|^{2}$ denotes elementwise magnitude squared.

As in , we initialize the algorithm with 50 truncated power iterations, at the end of which the relative root mean squared error (rmse) is $0.35$. Another pass through the data using TWF gets it down to $0.15$. Increasing the number of passes to 2, 5 and 10 achieves $7.1\times 10^{-2}$, $2.3\times 10^{-2}$ and $6.6\times 10^{-3}$ respectively. Making one pass through the data using ITWF gets it down from $0.35$ to $8.2\times 10^{-3}$. Increasing the number of passes to 2, 5 and 10 achieves $3.0\times 10^{-4}$, $1.4\times 10^{-8}$ and $9.0\times 10^{-16}$ respectively.

Remark: The sensing matrices in this example are highly structured (DFT, diagonal), due to which computing the sum of the $n$ gradients for one mask can be accomplished with $O(n\log n)$ computations instead of $O(n^{2})$, by utilizing FFT algorithms. Hence, the motivation for ITWF that one iteration can be made $m$-times cheaper does not hold in this case. Hence we consider an increment to be the set of measurements corresponding to one mask ($L$ measurements) instead of one measurement. This means that each iteration of ITWF is $L$-times cheaper than that of TWF where $L$ is the total number of masks; thus one iteration of TWF is computationally equivalent to $L$ iterations of ITWF.

4 Example 4: Noisy measurements

In this example, we generate noisy measurements as follows. As before, the sensing vectors are chosen to be i.i.d. $\mathcal{N}(0,I)$. The measurements are generated according to (9). Theorem 2 effectively says that ITWF achieves

where SNR is the signal-to-noise ratio $\frac{\sum_{i=1}^{m}(\bm{a}_{i}^{T}\bm{x})^{4}}{\sum_{i=1}^{m}\eta_{i}^{2}}$, which is approximately $\frac{3m\|\bm{x}\|^{4}}{\|\bm{\eta}\|^{2}}.$

Under the Poisson model, since $\sum_{i=1}^{m}\eta_{i}^{2}\approx\sum_{i=1}^{m}|\bm{a}_{i}^{T}\bm{x}|^{2}\approx m\|\bm{x}\|^{2}$, we refer to $3\|\bm{x}\|^{2}$ as the SNR. Comparing the final relative MSE of TWF and ITWF at various values of SNR Figure 4 shows that the final relative MSE (LHS of (10)) of ITWF is in fact nearly equal to that of ITWF at all values of SNR.

Remarks and Future Directions

While loosely bounding the constants arising during the analysis provides a recommendation for the step-size (for the random Gaussian model) which is $\approx\frac{0.00016}{n}$, this can be slightly pessimistic. In our numerical experiments, we find that larger step sizes also work.

Developing incremental methods that are able to extract the benefits of truncation for other problems such as sparse phase retrieval, low-rank matrix recovery from linear/quadratic measurements would be highly interesting.

References

Appendix

The proof requires introducing some constants, and we collect their definitions before proceeding to the proof.

where $\xi\sim\mathcal{N}(0,1)$ and $\gamma,$ $\alpha_{x}$, $\alpha_{z}^{\text{ub}}$ and $\alpha_{z}^{\text{lb}}$ can be chosen to be $5$, $0.3$, $5$ and $5$ respectively.

The term $\delta_{\text{init}}$ appearing in the proofs can be considered to be equal to 0.1 for concreteness.

The intermediate lemmas and propositions will establish statements that hold for all $\bm{z}$ that are in a neighborhood of $\bm{x}$ as follows:

Note that the initialization stage can be used to guarantee that $\bm{z}^{(0)}$ is in this neighborhood of $\bm{x}$.

Then, to prove the theorem, it suffices to prove Proposition 1. The reason why proving this proposition suffices is as follows. Consider the statement (1) in Proposition 1. We can take expectation of the RHS with respect to $i_{t-1}$, and then apply Proposition 1 again to get a similar relation for the previous iteration. Continuing this process till we arrive at the initialization point $\bm{z}^{(0)}$, and noting that the initialization stage guarantees that $\text{dist}(\bm{z}^{(0)},\bm{x})$ is at most a small constant times $\|\bm{x}\|$, we arrive at the statement of Theorem 1. This concludes the proof of Theorem 1.∎

Let $\bm{h}$ denote $\bm{z}^{(t)}-\bm{x}$, and $\mathcal{E}^{i}_{t}=\mathcal{E}^{i}_{1,t}\cap\mathcal{E}^{i}_{3}$ denote the overall truncation. The ITWF update (7), after some simple algebraic manipulations, gives us {IEEEeqnarray}rCl E_i_t[dist^2(z^(t+1),x)] & = ∥h∥^2 - 4μm ∑_i=1^m |a_i^Th|^2(1 + aiTxaiTz(t))1_E^i_t + 4μ2m∑_i=1^m ∥a_i∥^2|a_i^Th|^2(1 + aiTxaiTz(t))^2 1_E^i_t

Lemmas 1, 2 and 3 show that for appropriate $\zeta_{1}$, $\zeta_{2}$ and $\zeta_{3}$ and any small $\delta>0$, the following two relations hold simultaneously for all $\bm{z}$, $\bm{h}$ and $\bm{x}$ with high probability: {IEEEeqnarray*}l 1m∑_i=1^m |a_i^Th|^2(1 + aiTxaiTz)1_E^i_t ≥(0.99 - ζ_1 - ζ_2 - ζ_3 - 3δ)∥h∥^2, and {IEEEeqnarray*}l 1m∑