mL-BFGS: A Momentum-based L-BFGS for Distributed Large-scale Neural Network Optimization
Authors: Yue Niu
Zalan Fabian
Sunwoo Lee
Mahdi Soltanolkotabi
Salman Avestimehr
Venue: Accepted by TMLR ()

Abstract: Quasi-Newton methods still face significant challenges in training large-scale neural networks due to additional compute costs in the Hessian related computations and instability issues in stochastic training.
A well-known method, L-BFGS that efficiently approximates the Hessian using history parameter and gradient changes, suffers convergence instability in stochastic training.
So far, attempts that adapt L-BFGS to large-scale stochastic training incur considerable extra overhead, which offsets its convergence benefits in wall-clock time.
In this paper, we propose mL-BFGS, a lightweight momentum-based L-BFGS algorithm that paves the way for quasi-Newton (QN) methods in large-scale distributed deep neural network (DNN) optimization. 
mL-BFGS introduces a nearly cost-free momentum scheme into L-BFGS update and greatly reduces stochastic noise in the Hessian, therefore stabilizing convergence during stochastic optimization.
For model training at a large scale, mL-BFGS approximates a block-wise Hessian, thus enabling distributing compute and memory costs across all computing nodes.
We provide a supporting convergence analysis for mL-BFGS in stochastic settings.
To investigate mL-BFGS's potential in large-scale DNN training, we train benchmark neural models using mL-BFGS and compare performance with baselines (SGD, Adam, and other quasi-Newton methods). 
Results show that mL-BFGS achieves both noticeable iteration-wise and wall-clock speedup.

Published in Transactions on Machine Learning Research (July/2023)
mL-BFGS: A Momentum-based L-BFGS for Distributed Large-Scale Neural Network Optimization
Yue Niu
yuentuGusc.edv
Deportment of Electrical and Computer Engineering University of Southern Calffrria
Zalan Fabian
afabion edu
Department of Electrionl and Computer Engineering University of Southern California
Sunwoo Lee
sunwool@inha.co oc.kr
Department Computer Science and Engineering, Inha University
Mahdi Soltanolkotabi
sollanol@usc.ccm edu
Department of Electrical and Computer Engineering
University of Southern California
Salman Avestimehr
anestime
Department of Electrical and Computer Engineering
University of Southern California
Reviewed on OpenReview: https: // openreview. net/ pertffomm? 9d=
Abstract
Quasi-Newton methods still face significant challenges in training large-scale neural networks due to additional compute costs in the Hessian related computations and instability issues in stochastic training. A well-known method, T-BFGS that efficiently approxinates the Hessian using history parameter and gradient changes, suffers convergence instability in stochastic Training. So far, atlempts that adapt L-BFGS to large-scale stochastic training incur considerable extra overhead, which offsets its convergence benefits in wall-clock time. In this paper, we propose mI-BFGS, lightweight. momentum-based L-BFGS algorithm that paves the way for quasi-Newton (QN) methods in largo-scale distributed deep neural network (DNN) optimization. mL-BFGS introduces nearly cost-free momentum scheme into L-BFGS update and greatly reduces stochastic noise in the Hessian, therefore stabilizing convergence during stochastic optimization. For model training at large scale, mL-BFGS approximates a block-wise Hessian, thus enabling distributing compute and memory costs across all computing nodes. We provide a supporting convergence analysis for mL-BFGS in stochastic settings. To investigate mL-BFGS's potential large-scale DNN training, we train benchmark neural models using mL-BFGS and compare performance with baselines (SGD, Adam, and other quasi-Newton methods). Results show that mL-BFGG achieves both noticeable iteration-wis and wall-clock speedup.
1 Introduction
In supervisod learning, a typical task is to minimize empirical risk function,
1 N
min L(O;X)
Σ 8(0:I,,M.
(1)
N
where 0 E d denote the parameters to be optimized, and X' represents training samples {x,, y
At present, stochastic gradient descent method (SGD) and its variants, such as Adam (Kingma Ba 2014 are the preferred methods to optimize parameters 0 due to their simplicity, especially for large-scale machine

---

Published in Transactions on Machine Learning Research (02l)
learning problems. Nevertheless, second-order quasi-Newton (QN) methods have also been extensively investigated due to their superior convergence over gradient descent (GD) in strongly convex optimization settings (Gao & Goldfarb 2019 Rodomanov & Nesterov 2021).
However, the Achilles heel of QN methods that is impeding their wide adoption for large-scale machine learn- ing problems is their substantial compute and memory costs. Specifically, these barriers stem from attaining second-order information, including computing and storing the Hessian, performing matrix inversion, etc. For large-scale neural networks, these operations pose daunting challenges to QN's implementations and significantly affect running time. Moreover, second-order methods are inherently hard to parallelize as the Hessian inversion usually involves many sequential steps, making it difficult to leverage large-scale distributed systems to partition computations and memory across multiple nodes. As a result, even though superior convergence performance is observed in strongly convex settings, it still remains unclear how to convey such benefits to large-scale model training in distributed settings.
Due to the prohibitive challenge above, approximation methods are getting increasing atlention that attain second-order information by formulating the Hessian inverse in different ways. These methods, to some extent, open the door for second-order methods to large-scale machine learning. Among them, two lines have proved very promising. The first one arises from the Fisher information matrix I (expectation of the Hessian under a negative log-likelihood loss). Methods such as KFAC (Martens & Grosse 2015 Ba et al. 2017 Pauloski et al., 2020) first approximate I and then simplify matrix inversion by decomposing into small submatrices. However, the considerable overhead for obtaining the empirical Fisher information and its inverse greatly neutralizes its faster per-iteration convergence promises. Another line of the approximation methods directly approximates the Hessian inverse via BFGS update. With pairs of history gradient and parameter changes, BFGS directly approaches the Hessian inverse with no additional costs on computing and storing the Hessian matrix. However, BFGS or its variant L-BFGS (Nocedal 1980) has not proved effi- cient in large-scale stochastic optimization, where convergence instability is commonly observed. Additional operations introduced in recent works stabilize the training but with substantial costs (Mokhtari & Ríbeiro, 2015 Moritz et al., 2016 Gower et al. 2016). For instance, Moritz el al. (2016) uses a separale large batch of inputs to compute the consistent gradient and parameter changes, which inevitably increase the wall-clock time. Morcover, moving to distributed systems, these methods are not amenable to efficient paralleizaaiin due to either direct matrix inverse (e.g., KFAC) or the iterative procedure in BFGS-like methods.
To simultaneously mitigate the challenges of compute and memory costs, and scalability in distributed systems, we propose mL-BFGS, a stable distributed BFGS variant that preserves the convergence promises of second-order methods, with modest compute and memory costs compared to other QN methods. mL- BFGS addresses the aforementioned barriers as follows. First, mI-BFGS introduces an almost cost-free momentum scheme into the BFGS update rule to stabilize the optimization. By using the momentum of the history statistics (parameter and gradient changes) to estimate the Hessian inverse, mL-BFGS smooths out the approximation with no needs of costly variance reduction methods (e.g., separate large batch size to estimate the Hessian). Hence, both stability and compute efficiency are achieved. For efficient parallelization and low memory footprint in distributed systems, mL-BFGS presents a more generic block- wise Hessian approximation with each block consisting of one or multiple layers. During optimization, each node only computes and stores statistics for one block rather than the full Hessian. As a result, mL-BFGS can perform the Hessian inverse and gradient conditioning in a distributed way with marginal communication and synchronization overhead.
Our theoretical analysis shows mL-BFGS effectively suppresses noise in the Hessian approximation and achieves stable convergence. Empirical evaluations show that, on benchmark datasets, CIFAR-10 and Ima- geNet, and models such as ResNet and Vision Transformer, mI-BFGS achieves a faster per-iteration con- vergence compared to SGD and Adam. Furthermore, due to the lightweight momentum based Hessian approximation, mL-BFGS needs a much shorter wall-clock time to achieve the target accuracy compared to SGD, Adam and other QN methods such as KFAC.
In summary, our main contributions are as follows:
1. We develop mL-BFGS, a distributed stochastic QN method for large-scale models, that achieves fast and stable convergence with low computational complexity.
2

---

Published in Transactions on Machine Learning Research (223)
2. We provide theoretical analyses that show mL-BFGS significantly mitigates adverse effects of stochastic noise the Hessian approximation and achieves stable convergence for stochastic optimization problems.
3. We provide complexity analyses that demonstrate mL-BFGS incurs much less complexity than other QN methods, leading to reductions in overall wall-clock training time.
4. Finally, we carry out comprchensive evaluations on various models and datasets that show mI-BFGS empirically delivers faster per-iteration and wall-clock convergence compared to SGD, Adam, and other second-order optimizers.
2 Preliminaries
The risk function X) in Eq I is usually optimized through form of gradient descent as:
0 = 0, H gt:
(2)
where denotes step size (learning rate) at iteration t and is gradient pre-conditioner. In stochastic training. gradients are evaluated on a mini-batch input X, S X, namely 9: = VeL(e, Xt).
If A is an identity matrix, the update above reduced to SGD, whereas F is a diagonal matrix, becomes an adaptive training algorithm such as Adagrad (Duchi et 2011) or Adam (Kingma & Ba 2014). To further improve convergence performance, esp. in ill-conditioned problem (Nesterov) 2003), it is desired to incorporate more second-order information into H as done in quasi-Newton (QN) methods.
A prime challenge in QN methods is the evaluation of A and in particular its inverse. A well-known Broy- den Fletcher Goldfarb Shanno (BFGS) algorithm Fletcher 2013) addresses the challenge by formulating the Hessian inverse as minimization problem:
A Ar f|-
min A
(3)
s.t.
A H.y= yk
Sk: A is symmetric,
where 8): 0 0:- denotes the parameter changes, and = gs - 9h-1 the gradient changes in two consecutive updates By imposing the secand condition during minimization, BFGS gradually attains the curvature information close to the real Hessian. Rodomanov Nesterov 2021 establishes that BFGS converges to the real Hessian by a greedy strategy of choosing (8k: yk).
Knowing AR-1, the current A is obtained via:
H (I PKVIS P -t- -1831((1))))5))) PESuSt.
(4)
where Pk: V/ at
Hence the Hessian inverse is constructed in an iterative manner with no need to compute the Hessian matrix.
To simplify computation in the Hessian-voctor product, Å BFGS is stored the form of a sequence of history vectors {y:} and {si}. The matrix-vector product fl -g. 9: is replaced by a sequence of fast vector- vector products shown in Algorithm 2 (See Appendix A.I). Furthermore, limited-memory BFGS,
L-BFGS (Nocedal 1980) is usually adopted that only uses several latest history vectors when approximating the Hessian inverse.
3 mL-BFGS
mL-BFGS consists of two crucial techniques to improve convergence stability and scalability of using L- BFGS in large-scale models. First, mL-BFFS introduces momentum into the Hessian approximation. The momentum- based design effectively reduces the adverse effects stochastic noise while without resorting to costly noiso-reduction techniques as current solutions (Moritz al. 2016). Sccond, in distributed training,
kk rather than used in the equation as parameter/ /gaddient used might he different from the Eg
3

