Compressive sensing with un-trained neural networks: Gradient descent finds the smoothest approximation

Introduction

Un-trained convolutional neural networks have emerged as highly successful tools for image recovery and restoration, for a variety of problems including denoising, compressive sensing, and inpainting [Uly+18, Jin+19, Vee+18, JH19, Hec19, HH19, Bos+20, Wan+20, HA20, Aro+20]. As opposed to trained convolutional neural networks, that learn an image prior from training data, un-trained convolutional networks act as an image prior without any training and solely based on the architecture of the network and the optimization procedure used to fit them.

The benefit of untrained networks was first observed in the Deep Image Prior (DIP) paper [Uly+18]. The key observation of [Uly+18] is that fitting a standard over-parameterized convolutional autoencoder (specifically, the U-net [Ron+15] or variations thereoff) to a single noisy/corrupted image, when combined with early stopping, yields excellent denoising, inpainting, and super-resolution performance. Subsequent literature has demonstrated that many elements of the architecture of a convolutional autoencoder—such as the encoder part—are irrelevant for this behavior to emerge. In particular the papers [HH19, HS20] highlight the critical role of convolutions with fixed convolutional kernels.

Un-trained convolutional networks are empirically most effective when the network is over-parametrized, meaning that is has more parameters than image pixels. This holds even though in this regime the neural network can in principle fit any image perfectly, including random noise. Therefore, further regularization is critical to performance in many applications. For instance denoising [Uly+18, HS20] critically requires early stopping, as without early stopping the noisy image is fitted perfectly and no noise is removed. However, perhaps surprisingly, for some inverse problems including inpainting [Uly+18] and compressive sensing, no further regularization is necessary! That is, a convolutional neural network, when fitted to compressive measurements from a single image (no other training data) can estimate the original image well, as illustrated in Figure 1. This phenomenon demonstrates an intriguing self-regularization capability in the context of compressive sensing.

The overarching goal of this paper is to study compressive sensing with un-trained convolutional generators theoretically in order to explain the above phenomenon. In particular, our goal is to understand (i) why for compressive sensing problems gradient descent can reconstruct a good signal estimate without any further regularization or additional training data and to (ii) prove that this is possible with a minimal number of measurements that is proportional to an appropriately defined notion of signal dimensionality.

Let $\hat{\mathbf{C}}$ denote the solution found by gradient descent. The signal estimate can then be calculated as $\hat{\mathbf{x}}=G(\hat{\mathbf{C}})$.

The generator $G$ is over-parameterized and can express any image $\mathbf{x}^{\ast}$, including unstructured noise. Nevertheless, typically no further regularization in the form of early stopping the optimization is necessary. We demonstrate this phenomenon in Figure 1. This figure shows that running gradient descent on the loss $\mathcal{L}(\mathbf{C})$ eventually yields an estimate that is very close to the original image. This is surprising because i) there is no additional training data and ii) even though the generator $G$ can fit any image, including noise, gradient descent still finds an image close to the original one.

2 Contributions

The main contribution of this paper is to show that un-trained convolutional image priors provably enable recovery of natural images from a few random linear measurements. This holds by simply running gradient descent until convergence—without any further regularization. More specifically, we show that fitting an over-parameterized convolutional network with fixed convolutions (via gradient descent) to random measurements of a smooth signal essentially recovers that signal. Furthermore, the required number of measurements is commensurate to how smooth the signal is with more measurements required when the signal has “high-frequency” components. In more detail:

We plot these trigonometric basis functions in Figure 2 and formally define them later on in Section 4. Note that the smaller $p$, the smoother the signal $\mathbf{x}^{\ast}$ is, thus $p$ is a measure of smoothness.

Our main result shows that the estimate $\mathbf{C}_{\infty}$, obtained by running gradient descent on the loss (2) until convergence, yields an output $G(\mathbf{C}_{\infty})$ which is very close to $\mathbf{x}^{\ast}$, i.e., $G(\mathbf{C}_{\infty})\approx\mathbf{x}^{\ast}$. This holds as soon as the number of measurements exceeds the degrees of smoothness present in the signal ($p$). Since natural images are approximately smooth, this results provides a theoretical explanation why compressive sensing on natural images with over-parameterized convolutional generators works so well (see [Vee+18, JH19, Hec19, Aro+20] for corresponding empirical results).

In a nutshell, our main insight is that the behavior of large over-parameterized neural networks is dictated by the spectral properties of their Jacobian mapping. For the convolutional generators considered in this paper, the associated Jacobian matrix has singular vectors that can be well approximated by the orthonormal trigonometric basis function and singular values that decay very quickly from the low-frequency to the high-frequency trigonometric basis functions. Specifically, the associated singular values decay approximately geometrically.

To prove our result, we first characterize the least-squares solution of a randomly sketched least-squares problem with a design matrix with a decaying spectrum. To prove the result for convolutional generators we show that this non-linear learning problem behaves like an associated linear model with the above spectral characteristics. We then conclude the proof for the corresponding convolutional generator, by showing that the solutions obtained by running gradient descent on the non-linear problem is close to that obtained by running gradient descent on the linear problem.

