Promoting Exploration in Memory-Augmented Adam using Critical Momenta
Authors: Pranshu Malviya
Goncalo Mordido
Aristide Baratin
Reza Babanezhad Harikandeh
Jerry Huang
Simon Lacoste-Julien
Razvan Pascanu
Sarath Chandar
Venue: Accepted by TMLR ()

Abstract: Adaptive gradient-based optimizers, notably Adam, have left their mark in training large-scale deep learning models, offering fast convergence and robustness to hyperparameter settings. However, they often struggle with generalization, attributed to their tendency to converge to sharp minima in the loss landscape. To address this, we propose a new memory-augmented version of Adam that encourages exploration towards flatter minima by incorporating a buffer of critical momentum terms during training. This buffer prompts the optimizer to overshoot beyond narrow minima, promoting exploration. Through comprehensive analysis in simple settings, we illustrate the efficacy of our approach in increasing exploration and bias towards flatter minima. We empirically demonstrate that it can improve model performance for image classification on ImageNet and CIFAR10/100, language modelling on Penn Treebank, and online learning tasks on TinyImageNet and 5-dataset. Our code is available at https://github.com/chandar-lab/CMOptimizer.

Published in Transactions on Machine Learning Research (06/2024)
Promoting Exploration in Memory-Augmented Adam using Critical Momenta
Pranshu Malviya
Mila Quebec AT Institute, Polytechnique Montreal
Gonçalo Mordido
goncabonentido@@@aall.coo
Mila Quebec AT Institute, Pokytechnique Montreal
Aristide Baratin
baratina @milla. quebec
Samsung SAIT AI Lab Montreal
Reza Babanezhad Harikandeh
bubumnezioad@gmmaill.c
Samsung SAIT Al Lab Montreal
Jerry Huang
Berry.huanng@@mmlla.. quebbe
Mila Quebec Al Institute, Université de Montréal
Simon Lacoste-Julien
slacoste Diro. wmwntreal.ccm
Mila Quebec AI Institule, Universilé de Montréal, Samsung SAIT AI Lab Montreal, Canada CIFAR AI Chuir
Razvan Pascanu
mapp@googgee.com
Google Deep
Sarath Chandar
sorath. chandar@miiaaaaaa quebec
Mila Quebec AI Institute, Polytechnique Montreal, Canada CIFAR AI Chair
Reviewed on OpenReview: https: // openreview. net/ forum? id= sHSkJqyyqg
Abstract
Adaptive gradient-based optimizers, notably Adam, have left their mark in training large-scale deep learning models, offering fast convergence and robustness to hyperparameter settings. However, they often struggle with generalization, attributed to their tendency to converge to sharp minima in the loss landscape. To address this, we propose a new memory-angmented version of Adam that encourages exploration towards flatter minima by incorporating buffer of critical momentum terms during training. This buffer prompts the optimizer to overshoot beyond narrow minima, promoting exploration. Through comprehensive analysis in simple settings, we illustrate the efficacy of approach in increasing exploration and bias lowards flatter minima. We empirically demonstrate that it can improve model performance for image classification on ImageNet and CIFAR10/ 11000, language modelling on Penn Treebank, and online learning tasks on TinylmageNet and 5-dataset. Our code is available at htttps://iithu: com/ chandar- 1annaar-1bb/CMppcimzae
1 Introduction
Deep learning models are often sensitive to the choice of optimizer used during training, which significantly influences convergence speed and the qualitative properties the minima to which the system converges (Chol ct al. 2019). Stochastic gradient descent (SGD) (Robbins & Monro 1951), SGD with momentum (Polyak) 1964], and adaptive gradient methods such as Adam (Kingma & Ba 2015 have been the most popular choices for training large-scale models.

---

Published in Transactions on Machine Learning Research (06/2024)
Adaptive gradient methods are advantageous they automatically adjust the learning rate per-coordinate basis, converging quickly with minimal hyperparameter luning by using information about the loss curvature. However, they are also known to achieve worse generalization performance than SGD (Wilson et al. 2017 Zhou et al. 2020 Zou et al., 20230. which several recent works suggest is due to the greater stability of adaptive optimizers (Zhou al. 2020 Wu et al. 2018a; Cohen al. 2022). This can lead the system converge to sharper minima than SGD, resulting worse generalization performance (Hochreiter
Schmidhuber]
1994
Keskar al. 2016
Dziugaite & Roy 2017 Neyshabur al. 2017 Chaudhari al.,
2017 lzmailov al. 2018 Kaur et al. 2023
However, similar to exploration in reinforcement learning, we hypothesize that equipping Adam with on exploration stralegy could improve performance by escaping sharp minima. Building upon the framework proposed by McRae et al. (2022), which maintains a buffer containing a limited history of gradients from previous iterations (called critical gradients or CG) during training, the goal is to allow the optimizer overshoot and escape sharp minima by adding inertia the learning process, as to control for the necessary width of the minima in order for the system to converge. However, we show that the original eemmory-nggenned adaptive optimizers proposed by McRae al. (2022), particularly Adam using CG (referred to as Admm+CG), suffer from gradient cancellation: phenomenon where new gradients have high directional variance and large norm around sharp minima. This leads to the aggregated gradient over the buffer to vanish, preventing the optimizer from escaping sharp minima, which is in agreement with the poor generalization performance presented by McRae et al. (2022).
We propose to instead store critical momenta (CM) during training, leading to new memory- auggeee-dggentedd version of Adam (Algorithm 1) that effectively escape sharp basins and converge to flat loss regions. Figure D Dillustrates the optimization trajectories, a toy 2D loss surface corresponding to the Goldstein Price (GP) function (Picheny et al. 2013), of Adam, Adam +GG, Adam+CM, and Adam combined with sharpness- aware minimization (Adam + SAM) (Foret et al., 2021) from different initialisations. We observe that while other optimizers converge to higher and often sharper loss regions, Adam+ CM is able to find the flat region that contains the global minimum.
Algorithm 1: Adam with Critical Momenta
Adam
Adam+SAM
Require: Initial paramcters 6g and moments mo, vo loss L, step size buffer me, capacity C, decay X for 1= 1,2, ... do
Adam+CG
Adam+CM
Sample mini-batch & compute loss gradient Update 1st moments m, with equation 4
Aggregate buffer moments ma my with equation Update 2nd moments eM with equation if buffer is not full then
Add my to else Priority(m.) > min(Priority(m,) then Replace smallest priority element with m,
end
Decay Priority (m,)) using
Update parameter 0, with equation 7 end for
Figure 1: (Left) Learning trajectories for different optimizers on the Goldstein-Pric loss function starting from different initial points. While the other optimizers get stuck in sub-optimal surfaces, Adam+CM explores a lower loss surface and is able to reach the global minimum. (Right) Pseudo-code for Adam with critical momenta (M).
The key contributions of our work are as follows:
2

---

Published in Transactions on Machine Learning Research (06/2024)
We introduce framework for promoting exploration in adaptíve optimizers |section 3p. We propose new memory-angmentee version of Adam, which stores and leverages a buffer of critical momenta from previous iterations during training.
We provide a theoretical convergence analysis of our method in simplified settings (subsection 3.2)
Using numerous examples and benchmarks, we illustrate how our method surpasses existing memory- augmented methods and promotes exploration towards flat minima (section 4).
We observe empirically an improvement in model performance in supervised and online learning settings (sectiin 5
2
Related work
To improve convergence speed and achieve better generalization deep learning models, numerous optimizers have been proposed. While SGD with momentum tends to show superior performance in particular scenarios, it usually requires careful hyperparameter tuning (Le et al 2011 On the other hand, adaptive optimization methods (Duchi et al., 2011 Hinton et al., 2012 Zeiler 2012), which adjust the learning rate for each parameter based on past gradient information to accelerate convergence, have reached state-of-the-art performance in many supervised learning problems while being more robust. to hyperparameter choice. In particular, Adam (Kingma & Baj 2015) combines momentun with an adaptive learning rate and has become the preeminent choice of optimizer across variety of models and tasks, particularly in large- scale deep learning models (Dozat, 2016 Vaswani et al. 2017). Several Adam variants have since been
proposed (Loshchilov & Hutter 2019 Zhuang et al., 2020
Granziol et al. 2020
Defazio & Jelassi 2022) to
tackle Adam's lack of generalization ability (Wu et al. 2018b Zhou et al. 2020 Zou et al. 2022).
2023
Cohen et al.
Converging to flat minima has been shown to be a viable way of indirectly improving generalization
performance (Hochreiter & Schmidhuber 1994 Keskar et al. 2016 Dziugaite & Roy 2017 Neyshabur et al.,
2017 Izmailov et al.] 2018 Kaur et al. 2023 Jiang et al.] 2020). For example, sharpness-awe aware minimization
(SAM) Foret et al. (2021) jointly maximizes model performance and minímizes sharpness within a specific ncighborhood during training. Since its proposal, SAM has been utilized in several applications, enhancing generalization in vision transformers (Dosovilskiy et al.) 2021 Chen el al. 2022), reducing quantization error [Liu et al., 2023), and improving model robustness (Mordido et al. [022].. Numerous methods have been proposed to further improve its gencralization performance, c.g. by changing the neighborhood shape (Kim) et al. 2022b) or reformulating the definition of sharpness (Kwon et al., 2021 Zhuang el al. 2022), and to reduce its cost, mostly focusing on alleviating the need for the double backward and forward passes required by the original algorithm (Du et al. 2022a Liu al. (022)..
Memory-augmentedd optimizers extend standard optimizers by storing gradient braseent-based information during training to improve performance. Hence, they present a trade-off between performance and memory usage. Different memory augmentation optimization methods have distinct memory requirements. For instance, stochastic accelerated gradient (SAG) (Roux et al.) 2012) and its adaptive variant, SAGA (Defazio et al.) 2014), require storing all past gradients to achieve a faster convergence rate. While such methods show great performance benefits, their large memory requirements often make them impractical the context of deep learning. On the other hand, one may only use a subset of past gradients, as proposed in limited-history BFGS (LBFGS) (Nocedal 1980). its online variant (oLBFGS) (Schraudolph et al. 2007), and stochastic dual coordinate ascent (SDCA) (Shalev-Shwartz & Zhang. 2013). Additionally, memory-augmenter frameworks with critical gradients (CG) use a fixed-sized gradient bulfer during training, which has been shown to achieve a good performance and memory trade-off for deep learning compared to the previous methods (MeRae et al., 2022).
In this work, we further improve upon CG by storing critical momenta instead of critical gradients, leading to better exploration of the loss surface by adaptive optimizers, particularly Adam.
3

