M EMORY A UGMENTED O PTIMIZERS FOR D EEP L EARNING
A BSTRACT
Popular approaches for minimizing loss in data-driven learning often involve an abstraction or an explicit retention of the history of gradients for efﬁcient parameter updates.
The aggregated history of gradients nudges the parameter updates in the right direction even when the gradients at any given step are not informative.
Although the history of gradients summarized in meta-parameters or explicitly stored in memory has been shown effective in theory and practice, the question of whether all or only a subset of the gradients in the history are sufﬁcient in deciding the parameter updates remains unanswered.
In this paper, we propose a framework of memory-augmented gradient descent optimizers that retain a limited view of their gradient history in their internal memory.
Such optimizers scale well to large real-life datasets, and our experiments show that the memory augmented extensions of standard optimizers enjoy accelerated convergence and improved performance on a majority of computer vision and language tasks that we considered.
Additionally, we prove that the proposed class of optimizers with ﬁxed-size memory converge under assumptions of strong convexity, regardless of which gradients are selected or how they are linearly combined to form the update step.
I NTRODUCTION
Gradient-based learning involves minimizing a scalar-valued function F ( θ ) with respect to a parameter vector θ ∈ R d using an iterative procedure.
When training over n examples of input-output pairs, ( x i , y i ) , the optimization problem boils down to solving min θ F ( θ ) , where
F ( θ ) = 1 n n (cid:88) i =
L i ( θ ) , with L i ( θ ) =
L ( M θ ( x i ) , y i ) , for some problem-dependent loss function L and a predictive model M parameterized by θ .
Stochastic Gradient Descent (SGD) (Robbins and Monro, 1951) is one common method used to tackle this problem, and is often preferred to full-batch Gradient Descent when the quantity of data required to train θ is large, since it can be more efﬁcient to measure a single component of the gradient (or a mini-batch of component gradients (Bottou, 1999)), and move in a noisy direction, than to compute a full gradient at each time step.
Several techniques have been proposed to further accelerate the convergence of SGD (Ghadimi and Lan, 2016; Allen-Zhu and Yuan, 2016; Zhou et al., 2018; Xu et al., 2017).
These include approaches that maintain a knowledge of previous gradients implicitly by summarizing them in a momentum buffer (Polyak, 1964b), and potentially adapting the learning rate based on the gradient statistics (Duchi et al., 2011; Hinton et al., 2012; Zeiler, 2012; Kingma and Ba, 2015).
Optimization techniques such as SGD with Momentum (Polyak, 1964a), AdaGrad (Duchi et al., 2011), RMSprop (Hinton et al., 2012), AdaDelta (Zeiler, 2012), and Adam (Kingma and Ba, 2015) maintain a set of buffers that track running moments of the gradients.
Such light-weight techniques have shown signiﬁcant application advantage in practice, whereas in theory, algorithms that store all of the gradients, like SAG (Roux et al., 2012) and SAGA (Defazio et al., 2014), can achieve better convergence.
A drawback of full-history methods is that the memory requirement linearly increases with the size of the data.
The spectrum of gradient-based approaches (Figure 1) that use a knowledge of the past gradients has Adam, RMSprop and SGD with Momentum on the one end and approaches like SAG and
SAGA on the other.
An optimization algorithm in the middle of this spectrum can use less memory than full-history methods while providing richer updates than adaptive-gradient methods.
Limited-Memory BFGS (Nocedal, 1980) and its online, stochastic form, oLBFGS (Schraudolph et al., 2007) aim to exploit this trade-off, but may not converge faster than SGD (Mokhtari et al., 2015), and are thus disadvantaged compared to accelerated SGD variants.
As a stepping stone towards reaping the advantages of both ends of the spectrum, we propose using memory to augment conventional algorithms.
This memory is used to store a small set of critical gradients (e.g., gradients with a larger (cid:96) 2 -norm) that are occasionally updated with newer gradients.
Speciﬁcally, rather than storing gradients for all the examples in the dataset, our proposed approach aims to ﬁnd a smaller set of gradients to store that will aid in the optimization process.
Through extensive experiments we show that this marginal increase in the memory requirement (Figure 1) provides accelerated convergence and even improved test performances in majority of the experiments.
In this work, we:
• Present a framework of memory-augmented optimizers † , § , ‡ compatible with popular gradient-based algorithms.
• Prove theoretically that such algorithms converge for smooth strongly convex objectives.
• Show that the proposed memory augmented optimizers can lead to faster convergence and better ﬁnal performance through exhaustive empirical study on eight different architectures among four different tasks (classiﬁcation, language modeling, natural language inference, and dialogue) and six different datasets.
• Demonstrate the memory augmented optimizers’ robustness to gradient selection heuristics, replacement strategies, or the aggregation techniques to summarize information in the memory.
M EMORY -A UGMENTED O PTIMIZERS
In this work, we propose augmenting standard optimizers with a ﬁxed size memory of past gradients.
The gradients are stored in the limited memory only when they are deemed critical as deﬁned by their (cid:96) 2 -norm.
Gradient Selection to the Memory Letting g t := ∇ θ L i ( θ t ) denote a component (or mini-batch) gradient at time-step t , we use (cid:107) g t (cid:107) 2 as a scalar indicator of the importance of the gradient, which serves as a proxy for the priority of a gradient to remain in the memory buffer g c .
In order to ensure that the buffer eventually ﬂushes out stale gradients, the proxy norms for gradients in the buffer are scaled down by a hyperparameter decay factor, denoted decay ∈ [0 , 1) .
The proposed approach maintains a gradient buffer of ﬁxed capacity C and stores the gradients selected by a chosen heuristic.
We refer our heuristic as the critical gradients , which stores the top C gradients by using the (cid:96) 2 -norm of g t to determine priority.
Gradients in the buffer are stored as tuples ( (cid:107) g t (cid:107) 2 , g t ) , akin to a key-value pair in a dictionary structure, where the key (cid:107) g t (cid:107) 2 is referred to as the proxy norm, and decayed at each time step t by a decay factor decay ∈ [0 , 1) .
At any iteration t during training, the gradient in a full capacity priority buffer with the smallest (cid:96) 2 proxy-norm is replaced by the current mini-batch gradient g t if (cid:107) g t (cid:107) 2 is greater than the smallest (cid:96) 2 proxy-norm in the buffer.
The multiplicative decay factor ensures that the buffer is frequently refreshed with more recent gradients.
Note that only the proxy-norm used for the heuristic rule is decayed by the decay factor decay ∈ [0 , 1) , while the gradient itself is not affected.
This technique of storing the critical gradients is general enough to be employed in any deep learning model, and can be easily be combined with many existing adaptive gradient optimization algorithms, such as those described in (Polyak, 1964b; Zeiler, 2012; Duchi et al., 2011; Hinton et al., 2012; Kingma and Ba, 2015; Lucas et al., 2019).
Critical Gradient Stochastic Gradient Descent Critical Gradient Stochastic Gradient Descent ( SGD C ) is the explicit integration of critical gradients into SGD.
Speciﬁcally, the iteration comes in the form
[Equation]
where g c is the set of critical gradients and aggr denotes an aggregation function which is used to combine the critical gradients with the current-iteration gradient.
We propose two possible functions for aggr including mean , the average of g t and all critical gradients, and sum , the addition of g t to the average of all critical gradients.
Mathematically, for a buffer of size C these are deﬁned as
[Equation]
In general, we ﬁnd optimization to be robust to the speciﬁc choice of aggregation function, but we do observe that adaptive learning rate methods (Adam, RMSprop) perform best with mean , while convential SGD (with or without Polyak momentum) performs best with sum .
Critical Gradients with Accelerated Methods Agarwal et al.
(2019) show that curvature in adaptive gradients can be estimated without loss of generality with a ﬁnite moving window of past gradients.
Following which we propose critical gradient extensions of SGD with Momentum (SGDM),
RMSprop, and Adam.
The retention of the critical gradients can be naturally extended to more ﬁrstorder optimization methods with the replacement of g t in those algorithms with aggr ( g t , g c ) .
The update steps for the corresponding C versions are as follows:
[Equation]
The general algorithm to integrate critical gradients into any optimizer is provided in Algorithm 1 in
Appendix A.
Replacing g t with aggr allows the optimization function to retain key gradients and optimize the parameters.
The acceleration provided by the topC gradients complements the adaptive learning rate methods by not only computing the learning rate with respect to g t but also with respect to the critical gradients from the past.
We observe that all C -extended optimizers in our empirical analysis and experiments in § 5 lead to faster convergence than their vanilla
P ROOF OF C ONVERGENCE
Consider the SGD C update in Equation (1), where aggr ( g t , g c ) outputs a linear combination of the critical gradients in memory g k ∈ g c , and the stochastic gradient g t , which is computed at the current parameters θ t .
When the staleness of gradients in the memory is bounded, we can prove convergence of this method through the lens of linear multi-step methods .
Formally, suppose there exists an integer
K > 0 such that g k ∈ g c at iteration t , implies t − k < K , and suppose the objective F : R d → R is twice-continuously differentiable, L -smooth, and µ -strongly convex, with 0 < µ ≤ L .
Assume that for all k , the stochastic gradient g k is a random vector satisfying
E [ g k ] = ∇ f
( θ k ) and E (cid:104) (cid:107) g k [Equation] f ( θ k ) (cid:107) 2 (cid:105) ≤ σ 2 , for some ﬁnite constant σ 2 .
Letting ζ k :
= g k − ∇ f ( θ k ) denote the gradient noise at iteration k , we take these gradient noise terms to be mutually independent.
For examples of functions satisfying these properties, see, e.g., (Nesterov, 2004; Bubeck, 2015).
Examples of typical tasks satisfying these assumptions are (cid:96) 2 -regularized logistic regression and (cid:96) 2 -regularized least-squares regression (i.e., ridge regression).
Taken together, these properties imply that the Hessian ∇ 2 F ( θ ) exists, and for all θ ∈ R d , the eigenvalues of ∇ 2 F ( θ ) lie in the interval [ µ, L ] .
In this case, let θ (cid:63) denote the unique minimizer of F .
To prove convergence, we view SGD C as a linear multi-step method, i.e., the parameters at the current time step are updated with a linear combination of gradients from the past K time steps.
From this view, we can describe the evolution of the algorithm according to a discrete-time linear dynamical system, the convergence of which is characterized by the spectral properties of the system matrix.
Our convergence theorem relies on the largest singular value of this system matrix, but can be strengthened to instead rely on the spectral radius of the system matrix by constructing an appropriate Lyapnuov function containing a time-invariant system matrix satisfying a particular
Integral Quadratic Constraint (IQC) (Lessard et al., 2016).
However, one downside of the IQC framework is that one must typically solve a semideﬁnite program (e.g., using a numerical solver) to obtain explicit convergence rates (Hu and Lessard, 2017), whereas Theorem 1 provides an analytical convergence rate.
Theorem 1 (Linear Convergence of SGD C ) .
Let
[Equation]
where ρ (Λ) are the singular values of the numerical matrix
[Equation]
with λ 0 = 1 − αη 0 and λ k = − αw k η k for k > 0 .
If the step-size α > 0 is chosen sufﬁciently small, such that q ( α,K ) < 1 , then
[Equation]
where the expectation on the left hand side is with respect to the joint distribution over the random noise variables ζ k , . .
, ζ 1 .
Theorem 1 shows that SGD C can be made to converge exponentially fast to a neighbourhood of the solution, the size of which is proportional to the variance of the stochastic gradients σ 2 and the step-size α 2 .
If K = 1 (i.e., the memory g c , at all times t , only contains the mostly recently computed gradient), then SGD C reduces to regular stochastic gradient descent, and the rate q ( α,K ) in
Theorem 1 reduces to the well known convergence rate of SGD for smooth strongly convex functions ( max η ∈ [ µ,L ] | 1 − αηL | ), and the variance bound reduces to the standard variance bound for SGD with a constant step-size ( α 2 σ 2 / (1 − q 2 ) ), see, e.g., (Bottou et al., 2018;
Assran and Rabbat, 2020).
For K > 1 , it may be possible to choose the step-size α and the aggregation weights { w k } to obtain accelerated convergence rates, faster than SGD.
For example, accelerated gradient methods such as Polyak and Nesterov momentum can be viewed as multi-step methods with K = 2 .
Note that although K appears in the numerator of the coefﬁcient of the variance term, it is also present in the rate term in the denominator, where faster convergence rates, smaller q ( α,K ) , directly lead to smaller variance bounds.
Note that expressing the convergence rate of a multi-step method as the roots of a polynomial, as we do in Theorem 1, is not new; see, e.g., Polyak (1964a).
It is also straightforward to extend Theorem 1 to replace the bounded variance assumption (common in the stochastic approximation setting) with other conditions, such as the Strong Growth Condition or the Weak Growth Condition (Bottou et al., 2018), which are more tailored to stochastic ﬁnite-sum models.
R ELATED W ORK
Past-Gradients-Summarizing Optimizers Several popular optimization algorithms incorporate the history of previous gradients through a summarizing variable.
In practice, this is often done by means of a decaying average, which is updated with each iteration.
SGD with Momentum (SGDM) (Polyak, 1964b) uses this summary variable as a means of stabilizing descent along one path even if the current gradient points are uninformative.
A variant of SGDM uses Nesterov momentum (NAG)
(Nesterov, 1983; Sutskever et al., 2013) and applies the velocity to the current gradient as a correction factor.
SGD C shares similarity with NAG except in SGD C the “momentum" is computed only using
C selected gradients.
In a different use of a summarizing variable, AdaGrad (Duchi et al., 2011) keeps a running sum of squared gradients which is used to dynamically adjust the learning rate.
AdaDelta (Zeiler, 2012) and RMSprop (Hinton et al., 2012) replace AdaGrad’s running sum with a decaying average of the squared gradients.
Adam (and AdaMax) (Kingma and Ba, 2015) uses a combination of momentum and a decaying-average-dependent learning rate.
In general, the C variants ( SGD C ,
Adam C , RMSprop C , SGDM C ) provide an additional “layer" to the base algorithms that ensures any parameter computed within the algorithm incorporates the critical gradients.
Although optimization using a ﬁxed size buffer has been in existence, the novelty of the C variants comes from adding a heuristic to keep only the gradients that provide the largest displacements towards the minimum.
Memory-Enhanced Optimizers
Our algorithm belongs to a class of optimization methods which are aware of a limited history of gradients from previous iterations.
One fellow algorithm from this class is Limited-History BFGS (LBFGS)
(Nocedal, 1980), an optimization technique which approximates the Hessian of a function using a limited history of gradients, then uses the estimated
Hessian to compute the descent direction for optimization.
An online version of LBFGS (oLBFGS), well-suited for a machine learning context, was proposed by Schraudolph et al.
These
LBFGS methods utilize a moving window of past gradients, rather than our idea to maintain “critical" gradients based on a well-deﬁned metric.
The storage of past gradients as a means to improve optimizer convergence is not a novel idea; both Stochastic Accelerated Gradient (SAG) (Roux et al., 2012) and SAGA (Defazio et al., 2014) do so in the context of accelerating stochastic gradient descent, while Stochastic Dual Coordinate
Ascent (SDCA) (Shalev-Shwartz and Zhang, 2013) uses previous gradients for coordinate ascent.
Empirically and theoretically, these techniques have been shown to yield good performance.
The
SAG and SAGA algorithms differ in their weight update step, but both involve storage of up to n gradients, where n is the number of training examples, which can be costly.
Our method relies on the storage of a ﬁxed number of gradients and uses a heuristic to decide whether to store a gradient in a training-dataset-size-independent memory.
Our class of optimizers can natively run in a batch-based training loop unlike originally presented SAG and SAGA.
Gazagnadou et al.
(2019) address this issue in their recent work.
We also extend our method beyond SGD to several other ﬁrst-order stochastic gradient descent methods.
SAG and SAGA belong to the family of variance-reduced gradient-based algorithms (J. Reddi et al., 2015).
Such algorithms theoretically and empirically outperform SGD, though with the trade-off of having a large cost in terms of memory or gradient evaluations.
Some algorithms in this framework attempt to bridge both costs, for instance a variant of Stochastic Variance Reduced
Gradient (SVRG) (Johnson and Zhang, 2013) stores intermediate gradients to avoid recomputing them later, and StochAstic Recursive grAdient algoritHm (SARAH) (Nguyen et al., 2017) uses a summarizing variable.
Our method attempts to recapture the beneﬁts of variance-reduced methods without signiﬁcant computation or memory overhead.
Optimizing Optimizers Our optimizer method also shares similarities with techniques which seek to optimize optimizers; that is, techniques which automatically learn the best way to accomplish the task of tuning the optimizer used in an outer-problem.
The architecture proposed by Metz et al.
(2020) uses an LSTM per tensor of the network paramters, with each LSTM being passed the gradient norms to inﬂuence the parameter update step, echoing our technique’s use of gradient norm as a critical metric to parameter updates.
In a similar vein, Andrychowicz et al.
(2016) also use LSTMs, which are fed the complete gradient, and use the LSTM’s recurrence as a means of implicitly accessing past gradients.
Li and Malik (2017) uses a reinforcement learning-based approach and deﬁnes a state-space which includes a recent history of gradients.
While our method equally maintains a ledger of past gradients, unlike this latter approach we use (cid:96) 2 -norm as a metric to pick out the critical gradients instead of arbitrarily keeping a recent history.
Motivation for explicit memory Motivation for explicit memory-augmented optimizers over the ones that maintain an implicit summarization of history comes from memory-augmented neural networks like Neural Turing Machines (NTMs) and their variants (Graves et al., 2014; 2016; Gulcehre et al., 2018).
While a simple recurrent architecture like LSTM (Hochreiter and Schmidhuber, 1997) integrates information in its cell state which is a single vector, NTMs maintain a memory matrix which stores a set of such cell state vectors.
This richer parameterization of memory helps NTMs in learning complex algorithmic tasks which are difﬁcult to learn for an LSTM.
Analogous to LSTMs and NTMs, SGDM and our C -extended optimizers maintain a single vector and a set of vectors respectively as their memory.
While our current formulation for memory-augmented optimizers uses heuristics to choose what information to store in the memory and how to use it, one could automatically learn both criteria and hence learn to optimize.
Motivation for identifying critical gradients
Likewise, recent work (Paul et al., 2021) has identiﬁed that samples with larger gradient norms can be used to identify a smaller set of training data that is important for generalization, and thereby prune the data as a pre-processing step.
Our proposed approach of augmenting optimizers with memory buffer of critical gradients can be seen as a softer version of data-pruning in that the buffer retains gradients through the heuristic, thereby performing updates using only the important gradients throughout the entire training, without discarding data.
E XPERIMENTS
We compare the proposed class of optimizers augmented with memory ( C variants) with their vanilla versions on a variety of deep learning architectures * .
To understand the performance on common deep learning datasets, we experiment with shallow/deep convolutional neural network architectures (CO) on CIFAR 10/100 (Krizhevsky et al., 2009) respectively; Bi-LSTM (BL), InferSent (I), and a text-based convolutional architecture (C) on the Stanford Natural Language Inference (SNLI)
(Bowman et al., 2015) dataset; LSTM on word level language modeling (LS) with the PennTreeBank (Marcus et al., 1993), and WikiText (Merity et al., 2016) datasets; RoBERTa-Base (RoB) and BiLSTM (BL) on language generation in dialogue task with MultiWoZ 2.0 dataset (Budzianowski et al., 2018).
Additionally, we preform analysis experiments using logistic regression (LR) and multi layer perceptrons (MLP) on the MNIST digit classiﬁcation dataset (LeCun and Cortes, 2010).
Across the experiments, we compare all the 8 optimizers averaged over 5 different runs with different seed values— Adam , SGD , SGDM , RMSprop , Adam C , SGD C , SGDM C , and RMSprop C ontasks — CIFAR100, CIFAR10, SNLI-I, SNLI-C, SNLI-BL, MWoZ-RoB, MWoZ-BL, PTB-LS, and
WIKITEXT-LS.
For a fair comparison, we follow the standard experimental protocol for all the experiments and tune all the hyperparameters of an optimizer individually for every task.
Our experimental results are aggregated from 5 independent runs, with the hyperparameters for each optimizer extensively tuned.
This involves tuning the learning rate in all optimizers, the topC and decay parameters in all C algorithms, and all optimizer-speciﬁc hyperparameters in both the vanilla versions and their C counterparts.
A full description of the values used to tune the various parameters, architecture and dataset details are in Appendix §E, §D, and §C respectively.
The best set of hyperparameters was selected based on which ones yielded the best validation performance (i.e. highest accuracy, highest BLEU score, or lowest perplexity). [Figure]
Figure 2(a) shows that the proposed C variants provided a consistent acceleration in convergence over their vanilla versions.
The accelerated convergence also resulted in improved test performance, as highlighted in Figure 2(b) (Exhaustive list is provided in Appendix §G).
In addition to the aforementioned improvements in vanilla-vs- C comparisons, the C optimizers yielded the best test performance across all 8 optimizers in 6/9 tasks.
Of the three remaining tasks, the C versions performed close to the vanilla versions.
In the case of MWoZ − BL RMSprop C the C version showed marginally accelerated convergence close to the optimal solution.
Furthermore, Figure 3 compares the validation performance after each epoch when trained with the different optimizers on the different tasks.
One immediate observation was that in language experiments – SNLI, PTB, MWoZ – the C versions stayed above most if not all of the vanilla optimizers.
While that separation is not as clear in the vision tasks, we observed that the pairwise comparison mostly put the C optimizer at an advantage in accelerated convergence.
A NALYSIS
Buffer
Staleness Bound A key assumption behind the theoretical convergence of SGD C is the existence of an upper bound on the buffer gradients’ staleness (i.e the number of iterations a gradient remains in the critical buffer).
We analysed the distribution of average number of steps a gradient stays in the memory for the different values of topC and decay in Figure 4.
We observe that tuning the parameters ensures that gradients in the buffer are refreshed frequently, allowing for a renewal of information.
Later, in analysing the effect of decay and topC on the performance in a task, we observe that the higher performance of the C methods correlate with the ﬁndings of this analysis.
Use of Critical Gradients Our Critical Gradient method utilizes a buffer which is ﬁlled by using a norm-based, “King of the Hill" sampling of incoming gradients, i.e. gradients always stored and sorted in decreasing order by (cid:96) 2 -norm, with the smallest norm always being the one removed.
We conduct ablation studies to probe the three assumptions of critical gradients using LSTM architecture with PTB dataset.
: (a) the use of norm as a controlling metric, (b) the removal of the smallest-norm [Figure] entry when adding to a full buffer, and (c) the update step depending on an ensemble of buffer gradients via the mean rather than on select gradients.
[Figure]
Addressing assumption (a) we test the use of maximal norm sampling (Figure 5(a)) by comparing with Cosine Similarity Sampling (CSS), Cosine Diversity Sampling (CDS), Mean-Norm Diversity
Sampling (MNDS), and random sampling via "coin toss".
Figure 5(b) depicts our algorithm alongside different gradient replacement methods, including Random Replacement (RR) and Norm-Controlled
Probabilistic Replacement (NCPR), as well as comparing against keeping a running First-In-Fist-Out (FIFO) queue and only using the smallest (" bottomC ") gradients.
Finally we compare the typical optimizer aggregation method with the buffer gradients minimum, maximum, and median norm (Figure 5(c)).
A full description of these techniques can be found in Appendix §F.9.
We observe in most cases, the sampling (a) and replacement (b) strategies for gradients in the buffer and all ablations converge regardless of these considerations.
While the Critical Gradients formulation is not unique as a high-performing implementation of a memory-augmented optimizer, it [Figure] remains conceptually straightforward, easy to implement, and the buffer properties are deterministic as compared to some of our ablation studies which were probabilistic.
More experiments need to be done to verify that other sampling and selection strategies indeed perform similar in all other tasks and architectures.
Experiments comparing the performance of mean-based aggregation against aggregation based on individual gradients show that aggregation based on the ensemble must be employed, and that the optimizers beneﬁt from the added information.
Effect of hyperparameters We study the performance’s sensitivity to decay and topC using
Adam C on different tasks.
[Figure]
architectures explicit evidence gradient staleness
The decay ∈ [0 , 1) parameter allows to tune the retention of critical gradient information from the past.
Setting decay = 0 reduces the priority of new gradients added to the buffer, making the
C variants behave like an average over a moving window of length topC .
On the other extreme, decay = 1 does not decrease the priorities of the past gradients and allows them to be in the buffer for a much longer period hindering the convergence of the optimizers.
Also, decay maintains the staleness of the gradients by exponentially decaying their priority in the memory buffer.
We experimented with Adam C optimizer for the different values of decay with every other parameter ﬁxed (Figure 6).
We observed that decay values closer to 1 did not have any advantage at best or leads to inferior performance on an unseen set at worst in most of the tasks.
[Equation] ) deﬁnes the size of the buffer that holds the past gradients.
When topC is set to 0 ,
C variants behave like the base algorithms and when topC is the size of the dataset, the optimizer becomes SAG-like.
Also, topC is directly proportional to the staleness of the gradients in the buffer.
A higher topC , though, guarantees better expressibility also becomes responsible for more stale gradients in the buffer.
Bengio et al.
(2020) also observes a similar issue in momentum used in
TD-learning and proposes a correction factor to alleviate the issue.
Here, the staleness is contained with the parameters topC and decay .
As an empirical evidence, we observe that the optimizers perform well on the validation set with lower topC values.
In most of the experiments, topC >was either not useful or was hurting the performance (also in Appendix F).
C ONCLUSION
We propose a general method to enhance the performance of an optimizer with an augmented memory buffer to maintain a limited set of critical gradients from the history.
The proposed memory mechanism, when integrated with several state-of-the-art optimizers, accelerated the learning in all cases and also improved their performance in several cases.
The C family of optimizers proposed in this paper are the ﬁrst steps towards designing smart optimizers that can learn to optimize.
Some of the immediate future research directions include learning what to store in the memory and learning how to integrate the information in memory to the gradient update steps. A CKNOWLEDGEMENTS
We would like to acknowledge Compute Canada and Calcul Quebec for providing computing resources used in this work.
The authors would also like to thank members of Chandar Research Lab,
Mila for helping with the code reviews and reviewing the manuscripts. SC is supported by a Canada CIFAR AI Chair and an NSERC Discovery Grant A G ENERAL C RITICAL G RADIENT A LGORITHM
We present the pseudocode for the Critical Gradient algorithm below.
Algorithm
Deﬁne #epochs E , #gradient updates B , model θ , learning rate α Deﬁne optimizer O ( g t , θ t , h ) , which produces a parameter update given gradient g t , parameter θ t and hyperparameters h Deﬁne aggregation function aggr ( g t , g c ) which outputs a linear combination of g t and the entries in g c
Initialize empty critical gradient buffer { g c } with decay rate d for s = 0 to E do for t = 0 to B do
Sample an input-output pair ( x i , y i )
g t ← ∇ θ L i ( θ t ) θ t +1 ← O
( aggr ( g t , g c ) , θ t , h ) if { g c } is not full then Add g t to { g c } elseif || g t || 2 > || min( { g c } )
|| 2 then
Remove min( { g c } ) from { g c } Add g t to { g c } end if end if for g c ∈ { g c } do || g c || 2 ← d · || g c || 2 end for end for end for
B C ONVERGENCE P ROOF
Proof.
Recall that if F : R d → R is twice continuously differentiable, then for all x, y ∈ R d
[Equation]
Letting r k := θ k − θ (cid:63) denote the suboptimality of the parameters at iteration k and applying (7) with y
= θ k and x = θ (cid:63) , we get that
[Equation]
where H k =
[Equation]
where ζ k denotes the gradient noise at iteration k .
Using this and bounded staleness assumption, we have that
[Equation] [Equation]
where { w t,k } K − 1 k =0 are non-negative scalars in the closed interval [0 , 1] used to linearly aggregate the gradients in memory with the most recent gradient g t .
For gradients g t − k
that are not stored in memory, the corresponding aggregation scalar w t,k is equal to 0 .
Without loss of generality, we take w t, 0 = 1 , since we can always incorporate it into the step-size α and re-scale the other weights w t, 1 , . . .
, w t,K − 1 accordingly.
Substituting (8) into (1) and subtracting θ (cid:63) from each side, we get that
[Equation]
Thus, we have that SGD C evolves according to the linear system  
[Equation]
  where A t = (cid:34) T
