Memory Augmented Optimizers for Deep Learning
Authors: Paul-Aymeric Martin McRae
Prasanna Parthasarathi
Mido Assran
Sarath Chandar
Venue: ICLR 2022 Poster (2022)
Keywords: Optimization for Deep learning
Memory augmented Optimizers

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.

Published as a conference paper at ICLR 2022
MEMORY LEARNING
AUGMENTED
OPTIMIZERS FOR DEEP
Paul-Aymeric McRae t!, Prasanna Parthasarathi 112, Mahmoud Assran and Sarath Chandar
Milla Quebee AI Institute, Canada
McGGil University, Canada
'fcole Polytechnique de Montréal, Canada
*Canada CIFAR Al Chair
ABSTRACT
Popular approaches for minimizing loss in data-driven learning often involve an abstraction or 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 informative. Although the history of gradients summarized in meta-pant patammeessssssss or explicitly stored in memory has been shown effective theory and practice, the question of whether all or only subset of the gradients in the history are sufficient in deciding the parameter updates remains unanswered. this paper, we propose 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 experiments show that the memory augmented extensions standard optimizers enjoy accelerated convergence and improved performance on majority of computer vision and language tasks that 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.
INTRODUCTION
Gradient-basse learning involves minimizing a ecalar-valued function F(0) with respect to a parame- ter vector 0 E Rd using iterative procedure. When training over examples of input-output pairs, (1i, Wi), the optimization problem boils down to solving ming F0)), where
1 F(0) = n Σ L,(0), with (0) = C y:)
for some
loss function C 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 often preferred to full-batch Gradient Descent when the quantity of data required to train 0 is large, since it can be more efficient to measure single component of the gradient (or a mini-batch of gradients (Bottou 1999)), ). and move a noisy direction, than to compute full gradient at each time step. Several techniques have been proposed further accelerate the convergence SGD (Ghadimi and Lan 2016 Allen-Zhu and Yuan 2016 Zhou et al.) (018)). These include approaches that maintain 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 al.
Hinton 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 set of buffers that track running moments of the gradients. Such light-weigh techniques
The two authors contributed equally this paper.

---

