Adaptive Importance Sampling for Finite-Sum Optimization and Sampling with Decreasing Step-Sizes

Functions $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ of the form:

are prevalent in modern machine learning and statistics. Important examples include the empirical risk in the empirical risk minimization framework, ¹¹1Up to a normalizing factor of $\frac{1}{N}$ which does not affect the optimization or the log-posterior of an exchangeable Bayesian model. When $N$ is large, the preferred methods for solving the resulting optimization or sampling problem usually rely on stochastic estimates of the gradient of $f$, using variants of stochastic gradient descent (SGD) [1] or stochastic gradient Langevin dynamics (SGLD) [2]:

where $\{\alpha_{t}\}_{t=1}^{T}$ is a sequence of step-sizes, and the index $I_{t}$ is sampled uniformly from $[N]$, making $N\nabla f_{I_{t}}(x)$ an unbiased estimator of the gradient of $f$. We use $\{x_{t}\}_{t=1}^{T}$ to refer to either sequence when we do not wish to distinguish between them. It is well known that the quality of the answers given by these algorithms depends on the (trace of the) variance of the gradient estimator, and considerable efforts have been made to design methods that reduce this variance.

In this paper, we focus on the importance sampling approach to achieving variance reduction. At each iteration, the algorithm samples $I_{t}$ according to a specified distribution $p^{t}$, and estimates the gradient using:

It is immediate to verify that $\hat{g}^{t}$ is an unbiased estimator of $g^{t}:=\nabla f(x_{t})$. By cleverly choosing the distributions $p^{t}$, one can achieve significant variance reduction (up to a factor of N) compared to the estimator based on uniform sampling. Unfortunately, computing the variance-minimizing distributions at each iteration requires the knowledge of the Euclidean norm of all the individual gradients $g_{i}^{t}:=\nabla f_{i}(x_{t})$ at each iteration, making it unpractical [3, 4].

Many methods have been proposed that attempt to construct sequences of distributions $\{p^{t}\}_{t=1}^{T}$ that result in efficient estimators [3, 4, 5, 6, 7, 8, 9]. Of particular interest to us, the task of designing such sequences was recently cast as an online learning problem with bandit feedback [10, 11, 12, 13]. In this formulation, one attempts to design algorithms with sub-linear expected static regret, which is defined as:

where $\Delta$ denotes the probability simplex in $\mathbb{R}^{N}$, and $c_{t}(p)$ is the trace of the covariance matrix of the gradient estimator (4), which is easily shown to be:

Note that the second term cancels in the definition of regret, and we omit it in the rest of our discussion. In this formulation, and to keep the computational load manageable, one has only access to partial feedback in the form of the norm of the $I_{t}^{th}$ gradient, and not to the complete cost function (5). Under the assumption of uniformly bounded gradients, the best result in this category can be found in [12] where an algorithm with $\tilde{O}(T^{2/3})$ static regret is proposed. A more difficult but more natural performance measure that makes the attempt to approximate the optimal distributions explicit is the dynamic regret, defined by:

Guarantees with respect to the expected dynamic regret are more difficult to obtain, and require that the cost functions $c_{t}(p)$ do not change too rapidly with respect to some reasonable measure of variation. See [14, 15, 16, 17, 18] for examples of such measures and the corresponding regret bounds for general convex cost functions.

In this work, we propose Avare, an algorithm that achieves sub-linear dynamic regret for both SGD and SGLD when the sequence of step-sizes $\{\alpha_{t}\}_{t=1}^{T}$ is decreasing. The name Avare is derived from adaptive variance minimization. Specifically, our contributions are as follows:

• •

  We show that Avare achieves $\mathcal{O}(T^{2/3})$ and $\mathcal{O}(T^{5/6})$ dynamic regret for SGD and SGLD respectively when $\alpha_{t}$ is $\mathcal{O}(1/t)$.

• •

  We propose a new mini-batch estimator that combines the benefits of sampling without replacement and importance sampling while preserving unbiasedness.

• •

  We validate empirically the performance of our algorithm and identify settings in which it leads to significant improvements.

