Stochastic Quasi-Gradient Methods: Variance Reduction via Jacobian Sketching

Introduction

We consider the problem of minimizing the average of a large number of differentiable functions

where $f$ is $\mu$–strongly convex and $L$–smooth. In solving (1), we restrict our attention to first-order methods that use a (variance-reduced) stochastic estimate of the gradient $g^{k}\approx\nabla f(x^{k})$ to take a step towards minimizing (1) by iterating

In the context of machine learning, (1) is an abstraction of the empirical risk minimization problem; $x$ encodes the parameters/features of a (statistical) model, and $f_{i}$ is the loss of example/data point $i$ incurred by model $x$. The goal is to find the model $x$ which minimizes the average loss on the $n$ observations.

Typically, $n$ is so large that algorithms which rely on scanning through all $n$ functions in each iteration are too costly. The need for incremental methods for the training phase of machine learning models has revived the interest in the stochastic gradient descent (SGD) method . SGD sets $g^{k}=\nabla f_{i}(x^{k})$, where $i$ is an index chosen from $[n]\overset{\text{def}}{=}\{1,2,\dots,n\}$ uniformly at random. SGD therefore requires only a single data sample to complete a step and make progress towards the solution. Thus SGD scales well in the number of data samples, which is important in several machine learning applications since there many be a large number of data samples. On the downside, the variance of the stochastic estimates of the gradient produced by SGD does not vanish during the iterative process, which suggests that a decreasing stepsize regime needs to be put into place if SGD is to converge. Furthermore, for SGD to work efficiently, this decreasing stepsize regime needs to be tuned for each application area, which is costly.

Stochastic variance-reduced versions of SGD offer a solution to this high variance issue, which improves the theoretical convergence rate and solves the issue with ad hoc stepsize regimes. The first variance reduced method for empirical risk minimization is the stochastic average gradient (SAG) method of Schmidt, Le Roux and Bach . The analysis of SAG is notoriously difficult, which is perhaps due to the estimator of gradient being biased. Soon afterwards, the SAG gradient estimator was modified into an unbiased one, which resulted in the SAGA method . SAGA maintains a matrix of the latest gradients computed for each datapoint $i$, and uses this matrix to construct a stochastic estimate of the gradient. The analysis of SAGA is dramatically simpler than that of SAG. Another popular method is SVRG of Johnson and Zhang (see also S2GD ). SVRG enjoys the same theoretical complexity bound as SAGA, but has a much smaller memory footprint. It is based on an inner-outer loop procedure. In the outer loop, a full pass over data is performed to compute the gradient of $f$ at the current point. In the inner loop, this gradient is modified with the use of cheap stochastic gradients, and steps are taken in the direction of the modified gradients. A notable recent addition to the family of variance reduced methods, developed by Nguyen et al , is known as SARAH. Unlike other methods, SARAH does not use an estimator that is unbiased in the last step. Instead, it is unbiased over a long history of the method.

A fundamentally different way of designing variance reduced methods is to use coordinate descent to solve the dual. This is what the SDCA method and its various extensions do. The key advantage of this approach is that the dual often has a seperable structure in the coordinate space, which in turn means that each iteration of coordinate descent is cheap. Furthermore, SDCA is a variance-reduced method by design since the coordinates of the gradient tend to zero as one approaches the solution. One of the downsides of SDCA is that it requires calculating Fenchel duals and their derivatives. This issue was later solved by introducing approximations and mapping the dual iterates to the primal space as pointed out in . This resulted in primal variants of SDCA such as dual-free SDCA . A primal-dual variant which enables the use of arbitrary minibatch strategies was developped by Qu et al , and is known as QUARTZ.

2 Gaps in our understanding of SAGA

Despite significant research into variance-reduced stochastic gradient descent methods for solving (1), there are still big gaps in our understanding of variance reduction. For instance, the current theory supporting the SAGA algorithm is far from complete.

SAGA with uniform probabilities enjoys the iteration complexity ${\cal O}((n+\tfrac{L_{\max}}{\mu})\log\tfrac{1}{\epsilon})$, where $L_{\max}\overset{\text{def}}{=}\max_{i}L_{i}$ and $L_{i}$ is the smoothness constant of $f_{i}$. While importance sampling versions of SAGA have proved in practice to produce a speed-up over uniform SAGA , a proof of this speed-up has been elusive. It was conjectured by Schmidt et al. that a properly designed importance sampling strategy for SAGA should lead to the rate ${\cal O}((n+\tfrac{\bar{L}}{\mu})\log\tfrac{1}{\epsilon})$, where $\bar{L}=\tfrac{1}{n}\sum_{i}L_{i}$. However, no such result was proved. This rate is achieved by, for instance, importance sampling variants of SDCA and QUARTZ . However, the analysis only applies to a more specialized version of problem (1) (e.g., one needs an explicit strongly convex regularizer). Second, existing minibatch variants of SAGA do not enjoy the same rate as that offered by methods such as SDCA and QUARTZ. Are the above issues with SAGA unavoidable, or is it the case that our understanding of the method is far from complete? Lastly, no minibatch variant of SAGA with importance sampling is known.

One of the contributions of this paper is giving positive answers to all of the above questions.

3 Jacobian sketching: a new approach to variance reduction

Our key contribution in this paper is the introduction of a novel approach—which we call Jacobian sketching—to designing and understanding variance-reduced stochastic gradient descent methods for solving (1). We refer to our method by the name JacSketch. We shall now briefly introduce some of the key insights motivating our our approach.

The starting point of our new approach is the following trivial observation: the gradient of $f$ at $x$ can be computed from the Jacobian ${\bf\nabla F}(x)$ by a simple linear transformation:

which is a stochastic gradient of $f$ at $x$. In other words, by performing a random linear transformation of the Jacobian, we have arrived at the classical stochastic estimate of the gradient. This approach does not suffer from the first issue mentioned above as the Jacobian is not needed at all in order to compute $\nabla f_{i}(x)$. Likewise, it does not suffer from the second issue; namely, the cost of computing the stochastic gradient is merely ${\cal O}(d)$, and we can avoid a costly pass through the data.For the purposes of this narrative it suffices to assume that stochastic gradients can be sampled at cost ${\cal O}(d)$.

This equation generalizes both (5) and (6). The left hand side contains the sketched system matrix $\mathbf{S}_{k}$ and the unknown matrix ${\bf J}$, and the right hand side contains a quantity we can measure (through a random linear measurement of the Jacobian, which we assume is cheap). Of course, the true Jacobian solves (8). However, in general, and in particular when $\tau\ll n$ which is the regime we want to be in for practical reasons, the system (8) will have infinite ${\bf J}$ solutions.

In doing so, we have built a learning mechanism whose goal is to maintain good estimates of the Jacobian throughout the run of method (2). These estimates can be used to efficiently estimate the gradient by performing a linear transformation similar to (5), but with ${\bf\nabla F}(x)$ replaced by the latest estimate of the Jacobian. In practice, it is important to design sketching matrices so that the Jacobian sketch ${\bf\nabla F}(x)\mathbf{S}_{k}$ can be calculated efficiently.

The “sketch-and-project” strategy (1.3) for updating our Jacobian estimate is analogous to the way quasi-Newton methods update the estimate of the Hessian (or inverse Hessian) . From this perspective, our method can be viewed as a stochastic quasi-gradient method.The term “quasi-gradient methods” was popular in the 1980s , and refers to algorithms for solving certain stochastic optimization problems which rely on stochastic estimates of function values and their derivatives. In this paper we give the term a different meaning by drawing a direct link with quasi-Newton methods.

Problem (1.3) admits the explicit closed-form solution (see Lemma B.1):

is a projection matrix, and $\dagger$ denotes the Moore-Penrose pseudoinverse.

The key insight of our work is to propose an efficient Jacobian learning mechanism based on ideas borrowed from recent results in randomized numerical linear algebra.

Having established our update of the Jacobian estimate, we now need to use this to form an estimate of the gradient. Unfortunately, using ${\bf J}^{k+1}$ in place of ${\bf\nabla F}(x^{k})$ in (5) leads to a biased gradient estimate (something we explore later in Section 2.5). To obtain an unbiased estimator of the gradient, we introduce a stochastic relaxation parameter $\theta_{\mathbf{S}_{k}}$ and use

