The Practicality of Stochastic Optimization in Imaging Inverse Problems

Introduction

Stochastic gradient-based optimization algorithms have been ubiquitous in real world applications which involve solving large-scale and high-dimensional optimization tasks, particularly in the field of machine learning , due to their scalability to the size of the optimization problems. In this work we study the practicality of stochastic gradient-based optimization algorithms in imaging inverse problems, which are also large-scale and high-dimension by nature. The class of problems we consider, with typical examples including image deblurring, denoising, inpainting, superresolution, demosaicing, tomographic image reconstruction, etc, can be generally formulated as the following:

One of the most typical examples of the data fidelity term in imaging inverse problems is the least-squares loss:

while we typically obtain a robust estimator of $x^{\dagger}$ via jointly minimizing the least-squares data-fidelity term with a structure-inducing regularization $g(\cdot)$ which encodes prior information we have regarding $x^{\dagger}$:

Traditionally, the imaging inverse problems are solved most often by minimizing the regularized least-squares via the deterministic first-order solvers, such as the proximal gradient descent , and its accelerated and primal-dual variants . The iterates of the proximal gradient descent for solving (4) can be written as:

For least-squares data-fidelity term $\triangledown f(x^{i})=\frac{1}{n}A^{T}(Ax^{i}-b)$. We denote the proximal operator as:

Stochastic gradient descent methods , which is based on randomly selecting one or a few function $f_{i}(x)$ in each iteration and compute an efficient unbiased estimate of the full gradient $\triangledown f(x)$ and perform the descent step, by nature are the ideal solvers for the generic composite optimization task (1), including the regularized least-squares (4). These type of methods are able to achieve scalablity to large-scale problems compared to the deterministic gradient methods in many modern machine learning applications.

In recent years, researchers have developed several advanced variants of stochastic gradient methods, namely, the variance-reduced stochastic gradient methods . In each iteration of these new stochastic algorithms, more delicated stochastic gradient estimator is computed, which can reduce the variance of the stochastic gradient estimator progressively, with small computational or storage overheads, and hence significantly improve the convergence rate of stochastic gradient methods. Most recently, with further combining the variance-reduction technique with the Nesterov’s momentum acceleration technique which was originally designed to accelerate the deterministic gradient methods , researchers have developed several \sayoptimal stochastic gradient algorithms which can provably achieve the worse-case optimal convergence rate for (1).

While having been a proven success both in theory and in machine learning applications, there are few convincing results so far in the literature which report the performance of the stochastic gradient methods in imaging applications (with the possible exception of tomographic reconstruction ). Can stochastic gradient methods significantly facilitate inverse problems as they did for machine learning? If not, why might stochastic optimization be inefficient for some inverse problems? How can we understand such failures? How could we help practitioners to characterize whether a given inverse problem is suitable for stochastic gradient methods or not? This work is aimed at answering these questions in a systematic way.

We first report surprisingly negative results of stochastic gradient methods in solving a space-varying image deblurring problem, which go against the conventional wisdom and common believe of the large-scale optimization and machine learning community. The first step of this work is to find out the key factor which determines the success or failure of stochastic gradient methods to be more advantageous than their deterministic counterparts for an imaging inverse problem. We start by a motivational analysis from known upper and lower complexity bounds for solving (1), demonstrating that in the worst case the acceleration given by stochastic gradient methods in terms of objective-gap convergence is dominated by this ratio. In the context of imaging inverse problem, it is more desirable to further study whether the acceleration provided by stochastic gradient methods in terms of estimation-error convergence is also dominated by this ratio. To show this, we provide a novel analysis for the estimation-error convergence rate of minibatch proximal SGD in solving linear inverse problems with regularization constraints, under expected smoothness and restricted strong-convexity condition. By comparing our result for minibatch proximal SGD with the deterministic proximal gradient descent in the same setting, we can confirm that this ratio of Lipschitz constants is indeed the key factor which can be used to characterize whether a given inverse problem is suitable or not for applying stochatic gradient methods. Hence we find strong theoretical evidences, that the computational speedup which stochastic gradient methods can bring over their deterministic counterparts, is dominantly related to the ratio of the Lipschitz constants of the full gradient and the minibatch stochastic gradients.

Based on our theoretical analysis, we propose to evaluate the limit of possible acceleration of a stochastic gradient method over its full gradient counterpart by measuring the Stochastic Acceleration (SA) factors which are based on the ratio of the Lipschitz constants of the minibatched stochastic gradient and the full gradient. We also discover that the SA factors are able to characterize the benefits of using randomized optimization techniques, and that not all imaging problems have large SA factors.

1.2 Understanding the relationship between the structure of inverse problems and stochastic acceleration

An immediate and crucial question to be answered is, \saywhat types of inverse problems favor stochastic gradient algorithms?. We provide tight and insightful lower and upper bounds for the stochastic acceleration factors for partition minibatch schemes, using standard tools in numerical linear algebra . Our lower bound results suggest that the SA factor we propose is directly related the ratio $\frac{\max_{i\in[n]}\|a_{i}\|_{2}^{2}}{\|A\|^{2}}$, which can be efficiently evaluated by practitioners. This ratio $\frac{\max_{i\in[n]}\|a_{i}\|_{2}^{2}}{\|A\|^{2}}$ is also directly related to the eigenspectrum of the Hessian matrix, when the measurements are relatively balanced, i.e. $\frac{\max\|a_{i}\|_{2}^{2}}{\frac{1}{n}\sum_{j\in[n]}\|a_{j}\|_{2}^{2}}=O(1)$, which is generally true for most of imaging inverse problems:

If such an inverse problem’s Hessian matrix has a fast decaying eigenspectrum, then it can be guaranteed to have large SA factors, and hence can be characterized as a suitable application for stochastic gradient methods.

And vice-versa: if such an inverse problems’s Hessian matrix has a slowly-decaying eigenspectrum, then it is guaranteed to have small SA factors and can be deemed as unsuitable for stochastic gradient methods, not matter how delicately we partition the data.

