Variance Reduction in SGD by Distributed Importance Sampling

Introduction

Many of the advances in Deep Learning from the past 5-10 years can be attributed to the increase in computing power brought by specialized hardware (i.e. GPUs). The whole field of Machine Learning has adapted to this reality, and one of the latest challenges has been to make good use of multiple GPUs, potentially located on separate computers, to train a single model.

One widely studied solution is Asynchronous Stochastic Gradient Descent (ASGD), which is a variation on SGD in which the gradients are computed in parallel, propagated to a parameter server, and where we drop certain synchronization barriers to allow the algorithm to run faster. This method was introduced by Bengio et al. (2003) in the context of neural language models, and extended to model-parallelism and demonstrated on a large scale by Dean et al. (2012).

One of the important limitations of ASGD is that it requires a lot of bandwidth to propagate the parameters and the gradients. Moreover, many theoretical guarantees are lost due to the fact that synchronization barriers are removed and stale gradients are being used. Some theoretical guarantees can still be made in the context of convex optimization (see Agarwal & Duchi (2011), Recht et al. (2011), Lian et al. (2015)), but any result from convex optimization applied to neural networks (highly non-convex) has to be used with fingers crossed.

In this paper we present a different principle for distributed training based on importance sampling. We demonstrate many interesting theoretical results, and show some experiments to validate our ideas. The reader should view these experiments as a proof of concept rather than as an appeal to switch from Asynchronous SGD to Importance Sampling SGD. In fact, we can imagine our method being a supplement to ASGD.

Throughout this paper, we will use the word stale to refer to the fact that certain quantities are slightly outdated, but usually not to the point of being completely unusable. These stale values are usually gradients computed from a set of parameters $\theta_{t}$ when some reference model is now dealing with parameters $\theta_{t+\Delta t}$.

In section 2 we will explain our distributed Importance Sampling SGD approach. In section 3 we revisit a classical result from the importance sampling literature and demonstrate a more general result that applies to high dimensions. We also present a technique that can be used to compute efficiently the gradient norms for all the individual members of a minibatch. In section 4 we discuss our particular implementation for distributed training. In section 5 we show experiments to illustrate both the reduction in variance and the increase in performance that it can bring.

The main contribution of this paper is to open the door via theoretical and experimental results to a novel approach to distributed training based on importance sampling, to focus the attention of the learner on the most informative examples from the learning point of view.

Scaling Deep Learning by Distributing Importance Sampling

One of the most important constraints on ASGD is the fact that it requires a large amount of bandwidth. Indeed, all the workers connecting to the parameter server are required to regularly fetch a fresh copy of the parameters, and all their computed gradients have to be pushed to the parameter server. For every minibatch processed by a worker computing a gradient, the memory size of that gradient vector is equal to the memory size of the parameter vector for the model (i.e. every parameter value gets a gradient value). Delaying synchronization can result in “stale” gradients, that is, gradients that are computed from a set of parameters that have been fetched from the parameter server too long ago to be relevant.

The approach that we are taking in this paper is to focus on the most “useful” training samples instead of giving equal attention to all the training set. Humans can learn from a small collection of examples, and a good tutor is able to pick examples that are useful for a student to learn the current lesson. This work can therefore be seen as a follow-up on the curriculum learning ideas (Bengio et al., 2009), where the model itself is used to figure out which examples are currently informative for the learner. The method that we present in this paper will incorporate that intuition into a training algorithm that is justified by theory rooted in importance sampling (Tokdar & Kass, 2010).

The approach of calibrating the importance sampling coefficients in order to minimize variance during SGD is also presented by Bouchard et al. (2015), in their method called “Adaptive Weighted SGD”, in which they adjust the coefficients by performing an intermediate gradient step to learn the best sampling proposal. They demonstrate how this can lead to improvements in convergence speed and generalization performance. In our paper, we show how an exact method can be used to get those optimal coefficients.

