Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence

Introduction

We solve the finite-sum optimization problem:

Stochastic gradient descent (SGD) (Robbins and Monro,, 1951; Nemirovski and Yudin,, 1978, 1983; Shalev-Shwartz et al.,, 2007; Nemirovski et al.,, 2009; Hardt et al.,, 2016), is the workhorse for training supervised machine learning problems that have the generic form (1).

The main parameter for guaranteeing the convergence of SGD is the step-size or the learning rate. In recent years, several ways of selecting the step-size have been proposed. Moulines and Bach, (2011); Needell et al., (2016); Needell and Ward, (2017); Nguyen et al., (2018); Gower et al., (2019) propose a non-asymptotic analysis of SGD with constant step-size for convex and strongly convex functions. For non-convex functions, such an analysis can be found in Ghadimi and Lan, (2013); Bottou et al., (2018). Using a constant step-size for SGD guarantees convergence to a neighbourhoood of the solution. A common technique to guarantee convergence to the exact optimum is to use a decreasing step-size (Robbins and Monro,, 1951; Ghadimi and Lan,, 2013; Gower et al.,, 2019; Nemirovski et al.,, 2009; Karimi et al.,, 2016). More recently, adaptive methods (Duchi et al.,, 2011; Liu et al.,, 2019; Kingma and Ba,, 2015; Bengio,, 2015; Vaswani et al., 2019b, ; Li and Orabona,, 2019; Ward et al.,, 2019) that adjust the step-size on the fly have become wide-spread and are particularly beneficial when training deep neural networks.

Contributions: Inspired by the classical Polyak step-size (Polyak,, 1987) commonly used with the deterministic subgradient method (Hazan and Kakade,, 2019; Boyd et al.,, 2003), we propose a novel adaptive learning rate for SGD. The proposed step-size is a natural extension of the Polyak step-size to the stochastic setting. We name it stochastic Polyak step-size (SPS). Although computing SPS requires knowledge of the $f_{i}^{*}$; we argue that this information is readily available for modern machine learning applications (for example, $f_{i}^{*}=0$ for most standard surrogate losses), making SPS an attractive choice for SGD.

In Section 3, we provide theoretical guarantees for the convergence of SGD with SPS in different scenarios including strongly convex, convex and non-convex smooth functions. Although SPS is provably larger than the typically used constant step-size, we guarantee its convergence to a reasonable neighborhood around the optimum. We note that in the modern machine learning tasks that we consider, it is enough to converge to a small neighbourhood and not the exact minimizer to get good generalization performance. We also establish a connection between SPS and the optimal step-size used in sketch and project methods for solving linear systems. Furthermore, in Appendix C, we provide convergence guarantees for convex non-smooth functions. We also show that by progressively increasing the batch-size for computing the stochastic gradients, SGD with SPS converges to the optimum.

Contributions: Our analysis of SGD with SPS does not require any of these additional assumptions for guaranteeing convergenceExcept for our analysis for non-convex smooth functions where the weak growth condition is used.. We also note that our theoretical results do not require the finite-sum assumption and can be easily adapted to the streaming setting.

In addition, unlike standard analysis for constant step-size SGD, the use of SPS requires an adaptive step-size that uses the loss and stochastic gradient estimates at an iterate, resulting in correlations. One of the main technical challenges in the proofs is to carefully analyze the SGD iterates taking these correlations into account. Furthermore, since we need to be adaptive to the Lipschitz constant, we can not use the descent lemma (implied by smoothness and SGD update). This makes the convex proof more challenging than the standard analysis.

Modern machine learning models such as non-parametric regression or over-parametrized deep neural networks are highly expressive and can fit or interpolate the training dataset completely (Zhang et al.,, 2017; Ma et al.,, 2018). In this setting, SGD with constant step-size can been shown to converge to the exact optimum at the deterministic rate (Schmidt and Roux,, 2013; Ma et al.,, 2018; Vaswani et al., 2019a, ; Vaswani et al., 2019b, ; Gower et al.,, 2019; Berrada et al.,, 2020).

Contributions: As a corollary of our theoretical results, we show that SPS is particularly effective under this interpolation setting. Specifically, we prove that SPS enables SGD to converge to the true solution at a fast rate matching the deterministic case. Moreover, SPS does not require the knowledge of any problem-dependent constants or additional computational overhead.

In Section 4, we experimentally validate our theoretical results via experiments on synthetic datasets. We also evaluate the performance of SGD equipped with SPS relative to the state-of-the-art optimization methods when training over-parameterized models for deep matrix factorization, binary classification using kernels and multi-class classification using deep neural networks. For each of these tasks, we demonstrate the superior convergence of the proposed method. The code to reproduce our results can be found at https://github.com/IssamLaradji/sps.

SGD and the Stochastic Polyak Step-size

The optimization problem (1) can be solved using SGD:

where example $i\in[n]$ is chosen uniformly at random and $\gamma_{k}>0$ is the step-size in iteration $k$.

Before explaining the proposed stochastic Polyak step-size, we first present the deterministic variant by Polyak (Polyak,, 1987). This variant is commonly used in the analysis of deterministic subgradient methods (Boyd et al.,, 2003; Hazan and Kakade,, 2019).

For convex functions, the deterministic Polyak step-size at iteration $k$ is the one that minimizes an upper-bound $Q(\gamma)$ on the distance of the iterate $x_{k+1}$ to the optimal solution: $\|x^{k+1}-x^{*}\|_{2}^{2}\leq Q(\gamma)$, where $Q(\gamma)=\|x^{k}-x^{*}\|^{2}-2\gamma\left[f(x^{k})-f^{*})\right]+\gamma^{2}\|g^{k}\|^{2}.$ That is,

Here $g^{k}$ denotes a subgradient of function $f$ at point $x^{k}$ and $f^{*}$ the optimum function value. For more details and a convergence analysis of the deterministic subgradient method, please check Appendix A.2. Note that the above step-size can be used only when the optimal value $f^{*}$ is known, however Boyd et al., (2003) demonstrate that $f^{*}=0$ for several applications (for example, finding a point in the intersection of convex sets, positive semidefinite matrix completion and solving convex inequalities).

It is clear that using the deterministic Polyak step-size in the update rule of SGD is impractical. It requires the computation of the function value $f$ and its full gradient in each iteration.

