Gradient Descent Provably Solves Nonlinear Tomographic Reconstruction

Abstract

In computed tomography (CT), the forward model consists of a linear Radon transform followed by an exponential nonlinearity based on the attenuation of light according to the Beer–Lambert Law. Conventional reconstruction often involves inverting this nonlinearity as a preprocessing step and then solving a convex inverse problem. However, this nonlinear measurement preprocessing required to use the Radon transform is poorly conditioned in the vicinity of high-density materials, such as metal. This preprocessing makes CT reconstruction methods numerically sensitive and susceptible to artifacts near high-density regions. In this paper, we study a technique where the signal is directly reconstructed from raw measurements through the nonlinear forward model. Though this optimization is nonconvex, we show that gradient descent provably converges to the global optimum at a geometric rate, perfectly reconstructing the underlying signal with a near minimal number of random measurements. We also prove similar results in the under-determined setting where the number of measurements is significantly smaller than the dimension of the signal. This is achieved by enforcing prior structural information about the signal through constraints on the optimization variables. We illustrate the benefits of direct nonlinear CT reconstruction with cone-beam CT experiments on synthetic and real 3D volumes. We show that this approach reduces metal artifacts compared to a commercial reconstruction of a human skull with metal dental crowns.

1. Introduction

Computed tomography (CT) is a core imaging modality in modern medicine (Food & Administration, 2023). X-ray CT is used to diagnose a wide array of conditions, plan treatments such as surgery or chemotherapy, and monitor their effectiveness over time. It can image any part of the body, and is widely performed as an outpatient imaging procedure.

CT systems work by rotating an X-ray source and detector around the patient, measuring how much of the emitted X-ray intensity reaches the detector at each angle. Because different tissues absorb X-rays at different rates, each of these measurements records a projection of the patient’s internal anatomy along the exposure angle. Algorithms then combine these projection measurements at different angles to recover a 2D or 3D image of the patient. This image is then interpreted by a medical professional (e.g. physician, radiologist, or medical physicist) to help diagnose, monitor, or plan treatment for a disease or injury.

CT scanners in use today typically consider the image reconstruction task as a linear inverse problem, in which the measurements are linear projections of the signal at known angles. Omitting measurement noise, we can write this standard linear measurement model as:

$${\hat{y}}_i=a_i^{T}x,$$ (1.1)

where $x$ is a vectorized version of the unknown signal (which commonly lies in 2D or 3D) and $a_i$ is a known, nonnegative measurement vector that denotes the weight each entry in the signal contributes to the integral ${\hat{y}}_i$ along measurement ray $i$. Computed over a set of regularly-spaced ray angles, this is exactly the Radon transform (Radon, 1917). This linear measurement model is quite convenient, as it enables efficient computations using the Fourier slice theorem—which equates linear projections in real-space to evaluation of slices through Fourier space—as well as strong recovery guarantees from compressive sensing (Kak & Slaney, 2001; Foucart & Rauhuti, 2013; Bracewell, 1990). This linear projection model is accurate for signals of low density, for which the incident X-rays pass through largely unperturbed.

However, consider the common setting in which the signal contains regions of density high enough to occlude X-rays, such as the metal implants used in dental crowns and artificial joints. Such high-density regions produce nonlinear measurements for which the Fourier slice theorem, and standard compressed sensing results, no longer hold. In practice, tomographic reconstruction algorithms that assume a linear projection as the measurement model produce streak-like artifacts around high-density regions, potentially obscuring otherwise measurable and meaningful signal.

To avoid such artifacts, in this paper we consider a nonlinear measurement model, which correctly models signals with arbitrary density. Equation (1.1) then becomes:

$$y_i=1-\text{exp}\left(-a_i^{T}x\right),$$ (1.2)

where the exponential nonlinearity accounts for occlusion and is due to the Beer-Lambert Law. In practice, the partial occlusions captured by eq. (1.2) are commonly incorporated into a linear model by inverting the nonlinearity, converting raw measurements $y_i$ from eq. (1.2) into processed measurements ${\hat{y}}_i=-\text{ln}\left(1-y_i\right)$ for which eq. (1.1) holds. Indeed, this logarithmic preprocessing step is built into commercial CT scanners (Fu et al., 2016), though some additional preprocessing for calibration and denoising is often performed before the logarithm. The logarithm is well-conditioned for $y_i\approx 0$ but becomes numerically unstable as $y_i$ approaches unity, which corresponds to total X-ray absorption. This is particularly problematic for rays that pass through high-density materials, such as metal, as well as for very low-dose CT scans that use fewer X-ray photons.

Instead, we study reconstruction through direct inversion of eq. (1.2) via iterative gradient descent. We optimize a squared loss function over these nonlinear measurements $y_i$, which is optimal for the case of Gaussian measurement noise—though extending this analysis to more realistic noise models is also of interest (Fu et al., 2016). By avoiding the ill-conditioned logarithm of a near-zero measurement, this approach is well-suited to CT reconstruction with low-dose X-rays as well as CT reconstruction with reduced metal artifacts. However, direct reconstruction through eq. (1.2) is more challenging than reconstruction through the linearized eq. (1.1), because the resulting loss function is nonconvex. While the linear inverse problem defined by eq. (1.1) can be solved in closed form by methods such as Filtered Back Projection (Radon, 1917; Kak & Slaney, 2001; Natterer, 2001), the nonlinear inverse problem defined by eq. (1.2) requires an iterative solution to a nonconvex optimization problem. We show that gradient descent with appropriate stepsize successfully recovers the global optimum of this nonconvex objective, suggesting that direct optimization through eq. (1.2) is a viable and desirable alternative to current methods that use eq. (1.1).

Concretely, we make the following contributions:

• We propose a Gaussian model of eq. (1.2) and show that gradient descent converges to the global optimum of this model at a geometric rate, despite the nonconvex formulation. These results hold with a near minimal number of random measurements. To prove this result we utilize and build upon intricate arguments for uniform concentration of empirical processes.

• We also extend our result to a compressive sensing setting to show that a structured signal can be recovered from far fewer measurements than its dimension. In this case prior information about the signal structure is enforced via a convex regularizer; our result holds for any convex regularizer. We show that the required number of measurements is commensurate to an appropriate notion of statistical dimension that captures how well the regularizer enforces the structural assumptions about the signal. For example, for an $s$-sparse signal and an ${\ell }_1$ regularizer our results require on the order of $s\text{ log}\left(n/s\right)$ measurements, where $n$ is the dimension of the signal. This is the optimal sample complexity even for linear measurements.

• We perform an empirical comparison of reconstruction quality in 3D cone-beam CT on both synthetic and real volumes, where our real dataset consists of a human skull with metal dental crowns. We show that direct reconstruction through eq. (1.2) yields reduction in metal artifacts compared to reconstruction by inverting the nonlinearity into eq. (1.1).

2. Problem Formulation

In practice, the measurement vectors $a_i$ are sparse, nonnegative, highly structured, and dependent on the rays $i$, as only a small subset of signal values in $x$ will contribute to any particular ray. These vectors correspond to the weights in a discretized ray integral (projection) along the $i$‘th measurement ray in the Radon transform (Radon, 1917). We use these real, ray-structured measurement vectors in our synthetic and real-data experiments.

In our theoretical analysis, we make two simplifying alterations to eq. (1.2): (1) we model $a_i$ as a standard Gaussian vector, where the Gaussian randomness is an approximation of the randomness in the choice of ray direction, and (2) we wrap the inner product $a_i^{T}x$ in a ReLU, to capture the physical reality that the raw integral of density along a ray, and the corresponding sensor measurement, must always be nonnegative. This nonnegativity is implicit in eq. (1.2) because $a_i$ represents a ray integral with only nonnegative weights as its entries, and the true density signal $x$ is also nonnegative; in our model eq. (2.1) we make nonnegativity explicit (subscript + denotes ReLU):

$$y_i=f\left(a_i^{T}x\right)\text{, where }f(\cdot )=1-\text{exp}\left(-{(\cdot )}_{+}\right)\text{.}$$ (2.1)

Here $y_i\in ℝ$ is a measurement corresponding to ray $i,x\in {ℝ}^n$ is the signal we want to recover, and $a_i\in {ℝ}^n$ are i.i.d. random Gaussian measurement vectors distributed as $𝒩\left(0,I_n\right)$. In this paper we consider a least-squares loss of the form

$$ℒ\left(z\right)=\frac{1}{2m}\sum _{i=1}^m{\left(y_i-f\left(a_i^{T}z\right)\right)}^2,$$ (2.2)

which is optimal in the presence of Gaussian measurement noise. However, in our analysis we focus on the noiseless setting. We minimize this loss using subgradient descent starting from $z_0=0_n$, with step size ${\mu }_t$ in step $t$. More specifically, the iterates take the form

$$z_t=z_{t-1}-{\mu }_t\nabla ℒ\left(z_{t-1}\right)=z_{t-1}-\frac{{\mu }_t}{m}\sum _{i=1}^m{a}_i{f}^{\prime }\left(a_i^T{z}_{t-1}\right)\left(f\left(a_i^T{z}_{t-1}\right)-y_i\right).$$

Here, we use the following subdifferential of $f$:

$$f^{\prime }(\cdot )=\{\begin{array}{ll}0,\hfill & \text{ if }\cdot <0\hfill \\ \frac{1}{2},\hfill & \text{ if }\cdot =0\hfill \\ \text{exp}\left(-\cdot \right),\hfill & \text{ if }\cdot >0\hfill \end{array}.$$

We also consider a regularized (compressive sensing) setting where the number of measurements $m$ is significantly smaller than the dimension $n$ of the signal. In this case we optimize the augmented loss function

$$ℒ\left(z\right)=\frac{1}{2m}\sum _{i=1}^m{\left(y_i-f\left(a_i^{T}z\right)\right)}^2+\lambda ℛ\left(z\right)$$ (2.3)

via subgradient descent. Here, $ℛ\left(z\right)$ is a regularizer enforcing a priori structure about the signal, with regularization weight $\lambda$. In our experiments, we use 3D total variation as $ℛ$, to encourage our reconstructed structure to have sparse gradients in 3D space.

We note that $ℒ$ is a nonconvex objective, so it is not obvious whether or not subgradient descent will reach the global optimum. Do the iterates converge to the correct solution? How many iterations are required? How many measurements? How does the number of measurements depend on the signal structure and the choice of regularizer? In the following sections, we take steps to answer these questions.

3. Global Convergence in the Unregularized Setting

Our first result shows that in the unregularized setting, direct gradient-based updates converge globally at a geometric rate. We defer the proof of Theorem 1 to Appendix B.

Theorem 1 Consider the problem of reconstructing a signal $x\in {ℝ}^n$ from $m$ nonlinear CT measurements of the form $y_i=1-e^{-{\left(a_i^{T}x\right)}_{+}}$, where the measurement vectors $a_i$ are generated i.i.d. $𝒩\left(0,I_n\right)$. We consider a least-squares loss as in eq. (2.2) and run gradient updates of the form

$$z_t=z_{t-1}-{\mu }_t\nabla ℒ\left(z_{t-1}\right)$$

starting from $z_0=0_n$ with ${\mu }_1=4\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}$ and ${\mu }_t=\mu e^{-5\parallel x{\parallel }_{{\ell }_2}}$ with $\mu \le c_0$ for $t>1$. Here, erfc is the complementary error function. As long as the number of measurements obeys

$$m\ge \frac{c_1{e}^{c_2\parallel x{\parallel }_{{\ell }_2}}}{\parallel x{\parallel }_{{\ell }_2}^2}n$$

then

$${\Vert z_t-x\Vert }_{{\ell }_2}^2\le {\left(1-\mu e^{-10\parallel x{\parallel }_{{\ell }_2}}\right)}^t\parallel x{\parallel }_{{\ell }_2}^2$$

holds with probability at least $1-5e^{-c_{3}n}-3e^{-\frac{m}{2}}$. Here, $c_0,c_1,c_2$, and $c_3$ are fixed positive numerical constants.

Theorem 1 answers some of the key questions from the previous section in the affirmative. Even though the nonlinear CT reconstruction problem is a nonconvex optimization, gradient descent converges to the global optimum, the true signal, at a geometric rate.

Further, the number of required measurements $m$ is on the order of $n$, the dimension of the signal, which is near-minimal even for a linear forward model. In Theorem 2 we prove global convergence with even fewer measurements in the compressive sensing setting, when some prior knowledge of the signal structure is enforced through a convex regularizer.

We note that the initial step size ${\mu }_1$ used in Theorem 1 is a function of the signal norm $\parallel x{\parallel }_{{\ell }_2}$, which is a priori unknown. However, we briefly describe how this quantity can be estimated from the available measurements. By averaging over the $m$ measurements, we have

$$\frac{1}{m}\sum _{i=1}^m{y}_i=\frac{1}{m}\sum _{i=1}^m\left(1-e^{-{\left(a_i^{T}x\right)}_{+}}\right)=1-\frac{1}{m}\sum _{i=1}^m{e}^{-g_{i+}\parallel x{\parallel }_{{\ell }_2}}$$

where $g_i$ are i.i.d. standard Gaussian random variables. Since $e^{-g_{i+}\parallel x{\parallel }_{{\ell }_2}}$ is a 1-Lipschitz function of $g_i$, this quantity concentrates around its mean

$$\mathbb{E}\left[e^{-g_{i+}\parallel x{\parallel }_{{\ell }_2}}\right]=\frac{1}{2}\left(1+\text{exp }\left(\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\text{ erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)\right).$$

We can invert this relationship to get a close estimate of $\parallel x{\parallel }_{{\ell }_2}$ from the average measurement value.

We also note that both the convergence rate and the number of measurements in Theorem 1 are exponentially dependent on $\parallel x{\parallel }_{{\ell }_2}$. This is natural because as $\parallel x{\parallel }_{{\ell }_2}$ increases towards infinity the measurements $y_i=1-e^{-{\left(a_i^{T}x\right)}_{+}}$ approach the constant value 1 and the corresponding gradient of the loss approaches zero. Intuitively, this corresponds to trying to recover a CT scan of a metal box; if the walls of the box become infinitely absorbing of X-rays, we cannot hope to see inside it. Nonetheless, for real and realistic metal components in our experiments (Section 5) we do find good signal recovery following this approach.

4. Global Convergence in the Regularized Setting

We now turn our attention to the regularized setting. Our measurements again take the form $y_i=1-e^{-{\left(a_i^{T}x\right)}_{+}}$ for $i=1,2,\dots ,m$, where $x\in {ℝ}^n$ is the unknown but now a priori “structured” signal. In this case we wish to use many fewer measurements $m$ than the number of variables $n$, to reduce the X-ray exposure to the patient without sacrificing the resolution of the reconstructed image or volume $x$. Because the number of equations $m$ is significantly smaller than the number of variables $n$, there are infinitely many reconstructions obeying the measurement constraints. However, it may still be possible to recover the original signal by exploiting knowledge of its structure. To this aim, let $ℛ:{ℝ}^n\to ℝ$ be a regularization function that reflects some notion of “complexity” of the “structured” solution. For the sake of our theoretical analysis we will use the following constrained optimization problem in lieu of eq. (2.3) to recover the signal:

$$\underset{z\in {ℝ}^n}{\text{min }}ℒ\left(z\right)=\frac{1}{2m}\sum _{i=1}^m{\left(y_i-f\left(a_i^{T}z\right)\right)}^2\text{ subject to }ℛ\left(z\right)\le ℛ\left(x\right).$$ (4.1)

We solve this optimization problem using projected gradient updates of the form

$$z_{t+1}={𝒫}_{𝒦}\left(z_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right)\right).$$ (4.2)

Here, ${𝒫}_{𝒦}\left(z\right)$ denotes the projection of $z\in {ℝ}^n$ onto the constraint set

$$𝒦=\left\{z\in {ℝ}^n:ℛ\left(z\right)\le ℛ\left(x\right)\right\}.$$ (4.3)

We wish to characterize the rate of convergence of the projected gradient updates eq. (4.2) as a function of the number of measurements, the available prior knowledge of the signal structure, and how well the choice of regularizer encodes this prior knowledge. For example, if we know our unknown signal $x$ is approximately sparse, using an ${\ell }_1$ norm for the regularizer is superior to using an ${\ell }_2$ regularizer. To make these connections precise and quantitative, we need a few definitions which we adapt verbatim from Oymak et al., (2017); Oymak & Soltanolkotabi (2017b); Soltanolkotabi (2019b).

Definition 1 (Descent set and cone) The set of descent of a function $ℛ$ at a point $x$ is defined as

$${𝒟}_{ℛ}\left(x\right)≔\left\{h:ℛ\left(x+h\right)\le ℛ\left(x\right)\right\}.$$

The cone of descent, or tangent cone, is the conic hull of the descent set, or the smallest closed cone ${𝒞}_{ℛ}\left(x\right)$ that contains the descent set, i.e. ${𝒟}_{ℛ}\left(x\right)\subset {𝒞}_{ℛ}\left(x\right)$.

The size of the descent cone ${𝒞}_{ℛ}$ determines how well the regularizer $ℛ$ captures the structure of the unknown signal $x$. The smaller the descent cone, the more precisely the regularizer describes the properties of the signal. We quantify the size of the descent cone using the notion of mean (Gaussian) width.

Definition 2 (Gaussian width) The Gaussian width of a set $𝒞\in {ℝ}^p$ is defined as:

$$\omega \left(𝒞\right)≔{\mathbb{E}}_g\left[\underset{z\in 𝒞}{\text{sup}}\langle g,z\rangle \right],$$

where the expectation is taken over $g~𝒩\left(0,I_p\right)$. Throughout we use ${ℬ}^n/{\mathbb{S}}^{n-1}$ to denote the the unit ball/sphere of ${ℝ}^n$.

We now have all the definitions in place to quantify how well the function $ℛ$ captures the properties of the unknown signal $x$. This naturally leads us to the definition of the minimum required number of measurements.

Definition 3 (minimal number of measurements) Let ${𝒞}_{ℛ}\left(z\right)$ be a cone of descent of $ℛ$ at $z$. We define the minimal sample function as

$$ℳ\left(ℛ,z\right)≔{\omega }^2\left({𝒞}_{ℛ}\left(z\right)\cap {ℬ}^n\right).$$

We shall often use the short hand $m_0=ℳ\left(ℛ,z\right)$ with the dependence on $ℛ,z$ implied. Here we define $m_0$ for an arbitrary point $z$, but we will apply the definition at the signal $x$.

We note that $m_0$ is exactly the minimum number of samples required for structured signal recovery from linear measurements when using convex regularizers. Specifically, the optimization problem

$$\text{arg }\underset{z}{\text{min}}\frac{1}{2m}\sum _{i=1}^m{\left(y_i-a_i^{T}z\right)}^2\text{ subject to }ℛ\left(z\right)\le ℛ\left(x\right),$$ (4.4)

succeeds at recovering the unknown signal $x$ with high probability from $m$ measurements of the form $y_i=a_i^{T}x$ if and only if $m\ge m_0$.¹ Given that in our Gaussian-approximated nonlinear CT reconstruction problem we have less information (we lose information when the input to the ReLU is negative), we cannot hope to recover structured signals from $m\le m_0$ when using (4.1). Therefore, we can use $m_0$ as a lower-bound on the minimum number of measurements required for projected gradient descent iterations eq. (4.2) to succeed in recovering the signal of interest. With these definitions in place we are now ready to state our theorem in the regularized/compressive sensing setting. We defer the proof of Theorem 2 to Appendix C.

Theorem 2 Consider the problem of reconstructing a signal $x\in {ℝ}^n$ from $m$ nonlinear CT measurements of the form $y_i=1-e^{-{\left(a_i^{T}x\right)}_{+}}$, where the measurement vectors $a_i$ are generated i.i.d. $𝒩\left(0,I_n\right)$. We consider a constrained least-squares loss as in eq. (4.1) and run projected gradient updates of the form in eq. (4.2) starting from $z_0=0_n$ with ${\mu }_1=4\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_2}{\sqrt{2}}\right)}$ and ${\mu }_t=\mu e^{-5\parallel x{\parallel }_{{\ell }_2}}$ with $\mu \le \frac{c_0}{{\left(1+\frac{n}{m}\right)}^2}$ for $t>1$. Here, erfc is the complementary error function. As long as the number of measurements obeys

