Apollo: An Adaptive Parameter-wised Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization
Authors: Xuezhe Ma
Venue:  (2021)
Keywords: Optimization
Stochastic Optimization
Nonconvex
Quasi-Newton
Neural Network
Deep Learning

Abstract: In this paper, we introduce Apollo, a quasi-newton method for noncovex stochastic optimization, which dynamically incorporates the curvature of the loss function by approximating the Hessian via a diagonal matrix. Algorithmically, Apollo requires only first-order gradients and updates the approximation of the Hessian diagonally such that it satisfies the weak secant relation. To handle nonconvexity, we replace the Hessian with its absolute value, the computation of which is also efficient under our diagonal approximation, yielding an optimization algorithm with linear complexity for both time and memory. Experimentally, through three tasks on vision and language we show that Apollo achieves significant improvements over other stochastic optimization methods, including SGD and variants of Adam, in term of both convergence speed and generalization performance.

Under review as a conference paper at ICLR 2021
APOLLO: AN ADAPTIVE PARAMETER-WISE DIAG- ONAL QUASI-NEWTON METHOD FOR NONCONVEX STOCHASTIC OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
In this paper, we introduce APOLLO, a quasi-Newton method for nonconvex stochastic optimization, which dynamically incorporates the curvature of the loss function by approximating the Hessian via diagonal matrix. Importantly, the up- date and storage the diagonal approximation of Hessian is as efficient as adaptive first-order optimization methods with linear complexity for both time and memory. To handle nonconvexity, we replace the Hessian with its rectified absolute value, which is guaranteed to be positive-definite Experiments on three tasks of vision and language show that APOLLO achieves significant improvements over other stochastic optimization methods, including SGD and variants of Adam, in terms of both convergence speed and generalization performance. The implementation of the algorithm is available anonymous link.
INTRODUCTION
Nonconvex stochastic optimization of core practical importance many fields of machine learning, in particular for training deep neural networks (DNNs). First-order gradient-based optimization algorithms, conceptually attractive due to their linear efficiency on both the time and memory complexity, have led tremendous progress and impressive successes. A number of advanced first-order algorithms have emerged over the years to pursue fast and stable convergence, among which stochastic gradient descent (SGD) (Robbins & Monro 1951;
al.) 1998), equipped with momentum (Rumelhart et al. 1985 Qian 1999 Bottou & Bousquet, 2008). has stood out for its simplicity and effectiveness across a wide range of applications (Hinton & Salakhutdinov 2006, Hinton et al. 2012 Graves, 2013). However, one disadvantage of SGD is that the gradients in different directions are scaled uniformly, resulting in limited convergence speed and sensitive choice of the learning rate, and thus has spawned a lot of recent interest in accelerating SGD from the algorithmic and practical perspectives.
Recently, many adaptive first-order optimization methods have been proposed to achieve rapid training progress with element- sealed learning rates, and we can only mention a few here due to space limits. In their pioneering work, Duchi et al. (2011) proposed AdaGrad, which scales the gradient by the square root of the accumulative square gradients from the first iteration. While AdaGrad works well for sparse settings, its performance significantly degrades for dense settings, primarily due to the monotonic increase of the accumulation. Subsequently several methods have been proposed with the intuition limit the accumulation in particular exponentially reduce the weight of earlier iterations. Notable works incorporating this method are RMSProp (Tieleman & Hinton| 2012). AdaDelta (Zeiler 2012). and Adam (Kingma & Ba 2015), among which Adam has become the default optimization algorithm across many deep learning applications because its fast convergence speed and relatively consistent selections of hyper-parameters (Ruder, 2016 Zhang et al. (2020). However, it has been observed that these adaptive optimization methods may converge to bad/suspicious local optima, resulting in worse generalization ability than their non-adaptive counterparts (Wilson et al. 2017), or fail converge due unstable and extreme learning rates (Luo et al.) 2019).
small window of past iterations, and
Duasi-Newton methods have been widely used in solving convex optimization problems, due to their efficient computation and fast convergence rate (Broyden] 1967, Dennis & Moré 1977). However, the stochastic, high-dimensional and nonconvex nature of many machine learning tasks, such as

---

Under review as a conference paper at ICLR 2021
training deep neural networks, has rendered many classical quasi-Newton methods ineffective and/or inefficient (Keskar & Berahas 2016 Wang et al., 2017: Yao et al., 2020). Indeed, in many natural language processing (NLP) and computer vision (CV) tasks (He et al. 2016 Ma & Hovy 2016 Luo et al. 2019), SGD (with momentum) is chosen as the optimizer, benefiting from its stable and efficient training and outstanding generalization.
In this work, we develop APOLLO, quasi-Newton method for nonconvex stochastic optimization simultaneously tackle the aforementioned challenges of stochastic variance, nonconvexity and inefficiency. Algorithmically. APOLLO dynamically incorporates the curvature of the objective function with diagonally approximated Hessian. It only requires first-order gradients and updates the approximation of the Hessian diagonally so that it satisfies parameter-wise version of the weak secant condition (Wolfe 1959). To handle nonconvexity, we replace the Hessian with its rectified absolute value, the computation of which is also efficient under our diagonal approximation, yielding an efficient optimization algorithm with linear complexity for both time and memory ($3). Experimentally, through three tasks on CV and NLP with popular deep neural networks, including ResNets (He et al.) 2016), LSTMs (Hochreiter & Schmidhuber 1997) and Transformers (Vaswani al. 2017), we demonstrate that APOLLO significantly outperforms SGD and variants of Adam, in terms of both convergence speed and generalization performance ($4)
BACKGROUNDS
In this section, we set up the notations on nonconvex stochastic optimization, briefly review the (quasi-) Newton methods, and discuss the problems of applying quasi-Newton methods to nonconvex stochastic optimization that we attempt to study in the rest of the paper.
2.1 NONCONVEX STOCHASTIC OPTIMIZATION
In this paper, we consider the following stochastic optimization problem: min f(0) E(; [)] ER
(1)
where :: Rd x R" + R is continuously differentiable (and possible nonconvex) function, 0 Rd denotes the parameter be optimized, R² denotes a random variable with distribution function P. and B[-] denotes the expectation w.r.t r. Intuitively, T incorporates noises in f. leading to stochastic objective function. special case of (T) that arises frequently in machine learning the empirical risk minimization problem:
N Σ 4(0) N
1
min ((0) BER
where 4 Rd R the loss function corresponds to the i-th data, and N the number of data samples that is assumed to be extremely large. Objective functions may also have other sources of noise than data subsampling, such dropout (Srivastava al. 2014 in deep neural networks.
Decoupled Parameters. In this work, consider setting of decoupled parameters: 0 (0(), 1 I, L}. Intuitively, under this setting the parameter 0 is decoupled into sequence of parameters serving different functionalitics. For example, neural network training the parameters of neural network can be naturally decoupled into the parameters of different layers modules.
2.2 NEWTON QUASI-NEWTON METHODS
Newton's 's method usually employs the following updates to solve II]
0++1 = 0, H gt
where 9r (0.) is the gradient at 0, and H, ((,;) is the Hessian matrix. The convergence rate of Newton's 's method is quadratic under standard assumptions (Nocedal & Wright 2006). How- ever, major challenges with this method are i) the expensive computation of the inverse Hessian at every iteration and the corresponding quadratic memory complexity; and ii) the limitation to convex functions (nonconvexity results in negative curvature of H, and misleads the update directions).
2

---

