Optimal mini-batch and step sizes for SAGA

Nidham Gazagnadou, Robert M. Gower, Joseph Salmon

Introduction

Consider the empirical risk minimization (ERM) problem:

where each fif_{i} is LiL_{i}-smooth and ff is μ\mu-strongly convex. Each fif_{i} represents a regularized loss over a sampled data point. Solving the ERM problem is often time consuming for large number of samples nn, so much so that algorithms scanning through all the data points at each iteration are not competitive. Gradient descent (GD) falls into this category, and in practice its stochastic version is preferred.

Stochastic gradient descent (SGD), on the other hand, allows to solve the ERM incrementally by computing at each iteration an unbiased estimate of the full gradient, ∇fi(wk)\nabla f_{i}(w^{k}) for ii randomly sampled in [n]:={1,2,…,n}[n]:=\{1,2,\ldots,n\} (Robbins & Monro, 1951). On the downside, for SGD to converge one needs to tune a sequence of asymptotically vanishing step sizes, a cumbersome and time-consuming task for the user. Recent works have taken advantage of the sum structure in Eq. 1 to design stochastic variance reduced gradient algorithms (Johnson & Zhang, 2013; Shalev-Shwartz & Zhang, 2013; Defazio et al., 2014; Schmidt et al., 2017). In the strongly convex setting, these methods lead to fast linear convergence instead of the slow O(1/t){\cal O}(1/t) rate of SGD. Moreover, they only require a constant step size, informed by theory, instead of sequence of decreasing step sizes.

In practice, most variance reduced methods rely on a mini-batching strategy for better performance. Yet most convergence analysis (with the Katyusha algorithm of Allen-Zhu (2017) being an exception) indicates that a mini-batch size of b=1b=1 gives the best overall complexity, disagreeing with practical findings, where larger mini-batch often gives better results. Here, we show both theoretically and numerically that b=1b=1 is not the optimal mini-batch size for the SAGA algorithm (Defazio et al., 2014).

Our analysis leverages recent results in (Gower et al., 2018), where the authors prove that the iteration complexity and the step size of SAGA, and a larger family of methods called the JacSketch methods, depend on an expected smoothness constant. This constant governs the trade-off between the increased cost of an iteration as the mini-batch size is increased, and the decreased total complexity. Thus if this expected smoothness constant could be calculated a priori, then we could set the optimal mini-batch size and step size. We provide simple formulas for computing the expected smoothness constant when sampling mini-batches without replacement, and use them to calculate optimal mini-batches and significantly larger step sizes for SAGA.

In particular, we provide two bounds on the expected smoothness constant, each resulting in a particular step size formula. We first derive the simple bound and then develop a matrix concentration inquality to obtain the refined Bernstein bound. We also provide substantial theoretical motivation and numerical evidence for practical estimate of the expected smoothness constant. For illustration, we plot in Figure 1 the evolution of each resulting step size as the mini-batch size grows on a classification problem (Section 5 has more details on our experimental settings).

with Lmax⁡:=max⁡i∈[n]Li{L_{\max}}:=\max_{i\in[n]}L_{i}, Lˉ:=1n∑i=1nLi\bar{L}:=\frac{1}{n}\sum_{i=1}^{n}L_{i} and ϵ>0\epsilon>0 is the desired precision.

This complexity bound, and others presented in Section 3.3 show that SAGA enjoys a linear speedup as we increase the mini-batch size until an optimal one (as illustrated in Figure 2). After this point, the total complexity increases. We use this observation to develop optimal and practical mini-batch sizes and step sizes.

The rest of the paper is structured as follows. In Section 2 we first introduce variance reduction techniques after presenting our main assumption, the expected smoothnes assumption. We highlight how this assumption is necessary to capture the improvement in iteration complexity, and conclude the section by showing that to calculate the expected smoothness constant we need evaluate an intractable expectation. Which brings us to Section 3 where we directly address this issue and provide several tractable upper-bounds of the expected smoothness constant. We then calculate optimal mini-batch sizes and step sizes by using our new bounds. Finally, we give numerical experiments in Section 5 that verify our theory on artificial and real datasets. We also show how these new settings for the mini-batch size and step size lead to practical performance gains.

Background

We can introduce variance reduced versions of SGD in a principled manner by using a sampling vector.

With a sampling vector we can re-write (1) through the following stochastic reformulation

where fv(w)f_{v}(w) is called a subsampled function. The stochastic Problem (2) and our original Problem (1) are equivalent :

Consequently the gradient ∇fv(w)\nabla f_{v}(w) is an unbiased estimate of ∇f(w)\nabla f(w) and we could use the SGD method to solve (2). To tackle the variance of these stochastic gradients we can further modify (2) by introducing control variates which leads to the following controlled stochastic reformulation:

That is, starting from a vector w0w^{0}, given a positive step size γ\gamma, we can iterate the steps

where vk∼Dv^{k}\sim{\cal D} are i.i.d. samples at each iteration.

2 The expected smoothness constant

In order to analyze stochastic variance reduced methods, some form of smoothness assumption needs to be made. The most common assumption is

for each i∈[n]i\in[n]. That is each fif_{i} is uniformly smooth with smoothness constant Lmax⁡L_{\max}, as is assumed in (Defazio et al., 2014; Hofmann et al., 2015; Raj & Stich, 2018) for variants of SAGAThe same assumption is made in proofs of SVRG (Johnson & Zhang, 2013), S2GD (Konečný & Richtárik, 2017) and the SARAH algorithm (Nguyen et al., 2017).. In the analyses of these papers it was shown that the iteration complexity of SAGA is proportional to Lmax⁡,L_{\max}, and the step size is inversely proportional to Lmax⁡.L_{\max}.