In order to develop a better understanding of compressive sensing with untrained priors, we also carry out compressive sensing experiments for accelerating magnetic resonance imaging (MRI). Our experiments corroborate our theoretical finding that simply iterating until convergence is effective. This also suggests that there is little or no benefit to additional regularization.

Our paper is organized as follows: We start by stating the convolutional architecture considered in this paper in Section 2. In Section 3 we study the reconstruction of a signal from few a measurements with a linear over-parameterized generator to form intuition. In Section 4 we state our main results for signal recovery with convolutional generators. Section 5 contains our numerical result for MRI imaging. We conclude the paper with related work and a brief proof sketch, all formal proofs are deferred to the Appendix.

Convolutional generators

This architecture is a two-dimensional version of the deep decoder [HH19]. The deep decoder in turn is a sub-set of the deep image prior [Uly+18] and the U-net [Ron+15], as commented on below.

The deep decoder with $d$ layers (typically, $d=4,5,6$) is defined as

Finally, as mentioned before, the deep decoder can be viewed as the relevant part of a convolutional generator to function as an image prior. It can be deduced from a convolutional autoencoder (such as the deep image prior [Uly+18] and the U-net [Ron+15]) by removing the encoder part, any skip connections, and most surprisingly, the trainable convolutional filters of spatial extent larger than one. As demonstrated in [HS20], the critical aspect for an un-trained deep image prior are the convolutions with fixed convolutional kernels, implemented here by the operator $\mathbf{U}$.

Signal recovery with over-parameterized linear generators

Our goal is to estimate the signal $\mathbf{x}^{\ast}$ based on the measurement $\mathbf{y}$. We estimate the signal $\mathbf{x}^{\ast}$ by first computing a coefficient estimate $\hat{\mathbf{c}}$ by minimizing the loss

via running gradient descent with sufficiently small step size until convergence. We then estimate the signal via $\hat{\mathbf{x}}=\mathbf{J}\hat{\mathbf{c}}$. Since gradient descent applied on a least-squares problem yields the minimum-norm solution, the estimate $\hat{\mathbf{c}}$ can equivalently be expressed as

In closed form, $\hat{\mathbf{c}}$ is given as

where ${(\mathbf{A}\mathbf{J})}^{\dagger}$ is the pseudo-inverse of $\mathbf{A}\mathbf{J}$, and $\mathbf{P}_{\mathbf{J}^{T}\mathbf{A}^{T}}$ is a orthogonal projection operator onto the range of ${(\mathbf{A}\mathbf{J})}^{T}$. Thus, the signal estimation error is

The following theorem characterizes this signal estimation error.

The proof, given in the appendix, relies on arguments from [Hal+11, Sec. 8 and Sec. 9] developed for approximating low-rank matrices through random sampling.

The theorem guarantees that the error in estimating the signal $\mathbf{x}^{\ast}$ from compressive measurements $\mathbf{y}=\mathbf{A}\mathbf{x}^{\ast}$ is small provided that two conditions are satisfied:

The signal $\mathbf{x}^{\ast}$ lies (approximately) in the span of the leading $O(m)$ singular vectors of $\mathbf{J}$, where $m$ is the number of linear measurements.

The singular values of the generator matrix $\mathbf{J}$ decay sufficiently fast (for example geometrically).

Here, we used that the first term in the right-hand-side of (1) is bounded by $1/\sigma_{m/3}^{2}{\left\|\mathbf{x}^{\ast}\right\|}_{2}^{2}$, using that $\mathbf{x}^{\ast}$ is in the span of the leading singular vectors, and that $\sum_{i>2m/3}\sigma_{i}^{2}\leq\frac{\gamma^{2m/3}}{1-\gamma}$, by the formula for a geometric series. The bound (8) is very small provided that $\gamma$ is slightly below one (since $\gamma^{m/3}$ decays exponentially)—thus guaranteeing almost perfect recovery of a signal that is aligned with the leading singular vectors of $\mathbf{J}$.

Main results for compressive sensing with convolutional generators

We are now ready to state our main results for compressive sensing with convolutional generators. We consider the non-linear least-squares objective

In the previous section we studied a linear generator with generator matrix $\mathbf{J}$ with quickly decaying spectrum. In this section we extend the insights from the previous section to the non-linear case by replacing the role of the generator matrix $\mathbf{J}$ with the Jacobian of the non-linear generator $G$, defined as $[\mathcal{J}(\mathbf{C})]_{ij}=\frac{\partial}{\partial c_{i}}[G(\mathbf{C})]_{j}$. In contrast to the linear case, however, the Jacobian changes across iterations of gradient descent. Nevertheless, we can account for these changes in the Jacobian in our analysis.

With those definition, we are now ready to state our main result.

