Compressed Sensing with Deep Image Prior and Learned Regularization

Introduction

Sparsity approaches have proven successful, but more complex models with additional structure have been recently proposed such as model-based compressive sensing and manifold models . Bora et al. showed that deep generative models can be used as excellent priors for images. They also showed that backpropagation can be used to solve the signal recovery problem by performing gradient descent in the generative latent space. This method enabled image generation with significantly fewer measurements compared to Lasso for a given reconstruction error. Compressed sensing using deep generative models was further improved in very recent work . Additionally a theoretical analysis of the nonconvex gradient descent algorithm was proposed by Hand et al. under some assumptions on the generative model.

Inspired by these impressive benefits of deep generative models, we chose to investigate the potential application of such methods for medical imaging, a canonical application of compressive sensing. A significant problem, however, is that all these previous methods require the existence of pre-trained models. While this has been achieved for various types of images, e.g. human faces of CelebA via DCGAN , it remains significantly more challenging for medical images . Instead of addressing this problem in generative models, we found an easier way to circumvent it.

Surprising recent work by Ulyanov et al. proposed Deep Image Prior (DIP), which uses untrained convolutional neural networks. In DIP-based schemes, a convolutional neural network generator (e.g. DCGAN) is initialized with random weights; these weights are subsequently optimized to make the network produce an output as close to the target image as possible. This procedure is unlearned, using no prior information from other images. The prior is enforced only by the fixed convolutional structure of the generator network.

Generators used for DIP are typically over-parameterized, i.e. the number of network weights is much larger compared to the output dimension. For this reason DIP has empirically been found to overfit to noise if run for too many iterations: The reconstruction error initially decreases and then plateaus, at approximately $500$ iterations, as the network fits the original image. Then, at roughly $10,000$ iterations, the error decreases further, as the network starts fitting the noise . Early stopping is a heuristic intended to terminate the optimization procedure within this plateau region, and avoid overfitting to noise. In this paper we theoretically prove that this overfitting phenomenon occurs with gradient descent for any signal and hence justify the use of early stopping and other regularization methods.

In Section 3 we propose DIP for compressed sensing (CS-DIP). Our basic method is as follows. Initialize a DCGAN generator with random weights; use gradient descent to optimize these weights such that the network produces an output which agrees with the observed measurements as much as possible. This unlearned method can be improved with a novel learned regularization technique, which regularizes the DCGAN weights throughout the optimization process.

In Section 4 we theoretically prove that DIP will fit any signal to zero error with gradient descent. Our result is established for a network with a single hidden layer and sufficient constant fraction over-parametrization. While it is expected that over-parametrized neural networks can fit any signal, the fact that gradient descent can provably solve this non-convex problem is interesting and provides theoretical justification for early stopping, a phenomenon justified empirically by Ulyanov et al. .

In Section 5 we empirically show that CS-DIP outperforms previous unlearned methods in many cases. While pre-trained or “learned” methods frequently perform better , we have the advantage of not requiring a generative model trained over large datasets. As such, we can apply our method to various medical imaging datasets for which data acquisition is expensive and generative models are difficult to train.

Background

A classical assumption made in compressed sensing is that the vector $x^{*}$ is $k$-sparse in some basis such as wavelet or discrete cosine transform (DCT). Finding the sparsest solution to an underdetermined linear system of equations is NP-hard in general; however, if the matrix $A$ satisfies conditions such as the Restricted Eigenvalue Condition (REC) or Restricted Isometry Property (RIP) , then $x^{*}$ can be recovered in polynomial time via convex relaxations or iterative methods. There is extensive compressed sensing literature regarding assumptions on $A$, numerous recovery algorithms, and variations of RIP and REC .

Compressed sensing methods have found many applications in imaging, for example the single-pixel camera (SPC) . Medical tomographic applications include x-ray radiography, microwave imaging, magnetic resonance imaging (MRI) . Obtaining measurements for medical imaging can be costly, time-consuming, and in some cases dangerous to the patient . As such, an important goal is to reduce the number of measurements while maintaining good reconstruction quality.

Aside from the classical use of sparsity, recent work has used other priors to solve linear inverse problems. Plug-and-play priors and Regularization by Denoising have shown how image denoisers can be used to solve general linear inverse problems. A key example of this is BM3D-AMP, which applies a Block-Matching and 3D filtering (BM3D) denoiser to an Approximate Message Passing (D-AMP) algorithm . AMP has also been applied to linear models in other contexts . Another related algorithm is TVAL3 which leverages augmented Lagrangian multipliers to achieve impressive performance on compressed sensing problems. In many different settings, we compare our algorithm to these prior methods: BM3D-AMP, TVAL3, and Lasso.