[Equation]
T t
[Equation]
··· T
[Equation]
K −
I ··· 0 ··················  ···
I (cid:35) , with T t, 0 =
I − αH t and T t,k = − αw t,k H t − k for k = 1 , . .
, K − 1 .
Unrolling the recursion, we have that
[Equation]
 
[Equation]
1 k
[Equation]
[Equation]
j
[Equation]
k
[Equation]

 [Equation]
[Equation]
from which it is clear that we may expect convergence properties to depend on the spectral properties of the matrices { A k } tk =1 .
The convergence rate of SGD C will therefore depend on the largest singular value of the matrices { A k } tk =1 .
From Polyak (1964a, Lemma 4), we have that
[Equation]
where ρ (Λ) are the singular values of the numerical matrix
[Equation]
with λ 0 = 1 − αη 0 and λ k = − αw k η k for k = 1 , . . .
, K − 1 .
Note that Polyak (1964a, Lemma 4) originally refers to the eigenvalues of a block matrix, however, in the case where the blocks are commutative, it is straightforward to extend the lemma to describe the singular values.
To see this, simply apply the original lemma to the eigenvalues of the matrix Λ (cid:62) Λ , where commutativity of the individual blocks ( T j, 0 , . .
, T j,K are all symmetric and therefore commutative) allows you to specify the eigenvalues of the new block matrix in terms of the eigenvalues of the original blocks.
Letting q ( α,K ) denote the right hand side of (10), we have that (cid:107) A j (cid:107) ≤ q ( α,K ) for all j = 1 , . . .
, t , and q ( α,K ) has an analytic deﬁnition as the roots of a polynomial.
Taking 2-norms, applying submultiplicativity of matrix norms, and using the bounded variance assumption along with the fact that the noise terms { ζ k } tk =0 are mutually independent and that the aggregation weights { w t,k } t ≥ 0 ,k ∈ [1 ,K ) are in the closed interval [0 , 1] for all t, k , we have that
[Equation]
where the expectation on the left hand side is with respect to the joint distribution over the random noise variables ζ k , . .
, ζ 1 .
By assumption, since the step-size α is chosen such that q ( α,K ) < 1 , equation (11) simpliﬁes as
[Equation]
where we have implicitly used the upper bound on the limit of a geometric sequence.
[Equation]
C D ATA SET DISTRIBUTIONS
Table 1 details the train-
[Figure]
D E XPERIMENT D ETAILS
We tested the proposed class of optimizers on different tasks that have potentially different loss surfaces in large parameter spaces.
To cast a wide net and ensure that we capture a plethora of neural architectures, we vary the network designs by using fully-connected, recurrent (Rumelhart et al., 1985), convolutional (LeCun et al., 1999), dropout (Srivastava et al., 2014), ReLU (Agarap, 2018) layers and a large Transformer language model architecture – RoBERTa-Base (Liu et al., 2019) across different tasks.
A brief description of the data and architectures used follows.
CIFAR-10 is an image classiﬁcation task with 10 classes.
We train a shallow Convolutional Neural Network (ConvNet) with dropout and ReLU activation using the different optimizers and their C variants
CIFAR- 100 is an image dataset similar in size to CIFAR- 10 but with 100 classes of images This task was approached using a Convolutional
Neural Network with three batch- normalized convolutional blocks and one fully-connected block with dropout applied between blocks
We experiment with the commonly-used language inference dataset, Stanford Natural Language
Inference (SNLI) dataset, to compare the performances of C variants training on three different text encoder architectures: 4 -layer convolutional network (ConvNetEncoder), 2 layer bi-directional encoder with 2 linear projection layers (InferSent) and a unidirectional LSTM with 2 recurrent layers (LSTMEncoder).
The classiﬁer on top of the representations learned by the encoder architectures is a 2 -layer fully-connected Multi-Layer Perceptron with dropout connection.
The Penn Tree Bank (PTB) is a syntax-annotated text corpus sourced from stories from the Wall
Street Journal.
We use this corpus for a word-level language modeling task using a 1 -layer LSTM with dropout.
Gradient clipping is employed to avoid exploding gradients.
We evaluate the model by measuring perplexity (PPL) in the validation set; lower PPL scores are preferred.
MultiWoZ 2 . 0 is a popular dataset that has human-to-human goal-oriented conversations on different topics.
The objective is to generate the next utterance conditioned on the history of utterances in the conversation.
We experiment with a very large language model architecture – RoBERTaBase – and a BiLSTM Sequence-to-Sequence architecture (Vinyals and Le, 2015) for the next utterance prediction.
RoBERTa was trained with CausalLM-Head while BiLSTM model was an encoder-decoder architecture with a Bi-LSTM encoder and LSTM with Attention decoder.
We use models with varying size of trainable parameters as shown in Table 2.
[Table]
We use the same hyperparameter initialization for comparisons with base optimization methods and tune the learning rate hyperparameter using grid search.
The primary objective of the experiments is to verify the consistency of convergence to better solutions across tasks and architectures.
All reported results are averaged over 5 different runs with different random seed values.
D. M ODEL H PARAMS
All experiments are averaged for 5 different runs with different seeds.
We use PyTorch 1.1 for the experiments and use their implementation of the base optimizers available in torch.optim .
The details of hyperparameters used for the model are in Table 3.
[Table]
D. R UNTIME S TATISTICS
We logged the approximate time for each epoch for different values of topC across the different models.
Although the results are populated from the experiments with Adam and its C variants, the results can be extended to the other optimizers and its variants.
These times are reported in Table 4.
[Table]
E H YPERPARAMETERS E. R ANGE OF H PARAMS
We present the range of hyperparameters used in our experiments (Table 5).
Optimizer-speciﬁc parameters were used on both their vanilla and C versions
[Table] [Table]
E. HP ARAMS OF THE B EST C ONFIGURATIONS
The learning rates and other hyperparameters of the optimizers used in the results reported in the paper are listed in Tables 6, 7, 8 and
[Table]
[Table]
[Table] [Table]
N/A N/A
N/A 0.99 SGDM C 0.
N/A
N/A
N/A SGDM 0.
N/A
N/A
N/A SGD C 0.
N/A
N/A N/A
N/A SGD 0.
N/A
N/A N/A
N/A
F A NALYSIS F. N OTE ON AGGREGATION
The aggregation function aggr lies at the core of integrating critical gradients into an algorithm.
In §2 we presented mean and sum function for aggregation.
The notable difference in these methods is the importance of the current gradient, as it is more heavily weighted with sum (which does not scale it) than with mean (which scales it by a factor of 1 topC +1 .
While our algorithms converged
[Table]
/A N/A
N/A N/A
BLSTMEncoder
Adam C 0.
N/A 0.
N/A Adam 0.
N/A 0.
N/A RMSprop C 0.
N/A N/A
N/A 0.99 RMSprop 0.
N/A N/A
N/A 0.99 SGDM C 0.
N/A
N/A
N/A SGDM 0.
N/A
N/A
N/A SGD C 0.
N/A
N/A N/A
N/A SGD 0.
N/A N/A
N/A
N/A using both aggregation methods, we found that SGD C and SGDM C demonstrated better validation performance using sum whereas RMSprop C and Adam C performed better using mean .
This is likely due to both RMSprop and Adam having adaptive learning rates, allowing them to be more robust to changes in the scaling of the gradients.
F. C ONVERGENCE IN C ONVEX T ASKS
[Figure]
To empirically validate our proof of convergence for the proposed methods on convex loss surfaces, we train all C optimizer variants on three tasks with Logistic Regression: (cid:96)
-regularized binary classiﬁcation with rcv1 (Lewis et al., 2004), (cid:96) 2 -regularized multi-class classiﬁcation on covtype (Dua and Graff, 2017), and non-regularized multi-class classiﬁcation on MNIST.
Figure 7 shows training losses for all optimizers converging towards a minimal training loss.
F. B UFFER S TALENESS B OUND
We present additional results demonstrating the boundedness of the staleness of stored gradients (Figure 8).
F. A DDITIONAL A BLATION E XPERIMENTS
We present additional side-by-side comparisons of optimizers in the vein of §6 (Figure 9).
[Figure] [Figure]
As C variants retain larger gradients during training, and noisy data points may have gradients with higher values of (cid:96) 2 -norm, we experiment with the C variants by training Logistic Regression on a synthetic binary classiﬁcation dataset where the labels are perturbed with probability P noise .
We sample 500 data points from a 7 -dimensional hypercube with the class means separated by 1 unit.
The data is split by 80 / 20 for training and evaluation.
We train the model with the optimizers forepochs on a ﬁxed learning rate of 1 × 10 − 3 , topC = 5 , decay = 0 .
5 and other parameters from the base model set to their defaults.
We compare the performance of base method and its C variant to study anomalies when the C variants are exposed to noise (Figure 10).
We observe that the models do not behave any differently than the base methods showcasing that the noise in the dataset does not affect the workings of the C variants even with a small value for topC .
This could be attributed to the decay parameter of the optimizer that linearly scales down the priority of the older gradients in the buffer leading to them getting replaced.
To construct the synthetic data set for the experiments, we used sklearn’s make_classification method.
The speciﬁc hyperparameters used to construct the dataset:
• n_samples = 500, n_features = 10, • n_informative=7, n_redundant=0, n_repeated=0, • n_classes=2, n_clusters_per_class=1, • weights=class_imbalance, ﬂip_y=noise, class_sep= 0.5, • hypercube=
True, shift=0.4, scale=1.0,
[Table]
[Figure]
The proposed variants did not show any anomaly on a noise induced dataset.
For the experiments on the decay the topC was 5 and decay 0.5.
The plots averaged 5 different runs different seeds.
F. DECAY
The optimizers’ sensitivity to decay (Figure 11) had similar trend in other experiments, where the optimizers showed better performance when decay was away from 0.99.
[Figure] [Figure]
TOP
[Figure]
The optimizers’ sensitivity to topC (Figure 12) had a similar trend to other experiments, where the optimizers showed better performance for lower values of topC .
F. g t VS g c
As a follow-up experiment to validate the non-informative gradients with higher values of topC , we observe the trend in the difference between average ( || g c || 2 ) and || g t || 2 for different values of topC .
We see that the gradients at each step get lower as topC increases.
This could be because of the stochasticity in aggr when computed with fewer gradients, which incidentally allow models to converge better.
Although storing all of the past gradients in memory has theoretical advantages, in practice we observe that lower topC provides better training signal through g c than g t (Figure 13(a),13(b)).
This explains the slightly lower performance of topC = 100 across models in SNLI task and the nonexistent to marginal improvements in MultiWoZ dataset with RoBERTa-Base model.
[Figure] [Figure]
The plots of g t vs g c is crucial to the results in that they provide an explanations for the faster convergence (as well as smaller improvements) of the C variants over the base methods (Figure 14).
Since the parameter updates are reminded of large displacement of the gradients via the aggr method, the model updates the parameters more frequently.
In the cases where g c is higher than g t , we observe better performance and no improvements in cases where the difference is not as signiﬁcant.
[Figure]
Further, we observe the signiﬁcance of the gradients as deﬁned in (cid:96) 2 -norm diminishes as topC increases.
This indicates that the optimizer needs only to store only a small set of critical gradients in order to improve the performance without incurring much time and memory overhead.
Employing lazy updates to cut down on the time is a useful future direction of research.
F. D ETAILS OF A BLATION E XPERIMENTS
We detail the various ablation experiments whose results are presented in §6.
• King of the Hill sampling: The default method for Critical Gradients.
Incoming gradients are added if their norm is larger than the priority of the smallest gradient, which is subsequently removed.
• Smallest gradients (or " bottomC "): Replaces the max-heap buffer with a min-heap • First-In-First-Out (FIFO):
Replaces the max-heap (priority queue) buffer with a queue which begins to dequeue when it reaches a capacity of topC .
• "Coin-Toss"/"random": an incoming gradient is added to the buffer or not with equal probability.
• Mean-Norm Diversity Sampling (MNDS): Incoming gradients are probabilistically added to the buffer with likelyhood proportional to its difference from the mean of the norms of gradients in the buffer.
• Cosine Diversity Sampling (CDS): Incoming gradients are probabilistically added to the buffer with likelyhood antiproportional to its cosine similarity (normalized dot-product of ﬂattened vectors).
• Cosine Similarity Sampling (CSS): Incoming gradients are probabilistically added to the buffer with likelyhood proportional to its cosine similarity.
• Random Replacement (RR): When the buffer is full and a new gradient is added, the gradient which gets removed gets selected at random.
• Norm-Controlled Probabilistic Replacement (NCPR): When the buffer is full and a new gradient is added, the gradient g k which gets removed gets selected with probability
[Table]
G T EST P ERFORMANCE
The complete test performance of the models is reported in Table 10.
H R EPRODUCIBILITY C HECKLIST
As per the prescribed Reproducibility Checklist we provide the information of the following:
• A clear description of the mathematical setting, algorithm and/or model :
We provide
details of models used
in
§D • Submission of source code : Source code for the proposed optimizers and its variants is provided as a zip.
The code used to train the models on the different data sets are open source GitHub repositories.
Other codes developed for the project are included in the zip.
• Description of the computing infrastructure used :
We used 50 NVIDIA V100 32GB GPUs in parallel hyper parameter search over the grid using wandb and submitit packages.
For the ﬁnal runs we used 1 NVIDIA V100 32 GB GPUs for every seed of every model.
• Average runtime for each approach :
The approximate training time for our use of Adam C accross tasks is reported in §D.2.
• Explanation of evaluation metrics used, with links to code : The metrics used for evaluation of the models are the popular ones.
For the ease of readers citations for the metrics are included in the paper.
• Relevant statistics of the datasets used : We provide the statistics of the datasets used in C. • Explanation of any data that were excluded, and all pre-processing steps : We train on a fraction of the covtype and rcv1 datasets instead of using the entire data.
We sampleddatapoints at random with seed set to 100 .
• Link to downloadable version of data
: The data sets used in the paper are from public repositories.
Links to the paper that proposes the data sets is included in the README.md ﬁles in the submitted repository.
[Table]