channels and with convolutional kernel $\mathbf{u}$ of the convolutional operator $\mathbf{U}$ and associated weights $\bm{\sigma}=[\sigma_{1},\ldots,\sigma_{n}]$. Here, $\xi\leq 1$ is arbitrary and $C_{\mathbf{u}}$ is a constant that only depends on the convolutional kernel $\mathbf{u}$. In order to estimate the signal, we fit the convolutional generator to the signal by running gradient descent starting from a random initialization $\mathbf{C}_{0}$ with i.i.d. $\mathcal{N}(0,\omega^{2})$, entries, $\omega\propto\frac{{\left\|\mathbf{y}\right\|}_{2}}{\sqrt{n}}$, and sufficiently small stepsize to the loss $\frac{1}{2}{\left\|\mathbf{A}G(\mathbf{C})-\mathbf{y}\right\|}_{2}^{2}$ until convergence. Then, with high probability, the reconstruction error with parameters $\mathbf{C}_{\infty}$ at convergence obeys

Theorem 2 establishes that a convolutional generator enables the reconstruction of a natural signal from a few linear measurements. To see this, note that a good model for a natural image is a smooth signal, i.e., a signal that can be well-approximated by few leading trigonometric basis functions. More concretely, Figure 4 in [SO01] shows that the power spectrum of a natural image (i.e., the energy distribution by frequency) decays rapidly from low frequencies to high frequencies.

Thus it is reasonably to assume that the signal $\mathbf{x}^{\ast}$ can be represented with few of the trigonometric basis function; for concreteness say that $\mathbf{x}^{\ast}$ lies in the span of $\mathbf{w}_{1},\ldots,\mathbf{w}_{m/3}$. Next, recall from Figure 3 that the weights associated with a triangular kernel decay geometrically (i.e., $\sigma_{i}^{2}=\gamma^{i}$ for some $\gamma\in(0,1)$). Thus, from the same argument as used for (8), the bound (13) established by the theorem yields that the reconstruction error is bounded by

Thus our theorem guarantees the recovery of a sufficiently smooth signal by optimizing over the range of the generator. In particular if the signal is $p$-smooth, i.e., lies in the span of $\mathbf{w}_{1},\ldots,\mathbf{w}_{p}$, then $O(p)$ measurements are sufficient to provide an accurate estimate.

Our main theorem from the previous section relies on two critical ingredients:

The finding from [HS20] that the leading singular vectors of the Jacobian of a two-layer deep decoder are approximately the trigonometric basis function throughout all iterations of gradient descent.

The weights $\sigma_{1},\ldots,\sigma_{n}$ associated with the trigonometric basis functions decaying sufficiently fast, specifically approximately geometric. That is required for gradient descent applied to fitting $m$ compressive measurements until convergence to (approximately) only fit the signal to the leading $O(m)$ trigonometric basis functions.

Those results extend to deeper networks as follows. First, as shown numerically in [HS20], the leading singular vectors of the Jacobian of a four-layer deep decoder are also close to the trigonometric basis functions, and change only little across iterations. Second, as shown in Figure 4, the singular values of a four-layer deep decoder also decay (at least) geometrically, and the spectrum changes only little across iterations. Thus, the implications of our theory continue to apply for deeper deep decoders.

Numerical experiments for magnetic resonance imaging

In the final part of our paper we consider accelerating magnetic resonance imaging (MRI), one of the major application of compressive sensing. MRI is a medical imaging technique where measurements of an object can only be taken in the Fourier domain, referred to as $k$-space. If the full $k$-space measurement is collected, an image of the object can be computed almost perfectly (up the noise inherent in the measurement process). In order to accelerate the imaging process, it is common to only collect a small part of the $k$-space, which corresponds to taking few linear Fourier measurements; or in the notation of our paper, a measurement matrix $\mathbf{A}$ with subsampled rows of the Fourier matrix.

In order to understand whether our main finding—that signal reconstruction from compressive measurements without further regularization is possible—applies in practice, we consider the problem of reconstructing an image from few k-space measurements. We consider reconstruction of an image from 8-fold undersampled k-space measurements from the fastMRI dataset, recently released by facebook and NYU [Zbo+18]. We reconstruct with a $d=5$ layer and highly over-parameterized deep decoder. Figure 5 shows the corresponding loss curves. It can be seen that early stopping at the optimal early stopping point gives only marginally better performance than when optimizing until convergence, and in addition the optimal early stopping point is unknown in practice (because we do not have access to a reconstruction from a full measurement).

Related literature

In this paper we focus on un-trained neural network for solving inverse problems. In contrast a large body of recent result concentrates on using trained deep convolutional neural networks for image recovery and reconstruction. Training based deep learning methods for solving inverse problems are either trained end-to-end for tasks like denoising [Bur+12, Zha+17], or are based on learning a generative image model (by training an autoencoder or GAN [HS06, Goo+14]) and then using the resulting image models to regularize problems such as compressed sensing [Bor+17, HV18, Hua+18], denoising [Hec+20], or phase retrieval [Han+18, SA18]. In contrast to un-trained network, where optimization is over the weights of the un-trained generator, in the aformentioned papers it is over the input of the (trained) network.

Our proof relies on relating the dynamics of gradient descent on an over-parameterized network to that of gradient descent on an associated linear network. This proof technique has been used in a variety of recent publication [Sol+18, Ven+19, Du+18, OS19, OS19a, Aro+19, Oym+19, Bas+19, Li+19]. Most related to our work is the recent paper [HS20] that shows that the deep decoder enables denoising. Neither of the publications, however, addresses compressive sensing or reconstruction from randomly sketched data, and most of our technical results are specific to this setup.