$$m\ge \frac{c_1{e}^{c_2\parallel x{\parallel }_{{\ell }_2}}}{\parallel x{\parallel }_{{\ell }_2}^2}{m}_0,$$

with $m_0$ denoting the minimal number of samples per Definition 3, then

$$z_t-x_{{\ell }_2}^2\le {\left(1-\mu e^{-10\parallel x{\parallel }_{{\ell }_2}}\right)}^t\parallel x{\parallel }_{{\ell }_2}^2$$

holds with probability at least $1-5e^{-c_3{m}_0}-3e^{-\frac{m}{2}}$. Here, $c_0,c_1,c_2$, and $c_3$ are fixed positive numerical constants.

Theorem 2 parallels Theorem 1, likewise showing fast geometric convergence to the global optimum despite nonconvexity. In this regularized setting, the sample complexity of our nonlinear reconstruction problem is on the order of $m_0$, the number of measurements required for linear compressive sensing. In other words, the number of measurements required for regularized nonlinear CT reconstruction from raw measurements is within a constant factor of the number of measurements needed for the same reconstruction from linearized measurements. This is the optimal sample complexity for this nonlinear reconstruction task. For instance for an $s$ sparse signal for which $m_0\propto s\text{ log}\left(n/s\right)$, the above theorem states that on the order of $s\text{ log}\left(n/s\right)$ nonlinear CT measurements suffices for our direct gradient-based approach to succeed. Finally, we would like emphasize that the above result is rather general as it applies to any type structure in the signal and can also deal with any convex regularizer.

5. Experiments

We support our theoretical analysis with experimental evidence that gradient-based optimization through the nonlinear CT forward model is effective for a wide range of signal densities, including signals that are dense enough that the same optimization procedure through the linearized forward model produces noticeable “metal artifacts.”

All of our experiments are based on the JAX implementation of Plenoxels (Sara Fridovich-Keil and Alex Yu et al., 2022), with a dense 3D grid of optimizable density values connected by trilinear interpolation. We use a cone-beam CT setup and optimize with mild total variation regularization. Our experiments do not focus on speed or measurement sparsity, though we fully expect our optimization objective to pair naturally with efficient ray sampling implementations and regularizers of choice.

5.1. Synthetic Data

Our synthetic experiments use a ground truth volume defined by the standard Shepp-Logan phantom (Shepp & Logan, 1974) in 3D, with the following modifications:

• We scale down the voxel density values by a factor of 4, to more closely mimic the values in our real cone-beam CT skull dataset.

• We adjust one of the ellipsoids to be slightly larger than standard (to make it more visible), and gradually increase its ground truth density to simulate a spectrum from soft tissue to bone to metal.

We simulate CT observations of this synthetic volume and then reconstruct using either the linearized forward model with the logarithm and eq. (1.1), or directly using eq. (1.2). We also use a small amount of total variation regularization, and constrain results to be nonnegative.

Results of this synthetic experiment are presented in Figure 1. As the density of the test ellipsoid increases, the linearized reconstruction experiences increasingly severe “metal artifacts,” while the nonlinear reconstruction continues to closely match the ground truth. PSNR values are reported over the entire reconstructed volume compared to the ground truth, where PSNR is defined as −10 log₁₀ (MSE) and MSE is the mean squared voxel-wise error.

Figure 1:

Synthetic experiments using the Shepp-Logan phantom, showing a slice through the reconstructed 3D volume. From top to bottom, we increase the density of the central test ellipsoid to simulate soft tissue, bone, and metal. Nonlinear reconstruction is robust even to dense “metal” elements of the target signal.

Note that this synthetic experiment does not include any measurement noise or miscalibration; the instability of the logarithm with respect to dense signals arises even when the only noise is due to numerical precision. We also note that even the densest synthetic “metal” ellipsoid we test is no denser than what we observe in our real CBCT skull dataset in Figure 2, with a real metal dental crown.

Figure 2:

Real experiments using a 3D human skull phantom with a metal dental crown; here we show a cross-section of the reconstructed volume. Note the streak artifact to the left of the metal crown (annotated in red) and the X artifact below it (annotated in purple) in the reference linearized reconstruction, which are not present in the nonlinear reconstruction.

5.2. Real Data

Our real data experiment uses a cone-beam CT phantom made from a human skull with metal dental crowns on some of the teeth. In Figure 2 we show slices of our nonlinear reconstruction compared to a reference commercial linearized reconstruction, as no ground truth is available for the real volume. The nonlinear reconstruction exhibits reduced metal-induced streak artifacts compared to the commercial reconstruction, highlighted in red and purple in the leftmost panel.

6. Related Work

Tomographic reconstruction.

The measurement model in eq. (1.2) is a discretized corollary of the Beer-Lambert Law that governs the attenuation of light as it passes through absorptive media. Inverting the exponential nonlinearity in this model recovers the Radon transform summarized by eq. (1.1), in which measurements are linear projections of the signal at chosen measurement angles. The Radon transform has a closed-form inverse transform, Filtered Back Projection (Radon, 1917; Kak & Slaney, 2001; Natterer, 2001), that leverages the Fourier slice theorem (Bracewell, 1990; Kak & Slaney, 2001). Filtered Back Projection is a well-understood algorithm that can be computed efficiently, and is a standard option in commercial CT scanners, but its reconstruction quality can suffer in the presence of either limited measurement angles or metal (highly absorptive) signal components (Fu et al., 2016).

Many methods exist to improve the quality of CT reconstruction in the limited-measurement regime, such as limited baseline tomography, which is of clinical interest because not all viewpoints may be accessible and every measurement angle requires exposing the patient to ionizing X-ray radiation. These methods typically involve augmenting the data-fidelity loss function with a regularization term that describes some prior knowledge of the signal to be reconstructed. Such priors include sparsity (implemented through an ${\ell }_1$ norm) in a chosen basis, such as wavelets (Foucart & Rauhuti, 2013; Chambolle et al., 1998), as well as gradient sparsity (implemented through total variation regularization) (Candes et al., 2006). Compressive sensing theory guarantees correct recovery with fewer measurements in these settings, as long as the true signal is well-described by the chosen prior (Foucart & Rauhuti, 2013). CT reconstruction with priors cannot be solved in closed form, but as long as the regularization is convex we are guaranteed that iterative optimization methods such as gradient descent, ISTA (Chambolle et al., 1998), and FISTA (Beck & Teboulle, 2009) will be successful.

Recently, reconstruction with even fewer measurements has been proposed by leveraging deep learning, through either neural scene representation (Rückert et al., 2022) or data-driven priors (Szczykutowicz et al., 2022). These methods may sacrifice convexity, and theoretical guarantees, in favor of more flexible and adaptive regularization that empirically reduces reconstruction artifacts in the limited-measurement regime. However, these methods are still based on the linear measurement model of eq. (1.1), making them susceptible to reconstruction artifacts near highly absorptive metal components. In some cases neural methods may reduce metal artifacts compared to traditional algorithms, but this reduction is achieved by leveraging strong and adaptive prior knowledge rather than the measurements of the present signal.

We propose to resolve these metal artifacts by reconstructing from raw nonlinear X-ray absorption measurements, rather than the preprocessed, linearized measurements produced by standard CT scanners. Our method may pair particularly well with new photon-counting CT scanners (Shikhaliev et al., 2005), which were approved by the FDA in 2021 (Food & Administration, 2021). These scanners measure raw X-ray photon counts, which should enable finer-grained noise modeling and correction as well as our method for principled reconstruction of signals with metal.

Signal reconstruction from nonlinear measurements.

There are a growing number of papers focused on reconstructing a signal from nonlinear measurements or single index models. Early papers on this topic focus on phase retrieval and ReLU nonlinearities (Oymak et al., 2018; Soltanolkotabi, 2017; Candes et al., 2015) and approximate reconstruction (Oymak & Soltanolkotabi, 2017a). These papers do not handle the compressive sensing/structured signal reconstruction setting. The paper (Soltanolkotabi, 2019a) deals with reconstruction from structured signals for intensity and absolute value nonlinearties but only achieves the optimal sample complexity locally. A more recent paper (Mei et al., 2018) deals with a variety of nonlinearities with bounded derivative activations. However, this paper does not handle non-differentiable activations and only deals with simple structured signals such as sparse ones. In contrast to the above paper, our activation is non-differentiable, we handle arbitrary structures in the signal, and our results apply for any convex regularizer.

7. Discussion

In this paper, we consider the CT reconstruction problem from raw nonlinear measurements of the form $y_i=1-e^{-a_i^{T}x}$ for a signal $x$ and random measurement weights $a_i$. Although this nonlinear measurement model can be easily transformed into a linear model via a logarithmic preprocessing step ${\hat{y}}_i=-\text{ln}\left(1-y_i\right)=a_i^{T}x$, and this transformation is common practice in clinical CT reconstruction, the logarithm is numerically unstable when the measurements approach unity. This occurs frequently in practice, notably when the signal $x$ contains metal and especially for low-dose CT scanners that reduce radiation exposure. In this setting, traditional linear reconstruction methods tend to produce “metal artifacts” such as streaks around metal implants. Reconstruction directly through the raw nonlinear measurements avoids this numerically unstable preprocessing, in exchange for solving a nonconvex nonlinear least squares objective instead of convex linear least squares.

We prove that gradient descent finds the global optimum in CT reconstruction from raw nonlinear measurements, recovering exactly the true signal $x$ despite the nonconvex optimization. Moreover, it converges at a geometric rate, which is considered fast even for convex optimization. This nonconvex optimization requires order $n$ measurements, where $n$ is the dimension of the unknown signal, the same order sample complexity as if we had reconstructed through a linear forward model.

We also extend our theoretical results to the compressive sensing setting, in which prior structural knowledge of the signal $x$, enforced through a regularizer, allows for reconstruction with far fewer measurements than the dimension of the signal. Our results in this setting again parallel standard results from the linear reconstruction problem, even though we consider a nonlinear forward model and optimize a nonconvex formulation.

