Year: 2021

Abstract:
Popular approaches for minimizing loss in data-driven learning often involve an abstraction or an explicit retention of the history of gradients for efficient 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 sufficient 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 fixed-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.

heading: 1 INTRODUCTION
text: Gradient-based learning involves minimizing a scalar-valued function F (θ) with respect to a parameter vector θ ∈ Rd using an iterative procedure. When training over n examples of input-output pairs, (xi, yi), the optimization problem boils down to solving minθ F (θ), where
F (θ) = 1
n n∑ i=1 Li(θ), with Li(θ) = L (Mθ(xi), yi) ,
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 efficient 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 significant 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.
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 `2-norm) that are occasionally updated with newer gradients. Specifically, rather than storing gradients for all the examples in the dataset, our proposed approach aims to find 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 final performance through exhaustive empirical study on eight different architectures among four different tasks (classification, 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.

heading: 2 MEMORY-AUGMENTED OPTIMIZERS
text: In this work, we propose augmenting standard optimizers with a fixed size memory of past gradients. The gradients are stored in the limited memory only when they are deemed critical as defined by their `2-norm.
Gradient Selection to the Memory Letting gt := ∇θLi(θt) denote a component (or mini-batch) gradient at time-step t, we use ‖gt‖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 gc. In order to ensure that the buffer eventually flushes 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 fixed 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 `2-norm of gt to determine priority. Gradients in the buffer are stored as tuples (‖gt‖2 , gt), akin to a key-value pair in a dictionary structure, where the key ‖gt‖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 `2 proxy-norm is replaced by the current mini-batch gradient gt if ‖gt‖2 is greater than the smallest `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 (SGDC) is the explicit integration of critical gradients into SGD. Specifically, the iteration comes in the form θt+1 = θt − α · aggr(gt, gc), (1) where gc 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 gt and all critical gradients, and sum, the addition of gt to the average of all critical gradients. Mathematically, for a buffer of size C these are defined as
mean(gt, gc) = 1
C + 1 (gt + ∑ gc∈gc gc) (2) sum(gt, gc) = gt + 1 C ∑ gc∈gc gc (3)
In general, we find optimization to be robust to the specific 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 finite 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 firstorder optimization methods with the replacement of gt in those algorithms with aggr (gt, gc). The update steps for the corresponding C versions are as follows:
mt = γ ·mt−1 + aggr(gt, gc) θt+1 = θt − α ·mt (4)
E[g2]t = 0.9 · E[g2]t−1 + 0.1 · aggr(gt,gc)2
θt+1 = θt − η√
E[g2]t + · aggr(gt,gc)
(5)
mt = β1mt−1 + (1− β1) · aggr(gt, gc) vt = β2vt−1 + (1− β2) · (aggr(gt, gc))2 m̂t = mt
1− βt1 v̂t =
vt 1− βt2
θt+1 = θt − η√ v̂t + m̂t
(6)
The general algorithm to integrate critical gradients into any optimizer is provided in Algorithm 1 in Appendix A.
Replacing gt 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 gt 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, non memory-augmented counterparts.

heading: 3 PROOF OF CONVERGENCE
text: Consider the SGDC update in Equation (1), where aggr(gt, gc) outputs a linear combination of the critical gradients in memory gk ∈ gc, and the stochastic gradient gt, 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 gk ∈ gc at iteration t, implies t− k < K, and suppose the objective F : Rd → R is twice-continuously differentiable, L-smooth, and µ-strongly convex, with 0 < µ ≤ L. Assume that for all k, the stochastic gradient gk is a random vector satisfying
E[gk] = ∇f(θk) and E [ ‖gk −∇f(θk)‖2 ] ≤ σ2,
for some finite constant σ2. Letting ζk := gk −∇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 `2-regularized logistic regression and `2-regularized least-squares regression (i.e., ridge regression). Taken together, these properties imply that the Hessian ∇2F (θ) exists, and for all θ ∈ Rd, the eigenvalues of ∇2F (θ) lie in the interval [µ,L]. In this case, let θ? denote the unique minimizer of F .
To prove convergence, we view SGDC 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 semidefinite 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 SGDC). Let
q(α,K) := sup {wk∈[0,1]},{ηk∈[µ,L]} ρ(Λ),
where ρ(Λ) are the singular values of the numerical matrix
Λ := [ λ0 λ1 · · · λK−1 0 1 · · · 0 · · · · · · · · · · · · · · · · · ·
0 0 · · · 1
] ∈ RK×K ,
with λ0 = 1− αη0 and λk = −αwkηk for k > 0. If the step-size α > 0 is chosen sufficiently small, such that q(α,K) < 1, then
E ‖θt+1 − θ?‖2 ≤ q(α,K)2t ‖θ1 − θ?‖ 2 + α2K
1− q(α,K)2 σ2,
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 SGDC 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 gc, at all times t, only contains the mostly recently computed gradient), then SGDC 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− q2)), 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 {wk} 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 coefficient 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 finite-sum models.

heading: 4 RELATED WORK
text: 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. SGDC shares similarity with NAG except in SGDC 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 (SGDC, AdamC, RMSpropC, SGDMC) provide an additional “layer" to the base algorithms that ensures any parameter computed within the algorithm incorporates the critical gradients. Although optimization using a fixed 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. (2007). These LBFGS methods utilize a moving window of past gradients, rather than our idea to maintain “critical" gradients based on a well-defined 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 fixed 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 first-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 benefits of variance-reduced methods without significant 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 influence 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 defines 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 `2-norm as a metric to pick out the critical gradients instead of arbitrarily keeping a recent history.
Memory-Augmented Neural Networks 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 difficult 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.