Finally note that regularizing linear models with gradient descent via early stopping has a rich history in the signal processing community. In the 50s, Landweber proposed to recover a signal from linear measurements via gradient descent [Lan51] which became known as the Landweber algorithm in the inverse problems community. Subsequent work in this literature proposed to early-stop the Landweber iterations (i.e., gradient descent) in order to regularize ill-posed inverse problems [TC85].

Proof sketch

In this section we provide a sketch of our argument. Our statement and formal proof pertains to the two-layer case, in this section we provide the sketch for the general case where $G(\bm{\theta})$ is a generic network with a $N$-dimensional parameter vector $\bm{\theta}$, and then comment on how this general proof strategy is particularized to the two layer case.

Given a measurement $\mathbf{y}$, we characterize the solution of running gradient descent with fixed step size $\eta$ on the nonlinear least-squares objective

starting from an initial point $\bm{\theta}_{0}$. The updates take the form

Relevant for the dynamics of gradient descent, however, are the corresponding sketched original and reference Jacobians, defined as

Since we chose $\mathbf{J}_{G}\approx\mathcal{J}_{G}(\bm{\theta}_{0})$, we also have $\mathbf{J}\approx\mathcal{J}(\bm{\theta}_{0})$.

To characterized the behavior of the gradient descent updates in (24), we relate the non-linear least squares problem to a linearized one in a ball around the initialization $\bm{\theta}_{0}$. This general strategy has been utilized in a number of recent publications [Sol+18, Du+18, Aro+19, OS19a, Oym+19, HS20]. We define the associated linearized least-squares problem as

Starting from the same initial point $\bm{\theta}_{0}$, the gradient descent updates of the linearized problem are

The iterates and residuals of the non-linear and linear updates are close throughout the entire run of gradient descent provided the following assumptions are satisfied:

The smallest and largest singular values of the generator reference Jacobian are lower and upper bounded by constants $\alpha$ and $\beta$, respectively.

The reference Jacobian approximates the Jacobian at initialization, i.e., for $\epsilon_{0}>0$,

where ${\left\|\cdot\right\|}$ is the standard operator (matrix) norm.

Within a radius $R$ around the initialization, the Jacobian varies by no more than $\epsilon$ in the sense that

Here, $\mathcal{B}_{R}(\bm{\theta}_{0})\coloneqq\{\bm{\theta}\colon{\left\|\bm{\theta}-\bm{\theta}_{0}\right\|}_{2}\leq R\}$ is the ball with radius $R$ around $\bm{\theta}_{0}$.

Under these assumptions, we establish that the residuals of the linear problem,

are close during the entire run of gradient descent, and most importantly for proving our result, that the iterates of the linear and non-linear problem are close, again during the entire run of gradient descent:

2 Inheriting the properties of the linear problem

Recall that our goal is to characterize the signal estimate $G(\bm{\theta}_{\infty})$ at convergence. We characterize this estimate by

characterizing the estimate $\hat{\mathbf{x}}=\mathbf{J}_{G}\bm{\theta}_{\infty}$ obtained by running the linear problem until convergence and

showing that this estimate is close to the original estimate, i.e., $\mathbf{J}_{G}\bm{\theta}_{\infty}\approx G(\bm{\theta}_{\infty})$.

In more detail, suppose that the assumption i-iii are satisfied for sufficiently small closeness parameters $\epsilon_{0}$ and $\epsilon$. Then, as discussed above, the iterates of the non-linear problem and the linear problem are close at any iteration, in particular at convergence. Since the Jacobians are also close, we can establish that $\hat{\mathbf{x}}=\mathbf{J}_{G}\bm{\theta}_{\infty}\approx G(\bm{\theta}_{\infty})$.

In more detail, we can bound the signal estimation error at convergence as

The first term is controlled by analyzing the linear case with Theorem 1 from Section 3. To control the second term we need a simple definition

With this definition in place we can proceed to bound the second term as follows

3 Concluding the proof sketch

The proof for the two-layer case is then concluded by analyzing the associated linear problem. In particular, we use that the matrix $\mathbf{J}_{G}$ has as its left-singular vectors the trigonometric basis function, and its spectrum are the associated weights $\sigma_{1},\ldots,\sigma_{n}$ specified in Section 4.

In order to extend this proof to a multi-layer deep decoder $G(\bm{\theta})$, all we need to do is to characterize the associated matrix $\mathbf{J}_{G}$, in particular its left-singular vectors and corresponding singular values.

Code

Code to reproduce the experiments is available at https://github.com/MLI-lab/cs_deep_decoder.

Acknowledgements