To avoid this, we propose the stochastic Polyak step-size (SPS) for SGD:

In addition to SPS, in some of our convergence results we require its bounded variant:

Here $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}>0$ is a bound that restricts SPS from being very large and is essential to ensure convergence to a small neighborhood around the solution. If $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}=\infty$ then $\text{SPS}_{\max}$ is equivalent to SPS.

We now briefly compare against the recently proposed stochastic variants of the Polyak step-size (Rolinek and Martius,, 2018; Oberman and Prazeres,, 2019; Berrada et al.,, 2020). In Section 3, we present a detailed comparison of the theoretical convergence rates.

In Rolinek and Martius, (2018), the L4 algorithm has been proposed showing that a stochastic variant of the Polyak step for SGD achieves good empirical results for training neural networks. However it has no theoretical convergence guarantees. The step-size is very similar to SPS (2) but each update requires an online estimation of the $f_{i}^{*}$ which does not result in robust empirical performance and requires up to three hyper-parameters.

In the ALI-G algorithm proposed by Berrada et al., (2020), the step-size is set as: $\gamma_{k}=\min\left\{\frac{f_{i}(x^{k})}{\|\nabla f_{i}(x^{k})\|^{2}+\delta},\eta\right\}$, where $\delta>0$ is a positive constant. Unlike our setting, their theoretical analysis relies on an $\epsilon$-interpolation condition. Moreover, the values of the parameter $\delta$ and $\eta$ that guarantee convergence heavily depend on the smoothness parameter of the objective $f$, limiting the method’s practical applicability. In Section 3, we show that as compared to Berrada et al., (2020), the proposed method results in both better rates and a smaller neighborhood of convergence. For the case of over-parameterized models, our step-size selection guarantees convergence to the exact solution while the step proposed in Berrada et al., (2020) finds only an approximate solution that could be $\delta$ away from the optimum. In Section 4, we also experimentally show that $\text{SPS}_{\max}$ results in better convergence than ALI-G.

2 Optimal Objective Difference

This is a very weak assumption. Moreover when (1) is the training problem of an over-parametrized model such as a deep neural network or involves solving a consistent linear system or classification on linearly separable data, each individual loss function $f_{i}$ attains its minimum at $x^{*}$, and thus $f_{i}(x^{*})-f_{i}^{*}=0.$ In this interpolation setting, it follows that $\sigma=0$.

Convergence Analysis

In this section, we present the main convergence results. For the formal definitions and properties of functions see Appendix A.1. Proofs of all key results can be found in the Appendix B.

If a function $g$ is $\mu$-strongly convex and $L$-smooth the following bounds hold: $\frac{1}{2L}\|\nabla g(x)\|^{2}\leq g(x)-\inf_{x}g(x)\leq\frac{1}{2\mu}\|\nabla g(x)\|^{2}.$ Using these bounds and by assuming that the functions $f_{i}$ in problem (1) are $\mu_{i}$-strongly convex and $L_{i}$-smooth, it is straight forward to see that SPS can be lower and upper bounded as follows:

where $L_{\max}=\max\{L_{i}\}_{i=1}^{n}$.

2 Sum of convex functions: strongly convex objective

In this section, we assume that all components $f_{i}$ are convex functions and that the objective function $f$ is $\mu$-strongly convex.

Let $f_{i}$ be $L_{i}$-smooth convex functions and assume that the objective function $f$ is $\mu$-strongly convex function. Then, SGD with $\text{SPS}_{\max}$ with $c\geq 1/2$ converges as:

where $\alpha:=\min\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\}$ and $L_{\max}=\max\{L_{i}\}_{i=1}^{n}$ is the maximum smoothness constant. The best convergence rate and the tightest neighborhood are obtained for $c=1/2$.

Note that in Theorem 3.1, we do not make any assumption on the value of the upper bound $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$. However, it is clear that for convergence to a small neighborhood of the solution $x^{*}$ (unique solution for strongly convex functions) $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ should not be very largeNote that neighborhood $\frac{2\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\sigma^{2}}{\mu\alpha}$ has $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ in the numerator and for the case of large $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$, $\alpha=\frac{1}{2cL_{\max}}$..

Another important aspect of Theorem 3.1 is that it provides convergence guarantees without requiring strong assumptions like bounded gradients or growth conditions. We do not use these conditions because SPS provides a natural bound on the norm of the gradients. In the following corollaries we make additional assumptions to better understand the convergence of SGD with $\text{SPS}_{\max}$.

In our first corollary, we assume that our model is able to interpolate the data (each individual loss function $f_{i}$ attains its minimum at $x^{*}$). This condition is satisfied for unregularized least-squares regression on a realizable dataset, or when using the squared-hinge loss on a linearly-separable dataset. The interpolation assumption enables us to guarantee the convergence of SGD with SPS, without an upper-bound on the step-size ($\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}=\infty$).

Assume interpolation ($\sigma=0$) and let all assumptions of Theorem 3.1 be satisfied. SGD with SPS with $c=1/2$ converges as:

We compare the convergence rate in Corollary 3.2 to that of stochastic line search (SLS) proposed in Vaswani et al., 2019b . In similar setting, SLS achieves the slower linear rate $\max\left\{1-\frac{\bar{\mu}}{L_{\max}},1-\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\bar{\mu}\right\}$, where $\bar{\mu}=\sum_{i=1}^{n}\mu_{i}/n$ is the average strong-convexity of the finite sum. In particular, according to Theorem 1 of Vaswani et al., 2019b , the convergence of SLS requires that at least one of the $f_{i}$’s is $\mu_{i}$-strongly convex implying that the objective function $f$ is strongly convex. This is a stronger assumption than the one we have in Theorem 3.1. We also note that $\bar{\mu}\leq\mu$.

In Berrada et al., (2020), ALI-G is analyzed under the strong assumption that all functions $f_{i}$ are $\mu$-strongly convex and $L$-smooth. For detailed comparison of SPS with ALI-G, see Appendix B.1.1.

An interesting outcome of Theorem 3.1 is a novel analysis for SGD with a constant step-size. In particular, note that if the bound in $\text{SPS}_{\max}$ is selected to be $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$, then using the lower bound of (5), it can be easily shown that our method reduces to SGD with constant step-size $\gamma_{k}=\gamma=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$. In this case, we obtain the following convergence rate.

