Title: TRAINING DEEP NETWORKS WITH STOCHASTIC GRA-DIENT NORMALIZED BY LAYERWISE ADAPTIVE SEC-OND MOMENTS

Abstract: We propose NovoGrad, an adaptive stochastic gradient descent method with layerwise gradient normalization and decoupled weight decay. In our experiments on neural networks for image classification, speech recognition, machine translation, and language modeling, it performs on par or better than well-tuned SGD with momentum, Adam, and AdamW. Additionally, NovoGrad (1) is robust to the choice of learning rate and weight initialization, (2) works well in a large batch setting, and (3) has half the memory footprint of Adam.

1. INTRODUCTION:

The most popular algorithms for training Neural Networks (NNs) are Stochastic Gradient Descent (SGD) with momentum (Polyak, 1964; Sutskever et al., 2013) and Adam (Kingma & Ba, 2015) . SGD with momentum is the preferred algorithm for computer vision, while Adam is the most commonly used for natural language processing (NLP) and speech problems. Compared to SGD, Adam is perceived as safer and more robust to weight initialization and learning rate. However, Adam has certain drawbacks. First, as noted in the original paper (Kingma & Ba, 2015) , the second moment can vanish or explode, especially during the initial phase of training. To alleviate this problem, a learning rate (LR) warmup (Goyal et al., 2017) is typically used. Adam often leads to solutions that generalize worse than SGD (Wilson et al., 2017) , and to improve Adam regularization, Loshchilov & Hutter (2019) proposed AdamW with decoupled weight decay.

Our motivation for this work was to find an algorithm which: (1) performs equally well for image classification, speech recognition, machine translation, and language modeling, (2) is robust to learning rate and weight initialization, (3) has strong regularization properties. We start with Adam, and then (1) replace the element-wise second moment with the layer-wise moment, (2) compute the first moment using gradients normalized by layer-wise second moment, (3) and decouple weight decay (similar to AdamW) from normalized gradients. The resulting algorithm, NovoGrad, combines SGD's and Adam's strengths. We applied NovoGrad to a variety of large scale problems -image classification, neural machine translation, language modeling, and speech recognition -and found that in all cases, it performs as well or better than Adam/AdamW and SGD with momentum.

2. RELATED WORK:

NovoGrad belongs to the family of Stochastic Normalized Gradient Descent (SNGD) optimizers (Hazan et al., 2015; Nesterov, 1984) . SNGD uses only the direction of the stochastic gradient to update the weights, and the step size does not depend on the magnitude of that gradient. Hazan et al. (2015) proved that the direction of the gradient was sufficient for convergence. Ignoring the gradient magnitude makes SNGD robust to vanishing and exploding gradients.

SNGD with layer-wise gradient normalization was introduced by Singh et al. (2015) . The method scales up small gradients, while keeping large gradients unchanged:

ĝl t = g l t • (1 + log(1 + 1 ||g l t || ))

where g l t is the gradient for the layer l at step t. A similar approach was proposed by Yu et al. ( 2018), who used layer-wise gradient normalization for SGD and Adam to alleviate vanishing and exploding gradients. They divide the stochastic gradient g l t by its norm ||g l t ||:

ĝl t = g l t ||g l t ||

NovoGrad is also closely related to the Normalized Direction-preserving Adam (ND-Adam), proposed by Zhang et al. (2017) . For each layer, ND-Adam first removes the projection of gradient g l t on the layer's weights w l t :

ḡl t = g l t -(g l t , w l t ) • w l t

Then, ḡl t is used to compute Adam 1 st and 2 nd (scalar) moments:

m l t = β 1 • m l t-1 + (1 -β 1 ) • ḡl t v l t = β 2 • v l t-1 + (1 -β 2 ) • ||ḡ l t || 2

Similar to Adam, the weights are updated with the 1 st moment re-scaled by the 2 nd moment:

wl t+1 = w l t -λ t • m l t v l t +

Adaptive methods like Adam generalize worse than SGD with momentum as was shown in Wilson et al. (2017) . For example, Keskar & Socher (2017) proposed to use Adam during the initial stage only and then switch to SGD. Luo et al. (2019) suggested to improve Adam regularization by limiting the factor1 √ vt to a certain range: limiting from above helps to decrease the training loss while limiting from below helps to generalize better. Loshchilov & Hutter (2019) showed that Adam's weak regularization is due to the fact that the 2 nd moment normalization effectively turns off L2-regularization. They proposed AdamW, which decouples the weight decay d • w t from the gradient and uses it directly in the weight update:

w t+1 = w t -λ t • ( m t √ v t + + d • w t )

Adam needs to store the 2 nd moment, and this doubles the optimizer memory compared to SGD with momentum. This affects large models like GPT-2 (Radford et al., 2019) with 1.5 billion parameters. Shazeer & Stern (2018) proposed the AdaFactor algorithm, which replaced the full 2 nd moment with moving averages of the row and column sums of the squared gradients. For a layer defined by an n × m matrix, this would reduce memory from O(n × m) to O(n + m). NovoGrad consumes the same amount of memory as SGD with momentum.

3. ALGORITHM:

NovoGrad is based on 3 ideas: (1) layer-wise 2 nd moments instead of 2 nd moment per each parameter, (2) gradients normalization with layer-wise 2 nd moments, (3) decoupled weight decay.

Let g l t be the stochastic gradient for layer l at step t. NovoGrad first computes the layer-wise 2 nd moment v l t using the norm ||g l t ||:

1 v l t = β 2 • v l t-1 + (1 -β 2 ) • ||g l t || 2 (1)

where 0 ≤ β 2 ≤ 1. We use much smaller β 2 than in Adam, usually in the range [0.2, 0.5].2 

The moment v l t is used to normalize the gradient g l t before calculating the first moment m l t . Similarly to AdamW, we decouple weight decay d • w t from the stochastic gradient, but we add it to normalized gradient before computing moment m l t :

EQUATION

Algorithm 1 NovoGrad Parameters: Initial learning rate λ 0 , moments β 1 , β 2 , weight decay d, number of steps T Weight initialization: t = 0, Initialize w 0 .

Moment initialization: t = 1, for each layer

l set v l 1 = ||g l 1 || 2 ; m l 1 = g l 1 √ v l 1 + d • w l 0 . while t ≤ T do λ t ← LearningRateU pdate(λ 0 , t, T ) (compute the global learning rate) for each layer l do g l t ← ∇ l L(w t ) v l t ← β 2 • v l t-1 + (1 -β 2 ) • ||g l t || 2 m l t ← β 1 • m l t-1 + ( g l t √ v l t + + d • w l t ) w l t+1 ← w l t -λ t • m l t

end for end while where 0 < β 1 < 1 is the momentum, typically in the same range as in SGD or Adam [0.9 -0.95].

The first moment can be also computed via an exponential moving average in Adam-like style:

m l t = β 1 • m l t-1 + (1 -β 1 ) • ( g l t v l t + + d • w l t )

Finally, weights are updated the same way as in SGD with momentum.

Similar to Adam, one can construct a counter-example for NovoGrad in the stochastic convex optimization settings (Wilson et al., 2017) . However, the "AMS-Grad" fix (Reddi et al., 2018) for Adam can also be applied in this case to guarantee NovoGrad convergence:

v l t = β 2 • v l t-1 + (1 -β 2 ) • ||g l t || 2 vl t = max(v l t-1 , v l t ) m l t = β 1 • m t-1 + ( g l t vl t + + d • w l t )

4. EXPERIMENTS WITH DEEP LINEAR NETWORKS:

Following (Andrew M. Saxe & Ganguli, 2013; Ian J. Goodfellow & Saxe, 2015) we will use Novo-Grad to train linear model composed of two linear layers w 1 , w 2 without any non-linearity. The model y = (w 1 • w 2 ) • x should output 1 when x = 1. This model is linear with respect to the inputs, but it is non-linear with respect to the weights, since they are factorized into the product of layers' weights. Training the model is equivalent to the minimization of the loss

L(w 1 , w 2 ) = (w 1 •w 2 -1) 2 (Ian J. Goodfellow & Saxe, 2015

). The loss is not convex, and its minima are located on the hyperbola w 1 w 2 = 1 (see Figure 1 ). Minima close to the points (-1, -1) and (1, 1) are good "flat" minima which generalize well. Minima close to the axes are "sharp" minima (Keskar et al., 2016) .

We trained the model with SGD with momentum, Adam, AdamW, and NovoGrad, using the same fixed learning rate,3 weight decay, and weights initialization. The model was trained for 500 steps.

Figure 2 shows the training trajectory and the zoomed-out area near the final point. All algorithms behave in a similar way: first the trajectory goes to the curve w 2 = 1/w 1 , and then follows the hyperbola towards (1, 1) or (-1, -1). During the first phase, training loss decreases, and during the second phase, generalization improves. SGD converges nicely toward (1, 1) but its trajectory is still slightly off of the optimal solution. Adam oscillates wildly around hyperbola w 2 = 1/w 1 , while AdamW behaves much better since weight decay decoupling significantly reduces oscillations.