This strategy indeed works, as we show in detail in this paper. Under appropriate conditions (on the stepsize $\alpha$, properties of $f$ and randomness behind the sketch matrices $\mathbf{S}_{k}$ and so on), the variance of $g^{k}$ diminishes to zero (e.g., see Lemma 3.10), which means that JacSketch is a variance-reduced method. We perform an analysis for smooth and strongly convex functions $f$, and obtain a linear convergence result (Theorem 3.6). We summarize our complexity results in detail in Section 1.7.

4 SAGA as a special case of JacSketch

Of particular importance in this paper are minibatch sketches, which are sketches of the form $\mathbf{S}_{k}={\bf I}_{S_{k}}$, where $S_{k}$ is a random subset of $[n]$, and ${\bf I}_{S_{k}}$ is a random column submatrix of the $n\times n$ identity matrix with columns indexed by $S_{k}$. For minibatch sketches, JacSketch corresponds to minibatch variants of SAGA. Indeed, in this case, and if ${\bf W}={\rm Diag}(w_{1},\dots,w_{n})$, we have ${\bf\Pi}_{\mathbf{S}_{k}}e=e_{S_{k}}$, where $e_{S}=\sum_{i\in S}e_{i}$ (see Lemma 4.7). Therefore,

In view of (11), and since ${\bf\Pi}_{\mathbf{S}_{k}}={\bf I}_{S_{k}}{\bf I}_{S_{k}}^{\top}$ (see Lemma 4.7), the Jacobian estimate gets updated as follows

Standard uniform SAGA is obtained by setting $S_{k}=\{i\}$ with probability $1/n$ for each $i\in[n]$, and letting $\theta_{\mathbf{S}_{k}}\equiv n$. SAGA with arbitrary probabilities is obtained by instead choosing $S_{k}=\{i\}$ with probability $p_{i}>0$ for each $i\in[n]$, and letting $\theta_{\mathbf{S}_{k}}\equiv\tfrac{1}{p_{i}}$. However, virtually all minibatching and importance sampling strategies can be treated as special cases of our general approach.

The theory we develop answers the open questions raised earlier. In particular, we answer the conjecture of Schmidt et al. about the rate of SAGA with importance sampling in the affirmative. In particular, we establish the iteration complexity $(n+\frac{4\bar{L}}{\mu})\log\tfrac{1}{\epsilon}.$ This complexity is obtained for different importance sampling distributions than that currently proposed in the literature for SAGA. In order to achieve this, we develop a new analysis technique which makes use of a stochastic Lyapunov function (see Section 5). That is, our Lyapunov function has a random element which is independent of the randomness inherited from the iterates of the method. This is unlike any other Lyapunov function used in the analysis of stochastic methods we are aware of. Further, we prove that SAGA converges with any initial matrix ${\bf J}^{0}$ in place of the matrix of gradients of functions $f_{i}$ at the starting point. We also show that our results give better rates for minibatch SAGA than are currently known, even for uniform minibatch strategies. We also allow for a family of completely new uniform minibatching strategies which were not considered in connection with SAGA before, and consider also SAGA with importance sampling for minibatchesFor some prior results on importance sampling for minibatches, in the context of QUARTZ, see . (based on a partition of $[n]$). Lastly, as a special case, our method recovers standard gradient descent, together with the sharp iteration complexity of $\frac{4L}{\mu}\log\tfrac{1}{\epsilon}$.

Our general approach also enables a novel reduced memory variant of SAGA as a special case. Let $\mathbf{S}_{k}=e_{S_{k}}$, and choose ${\bf W}={\bf I}.$ Since ${\bf\Pi}_{\mathbf{S}_{k}}e=e_{S_{k}}$, the formula for $g^{k}$ is the same as in the case of SAGA, and is given by (16). What is notably different about this sketch (compared to ${\bf I}_{S_{k}}$) is that, since ${\bf\Pi}_{{\bf I}_{S_{k}}}=\frac{1}{|S_{k}|}e_{S_{k}}e_{S_{k}}^{\top},$ the update of the Jacobian estimate (39) is given by

Thus, the same update is applied to all the columns of ${\bf J}^{k}$ that belong to $S_{k}$. Equivalently, this update can be written as

In particular, if $S_{k}$ only ever picks sets which correspond to a partition of $[n]$, and we initialize ${\bf J}^{0}$ so that all the columns belonging to the same partition are the same, then they will be the same within in each partition for all $k$. In such a case, we do not need to maintain all the identical copies. Instead, we can update and use a condensed/compressed version of the Jacobian, with one column per partition set only, to reduce the total memory usage. This method, with non-uniform probabilities, is analyzed in our framework in Section 5.6.

5 Sketch and project

It has long been known, and was explored in detail by Needell, Srebro and Ward , that the randomized Kaczmarz method is a specific instantiation of SGD, applied to a suitable least-squares type function. In the context of sketch and project methods with arbitrary sketching matrices $\mathbf{S}_{k}$, this was explored by Richtárik and Takáč , who also demonstrated that the sketch and project method, and hence also our Jacobian learning iteration (1.3), can be interpreted as stochastic gradient descent applied to a suitable stochastic optimization problem. Therefore, and quite surprisingly:

In our Jacobian sketching framework, variance reduction is obtained by applying SGD to the problem of learning the Jacobian. So, our method uses SGD in two different ways: as a method for performing the step toward minimizing the loss (this is standard), and as a method for learning the Jacobian which is then used to lower the variance of the search direction (this is our new insight).

As a follow up to , Gower and Richtárik further extended their analysis in to arbitrary consistent linear systems (i.e., beyond systems with a single solution, such as (7)). Therein they show that the sketch and project method converges linearly to the projection of the starting iterate onto the solution space of the system, and also uncover a dual interpretation of the method as stochastic dual subspace ascent. Related ideas were later used to design stochastic algorithms for inverting matrices and computing the pseudoinverse of a rectangular matrix . For a compendium of some of the above papers on sketch and project, see also .

An accelerated (in the sense of Nesterov) sketch and project method was proposed and analyzed in . However, the analysis was restricted to a weak type of convergence. This was remedied by Tu et al. for positive definite systems and a special class of sketchings, by Richtárik and Takáč for general linear systems and general sketchings, and further extended to Euclidean setting and applied to matrix inversion and quasi-Newton updates by Gower et al. . A sketch and project method with the heavy ball momentum was studied in .

6 Controlled stochastic reformulation

Loosely motivated by , we shall explore an alternative narrative to the sketch-and-project motivation described above. In particular, the development of JacSketch can instead be motivated through the lens of controlled stochastic reformulations of (1).

Let us now very briefly outline the main idea. First, we will use the distribution ${\cal D}$ from which the sketching matrices are drawn to define a stochastic optimization reformulation of problem (1). That is, we write $f$ as an expectation over some carefully constructed functions $f_{\mathbf{S}}(x)$ instead, where the expectation is taken over $\mathbf{S}\sim{\cal D}$. We then add a “smart” zero function, also of the form of an expectation of some functions over ${\cal D}$, to this reformulation. However, this zero perturbation depends on ${\bf J}^{k}$. While this does not change the objective function, it affects the stochastic gradients in a positive way: it reduces their variance. We then apply an SGD step to this perturbed (or “controlled”) reformulation, followed by an update of the Jacobian (through sketch and project). This is iterated until convergence, and results in JacSketch. This alternative narrative is provided in Section 2.

7 Summary of complexity results

All convergence results obtained in this paper are summarized in Table 1.

Theorem 3.6 is our most general result, allowing for any(unbiased) sketch $\mathbf{S}$ (see (15)), and any weight matrix ${\bf W}\succ 0$. The resulting iteration complexity given by this theorem is

and is also presented in the first row of Table 1. This result depends on two expected smoothness constants ${\cal L}_{1}$ (measuring the expected smoothness of the stochastic gradient of our stochastic reformulation; see Assumption 3.1) and ${\cal L}_{2}$ (measuring the expected smoothness of the Jacobian; see Assumption 3.2). The complexity also depends on the stochastic condition number $\kappa$ (see (48)) and the sketch residual $\rho$ (see (37) and (55)). We devote considerable effort to give simple formulas for these constants under some specialized settings (for special combinations of sketches $\mathbf{S}$ and weight matrices ${\bf W}$). In fact, the entire Section 4 is devoted to this. In particular, all rows of Table 1 where the last column mentions Theorem 3.6 arise as special cases of the general iteration complexity in the first row.