2 Compressed Sensing: Learned Approaches

While sparsity in some chosen basis is well-established, recent work has shown better empirical performance when neural networks are used . This success is attributed to the fact that neural networks are capable of learning image priors from very large datasets . There is significant recent work on solving linear inverse problems using various learned techniques, e.g. recurrent generative models and auto-regressive models . Additionally approximate message passing (AMP) has been extended to a learned setting by Metzler et al. .

Bora et al. is the closest to our set-up. In this work the authors assume that the unknown signal is in the range of a pre-trained generative model such as a generative adversarial network (GAN) or variational autoencoder (VAE) . The recovery of the unknown signal is obtained via gradient descent in the latent space by searching for a signal that satisfies the measurements. This can be directly applied for linear inverse problems and more generally to any differentiable measurement process. Recent work has built upon these methods using new optimization techniques , uncertainty autoencoders , and other approaches . The key point is that all this prior work requires pre-trained generative models, in contrast to CS-DIP. Finally, there is significant ongoing work to understand DIP and develop related approaches .

Proposed Algorithm

Our approach is to find a set of weights for the convolutional network such that the measurement matrix applied to the network output, i.e. $AG(z;w)$, matches the measurements $y$ we are given. Hence we initialize an untrained network $G(z;w)$ with some fixed $z$ and solve the following:

This is, of course, a non-convex problem because $G(z;w)$ is a complex feed-forward neural network. Still we can use gradient-based optimizers for any generative model and measurement process that is differentiable. Generator networks such as DCGAN are biased toward smooth, natural images due to their convolutional structure; thus the network structure alone provides a good prior for reconstructing images in problems such as inpainting and denoising . Our finding is that this applies to general linear measurement processes. Furthermore, our method also directly applies to any differentiable forward operator $A$. We restrict our solution to lie in the span of a convolutional neural network. If a sufficient number of measurements $m$ is given, we obtain an output such that $x^{*}\approx G(z;w^{*})$.

Note that this method uses an untrained generative model and optimizes over the network weights $w$. In contrast previous methods, such as that of Bora et al. , use a trained model and optimize over the latent $z$-space, solving $z^{*}=\operatorname*{arg\,min}_{z}\|y-AG(z;w)\|^{2}$. We instead initialize a random $z$ with Gaussian i.i.d. entries and keep this fixed throughout the optimization process.

In our algorithm we leverage the well-established total variation regularization , denoted as $TV(G(z;w))$. We also propose an additional learned regularization technique, $LR(w)$; note that without this technique, i.e. when $\lambda_{L}=0$, our method is completely unlearned. Lastly we use early stopping, a phenomenon that will be justified theoretically in Section 4.

Thus the final optimization problem becomes

The regularization term contains hyperparameters $\lambda_{T}$ and $\lambda_{L}$ for total variation and learned regularization: $R(w;\lambda_{T},\lambda_{L})=\lambda_{T}TV(G(z;w))+\lambda_{L}LR(w)$. Next we discuss this $LR(w)$ term.

2 Learned Regularization

Without learned regularization CS-DIP relies only on linear measurements taken from one unknown image. We now introduce a novel method which leverages a small amount of training data to optimize regularization. In this case training data refers to measurements from additional ground truth of a similar type, e.g. measurements from other x-ray images.

To leverage this additional information, we pose Eqn. 3 as a Maximum a Posteriori (MAP) estimation problem and propose a novel prior on the weights of the generative model. This prior then acts as a regularization term, penalizing the model toward an optimal set of weights $w^{*}$.

In this setting we want to find a set of weights $w^{*}$ that maximizes the log posterior on $w$ given $y$, i.e.,

This gives us the learned regularization term

where the coefficient $\lambda_{L}$ in Eqn. 4 controls the strength of the prior.

In the previous section, we introduced the learned regularization term $LR(w)$ defined in Eqn. 5. However we have not yet learned values for parameters $(\mu,\Sigma)$ that incorporate prior knowledge of the network weights. We now propose a way to estimate these parameters.