---

Published in Transactions on Machine Learning Research (023)
mL-BFGS allows block-wise approximation along arbitrary diagonal blocks. mL-BFGS can flexibly assign each computing node block that comprises multiple layers to distribute the workload. We describe the method in detail below.
3.1 Momentum-based Hessian: Reduce Effects of Stochastic Noise
While momentum is widely used in first-order methods, it is rarely explored in the second-order domain. Surprisingly, we find that momentum is also a crucial component for a stable Hessian approximation. Fur- thermore, compared to other noise-reduction methods, the momentum-bant based design is almost cost-free.
In this paper, we apply momentum to past parameters and gradients
0 MM. 3- MM + (1-3)0t-
(5)
g Mg. B -MM.. + 8)g
where b,, 9t denotes parameters and gradients at t-th itcration, 8 is the momentum coefficient.
Following the BFGS update rule in Eq 33 and assuming that A is updated for every T mini-batch iterations, and ge are obtained as
8 Me Mewro
yk = Mg MMNT
(6)
This simple but effective technique works surprisingly well when gradients are noisy. To intuitively show improvements of using momentum over the vanilla I.-BGGS we visualize the stochastic optimization of simple quadratic loss function, L = |10|2. We simulate gradients in stochastic scttings at itcration as = 0r + ne, where nı, denotes the slochaslic noise. We model the noise as i.i.d. Gaussian, namely nu ~N, o²). Figure shows optimization trajectories for SGD, vanilla L-BFGS, L-BFGS with momentum and the exact 2nd-order method given o 0.2. With the same initial point, we observe that L--FGS with momentum is as fast as the exact 2nd-order method, and much faster than SGD. On the other hand, vanilla L-BFGS obviously suffers convergence issues due to noisy yk-
Analysis- We use quadratic function as show- case to analyze important theoretical properties of the momentum-based design:
Now Loes Optimel exact 2nd-order
SGD
vanilla L-BFGS
((0) -(0.)
0o) (0 B(0 0u)
momentum L-BFGS
The following lemmas and theorem establish theo- retical improvement of using momentum over vanilla. L-BFGS. Specifically, Lemma [ shows that momen- tum in Eq 5 significantly reduces stochastic noise in gradient changes yk- Temma 2 further provides a theoretical guarantee that using momentum can still obtain the same Hessian approximation as the vanilla L-BFGS algorithm in the noise-free case.
Finally, Theorem D shows that in stochastic settings, the achievable loss by L-BFGS is lower bounded by value depending on noise variance in yk- Com- bined with Lemma | and 2 it is obvions that the momentum-based approximation achieves a much lower bound compared the vanilla version. Such conclusion is aligned with our empirical observa- tions in Figure
Figure 1: Optimization using SGD, vanilla L-BFGS, L-BFGS with momentum (8 = 0.9) and exact 2nd-order method. Vanilla L-BFGS fails to find the desirable optimization path. However, momentum reduces noise in gradients and significantly stabilizes the optimization.
Lemma 1. Given L(O), and assuming 9: in iteration t contains Gaussian noise with zero mean and Ell[n, 1 < e then the noise variance in Mg. in Eq 55 is reduced to it3
1-9 2 as Furthermore, the
moise variance in ys is reduced to 4(1-0) 2
1+3

---

Published in Transactions on Machine Learning Research (223)
Lemma 2. Considering ((0) above, with no noise involved, Eq 6 obtoins the some Hession approzimation the vonilla I-BFGSS
Theorem 1. Considering L(0) with AT B AT where A 2 >>
0, and noise variance in yk is bounded by 22a. Assuming during the Hessian estimation, we always choose S yk such that (8k. yv) 2 a | with a > 2, then ((O,) 2 (a (a-2) at any iteration
2A
Remark 1. (8k, yn) 2 a |ksk| with 2 ensures positive-definitnnsss in the Hessian approximation.
3.2 Block-wise Approximation: Improve Memory Efficiency in Distributed Training
With the L-BFGS update rule, we need to store sufficient number of history vectors 82, y to obtain a good Hessian approximation. Given model parameters E Rd and supposing M history vectors are needed, it required (22Md) memory space. While it is much smaller than storing the full Hessian (O(d)) with M < d, it can still be critical bottleneck when training large models.
Furthermore, when it comes to distributed training, simply adopting data parallelism (DP) unfortunately does not help relieve memory pressure on each node, as parameters and gradients are duplicated on each node in DP. naive implementation of the L-BFGS update needs to maintain a copy of {Sk. yı h on each node. As result, more nodes do not help reduce per-node memory costs in the Hessian approximation. more efficient design is needed to reduce per-node memory costs when more nodes are available.
To improve memory efficiency, propose approximation methods such as KFAC (Ba et al. 2017), our block-wise approximation method allows each diagonal Hessian block to capture curvature information across multiple layers. Such a design enjoys a more flexible configuration depending on the size of each layer's parameters and each node's capacity. For instance, as shown in Figure 2 if node 1 and 2 have similar memory capacities, and the sizes of layers 1-2 and 3-4 are also close, we can approximate the Hessian blocks with layers 1-2 and 3-4 in node and 2, respectively. Otherwise, we can group layers in different way and ensure node and 2 share similar memory costs.
block-diagonal Hessian approximation. Unlike diagonal
During training, for the Hessian approximation, the i-th node only collects parameters and gradients corre- sponding to its block and computes and stores vectors { vi based on Eq 6 When applying gradient conditioning, each node performs the Hessian-vector product on the corresponding layers, followed by up- dating parameters. And then, all nodes invoke a All-Gather operation to send/ /collect updated parameters to/from other nodes
node
node
Hessian inverse matrix
blockt: layer 1,2
block2: layer 3,4
9:
9
96
1.2
1.2
Cond.
Cond.
y
Update
All- Gather
Update
01.3
03+1
0++1
A distributed block-wise approximation. Rach node can compute the Hessian inverse of a diagonal black comprising one or more layers. Unlike diagonal approximation methods such as KFAC, mL-BFGS allows much more flexible configurations depending on the size of each layer's parameters and each node's capacity.
Figure 2:
Analysis Given p nodes and assuming model's parameters can be evenly distributed among these nodes, then each node only needs O 2M d) memory to store the history vectors. If more nodes are available, per-
5

---

Published in Transactions on Machine Learning Research (023)
node memory costs can be further reduced. On the other hand, it is easy show that the Hessian inverse comprising block-diagonal matrices is positive-definite if each of these blocks is positive-definte, therefore ensuring stable convergence with positive-definite diagonal Hessian blocks.
Lemma 3. In k-th Hessian update, if block i(i 1,.. .p), fg is bounded by 111 fij E'1, then min 5 Hx 3 m
3.3 Put All Together
With the momentum-based design and block-wise approximation, we present the whole algorithm, mL- BFGS, as shown in Alg [ For cach paramcter block at itcration t, we compute momentum of paramoters and gradients, and store them Ma, and Mg. (line 6). For every iterations, we compule s yi, and update the Hessian approximation (line 13-18). It is worth noting that in UpdateHessian, the Hessian inverse is not explicitly computed. Instead, we only push s v to the history vector buffer. If the buffers are full, use SGD to conduct a warmup training as the
will pop the oldest vectors. In the first 27' iverations, Hessian inverse is not available yet (line 8). After the initial warmup training, we will pre-condition gradients gr before applying updates to 0 (line 10-11). At the end of each iteration, All- Alther is called to update parameters on all nodes (line 20).
Algorithm mI-BFGS algorithm (T, M, parameter block {0')=-) Initialize Oo, Ma. 0g Mg = %
2: for 1= 1, max iter do
3: Randomly choose mini-batch input X, E X
4: Perform model forward and backward, compule gradients Ir given X,
5: for each parameter block do
M 3 M + - (3 0; My 9:
BBM 2. +(1 - B) gi
7:
if t 5 2T then
Warmup with SGD: 0'4,1 0: 17: gi else
10:
Pre-condition: A0; = Hij gi (Algorithm
11:
0-n 0: Th. A0;
12:
end
if
13:
if 17T == 0 and T then
14:
k $ MMo M yi M₁ M
15:
g:
16:
Apply damping: D. T y (1 () 8
17:
Update Hessian: Hi = UpdateHesssian( 11 9 M)
18:
end if
19:
end for
All-Gather {0++)]² across all nodes.
20:
21: end for
Hessian damping For non-convex QN optimization, damping is a common technique that ensures the positive definiteness of the Hessian Al-Baali et al. 2014 Al-Baali Grandinetti] 2017 Martens Grosse 2015). Furthermore, even with momentum, stochastic training can still cause undesirable fluctuation in the Hessian approximation. To preserve the positive of the Hessian, we adopt an adaptive damping scheme as
ú: = y: 7) 8i.
(7)

---

Published in Transactions on Machine Learning Research /023)
with 7 obtained
min 1-th, To)
I-/ 2=-1
min
To)
# oH >1
70
otherwise
where *: vs OL and ои are the lower and upper thresholds for restraining eigenvalues in f and To < constant coefficient.
o,"
It is easy to show that tt is bounded between [0,, TH), so that the positive definiteness of the Hessian .f
approximation is preserved during training (See A.2.4 for the proof).
4 Theoretical Guarantees
In this section, we first prove that mL-BFGS achieves linear convergence rate under non-convex settings with proper assumptions. Then, in the second part of this section, we delve into the compute and memory costs in mL-BFGS, and show its benefits in wall-clock convergence compared to other baseline methods: stochastic L-BFGS, and KFAC.
4.1 Convergence Analysis
We assume the risk function C satisfies the following conditions:
AS 1. C(0) is twice continuously differentiable.
AS 2. (0) is A-smooth for i N, A > 0: vo,, 02, (02) - V4(0))I = A ||0, - -0111-
AS 3. ((0) is X-PL: it satisfies Polyak-Lojassewicc (PI) condition for a constant À > 0: || 2 AC
The smooth condition in AS commonly used in analyzing convergence in practical optimization. Tn addition, noting that compared to typical strong convexity assumptions, the PL condition in AS 3 applies to more general setting (Polyak] 1963). Strong convexity implies the PL condition, but not vice versa. In we relax the constraint and only require the gradient variance to be lower bounded.
AS 3
With the assumptions, we present the convergence theorem as follows. Proofs are deferred to Appendix A.2
Theorem 2. Assume AS i hold at each iteration t of mL-BFGa with mini-batch input X, where each sample is randomly sampled from X' with replacement, then the expectation of L(O.) salisfies
Ex, [C(,)) S (C
where 0:-1 01-1=1-144 1-111-1-11 i & and denotes the lower and upper bound of the ft.
By choosing Th-1 such that at-1 1, mL-BFGS converges at linear rate. The convergence rate matches the best rate in stochastic QN optimizations. It is worth mentioning that no convergence rate beyond linear is observed in stochastic I-BFGS optimizations. Theorem 2 is not aimed to push the theoretical convergence limit, Instead, it investigates the effects of momentum in the Hessian and block-diagonal approximation. In addition, mentioned carlier, Thcorem [ also applics convex scttings, as strong convexity implics [[]]()) is lower bounded by L(0) for an appropriate and hence AS holds.
4.2 Complexity Analysis
In this section, we analyse the compute and memory cost of SGD, KFAC Ba et al. 2017): stochastic L-BFGS (we call it sL-BFGS) (Chang et al. 2019) and mL-BFGS. As the main motivation, reducing the complexities of QN methods crucial for their meployment in real large-scale neural network optimization.
Given a model with parameter 0 E Rdd we use Co. and Mp. to represent the compute and memory cost of forward/ forward/backwar backward (Fwd/Bar (wddd pass with a batch size of b 1. Furthermore, Copt denotes the compute cost of model updates (Opt) which consists of gradient reduction, computing and apply the update A0.

---

Published in Transactions on Machine Learning Research (02y)
Table 1 summarizes the total compute and memory cost of SGD, KFAC, sL-BFGS and mL-BFGS in general distributed system with p workers. Compared to SGD, during the forward and backward passes, mL-BFGS needs to additionally compute Ma, MMg, for which the complexity increases linearly with model size (d). The main extra compute mI-BFGS introduces is the Hossian-voctor product, in which we need to M and {yi)== as shown in Alg A.1 The complexity increases linearly with the number iterate over {=:);=1
of history vectors and model size (2Md). However, it only adds average cost of O ( on each worker. Compared to (bCC)s) complexity in forward and backward passes, such costs are relatively marginal.
As comparison, KFAC adds significant additional computations through 1) possible multiple backward passes to updale factors (ydOns with 7 2 1), 2) matrix inversion (E + (4)))) for every T iterations, and 3) Matrix-vector products E4 + 4. )d,). On the other hand, sI-BFGS also resorts to computation- intensive operations including full-batch gradients and separate large batch to estimate the Hessian, which respectively adds amortized costs of ((6C)) and #bn With data parallelism, it requires each worker to locally perform gradient conditioning, which adds another cost of O(2Md) in total.
As for memory usage, mL-BBFGS mainly needs ((2Md) to store history vectors. The amortized costs of each worker are O( (2Md). In practice, M is set to be 10 ~ 20, which ensures that memory usage is managcable in mL.BFFS. sT-BFGS needs (Mdd storage for 8,, yi in total, and amortized cost of by Mn.) for additional backward passes. While KFAC needs 0(2 (1) + (4)) to store sub-matrices and their inverse, where the actual memory footprint hinges on model architectures.
Table 1: Computations and Memory in SGD, KFAC, BF-BFGS and mL-BFGS
SGD
KFAC
sL-BFGS
mL-BFGS
Per-Node Computation
Fwd&&Bwd
O(bC())
O(d+nbCa, + $ Σ 02 (()))) o(d+2 CbnC Old O(d+2Md)
O (CB) O(d + 2Md)
Opt
((d)
Old + 2 Σ (4 4. )d.
Per-Node Memory
Fwd&&Bwd
O(bMa)
Old + bMa) (d) Old + (+)
Old + bMr + tbu Ma) O(d + 2Md)
O - bM
Opt
Old + 2M4
b: per-node batch size. bu: batch size for the Hessian approx. workerss T: Hessian updale period.
de #params i-th layer. 4; #input neurons i-th layer. M: max #history vectors.
Table 2 lists detailed amortized costs of different optimizers on PeeNet-50/ImageNee in distributed system. Compared to KFAC and sI.BBFS, mI-BFGS significantly reduces compute costs in approximating the Hessian and gradient conditioning. Due to the efficient distributed design, memory consumption on each worker is also reduced compared to sL-BFGS.
Table 2: Computation (MACs) and memory costs of different optimizers on ResNeet500/lmaeeNet (b = 64, T' = 20; M = 10;
by = 1024, y = L, p 8). "B" denotes billion, and "M" denotes million.
SGD
KFAC
sL-BFGS
mL-BFGS
Fwd/Bwd
769B
Å Compute Opt. Å Memory
570B
1414B
52M
26M
156B
4B
546M
308M
4B
520M
5 Experiments
We conduct various experiments on computer vision (CV) problems involving datasets such as CIFAR-10, CIFAR-100 and ImageNet. We choose SGD and Adam as the main baselines since widely used in these tasks and maintains the best training performance. We also compare with another quasi-Newton method,
8

---

Published in Transactions on Machine Learning Research (023)
KFAC, in large-scale model training. We separately tune hyperparameters for each optimizer to ensure it achieves the best validation accuracy.
We use single GPU server with 8 Nvidia Quadro RTX 5000 GPUs to simulate a distributed system, where each GPU used as a worker to perform forward and backward passes, and model updates. Furthermore, each worker is also assigned with one Hessian block compute the Hessian inverse and gradient conditioning. The current implementation based on PyTorrch. We set lower and upper thresholds of damping oL, oh be 0.01, 1.5 in all experiments to smooth the Hessian approximation.
5.1 Experiments on CIFAR-10/CIFARE
We first evaluate mL-BFGS on two small-scale problems: CIFAR-10 and CIFAR-100, and demonstrate the convergence advantage of mL-BFGS compared to SGD and Adam. The models used are ResNet-18 and DefT-Tiny Touvron al.] (2021), where DeiT-Tiny an efficient image transformer with layers, 3 atlention heads, and hidden and MLP dimension of 192.
For ResNet18, divide it into 4 blocks such that each block consists of resblocks (He et al. 2016 The lincar layer for classification is packed into the last block. For DeiT-Tiny, due to the small model size, we choose to approxinate the whole Hessian.
Hyperparameters are tuned to achieve the best validation accuracy. Details are provided in Appendix A.4.1
Figure shows training loss. We observe that mL-BFGS achieves a much faster convergence rate compared to SGD and ADAM. Table lists the validation accuracy on CIFAR-10 and CIFAR-100. We note that mL-BFGS also achieves similar accuracy as SGD. On the other hand, in Figure 3 although ADAM has convergence rate close to mL-BFGS, the validation accuracy are much lower in all experiments compared to mL-BFGS. Therefore, we obverse mL-BFGS not only deliver faster convergence, but also achieves good generalization performance.
0.40
08
08
ADAM
mLBFGS
mLBFGS
0.50
0.6
0.6
Loss
0.20
0.4
0.4
0.3
0.05
011
0.1
56.
10c
15K
20k
25K
30K
5<
10K
15k
20K
25k
30C
16.
65
8K
10c
Iterations
Iteraaionns
Iterations
(a) 18eNet-18/CIFAR-II
(b) ResNet-
(c) COIT-TinyCIFAAR-AA
Figure 3: Training loss mL-BFGS, SGD, ADAM ResNet-18 and Deit-Tiny CIFAR-10/100. mL-BFGS delivers faster convergence than SGD and ADAM.
Table 3: Validation accuracy of ResNet-18 and Deit-Tiny on CIFAR-10/100 using SGD, ADAM, and mL-BFGS.
1esNet-18/CIFAR
ResNet-18//CIFA 100
DeIT-Tny/CIFARE 100
SGD 94.1 0.1
ADAM 92.7 0.1
mL-BFGS
SGD 75 0.15
ADAM
mL-BFGS 0.14
SGD 80.5 = 0.2
ADAM 75.34.33
mL-BFGS 79.9 0.2
$3.9 0.1
72.2 = 0.16

---

Published in Transactions on Machine Learning Research (22l)
5.2 Experiments ImageNet
ImageNet has been the gold standard for evaluating the performance of optimizers. It consists of ~1.2M training and ~50K test images, categorized into 1000 classes. We follow the standard data pre-processing procedure, where cach image is first resized to 256 x 256, and randomly cropped to 224 x 224 and flipped horizontally. Each image is then normalized using pre-computed mean and variance.
ResNet-50 When approximating the Hessian, we divide ResNet-50 into 8 blocks such that each block consists of resblocks. Similar to ResNet-18 in CIFAR-10, the linear layer is packed into the last block. Figure shows ileration-wise convergence ResNet-50 using SGD, Adam, KFAC and mI-BFGS. Detailed hyperparameter settings provided in Appendix A.4.2 Compared to Adam and SGD, mL-BFGS enjoys much faster per-itcration convergence. Such fast convergence is also reflected in the validation dataset (Figure 4b Furthermore, it also generalizes well on the validation set, and finally reaches comparable validation accuracy to SGD.
SGD
SGD
SGD
55
KFAC mi BFGS
KFAC mLBFGS
KFAC
Erior (%)
mLBFGS
00
3:44
Emor
500
13
10.00
20:00
2004
Itorations
bpocc
TIme(h)
(a) Training loss
Validation
(c) Wall-clock convergence
Figure 4: Training loss and validation error of mL-BFGS, SGD, Adam and KFAC ImageNet using ResNet-50. mL-BFGS
delivers faster iteration-wist and real-time convergence compared SGD and Adam. We plot the mean and standard with different random seeds
Table Validation accuracy of ResNet-50 on ImageNet using SGD, ADAM, KFAC, and BIGBFGS.
SGD
ADAM
BF-BFGS
KFAC
Acc
Time/epoch 7.8min
Acc 73.95 + 0.1
Time /mee/eooch
Acc 74.6 + 0.13
Time/epoch 7.9min
Acc 74.5 0.1
Time/epoch 17min
74.9 0.11
7.8min
The benefit of mL-BFGS is even more striking in terms of wall-clock time. As listed in Figure 4c and Table due to light compute costs, the per-epoch runtime of mL-BFGS is almost the same as SGD and Adam. On the other hand, for KFAC, while it delivers fast per-iteration convergence compared to SGD and ADAM, the wall-clock performance is significantly diminished by its additional compute costs. The per-epoch runtime is > 2x more than mL.BBGS.
5.3 Ablation Study: The Effects of Momentum and Damping
In this section, we give more insight into the effects of momentum and damping used in mL-BFGS. To this end, ablate two critical components in mL-BFGS: momentum and damping in the Hessian approximation, and then the ablated version to train ResNet-18 CIFAR-10. We focus CIFAR-10 10 since we observed more convergence instability on this dataset compared to others.
Figure 5 shows convergence using the ablated mL-BFGS with only momentum (black), with only damping (purple), and with no momentum or damping (red). Due to stochastic noise, the ablated version of mL-BFGS without patpenntu/dmaapiiii (vanilla L-BFGS) diverges casily the early stages. With momentum (black), the whole optimization significantly stabilized. However, it still fails to converge when there radical change in the loss landscape (for example, when learning rate decays). With damping (purple), the Hessian
10

---

Published in Transactions on Machine Learning Research (22y)
approximation is effectively restrained, especially when such sudden changes in the loss landscape happen. Tt is interesting to observe that while damping prevents divergence, the whole training is still largely affected by stochastic noise. Notable fluctuation in the loss is commonly observed during training. As a comparison, the complete mL-BFGS (blue) effectively addresses these issucs achieving much more stable convergence.
10P
Diverges easily
Abnormal Hessian
Large stochastic noise
Training loss
20-1
Momentum & Damping
10a2
Momentum only
Damping only
No Momentum/Damping
10k
20k
30k
Iterations
Figure Ablation study mI-BFGS on ResNee-18/CIFA C- (batch size: 256). Vanilla L-BFGS (red) diverges easily. Damping-only scheme (purple) cannot effectively suppress stochastic noise. Momentum-only scheme (black) still diverges with radical changes in the Hessian. While damping and momentum combined (blue) achieve a very smooth and stable optimization.
6 Related Works
While SGD is widely used in many machine learning tasks, other forms of optimization have also been investigated extensively in the past years. Among these attempts, designing optimizers with preconditioned gradients is one of the most promising areas.
Adaptive methods such as Adam, AdaGrad, AdaDelta (Kingma & Ba. 2014 Duchi et al. 2011 Zeiler 2012)
construct a diagonal matrix by incorporating knowledge from the past gradients. Such diagonal matrix adaptively adjusts the learning rale for each parameter. For instance, AdaGrad uses a large learning rate for those irrelevant features (small gradients), and small learning for those relevant ones (large gradients). Adams further uses the first and second moment of gradients to adjust the learning rate.
Besides diagonal preconditioning matrices, constructing block diagonal matrices or even full matrices has received increasing attention in recent years. Methods such as Shampoo (Gupta al. 2018) approximate the full matrix version of AdaGrad as block-diagonal matrix to incorporate more curvature information during optimization. Similarly, well-known KFAC method [Martens & Grosse 2015
approximates the Fisher information matrix a block-diagonal matrix. T-BFGS methods Moritz et al. (2016) the other hand directly construct the full Hessian matrix as preconditioner during optimization. These preconditioners have been empirically proved to achieve fast convergence compared to SGD, as well as adaptive mothods. Authors in Goldfarb et al. 2020 further adopt L-BEGS to efficiently compute matrix inversion in KFAC.
Variance reduction in the preconditioning matrix is also crucial to ensure stable optimization. The algorithm in Moritz et al. (2016) adopts separate large batch of data to estimate current curvature. On the other hand, VITE Lucchi et al. (2015) chooses to use pivot parameter together with full- gradients to reduce variance in the Hessian approximation.
7 Conclusion
In this paper, we propose mL-BFGS, a quasi-Newton method that simultancously mitigates computation and convergence instability barriers in second-order methods, as well as scalability issues in distributed systems. By introducing momentum and damping into the Hessian update, mL-BFGS obviates the need for highly costly estimation on the Hessian. Approximation along diagonal blocks further reduces memory and compute costs in distributed systems. Empirical analyses on CV models, such ResNet-50 and Vision Transformer show that mL-BFGS achieves faster convergence, and reaches similar accuracy compared to SGD.
11

---

Published in Transactions on Machine Learning Research (22l)
Acknowledgements
This material is based upon work supported by Defense Advanced Research Projects Agency (DARPA) under Contract FASTNICS HR001120C00S8. The views, opinions, and/or findings expressed are those of the author(s) and should not be interpreted as representing the official views or policies of the Department of Defense or the U.S. Government.
References
Mchiddin Al-Baali and Lucio Grandinctti. Improved damped quasi-newton methods for unconstrained oplimization. Pacific Journal of Oplimizalion (To appear), 2017.
Mchiddin Al-Baali, Lucio Grandinetti, and Ornella Pisacane. Damped techniques for the limited memory bfgs method for large-scale optimization. Journal of Oplimization Theory and Applications, (2)(2):688 699, 2014.
Jimmy Ba, Roger Grosse, and James Martens. Distributed socond-order optimization using kronecker- factored approximations. In International Conference on Leurning Represenlations, 2017.
Richard H Byrd and Jorge Nocedal. A tool for the analysis of quasi-newton methods with application to unconstrained minimization. SIAM Journal on Numerical Analysis, 26(3):727-775 1989.
Daqing Chang, Shiliang Sun, and Changshui Zhang. An accelerated linearly convergent stochastic 1-bfgs algorithm. IEFE transactions on neural networks and learning systems, 30(11):3338 3333-3366, 2019.
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. Practical methods of optimization. John Wiley & Sons, 2013.
Wenbo Gao and Donald Goldfarb. Quasi-newton methods: superlincar convergence without line scarches for self-concordant functions. Oplimization Methods and Software, 34(1):194 217, 2019.
Donald Goldfarb, Yi Ren, and Achraf Bahamon. Practical quasi-newton mothods for training deep neural networks. Xiviv preprinl xrx::22006.0 08877, 2020.
Robert Gower, Donald Goldfarb, and Peter Richtárik. Stochastic block bfgs: Squeezing more curvature out of data. In International Conference on Machine Learning, pp. 1869-1878. PMLR, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
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. ar preprinl
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization.
arXiv:1412.6900 2014.
Aurelien Lucchi, Brian McWillians, and Thomas Hofmann. A variance reduced stochastic newton method. arXiv preprint orx:::5503 08316, 2015.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curva- ture. In International conference on machine learning, pp. 2408 2417. PMLR, 2015.
Aryan Mokhtari and Alejandro Ribciro. Global convergence of online limited memory bfgs. The Journal of Machine Learning Research, 16(1):3151 1-111111111,,,,, 2015.
Philipp Moritz, Robert Nishihara, and Michael Jordan. linearly-convergent stochastic I-bfgs algorithm. In Arlificial Intelligence and Slatistics, pp. 249 258. PMLR, 2016.
12

---

Published in Transactions on Machine Learning Research (223)
Yurii Nesterov. Introductory lectures conver optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Jorge Nocedal. Updating quasi-newtor matrices with limited storage. Mathematics of computation, 35 (5(1) 773-782, 1980.
J. Gregory Pauloski, Zhao Zhang, Lei Huang, Weijia Xu, and lan Foster. Convolutional Neural Network Training with Distributed K-FAC. In Proceedings of the International Conference for High Performance Computing, Nehworking, Storge and Analysis, SC '20 TEFE Press, 2020. ISBN 97817281 (781178119986. doi:
00.5555/3433701.343333
Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal vychislitel noi matematiki i matematicheskot fialki, 3(4):643 653, 1963.
Anton Rodomanov and Yurii Nesterov. Greedy quasi-newton methods with explicit superlinear convergence. SIAM Journal on Optimization, 31(1):785 811, 2021.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference machine learning, pp. 10347-10357. PMLR, 2021.
Matthew D) Zeiler. Adadelta: an adaptive learning rate method. ar preprint or 1xx::12112.500,, 2012.
A
Appendix
The appendix is arranged as follows: In Sec A.1
we include the Hessian-vector vector product used in mL-BFGS. See A.2 and A.3 we provide proofs of lemmas and theorems in the paper. In See A.4 we list hyperparameters used in the experiments.
A.1 Hessian-Vector Product in L-BFGS
Algorithm Hessian- Vecttrr Vector in L-BEGS
Input: gt, [ M {
M
Output: 9t
for ..... .M- - 1 do
s7 Y:
3: end for
4: for = 0. ... M do
$ 1
5: a;
g. = g: 0i 9M-i-1
end for
8: = Ho 9: { (Hg
T UM 17
for =
-.... ... M do
8 v
10:
Pt
11: g. = g: (aas-i-1 - Bi). 8:
12: end for
A.2 Proof of Lemmas and Theorems
A.2.1 Proof of Lemma
Proof. Let Mn, denotes the stochastic noise in Mg then it can be written as:
13

---

Published in Transactions on Machine Learning Research JJuly20233
Mn. BMn. (1 3)m, n' + (1-3) 3) Ei=1 n(
Since n2 for i= 0, ,t are independent, therefore
E [lMM. 12] =B2 ||00| [2]
Ei= ()2 E ||m:11²] -1
Since ||mr11?] bounded by e², therefore
(1-9)² E Mm. 12] 3 Ei= (1-3)11122001) 3)2 =22(3at
T-
(1 (321))
It is obvious that lim E [IM 1P] c.
Furthermore, noise variance in yı is bounded by lim E [|Mms Mm 12]
4(1-8) e. 1+3
A.2.2 Proof of Lemma
Proof. First, consider the case with T = 1, then
8n Mo. MMa -11 () - 0o)
yo Mg. MMg (8 go)
Since 9u B(0 0o) in L-BFGS update, then it is also valid that yo Bsu- Assume for (I: 1)th approximation, y/ = Ma. MM = Bss = B(Mo. -Ma Expand Mo, we can easily get yk = Mg -a-1)(1-33)) B() -o.-1). Then, for kth approximation, we have
B8 = B(Mom - Mo,) = B(1 (0x+1 - Mo.)
Further expand Ma get
Bss+1 (1-33)((0++) 0% + B(0 -Mo_.))
Similarly, since gk+1 ga = B(0++1 - -Os), we have
B8s+1 = - () 9k) + B(Mg - -Mg ,) yk
Consider the case with T> 1, yk can be written as
yv I Ma MMg
T-1
According to the result with T 1, ys can be further written as
y =T= B(Mg M0 -o BET (Me -M Bsk.
Therefore, it is equivalent to use the momentum based scheme to estimate the Hessian without losing accuracy compared to naive L-BFGS.
14

---

Published in Transactions on Machine Learning Research (02l)
A.2.3 Proof of Theorem
Proof. First, we show that (ék:Me) is well bounded, so that most iterates using BFGS goes toward desirable [8..51]
direction according to Byrd & Nocedal 1989
Given (k) 2 | we have (#a.91))
Let va denotes gradient changes with no noise added, then yı 0% + 72 72--1 Since noise variance in y is bounded by 2c?, we have E
|| S e
Then,
E
(ax 201
E
(a).va) (8))
(en ny
(1) (8)
X |
(8k.36)
|
According the two inequalities above, we have E
[s (ok,0k)
a+2
As for the upper bound, first have E [88k yn)] = E K8RK ok + T2g Th 1)] S (Sk, wk) + 2e ||
Because H 3 AI, then (sk, ok) < A || Combined with (Sk, yk) ac | we have
|| 2 E, (sk, | (A |s + 2e)
(8)
With W3 3 we have
(sk, vk) 2 X+A
||sx|| 2 + x+x |
Rearrange this incquality, we have
||
A=A 2
sk
S الد
||
Due to yb = ok + n2 - 72-11 we can convert the equation above to
E Iy 3 sx|l²] [1]]
< ( A
||a||| 26)
Expand the equation, we have
F |1yx 1|2 - (A A) (8K, +((1(t)))) (2 ssall]]]
< {( || + 2(A A) /|///1
4e2
Divide by (8k, 8k) on both sides, we have
(A+A (M).8k)
(M).
According Eq we simplify the expectation as
ykk 3k)
(A + A)(A |s
+ 2e)
(2e + A 2e+AA|sxl111
|sxll)
E
(82, 8x)
|
(8x, 8K)
2(A + A)e
462
2Ae
2)e
A(A + X) +
2 +
AA
|| |I
||s8
||
4Ac
4e2
2
||
4A
4A2
2
A)2
(a 2)²
15