Let all assumptions of Theorem 3.1 be satisfied. SGD with $\text{SPS}_{\max}$ with $c=1/2$ and $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{L_{\max}}$ becomes SGD with constant step-size $\gamma\leq\frac{1}{L_{\max}}$ and converges as:

If we further assume interpolation ($\sigma=0$), the iterates of SGD with constant step-size $\gamma\leq\frac{1}{L_{\max}}$ satisfy:

3 Sum of convex functions

Here, we derive the convergence rate when all component functions $f_{i}$ are convex without any strong convexity and obtain the following theorem.

Assume that $f_{i}$ are convex, $L_{i}$-smooth functions. SGD with $\text{SPS}_{\max}$ with $c=1$ converges as:

Here $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}$ and $\bar{x}^{K}=\frac{1}{K}\sum_{k=0}^{K-1}x^{k}$.

Analogous to the strongly-convex case, the size of the neighbourhood is proportional to $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$. When interpolation is satisfied and $\sigma=0$, we observe that the unbounded variant of SPS with $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}=\infty$ converges to the optimum at a $O(1/K)$ rate. This rate is faster than the rates in Vaswani et al., 2019b ; Berrada et al., (2020) and we refer the reader to the Appendix for a detailed comparison. As in the strongly-convex case, by setting $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{max}}$, we obtain the convergence rate obtained by constant step-size SGD.

4 Consistent Linear Systems

5 Sum of non-convex functions: PL Objective

We first focus on a special class of non-convex functions that satisfy the Polyak-Lojasiewicz (PL) condition (Polyak,, 1987). The PL inequality is a generalization of strong-convexity and is satisfied for matrix factorization (Sun and Luo,, 2016) or when minimizing the logistic loss on a compact set (Karimi et al.,, 2016). In particular, we assume that function $f$ satisfies the PL condition but do not assume convexity of the component functions $f_{i}$.

Assume that function $f$ satisfies the PL condition (7), and let $f$ and $f_{i}$ be smooth functions. SGD with $\text{SPS}_{\max}$ with $c>\frac{L_{\max}}{4\mu}$ and $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\geq\frac{1}{2cL_{\max}}$ converges as:

where $\nu=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\left(\frac{1}{\alpha}-2\mu+\frac{L_{\max}}{2c}\right)\in(0,1]$ and $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}$.

Under the interpolation setting, $\sigma=0$, and SPSmax converges to the optimal solution at a linear rate. If $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\min\left\{\frac{1}{2cL_{\max}},\frac{2c}{4\mu c-L_{\max}}\right\}$ using the lower bound in (5), the analyzed method is SGD with constant step-size and we obtain the following corollary.

Assume that $f$ satisfies the PL condition (7), and let $f$ and $f_{i}$ be smooth functions. SGD with constant step-size $\gamma_{k}=\gamma\leq\frac{\mu}{L_{\max}^{2}}$ converges as:

To the best of our knowledge this is the first result for the convergence of SGD for PL functions without assuming bounded gradient or bounded variance or interpolation (for more details see results in Karimi et al., (2016) and discussion in Gower et al., (2019)). In the interpolation case, we obtain linear convergence to the optimum with a constant step-size equal to that used in Vaswani et al., 2019a ; Lei et al., (2019).

6 General Non-Convex Functions

In this section, we assume a common condition used to prove convergence of SGD in the non-convex setting (Bottou et al.,, 2018).

Let $f$ and $f_{i}$ be smooth functions and assume that there exist $\rho,\delta>0$ such that the condition (8) is satisfied. SGD with SPSmax with $c>\frac{\rho L}{4L_{\max}}$ and $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\max\left\{\frac{2}{L\rho},\bar{\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}}\right\}$ converges as:

where $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\},\quad\zeta=\left(\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}+\alpha\right)-\rho\left(\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}-\alpha+L\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}^{2}\right)\quad$ and

From the above theorem, we observe that SGD with SPS results in $O(1/K)$ convergence to a neighborhoud governed by $\delta$. For the case that $\delta=0$, condition (8) reduces to the strong growth condition (SGC) used in several recent papers (Schmidt and Roux,, 2013; Vaswani et al., 2019b, ; Vaswani et al., 2019a, ). It can be easily shown that functions that satisfy the SGC condition necessarily satisfy the interpolation property (Vaswani et al., 2019a, ). In the special case of interpolation, SGD with SPS is able to find a first-order stationary point as efficiently as deterministic gradient descent. Moreover, for $c\in\left(\frac{\rho L}{4L_{\max}},\frac{\rho L}{2L_{\max}}\right]$, the lower bound $\frac{1}{2cL_{\max}}$ of SPS lies in the range $\left[\frac{1}{\rho L},\frac{2}{\rho L}\right)$ and thus the step-size is larger than $\frac{1}{\rho L}$, the best constant step-size analyzed in this setting (Vaswani et al., 2019a, ).

7 Additional Convergence Results

In Appendix C, we prove a $O(1/\sqrt{K})$ convergence rate for non-smooth convex functions. Furthermore, similar to Schmidt et al., (2011), we propose a method to increase the mini-batch size for evaluating the stochastic gradient and guarantee convergence to the optimal solution without interpolation.

Experimental Evaluation

We validate our theoretical results using synthetic experiments in Section 4.1. In Section 4.2, we evaluate the performance of SGD with SPS when training over-parametrized models. In particular, we compare against state-of-the-art optimization methods for deep matrix factorization, binary classification using kernel methods and multi-class classification using standard deep neural network models.

2 Experiments for over-parametrized models