We also compare linearized and nonlinear CT reconstruction experimentally in the setting of 3D cone-beam CT, using both a synthetic 3D Shepp-Logan phantom for which we know the ground truth volume as well as a real human skull phantom with metal dental crowns. In both cases, we find that nonlinear reconstruction reduces metal artifacts compared to linearized reconstruction, whether that linearized reconstruction is done by gradient descent or a commercial algorithm.

Our work is a promising first step towards higher-quality CT reconstruction in the presence of metal components and low-dose X-rays, offering both practical and theoretical guidance for trustworthy reconstruction. Future work may extend our results both theoretically and experimentally to consider more realistic measurement noise settings such as Poisson noise, which is particularly timely given the emergence of new photon-counting CT scanners.

A. APPENDIX

B. Proof of Theorem 1

In the following subsections we prove Theorem 1, beginning with a proof outline.

B.1. Proof Outline

The proof consists of the following four steps.

Step I: First iteration: Welcome to the neighborhood.

In the first step, we show that as long as the number of measurements is sufficiently large $\left(m\ge 5n\right)$, the first iteration obeys

$$\parallel z_1-{x\parallel }_{{\ell }_2}\le \frac{1}{4}\parallel x{\parallel }_{{\ell }_2}$$ (B.1)

with probability at least $1-2e^{-\frac{m}{2}}-e^{-cn}$, for a constant $c$. We prove this in Appendix B.2.

Step II: Local pseudoconvexity.

In the second step, we show that our nonconvex objective function is locally strongly pseudoconvex inside this local neighborhood of eq. (B.1). Specifically, we show the correlation inequality

$$\nabla ℒ{(z)}^T\left(z-x\right)\ge \alpha \parallel z-x{\parallel }_{{\ell }_2}^2$$ (B.2)

holds with $\alpha =\frac{1}{2}{e}^{-\left(5\parallel \mathbf{x}{\parallel }_{{\ell }_2}+2\right)}$, with probability at least $1-4e^{-n}$ as long as the number of measurements $m$ is at least $\frac{c_1}{e^{c_2∥\mathbf{x}{\ell }_{{\ell }_2}∥\mathbf{x}{\parallel }_{{\ell }_2}^2}}n$ for constants $c_1$ and $c_2$.

We compute the $\alpha$ in Equation (B.2) in two cases, where case 1 corresponds to $\frac{{\mathbf{x}}^T(\mathbf{z}-\mathbf{x})}{\parallel \mathbf{x}{\parallel }_{{\ell }_2}\parallel \mathbf{z}-\mathbf{x}{\parallel }_{{\ell }_2}}\ge -0.6$ and case 2 corresponds to $\frac{{\mathbf{x}}^T(\mathbf{z}-\mathbf{x})}{\parallel \mathbf{x}{\parallel }_{{\ell }_2}\parallel \mathbf{z}-\mathbf{x}{\parallel }_{{\ell }_2}}<-0.6$. These cases are analyzed in Appendix B.3 and Appendix B.4, respectively. Combining cases 1 and 2, we have that the lower bound on $\alpha$ from case 2 lower bounds the bound from case 1, as shown in Figure 3, so we use that bound in eq. (B.2). The sample complexity in eq. (B.2) is the maximum over the sample complexities of cases 1 and 2, which are $m\ge \frac{cn}{{\alpha }^2\parallel x{\parallel }_{{\ell }_2}^2}$ and $m\ge \frac{Cn}{(\sqrt{2\alpha }-\sqrt{\alpha }{)}^2}$, respectively, for constants $c$ and $C$.

Figure 3:

Case 2 provides a lower bound on the expected correlation for both cases.

Step III: Smoothness.

In the third step, we show

$$\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}\le L\parallel z-x{\parallel }_{{\ell }_2}$$ (B.3)

holds with probability at least $1-e^{-\frac{1}{2}\left(m+n\right)}$. , for a constant $L$. This smoothness condition is proved in Appendix (B.5).

Step IV: Completing the proof via combining Steps I-III.

Finally, we combine the first three steps into a complete proof. At the core of the proof is the following lower bound on the correlation between the loss gradient and the error vector

$$\nabla ℒ{(z)}^T\left(z-x\right)\ge A\parallel z-x{\parallel }_{{\ell }_2}^2+B\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}^2$$ (B.4)

for positive constants $A$ and $B$. This starting point is similar to the proof of Lemma 7.10 in Candes et al. (2015), with some modifications necessary for the nonlinear CT reconstruction problem. To prove eq. (B.4) we first combine eq. (B.2) and eq. (B.3) to conclude that

$$\begin{gathered} \nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{\alpha }{2}\parallel z-x{\parallel }_{{\ell }_2}^2+\frac{\alpha }{2}\parallel z-x{\parallel }_{{\ell }_2}^2 \\ \ge \frac{\alpha }{2}\parallel z-x{\parallel }_{{\ell }_2}^2+\frac{\alpha }{2L^2}\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}^2. \end{gathered}$$

Thus eq. (B.4) holds with $A=\frac{\alpha }{2}$ and $B=C$ for a new constant $C$, with probability at least $1-4e^{-n}-e^{-\frac{1}{2}\left(m+n\right)}-2e^{-\frac{m}{2}}-e^{-cn}$, for a constant $c$ (by a union bound over the first three steps of the proof). Using eq. (B.4) with adequate choice of stepsize suffices to prove geometric convergence.

$$\begin{gathered} \parallel z_{t+1}-{x\parallel }_{{\ell }_2}^2={\Vert z_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right)-x\Vert }_{{\ell }_2}^2 \\ \stackrel{\left(a\right)}{=}{\Vert z_t-x\Vert }_{{\ell }_2}^2-2{\mu }_{t+1}\nabla ℒ{\left(z_t\right)}^T\left(z_t-x\right)+{\mu }_{t+1}^2{\Vert \nabla ℒ\left(z_t\right)\Vert }_{{\ell }_2}^2 \\ \stackrel{\left(b\right)}{\le }\left(1-2{\mu }_{t+1}A\right){\Vert z_t-x\Vert }_{{\ell }_2}^2+{\mu }_{t+1}\left({\mu }_{t+1}-2B\right){\Vert \nabla ℒ\left(z_t\right)\Vert }_{{\ell }_2}^2 \\ \stackrel{\left(c\right)}{\le }\left(1-2{\mu }_{t+1}A\right){\Vert z_t-x\Vert }_{{\ell }_2}^2. \end{gathered}$$

In (a) we expand the square. In (b) we apply eq. (B.4). In (c) we choose ${\mu }_{t+1}\in \left(0,2B\right]$, making the second term negative. Applying this relation inductively over $T$ steps of gradient descent yields geometric convergence with rate $1-2\mu A$, provided that the step size ${\mu }_t$ is less than $2B$ for $t>1$.

The number of measurements $m$ required in theorem 1 is the maximum over the number of samples required for each of the first two steps of the proof. For the first step, $m\ge 5n$ measurements are sufficient to reach our neighborhood of radius $\frac{1}{4}\parallel x{\parallel }_{{\ell }_2}$. For case 1 in the second step, $m\ge \frac{\overline{c_{1}n}n}{{\alpha }^2\parallel x{\parallel }_{{\ell }_2}^2}$ measurements are sufficient for correlation concentration. For case 2 in the second step, $m\ge \frac{c_{2}n}{{(\sqrt{2\alpha }-\sqrt{\alpha })}^2}$ measurements are sufficient for concentration. Maximizing over these bounds yields the result in theorem 1 (note that the constants change during the maximization).

B.2. First Step: Welcome to the Neighborhood

We first consider what happens in expectation when we take our first gradient step starting from an initialization at $z_0=0$. The expectation is over the randomness in the Gaussian measurement vector $a$. We have

$$\begin{gathered} {\mathbb{E}}_a\left[z_1\right]={\mathbf{0}}_n-{\mu }_1{\mathbb{E}}_a\left[{\nabla ℒ\left(z;a\right)|}_{z=0}\right] \\ =-{\mu }_1{\mathbb{E}}_a\left[a{f}^{\prime }\left(0\right)\left(f\left(0\right)-y\right)\right] \\ \stackrel{\left(a\right)}{=}-{\mu }_1{\mathbb{E}}_a\left[-\frac{1}{2}ay\right] \\ \stackrel{\left(b\right)}{=}\frac{{\mu }_1}{2}{\mathbb{E}}_a\left[a\left(1-\text{exp}\left(-{\left(a^{T}x\right)}_{+}\right)\right)\right] \\ \stackrel{\left(c\right)}{=}-\frac{{\mu }_1}{2}{\mathbb{E}}_a\left[a\text{ exp}\left(-{\left(a^{T}x\right)}_{+}\right)\right)] \\ \stackrel{\left(d\right)}{=}-\frac{{\mu }_1}{2}{\mathbb{E}}_a\left[\frac{x{x}^T}{\parallel x{\parallel }_{{\ell }_2}^2}a\text{ exp}\left(-{\left(a^{T}x\right)}_{+}\right)\right)] \\ \stackrel{\left(e\right)}{=}-\frac{{\mu }_1}{2}\frac{x}{\parallel x{\parallel }_{{\ell }_2}}{\mathbb{E}}_g\left[g\text{ exp}\left(-g_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)] \\ \stackrel{\left(f\right)}{=}\frac{{\mu }_1}{4}\text{ exp }\left(\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\text{erfc}\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)x \\ \stackrel{\left(g\right)}{=}x. \end{gathered}$$

In (a) we evaluate $f\left(0\right)=0$ and $f^{\prime }\left(0\right)=\frac{1}{2}$, where for the latter we use $\frac{1}{2}$ as the sub-differential even though $f$ is nondifferentiable at 0 due to the non-differentiability of ReLU (this choice is also justified as it is the expected gradient of $f$ around a small random initialization around 0). In (b) we plug in the value of the measurement $y$. In (c) we use linearity of expectation and evaluate the first term, ${\mathbb{E}}_a\left[a\right]=0_n$. In (d) we separate the leading $a$ into components parallel and orthogonal to $x$, and evaluate the expectation of the orthogonal term to zero. In (e) we replace $a^{T}x$ with $g\parallel x{\parallel }_{{\ell }_2}$ for a scalar Gaussian $g$, as these have the same distribution. In (f) we evaluate the remaining expectation. In (g) we choose ${\mu }_1=4\text{exp}\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc}\left(\frac{\parallel x{\parallel }_2}{\sqrt{2}}\right)}$ so that in expectation, our first step exactly recovers the signal $x$.