---

Published in Transactions on Machine Learning Research /223)
With the well-bounded Hessian approximation, show that there is an lower bound for the optimization with noise involved.
With (sk: ya) 2 ar || and E [|m:112] S e², we get (8k, yk) = (Sk, wk + r2z+1 ni)) S (8k, wk) + 2e |
Then (86, xk) 2 (a - 2) ||
With the upper bound of H, we have (a - 2) || S SAl ||
With the lower bound of H, we have L(O,) L(0;-1) 2 gh+11 (-8k(())))))))))))))))() 8k) +
||
We focus on near-optimal region, where [(0++1) = and g++1 = 0, then we get ((0)) 2
||
(a-2)
e².
21
Therefore, the achievable loss by L-BFGS with noise is lower bounded by (a 2a-2)²2 2
21²
A.2.4 Proof of Hessian Damping in Eq
Proof. According to Eq (7) ss Ú. sf (Ty. (1=r)s;) (UT + T)ST st. where μ
8 y:. 8. 8;
For S oT, two cases need to be considered:
If T = To: then 1- To: and jut + 1 T 2. T-p
If r 1- then μT + -
Therefore, when μ S o. .. 2 0
For or. < < oH:
We can write μT + 1 - prmo 1 - To- It is easy to show that
ITo + 1- To oz (1 -0)-o)((11 -o)
pTo +1 oo S (1 0H)(1 - To) <0
Therefore, when < < on, o1 8: < < < <on
For 2 σH, similarly two cases might arise:
If To, then -=-l 1-1 2 Tos and HI + TH.
1=-1 then pr + 0H- If == p-1
Therefore, when l 2 σh, x ÿ S .8
In summary, σL s d < σh-
A.2.5 Proof of Lemma
Proof. For any arbitrary vector x o 0, x² fly x can be written as
(x!". -.g..... ".." (i (x'". x Y El xt" f x
where x² sub-voctor corresponding to block i.
Let = ming and E max Ei, therefore, & $ TT. Ål. S E.
16

---

Published in Transactions on Machine Learning Research (023)
A.3 Proof of Theorem 2
To prove Theorem 2 we need several lemmas to bound the Hessian approximation and the gradient variance of the objective function.
A.3.1 Some Lemmas for Theorem 2
Lemma 4. Given damping scheme in Eg 17), at the k-th Hessian update, fl during the optimization is bounded by ŞI 3 Hg 1 EI, where = (M 1) az and 1
Proof. Lower bound: A initialized as Ho
s go 96 90
I. According to the damping scheme A.2.4 there
exists Ho(O) L such that Do Ho 80-
Therefore,
8. 7 en Hoeo 9 90
(s 11 ( (2 20)
1 1T.
% Ho- Home
[0] Ho No- 1177aa) ma)
Then for & 2 1, assuming Ma-1 2 an
17 hold, based on Eq (4),
Mr = (T pEiL Hx-1 -p: +
ot yu
Because (I- (4)))))))) HR-I(I- prDN)s() is positive definite, can bound Hi as: Hy syal mvah of Uo. H..
11 OF
Therefore, lower bound of AI, {= 1 a8
Upper bound: Since H) 2 o,, we can gel
To vo
1
(a / 80)
I 1
96 9o
(o (0 Ho- a6
Similarly, for & 2 1, we assume Hz-1 = 0z
hold.
For the first part in Eq 44, let Q (1 ( Then for V x 0 x"' Q" Q² fl- 1Qx (Qx)x.fls-1-(0c)) 5
oL x0" (QTQ)=
Let P= = 2. BEST 5 BLo
besT+ani then Q=Q=1+P. 0
Because P is rank-1 matrix with eigenvalue -1 and eigenvector Sr, we have xf Q" flg S bx
a6
(I P)+ x < LxT
For the second part in Eq (4,, we directly gel Therefore, aT Ag az
Llel s" ya o HE
3 -I
01.
xT x + o1 xT
x
k+1 og
an x
In mL-BFGS, & is at most M, which is the length of history vector, therefore E = (M + 1) In summary, GN
1 3 Ay 3 (M+1)
Lemma 5. Assume AS holds, then loss function L(O) is at least A-smooth. At iteration t with mani-batch impul S,, where each sample is randomly sampled from X with replacement, gradient (L(,;; S.) satisfies
Es.1IVC(00:S0111 s 2A -())
Proof. Smoothness of L(().
Given AS hold, we have (0,) ve(02)) & 0 - 0z11 For C, have VL(0,) VL(02) 1 IN (4())) V4(02).
17

---

Published in Transactions on Machine Learning Research (023)
Then,
N Ve((1)) V2((22
[VC())) - V2(02))
N
Ve:(01) ve,(02))
N
N
A
||0 02ll
1 021
"TI' indicates triangle inequality. Therefore, C is least A-smooth.
Gradient variance bound.
Es]1)VVC((::8))) Es, V ((0)), (EE
Expand summation and regroup,
Es.IIVC((@::50)1 Es, Ee + E E Ej (V4(0.), (e,(0.))
Take expectation each sample,
Es, [s 3 Em= [VV(())))) E 21=1, Ex, m; (V)) (2,(0.))
Because æ; and æy are independent, the second part can be simplified as,
( + E. Li,,, [[.)) a Σ. ES [L(,,)]
With further simplification, we get
Es b Ez, VI. بيه |V
Given AS 2 we have [[VE,()))1 < 2A (0,) and |V( < 2A L(0.)
Therefore,
1
b- 2AC(O))
Es. Es.
<
bZ Σ E.. 224, (0,)
b
+ b b- 1-1 2.C(0.)
2A
b
b
2A (((,))
A.3.2 Proof of Theorem 2
Proof. Given AS 2 and Lemin 5 L(O.) can be bounded by L(O.) < L(0,-1) + 10 0,-1)
/ ||0, 0,-1 for 0
18

---

Published in Transactions on Machine Learning Research /223)
In mL-BFGS, 0, 0 - V Therefore, we upper bound ((O,)
((0,) SC(0)-1) 7/-1 VA
+ A
[[
SC(0--1) 1 VC(0.-1) V
2 A= 1--1 2
[[(00-1;
Since Hg is independent with S,-1, we take expectation w.r.t S:-1 and St:
Es.(2(0)))6-11 (C(0--1)) 1/--1 V.C.
+ + A=
Es
= -[(6,-1) - 1/--1 V[(0--))
+ n A=
Es. [|
According AS
8
we have
I(0,-1) S [V
According to Lemma 5
we have
Es. 1.s.:llVL(L..... 11 2AL(O.-1)
Therefore, Es, ((( -. A3 After simply regrouping, can get
Es,
Apply total expectation rule w.r.t St, we have
Ess. [C())) S (1 - 12 1 (C(0,-1)
A.4 Hyperparameter Settings
A.4.1 CIFAR-10/CIFAR-100
For ResNet-18 and DeiT, detailed settings shown in Table 5 The batch size is set to be 256 for all optimizers. For the learning rate, use cosine annealing scheduling strategy with minimum learning rate of 0.0001. Maximum epoch is 150 for ResNet--18 and 50 for DeiT.
A.4.2 ImageNet
For ResNet50, detailed settings are shown in Table G The batch size is set be 512 for all optimizers. For the learning rate, we used cosine annealing scheduling strategy with minimum learning rate of 0.0001. Maximum epoch is 100.
19

---

Published in Transactions on Machine Learning Research (02y)
Table 5: Hyperparameters for SCD, Adam, mL-BFGS on CIEAR-I00/IPAAR10000
Optimizer
b
Ir
momentum
wd
T0
B
SGD
256
0.1
0.9
le-4
Adam
256
0.01
0.9/0.999
le-2
mI-BFGS
256
0.1
0.9
2e-4
0.99
0.999
50
10
momentum for the Hossian; Tu: initial damping: T: frequency for updating the Hossian; M: length of history vector (s.v)
Table 6: Hyperparametens for SGD, Adam, KFAC and mL-BFGS of ResNet50 on ImageNet
Optimizer
b
lr
momentum
wd
To
B
M
SGD
512
0.1
0.9
le-4
Adam
512
0.07
0.9/0.999
3e-2
KFAC
512
0.06
0.9
3c-2
50
mI-BFGS
256
0.1
0.9
2e-4
0.99
0.999
50
10
f: momentum for the Hessian; initial damping; T: frequency for updating the Hesslan; M: length of history vector (8. y)
30
1 Rlock
Block
2 Blocks
Blocks
2.5
4 Blocks
Blocks
2.0
(rror(%)
70
Loss
10
03
23
100
Iterations
Epoch
(a) Training loss
(b) Validation accuracy
Figure 6: Training loss and validation accuracy using mL-BFGS on ResNet-8//IIAAR1100 with different block sizes.
A.5
Ablation Study: Impact of Granularity of Block-wise Approximation
We study the impact of granularity of block-wise approximation in this section. We train ResNet-18 on CIFAR-100 using mL-BFGS with a different number of Hessian blocks: 1, 2, 4 blocks. mL-BFGS with one block approximates the whole Hessian matrix, while mL-BFGS with two blocks approximates two diagonal Hessian blocks with each one consisting of reshlocks in ResNet-18. At last, mI-BFGS with four blocks is the default setting in our CIFAR-106 experiment, where each block consists of 2 resblocks. As shown in Figure 6 increase the size of the blocks, mL-BFGS converges faster and achieve better generaliza- tion performance. This observation aligns with the argument that mL-BFGS estimate more accurate curvature information by approxinating large Hessian block, improving training performance.
A.6 Justification of PL-condition in Real Neural Networks
Figure 7 shows gradient norm and loss when training ResNet-18 on CIFAR-100. We can observe that with appropriate parameter X, the Pl-condition can be easily satisfied.
20

---

Published in Transactions on Machine Learning Research /02))
Loss
Gradlent
Gradient Norm
Iterations
Figure 7: The relation gradient norm and loss during training (ResNet-11 on CIFAR-100)
A.7 Other Second-order Methods CIFAR-10/100
08
SGD
322
ADAM
07-
ADAM
mLDFGS
mLBFGS
530
KFAC
06
KFAC
025
65
LOSS
Iterations
(a) Training loss CIFAR-10
(b) Training loss CIFAR-100
Figure 8: Training loss using mL-BFGS, KFAC, Adam and SGD on CIFAR-10/100.
Figure 8 show training loss of using mL-BFGS, KFAC, Adam, and SGD on CIFAR-10/100. Compared
to the second-order method KFAC, mL-BFGS converges faster as it can estimate more accurate curvature information by approximating large Hessian blocks. On the validation dataset, mL-BFGS achieves higher accuracy than KFAC (KFAC CIFAR-10: 93.1 0.1, CIFAR-100: 73.2 0.13).
21


Reviews:
[{"review_id": "KLPhvKsdbD", "reviewer": "Reviewer_xAnx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The proposed method is novel and new (to my knowledge). And it is an important step for the traditional L-BFGS algorithm to scale.\n- The theoretical contribution is sufficient, the authors provide analysis to both convergence analysis and memory usage.\n\nWeaknesses\n\n- In the original [ResNet paper](https://arxiv.org/abs/1512.03385), momentum-based SGD with decaying learning rate is used, in order to reach target accuracy for both CIFAR and ImageNet. I'm not sure if the current learning rate schedule (cosine annealing scheduling strategy) without momentum for SGD is comparable to that, and why authors choose this schedule.\n- The experiments are small-scale. Despite DeiT-Tiny is a transformer model, its size is much smaller than ResNet-18, is there a reason why the authors choose DeiT-Tiny instead of DeiT-Base or DeiT-Small? Is it because of the cost for memory consumption?\n- Last but not least, I wonder if there's any particular reason second-order optimizer should be preferred than the first-order optimizer for model training, especially distributed training in general?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed mL-BFGS, a distributed stochastic Quasi-Newton for large-scale model training. mL-BFGS contains two main designs: (1) momentum for Hessian approximation, which reduces the stochastic noise; (2) a block-wise approximation method that allows distributed storage for model layers and thus enables efficient distributed training. The authors give convergence proof and complexity analysis on the memory. Empirically, comparison between mL-BFGS and other algorithms (SGD, Adam, etc.) is done on CIFAR and ImageNet.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please consider clarifying Weakness 1 and 2.\n- Please consider adding some text to answer Weakness 3.", "broader_impact_concerns": "There is no foreseeable ethical implication.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "iu4atpl3wl", "reviewer": "Reviewer_fNCc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. First of all, there is no explicit mention of what \"m\" stands for in mL-BFGS to begin with.  I think it is for \"momentum\" but am not sure.  It would be helpful to make it clear.\n\n2. I wonder what the convergence rate is without the assumption of the PL-condition.  Usually the convergence rate of SGD or its variants is given as  $O(\\frac{1}{\\sqrt{T}})$ under the assumptions on Lipschitz continuous gradients, its unbiasedness and bounded variance for non-convex objective functions.  It would be helpful to show the convergence rate of mL-BFGS under the same assumptions. PL-condition seems a bit strong to me.\n\n3. In Eq.5, do $\\theta$ and $g$ have to share the same momentum parameters $\\beta$ ?\n\n4. The block-wise approximation of the Hessian matrix using model parallelization is a very popular technique in the ML community and other Quasi-Newton methods can also readily use it.  I am not sure this can be considered a unique part of the design of mL-BFGS.  Also,  it would be helpful to analyze the impact of granularity of block-wise approximation to the performance.\n\n5. I am bit confused by the loss curves in Fig 3 and accuracies in Tab. 3 for CIFAR-10/100.  Do they show that in order to get the best accuracy, mL-BFGS has to run as many iterations as SGD and Adam?  If the required numbers of iterations are the same for the best accuracy, then what is the advantage of using mL-BFGS since it requires more computations and memory?  I would like to see the iterations used to get the accuracies reported in Tab. 3 under SGD, Adam and mL-BFGS.\n\n6.  I have a concern that mL-BFGS may not be large-batch friendly.  Large scale distributed optimization with large number of learners (say, 64 or 128) relies on large batch size. Given the extra memory used by mL-BFGS, the local affordable batch size on each learner may be smaller than SGD, if not much smaller.\n\n\n7.  I would suggest the authors also verify the effectiveness of  mL-BFGS on models using RNN/LSTM architectures.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a momentum-based L-BFGS algorithm, mL-BFGS,  for large scale distributed optimization of deep neural networks.   There are two major algorithmic contributions. One is introducing momentum to the stochastic L-BFGS and also with damping.  The other is model parallelization across learners in the distributed setting to reduce the computational and memory costs, which amounts to a block-wise Hessian approximation. mL-BFGS is a light weight second-order stochastic optimizer.  It takes into account the curvature information of the optimization landscape and therefore is supposed to be faster in convergence.  The authors give the convergence rate, analyze its complexity and show its effectiveness by comparing to existing techniques on various datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to my questions and concerns above.", "broader_impact_concerns": "I don't have concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EJSVAbxBWw", "reviewer": "Reviewer_CWvY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n-\tThe motivation is clear.\n\n-\tThe proposed method appears to be technically sound\n\n-\tThe paper is well-written and easy to follow\n\n-\tThe paper provides both theoretical and empirical analysis\n\nWeaknesses\n\n-\tThe meaning of the theoretical results is a bit unclear.\n\n-\tMissing closely related baselines in experiments.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new second-order optimization method for improving training stability and reducing the computation and memory costs. Specifically, based on L-BFGS, the authors introduce momentum and damping to reduce stochastic noise and stabilize Hessian estimation. They further introduce a block-wise approximation of Hessian to improve memory efficiency in distributed training. Experiments on CIFAR and ImageNet demonstrate the effectiveness of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "-\tWhen comparing to the standard L-BFGS, the authors compare convergence lower bounds (Theorem 1). Can the authors explain why considering lower bounds rather than upper bounds, as upper bounds seem more commonly used in the literature?\n\n-\tIt is unclear to me how Theorem 2 shows the effect of momentum in the Hessian and block-diagonal approximation.\n\n-\tIn the CIFAR experiments, it will be better to add some second-order baselines. Besides, since the proposed method is a variant of L-BFGS, it seems reasonable to add it as a baseline.\n", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is reviewed by three expert reviewers and all three agree that the paper can be accepted. That said, all three reviewers are "leaning accept" as such, this makes it hard to accept the paper as is. I kindly ask the authors to go over all reviewer concerns one final time and address any remaining issues that remains. After this, I believe the paper will be in great shape for publication at TMLR.
Decision: Accept (poster)