Gradient descent. As a starting point, in row 3 we highlight that one can recover gradient descent as a special case of JacSketch with the choice $\mathbf{S}={\bf I}$ (with probability 1) and ${\bf W}={\bf I}$. We get the rate $\tfrac{4L}{\mu}\log\tfrac{1}{\epsilon}$, which is tight.

SAGA with uniform sampling. Let us now focus on a slightly more interesting special case: row 5. We see that SAGA with uniform probabilities appears as a special case, and enjoys the rate $(n+\tfrac{4L_{\max}}{\mu})\log\tfrac{1}{\epsilon}$, recovering an existing result.

SAGA with importance sampling. Unfortunately, the generality of Theorem 3.6 comes at a cost: we are not able to obtain an importance sampling version of SAGA as a special case which would have a better iteration complexity than uniform SAGA. This will be remedied by our second complexity theorem, which we shall discuss later below.

Minibatch SAGA. Rows 9-13 correspond to minibatch versions of SAGA. In particular, row 9 contains a general statement (albeit still a special case of the statement in row 1), covering virtually all minibatch strategies. Rows 10-13 specialize this result to two particular minibatch sketches (i.e., $\mathbf{S}={\bf I}_{S}$), each with two choices of ${\bf W}$. The first sketch corresponds to samplings $S$ which choose from among all subsets of $[n]$ uniformly at random. This sampling is known in the literature as $\tau$-nice sampling . The second sketch corresponds to $S$ being a $\tau$–partition sampling. This sampling picks uniformly at random subsets of $[n]$ which form a partition of $[n]$, and are all of cardinality $\tau$. Notice that the complexities in rows 10 and 11 are comparable (each can be slightly better than the other, depending on the values of the smoothness constants $\{L_{i}\}$). On the other hand, in the case of $\tau$–partition, the choice ${\bf W}={\rm Diag}(L_{i})$ is better than ${\bf W}={\bf I}$: the complexity in row 13 is better than that in row 12. This is because $\max_{C\in{\rm supp}(S)}\frac{1}{\tau}\sum_{i\in C}L_{i}\leq L_{\max}.$

Optimal minibatch size for SAGA. Our analysis for mini-batch SAGA also gives the first iteration complexities that interpolate between the $(n+\frac{4L_{\max}}{\mu})\log\tfrac{1}{\epsilon}$ complexity of SAGA and the $\frac{4L}{\mu}\log\tfrac{1}{\epsilon}$ complexity of gradient descent, as $\tau$ increases from $1$ to $n$. Indeed, consider the complexity in rows 10, 11 and 13 for $\tau=1$ and $\tau=n.$ Our iteration complexity of mini-batch SAGA is the first result that is precise enough to inform an optimal mini-batch size (see Section 6.2). In contrast, the previous best complexity result for mini-batch SAGA interpolates between $(n+\frac{4L_{\max}}{\mu})\log\tfrac{1}{\epsilon}$ and $\frac{4L_{\max}}{\mu}\log\tfrac{1}{\epsilon}$ as $\tau$ increases from $1$ to $n$, and thus is not precise enough as to inform the best minibatch size. We make a more detailed comparison between our results and in Section 4.7.

Specialized theorem.

We now move to the second main complexity result of our paper: Theorem 5.2. The general complexity statement is listed in row 2 of Table 1:

Gradient descent. As a starting point, we point out that just like Theorem 3.6, Theorem 5.2 also recovers the correct complexity of gradient descent as a special case (this is when $S=[n]$ with probability 1); this can be seen in row 4. Indeed, in this case we have $S=[n]$ with probability 1 (hence, $p_{[n]}=1$), ${\rm supp}(S)=\{[n]\}$, $\tau=n$ and $L_{[n]}=L$. Hence, (19) specializes to $\frac{4L}{\mu}\log\frac{1}{\epsilon}$.

SAGA with importance sampling. The first remarkable special case of (19) is summarized in row 8, and corresponds to SAGA with importance sampling. The complexity obtained, $(n+\tfrac{4\bar{L}}{\mu})\log\tfrac{1}{\epsilon}$, answers a conjecture of Schmidt et al. in the affirmative. In this case, the support of $S$ are the singletons $\{1\}$, $\{2\},\dots,\{n\}$, $p_{\{i\}}=p_{i}$ for all $i$, $\tau=1$ and $L_{\{i\}}=L_{i}$. Optimizing the complexity bound over the probabilities $p_{1},\dots,p_{n}$, we obtain the importance sampling $p_{i}=\frac{\mu n+4\tau L_{i}}{\sum_{j}\mu n+4\tau L_{j}}.$

Minibatch SAGA with importance sampling. In row 14 we state the complexity for a minibatch SAGA method with importance sampling. This is the first result for this method in the literature. Note that by comparing rows 13 and 14, we can conclude that the complexity of minibatch SAGA with importance sampling is better than for minibatch SAGA with uniform probabilities. Indeed, this is becauseWe prove inequality (20) in the appendix; see Lemma A.1.

8 Outline of the paper

We present an alternative narrative motivating the development of JacSketch in Section 2. This narrative is based on a novel technical tool which we call controlled stochastic optimization reformulations of problem (1). We then develop a general convergence theory of JacSketch in Section 3. This theory admits practically any sketches $\mathbf{S}$ (including minibatch sketches mentioned in the introduction) and weight matrices ${\bf W}$. The main result in this section is Theorem 3.6. In Section 4 we specialize the general results to minibatch sketches. Here we also compute the various constants appearing in the general complexity result for JacSketch for specific classes of minibatch samplings. In Section 5 we develop an alternative theory for JacSketch, one based on a novel stochastic Lyapunov function. The main result in this section is Theorem 5.2. Computational experiments are included in Section 6.

9 Notation

We will introduce notation when and as needed. If the reader would like to recall any notation, for ease of reference we have a notation glossary in Section D. As a general rule, all matrices are written in upper-case bold letters. By $\log t$ we refer to the natural logarithm of $t$.

Controlled Stochastic Reformulations

In this section we provide an alternative narrative behind the development of JacSketch; one through the lens of what we call controlled stochastic reformulations. These reformulations are a novel technical tool enabling us to view JacSketch from a novel perspective.

It will be useful to formalize the condition mentioned in Section 1.3 which leads to $g^{k}$ being an unbiased estimator of the gradient.

Let ${\bf W}\succ 0$ be a weighting matrix and let ${\bf D}$ be the distribution from which the sketch matrices $\mathbf{S}$ are drawn. There exists a random variable $\theta_{\mathbf{S}}$ such that

When this assumption is satisfied, we say that $(\mathbf{S},\theta_{\mathbf{S}},{\bf W})$ constitutes an “unbiased sketch”, and we call $\theta_{\mathbf{S}}$ the bias-correcting random variable. When the triple is obvious from the context, sometimes we shall simply say that $\mathbf{S}$ is an unbiased sketch.

The first key insight of this section is that besides producing unbiased estimators of the gradient, unbiased sketches produce unbiased estimators of the loss function as well. Indeed, by simply observing that $f(x)=\frac{1}{n}\left<F(x),e\right>$, we get

In other words, we can rewrite the finite-sum optimization problem (1) as an equivalent stochastic optimization problem where the randomness comes from ${\cal D}$ rather than from the representation-specific uniform distribution over the $n$ loss functions:

The stochastic optimization problem (22) is a stochastic reformulation of the original problem (1). Further, the stochastic gradient of this reformulation is given by

With these simple observations, our options at designing stochastic gradient-type algorithms for (1) have suddenly broadened dramatically. Indeed, we can now solve the problem, at least in principle, by applying SGD to any stochastic reformulation:

But now we have a parameter to play with, namely, the distribution of $\mathbf{S}$. The choice of this parameter will influence both the iteration complexity of the resulting method as well as the cost of each iteration. We now give a few examples of possible choices of ${\cal D}$ to illustrate this.