Because in practice we do not have access to an expectation over infinite measurements, we also care about the concentration of this first gradient step. This first-step concentration determines the local neighborhood around the signal $x$ that our gradient descent will operate within for the remaining iterations.

$$\begin{gathered} -{\mu }_1\nabla ℒ\left(z=0\right)=\frac{{\mu }_1}{2m}\sum _{i=1}^m{a}_i\left(1-\text{exp}\left(-{\left(a_i^{T}x\right)}_{+}\right)\right) \\ \stackrel{\left(a\right)}{=}\frac{{\mu }_1}{2m}\sum _m^{i=1}\frac{x{x}^T}{\parallel x{\parallel }_{{\ell }_2}^2}{a}_i\left(1-\text{exp}\left(-{\left(a_i^{T}x\right)}_{+}\right)\right)+\frac{{\mu }_1}{2m}\sum _m^{i=1}\left(⫿-\frac{x{x}^T}{\parallel x{\parallel }_{{\ell }_2}^2}\right)a_i\left(1-\text{exp}\left(-{\left(a_i^{T}x\right)}_{+}\right)\right) \\ \stackrel{\left(b\right)}{=}\frac{{\mu }_1}{2m}\sum _{i=1}^m{g}_i\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sum _{i=1}^m\left(⫿-\frac{x{x}^T}{\parallel x{\parallel }_{{\ell }_2}^2}\right)a_i\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right) \\ \stackrel{\left(c\right)}{=}\frac{{\mu }_1}{2m}\sum _m^{i=1}{g}_i\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _m^{i=1}{\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)}^2{a}_{\perp }} \\ \stackrel{\left(d\right)}{=}\frac{{\mu }_1}{2m}\sum _m^{i=1}{\hat{g}}_i\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _m^{i=1}{\tilde{g}}_i{a}_{\perp }.} \end{gathered}$$

In (a) we separate $a_i$ into components parallel and orthogonal to $x$. In (b) we replace $a_i^{T}x$ with $\parallel x{\parallel }_{{\ell }_2}{g}_i$ for a standard scalar Gaussian $g_i$, as these have the same distribution. In (c) we simplify the second term by rewriting it with a standard Gaussian vector $a_{\perp }$ with $n-1$ free dimensions (constrained to be orthogonal to $x$), and $a_{\perp }$ independent of $g_i$ for all $i$. In (d) we write ${\hat{g}}_i≔g_i\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)$ and ${\tilde{g}}_i≔{\left(1-\text{exp}\left(-{\left(g_i\right)}_{+}\parallel x{\parallel }_{{\ell }_2}\right)\right)}^2$.

Note that the first term has expectation $\frac{{\mu }_1}{4}\text{ exp }\left(\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\text{ erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)x=x$ computed above, and the second term has expectation zero because ${\tilde{g}}_i$ and $a_{\perp }$ are independent, and $a_{\perp }$ is mean zero. The first term is aligned with the signal $x$ but with a scaling factor $\frac{\mu }{2m\parallel x{\parallel }_{{\ell }_2}}{\sum }_{i=1}^m{\hat{g}}_i$; we can bound the deviation of this scaling from its mean using Hoeffding’s concentration bound for sums of sub-Gaussian random variables. We use the definition of sub-Gaussianity provided by Definition 2.2 in Wainwright (2019), for which ${\hat{g}}_i$ has sub-Gaussian parameter 1. We omit the leading constant $\frac{\mu }{2m\parallel x{\parallel }_{{\ell }_2}}$, and apply the Hoeffding bound as presented in Proposition 2.5 in Wainwright (2019) to conclude that

$$\left|\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\right|\le s$$

with probability at least $1-2e^{-\frac{s^2}{2m}}$.

The second term is a nonnegative scaling $\frac{\mu }{2m}\sqrt{{\sum }_{i=1}^m{\tilde{g}}_i}$ times a standard $n$-dimensional Gaussian vector $a_{\perp }$ with $n-1$ degrees of freedom (because it is orthogonal to $x$). Since ${\tilde{g}}_i$ is bounded in [0, 1], we can upper bound this scaling by $\frac{\mu }{2\sqrt{m}}$. Then the norm of the random vector $a_{\perp }$ can also be bounded using Exercise 5.2.4 in Vershynin (2018), to conclude that

$${\Vert a_{\perp }\Vert }_{{\ell }_2}\le \sqrt{\mathbb{E}\left[{\Vert a_{\perp }\Vert }_{{\ell }_2}^2\right]}+t\stackrel{\left(a\right)}{=}\sqrt{n-1}+t\stackrel{\left(b\right)}{\le }\left(1+t\right)\sqrt{n}$$

with probability at least $1-e^{-c{t}^{2}n}$ for a constant $c$, where in (a) we evaluate the expectation of a Gaussian norm with $n-1$ degrees of freedom and in (b) we upper bound $n-1$ by $n$ and do a change of variables to replace $t$ with $t\sqrt{n}$. Putting these together with a union bound, we have that

$$\begin{gathered} {\Vert z_1-\mathbb{E}\left[z_1\right]\Vert }_{{\ell }_2}={\Vert z_1-x\Vert }_{{\ell }_2}\le \frac{{\mu }_1}{2m}\left|\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\right|+\frac{{\mu }_1}{2\sqrt{m}}{\Vert a_{\perp }\Vert }_{{\ell }_2} \\ \le \frac{{\mu }_1}{2}\left(\frac{s}{m}+\left(1+t\right)\sqrt{\frac{n}{m}}\right) \\ =2\text{ exp}\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc}\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(\frac{s}{m}+\left(1+t\right)\sqrt{\frac{n}{m}}\right) \\ \stackrel{\left(a\right)}{=}2\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc}\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(s+\left(1+t\right)\sqrt{\frac{n}{m}}\right) \end{gathered}$$

with probability at least $1-2e^{-\frac{s^{2}m}{2}}-e^{-c{t}^{2}n}$, where in (a) we change variables and replace $s$ with $sm$, and $c$ is a constant. If we choose $s=t=1$, this simplifies to

$$\parallel z_1-{x\parallel }_{{\ell }_2}\le 2\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(1+2\sqrt{\frac{n}{m}}\right)$$

with probability at least $1-2e^{-\frac{m}{2}}-e^{-cn}$, for a constant $c$. For our first step to lie within a distance of $\frac{1}{4}\parallel x{\parallel }_{{\ell }_2}$, we need the number of measurements to satisfy

$$m\ge \frac{n}{{\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{16}\text{ exp }\left(\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\text{ erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)-\frac{1}{2}\right)}^2}.$$

The denominator in this expression is lower bounded by 0.2 for all $\parallel x{\parallel }_{{\ell }_2}$, so we can also guarantee the first step concentration to this neighborhood using $m\ge 5n$ measurements.

B.3. Correlation Concentration: Case 1

Consider the correlation

$$\nabla ℒ{(z)}^T\left(z-x\right)=\frac{1}{m}{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-a_i^{T}z}\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-a_i^{T}z}\right)\left(a_i^{T}h\right)$$

where $h≔z-x$. Now note that if we have $a_i^{T}x\ge 0$ and $a_i^{T}h\ge 0$ it implies $a_i^{T}x+a_i^{T}h=a_i^{T}z\ge 0$. Thus we have

$${\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}\ge {\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^{T}h\ge 0\right\}}.$$

Using the above we can conclude that

$$\begin{gathered} \nabla ℒ{\left(z\right)}^T\left(z-x\right)\stackrel{\left(a\right)}{\ge }\frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^{T}h\ge 0\right\}}{e}^{-a_i^{T}z}\left(e^{-a_i^{T}x}-e^{-a_i^{T}z}\right)\left(a_i^{T}h\right) \\ \stackrel{\left(b\right)}{=}\frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^{T}h\ge 0\right\}}{e}^{-2a_i^{T}x}{e}^{-a_i^{T}h}\left(1-e^{-a_i^{T}h}\right)\left(a_i^{T}h\right). \end{gathered}$$

In (a) we plug in the indicator inequality, and accordingly remove the now-superfluous ReLU. In (b) we use the $h$ notation, and regroup terms. To continue, we divide both sides by $\parallel h{\parallel }_{{\ell }_2}^2$ and use the notation $\hat{h}=\frac{h}{\parallel h{\parallel }_{{\ell }_2}}$.

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^T\hat{h}\ge 0\right\}}{e}^{-2a_i^{T}x}\frac{e^{-\parallel h{\parallel }_{{\ell }_2}{a}_i^T\hat{h}}\left(1-e^{-\parallel h{\parallel }_{{\ell }_2}{a}_i^T\hat{h}}\right)}{\parallel h{\parallel }_{{\ell }_2}}{a}_i^T\hat{h}.$$

To continue note that the function

$$g\left(x,s\right)=\frac{e^{-sx}\left(1-e^{-sx}\right)}{s}$$

has non-positive derivative as

$$\frac{\partial g}{\partial s}=\frac{e^{-2sx}\left(2sx-e^{sx}\left(sx+1\right)+1\right)}{s^2}\le 0$$

for all values of $s$ and $x$. This implies that $g\left(x,s\right)$ is a non-increasing function of $s$. Thus, we can conclude that

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{1}{rm}\sum _{i=1}^m{1}_{\left\{a_i^{T}x\ge 0\right\}}{1}_{\left\{a_i^T\hat{h}\ge 0\right\}}{e}^{-2a_i^{T}x}{e}^{-r{a}_i^T\hat{h}}\left(1-e^{-r{a}_i^T\hat{h}}\right)a_i^T\hat{h}.$$

Thus we can focus on lower bounding