In this section, we consider training over-parameterized models that (approximately) satisfy the interpolation condition. Following the logic of the previous section, we evaluate the performance of both the SPS and SPSmax variants with $f_{i}^{*}=0$. Throughout our experiments, we found that SPS without an upper-bound on the step-size is not robust to the misspecification of interpolation and results in large fluctuations when interpolation is not exactly satisfied. For SPSmax, the value of $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ that results in good convergence depends on the problem and requires careful parameter tuning. This is also evidenced by the highly variable performance of ALI-G (Berrada et al.,, 2020) that uses a constant upper-bound on the step-size. To alleviate this problem, we use a smoothing procedure that prevents large fluctuations in the step-size across iterations. This can be viewed as using an adaptive iteration-dependent upper-bound $\gamma^{k}_{{}_{\mathsf{\scriptscriptstyle b}}}$ where $\gamma^{k}_{{}_{\mathsf{\scriptscriptstyle b}}}=\tau^{b/n}\,\gamma^{k-1}$. Here, $\tau$ is a tunable hyper-parameter set to $2$ in all our experiments, $b$ is the batch-size and $n$ is the number of examples. We note that using an adaptive $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ can be easily handled by our theoretical results. A similar smoothing procedure has been used to control the magnitude of the step-sizes when using the Barzilai-Borwein step-size selection procedure for SGD (Tan et al.,, 2016) and is related to the “reset“ option for using larger step-sizes in Vaswani et al., 2019b . We set $c=1/2$ for binary classification using kernels (convex case) and deep matrix factorization (non-convex PL case). For multi-class classification using deep networks, we empirically find that any value of $c\geq 0.2$ results in convergence. In this case, we observed that across models and datasets, the fastest convergence is obtained with $c=0.2$ and use this value.

We compare our methods against Adam (Kingma and Ba,, 2015), which is the most common adaptive method, and other recent methods that report better performance than Adam: (i) stochastic line-search (SLS) (Vaswani et al., 2019b, ) (ii) ALI-G (Berrada et al.,, 2020)With ALI-G we refer to the method analyzed in Berrada et al., (2020). This is SGD with step-size the one described in Section 2. We highlight that the experiments in Berrada et al., (2020) used momentum on top of the analyzed method but without any convergence quarantees. To ensure a fair comparison with SPS, we do not use such momentum. (iii) rectified Adam (RADAM) (Liu et al.,, 2019) (iv) Look-ahead optimizer (Zhang et al.,, 2019). We use the default learning rates and momentum (non-zero) parameters and the publicly available code for the competing methods. All our results are averaged across $5$ independent runs.

Next, we compare the optimizers’ performance in the convex, interpolation regime. We consider binary classification using RBF kernels, using the logistic loss without regularization. The bandwidths for the RBF kernels are set according to the validation procedure described in Vaswani et al., 2019b . We experiment with four standard datasets: mushrooms, rcv1, ijcnn, and w8a from LIBSVM (Chang and Lin,, 2011). Figure 2 shows the training loss on the mushrooms and ijcnn for the different optimizers. Again, we observe the strong performance of SPS compared to the other optimizers.

We benchmark the convergence rate and generalization performance of SPS methods on standard deep learning experiments. We consider non-convex minimization for multi-class classification using deep network models on the CIFAR10 and CIFAR100 datasets. Our experimental choices follow the setup in Luo et al., (2019). For CIFAR10 and CIFAR100, we experiment with the standard image-classification architectures: ResNet-34 (He et al.,, 2016) and DenseNet-121 (Huang et al.,, 2017). For space concerns, we report only the ResNet experiments in the main paper and relegate the DenseNet and MNIST experiments to Appendix E. From Figure 2, we observe that SPS results in the best training loss across models and datasets. For CIFAR-10, SPS results in competitive generalization performance compared to the other optimizers, whereas for CIFAR-100, its generalization performance is better than all optimizers except SLS. Note that ALI-G, the closest related optimizer results in worse generalization performance in all cases. We note that SPS is able to match the performance of SLS, but does not require an expensive back-tracking line-search or additional tricks.

For this set of experiments, we also plot how the step-size varies across iterations for SLS, SPS and ALI-G. Interestingly, for both CIFAR-10 and CIFAR-100, we find that step-size for both SPS and SLS follows a cyclic behaviour - a warm-up period where the step-size first increases and then decreases to a constant value. Such a step-size schedule has been empirically found to result in good training and generalization performance (Loshchilov and Hutter,, 2017) and our results show that SPS is able to simulate this behaviour.

Conclusion

We proposed and theoretically analyzed a stochastic variant of the classical the Polyak step-size. We quantified the convergence rate of SPS in numerous settings and used our analysis techniques to prove new results for constant step-size SGD. Furthermore, via experiments on a variety of tasks we showed the strong performance of SGD with SPS as compared to state-of-the-art optimization methods. There are many possible interesting extensions of our work: using SPS with accelerated methods, studying the effect of mini-batching and non-uniform sampling techniques and extensions to the distributed and decentralized settings.

Nicolas Loizou and Sharan Vaswani acknowledge support by the IVADO Postdoctoral Funding Program. Issam Laradji is funded by the UBC Four-Year Doctoral Fellowships (4YF). This research was partially supported by the Canada CIFAR AI Chair Program and by a Google Focused Research award. Simon Lacoste-Julien is a CIFAR Associate Fellow in the Learning in Machines & Brains program.

The authors would like to thank Frederik Kunstner for help with the convex proofs, and Aaron Defazio for fruitful discussions and feedback on the manuscript.

References

Appendix A Technical Preliminaries

Let us present some basic definitions used throughout the paper.

A.2 The Deterministic Polyak step-size

where $\gamma_{k}$ is the step-size (learning rate) and $g^{k}$ is any subgradient of function $f$ at point $x^{k}$.

Let $f$ be convex function. Let $\gamma_{k}=\frac{f(x^{k})-f(x^{*})}{\|g^{k}\|^{2}}$ be the step-size in the update rule of subgradient method. Here $f(x^{*})$ denotes the optimum value of function $f$. Let $G>0$ such that $\|g^{k}\|^{2}<G^{2}$. Then,

where $f^{k}_{*}=\min\{f(x^{i}):i=0,1,\dots,k\}$.

where the last line follows from the definition of subgradient:

which is precisely the step-size that minimize the right hand side of (13). That is,

. By using this choice of step-size in (13) we obtain:

From the above note that $\|x^{k}-x^{*}\|^{2}$ is monotonic function. Now using telescopic sum and by assuming $\|g^{k}\|^{2}<G^{2}$ we obtain:

Let us define $f^{k}_{*}=\min\{f(x^{i}):i=0,1,\dots,k\}$ then: $[f^{k}_{*}-f(x^{*})]^{2}\leq\frac{G^{2}\|x^{0}-x^{*}\|^{2}}{k+1}$ and