---

Published in Transactions on Machine Learning Research (06/2024)
Adam+CM Adam+CG
Adam+CM
Adam+CG
Buffer quantit
Buffer quantities
Buffer mean
Buffer mean
New gradient
New gradient
Figure 2: First 10 steps of Adam + CG and Cd trajectorics on Ackley loss surface. Coloured diamonds represent the final points reached by the optimizers. Gradient cancellation is observed in Adam+ CG as buffer mean and new gradients cancel each other out, yielding a small update. Conversely, Adam+ escapes sub-optimal minima and converges near the global minimum.
3 Memory-augmented Adam
We build upon the menory-angmented framework presented by McRac et al. 2022 and focus Adam in supervised learning setting. Adam (Kingma & Ba 2015) has standard parameler updates
m: B1m--1 (1 31)g
828 (1 B2)9..
(1)
my
1t
me
m
Oc+1 0g V0, +
(2)
3
3
where Og denotes the model parameter iteration 9t the loss gradient on the current mini-batch, is the learning rate, 82 E [0, 1) are the decay rates for m. and ut.
Critical gradients (CG). To memory-augment Adam, McRae el al. 2022 introduce fixed-size buffer of priority gradients 9. maintained in memory during training, and apply an aggregation function over this buffer to modify the moment updates (Equation
mi Bim + (1 - 1:, (1))ggr((()..))
-20 + (1 (b)aggr((9)) B2)
(3)
The gradient norm is used as the selection criterion for the buffer. The buffer takes the form of dictionary where the key-value pairs are (Ilgell2, 9e); additionally, the priority keys are decayed at each iteration by a decay factor (0, 1) to encourage buffer update. Thus, at each iteration t, if the norm l|gell2 the current gradient is larger than the smallest priority key in the buller, the corresponding critical gradient gets replaced by St in the buffer. A standard choice of aggregation function adds St the average of the critical gradients in the buffer.
The gradient cancellation problem. However, combining Adam with critical gradients has its pitfalls. We hypothesize that with CG, while the buffer gradients can promote exploration initially (Figure 1), the parameters remain fixed within sharp regions due lo gradient cuncellation. Gradient cancellation primarily occurs when existing buffer gradients are quickly replaced by high maggitage gradients when the parameters are near a sharp basin. As a result, the buffer quickly converges high variance gradients whose mean goes zero, allowing learning converge. Intuitively, the parameters bounce back and forth off the sides and bottom of the sharp basin: whenever the parameters try to escape the basin, the outgoing gradient.
4

---

Published in Transactions on Machine Learning Research (06/2024)
gets cancelled by incoming gradients in the buffer. Figure illustrates this phenomenon toy surface, by showing the buffer gradients (thin blue lines) and their means (black arrow) well as the new gradient. (green arrow), within sharp basins where Ad gets stuck. Additional plots are found in Appendix A.1
3.1 Critical momenta (CM)
As gradient cancellation hinders the ability of Adam+CG to escape sharp minima, our approach addresses this by leveraging a buffer m. of crilical momenla me during training. Like McRae et al. (2022), we use the gradient 12-norm as priority criterion The buffer a dictionary of key-value pairs (1ge) 2. me) with a factor 1) with which the values are decayed at each iteration. The integration with critical momenta leads to new algorithm, Ad,m+CM, defined by moment updates:
my -1mm-1 (1 11-B3 Bi)gt,
mi M aggr (mg, m.).
(4)
"/2 320, + (1 - -2) aggr (m4: m²)²
(5)
where aggr is the addition of the current momentum to the average of all critical momenta:
aggr(m,, m²) = me + C Σ
mc.
(6)
m. Em.
Finally, the Adam+CCM update rule given by
M 3
M
m
Ott1 0:
(7)
32
The pseudo-code of Adam+CM given in Algorithm
While at a sharp minimum, the elements of the buffer will still be quickly replaced (Figure 1 due to the inertia in the momentum terms the variance will stay low. Moreover, the fact that gradients quickly change direction will lead to the new momentum terms being smaller and hence have a smaller immediate influence on the aggregate value of the buffer. This allows the overshooting effect to still happen, enabling the exploration effect and helping to learn to escape sharp minima. Furthermore, the larger the size of the buffer, the stronger the overshooting effect will be, and the wider the minimum needs to be for learning to converge. That is because learning needs stay long enough the basin of a minimum to fill up most of the buffer in order turn back to the minimum that it jumped over and for the optimizer to converge. We observe this empirically in Tigure 9 and Appendix A.2.2
10a
3.2 Convergence analysis
10-1
Convergenca
10
In this section, we discuss the convergence of gradient- based algorithms with critical momenta under simplifying assumptions. Leaving aside the specificity of Adam, we consider the simplest variant characterized by the following update,
10'1
Momentum
10-+
[rrical Mornntta
10
Classical Momentumn (beta-0.9)
Critcal Momenta (beta=0.99
10-4
10:
102
10
10a
10a
Condition number
D++1 02 - a aggr aggr(m;; m.)
(8)
Figure 3: Quadratic convergence rates (1 p²) of clas- sical momentum and critical momenta. Solid curves indicate that both and / were optimized to obtain p", while dashed lines indicate that p² obtained with 0.9. Critical momenta converges for a wide range of condition numbers both cases.
In what follows, we restrict our attention quadratic losses and assume that the optimum lics 0* for simplicity, so that L(0) t HO. We also make the following assumptions: (i) as in McRae et al. (2022), we assume a fixed bound on the staleness of momenta, i.c., there is an integer K > 0 such that
We do use the alternative [|me]]a since the buffer will not get updated fast enough using this criterion.
5

---

Published in Transactions on Machine Learning Research (06/2024)
at each iteration t, m me implies k I K; (ii) further assume that K coincides with the buffer size and that at iteration i, mc contains ezaclly the momenta m4-1," -C
Under these assumptions, the convergence can be analyzed using standard spectral analysis for multistep linear systems. Considering Vi+1 = 18 Br, my mg-1.... M,-C41l, it straightforward to show that the dynamics can be cast as linear dynamical system V++1 AV, where the matrix A depends on the Hessian HI, the learning rate a and the momentum parameter 3, and takes the form:
I aH
0
a(3
T
0
0
H
0
31
A =
0
0
I
: :
: 0
-.
Assuming L-smoothness and j-strongly convexity, the worst-case bound on the convergence rate is obtained by maximizing the spectral radius (largest singular value) p(A) over all admissible (Lessard et al., 2016) Let denote the eigenvalues of the Hessian matrix H. Note that the singular values of the block matrix A are the same as those of the matrices As, obtained by collapsing the blocks of by replacing I by 1 and H by h. This leads to the simplified problem:
Ma, max p(As) NE(u,L)
The optimal convergence rate obtained by tuning these parameters minimize the spectral radius, i.e., p² mi p(a, 3). Similar to Zhang al. (2019), we solve this problem numerically using standard solvers to compute the eigenvalues and minimize the worse spectral radius; we also compare with the convergence rate of critical momenta with classical momentum diagonal quadratic functions in Figure 3 We plot the convergence rates of classical momentum and critical momenta with C = 5 on a quadratic loss function for two variants: (i) Optimal a (solid curves), (ii) optimal a with fixed B = 0.9. In both variants, we can see that critical momenta converge for a wide range of condition numbers.
Following the litcrature [Kurr et al. 2023), we use the maximum eigenvalue (Amaz) of the Hossian H as the indicator of sharpness in the rest of the paper: as hmar increases, the surface becomns sharper.
4 Insights from toy examples
In this section, we use loy lasks to empirically validate our hypothesis by analyzing and comparing various combinations of Adam with memory augmentation and sharpness-aware minimization.
Critical momenta promote exploration. We first compare the optimization trajectories of Adam+CM with Adam, AdM,+SAM, and Adam+CG, on interpretable, non-convex 2D loss surfaces. We also include the double combination of Adam with SAM and CM. To complement the Goldstein- Price function in Figure I we consider the Ackley function (Ackley) 1987 (see equation 11 Appendix A.2.1 for the explicit formula), which contains a ncarly flat outer region with many sharp minima and a large hole at the center with the global minimum at (0,0).
We minimize the Ackley function for 10 different initialisation seeds and compare the trajectories of the different optimizers. We run each model for 500 steps and reduce the learning rate by factor of 10 at the t5 step. To get the best-performing setup, we perform a grid search over the hyper-parameters for each optimizer. Figure 4 shows the training curves (left) and optimization trajectories (right) of the different optimizers, for the same initialisation (black square). We observe that, here, only the CM variants are able explore the loss surface, resulting in lower loss solution. Additional trajectorics with various different seeds for both the Ackley and Goldstein- Price loss surfaces are shown in Appendix A.2.1 Figure 14 and Figure 13).

---

Published in Transactions on Machine Learning Research (06/2024)
Adam
Adam+DG
Adam
Train Loss
3
Adam+CG
Adam+CM
Adam+SAM
500
Epochs
Figure 4: Training loss curves (left, averaged across 10 seeds) and learning trajectories (right, one seed) for different optimizers on the Ackley loss surface. While the other optimizers get stuck in sub-optimal minima near the initialisation point (black square), both CM variants explore and find the lower loss surface near the global solution (black diamond).
Figure 5: Optimization trajectory of Adam (left), Adam+CG (middle), and Adam+CM (right) on toy ID function with flat and sharp minimum with increasing sharpness (across columus), for different initialisation points (across rows). Green backgrounds indicate that the optimizer escapes the sharper minimum while red backgrounds indicate otherwise. The vertical line indicates the final point in cach sub-figure. We observe that Adam mostly converges to the minimum closest the initial point. Adam+CM converges to the flatter minimum for different initial points and degrees of sharpness more often than ACam+CG.
Critical momenta reduce sharpness. We now want to compare specifically the implicit bias of the different optimizers towards flat regions of the loss landscape. We first examine the solutions of optimizers trained on the Goldstein-PPen and Levy functions (Laguna & Marti 2005) (see Appendix A.2.1p Both of these functions contain several local minima and one global minimum. We evaluate the solutions based on the final loss and sharpness, averaged across 20 seeds. As a simple proxy for sharpness, we compute the highest eigenvalue of the loss Hessian.
Results in Table show that Adam+CM finds flatter solutions with lower loss value compared to Adam, Adam+CG, and Adam+SAM in both examples. Furthermore, Adam and Adam+SAM reach almost equal loss values for the Levy function with negligible difference in sharpness, but for the GP function, Adam+SAM converges to a sub-optimal minimum with lower sharpness.
where we show that even with small buffer size, Adam+CM can escape sharper minima and explore lower loss regions than other optimizers. The results also suggest that a controlled setting, the larger buffer size helps find flatter minimum. To further investigate the escaping abilities of the various optimizers, we consider the
We analyze how the buffer size controls the amount of exploration empirically in Appendix A.2.1 [p.peddi(A.22111

---

Published in Transactions on Machine Learning Research (06/2024)
Ad C-20
Adam+GB C-20
AdM
0.7
Adam
Optimizers
Loss
Sharpness
0.8
Adam
0.86
1.49
Aadam+-SAl
3.14
1.43
GP
Adam+CCC CG
0.85
1.51
0.3
Adam CMM
0.81
1.36
Adam
13.87
65.65
AdAm+SSAM
13.87
65.62
Levy
Adam+GG
13.61
64.45
Adam-CM
12.50
62.53
Sharpnees coefficient
100
Table 1: Loss vs sharpness of the solutions of different optimizers for loy loss surfaces. The buffer decay is set to 0.99 for these experiments. Adam+CM is able to find solutions that are both flatter and deeper (lower loss) than other optimizers in this setting.
Figure Escape ratio (number of times the opti- mizer escapes the sharp minimum to reach the global minimum out of 50 runs) in the 10-D toy example (equation 99,, for different values of the sharpness co- efficient. Adam CM shows a higher ability to escape sharp minima in this setting.
following class of functions on R.:
D fa(x) = Σ ((an d=1
( 11)))
(9)
where is sharpness cocfficient. Each function in this class has two global minima: flat minimum at the origin and sharper minimum at (I...1). 1).
Figure 5] shows optimization trajectories in the one-dimensional case for various values of the sharpness coefficient & E {5, 10, 100} (across columns) and initial point x E {-2, 2, (across rows). We can see that Adam mostly converges to the minimum closest to the initial point. Adam+ +CM converges to the flatter minimum for different initial points and degreos of sharpness more often than Adam+CG. Additional plots are shown in Appendix A.3 for various values of the hyperparameters.
In Figure 7 (left), we compare the escape ratio for different optimizers on Equation 9 with d (similar to Figure 6). We can see that for C = 10, with an exception at = 2, Adam+CMM consistently finds the minimum at æ = more often than other optimizers. Interestingly, as 8 increases, Adam++GG is outperformed by both Adam and Adam+SAM, which indicates that Adam+CG is more susceptible to getting stuck at sharper minima in this example. Figure 7 (right) shows that the tendency of Adam+CM to escape minima is dependent on such that, larger value of results in convergence flatter minimum more often.
A Analyzing with D 10, we uniformly sample 50 unique initial points in -,4,5 Of these runs, count. the number of times optimizer finds the flat minimum at the origin by escaping the sharper minimun. Figure 6 reports the escape ratio for different values of the sharpness coefficient. We observe that Adam+CM (with buffer capacity C' = 20) has a higher escape ratio than others as the sharpness increases. We replicate this experiment with various values of the buffer capacity (see Figure 7
5 Experimental results
Results from the previous section show that our proposed approach exploration adaptive optimizers (Adam+CMP less sensitive the sharpness of the loss surface, less prone gradient cancellation, and finds flatter solutions on variety of loss surfaces. The goal of this section is to evaluate our method empirically complex models and with the following benchmarks:
8

---

Published in Transactions on Machine Learning Research (06/2024)
1.0
12M topC-10
1.0
Adum+CM tep0-5
top0-10 Adam- SAM
Adam+CM top
Adam+CM lopo=20
Adom
Adam+C tapCn40
Excape rello
02
0.2
100
100
Sharpness coefficient
Figure 7: (Left) Escape ratio (number of times when optimizer reaches the minima at x 0 out of total 50 runs) for different sharpness coefficient (s) for minima at x = 1 in 1D toy example shown in Figure 5) (Right) Escape ratio vs sharpness coefficient (s) for different C.
The Penn Treebank (PTB) (Marcus al. 1993]: It part-of-speech (POS) tagging task where a model must determine what part of speech (ex. verb, subject, direct object, etc.) every element of a given natural language sentence consists of.
CIFAR10 (Krizhevsky et al., 2009): is multiclass image classification task the training and validation set consists of images that separated into 10 distinct classes, with each class containing an equal number of samples across the training and validation sets. The goal is to predict the correct image class, provided as an annotated label, given a sample from the sets.
CIFAR100: It is the same task as CIFARIO, except that images are now separated into 100 classes with equal number of samples within each class.
ImageNet (Deng et al, )009) It a large-scale dataset consisting of images separated into 1000 distinct classes, The objective of the task is the same as CIFAR10 and CIFAR100, which to classify images from the training and validation set correctly into class with annotated label.
All results presented in this section are averaged across five seeds.
5.1 Language modelling
Starting with a langnage-based task, a single-layer long short-term memory network (LSTM) (Hochreiter & Schmidhuber, 1997) is trained on the PTB dataset. We evaluate the performance by reporting the validation perplexity on held-out set. We train the model for 50 epochs (similar to McRae et al. (2022)) and we reduce the learning at the 25th epoch by dividing by 10. The results are reported after performing grid search over corresponding shows that Adam+ CM outperforms other optimizers in terms of best validation perplexity achieved during training. We also note that AdA is the socond-best optimizer and achieves better performance than its CM variant.
parameerassssss. Details of this grid search present in Appendix TTall [Table 2
5.2 Image classification
Next, we evaluate Ad on different model sizes for image classification. We train ResNet34 models Hc et al. 2016) on 11FAR10/100 datasets. We train the models for 100 epochs where we reduce the learning at. the 50th epoch by dividing it by 10. We also train an EfficientNet-B0 model Tan & Le 2019) from scratch on ImageNet. We used a publicly available EfficientNet implementatior in PyTorch (aaske et al. 2019).

---

Published in Transactions on Machine Learning Research (06/2024)
Language modelling Validation Perplexity (4)
Image classification
Validation Accuracy % (t)
Optimizers
PTB
CIFARI0
CIFAR100
ImageNet
Adam
179.4
93.9
70.7
67.8
Adam+CG
174.445.5
93.840.4
71.0+..3
69.7
ASA
168.741.9
93.7
70.5
65.440.2
Adam+CM (ours)
163.2
94.040.3
71.2+0.3
71.740.2
Adam+SAM+CM (ours)
176.2
94.410.4
69.7.0.3
71.3
Table 2: Comparison of performance in terms of best validation perplexity and accuracy (%) achieved by the existing baselines with Cd and its SAM variant on language modelling and image classification tasks. Overall, CM outperforms the baselines in all four datascts.
Adam
10a
AdM
Adam-SAM
Distonce
20-1
Adamm
Adam+OG
AdM
Adam+SAM-CH
Epochs
Epochs
Adam+ CG
Adam- + CM
Cosine cimitarty
0.20
Adam
Adam+ CM
0.15
10
0.10
0.05
60
Epochs
Figure 8: Sharpness (top-left), distance (top-right) buffer variance (bottom-left), and cosine similarity (bottom-right) in buffer elements of the optimizers on CIFAR100. These results indicate that buffer elements in Adam+CM agree more with each other and have lower sharpness than Adam+CG.
with a weight decay (Loshchilov & Hutter 2019) of 10-4 and a learning rate scheduler where the initial learning rate reduced by a factor of 10 every 30 epochs. We provide additional details about the grid search, datasets, and models in Appendix A.2.2
The best validation accuracy achieved by models for image classification tasks is reported in Table We again observe that Adam+CM achieves better performance than the other optimizer baselines. Moreover, Adam+SAM+CH yielded the best validation accuracy for CIFAR10, while Adam+CM performed the best on the other two datasets.
Figure 8 corroborates the claim that Adam+CM finds a flatter surface containing the global minimum, with the top-left plot showing lower sharpness compared to Adam or Adam+SAM. The top-right plot further reveals the greater distance travelled by parameters during training, indicating that using CM promotes more exploration than the other optimizers. We also note that the distance travelled by parameters in SAM is similar to that of Adam. This similarity is also observed with Adam+ and Adam+SAM+CN as their
10

---

Published in Transactions on Machine Learning Research (06/2024)
64
128
Validation accuracy
90
512
256
Validation accuracy
90
80
80
70
70
Resnet34 CIFAR10
Resnet34 CIFAR100
Resnet34 CIFAR10
Resnet34 CIFAR100
Figure 10: Best validation accuracy for different. batch sizes (left) and buffer size (right) using Adam+ CMM on CIFAR10/100, with hyperparameters grid search. The default batch size of 64 and buffer size works best for these benchmarks.
distance curves overlap. Figure (bottom-left) shows that buffer elements stored by Adam+CM have lower variance during training compared to Adam+CG. To compare agreement among buffer quantáties, take the clement with the highest norm within the buffer, compute the cosine similarities with other elements in the buffer, and take the mean of these sinilarities.
Figure (bottom-right) shows the agreement in Adam+-MM remains higher than in Adam+CG, indicating the aggregation of buffer cle- ments Adam+ CM will more oflen result a non-zero quantily the desired direction. On the other hand, high variance and disagree- ment among elements in the Adam+ CG buffer may cause gradient cancellation during aggregation and result in Adam-like behavior. These results are associated with the best-performing set for the given optimizers. They highlight that it is the property of Adam+CM that facilitates efficient exploration and finding flatter surfaces to achieve the best performance.
10
C=20
Sharpness
Figure 9 shows the final sharpness metric for different buffer sizes recorded for CIFARI0/100 experiments with default hyperparameter setup. It is clear that using large buffer size further reduce the sharpness of the solution in such complex settings.
Resnet34 CIFAR10
Resnet34 CIFAR100
Figure 9: Sharpness for different buffer sizes using Adam+CM CI- FAR10/100, with other hyperparam- eters fixed. Using larger buffers re- sults in lower sharpness even for high- dimensional models.
Next, we show how batch size influences the performance of Adam+CM for CIFAR1O/IH experiments. In both datasets, we found that the default batch size of 64 performs best in terms of validation accuracy. We provide the validation accuracy achieved by the best-performing setups for other batch sizes in Figure 10 (left). We also perform similar experiment but for different buffer sizes CIFAR10/100 experiments. We found that the default buffer size of performs best in terms of validation accuracy. We provide the best-performing sctups in Figure 10 (right) and also observe that smaller buffer size results have a slightly higher standard deviation.
5.3 Using non-adaptive optimizer
Next, we show the advantage of integrating exploration capabilities in SGD. To implement SGD+CG, we follow McRae al. (2022) and replace ge with aggregation of g, and the buffer gradients g, for parameter
11

---

Published in Transactions on Machine Learning Research (06/2024)
Optimizers
CIFAR10
CI
Accuracy
Speed-up
Accuracy
Speed-up
SGD
93.0=0.4
1x
69. 69.640.7
1x
SGD+CG
93.1 93.1 00.4
1.2x
69.610.3
1.2x
SGD+CM (ours)
93.1-0.1
1.3x
69.6+0.7
1.2x
Table 3: Comparison of performance in terms of best validation accuracy (%) and speed-up by SGD+CM CIFARIO and CIFAR100. Although the performance remains the same, there speed-up by both memory-augmentedd optimizers.
updates:
0+11 1) (9:
1 C
Σ gc)
(10)
g-Eg.
Similarly, for SGD+CM, we aggregate the momentum using equation and update the parameters using equation We perform a hyperparameter grid search on CIFARI0 1100 experiments to compare SGD, SGD+ CG with SGD+CM and report the results in [Tall 3 Apart from validation accuracy, we also report the speed-up in the first 50 epochs of the training process and observe that although generalization performance remains the same, there is speed-up in both SGD+CG and SGD+CM.
5.4 Online learning
We also evaluate the methods an online learning sctup where a model with a limited capacity is trained on a stream of lasks. Here, the goal adapt to the latest task dataset by maximizing its performance. Therefore, it is important for an optimizer escape task-specific solutions as the loss landscape changes once new task dataset arrives. We hypothesize that with its greater exploration capabilitios, Adam+ CM will be able to do better than other baselines. We evaluate the optimizers on the following benchmarks:
Tinylmagenet Zhang et al. 2019: : This dataset created by partitioning its 200 classes into 40
5-way classification tasks. The implementation of TinyImagenet is based on Gupta et al. (2020) where 4-layer CNN model trained.
5-daaasstt Mehht et al. 2023): It consists of five different 10-way image classification tasks: CIFAR1O, Le LeCun 1998)
MashioonMNISST Xiao et al. (2017). SVHN Netzer el al. 2001 and MNIST
Bulatov [0011) The implementation is based on Mehta et al. (2023) where ResNet18 He et al.
(2016) model is trained.
We report their learning accuracics, which is the average validation accuracy for each task when the model is trained that specific task. More details about the models, datasets, and hyper-parameter grid search are provided in Appendix A.2.3
we observe that Adam+CM achieves better performance than Adam. Morcover, while Adam+SAM performs similarly Adam, Adam+SAM-CM re- sults in the best accuracy on both bench- marks. One possible reason for this is that, while SAM primarily focuses on im- proving the flatness of a solution, CM is more effective in exploring the loss surface and finding a basin containing more gen- eralizable solution. Our findings validate the hypothesis that CM can be applied in
In Table 4
Optimizers
TinyImagenet
5-dataset
Adam
62.9
88.2
Adam+CG
62.8
88.240.5
Adam+SAM
62.8
88.3
Adam+CM
63.4
88.4.0.6
Adam+SAM+CM
63.940.7
88.540.5
Table 4: Performance comparison in terms of best learning accu- racy (%) of the existing basclines with Adam+ CM and its SAM variant on online learning benchmarks.
12

---

Published in Transactions on Machine Learning Research (06/2024)
online learning settings to promote exploration. Moreover, it complements existing baselines and can be seamlessly integrated for performance improvements.
6 Conclusion
This work presents a framework for enhancing exploration in adaptive optimizers. We introduce CM, novel memory-augmenteed variant of Adam that incorporates a buller of critical momenta and adapls the parameters update rule using an aggregation function. Our analysis demonstrates that Adam+CM effectively mitigates the limitations of existing menory-augmented adaptive optimizers, fostering exploration toward flatter regions of the loss landscape. Empirical evaluations showcase that Adam+CM outperforms Adam, SAM, and CG across standard supervised and online learning tasks. An exciting future direction is the application of reinforcement learning settings, where knowledge transfer without overfitting on a single task is crucial. Furthermore, our findings hint at the potential of CM to capture higher-order dynamics of the loss surface, warranting further investigation. Even though the quadratic assumption in our theoretical analysis can be (slightly) relaxed by assuming that the loss Hessian matrices computed along the optimization trajectory all commute with each other, analyzing broader classes of functions is worth exploring in the future.
Acknowledgments
This research was supported by Samsung Electronics Co., Ltd. through a Samsung/ Mila collaboration grant, and was enabled in part by compute resources provided by Mila, the Digital Research Alliance of Canada, and NVIDIA. Pranshu Malviya is partially supported by the Merit scholarship program for forcign students (PBEEE) by Fonds de Recherche du Québec Nature el lechnologies (FRQNT). Gonçalo Mordido was supported by an FRQNT postdoctoral scholarship (PBEEE) during part of this work. Simon Lacoste-Julien is a CIFAR Associate Fellow in the Learning Machincs & Brains program and supported by NSERC Discovery Grants. Sarath Chandar is supported by the Canada CIFAR AT Chairs program, the Canada Research Chair in Lifelong Machine Learning, and the NSERC Discovery Grant.
References
David H Ackley. The model. In A Connectionist Machine for Genetic Hillclimbing, pp. 29-70. Springer, 1987.
Yaroslav Bulatov. Notmnist datasct. Technical report, Google (Books/OCR), 2011. URL http:///ryyooolwwwbb html
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In International Conference on Learning Representations, 2017.
Xiangning Chen, Cho-Jui Hsich, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representations, 2022.
Dami Choi, Christopher Shallue, Zachary Nado, Jachoon Lee, Chris J Maddison, and George Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv: 991..05446, 2019.
Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. Xiv preprint. ar Xiv:2200 14484, 2022.
Aaron Defazio and Samy Jelassi. Adaptivity without compromise: momentumized, adaptive, dual averaged gradient method for stochastic optimization. Journal of Machine Learning Research, 2022.
Aaron Defazio, Francis Bach, and Simon Lacosto-Julien. SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances in Neural Information Processing Systems, 2014.
13

---

Published in Transactions on Machine Learning Research (06/2024)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In TEEE Conferencee Computer Vision and Pallern Recognation, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Timothy Dozat. Incorporating Nesterov momentum into Adam. In International Conference on Learning Representations workshop, 2016.
Jiawei Du, Zhou Daquan, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. In Advances in Newral Information Processing Systems, 2022a.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent. Tan. Efficient sharpness-awaie aware minimization for improved training of neural networks. In International Conference Learning Representations, 2022b.
John Duchi, Flad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 2011.
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Conference on Uncertainty Artificial Intelligence, 2017.
Pierre Forct, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpncss-awaro minimization for efficiently improving generalization. In International Conference on Learning Represenlations, 2021.
Diego Granziol, Xingchen Wan, Samuel Albanic, and Stephen Roberts. Iterative averaging in the quest for best Lest. error. Xiviv preprint arXiv:2003.012 2020.
Gunshi Gupta, Karmesh Yadav, and Lian Paull. Look-ahead meta learning for continual learning. Advances in Neural Information Processing Systems, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Pattern Recognition, 2016.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent, 2012.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in Neural Information Processing Systems, 1994.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 131735-7780 1997.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Conference on Uncertainty in Artificial Intelligence, 2018.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization mcasures and where to find them. In International Conference on Learning Representations, 2020.
Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and general- ization. In Proceedings of Machine Learning Research, volume 187, pp. 51 65. PMLR, 03 Dec 2023.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On largo-batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2016.
Junhyung Lyle Kim, Gauthier Gidel, Anastasios Kyrillidis, and Fabian Pedregosa. When momentum extragradient optimal? polynomial-based analysis. Xiv preprint 22v 11 1.4659, 2022a.
14

---

Published in Transactions on Machine Learning Research (06/2024)
Minyoung Kim, Da Li, Shell X. Hu, and Timothy Hospedales. Fisher SAM:
Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, 2022b.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Alex Krishevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpness-awen aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, 2021.
Manuel Laguna and Rafael Marti. Experimental testing of advanced scatter search designs for global optimization of multimodal functions. Journal of Global Oplimizalion, 33:235-220 255, 2005.
Quoe V Le, Jiquan Ngiam, Adam Coates, Abhik Lahiri, Bobby Prochnow, and Andrew Y Ng. On optimization methods for deep learning. In International Conference on Machine Learning, 2011.
Yann LeCun. The mnist database of handwritten digits, 1998. URL httphtttp/aaae lecun. cww. hoom/exbbmmist..
Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of optimization algorithms via
integral quadratic constraints. SIAM Journal Optimization, 26(1):57-955 2016. doi: 10.1137/15M1009599..
URL httph:t:///// :/ttps://doi.. org/10. 1137/15M100959
Ren Liu, Fengmiao Bian, and Xiaoqun Zhang. Binary quantized network training with sharpness-aware minimization. Journal of Scientifie Compuling, 2023.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness- aware minimization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Tya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Mitchell P. Marcus, Mary Ann Marcinkiewicz, and Beatrice Santorini. Building large annotated corpus of english: The penn treebank. Compul. Linguist., 19(2):313 330, jun 1993. ISSN 0891-2017.
Paul-Aymeric Martin McRac, Prasanna Parthasarathi, Mido Assran, and Sarath Chandar. Memory augmented optimizers for deep learning. In International Conference on Learning Representations, 2022.
Sanket Vaibhav Mehta, Darshan Patil, Sarath Chandar, and Emma Strubell. An empirical investigation of the role of pre-traiming in lifelong learning. Journal of Machine Learning Research, 24(214):1 2023. html
URL http://jmlr.
Gonçalo Mordido, Sarath Chandar, and François Leduc-Primeau. Sharpness-aware training for accurate inference on noisy DNN accelerators. arXiv preprinl arXiv:2211111661 2022.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning, 2011.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in Neural Information Processing Systems, 2017.
Jorge Nocedal. Updating quasi-Newton matrices with limited storage. Mathematics of computation, 1980.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 2019.
Victor Picheny, Tobias Wagner, and David Ginsbourger. A benchmark of kriging based infill criteria for noisy optimization. Structural and multidisciplinary oplimization, 48:607-615 -6607-626, 2013.
triging-based
15

---

Published in Transactions on Machine Learning Research (06/2024)
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathemalics and Mathemalical Physics, 1964.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, 1951.
Nicolas Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an exponential convergence rate for finite training sets. Advances in Neural Information Processing Systems, 2012.
Nieol N Schraudolph, Jin Yu, and Simon Günter. A stochastic quasi-newton method for online convex optimization. In Conference on Artificial Intelligence and Statistics, 2007.
Shai Shalev-Shwartz and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Research, 2013.
Mingxing Tan and Quoc Le. EfficientNet: Rethinking model scaling for convolutional neural networks. In International Conference on Machine Learning, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Aduances in Neural Information Processing Systems, 2017.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, 2017.
Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over pver-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems, 2018a.
Lei Wu, Chao Ma, et al. How SGD selects the global minima in over-parameteerieed learning: A dynamical stability perspective. Advonces in Neural Information Processing Systems, 31, 2018b.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint 10X.::0788.077777 2017.
Matthew D) Zeiler. AdaDelta: An adaptive learning rate method. arXiv preprint arXi: 2012.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. Advances in neural information processing systems, 2019.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why SGD generalizes better than Adam in deep learning. Advances in Neural Information Processing Systems, 2020.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. AdaBelief optimizer: Adapting stepsizes by the belief in observed gradients. In Advances in Neural Information Processing Systems, 2020.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training, In Interna- tional Conference on Learning Representations, 2022.
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in lcarning neural networks with proper regularization. In The Eleventh International Conference on Learning Representations, 2023.
16

---

Published in Transactions on Machine Learning Research (06/2024)
A
Appendix
In this section, we provide the details and results not present in the main content. In section A. more evidence of gradient cancellation in a loy example. In section A.2 we describe the implementation details including hyper-per hyper-parameters parameters values used our experiments. All experiments were executed on an
we report
NVIDIA A100 Tensor Core GPUs machine with 40 GB memory.
A.1 Gradient cancellation in CG
When we track the trajectory along with the gradient directions of the buffer in the Adam++G optimizer on the Ackley+ Rosenbrock function (defined in the next section), we found that CG gets stuck in the sharp minima (see Figure 11p. This is because of the gradient cancellation problem discussed earlier.
Figure 11: Three consccutive steps in training Adam+CG on A Rosenbrock function with C 5. The white lines indicate the gradient directions in the buffer. Since four of these gradients point in opposite directions, their mean would be small and cancel the current gradient update out. As result, parameters get stuck in the sharp minima.
Next, provide additional analyses similar to Figure We compare the trajectories and gradient directions of Adam+CM and AdG for six consecutive steps on the Goldstein-Prico function in Figure 12 Similar to our observations on the Ackley function, we find that the buffer quantities in CG are dominated by directions higher sharpness (towards the top-left direction and bottom-right direction). Even when the new gradient direction aligns with the buffer mean in CG, it remains stuck in the sharp surface. On the other hand, the buffer directions do not drastically change in CM and therefore it finds the surface that contains the global minima.
A.2 Implementation details and other results
A.2.1 Toy examples
We evaluate our optimizer on the following test functions in Section
1. Ackley function:
(11) The global minimum is present at (0,0). In Figure 13 we visualize the trajectories of Adam, Adam+SAM, Adam+CG, and Adam+CM for different initialisation points. While other optimizers may get stuck at nearby local minima, Adam + CM benefits from more exploration and finds a lower-loss surface that may contain the global minima.
f(x, v) 20 -00exp(-002200.00 + v2)) exp(O.5(cos 2wr cos (0s)) + e 20
2. Soldstein-Price function:
f(x. v) 2.427 log
(z + 119 14r + 3x 14y 6xy + 3y
(12)
x [30 + (2x 318 - 32r + 12n + 48y 360y + 27y²) 8.693]
17

---

Published in Transactions on Machine Learning Research (06/2024)
AdsRIcM AdaenCC cunntiis
AcamiCM
Acam SM
Sutter quartille
Buffer quarddes
Burforment men gradint
Gudler moon
Buffor moan
New gradiett
gradiant
Adamn DXt
AdamicM
AdamiCM
Adla-DD
Maniic
Bulkr
Seheer
sxarolic
Six consecutive steps of Adam+CG and Adam+-CM trajectories on Goldstein Price function function. Buffer quantities in Adam+CG are dominated by directions of higher sharpness (towards top-left direction and bottom-right direction) and the new gradient directions have high variance, yielding a small update. Conversely, Adam+CM escapes sub-optimal minima and converges near the global minimum.
Figure 12:
The global minimum present at [0, 1j2. In Figure 14, we visualize the trajectories of Adam, Adam+SAM, Adam+CG, and Adam+CCM CM for different initialisation points. While other optimizers may get stuck at sub-optimal loss surface, Adam+CM benefits from more exploration and finds the global minimum.
Levy function:
d-1 (w, - 1 10sin?(Tm) + 1)] + (wra - 1) sin sin? (2mmaa)
[(11 x2) = sin? (TT))
(13)
where w, 1 ببي and d is the number of variables. The global minimum is present (1, 1).
+ function:
f(x,y) 0.05(1 - (0 0.05 - 22) 2:::://==2):...fxxppppooo -
(14)
- cxp(O05(cos 2xx + +os cos + +
The global minimum is present at (2.1,3.9).
we also compared different optimizers in terms of the sharpness of the solution reached by each optimizer on different functions. Table 5 we perform a similar experiment and compare sharpness across different values of hyperparameters for Adam+CM We observe that Adam+CM able to find flatter and lower-loss solution as buffer size increases. This is consistent with complex tasks in Figure 9
In Figure ]
We also provide comparison with another existing method called Momentum Extragradient (MEG) (Kim ct al. 2022ah and Adam+CM on the Ackley and Goldstein-Price functions in Figure 15) We observe that Adam CM has better exploration capabilities than MEG. When evaluated on CIFARIO dataset, Adam+MEG achieves the accuracy of 93.8 % which is lower than both Adam+CM (94.040.3 %) and Adam+SAM+CM (94.4+0.4 %).
In Figure 4 we compared different variants of Adam on Ackley function by running 500 steps. We observed that, unlike other baselines, Adam+CM and Adam+SAM+ON escape local minima, explore the loss surface,
18

---

Published in Transactions on Machine Learning Research (06/2024)
Optimizers
Loss
Sharpness
CG
5
0.7
14.40
64.72
Cd
5
0.99
14.41
64.67
Ad
20
0.99
13.61
64.45
Adam+CM
5
0.7
12.90
63.74
Adam+CMM
5
0.99
13.02
63.83
Adam+CM
20
0.99
12.50
62.53
Table 5: Loss vs sharpness the solutions of Adam+CM for different hyperparameters buffer sizes on Levy function. Adam+CM is able to find solutions that both flatter and deeper (lower loss) bufler size increases. Morcover, Adam+ CM always outperforms Adam+CG even with smaller buffer size.
Adam+CM
Adam+OM
Adam+CG
Adam+CG
AdMM+SAM+OM
Adam+SAM-CM
Adam+SAM
Adam+SAM
Adam+
Adam+CM
Adam+CG
Adam+CG
Adam
Adam
Adam+ SAM+CM
Adam+SAM-CM
Adam+SAM
Adam*SAM
Figure 13: Optimization trajectories of various optimizers on the Ackley loss surface for different initial points (black square). The black diamond indicates the global minimum.
and navigate towards the region containing global minima. In this experiment, reducing the learning rate helps settling down to the global minimum and stopping further exploration. As an ablation study, we perform the same experiment but without learning rate decay. We plot the loss curve in Figure 16 and observe that Adam+CM oscillates near the global minima.
19

---

we provide a summary of all datasets and deep learning models used in the experiment Published in Transactions on Machine Learning Research (06/2024)
Adam+OG
Adam+OG
Adam
Adam
Adam+SAM
Adam+SAM
Adam+CM
Adam+CM
Adam+CG
Adam+CG
Adam
Adam
Adam+ SAM
Adam SAM
Figure 14: Optimization trajectories of various optimizers the Goldstein- Colice loss surface for different initial points (black square). The black diamond indicates the global minimum..
AdamHCM
Adam
Adam+BAM
Adam-MEG
Adam+OM
Adam
Adami
Figure 15: Comparing adaptation trajectorics of Adam+MEG with Adam, Adam+ and Adam+ Overall, we observe that Adam+CM exhibits better exploration the loss landscape.
A.2.2 Deep learning experiments
In Table 6 and Table
from Section
20

---

Published in Transactions on Machine Learning Research (06/2024)
3
Adam
Adam+CG
Adam+SAM
SAM-CM
Epochs
Figure 16: Loss (averaged across 10 seeds) obtained and learning trajectories for different optimizers with constant learning rate on the Ackley loss surface. We observe that while other optimizers get stuck at sub-optimal minima, Adam+CM and AdMM+SAM+CM achieve lower loss values and oscillates near the global minima.
Dataset
Train set
Validation set
PTB
890K
70K
CIFARIO
40K
10K
CIFAR100
40K
10K
ImageNet
1281K
50K
Table Dataset details
Model
Number of parameters
LSTM
20K
ResNet34
22M
EfficientNeetB0
5.3M
Table Model details
For all toy examples experiments based on Equation 9 the lcarning rate is set as 0.05. The learning rate is set as 0.1 for other loy examples. Unless specified in the experiment description, the default set of other hyperparameters in all our experiments is {B1, B2, C, A, p} {0.9, 0.99, 5, 0.7, 0.05] except in CIFAR10/ 100 experiments where 32 is set to 0.999. The default values of and are decided based on the suggested values from McRae al. (2022) and p basod on Foret et al. (2021) For the results in Figure and Figure C and
À are set. Lo 20 and 0.99.
For supervised learning experiments, provide the details pareame grid-search in Table 8 and the best settings for all experiments and Table 9 and Table 10 In these tables, we report the hyperparameter set for each optimizer as follows:
Adam: {a}
Adam+CG: (a, B1, B2, C2
Adam+SAM: {a, 31, 32
Adam+CM: {a, 3, 32, C,
Adam+SAM+CM: (2, C.
21

---

Published in Transactions on Machine Learning Research (06/2024)
CM essentially uses the equivalent cost of computation and space complexity CG. However, after performing hyperparameter grid search over the buffer size, observed that CM consistently requires smaller buffer size than CG and obtains better performance with (C 5, A = 0.7).
Hyper-pane parameter
Set
lr
{0.1, 0.01, 0.001, 0.0001} {0.9, 0.99, 0.999) (0.99, 0.999, 0.9999) (5. 20) {0.7, 0.99}
B
B2
C
X
P
{0.01, 0.05, 0.1}
Table Details on grid search on hyper-parameter setting.
Optimizers
PTB LSTM
Adam
{0.001, 0.9, 0.99]
Adam+CGG
{0.001, 0.9, 0.999, 5, 0.7)
Adam+SAM
{0.001, 0.9, 0.9, 0.01]
Adam+CCM
(0.001, 0.9, 0.999, 5, 0.7)
AdMm+SAM+CM
{0.001, 0.9, 0.999, 5, 0.7, 0.1)
Table 9: Best hyperparameter settings for different optimizers on PTB.
Optimizers
CIFAR10
CIFAR100
ImageNett
Adam
[0.001, 0.9, 0.999)
(0.001, 0.9, 0.99)
{0.0001, 0.9, 0.99)
Adam+-GG
{0.0001, 0.9, 0.999, 20, 0.7}
(0.001, 0.9, 0.99, 20, 0.7)
(0.0001, 0.9, 0.99, 5, 0.7)
Adam+SAM
{0.0001, 0.9, 0.99, 0.05)
{0.001, 0.9, 0.999, 0.05]
{0.0001, 0.9, 0.99, 0.05]
Adan+CMM
(0.0001, 0.9, 0.999, 0.7]
(0.001, 0.9, 0.9999, 5, 0.7)
{0.0001, 0.9, 0.99, 5, 0.7)
Adam+SAM-CM
(0.0001, 0.9, 0.99, 5, 0.7, 0.05)
{0.001, 0.9, 0.99, 5, 0.7, 0.1]
{0.0001, 0.9, 0.99, 5, 0.7, 0.05)
Table 10: Best hyperparameter settings for different optimizers on image classification benchmarks.
A.2.3 Online learning setup
The online learning experiments performed the following benchmarks:
Tinylmagenet:
This dataset is created by partitioning its 200 classes into 40 5-way classification tasks. The implementation of Tinylmagenet is based on Gupta et al. (2020 where a 4-layer CNN model is trained.
5-dataset: It consists of five different 10-way image classification tasks: CIFAR10, MNIST
LeCun
(1998), Fashion-MNIST Xiao et al. (2017). SVHN Netzer al. (2011), and notMNIST Bulatov 2011
The implementation is based on Mchta et al. (2023) where ResNet18 He et al. (2016) model is
trained.
For online learning experiments, we provide the details on hyper-parameter grid-search best settings for all experiments in Table 13
Table 12 and the
A.3 Sensitivity analysis
Following experiments in Figure 5 we fix decay 0.7 in Figure and vary C. We perform a similar
experiment with decay= 0.99 and plot them in Figure 18 In both these figures, the observation remains
22

---

Published in Transactions on Machine Learning Research (06/2024)
Optimizers
GP
Ackley
Adam
(0.1, 0.9, 0.99)
{0.1, 0.9, 0.99) {0.1, 0.9, 0.99, 5, 0.7}
Adam+CG
{0.1, 0.9, 0.99, 5, 0.7}
Adam+SSMM
{0.1, 0.9, 0.99, 0.01)
(0.1, 0.9, 0.99, 0.01]
Adam+CM
{0.1, 0.9, 0.999, 20, 0.99}
{0.1, 0.9, 0.999, 20, 0.99}
Adam-SAAe CAMM+M
{0.1, 0.9, 0.999, 20, 0.99, 0.05)
{0.1, 0.9, 0.999, 20, 0.99, 0.05)
Table 11: Best hyperparameter settings for different optimizers on toy examples.
Slyper-paramtten
Values
step size (n)
{0.01, 0.001, 0.0001, 0.00001} {0.9, 0.99) {0.99, 0.999, 0.9999}
81
B2
Table 12: Details the hyper-parameter grid search used for the online learning experiments.
Optimizer
Tiny Imagentt
5-dataset
Adam
{0.0001, 0.9, 0.999)
{0.0001, 0.9, 0.999)
Adam+CG
{0.0001, 0.9, 0.999}
{0.0001, 0.9, 0.999)
Adam+SAM
(0.0001, 0.9, 0.9999)
{0.001, 0.9, 0.999]
Adan+CM
(0.0001, 0.9, 0.9999)
{0.001, 0.9, 0.9999}
SAM
{0.0001, 0.9, 0.99}
(0.00001, 0.9, 0.999}
Table 13: Best hyper- prper-parameter settings for online learning experiments.
the same that is Adam+CM converges to flatter minima for different initial points and degrees of sharpness. We also observe that C plays an important role in convergence to flatter minima in both Adam+CC and Adam+CMM
A.3.1 m-sharppess
There have been various definitions of sharpness used in optimizatiin literature. Specifically, Foret et al. (021) employs m-sharpness to indicate the sharpness of the loss landscape and demonstrates its correlation with generalization performance.
The m-sharpness is defined as:
1 Σ MED
max Σ - JellaSr
(15)
SEM
where D represents the training dataset, which is composed of mini-batches M of size and is hyper-parameter set to 0.05 by default.
In Figure 19 we also monitor the m-sharpness during the training of CIFAR1O (left) and CIFAR100 (right)
with the same learning rate of 0.0001 and fixed hyperparameter setups. We compare the baseline optimizers with Adam+ and observe that Adam CM exhibits lower m-sharpness on both datasets, which consistent with our obscrvations Figure
A.3.2 Learning rate and sharpness
In this section, we compare generalization performances and empirical sharpness of the solutions obtained using the Adam optimizer with different learning rates on CIFAR10. We keep the other hyper-parameters fixed their default values for fair comparison. We show that even when increasing the learning rate decreases the overall sharpness (Figure 20), the resulting minima is sub-optimal (Figure 14).
23

---

Published in Transactions on Machine Learning Research (06/2024)
5-
5-
-35
3.0
-35
3/3
4-
6-
-355
3.6
-3.5
06
-3.5
3.6
2.5
3.5
3/5
-3.5
2-
13
0-
0-
2.5
0.0
2.5
2.5
0.
23
355 3.0 2.5
2.5 0.0
2.5
255
0.
3.5
3.5 0.0
25
Figure 17: Following Figure 5 compare trajectorics of Adam+CC (left columns) and Adam- CM (right column) on Equation 9 with d = 1 where À is set to 0.7 and C: (i) 5 (first row), (ii) 10 (second row) and (iii) 20 (third row). For different initial points and degrees of sharpness, Adam+Ch converges to flatter minima more often than Adam+CG.
24

---

Published in Transactions on Machine Learning Research (06/2024)
5-
5-
-35
3.0
-35
3/3
4-
6-
-355
3.6
-3.5
06
-3.5
3.6
2.5
3.5
3/5
-3.5
2-
13
0-
0-
2.5
0.0
2.5
2.5
0.
23
355 3.0
2.5
2.5 0.0
255
255
0.6
2.5
3.5 0.0
25
Figure 18: Following Figure 5 compare trajectorics of Adam+CC (left columns) and Adam- CM (right column) on Equation 9 with d = where À is set to 0.99 and C: (i) (first row), (ii) 10 (second row) and (iii) 20 (third row). For different initial points and degrees of sharpness, Adam+Ch converges to flatter minima more often than Adam+CG.
25

---

Published in Transactions on Machine Learning Research (06/2024)
10
Adam
Adam+CG
Adam-CM
m-Sharpness
m-Sharpness
10-3
Adam
Adam+CG
Adam-SAM
40
80
100
20
80
100
Epochs
Epochs
Figure 19: m-sharpness computed upon training ResNet34 CIFAR10 (left) and CIFAR100 (right) with different optimizers on fixed hyper-parameter setup. In both cases, Adam+CM has lower m-sharpness than other baselines.
I-0.1
k-0.0001
10
Learning rate Validation Accuracy
Sharpness
0.01
93.5+0.3
0.001
70.2
10
0.0001
93.4
10
0
20
40
80
80
Epochs
Table 14: Validation Accuracy of Adam with different learning rates CIFAR10. The best result was obtained using learning rate of 0.001.
Figure 20: Sharpness obtained using Adam with differ- ent learning rates on CIFARIO. These results indicate that higher learning rate results in lower sharpness.
26


Reviews:
[{"review_id": "EC0XLXLdhO", "reviewer": "Reviewer_q45D", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "$\\textbf{Strength}$\n\n$\\textbf{S.1}$. The proposed critical momenta is a direct enhancement on the critical gradient version of the memory-augmented Adam. The motivation is strong with very intuitive explanation of the issues, i.e., gradient cancellation, caused by critical gradient. The approach of resolving the issue and hence improve the performance is very well presented and novel.\n\n$\\textbf{S.2}$. There are a considerable amount of very intuitive and insightful explanation on how to use memory buffer to better leverage the curvature, i.e., second-order information, what kind of issues existed in pre-existing literature of memory-augmented optimizers, and how to address the issues. It does an excellent job on educating the audience and presenting the novel idea of using critical momenta.\n\n$\\textbf{S.3}$. The brief technical analysis is very insightful and convincing in demonstrating the proposed approach given in convex context. Meanwhile, there are a lot of insightful and intuitive empirical analysis to further explain and present the idea.\n\n$\\textbf{S.4}$. The main empirical study is very extensive. While the technical insight mainly focuses on convex case, the main experiment focuses mainly on up-to-date relevant nonconvex problems in supervised learning, such as token classification in language modeling and image classification, as well as online learning setting.\n\n\n\n$\\textbf{Weakness}$:\n\n$\\textbf{W.1}$. Compared to the empirical analysis, the technical analysis, especially regarding the convergence analysis, is rather weak. There lacks a formal convergence analysis in the proposed strongly convex case, and whether a convergence can be achieved or not in a (specific or general) nonconvex setting is unknown. \n\n$\\textbf{W.2}$. While the technical analysis reveals how the curvature is better captured through CM and how the 'flatter sharpness' is realized. It is lacking a formal analysis, on defining i) the relevant metrics, such as 'how flat is flatter?' and 'what does flatness mean' ii) how the curvature is better captured or captured in a certain way, resulting an improved behavior over CG version.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a variant of memory-augmented optimizer (McRae et al. 2022), with replacement of gradient by the momentum and applied specifically in Adam-class algorithms. The main idea is to use a memory buffer to store the past history of momentums, aka `critical momenta`, which alleviate the issues of `gradient cancellation` induced by `critical gradient` proposed in McRae et al. (2022). This allows the algorithm to escape sharp local curvature and seek for flatter neighborhood of minima.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "$\\textit{Note}$: this section covers both questions and requested changes.\n\n$\\textbf{Q.1}$. In McRae et al. (2022)., author explicitly explains that using memory buffer with past gradient captures the curvature. Is there anyway besides the specific/simple example showing `section 3.2` to generalize the idea? or if author had different ideas of how the CM helps rather than better capturing the curvature, would it be possible to explain it further with generalized case?\n\n$\\textbf{Q.2}$. Compared to McRae et al. (2022)., Can I assume the equivalent cost of computation and space complexity? Is it possible to add complexity analysis to showcase the merit of the proposed CM method?\n\n$\\textbf{Q.3}$. In `section 5.3`, authors briefly mention applying CM with SGD. Would it be possible (or appropriate) to address the different memory-augmented SGD variants in more details? It seems that the current version of paper has no mention of these elsewhere.\n\n$\\textbf{Q.4}$. Given the emphasis of 'sharpness' of this paper, would it be better to give more explanation and/or define the sharpness more formally, so to benefit the general audience? Also, for 'm-sharpness' seen in Appendix, it is clear that the definition/idea was from pre-existing literature, and could be very beneficial to add more context in the current version of the paper.", "broader_impact_concerns": "I do not see a `broader impact statement` being enclosed in the current version of the paper. However, as I see that the paper mainly concerns a very specific improvement in enhancing the performance of a very well known optimizer, i.e., Adam, and thus I have no concerns of any negative impact.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "b3N7iuGcud", "reviewer": "Reviewer_kzdk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenghts:\n- The idea of CM is simple yet effective by applying same techniques as CG into the first-order momentum term of Adam instead of the gradient.\n- The paper points out the drawback of gradient cancellation when using Adam with CG and propose CM to address that.\n- The authors do a good job of illustrating how Adam with CM can lead to better stationary points compared to existing variants by looking at several toy examples in Section 4.\n- Numerical experiments are very extensive, there are many ablation studies on effect of key parameters such as buffer size vs local sharpness.\n\nWeakness:\n- The convergence analysis only considers the quadratic losses which is somewhat simple. I hope to see analyses for broader classes of function.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new algorithm within the memory-augmented framework applied to Adam called Adam with Critical Momenta (Adam + CM). The effect of CM is shown in several toy examples showing that Adam with CM is able to find stationary points with flatter surface than existing variants of Adam. Numerical experiments are provided to illustrate the advantage of CM in different benchmarks including language modeling, image classification, and online learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I think the equation references in Algorithm 1 can be improved. E.g.: Update 1nd moment should point to equation 4. Update parameter $\\theta_t$ should point to equation 7?\n- I think it is worth to explicitly define $h$ (probably eigenvalues of $H$) in section 3.2\n- I suggest having the same order and color for 2 plots in Figure 4.\n- For completeness, I would explicitly describe how CG and CM are integrated into SGD. Even though it can be infered from how CG is applied to Adam.\n- There are few cases where combing CM with SAM does not perform as well, or maybe worse than their individual technique. I wonder if there is any intuition on that.\n- In several experiments the variant is a bit high which includes the interval of other variant (despite higher mean) so we can't really conclude Adam+CM is indeed better. E.g: CIFAR10, CIFAR100 in Table 2, 5-dataset in Table 4.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Wb3GfFFzT2", "reviewer": "Reviewer_EuAX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1) The paper presents an interesting algorithm of improving generalization by escaping sharp minimas. Adam is well known to generalize worse than SGD, and improving it's generalization is critical. \n2) The experiments on toy examples and large datasets show the applicability of this method well. \n\nWeakness / Questions:\n\n1) \"Gradient cancellation primarily\noccurs when existing buffer gradients are quickly replaced by high-magnitude gradients when the parameters\nare near a sharp basin. \" this argument does not seem to be generic. There is no guarantee that at sharp minima, the current gradient will oppose the existing buffer accumulation. It depends on the direction of divergence and may vary for different loss curves. I could counter-argue that the current gradient will be in the same direction as the existing buffer accumulation facilitating escape from sharp minimas. Curves in Figure-2 seem very specific to the Ackley loss surface. \n\n2) I wonder methods like momentum and ADAM already take in the trajectory history into account. So, why would there be a need for additional exploration. Methods like SAM takes a flatter direction step from the current iterate and that is understandable. Infact there are works that explore extragradient exploration in top of momentum (https://openreview.net/forum?id=ZLVbQEu4Ab). Why shouldn't these methods be enough? Furthermore, just averaging weights in the trajectory are well known to improve generaliaztion (https://arxiv.org/abs/1803.05407). Having just model averging on top of Adam should work well. \n\n3) Figure-4  shows signs that learning rate decay was used at around 250 epoch for Adam+CM, Adam+SAM+CM, Adam+CG, which is known to decrease the training loss abruptly. Infact, before 250 epoch, just Adam+SAM seem to have a lower loss. Can the authors confirm this? And if so, a fair comparison should be done without learning rate decay. \n\nQuestions:\n\n1) Will increasing the learning rate in Adam help to escape sharp minimas. For plain Adam (without an explorations) with step size η and\nβ1 = 0.9, this stability threshold is 38/η  (https://arxiv.org/pdf/2207.14484.pdf).  Can lr be increased enough to match the sharpness levels in Figure-8?\n\n2) How does the generalization of simple SGD compare with the generalization capabilities of exploration based Adam methods? I did not find any comparison with this.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Thes paper provides an improvement over adam+cg and proposes adam+cm (critical momenta). The main innovation is that adam+cg, which is an exploratory step built on top of adam to escape sharp minimas, can have the drawback that when it encounters a large gradient in a sharp minima, it cancels with existent gradient buffers and provides no exploration. To improve over this, the authors apply the buffering on the momenta term instead of the gradient. Through analyzing trajectories on non-convex toy examples, they show that adam+cm can escape sharp minimas whereas other variants of adam + exploration strategies fail to do so. Furthermore, this method has a faster convergence rate on quadratic losses.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like the authors to answer the 3 weakness points raised. Further answering the questions 1 and 2 are critical.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: ### Comment on the recommendation
The reviewers agree that the work is well-motivated, the paper is clearly written, and the claims are supported by high-quality empirical results. (Please also see **Claims And Evidence**.)

The authors addressed the major concerns raised by the reviewers with the replies and the revised manuscript.

The contributions are highly relevant to the TMLR community, presented in a good quality. Hence, I recommend acceptance.

### Requests to the author for the final version
* Please perform minor revisions based on the revised version uploaded on Apr 25.
* I suggest that the authors briefly mention the limitation of the theory part, e.g., in Section 6 as this was the biggest concern of the reviewers.
* Please consider improving Figure 2. The contours are nearly invisible to me.
* I think Figure 1 should be for Table 1 in the sentence "Results in Figure 1 show [...]" on p.7.
* Please add punctuation marks after equations.
* Typo: worse-case --> worst-case
Decision: Accept (poster)