But as was shown in (Gower et al., 2018), we can set a much larger step size by making use of the smoothness of the subsampled functions fv.f_{v}. For this Gower et al. (2018) introduced the notion of expected smoothness, which we extend here to all sampling vectors and control variates.

Note that we refer to any positive constant L{\cal L} that satisfies (7) as an expected smoothness constant. Indeed L→∞{\cal L}\rightarrow\infty is a valid constant in the extended reals, but as we will see, the smaller L{\cal L}, the better for our complexity results.

Gower et al. (2018) show that the expected smoothness constant plays the same role that Lmax⁡L_{\max} does in the previously existing analysis of SAGA, namely that the step size is inversely proportional to L{\cal L} and the iteration complexity is proportional to L{\cal L} (see details in Theorem 1). Furthermore, by assuming that ff is LL–smooth, the expected smoothness constant is bounded

as was proven in Theorem 4.17 in (Gower et al., 2018). Also, the bounds Lmax⁡L_{\max} and LL are attained when using a uniform single element sampling and a full batch, respectively. And as we will show, the constants Lmax⁡L_{\max} and LL can be orders of magnitude apart on large dimensional problems. Thus we could set much larger step sizes for larger mini-batch sizes if we could calculate L{\cal L}. Though calculating L{\cal L} is not easy, as we see in the next lemma.

If the sampling has a very large combinatorial number of possible realizations — for instance sampling mini-batches without replacement — then this expectation becomes intractable to calculate. This observation motivates the development of functional upper-bounds of the expected smoothness constant that can be efficiently evaluated.

3 Mini-batch without replacement: b𝑏b–nice sampling

Now we will choose a distribution of the sampling vector vv based on a mini-batch sampling without replacement. We denote a mini-batch as B⊆[n]B\subseteq[n] and its size as b=∣B∣b=|B|.

SS is a bb-nice sampling if SS is a set valued map with a probability distribution given by

where \mathds1S\mathds{1}_{S} denotes the indicator function of the random set SS. Now taking expectation in (9) gives

using ∣{B⊆[n]:∣B∣=b∧i∈B}∣=(n−1b−1)|\{{{B\subseteq[n]:|B|=b\wedge i\in B}}\}|=\binom{n-1}{b-1}.

Here we are interested in the mini-batch SAGA algorithm with bb-nice sampling, which we refer to as the bb-nice SAGA. In particular, bb-nice SAGA is the result of using bb-nice sampling, together with a linear model for the control variate zv(w)z_{v}(w). Different choices of the control variate zv(w)z_{v}(w) also recover popular algorithms such as gradient descent, SGD or the standard SAGA method (see Table 1 for some examples).

A naive implementation of bb-nice SAGA based on the JacSketch algorithm is given in Algorithm 1We also provide a more efficient implementation that we used for our experiments in the appendix in Algorithm 2..

Upper Bounds on the Expected Smoothness

To determine an optimal mini-batch size b∗b^{*} for bb-nice SAGA, we first state our assumptions and provide bounds of the smoothness of the subsampled function. We then define b∗b^{*} as the mini-batch size that minimizes the total complexity of the considered algorithm, i.e., the total number of stochastic gradients computed. Finally we provide upper-bounds on the expected smoothness constant L{\cal L}, through which we can deduce optimal mini-batch sizes. Many proofs are deferred to the supplementary material.

We consider that the objective function is a GLM with quadratic regularization controlled by a parameter λ>0\lambda>0:

We assume that the second derivative of each ϕi\phi_{i} is uniformly bounded, which holds for our aforementioned examples.

For a batch B⊆[n]B\subseteq[n], we rewrite the subsampled function as

and its second derivative is thus given by

where IdI_{d} denotes the identity matrix of size dd.

For a symmetric matrix MM, we write λmax⁡(M)\lambda_{\max}(M) (resp. λmin⁡(M)\lambda_{\min}(M)) for its largest (resp. smallest) eigenvalue. Assumption 1 directly implies the following.

Let B⊂[n]B\subset[n], and let AB=[ai]i∈BA_{B}=[a_{i}]_{i\in B} denote the column concatenation of the vectors aia_{i} with i∈B.i\in B. The smoothness constant of the subsampled loss function 1∣B∣∑i∈Bϕi(ai⊤w)\frac{1}{|B|}\sum_{i\in B}\phi_{i}(a_{i}^{\top}w) is given by

Proof. The proof follows from Assumption 1 as

Combined with (11), we get that fBf_{B} is (LB+λ)(L_{B}+\lambda)-smooth.

Another key quantity in our analysis is the strong convexity parameter.

The strong convexity parameter is given by

Since we have an explicit regularization term with λ>0\lambda>0, ff is strongly convex and μ≥λ>0.\mu\geq\lambda>0.

We additionally define LiL_{i}, resp. LL, as the smoothness constant of the individual function ϕi(ai⊤w)\phi_{i}(a_{i}^{\top}w), resp. the whole function 1n∑i=1nϕi(ai⊤w)\frac{1}{n}\sum_{i=1}^{n}\phi_{i}(a_{i}^{\top}w). We also recall the definitions of the maximum of the individual smoothness constants by Lmax⁡:=max⁡i∈[n]Li{L_{\max}}:=\max_{i\in[n]}L_{i} and their average by Lˉ:=1n∑i=1nLi\bar{L}:=\frac{1}{n}\sum_{i=1}^{n}L_{i}. The three constants satisfies

The proof of (13) is given in Lemma 10 in the appendix.

2 Path to the optimal mini-batch size

Our starting point is the following theorem taken from combining Theorem 3.6 and Eq. (103) in (Gower et al., 2018)Note that λ\lambda has been added to every smoothness constant since the analysis in Gower et al. (2018) depends on the (L+λ)(L+\lambda)-smoothness of ff and the (LB+λ)(L_{B}+\lambda)-smoothness of the subsampled functions fBf_{B}..