Let $\mathbf{S}$ be equal to ${\bf I}$ (or any other $n\times n$ invertible matrix) with probability 1 and let ${\bf W}\succ 0$ be chosen arbitrarily. Then $\theta_{\mathbf{S}}\equiv 1$ is bias-correcting since

With this setup, the SGD method (24) becomes gradient descent:

Let $S_{k}=\{i_{k}\}$ be picked at iteration $k$. Then the SGD method (24) becomes SGD with non-uniform sampling:

Note that with this setup, and when $p_{i}=1/n$ for all $i$, the stochastic reformulation is identical to the original finite-sum problem. This is the case because $f_{e_{i}}(x)=f_{i}(x)$.

Let $\mathbf{S}=e_{S}=\sum_{i\in S}e_{i}$, where $S=C\subseteq[n]$ with probability $p_{C}$. Let ${\bf W}={\bf I}$. Assume that the cardinality of the set $\{C\subseteq[n]\;:\;C\in{\rm supp}(S),\;i\in C\}$ does not depend on $i$ (and is equal to $c_{1}>0$). Then $\theta_{e_{S}}=1/(c_{1}p_{S})$ is bias-correcting since

Note that ${\bf\Pi}_{e_{S}}e=e_{S}$. Assume that set $S_{k}$ is picked in iteration $k$.Then the SGD method (24) becomes minibatch SGD with non-uniform sampling:

Finally, note that gradient descent (25) is a special case of (27) if we set $p_{[n]}=1$ and $p_{C}=0$ for all other subsets $C$ of $[n]$. Likewise, SGD with non-uniform probabilities (26) is a special case of (27) if we set $p_{\{i\}}=p_{i}>0$ for all $i$ and $p_{C}=0$ for all other subsets $C$ of $[n]$.

2 The controlled stochastic reformulation

Though SGD applied to the stochastic reformulation can generate several known algorithms in special cases, there is no reason to believe that the gradient estimates $g^{k}$ will have diminishing variance (excluding the extreme case such as gradient descent). Here we handle this issue using control variates, a commonly used tool to reduce variance in Monte Carlo methods .

Given a random function $z_{\mathbf{S}}(x)$, we introduce the controlled stochastic reformulation:

is an unbiased estimator of the gradient $\nabla f(x)$, we can apply SGD to the controlled stochastic reformulation instead, which leads to the method

Reformulation (22) and method (24) is recovered as a special case with the choice $z_{\mathbf{S}}(x)\equiv 0$. However, we now have the extra freedom to choose $z_{\mathbf{S}}(x)$ so as to control the variance of this stochastic gradient. In particular, if $\nabla z_{\mathbf{S}}(x)$ and $\nabla f_{\mathbf{S}}(x)$ are sufficiently correlated, then (29) will have a smaller variance than $\nabla f_{\mathbf{S}}(x).$ For this reason, we choose a linear model for $z_{\mathbf{S}}(x)$ that mimicks the stochastic function $f_{\mathbf{S}}(x).$

We collect this observation that (32) is unbiased in the following lemma for future reference.

If $\mathbf{S}$ is an unbiased sketch (see Definition 2.1), then

Now it remains to choose the matrix ${\bf J}$, which we do by minimizing the variance of our gradient estimate.

3 The Jacobian estimate, variance reduction and the sketch residual

and we have used the weighted Frobenius norm with weight matrix ${\bf B}$ (see (10)).

For most distributions ${\cal D}$ of interest, the matrix ${\bf B}$ is positive definiteExcluding such trivial cases as when $\mathbf{S}$ is an invertible matrix and $\theta_{\mathbf{S}}=1$ with probability one, in which case ${\bf B}=0$. Letting $v_{\mathbf{S}}\overset{\text{def}}{=}({\bf I}-\theta_{\mathbf{S}}{\bf\Pi}_{\mathbf{S}})e$, we can bound the largest eigenvalue of matrix ${\bf B}$ via Jensen’s inequality as follows:

Combined with (34), we get the the following bound on the variance of $\nabla f_{\mathbf{S},{\bf J}}$:

Let us now return to the identity (34) and its role in choosing ${\bf J}$. Minimizing the variance in a single step is overly ambitious, since it requires setting ${\bf J}={\bf\nabla F}(x)$, which is costly. So instead, we propose to minimize (34) iteratively. But first, to make (34) more manageable, we upper-bound it using a norm defined by the weight matrix ${\bf W}$ as follows

is the largest eigenvalue of ${\bf W}^{1/2}{\bf B}{\bf W}^{1/2}$. We refer to the constant $\rho$ as the sketch residual, and it is a key constant affecting the convergence rate of JacSketch as captured by Theorem 3.6. The sketch residual $\rho$ represents how much information is “lost” on average due to sketching and due to how well ${\bf W}^{-1}$ approximates ${\bf B}$. We develop formulae and estimates of the sketch residual for several specific sketches of interest in Section 4.5.

Consider the setup from Example 2.2 (gradient descent). That is, let $\mathbf{S}$ be invertible with probability one and let $\theta_{\mathbf{S}}=1$ be the bias-reducing variable. Then ${\bf\Pi}_{\mathbf{S}}e=e$ and hence ${\bf B}=0$, which means that $\rho=0$.

We have switched from the ${\bf B}$ norm to a user-controlled ${\bf W}^{-1}$ norm because minimizing under the ${\bf B}$ norm will prove to be impractical because ${\bf B}$ is a dense matrix for most all practical sketches. With this norm change we now have the option to set ${\bf W}$ as a sparse matrix (e.g., the identity, or a diagonal matrix), as we explain in Remark 2.8 further down. However, the theory we develop allows for any symmetric positive definite matrix ${\bf W}$.

We can now minimize (36) iteratively by only using a single sketch of the true Jacobian at each iteration. Suppose we have a current estimate ${\bf J}^{k}$ of the true Jacobian and a sketch of the true Jacobian ${\bf\nabla F}(x^{k})\mathbf{S}_{k}$. With this we can calculate an improved Jacobian estimate using a projection step

the solution of which, as it turns out, depends on ${\bf\nabla F}(x^{k})$ through its sketch ${\bf\nabla F}(x^{k})\mathbf{S}_{k}$ only. That is, we choose the next Jacobian estimate ${\bf J}^{k+1}$ as close as possible to the true Jacobian ${\bf\nabla F}(x^{k})$ while restricted to a matrix subspace that passes through ${\bf J}^{k}$. Thus in light of (36), the variance is decreasing. The explicit solution to (38) is given by

See Lemma B.1 in the appendix for the proof. Note that, as alluded to before, ${\bf J}^{k+1}$ depends on ${\bf\nabla F}(x^{k})$ through its sketch only. Note that (39) updates the Jacobian estimate by re-using the sketch ${\bf\nabla F}(x^{k})\mathbf{S}_{k}$ which we also use when calculating the stochastic gradient (32).

Note that (39) gives the same formula for ${\bf J}^{k+1}$ as (11) which we obtained by solving (1.3); i.e., by projecting ${\bf J}^{k}$ onto the solution set of (8). This is not a coincidence. In fact, the optimization problems (1.3) and (38) are mutually dual. This is formally stated in Lemma B.1 which can be found in the appendix. In the context of solving linear systems, this was observed in . Therein, (1.3) is called the sketch-and-project method, whereas (38) is called the constrain-and-approximate problem. In this sense, the Jacobian sketching narrative we followed in Section 1.3 is dual to the Jacobian sketching narrative we are pursuing here.

Loosely speaking, the denser the weighting matrix ${\bf W}$, the higher the computational cost for updating the Jacobian using (39). Indeed, the sparsity pattern of ${\bf W}$ controls how many elements of the previous Jacobian estimate ${\bf J}^{k}$ need to be updated. This can be seen by re-arranging (39) as

4 JacSketch Algorithm

Combining formula (32) for the stochastic gradient of the controlled stochastic reformulation with formula (39) for the update of the Jacobian estimate, we arrive at our JacSketch algorithm (Algorithm 1).

From the point of view of the controlled stochastic reformulation, JacSketch can also be written in the form of Algorithm 2.