$$\frac{1}{rm}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^T\hat{h}\ge 0\right\}}{e}^{-2a_i^{T}x}{e}^{-r{\left(a_i^T\hat{h}\right)}_{+}}\left(1-e^{-r{\left(a_i^T\hat{h}\right)}_{+}}\right){\left(a_i^T\hat{h}\right)}_{+}$$

over the set

$$\left\{\hat{h}\in {\mathbb{S}}^{n-1}:\frac{x^T\hat{h}}{\parallel x{\parallel }_{{\ell }_2}}\ge \rho \right\},$$

where we reintroduce superfluous ReLUs around $a_i^T\hat{h}$ as it will be convenient in the next steps. To continue, note that the function $f\left(h\right)=e^{-rh}\left(1-e^{-rh}\right)h$ has derivative

$$f^{\prime }\left(h\right)=e^{-2rh}\left(-1-e^{rh}\left(rh-1\right)+2rh\right).$$

It is easy to verify numerically that for $h\ge 0$ this gradient is maximized around $h_{\text{max}}\approx \frac{0.402673}{r}$, with maximum value $f^{\prime }\left(h_{\text{max}}\right)\approx 0.312334$. Thus, for all $h\ge 0$

$$f^{\prime }\left(h\right)\le \frac{1}{3}.$$

As a result the function $g\left(h\right)=e^{-h_{+}}\left(1-e^{-h_{+}}\right)h_{+}$ is a $\frac{1}{3}$-Lipschitz function of $h$. Thus for any $h_1\in {\mathbb{1}}^d$ and $h_2\in {\mathbb{1}}^d$ we have

$$\left|e^{-{\left(a_i^T{h}_2\right)}_{+}}\left(1-e^{-{\left(a_i^T{h}_2\right)}_{+}}\right){\left(a_i^T{h}_2\right)}_{+}-e^{-{\left(a_i^T{h}_1\right)}_{+}}\left(1-e^{-{\left(a_i^T{h}_1\right)}_{+}}\right){\left(a_i^T{h}_1\right)}_{+}\right|\le \frac{1}{3}\left|a_i^T\left(h_2-h_1\right)\right|.$$

We now define the random variable ${𝒳}_i\left(\hat{h}\right)={\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{e}^{-2a_i^{T}x}{e}^{-{\left(a_i^T\hat{h}\right)}_{+}}\left(1-e^{-{\left(a_i^T\hat{h}\right)}_{+}}\right){\left(a_i^T\hat{h}\right)}_{+}$ and note that the Lipschitzness of $g$ implies that for any $h_1,h_2\in {\mathbb{1}}^d$ we have

$$\begin{gathered} \left|{𝒳}_i\left(h_2\right)-{𝒳}_i\left(h_1\right)\right| \\ ={\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{e}^{-2a_i^{T}x}\left|e^{-{\left(a_i^T{h}_2\right)}_{+}}\left(1-e^{-{\left(a_i^T{h}_2\right)}_{+}}\right){\left(a_i^T{h}_2\right)}_{+}-e^{-{\left(a_i^T{h}_1\right)}_{+}}\left(1-e^{-{\left(a_i^T{h}_1\right)}_{+}}\right){\left(a_i^T{h}_1\right)}_{+}\right| \\ \le \frac{1}{3}\left|a_i^T\left(h_2-h_1\right)\right|. \end{gathered}$$

Since $a_i^T\left(h_2-h_1\right)$ is a sub-Gaussian random variable with sub-Gaussian norm on the order of ${\Vert h_2-h_1\Vert }_{{\ell }_2}$, we have that

$${\Vert {𝒳}_i\left(h_2\right)-{𝒳}_i\left(h_1\right)\Vert }_{{\psi }_2}\le \frac{c}{2}{\Vert h_2-h_1\Vert }_{{\ell }_2}$$

for some constant $c$. We use $c$ to denote any universal constant; note that this constant may vary between different lines. Thus using the centering rule for sub-Gaussian random variables, the centered processes ${\overline{𝒳}}_i\left(h\right)≔{𝒳}_i\left(h\right)-\mathbb{E}\left[{𝒳}_i\left(h\right)\right]$ obey

$${\Vert {\overline{𝒳}}_i\left(h_2\right)-{\overline{𝒳}}_i\left(h_1\right)\Vert }_{{\psi }_2}\le c{\Vert h_2-h_1\Vert }_{{\ell }_2}.$$

Using the rotational invariance property of sub-Gaussian random variables, this implies that the stochastic process

$$𝒳\left(h\right)≔\frac{1}{m}\sum _{i=1}^m{𝒳}_i\left(h\right)-\mathbb{E}\left[{𝒳}_i\left(h\right)\right]$$

has sub-Gaussian increments. That is,

$${\Vert 𝒳\left(h_2\right)-𝒳\left(h_1\right)\Vert }_{{\psi }_2}\le \frac{c}{\sqrt{m}}{\Vert h_2-h_1\Vert }_{{\ell }_2}.$$

Thus using Exercise 8.6.5 of Vershynin (2018) we can conclude that

$$\underset{\parallel \hat{h}{\parallel }_{{\ell }_2}=1}{\text{sup}}\frac{\left|𝒳\left(\hat{h}\right)\right|}{r}\le \frac{c}{\sqrt{m}r}\left(\sqrt{n}+u\right)$$

holds with probability at least $1-2e^{-u^2}$. Thus, we conclude that for all $\mathbf{h}$ obeying $\parallel \mathbf{h}{\parallel }_{{\ell }_2}\le r$ we have

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{1}{r}\mathbb{E}\left[1_{\left\{a^{T}x\ge 0\right\}}{1}_{\left\{a^T\hat{h}\ge 0\right\}}{e}^{-2a^{T}x}{e}^{-r{a}_i^T\hat{h}}\left(1-e^{-r{a}_i^T\hat{h}}\right)a_i^T\hat{h}\right]-\frac{c}{r}\frac{\sqrt{n}}{\sqrt{m}}$$

with probability at least $1-2e^{-n}$, for a constant $c$.

We can estimate and lower bound the expectation above using a numerical average over many (50000) two-dimensional Gaussian samples, with the two dimensions corresponding to ${\mathbf{a}}^T\stackrel{\mathbf{ˆ}}{\mathbf{h}}$ and $\frac{{\mathbf{a}}^T\mathbf{x}}{\parallel \mathbf{x}{\parallel }_{{\ell }_2}}$, minimizing over all correlations between $\mathbf{x}$ and $\mathbf{h}$ at least $\rho$ (i.e. all correlations in this case). We arrive at the following lower bound, for $r=\frac{1}{4}\parallel \mathbf{x}{\parallel }_{{\ell }_2}$ and $\rho =-0.6$.

$$\frac{1}{r}\mathbb{E}\left[{\mathbb{1}}_{\left\{a^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a^T\hat{h}\ge 0\right\}}{e}^{-2a^{T}x}{e}^{-r{a}_i^T\hat{h}}\left(1-e^{-r{a}_i^T\hat{h}}\right)a_i^T\hat{h}\right]\ge e^{-\sqrt{10\parallel x{\parallel }_{{\ell }_2}+7}}.$$

This bound is illustrated in a “proof by picture” in Figure 4.

Figure 4:

Lower bound on the expected correlation in case 1.

B.4. Correlation Concentration: Case 2

In this case we will focus on controlling the correlation inequality in the region where

$$\left\{h\in {\mathbb{1}}^n:\parallel h{\parallel }_{{\ell }_2}\le r\text{ and }\frac{x^{T}h}{\parallel x{\parallel }_{{\ell }_2}\parallel h{\parallel }_{{\ell }_2}}\le \rho \right\}.$$

Consider the correlation

$$\begin{gathered} \nabla ℒ{(z)}^T\left(z-x\right)=\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-a_i^{T}z}\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-a_i^{T}z}\right)\left(a_i^{T}h\right) \\ \stackrel{\left(a\right)}{\ge }\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-a_i^{T}z}\left(e^{-a_i^{T}x}-e^{-a_i^{T}z}\right)\left(a_i^{T}h\right) \\ \stackrel{\left(b\right)}{=}\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a_i^{T}h\ge -a_i^{T}x\right\}}{e}^{-2a_i^{T}x}{e}^{-a_i^{T}h}\left(1-e^{-a_i^{T}h}\right)\left(a_i^{T}h\right) \\ \stackrel{\left(c\right)}{\ge }\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\ge a_i^{T}h\ge -a_i^{T}x\right\}}{e}^{-2a_i^{T}x}{e}^{-a_i^{T}h}\left(1-e^{-a_i^{T}h}\right)\left(a_i^{T}h\right). \end{gathered}$$

In (a) we provide a lower bound by introducing an additional indicator function on ${\mathbf{a}}_i^T\mathbf{x}$, which allows us to remove the ReLU. In (b) we use the notation $\mathbf{h}≔\mathbf{z}-\mathbf{x}$, and combine terms. In (c) we again provide a lower bound by adding an indicator to restrict ${\mathbf{a}}_i^T\mathbf{h}\le 0$. By flipping the sign of $\mathbf{h}$ we can alternatively lower bound

$$\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\le a_i^{T}h\le a_i^{T}x\right\}}{e}^{-2a_i^{T}x}{e}^{a_i^{T}h}\left(e^{a_i^{T}h}-1\right)\left(a_i^{T}h\right)$$

over the set

$$\left\{h\in {\mathbb{1}}^n:\parallel h{\parallel }_{{\ell }_2}\le r\text{ and }\frac{x^{T}h}{\parallel x{\parallel }_{{\ell }_2}\parallel h{\parallel }_{{\ell }_2}}\ge -\rho \right\}.$$

To this aim note that for $s\ge 0$ we have $e^s\ge 1$ and $\left(e^s-1\right)s\ge s^2$. Thus,

$$\nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{1}{m}{\sum }_{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\le a_i^{T}h\le a_i^{T}x\right\}}{e}^{-2a_i^{T}x}{\left(a_i^{T}h\right)}^2.$$

To continue, we introduce the notation $\stackrel{\mathbf{ˆ}}{\mathbf{h}}=\frac{h}{\parallel \mathbf{h}{\parallel }_{{\ell }_2}}$, and divide both sides by $\parallel \mathbf{h}{\parallel }_{{\ell }_2}^2$. Thus we have