NovoGrad is the most stable out of four algorithms. It exhibits better generalization and closely follows the minima curve because normalized gradients prevent trajectory from going far from it. We also found that NovoGrad is more robust than other algorithms to the choice of learning rate, weight decay, and weight initialization (see for details Appendix A). 

5.1. IMAGE CLASSIFICATION:

We used ResNet-50 v2 (He et al., 2016) for ImageNet classification task (Russakovsky et al., 2015) .

We trained this model with three optimizers: SGD with momentum (SGD), AdamW, and NovoGrad.

All models have been trained with the batch size of 1024 for 100 epochs. We used quadratic LR decay for SGD with momentum and cosine decay (Loshchilov & Hutter, 2016) for AdamW and NovoGrad. We could not find any training recipe for ResNet-50 with AdamW, so we report the best accuracy we achieved after extensive hyper-parameter search. We used only standard data augmentation methods: re-size, flip, random crop, and did not employ any additional training tricks (He et al., 2018) . The single-crop validation accuracy for each algorithm is reported in Table 1 . Hazan et al. (2015) showed that large batch size is beneficial for SNGD convergence, which motivated us to explore NovoGrad for large batch training. We trained ResNet-50 v2 with batch sizes of 8K and 32K. To compare with the previous methods, we train the model for 90 epochs using cosine LR decay. To emulate a large batch, we used a mini-batch of 128 per GPU and accumulated gradients from several mini-batches before each weight update. To establish the baseline for NovoGrad training with batch 32K we first used the method similar to proposed in Goyal et al. (2017) : scaling the learning rate linearly with the batch size and using LR warmup. This method gives top-1=75.09% and top-5=92.27%. We found that we get much better results when we increase both the learning rate λ and the weight decay d to improve the regularization (see Table 2 ).

For comparison, we took 3 methods, which (1) use fixed batch size during training and (2) do not modify the original model. All 3 methods employ SGD with momentum. The first method (Goyal et al. (2017) ) scales LR linearly with batch size and uses the LR warmup to stabilize the initial training phase. The second method (You et al. (2018) ) combines warmup with Layer-wise Adaptive Rate Scaling (LARS) (You et al., 2017) . The last method (Codreanu et al. (2017) ) uses warmup and dynamic weight decay (WD). NovoGrad outperformed all other methods without using any additional techniques like LR warmup (Goyal et al., 2017) , dynamic weight decay, special batch normalization initialization, etc. Using warm-up (500 steps) we slightly improved top1 accuracy to 75.99% and top5 to 92.72%. We conducted experiments with Jasper-10x5 (Li et al. (2019) ), a very deep convolutional neural acoustic model, on the LibriSpeech speech recognition task (Panayotov et al., 2015) . Jasper was trained with SGD with momentum (SGD), Adam and NovoGrad for 400 epochs with a batch of 256, polynomial LR decay, and Layerwise Adaptive Rate Clipping (LARC). 5 We found that NovoGrad yields lower Word Error Rates (WER) comparing to SGD, especially for the long runs. The model and training parameters are described in Li et al. (2019) . 

5.2.1. LARGE BATCH TRAINING FOR ASR:

We trained Jasper10x5 with batch sizes of 512, 4K, 8K, 16K and 32K on LibriSpeech. In all cases, we trained the model for 400 epochs. For batch size up to 8K we scaled LR linearly with the batch size and used LR warmup. To scale batch to 16K and 32K we also increased weight decay (see Table 5 ). The batch 16K leads to WER comparable to the baseline. Batch 32K has higher WER due to the smaller number of training steps (9 weights updates per epoch). Figure 3 shows WER on dev-clean during training for different batch sizes.

5.3. LANGUAGE MODELING:

We trained Transformer-XL (Dai et al., 2019) , the state-of-the-art LM architecture on the wordlevel WikiText-103 (Merity et al., 2016) benchmark. For all the experiments we used a 16-layer (Sennrich et al., 2015) . For Adam and AdamW we used dropout of P drop = 0.3 and for NovoGrad we used P drop = 0.2. We trained all algorithms with mixed-precision (Micikevicius et al., 2017) for 100K steps (approximately 150 epochs) with a 4K steps warmup on batches of up to 490K source and target tokens obtained via gradient accummulation (Ott et al., 2018) with cosine learning rate annealing. We did not use checkpoint averaging, all the results are reported for the last checkpoint in the corresponding run. Training the model is equivalent to the minimization of the loss L(w 1 , w 2 ) = (w 1 • w 2 -1) 2 . The loss is not convex, and its minima are located on the curve: w 1 w 2 = 1. Minima close to the points (-1, -1) and (1, 1) are good "flat" minima which generalize well. Minima close to axes are "sharp" minima with bad generalization (see (Keskar et al., 2016) ). The 2D-contour plot of the loss function shown on Figure 5 . (1, 1) are good "flat" minima, and solutions near axes are "sharp" minima.

We will study how the behavior of each algorithm depends on learning rate, weight decay and initialization. We will train the model with each optimizer for 500 steps using the same learning rate, weight decay, and weights initialization. To use the same learning rate for all optimizers, we will use the "gradient averaging" for NovoGrad. We will also use the version of SGD with "gradient averaging" (similar to Adam): m t = β • m t-1 + (1 -β) • g t . For fixed learning rate this SGD version is equivalent to the regular SGD with momentum.

Training trajectories for the baseline (fixed learning rate 0.2, weight decay 0.1, and β 1 = 0.95, β 2 = 0.5.) are shown on the Figure 6 . All algorithms first go to the curve w 2 = 1/w 1 , and then slide along hyperbola towards (1, 1) or (-1, -1). SGD is slightly off with respect to the optimal solution. Adam oscillates wildly around line w 2 = 1/w 1 . AdamW behaves better since weight decay decoupling significantly reduces osculations. NovoGrad is the most stable out of four algorithms, it also shows much better generalization than other algorithms and converges to (1, 1) closely following the minima curve.

Next, we increased learning rate from 0.2 to 1.0 while keeping weight decay equal to 0. Similarly, when we increased weight decay from 0.1 to 0.5 while keeping learning rate 0.2, all algorithms except NovoGrad diverge, while NovoGrad demonstrates high robustness to the weight decay choice (see Figure 8 ).

Finally, we started training from different initial point. SGD and NovoGrad are most robust with respect to the initialization, while AdamW diverge (see Figure 9 ). 

:

We use L2-norm ||g l t || to compute v l t . It would be interesting to see how L1 or L∞ norms perform.

If β2 = 0, then v l t = ||g l t || 2, and NovoGrad becomes layer-wise NGD with decoupled weight decay.

We will use the gradient averaging for SGD first moment as in Adam: mt = β • mt-1 + (1 -β) • gt to use the same LR for all optimizers.

Training was done in OpenSeq2Seq(Kuchaiev et al., 2018): https://github.com/NVIDIA/ OpenSeq2Seq/blob/master/example_configs/image2label/resnet-50-v2-mp.py.

https://github.com/NVIDIA/apex/blob/master/apex/parallel/LARC.py.

https://github.com/cybertronai/transformer-xl

Figures and Tables:

Figure 1: 2D-contour plot of surface y = (w 1 • w 2 -1) 2 of linear model with two layers. The loss functions has many global minima are located on hyperbola w 2 = 1/w 1 . Solutions near (-1, 1) and(1, 1) are good "flat" minima, and solutions near axes are "sharp" minima.

Figure 2: Deep Linear Network with two layers: training with SGD, Adam, and NovoGrad

Figure 3: ASR, large batch training. Jasper-10x5 trained with NovoGrad on LibriSpeech.

Figure 4: LM. Transformer-XL trained with Adam and NovoGrad on WikiText-103.

TRAINING OF DEEP LINEAR NETWORKS Following (Andrew M. Saxe & Ganguli, 2013; Ian J. Goodfellow & Saxe, 2015) we will use Novo-Grad to train linear NN model composed of two linear layers w 1 , w 2 without any non-linearity. The model y = (w 1 • w 2 ) • x should output 1 when x = 1. The model is linear function of the inputs, but it is non-linear function of the weights, since they are factorized into the product of layers' weights.

Figure 5: 2D-contour plot of surface y = (w 1 • w 2 -1) 2 of linear model with two layers. The loss functions has many global minima located on hyperbola w 2 = 1/w 1 . Solutions near (-1, -1) and(1, 1) are good "flat" minima, and solutions near axes are "sharp" minima.

1. Training trajectories are shown on the Figure 7: SGD and Adam diverge. Only AdamW and NovoGrad converge.

Figure 6: DLN training -baseline: learning rate 0.2, weight decay 0.1.