Published as a conference paper at ICLR 2022
have shown significant application advantage in practice, whereas in theory, algorithms that store all of the gradients, like SAG (Roux et al. (012) and SAGA (Defazio et al.] (014)) 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-basets approaches (Fig- ure [ 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 mem- ory than full-history methods while providing richer updates than adaptive-gradie methods.
Miniaat
More
Memarttss
Kemoor
kde2zed
Momantus
520
Limited-Memory BFGS (Nocedal 1980) and its online, stochastic form, oLBFGS (Schraudolph et al. al.] 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, propose using memory to augment conventional optimizers. This memory is used to store small set of critical gradients (e.g., gradients with larger 4--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 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.
Figure 1: Common optimization algorithms placed over the spectrum of memory requirements. The proposed algorithms the c variants have a fixed sized memory to store gradients during training.
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, replace- ment, or the aggregation techniques to summarize information in the memory.
2 MEMORY-AUGMENTEE OPTIMIZERS
In this work, we propose augmenting standard optimizers with fixed size memory of past gradients. The gradients are stored in the limited memory only when they are deemed critical as defined by their fz-norm.
Gradient Selection to the Memory Letting 9: = (B,) denote a component (or mini-batch gradient at time-step t, we use l[gel/o as scalar indicator of the importance of the gradient, which serves proxy for the priority of a gradient to remain in the memory buffer Jc- In order to ensure that the buffer eventually flushes out stale gradients, the proxy norms for gradients in the buffer are scaled down by hyperparameter decay factor, denoted decay E [0, L).
The proposed approach maintains a gradient buffer of fixed capacity 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 fo-morm of 9e to determine priority. Gradients in the buffer are stored as tuples (I19c ll2 9t), akin to key-value pair in dictionary structure, where the key || is referred to as the proxy norm, and decayed at each time step t by a decay factor decay E [0, 1). At any iteration t during training, the gradient in full capacity priority buffer with the smallest &g proxy-norm is
2

---

Published as a conference paper at ICLR 2022
replaced by the current mini-batch gradient 9t if 11 ll is greater than the smallest &2 proxy- in the buffer. The multiplicative decay factor ensures that the buffer is frequently refreshed with more recent gradients. Note that only the proxy-norn used for the heuristic rule is decayed by the decay factor decay E [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 al. 2012 Kingma and Ba 2015).
Critical Gradient Stochastic Gradient Descent Critical Gradient Stochastic Gradient Descent (SGDc) is the explicit integration of critical gradients into SGD. Specifically, the iteration comes as,
8141
a aggr(gt: ge).
where Be is the set of critical gradients and aggr denotes 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 Sr and all critical gradients, and sum, the addition of gi to the average of all critical gradients. Mathematically, for buffer of size C' these are defined as
mean(yt: ge) - C (9:
Σ
go)
(2)
sum(ge, ge) - Sr + C
Σ ge
(3)
BrEge
9.Egr
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 finite moving window of past gradients. Following which we propose critical gradient extensions of SGD with Momentum (SGDM) 4, RMSprop (5q(5). and Adam (Fq The retention of the critical gradients can be naturally extended to more first-order optimization methods with the replacement of g in those algorithms with aggr (gt, 9c). The update steps for the corresponding c versions
follows:
m. 7 m.-1 aggr(g..g.)
m. 3 + (1 (:) aggr(si, ge)
(4)
8.41 a me
2, 820-1 (2) (aggr(9s, gc))²
m,
m, - 8
Elg"}. = 0.9 Elg" 0.1 aggz(g).g) ga)
6, it
35
p.n 0: VElg+,+c aggr(g..g)
7
0.41 0:
me
Vo, +
The general algorithm to integrate critical gradients into any optimizer is provided in Algorithm ii Appendix A Replacing 9: 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 9: but also with respect to the critical gradients from the past. We observe that all c-extended optimizers in our experiments $ lead to faster convergence than their vanilla, non memory-augmente counterparts.
PROOF OF CONVERGENCE
Consider the SGDc update in Equation T where aggr(gi, ge) outputs a linear combination the critical gradients in memory 9& E ge. and the stochastic gradient gt, which is computed at the current parameters 0, When the staleness of gradients the memory is bounded, we can prove convergence of this method through the lens of linear multi-step methods. Formally, suppose there exists integer
3

---

Published as a conference paper at ICLR 2022
K > such that Ok 9. at iteration t, implies t < K, and suppose the objective F R + R is wrice-continuously differentiable, L-smooth, and p-strongly convex, with 0 < μ S L. Assume that for all k, the stochastic gradient 9k is random vector satisfying
E Bigx] = f(0x) and E |1g
f((+))) so²
for some finite constant 0². Letting ÇA = 9h f(0x) denote the gradient noise iteration 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 fz-regularized logistic regression and Ez-regularized least-squares regression. Taken together, these properties imply that the Hessian V F()) exists, and for all E R the eigenvalues of F() lie in the interval [M, Lj. In this case, let 0* denote the unique minimizer of F.
To prove convergence, we view SGDe 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 instead rely the spectral radius of the system matrix by constructing appropriate Lyapnuov function containing a time-invariant system matrix satisfying particular Integral Quadratic Constraint (IQC) (Lessard et al. 2016). However, one downside of the IQC framework is that one must typically solve semidefinite program to obtain explicit convergence rates (Hu and Lessard (refer Appendix $B for proof) provides an analytical convergence rate.
2017 whereas Theorem !
Theorem (Linear Convergence of SGDc). Let
sup K). (w1E(0.1]) Em EI
P(A),
where p(A) are the singular values of the numerical matrix
Ag
21
Ax-
A
1
E RRRK K
with Àu =1 anju and A -0 for ls >> O. If the step-size >> O is chosen sufficiently small, such that (),,K) < 1, then
a K
E [10441 - 1-112 < [(.,)) 100 8+112+ +
2²
((,,,K
where the expectation the left hand side is with respect to the joint distribution over the random noise variables ......,
Theorem I shows that SGDg can be made converge exponentially fast to a neighbourhood of the solution, the size of which is proportional to the variance of the stochastic gradients o² and the step-size a². If K
(i.e., the memory 9c. at all times i, only contains the mostly recently computed gradient), then SGDe reduces to regular stochastic gradient descent, and the rate ((a in Theorem reduuces the well known convergence rate of SGD for smooth strongly convex functions (max,ee(uL) [1 Ln) and the variance bound reduces to the standard variance bound for SGD with constant step-size (a2o²/(0 /01 q²)), see, e.g., (Bottou et al. 2018 Assran and Rabbat 2020). For K > 1, it may be possible to choose the step-size a and the aggregation weights [wx]] 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 K() directly lead to smaller variance bounds. Note that expressing the convergence rate of a multi-step method as the roots of polynomial, as we do in Theorem is not new; see, e.g.. Polyak (194aa).
4 RELATED WORK
Soadients-Summarrinn Optimizers Several popular optimization algorithms incorporate
the history of previous gradients through a summarizing variable. In practice, this is often done by
4

---

Published as a conference paper at ICLR 2022
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 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. DC shares similarity with NAG except in SGDc the "momentum" is computed only using C selected gradients. In different use of summarizing variable, AdaGrad (Duchi et al., 2011) keeps 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 s running sum with a decaying average of the squared gradients. Adam (and AdaMax) (Kingma and Ba 2013) uses a combination of momentum and a decaying-averan aveeaaaeeeegeeeeeeeen learning rate. In general, the a 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 variants comes from adding a heuristic to keep only the gradients that provide the largest displacements towards the minimum.
Memory-Enhen 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-Histony BFGS (LBFGS) (Nocedal 1980), an optimization technique which approximates the Hessian of function using 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 SSchraudolph al. (0007) 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 means to improve optimizer convergence is not a novel idea; both Stochastic Accelerated Gradient (SAG) (Roux al. 2012) and SAGA (Defazio et al. 2014) do so in the context of accelerating stochastic gradient descent, while Stochastic Dual Coordinate Ascent Shalev-Shwart 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 is the number of training examples, which can be costly. Our method relies on the storage of fixed number of gradients and uses a heuristic to decide whether to store a gradient in a dataset-size- independent memory. Our class of optimizers can natively run in 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 optimization methods.
SAG and SAGA belong to the family of variance-reduced gradient- algorithms (J. Reddi eet 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 L.STM being passed the gradient norms to influence the parameter update step, echoing our technique's use of gradient norm as 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 's recurrence past gradients. [i and Malik (2017) uses a reinforcement learning based approach and defines a state-space which includes recent history of gradients. While our method equally maintains ledger of past gradients, unlike this latter approach we use fnrnr as a metric to pick out the critical gradients instead of arbitrarily keeping a recent history.
means of implicitly accessing
Motivation for explicit memory Motivation for explicit ones that maintain implicit summarization of history comes from memory- -augentd neural
optimizers over the
5

---

Published as a conference paper at ICLR 2022
networks like Neural Turing Machines (NTMs) and their variants (Graves al. 2014 2016 Gulcehre jt al. 2018). While a simple recurrent architecture like LSTM (Hochreiter and Schmidhuber 1997) integrates information in its cell state which is single vector, NTMs maintain 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 learn for LSTM. Analogous L.TTMs and NTMs, SGDM and our c-extended optimizers maintain single vector and a set of vectors respectively their memory. While our current formulation for memnory-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 optimize.
Motivation for identifying critical gradients Likewise, recent work (Paul et al. 2021 has identi- fied that samples with larger gradient norms can be used to identify smaller set of training data that is important for generalization, and thereby prune the data pre-processing step. Our proposed approach augmenting optimizers with memory buffer of critical gradients can be seen as a softer version of mata-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.
5 EXPERIMENTS
SNIL-C RMSoropy WITEXXE
SNLL-C RMSprope WIKITEX Y REPEORR
ENL IBI SCA. SO
OTE- Adam.
CIFARTE &O PTE ABamT
ON PTB RIMSarOO:
OF DOLL SO CORE
SALLE SODM,
SHLE SOUNC
SNU MHE AdMIC
TL MWoz ROB ARSSPOOS WIKITEET FIE SER
MWoz MIXO BOT SERE
WIKITEXT. COE
WIKITE SOB ARRIC
PTB ABRIN
CEARIO SOAD
SC CIFARID CO SODE
SA SEOORC
MIVOZ B. SEDS
MWGZ EL RMSORRD
MWGZZ BL R
2 3GONE
MINSZ EL SA6GS
MWOZ SL
SMUESL SCOM
SAL AONITS
GEAMS
CIFARLOC 5damc SH CIFAR1OD RMSPrOpe DIAAR SCOMC
OFAR10D CO RMSSPOO:
UFARTO Adam
2x
4x
Acce eration
Imprevement
(a) Acceleration in Convergence
(b) Improvement over the Test performance
Figure 2: (a) Showcases the acceleration in convergence (the unit being the number of gradient updates with batch-size fixed) using a validation set provided by the c variants, compared pairwise (eg: Adam Adamc), the different (bs denotes the difference in performance of the models trained with variants with respect to its corresponding base method on the test set. A positive value indicates that & variant performed better than its base method in that task. The side-by-side comparison allows us to observe that the proposed optimizers enabled accelerated convergence to better test results. Of the tasks considered, the e optimizers performed the best in 6 tasks. Further, (b) shows that version of the base optimizers enhanced the performance to as much as 23%.
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 al. 2009) respectively; Bi-L.TTM (BL., InferSent (I), and text- based convolutional architecture (C) on the Stanford Natural Language Inference (SNLI) 2015) dataset; LSTM on word level language modeling (LS) with the PennTreeBank (Marcus et al. 1993). and WikiText (Merity et al. 2016) datasets; RoBERTa-Bass (RoB) and Bi- LSTM (BL) on language generation in dialogue task with MultiWoZ 2.0 dataset (Budzianowski 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, compare all the 8 optimizers averaged over 5 different with different
(Bowman et al.
et al. 2018). Additionally,

---

Published as a conference paper at ICLR 2022
seed values Adam, SGD, SGDM, RMSprop, Adame, SGDc, SGDMc. and RMSprope tasks 100, CIFARIO, SNLI-I, SNLI-C, SNLI-BL, MWoZ-RoB. MWoZ-BL, PTB-LS, and WIKITEXT-LS" For 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 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 counterparts. A full description of the values used to tune the various parameters, architecture and dataset details are in Appendix şE] $D and reeppectivel 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 с variants provided consistent acceleration in convergence over their vanilla versions. The accelerated convergence also resulted improved test performance, as highlighted in Figure 2(b (Exhaustive list is provided in Appendix { In addition to the aforementioned improvements 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 MWoZZ BL RMSprope the version showed marginally accelerated convergence close to the optimal solution. Furthermore. Figure 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 versions stayed above most if not all of the vanilla optimizers. While that separation is clear in the vision tasks, we observed that the pairwise comparison mostly put the a optimizer at an advantage in accelerated convergence.
vanilla-vS-c: comparisons, the c optimizers yielded the best test
54
52
ALNT
2coM:
SCOM
Vel. Accuracy
Accurscy
20
Vsl. 28
500ML
Epochs
Epochs
(a) SNLI/Infereen
(b) SNLIBLSTM
(c) SNLI/ConvNee
(d) MWoZ/BiLSTM
850
57.
750
Vol. Accuratty
95.0
600
52.5
500
30.0
475
400
120
332
100L
SCOM
23C
10
22
opochh
(e) CIFARIO/ConNNet
(I) CIFARIOO/ConvNet
(g) PTB/LSTM
(h) 2/LT
Figure The validation performance comparison of the proposed optimizers and their vanilla versions showed that, in general. the proposed optimizers showed accelerated convergence than their counterparts. Also, the tasks SNLI-I (a), SNLI-BL (b) and MWoZ-BL (d) where the methods are not the optimal - only falls short by negligible margin. For ease of visualization, confidence intervals across seeds are suppressed.
"Code to reproduce results: http://hhh....oooooomaaaiiiiaaad Optimizer package: http://ggglaaccccccmmmmaallloinn
7

---

Published as a conference paper at ICLR 2022
6 ANALYSIS
Buffer Staleness Bound A key assumption behind the theoretical convergence of SGDe is the existence of an upper bound the buffer gradients' staleness (i.e the number of iterations 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 opp and decay in Figure 4 We observe that tuning the parameters ensures that gradients in the buffer are refreshed frequently, allowing for renewal of information. Later, in analysing the effect of decay and topC on the performance observe that the higher performance of the c methods correlate with the findings of this analysis.
task, we
Use of Critical Gradients Our Critical Gradient method utilizes 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 mn-nrmm 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 amallest-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.
250
250
250
0.7
0.7
0.7
200
200
0.9
200
0.9
Count
150
0.99
150
0.99
Court
150
0.99
100
100
100
50
50
50
60
80
80
Age (Steps)
Age (Steps)
Age (Steps)
(a) topC= =
(b) topC
(c) topC = 20
Figure 4: Histograms showing ages of the buffer as recorded the end of every epoch, across five seeds for several decay values, for a MLP on MNIST. 0.99 decays at a much slower rate, keeping stale gradients in the memory, while 0.7 replaces the gradients frequently. Also, the size of the buffer (topC) ), as expected, correlates directly with retention of stale gradients.
Addressing assumption (a) we test the use of maximal norm sampling (Figure (aa) by comparing with Cosine Similarity Sampling (CSS), Cosine Diversity Sampling (CDS), Mcan-Norm Diversity Sampling (MNDS), and random sampling via "coin toss". Figure ((b)ddppcct our algorithm alongside different gradient replacement methods, including Random Replacement (RR) and Nont-Controlled Probabilistic Replacement (NCPR), as well comparing against keeping running First-In-Fist. Futt-OOut (FIFO) queue and only using the smallest ('bottom"" 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 found in Appendix $9
600
MNDS
600
AIF
600
median
CDS
500
NCPR
min
500
random
RR
500
max
400
CSS
g 400
bottomC
g
400
mean
300
KotH
300
KotH
300
200
200
20
20
(a)
(b)
(C)
Figure (a) compares gradient selection techniques based on different metrics, (b) compares different replacement strategies when flushing out a gradient from memory and (c) compares choice of aggregating the gradients the buffer. The algorithm converges all the different choices.
8

---

Published as a conference paper at ICLR 2022
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 remains conceptually straightforward, easy to implement, and the buffer properties are deterministic compared some of our ablation studies which were probabilistic. Although, rigorous analysis is necessary to verify the generality of this result.
nigh-performing implementation of memory-
optimizer, it
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 s sensitivity to decay and topC using Adame on different tasks.
DanfAC
Conshet
Mgp.
(a) SNLI decay
(b) SNLI topc
(c) CIFAR-100, decay (d) CIFAR-100, topC
Figure 6: Adame with different architectures on SNLI task. Also, the usefulness of explicit memory aligns with the evidence on gradient staleness from Figure
We observe that non-trivial choices of decay and topC provided best performance for
The decay E [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 variants behave like an average over moving window of length topC. On the other extreme, decay- = 1 does not decrease the priorities of the past gradients and allows them to in the buffer for much longer period hindering the convergence of the optimizers. Also, decay maintains the staleness the gradients by exponentially decaying their priority in the memory buffer. We experimented with Adame optimizer for the different values of decay with every other parameter fixed (Figure(6). We observed that decay values closer to did not have any advantage at best or leads to inferior performance on an unseen set at worst in most of the tasks.
topC E [0, oo) defines the size of the buffer that holds the past gradients. When topC is set to 0, a variants behave like the base algorithms and when topC the size of the dataset, the optimizer becomes SiG-like. Also, topC is directly proportional to the staleness of the gradients in the buffer. 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
7 CONCLUSION
We propose a general method to enhance the performance of an optimizer with an augmented memory buffer maintain a limited set of critical gradients from the history. The proposed memory mechanism, when integrated with several state-of oiite-of-lie 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.
Acknowledgements 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 Canada CIFAR AI Chair and an NSERC Discovery Grant.
9

---

Published as a conference paper at ICLR 2022
REFERENCES
Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv, 2018.
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning. pages 102-110. PMLR, 2019.
Zeyuan Allen-Zhu and Yang Yuan. Improved svrg for non-stroogglyconvex sum-of-non-convex objectives. In International conference on machine learning, pages 1080-1089. PMLR, 2016.
Marein Andrychowicz, Misha Denil, Sergio Gómez Colmenarejo, Matthew Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. In NeurIPS, 2016.
Mahmoud Assran and Michael Rabbat. On the convergence of nesterov's accelerated gradient method in stochastic settings. arXiv preprint arXiv:2002.12414 2020.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Correcting momentum in temporal difference learning. Deep Learning Workshop, NeurlPS, 2020.
Léon Bottou. On-line learning and stochastic approximations. On-Line Learning in Neural Networks, 1999.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, (0(2):223-311 2018.
Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 632-642, Lisbon, Portugal, September 2015. Association for Computational Linguistics. doi: 10. 10.18653/1//D15-107. URL https www.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 8(3-4):231-357, 2015.
Pawel Budzianowski, Tsung- Tsinn Wen, Bo- Tseng, Inigo Casanueva, Stefan Ultes, Osman Ramadan, and Milica Gašić. Multiwoz-a large-scale multi-domain wizard-of-oz dataset for task-oriented dialogue modelling. In EMNLP. 2018.
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. In NeurlPS, 2014.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 2011.
Nidham Gazagnadou, Robert Gower, and Joseph Salmon. Optimal mini-batch and step sizes for SAGA. In ICML, 2019.
Saeed Ghadimi and Guanghui Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming, 156(1-2):59-99 2016.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv, 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 2016.
Caglar Gulcehre, Sarath Chandar, Kyunghyun Cho, and Yoshua Bengio. Dynamic neural turing machine with continuous and discrete addressing schemes. Neural computation, 2018
Geoffrey Hinton, Nitish Srivastava. and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent, 2012.
10

---

Published as a conference paper at ICLR 2022
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
Bin Hu and Laurent Lessard. Dissipativity theory for nesterov's accelerated method. In International Conference on Machine Learning, pages 1549-1557. PMLR, 2017.
Sashank J. Reddi, Ahmed Hefny, Suvrit Sra, Barnabas Poczos, and Alexander Smola. On variance reduction in stochastic gradient descent and its asynchronous variants. NeurIPS, 2015.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In NeurlPS, 2013.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR, 2015.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, In Citeseer, 2009.
Yann LeCun and Corinna Cortes. 2010.
MNIST
handwritten digit database.
Yann LeCun, Patrick Haffner, Léon Bottou, and Yoshua Bengio. Object recognition with gradient- based learning. In Shape, contour and grouping in computer vision. Springer, 1999.
Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of optimization algorithms via integral quadratic constraints. SIAM Journal Optimization, 26(1):57-95, 2016.
David D. Lewis, Yiming Yang, Tony G. Rose, and Fan Li. Rev1: new benchmark collection for text categorization research. JMLR, 2004.
Ke Li and Jitendra Malik. Learning optimize. In ICLR, 2017.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv, 2019.
Mitchell Marcus, Mary Ann Mareinkiewicz, and Beatrice Santorini. Building large annotated
corpus of english: The penn treebank. Computational Linguistics, 1993.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: (eee::609.07770 07843, 2016.
Luke Metz, Niru Maheswaranathan, C. Daniel Freeman, Ben Poole, and Jascha Soch-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. arXiv, 2020.
Aryan Mokhtari, Alej, and ro Ribeiro. Global convergence of online limited memory bfgs. Journal of Machine Learning Research, 2015.
Y. E. Nesterov. A method for solving the convex programming problem with convergence rate o11kk)). Dokd. Akad. Nauk SSSR, 1983.
Yurii Nesterov. Introductory lectures convex optimization: basic course. Khuwer Academic Publishers, pages 71-81, 2004.
Lam M. Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáč. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In ICML, 2017.
Jorge Nocedal. Updating quasi-newton matrices with limited storage. Mathematics of Computation, 1980.
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning data diet: Finding important examples early in training. NeurIPS, 2021.
Boris T. Polyak. Some methods speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964a.

---

Published as a conference paper at ICLR 2022
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 1964b.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 1951.
Nicolas Roux, Mark Schmidt, and Francis Bach. stochastic gradient method with an exponential convergence __rate for finite training sets. In NeurlPS, 2012.
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning internal representations by error propagation. Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
Nicol N. Schraudolph, Jin Yu, and Simon Günter. A stochastic quasi-newton method for online convex optimization. In AISTATS, 2007.
Shai Shalev-Shwards and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. JMLR, 2013.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way prevent neural networks from overfitting. JMLR, 2014.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In ICML, 2013.
Oriol Vinyals and Quoe Le. A neural conversational model. arXiv, 2015.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv, 2012.
Dongruo Zhou, Pan Xu, and Quanquan Gu. Stochastic nested variance reduction for nonconvex optimization. arxiv preprint 1rX1v:1806.0781,, 2018.
12

---

Published as a conference paper at ICLR 2022
A
GENERAL CRITICAL GRADIENT ALGORITHM
We present the pseudocode for the Critical Gradient algorithm below.
Algorithm 1 Critical Gradients Optimization
Define #epochs E, #gradient updates B3 model 0, learning rate
Define optimizer Ogge, Oc, ht). which produces parameter update given gradient gt. parameter O and
hyperparameters
Define aggregation function aggr(gt, ge) which outputs a linear combination of g: and the entries in ge Initialize empty critical gradient buffer {ge} with decay rate d for 0 to E do
do Sample an input-output pair (Ir,, M6)
=
9t + Ve (0.)
01+1 O(aggz(9, ge); 0,, h)
if {ge is full then
Add ge to {ge}
else
if 1lg//12 > mim((gg)) then
Remove min((ge) from (ge)
Add gi to {ga}
end if end if
for
g. {ge) do
9-112 l
d 11g-112
end for
end for
end for
B CONVERGENCE PROOF
Proof. Recall that if F Rd + R is twice continuously differentiable, then for all a, R
VF VF(c) + ( x))dt (y (-))..
Letting Fkk = 02 8* denote the suboptimality of the parameters iteration and applying 7 with
0r and 8", we get that
where H = V2 5(0" + trx)dt,
gl Hgrk + Sk.
where &L denotes the gradient noise at iteration k. Using this and bounded staleness assumption, we have that
K-1
K-1
aggr((:;9)) = Σ Wr = Σ
WHA -P Ge
(8)
K=
where {wr, are non-negative scalars in the closed interval [0, 1] used to linearly aggregate the gradients in memory with the most recent gradient 9t- For gradients Jt-k that not stored in memory, the corresponding aggregation scalar W,,k is equal to 0. Without loss of generality, we take wh.0 since we can always incorporate it into the step-size a and re-scale the other weights Wt,1, Wt,K accordingly. Substituting (8) into I and subtracting 0* from each side, we get that
T =
(Ha + G(-R).
k=
Thus, we have that SGDe evolves according to the linear system
Mi we
1
Tu
T
T-K-1
At
where At
Tt11- K
13

---

Published as a conference paper at ICLR 2022
with Ti.0 have that
I - aHl and Ti, -owi H for k = .,....K K 1. Unrolling the recursion,
7++1
01
Diso "
Lizo
0
0
from which clear that we may expect convergence properties to depend on spectral properties of the matrices {Aa A The convergence rate of SGDe will therefore depend on the largest singular value of the matrices [AA]i== From Polyak 1964a Lemma 4), we have that
|| $
sup P(A),
(10)
10 hi E|p.41
where p(A) are the singular values of the numerical matrix
Aa
AK-1
E K
0
with Àg 1 onjo and d -00 for k I,.... -- 4) originally refers to the eigenvalues of 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 cigenvalues of the matrix A A, where commutativity of the individual blocks (7..0..... TjKK 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.
Note that Polyak (164aa Lemma
Letting (() denote the right hand side of (10). we have that ||A I ((,KK) for all j = .....,t, and K() has analytic definition as the roots of a polynomial. Taking 2-norms, applying submul- tiplicativity of matrix norms, and using the bounded variance assumption along with the fact that the noise terms {Gx )= are mutually independent and that the aggregation weights (wo,, in the closed interval [0, 1] for all
have that
E |ree1l < L.
|10 0-
+
_
2(t 2(t-j) (a,k)
(11)
where the expectation on the left hand side is with respect the joint distribution over the random noise variables , By assumption, since the step-size chosen such that ((. < 1, equation
simplifies as
E Elre+ 2| < (a. 2 1|0 - 0-1 + a² K
(a,
where we have implicitly used the upper bound on the limit of a geometric sequence.
C DATA SET DISTRIBUTIONS
Table details the train-valide valid-test splits of all the data sets used in the experiments.
Dataset coviype
PTrain
#Yalid
#Test
NIA
nevl
5000
N/A
N/A
MNIST
S0K
10K
IOK
CIFAR-I0
1OK
CIFAR-I00
10K
SNLJI
SS0K
JOK
1OK
WikTTea
2M
21 3K
241K
B9OK
MoldWo
115K (1.5M)
20K (200K)
20K (200K)
Table 1: The splits of the data sets in number of samples. For WikiText, Penn TreeBank and MultiWoZ the models trained language modeling objective, the splits are given in number of tokens. For MultiWoZ, the number utterances provided with the number of tokens parentheses.
14

---

Published as a conference paper at ICLR 2022
D EXPERIMENT DETAILS
We tested the proposed class of optimizers on different tasks that have potentially different loss surfaces in large parameter spaces. To cast wide and ensure that capture 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 al. 2014), ReLU (Agarap) (018) layers and large Transformer language model architecture RoBERTa Ra-BBase (Liu et al. 2019) across different tasks. brief deseription of the data and architectures used follows.
CIFAR-10 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-T00 is image dataset similar size CIFAR-10 but with 100 classes of images. This task was approached using 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 (Conv/etfEncoder)), 2 layer bi-directional encoder with linear projection layers (InferSent) and a unidirectional LSTM with 2 recurrent layers (LSTMEncoder). The classifier top of the representations learned by the encoder architectures is 2-layer fully-connected Multi-Layer Perceptron with dropout connection.
The Penn Tree Bank (PTB) is vyntax-annotale 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 employed to avoid exploding gradients. We evaluate the model by measuring perplexity (PPL.) in the validation set; lower PPL. scores are preferred.
goal-oriented conversations 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 RoBERTa- Base and BiLSTM Sequence-to-Sequenc architecture (Vinyals and Le 2013) for the next utterance prediction. RoBERTa was trained with CausalLM M--aaddd while BiLSTM model was an encoder-decoder architecture with a Bi-LSTM encoder and LSTM with Attention decoder.
MultiWoZ. 2.0 popular dataset that has
We use models with varying size of trainable parameters as shown in Table
Table 2: We experimented with models with different sizes for a comprehensive study of the proposed variants.
#Params
Dataset(s)
Logistic Regression
8K (55//4))
MNIST (rev l/covtype)
NeuralNetwork
25K
MNIST
ConvNet
CIFAR-10 (CIFAR-10)
1.STM
20K (600K)
PTB (WikiText)
LSTMEncoder
SNLI
InferSent
1.2M
SNLI
ConnNerEncoder
3.
SNLI
RoBERTa-Basc
125M
MuliWoZ
Bi-LSTM
S74K
MutaWoZ
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 to verify the consistency of convergence to better solutions across tasks and architectures. All reported results are averaged over different runs with different random seed values.
D.1 MODEL HPARAMS
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. torch.opt.i im. The details of hyperparameters used for the model are in Table
15

---

Published as a conference paper at ICLR 2022
Table 3: Architecture details for models used in experiments.
Model
Dataset
#Layers
#Hidden
Rel.U/Drooout
Log. Reg.
caviype
NIA
No/No
Log. Reg.
rel
N/A
No/No
Reg.
MNIST
N/A
No/No
Log. Reg.
synthetic
N/A
No/No
MLP
CNN
CIFAR-10
120
Yes/Yes
CNN
CIFAR- 100
4096
Yes
LSTTM
FTB
128
NiNNo
LSTM
WikiText
128
NoNNo
ConvNetEne.
SNLI
Yes/Yes
LST MMSEeddrrrr
SNLI
200
No/Yes
InferSent
RoBERTa-Baal Buse
MalaWoZ
768
YesYees
Bi LSTM Altn
MalaWoZ
200
No/No
D.2 RUNTIME STATISTICS
We logged the approximate time for each epoch for different values of topC across the different models. Although the results populated from the experiments with Adam and its a variants, the results can be extended to the other optimizers and its variants. These times are reported in Table
Table 4: Approximate time taken for epoch of training the models in the public repositories on the different data sets. The time is clocked in minutes. The time per epoch increases linearly with increase in topC over the (B) Jase method. Although for lower values of topC the time per epoch is smaller, there is still room for improvement to bring down the time with smart update rules as discussed in $7
Model
C100
MNIST
Log.Reg
0.3
0.4
0.5
0.6
0.9
.5
Neural-Net
0.3
0.6
2.3
PTB
LSTM
0.08
0.4
0.6
6
WikiTest
LSTM
0.6
1.4
2.6
10
0.5
CIFAR-100
CNN
0.75
1.5
2
12
LSTMEnc.
1
SNLI
InferSent
2
50
ConsNet
1
65
MuliWW7
RoBERTa
15
35
NA
BILSTM
3
10
E HYPERPARAMETERS
E. RANGE OF HPARAMS
We present the range of hyperparameters used in experiments (Table 5). Optimizer-specific parameters were used on both their vanilla and versions.
Table 5: The hyperparameter configurations of the different optimizers in the experiments. Hparam
Choices
Adam 31
(0.9, 0.99, 0.999)
Adam 82
(0.99, 0.999, 0.9999]
RMSprop a
{0.9, 0.99, 0.999)
SGDM momentum
[0.9, 0.99, 0.999)
learning rate
{0.1, 0.01, 0.000, 0.00001
decay
{0.7, 0.9, 0.99)
topC
{5, 10, 20)
16

---

Published as a conference paper at ICLR 2022
E.2 HPARAMS OF THE BEST CONFIGURATIONS
The learning rates and other hyperparameters of the optimizers used in the results reported in the paper listed in Tables 6 and
Table 6: Best hyperparameter configurations for CIFAR-I0 and CIFAR-100 image classification tasks.
Optimizer
Leaming
Topic
Adame:
NIA
NIA
Adam
0001
N/A
NAA
N/A
0.9
0.99
NiA
SMPore
00001
NIA
NIA
N/A
Shallow
RMSprop
CIFAR-10
ConvNet
SGDMc
00011
NIA
NA
NGA
SGDM
0.011
NA
NI
0,9
NIA
NA
NIA
SGDe
0.01
0,99
N/A
NIA
N/A
NIA
Adamg
0.0001
20
0.7
NA
Q
6,999
NIA
Adam
0.000
N/A
NIA
N/A
09
0.9999
NIA
RMSpropG
0.0001
20
0.99
CIEAR 100
ConvNet
Deep
RMSрrop
SGDMc
0.011
20
0.7
0.9
N/A
N/A
NIA
SGDM
0.11
N/A
N/A
0.9
N/A
N/A
NIA
SGDc
SGD
0.1
N/A
N/A
Table 7: Best hyperparameter configurations for language modelling tasks WikiText and PennTree--
Bank datasets.
Dutaset
Model
Optimizer
LR
Momenium
Adamg
Adam
0.0001
0
M/A
0.9
0.999
N/A
RMSSproP
0.0011
20
0.9
V/A
NIA
N/A
0.9
RMSşcop
0.0011
N/A
SGDMc
0.9
N/A
SGIM
0.1
0.9
NIA
N/
N/A
SGDMc
0.1
2
0.95
NA
N/A
N/A
N/A
Adamgg
Adam
0.0011
0
N/A
0.9
(.9999
N/A
RMSprrpp
0.0811
0.9
N/A
N/A
N/A
RMSpcop
0.0811
SGDMc
0.9
N/A
N/A
SGDM
0.1
0
0.9
NIA
N/A
N/A
SGDMc
0.1
0.95
N/A
N/A
N/A
N/A
Table 8: Best hyperparameter configurations for text generation for dialogue on Multi MultiWoZ2.0 dataset.
Dataset
Model
Optimizer
I.
topc
decay
Momentum
91
Ag
Adamç:
Adam
000001
N/A
0.99
0.999
N/A
RMSpropc
0.0011
0.7
NIA
N/A
N/A
0.99
RoRERTT-Rase
SGDMc
RMSpnop
0.00011
0
NIA
NAA
NA
SGDM
0.01
0.9
NA
NA
NIA
SGDc
0.1
0.7
NIA
N/A
MIA
N/A
SGD
NA
Adame
Adua
0.001
N/A
0.99
0.999
N/A
RMSpropc
0.001
0.7
NIA
N/A
N/A
0.9
Ri-LSTM
SGDMc
RMSppop
0.001
0
N/A
NA
NA
NAA
SGDM
0.1
0.9
NA
NA
N/A
SGDc
0.1
0.7
NIA
N/A
M/A
N/A
SGD
N/A
F ANALYSIS
F.1 NOTE 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 s sum (which does not While algorithms converged scale it) than with mean (which scales by factor of 1opC4I-
17

---

Published as a conference paper at ICLR 2022
Table 9: Best hyperparameter configurations for text classification for language inference SNLI
dataset.
Dalaset
Model
Oplimizer
LR
Momenium
Adamc:
0.0001
0.7
NIA
Adam
0.001
NIA
0.99
0.999
NIA
RMSPropp
0.0001
0.7
NiA
N/A
N/A
0.99
RMSpcap
N/A
SGDMc
8.81
0.7
SGDM
0.1
0.9
N/A
N/
NIA
SGDc
0,7
N/A
N/A
N/A
NIA
Adame
0000
0.7
NIA
Adam
0.001
NIA
0.99
0.990
NIA
RMSprrpp
0.0001
0.7
NIA
N/A
0.99
Comenn
SGDMc
0.01
0.7
N/A
SGDM
0.1
0.9
NVA
NA
NIA
SGDc
0.7
NIA
N/A
NIA
Adame
0,000
0.7
NIA
NIA
Adam
0.001
0
N/A
0.99
0,990
NIA
RMSpropp
0.0001
0.7
NIA
0.99
RMSpcup
SGDMc
0.01
0.7
N/A
NIA
SGDM
0.1
0.9
N/A
N/A
NIA
SGDG
0.7
NIA
N/A
NIA
using both aggregation methods, we found that SGDe and SGDMc demonstrated better validation performance using sum whereas RMSprope and Adamc performed better using mean. This is likely due to both RMSprrp and Adam having adaptive learning rates, allowing them to be more robust changes in the scaling of the gradients.
F.2
CONVERGENCE IN CONVEX TASKS
SGDc
10
Adamc
Adamc
0.08
Adamç
SGDc
0.06
SGDMc
Training Loss
Training Loss
Training Loss
SGDMc
SGDMc
SGDc
RMSpropc
RMSpropc
0.04
RMSpropc
0.6
0.4
0.05
10
20
20
10
20
epoch
epoch
epoch
(a) covtype
(b) MNIST
(c) revl
Figure 7: The training on the three different convex datasets show that the variants of the base optimizers empirically converge.
To empirically validate our proof of convergence for the proposed methods on convex loss surfaces, we train all optimizer variants on three tasks with Logistic Regression: fz-regularized binary classification with rcv/ (Lewis al. 2004). fz-regularized multi-class classification on covtype (Dual and Graff, 2017), and non-regularized multi-class classification on MNIST. Figure shows training losses for all optimizers converging towards minimal training loss.
F.3 BUFFER STALENESS BOUND
We present additional results demonstrating the boundedness of the staleness of stored gradients (Figure 8)
F.4
ADDITIONAL. ABLATION EXPERIMENTS
We present additional side-by-side comparisons of optimizers in the vein of $ (Figure 9).
18

---

Published as a conference paper at ICLR 2022
250
250
250
0.7
200
200
200
0.99
0.99
0.99
150
Count
150
Count
150
100
100
100
50
50
50
Age (Steps)
Age (Steps)
Age (Steps)
(a) topC 5
(b) topC 10
(c) t topC 20
Figure 8: Histograms showing ages of the buffer as recorded the end of every epoch, across five seeds, for Logistic Regression on MNIST.
60
60
Val. Accuracy
Val. Accuracy
Vol. Accuracy
50
MNDS
50
FIFO
50
CDS
NCPR
max
40
CSS
40
bottomC
min
random KotH
RR
median
30
KotH
30
sum
50
100
50
100
0
50
100
epoch
epoch
epoch
(a)
(b)
(c)
Figure Additional ablation experiments performed on CIFAR-10
F.5 ROBUSTNESS TO NOISE
As variants retain larger gradients during training, and noisy data points may have gradients with higher values of norm,, experiment with the a variants by training Logistic Regression synthetic binary classification dataset where the labels are perturbed with probability Pnoise-
We sample 500 data points from 7-dimensional hypercuhe with the class means separated by unit. The data is split by 80/20 for training and evaluation. We train the model with the optimizers for 20 epochs on fixed learning rate of 1 10-3 topC decay 0.5 and other parameters from the base model set to their defaults.
We compare the performance of base method and its a variant study anomalies when the a variants are exposed to noise (Figure I0). We observe that the models do not behave any differently than the base methods showcasing that the noise the dataset does not affect the workings the a 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 synthetic data set for experiments, we used sklearn' make_classiiicco method. The specific hyperparameters used to construct the dataset:
samples 500, n_ features 10,
infformativ=77,.. r_rennddandddn. repeated=0),
n_classe==2, suinesse___cuussessssp___asss__
beightsoclas_jmmaalnnee flip_y=noise, class sesss_sep=000 0.5,
hypercube=True, shift=0.4, shift:0:0 scale=1.0, shuffle=True, random 3taddmm_tttte11033
19

---

Published as a conference paper at ICLR 2022
Adamc
0.8
0.8
0.8
0.8
Adam
F1
0.7
[1score
0.7
F1 F1score
0.7
score
0.7
SGDc
0.6
SGDMc
0.6
RMSpropc
0.6
SGD
SGDM
RMSprop
0.5
0.0
0.5
10
0.0
0.5
10
0.0
0.5
10
0.0
0.5
1.0
Phebe
(a)
(b)
(c)
(d)
Figure 10: Proiie is increased from 0 to in increments of 0. indicating 10% more injection of noise. The proposed variants did not show any anomaly on noise induced dataset. For the experiments on the decay the topC was 5 and decay 0.5. The plots are averaged over 5 different runs with different seeds.
F.6 DECAY
The optimizers' sensitivity to decay (Figure [1] had similar trend in other experiments, where the optimizers showed better performance when decay was away from 0.99.
VAl Accurery
Val Actunccy
22
Log. Reg.
MLP
20
ConviNt
0.99
0.05
decey
decay
(a) HolliWWZ-BiiSTMM
(b) CIFAR-10
(c) MNIST
Ro3E
Valitatton Perploxtty
LSTM
240
LSIM
Val
2.5
280
200
0.2
0.2
180
decay
(d) RoBBRoZ
(e) PTB
(I) WikiText
Figure 11: Experiments varying decay that across tasks. Graphs in the top row measure BLEU score or accuracy (where a higher value is desired) and those in the bottom row measure perplexity (where lower value is desired)
F.7 TOPC
The optimizers' sensitivity to topC (Figure 12) had a similar trend to other experiments, where the optimizers showed hetter performance for lower values of topC.
F.8
go VS 9a
As follow-up experiment to validate the non-informative gradients with higher values of topC. we observe the trend in the difference between ave ragee(1g)ll))) and ||gel12 for different values of topC. We see that the gradients 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 the past gradients in memory has theoretical advantages, in practice observe that lower topC provides better training signal through ge than Sr (Figure This explains the slightly lower performance of topC =100 across models SNLI task and the nonexistent marginal improvements in MultiWoZ dataset with RoBERTa Ta-Base model.
20

---

Published as a conference paper at ICLR 2022
ACCurac
21.5
86
Medel
ConvNet
200
0.05
0.1
0.2
docay
(a) MultiWo7.
(h) CIFAR-10
Valdeton Perpletty
210
LSTM
Accurety
MLP
Log Reg.
190
90.0
topC
topC
(c) WikiText
(d) MNIST
Figure 12:
Experiments on varying topC that across tasks. Top row: (a) MultiWoz BiLSTM
(BLEU Score) (b) CIFAR-10 Convnet (Accuracy). Bottom row: (c) WikiText LSTM (Perplexity) (d) MNIST Log. Reg. and MLP (Accuracy).
0.10
10
100
0.00
100
0.00
100000
steps
steps
(a) CIFAR-100
(b) SNLI
Figure 13: The difference between ave average( (1gg|12) and ||ge shown in plots (a) and (b) indicate that increasing topC could hurt the performance. These plots correlate experiments in $6
21

---

Published as a conference paper at ICLR 2022
The plots of St 9e is crucial to the results in that they provide explanations for the faster convergence (as well as smaller improvements) of the variants over the base methods (Figure Since the parameter updates reminded of large displacement of the gradients via the aggr method, the model updates the parameters more frequently. In the cases where Je is higher than St, we observe better performance and no improvements in cases where the difference is not as significant.
14
0.10
5
0.10
20
0.05
0.05
20
0.0
100
0.00
100
0.00
100
50000
100000
150000
50000
100000
150000
50000
100000
150000
(a) SNL.I ConvNetEne.
(b) SNLI I. MTNnondeer
(c) SNLI Infersent
0.050
0.002
30
0.000
20
0.000
-0 0.001
10000
15000
2000
4000
1020
5000
10000
2005
2000
stepm
(d) CIFAR-10
(e) PTB
(1) Roberta-MWoZ
Figure 14: Additional results on St vs Jc, explaining rationale behind the better performances of lower values of topC across all the experiments. As the variance in aggr get diminished with increase in topC the model does not gain much with increase in topc.
Further, we observe the significance of the gradients as defined in & fo-norm diminishes as topC increases. This indicates that the optimizer needs only to store only small set of critical gradients in order to improve the performance without incurring much time and memory overhead. Employing lazy updates down the time useful future direction of research.
F.9 DETAILS OF ABLATION EXPERIMENTS
We detail the various ablation experiments whose results are presented in $G
King 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 min-heap
First-In-First-Our (FIFO): Replaces the max-heap (priority queue) buffer with a queue which begins dequeue when it reaches capacity of topC.
"Coin-Toss"/1 randdom""" an incoming gradient 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 its difference from the mean of the norms of gradients in the buffer.
Cosine Diversity Sampling (CDS): Incoming gradients 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 full and new gradient is added, the gradient which gets removed gets selected at random.
22

---

Published as a conference paper at ICLR 2022
Norm-Controlles Probabilistic Replacement (NCPR): When the buffer is full and gradient is added, the gradient Ik which gets removed gets selected with probability l
2, llailla"
G TEST PERFORMANCE
The complete test performance the models reported in Table I0
H REPRODUCIBILITY CHECKLIST
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 DD
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 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 final runs we used I 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 reported in $D2
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 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 fraction of the covtype and revl 1 datasets instead of using the entire data. We sampled 5000 datapoints at random with seed set 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.nd files in the submitted repository.
23

---

Published as a conference paper at ICLR 2022
Table 10: Complete test results on the different tasks.
Model
Dataset
Metric
Optimizer
Performance
KabLRT
StWa
FPL
Adaı
2.77 1: 0:0
Adamc
RoBERT,
MWaZ
PPL.
RMSprop
2.56
RaBERTa
MWaZ
PPL
RMSPropc
2.42 0.0
RaBERTa
MWaZ
PPL
SGDM
2.46 0.0
PPL
SGDMc
RaBERTa
MWaZ
PPL
SGD
2.62
RaBERTa
MWAZ
PPL
SGDc
2.41 0.0
Model
Dalaset
Metric
Optimizer
Performance
BILSIM
StWoZ
BLEU
Adam
20.64 a 0.19
BILSTM
Adamc:
20.94
RMSprop
BiLSTM
MWaZ
BLEU
RMSPrope
21.21
BILSIM
StWaZ
BLEU
SGDM
19.31 +: 0.27
BILSIM
MWWZ
BLEU
SGDMc
19.58
BILSTM
MWaZ
BLEU
SGDg
16.35 0.24
Model
Metric
Optimizer
LSTM
PTB
PPI.
Adaam
153.54
1.65
LSTM
FIB
PPL
Adamc
133.28
0.61
RMSprop
RMSpropC
LSTM
PTB
PPL
SCDM
139.95
0.64
LSTM
PTB
PPI
SGDMc
132.22 1.1
PPL
SGD
1861
LSTM
PTB
PPL
SGDe
295.33
Model
Metric
Optimizer
Performance
LSTM
WixTlexi
PPL
Adam
170.3 0:
LSTM
WikiText
PPL
Adame
169.23
1.06
LSTM
WakTTest
PPI
RMSprop
180.42
1.64
IMSpropp
LSTM
WikeTex
PPL
SGDM
166.82
LSTM
PPI
SGDMc
156.75
1.32
LSTM
WakTText
PPL
SGD
461.81 12.68
WikeTest
PPL
Model
Dataset
Metric
Optimizer
Performance
Accuracy
InderSeat
SNLI
Accuracy
Adame
79.03 0.18
InterSent
SNLI
Accuracy
RMSprop
77.91 0.35
InferSent
SNLI
Accuracy
RMSprope:
79.07 0.28
Accuracy
SGDM
79.15
InterSent
SNLI
Accuracy
SGDMc
78.52 0.25
InterSent
SNLI
Accuracy
SGD
78.3 :1 0.16
InderSeat
Accuracy
SGDe
Model
Dataset
Matric
Optimizer
Performance
ConsNer
Accuracy
Adam
0.0
Accuracy
Adame
72.77
ComWNrr
SNL
Accuracy
RMSprop
74.69 0.36
ConvNer
SNLI
Accuracy
RMSpropC
77.49 0.24
ConvNet
SNLI
Accuracy
SGDM
76.79
Accuracy
SGDMc
78.06
ConvNer
SNLI
Accuracy
SCD
78.5 0.1
ConvNet
SNLI
Accuracy
SGDc
78.39 0.24
Model
Dataset
Metric
Optimizer
Performance
LSTMEnc
SXLI
Accuracy
Adamm
76.79 0.38
Accuracy
Accuracy
RMSprop
76.37.2
LSTMEne
SNLI
Acsuray
RMSpropC
77.48
LSTMEBc
SNLI
Accuracy
SGDM
77.6 0.22
L.TMFFnc
Accuracy
SGDMc
LSTMEne
SNLI
Accuracy
SGD
74.92 2:2:::10
LSTMFBne
SNLI
Accuracy
SGDc
76.81 2: 0.59
Model
Optimizer
Performance
ComNett
CIF
Accuracy
Adam
51.71 0.S5
ComNet
CIFAR 100
Accuracy
Adamc:
53.55 1: 0.42
Accuracy
RMSprop
Cum-Net
CIFAR100
Accuracy
RMSprope
52.93
ConvNet
CIARRO0
Accuracy
SGDM
56.03 0.49
Com-Net
CIEAR100
Accuracy
SGDMc
59.07 0.2
CumNNet
Accuracy
SGD
55.37
Con-NNt
CIFAR100
Accuracy
SGDe
38.0 0.37
Model
Metric
Optimizer
Performance
ConvNer
CIFARIO
Accuracy
Adam
64.26 0.47
Com-Net
CIFARIO
Accuracy
Adamg
63.72 4: 0.57
Com-Net
CIEARIO
Accuracy
RMSprop
Accuracy
RMSproPC
Con-NNer
CIARROO
Accuracy
SGDM
64.7 0.69
ComwNer
CIFAR10
Accuracy
SGDMc
63.68 2 0.77
ConvNer
CIFAR1O
SGD
Accuracy
SGDe
24


Reviews:
[{"review_id": "uJuixu-wR-C", "reviewer": "Reviewer_tj5C", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- A general framework which can accommodate existing first order optimizers to trade off memory with faster convergence.\n- Convergence proof for smooth strongly convex functions \n- Solid set of experiments\n  - Sensitivity to hparams\n  - Alternatives to l_2 norm for deeming gradients critical\n\n\nAreas of improvements and questions:\n- Either make figure 1 vertical or at least make the text vertical so that its easier to read.\n- In Figure 2, I see that several tasks degrade convergence. Is there a way to make sure this method does no worse than before? I will increase my score if this is done.", "summary_of_the_paper": "Bridging the advances from theoretical and empirical side, this paper proposes a framework for memory augmented optimizers, a strategy which can improve any off-the-shelf first order optimizers trading off increased memory with faster convergence. ", "summary_of_the_review": "Overall good paper, some areas of improvements", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "ns-xodFkZRM", "reviewer": "Reviewer_XUqb", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The critical gradient strategy seems fairly interesting and a promising way to accelerate training. The paper was fun to read and well written. The experimental evaluation is fairly comprehensive and has interesting analysis. The main concerns I have are:\n* Although the convergence speed in terms of iterations is faster, the wall clock time relative to baselines especially with large buffer sizes seems much higher Appendix D2 Table 4. The overhead seems to stem from the book keeping, replacement and aggregation strategy. Is there a convergence speed comparison Figure 2 with wall clock times as opposed to number of gradient updates?  \n* In Analysis there is nice study on the staleness of the gradients in the buffer. Can the staleness bound just be enforced instead of relying on it being empirically the case?  \n* If the critical gradients are finally aggregated either via sum or averaging there can be easier strategies for doing the same aggregation in a more online fashion without additional memory. Especially given that there is a decay term when doing gradient selection for replacement. Was there an attempt to try more online approaches?\n", "summary_of_the_paper": "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. \n", "summary_of_the_review": "The proposed method of only storing critical gradients in a limited size memory buffer seems interesting. The experimental evaluation is comprehensive. However, the method only seems to improve in terms of gradient steps but not wall clock time to converge due to bookkeeping overheads.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Kemf0XTgT5u", "reviewer": "Reviewer_KE2X", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros\n* The selection of the critical gradients directly depends on the $l_2$ norm of the gradients and the aggregation function is simple enough, so the idea itself looks reasonable and the implementation of the algorithm is straightforward. Furthermore, the optimizer could be combined with other popular first-order optimizers which can make it more useful.\n* The related work section is well-written and the number of citations is adequate. I feel comfortable reading through the paper and the points are obvious to me. The comparison between the proposed method and other algorithms on memory requirements and performance improvement is clear.\n* The experiment results are strong. Figure 2 gives a nice summary about how the proposed memory-augmented optimizer influences performance of the original optimizer.\n\nCons\n* Although it seems natural to choose gradients of large norm as the critical ones. I did not find any motivation or explanation for that. There are several related works trying to utilize the information of historical gradients, but it is not clear how saving the critical gradients by looking at their norm helps improve the performance. More specifically, In Figure 5 (a) and (b), it seems the choices of selection techniques and replacement methods do not influence the results a lot. I would be rather interested if the author could give some insights about the reason why such selection mechanisms work well.\n\n* The theoretical proof looks merely fine to me. It is not clear how the choice of $K$ will influence the bound. It would be more obvious and convincing if the authors could give some examples rather than simply claim it may be possible to accelerate convergence rates when $K > 1$. I guess such acceleration might depend on specific conditions and these conditions could show some improvements and limitations of the proposed algorithm. \n\nComments\n* Maybe there is a limitation for the number of pages, but I personally feel it will make the paper look more complete if Algorithm 1 goes to the main body.\n* The description of $A_t$ at the bottom of page 13 seems to be wrong. Shouldn't the blocks of identity matrices be moved one block to their left?", "summary_of_the_paper": "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.", "summary_of_the_review": "In general, the paper is well-written and the experimental results look strong. However, the lack of clear motivation and explanation of the proposed method downgrades my rating to the paper. Furthermore, the theoretical part could be more detailed as the current one does not provide a convincing improvement over vanilla SGD method. I hope the authors will answer these questions in the future.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HGI-uTcyPOV", "reviewer": "Reviewer_v5GJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- analysis on strongly convex smooth objectives\n- good empirical evaluation\n- strong results on evaluated datasets\n\nWeaknesses\n\n- nowhere near SotA validation on CIFAR 10/100? Validation should be around 75% looking at wideresnets, results, while this is an optimisation paper, I feel like this makes it hard to compare/more vulnerable to accidental cherry picking\n- architectures only visible in supplimentary, and appendix listing them for analysis would help the reader to have all information in one place\n- strongly convex smooth is a strong assumption to place on convergence, although the empirical performance alleviates this to some degree\n- I'm missing some comparison to other discussions of gradient staleness e.g. from the Asychronous Distributed SGD community http://proceedings.mlr.press/v80/damaskinos18a/damaskinos18a.pdf \n\nQuestions:\n\n- do you have any explanations of the cases where you regress in performance/speed? violation of the assumptions, high gradient variance or anything else?", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, I like this paper as a simple to implement empirical improvement to other algorithm that just \"makes sense\": very large gradients should probably have an impact for a longer time BUT not all at once so as to not dominate fully.  It gives good but not overwhelming empirical evidence (difficult to compare to SotA due to architecture choice and different tuning) and some theoretical justification (albeit not applicable to the settings evaluated).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a general method to enhance the performance of first-order optimizers. The main idea is to use a memory buffer to maintain a limited set of critical gradients from recent history. Namely, gradients with large l2 norm. The paper includes a convergence proof on strongly convex smooth objectives. Experimental results are reported for several architectures in vision and language tasks. When integrated with several commonly used optimizers (SGD, SGDM, RMSProp, Adam), the method shows an improvement in terms of learning speed as well as improved performance (in almost all cases). Several ablations are performed to show strong robustness to hyperparameters introduced.

The paper is very well written and easy to follow. The proposed method is simple and effective. The empirical evaluation is strong and the ablation studies exhaustive and convincing, as pointed out by all four reviewers.

The authors provided a solid rebuttal addressing many questions raised by the reviewers.

Reviewer KE2X is the only reviewer recommending to reject the paper, pointing out that the paper lacks a clear motivation on why the critical gradients are selected based on the l2 norm. In their response the authors provided a connection with recent methods designed to find important examples when training neural networks. In the discussion (not visible to the authors). Reviewer v5GJ stated that this criterion is intuitive and acceptable given the strong empirical performance reported. The AC agrees with this view.

Reviewer KE2X also points out that the theoretical results do not provide a convincing improvement over vanilla SGD. The authors acknowledge this point, adding that the convergence results apply to a wide variety of critical gradient methods and aggregation strategies, so it is natural to expect that improvements will depend on the specific conditions. While the AC agrees that having stronger theoretical results would improve the paper, as it stands it is certainly above acceptance threshold.

Overall the paper makes a solid contribution with a method that is simple and effective. It will likely inspire other alternative methods in the future. The AC recommends accepting the work.
Decision: Accept (Poster)