While the spectral properties of the forward operator fundamentally controls the suitability of stochastic proximal gradient methods for an inverse problem, we know that for some inverse problems, different choices of partition can lead to different convergence rates for stochastic gradient algorithms in practice. One of our lower bounds for SA factors demonstrates that:

If a partition scheme generates minibatches which have low local coherence structure, i.e. the measurements within minibatches are less correlated to each other, then it is superior to other partition schemes which have high local coherence structure.

The SA factors and the lower bounds we propose provide for the practitioners efficient ways to check whether they should use stochastic proximal gradient techniques or classical deterministic proximal gradient methods to solve a given inverse problem, and also compare between different partition minibatch schemes and choose the best one among them in practice.

1.3 Breaking the computational bottleneck of expensive/multiple proximal operators for momentum SGD

Another factor in imaging applications which significantly affects the SGD-type methods’ actual performance is the frequent calculation of the costly proximal operators for the regularization terms, such as the TV semi-norm – SGD methods need to calculate these much more frequently than full gradient methods. Moreover most of the fast SGD methods can not cope with more than one non-smooth regularization term . To overcome these issues we propose an accelerated primal-dual SGD (Acc-PD-SGD) algorithm based on the primal-dual hybrid gradient framework , as a side-contribution. The proposed Acc-PD-SGD algorithm is able to efficiently handle (1) regularization with a linear operator, (2) multiple regularization terms, while (3) maintaining Nesterov-type accelerated convergence speed in practice.

2 Outline

Now we set out the rest of the paper. We start by presenting in section 2 a surprising negative result of state-of-the-art stochastic gradient methods in a space-varying image deblurring task. In section 3 we describe our notations and definitions which will be frequently used throughout the paper. Then in section 4, we provide theoretical analysis regarding the limitation of stochastic optimization algorithms, and particularly our novel analysis of minibatch SGD and the theory-inspired SA factors. In section 4, we also present bounds for the SA factors with respect to the spectral property of the forward operator, and hence derive a condition for an inverse problem to be a suitable application of stochastic gradient methods. In section 5 and 6, we present the accelerated primal-dual SGD algorithm and the numerical experiments. Final remarks appear in section 7, while we include the proofs of our theoretical results in the appendix.

A motivating example

Image deblurring is an important type of imaging inverse problems and has been studied intensely during the recent decades. For uniform deblurring, due to the cyclic structure of the deconvolution, FFT-based ADMMThe computationally demanding sub-problems of alternating direction method of multipliers (ADMM) in this case can be solved with an efficient matrix inversion by FFT due to the cyclic structure of the uniform deconvolution. variants have shown to be remarkably efficient when compared to classic gradient-based solvers such as FISTA . Such techniques, although being computationally efficient, are specifically tailored to a restricted range of problems where the observation models are diagonalizable by a DFT. For image deblurring, it is often not realistic to assume that imaging devices induce a uniform blur . If the blurring is different across the image, then the efficient implementation of ADMM is not applicable in general. Then standard ADMM and deterministic gradient methods such as FISTA can be computationally expensive. It is therefore natural to ask: can stochastic gradient methods offer us a more efficient solution?

We start by a simple space-varying deblurring example where a part (sized 256 by 256) of the \sayKodim04 image from Kodak Lossless True Color Image Suite is blurred with a space-varying blur kernel which imposes less blurring at the center but increasingly severe blurring towards the edge. For the shape of the blur kernel, we choose the out-of-focus kernel provided in . We also add a small amount of noise to the blurred image.

We test the effectiveness of several algorithms by solving the same TV-regularized least-squares problem, to get an estimation of the ground truth image. The algorithms we test in the experiments include the accelerated full gradient method FISTA , proximal SGD , the proximal SVRG and its accelerated variant, Katyusha algorithm which has achieved optimal convergence rate in theory for (1). Perhaps surprisingly, on this experiment we report a negative result in Fig.1 for all these randomized algorithms. The most efficient solver in this task is the full gradient method FISTA in terms of wall clock time and number of epochs (datapasses). The state-of-the-art stochastic gradient methods with Nesterov’s acceleration even cannot beat FISTA in terms of epoch counts. For all the randomized algorithms we choose a minibatch size which is 10 percent of the total data size. For stochastic gradient methods, a smaller minibatch size in this case did not provide better performance in datapasses and significantly slowed down running time due to the multiple calls on the proximal operator.

Notations and definitions

The smoothness conditions of the full batch $f(x)$ and minibatches $f_{S_{k}}(x)$ are formally described as the following:

(Smoothness of the Full-Batch and the Mini-Batches.) $f(\cdot)$ is $L_{f}$-smooth and each $f_{S_{k}}(\cdot)$ is $L_{b}$-smooth, that is:

We refer to and [46, Theorem 2.1.7] for details. In this paper we mainly consider two types of minibatch schemes, the partition minibatch sampling and random with-replacement sampling:

Limitations of stochastic optimization

The previous deblurring example appears to be contrary to the popular belief among the stochastic optimization community and the experience of machine learning practitioners, that stochastic gradient methods are much faster in terms of iteration complexity than deterministic gradient methods in solving large scale problems. To be specific – to achieve an objective gap suboptimality of $F(x)-F(x^{\star})\leq\varepsilon$, optimal stochastic gradient methods needs only $\Theta\left(n+\sqrt{n{L}/{\varepsilon}}\right)$ evaluations of $\triangledown f_{i}$ where $L$ denotes the gradient Lipschitz constant of $f_{i}$, see e.g. , while $\Theta\left(n\sqrt{{L}/{\varepsilon}}\right)$ are needed for optimal full gradient methods . Where is the loophole?

It is often easily ignored that the complexity results above are derived under different smoothness assumptions. For the convergence bound for full gradient, the full smooth part of the cost function $f(.)$ is assumed to be $L$-smooth, while for the case of stochastic gradient, every individual function $f_{i}(.)$ is assumed to be $L$-smooth. Now we can clearly see the subtlety: to compare these complexity results and make meaningful conclusions, one has to assume that these two Lipschitz constants are roughly the same. While this can be true, and is true for many problems, there are exceptions – image deblurring is one of them.

