We are interested in the unconstrained minimization problem:

where $\xi$ is a real random vector with distribution $P$. In machine learning, $x$ represents the parameters of the prediction rule, $\xi$ represents the pair of input/output variables, $f$ the loss function, and $F$ the risk.

In particular, we focus on the sample average approximation to (1):

where $(\xi_{i})_{i=1}^{n}$ are i.i.d. samples from $P$. This is also known as empirical risk minimization in machine learning.

Defining $f_{i}(x)\coloneqq f(x,\xi_{i})$ for $i\in[n]$, and letting $i$ be a uniformly distributed random variable on $[n]$, we can rewrite (2) as:

where the expectation is taken with respect to $i$.

We make the following assumptions on the function $F$ and the component functions $f_{i}$:

Note that by Assumption 3 and the triangle inequality, $F$ is also $L$-smooth, and we define its condition number by $\kappa=L/\mu$.

Let $x^{*}$ be the unique minimizer of $F$. Our goal is to design the fastest algorithm that outputs a point $x\in\mathbb{R}^{d}$ such that $\left\lVert x-x^{*}\right\rVert_{2}^{2}<\varepsilon$ for a given accuracy $\varepsilon>0$. Following (Agarwal & Bottou, 2015), we assume that we have access to a gradient oracle that takes as input a point $x\in\mathbb{R}^{d}$ and an index $i\in[n]$ and returns $\nabla f_{i}(x)$. We then measure the complexity of an algorithm by counting the number of oracle calls needed to achieve a desired accuracy.

The most straightforward way to solve (3) is to ignore the particular structure of $F$, and simply run gradient descent or accelerated gradient-descent (Nesterov, 1983) on $F$. Under our assumptions and definition of complexity, gradient descent converges at the rate of $O(n\kappa\log{(1/\varepsilon)})$, while accelerated gradient descent converges at the rate of $O(n\sqrt{\kappa}\log{(1/\varepsilon)})$ (Nesterov, 2018). When $n$ is large, this becomes prohibitively expensive.

One solution to this is to view problem (3) in its expectation form, and use the stochastic approximation of gradient descent. This yields stochastic gradient descent (SGD) which estimates $\nabla F(x)$ by $\nabla f_{i}(x)$ for $i$ uniformly distributed on $[n]$ in the gradient descent update (Robbins & Monro, 1951; Nemirovsky & Yudin, 1983; Nemirovski et al., 2009; Moulines & Bach, 2011). The complexity of SGD under our assumptions is known to have two regimes. Denote by $\sigma^{2}$ the variance of the gradient estimator of SGD at the minimizer $x^{*}$. For $\varepsilon>\sigma^{2}/L\mu$, SGD converges at the fast rate $O(\kappa\log{(1/\varepsilon)})$, while for $\varepsilon\leq\sigma^{2}/L\mu$, SGD suffers from the sublinear rate $O(\sigma^{2}/\mu^{2}\varepsilon)$ (Needell et al., 2014; Nguyen et al., 2018; Gower et al., 2019). When $\sigma^{2}$ is large, SGD takes a long time to converge to even moderately accurate solutions.

While SGD does take into account the expectation form of the objective (3), it fails to adapt to the fact that the expectation is taken over a discrete distribution which gives rise to a finite-sum structure. The last decade has seen the development of many new methods that leverage this additional structure, including (but not limited to) SAG (Roux et al., 2012; Schmidt et al., 2017), SAGA (Defazio et al., 2014), SVRG (Johnson & Zhang, 2013), and SARAH (Nguyen et al., 2017). All of these methods converge at the fast linear rate of $O((n+\kappa)\log{(1/\varepsilon)})$. Further work in this direction led to the development of lower-bounds for the optimization of finite-sum functions under our assumptions and the oracle model of complexity (Agarwal & Bottou, 2015; Woodworth & Srebro, 2016; Arjevani, 2017; Lan & Zhou, 2018), and, similar to the deterministic case, accelerated methods have been developed that closely match them (Lin et al., 2018; Allen-Zhu, 2018; Lan et al., 2019; Song et al., 2020). The core innovation behind this set of algorithms is the design of more efficient gradient estimators using carefully designed control variates, earning them the name of variance-reduced methods.

Despite all this progress, SGD remains the algorithm of choice in practice for some machine learning problems where it empirically outperforms variance-reduced methods (Defazio & Bottou, 2019). From an optimization perspective, there are a few explanations to this observation.

One is the phenomenon of interpolation (Ma et al., 2018; Vaswani et al., 2019a, b): if $\sigma^{2}=0$, then SGD converges linearly for any $\varepsilon>0$ without any variance reduction. A second possible explanation is that in many cases we are content with a low accuracy solution $\varepsilon>\sigma^{2}/L\mu$, in which case SGD converges linearly as well. In either of these scenarios, SGD and variance-reduced methods have the same iteration complexity, but variance-reduced methods require on average three times as many gradient evaluations per iteration (SVRG/SARAH).