Assume we have a set of measurements $S_{Y}=\{y_{1},y_{2},\cdots,y_{Q}\}$ from $Q$ different images $S_{X}=\{x_{1},x_{2},\cdots,x_{Q}\}$, each obtained with a different measurement matrix $A$. For each measurement $y_{q},q\in\{1,2,...,Q\},$ we run CS-DIP to solve the optimization problem in Eqn. 3 and obtain an optimal set of weights $W^{*}=\{w^{*}_{1},w^{*}_{2},\cdots,w^{*}_{Q}\}$. Note that when optimizing for the weights $W^{*},$ we only have access to the measurements $S_{Y}$, not the ground truth $S_{X}$.

2.2 Discussion of Learned Regularization

The proposed CS-DIP does not require training if no learned regularization is used, i.e. if $\lambda_{L}=0$ in Eqn. 3. This means that CS-DIP can be applied only with measurements from a single image and no prior information of similar images in a dataset.

Our next idea, learned regularization, utilizes a small amount of prior information, requiring access to measurements from a small number of similar images (roughly $10$). In contrast, other pre-trained models such as that of Bora et al. require access to ground truth from a massive number of similar images (tens of thousands for CelebA). If such a large dataset is available, and if a good generative model can be trained on that dataset, we expect that pre-trained models would outperform our method. Our approach is instead more suitable for reconstructing problems where large amounts of data or good generative models are not readily available.

Theoretical Results

We focus on generators consisting of a single hidden-layer ReLU network with $k$ inputs, $d$ hidden units, and $n$ outputs. Using $w=(W,V)$ the generator model in this case is given by

with step size $\eta=\frac{\bar{\eta}}{\left\|y\right\|^{2}}\frac{8m}{4m+d}$ where $\bar{\eta}\leq 1$. Assuming that $d\geq Cm,$ with $C$ a fixed numerical constant, then

holds for all $\tau$ with probability at least $1-5e^{-m/2}-e^{-d/2}-e^{-4d^{\frac{2}{3}}m^{\frac{1}{3}}}$.

Our theoretical result shows that after many iterative updates, gradient descent will solve this non-convex optimization problem and fit any signal $y$, if the generator network is sufficiently wide. This occurs as soon as the number of hidden units $d$ exceeds the signal size $n$ by a constant factor. Our theorem directly applies to many compressed sensing measurement matrices, in particular any matrix obtained by subsampling the rows of an orthonormal matrix (e.g. sub-sampling a Fourier matrix). This is possible because, for any such orthonormal matrix, $AV$ has the same distribution as a Gaussian matrix with i.i.d. entries. This result demonstrates that early stopping is necessary for DIP-based methods to be successful; otherwise the network can fit any signal, including one that is noisy.

Our proof builds on theoretical ideas from Oymak et al. which provide a general framework for establishing global convergence guarantees for overparameterized nonlinear learning problems based on various properties of the Jacobian mapping along the gradient descent trajectory. While our proof leverages relevant prior work , our argument is quite specialized and intricate with new techniques such as Gordon’s Lemma. This allows us to have moderate network overparameterization that is only linear in the number of measurements, contrary to other results in the literature which require a significant amount of overparameterization. Ultimately we combine tools from empirical process theory, random matrix theory, and matrix algebra to show that, starting from a random initialization, the Jacobian mapping across all iterates has favorable properties with high probability, hence facilitating convergence to a global optima.

Experiments

Datasets: We use our algorithm to reconstruct both grayscale and RGB images. For grayscale we use the first 100 images in the test set of MNIST and also 60 random images from the Shenzhen Chest X-Ray Dataset , downsampling a $512\times 512$ crop to $256\times 256$ pixels. For RGB we use retinopathy images from the STARE dataset with $512\times 512$ crops downsized to $128\times 128$ pixels.

Baselines: We compare our algorithm to state-of-the-art unlearned methods such as BM3D-AMP , TVAL3 , and Lasso in a DCT basis . We also evaluated the performance of Lasso in a Daubechies wavelet basis but found this performed worse than Lasso - DCT on all datasets. Thus hereon we refer to Lasso - DCT as “Lasso” and do not include results of Lasso - Wavelet. We used sci-kit learn for the implementation of Lasso and code provided by the original authors for BM3D-AMP and TVAL3. A standard grid search was performed over each baseline to tune hyperparameters.

Metrics: To quantitatively evaluate the performance of our algorithm, we use per-pixel mean-squared error (MSE) between the reconstruction $\hat{x}$ and true image $x^{*}$, i.e. $\frac{\|\hat{x}-x^{*}\|^{2}}{n}$. Note that because these pixels are over the range $$, it’s possible for MSE to be greater than$1$.