Compared to ASGD, our approach can be used to alleviate some of the communication costs. Instead of communicating the gradients on minibatches, the workers communicate one floating-point number per training sample. In a situation where the parameters can be of size ranging from 100 MB to 1GB, this cuts down the network transfers significantly. The parameters still have to be sent on the network to update the workers, however, but that cost can be amortized over a long period if the algorithm turns out to be robust to the use of older parameters in order to select the important samples. Our experiments confirm that hypothesis.

Importance Sampling in theory

Importance sampling is a technique used to reduce variance when estimating an integral of the form

through a Monte-Carlo estimate based on samples drawn from $p(x)$. Here $f(x)$ can only take on real values, but $x$ can be anything as long as it’s compatible with the probability density function $p(x)$.

It relies on a sampling proposal $q(x)$, for which $0<q(x)$ whenever $0<p(x)$, and the observation that

Since all the quantities in the following empirical sum are independent,

we can directly verify they are unbiased and then try to minimize their variance. The unbiasedness follows directly from equation (1), and with a little work can prove that that the variance is minimized when

2 Extending beyond a single dimension

Minimizing the variance is a well-defined objective in one dimension, but when going to higher dimensions we have to decide what we would like to minimize.

For our application, a natural choice of objective function (Bouchard et al., 2015) would be the trace of the covariance matrix of the proposal distribution, $\operatorname{Tr}(\Sigma)$, because it corresponds to the sum of all the eigenvalues of $\Sigma$, which is a positive semi-definite matrix. It also corresponds to sum of all the variances for each individual component of the gradient vector. We can also imagine minimizing $\left\|\Sigma(q)\right\|_{\textrm{F}}^{2}$, but in this case this would yield a different $q^{*}$ for which we do not know of an analytical form.

A nice consequence of our choice is that, when $d=1$, this $\operatorname{Tr}(\Sigma)$ will get back the classic result from the importance sampling literature. This is an pre-requisite for any general result.

The context requires that $q(x)>0$ whenever $p(x)>0$. We know that the importance sampling estimator

Let $\Sigma(q)$ be the covariance of that estimator, where we include $q$ in the notation to be explicit about the fact that it depends on the choice of $q$.

Then the trace of $\Sigma(q)$ is minimized by the following optimal proposal $q^{*}$ :

Note that in theorem 1 we refer to a general function $f$. It should be understood by the reader that we are really interested in the particular situation in which $f$ represents the gradient of a loss function with respect to the parameters of a model to be trained. However, since our results are meant to be more general than that, we tried to avoid contaminating them with those specific details, and decided to stick with $f(x)$ instead of talking about $\nabla_{\theta}\mathcal{L}(x_{n})$.

Then we have that the trace of the covariance of the importance sampling estimator is given by

To sample from $q(x)$ we just normalize the probability weights

and we sample from a multinomial distribution with argument $(\omega_{1},\ldots,\omega_{N})$ to pick the corresponding element in $\mathcal{D}$.

3 Dealing with minibatches

To apply the principles of ISSGD, we need to be able to evaluate $\left\|g(x_{n})\right\|_{2}$ efficiently for all the elements of the training set, where $g$ here is the gradient of the loss with respect to all the parameters of the model.

In the current landscape of machine learning, using minibatches is a fact of life. Any training paradigm has to take that into consideration, and this can be a challenge when one considers that the gradient for a single training sample is as big the parameters themselves. This fact is generally not a problem since the gradients are aggregated for all the minibatch at the same time, so the cost of storing the gradients is comparable to the cost of storing the model parameters.

In this particular case, what we need is a recipe to compute the gradient norms directly, without storing the gradients themselves. The recipe in question, formulated here as proposition 1, was published by Goodfellow (2015) slightly prior to our work. It applies to the fully-connected layers, but unfortunately not to convolutional layers.