Consider the iterates wkw^{k} of Algorithm 1. Let the step size be given by

Through Theorem 1 we can now explicitly see how the expected smoothness constant L{\cal L} controls both the step size and the resulting iteration complexity. This is why we need bounds on L{\cal L} so that we can set the step size. In particular, we will show that the expected smoothness constant is a function of the mini-batch size bb. Consequently so is the step size, the iteration complexity and the total complexity. We denote KtotalK_{\text{total}} the total complexity defined as the number of stochastic gradients computed, hence with (15),

As we have shown in Lemma 1, computing a precise bound on L{\cal L} can be computationally intractable. This is why we focus on finding upper bounds on L{\cal L} that can be computed, but also tight enough to be useful. To verify that our bounds are sufficiently tight, we will always have in mind the bounds L≤L≤Lmax⁡L\leq{\cal L}\leq L_{\max} given in (8). In particular, after expressing our bounds of L=L(b){\cal L}={\cal L}(b) as a function of bb,we would like the bounds (8) to be attained for L(1)=Lmax⁡{\cal L}(1)=L_{\max} and L(n)=L.{\cal L}(n)=L.

3 Expected smoothness

All bounds we develop on L{\cal L} are based on the following lemma, which is a specialization of (1) for bb-nice sampling.

For the bb-nice sampling, with b∈[n]b\in[n], the expected smoothness constant is given by

Let SS the bb-nice sampling as defined in Definition 3 and let v=nb∑j∈Sejv=\frac{n}{b}\sum_{j\in S}e_{j} be its corresponding sampling vector. Note that

Taking the maximum over all i∈[n]i\in[n] gives the result. ∎

The first bound we present is technically the simplest to derive, which is why we refer to it as the simple bound.

For a bb-nice sampling SS, for b∈[n]b\in[n], we have that

The proof, given in Section A.2, starts by using the that LB≤1b∑j∈BLjL_{B}\leq\frac{1}{b}\sum_{j\in B}L_{j} for all subsets BB, which follows from repeatedly applying Lemma 8 in the appendix. The remainder of the proof follows by straightforward counting arguments. ∎

The previous bound interpolates, respectively for b=1b=1 and b=nb=n, between Lmax⁡L_{\max} and Lˉ.\bar{L}. On the one hand, we have that Lsimple⁡(b){\cal L}_{\operatorname{simple}}(b) is a good bound for when bb is small, since Lsimple⁡(1)=Lmax⁡{\cal L}_{\operatorname{simple}}(1)=L_{\max}. Though Lsimple⁡(b){\cal L}_{\operatorname{simple}}(b) may not be a good bound for large bb, since Lsimple⁡(n)=L‾≥L{\cal L}_{\operatorname{simple}}(n)=\overline{L}\quad{\geq}\quad L, thanks to (13). Thus Lsimple⁡(b){\cal L}_{\operatorname{simple}}(b) does not achieve the left-hand side of (8). Indeed L‾\overline{L} can be far from LL. For instanceWe numerically explore such extreme settings in Section 5., if f(w)=1n∑i∈[n]12(ai⊤w−bi)2f(w)=\frac{1}{n}\sum_{i\in[n]}\frac{1}{2}(a_{i}^{\top}w-b_{i})^{2} is a quadratic function, then we have that L‾=1n\mboxTr(AA⊤)\overline{L}=\frac{1}{n}\mbox{Tr}\left(AA^{\top}\right) and L=1nλmax⁡(AA⊤)L=\frac{1}{n}\lambda_{\max}(AA^{\top}). Thus if the eigenvalues of AA⊤AA^{\top} are all equal then L‾=dL\overline{L}=dL. Alternatively, if one eigenvalue is significantly larger than the rest then L‾≈L\overline{L}\approx L.

Due to this shortcoming of Lsimple⁡,{\cal L}_{\operatorname{simple}}, we now derive the Bernstein bound. This bound explicitly depends on LL instead of Lˉ\bar{L}, and is developed through a specialized variant of a matrix Bernstein inequality (Tropp, 2012, 2015) for sampling without replacement in Appendix C.

The expected smoothness constant is upper bounded by

Checking again the bounds of LBernstein(b){\cal L}_{\text{Bernstein}}(b), we have on the one hand that LBernstein(1)=(1+43log⁡d)Lmax⁡≥Lmax⁡,{\cal L}_{\text{Bernstein}}(1)=\left(1+\tfrac{4}{3}\log d\right)L_{\max}\geq L_{\max}, thus there is a little bit of slack for bb small. On the other hand, using 1nLmax⁡≤L\tfrac{1}{n}L_{\max}\leq L (see Lemma 10 in appendix), we have that

which depends only logarithmically on dd. Thus we expect the Bernstein bound to be more useful in the large dd domains, as compared to the simple bound. We confirm this numerically in Section 5.1.

The simple bound is relatively tight for bb small, while the Bernstein bound is better for large bb and large dd. Fortunately, we can obtain a more refined bound by taking the minimum of the simple and the Bernstein bounds. This is highlighted numerically in Section 5.

Next we propose a practical estimate of L{\cal L} that is tight for both small and large mini-batch sizes.

Indeed Lpractical(1)=Lmax⁡{\cal L}_{\text{practical}}(1)=L_{\max} and Lpractical(n)=L,{\cal L}_{\text{practical}}(n)=L, achieving both limits of (8). The downside to Lpractical(b){\cal L}_{\text{practical}}(b) is that it is not an upper bound of L{\cal L}. Rather, we are able to show that Lpractical(b){\cal L}_{\text{practical}}(b) is very close to a valid smoothness constant, but it can be slightly smaller. Our theoretical justification for using Lpractical(b){\cal L}_{\text{practical}}(b) comes from a mid step in the proof of the Bernstein bound which is captured in the next lemma.