For the case where the minibatch size is $1$, we can denote the smoothness constants of $f(\cdot)$ and $f_{i}(\cdot)$ as $L_{f}$ and $L_{b}$ respectively, we illustrate here some extreme examples for the two smoothness constants to demonstrate this possible dramatic difference:

Let $f(x)=\frac{1}{2n}\|Ax-b\|_{2}^{2}=\frac{1}{n}\sum_{i=1}^{n}\frac{1}{2}(a_{i}^{T}x-b_{i})^{2}:=\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)$.

(1) If $a_{1}=a_{2}=a_{3}...,=a_{n-1}=a_{n}$, then $L_{f}=L_{b}$.

(2) If $A=I$, then $L_{f}=\frac{1}{n}L_{b}$.

Now we turn to our analysis. Given a minibatch partition $[S_{1},S_{2},...,S_{K}]$ such that $S_{1}\cup S_{2}\cup...\cup S_{K}=[n]$ and:

while $k\in[K]$. In order to identify the potential of a certain optimization problem to be more efficiently solved using stochastic gradient methods, we start by deriving a motivating theorem comparing the convergence of the optimal full gradient methods as well as the optimal stochastic gradient methods.

We consider comparing two classes of algorithm: the optimal deterministic gradient methods which meet the deterministic gradient-complexity lower bound [46, Theorem 2.16] presented in Theorem A.1 and the optimal stochastic gradient methods which are able to match the stochastic gradient-complexity lower bound [66, Theorem 7] presented in Theorem A.2. The FISTA algorithm and the Katyusha algorithm are typical instances from these two classes of algorithms.

It is known that the FISTA algorithm satisfies this definition with $C_{1}=4$ . We also define the class for optimal stochastic gradient methods:

for some positive constants $C_{2}$ and $C_{3}$.

Note that the accelerated stochastic variance-reduced gradient methods such as Katyusha , MiG and Point-SAGA satisfy this definition with different constants of $C_{2}$ and $C_{3}$.

Now we are ready to present the motivational theorem, which follows from simply combining the existing convergence results of the lower bounds for the stochastic and deterministic first-order optimization .

for some positive constant $c_{0}$ which does not depend on $L_{b}$, $L_{f}$ and $K$.

2 An in-depth analysis of minibatch SGD for linear inverse problems

In the previous subsection, we have provided a preliminary motivational analysis, which demonstrates that the speedup of stochastic gradient methods (with data-partition minibatches) over their deterministic counterparts in terms of objective gap convergence are at the worst case controlled by the ratio of Lipschitz constants of the stochastic gradient and full gradient, for the case of unregularized smooth optimization. Such analysis, although motivational, is restrictive in some aspects: in imaging inverse problems we usually consider non-smooth regularization, and we are more concerned with the convergence rates of optimization algorithms regarding estimation error. In this subsection, for the case where the linear measurements are noiseless (i.e. $\|w\|_{2}=0$), we provide a novel convergence rate analysis of minibatch SGD on solving constrained least-squares, which is a subclass of the regularized least-squares (4). By comparing our rate of minibatch SGD with the best known result on deterministic proximal (projected) gradient descent (PGD) in the same setting, we confirm that the ratio of the Lipschitz constants of stochastic gradient and full gradient is indeed the key to characterize the practicality of stochastic optimization for a given inverse problem.

The constrained least-squares objective is written as the following:

where the constraint set $\mathcal{K}$ is enforced as regularization, and the indicator function is used as the regularization to utilize the prior knowledge for better estimation:

One typical example would be the total-variation (TV) semi-norm constraint in imaging applications such as inpainting and deblurring , using an efficient TV-projection operator such as the one developed by .

We restrict ourselves to make the convergence rate comparison of minibatch SGD and deterministic PGD on constrained least-squares mainly due to the fact that the restricted strong-convexity , which is essential for showing estimation-error convergence of the iterates, when applicable, is valid globally in this case since all descent directions are restricted within a tangent cone of the constraint set, as we will see. While for generic regularizers, such necessary restricted strong-convexity condition can only be valid locally . Such an issue will make the desired accurate convergence rate comparison on the estimation error hopeless under the currently known framework for analyzing first-order methods, unless strong extra assumptions are made.

In this subsection, we also study the stochastic acceleration in the case where random with-replacement sampling is used instead of partitioning. Both random with-replacement minibatch scheme and the data-partitioning minibatch scheme are standard choices for stochastic gradient methods. Our analysis for minibatch SGD cover both data-partition minibatch schemes (Def. 3.2) and random with-replacement minibatch schemes (Def. 3.3).

Unlike the analysis of data-partitioning sampling, a major difficulty for the analysis of random with-replacement sampling is that, the step-size choices suggested by the existing convergence results of minibatch proximal stochastic gradient methods can be highly suboptimal, which lead to conservative convergence rate guarantees. Fortunately, there has been recent progress identifying near optimal step-size choices for minibatch stochastic gradient descent and SAGA algorithms for minimizing strongly convex and smooth objective functions. However, these existing results cannot be directly applied in inverse problems, mainly due to the following reasons:

Firstly, these results are only for smooth optimization, while we often use non-smooth regularization in inverse problems, such as sparsity-inducing norms. It is unclear whether such large step-size choices are still allowed in the proximal setting.

Secondly, these results require the objective function to be strongly-convex, which is not satisfied in general for inverse problems.

Due to these obstacles, the first step we should take is to extend the analysis of to linear inverse problems with non-smooth regularization.

We denote $\mathcal{P}_{\mathcal{K}}$ as the orthogonal projection on to the set $\mathcal{K}$ and $\eta$ denotes the step size. We write down the minibatch SGD algorithm using with-replacement random sampling as the solver for (17):

In contrast to (12), we introduce the notion of expected smoothness proposed by , which will be invloved in our analysis