Figure 7: DLN training -increased learning rate 0.2 → 1.0: SGD and Adam diverges, AdamW and NovoGrad converge.

Figure 8: DLN training -increased weight decay 0.1 → 0.5: All algorithms except NovoGrad diverge.

Figure 9: DLN training -"bad" initialization: AdamW diverges.

ImageNet classification -ResNet-50 v2, batch 1024, top-1 and top-5 accuracy(%).

Table content: <table><tr><td colspan="7">Optimizer LR policy Init LR Weight decay Epochs top-1,% top-5,%</td></tr><tr><td>SGD</td><td>poly (2)</td><td>0.400</td><td>0.0001</td><td>100 200</td><td>76.38 76.33</td><td>93.08 92.96</td></tr><tr><td>AdamW</td><td>cosine</td><td>0.002</td><td>0.120</td><td>100 200</td><td>76.36 76.48</td><td>93.01 92.94</td></tr><tr><td/><td/><td/><td/><td>100</td><td>76.94</td><td>93.41</td></tr><tr><td>NovoGrad</td><td>cosine</td><td>0.007</td><td>0.002</td><td>200</td><td>77.74</td><td>93.70</td></tr><tr><td/><td/><td/><td/><td>300</td><td>77.65</td><td>93.62</td></tr><tr><td colspan="7">NovoGrad outperformed both AdamW and SGD with the top-1 accuracy of 76.94% after 100</td></tr><tr><td colspan="7">epochs. SGD and Adam accuracy remained under 76.5% if we trained for 200 epochs instead,</td></tr><tr><td colspan="7">while NovoGrad accuracy improved to 77.74%. NovoGrad demonstrated powerful regularization</td></tr><tr><td colspan="7">capabilities: training for 100 additional epochs kept top-1=77.65% without overfitting. Note that</td></tr><tr><td colspan="7">this is "vanilla" ResNet-50, without sophisticated data augmentation or model tweaking.</td></tr><tr><td colspan="5">5.1.1 LARGE BATCH TRAINING FOR IMAGE CLASSIFICATION</td><td/><td/></tr></table>

Large batch training with NovoGrad -ImageNet, ResNet-50 v2, 90 epochs, accuracy(%).

Table content: <table><tr><td colspan="5">Batch Init LR Weight decay Top-1,% Top-5,%</td></tr><tr><td>1K</td><td>0.070</td><td>0.002</td><td>76.86</td><td>93.31</td></tr><tr><td>8K</td><td>0.016</td><td>0.006</td><td>76.64</td><td>93.14</td></tr><tr><td>32K</td><td>0.026</td><td>0.010</td><td>75.78</td><td>92.54</td></tr></table>

Large batch training comparison -ImageNet, ResNet-50v 2, top-1 accuracy(%) .

Table content: <table><tr><td>Optimizer</td><td>Reference</td><td>Bag of Tricks</td><td colspan="3">Epochs B=8K B=32K</td></tr><tr><td>SGD</td><td>Goyal et al. (2017)</td><td>warmup</td><td>90</td><td>76.26</td><td>72.45</td></tr><tr><td>SGD</td><td>You et al. (2018)</td><td>warmup, LARS</td><td>90</td><td>75.30</td><td>75.40</td></tr><tr><td>SGD</td><td colspan="2">Codreanu et al. (2017) warmup, multi-step WD</td><td>100</td><td>76.60</td><td>75.31</td></tr><tr><td>NovoGrad</td><td/><td>-warmup</td><td>90 90</td><td>76.64 -</td><td>75.78 75.99</td></tr><tr><td colspan="2">5.2 SPEECH RECOGNITION</td><td/><td/><td/><td/></tr></table>

Speech recognition -Jasper-10x5, LibriSpeech, 400 epochs, WER (%).

Table content: <table><tr><td>Optimizer</td><td>Dev clean other clean other Test</td></tr><tr><td>Adam</td><td>13.20 31.71 13.36 32.71</td></tr><tr><td>SGD</td><td>3.91 12.77 3.98 12.79</td></tr><tr><td>NovoGrad</td><td>3.64 11.89 3.86 11.95</td></tr></table>

Large batch training with NovoGrad -Jasper-10x5, LibriSpeech, 400 epochs, WER (%).