While these reasons can be used to explain the superior performance of SGD in some settings, there is evidence that there are scenarios where variance-reduction might be useful even when only moderate accuracy is needed. In particular, (Goyal et al., 2018) showed that reducing the variance of the gradient estimator by increasing the batch size leads to optimization gains in the training of deep neural networks for ImageNet. Even in such cases however, the user is faced with a new problem. If we use variance-reduced methods from the beginning of the optimization, we make progress three times more slowly than SGD in the initial phase, which is often the dominating phase. Ideally, we would like to pay the additional computational cost for variance-reduction only when we need it, but it is not clear how we can detect when progress becomes constrained by the variance of the gradient estimator, although some advances have been made in this direction, see for example (Pesme et al., 2020).

A pressing question that arises from our discussion is therefore: Can we develop a variance-reduced algorithm that uses exactly the same number of gradient evaluations as SGD ? Such an algorithm would enjoy the fast rate of SGD in the initial phase of optimization, while automatically performing variance-reduction when needed. This question is the main motivation for this paper.

Before closing this section, we give one additional motivation for our work. As previously mentioned, the core idea of variance-reduced methods is the use of carefully crafted control variates to reduce the variance of the gradient estimator. There are, however, other ways to reduce the variance of a Monte Carlo estimator. In particular, one can use importance sampling. There is a large literature on such methods. Initial attempts focused on using a fixed distribution throughout the optimization based on prior knowledge about the functions $f_{i}$ (Needell et al., 2014; Zhao & Zhang, 2015; Csiba & Richtárik, 2018), and were able to show slightly favorable rates compared to uniform sampling. A more recent line of work attempts to adaptively design the distributions (Schaul et al., 2015; Loshchilov & Hutter, 2015; Alain et al., 2016; Bouchard et al., 2016; Canevet et al., 2016; Stich et al., 2017; Katharopoulos & Fleuret, 2018; Johnson & Guestrin, 2018), although the methods developed in these works mostly rely on heuristics and do not come with theoretical guarantees. Finally, in an attempt to put adaptive importance sampling methods on a firmer theoretical ground, a recent set of papers embed the problem of designing the distributions in an online learning framework and propose methods that achieve sublinear regret (Namkoong et al., 2017; Salehi et al., 2017; Borsos et al., 2018, 2019; El Hanchi & Stephens, 2020). These analyses however are still not satisfactory as they assume that the gradient norms of the functions $f_{i}$ are uniformly bounded on $\mathbb{R}^{d}$ which does not hold, for example, when $F$ is strongly-convex. The second question we ask here is therefore: Can we develop a variance-reduced algorithm based on importance sampling and which provably outperforms SGD ?

Contributions: In this paper, we give positive answers to both questions. In particular, our contributions are as follows:

• •

  We propose stochastic reweighted gradient descent (SRG), a variance-reduced algorithm for the optimization of finite-sums based on importance sampling.

• •

  As oppose to SAGA/SAG which require $O(nd)$ memory, SRG only requires $O(n)$ memory.

• •

  Unlike SVRG/SARAH which require on average three gradient evaluations per iteration, and like SGD, SRG requires a single gradient evaluation per iteration.

• •

  We develop a specialized red-black tree that allows an efficient implementation of SRG, incurring an overhead of only $O(\log{n})$ operations per iteration compared to SGD.

• •

  We show that SRG provably outperform SGD under our assumptions. Let $\sigma^{2}_{*}$ be the minimal variance achievable through importance sampling at the minimizer $x^{*}$. Then we show that SRG has the same convergence as SGD, but with $\sigma^{2}$ replaced with $\sigma_{*}^{2}$, which can be up to $n$ times smaller.

Before introducing our proposed algorithm, let us first introduce some notation. For an iteration number $k\in\mathbb{N}$, SGD performs the following update to minimize (3):

for a step size $\alpha_{k}>0$ and a random index $i_{k}$ drawn uniformly from $[n]$ and independently at each iteration. The idea behind importance sampling is to instead sample the index $i_{k}$ according to a given distribution $p_{k}$ on $[n]$, and to perform the update:

where $p_{k}^{i}$ is the $i^{th}$ component of the vector $p_{k}\in\Delta$, and $\Delta$ is the probability simplex in $\mathbb{R}^{n}$. It is immediate to verify that the importance sampling estimator of the gradient is unbiased as long as $p_{k}>0$. The challenge is to design a sequence $\{p_{k}\}_{k=0}^{\infty}$ that produces more efficient gradient estimators than the ones produced by uniform sampling.

Our proposed method is a simple modification of the one given by (El Hanchi & Stephens, 2020), and is given in Algorithm 1. The algorithm follows from the following reasoning. The variance of the importance sampling gradient estimator is given by, up to an additive constant:

Ideally we would like to choose $p_{k}\in\Delta$ so that this variance is minimized, but this is not feasible as we do not have access to the gradient norms $\left\lVert\nabla f_{i}(x_{k})\right\rVert_{2}$. Inspired by control variates methods, and in particular SAG/SAGA, we instead maintain a table $(g_{k}^{i})_{i=1}^{n}$ that aims to track the component gradients $(\nabla f_{i}(x_{k}))_{i=1}^{n}$. We then use this table to construct an approximation of the true variance:

Finally, we choose $p_{k}$ so as to minimize this quantity. Since this is only an approximation, we perform the minimization over the restricted simplex $\Delta(\varepsilon_{k})=\{p\in\Delta\mid p\geq\varepsilon_{k}\}$ for a given $\varepsilon_{k}\in(0,1/n]$. This enforces $p_{k}>0$, while making sure the error on the approximation of the variance is taken into account. The difference between our method and the one proposed in (El Hanchi & Stephens, 2020) is in the way the table $(g_{k}^{i})_{i=1}^{n}$ is updated: we add a bernoulli random variable to determine whether the update occurs or not. This step is crucial for the analysis as we will see, although in practice it might be skipped.

The goal of the next two sections will be to (i) show that SRG can be efficiently implemented. (ii) analyse the convergence of SRG and show that it outperforms SGD.

In this section, we show how Algorithm 1 can be efficiently implemented, leading to a per-iteration cost that is competitive with SGD.

First, note that while we use the table $(g_{k}^{i})_{i=1}^{n}$ in the presentation of Algorithm 1, in practice it is enough to store the gradient norms $(\left\lVert g_{k}^{i}\right\rVert_{2})_{i=1}^{n}$ only. SRG therefore requires only $O(n)$ memory compared to SAG/SAGA which require $O(nd)$ memory in general.

With memory issues out of the way, we turn to the main bottleneck in SRG which is sampling from:

The following lemma, taken from (El Hanchi & Stephens, 2020), sheds some light on the solution:

In this section, we analyze the convergence of SRG, and show that it outperforms SGD under our assumptions. Two key constants are helpful in contrasting the convergence of SRG and SGD. Recall the definition of $\sigma^{2}(x_{k},p_{k})$ in (4), and define:

It is well known that the convergence of SGD depends on $\sigma^{2}$ (Gower et al., 2019). We here show that SRG is able to reduce this dependence to $\sigma_{*}^{2}$, which can be up to $n$ times smaller than $\sigma^{2}$.

In this section, we experiment with SRG on $\ell_{2}$-regularized logistic regression problems. In this case, the functions $f_{i}$ are given by:

where $y_{i}\in\{0,1\}$ is the label of data point $a_{i}\in\mathbb{R}^{d}$. It is not hard to show that each $f_{i}$ is convex and $L_{i}=0.25\left\lVert a_{i}\right\rVert_{2}^{2}+\mu$ smooth. Their average $F$ is also $\mu$-strongly convex.

We experiment with four datasets from LIBSVM (Chang & Lin, 2011): ijcnn1, w8a, mushrooms, and phishing. We start by normalizing the data points $(a_{i})_{i=1}^{n}$ so that $\left\lVert a_{i}\right\rVert_{2}=1$ for all $i\in[n]$. This makes all the $f_{i}$ $L$-smooth with $L=0.25+\mu$. As is standard for regularized logistic regression, we take $\mu=\frac{1}{n}$. We conduct two sets of experiments: one set using a constant step size, and another using decreasing step sizes. In both cases we evaluate the performance of the algorithms by tracking the relative error $\frac{\left\lVert x_{k}-x^{*}\right\rVert_{2}^{2}}{\left\lVert x_{0}-x^{*}\right\rVert_{2}^{2}}$. For all experiments, we sample the indices without replacement using a batch size of $128$. Finally, we use the same step size across all algorithms for each dataset to fairly compare their performance. We give more details on the step sizes we use in the supplementary material.

The results of the experiments with constant step-size can be seen on the top row of Figure 1. We compare SRG with SGD on the one hand and SVRG on the other. We observe that SRG consistently outperforms SGD, reaching solutions that are about an order of magnitude more accurate. We also notice that SRG extends the range of accuracies on which SGD dominates SVRG. Eventually, SVRG outperforms both SGD and SRG, but our results show that SRG allows reaching solutions of moderate accuracy more quickly.

The bottom row of Figure 1 shows the results of the experiments where we use decreasing step-sizes. Here we compare SRG with SGD as well as SGD with random shuffling. SGD with random shuffling has attracted a lot of attention recently, and many analyses have been developed that show its superior performance, see for example (Ahn et al., 2020). Our results confirm this as SGD with shuffling outperforms both SGD and SRG. We still see however that SRG converges slightly faster than SGD, as our analysis suggests.

While the empirical results we present here are positive but not impressive, we would like to point out that the improvement of SRG over SGD depends directly on the ratio $r=\sigma^{2}/\sigma^{2}_{*}$. In particular, if $\sigma_{*}^{2}=\sigma^{2}$, then we should not expect any improvement. For the experiments in Figure 1, $r$ is in the range $[5,10]$. One can construct artificial examples in which this ratio is very large, which we do in the supplementary material, leading to significant improvements. It is however still unclear how often such cases are encountered in practice. We hope to explore this question more in future work.