We would like to point out that while the decreasing step size requirement might seem restrictive, we argue that for SGD and SGLD, it is the right setting to consider for variance reduction. Indeed, it is well known that under suitable technical conditions, both algorithms converge to their respective solutions exponentially fast in the early stages. Variance reduction is primarily useful at later stages when the noise from the stochastic gradient dominates. In the absence of control variates, one is forced to use decreasing step-sizes to achieve convergence. This is precisely the regime we consider.

It is easy to see that the cost functions (5) are convex over the probability simplex. A celebrated algorithm for convex optimization over the simplex is entropic descent [19], an instance of mirror descent [20] where the Bregman divergence is taken to be the relative entropy. A slight modification of this algorithm for online learning is the EXP3 algorithm [21] which mixes the iterates of entropic descent with a uniform distribution to avoid the assignment of null probabilities. See [22, 23, 24] for a more thorough discussion of online learning (also called online convex optimization) in general and variants of this algorithm in particular.

Since we are working over the simplex, the EXP3 algorithm is a natural choice. This is the approach taken in [10] and [11], although strictly speaking neither is able to prove sub-linear static regret bounds. The difficulty comes from the fact that the norm of the gradients of the cost functions (5) explode to infinity on the boundary of the simplex. This is amplified by the use of stochastic gradients which grow as $1/p_{min}^{5}$ in expectation, making it very difficult to reach regions near the boundary. Algorithms based on entropic descent for dynamic regret problems also exist, including the fixed share algorithm and projected mirror descent [25, 26, 27, 28]. Building on these algorithms, and by artificially making the cost functions strongly-convex to allow the use of decreasing step-sizes, we were only able to show $\tilde{O}(T^{7/8})$ dynamic regret using projected mirror descent for SGD with $O(1/t)$ decreasing step sizes and uniformly bounded gradients.

The approach taken in [12] is qualitatively different and is based on the follow-the-regularized-leader (FTRL) scheme. By solving the successive minimization problems stemming from the FTRL scheme analytically, the authors avoid the above-mentioned issue of exploding gradients. The rest of their analysis relies on constructing an unbiased estimate of the cost functions (5) using only the partial feedback $\left\lVert g_{I_{t}}^{t}\right\rVert_{2}^{2}$, and probabilistically bounding the deviation of the estimated cost functions from the true cost functions using a martingale concentration inequality. The final result is an algorithm that enjoys an $\tilde{O}(T^{2/3})$ static regret bound.

Our approach is similar to the one taken in [12] in that we work directly with the cost functions and not their gradients to avoid the exploding gradient issue. Beyond this point however, our analysis is substantially different. While we are still working within the online learning framework, our analysis is more closely related to the analysis of variance-reduced stochastic optimization algorithms that are based on control variates. In particular, we study a Lyapunov-like functional and show that it decreases along the trajectory of the dynamics defined by the algorithms. This yields simpler and more concise proofs, and opens the door for a unified analysis.

Most of the literature on online convex optimization with dynamic regret relies on the use of the gradients of the cost functions [14, 15, 16, 17, 18]. However, as explained in the previous section, such approaches are not viable for our problem. Unfortunately, the regret guarantees obtained from the FTRL scheme for static regret used in [12] do not directly translate into guarantees for dynamic regret.

We start by presenting the high level ideas that go into the construction of our algorithm. We then present our algorithm in explicit form, and discuss its implementation and computational complexity.

In this section, we prove a sub-linear dynamic regret guarantee for our proposed algorithm when used with SGD and SGLD with decreasing step-sizes. We present our results in a more general setting, and show that they apply to our cases of interest. We start by stating our assumptions:

The bounded gradients assumption has been used in all previous related work [10, 11, 12, 13], although it is admittedly quite strong. The smoothness assumption is standard in the study of optimization and sampling algorithms. Note that we chose to state our assumptions using index-independent constants to make the presentation clearer and since this does not affect our derivation of the sequence $\{\varepsilon_{t}\}_{t=1}^{T}$.