5 A window into biased estimates and SAG

We will now take a small detour from the main flow of the paper to develop an alternative viewpoint of Algorithm 1 and also make a bridge to biased methods such as SAG .

suggests that $\hat{g}^{k}=\frac{1}{n}{\bf J}^{k+1}e$, where ${\bf J}^{k+1}\approx{\bf\nabla F}(x^{k})$ would give a good estimate of the gradient. To decrease the variance of $\hat{g}^{k}$, we can also use the same update of the Jacobian estimate (39) since

The issue with using $\frac{1}{n}{\bf J}^{k+1}e$ as an estimator of the gradient is that it decreases the variance too aggressively, neglecting the bias. However, this can be fixed by trading off variance for bias. One way to do this is to introduce the random variable $\theta_{\mathbf{S}}$ as a stochastic relaxation parameter

If $\theta_{\mathbf{S}}$ is bias correcting, we recover the unbiased SAGA estimator (13). By allowing $\theta_{\mathbf{S}}$ to be closer to one, however, we will get more bias and lower variance. We leave this strategy of building biased estimators for future work. It is conceivable that SAG could be analyzed using reasonably small modifications of the tools developed in this paper. Doing this would be important due to at least four reasons: i) SAG was the first variance-reduced method for problem (1), ii) the existing analysis of SAG is not satisfying, iii) one may be able to obtain a better rate, iv) one may be able to develop and analyze novel variants of SAG.

Convergence Analysis for General Sketches

In this section we establish a convergence theorem (Theorem 3.6) which applies to general sketching matrices $\mathbf{S}$ (that is, arbitrary distributions ${\cal D}$ from which they are sampled). By design, we keep the setting in this section general, and only deal with specific instantiations and special cases in Section 4.

We first formulate two expected smoothness assumptions tying together $f$, its Jacobian ${\bf\nabla F}(x)$ and the distribution ${\cal D}$ from which we pick sketch matrices $\mathbf{S}$. These assumptions, and the associated expected smoothness constants, play a key role in the convergence result.

Our first assumption concerns the expected smoothness of the stochastic gradients $\nabla f_{\mathbf{S}}$ of the stochastic reformulation (22).A similar relation to (43) holds for the stochastic optimization reformulation of linear systems studied by Richtárik and Takáč . Therein, this relation holds as an identity with ${\cal L}_{1}=1$ (see Lemma 3.3 in ). However, the function $f_{\mathbf{S}}$ considered there is entirely different and, moreover, $f(x^{*})=0$ and $\nabla f_{\mathbf{S}}(x^{*})=0$ for all $\mathbf{S}$.

There is a constant ${\cal L}_{1}>0$ such that

It is easy to see from (23) and (32) that

Our second expected smoothness assumption concerns the Jacobian of $F$.

There is a constant ${\cal L}_{2}>0$ such that

where the norm is the weighted Frobenius norm defined in (10).

Therefore, (45) can be equivalently written in the form

which suggests that the above condition indeed measures the variation/smoothness of the Jacobian under a specific weighted Frobenius norm. To the best of our knowledge, the above expected smoothness conditions are new, and have not been considered in the literature before.

2 Stochastic condition number

By the stochastic condition number associated with ${\bf W}$ and ${\cal D}$ we mean the constant defined by

In the next lemma we show that $0\leq\kappa\leq 1$ for all distributions ${\cal D}$ for which the expectation (48) exists.

For all distributions $\mathcal{D},$ we have the bounds $0\leq\kappa\leq 1.$

In our convergence theorem we will assume that $\kappa>0$. This can be achieved by choosing a suitable distribution ${\cal D}$ and it holds trivially for all the examples we develop. The condition $\kappa>0$ essentially says that the distribution is sufficiently rich. This condition number was first proposed in in the context of randomized algorithms for solving linear systems. We refer the reader to that work for details on sufficient assumptions about ${\cal D}$ guaranteeing $\kappa>0$. Below we give an example.

Let ${\bf W}\succ 0$, and let ${\cal D}$ be given by setting $\mathbf{S}=e_{i}$ with probability $p_{i}>0$. Then

3 Convergence theorem

Our main convergence result, which we shall present shortly, holds for $\mu$-strongly convex functions. However, it turns out we can establish the result for a somewhat larger family of functions. This family is described next.

We are now ready to present the main result of this section.

Let ${\bf W}\succ 0$. Let $f$ satisfy Assumption 3.5. Let Assumption 2.1 be satisfied (i.e, $\mathbf{S}$ is an unbiased sketch and $\theta_{\mathbf{S}}$ is the associated bias-correcting random variable). Let the expected smoothness assumptions be satisfied: Assumption 3.1 and Assumption 3.2. Assume that $\kappa>0$. Let the sketch residual be defined as in (37), i.e,

If we choose $\alpha$ to be equal to the upper bound in (53), then

Recall that the iteration complexity expression from (55) is listed in row 1 of Table 1.

The Lyapunov function we use is simply the sum of the squared distance between $x^{k}$ to the optimal $x^{*}$ and the distance of our Jacobian estimate ${\bf J}^{k}$ to the optimal Jacobian ${\bf\nabla F}(x^{*}).$ Hence, the theorem says that both the iterates $\{x^{k}\}$ and the Jacobian estimates $\{{\bf J}^{k}\}$ converge.

4 Projection lemmas and the stochastic condition number κ𝜅\kappa

In this section we collect some basic results on projections. Recall from (12) that ${\bf\Pi}_{\mathbf{S}}=\mathbf{S}(\mathbf{S}^{\top}{\bf W}\mathbf{S})^{\dagger}\mathbf{S}^{\top}{\bf W}$ and from (46) that ${\bf H}_{\mathbf{S}}=\mathbf{S}(\mathbf{S}^{\top}{\bf W}\mathbf{S})^{\dagger}\mathbf{S}^{\top}$.

Proof: Using the pseudoinverse property ${\bf A}^{\dagger}{\bf A}{\bf A}^{\dagger}={\bf A}^{\dagger}$ we have that

and as a consequence (56) holds. Moreover,

Finally, taking expectation over (58) and (59) gives (57). ∎

By taking expectations in ${\cal D}$, we get

where in the last step we used the estimate

5 Key lemmas

We first establish two lemmas. The first lemma provides an upper bound on the quality of new Jacobian estimate in terms of the quality of the current estimate and function suboptimality. If the second term on the right hand side was not there, the lemma would be postulating a contraction on the quality of the Jacobian estimate.

Let Assumption 3.2 be satisfied. Then iterates of Algorithm 1 satisfy

Proof: Subtracting ${\bf\nabla F}(x^{*})$ from both sides of (39) gives

Taking norms on both sides, then expectation with respect to $\mathbf{S}_{k}$ and then using Lemma 3.8, we get

We now bound the second moment of $g^{k}$. The lemma implies that as $x^{k}$ approaches $x^{*}$ and ${\bf J}^{k}$ approaches ${\bf\nabla F}(x^{*})$, the variance of $g^{k}$ approaches zero. This is a key property of JacSketch which elevates it into the ranks of variance-reduced methods.

Let $\mathbf{S}$ be an unbiased sketch. Let Assumption 3.1 be satisfied (i.e., assume that inequality (43) holds for some ${\cal L}_{1}>0$). Then the second moment of the estimated gradient is bounded by

Proof: Adding and subtracting $\tfrac{\theta_{\mathbf{S}_{k}}}{n}{\bf\nabla F}(x^{*}){\bf\Pi}_{\mathbf{S}_{k}}e$ in (13) gives

Taking norms on both sides and using the bound $\|a+b\|_{2}^{2}\leq 2\|a\|_{2}^{2}+2\|b\|_{2}^{2}$ gives

In view of Assumption 3.1 (combine (43) and (44)), we have

If we now let $v={\bf W}^{1/2}(\theta_{\mathbf{S}_{k}}{\bf\Pi}_{\mathbf{S}_{k}}-{\bf I})e$ and ${\bf M}=({\bf J}^{k}-{\bf\nabla F}(x^{*})){\bf W}^{-1/2}$, then we can continue:

where in the last step we have used the assumption that $\theta_{\mathbf{S}_{k}}$ is bias-correcting:

It now only remains to substitute (66) and (67) into (65) to arrive at (64). ∎

6 Proof of Theorem 3.6

With the help of the above lemmas, we now proceed to the proof of the theorem. In view of the strong convexity assumption (50), we have

By using the relationship $x^{k+1}=x^{k}-\alpha g^{k}$, the fact that $g^{k}$ is an unbiased estimate of the gradient $\nabla f(x^{k})$, and using one-point strong convexity (69), we get

Next, applying Lemma 3.10 leads to the estimate

We now choose $\alpha$ so that $\text{I}\leq 0$ and $\text{II}\leq 1-\alpha\mu$, which can be written as

Minibatch Sketches

In this section we focus on special cases of Algorithm 1 where one computes $\nabla f_{i}(x^{k})$ for $i\in S^{k}$, where $S^{k}$ is a random subset (mini-batch) of $[n]$ chosen in each iteration according to some fixed probability law. As we have seen in the introduction, this is achieved by choosing $\mathbf{S}_{k}={\bf I}_{S_{k}}$.

where $\sum_{C\subseteq[n]}p_{C}=1$ and $p_{C}\geq 0$ for all $C$.

We refer the reader to for a background reading on samplings and their properties.

The support of a sampling $S$ is the set of subsets of $[n]$ which are chosen by $S$ with positive probability: ${\rm supp}(S)\overset{\text{def}}{=}\{C\;:\;p_{C}>0\}$. We say that $S$ has uniform support if

for all $i,j\in[n]$. In such a case we say that the support is $c_{1}$–uniform.

To illustrate the above concepts, we now list a few examples with $n=4$.

The sampling defined by setting $p_{\{1,2\}}=p_{\{3,4\}}=0.5$ is non-vacuous, proper, $2$–uniform ($p_{i}=0.5$ for all $i$ and $|S|=2$ with probability 1), and has $1$–uniform support. If we change the probabilities to $p_{\{1,2\}}=0.4$ and $p_{\{3,4\}}=0.6$, the sampling is no longer uniform (since $p_{1}=0.4\neq 0.6=p_{3}$), but it still has $1$–uniform support, is proper and non-vacuous. Hence, a sampling with uniform support need not be uniform. On the other hand, a uniform sampling need not have uniform support. As an example, consider sampling $S$ defined via $p_{\{1\}}=0.4$, $p_{\{2,3\}}=p_{\{3,4\}}=p_{\{2,4\}}=0.2$. It is uniform (since $p_{i}=0.4$ for all $i$). However, while element $1$ appears in a single set of its support, elements $2,3$ and $4$ each appear in two sets. So, this sampling does not have uniform support.

A uniform sampling need not be $\tau$–uniform for any $\tau$. For example, the sampling defined by setting $p_{\{1,2,3,4\}}=0.5$, $p_{\{1,2\}}=0.25$ and $p_{\{3,4\}}=0.25$ is uniform (since $p_{i}=0.75$ for all $i$), but as it assigns positive probabilities to sets of at least two different cardinalities, it is not $\tau$–uniform for any $\tau$.

Further, the sampling defined by setting $p_{\{1,2\}}=1/6$, $p_{\{1,3\}}=1/6$, $p_{\{1,4\}}=1/6$, $p_{\{2,3\}}=1/6$, $p_{\{2,4\}}=1/6$, $p_{\{3,4\}}=1/6$ is non-vacuous, $2$–uniform ($p_{i}=1/2$ for all $i$ and $|S|=2$ with probability 1), and has $3$–uniform support. The sampling defined by setting $p_{\{1,2\}}=1/3$, $p_{\{2,3\}}=1/3$, $p_{\{3,1\}}=1/3$ is non-vacuous, proper, $2$–uniform ($p_{i}=2/3$ for all $i$ and $|S|=2$ with probability 1) and has $2$–uniform support.

Note that a sampling with uniform support is necessarily proper as long as $c_{1}>0$. However, it need not be non-vacuous. For instance, the sampling $S$ defined by setting $p_{\emptyset}=1$ has –uniform support and is vacuous. From now on, we only consider samplings with the following properties.

$S$ is non-vacuous and has $c_{1}$–uniform support with $c_{1}\geq 1$.

Note that if $S$ is a non-vacuous sampling with $1$–uniform support, then its support is necessary a partition of $[n]$. We shall pay specific attention to such samplings in Section 5 as for them we can develop a stronger analysis than that provided by Theorem 3.6.

2 Minibatch sketches and projections

In the next result we describe some basic properties of the projection matrix ${\bf\Pi}_{\mathbf{S}}=\mathbf{S}(\mathbf{S}^{\top}{\bf W}\mathbf{S})^{\dagger}\mathbf{S}^{\top}{\bf W}$ associated with a minibatch sketch $\mathbf{S}$.

${\bf\Pi}_{\mathbf{S}}={\bf I}_{S}{\bf I}_{S}^{\top}$. This is a diagonal matrix with the $i$th diagonal element equal to 1 if $i\in S$, and if $i\notin S$.

${\bf\Pi}_{\mathbf{S}}e=e_{S}\overset{\text{def}}{=}\sum_{i\in S}e_{i}.$

The stochastic condition number defined in (48) is given by $\kappa=\min_{i}p_{i}$

Let $S$ satisfy Assumption 4.6. Then the random variable

This follows by noting that ${\bf I}_{S}^{\top}{\bf W}{\bf I}_{S}$ is the $|S|\times|S|$ diagonal matrix with diagonal entries corresponding to $w_{i}$ for $i\in S$, which in turn can be used to show that $({\bf I}_{S}^{\top}{\bf W}{\bf I}_{S})^{-1}{\bf I}_{S}^{\top}{\bf W}={\bf I}_{S}^{\top}$.

This follows from (i) by taking expectations of the diagonal elements of ${\bf\Pi}_{\mathbf{S}}$.

where the last equation follows from the assumption that the support of $S$ is $c_{1}$–uniform. ∎

The following simple observation will be useful in the computation of the constant ${\cal L}_{1}$. The proof is straightforward and involves a double counting argument.

Let $S$ be a sampling satisfying Assumption 4.6. Moreover, assume that $S$ is $\tau$–uniform. Then $\frac{|{\rm supp}(S)|}{c_{1}}=\frac{n}{\tau}$. Consequently, $\kappa=p_{1}=p_{2}=\dots=p_{n}=\frac{\tau}{n}=\frac{c_{1}}{|{\rm supp}(S)|}$, where $\kappa$ is the stochastic condition number associated with the minibatch sketch $\mathbf{S}={\bf I}_{S}$.

3 JacSketch for minibatch sampling = minibatch SAGA

As we have mentioned in Section 1.4 already, JacSketch admits a particularly simple form for minibatch sketches, and corresponds to known and new variants of SAGA. Assume that $S$ satisfies Assumption 4.6 and let ${\bf W}={\rm Diag}(w_{1},\dots,w_{n})$. In view of Lemma 4.7(vi), this means that the random variable $\theta_{\mathbf{S}}=\frac{1}{c_{1}p_{S}}$ is bias-correcting, and due to Lemma 4.7(ii), we have ${\bf\Pi}_{\mathbf{S}_{k}}e=e_{S_{k}}=\sum_{i\in S_{k}}e_{i}$. Therefore,

By Lemma 4.7(i), ${\bf\Pi}_{\mathbf{S}_{k}}={\bf I}_{S_{k}}{\bf I}_{S_{k}}^{\top}$. In view of (11), the Jacobian estimate gets updated as follows

The resulting minibatch SAGA method is formalized as Algorithm 3.

Below we specialize the formula for $g^{k}$ to a few interesting special cases.

Standard uniform SAGA is obtained by setting $S_{k}=\{i\}$ with probability $1/n$ for each $i\in[n]$. Since the support of this sampling is $1$–uniform, we set $c_{1}=1$. This leads to the gradient estimate

However, we can use non-uniform probabilities instead. Let $S_{k}=\{i\}$ with probability $p_{i}>0$ for each $i\in[n]$. Since the support of this sampling is 1–uniform, we have $c_{1}=1$. So, the gradient estimate has the form