Table content: <table><tr><td colspan="4">Batch Init LR Warmup Weight decay</td><td>Dev clean other clean other Test</td></tr><tr><td>512</td><td>0.015</td><td>-</td><td>0.001</td><td>3.58 11.30 3.85 11.29</td></tr><tr><td>4K</td><td>0.03</td><td>0.05</td><td>0.001</td><td>3.66 11.67 3.92 11.68</td></tr><tr><td>8K</td><td>0.06</td><td>0.05</td><td>0.001</td><td>3.69 11.76 3.96 11.75</td></tr><tr><td>16K</td><td>0.06</td><td>0.05</td><td>0.003</td><td>3.67 11.03 3.94 11.19</td></tr><tr><td>32K</td><td>0.06</td><td>0.08</td><td>0.004</td><td>4.01 11.73 4.14 11.89</td></tr></table>

LM. Transformer-XL trained on WikiText-103 with batch size 256, sequence length 512.

Table content: <table><tr><td colspan="6">Optimizer Tokens Init LR Weight decay Val PPL Test PPL</td></tr><tr><td>Adam</td><td>12B</td><td>0.00025</td><td>-</td><td>23.84</td><td>25.40</td></tr><tr><td>AdamW</td><td>12B</td><td>0.00025</td><td>0.001</td><td>23.64</td><td>25.06</td></tr><tr><td>NovoGrad</td><td>12B 20B</td><td>0.01 0.01</td><td>0 0</td><td>20.53 19.89</td><td>21.26 20.65</td></tr></table>

batch 490K tokens, 150 epochs,  no checkpoint averaging. Tokenized BLEU and detokenized SacreBLEU on WMT'14 (newstest14). We tested NovoGrad on deep models for image classification, speech recognition, translation, and language modeling. In all experiments, NovoGrad performed equally or better than SGD and Adam/AdamW. NovoGrad is more robust to the initial learning rate and weight initialization than other methods. For example, NovoGrad works well without LR warm-up, while other methods require it. NovoGrad performs exceptionally well for large batch training, e.g. it outperforms other methods for ResNet-50 for all batches up to 32K. In addition, NovoGrad requires half the memory compared to Adam.

Table content: <table><tr><td colspan="5">Optimizer Init LR Weight decay SacreBLEU TokenBLEU</td></tr><tr><td>Adam</td><td>0.0006</td><td>-</td><td>28.26</td><td>28.71</td></tr><tr><td>AdamW</td><td>0.0006</td><td>0.005</td><td>28.24</td><td>28.72</td></tr><tr><td>NovoGrad</td><td>0.04</td><td>0.0001</td><td>28.80</td><td>29.35</td></tr><tr><td>6 CONCLUSION</td><td/><td/><td/><td/></tr><tr><td colspan="5">We propose NovoGrad -an adaptive SGD method with gradients normalized by the layer-wise</td></tr><tr><td colspan="3">second moment and with decoupled weight decay.</td><td/><td/></tr></table>

Peer review (iclr2020):
The authors present a variation of Adam that combines layerwise normalization and weight decay decoupling. They test quite extensively their algorithm against other optimization methods on varied tasks. From a theoretical point of view, the contribution seems very incremental. The authors acknowledge that weight decay decoupling is already present in the literature, so it appears that the only contribution is the layer-wise normalization. No justification is proposed as to why this kind of normalization would either accelerate the convergence or leads to better generalization. Proof of convergence even in a deterministic convex setting are missing, so the reader has to extrapolate correctness from previous work on adaptive gradient descent. On the other hand, the proposed algorithm is tested on a great variety of tasks, using state of the art models. The reported performance for the benchmarks (checked for Resnet and Transformer-XL) are on par with what can be found in the literature. The proposed method outperform consistently the other optimizers in terms of generalization performance. I have some concerns regarding section 4: “SGD converges nicely toward (1, 1) but its trajectory is still slightly off of the optimal solution”. It is unclear to me what the reader should understand. Does it converge to the optimal solution? If yes, why should we expect the trajectory to follow the hyperbola? Using the same learning rate for all methods is a bit odd. Why not search for the best learning rate for each optimizer and report its performance? It seems that the oscillation of some of the optimizers could be fixed by using a smaller learning rate. Also, it can be seen in section 5 that Adam consistently needs a smaller learning rate than Novograd. If the difference between NovoGrad and AdamW lies in the layer-wise second moment, it is unclear to me why their performance should differ on this task, as each layer has only one weight. It would be great if the authors could clarify this point. As a conclusion, I am a bit conflicted regarding this paper. The motivation for this modified version of AdamW are unclear, but the empirical results are convincing and rigorous. The authors made a great effort in testing in a variety of different settings. I’m leaning toward accepting this paper, to give the community a chance of testing and, maybe, adopting it.