Under review as a conference paper at ICLR 2021
A standard alternative Newton's method is class of quasi-Newton methods, which have been widely used in solving convex deterministic optimization:
0 = 8: - 1/: B, 9:
(4)
where 7/s is the stepsize (a.k.a learning rate), B, is an approximation to the Hessian matrix (2,( Oe, which is updated based on the well-known secant equation:
B 1 argmin | - BBBll
B
S.t. 14 = y/: (secant equation)
where 8, = 0 - 02 and gt+1 9t. B++1 is, in the sense of some matrix norm, the closest to B among all symmetric matrices that satisfy the secant equation. Each choice of the matrix norm results in different update formula, such as DFP (Davidon 1991 Fletcher 1987) and BFGS (Broyden 1970, Fletcher 1970; Goldfarb 1970; Shanno 1970). The popularity of this method due the fact that only the gradient of the objective function is required at each iteration. Since no second derivatives (Hessian) are required, quasi-News Newton method is sometimes more efficient than Newton's method, especially when the computation of Hessian is expensive. To further reduce memory cost, one seminal work is the limited memory BFGS (L-BFGS) (Liu & Nocedal 1989; Byrd et al. 1995) that achieves desirable linear computational and memory complexity by approximating the Hessian as a series of sum of first order information from previous iterations.
2.3 PROBLEMS OF QUASI-NEWTON METHODS
Despite their impressive successes on convex deterministic optimization, quasi-Newton methods suffer from their own problems in the scenario of large-scale nonconvex stochastic optimization.
Stochastic Variance. One main challenge that prevents quasi-Newton methods from being applied to nonconvex stochastic optimization (I is the variance introduced by the stochastic nature of the problem. At each iteration, only the stochastic gradient ge is available, which is an unbiased estimation of the gradient V (O,) and may lead to an erroneous approximation of Hessian (Byrd et al. 2011).
Nonconvexity. Another key challenge in designing stochastic quasi-Newton methods for noncon- vex problem lies in the difficulty of preserving the positive-definiteess of B, in (5). due to the nonconvexity of the objective function. What is worse is that performing line search is infeasible in the stochastic setting, due to the presence of noise in the stochastic gradients (Wang et al.) 2017).
Computational and Memory Efficiency. Even though quasi-Newton methods are more efficient than Newton's method, the time and memory complexities are still relatively large compared with adaptive first-order methods. For instance, L-BFGS requires to store first-order information from m previous iterations with commonly m 2 5, which is still too expensive for deep neural networks con- taining millions of parameters. Moreover, adapting quasi-Newton methods to nonconvex stochastic optimization probably introduces additional computation, further slowing down these methods.
Due to the aforementioned problems, no quasi-Newton methods (to our best knowledge) designed for nonconvex optimization consistently outperform adaptive first-order algorithms w.r.t convergence speed and generalization performance. The main goal of this work is to algorithmically design and experimentally demonstrate novel quasi-Newton method, in hope of improving the convergence speed and generalization performance of nonconvex stochastic optimization eventually.
ADAPTIVE PARAMETER- WISE DIAGONAL QUASI-NEWTON
With the end goal of designing an efficient quasi- Meastn method to solve the problem in (T) mind.
we first propose to approximate the Hessian with diagonal matrix, whose elements are determined
by the variational approach subject to the parameter-wise weak secant equation ($.. Then, we
explain our stepsize bias correction technique to reduce the stochastic variance in $3.2 To handle nonconvexity, we directly use the rectified absolute value of the diagonally approximated Hessian as the preconditioning of the gradient ($3.3). The initialization technique of APOLLO allows us to eliminate one
($3.4). At last, we provide a theoretical analysis of APOLLO'S convergence in online convex programming ($3.5.. The pseudo-code is shown in Algorithm
3

---

Under review as a conference paper at ICLR 2021
3.1 QUASI-NEWTON METHODS WITH DIAGONAL HESSIAN APPROXIMATION
As discussed in Bordes et al. 2200), designing an efficient stochastic quasi-Newton algorithm involves careful trade-off between the sparsity of the approximation matrix B, and the quality of its approximation of the Hessian H, and diagonal approximation is a reasonable choice (Becker et al. 1988). If B is chosen to be diagonal matrix satisfy the secant equation in (5). one can obtain an updating formula similar to the SGD-QN algorithm (Bordes ct al. 2009).
An alternative of the secant equation in the updating formula (5. as first introduced by Nazareth (195),, is the weak secant equation (Dennis & Wolkowicz, 1993):
B++1 argmin | - Bell
B s.t. $f B141s = sf y: (weak secant equation)
(6)
The motivation of using the weak secant condition in diagonal quasi-Newtor method is straight- forward: the standard mean-value theorem might not necessarily hold for vector-valued functions expressed in the secant equation, B 8 = Yt R (0,))) Thus, we do not know whether there exists vector 9 E Rd such that yr = f())s (Dennis & Moré 1977). On the other hand, the Taylor theorem ensures that there exists such 0 that s Ut = 8 V f()).. leading to the reasonable assumption of the weak secant condition (6).
Based on the variational technique (Zhu et al.
1999), the solution of (6) with Frobenius norm is: $f yr
sf B,S, Diag(s?)
A t Bi41 B,
where of is the element wise square vector of st, Diag(s?) is the diagonal matrix with diagonal elements from vector sy. and || I|a the 4-norm of vector.
Parameter- Wise Weak Secant Condition. However, in optimization problems with high- dimensional parameter space, such as training deep neural networks with millions of parameters, the weak secant condition might be too flexible to produce good Hessian approximation. In the setting of decoupled parameters ($ (.2.1) we propose a parameter-wiss version the weak secant equation to achieve a trade-off between the secant and weak secant conditions: for each parameter ((l)) 0, we update B corresponding to g(t) by solving (6) individually. Remarkably, the secant condition restricts B with an equation of dddimenniionl vector, while the weak secant condition relaxes it with 1-dimensional scalar. The parameter-wise weak secant condition expresses the restriction as a I-dimension vector (1 < < d), resulting in reasonable trade-off. The updating formula is the same as 7 for each parameter-wise B.
3.2 STEPSIZE BIAS CORRECTION
To mitigate the stochastic variance problem in stochastic quasi-Newton methods, APOLLO utilizes stepsize bias correction on the stochastic gradients at each step f. Concretely, comparing the parameter updating formulas of Newton's method (3) and quasi-Newton methods (4), we observe that the optimal stepsize equals to 1 if the Hessian approximation B, and the stochastic gradient 9: are close to the exact Hessian H; and gradient V ((.), respectively. Inspired by this, we correct the stepsize bias in the stochastic gradient ge by replacing it with a corrected gradient g! = Ne9t- Together with the corresponding corrected y/ = g(+1 gt = (elli, we correct the updating term of B, in (7 by:
A' 2/ xf B,st Diag(s?)
||
(8)
df y1 + df Bed. Diag(d?)
[ddll]
where de - B, 9: is the update direction. Based previous studies, incorporating exponential moving averages (EMVs) for the stochastic gradients significantly reduces the vari- ance (Kingma & Bal 2015 Yao al. 2020)) We follow these work and apply exponential moving average to St. together with the initialization bias correction:
3(1 - 8t)
1-8 3
Mc+1
1 --4 M4 +
1 gt+ 9c+1
(9)
4

---