R. Heckel is partially supported by NSF award IIS-1816986 and acknowledges support of the NVIDIA Corporation in form of a GPU. M. Soltanolkotabi is supported by the Packard Fellowship in Science and Engineering, a Sloan Research Fellowship in Mathematics, an NSF-CAREER under award #1846369, the Air Force Office of Scientific Research Young Investigator Program (AFOSR-YIP) under award #FA9550-18-1-0078, an NSF-CIF award #1813877, DARPA under the Learning with Less Labels (LwLL) and Fast Network Interface Cards (FastNICs) program, and a Google faculty research award.

References

Appendix A Proof of Theorem 1

The statement follows from the following more general result.

To see this, note that with $p=k/2$ and $u=\sqrt{p}$, the proposition guarantees that with probability at least $1-3e^{-p}$,

Noting that $m=3/2k$, $\hat{\mathbf{x}}=\mathbf{J}\hat{\mathbf{c}}$ and $\mathbf{x}^{\ast}=\mathbf{J}\mathbf{c}^{\ast}$ concludes the proof.

By the characterization (6), our goal is to upper bound

Our proof relies on arguments from [Hal+11, Sec. 8 and Sec. 9] developed for approximating low-rank matrices through random sampling.

We start by partitioning the right-singular vectors of ${\mathbf{J}}^{T}$ into two blocks $\mathbf{V}_{1}$ and $\mathbf{V}_{2}$ containing $k$ and $n-k$ columns, respectively.

Note that both matrices are standard Gaussian, and, because they are non-overlapping sub-matrices of $\mathbf{V}\mathbf{A}$, they are also stochastically independent. Moreover, $\Omega_{1}$ has full row-rank with probability one.

Next, we record a useful property from [Hal+11, Prop. 8.4]: For a unitary matrix $\mathbf{U}$ any matrix $\mathbf{M}$,

To see that the identity (19) holds, first note that the matrix $\mathbf{P}={\mathbf{U}}^{T}\mathbf{P}_{\mathbf{M}}\mathbf{U}$ is an orthogonal projection operator because it is Hermitian an $\mathbf{P}^{2}=\mathbf{P}$. Moreover,

Since the range determines the orthogonal projector onto its range, we have that $\mathbf{P}={\mathbf{U}}^{T}\mathbf{P}_{\mathbf{M}}\mathbf{U}=\mathbf{P}_{{\mathbf{U}}^{T}\mathbf{M}}$, concluding the proof of (19). Next, let

be the full singular value decomposition of ${\mathbf{J}}^{T}$, including the singular vectors $\mathbf{U}_{d-n}$ multiplying with zero singular values. Applying the identity (19) and that ${\mathbf{U}}^{T}\mathbf{U}$ we proceed as

where the second-to-last inequality follows from [Hal+11, Last ineq in Sec. 9.2]. Finally, the last inequality holds with the probability specified in the proposition because by [Hal+11, Last inequality in Sec. 10.3], for $p\geq 4$ and $u>0$,

This concludes the proof of the proposition.

Appendix B Proof of Theorem 2

The result stated in the main text (Theorem 2) is obtained from a slightly more general result which applies beyond convolutional networks. Specifically, we consider neural network generators of the form

Our results depend on the largest and smallest eigenvalue of $\bm{\Sigma}(\mathbf{U})$ denoted by $\sigma_{n}^{2}$ and ${\left\|\mathbf{U}\right\|}^{2}$ and in particular a condition number denoted by $\kappa$ formally defined as

With these definitions in place we are now ready to state our result about neural generators.

via running gradient descent with iterations $\mathbf{C}_{t+1}=\mathbf{C}_{t}-\eta\nabla\mathcal{L}(\mathbf{C}_{t})$, starting from $\mathbf{C}_{0}$ with i.i.d. $\mathcal{N}(0,\omega^{2})$ entries, $\omega=\frac{\xi{\left\|\mathbf{y}\right\|}_{2}}{2\sqrt{n}{\left\|\mathbf{U}\right\|}}$, and step size obeying $\eta\leq\frac{m}{4n{\left\|\mathbf{U}\right\|}^{2}}$. Then, with probability at least $1-ne^{-k^{2}}-2e^{-\frac{m}{2}}-\delta$, for all iterations $t$,

Theorem 2 follows directly from Theorem 3 by noting that for $\mathbf{U}$ a circulant matrix (implementing a convolution), as found in [HS20], the left singular vectors of $\bm{\Sigma}(\mathbf{U})$ are given by the trigonometric basis functions in (10) and the singular values are given by (11).

Appendix C The dynamics of linear and nonlinear least-squares

Theorem 3, proven below, builds on a result on the dynamics of a general non-linear least squares problem that is stated and discussed in this section. Consider a nonlinear least-squares fitting problem of the form

To solve this problem, we run gradient descent with a fixed stepsize $\eta$, starting from an initial point $\bm{\theta}_{0}$, with updates of the form

The associated linearized least-squares problem is defined as

To show that the non-linear updates (24) are close to the linearized iterates (26), we make the following assumptions:

We assume the singular values of the reference Jacobian obey for some $\alpha,\beta$

Furthermore, we assume that the Jacobian mapping associated with the nonlinear model $f$ obeys

We assume the reference Jacobian and the Jacobian of the nonlinearity at initialization $\mathcal{J}(\bm{\theta}_{0})$ are $\epsilon_{0}$-close in the sense that