Consider a multi-layer perceptron (MLP) applied to minibatches of size $N$, and with loss $\mathcal{L}=\mathcal{L}_{1}+\ldots+\mathcal{L}_{N}$, where $\mathcal{L}_{n}$ represents the loss contribution from element $n$ of the minibatch.

Let $(W,b)$ be the weights and biases at any particular fully-connected layer so that $XW+b=Y$, where $X$ are the inputs to that layer and $Y$ are the outputs.

The gradients with respect to the parameters are given by

where the values $\left(\frac{\partial\mathcal{L}_{n}}{\partial W},\frac{\partial\mathcal{L}_{n}}{\partial b}\right)$ refer to the particular contributions coming from element $n$ of the minibatch. Then we have that

where the notation $X[n,:]$ refers to row $n$ of $X$, and similarly for $\frac{\partial\mathcal{L}}{\partial Y}[n,:]$.

That is, we have a compact formula for the Euclidean norms of the gradients of the parameters, evaluated for each $N$ elements of the minibatch independently.

Note that proposition 1 applies to MLPs that have all kinds of activation functions and/or pooling operations, as long as the parameters $(W,b)$ are not shared between layers. We can ignore the activation functions when applying proposition 1 because the activation functions do not have any parameters, and the linear operation part (matrix multiplication plus vector addition) simply uses whatever quantities are backpropagated without knowing what comes after in the sequence of layers.

Despite the fact that convolutions are linear operations (in the mathematical sense), this formula fails to apply to convolutions because of their sparsity patterns and their parameter sharing.

In a situation where we face convolutional layers along with fully-connected layers, proposition 1 applies to the fully-connected layers. For our purpose of performing importance sampling, this is not satisfying because we would have to find another way to compute the gradient norms for all the parameters. One might suggest to abandon the plan of achieving optimal importance sampling and simply ignore the contributions of sparsely-connected layers, but we do not investigate this strategy in this paper.

Distributed implementation of ISSGD

We can see from corollary 1 that the expected trace of the covariance matrix over the whole training set is given by

The constant $\left\|g^{\textsc{true}}\right\|_{2}^{2}$ does not depend on the choice of $q$ so we will leave it out of the current discussion. Refer to section B.2 for more details about it.

The oracle allows us to achieve the ideal Importance Sampling SGD, and this quantity becomes

In this situation, we are using $q_{\textsc{ideal}}$ as notation instead of $q^{*}$. This is because we will want to contrast this situation with $q_{\textsc{unif}}$ and $q_{\textsc{stale}}$ that we will define shortly.

When performing SGD training with $q_{\textsc{ideal}}$, we can plot those values of equation (7) as we go along, and we can compare at each time step the $\operatorname{Tr}(\Sigma(q_{\textsc{ideal}}))$ with the value of $\operatorname{Tr}(\Sigma(q_{\textsc{unif}}))$ that we would currently have if we were using uniform sampling to construct the minibatches. The latter are given by

In figure 4 we will see those quantities compared during an experiment where we do not have access to an oracle, but where we can still evaluate what would have been the $\operatorname{Tr}(\Sigma(q_{\textsc{ideal}}))$ that we would have had if we had an oracle. This is relatively easy to evaluate by using equation (7).

2 Implementing the oracle using multiple machines

There are degrees of staleness, and the usefulness of a weight computed 5 minutes ago differs greatly from that of a weight computed 2 seconds ago. We refer to $q_{\textsc{stale}}$ as the proposal that is based on all the weights from the previous iteration. It serves its role as pessimistic estimator, which is generally worse than what we are actually using. It is also easier to compute because we can get it from values stored in the database without having to run the model on anything more.

We know for a fact that $\operatorname{Tr}(\Sigma(q_{\textsc{ideal}}))$ is the lower bound on all the possible $\operatorname{Tr}(\Sigma(q))$. When the weights are not in a horrible state due to excessive staleness, we generally observe experimentally that the following inequality holds:

3 Exact implementation vs relaxed implementation