with Ni:=1b∑j∈Siajaj⊤−1bb−1n−1∑j∈[n]∖{i}ajaj⊤N_{i}:=\frac{1}{b}\sum_{j\in S^{i}}a_{j}a_{j}^{\top}-\frac{1}{b}\frac{b-1}{n-1}\sum_{j\in[n]\setminus\{i\}}a_{j}a_{j}^{\top}.

Lemma 3 shows that the expected smoothness constant is upper-bounded by Lpractical(b){\cal L}_{\text{practical}}(b) and an additional term. In this additional term we have the largest eigenvalue of a random matrix. This matrix is zero in expectation, and we also find that its eigenvalues oscillate around zero. Indeed, we provide extensive experiments in Section 5 confirming that Lpractical(b){\cal L}_{\text{practical}}(b) is very close to L{\cal L} given in (17).

Optimal Mini-Batch Sizes

Now that we have established the simple and the Bernstein bounds, we can minimize the total complexity (16) in the mini-batch size.

For instance for the simple bound, given ϵ>0\epsilon>0 and plugging in (18) into (16) gives

where gsimple(b):=4(nLˉ−Lmax⁡+(n−1)λ)μ(n−1)b+4n(Lmax⁡−Lˉ)μ(n−1),g_{\text{simple}}(b):=\tfrac{4\left(n\bar{L}-L_{\max}+(n-1)\lambda\right)}{\mu(n-1)}b+\tfrac{4n\left(L_{\max}-\bar{L}\right)}{\mu(n-1)}, and h(b):=−4(Lmax⁡+λ)μ(n−1)b+n(1+1n−14(Lmax⁡+λ)μ).h(b):=-\tfrac{4(L_{\max}+\lambda)}{\mu(n-1)}b+n\left(1+\tfrac{1}{n-1}\tfrac{4(L_{\max}+\lambda)}{\mu}\right).

The right-hand side term h(b)h(b) is common to all our bounds since it does not depend on L{\cal L}. It linearly decreases from h(1)=n+4(Lmax⁡+λ)μh(1)=n+\frac{4(L_{\max}+\lambda)}{\mu} to h(n)=nh(n)=n.

We note that gsimple(b)g_{\text{simple}}(b) is a linearly increasing function of bb, because Lmax⁡≤nLˉL_{\max}\leq n\bar{L} (as proven in Lemma 10). One can easily verify that gsimple(b)g_{\text{simple}}(b) and h(b)h(b) cross, as presented in Figure 2, by looking at initial and final values:

At b ⁣= ⁣1b\!=\!1, gsimple(1)=4μ(Lmax⁡+λ)=h(1)−ng_{\text{simple}}(1)=\tfrac{4}{\mu}(L_{\max}+\lambda)=h(1)-n. So, gsimple(1)≤h(1)g_{\text{simple}}(1)\leq h(1).

At b ⁣= ⁣nb\!=\!n, gsimple(n)=4(Lˉ+λ)μn=4(Lˉ+λ)μh(n)g_{\text{simple}}(n)=\tfrac{4(\bar{L}+\lambda)}{\mu}n=\tfrac{4(\bar{L}+\lambda)}{\mu}h(n). Since Lˉ≥μ\bar{L}\geq\mu, we get gsimple(n)≥h(n)g_{\text{simple}}(n)\geq h(n).

Consequently, solving gsimple(b)=h(b)g_{\text{simple}}(b)=h(b) in bb gives the optimal mini-batch size

For the Bernstein bound, plugging (19) into (16) leads to

The function gBernsteing_{\text{Bernstein}} is also linearly increasing in bb and its initial and final values are

At b ⁣= ⁣1b\!=\!1, gBernstein(1)=(1+43log⁡d)4Lmax⁡μ+4λμ.g_{\text{Bernstein}}(1)=(1+\frac{4}{3}\log d)\tfrac{4L_{\max}}{\mu}+\tfrac{4\lambda}{\mu}.

At b ⁣= ⁣nb\!=\!n, gBernstein(n) ⁣= ⁣n4(2L+λ)μ ⁣+ ⁣163μ(Lmax⁡ ⁣+ ⁣λ)log⁡(d)g_{\text{Bernstein}}(n)\!=\!n\tfrac{4(2L+\lambda)}{\mu}\!+\!\tfrac{16}{3\mu}{(L_{\max}\!+\!\lambda)}\log(d). Since L≥μL\geq\mu, we get gBernstein(n)≥h(n)g_{\text{Bernstein}}(n)\geq h(n).

Yet, it is unclear whether gBernstein(1)g_{\text{Bernstein}}(1) is dominated by h(1)h(1). This is why we need to distinguish two cases to minimize the total complexity, which leads to the following solution

Numerical Study

All the experiments were run in Julia and the code is freely available on https://github.com/gowerrobert/StochOpt.jl.

In Figure 4 we see that Lpractical{\cal L}_{\text{practical}} is arbitrarily close to L{\cal L}, making it hard to distinguish the two line plots. This was the case in many other experiments, which we defer to Section E.1. For this reason, we use γpractical\gamma_{\text{practical}} in our experiments with the SAGA method.

Furthermore, in accordance with our discussion in Section 3.3, we have that Lsimple{\cal L}_{\text{simple}} and LBernstein{\cal L}_{\text{Bernstein}} are close to L{\cal L} when bb is small and large, respectively. In Section E.2 we show, by applying ridge and regularized logistic regression to publicly available datasets from LIBSVMhttps://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/ and the UCI repositoryhttps://archive.ics.uci.edu/ml/datasets/, that the simple bound performs better than the Bernstein bound when n≫dn\gg d, and conversely for dd slightly smaller than nn, larger than nn or when scaling the data.

2 Related step size estimation