Let $C_{1},\dots,C_{q}$ be nonempty subsets of forming a partition $[n]$. Let $S_{k}=C_{j}$ with probability $p_{C_{j}}>0$. The support of this sampling is $1$–uniform, and hence we can choose $c_{1}=1$. This leads to the gradient estimate

Let $S_{k}$ be chosen uniformly at random from all subsets of $[n]$ of cardinality $\tau\geq 2$. That is, $\mathbf{S}_{k}$ is the $\tau$-nice sampling, and the probabilities are equal to $p_{S_{k}}=1/{n\choose\tau}$. This sampling has $c_{1}$–uniform support with $c_{1}={n-1\choose\tau-1}=\frac{\tau}{n}{n\choose\tau}$. Thus, $nc_{1}p_{S_{k}}=\tau$, and we have

Consider the same situation as in Example 4.12, but with $\tau=n$. That is, we choose $S_{k}=[n]$ with probability $1$, and $c_{1}=1$. Then

Here we compute the expected smoothness constants ${\cal L}_{1}$ and ${\cal L}_{2}$ in the case of $\mathbf{S}$ being a minibatch sketch $\mathbf{S}={\bf I}_{S}$, and assuming that $f$ is convex and smooth. We first formalize the notion of smoothness we will use.

The above assumption is somewhat non-standard. Note that, however, if we instead assume that each $f_{i}$ is convex and $L_{i}$-smooth, then the above assumption holds for $L_{C}=\frac{1}{|C|}\sum_{i\in C}L_{i}$. In some cases, however, we may have better estimates of the constants $L_{C}$ than those provided by the averages of the $L_{i}$ values. The value of these constants will have a direct influence on ${\cal L}_{1}$ and ${\cal L}_{2}$, which is why we work with this more refined assumption instead.

where in the last step we used the fact that $\sum_{i=1}^{n}\nabla f_{i}(x^{*})=n\nabla f(x^{*})=0.$ ∎

where $L_{i}=L_{\{i\}}$. If moreover, $S$ is $\tau$-uniform, thenNote that $c_{1}=|\{C\in{\rm supp}(S)\;:\;1\in C\}|$, and hence ${\cal L}_{1}$ has the form of a maximum over averages.

where in this last inequality we have used convexity of $f_{i}$ for $i\in[n]$. Since

the formula for ${\cal L}_{1}$ now follows by comparing (86) to (43). In order to establish the formula for ${\cal L}_{2}$, we estimate

The specialized formulas (85) for $\tau$–uniform sampling follow as special cases of the general formulas (84) by applying Lemma 4.8. ∎

In the next result we establish some inequalities relating the quantities $L$, $L_{\max}$, $L_{C}$ and $L^{{\cal G}}_{\max}.$ In particular, the results says that for a certain family of samplings $S$ (the same for which we have defined the quantity $L^{{\cal G}}_{\max}$ in (85)), the expected smoothed constant $L^{{\cal G}}_{\max}$ is lower-bounded by the average of $L_{C}$ over $C\in{\cal G}={\rm supp}(S)$, and upper-bounded by $L_{\max}$.

Let $S$ be a $\tau$–uniform sampling ($\tau\geq 1$) with $c_{1}$–uniform support ($c_{1}\geq 1$). Let ${\cal G}={\rm supp}(S)$. Then

The last inequality holds without the need to assume $\tau$–uniformity.

Proof: Using the fact that $S$ has $c_{1}$–uniform support, and utilizing a double-counting argument, we observe that $\sum_{C\in{\cal G}}|C|f_{C}(x)=c_{1}\sum_{i=1}^{n}f_{i}(x)$. Multiplying both sides by $\frac{1}{nc_{1}}$, and since $|C|=\tau$ for all $C\in{\cal G}$, we get $\frac{\tau|{\cal G}|}{c_{1}n}\frac{1}{|{\cal G}|}\sum_{C\in{\cal G}}f_{C}(x)=\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)=f(x).$ To obtain (88), it now only remains to use the identity

which was shown in Lemma 4.8. The first inequality in (89) follows from (88) using standard arguments (identical to those that lead to the inequality $L\leq\bar{L}$).

Let us now establish the second inequality in (89). Define $L^{{\cal G}}_{i}\overset{\text{def}}{=}\frac{1}{c_{1}}\sum_{C\in{\cal G}\;:\;i\in C}L_{C}$. Again using a double-counting argument we observe that $\tau\sum_{C\in{\cal G}}L_{C}=c_{1}\sum_{i=1}^{n}L^{{\cal G}}_{i}.$ Multiplying both sides of this equality by $\frac{|{\cal G}|}{c_{1}n}$ and using identity (90), we get $\frac{1}{|{\cal G}|}\sum_{C\in{\cal G}}L_{C}=\frac{1}{n}\sum_{i=1}^{n}L^{{\cal G}}_{i}\leq\max_{i}L^{{\cal G}}_{i}=L^{{\cal G}}_{\max}.$ We will now establish the last inequality by proving that $L^{{\cal G}}_{i}\leq L_{\max}$ for any $i$:

Note that we did not need to assume $\tau$–uniformity to prove that $L^{{\cal G}}_{\max}\leq L_{\max}$. ∎

5 Estimating the sketch residual ρ𝜌\rho

In this section we compute the sketch residual $\rho$ for several classes of samplings $S$. Let ${\cal G}={\rm supp}(S)$. We will assume throughout this section that $S$ is non-vacuous, has $c_{1}$–uniform support (with $c_{1}\geq 1$), and is $\tau$–uniform.

Further, we assume that ${\bf W}={\rm Diag}(w_{1},\dots,w_{n})$, and that the bias-correcting random variable $\theta_{\mathbf{S}}$ is chosen as $\theta_{\mathbf{S}}=\tfrac{1}{c_{1}p_{S}}=\tfrac{|{\cal G}|}{c_{1}}$ (see (75) and Lemma 4.8). In view of the above, since ${\bf\Pi}_{{\bf I}_{C}}e=e_{C}$, the sketch residual is given by

where the last equality follows by permuting the multiplication of matrices within the $\lambda_{\max}.$

In the following text we calculate upper bounds for $\rho$ for $\tau$–partition and $\tau$–nice samplings. Note that Theorem 3.6 still holds if we use an upper bound of $\rho$ in place of $\rho$.

If $S$ is the $\tau$–partition sampling, then

Proof: Using Lemma 4.8, and since $c_{1}=1$, we get $\frac{|{\cal G}|}{c_{1}^{2}}=\frac{n}{\tau}$. Consequently,

where $w_{C}=\sum_{i\in C}w_{i}e_{i}$ and we used that $-{\bf W}^{1/2}ee^{\top}{\bf W}^{1/2}$ is negative semidefinite. When ${\bf W}={\bf I}$, the above bound is tight. By Gershgorin’s theorem, every eigenvalue $\lambda$ of the matrix is bounded by at least one of the inequalities $\lambda\leq\sum_{i\in C}w_{i}$ for $C\in{\cal G}$. Consequently, from (93) we have that $\rho\leq\frac{n}{\tau}\max_{C\in{\cal G}}\sum_{i\in C}w_{i}.$ ∎

Next we give an useful upper bound on $\rho$ for a large family of uniform samplings (for proof, see Appendix C).

Let ${\cal G}$ be a collection of subsets of $[n]$ with the property that the number of sets $C\in{\cal G}$ containing distinct elements $i,j\in[n]$ is the same for all $i,j$. In particular, define

Now define a sampling $S$ by setting $S=C\in{\cal G}$ with probability $\frac{1}{|{\cal G}|}$. Moreover, assume that the support of $S$ is $c_{1}$–uniform. Consider the minibatch sketch $\mathbf{S}={\bf I}_{S}$.

If ${\bf W}={\rm Diag}(w_{1},\ldots,w_{n})$, then

Note that as long as $\tau\geq 2$, the $\tau$–nice sampling $S$ satisfies the assumptions of the above theorem. Indeed, ${\cal G}$ is the support of $S$ consisting of all subsets of $[n]$ of size $\tau$, $|{\cal G}|={n\choose\tau}$, $c_{1}={n-1\choose\tau-1}$, and $c_{2}={n-2\choose\tau-2}$. As a result, bound (95) simplifies to