This kind of relaxation is analogous to how ASGD discards the synchronization barriers to trade away correctness to gain performance. However, in the case of ISSGD, stale probability weights may lead to more variance but we will always get an unbiased estimator of the true gradient, even when we get rid of all the synchronization barriers.

In the appendix we discuss three aspects of how training can be adapted to be more practical and robust. In section B.1 we discuss the possibility of using only a subset of the probability weights, filtering them based on how recently they have been updated. In section B.2 we discuss how to approximate $\left\|g^{\textsc{true}}\right\|_{2}^{2}$, which is not a quantity that we absolutely need to compute for perform training, but which is something that we like to monitor to assess the benefits of using ISSGD instead of regular SGD. In section B.3 we add a smoothing constant to the probability weights in order to make training more robust to sudden changes in gradients.

Experimental results

We evaluated our model on the Street View House Numbers (SVHN) dataset from Netzer et al. (2011). We used the cropped version of the dataset (sometimes referred to as SVHN-2), which contains about 600,000 32x32 RGB images of house number digits from Google Street View.

Since there is no standard validation set, we randomly split 5% of the data to form our own validation set. Since our per-example gradient norm computation (from section 3.3) does not work with parameter sharing models (such as RNNs and Convnets), we consider the permutation invariant version of the SVHN task, in which the model is forced to discard the spatial structure of the pixels. While the permutation-invariant task is not practically relevant (as the spatial structure of the pixels is useful), it is commonly used as a testbed for studying fully connected neural networks (Goodfellow et al., 2013; Srivastava et al., 2013).

This is not meant to be a paper about exploring a variety of models, and since we stick with the permutation-invariant task this already limits our ability to use more interesting models. In any case, we picked an MLP with 4 hidden layers, each with 2048 hidden units and with a ReLU at its output (except for a softmax at the final layer). We are very much aware that a convolutional model would perform better.

We have used Theano ((Bergstra et al., 2010; Bastien et al., 2012)) to implement the model, and Redis as a database solution. The master and workers are each equipped with a k20 GPU.

2 Reduced training time and better prediction error

We compare in figure 2 the training loss for a model trained with ISSGD (in green) and regular SGD (in blue). We used 3 workers to help with the master. In the case of regular SGD, we also used a worker in the background to be able to compute statistics as we go along without imposing that burden on the process training the model. To make sure that the results are not due to the random initialization of parameters, we ran this experiment 50 times. We report here the median (thicker line), and the quartiles 1 and 3 above and below (thinner lines). This represents a “tube” into which half of the trajectories fit.

In all the figures from this section, we always compare the same two sets of hyperparameters. On the left we always have a setting where the learning rate is higher (0.01) and where we smoothe the probability weights by adding a constant (+10.0) to them (see section B.3 in the appendix for more explanations on this technique). On the right we always have a setting where the learning rate is smaller (0.001) and where the smoothing constant is also smaller (+1.0).

In figure 2 we can see that in both cases ISSGD minimizes the train loss more quickly than regular SGD, and it actually reaches 0.0. This obviously corresponds to overfitting, but since we are presenting here an optimization method, it seems natural to celebrate the fact that it can minimize the objective function faster and better.

In figure 3 we show the test prediction error. These results are not so easy to interpret, and we see that faster convergence does not always lead to a better generalization error. This suggests that regular SGD benefits here from a kind of regularization effect.

We also report in table 1 what are the final prediction errors for both methods (averaged over the last 10% of the timesteps plotted). We picked the set of hyperparameters that had the best validation prediction error and reported the test prediction errors. Unsurprizingly, this corresponds to using the result from figure 3(a) for ISSGD and figure 3(b) for regular SGD. The final values are very similar for the two methods.

3 Variance reduction

Here we look at the values of values of $\operatorname{Tr}(\Sigma(q))$ during the ISSGD training from the previous section (which led to figure 2 and figure 3).