For more details and slightly different analysis check Polyak, (1987) and Boyd et al., (2003). In Hazan and Kakade, (2019) similar analysis to the above have been made for the deterministic gradient descent ($g^{k}=\nabla f(x^{k})$) under several assumptions. (convex, strongly convex , smooth).

Appendix B Proofs of Main Results

In this section we present the proofs of the main theoretical results presented in the main paper. That is, the convergence analysis of SGD with $\text{SPS}_{\max}$ and SPS under different combinations of assumptions on functions $f_{i}$ and $f$ of Problem (1).

First note that the following inequality can be easily obtained by the definition of $\text{SPS}_{\max}$ (3):

We use the above inequality in several parts of our proofs. It is the reason that we are able to obtain an upper bound of $\gamma_{k}^{2}\|\nabla f_{i}(x^{k})\|^{2}$ without any further assumptions. For the case of SPS (2), inequality (17) becomes equality.

From convexity of functions $f_{i}$ it holds that $-\langle x^{k}-x^{*},\nabla f_{i}(x^{k})\rangle+f_{i}(x^{k})-f_{i}(x^{*})\leq 0$, $\forall i\in[n]$. Thus,

By taking expectation condition on $x^{k}$

From strong convexity of the objective function $f$ we have that $f(x^{k})-f(x^{*})-\langle x^{k}-x^{*},\nabla f(x^{k})\rangle\leq-\frac{\mu}{2}\|x^{k}-x^{*}\|^{2}$. Thus, we obtain:

Taking expectations again and using the tower property:

Recursively applying the above and summing up the resulting geometric series gives:

Let $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}$ then,

From definition of $\alpha$ is clear that having small parameter $c$ improves both the convergence rate $1-\mu\alpha$ and the neighborhood $\frac{2\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\sigma^{2}}{\mu\alpha}$. Since we have the restriction $c\geq\frac{1}{2}$ the best selection would be $c=\frac{1}{2}$. ∎

In the next corollary, in order to compare against the results for ALI-G from Berrada et al., (2020), we make the strong assumption that all functions $f_{i}$ have the same properties. We note that such an assumption in the interpolation setting is quite strong and reduces the finite-sum optimization to minimization of a single function in the finite sum.

Let all the assumptions in Theorem 3.1 be satisfied and let all $f_{i}$ be $\mu$-strongly convexThis is a much stronger assumption than assuming that functions $f_{i}$ are convex and $f$ is strongly convex, which is the main assumption of Theorem 3.1. Nevertheless, assuming that all $f_{i}$ are $\mu$-strongly convex functions implies that the objective function $f$ is $\mu$-strongly convex and that functions $f$ are convex. Thus, Theorem 3.1 still holds. and $L$-smooth. SGD with $\text{SPS}_{\max}$ with $c=1/2$ converges as:

For the interpolated case ($\sigma=0$) we obtain the same convergence as Corollary 3.2 with $L_{\max}=L$.

Note that, the result of Corollary B.1 is obtained by substituting $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\overset{\eqref{NewBounds}}{=}\frac{1}{2c\mu}\overset{c=\frac{1}{2}}{=}{\frac{1}{\mu}}$ into (6).

For the setting of Corollary B.1, Berrada et al., (2020) show the linear convergence to a much larger neighborhood than ours and with slower rate. In particular, their rate is $1-\frac{\mu}{8L}$ and the neighborhood is $\frac{8L}{\mu}(\frac{\epsilon}{L}+\frac{\delta}{4L^{2}}+\frac{\epsilon}{2\mu})$ where $\delta>2L\epsilon$ and $\epsilon$ is the $\epsilon$-interpolation parameter $\epsilon>\max_{i}[f_{i}(x^{*})-f_{i}^{*}]$ which by definition is bigger than $\sigma^{2}$. Under interpolation where $\sigma=0$, our method converges linearly to the $x^{*}$ while the algorithm proposed by Berrada et al., (2020) still converges to a neighborhood that is proportional to the parameter $\delta$.

B.2 Proof of Theorem 3.4

Let $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}$ and recall that from the definition of $\text{SPS}_{\max}$ (3) we obtain:

From the above if $\alpha=\frac{1}{2cL_{\max}}$ then the step-size is in the regime of the stochastic Polyak step (5). In the case that $\alpha=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ then the analyzed method becomes the constant step-size SGD with stepsize $\gamma_{k}=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$.

Since $c>\frac{1}{2}$ it holds that $\left(2-\frac{1}{c}\right)>0$. Using (21) into (20) we obtain:

where in the last inequality we use that $\alpha\left(2-\frac{1}{c}\right)\left[f_{i}(x^{*})-f_{i}^{*}\right]>0$.

By taking expectation condition on $x^{k}$ and dividing by $\alpha\left(2-\frac{1}{c}\right)$:

Taking expectation again and using the tower property:

Summing from $k=0$ to $K-1$ and dividing by $K$:

Let $\bar{x}^{K}=\frac{1}{K}\sum_{k=0}^{K-1}x^{k}$, then:

At this point we highlight that $c=1$ is selected to simplify the expression of the upper bound in (B.2). This is not the optimum choice (the one that makes the rate and the neighborhood of the upper bound smaller). In order to compute the optimum value of $c$ one needs to follow similar procedure to Gower et al., (2019) and Needell et al., (2016). In this case $c$ will depend on parameter $\sigma$ and the desired accuracy $\epsilon$ of convergence.

However as we show bellow having $c=1$ allows SGD with SPS to convergence faster than the ALI-G algorithm (Berrada et al.,, 2020) and the SLS algorithm (Vaswani et al., 2019b, ) for the case of smooth convex functions.

Similar to the strongly convex case let us compare the above convergence for smooth convex functions with the convergence rates proposed in Vaswani et al., 2019b and Berrada et al., (2020).

For the smooth convex functions, Berrada et al., (2020) show the linear convergence to a much larger neighborhood than ours and with slower rate. In particular, their rate is $\frac{1}{K}\left(\frac{2L}{1-\frac{2L\epsilon}{\delta}}\right)$ and the neighborhood is $\frac{\delta}{L(1-\frac{2L\epsilon}{\delta})}$ where $\delta>2L\epsilon$ and $\epsilon$ is the $\epsilon$-interpolation parameter $\epsilon>\max_{i}[f_{i}(x^{*})-f_{i}^{*}]$ which by definition is bigger than $\sigma^{2}$. Under interpolation where $\sigma=0$, our method converges with a $O(1/K)$ rate to the $x^{*}$ while the algorithm proposed by Berrada et al., (2020) still converges to a neighborhood that is proportional to the parameter $\delta$.