Let $\mathcal{D}$ be the distribution where the random subsampling operator $S$ is drawn from, and we denote $S_{o}$ as the corresponding index set subselected by $S$, the expected smoothness of the minibatches is defined as:

If we use a data-partition minibatch, we have $L_{e}\leq L_{b}$ in (12), as shown in .

2.2 Preliminaries for the analysis of minibatch SGD

We next provide a general theoretical framework for the analysis of minibatch SGD with the restricted strong convexity .

Cone $\mathcal{C}$ is the smallest cone at point $x^{\dagger}$ which contains the shifted set $\mathcal{K}-x^{\dagger}$:

The restricted strong-convexity constant $\mathcal{\mu}_{c}$ is the largest non-negative value which satisfies:

If the measurement system is noiseless, i.e. $\|w\|_{2}=0$, we expect the estimator (17) to be exact: $x^{\star}=x^{\dagger}$, if not, we expect the estimator to be robust to noise: $\|x^{\star}-x^{\dagger}\|_{2}\leq\frac{2\|w\|_{2}}{\epsilon}$. The success of exact/robust estimation is completely dependent on the null-space property of $A$ and the tangent cone $\mathcal{C}\in\mathcal{K}-x^{\dagger}$ on $x^{\dagger}$. In short, for the first scenario the necessary condition for exact recovery is $\|Az^{\prime}\|_{2}>0$ for any normalized vector $z^{\prime}\in\mathcal{C}$ [17, Proposition 2.1]; for the second scenario the necessary condition for robust recovery is $\|Az\|_{2}\geq\epsilon\|z\|_{2}$ for any $z\in\mathcal{C}$ [17, Proposition 2.2]. This relationship between the null space property of $A$ and the constraint on $x^{\dagger}$ is fully captured by the restricted strong convexity property. If the restricted strong convexity condition is valid for (17), we know that $x^{\star}$ provides reliable and robust estimation for $x^{\dagger}$.

2.3 Convergence result of minibatch SGD

Using the expected smoothness result and restricted strong-convexity condition, we are able to derive the following convergence rate for minibatch proximal SGD under uniform random with-replacement minibatch scheme. If we set the step size of the minibatch SGD to be $\frac{1}{L_{e}}$, we can have the following linear convergence result up to a statistical accuracy:

Suppose that $\|w\|_{2}=0$, let the step size of the minibatch SGD algorithm $\eta=\frac{1}{L_{e}}$. The expected error of the update by the $i$-th iteration obeys:

We include the proof of this convergence theorem in Appendix C. The convergence result of the deterministic projected gradient descent under restricted strong-convexity is well-studied in the literature, and we present it here for comparison, while the proof is simple, following a similar procedure to that in, e.g. :

Suppose that $\|w\|_{2}=0$, let the step-size of projected gradient descent algorithm be $\eta=\frac{1}{L_{f}}$, the estimation error of the update by the $i$-th iteration obeys:

We can compare our convergence result of minibatch SGD in Theorem 4.7 with the result for deterministic projected gradient descent in Theorem 4.8. To guarantee an estimation accuracy $\|x^{N}-x^{\dagger}\|_{2}\leq\varepsilon$, the deterministic proximal gradient descent needs:

iterations, while the minibatch SGD needs:

Hence the iteration complexity of minibatch SGD is $O(\Upsilon_{e})$-times smaller than the projected gradient descent, where:

For the data-partition minibatch scheme where we partition the least-squares loss function in to $K=\frac{n}{m}$ minibatches, we know that $L_{e}\leq L_{b}$, as shown in [25, proposition 3.7]. Hence we have:

which demonstrates that for data-partition minibatch scheme, the acceleration of minibatch SGD can offer over its deterministic counterpart, is dominated by the ratio of the Lipschitz constants of the full gradient and minibatch stochastic gradient.

3 Evaluating the limitation of SGD-type algorithms

We introduce a metric called the Stochastic Acceleration (SA) factor based on our theoretical analysis of minibatch SGD in the previous section. The curve for SA factor as a function of the minibatch number $K$ (for a given minibatch pattern) is able to provide a way of characterizing inherently whether for a given inverse problem and a certain partition minibatch sampling scheme, randomized gradient methods should be preferred over the deterministic full gradient methods or not.

For a given disjoint partition minibatch index $[n]=S_{1}\cup S_{2}\cup...\cup S_{K}:=\bar{S}$ where $S_{i}\cap S_{j}=\emptyset,\forall i\neq j\in[K]$, with corresponding subsampling operators $[S^{1},...S^{K}]$, the Stochastic Acceleration (SA) factor is defined as:

We next evaluate the SA factors for the least squares loss function $f(x)=\frac{1}{2n}\|Ax-b\|_{2}^{2}$ with different types of forward operator. We use the without-replacement sampling (data-partitioning) which are most applied in practice. In this case we have

In Figure 3, we demonstrate the SA factors for these 5 problem instances as a function of the number of minibatches along with the empirical acceleration observed when solving these problems. From the result demonstrated in the Figure 3 we find that indeed the stochastic methods have a limitation on some optimization problems like deblurring and inverse problems with random matrices, where we see that the curve for the SA factor of such problems stays low and flat even when we increase the number of minibatches. For the machine learning datasets and X-ray CT imaging, the SA factor increases rapidly and almost linearly as we increase the number of minibatches, which is in line with observations in machine learning on the superiority of SGD and also the observation in CT image reconstruction of the benefits of using the ordered-subset methods which are similar to stochastic gradient methods.

The curves for the SA factor on the first figure qualitatively predict the empirical comparison resultWe compare the objective-gap convergence of FISTA and Katyusha for a fixed number of datapasses (epochs). of the Katyusha and FISTA algorithms shown on the second, where we observe that Katyusha offers no acceleration over the FISTA on either the deblurring or the Gaussian random inverse problem, but significantly outperforms FISTA on the other cases. Indeed, positive results for applying SGD-type algorithms on these problems are well-known already . Hence we have shown that the SA factor we propose is useful in characterizing whether an inverse problem is inherently a suitable candidate for stochastic gradient methods for a given partition.