We would like to compare the values of $\operatorname{Tr}(\Sigma(q))$ for $(q_{\textsc{ideal}},q_{\textsc{stale}},q_{\textsc{unif}})$. Note that $q_{\textsc{unif}}$ does not mean here that we trained with the regular SGD (that assigns the same probability to each training example). It means that, during ISSGD training, we can report the value of $\operatorname{Tr}(\Sigma(q))$ that we would get if we performed the next step with regular SGD. In figure 4, we refer to this as “SGD, ideal”. We compare it to “ISSGD, ideal”, which corresponds to the best possible situation for our method, $\operatorname{Tr}(\Sigma(q_{\textsc{ideal}}))$, which is not necessarily achieved in practice.

In section B.3 of the Appendix we describe how we add a constant to the probability weights in order to make the method more robust. We are trading away potential gains to make training more stable.

On both plots of figure 4 we show the “ideal” measurements that we would get with exact probability weights, and we compare with the “stale” measurements that we get with probability weights used in the actual experiments, which are all stale to varying degrees. On those stale curves, we show the effects of using the actual additive constant to the probability weights, and the effects of using an alternate one. Bear in mind that, in both cases, the validation loss reached its minimum in around 30 minutes, and these plots are shown for 2.5 hours. Also, these are the $\operatorname{Tr}(\Sigma(q))$ with respect to the gradient on the training set. One naturally expects that gradient to converge to 0.0 during the overfitting regime.

Note that in figure 4 we report the square root of those values in order to have it be on the same scale and the gradients themselves (this is analogous to reporting $\sigma$ instead of $\sigma^{2}$).

Future work

One of the constraints that we are facing is that proposition 1 works with models with only fully-connected layers. This rules out all the convolutional neural networks, which are very popular and very useful.

One alternative would be to use an approximate formula for the individual gradient norms for convolutional layers. Either something naive (such as applying proposition 1 without proper justification), or possibly even ignoring the contributions from those layers. This would yield an importance sampling scheme that would be of lesser quality, but it would also be hard to evaluate how much we actually suffer for that.

We have avoided direct comparisons with ASGD in this paper because we are not currently in possession of a good production-quality ASGD implementation. We would certainly like to see how ASGD could be combined with ISSGD, whether this would create positive interactions or whether the two methods would impede each other.

Note that there are alternative ways to combine our method with ASGD, and they are not equally promising. Our recommendation would be to get rid of the master/workers distinction and have only workers (or “peers”) along with a parameter server (or shared memory, or whatever synchronization method is used to aggregate the gradients and parameters). Whenever a gradient contribution is computed, the importance weights can be obtained at the same time. These can be shared in the same way that the gradients are shared, so that all the workers are able to use the importance weights to run ISSGD steps.

Conclusion

We have introduced a novel method for distributing neural network training by using multiple machines to search for the most informative examples to train on. This method led to significant improvements in training time on permutation invariant SVHN. Our results demonstrated that importance sampling reduced the variance of the gradient estimate, even in the distributed setting where the importance weights are not exact. One area for future work is extending this method to models that use parameter sharing (such as convnets and RNNs), either by finding a new formula for per-example gradient norms or by finding an approximation to the gradient norm that is easy to compute. Finally, much of the most successful work on data parallel distributed deep learning has used a variant of Asynchronous SGD. It would be useful to understand exactly how our method compares with Asynchronous SGD and to see if further improvement is gained by using both approaches simultaneously.

The authors would like to acknowledge the support of the following agencies for research funding and computing support: NSERC, Calcul Québec, Compute Canada, the Canada Research Chairs and CIFAR, and the Nuance Foundation.

The authors would like to acknowledge the stimulating discussions with Ian Goodfellow, Zack Lipton, Kari Torkkola, Daniel Abolafia, and Ethan Holly.

We would also like to thank the developers of Theano. http://deeplearning.net/software/theano/

References