$$\begin{gathered} \frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{\left(z\right)}^T\left(z-x\right)\ge \frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\le a_i^T\hat{h}\le a_i^{T}x/\parallel h{\parallel }_{{\ell }_2}\right\}}{e}^{-2a_i^{T}x}{\left(a_i^T\hat{h}\right)}^2 \\ \ge \frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\{0\le a_i^T\hat{h}\le \frac{a_i^{T}x}{r}}{e}^{-2a_i^{T}x}{\left(a_i^T\hat{h}\right)}^2. \end{gathered}$$

Thus it suffices to lower bound

$$\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\le a_i^T\hat{h}\le \frac{a_i^{T}x}{r}\right\}}{e}^{-2a_i^{T}x}{\left(a_i^T\hat{h}\right)}^2$$

over

$$\left\{\hat{h}\in {\mathbb{S}}^{n-1}:\frac{x^T\hat{h}}{\parallel x{\parallel }_{{\ell }_2}}\ge -\rho \right\}.$$

To continue using Jensen’s inequality we have

$$\begin{gathered} \frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{\left(z\right)}^T\left(z-x\right)\ge {\left(\frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{0\le a_i^T\hat{h}\le \frac{a_i^{T}x}{r}\right\}}{e}^{-a_i^{T}x}{a}_i^T\hat{h}\right)}^2 \\ \ge {\left(\frac{1}{m}\sum _m^{i=1}{\mathbb{1}}_{\left\{a_i^{T}x\ge 0\right\}}𝒮\left(a_i^T\hat{h};\frac{a_i^{T}x}{r}\right)e^{-a_i^{T}x}\right)}^2 \end{gathered}$$

where we have defined the function

$$𝒮\left(v;w\right)=\{\begin{array}{ll}0\hfill & v<0\hfill \\ v\hfill & 0\le v\le \frac{w}{2}\hfill \\ w-v\hfill & \frac{w}{2}\le v\le w\hfill \\ 0\hfill & v\ge w\hfill \end{array}$$

which is a 1-Lipschitz function of $v$. We now define the random variable ${𝒳}_i\left(\stackrel{\mathbf{ˆ}}{\mathbf{h}}\right)=1_{\left\{{\mathbf{a}}_i^T\mathbf{x}\ge 0\right\}}𝒮\left({\mathbf{a}}_i^T\stackrel{\mathbf{ˆ}}{\mathbf{h}};\frac{{\mathbf{a}}_i^T\mathbf{x}}{r}\right)e^{-{\mathbf{a}}_i^T\mathbf{x}}$ and note that the Lipschitzness of $𝒮$ implies that for any ${\mathbf{h}}_1,{\mathbf{h}}_2\in {\mathbb{R}}^d$ we have

$$\left|{𝒳}_i\left(h_2\right)-{𝒳}_i\left(h_1\right)\right|\le \left|a_i^T\left(h_2-h_1\right)\right|$$

Since ${\mathbf{a}}_i^T\left({\mathbf{h}}_2-{\mathbf{h}}_1\right)$ is a sub-Gaussian random variable with sub-Gaussian norm on the order of ${∥{\mathbf{h}}_2-{\mathbf{h}}_1∥}_{{\ell }_2}$, we have that

$${\Vert {𝒳}_i\left(h_2\right)-{𝒳}_i\left(h_1\right)\Vert }_{{\psi }_2}\le \frac{c}{2}{\Vert h_2-h_1\Vert }_{{\ell }_2}$$

for some constant $c$. We use $c$ to denote any universal constant; note that this constant may vary between different lines. Using the centering rule for sub-Gaussian random variables, the centered processes ${\stackrel{-}{𝒳}}_i\left(\mathbf{h}\right)≔{𝒳}_i\left(\mathbf{h}\right)-\mathbb{E}\left[{𝒳}_i\left(\mathbf{h}\right)\right]$ obey

$${\Vert {\overline{𝒳}}_i\left(h_2\right)-{\overline{𝒳}}_i\left(h_1\right)\Vert }_{{\psi }_2}\le c{\Vert h_2-h_1\Vert }_{{\ell }_2}.$$

Using the rotational invariance property of sub-Gaussian random variables, this implies that the stochastic process

$$𝒳\left(h\right)≔\frac{1}{m}\sum _{i=1}^m{𝒳}_i\left(h\right)-\mathbb{E}\left[{𝒳}_i\left(h\right)\right]$$

has sub-Gaussian increments. That is,

$${\Vert 𝒳\left(h_2\right)-𝒳\left(h_1\right)\Vert }_{{\psi }_2}\le \frac{c}{\sqrt{m}}{\Vert h_2-h_1\Vert }_{{\ell }_2}.$$

Thus using Exercise 8.6.5 of Vershynin (2018) we can conclude that

$$\underset{\parallel \hat{h}{\parallel }_{{\ell }_2}=1,{\text{cos}}^{-1}\left(\frac{x^T\hat{h}}{\parallel x{\parallel }_{{\ell }_2}}\right)\le \delta }{\text{sup}}\left|𝒳\left(\hat{h}\right)\right|\le \frac{c}{\sqrt{m}}\left(\sqrt{n}{e}^{-\frac{n{\text{ cos}}^2\left(\delta \right)}{2}}+u\right)$$

holds with probability at least $1-2e^{-u^2}$. In the last line we used the fact that the surface area of a spherical cap with distance at least $ϵ$ away from the center is bounded by $e^{-n\frac{{ϵ}^2}{2}}$. By using $u=\sqrt{n}$, this implies that

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge {\left(\mathbb{E}\left[{\mathbb{1}}_{\left\{a^{T}x\ge 0\right\}}𝒮\left(a^T\hat{h};\frac{a_i^{T}x}{r}\right)e^{-a^{T}x}\right]-c\frac{\sqrt{n}}{\sqrt{m}}\right)}^2$$

holds with probability at least $1-2e^{-n}$, for a constant $c$.

We can estimate and lower bound the expectation above using a numerical average over many (50000) two-dimensional Gaussian samples, with the two dimensions corresponding to ${\mathbf{a}}^T\stackrel{\mathbf{ˆ}}{\mathbf{h}}$ and $\frac{{\mathbf{a}}^T\mathbf{x}}{\parallel \mathbf{x}{\parallel }_{{\ell }_2}}$, minimizing over all correlations between $\mathbf{x}$ and $\mathbf{h}$ at most $\rho$ (i.e. all correlations in this case). We arrive at the following lower bound, for $r=\frac{1}{4}\parallel \mathbf{x}{\parallel }_{{\ell }_2}$ and $\rho =-0.6$.

$$\mathbb{E}{\left[{\mathbb{1}}_{\left\{a^{T}x\ge 0\right\}}𝒮\left(a^T\hat{h};\frac{a_i^{T}x}{r}\right)e^{-a^{T}x}\right]}^2\ge e^{-\left(5\parallel x{\parallel }_{{\ell }_2}+2\right)}.$$

This bound is illustrated in a “proof by picture” in Figure 5

Figure 5:

Lower bound on the expected correlation in case 2.

B.5. Bounding the Gradient Norm

Consider the gradient and note that

$$\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}=\underset{u\in {\mathbb{S}}^{n-1}}{\text{sup}}\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-{\left(a_i^{T}z\right)}_{+}}\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-{\left(a_i^{T}z\right)}_{+}}\right)\left(a_i^{T}u\right),$$

where ${\mathbb{S}}^{n-1}$ denotes the set of all real $n$-dimensional unit-norm vectors. To continue, we use the Cauchy-Schwarz inequality:

$$\begin{gathered} \parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}\le \sqrt{\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-2{\left(a_i^{T}z\right)}_{+}}{\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-{\left(a_i^{T}z\right)}_{+}}\right)}^2}\sqrt{\underset{u\in {\mathbb{S}}^{n-1}}{\text{sup}}\frac{1}{m}\sum _{i=1}^m{\left(a_i^{T}u\right)}^2} \\ =\sqrt{\frac{1}{m}\sum _{i=1}^m{\mathbb{1}}_{\left\{a_i^{T}z\ge 0\right\}}{e}^{-2{\left(a_i^{T}z\right)}_{+}}{\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-{\left(a_i^{T}z\right)}_{+}}\right)}^2}\frac{\parallel A\parallel }{\sqrt{m}} \\ \le \sqrt{\frac{1}{m}\sum _{i=1}^m{\left(e^{-{\left(a_i^{T}x\right)}_{+}}-e^{-{\left(a_i^{T}z\right)}_{+}}\right)}^2\frac{\parallel A\parallel }{\sqrt{m}}} \\ \stackrel{\left(a\right)}{\le }\sqrt{\frac{1}{m}\sum _{i=1}^m{\left(a_i^T\left(z-x\right)\right)}^2\frac{\parallel A\parallel }{\sqrt{m}}} \\ \le \frac{\parallel A{\parallel }^2}{m}\parallel z-x{\parallel }_{{\ell }_2}, \end{gathered}$$

where $\parallel \mathbf{A}\parallel$ is the operator norm of a matrix comprised by stacking the vectors ${\mathbf{a}}_i$, and in (a) we used the fact that the function $f\left(z\right)=e^{-\left(z\right)+}$ is 1-Lipschitz. Finally, using the fact that $\parallel \mathbf{A}\parallel \le 2(\sqrt{m}+\sqrt{n})$ with probability at least $1-e^{-0.5(m+n)}$, we conclude that

$$\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}\le 8\left(1+\frac{n}{m}\right)\parallel z-x{\parallel }_{{\ell }_2}≔L\parallel z-x{\parallel }_{{\ell }_2}$$

with probability at least $1-e^{-0.5(m+n)}$, where $L$ is a constant since we have the number of measurements at least a constant times the number of unknowns. This completes the proof of smoothness of the gradient towards the global optimum.

C. Proof of Theorem 2

The general strategy of the proof is similar to Theorem 1 but requires delicate modifications in each step. Concretely, we have the following four steps.

Step I: First iteration: Welcome to the neighborhood.

In the first step we show that the first iteration obeys

$$\Vert z_1-x_{{\ell }_2}\Vert \le \frac{1}{4}\parallel x{\parallel }_{{\ell }_2}$$