Different bounds on L{\cal L} also give different step sizes (14). Plugging in our estimates Lsimple{\cal L}_{\text{simple}}, LBernstein{\cal L}_{\text{Bernstein}} and Lpractical{\cal L}_{\text{practical}} into (14) gives the step sizes γsimple\gamma_{\text{simple}}, γBernstein\gamma_{\text{Bernstein}} and γpractical\gamma_{\text{practical}}, respectively. We compare our resulting step sizes to γL\gamma_{{\cal L}} where L{\cal L} is given by Eq. 17 and to the step size given by Hofmann et al. (2015), which is γHofmann(b)=K2Lmax⁡(1+K+1+K2)\gamma_{\text{Hofmann}}(b)=\frac{K}{2L_{\max}(1+K+\sqrt{1+K^{2}})}, where K:=4bLmax⁡nμ.K:=\frac{4bL_{\max}}{n\mu}. We can see in Figure 4, that for b=1b=1, all the step sizes are approximately the same, with the exceptions of the Bernstein step size. For b>5b>5, all of our step sizes are larger than γHofmann(b)\gamma_{\text{Hofmann}}(b), in particular γpractical(b)\gamma_{\text{practical}}(b) is significantly larger. These observations are verified in other artificial and real data examples in Sections E.3 and E.4.

3 Comparison with previous SAGA settings

Here we compare the performance of SAGA when using the mini-batch size and step size b=1,γDefazio:=1/3(nμ+Lmax⁡)b=1,\gamma_{\text{Defazio}}:=1/{3(n\mu+L_{\max})} given in (Defazio et al., 2014), b=20b=20 and γHofmann=20/nμ\gamma_{\text{Hofmann}}={20}/{n\mu} given in Hofmann et al. (2015), to our new practical mini-batch size bpractical=⌊1+μ(n−1)4(L+λ)⌋b_{\text{practical}}=\left\lfloor 1+\frac{\mu(n-1)}{4(L+\lambda)}\right\rfloor and step size γpractical\gamma_{\text{practical}}. Our goal is to verify how much our parameter setting can improve practical performance. We also compare with a step size γgridsearch\gamma_{\text{gridsearch}} obtained by grid search over odd powers of 22. These methods are run until they reach a relative error of 10−410^{-4}.

We find in Figure 5 that our parameter settings (γpractical,bpractical)(\gamma_{\text{practical}},b_{\text{practical}}) significantly outperforms the previously suggested parameters, and is even comparable to grid search. Finally, In we show in Section E.5 that the settings (γHofmann,b=20)(\gamma_{\text{Hofmann}},b=20) can lead to very poor performance compared to our settings.

4 Optimality of our mini-batch size

In the last experiment, detailed in Section E.6, we show that our estimation of the optimal mini-batch size bpracticalb_{\text{practical}} leads to a faster implementation of SAGA. We build a grid of mini-batch sizesOur grid is {2i,i=0,…,14}\{2^{i},i=0,\dots,14\}, with 216,2182^{16},2^{18} and nn being added when needed. and compute the empirical total complexity required to achieve a relative error of 10−410^{-4}, as in Section 5.3.

In Figure 6 we can see that the empirical complexity when using bpracticalb_{\text{practical}} is almost the same as one for the optimal mini-batch size calculated through grid search. Yet, our bpracticalb_{\text{practical}} being always larger than the mini-batch obtained by grid search, it leads to a faster algorithm for two reasons. Firstly, the corresponding step size is larger and secondly, and secondly, computing the stochastic gradients in parallel improves the running time. What is even more interesting, is that bpracticalb_{\text{practical}} always predicts a regime change, where using a larger mini-batch size results in a much larger empirical complexity.

Conclusions

We have explained the crucial role of the expected smoothness constant L{\cal L} in the convergence of a family of stochastic variance-reduced descent algorithms. We have developed functional upper-bounds of this constant to build larger step sizes and closed-form optimal mini-batch values for the bb-nice SAGA algorithm. Our experiments on artificial and real datasets showed the validity of our upper-bounds and the improvement in the total complexity using our step and optimal mini-batch sizes. Our results suggest a new parameter setting for mini-batch SAGA, that significantly outperforms previous suggested ones, and is even comparable with a grid search approach, without the computational burden of the later.

Acknowledgements

This work was supported by grants from DIM Math Innov Région Ile-de-France (ED574 - FMJH) and by a public grant as part of the Investissement d’avenir project, reference ANR-11-LABX-0056-LMH, LabEx LMH, in a joint call with Gaspard Monge Program for optimization, operations research and their interactions with data sciences.

References

Appendix A Proofs of the Upper Bounds of ℒℒ{\cal L}

Since the fif_{i}’s are convex, each realization of fvf_{v} is convex, and it follows from equation 2.1.7 in (Nesterov, 2014) that

Taking expectation over the sampling gives

where in the last equality the full gradient vanishes because it is computed at optimality. The result now follows by comparing the above with the definition of expected smoothness in (7). ∎

A.2 Proof of the simple bound

To derive this bound on L{\cal L} we use that

which follows from repeatedly applying Lemma 8. For b≥2b\geq 2, it follows from Equation 17 and Equation 25 that

Using a double counting argument we can show that

We also verify that this bound is valid for 11-nice sampling. Indeed, we already have that in this case L=Lmax⁡{\cal L}=L_{\max}. ∎

A.3 Proof of the Bernstein bound

To start the proof of Theorem 3, we re-write the expected smoothness constant as the maximum over an expectation. Let SiS^{i} be a (b−1)(b-1)-nice sampling over [n]∖{i}.[n]\setminus\{i\}. We can write

One can come back to the definition of the subsample smoothness constant Equation 12 and interpret previous expression as an expectation of the largest eigenvalue of a sum of matrices. This insight allows us to apply a matrix Bernstein inequality, see Theorem 7, to bound L{\cal L}.