4 Local coherence structure, eigenspectrum, and stochastic acceleration

From the results we have obtained so far, we now go deeper to investigate the relationship of the SA factor and the structure of the forward operator of the inverse problem.

Subsequently we will assume that each partition has an equal size $m$ for the simplicity of presentation. We will find the following definition of the local-accumulated-coherence to be useful.

Give a partition $[n]=S_{1}\cup S_{2}\cup...\cup S_{K}:=\bar{S}$ for $A=[a_{1};a_{2};...;a_{n}]$ which

Our definition of the local cumulative coherence is similar to and more general than the one presented in and related works, but we do not require the rows to be normalized and the summation includes the term $|\langle a_{j},a_{j}\rangle|$. The local-accumulated-coherence captures the the correlation characteristic between the linear measurements within each partitioned minibatches. As we will see, if a partition have a smaller local accumulated coherence than another partition, then it typically can have a better SA factor. Our analysis in this subsection is based on the Gersgorin disk theorem:

The Gersgorin disk theorem relates a square symmetric matrix’s eigenvalues with its entries, and links to the gradient-Lipschitz constant $L_{b}$ – which in the least-squares context can be written as:

The SA factor for any linear inverse problem with $f(x)=\frac{1}{2n}\|Ax-b\|_{2}^{2}$ is lower bounded as:

We have derived partition independent lower bounds for $\Upsilon(A,\bar{S},K)$. However, these lower bounds, by definition have to cover the worst case of partition. Hence for some inverse problems which admit inferior partitions, these may be crude lower bounds. It is therefore insightful to derive a lower bound for the case where we randomly partition the data. We provide the following lower bound using the Matrix Chernoff inequality and the union bound, following a similar argument by [40, Proposition 3.3]. We present the proof in Appendix G.

If $\bar{S}$ is a uniform random partition, then for $K\in\left[\frac{\|A\|^{2}}{\|A^{T}\|_{1\rightarrow 2}^{2}},\min(n,d)\right]$, the following lower bounds hold:

with probability at least: $1-d^{2}\left(\frac{e}{\delta}\right)^{\delta}$.

This lower bound for random partition scheme again demonstrates the strong relationship between SA factor and the ratio $\frac{\|A^{T}\|_{1\rightarrow 2}^{2}}{\|A\|^{2}}$ which is controlled by the eigenspectrum $\frac{\sum_{i=1}^{d}\sigma(A^{T}A,i)}{\sigma(A^{T}A,1)}$. Note that due to the Matrix Chernoff inequality , this theorem holds with a probability $1-d^{2}\left(\frac{e}{\delta}\right)^{\delta}$, which is dimension-dependent, and meanwhile we also need to note that the theorem covers only the regime where $K$ is sufficiently large. For a fixed dimension $d$, the smaller the ratio $\frac{\|A^{T}\|_{1\rightarrow 2}^{2}}{\|A\|^{2}}$, or rather, the faster the eigenspectrum decays, the larger SA factors will be for the number of minibatch $K$ within the range $\left[\frac{\|A\|^{2}}{\|A^{T}\|_{1\rightarrow 2}^{2}},\min(n,d)\right]$.

To be more specific, if we demand here $1-d^{2}\left(\frac{e}{\delta}\right)^{\delta}\geq 0.9$, we can compute for a given $\delta$ the maximum allowed $d$ for the Theorem 4.14 to hold with this probability via the following bound:

We list the values of $\sqrt{\frac{1}{10\left(\frac{e}{\delta}\right)^{\delta}}}$ for a range of $\delta$ in table 1.

Moreover, as we will show, qualitatively this works but using a smaller $\delta$ seems to better describe what we see in the SA factor for random partition. On the other hand, this result requires the number of minibatches $K$ to be sufficiently large ($K\in\left[\frac{\|A\|^{2}}{\|A^{T}\|_{1\rightarrow 2}^{2}},\min(n,d)\right]$), however in our numerical experiments we can observe that for small values of $K$ the lower bound $\alpha_{r}(A,K,\delta)$ still provides a reasonably good estimate. Similar restrictions occur in [40, Proposition 3.3] which is also based on the Matrix Chernoff inequality. Whether such a restriction can be technically removed is an interesting open question.

Meanwhile, we can also have an upper bound for the SA factor, independent of the partition $\bar{S}$, in terms of the eigenspectrum of the Hessian matrix $A^{T}A$. This upper bound can be derived from a standard result [28, Theorem 4.3.15] using the fact that the matrix $(S^{k}A)^{T}S^{k}A$ and $S^{k}A(S^{k}A)^{T}$ share the same non-zero eigenvalues. We denote the $i$-th large eigenvalue of a Hermitian matrix $H$ as $\sigma(H,i)$, and the upper bound is written as the following:

The SA factor for any linear inverse problem with $f(x)=\frac{1}{2n}\|Ax-b\|_{2}^{2}$ is upper bounded as:

We include the proof in Appendix F for completeness. The upper bound (45) suggests that, if the Hessian matrix $A^{T}A$ has slowly-decaying eigenvalues at the tail, it indeed typically cannot have a large SA factor, no-matter how delicately we partition the forward operator $A$. The upper bound and the lower-bounds jointly suggest that, having a fast-decaying eigenspectrum of the Hessian is a sufficient and necessary condition for an inverse problem to have good SA factors.

A key result in this analysis is the lower bound which relates the stochastic acceleration factor with the local coherence/correlation structure of the given partition:

An immediate conclusion we can have is that, for a given linear inverse problem with forward operator $A$, and a partition $\bar{S}$, the smaller the local accumulated coherence is, the larger the SA factor will be. For some inverse problems, judiciously choosing the partition for minibatches is important – good choices of partitioning can have small local coherence and hence lead to larger SA factors in practice.

In Figure 5 we present a simulation result where we generate 4 random matrices of the same size $n=200,d=1000$ with different distributions and check the relationship of their SA factors and the eigenspectrum of their Hessian matrices. The forward operators we generate are:

(1) random Gaussian matrix with each entry drawn from a Gaussian distribution with zero-mean and unit-variance;

(3) random Gaussian matrix with each entry drawn from a Gaussian distribution with 0.25-mean and unit-variance;

(4) random matrix with each entry drawn from a uniform distribution supported on the interval $$.

From the experimental result we can observe that, these four forward operators have very different decay-rates on their Hessians’ eigenspectrum, and correspondingly, very different SA factors. The case (4) has the fastest decay-rate, and has the largest SA factors and it grows almost linearly as the number of minibatches increases. The case (1) has the slowest decay-rate on the eigenspectrum, and correspondingly, it has the worst SA factors among the 4 cases. This numerical result is in broad agreement with our analysis.

We also test our lower bound $\alpha_{r}(A,K,\delta)$ for all the examples we have considered. We first compute the lower bound estimate $\alpha_{r}(A,K,\delta)$ by (42) for different forward operators, with the choice of $\delta=15$ which is sufficient for the lower bound to hold with probability at least $0.9$ for all these forward operators. We present the result in Figure 6(b) and compare it with the SA factors presented in Figure 6(a). We find that our theoretically justified lower bound is still able to distinguish well whether a given inverse problem is suitable or not for stochastic optimization, but seems to be very conservative for the choice of $\delta$. Interestingly, if we heuristically reduce $\delta$ from $15$ to $2$, then we can actually obtain a much better lower bound estimate for the SA factors, as shown in Figure 6(c).

4.2 Concluding remark

In short, the take-home message of our analysis in this subsection contains following aspects:

Firstly, for linear inverse imaging problems, in order to have a good SA factors, the forward operator $A$ should have a small ratio of $\frac{\|A^{T}\|_{1\rightarrow 2}^{2}}{\|A\|^{2}}$, which essentially means that the Hessian matrix $A^{T}A$ should have a relatively fast-decaying eigenspectrum.

Meanwhile, optimizing the choice of partitioning can be crucial in some inverse problems – a judiciously chosen partition scheme can significantly improve the SA factor if the forward operator has local incoherence structure.

The lower bounds, particularly the one for the random partitionIf one wishes to use the relaxed lower bound estimate, the only need is to compute $\frac{\|A\|^{2}}{\|A^{T}\|_{1\rightarrow 2}^{2}}$ which can be very efficiently obtained. (42) can be readily applied by practitioners to conveniently evaluate whether a given inverse problem is suitable for using stochastic gradient methods as solvers. The first lower bound (38) can also be used to compare between different partition schemes and select the best one among themOne may also achieve this goal by directly computing $L_{f}$ and $L_{b}$ with the power method or its accelerated and stochastic variants for each of the compared partition schemes.. However, finding the best partitioning for a given inverse problem is a combinatorial problem, and we do not currently have any generic scheme which is guaranteed to solve this for arbitrary inverse problems – we leave this as an open problem for future work. Note that, in this section we have mainly focused on the data-partition minibatch scheme, while we also find that similar results is also valid for random with-replacement minibatch schemes, and we refer the readers to Appendix D for details.

A practical accelerated stochastic gradient method for imaging

So far we have considered the role of the Lipschitz constants (and associated step-sizes) in determining the potential advantages of stochastic over deterministic gradient methods in inverse problems. However there are other aspects that complicate the analysis. The most obvious one is that stochastic gradient methods formulated in the primal domain need to calculate the proximal operator many more times than full gradient methods and hence slow down dramatically the run time. There are also scenarios, (see e.g. ), where more than one non-smooth regularization term may be desired, where most of the existing fast stochastic methods such as Katyusha are inapplicable. Here we present a new SGD formalism that aims to mitigate this problem.

We consider in this section the following optimization problem:

The most famous algorithm for solving this saddle-point problem is the primal-dual hybrid gradient (PDHG, also known as the Chambolle-Pock algorithm) , which interleaves the update of the primal variable $x$ and the dual variable $y$ throughout the iterates. With this reformulation the linear operator $D$ and the function $g(.)$ are decoupled and hence one can divide-and-conquer the expensive TV-proximal operator with the primal-dual gradient methods. The stochastic variant of the PDHG for the saddle-point problem (48) has been very recently proposed by Zhao $\&$ Cevher [68, Alg.1, \saySPDTCM] and shown to have state-of-the-art performance when compared to PDHG, stochastic ADMM and stochastic proximal averaging . We find out that, the SPDTCM method is an efficient and practical optimization algorithm for imaging problems with regularization terms which are coupled with linear operators. Additionally, since the effect of acceleration given by Nesterov’s momentum appears to be very important, we also need to consider a way to ensure that our method is acceleratedWe are aware of the recent attempt to develop accelerated SPDTCM algorithm. However we find this accelerated method is not as practical as SPDTCM and our method in imaging problems since the step-size choices depend on unknown parameters ($\rho$ and $\rho^{\prime}$ in [69, Section 2]) in unconstrained minimization tasks, which need careful tuning manually for different imaging inverse problems in practice. Sub-optimal choices may lead to compromised convergence rates and even diverging behaviors. . Since the SPDTCM method does not have Nesterov-type acceleration, we propose a variant of it which adopts the outerloop acceleration scheme given by the Katyusha-X algorithm , which is a simplified variant of the Katyusha algorithm . We observe that such a momentum step is important for the stochastic primal-dual methods in this application. We present our method as Algorithm 2. One can directly choose the same step-size sequences $\alpha_{(.)},\eta_{(.)},\theta_{(.)}$ as suggested in [68, Section 2.3].

Numerical experiments

We plot the estimation error $\log_{10}\|x-x^{\dagger}\|_{2}^{2}$ in Figure 7 for each algorithm, where $x^{\dagger}$ denotes the ground truth image. We observe a roughly $4\times$ improvement in run time compared to FISTA since our algorithm can avoid the heavy cost of the TV proximal operatorFor the computation of the TV proximal-operator for FISTA algorithm, we use the popular implementation from the UnLocBox toolbox . while maintaining the fast convergence provided by Nesterov-type momentum and randomization. We also report a significant improvement over the SPDTCM algorithm both in time and iteration complexity. In terms of datapasses (number of epochs), the SPDTCM does not show any advantage over the deterministic method FISTA, while our Acc-PD-SGD with 10 minibatches is able to achieve 2 times acceleration over FISTA.