6 Calculating the iteration complexity for special cases

In this section we consider minibatch SAGA (Algorithm 3) and calculate its iteration complexity in special cases using Theorem 3.6 by pulling together the formulas for ${\cal L}_{1},{\cal L}_{2},\kappa$ and $\rho$ established in previous sections. In particular, assume $S$ is $\tau$–uniform and has $c_{1}$–uniform support with $c_{1}\geq 1$. In this case, formula (85) for ${\cal L}_{1},{\cal L}_{2}$ from Lemma 4.16 applies and we have ${\cal L}_{1}=L^{{\cal G}}_{\max}$ and ${\cal L}_{2}=\tau\max_{i}\left\{\frac{L_{i}}{w_{i}}\right\}$.

Moreover, by Lemma 4.8, $\kappa=\tfrac{\tau}{n}$. By Theorem 3.6, if we use the stepsize

then the iteration complexity is given by

Complexity (100) is listed in line 9 of Table 1. The complexities in lines 3, 5 and 10–13 arise as special cases of (100) for specific choices of $S$:

In line 3 we have gradient descent. This arises for the choice ${\bf W}={\bf I}$ and $S=[n]$ with probability 1. In this case, $\tau=n$, $L^{{\cal G}}_{\max}=L$ and $\rho=0$. So, (100) simplifies to

In line 5 we have uniform SAGA. We choose ${\bf W}={\bf I}$ and $S=\{i\}$ with probability $1/n$. We have $\tau=1$ and $L^{{\cal G}}_{\max}=L_{\max}$. In view of Theorem 4.18, $\rho\leq n$. So, (100) simplifies to

In line 10 we choose ${\bf W}={\bf I}$ and $S$ is the $\tau$-nice sampling. In this case, Theorem 4.19 says that $\rho=\frac{n}{\tau}\frac{n-\tau}{n-1}$ (see (98)). Therefore, (100) reduces to

In line 11 we choose ${\bf W}={\rm Diag}(L_{i})$ and $S$ is the $\tau$-nice sampling. Theorem 4.19 says that $\rho\leq\tfrac{n-\tau}{\tau}\left(\tfrac{n-2}{n-1}L_{\max}+\tfrac{n}{n-1}\bar{L}\right)$ (see (97)). Therefore, (100) reduces to

To simplify the above expression, one may further use the bound $\tfrac{n-2}{n-1}L_{\max}+\tfrac{n}{n-1}\bar{L}\leq L_{\max}+\bar{L}$. In Table 1 we have listed the complexity in this simplified form. Whether (103) or (104) is better depends on the constants $\{L_{i}\}$. Indeed, when there exists $i$ such that $L_{i}\gg L_{j}$, for $j\neq i$ then

thus (104) is smaller than (103). On the other extreme, when $L_{i}=L_{j}$ for all $i,j$, then

so long as $n\geq 1$. In this case (104) is larger than (103).

In line 12 of Table 1 we let ${\bf W}={\bf I}$ and $S$ is the $\tau$-partition sampling. In view of Theorem 4.18, $\rho\leq\tfrac{n}{\tau}\tau=n$ and hence (100) reduces to

In line 13 of Table 1 we let ${\bf W}={\rm Diag}(L_{i})$ and $S$ is the $\tau$-partition sampling. In view of Theorem 4.18, $\rho\leq\frac{n}{\tau}\max_{C\in{\cal G}}\sum_{i\in C}L_{i}$ and hence (100) reduces to

Note that the bound in (106) is better than (105) because $\max_{C\in{\cal G}}\sum_{i\in C}L_{i}\leq\tau L_{\max}.$

7 Comparison with previous mini-batch SAGA convergence results

Recently in , a method that includes a mini-batch variant of SAGA was proposed. This work is the most closely related to our minibatch SAGA, and was developed independently from ours. The methods described in can be cast in our framework. In the language of our paper, in the authors update the Jacobian estimate according to (77), where $S_{k}$ is sampled according to a uniform probability with $p_{i}=\tau/n,$ for all $i=1,\ldots,n.$ What do differently is that instead of introducing the bias-corecting random variable $\theta_{\mathbf{S}}$ to maintain an unbiased gradient estimate, the gradient estimate is updated using the standard SAGA update (78) and this sampling process is done independently of how $S_{k}$ is sampled for the Jacobian update. Thus at every iteration a gradient $\nabla f_{i}(x^{k})$ is sampled to compute (78), but is then discarded and not used to update the Jacobian update so as to maintain the independence between ${\bf J}^{k}$ and $g^{k}.$ By introducing the bias-correcting random variable $\theta_{\mathbf{S}}$ in our method we avoid the data-hungry strategy used in .

The analysis provided in shows that, by choosing the stepsize appropriately, the expectation of a Lyapunov function similar to (52) is less than $\epsilon>0$ after

iterations, where $K\overset{\text{def}}{=}\frac{4L_{\max}}{\mu}$. When $\tau=1$ this gives an iteration complexity of $O(n+K)\log\frac{1}{\epsilon},$ which is essentially the same complexity as the standard SAGA method. The main issue with this complexity is that it decreases only very modestly as $\tau$ increases. In particular, on the extreme end when $\tau=n$, since $K\geq 4$, we can approximate $(1+K)^{2}\approx 1+K^{2}$ and the resulting complexity (107) becomes

Yet we know that $\tau=n$ corresponds to gradient descent, and thus the iteration complexity should be $O(\tfrac{L}{\mu}\log(1/\epsilon)),$ which is what we recover in the analysis of all our mini-batch variants. In Figures 3(a), 3(b) and 3(c) in the experiments in Section 6 we illustrate how (107) descreases very modestly as $\tau$ increases.

A Refined Analysis with a Stochastic Lyapunov Function

In this section we perform a refined analysis of JacSketch applied with a minibatch sketch $\mathbf{S}={\bf I}_{S}$ for a special class of samplings $S$ which pick uniformly at random from a partition of $[n]$ into sets of size $\tau$. This is only possible when $n$ is a multiple of $\tau$.

In the terminology introduced in Section 4.1, a $\tau$–partition sampling is non-vacuous, proper and $\tau$–uniform. Its support is a partition of $[n]$, and is $1$–uniform. It satisfies Assumption 4.6.

Restricting our attention to $\tau$–partition samplings will allow us to perform a more in-depth analysis of JacSketch using a stochastic Lyapunov function. Unlike Theorem 3.6, and as explained in Section 1.7, our main result in this section (Theorem 5.2) is capable of obtaining the conjectured rate $O((n+\tfrac{\bar{L}}{\mu})\log\tfrac{1}{\epsilon})$ for SAGA with importance sampling.

One of the key reasons why we restrict our attention to $\tau$-partition samplings is the fact that

for $C_{1},C_{2}\in{\cal G}$. Recall from Lemma 4.7 that if ${\bf W}={\bf I}$, then ${\bf\Pi}_{{\bf I}_{C}}={\bf I}_{C}{\bf I}_{C}^{\top}$. Consequently, for $C_{1},C_{2}\in{\cal G}$ we have

This orthogonality property will be fundamental for controlling the convergence of the gradient estimate in Lemma 5.3.

Recall from (32) that the stochastic gradient of the controlled stochastic reformulation (28) of the original finite-sum problem (1) is given by

provided that we use the minibatch sketch $\mathbf{S}={\bf I}_{S}$ and bias-correcting variable $\theta_{\mathbf{S}}=\theta_{{\bf I}_{S}}=1/p_{S}$ given by Lemma 4.7(vi). This object will appear in our Lyapunov function, evaluated at $x=x^{*}$ and ${\bf J}={\bf J}^{k}$. We are now ready to present the main result of this section.

$\mathbf{S}$ be a minibatch sketch (i.e., $\mathbf{S}={\bf I}_{S}$)We can alternatively set $\mathbf{S}=e_{S}$ and the same results will hold. , where $S$ is a $\tau$–partition sampling with support ${\cal G}={\rm supp}(S)$,

$f_{C}\overset{\text{def}}{=}\tfrac{1}{|C|}\sum_{i\in C}f_{i}$ be $L_{C}$–smooth and $\mu$–strongly convex (for $\mu>0$) for all $C\in{\cal G}$,