B.3 SPS on Methods for Solving Consistent Linear Systems

Recently several new randomized iterative methods (sketch and project methods) for solving large-scale linear systems have been proposed (Richtárik and Takác,, 2020; Loizou and Richtárik, 2020b, ; Loizou and Richtárik, 2020a, ; Gower and Richtárik,, 2015). The main algorithm in this literature is the celebrated randomized Kaczmarz (RK) method (Kaczmarz,, 1937; Strohmer and Vershynin,, 2009) which can be seen as special case of SGD for solving least square problems (Needell et al.,, 2016). In this area of research, it is well known that the theoretical best constant step-size for RK method is $\gamma=1$.

As we have already mentioned in Section 3.4, given the consistent linear system

Richtárik and Takác, (2020) provide a stochastic optimization reformulation of the form (1) which is equivalent to the linear system in the sense that their solution sets are identical. That is, the set of minimizers of the stochastic optimization problem ${\cal X}^{*}$ is equal to the set of solutions of the stochastic linear system ${\cal L}:=\{x:\mathbf{A}x=b\}$.

In particular, the stochastic convex quadratic optimization problem proposed in Richtárik and Takác, (2020), can be expressed as follows:

Here the expectation is over random matrices ${\bf S}$ drawn from an arbitrary, user defined, distribution ${\cal D}$ and $f_{{\bf S}}$ is a stochastic convex quadratic function of a least-squares type, defined as

For solving problem (27), Richtárik and Takác, (2020) analyze SGD with constant step-size:

where $\nabla f_{{\bf S}_{k}}(x^{k})$ denotes the gradient of function $f_{{\bf S}_{k}}$. In each step the matrix ${\bf S}_{k}$ is drawn from the given distribution ${\cal D}$.

The above update of SGD is quite general and as explained by Richtárik and Takác, (2020) the flexibility of selecting distribution ${\cal D}$ allow us to obtain different stochastic reformulations of the linear system (26) and different special cases of the SGD update. For example the celebrated randomized Kaczmarz (RK) method can be seen as special cases of the above update as follows:

Randomized Kaczmarz Method: Let pick in each iteration the random matrix $\mathbf{S}=e_{i}$ (random coordinate vector) with probability $p_{i}=\|\mathbf{A}_{i:}\|^{2}/\|\mathbf{A}\|_{F}^{2}$. In this setup the update rule of SGD (29) simplifies to

Many other methods like Gaussian Kacmarz, Randomized Coordinate Descent, Gaussian Decsent and their block variants can be cast as special cases of the above framework. For more details on the general framework and connections with other research areas we also suggest (Loizou and Richtárik,, 2019; Loizou,, 2019).

As we will see in the next Theorem, using the special structure of the stochastic reformulation (27), SPS (2) with $c=1/2$ takes the following form:

which is the theoretically optimal constant step-size for SGD in this setting (Richtárik and Takác,, 2020). This reduction implies that SPS results in an optimal convergence rate when solving consistent linear systems. We provide the convergence rate for SPS in the next Theorem.

Though a straight forward verification of the optimality of SPS, we believe that this is the first time that SGD with adaptive step-size is reduced to constant step-size when is used for solving linear systems. SPS does that by obtaining the best convergence rate in this setting.

Let $\mathbf{A}x=b$ be a consistent linear system and let $x^{*}$ is the projection of vector $x$ onto the solution set ${\cal X}^{*}={\cal L}$. Then the SGD with SPS (2) with $c=1/2$ for solving the stochastic optimization reformulation (27) satisfies:

Let us select $\gamma_{k}$ such that the RHS of inequality (33) is minimized. That is, let us select:

Substitute this step-size to (33) we obtain:

By taking expectation with respect to ${\bf S}_{k}$ and using quadratic growth inequality (31):

Taking expectation again and by unrolling the recurrence we obtain (32). ∎

We highlight that the above proof provides a different viewpoint on the analysis of the optimal constant step-size for the sketch and project methods for solving consistent liner systems. The expression of Theorem B.3 is the same with the one proposed in Richtárik and Takác, (2020).

B.4 Proof of Theorem 3.6

By the smoothness of function $f$ we have that

Combining this with the update rule of SGD we obtain:

and by taking expectation condition on $x^{k}$:

Let $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}$. Then,

Using $\gamma_{k}\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ and by taking expectations again:

By having $\nu\in(0,1]$ and by recursively applying the above and summing the resulting geometric series we obtain:

In the above result we require that $0<\nu=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\left(\frac{1}{\alpha}-2\mu+\frac{L_{\max}}{2c}\right)\leq 1$. In order for this to hold we need to make extra assumptions on the values of $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ and parameter $c$. This is what we do next.

Let us divide the analysis into two cases based on the value of parameter $\alpha$. That is:

(i) If $\frac{1}{2cL_{\max}}\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ then,

By preliminary computations, it can be easily shown that $\nu>0$ for every $c\geq 0$. However for $\nu\leq 1$ we need to require that $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{\left(\frac{1}{\alpha}-2\mu+\frac{L_{\max}}{2c}\right)}$ and since we are already assume that $\frac{1}{2cL_{\max}}\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ we need to force

to avoid contradiction. This is true only if $c>\frac{L_{\max}}{4\mu}$ which is the assumption of Theorem 3.6.

(ii) If $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$ then,

Note that if we have $c>\frac{L_{\max}}{4\mu}$ (an assumption of Theorem 3.6) it holds that $\nu\leq 1$. In addition, by preliminary computations, it can be shown that $\nu>0$ if $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\frac{2c}{4\mu c-L_{\max}}$. Finally, for $c>\frac{L_{\max}}{4\mu}$ it holds that $\frac{1}{2cL_{\max}}\leq\frac{2c}{4\mu c-L_{\max}}$, and as a result $\nu>0$ for all $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$.

By presenting the above cases on bound of $\nu$ we complete the proof. ∎