We assume that within a radius $R$ around the initialization, the Jacobian varies by no more than $\epsilon$ in the sense that

where $\mathcal{B}_{R}(\bm{\theta}_{0})\coloneqq\{\bm{\theta}\colon{\left\|\bm{\theta}-\bm{\theta}_{0}\right\|}\leq R\}$ is the ball with radius $R$ around $\bm{\theta}_{0}$.

Under these assumptions i) the difference of the nonlinear iterative updates (24) and the linear iterative updates (26) is bounded, and ii) the difference of the linear and non-linear residuals, defined as

are close throughout the entire run of gradient descent; both in the proximity of the initialization.

Here, ${\mathbf{J}}^{\dagger}$ is the pseudo-inverse of $\mathbf{J}$. We run gradient descent with stepsize $\eta\leq\frac{1}{\beta^{2}}$ on the linear and non-linear least squares problem, starting from the same initialization $\bm{\theta}_{0}$. Then, for all iterations $t$,

the non-linear residual converges geometrically

the residuals of the original and the linearized problems are close

the parameters of the original and the linearized problems are close

and finally, the parameters are not far from the initialization

The above theorem formalizes that in a (small) radius around the initialization, the non-linear problem behaves similarly as its linearization. Thus to characterize the dynamics of the nonlinear problem, it suffices to characterize the dynamics of the linearized problem. This is the subject of our next theorem, which is a standard results on the iterates of least squares, see [HS20, Thm. 5] for the proof.

Moreover, using a step size satisfying $\eta\leq\frac{1}{\sigma_{1}^{2}}$, the linearized iterates (26) obey

In the next section we show we can combine these two general theorems to provide guarantees for compressed sensing using general neural networks.

The proof is by induction. We note that the base case $t=0$ is trivially true. We suppose the statement, in particular the bounds (34), (35), (36), (37), and (38) hold for all iterations $\tau\leq t-1$. We then show that those relations continue to hold for iteration $t$ in five steps: In Step I, we show that a weaker version of (38) holds, specifically that ${\left\|\bm{\theta}_{t}-\bm{\theta}_{0}\right\|}_{2}\leq R$. This guarantees that we can work with our assumptions; those require the iterates to be sufficiently close to the initial values. In Step II we show that the nonlinear residual decreases at a geometric rate proving (34). In Steps III and IV we show that the residuals and the coefficients of the linear and non-linear problem are close, respectively. Finally, in Step V we utilize Steps I-IV to complete the proof by showing that the iterates of the non-linear problem are close to its initialization (i.e., equation (38)).

Before we start, we note that under our assumption, the residual of the linear problem converges linearly. Specifically, by the updates of the linear problem (26), we have that

Using that the smallest singular values of $\mathbf{J}{\mathbf{J}}^{T}$ is lower bounded by $2\alpha^{2}$, this guarantees that

establishing linear convergence of the linear problem.

We start by using a coarse argument that establishes $\bm{\theta}_{t}\in\mathcal{B}_{R}(\bm{\theta}_{0})$. First note that by the triangle inequality and the induction assumption (38) we have

So to prove ${\left\|\bm{\theta}_{t}-\bm{\theta}_{0}\right\|}_{2}\leq R$ it suffices to show that ${\left\|\bm{\theta}_{t}-\bm{\theta}_{t-1}\right\|}_{2}\leq R/2$. To this aim note that

Here, (ii) follows from the fact that $\frac{1}{2}\leq\frac{\beta^{2}}{\alpha^{2}}$ and inequality (i) follows from Assumptions 1-3, the induction hypothesis (36), ${\left\|\widetilde{\mathbf{r}}_{\tau-1}\right\|}\leq{\left\|\mathbf{r}_{0}\right\|}$, and the bound

To continue we use the fact that $\eta\leq\frac{1}{\beta^{2}}$ in (C.1) to conclude that

The last inequality follows by definition of $R$ in (33), and concludes the proof of Step I.

Step II: Geometric decay of non-linear iterate.

Since the linear residuals converge linearly and the Jacobian of the non-linear problem is close the Jacobian of the linear problem, $\mathbf{J}$, the non-linear problem also converges linearly. To see this, with $\mathcal{J}(\mathbf{a},\mathbf{b})=\int_{0}^{1}\mathcal{J}(s\mathbf{b}-(1-s)\mathbf{a})ds$, we have that, by the mean value theorem

where in the last equality we defined the matrices $\mathbf{B}_{1}$ and $\mathbf{B}_{2}$ accordingly for notational convenience. This implies that

For inequality (ii) we used the assumption $2\beta(\epsilon_{0}+\epsilon)\leq\alpha^{2}$, and for inequality (i) we used the bound

where the last inequality follows from our assumptions, and using that, by the triangle inequality and assumptions 2 and 3, we have

where in the last inequality we used the induction hypothesis (34). This completes the proof of the bound (34) for iteration $t$ concluding Step II.

Step III: Original and linearized residuals are close.

In this step, we bound the deviation of the residuals of the original and linearized problem defined as