For the proof of Theorem 3, we first need the two following results.

This results follows using a double-counting argument at the fourth line of the computation.

We then introduce another two lemmas which give a first intermediate bound.

where in the first inequality we add and remove the mean and then apply Lemma 8. In the second equality we explicit the mean with Lemma 4 and in the last inequality we use again Lemma 8 for the left-hand side term. Finally, we multiply by UU on both sides of the inequality. ∎

We recall the following lemma used to introduced the practical estimate given by

The result comes from applying re-writing L{\cal L} as an expectation of the largest eigenvalue of a sum of matrices. Then we apply Lemma 5 and then taking the maximum over all i∈[n]i\in[n]. Thus, we have

with N:=1b∑j∈Siajaj⊤−1bb−1n−1∑j∈[n]∖{i}ajaj⊤N:=\frac{1}{b}\sum_{j\in S^{i}}a_{j}a_{j}^{\top}-\frac{1}{b}\frac{b-1}{n-1}\sum_{j\in[n]\setminus\{i\}}a_{j}a_{j}^{\top}.

Using this notation, the matrix NN which can be further decomposed as

where we have encoded the sampling SiS^{i} using unit coordinate vectors. The matrices M1,…,Mb−1M_{1},\ldots,M_{b-1} are sampled without replacement from the set

Now let X1,…,XbX_{1},\ldots,X_{b} be matrices sampled with replacement from (34) and let Xk:=1b∑j∈[n]∖{i}(zjk−1n−1)ajaj⊤X_{k}:=\frac{1}{b}\sum_{j\in[n]\setminus\{i\}}\left(z_{j}^{k}-\frac{1}{n-1}\right)a_{j}a_{j}^{\top} and Y:=∑k=1b−1XkY:=\sum_{k=1}^{b-1}X_{k} thus the vectors zkz^{k} are sampled with replacement from {e1,…,ei−1,ei+1,…,en}.\{e_{1},\ldots,e_{i-1},e_{i+1},\ldots,e_{n}\}. Consequently

We are now in a position to apply the Bernstein matrix inequality. To this end we have

Let k∗k^{*} be the unique index such that zk∗k=1.z^{k}_{k^{*}}=1. We have a uniform bound of the largest eigenvalue of our XkX_{k}

where we applied the Lemma 9 in the first inequality.

Summing in (36), taking the largest eigenvalue and applying Lemma 9 results in

Considering Equations 35 and 37 and applying the matrix Bernstein concentration inequality in Theorem 7 we get

Taking the maximum over ii and using L[n]∖{i}≤nn−1LL_{[n]\setminus\{i\}}\leq\frac{n}{n-1}L we have that

Combining the above result with (32) leads us to

where in the second inequality we used the inequality 2ab≤a+b\sqrt{2ab}\leq a+b. ∎

Appendix B Linear Algebra Tools

This appendix is dedicated to the presentation of useful results to manipulate more easily the smoothness constants.

Let us recall some useful spectral results on Hermitian and positive semi-definite matrices.

whenever i,j≥1i,j\geq 1 and i+j−1≤n.i+j-1\leq n\enspace.

Moreover, as a direct consequence of the variational characterization of eigenvalues, namely

we have an inequality between the maximum diagonal term of a positive semi-definite matrices and its maximum eigenvalue.

The following lemma is a direct consequence of Weyl’s inequality for i=j=1.i=j=1.

Lastly, we present a result arising from previous lemma.

where the first inequality stems from Lemma 8 and the second from B⪰0.B\succeq 0. ∎

B.2 Basic properties of the smoothness constants

The complexity results of Gower et al. (2018) depends on smoothness constants defined in Section 3.1. Here are some inequalities giving an idea of the order of those constants.

Let ∅≠B⊆[n]={1,…,n}\emptyset\neq B\subseteq[n]=\{1,\dots,n\} a batch set drawn randomly without replacement. The following inequalities hold

One directly gets that Li≤max⁡j=1,…,nLj=Lmax⁡L_{i}\leq\max_{j=1,\ldots,n}L_{j}={L_{\max}}.

This inequality states that the smoothness constant LBL_{B} of the averaged function fBf_{B} is upper bounded by the average of the corresponding smoothness constants LiL_{i}, over the batch BB. The proof consists in ∣B∣|B| repetitive calls of Lemma 8.

(a)(a) Direct implication of (ii) for B=[n]B=[n].

(c)(c) Let us first recall the matrix formulation of our smoothness constants:

Dividing the above by nn on both sides gives

Appendix C Matrix Bernstein Inequality: Sampling Without Replacement

In this appendix, we present the matrix Bernstein inequality for independent Hermitian matrices from Tropp (2015). We also provide another version of this theorem for matrices sampled without replacement and prove it as explicitly as possible, taking our inspiration from Tropp (2011). The proof is based the possibility of transferring the results from sampling with to without through the inequality (50) due to Gross & Nesme (2010). The exact same work can be done for the tail bound, which is for instance used in Bach (2013).

We first present Theorem 4 which gives a Bernstein inequality for a sum of random and independent Hermitian matrices whose eigenvalues are upper bounded. If the matrices XkX_{k} are sampled from a finite set X{\cal X}, one can interpret this random sampling of independent matrices as a random sampling with replacement.

Consider a finite sequence {Xk}k=1,…,n\{X_{k}\}_{k=1,\dots,n} of nn independent, random, Hermitian matrices with dimension dd. Assume that

Let v(SX)v(S_{X}) be the matrix variance statistic of the sum:

This theorem is the one we extend in Theorem 7 to the case when the random matrices XkX_{k} are sampled without replacement from a finite set X{\cal X}. We drew our inspiration from the proof of the matrix Chernoff inequality in Tropp (2011) and the one of the matrix Bernstein tail bound in Bach (2013), both in the case of sampling without replacement.