The expression of Corollary 3.7 is obtained by simply use $c=\frac{L_{\max}}{2\mu}$ in the case (ii) of the above proof. In this case we have $\gamma\leq\frac{\mu}{L_{\max}^{2}}$ and $\nu=1-\mu\gamma$.

B.5 Proof of Theorem 3.8

By the smoothness of function $f$ we have that

Combining this with the update rule of SGD we obtain:

By taking expectation condition on $x^{k}$:

Since $0<\alpha\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ we have that $\left(\frac{\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}}{2}-\frac{\alpha}{2}+\frac{L\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}^{2}}{2}\right)>0$. Thus, we are able to use (8):

By rearranging and taking expectation again:

By summing from $k=0$ to $K-1$ and dividing by $K$:

In the above result we require that $\zeta=\left(\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}+\alpha\right)-\left(\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}-\alpha+L\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}^{2}\right)\rho>0$. In order for this to hold we need to make extra assumptions on the values of $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ and parameter $c$. This is what we do next.

Let us divide the analysis into two cases. That is:

(i) If $\frac{1}{2cL_{\max}}\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ then $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}=\frac{1}{2cL_{\max}}$ and

By solving the quadratic expression of $\zeta$ with respect to $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$, it can be easily shown that $\zeta>0$ if

To avoid contradiction the inequality $\frac{1}{2cL_{\max}}<\bar{\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}}$ needs to be true, where $\bar{\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}}$ is the above upper bound of $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$. This is the case of $c>\frac{L\rho}{4L_{\max}}$ which is the assumption of Theorem 3.8.

(ii) If $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$ then $\alpha=\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}=\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ and

In this case, by preliminary computations, it can be shown that $\zeta>0$ if $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\frac{2}{L\rho}$. For $c>\frac{L\rho}{4L_{\max}}$ it also holds that $\frac{1}{2cL_{\max}}<\frac{2}{L\rho}$.

Let us now present an extra theoretical result in which we assume that the step-size $\gamma_{k}$ and the stochastic gradient $\nabla f_{i}(x^{k})$ in each step are not correlated. Such assumption has been recently used to prove convergence of SGD with the AdaGrad step-size (Ward et al.,, 2019) and for the analysis of stochastic line search in Vaswani et al., 2019b . From a technical viewpoint, we highlight that for the proofs in the non-convex setting, we use the lower and upper bound of SPS rather than its exact form. This is what allows us to use this independence.

Let us state the main Theorem with the extra condition of independence and present its proof.

Let $f$ and $f_{i}$ be smooth functions and assume that there exist $\rho,\delta>0$ such that the condition (8) is satisfied. Assuming independence of the step-size $\gamma_{k}$ and the stochastic gradient $\nabla f_{i_{k}}(x^{k})$ at every iteration $k$, SGD with SPSmax with $c>\frac{\rho L}{4L_{\max}}$ and $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\max\left\{\frac{2}{L\rho},\frac{1}{\sqrt{\rho c\,LL_{\max}}}\right\}$ converges as:

where $\beta_{1}=1-\frac{\rho c\,LL_{max}\,\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}^{2}}{2}$ and $\beta_{2}=1-\frac{\rho L\,\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}}{2}$, $\alpha=\min\left\{\frac{\beta_{1}}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\,\beta_{2}\right\}$.

By the smoothness of function $f$ we have that

Combining this with the update rule of SGD we obtain:

Taking expectations with respect to $i_{k}$ and noting that $\gamma_{k}$ is independent of $\nabla f_{i}(x^{k})$ yields:

By rearranging and taking expectations again:

By summing from $k=0$ to $K-1$ and dividing by $K$:

In the above result we require that $\alpha=\left(\min\left\{\frac{1}{2cL_{\max}},\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\right\}-\frac{L\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}^{2}}{2}\rho\right)>0$. In order for this to hold we need to make extra assumptions on the values of $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ and parameter $c$. This is what we do next.

Let us divide the analysis into two cases. That is:

(i) If $\frac{1}{2cL_{\max}}\leq\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}$ then,

By preliminary computations, it can be easily shown that $\alpha>0$ if $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\frac{1}{\sqrt{L\rho cL_{\max}}}$. To avoid contraction the inequality $\frac{1}{2cL_{\max}}<\frac{1}{\sqrt{L\rho cL_{\max}}}$ needs to be true. This is the case of $c>\frac{L\rho}{4L_{\max}}$ which is the assumptions of Theorem B.5.

(ii) If $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\leq\frac{1}{2cL_{\max}}$ then,

In this case, by preliminary computations, it can be shown that $\alpha>0$ if $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}<\frac{2}{L\rho}$. For $c>\frac{L\rho}{4L_{\max}}$ it also holds that $\frac{1}{2cL_{\max}}<\frac{2}{L\rho}$.

Appendix C Additional Convergence Results

In this section we present some additional convergence results. We first prove a $O(1/\sqrt{K})$ convergence rate of stochastic subgradient method with SPS for non-smooth convex functions in the interpolated setting. Furthermore, similar to Schmidt et al., (2011), we propose a way to increase the mini-batch size for evaluating the stochastic gradient and guarantee convergence to the optimal solution without interpolation.

In all of our previous results we assume that functions $f_{i}$ are smooth. As a result, in the proofs of our theorems we were able to use the lower bound (5) of SPS. In the case that functions $f_{i}$ are not smooth using this lower is clearly not possible. Below we present a Theorem that handles the case of non-smooth function for the convergence of stochastic subgradient methodNote that for non-smooth functions, it is required to have stochastic subgradient method instead of SGD. That is, in each iteration we replace the evaluation of $\nabla f_{i}(x)$ with its subgradient counterpart $g_{i}(x)$. For this result we require that a constant $G$ exists such that $\|g_{i}(x)\|^{2}<G^{2}$ for each subgradient of function $f_{i}$. This is equivalent with assuming that functions $f_{i}$ are $G$-Lipschitz. To keep the presentation simple we only present the interpolated case. Using the proof techniques from the rest of the paper one can easily obtain convergence for the more general setting.

where $\bar{x}^{K}=\frac{1}{K}\sum_{k=0}^{K-1}x^{k}.$

The proof is similar to the deterministic case (see Theorem A.4). That is, we select the $\gamma_{k}$ that minimize the right hand side of the inequality after the use of convexity.