Under review as a conference paper at ICLR 2021
Algorithm 1: APOLLO, our proposed algorithm for nonconvex stochastic optimization. All operations on vector are element-wise wise. Good default settings are B 0.9 and = le-t.
Initial: mo. do, Bn < 0, 0,
// Initialia mo, do, Bg to zero
while t {O...., } do
for 0 E (0' ..... 0%) do
gc+1 + (((+))
// calculate gradient at step t
m++1 - T T-pett + 1-8
A(1-3')
T-
// Update blas-corrected moving
+ df (m m B,d.
1/ Calculate coefficient B update
([della+e))
B++1 + B₁ 1 a
- -Dagg(d))
// Update diagonal Hessian // Handle nonconvexity
D++1 t- rectify (Brrly
d+1 t- D241me+
// Calculate update direct on // Update parameters
O.t1 c Os - 1+11d++1
end
end
return 0T
where 0 < 3 < 1 is the decay rate of EMV and yt in (8) is written M4+1 m,- Note that do not apply moving average methods to the approximated Hessian, though the diagonal matrix is easier to be explicitly formed to average than full matrices. Further improving APOLLO by investigating the moving average of the diagonal matrix B, might be a interesting direction of future work.
3.3 RECTIFIED ABSOLUTE VALUE OF HESSIAN FOR NONCONVEXITY
To guarantee convergence, quasi-Newton methods require the approximated Hessian matrix B, to be positive definite at each step. The common strategy in previous studies is to solve the updating formula in (5) by restricting the candidate matrix B to be symmetric positive definite. It is known that the BFGS update preserves the
of Bu+1 as long as the curvature condition sf yt 0 holds, which can be guaranteed for strongly convex problem. For nonconvex problem, the curvature condition can be satisfied by performing line search, which is, however, expensive or even infeasible in stochastic setting, because the exact function values and gradient information are unavailable. Wang et al. (2017) proposed the stochastic damped L-BFGS (SdLBFGG) method that implicitly generates a positive definite matrix without line search. However, it usually requires large history size (m 2 100) to guarantee convergence, which is infeasible for large-scale optimization.
To handle nonconvexity, we adopt a different strategy that does not require the solution of B, in (5) to be positive definite. Intuitively, search for B, that is a good approximation of the real Hessian, which is not necessarily positive definite in nonconvex problem. When we use B, as preconditioning to calculate the update direction, we use its absolute value:
B V B B.
(10)
where V- is the positive definite square root of matrix. The motivation of absolute value is straight- forward: for dimensions with large absolute values of curvature, the objective function could be very sharp and we would prefer to take relatively smaller steps than those flatter dimensions. Since APOLLO formulate B, as a diagonal matrix, the cost of computing [:]] is marginal.
Rectified Absolute Value of B, For nonconvex objective functions, there exist inflection points whose curvatures are zero. To prevent the steps from becoming arbitrarily large, we rectify the absolute value of B, with convexity byper-parameter d:
Dr = rectify ( [B,, o) = max( ()))
(11)
where the rectify(;, function is similar to the rectified linear unit (ReL.U) (Nair & Hinton 2010 with threshold set to The update direction in ( is then d, D mm-
AdaHessian Yao et al. 2020) used an idea similar to the absolute values of B, to handle nonconvexity, where the root mean square averaging is applied to compute the Hessian diagonal. Different from APOLLO, AdaHessian requires second-order information to compute the Hessian matvec oracle and approximate the Hessian diagonal using Hutchinson's 's method, which is significantly more costly.
5

---

Under review as a conference paper at ICLR 2021
3.4 INITIALIZATION
The rectified D, in (1I introduces one more hyper-parameter o, limiting the application of APOLLO in practice. In this section, show that the zero initialization approach in APOLLO, which initializes the moving average of gradient ma. the parameter update direction dg and the diagonal approximation of Hessian Bo as (vector of) zeros, leads to coupled stepsize 77 and convexity o, allowing us to hyper-parameter parameter of 17 or
eliminate one hyper-
Coupled Stepsize and Convexity o. With the zero initialization of mo, do, and Bo, we have the following theorem to illustrate the relation between 1j and a (details in Appendix A
Theorem 1. Given zero initialization of my, du, and Bo and a fixed parameter intialization Og- Suppose that we have two sets of inpper-paraneter 17. and n' o' with the same ratio: = 2. Then the convergence trajectories of these two sets of hyper-parametter are exactly the same:
0: == VI (I,...,T):
(12)
where Or and 02 are the parameters of (n, 0) and (n', d') at iteration I, respectively.
From Theorem 1 we observe that and are coupled with each other and in practice we only need to tune one of them, leaving the other fixed. Therefore, in our experiments ($4). we fix a 1 and tune 17 on different problems.
Learning Rate Warmup for APOLLO As discussed in Kingma & Ba 2015) zero initialization
leads to estimations biased towards zero in the initial iterations. For the moving average m,, this bias can be corrected by dividing the bias-correction term (9). For de and B,, however, we cannot derive such bias correction terms. Fortunately, simple linear warmup heuristic of 1) at the beginning iterations achieves remarkablely stable training in our experiments.
3.5 CONVERGENCE ANALYSIS
We analyze the convergence of APOLLO using the online learning framework (Zinkevich 2003) for sequence of convex cost functions fi (0), f2(
Theorem 2. Let {0r} be the sequence obtained from Algorithm 1, m = 3/ 3 1 Vt E [T].
Suppose |[ae|22 < G, |0; 0e S D, Vt, t' [T]. For Og generated with the APOLLO algorithm,
we have the following bound on the regret:
D
37G
Rr -
2(1 - (2VT
(13)
2(1
11e
t=1
The following result falls as an immediate corollary of the above result. Corollary 2.1. Suppose B, --'=1,11 1''''1 in Theorem we have
g²
GrG
RT S 21)(1 3) 22)2
2(1 3) (2VT
-1)
(14)
For the convex case, Theorem lmplies the regret of APOLLO is upper bounded by O(VT) (proof in Appendix B). The conditions for Corollary 2.1 as in Reddi et al. al. (2188) can be relaxed to Br 3/t
and still ensures regret of O(VT).
4 EXPERIMENTS
To evaluate APOLLO, we conduct experiments on four benchmark datasets across three tasks of CV and NI.P that are commonly used to evaluate optimization algorithms: CIFAR-10 (Krizhevsky Hinton] 2009) and ImageNet (Deng et al. 2009) for image classification; One Billion Words (Chelba
ct al. 2013) for language modeling; and WMT 2014 English-German for neural machine translation.
The four baseline methods compare with are SGD with momentum (Bottou & Bousquet] 2008).
Adam (Kingma& & Ba 2015). Rectified Adam (RAdam) (Liu et al. 2020). and AdaHessian ( Yao et al. decouple weight decays in Adam. RAdam and AdaHessian in all the experiments. For each experiment, we report the average over 5 runs. More
2020). Following Loshchilov & Hutter (2019),
detailed descriptions, results and analysis of the conducted experiments provided in Appendix C

---

Under review as a conference paper at ICLR 2021
SGD
Adam
RAdam
Adam-ad|
RAdam-adj
AdaHassian
Apollo accuracy (milestone)
(milestone)
accuracy (milestonee)
Training (milestone)
50
2.5
72
80
1.5
56
02
1.0
0.0
700
0.00
400
Epochs
Training (cosine)
accuracy
Training (cosine)
Test accuracy (cosine)
50
2.5
72
0.3
02
80
1.5
1.0
0.0
(90
0.00
400
Fpochs
CIFAR-10
Fporhs
Fpochs
ImageNet
Fpochs
Figure I: Training loss and test accuracy of ResNet- 10 on CIFAR-10 and ResNeXt-50 on ImageNet, with two schedule strategies of learning rate decay,
Adam
Apallo
0.5
0.5
0.5
0.5
0.4
0.4
0.4
0.4
0.2
0.1
6.1
0.1
0.0
20
40
80
100
120
140
160
0.0
20
62
80
100
120
142
160
0.00
20
40
62
80
100
120
140
160
20
60
80
100
120
140
150
03
90
90
85
85
80
80
73
75
20
40
80
300
120
340
360
19g
20
62
80
100
120
100
20
40
80
100
120
340
160
70g
20
60
80
100
120 300
0.2
1.0
2.0
10.0
Figure 2: SGD, Adam, RAdam and APOLLO with different learning rates on CIFAR-10.
4.1 IMAGE CLASSIFICATION
We begin our experiments with an evaluation of the convergence and generalization perfor- mance on image classification. We use ResNet- 11 for CIFAR-10 and ResNext-50 (Xie et al. > for ImageNet, respectively (details in Ap- pendix C.I). The results on CIFAR- 10 and Im- ageNet are presented in Figure I and Table I together with the four baselines. For each opti- mizer, we use two scheduling strategies of learn- ing rate decay: milestone that decays the learn- ing rate at the end of some predefined epochs; and ii) cosine annealing schedule proposed in Loshchilov & Hutter (2077).
Table 1: Test Acc. CIFAR-10 and ImageNet.
CIFAR-10
ImageNet milestone
Method
milestone
cosine
cosine
SGD
93.91
94.53
77.19
78.17
Adam RAdam
91.41
91.56
71.72
71.19
91.80
91.88
72.37
71.64
Adam-adj
93.74
94.24
76.86
77.54
RAdam-adj
93.88
94.38
76.91
77.68
AdaHHessian
93.97
94.48
APOLLO
94.20
94.60
77.90
78.54
Though Loshchilov & Hutter (001)) claimed that the optimal settings of the learning rate and weight decay factor in Adam with decoupled weight decay is more independent than the original Adam, we observed that the strength of weight decay regularization is still co-related with the learning rate. Since the strength of regularization has significant effect on both the performance of convergence and generalization, we adjust the weight decay rates for Adam and RAdam, named Adam-odj and R RAdam-djj, so that the they have similar strength of regularization for fair comparison. More detailed analysis on the effect of weight decay is provided in Appendix D.2
ResNee- 1 10 is modificd (small) version of ResNet-18 adapt the image size 32 32 in CIFAR-10.
7

---

Under review as a conference paper at ICLR 2021
Adam
RAdam
Apollo
Training
Test PP!
Method
PPL
100
50
43
Adam
36.68
80
39
RAdam
36.20
APOLLO
31.94
50k
10th
150c
200K
2506
2M
6M
10 12
15
Updates
Figure 3: Language modeling (LSTMs) One Billion Words.
Table 2: Test PPL.
From Figure 1 and Table 1 first see that Adam and RAdam, without adjusting the weight decay rates, converge much faster than other optimization methods, while obtaining significantly worse classification accuracy. After adjusting the weight decay rates, the test accuracy of Adam-adj and RAdam-adj remarkably improves, with rapid decline of convergence speed. This suggests that the fast convergence speed of Adam and RAdam results from relatively weak regularization. Thus, the effect of regularization strength needs to be considered when we analyze the performance of different optimization methods. Second, APOLLO outperforms the three first-order methods (SGD, Adam and RAdam) on both the convergence speed and classification accuracy. demonstrating its effectiveness on training the ResNet architectures based on convolutional neural networks (CNNs) (LeCun et al.) 1989). Comparing with AdaHessian, APOLLO obtains better test accuracy with similar convergence speed. Note that AdaHessian requires second-order information and significantly more costly (detailed comparison in Appendix D.3) and we omit it from the following experiments.
Robustness to Learning Rate Change. Besides performance improvements, we also investigate the robustmess of different optimization methods the change learning rate. For each optimizer, we use the learning rate in the previous experiment (Table ( as the base, i.c. 0.1 for SGD, 0.001 for Adam and R. RAdam, and 0.5 for APOLLO. Then, we explore different learning rates that are a times of the base learning rate, with E (0.2, 1.0, 2.0, 10.0). To eliminate the impact of weight decay, we adjust the weight decay rates according to the factor a. Experimental results with ResNet- 110 on CIFAR-10 are summarized in Figure 2 After correcting the impact of weight decay, all the optimization methods, except SGD with = 10.0, achieves consistent model performance (highly overlapped training and test curves with each other), while APOLLO slightly improves the robustness of model training over the three baseline methods.
4.2 LANGUAGE MODELING
To evaluate APOLLO on Recurrent Neural Networks (RNNs) that are applied in wide range of NLP tasks (Graves 2013), we conduct experiments the One Billion Words dataset, using two-layer LSTM network for language modeling (details Appendix C.2)
Figure and Table illustrate the perplexity (PPL) of training and test for each optimization method. As shown in Figure3] although APOLLO is slower than Adam and RAdam the first few updates, its convergence is much faster after that. On generalization performance, APOLLO achieves significant improvements (more than 4.0 PPL. points on test data) over Adam and RAdam. This demonstrates the effectiveness of APOLLO on training LSTM-based neural architectures.
Training Stability. From the middle plot Figure 3 we see that the training losses of Adam and RAdam may suddenly increase. This occurred in all the runs of experiments using Adam and RAdam, and we selected these successfully converged - the loss went back to normal after some updates, and discarded those failed to converged - the model crashed due to loss numerical overflow. The models optimized with APOLLO never suffered from this issue, demonstrating the convergence stability of it.
4.3 NEURAL MACHINE TRANSLATION
Table 3: BLEU on WMT-14.
To evaluate APOLLO on Attention-based Transformer architec- ture (Vaswani et al. 2017), we train the Transformer-base model on the WMT2014 English-German (EN-DE) dataset (around 4.5M sentence pairs). We use the same data preprocessing steps as in Ma et al, (019) (details in Appendix C3 From Table [3] APOLLO obtains improvements over all the baseline methods. For each experiment, report the mean and standard variance over 5 runs.
Method
BLEU
SGD
26.59
Adam RAdam
27.84
20.15
APOLLO
28.34
8

---

Under review as a conference paper at ICLR 2021
5 RELATED WORK
Stochastic Quasi-Newton Methods. In the literature of (nonconvex) stochastic quasi-Newtur
methods, several algorithms have been developed recently for large-scale machine learning problems:
oLBFGS (Schraudolph et al. 2007] Mokhtari & Ribeiro, ) RES (Mokhtari & Ribeiro 2014).
SFO (Sohl-Dickstein al.) 2014), SQN (Byrd al. 2016), SdLBFGS (Wang et al.) 201 7, and
AdaQN (Keskar & Berahas 2016). among which only SILBFGS and AdaQN are designed to solve nonconvex optimization problems. The SdLBFGS algorithm carefully controls the quality of modified BFGS updates to preserve the cositie-defiinitensss of B, in (5) without line search. AdaQN shares a similar idea but is specifically designed for RNNs by refining the initial L-BFGS scaling, step acceptance control and choice of curvature information matrix, and adopting the SQN framework (Byrd et al. 2016). Different from these two methods, APOLLO does not require B, in (5) to be positive definite, but replacing it with its rectified absolute value to handle nonconvexity. Moreover, both SdLBFGS and AdaQN use the updating formula similar to L-BFGS and require even larger history size (commonly 2 100) to guarantec convergence, preventing them from being applicd to large-scale optimization. For comprehensive comparison, we conducted experiments with small toy CNN models to compare SdLBFGS with Apollo (details in Appendix E.
Adaptive First-Order Methods. From the diagonal approximation of Hessian, APOL.OO is also related to those diagonally-scale RMSProp (Tilleman & Hinton| (012)) AdaDelta [Zeiler] (012), and Adam (Kingma & Ba 2015
first-order algorithms, such as AdaGrad (Duchi et al. 201 1).
Subsequently, a number of techniques have emerged to theoretically justify and algorithmically im- prove Adam, including AMSGrad (Reddi et al. 2018), AdaBound (Luo et al. 2019) and RAdam (Liu et al., 2020). The main difference is that the diagonal preconditioning in APOLLO is directly derived from the quasi-Newtor updating formula (6). In terms of memory efficiency, Anii et al. (2019) and Chen et al. (2020) further reduces the memory cost adaptive methods, and Agarwal et al. (2019)
proposed efficient method for full-matrix adaptive regularization.
Stochastic Second-Order Hessian Methods. Stochastic Second-Order Hessian-Free meth-
ods (Martens) 2010 Martens & Sutskever] 201T) implicitly solve quadratic models using matrix -vector products. [Dauphin et al. (2014) argued the existence of saddle points and proposed method to rapidly escape them. K-FAC (Martens & Grosse 2015) computes a second-order step by constructing an invertible approximation of the Fisher information matrix in an online fashion. Shampoo (Gupta) jet (018)) approximates the Fisher information matrix using low-rank decomposition. Recently, Yao et al. (2020) proposed AdaHessian, which approximates the Hessian diagonal using Hutchinson's 's method. These second-order methods differ from A APOLLO mainly in the request of second-order information of the objective function at each iteration.
CONCLUSION AND EXTENSIONS
We have introduced APOLLO, simple and computationally efficient quasi-Newton algorithm for nonconvex stochastic optimization. This method is aimed towards large-scale optimization problems in the sense of large datasets and/or -igh-dimensoonal parameter spaces.
We briefly outline a few extensions to APOLLO that we want to explore in future work.
Parameter- Wise Gradient Clipping. The standard gradient clipping method (Pascanu et al. 2013) is to clip the gradients based on the norm computed over gradients of all the parameters together. modification of gradient clipping to properly apply it to APOLLO is clip the gradient of each parameter individually based on its own norm. Preliminary results are provided in Appendix D.4]
Decoupled Weight Decay in APOLLO. (Loshchilov & Hutter 2019) demonstrated that L2 regu- larization is not identical to weight decay for adaptive gradient methods and proposed Adam with decoupled weight decay (AdamW). The application of decoupled weight decay to APOLLO is slightly different from AdamW as APOLLO memorizes the update direction of the last iteration de to update the diagonal Hessian. The algorithm of APOLLO with decoupled weight decay is in Appendix F
A APOLLO is Scale- Variant. An important advantage of adaptive optimization methods, including Adam and its variants, is that they are inherently scale-invariant - invariant with the scale of the objective function. The property of scale- invarinnvee yields more consistent parameeaassserss of these adaptive methods than SGD across different machine learning tasks. Unfortunately, APOLLO does not hold the property of scale-in invarinnaee Developing a scale-invarian version of APOLLO might be potential direction for future work.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning. pp. 102-110, 2019.
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. In Advances in Neural Information Processing Systems, pp. 9749-9758, 2019.
Sue Becker, Yann Le Cun, et al. Improving the convergence of back-propagation learning with second order methods. In Proceedings of the 1988 connectionist models summer school, pp. 29-37, 1988.
Antoine Bordes, Léon Bottou, and Patrick Gallinari. SGD-QN: Careful quasi-newton stochastic gradient descent. The Journal of Machine Learning Research, 10:1737-1754, 2009.
Léon Bottou and Olivier Bousquet. The tradeoffs of large scale learning. Advances in neural information processing systems, pp. 161-168, 2008
Charles G Broyden. Quasi-newton methods and their application to function minimisation. Mathe- matics of Computation, (99)):68-33)) 1967.
Charles George Broyden. The convergence of class of double-rank minimization algorithms. IMA Journal of Applied Mathematics, 6(1):76-90, 1970.
Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. A limited memory algorithm for bound constrained optimization. SIAM Journal on scientific computing, 16(5):1190-150)) 1995.
Richard H Byrd, Gillian M Chin, Will Neveitt, and Jorge Nocedal. On the use of stochastic hessian information in optimization methods for machine learning. SIAM Journal on Optimization, 21(3): 977-995, 2011.
Richard H Byrd, Samantha L Hansen, Jorge Nocedal, and Yoram Singer. A stochastic quasi-newton method for large-scale optimization. SIAM Journal on Optimization, 26(2):1008-103 2016.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3000 2013.
Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. In International Conference on Learning Representations, 2020.
Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in high-dimensional non-convex optimization. In Advances in neural information processing systems, pp. 2933-2944 2014.
William C Davidon. Variable metric method for minimization. SIAM Journal on Optimization, 1(1):
1-17, 1991.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248 255.255. leee, 2009.
John E Dennis, and Jorge J Moré. Quasi-newton methods, motivation and theory. SIAM review, 19 (1):46-8) 89, 1977.
John E Dennis, Jr and Henry Wolkowiez. Sizing and least-change secant methods. SIAM Journal on Numerical Analysis, 30(5):1291-1314 1993.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Roger Fletcher. A new approach to variable metric algorithms. The computer journal, 13(3), 1970.
Roger Fletcher. Practical methods of optimization. John Wiley & Sons, 1987.
10

---

Under review as a conference paper at ICLR 2021
D Goldfarb. A family of variable metric updates derived by variational means. Mathematics of Computation, 24(109):23- 22,)), 1970.
Alex Graves. Generating sequences with recurrent neural networks. arXiv preprint:1388.085000 1308.0800, 2013.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- 2018.
tion. arXiv preprini
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IFEE Signal processing magazine, 29()):2-97, 2012.
Geoffrey E Hinton and Ruslan R Salakhutdinov. Reducing the dimensionality of data with neural networks. science, 3 313(578)):504-500 2006.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1780-1780, 1997.
Nitish Shirish Keskar and Alhert S Berahas. AdaQN: An adaptive quasi-newton algorithm for training In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 1-16. Springer, 2016.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 144):541-551, 1989.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-233) 1998.
Yingkai Li and Huidong Liu. Implementation of stochastic quasi-newton's 's method in pytorch. arXiv 2018
preprint
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1-3):503-528, 1989.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representationss 2020.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, In International Confer- ence on Learning Representations, 2019.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. International Conference on Learning Representations, 2019.
Xuezhe Ma and Eduard Hovy. End-to-end sequence labeling via bi-directional LSTM-CNNs-CRE In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1064 1074, Berlin, Germany, August 2016.

---

Under review as a conference paper at ICLR 2021
Xuezhe Ma, Chunting Zhou, Xian Li, Graham Neubig, and Eduard Hovy. Flowseq: Non- autoregressive conditional sequence generation with generative flow. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing, pp. 4282-42228 4292, Hong Kong, November 2019.
James Martens. Deep learning via hessian-free optimization. In Proceedings of the 27th International Conference on International Conference on Machine Learning, pp. 735-742, 2010.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417, 2015.
James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian ffrce optimization. In Proceedings of the 28th international conference machine learning (ICML-I I), pp. 1033- 1040. Citeseer, 2011.
Aryan Mokhtari and Alejandro Ribeiro. Res: Regularized stochastic bfgs algorithm. IEEE Transac- tions on Signal Processing, 62(23))6089-6104, 2014.
Aryan Mokhtari and Alejandro Ribeiro. Global convergence of online limited memory bfgs. The Journal of Machine Learning Research, 16(1)33111-1181, 2015.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In ICML, 2010.
JL Nazareth. If quasi-newton then why not quasi-cauchy. SIAG/Opt Views-and-news, 6:11 1995.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Science & Business Media, 2006.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. FairSeq: fast, extensible toolkit for sequence modeling. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations)) pp. 48 28-53,200 2019.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. International conference on machine learning, pp. 1310-1318, 2013.
Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks, 12(1): 145-151, 1999.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint drXiv:1609.04747,220 2016.
David E Rumelhart, Geoffrey Hinton, and Ronald J Williams. Learning internal representations by eтror propagation. Technical report, California Uniy San Diego La Jolla Inst for Cognitive Science, 1985.
Nicol N Schraudolph, Jin Yu, and Simon Günter. A stochastic quasi-newton method for online convex optimization. In Artificial intelligence and statistics, pp. 436-443, 2007.
David F Shanno. Conditioning of quasi-newton methods for function minimization. Mathematics of computation, 24(111):647-656 1970.
Jascha Sohl-Dickstein. Ben Poole, and Surya Ganguli. Fast large-scale optimization by unifying stochastic gradient and quasi-newton methods. International Conference on Machine Learning, pp. 604-612, 2014.
12

---

Under review as a conference paper at ICLR 2021
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1988, 2014.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe. Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826, 2016.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2), 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Xiao Wang. Shiqian Ma. Donald Goldfarb, and Wei Liu. Stochastic quasi-newton methods for nonconvex stochastic optimization. SIAM Journal on Optimization, 27(2):927-956, 2017.
Ashia C. Wilson, Rebecca Roelofs, Mitchell Stem, Nati Srebro, and Benjamin Recht. The marginal
value of adaptive gradient methods in machine learning. In Advances in neural information processing systems, pp. 4148-44 14884118,, 2017.
Philip Wolfe. The secant method for simultancous nonlinear equations. Communications of the ACM, 2(12): 1959.
Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1492-1500, 2017.
Zhewei Yao, Amir Gholami, Sheng Shen, Kurt Keutzer, and Michael W Mahoney. ADAHESSIAN: An adaptive second order optimizer for machine learning. arXiv preprint arXiv:2006.0077 2020.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint: preprint: 1212.5701, 2012.
Jingzhao Zhang, Sai Pranceth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33, 2020.
Mingfa Zhu, John Lawrence Nazareth, and Henry Wolkowicz. The quasi-cauchy relation and diagonal updating. SIAM Journal on Optimization, 9(4)11122-204, 1999.
Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference machine learning (ieml-03), pp. 928-9398-69999 2003.
13

---

Under review as a conference paper at ICLR 2021
APPENDIX: APOLLO: AN ADAPTIVE PARAMETER- WISE DIAGONAL QUASI-NEWTON METHOD FOR NONCONVEX STOCHASTIC OPTIMIZATION
A COUPLED STEPSIZE AND CONVEXITY
Before proving Theorem we first define the notations.
Let d' be the ratio of the two sets of learning rates. Let d, and B; be the corresponding terms of parameter 8% at step for (r', a').
Proof of Theorem
Proof. Induction on the step of updates t, we attempt to prove that at each step
my m/, Bj = aB: and 0, 6', 14
(15)
Initial step: when t= 1, since 0g = 8/ we have my m/. With do d' = and 8
we have
Bi Bi = and:
Dj recify reiif((B1,0) (B
D recify recify(B(,,0)) d'
Then, Dí a and
0f = ==00- 10- 1 Df-'D''mm m² = 0o - ppa(Dii/a)m) 0o - -nDi'm = =0.
Thus, the statement (15) true.
Induction on step t: Suppose that the statement (55 is true for all the previous t steps. Now we prove the case 1. From the inductive assumption and D we have,
B₁ aB
d, de and m4+1 m
From 8
d. s. d.B.ddd
Bidd (
B11 B,
1d,
=aB; (e y/: + (a (
di)
dTy, yt de B,d, Diag(d?)
= =aB, -a
||d 12
aBB+11.
Then,
D'41 recify( 0')1100)
veiiy(aaBtt1.0
arecify ())
= aDD+1
and we have 0'+1 Oz4-1-
Finally, to sum up with the induction, have proven Theorem
14

---

Under review as a conference paper at ICLR 2021
B CONVERGENCE ANALYSIS
Proof of Theorem
T Proof. Let 0 = argmin Σ f(),, where F is the feasible set of A. From the Algorithm
we have
BEF 1=1
0.41- O" 0, 0" Dr mm, Then, we have the following:
||D| - -0")12 1 - -0") - Dr 1/2 mel 1
( - 8²)1 + 11m: 2 11 mmm! (0: - -0")
11,,//(0, 0²) 01) + 11m=11-
1/2 m2 (3 (1 Be)g:)((0)) 0*)
Then, have
of (0. 0*)
1
|1)//(00 (0,- 11
03
21 (1 3,)
+ 2(1 B.)
1/
11 1/2 m2l1
Br
m T (0, 0*)
1 Be
|ID|//2 (0, - -02) D (0c+1 - 10)1
2m (1 8,)
1/1
D2
1/2
m2112 +
1/8
m+ 112
33
00:
10
2(1 3.)
2(1
3.)
21
Be)
Using the standard approach of bounding the regret at each step with convexity of the functions T
(f:)T- we have the following bound of R Σ fe(O.) fe(0)):
T
f.(0.) fe(0%) S
Σ
T D6
(0. 0²)
1
[ID| (0,-0'11- ID (0 - 00)
211:(1 - {e)
(16)
1).
|| 1/2 1/2 m2
1i 2(1 - 3e)
||+ 1112
2(1 - 31)
B8
11: - 11
27/ 3e)
From the triangle inequality, we have
1ID/ (0, - 03
|1D (0.+1 - -0+1-00113 < (2) 02+1)ll2 = "1 1/2 m2 (17)
To sum up with [6) and (17) we have
T
RT <
1
D.
1/2 1/2mm ml
27p
B.)
t=1
2(1 11
|D, 1/2
m2113
m+
11
111
| 112
B1)
2(1 3e)
(18)
110 - 0²1
21 - 31)
T
T
32
Tl
D 1/2
m
11 +
2(1 T/c
mm 112 +
10
001
1
3r
3.)
2 (1
3,)
t=1
t=1
15

---

Under review as a conference paper at ICLR 2021
Since the element of D, is rectified by i.e. Di,3 1, and llme|la 1/ S G, < < 1, we have
1
1/2
m211
U/
m 1
2
| 12
11
m
Bt
2(1 3r)
Br
2(1 - 3,)
3G
Σ 2(1 - 3))
Σ
1=1
(19)
3n/ T
2(1 - 3,) 1=1 Vt
31
< 2(1 31)
(2VT 1)
The last inequality is due to the following upper bound:
T
Σ
1= Vt = 2VT
In addition, as 110, - 10,-00112 S
and 3e < B < 1, we have
32
D2
T 38
00:
8 11
Σ
(20)
2m7
Bi)
2(1
3)
171
(=1
t=1
Finally, to sum up with 18
19 and
20
we have
D
T
82
3r (2VT- - 1)
R3 S
2(1 - 3)
11
2(1 - 3)
t=1
Proof of Corollary 2.1
Proof. Since Be = =3 by sum of arithmetico writhnetico-geometric series we have
Σ 1
2((-11)
1
(21)
(1-
Plugging 21 into 20
have
32
3nG
Rr S 21y(1 3)(1 x2)2
(2VT
1)
2(1 3)
C EXPERIMENTAL DETAILS
C.I IMAGE CLASSIHICATION
For CIFAR-10 dataset, we use the ResNet-110 architecture in the public implementation Note that ResNet- 10 is modified version of ResNet-18 (He et al. 2016) to adapt the small image size 32 32 in CIFAR-10 10, and much smaller than ResNet-18. The number of parameters for ResNet- 10 and ResNet-18 18 are 1.73 M and 1 11.69 M, respectively. For ImageNet, we used the neural architecture of ResNeXt-50 (Xie et al., 2017). Batch size 128 for CIFAR-10 10 and 256 for ImageNet. For each optimizer, used two learning rate decay strategies. First, we train the model CIFAR- 10 for 164 epochs and decay the learning rate at the end of 80-th and 120-th epochs by 0. For ImageNet, train the model for 120 epochs and decay the learning rate at the end of 40-th and 80-th
https:: https::///iithubbccom//eeee http://thtw....///e.//////ttttt/ assi Fication
16

---

Under review as a conference paper at ICLR 2021
epochs by 0.1. Second, we also used the cosine annealing schedule (oshhhhlovv & Hutter 2017). For the cosine annealing schedule, we train a CIFAR- 10 model for 200 epochs and a ImageNet model for 120 epochs. The start learning rate is 0.1 for SGD, 0.001 for Adam and RAdam, 0.15 for AdaHessian and 0.5 for APOLLO. The weight decay rates for SGD and APOLLO are 5e on CIFAR-10 and 2 on ImageNet. For AdaHessian, the weight decay rate is le-A on CIFAR-10 and 2e-4 on ImageNet. For Adam and RAdam, we adjusted the weight decay rates based on the ratio of the learning rate, i.e. multiplying 500 to the original decay rates] For APOLLO learning rates are warmed up linearly from 0.01 to 0.5 in the first 100 updates. For AdaHessian, we performed grid search for hyper-parameters with 1) E (0.1,0.15, 0.2), c E fle-1, 1e 1e 11e-") and weight decay E on the generalization accuracy (7) = 0.15, = le_?, weight decay = 1e- In addition, we observed that warmup is also crucial for AdaHessian. Thus, we linearly warmed up the learning rates from 0.001 to 0.1 in the first 100 updates, similar to APOLLO. The implementation of AdaHessian is based on the public implementation" Random cropping and random horizontal flipping are applied to training data. For each experiment, we conduct training on one NVIDIA Tesla V100 GPU for CIFAR-10 and eight NVIDIA Tesla V100 GPUs for ImageNet.
le 2e 3e
-3 4e ). and selected the optimal hyper-parameters based
C.2 LANGUAGE MODELING
One Billion Words dataset (Chelba et al. 2013) is publicly available benchmark for measuring progress of language modeling. It contains about 0.8 billion tokens with a vocabulary of 793,471 words, including sentence boundary markers. Different from Liu et al. (2020) which shrinks the vocabulary to about 0.64 million words, we used the standard vocabulary For the language model, we used two-layer LSTM with 2048 hidden states with adaptive softmax and 300-dimensional word embeddings as input. The cut-offss of the adaptive softmax are set to (60000, 100000, 640000], which different from Liu et al.] (2020). Dropout (Srivastava et al., 2014) is applied to each layer with drop rate of 0.1. The start learning rate is set to 0.001 for Adam and RAdam, and 10.0 for APOLLO. Gradient clips with 1.0 are applied to all the optimization methods. Each model is trained for 20 epochs, and the learning rate decays at the end of the 12-th and 18-th epochs by decay rate 0.1. LSTMs are unrolled for 20 steps without resetting the LSTM states and the batch size is set to 128. Every models is trained on one NVIDIA Titan RTX GPU.
C.3 NEURAL MACIINE TRANSLATION
Our experiments on WMT 2014 English-German are based on the Transformer-base model (Vaswani ct al. 2017), with implementation from the FairSeq package (Ot et al.] (2019). This dataset contains 4.5M parallel sentence pairs for training. We following the standard setting (Vaswani al. 2017).
using Newstest2013 as the validation set and Newstest2014 as the test set. The dataset is pre- pre-processed following (Ma et al. 2019). using the scripts from FairSeq package Specifically, embedding with 512 dimension and 6-layer encoder/decoder with 8 multi-head attention and 2048 feed-forward dimensions. We apply 0.1 label smoothing (Szegedy et al.] (0166). and perform totally 500, 000 updates to train each model. For Adam and RAdam, we use start learning rate 0.0005. For Adam we set (0.9, 0.98). while for RAdam we set 3 = (0.9, 0.999). For SGD and APOLLO, the start learning rates is 0.1 and 10.0, respectively. The momentum of SGD is 0.9. For learning rate scheduling, we applied linear warm up the learning rate for SGD, Adam APOLLO - 4000 updates for Adam and 1000 updates for SGD and APOLLOO. For RAdam, we did not apply warm up because RAdam inherently designed to avoid it. After learning rate warming up, we applied the inverse square root decay (Vaswani et al. 2017) to Adam. For SGD, RAdam and APOLLO, we decayed the learning rate at the 250, 000 and 450, 000 updates by decay rate O. Gradient clips with 1.0 are applied to all the optimization methods, and the dropout ratio are set to 0.1. Weight decay rates 1e-4 for Adam and RAdam, le 6 for SGD, and 1 for APOLLO. The decoding beam size is set to 5, and the checkpoints of the last 10 epochs averaged before evaluation. For each experiment, we conducted distributed training across eight NVIDIA Tesla V100 GPUs with maximum batch size as 8192 tokens per GPU (totally 8192 tokens per batch).
word
Weight decay adjustment: Todom f/ Yapolla * Nappolio where 7 is the weight decay factor.
an
httt::::/qithuub.co///raa http:///.a..aaaaoooooooooa/a/aa/
17

---

Under review as a conference paper at ICLR 2021
Table 4:
Classification accuracy on CIFAR-10 and ImageNet. For each experiment, we report the mean and standard variance over 5 runs.
CIFAR-10
ImageNet
Method
milestone decay
cosine annealing
milestone decay
cosine annealing
SGD
93.91+00.07
94.53+0.27
77.19+0.07
78. 78.17=0.06
Adam
91.41+0..00
91.56:1:.19
71.72++0.13
71.19=0.10
RAdam
91.80+0.04
91.88=0015
72.37+0008
71.64+0.14
Adam-adj
93.74+0.15
94.2410.099
76.86+0.006
77.54=0.16
RAdam-adj
93.88+0.11
94.38+0.25
76.91
77.68=0.08
AdaHessian
93.97-10.22
94.48
APOLLO
94.20-0012
94.60
77.90
78.54+00.9
Text 800
Figure 4: Training loss and test accuracy of ResNet-1 10 on CIFAR-10 with various rates of weight decay, with two schedule strategies of learning rate decay,
D DETAILED EXPERIMENTAL RESULTS
In this section, we report the detailed experimental results in Section 4 and the results the investigation of the effect of weight decay.
D.1 DETAILED RESULTS ON IMAGE CLASSIFICATION
Table 4 iillstrates the details of the experimental results Image Classification. For each experiment, we report the mean values with corresponding standard deviations over runs.
D.2 EFFECT OF WEIGIT DECAY RATE ON OPTIMIZATION
To further investigate the effect of weight decay rate converge speed and generalization perfor- mance for different optimization methods. we conduct experiments on CIFAR-10 ResNet- 110 with range of weight decay rates from 1 to 9e-1. As discussed before, for Adam and R Adamm, we adjust the weight decay rate based on the ratio of the learning rate such that the they have similar strength of regularization for fair comparison:
Nadam * Maddm = Tapolto TopooHo
(22)
where 7 is the weight decay rate and 17 is the learning rate.
Figure 4 shows the convergence of different optimization methods with various rates of weight decay, together with the classification accuracy. APOLLO achieves improvements over all the three baselines convergence speed with different rates of weight decay. For classification accuracy, APOLLO obtains the best accuracy when the weight decay rate is larger than 4e -4 When the weight decay rate is decreasing, SGD obtains the best accuracy, while APOLLO achieves comparable performance.
18

---

Under review as a conference paper at ICLR 2021
D.3 COMPARISON ON TRAINING SPEED AND MEMORY COST
In this section, we compare the training speed and memory between SGD, Adam, AdaHessian and APOLLO. summarzzss the comparison cost of a single iteration of update. For fair comparison, experiments of CIFAR-10 and One Billion Words are conduced on single NVIDIA TITAN RTX GPU, while experiments of ImageNet and WMT are performed with distributed training on 8 NVIDIA Tesla V100 GPUs.
Table Comparison between different optimization methods on training speed and memory cost.
CIFAR-10
ImageNet
1BW
WMT-14
Cost (XSGD)
Speed
Memory
Speed
Memory
Speed
Memory
Speed
Memory
SGD
1.00
1.00
1.00
1,00
1.00
1.00
1.00
1.00
Adam
1.16
1.01
1.01
1.03
1.19
1.34
1.13
1.04
Apollo
1.42
1.01
1.23
1.05
1.49
1.62
1.19
1.06
AdaHessian
5.76
2.12
11.78
2.51
3.51
2.78
8.46
2.47
From T we that the second-order AdaHessian requires much more computational resource than first-order methods on both time and memory. In addition, the slow-down of AdaHessian becomes more significant for larger-scale models with distributed training across multiple GPUS, such as ResNext-50 on ImageNet and Transformer on WMT.
D.4 EXPERIMENTS ON PARAMETER- WISE GRADIENT CLIPPING
In this section, we provide some preliminary results on parameter-wise gradient clipping, a modifica- tion of the standard gradient clipping that is inherently proper to APOLLO. Parameter wise gradient clipping is to clip the gradient of each parameter individually based on its own norm. It can be regarded trade-off between gradient clipping by global norm and by each value.
We conducted two groups of experiments to compare with the standard gradient clipping method language modeling and neural machine translation. The experimental settings for standard gradient clipping are exactly the same as in section [ where we clipped the gradient by global norm 1.0 for each model. For parameter-wise gradient clipping, we clipped each parameter by 0.5 for the LSTM model in language modeling, and 0.1 for the Transformer base model in NMT.
Table 6: Comparison between APOLLO with standard and parameter-wise gradient clipping One Billion Words and WMT-14. We report the mean and standard variance over runs.
IBW
WMT-14
Standard
28.34+0.10 31.94+0.009
Parameter-wise
31.75+t0.10
28.39
Table 6 lists the preliminary results. On both the two groups of experiments, parameter-wiss gradient clipping slightly outperforms the standard one.
E EXPERIMENTS WITH SMALL TOY CNN MODELS
In this section, we provide the comparison between SdLBFGS (Wang et al. 2017) and APOLLO CIFAR-10 10 dataset with a small toy CNN mode The implementation of SdLBFGS is based on the public Py PoTcc ( which includes two important modifications to the original SdLBFGS algorithm: identity matrix initialization and direction normalization (Li & Liu 2018). For each optimizer, we train the CNN model for 50 epochs with batch size equals to 64. After each epoch, the learning rate is decayed by the rate 0.95, For the start learning rate for each optimizer, we performed
19

---

Under review as a conference paper at ICLR 2021
search in wide range: 7) (1.0, 0.5, 0.2, 0.1, 0.05, 0.01, 0.005). and select the obtains the optimal performance. The final start learning rate for both SdLBFGS and APOLLO is 0.1. Following [i & Liu (2018), the memory size of SdLBFGS is set 100. For APOLLO, linearly warmed up the learning rate from 0.01 in the first 10 updates. For other hyper-parameters of each optimizer, we choose the default value.
SdLBFGS
Apollo
Training loss
Test accurocy
1.7%
1.80
50
0.7%
640 10
Figure 5: Training loss and test accuracy of SdLBFGS and APOLLO on CIFAR-10 with the small toy CNN model.
we see that APOLLO convergences faster than SaLBFGS and obtains comparable test accuracy. Note that APOLLO is much faster (more than 10 times for one iteration) than SdLBFGS and consumes much less memory (SdLBFGS stores 100 previous update directions).
From Figure 5
F APOLLO WITH DECOUPLED WEIGHT DECAY
Algorithm APOLLO with weight decay (1.2/Decoupled)
Initial: mg. do. Bu < 0, 0, while / {0, .g.....7. T} do
// Initialize mo, do, Bo to zero
for 0 E (0 ..... 06) do
9++1 < +20,
// Calculate gradient at step t
A(1-9')
M++1 & 1-ge-1 my + 198--9941
// Update bias-corrected moving
+ d_ (m -m))+mm)++ B,d:
// calculate coefficient of B update
B4+1 +- B₁ B--08 a Diag(d?) (d?)
// Update diagonal Hessian // Handle nonconvexity
D(+1 rectify( (B++1, 1)
d+1 t
// Calculate update direction
f.t + O₁ - m+41d4+1
// Update parameters
end
end
return 0T
20


Reviews:
[{"review_id": "sal3hhQ5HiT", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is clearly written and easy to understand. The authors give a clear introduction to Newton and quasi-Newton methods, and summarize three main drawbacks of quasi-Newton methods for nonconvex stochastic optimization which is the real case for many practical problems. The three main challenges are stochastic variance, nonconvexity and inefficiency.\n\nAs far as I known, there is no quasi-Newton methods that could solve the above challenges simultaneously, and this paper give a solution that considering all the three aspects.\n\nFor efficiency, the computational and memory efficiency are both considered. Instead of depending on first order informations from m previous iterations, the proposed method approximates the Hessian matrix by  considering  the diagonal parameters of B_t which is more memory and computational efficient. \n\nTo solve the stochastic variance, the authors propose stepsize bias correction to work better in stochastic gradient descent framework. And exponential moving average to g_t is adopted to make the gradient more stable.\n\nAnd to better support nonconvexity, a rectified trick is used in the constraint B_t. What’s more, in order to produce a better Hessian approximation, a parameter-wise weak secant condition is used.\n\nAll the above ideas look reasonable, and this give a full workable solution to train networks in variance domains.\n\nThe authors conduct several experiments in Image Classification, Language Modeling and Neural Machine Translation, and experiments show the effectiveness of the proposed method. The proposed method has advantage in all tasks, and for Language Modeling, the improvements seem to be very large.\n\nQuestions:\n(1) What are the effects of different factors such as “stepsize bias correction”, “rectified trick” and “a parameter-wise weak secant”.\n(2) Why the improvements for Language Modeling are larger than for Neural Machine Translation and Image Classification?\n\nThe authors are strongly suggested to test more network structures on challenging tasks and try to achieve new state of the art results which could make the method more convincing. And more theoretical analysis is suggested to better understand the proposed method.\n\n==========\n\nI would like to change my rating from 9 to 5. The paper proposes an interesting idea and does achieve good results on several datasets. However,  after reading all the comments and feedbacks, I notice that the comparisons are not convincing enough, and I have some concerns about the performance of the proposed method on more general and challenging tasks.", "extra_scores": {}, "extra_text": {}}, {"review_id": "TNub_MjNUxB", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**:\n\nThe paper proposes a Quasi-Newton inspired optimization algorithm for Stochastic Optimization named APOLLO. It adjusts a previously known update formula to better suit Deep Learning by using 1) a layer-wise diagonal approximation to the Hessian, 2) an exponential average of gradients to address the noise. Overall the algorithm shows promising results on the assigned experiments.\n\n\n**Reasons for score**:\n\nI believe the paper has a lot of potential but requires some additional discussion and therefore weakly reject the paper in its current form. My main concern lies in the (mostly lack off) comparison to SGD with momentum followed by the choice of matrix norm used for the update. \n\n\n**Pros**:\n\n+ The authors identify the main problems of QN methods in the DL setting and provides a good argument for using a diagonal approximation instead. \n\n+ The presentation of the theory and algorithm is straightforward and easy to follow with references acknowledging preceding work.\n\n+ I am positively surprised by the efficacy of the presented algorithm given that the diagonal approximation is only chosen to satisfy the weak secant equation. It seems to strike a good trade-off between a light constraint but a restrictive model for the Hessian.\n\n\n**Concerns**:\n\n- In my opinion, the most important optimizer to compare to is SGD with momentum since there we would see the effect of the proposed preconditioning vs none. It is also the second best in the image classification (sometimes surpassing APOLLO). Can you explain why iti is not present in the other experiments?  \n\n- Eq. 5 and 6 should preferably highlight the usage of matrix norm by changing $|| B - B_t || $->$ || B - B_t ||_W$ (for example) and clarify W=I is used.\n\n- The authors mention that the solution to the constrained optimization in Eq. 5 can result in different popular QN algorithms depending on the matrix norm, which often plays an important role in traditional unconstrained optimization. APOLLO uses the Frobenius norm (W=I) which I would like some further comments on. Did you consider using other norms like $W=B_{t-1}$ or instead estimating the inverse Hessian like BFGS (s <-> y)? \n\n- How does the momentum influence the performance of APOLLO? Is a high beta required or does it also work for beta=0 so the batch gradient is used? Given that all the gradients in the update are replaced with the momentum term the algorithm becomes reminiscent of an Adagrad update (with gradient replaced by momentum), but with additional \"forgetting\" (the lack of which has plagued Adagrad in non-sparse settings and lead to the development of several new optimizers). \n\n- How sensitive is the algorithm to the warmup? This parameter is not mentioned in the language modeling experiment. Could not all of the algorithms benefit from a warmup strategy? \n\n- The experiments are varied and overall look good but with a relatively small margin in the case of image classification. Given the matter of weight-decay it is difficult to interpret if the improvement over (R)Adam is due to the algorithm or ill-tuned regularization. I think this is handled in a suitable way as outlined in the appendix but it makes the comparison more difficult. \n\n\n--------------\n\n**Post rebuttal**\n\nI have considered the revised paper, rebuttal and feedback+rebuttal of fellow reviewers and in the end decided to leave my score unchanged. Below is a summary of my reasoning.\n\n----\n\nThe Rebuttal has addressed many concerns and the revised edition has further strengthened the paper in many ways but unfortunately lack in the empirical evaluation. \nIt seems like the warmup of Apollo requires a [start, end] learning rate as well as the increase rate compared to the single learning rate of say SGD. It is not clear that the additional overhead of tuning these parameters could not be used to further improve the training schedule of the baseline for better performance (particularly in the CV). At the provided [[link]](https://github.com/bearpaw/pytorch-classification/blob/master/TRAINING.md) for the CV task it looks like the weight decay for CIFAR-10 with the Resnet-110 architecture is set to $10^{-4}$ (not $5\\cdot 10^{-4}$) for which $\\eta=0.1$ (SGD+M) was good, meaning that the used $\\eta$ is not necessarily optimal for the higher weight decay. \n\nThe results on the language modeling task are impressive but for the algorithm to be accepted as a particularly good algorithm for RNNs it should compare to the more elaborate SotA algorithms for this particular task (AdaQN was proposed). \n\n\n\n----\nSome points that would be good to address regardless of outcome (no influence on my decision):\n- Does the algorithm work in the $\\beta=0$ setting? SGD with momentum reverts to SGD and Adam reverts to RMSProp which both are competitive optimizers. Does that also hold for Apollo?\n- How are the values in Table 5 (D.3) calculated? Depending on the implementation of Apollo it looks like 3-4 parameter-sized vectors are required per update ($g$,$m$,$B$,$d$) which in the case of 4 is twice the amount of SGD with momentum, yet the memory is only a few % larger.\n- In algorithm 2 you should replace $\\lambda$ with $\\gamma$ to be consistent with the rest of the paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "LwKvsrK7R-A", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents the optimization method Apollo, a quasi-Newton method that relies on a parameter-wise version of the weak secant condition to allow for a diagonal approximation of the Hessian. Additionally, the issue of a potentially non-PSD approximation is addressed by replacing the approximation with a rectified absolute value. While the combination of techniques is interesting, my main hesitation comes from the limited discussion concerning other quasi-Newton methods for the same problem setting.\n\nTo begin, a much more significant overview of the distinctions between this work and those of AdaQN and SdLBFGS is certainly warranted, as few details are provided to explain how they differ from the methods in this paper. For example, the comment made about AdaQN is that it \"shares a similar idea but specifically designed for RNNs.\" First, I am unsure why the implication is that this somehow weakens the merit of AdaQN as compared to the Apollo method, especially since the authors themselves evaluate their Apollo method on RNNs for language modeling in Section 4.2. In fact, the authors do not even compare with AdaQN in the RNN experiments, choosing instead only to run against Adam and RAdam. This brings us to a key issue with the paper: why are there no comparisons to any other quasi-Newton methods, for any setting (RNN or otherwise)? Since AdaQN is designed for RNNs, it is perfectly suited as a method to compare with in the language modeling tasks, which exhibit the most notable claimed improvement for Apollo over the adaptive first-order methods.\n\nAs for other related methods such as AdaHessian, I agree that there is a distinction between quasi-Newton methods and second-order Hessian-free methods in terms of the information that is accessed. However, just because second-order information is invoked (through Hessian vector products) does not mean by default that the method is \"significantly more costly\" than these quasi-Newton methods, as is claimed in the paper. Hessian vector product-based methods are desirable precisely because the computational cost is comparable to first-order methods, and here too additional comparison is needed.\n\nOverall, the previous works on quasi-Newton methods for stochastic non-convex optimization have not been sufficiently addressed or compared to, particularly given how those works may also handle the issue of preserving positive-definiteness of B_t.\n\nSmall comments:\n.- \"we demonstrate that Apollo significantly outperforms SGD and variants of Adam\"\nThis is overstated, as the only notable improvement claimed by the paper is for language modeling (the others, particularly for image classification tasks, are modest improvements at best).\n\n.- \"Newton's method usually employs the following updates to solve (1)\"\nIt should be clarified that convexity is important when trying to use (plain) Newton's method to solve problems such as (1).\n\n.- \"unnecessarily\" -> \"not necessarily\"\n\n.- Related work on Hessian-free methods that consider absolute value-based transformations of the Hessian:\n\nDauphin, Yann N., Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. \"Identifying and attacking the saddle point problem in high-dimensional non-convex optimization.\" In Advances in Neural Information Processing Systems, pp. 2933-2941. 2014.\n\n.- Related works, in addition to Gupta et al. (2018), in terms of memory-efficient adaptive first-order methods:\n\nNaman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. \"Efficient full-matrix adaptive regularization.\" In International Conference on Machine Learning, pp. 102-110. 2019.\n\nRohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. \"Memory Efficient Adaptive Optimization.\" In Advances in Neural Information Processing Systems, pp. 9749-9758. 2019.\n\nXinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. \"Extreme Tensoring for Low-Memory Preconditioning.\" In International Conference on Learning Representations. 2020.", "extra_scores": {}, "extra_text": {}}, {"review_id": "BFU17FoqZ9O", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work considers a layer-wise weak secant equation to update the Hessian approximation and  train deep learning models through a stochastic quasi-Newton-like update. The major idea is to construct diagonal approximations so the computational cost is low, and the idea is to some extent similar to adagrad in modifying each coordinate individually with different weights.\n\nThe idea itself is interesting, but to fully demonstrate the effectiveness of the proposed algorithm, there should be at least some comparisons with methods of similar flavor.\nMore specifically, the experiment (section 4) only considered first-order baselines (SGD and Adam-type methods), while later in literature review (section 5), the authors mentioned that there are actually existing stochastic Quasi-Newton methods and adaptive first-order methods that can all be seen as using a sort of diagonal approximation of the Hessian. State of the art in those directions should have been included in the comparison to demonstrate that the proposed work is indeed desirable.\nMoreover, Adam/RAdam are shown to be the worst baseline in the first experiment, while in the later 2 experiments those are not excluded, but instead are the only baselines kept for comparison. This is quite suspicious to me and the authors didn't explain why other baselines are discarded at all.\n\nThere are also discrepancy  in the text description and the algorithm presentation.\nEspecially, the update of alpha in Algorithm 1 is different from (8). Although the motivation is understandable, there should be some text explanation of this change in the algorithm.\n\nNote that Section 3.2 totally doesn't make sense, as the normal Newton method also needs a step size to ensure convergence. There are plenty of examples showing that Newton without line search could fail to converge, even for convex problems. Therefore, the motivation is quite weak. Also note that the update of y'_t is wrong, as g'_{t+1} and g'_t use different eta, so y'_t is not a simple scaling of y_t, although in Algorithm 1 it is anyway not calculated in this manner.\n\nSome minor issues:\n- Johnson & Zhang 2013 didn't use the kind of moving average of  gradient described in (9), thus this citation is not quite correct.\n- The learning rate warmup strategy in the end  of section 3  is not quite satisfactory. The text only mentioned that some heuristic is applied, but didn't describe it there nor in the experiment. Relevant information is only available in the appendix.\n- There are many minor typos. The authors should have proofread the manuscript once more before submision. \n\n=====Post Rebuttal======\n\nI appreciate the authors' responses and the revision of the manuscript. My point about Section 3.2 is not that the correction doesn't make sense, but that the reasoning is not quite convincing from an optimization point of view. I would suggest the authors to simply say what they have replied me instead of trying to link this part to Newton's method.\nI wonder why the warm up strategy is only applied to the proposed method but not others if that is effective. Isn't it an unfair comparison?\nI also appreciate the new experiments, but decided to retain the score unchanged. My major concern is that if the additional cost of other solvers is the main issue, then probably it would be better for the authors to directly show the training time as well so that the comparison can be straightforward.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Dear authors,

I took your concerns into account, and I also understand the whole crazy situation around the COVID-19. Many of the reviewers have families (e.g., in US, many kids are now homeschooled, and there are no good daycare solutions as well). I do not plan to list all the good parts of the paper and list weaknesses that are already mentioned and visible to you. Hence, let me focus on my concerns about this paper (and I hope you could find them interesting and they will help you to improve your paper).

+ I personally find the use of 2nd order method in DNN a way to improve many inefficiencies of ADAM/SGD, .... and using diagonal scaling is one way to do it. 

-- I personally find some sections not very motivated and explanatory. E.g. Section 3.2 is just telling half of the story and is missing some details to give the reader the full understanding.

-- The fact that B_t  is not necessary >0, it makes intermediate sense to use some kind of \max\{B_i, \sigma\}  to have the "scaling" to be $\succ 0$.
Note that there are also SR1 methods that would guarantee the matrix to be not necessary pd.

-- Your main motivation was non-convex problems, but the only theorem in the main paper was for convex loss only, right? In this case, I guess there is no issue with B_t to have some coordinates <0, right?

Overall, I find the topic interesting and would like to see an updated paper in some of the top ML venues, but right now I cannot recommend it for acceptance!


Decision: Reject