C.2 Technical random matrices prerequisites

Before proving Theorem 7, which extends the matrix Bernstein inequality to sampling without replacement, we need to introduce the key tools of the matrix Laplace transform technique. This technique is precious to prove tail bounds for sums of random matrices such as Chernoff, Hoeffding or Bernstein bounds, as presented in (Tropp, 2012).

Here, ∥⋅∥\left\|\cdot\right\| denotes the spectral norm, which is defined for any Hermitian matrix HH by

We also introduce the moment generating function (mgf) and the cumulant generating function (cgf) of a random matrix, which are essential in the Laplace transform method approach.

These expectations may not exist for all values of θ\theta.

Let XX be a random Hermitian matrix. Then

This proposition is an adaptation of the Laplace transform method to obtain a bound of the expectation of the maximum eigenvalue of a random Hermitian matrix. Contrary to the tail bounds, there is no exact analog of the expectation bounds in the scalar setting.

Fix a positive number θ\theta. Because λmax⁡(⋅)\lambda_{\max}(\cdot) is a positive-homogeneous map, we have

where in the third line we used the Jensen’s inequality, in the fourth one the spectral mapping theorem and in the last line the domination by the trace of a positive-definite matrix. ∎

Let HH be a fixed Hermitian matrix with dimension dd. The function

is a concave map on the the convex cone of d×dd\times d positive-definite matrices.

Let HH be a fixed Hermitian matrix with dimension dd. Let XX be a random Hermitian matrix of same dimension. The following inequality holds

is a concave map on the the convex cone of d×dd\times d positive-definite matrices.

where the inequality comes from the application of Theorem 5 and Jensen’s inequality. ∎

Let XX a random Hermitian matrix such that

C.3 Extended results for sampling without replacement

This section is dedicated to the main result, Lemma 13, needed for transferring results from sampling with to without replacement. This lemma is actually the matrix version of a classical result from Hoeffding (1963). We then combine it with previous results of Section C.2 to produce a new master bound in Theorem 6, which is the key inequality of the proof of Theorem 7.

The left-hand side equality directly arises from Definition 6 and the fact that the trace commutes with the expectation because it is a linear operator. For the right-hand side inequality, see the proof in Gross & Nesme (2010). ∎

Consider two finite sequences, of same length nn, {Xk}k=1,…,n\{X_{k}\}_{k=1,\dots,n} and {Yk}k=1,…,n\{Y_{k}\}_{k=1,\dots,n} of Hermitian random matrices of same size sampled respectively with and without replacement from a finite set X{\cal X}. Then

This theorem is a modified version of Theorem 3.6.1 in Tropp (2015) for a sum of matrices sampled without replacement.

Consider two finite sequences, of same length, {Xk}\{X_{k}\} and {Yk}\{Y_{k}\} of Hermitian random matrices of same size sampled respectively with and without replacement from a finite set X{\cal X}. Let θ\theta a positive number.

where we used successively Proposition 2, Lemma 13 and Lemma 11. First, we use the expectation bound for the maximum eigenvalue. We then use the main result of Gross & Nesme (2010) and invoked in Tropp (2011) to extend the matrix Chernoff bound for matrices sampled without replacement. This lemma allows us to transfer our results to sampling with replacement. And finally, we then apply the subadditivity of matrix cgfs to get the desired result. ∎

C.4 Bernstein inequality for sampling without replacement

The following theorem is almost the same than Theorem 4, but in the case of matrices sampled without replacement from a finite set. The proof stems from results established in previous Sections C.3 and C.2.

Let X{\cal X} be a finite set of Hermitian matrices with dimension dd such that

Sample two finite sequences, of same length nn, {Xk}k=1,…,n\{X_{k}\}_{k=1,\dots,n} and {Yk}k=1,…,n\{Y_{k}\}_{k=1,\dots,n} uniformly at random from X{\cal X} respectively with and without replacement such that

Let v(SX)v(S_{X}) be the matrix variance statistic of the second sum

Consider X{\cal X} a finite set of Hermitian matrices of dimension dd such that

Sample two finite sequences, of same length, {Xk}\{X_{k}\} and {Yk}\{Y_{k}\} uniformly at random from X{\cal X} respectively with and without replacement such that

The {Xk}\{X_{k}\} matrices are thus independent. Introduce the sums SX=∑k=1nXkS_{X}=\sum_{k=1}^{n}X_{k} and SY=∑k=1nYkS_{Y}=\sum_{k=1}^{n}Y_{k}. Let us bound the expectation of the largest eigenvalue of the latter

Appendix D Miscellaneous

Appendix E Additional Experiments

As described in Section 5, we compute our the simple and Bernstein bounds, our practical estimate and the true L{\cal L} for ridge regression applied to small artificial datasets: uniform (n=24,d=50)(n=24,d=50), staircase eigval (n=d=24)(n=d=24) and alone eigval (n=d=24)(n=d=24). Figure 7 shows first that the practical estimate is a very close approximation of L{\cal L}. On the one hand, we observe in Figure 7(a) that the Bernstein bound performs poorly since the feature dimension is very small d=50d=50. On the other hand, Figure 7(c) shows a regime change for b≈10b\approx 10, which highlight the usefulness of combining our bounds to approximate the expected smoothness constant. Finally, we observe that for the alone eigval dataset Figure 7(b), which has one very large eigenvalue far from the rest of the spectrum, the simple bound matches L{\cal L} because the gap between Lˉ\bar{L} and LL shrinks. Indeed, in this configuration Lˉ≈L≈Lmax⁡n\bar{L}\approx L\approx\frac{L_{\max}}{n}. When the spectrum is more concentrated, like for staircase eigval, we get a significant gap between Lˉ\bar{L} and LL as shown in Figure 7(c), where the simple bound is far from L{\cal L} when b=nb=n.