Implementation: To find a set of weights $w^{*}$ that minimize Eqn. 3, we use PyTorch with a DCGAN architecture. Our network has depth 7 and uses convolutional layers with ReLU activations. We use the RMSProp optimizer with learning rate $10^{-3}$, momentum $0.9$, and $1000$ update steps for every set of measurements. These parameters are the same across all datasets. We initialize one random measurement matrix $A$ for each image.

More implementation details can be found in the appendix, such as hyperparameter search, network initializations, and early stopping criterion. Code for these experiments is available in our GitHub repository: github.com/davevanveen/compsensing_dip.

2 Experimental Results

We first evaluate the benefits of learned regularization by comparing our algorithm with and without learned regularization, i.e. $\lambda_{L}=100$ and $\lambda_{L}=0$, respectively, while all other parameters across this comparison are held constant. The latter setting of $\lambda_{L}=0$ is an unlearned method, as we are not leveraging ($\mu,\Sigma$) from a specific dataset. In the former setting of $\lambda_{L}=100$, we first learn ($\mu,\Sigma$) from a particular set of ten x-ray images; we then evaluate on a different set of x-ray images. We compare these two settings with varying noise and different number of measurements.

2.2 Results: Unlearned CS-DIP

For the remainder of this section, we evaluate our algorithm in the noiseless case without learned regularization, i.e. when $\eta=0$ in Eqn. 1 and $\lambda_{L}=0$ in Eqn. 3. Hence CS-DIP is completely unlearned; as such, we compare it to other state-of-the-art unlearned algorithms on various datasets and with different measurement matrices.

MNIST: In Figure 1(b) we plot reconstruction error with varying number of measurements $m$ of $n$ = 784. This demonstrates that our algorithm outperforms baselines in almost all cases. Figure 2(b) shows reconstructions for 75 measurements, while remaining reconstructions are in the appendix.

Chest x-rays: In Figure 1(a) we plot reconstruction error with varying number of measurements $m$ of $n$ = 65536. Figure 2(a) shows reconstructions for 2000 measurements; remaining reconstructions are in the appendix. On this dataset we outperform all baselines except BM3D-AMP for higher $m$. However for lower $m$, e.g. when the ratio $\frac{m}{n}\leq 3\%$, BM3D-AMP often doesn’t converge. This finding seems to support the work of Metzler et al. : BM3D-AMP performs well on higher $m$, e.g. $\frac{m}{n}\geq 10\%$, but recovery at lower sampling rates is not demonstrated.

Comparison to pre-trained DCGAN: The approach of Bora et al. similarly employs a DCGAN but is pre-trained over a large dataset. As expected, this method outperforms ours for lower number of measurements $m$; however, as $m$ increases, the pre-trained method’s performance saturates while our algorithm continues to improve and consequently outperform its pre-trained counterpart, per Figure 4(a) in the appendix. This can be attributed to our method optimizing over the weights $w$ as opposed to the latent space $z$, allowing for a more expressive network capable of reconstructing complicated signals e.g. medical images. The method of Bora et al. is comparatively less expressive; indeed it is only capable of reconstructing simple images, such as MNIST or aligned CelebA .

Additional experiments: In the appendix we further demonstrate our algorithm (1) using a Fourier measurement process for $A$ instead of a Gaussian i.i.d. matrix, (2) on RGB retinopathy images, and (3) in the presence of additive noise. We also perform a runtime analysis.

Conclusion

We demonstrate how Deep Image Prior (DIP) can be generalized to solve any differentiable linear inverse problem, in many cases outperforming state-of-the-art unlearned methods. We further propose learned regularization which enforces a learned Gaussian prior on the network weights. This prior reduces reconstruction error, particularly for noisy or compressed measurements. Lastly we prove that the DIP optimization technique can fit any signal given a sufficiently wide single-layer network. This provides theoretical justification for regularization methods such as early stopping.

References

Appendix A Implementation Details

Hyperparameter search: After a standard grid search procedure, we set the TV hyperparameter $\lambda_{T}=0.01$, which aids in providing a sharper reconstruction of the image’s high frequency components. For the learned regularization experiments in Section 5.2.1, a similar grid search was performed to set $\lambda_{L}=100$. The criteria for selecting a hyperparameter is one that provides lowest error with the observed measurements, i.e. without observing ground truth. We tune hyperparameters on a random set of ten images; this set is disjoint from the images used for evaluation.