heading: 5 EXPERIMENTS
text: 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 classification dataset (LeCun and Cortes, 2010). Across the experiments, we compare all the 8 optimizers averaged over 5 different runs with different
*The code to reproduce the experiments is submitted as supplementary and will be released upon acceptance.
seed values— Adam, SGD, SGDM, RMSprop, AdamC, SGDC, SGDMC, and RMSpropC on 9 tasks — 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-specific 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 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 RMSpropC 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.

heading: 6 ANALYSIS
text: Buffer Staleness Bound A key assumption behind the theoretical convergence of SGDC 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 findings of this analysis.
Use of Critical Gradients Our Critical Gradient method utilizes a buffer which is filled by using a norm-based, “King of the Hill" sampling of incoming gradients, i.e. gradients always stored and sorted in decreasing order by `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 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.
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
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 benefit from the added information.
Effect of hyperparameters We study the performance’s sensitivity to decay and topC using AdamC on different tasks.
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 AdamC optimizer for the different values of decay with every other parameter fixed (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.
topC ∈ [0,∞) defines 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 > 20 was either not useful or was hurting the performance (also in Appendix F).

heading: 7 CONCLUSION
text: 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 first 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.

heading: A GENERAL CRITICAL GRADIENT ALGORITHM
text: We present the pseudocode for the Critical Gradient algorithm below.
Algorithm 1 Critical Gradients Optimization Define #epochs E, #gradient updates B, model θ, learning rate α Define optimizer O(gt, θt, h), which produces a parameter update given gradient gt, parameter θt and hyperparameters h Define aggregation function aggr(gt, gc) which outputs a linear combination of gt and the entries in gc Initialize empty critical gradient buffer {gc} with decay rate d for s = 0 to E do
for t = 0 to B do Sample an input-output pair (xi, yi) gt ← ∇θ Li(θt) θt+1 ← O(aggr(gt, gc), θt, h) if {gc} is not full then
Add gt to {gc} else
if ||gt||2 > ||min({gc})||2 then Remove min({gc}) from {gc} Add gt to {gc}
end if end if for gc ∈ {gc} do ||gc||2 ← d · ||gc||2
end for end for
end for

heading: B CONVERGENCE PROOF
text: Proof. Recall that if F : Rd → R is twice continuously differentiable, then for all x, y ∈ Rd ∇F (y) = ∇F (x) + ∫ 1 0 ∇2F (x+ t(y − x))dt (y − x). (7) Letting rk := θk − θ? denote the suboptimality of the parameters at iteration k and applying (7) with y = θk and x = θ?, we get that
gk = Hkrk + ζk, where Hk = ∫ 1 0 ∇2f(θ? + trk)dt,
where ζk denotes the gradient noise at iteration k. Using this and bounded staleness assumption, we have that
aggr(gt,Gt) = K−1∑ k=0 wt,kgt−k = K−1∑ k=0 wt,k(Ht−krt−k + ζt−k), (8)
where {wt,k}K−1k=0 are non-negative scalars in the closed interval [0, 1] used to linearly aggregate the gradients in memory with the most recent gradient gt. For gradients gt−k that are not stored in memory, the corresponding aggregation scalar wt,k is equal to 0. Without loss of generality, we take wt,0 = 1, since we can always incorporate it into the step-size α and re-scale the other weights wt,1, . . . , wt,K−1 accordingly. Substituting (8) into (1) and subtracting θ? from each side, we get that
rt+1 = rt − α K−1∑ k=0 wt,k(Ht−krt−k + ζt−k).
Thus, we have that SGDC evolves according to the linear system rt+1rt. . .
rt+1−K
 = At  rtrt−1.
.
. rt−K
− α ∑K−1k=0 wt,kζt−k0.
. . 0
 where At = [Tt,0 Tt,1 · · · Tt,K−10 I · · · 0· · · · · · · · · · · · · · · · · · 0 0 · · · I ] ,
with Tt,0 = I − αHt and Tt,k = −αwt,kHt−k for k = 1, . . . ,K − 1. Unrolling the recursion, we have that rt+1 rt . .
. rt+1−K
 = (At · · ·A1)  θ1 − θ? 0 . .
. 0
 − α  ∑K−1 k=0 wt,kζt−k 0 . .
. 0
 − α t−1∑ j=1 (At · · ·Aj+1)  ∑K−1 k=0 wj,kζj−k 0 . .
. 0
 , (9)
from which it is clear that we may expect convergence properties to depend on the spectral properties of the matrices {Ak}tk=1. The convergence rate of SGDC will therefore depend on the largest singular value of the matrices {Ak}tk=1. From Polyak (1964a, Lemma 4), we have that
‖Aj‖ ≤ sup {wk∈[0,1]},{ηk∈[µ,L]} ρ(Λ), (10)
where ρ(Λ) are the singular values of the numerical matrix
Λ := [ λ0 λ1 · · · λK−1 0 1 · · · 0 · · · · · · · · · · · · · · · · · ·
0 0 · · · 1
] ∈ RK×K ,
with λ0 = 1 − αη0 and λk = −αwkη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 Λ>Λ, where commutativity of the individual blocks (Tj,0, . . . , Tj,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 ‖Aj‖ ≤ q(α,K) for all j = 1, . . . , t, and q(α,K) has an analytic definition 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 {wt,k}t≥0,k∈[1,K) are in the closed interval [0, 1] for all t, k, we have that
E ‖rt+1‖2 ≤ q2t(α,K) ‖θ1 − θ ?‖2 + α2σ2K t∑ j=1 q 2(t−j) (α,K) , (11)
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) simplifies as
E ‖rt+1‖2 ≤ q(α,K)2t ‖θ1 − θ?‖ 2 + α2K
1− q(α,K)2 σ2,
where we have implicitly used the upper bound on the limit of a geometric sequence.

heading: C DATA SET DISTRIBUTIONS
text: Table 1 details the train-valid-test splits of all the data sets used in the experiments.

heading: D EXPERIMENT DETAILS
text: 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 classification 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 classifier 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.
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.

heading: D.1 MODEL HPARAMS
text: 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.

heading: D.2 RUNTIME STATISTICS
text: 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.

heading: E HYPERPARAMETERS

heading: E.1 RANGE OF HPARAMS
text: We present the range of hyperparameters used in our experiments (Table 5). Optimizer-specific parameters were used on both their vanilla and C versions.

heading: E.2 HPARAMS OF THE BEST CONFIGURATIONS
text: 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 9.

heading: F ANALYSIS

heading: F.1 NOTE ON AGGREGATION
text: 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 1topC+1 . While our algorithms converged
using both aggregation methods, we found that SGDC and SGDMC demonstrated better validation performance using sum whereas RMSpropC and AdamC 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.

heading: F.2 CONVERGENCE IN CONVEX TASKS
text: 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: `2-regularized binary classification with rcv1 (Lewis et al., 2004), `2-regularized multi-class classification on covtype (Dua and Graff, 2017), and non-regularized multi-class classification on MNIST. Figure 7 shows training losses for all optimizers converging towards a minimal training loss.

heading: F.3 BUFFER STALENESS BOUND
text: We present additional results demonstrating the boundedness of the staleness of stored gradients (Figure 8).

heading: F.4 ADDITIONAL ABLATION EXPERIMENTS
text: We present additional side-by-side comparisons of optimizers in the vein of §6 (Figure 9).

heading: F.5 ROBUSTNESS TO NOISE
text: As C variants retain larger gradients during training, and noisy data points may have gradients with higher values of `2-norm, we experiment with the C variants by training Logistic Regression on a synthetic binary classification dataset where the labels are perturbed with probability Pnoise. 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 for 20 epochs on a fixed 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 specific 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, flip_y=noise, class_sep= 0.5,
• hypercube=True, shift=0.4, scale=1.0, shuffle=True, random_state=1403

heading: F.6 DECAY
text: 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.

heading: F.7 TOPC
text: 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.8 gt VS gc
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(||gc||2) and ||gt||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 gc than gt (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.
The plots of gt vs gc 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 gc is higher than gt, we observe better performance and no improvements in cases where the difference is not as significant.
Further, we observe the significance of the gradients as defined in `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.

heading: F.9 DETAILS OF ABLATION EXPERIMENTS
text: 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 flattened 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 gk which gets removed gets selected with probability ||gk||2∑
gi∈gc ||gi||2 .

heading: G TEST PERFORMANCE
text: The complete test performance of the models is reported in Table 10.

heading: H REPRODUCIBILITY CHECKLIST
text: 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 final 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 AdamC 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 sampled 5000 datapoints 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 files in the submitted repository.

Review summaries:
The paper proposes a memory-augmented optimizer to make use of the history information of the gradients. The motivation is that the aggregation of past gradients nudges parameter updates in a better direction. However, it remains to answer how much historical information of the gradients is sufficient. The authors study and analyze the memory requirement of several common optimization algorithms. Based on that, the authors propose a new optimizer that retains a limited gradient history and those gradients are selected by their importance. This method could be built upon existing popular optimizers like Adam and SGD and enjoys accelerated convergence and improved performance.

The paper proposes a framework for memory-augmented gradient descent optimizers which only keep a limited a limited buffer of gradient history and can be integrated with existing optimizers. The method proposed in the paper is centered around "critical gradients" with large l2 norm and only retaining them in the limited buffer. Experimental results show the faster convergence in terms of iterations compared to standard SGD algorithms without critical gradient buffers.  

The paper proposes an optimisation framework that uses K previous gradients stored in a buffer to compute an aggregated proxy-gradient for the  use in other gradient frameworks (studied are SGD, SGDM, RMSProp, Adam) and evaluates it on CIFAR 10/100 as well as NLP tasks. The paper also presents a convergence proof on strongly convex smooth objectives. In almost all empirical evaluations the method shows an improvement in terms of iteration and performance and ablations are performed to show strong robustness to hyperparameters introduced.