Specifically, we use the induction hypothesis together with the fact that based on Step I we have $\bm{\theta}_{t-1},\bm{\theta}_{t}\in\mathcal{B}_{R}(\bm{\theta}_{0})$, to show that

Before we prove this however note that for $x\leq 1/2$ we have $(1-x)^{t-1}t\leq\frac{1}{e(\ln 2)x}$ for all $t\geq 0$. Now using this identity with $x=\eta\alpha^{2}\leq\frac{\alpha^{2}}{\beta^{2}}\leq\frac{1}{2}$ in (47) we conclude that

completing the proof of (36) for iteration $t$. Thus, all that remains in this step is to establish (47). To this aim note that from the formulas for the linear and non-linear residuals in (41) and (43), we have that

Thus for $\mathbf{e}_{t}=\widetilde{\mathbf{r}}_{t}-\mathbf{r}_{t}$ we have, with the same notation as in step II,

where the last inequality follows from ${\left\|\mathbf{B}_{1}\mathbf{B}_{2}-\mathbf{J}{\mathbf{J}}^{T}\right\|}\leq 2\beta(\epsilon_{0}+\epsilon)$, by (44), and from using the fact that ${\left\|\mathbf{r}_{t-1}\right\|}_{2}\leq(1-\eta\alpha^{2})^{t-1}{\left\|\mathbf{r}_{0}\right\|}_{2}$ which holds based on Step II. Finally, plugging in the induction hypothesis ${\left\|\mathbf{e}_{t-1}\right\|}_{2}\leq c\xi^{t-2}(t-1){\left\|\mathbf{r}_{0}\right\|}_{2}$ with $\xi:=1-\eta\alpha^{2}$ and $c:=2\eta\beta(\epsilon_{0}+\epsilon)$ in the above we conclude that

This concludes the proof of the bound (47) for iteration $t$, finishing Step III.

Step IV: Original and linearized parameters are close:

The difference between the parameter of the original iterate $\bm{\theta}$ and the linearized iterate $\widetilde{\bm{\theta}}$ obey

Here, (i) follows from (45) combined with Assumption 1 and (ii) follows from (47) established in step III. We now proceed by using the formulas for low-order polylogarithms to conclude that

This concludes the proof of (37) for iteration $t$, completing Step IV.

Step V: Proof of (38):

Here, inequality (ii) follows from the definition of $R$ in equation (33). Moreover, inequality (i) follows from the bound (37), which we just proved, and the fact that, from equation (40) in Theorem 2,

This concludes the proof of (38) for iteration $t$, completing the proof of Step V and the entire theorem.

Appendix D Proofs for neural network generators (proof of Theorem 3)

The proof of Theorem 3 relies on the fact that, in the overparameterized regime, the non-linear least squares problem is well approximated by an associated linearized least-squares problem. Studying the associated linear problem enables us to prove the result.

We apply Theorem 4, which ensures that the associated linear problem is a good approximation of the non-linear least squarest problem, with the non-linear function

Here, expectation is with respect to $\mathbf{C}$ with iid $\mathcal{N}(0,\omega^{2})$ parameters, and not with respect to $\mathbf{A}$. We apply Theorem 4 with

We next verify that the conditions of Theorem 4 are satisfied (specifically, Assumptions 1, 2, 3) by applying a series of Lemmas.

hold with probability at least $1-2e^{-\frac{\eta^{2}}{2}m}$ which with $\eta=1$ in turn implies that for $m\leq\frac{n}{9}$ we have

holds with probability at least $1-2e^{-\frac{m}{2}}$. See [Ver12, Corollary 5.35] for a proof of this standard result.

We start with bounding the initial residual by applying the following lemma.

With this lemma in place, the initial residual can be upper bounded as follows

Here (i) holds with probability at least $1-e^{-\frac{m}{2}}$ using the fact that $\mathbf{A}$ has i.i.d. Gaussian entries that are independent of $G(\mathbf{C}_{0})$, and for (ii) we used that, by Lemma 1,

where (i) follows from $\omega=\frac{\xi{\left\|\mathbf{y}\right\|}_{2}}{\beta\sqrt{m}}=\frac{\xi{\left\|\mathbf{y}\right\|}_{2}}{2\sqrt{n}{\left\|\mathbf{U}\right\|}}$ and for (ii) we used the fact that $\xi\leq\frac{1}{\sqrt{2\log(2n/\delta)}}$.

Verifying Assumption 1:

We next show that the norm of the reference Jacobian and the Jacobian are bounded, with the lemma below.

By Lemma 2, with ${\left\|\mathbf{v}\right\|}_{2}=1$,

where the last inequality follows from Lemma 2, with ${\left\|\mathbf{v}\right\|}_{2}=1$, and by using that, with high probability, ${\left\|\mathbf{A}\right\|}\leq 2\sqrt{n/m}$ per (48). Analogously, we obtain ${\left\|\mathcal{J}(\mathbf{C})\right\|}\leq\beta$, for all $\mathbf{C}$, with high probability. This completes the verification of Assumption 1.

Verifying Assumption 2:

To verify the assumption, we first state a concentration lemma from [HS20].