Initalizations: The measurement matrix $A$ is initialized at random for each sample. Similarly network input $z$ in Eqn. 3 is initialized with random Gaussian i.i.d. entries and then held fixed as we optimize over network weights $w$. We set the dimension of $z$ to be $128$, a standard choice for DCGAN architectures. For a sufficient number of pixels $n$, i.e. for chest x-ray and retinopathy images, different initializations of $z$ do not affect performance. However for smaller $n$, i.e. for MNIST images, performance can vary with different initializations of $z$.

Early stopping: We stop after $1000$ iterations in all experiments. Similar to Figure 2 in Ulyanov et al. , we found MSE to decrease initially and then plateau until roughly $10,000$ iterations, at which point the network overfits to noise. Hence we terminate the optimization procedure within this plateau region to avoid overfitting. This early stopping technique is common in DIP methods, hence our motivation to justify it theoretically in Section 4.

Appendix B Additional Experiments

Fourier measurement process: All other experiments used a measurement matrix $A$ containing Gaussian i.i.d. entries. We now consider the case where the measurement matrix is a subsampled Fourier matrix. For a 2D image $x$ and a set of indices $\Omega$, the measurements we receive are given by $y_{(i,j)}=[\mathcal{F}(x)]_{(i,j)},(i,j)\in\Omega$, where $\mathcal{F}$ is the 2D Fourier transform. We choose $\Omega$ to be indices along radial lines, as shown in Figure 13 of the appendix; this choice of $\Omega$ is common in literature and MRI applications . While Fourier subsampling is common in MRI applications, we use it here on images of x-rays simply to demonstrate that our algorithm performs well with different measurement processes.

In Figure 3(b), we compare our algorithm to baselines on the x-ray dataset for $\{3,5,10,20\}$ radial lines in the Fourier domain, which corresponds to $\{381,634,1260,2500\}$ Fourier coefficients, respectively. Quantitatively we outperform all baselines. Qualitative reconstructions can be found in Figure 14.

Retinopathy: We plot reconstruction error with varying number of measurements $m$ of $n$ = 49152 in Figure 3(a). On this RGB dataset we quantitatively outperform all baselines except BM3D-AMP on higher $m$; however, even at these higher $m$, patches of green and purple pixels corrupt the image reconstructions as seen in Figure 10. Similar to x-ray for lower $m$, BM3D-AMP often fails to produce anything sensible. All retinopathy reconstructions are located in the appendix.

Robustness to noise: In Figure 4(b) we demonstrate that our algorithm is robust to additive noise, i.e. when $\eta\neq 0$ in Eqn. 1, achieving similar behavior to baselines.

Runtime: We demonstrate runtimes for all algorithms on the x-ray dataset in Table 2. While our algorithm is faster in most cases, we acknowledge this is not a fair comparison as baselines do not have the benefit of running GPU. Meanwhile our algorithm was run on a NVIDIA GTX 1080-Ti. Ultimately this demonstrates that our algorithm executes in a reasonable amount of time, which can be an issue with DIP methods employing a U-net architecture.

Appendix C Proof of Section 4: Theoretical Justification for Early Stopping

In this section we prove our theoretical result in Theorem 4.1. We begin with a summary of some notations we use throughout in Section C.1. Next, we state some preliminary calculations in Section C.2. Then, we state a few key lemmas in Section C.3 with the proofs deferred to Appendix D. Finally, we complete the proof of Theorem 4.1 in Section C.4. We note that when $V$ has i.i.d. Gaussian entries and $A$ contains orthonormal rows, $AV$ also has i.i.d. Gaussian entries. Therefore without loss of generality we carry out the proof with $A=I$ and $m=n$. The result stated in the theorem simply follows by replacing $V$ in our proof with $AV$.

C.2 Preliminaries

In this section we carryout some simple calculations yielding simple formulas for the gradient and Jacobian mappings. We begin by noting we can rewrite the gradient descent iterations in the form

is the Jacobian mapping associated to the network and

is the misfit or residual vector. Note that

C.3 Lemmas for controlling the spectrum of the Jacobian and initial misfit

In this section we state a few lemmas concerning the spectral properties of the Jacobian mapping, its perturbation and initial misfit of the model with the proofs deferred to Appendix D.

holds with probability at least $1-2e^{-n}$.

holds with probability at least $1-e^{-n/2}-e^{-d/2}$.

C.4 Proof of Theorem 4.1

Consider a nonlinear least-squares optimization problem of the form