We also report that in this experiment, if we further increase the number of subsets of SPDTCM and Acc-PD-SGD, we do not observe faster convergence for these algorithms. In other words, no matter how we increase the number of subsets, this 2-time acceleration (in terms of number of datapasses) is the limit of our algorithm – such a trend is successfully predicted by the SA factor shown in the Figure 3, where we can see that the curve of the SA factor for deblurring task goes flat instead of increasing after the number of minibatches $K>10$.

2 X-Ray computed tomography image reconstruction experiment

In the first experiment, we have demonstrated the superior performance of the proposed Acc-PD-SGD algorithm compared to the deterministic algorithm FISTA, and the state-of-the-art stochastic primal-dual gradient method SPDTCM on a space-varying deblurring problem, although it is not inherently favorable for the application of stochastic gradient methods. In this subsection, we turn to another imaging inverse problem – the computed tomography image reconstruction. As suggested by the curve of the SA factor, the X-ray CT image reconstruction is a nice application for stochastic gradient methods, where we expect them to achieve significant speed-ups over the deterministic methods.

From the convergence results of the iterative algorithms in Figure 8 we can observe that for this experiment, the stochastic methods SPDTCM and Acc-PD-SGD converges significantly faster than the full gradient method FISTA both in terms of number of datapasses and wall-clock time. Meanwhile, we can also see that, our proposed method Acc-PD-SGD converges faster than the SPDTCM which does not use Katyusha-X momentum for acceleration.

Conclusion

In this work we began by investigating the practicability of the state-of-the-art stochastic gradient methods in imaging inverse problems. We first presented a surprisingly negative result on existing SGD-type methods on image deblurring, as a motivational example. To understand the limitation of stochastic gradient methods in inverse problems, we have provided a novel analysis for the estimation-error convergence rate of minibatch SGD in the setting of linear inverse problem with constraints as regularization, under restricted strong-convexity and expected smoothness conditions. Based on our theoretical analysis, we have proposed the SA factor to evaluate the possible computational advantage of using stochastic techniques for a given task. Then we went further and made an in-depth analysis for understanding how the inherent structure of the forward measurement model can contribute to the practicality of the stochastic gradient methods with partition minibatch schemes.

We also derived lower and upper bounds of the SA factor. From the theoretical results, we find out that, if an linear inverse problem has a small ratio of $\frac{\|A^{T}\|_{1\rightarrow 2}^{2}}{\|A\|^{2}}$, which means the Hessian matrix $A^{T}A$ has fast-decaying eigenspectrum, then it typically admit good SA factors, hence can be rapidly solved by stochastic gradient methods. Our analysis also suggests that, excellent partition schemes typically have low local-accumulated-coherence, which essentially means the measurements within one minibatch are mutually less correlated. Using our SA factor, jointly with the derived lower bounds, the practitioners can easily identify whether they should use stochastic gradient or deterministic gradient algorithms for given inverse problems, and evaluate the potential of given partition schemes.

While our results are mainly for linear inverse problems with least-squares data-fidelity terms and convex regularizers, we believe that they also can be extended and give insights to non-linear inverse problems since one can construct majorizing linearized subproblems (proximal Newton-steps) and solve these subproblems with deterministic or stochastic proximal gradient methods. Our results can also be extended for understanding and analyzing the limitations of stochastic gradient-based methods with the plug-and-play priors and regularization-by-denoising schemes in imaging inverse problems, which we leave as a future direction.

Although we have concentrated on stochastic gradient methods vs deterministic gradient methods, there are other considerations that might affect the choice of whether to go stochastic. For example, if an inverse problem can be effectively preconditioned by simple preconditioners (such as diagonal preconditioners), then the potential benefit of stochastic methods over deterministic methods may possibly be reduced, since the preconditioned forward operator may not have as fast-decaying spectrum as the original one. Moreover, if the forward operator can be implemented with a fast transform such as the FFT, for example in MRI image reconstruction tasks, the deterministic gradient methods are usually much more favored since they can benefit from the fast operation while current stochastic gradient methods cannot.

Finally, as a side contribution, we propose the Accelerated Primal-Dual SGD to cope with multiple regularizers (potentially) with a linear operator while maintaining the fast convergence, and demonstrate its effectiveness via experiments on space-varying deblurring and X-Ray CT image reconstruction. Although we have not yet done the theoretical convergence analysis of the Acc-PD-SGD algorithm, we believe it provides insights for the algorithmic design of fast stochastic gradient methods tailored specifically for imaging inverse problems, from understanding the inherent limitation, to the practical algorithmic framework.

Appendix A Lower bounds for deterministic gradient and stochastic gradient optimization

We present some well-known lower bounds for first-order optimization. We start by the lower-bound derived by Nesterov [46, Theorem 2.16]:

Such a lower-bound suggests that there exists at least one $L$-smooth convex function on which any first-order method cannot converge faster than $O(1/k^{2})$ for a limited number of iterations which $1\leq k\leq\frac{1}{2}(d-1)$.

For the stochastic gradient-based optimization, several researchers have derived important lower-bounds for optimizing the finite-sum objective with stochastic gradient oracle $\triangledown f_{i}(.)$, and we present here a typical well-known result:

number of stochastic gradient evaluations are needed.

Appendix B The Proof of Theorem 4.3

The proof of this theorem is straight forward and is based on combining the existing results since we have assumed the dimension $d$ is large enough for the lower-bound for stochastic gradient oracles to hold on a domain:

calls of the stochastic gradient oracle $\triangledown f_{i}()$. In other words, for this worst case function, if we run any stochastic gradient method with only $Ks$ calls on the stochastic gradient oracle such that:

where the constant $C_{1}$ is independent of $L_{b}$, $L_{f}$ and $K$. Combining these two bounds we can have:

Appendix C The proof of Theorem 4.7: convergence of minibatch SGD on constrained Least-squares

Due to the definition of expected smoothness (21), we have:

Now set $y=x^{\dagger}$, and since $\langle\triangledown f(x^{\dagger}),x-x^{\dagger}\rangle=0$ at the noiseless case, we have:

Note that by definition $b=Ax^{\dagger}+w$, and since it is assumed here $\|w\|_{2}=0$, we have:

Now we present the complete proof of Theorem 4.7:

For iteration index $i$ of minibatch SGD we have the following:

Line (e) uses the Jensen’s inequality, line (f) is due to Lemma C.1, while inequality (g) holds because of the restricted strong-convexity condition (Def. 4.6), and we choose $\eta\leq\frac{2}{L_{e}}$. Then because the subsampling in each iteration is independent from the previous error vector, by the tower rule we yield:

Thus we finish the proof by choosing $\eta=\frac{1}{L_{e}}$.

Appendix D Estimating the stochastic acceleration for random with-replacement sampling schemes

The notion of stochastic acceleration factor can also be extended to the case where we use random with-replacement sampling scheme. For random with-replacement minibatch scheme, [25, Proposition 3.8] shows that the expected smoothness constant $L_{e}$ in (21) can be upper bounded by:

Recall our Remark 4.9 comparing deterministic gradient descent and minibatch SGD on constrained least-squares. Denote here $L_{f}=\frac{1}{n}\|A^{T}A\|_{2}$. If $\|w\|=0$ and, to guarantee an estimation accuracy $\|x_{K}-x^{\dagger}\|_{2}\leq\varepsilon$, the deterministic proximal gradient descent needs:

iterations, while the minibatch SGD needs:

We name $\Upsilon_{e}(A,m)$ as the expected SA factor for the random with-replacement sampling scheme.

From the expression of $\Upsilon_{e}(A,m)$ we can find that the key factor which influences the advantage of stochastic gradient over determinstic gradient is again the ratio $\frac{\|A^{T}\|_{1\rightarrow 2}^{2}}{\|A\|^{2}}$ which also occurs in our lower bounds for data-partition minibatch schemes in Theorem 4.13 and 4.14.

Our results in this section regarding random with-replacement sampling have several restrictions. So far we can establish results only for minibatch SGD algorithm and meanwhile $\Upsilon_{e}(A,m)$ is derived based on comparing this result with the iteration complexity of proximal gradient descent. If we use $\Upsilon_{e}(A,m)$ to measure whether an inverse problem is suitable for stochastic gradient algorithms using a random with-replacement minibatch scheme, we are implicitly making the conjecture that all these minibatch proximal stochastic gradient methods including variance-reduced methods such as SVRG, SAGA and Katyusha, admit the large step-size choices $O(1/L_{e})$ based on the expected smoothness constant $L_{e}$, which has not been shown yet in the literature.

In Figure 12, we plot the expected SA factor $\Upsilon_{e}(A,m)$ for a range of inverse problems which we have considered so far. Similar to the results shown in Figure 3, we observe that, the space-varying deblurring example and the zero-mean random Guassian example have the worst expected SA factors.

Appendix E The proof of Theorem 4.13

Now we present the proof of Theorem 4.13.

If we set $H=S^{k}A(S^{k}A)^{T}$ for some $k\in[K]$, the top eigenvalue of $S^{k}A(S^{k}A)^{T}$ is no larger than the largest value within the set $G(S^{k}A(S^{k}A)^{T})$ which we denote here as $G_{\max}(S^{k}A(S^{k}A)^{T})$. We have the following relationship:

By definition of the SA factor, we can write:

and hence we can have a relaxed lower bound for $\Upsilon(A,\bar{S},K)$:

Suppose that for some positive constant $\rho$ we have:

and hence we can further lower bound $\Upsilon(A,\bar{S},K)$ by the cumulative eigenspectrum of the Hessian:

Appendix F The proof of Theorem 4.16

We present the proof of Theorem 4.16 here.

If we set $H_{m}=S^{k}A(S^{k}A)^{T}$, then we have:

Now we use the fact that $S^{k}A(S^{k}A)^{T}$ and $(S^{k}A)^{T}S^{k}A$ share the same non-zero eigenvalues, and meanwhile $AA^{T}$ and $A^{T}A$ also shares the same non-zero eigenvalues, we can have the following bound:

Then by the definition of $\Upsilon(A,\bar{S},K)$ we can obtain the upper bound.

Appendix G The proof of Theorem 4.14

We now present the proof of Theorem 4.14:

Suppose we randomly permute the index $[n]$ and generate the partition index $[S_{1},S_{2},...,S_{K}]$. If we pick arbitrarily a number $k\in[K]$ where $S^{k}$ is the subsampling matrix, by the Matrix Chernoff inequality we have the following relationship:

for any $\delta_{0}>0$. Now by choosing $\delta_{0}=\delta\cdot\frac{K\|A^{T}\|_{1\rightarrow 2}}{\|A\|^{2}}$ we can have following:

with probability at least $P^{\prime}$ where:

(This is because we restrict here $K\geq\frac{\|A\|^{2}}{\|A^{T}\|_{1\rightarrow 2}^{2}}$.) Now by applying the union bound over all possible choices of $k$ and since we assume here $K\leq\min(n,d)$, we have, with probability at least $1-d^{2}\cdot\left[\frac{e^{\delta}}{\delta^{\delta}}\right]$:

Acknowledgment

We acknowledge the support from H2020-MSCA-ITN 642685 (MacSeNet) , ERC Advanced grant 694888, C-SENSE and a Royal Society Wolfson Research Merit Award. We thank Alessandro Foi, Vladimir Katkovnik, Cristovao Cruz, Enrique Sanchez-Monge, Zhongwei Xu, Jingwei Liang, Derek Driggs, Alessandro Perelli, Mikey Sheehan and Jonathan Mason for helpful discussions.

References