Using the subgradient counterpart of SPS (2) with $c=1$, that is, $\gamma_{k}=\frac{f_{i}(x^{k})-f_{i}(x^{*})}{\|g_{i}^{k}\|^{2}}$Recall that in the interpolation setting it holds that $f_{i}^{*}=f^{*}=f(x^{*})$. we obtain:

Taking expectation again and using the tower property:

By rearranging, summing from $k=0$ to $K-1$ and dividing by $K$:

Taking square roots and using Jensen’s inequality:

where $\bar{x}^{K}=\frac{1}{K}\sum_{k=0}^{K-1}x^{k}.$ ∎

C.2 Increasing Mini-batch Size

We propose a way to increase the mini-batch size for evaluating the stochastic gradient and guarantee convergence to the optimal solution without interpolation. We present two main Theorems. In the first Theorem we assume that functions $f_{i}$ of problem (1) are $\mu_{i}$-strongly convex functions and in the second that each function $f_{i}$ satisfies the PL condition (10) with $\mu_{i}$ parameter.

Let us have the same assumptions as in Theorem 3.1 and let all $f_{i}$ be $\mu_{i}$-strongly convex functions. Then SGD with SPS, and increasing the batch-size progressively such that the batch-size $b_{k}$ at iteration $k$ satisfies:

Following the proof of Theorem 3.1 for the batch $b$. From Equation B.1,

By strong-convexity of all $f_{i}$, $\forall i\in[n]$ the minibatch function $f_{b}$ is $\mu_{b}$-strongly convex and it holds that:

where in the last inequality we use that $\mu_{\min}=\min\{\mu_{i}\}_{1=1}^{n}\leq\mu_{i}\leq\mu_{b}$. By the assumption that the gradients at the optimum have bounded variance, from Harikandeh et al., (2015); Lohr, (2019):

If we set the batch-size in iteration $k$ such that,

Following the remaining proof of Theorem 3.1,

For SPS it holds that $\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}=\infty$. Thus,

Assume that all functions $f_{i}$ satisfy the PL inequality (10) and let $f$ and $f_{i}$ be smooth functions. Then SGD with SPSmax, and increasing the batch-size progressively such that the batch-size $b_{k}$ at iteration $k$ satisfies:

where $v=1-\gamma_{{}_{\mathsf{\scriptscriptstyle b}}}\left(\frac{1}{\alpha}-2\mu+\frac{L_{\max}}{2c}\right)\in(0,1)$.

Following the proof of Theorem 3.6, from Equation 38,

Similar to the proof of Theorem C.2, since each function $f_{i}$ is PL,

Following the remaining proof of Theorem 3.6,

Following Mező and Baricz, (2017), we first start by presenting the definition of the Lambert $W$ function and its recent generalization, the $r$-Lambert function.

The inverse of the function on the left-hand side of the above equation ($xe^{x}$) is called the Lambert $W$ function and is denoted by $W$. In general, $W$ is a multivalued function on complex numbers. But for $a\geq 0$, there is a unique real solution (called the principal branch) and this is the one that we will consider in this paper, i.e. the unique real solution of (67) for $a\geq 0$ is given by $x=W(a)$. We note that $W(a)$ is a strictly increasing function for $a\geq 0$ with $W(0)=0$, and that there are efficient numerical routines to compute it (Corless et al.,, 1996).

The $r$-Lambert function is a direct generalization of the Lambert $W$ function first proposed by Mező and Baricz, (2017). It is used to express the solution of the transcendental equation

From the above definitions, it is clear that the classical Lambert $W$ function is a special case of the $r$-Lambert function when $r=0$. In this case we write $W_{0}$.

Manipulating the transcendental equation (68), we can easily solve the slightly more general equation as given in the following Theorem.

the binary log-loss. We show that in this case that $f_{i}^{*}$ can be computed in closed form expression using the $r$-Lambert function.

the binary exponential loss. As mentioned by Bartlett et al., (2006) this loss appears in the Adaboost algorithm, amongst others. In this case, the Lambert $W$ function is used.

where $A_{i}$ is the input feature vector for the $i^{th}$ datapoint while $b_{i}\in\{-1,1\}$ is its label, and $\lambda$ is the regularization parameter.

Note that by following the notation of the rest of the paper, in (70) we have

Note that $\log\left(1+e^{-\alpha\langle\hat{x},z_{i}\rangle}\right)$ is decreasing as $\langle\hat{x},z_{i}\rangle$ increases. Thus, by Cauchy-Schwartz inequality, $\inf_{\hat{x}}$ is reached when $\hat{x}=\frac{z_{i}}{\|z_{i}\|}$ and equation (71) takes the following form:

$g(\alpha)$ is a strongly convex function of $\alpha$, and we can find its global minimum by setting its gradient to zero:

Let $c=\|z_{i}\|$, then by rearranging the last equation, we obtain:

Using Theorem D.3, the solution of the above equation can be expressed as follows (using $r=1$):

Thus to get a closed form expression for $f_{i}^{*}$, one can plug $\alpha^{*}$ in (72), i.e. $f_{i}^{*}=g(\alpha^{*})$.

D.2 Binary Exponential Loss

with the same notation as for the logistic regression problem (70).

As for the logistic regression derivation, defining $z_{i}:=b_{i}A_{i}$, and letting $x=\alpha\hat{x}$ with $\|\hat{x}\|=1$, then we obtain the following:

As for the logistic regression, note that by Cauchy-Schwartz inequality, $\inf_{\hat{x}}$ is reached when $\hat{x}=\frac{z_{i}}{\|z_{i}\|}$. Thus,

Again, $g(\alpha)$ is a strongly convex function of $\alpha$, and we can find its global minimum by setting its gradient to zero:

Let $c=\|z_{i}\|$, then by rearranging the last equation, we obtain:

Using Theorem D.3, the solution of the above equation can be expressed as follows (using $r=0$):

where $W_{0}$ is the classical Lambert $W$ function.

Thus to get a closed form expression for $f_{i}^{*}$, one can plug $\alpha^{*}$ in (76), i.e. $f_{i}^{*}=g(\alpha^{*})$.

Appendix E Additional Experiments

Following the experiments presented in the main paper, we further evaluate the performance of SGD with SPS when training over-parametrized models.