with high probability as long as $m\ge c{m}_0$. We prove this in subsection C.1.

Step II: Local pseudoconvexity.

In this step we prove that the loss function is locally strongly pseudoconvex. Specifically we show that for all $z\in 𝒦$ that also belong to the local neighborhoood $𝒩\left(\mathbf{x}\right)≔\left\{\mathbf{z}\in {\mathbb{R}}^n:\parallel \mathbf{z}-\mathbf{x}{\parallel }_{{\ell }_2}\le \frac{1}{4}\parallel \mathbf{x}{\parallel }_{{\ell }_2}\right\}$ we have

$$\langle \nabla ℒ\left(z\right),z-x\rangle \ge \alpha \parallel z-x{\parallel }_{{\ell }_2}^2$$

with high probability as long as

$$m\ge \frac{c_1{e}^{c_2\parallel x{\parallel }_{{\ell }_2}}}{\parallel x{\parallel }_{{\ell }_2}^2}{m}_0.$$

Here, the value of $\alpha$ is the same as in Theorem 1 (see Appendix B.1). We prove this in subsection C.2 by again considering two cases.

Step III: Local smoothness.

We also use the fact that the loss function is locally smooth, that is,

$$\parallel \nabla ℒ\left(z\right){\parallel }_{{\ell }_2}\le 8\left(1+\frac{n}{m}\right)\parallel z-x{\parallel }_{{\ell }_2}≔L\parallel z-x{\parallel }_{{\ell }_2}$$

holds with probability at least $1-e^{-0.5\left(m+n\right)}$.per Section B.5.

Step IV: Completing the proof via combining steps I-III.

In this step we show how to combine the previous steps to complete the proof of the theorem. First, note that by the first step the first iteration will belong to the local neighborhood $𝒩\left(\mathbf{x}\right)$ and thus belongs to the set $𝒦\cap 𝒩\left(\mathbf{x}\right)$. Next, note that

$$\begin{gathered} \Vert z_{t+1}-x{\Vert }_{{\ell }_2}=\Vert {𝒫}_{𝒦}\left(z_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right)\right)-x{\Vert }_{{\ell }_2} \\ =\Vert {𝒫}_{𝒫}\left(h_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right)\right){\Vert }_{{\ell }_2} \\ \le \Vert {𝒫}_{𝒞}\left(h_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right)\right){\Vert }_{{\ell }_2} \\ \le \Vert h_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right){\Vert }_{{\ell }_2}. \end{gathered}$$

Squaring both sides and using the local psudoconvexity inequality from Step II we conclude that

$$\begin{gathered} \Vert z_{t+1}-x{\Vert }_{{\ell }_2}^2\le \Vert h_t-{\mu }_{t+1}\nabla ℒ\left(z_t\right){\Vert }_{{\ell }_2}^2 \\ =\Vert h_t{\Vert }_{{\ell }_2}^2-2{\mu }_{t+1}\langle h_t,\nabla ℒ\left(z_t\right)\rangle +{\mu }_{t+1}^2\Vert \nabla ℒ\left(z_t\right){\Vert }_{{\ell }_2}^2 \\ \le \Vert h_t{\Vert }_{{\ell }_2}^2-2{\mu }_{t+1}\alpha \Vert h_t{\Vert }_{{\ell }_2}^2+{\mu }_{t+1}^2\Vert \nabla ℒ\left(z_t\right){\Vert }_{{\ell }_2}^2. \end{gathered}$$

Next we use the local smoothness from Step III to conclude that

$$\begin{gathered} \Vert z_{t+1}-x{\Vert }_{{\ell }_2}^2\le \Vert h_t{\Vert }_{{\ell }_2}^2-2{\mu }_{t+1}\alpha \Vert h_t{\Vert }_{{\ell }_2}^2+{\mu }_{t+1}^2\Vert \nabla ℒ\left(z_t\right){\Vert }_{{\ell }_2}^2 \\ \le \Vert h_t{\Vert }_{{\ell }_2}^2-2{\mu }_{t+1}\alpha \Vert h_t{\Vert }_{{\ell }_2}^2+{\mu }_{t+1}^2{L}^2\Vert h_t{\Vert }_{{\ell }_2}^{2}NEW=\left(1-{\mu }_{t+1}\alpha \right)\Vert h_t{\Vert }_{{\ell }_2}^2 \\ =\left(1-{\mu }_{t+1}\alpha \right)\Vert z_t-x{\Vert }_{{\ell }_2}^2, \end{gathered}$$

where in the last line we used the fact that ${\mu }_{t+1}\le \frac{\alpha }{L^2}$, completing the proof.

C.1. Proof of Step I

The beginning of this proof is the same as the unregularized version where we note that

$$\begin{gathered} {\Vert z_1-x\Vert }_{{\ell }_2}={\Vert {𝒫}_{𝒦}\left(\frac{{\mu }_1}{2m}\sum _{i=1}^m{\hat{g}}_i\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _{i=1}^m{\tilde{g}}_i{a}_{\perp }}\right)-x\Vert }_{{\ell }_2} \\ =\parallel {𝒫}_{𝒫}{\left(\frac{{\mu }_1}{2m}\sum _{i=1}^m{\hat{g}}_i\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _{i=1}^m{\tilde{g}}_i{a}_{\perp }}-x\right)\parallel }_{{\ell }_2} \\ =\parallel {𝒫}_{𝒫}{\left(\frac{{\mu }_1}{2m}\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _{i=1}^m{\tilde{g}}_i{a}_{\perp }}\right)\parallel }_{{\ell }_2} \\ \le \parallel {𝒫}_{𝒞}{\left(\frac{{\mu }_1}{2m}\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _{i=1}^m{\tilde{g}}_i{a}_{\perp }}\right)\parallel }_{{\ell }_2} \\ \le \underset{v\in 𝒞\cap {\mathbb{S}}^{n-1}}{\text{sup}}{v}^T\left(\frac{{\mu }_1}{2m}\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\frac{x}{\parallel x{\parallel }_{{\ell }_2}}+\frac{{\mu }_1}{2m}\sqrt{\sum _{i=1}^m{\tilde{g}}_i{a}_{\perp }}\right) \\ \le \frac{{\mu }_1}{2}\left|\frac{1}{m}\sum _{i=1}^m\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\right|+\frac{{\mu }_1}{2\sqrt{m}}\underset{v\in 𝒞\cap {\mathbb{S}}^{n-1}}{\text{sup}}{v}^T{a}_{\perp }. \end{gathered}$$

Now similar to the unregularized case we have

$$\begin{gathered} \frac{{\mu }_1}{2m}\left|\sum _m^{i=1}\left({\hat{g}}_i-\mathbb{E}\left[\hat{g}\right]\right)\right|+\frac{{\mu }_1}{2\sqrt{m}}\underset{v\in 𝒞\cap {\mathbb{S}}^{n-1}}{\text{sup}}{v}^T{a}_{\perp }\le \frac{{\mu }_1}{2}\left(\frac{s}{m}+\left(1+t\right)\sqrt{\frac{m_0}{m}}\right) \\ =2\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(\frac{s}{m}+\left(1+t\right)\sqrt{\frac{m_0}{m}}\right) \\ \stackrel{\left(a\right)}{=}2\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(s+\left(1+t\right)\sqrt{\frac{m_0}{m}}\right) \end{gathered}$$

with probability at least $1-2e^{-\frac{s^{2}m}{2}}-e^{-c{t}^2{m}_0}$, where in (a) we change variables and replace $s$ with $sm$, and $c$ is a constant. If we choose $s=t=1$, this simplifies to

$${\Vert z_1-x\Vert }_{{\ell }_2}\le 2\text{ exp }\left(-\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\frac{1}{\text{erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)}\left(1+2\sqrt{\frac{m_0}{m}}\right)$$

with probability at least $1-2e^{-\frac{m}{2}}-e^{-c{m}_0}$, for a constant $c$. For our first step to lie within a distance of $\frac{1}{4}\parallel \mathbf{x}{\parallel }_{{\ell }_2}$, we need the number of measurements to satisfy

$$m\ge \frac{m_0}{{\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{16}\text{ exp }\left(\frac{\parallel x{\parallel }_{{\ell }_2}^2}{2}\right)\text{ erfc }\left(\frac{\parallel x{\parallel }_{{\ell }_2}}{\sqrt{2}}\right)-\frac{1}{2}\right)}^2}.$$

The denominator in this expression is lower bounded by 0.2 for all $\parallel \mathbf{x}{\parallel }_{{\ell }_2}$, so we can also guarantee the first step concentration to this neighborhood using $m\ge 5m_0$ measurements.

C.2. Proof of Step II

The proof of this step is virtually identical to that of the unregularized case. The only difference is that when we apply Exercise 8.6.5 of Vershynin (2018) $n$ is replaced with $m_0$ (indeed this exercise is stated with $m_0$). As a result in the two cases we conclude

Case I: In this case using the above yields

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge \frac{1}{r}\mathbb{E}\left[{\mathbb{1}}_{\left\{a^{T}x\ge 0\right\}}{\mathbb{1}}_{\left\{a^T\hat{h}\ge 0\right\}}{e}^{-2a^{T}x}{e}^{-r{a}_i^T\hat{h}}\left(1-e^{-r{a}_i^T\hat{h}}\right)a_i^T\hat{h}\right]-\frac{c}{r}\frac{\sqrt{m_0}}{\sqrt{m}}$$

holds with probability at least $1-2e^{-m_0}$, for a constant $c$. Case II: In this case using the above yields

$$\frac{1}{\parallel h{\parallel }_{{\ell }_2}^2}\nabla ℒ{(z)}^T\left(z-x\right)\ge {\left(\mathbb{E}\left[1_{\left\{a^{T}x\ge 0\right\}}𝒮\left(a^T\hat{h};\frac{a_i^{T}x}{r}\right)e^{-a^{T}x}\right]-c\frac{\sqrt{m_0}}{\sqrt{m}}\right)}^2$$

holds with probability at least $1-2e^{-m_0}$, for a constant $c$.

Thus the remainder of the proof is identical and the only needed change in this entire step is to replace $n$ with $m_0$.