Fix a point $\theta_{0}$. We have that $\sigma_{\min}\left(\mathcal{J}(\theta_{0})\right)\geq 2\alpha$.

Under these assumptions we can state the following theorem from .

Then, picking constant learning rate $\eta\leq\frac{1}{\beta^{2}}$, all gradient iterations obey the followings

We shall apply this theorem to the case where the parameter is $W$ and the nonlinear mapping is given by $V\phi\left(Wz\right)$ and $\phi=ReLU$. All that is needed to be able to apply this theorem is check that the assumptions hold. Per the assumptions of the theorem we use

To this aim note that using Lemma C.1 Assumption 1 holds with

with probability at least $1-2e^{-n}$. Furthermore, by Lemma C.3 Assumption 3 holds with

with probability at least $1-e^{-n/2}$. All that remains for applying the theorem above is to verify Assumption 2 holds with high probability

In the above we have used Lemma C.4 to conclude that $\left\|V\phi\left(Wz\right)-y\right\|\leq 3\left\|y\right\|$ holds with probability at least $1-e^{-n/2}-e^{-d/2}$. Thus, using Lemma C.2 all that remains is to show that

holds with $R=60\sqrt{n}$ and with probability at least $1-e^{-n/2}-e^{-\frac{(120)^{\frac{2}{3}}}{6}d^{\frac{2}{3}}n^{\frac{1}{3}}}\geq 1-e^{-n/2}-e^{-4d^{\frac{2}{3}}n^{\frac{1}{3}}}$. The latter is equivalent to

which holds as long as $d\geq 4.3\times 10^{15}n$. Thus with $d\geq Cn$ then Assumptions 1, 2, and 3 holds with probability at least $1-5e^{-n/2}-e^{-d/2}-e^{-4d^{\frac{2}{3}}n^{\frac{1}{3}}}$. Thus, Theorem C.5 holds with high probability. Applying Theorem C.5 completes the proof.

Appendix D Proof of Lemmas for the Spectral Properties of the Jacobian

We prove the result for $\nu=1$, the general result follows from a simple re-scaling. Define the vectors

To bound the minimum eigenvalue we state a result from .

Thus we can take $h=\sqrt{6}$. Therefore, using Theorem D.1 with $\delta=2e^{-n}$ we can conclude that

holds with probability at least $1-2e^{-n}$ as long as

Plugging this into (D.1) we conclude that with probability at least $1-2e^{-n}$

D.2 Proof of Lemma C.2

We prove the result for $\nu=1$, the general result follows from a simple rescaling. Based on (C.2) we have

To continue further note that by Gordon’s lemma we have

with probability at least $1-e^{-t^{2}/2}$. In particular using $t=\sqrt{n}$ we conclude that

with probability at least $1-e^{-n/2}$. To continue further we state a lemma controlling the size of $|\mathcal{S}(W)|$ based on the size of the radius $R$.

holds with probability at least $1-e^{-\frac{\left(2dR\right)^{\frac{2}{3}}}{6}}$.

Combining (12) together with (14) (using $s=3\left(2dR\right)^{\frac{2}{3}}$) and Lemma D.2 we conclude that

holds with probability at least $1-e^{-n/2}-e^{-\frac{\left(2dR\right)^{\frac{2}{3}}}{6}}$.

D.3 Proof of Lemma D.2

To prove this result we utilize two lemmas from . In these lemmas we use $|v|_{m-}$ to denote the $m$th smallest entry of $v$ after sorting its entries in terms of absolute value.

Combining the latter two lemmas with $m=\lceil\left(2dR\right)^{\frac{2}{3}}\rceil$ we conclude that when

then with probability at least $1-e^{-\frac{\left(2dR\right)^{\frac{2}{3}}}{6}}$ we have

D.4 Proof of Lemma C.3

We prove the result for $\nu=1$, the general result follows from a simple rescaling. Using (C.2) we have

The proof is complete by using standard concentration results for the spectral norm of a Gaussian matrix that allow us to conclude that

holds with probability at least $1-e^{-n/2}$.

D.5 Proof of Lemma C.4

To continue further let us consider one entry of $V\phi\left(Wz\right)$ and note that it has the same distribution as

with probability at least $1-e^{-n/2}$. Furthermore, note that

holds with probability at least $1-e^{-n/2}-e^{-d/2}$. Combining the latter with (15) we conclude that

holds with probability at least $1-e^{-n/2}-e^{-d/2}$.