We also report the influence of changing the value of the regularization parameter λ\lambda. Figure 8 shows that this parameter has little impact on the general shape of the bounds and of L{\cal L}.

Finally, we study the impact of scaling or standardizing (i.e., removing the mean and dividing by the standard deviation for each feature) our artificial datasets. In order not to benefit from the diagonal shape of the alone eigval and staircase eigval datasets we also give examples of the bounds of L{\cal L} after a rotation of the data. The rotation aims at preserving the spectrum while erasing the diagonal structure of the covariance matrix AA⊤AA^{\top}. This rotation procedure consists in transforming AA into Q⊤AQQ^{\top}AQ, where QQ is the orthogonal matrix given by the QR decomposition of a random squared matrix (with dimension the same as the one of AA) with uniformly random coefficients MM, such that M=QRM=QR.

We observe in Figure 9 that rotations do not affect our estimates of L{\cal L}, because they preserve the spectrum. Scaling non-diagonal datasets does not change the general shape neither. As predicted, scaling diagonal matrices leads to a particular case where the spectrum of the covariance matrix is flattened and for all i∈[n]i\in[n], Li≈Lmax⁡≈LˉL_{i}\approx L_{\max}\approx\bar{L}. This is why we get a flat simple bound in Figures 9(c) and 9(g). Even after those different types of preprocessing (rotation and scaling) and with different values of λ\lambda, we end up with the same strong observation that the practical estimate is a very sharp approximation of the expected smoothness constant.

E.2 Experiment 1: estimates of the expected smoothness constant for real datasets

In what folows, we also used publicly available datasets from LIBSVMhttps://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/ provided by Chang & Lin (2011) and from the UCI repositoryhttps://archive.ics.uci.edu/ml/datasets/ provided by Dheeru & Karra Taniskidou (2017). We applied ridge regression to the following datasets: YearPredictionMSD (n=515,345,d=90)(n=515,345,d=90) from LIBSVM and slice (n=53,500,d=384)(n=53,500,d=384) from UCI. We also applied regularized logistic regression for binary classification on ijcnn1 (n=141,691,d=22)(n=141,691,d=22), covtype.binary (n=581,012,d=54)(n=581,012,d=54), real-sim (n=72,309,d=20,958)(n=72,309,d=20,958), rcv1.binary (n=697,641,d=47,236)(n=697,641,d=47,236) and news20.binary (n=19,996,d=1,355,191n=19,996,d=1,355,191) from LIBSVM. When a test set was available, we concatenated it with the train set to have more samples.

One can observe in Figure 10, that for unscaled datasets the Bernstein bound performs better than the simple bound, except for YearPredictionMSD (n=515,345,d=90)(n=515,345,d=90) and covtype.binary (n=581,012,d=54)(n=581,012,d=54). From Figure 11, we observe that after feature-scaling, the Bernstein bound is always a tighter upper bound of L{\cal L} than the simple bound.

E.3 Experiment 2: step size estimates for artificial datasets

In this section we give the step sizes estimate corresponding to the expected smoothness constant, the simple and Bernstein upper-bounds and the practical estimate for our small artificial datasets. In Figure 12, we show that the practical step size estimate is larger than all others. Moreover, for except for small value sof bb, our γsimple\gamma_{\text{simple}} or γBernstein\gamma_{\text{Bernstein}} estimates are in most cases larger than the one proposed in (Hofmann et al., 2015).

E.4 Experiment 2: step size estimates for real datasets

Here we show the step sizes estimate corresponding to the simple and Bernstein upper-bounds and the practical estimate for real datasets detailed in Section E.2. On these real data, unscaled in Figure 13 and scaled in Figure 14, we see that the gap between our step size estimates and γHofmann\gamma_{\text{Hofmann}} is even larger. We observe in Figure 13, accordlingly to previous remarks in Section E.2, that Bernstein bound leads to larger step sizes than the simple one, except for the unscaled YearPredictionMSD and covtype.binary datasets. Yet, as noticed before, Figure 14 seems to show that scaling the data leads to γBernstein\gamma_{\text{Bernstein}} larger than γsimple\gamma_{\text{simple}}.

E.5 Experiment 3: comparison with previous SAGA settings

In this section we provide more example of the performance of our practical settings compared to previously known SAGA settings. In Figures 16, 17, 18, 19, 20 and 21 we run our experiments on real datasets introduced in detail in Section E.1. SAGA implementations are run until the suboptimality reaches a relative error (f(wk)−f(w∗))/(f(w0)−f(w∗))(f(w^{k})-f(w^{*}))/(f(w^{0})-f(w^{*})) of 10−410^{-4}, except in some cases where the Hofmann’s runs exceeded our maximal number of epochs like in Figure 17. In Figure 19, the curves corresponding to Hofmann’s settings are not displayed because they achieve a total complexity which is too large. Figure 15 shows an example of such a configuration.

These experiments show that our settings (bpractical,γpracticalb_{\text{practical}},\gamma_{\text{practical}}) most of the time outperforms whether the classical (b=1,γDefaziob=1,\gamma_{\text{Defazio}}) or the (b=20,γHofmannb=20,\gamma_{\text{Hofmann}}) settings both in terms of epochs and running time.

E.6 Experiment 4: optimality of the mini-batch size

We always observe a change of regime in the empirical complexity. For small values of bb, the complexity is of the same order of magnitude, then, for values greater than the empirical optimal mini-batch size, the complexity explodes. This experiment shows that our optimal mini-batch size bpracticalb_{\text{practical}} correctly designates the largest mini-batch achieving the best complexity as large as possible, without reaching the regime where the total complexity explodes.