Finally, Assumption 3 should really be derived from more basic assumptions, and is only stated to allow for a unified analysis. Note that in the optimization case this is a very mild assumption since we are generally only interested in convergent sequences, and for any such sequence with reasonably fast convergence this assumption holds. The following proposition shows that Assumption 3 holds for our cases of interest. All the proofs for this section can be found in appendix B.

We now state a proposition that relates the optimal function value for the problem (13) over the restricted simplex, with the optimal function value over the whole simplex. Its proof is taken from ([12], Lemma 6):

The following lemma gives our first bound on the regret per time step:

The expectation in the first term of the above lemma is highly reminiscent of the first term of the Lyapunov function used to study the convergence of SAGA first proposed in [30] and subsequently refined in [31] (with $g_{i}^{t}$ replaced by $g_{i}^{*}$, the gradient at the minimum). Inspired by this similarity, we prove the following recursion:

The natural thing to do at this stage is to unroll the above recursion, replace in Lemma 2, sum over all time steps, and minimize the obtained regret bound over the choice of the sequence $\{\varepsilon_{t}\}_{t=1}^{T}$. However, even if we can solve the resulting optimization problem efficiently, the solution will still depend on the constants $G$ and $L$ and on the initial error due to the initialization of the $h_{i}$s, both of which are usually unknown. Here instead, we make an asymptotic argument to find the optimal decay rate of $\{\varepsilon_{t}\}_{t=1}^{T}$, propose a sequence that satisfies this decay rate, and show that it gives rise to the dynamic regret bounds stated.

Denote by $\varphi(t)$ the expectation in the first term of the upper bound in Lemma 2. Suppose we take $\varepsilon_{t}$ to be of order $t^{-\beta}$. Looking at the recursion in Lemma 3, we see that to control the positive term, we need the negative term to be of order at least $t^{-\delta}$, so that $\varphi(t)$ cannot be smaller than $t^{\beta-\delta}$. The bound of Lemma 2 is therefore of order $t^{2\beta-\delta}+t^{-\beta}$. The minimum is attained when the exponents are equal so we have: $2\beta-\delta=-\beta\implies\beta=\frac{\delta}{3}$

We are now ready to guess the form of $\varepsilon_{t}$. Matching the positive term in Lemma 3, we consider the following sequence:

For a free parameter $C$ satisfying $C\geq N$ to ensure $\varepsilon_{1}\leq 1/N$. With this choice of the sequence $\{\varepsilon_{t}\}_{t=1}^{T}$, we are now finally ready to state our main result of the paper:

Note that since $C\geq N$, we have $t_{0}\leq(2^{3/\delta}-1)N+1$, so $t_{0}$ is bounded by a constant independent of $T$. Furthermore, setting $C=2N$ makes the theorem hold for all $T\in\mathbb{N}$. In practice however, it might be beneficial to set $C=N$ to overcome a bad initialization of the $h_{i}$s.

Combining Theorem 1 with proposition 1 we obtain the following corollary:

In most practical applications, one uses a mini-batch of samples to construct the gradient estimator instead of just a single sample. The most basic such estimator is the one formed by sampling a mini-batch of indices $S_{t}=\{I_{t}^{1},\dots,I_{t}^{m}\}$ uniformly and independently, and taking the sample mean. This gives an unbiased estimator, whose variance decreases as $1/m$. A simple way to make it more efficient is by sampling the indices uniformly but without replacement. In that case, the variance decreases by an additional factor of $(1-(m-1)/(N-1))$. For $m\ll N$, the difference is negligible, and the additional cost of sampling without replacement is not justified. However, when using unequal probabilities, this argument no longer holds, and the additional variance reduction obtained from sampling without replacement can be significant even for small $m$.

For our problem, besides the additional variance reduction, sampling without replacement allows a higher rate of replacement of the $h_{i}$s, which is directly related to our regret through the factor in front of the second term of Lemma 3, whose proper generalization for mini-batch sampling is $(1-\min_{i\in[N]}\pi_{i}^{t})$, where $\pi_{i}^{t}=P(i\in S_{t})$ is the inclusion probability of index $i$. This makes sampling without replacement desirable for our purposes. Unfortunately, unlike in the uniform case, the sample mean generalization of (4) is no longer unbiased. We propose instead the following estimator:

We summarize some of its properties in the following proposition:

The proposed estimator is therefore unbiased and its variance decreases super-linearly in $m$ (by (b) and (d)). Although we were unable to prove a regret bound for this estimator, (c) suggests that it is still reasonable to use our algorithm. To take into account the higher rate of replacement of the $h_{i}$s, we propose using the following $\varepsilon_{t}$ sequence, which is based on the mini-batch equivalent of Lemma 3 and the inequality $\min_{i\in[N]}\pi_{i}^{t}\geq m\varepsilon_{t}$ [32, 33, 34]:

In this section, we present results of experiments with our algorithm. We start by validating our theoretical results through a synthetic experiment. We then show that our proposed method outperforms existing methods on real world datasets. Finally, we identify settings in which adaptive importance sampling can lead to significant performance gains for the final optimization algorithm.

In all experiments, we added $l_{2}$-regularization to the model’s loss and set the regularization parameter $\mu=1$. We ran SGD (2) with decreasing step sizes $\alpha_{t}=\frac{m}{2NL+m\mu t}$ where $m$ is the batch size following [35]. We experimented with 3 different samplers in addition to our proposed sampler: Uniform, Multi-armed bandit Sampler (Mabs) [11], and Variance Reducer Bandit (Vrb) \citesuppDBLP:conf/colt/Borsos0L18. The hyperparameters of both Mabs and Vrb are set to the ones prescribed by the theory in the original papers. For our propsoed sampler Avare, we use the epsilon sequence given by (20) with $C=N$, $\delta=1$, and initialized $h_{i}=0$ for all $i\in[N]$. For each sampler, we ran SGD 10 times and averaged the results. The shaded areas represent a one standard deviation confidence interval.

To validate our theoretical findings, we randomly generated a dataset for binary classification with $N=100$ and $d=10$. We then trained a logistic regression model on the generated data, and used a batch size of 1 to match the setting of our regret bound. The results of this experiment are depicted in figure 1, and show that Avare outperforms the other methods, achieving significantly lower dynamic regret and faster convergence.

For our second experiment, we tested our algorithm on three real world datasets: MNIST, IJCNN1 [36], and CIFAR10. We used a softmax regression model, and a batch size of 128 sampled with replacement. The results of this experiment can be seen in figure 2. The third column shows the relative error which we define as $\left[c_{t}(p^{t})-\min_{p\in\Delta}c_{t}(p)\right]/\min_{p\in\Delta}c_{t}(p)$. In all three cases, Avare achieves significantly smaller dynamic regret, with a relative error quickly decaying to zero. The performance gains in the final optimization algorithm are clear in both MNIST and IJCNN1, but are not noticeable in the case of CIFAR10.

To determine the effectiveness of non-uniform sampling in accelerating the optimization process, previous work [4, 11] has suggested to look at the ratio of the maximum smoothness constant and the average smoothness constant. We show here that this is the wrong measure to look at when using adaptive probabilities. In particular, we argue that the ratio of the variance with uniform sampling at the loss minimizer to the optimal variance at the loss minimizer is much more informative of the performance gains achievable through adaptive importance sampling. For each dataset, these ratios are displayed in Table 1, supporting our claim. We suspect that for large models capable of fitting the data nearly perfectly, our proposed ratio will be large since many of the per-example gradients at the optimum will be zero. We therefore expect our method to be particularly effective in the training of models of this type. We leave such experiments for future work. Finally, in appendix D, we propose an extension of our method to constant step-size SGD, and show that it preserves the performance gains observed when using decreasing step-sizes.

Our work develops a new method for variance reduction for stochastic optimization and sampling algorithms. On the optimization side, we expect our method to be very useful in accelerating the training of large scale neural networks, particularly since our method is expected to provide significant performance gains when the model is able to fit the data nearly perfectly. On the sampling side, we expect our method to be useful in accelerating the convergence of MCMC algorithms, opening the door for the use of accurate Bayesian methods at a large scale.

This research was supported by an NSERC discovery grant. We would like to thank the anonymous reviewers for their useful comments and suggestions.