To show that (51) implies the condition in (29), we use the following lemma.

Using this inequality, as well as that ${\left\|\mathbf{A}\right\|}\leq 2\frac{\sqrt{n}}{\sqrt{m}}$, per (48), we get

as desired. This concludes the proof of Assumption 2.

This part of the proof also specifies our choice of the reference Jacobian $\mathbf{J}=\mathbf{A}\mathbf{J}_{G}$ as a matrix that is $\epsilon_{0}$ close to the Jacobian at initialization, $\mathcal{J}(\mathbf{C}_{0})$, and that exists by Lemma 4 above.

Verifying Assumption 3:

Verification of the assumption requires us to control the perturbation of the Jacobian matrix around a random initialization. We begin with the following lemma from [HS20].

Let $\mathbf{C}_{0}$ be a matrix with i.i.d. $\mathcal{\mathcal{missing}}{N}(0,\omega^{2})$ entries. Then, for all $\mathbf{C}$ obeying

with probability at least $1-ne^{-\frac{1}{2}\widetilde{R}^{4/3}k^{7/3}}$.

In order to verify Assumption 3, first note that the radius in the theorem, defined in equation (33), obeys

Here, (i) follows from the fact that ${\left\|\mathbf{J}^{\dagger}\mathbf{r}_{0}\right\|}_{2}\leq\frac{1}{\alpha\sqrt{2}}{\left\|\mathbf{r}_{0}\right\|}_{2}$, and using that $\epsilon_{0}+\epsilon\leq 2\epsilon=\frac{1}{8}\xi\frac{\alpha^{4}}{\beta^{3}}\leq\frac{1}{8}\xi\frac{\alpha^{4}}{\beta^{3}}$ (ii) from $\beta\geq\alpha$ and from the bound on the initial residual (49), (iii) from $\omega=\frac{\xi{\left\|\mathbf{y}\right\|}}{\beta\sqrt{m}}$ and finally (iv) follows from the assumption (21) which is equivalent to

For this choice of radius by Lemma 5 and by using ${\left\|\mathbf{A}\right\|}\leq 2\frac{\sqrt{n}}{\sqrt{m}}$ (per (48)) we have

where in (i) we used (21). Therefore, Assumption 3 holds with high probability by our choice of $\epsilon=\frac{\xi}{30}\frac{\alpha^{4}}{\beta^{3}}$.

Concluding the proof of Theorem 3:

To begin, let $\mathbf{c}^{\ast}$ be a solution to the optimization problem

To complete the proof of Theorem 3 let us consider the linearized optimization problem which takes the form

Here, $\mathbf{c}$ is the vectorized version of $\mathbf{C}$, with a slight abuse of notation. With this notation, we conclude the proof as

The bound (53) follows from Theorem 1 by noting that $\mathbf{w}_{1},\ldots,\mathbf{w}_{n}$ are the left singular vectors of $\mathbf{J}_{G}$ with associated singular values $\sigma_{1}\geq\ldots\geq\sigma_{n}$ (because $\mathbf{J}_{G}{\mathbf{J}}^{T}_{G}=\bm{\Sigma}(\mathbf{U}))$.

It remains to prove the bound (52). With $\mathcal{J}_{G}(\mathbf{a},\mathbf{b})=\int_{0}^{1}\mathcal{J}_{G}(s\mathbf{b}-(1-s)\mathbf{a})ds$, at $t=+\infty$,

In the above (i) follows from ${\left\|\mathbf{J}_{G}(\mathbf{C})\right\|}\leq{\left\|\mathbf{U}\right\|}=\frac{\sqrt{m}}{2\sqrt{n}}\beta$ (recall that $\beta=2\frac{\sqrt{n}}{\sqrt{m}}{\left\|\mathbf{U}\right\|}$) and from the bound

Moreover, (ii) follows from $m\leq n$ and ${\left\|\mathbf{c}_{\infty}\right\|}_{2}\leq{\left\|\mathbf{c}^{\ast}\right\|}_{2}\leq\frac{{\left\|\mathbf{x}^{\ast}\right\|}_{2}}{\sigma_{\min}(\mathbf{J}_{G})}$. We can now apply Theorem 4 equation (37) to bound the first term on the right-hand-side above to obtain

Here, (i) follows from ${\left\|\mathbf{r}_{0}\right\|}_{2}\leq 3{\left\|\mathbf{y}\right\|}_{2}=3{\left\|\mathbf{A}\mathbf{x}^{\ast}\right\|}_{2}\leq 6{\left\|\mathbf{x}^{\ast}\right\|}$, where we used (49) combined with the fact that ${\left\|\mathbf{A}\mathbf{x}^{\ast}\right\|}_{2}\leq 2{\left\|\mathbf{x}^{\ast}\right\|}_{2}$. Moreover, (ii) follows from $\frac{\beta}{\alpha}\geq 1$ and $\epsilon_{0}\leq\epsilon$ and finally (iii) from the choice $\epsilon=\frac{\xi}{16}\frac{\alpha^{4}}{\beta^{3}}$. This concludes the proof of the bound (52) and the proof of the theorem.