Appendix A Importance sampling in theory

The context requires that $q(x)>0$ whenever $p(x)>0$. We know that the importance sampling estimator

Let $\Sigma(q)$ be the covariance of that estimator, where we include $q$ in the notation to be explicit about the fact that it depends on the choice of $q$.

Then the trace of $\Sigma(q)$ is minimized by the following optimal proposal $q^{*}$ :

This proof is almost exactly the same as the well-known result in one dimension involving Jensen’s inequality. Everything follows from the decision to minimize $Tr\left(\Sigma\right)$ and the choice of $q^{*}$. Nevertheless, we include it here so the reader can get a feeling for where $q^{*}$ comes into play.

When sampling from $q(x)$ instead of $p(x)$, we are looking at how the unbiased estimator

which has mean $\mu$ and covariance $\Sigma(q)$. We make use of the fact that the trace is a linear function, and that $\operatorname{Tr}(\mu\mu^{T})=\left\|\mu\right\|_{2}^{2}$. The trace of the covariance is given by

There is nothing to do about the $\left\|\mu\right\|_{2}^{2}$ term since it does not depend on the proposal $q(x)$. Using Jensen’s inequality, we get that

which is the minimal value achievable, so $q^{*}$ is indeed the best proposal in terms of minimizing $Tr(\Sigma(q))$. ∎

Note also that the single-dimension equivalent, mentioned in section 3.1, is a direct corollary of this proposition, because the Euclidean norm turns into the absolute value.

Then we have that the trace of the covariance of the importance sampling estimator is given by

To sample from $q(x)$ we just normalize the probability weights

and we sample from a multinomial distribution with argument $(\omega_{1},\ldots,\omega_{N})$ to pick the corresponding element in $\mathcal{D}$.

We start from equation (13), which applies to a general proposal $q$. In fact, we make it to equation (A.1) still without making assumptions on $q$. At that point we can look at the normalizing constant of $q$, which is equal to

A.2 Dealing with minibatches

Consider a multi-layer perceptron (MLP) applied to minibatches of size $N$, and with loss $\mathcal{L}=\mathcal{L}_{1}+\ldots+\mathcal{L}_{N}$, where $\mathcal{L}_{n}$ represents the loss contribution from element $n$ of the minibatch.

Let $(W,b)$ be the weights and biases at any particular fully-connected layer so that $XW+b=Y$, where $X$ are the inputs to that layer and $Y$ are the outputs.

The gradients with respect to the parameters are given by

where the values $\left(\frac{\partial\mathcal{L}_{n}}{\partial W},\frac{\partial\mathcal{L}_{n}}{\partial b}\right)$ refer to the particular contributions coming from element $n$ of the minibatch. Then we have that

where the notation $X[n,:]$ refers to row $n$ of $X$, and similarly for $\frac{\partial\mathcal{L}}{\partial Y}[n,:]$.

That is, we have a compact formula for the Euclidean norms of the gradients of the parameters, evaluated for each $N$ elements of the minibatch independently.

The usual backpropagation rules give us that

All the backpropagation rules can be inferred by analyzing the shapes of the quantities involved and noticing that only one answer can make sense. If we focus on $\mathcal{L}_{n}$ for some $n\in\{1,\ldots,N\}$, then we can see that

Note here that $X[n,:]^{T}\frac{\partial\mathcal{L}}{\partial Y}[n,:]$ is the outer product of two vectors, which yields a rank-1 matrix of the proper shape for $\frac{\partial\mathcal{L}_{n}}{\partial W}$. Similarly, we have that $X[n,:]X[n,:]^{T}=\|X[n,:]\|_{2}^{2}$ is a 1x1 matrix, which can be treated as a real number in all situations.

The conclusion for $\left\|\frac{\partial\mathcal{L}_{n}}{\partial b}\right\|_{2}^{2}$ follows automatically from taking the norm of the corresponding expression in equation (18). Some more work is required for $\left\|\frac{\partial\mathcal{L}_{n}}{\partial W}\right\|_{2}^{2}$. We will make use of the three following properties of matrix traces.

$\left\|A\right\|_{\textrm{F}}^{2}=\operatorname{Tr}(AA^{T})$

$\operatorname{Tr}(ABC)=\operatorname{Tr}(BCA)=\operatorname{Tr}(CAB)$

One might wonder why we are interested in computing the Frobenius norm of the matrix $\frac{\partial\mathcal{L}_{n}}{\partial W}$ instead of its L2-norm. The reason is that when running SGD we serialize all the parameters as a flat vector, and it is the L2-norm of that vector that we want to compute. We flatten the matrices, and the following equality reveals why this means that we want the Frobenius norms of our matrix-shaped parameters :

Appendix B Distributed implementation of ISSGD

We have tried training without that staleness threshold and it is hard to see a difference. Adding more workers naturally lowers the average staleness of probability weights, because more workers can update them more frequently. If it were not the cost of communicating the model parameters, we could argue that a sufficiently large number of workers would simulate an oracle perfectly.

To report values of $\operatorname{Tr}(\Sigma(q))$, we need to be able to compute the actual expected gradient over the whole training set. We refer to that quantity as the true gradient $g^{\textsc{true}}=\frac{1}{N}\sum g(x_{n})$, but we never really compute it due to practical reasons. This would entail reporting the gradient for each chunk of the training set and aggregating everything. This is precisely the kind of thing that we avoid with ISSGD.

Instead we compute the gradients of the parameter for each minibatch, and we report the L2-norm of those. We then average those values. This produces an upper-bound to the actual value of $\left\|g^{\textsc{true}}\right\|_{2}$.

One important thing to note is that the equations (7), (8) and (9) each have the $\left\|g^{\textsc{true}}\right\|_{2}^{2}$ term, so any approximation of that term, provided that it is the same for all three, will not alter the respective order of $\operatorname{Tr}(\Sigma(q_{\textsc{ideal}})),\operatorname{Tr}(\Sigma(q_{\textsc{stale}})),\operatorname{Tr}(\Sigma(q_{\textsc{unif}}))$

Moreover, when are getting close to the end of the training, we should have that $\left\|g^{\textrm{true}}\right\|_{2}$ is getting close to zero. That is, the gradient is zero when we are close to an optimum. This does not meant that the individual gradients are all zero, however. But when our upper-bound on $\left\|g^{\textsc{true}}\right\|_{2}$ is getting close to being insignificant, then we can tell for sure that our values computed for the three $\operatorname{Tr}(\Sigma(q))$ are very close to their exact values.

B.3 Smoothing probability weights

Sometimes we can end up with probability weights that fluctuate too rapidly. This can lead to some problems in a situation where one training sample $x_{n}$ is assigned a small probability weight $\epsilon$, when compared to the other probability weights. Things normally balance out because $x_{n}$ has a probability proportional to $\epsilon$, and when it gets selected its gradient contribution $g(x_{n})$ gets scaled by $1/\epsilon$. The resulting contribution is a gradient of norm $\approx 1$.

However, when that gradient changes quickly (and probability weight along with it), it is possible to get into a situation where the gradient computed on the master is now much larger (due to the parameters having changed), and it still gets divided by $\epsilon$ when selected. This does not affect the bias, but it affects the stability of the method in the long term. When running for an indefinitely long period, it is dangerous to having a time bomb in the algorithm that has a small probability of ruining everything.

We had some ideas for using an adaptive method to compute this smoothing constant, but this was not explored due to the large number of other hyperparameters to study. One suggestion was to look at the entropy of the distribution of the $\left\{\omega_{n}\right\}_{n=0}^{\infty}$ that determine which training sample are going to be used. With a smoothing constant sufficiently large, we can bring this entropy down to any target level (or down to regular SGD when that constant is infinite).