Original: SLIM-QN: A S TOCHASTIC , L IGHT , M OMENTUMIZED Q UASI -N EWTON O PTIMIZER FOR D EEP N EURAL N ETWORKS
Revised: SLIM-QN: A S TOCHASTIC , L IGHT , M OMENTUMIZED Q UASI -N EWTON O PTIMIZER FOR D EEP N EURAL N ETWORKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose SLIM-QN, a light stochastic quasi-Newton optimizer for training large-scale deep neural networks (DNNs).
Revised: We propose SLIM-QN, a light stochastic quasi-Newton optimizer for training large-scale deep neural networks (DNNs).

Original: SLIM-QN addresses two key barriers in existing second-order methods for large-scale DNNs: 1) the high computational cost of obtaining the Hessian matrix and its inverse in every iteration (e.g. KFAC); 2) convergence instability due to stochastic training (e.g. L-BFGS).
Revised: SLIM-QN addresses two key barriers in existing second-order methods for large-scale DNNs: 1) the high computational cost of obtaining the Hessian matrix and its inverse in every iteration (e.g. KFAC); 2) convergence instability due to stochastic training (e.g. L-BFGS).

Original: To tackle the ﬁrst challenge, SLIM-QN  directly approximates the Hessian inverse using past parameters and gradients, without explicitly constructing the Hessian matrix and then computing its inverse.
Revised: To tackle the ﬁrst challenge, SLIM-QN uses the BFGS update rule that directly approximates the Hessian inverse using past parameters and gradients, without explicitly constructing the Hessian matrix and then computing its inverse.

Original: To achieve stable convergence, SLIM-QN introduces momentum in Hessian updates together with an adaptive damping mechanism.
Revised: To achieve stable convergence, SLIM-QN introduces momentum in Hessian updates together with an adaptive damping mechanism.

Original: We provide rigorous theoretical results on the convergence of SLIM-QN in a stochastic setting.
Revised: We provide rigorous theoretical results on the convergence of SLIM-QN in a stochastic setting.

Original: We also demonstrate that SLIM-QN has much less compute and memory overhead compared to existing second-order methods.
Revised: We also demonstrate that SLIM-QN has much less compute and memory overhead compared to existing second-order methods.

Original: To better understand the limitations and beneﬁts of SLIM-QN, we evaluate its performance on various datasets and network architectures.
Revised: To better understand the limitations and beneﬁts of SLIM-QN, we evaluate its performance on various datasets and network architectures.

Original: For instance on large datasets such as ImageNet, we show that SLIM-QN achieves near optimal accuracy 1 . 5 × faster when compared with SGD ( 1 . 36 × faster in wall-clock time) using the same compute resources.
Revised: For instance on large datasets such as ImageNet, we show that SLIM-QN achieves near optimal accuracy 1 . 5 × faster when compared with SGD ( 1 . 36 × faster in wall-clock time) using the same compute resources.

Original: We also show that SLIM-QN can readily be applied to other contemporary non-convolutional architectures such as Transformers.
Revised: We also show that SLIM-QN can readily be applied to other contemporary non-convolutional architectures such as Transformers.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Second-order methods have been extensively investigated in the classical convex optimization literature.
Revised: Second-order methods have been extensively investigated in the classical convex optimization literature.

Original: Indeed variants such as quasi-Newton methods are known to deliver faster convergence than gradient descent (GD) and achieve the best overall run time for many tasks (Gao & Goldfarb, 2019;
Revised: Indeed variants such as quasi-Newton methods are known to deliver faster convergence than gradient descent (GD) and achieve the best overall run time for many tasks (Gao & Goldfarb, 2019;

Original: Rodomanov & Nesterov, 2021).
Revised: Rodomanov & Nesterov, 2021).

Original: However, the Achilles heel of second-order methods that has impeded their wide adoption for large-scale machine learning problems is their substantial compute and memory cost, rendering them less favorable than popular stochastic ﬁrst-order methods such as
Revised: However, the Achilles heel of second-order methods that has impeded their wide adoption for large-scale machine learning problems is their substantial compute and memory cost, rendering them less favorable than popular stochastic ﬁrst-order methods such as

Original: SGD (Saad, 1998) and its variants (Duchi et al., 2011; Kingma & Ba, 2014).
Revised: SGD (Saad, 1998) and its variants (Duchi et al., 2011; Kingma & Ba, 2014).

Original: The aforementioned barriers stem from computing second-order information (loss Hessian w.r.t model parameters and its inverse), which typically dominates run-time, especially for large-scale models such as ResNet (He et al., 2016) and
Revised: The aforementioned barriers stem from computing second-order information (loss Hessian w.r.t model parameters and its inverse), which typically dominates run-time, especially for large-scale models such as ResNet (He et al., 2016) and

Original: Vision Transformer (Dosovitskiy et al., 2020a), on large datasets such as ImageNet (Deng et al., 2009).
Revised: Vision Transformer (Dosovitskiy et al., 2020a), on large datasets such as ImageNet (Deng et al., 2009).

Original: To mitigate these issues, approximation methods have been developed to either approximate Fisher information matrix (expected Hessian matrix under negative log-likelihood loss)
Revised: To mitigate these issues, approximation methods have been developed to either approximate Fisher information matrix (expected Hessian matrix under negative log-likelihood loss)

Original: (Martens & Grosse, 2015; Ba et al., 2016;
Revised: (Martens & Grosse, 2015; Ba et al., 2016;

Original: Pauloski et al., 2020) or directly approximate the Hessian inverse (Fletcher, 2013; Mokhtari & Ribeiro, 2015; Moritz et al., 2016; Gower et al., 2016; Gao & Goldfarb, 2018).
Revised: Pauloski et al., 2020) or directly approximate the Hessian inverse (Fletcher, 2013; Mokhtari & Ribeiro, 2015; Moritz et al., 2016; Gower et al., 2016; Gao & Goldfarb, 2018).

Original: A prominent example of the ﬁrst category is
Revised: A prominent example of the ﬁrst category is

Original: KFAC (Martens & Grosse, 2015; Ba et al., 2016) which utilizes a gradient conditioner based on
Revised: KFAC (Martens & Grosse, 2015; Ba et al., 2016) which utilizes a gradient conditioner based on

Original: Fisher information and also approximates the matrix as Kronecker product of small sub-matrices, therefore simpliﬁes matrix inversion.
Revised: Fisher information and also approximates the matrix as Kronecker product of small sub-matrices, therefore simpliﬁes matrix inversion.

Original: An example from the second category is L-BFGS (Nocedal, 1980; Liu & Nocedal, 1989) which aims to directly approximate the Hessian inverse by iterating over the past parameter and gradient changes, without explicitly constructing the Hessian matrix itself.
Revised: An example from the second category is L-BFGS (Nocedal, 1980; Liu & Nocedal, 1989) which aims to directly approximate the Hessian inverse by iterating over the past parameter and gradient changes, without explicitly constructing the Hessian matrix itself.

Original: These methods, while to some extent alleviate the compute and memory costs of second order methods, still experience performance or convergence issues when used for training on large-scale models.
Revised: These methods, while to some extent alleviate the compute and memory costs of second order methods, still experience performance or convergence issues when used for training on large-scale models.

Original: For instance, even though KFAC has faster per-iteration convergence compared to SGD, this gain is often signiﬁcantly neutralized by the need to run backward passes multiple times to estimate
Revised: For instance, even though KFAC has faster per-iteration convergence compared to SGD, this gain is often signiﬁcantly neutralized by the need to run backward passes multiple times to estimate

Original: Fisher information in each mini-batch iteration and then perform costly matrix inversion (Pauloski et al., 2020).
Revised: Fisher information in each mini-batch iteration and then perform costly matrix inversion (Pauloski et al., 2020).

Original: Similarly, stochastic BFGS variants (Mokhtari & Ribeiro, 2015; Moritz et al., 2016;
Revised: Similarly, stochastic BFGS variants (Mokhtari & Ribeiro, 2015; Moritz et al., 2016;

Original: Gower et al., 2016) have yet to be proven efﬁcient in training large-scale models such as ResNet in practice.
Revised: Gower et al., 2016) have yet to be proven efﬁcient in training large-scale models such as ResNet in practice.

Original: A key challenge is that computationally intensive techniques are required to address convergence instability issues in this literature (Mokhtari & Ribeiro, 2015; Moritz et al., 2016), which unfortunately offset the compute beneﬁts.
Revised: A key challenge is that computationally intensive techniques are required to address convergence instability issues in this literature (Mokhtari & Ribeiro, 2015; Moritz et al., 2016; Chang et al., 2019), which unfortunately offset the compute beneﬁts.

Original: More recently, authors in (Goldfarb et al., 2020) propose to lessen the computation burden of matrix inversion in KFAC via BFGS-like updates with promising results on simple architectures.
Revised: More recently, authors in (Goldfarb et al., 2020) propose to lessen the computation burden of matrix inversion in KFAC via BFGS-like updates with promising results on simple architectures.

Original: However the efﬁcacy of this approach is yet to be demonstrated on practical DNNs and large-scale datasets.
Revised: However the efﬁcacy of this approach is yet to be demonstrated on practical DNNs and large-scale datasets.

Original: To simultaneously mitigate the computation and instability barriers in second-order methods, we propose SLIM-QN, a stochastic light stable BFGS-like method that achieves convergence advantages of second-order methods, while only using modest compute and memory cost compared to other techniques such as KFAC.
Revised: To simultaneously mitigate the computation and instability barriers in second-order methods, we propose SLIM-QN, a stochastic light stable BFGS-like method that achieves convergence advantages of second-order methods, while only using modest compute and memory cost compared to other techniques such as KFAC.

Original: SLIM-QN addresses the barriers in second-order methods in two ways.
Revised: SLIM-QN addresses the barriers in second-order methods in two ways.

Original: To reduce compute cost while maintaining fast convergence of second-order methods, SLIM-QN introduces momentum into the Hessian update.
Revised: To reduce compute cost while maintaining fast convergence of second-order methods, SLIM-QN introduces momentum into the Hessian update.

Original: By utilizing momentum on past parameter and gradient changes, SLIM-QN smooths out the Hessian approximation without incurring costly variance reduction methods (e.g. a separate large batch size to estimate the Hessian).
Revised: By utilizing momentum on past parameter and gradient changes, SLIM-QN smooths out the Hessian approximation without incurring costly variance reduction methods (e.g. a separate large batch size to estimate the Hessian).

Original: Furthermore, to ensure stable convergence, SLIM-QN uses an adaptive damping mechanism to adjust gradient changes so as to guarantee positive deﬁniteness of the approximated Hessian inverse in stochastic settings.
Revised: Furthermore, to ensure stable convergence, SLIM-QN uses an adaptive damping mechanism to adjust gradient changes so as to guarantee positive deﬁniteness of the approximated Hessian inverse in stochastic settings.

Original: This adaptive damping scheme effectively restrains abnormal eigenvalues in the Hessian inverse and steers the optimization trajectory towards desirable directions.
Revised: This adaptive damping scheme effectively restrains abnormal eigenvalues in the Hessian inverse and steers the optimization trajectory towards desirable directions.

Original: SLIM-QN delivers faster convergence compared to SGD in various models and datasets.
Revised: SLIM- QN delivers faster convergence compared to SGD in various models and datasets.

Original: The convergence advantage is even more striking on large-scale models and datasets.
Revised: The convergence advantage is even more striking on large-scale models and datasets.

Original: Furthermore, due to its simplicity, SLIM-QN enjoys much better wall-clock convergence gains than other second-order methods such as KFAC.
Revised: Furthermore, due to its simplicity, SLIM-QN enjoys much better wall-clock convergence gains than other second-order methods such as KFAC.

Original: In summary, our main contributions are as follows:
Revised: In summary, our main contributions are as follows:

Original: We develop SLIM-QN, a stochastic quasi-Newton algorithm targeting large-scale models, that achieves both fast and stable convergence and low computation complexity, via introducing momentum and damping into the Hessian updates.
Revised: We develop SLIM-QN, a stochastic quasi-Newton algorithm targeting large-scale models, that achieves both fast and stable convergence and low computation complexity, via introducing momentum and damping into the Hessian updates.

Original: We provide a rigorous analysis for SLIM-QN showing that this algorithm converges at a linear rate for stochastic optimization problems.
Revised: We provide a rigorous analysis for SLIM-QN showing that this algorithm converges at a linear rate for stochastic optimization problems.

Original: We provide complexity analysis that demonstrates that SLIM-QN is lighter than other second-order methods leading to reductions in overall wall-clock training time.
Revised: We provide complexity analysis that demonstrates that SLIM-QN is lighter than other second-order methods leading to reductions in overall wall-clock training time.

Original: Finally, we carry out comprehensive evaluations on various models and datasets that show that SLIM-QN delivers faster convergence compared to SGD, especially for large datasets such as ImageNet.
Revised: Finally, we carry out comprehensive evaluations on various models and datasets that show that SLIM-QN delivers faster convergence compared to SGD, especially for large datasets such as ImageNet.

Original: For instance, to reach near optimal accuracy when training ResNet-50 on ImageNet, SLIM-QN is 1 . 5 × than SGD ( 1 . 36 × faster in wall clock time).
Revised: For instance, to reach near optimal accuracy when training ResNet-50 on ImageNet, SLIM-QN is 1 . 5 × than SGD ( 1 . 36 × faster in wall clock time).

Original: Furthermore, when training Vision Transformer models, SLIM-QN also achieves faster convergence and higher accuracy.
Revised: Furthermore, when training Vision Transformer models, SLIM-QN also achieves faster convergence and higher accuracy.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this paper we consider a typical empirical loss minimization problem of the form min θ L ( θ , X ) :
Revised: In this paper we consider a typical empirical loss minimization problem of the form min θ L ( θ , X ) :

Original: = 1 N (cid:80) Ni =1 (cid:96) ( θ , x i ) , where θ denotes the parameters of the model to be optimized, and X = { x i } Ni = are the training data where x i consists of both features and labels.
Revised: = 1 N (cid:80) Ni =1 (cid:96) ( θ , x i ) , where θ denotes the parameters of the model to be optimized, and X = { x i } Ni = are the training data where x i consists of both features and labels.

Original: The loss function is typically minimized through some variant of gradient descent, that uses the local gradients g t = ∇ θ L ( θ t ) directly to update the model parameters via iterates of the form 
Revised: The loss function is typically minimized through some variant of gradient descent, that uses the local gradients g t = ∇ θ L ( θ t ) directly to update the model parameters via iterates of the form [Equation]

Original: [Equation]
Revised: [Equation]

Original: (1) where η t denotes the step size (learning rate) at iteration t .
Revised: ) where η t denotes the step size (learning rate) at iteration t .

Original: However, such GD updates are typically slow especially for ill-conditioned problems (Nesterov, 2003).
Revised: However, such GD updates are typically slow especially for ill-conditioned problems (Nesterov, 2003).

Original: To speed up the convergence, often second-order methods are used.
Revised: To speed up the convergence, often second-order methods are used.

Original: In particular, Quasi-Newton (QN) methods ﬁnd an approximate
Revised: In particular, Quasi-Newton (QN) methods ﬁnd an approximate

Original: Hessian inverse ˆ H − 1 to pre-condition the gradient vector and apply the following update to minimize the loss:
Revised: Hessian inverse ˆ H − 1 to pre-condition the gradient vector and apply the following update to minimize the loss:

Original: [Equation]
Revised: [Equation]

Original: In stochastic training, the gradient vector is evaluated on a mini-batch input S t ⊆
Revised: In stochastic training, the gradient vector is evaluated on a mini-batch input S t ⊆

Original: X , namely g t = ∇ θ L ( θ t , S t ) .
Revised: X , namely g t = ∇ θ L ( θ t , S t ) .

Original: If ˆ H − 1 is the identity matrix, the update above reduces to SGD, whereas if ˆ H −is a diagonal matrix, it reduces to adaptive training algorithms such as Adagrad (Duchi et al., 2011) or
Revised: If ˆ H − 1 is the identity matrix, the update above reduces to SGD, whereas if ˆ H −is a diagonal matrix, it reduces to adaptive training algorithms such as Adagrad (Duchi et al., 2011) or

Original: Adam (Kingma & Ba, 2014).
Revised: Adam (Kingma & Ba, 2014).

Original: However, to incorporate more curvature information in the optimization process, it often requires approximating ˆ H − 1 with a full symmetric matrix.
Revised: However, to incorporate more curvature information in the optimization process, it often requires approximating ˆ H − 1 with a full symmetric matrix.

Original: A prime challenge in QN methods is the evaluation of ˆ H and in particular its inverse.
Revised: A prime challenge in QN methods is the evaluation of ˆ H and in particular its inverse.

Original: To address this challenge the well-known Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm has been proposed.
Revised: To address this challenge the well-known Broyden–Fletcher–Goldfarb–Shanno (BFGS) algorithm has been proposed.

Original: BFGS approaches the Hessian inverse as a minimization problem:
Revised: BFGS approaches the Hessian inverse as a minimization problem:

Original: [Equation]
Revised: [Equation]

Original: where s k = θ k − θ k − 1 denotes the parameter changes, and
Revised: where s k = θ k − θ k − 1 denotes the parameter changes, and

Original: y k = g k − g k − 1 the gradient changes in two consecutive updates 1 .
Revised: y k = g k − g k − 1 the gradient changes in two consecutive updates 1 .

Original: Knowing ˆ H − 1 k − 1 from the previous update, the current ˆ H − 1 is obtained via:
Revised: Knowing ˆ H − 1 k − 1 from the previous update, the current ˆ H − 1 is obtained via:

Original: UpdateHessian:
Revised: UpdateHessian:

Original: [Equation]
Revised: [Equation]

Original: where ρ k = 1 y Tk s k .
Revised: where ρ k = 1 y Tk s k .

Original: Therefore, ˆ H − 1 is constructed in an iterative manner without explicitly computing the Hessian itself.
Revised: Therefore, ˆ H − 1 is constructed in an iterative manner without explicitly computing the Hessian itself.

Original: Given such an update rule, methods such as Greedy BFGS (Rodomanov & Nesterov, 2021) show ˆ H − 1 can converge to the real Hessian at a linear rate.
Revised: Given such an update rule, methods such as Greedy BFGS (Rodomanov & Nesterov, 2021) show ˆ H − 1 can converge to the real Hessian at a linear rate.

Original: In real-world problems, θ usually consists of millions of parameters.
Revised: In real-world problems, θ usually consists of millions of parameters.

Original: As a result, it is infeasible to store the whole
Revised: As a result, it is infeasible to store the whole

Original: ˆ H − 1 k matrix with O ( | θ | 2 ) memory cost.
Revised: ˆ H − 1 k matrix with O ( | θ | 2 ) memory cost.

Original: To reduce memory footprint and simplify computation, ˆ H − 1 in BFGS is stored in the form of a sequence of history vectors { y i } and { s i } .
Revised: To reduce memory footprint and simplify computation, ˆ H − 1 in BFGS is stored in the form of a sequence of history vectors { y i } and { s i } .

Original: By exploiting the Hessian update formula in equation 2, the matrix-vector product ˆ H − 1 k · g t necessary to pre-condition the gradient can be replaced by a sequence of fast vector-vector products as shown in
Revised: By exploiting the Hessian update formula in equation 2, the matrix-vector product ˆ H − 1 k · g t necessary to pre-condition the gradient can be replaced by a sequence of fast vector-vector products as shown in

Original: Algorithm 1.
Revised: Algorithm 1.

Original: Furthermore, to limit memory and compute costs, a limited-memory version of BFGS,
Revised: Furthermore, to limit memory and compute costs, a limited-memory version of BFGS,

Original: L-BFGS (Nocedal, 1980) is proposed that only uses the latest M history vectors when approximating the Hessian inverse.
Revised: L-BFGS (Nocedal, 1980) is proposed that only uses the latest M history vectors when approximating the Hessian inverse.

Original: Input: g t , { y i } Mi =1 , { s i } Mi =

Revised: Input: g t , { y i } Mi =1 , { s i } Mi =

Original: Output: g t 1: for i = 0 , · · · , M − 1 do 2: ρ i =
Revised: Output: g t 1: for i = 0 , · · · , M − 1 do 2: ρ i =

Original: s
Revised: s

Original: Ti · y i 3: for i = 0 , · · · , M − 1 do 4: α i =
Revised: Ti · y i 3: for i = 0 , · · · , M − 1 do 4: α i =

Original: s T M −
Revised: s T M −

Original: i − 1 g t ρ M − i −5: g t =
Revised: i − 1 g t ρ M − i −5: g t =

Original: g t − α i · y M − i −6:
Revised: g t − α i · y M − i −6:

Original: g t =
Revised: g t =

Original: ˆ H − 1 0 · g t (cid:46)
Revised: ˆ H − 1 0 · g t (cid:46)

Original: ˆ H − 1= s TM − 1 y M − 1 y TM − 1 y M − 1 · I 7: for i = 0 , · · · , M − 1 do 8: β i =
Revised: ˆ H − 1= s TM − 1 y M − 1 y TM − 1 y M − 1 · I 7: for i = 0 , · · · , M − 1 do 8: β i =

Original: y Ti g t ρ
Revised: y Ti g t ρ

Original: i 9: g t =
Revised: i 9: g t =

Original: g t + ( α M − i − 1 − β i ) · s i
Revised: g t + ( α M − i − 1 − β i ) · s i

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original:  SLIM-QN
Revised:  SLIM-QN

Original: While BFGS achieves faster convergence compared to GD in full-batch training, it still suffers convergence instability in the stochastic setting, especially for large-scale DNNs.
Revised: While BFGS achieves faster convergence compared to GD in full-batch training, it still suffers convergence instability in the stochastic setting, especially for large-scale DNNs.

Original: Speciﬁcally, Naive
Revised: Speciﬁcally, Naive

Original: BFGS, which simply uses parameter θ t and gradients g t at each iteration to calculate s k and y k , suffers from severe instability due to stochastic noise introduced by mini-batch training (See the[Table] [Table]
Revised: BFGS, which simply uses parameter θ t and gradients g t at each iteration to calculate s k and y k , suffers from severe instability due to stochastic noise introduced by mini-batch training (See the

Original: [Table]
Revised: [Table]

Original: ablation study in Sec. 5.2).
Revised: ablation study in Sec. 5.2).

Original: To stabilize the optimization, a common solution is to use a separate large batch of data when estimating s k and
Revised: To stabilize the optimization, a common solution is to use a separate large batch of data when estimating s k and

Original: y k
Revised: y k

Original: (Moritz et al., 2016) in order to reduce stochastic noise.
Revised: (Moritz et al., 2016; Chang et al., 2019) in order to reduce stochastic noise.

Original: However, this dramatically increases the computation cost and negates performance gains in wall-clock time.
Revised: However, this dramatically increases the computation cost and negates performance gains in wall-clock time.

Original: In order to reduce variance in ˆ H − 1 without incurring such expensive additional computation as in (Moritz et al., 2016), we propose a novel quasi-Newton method, SLIM-QN which introduces momentum and damping to the Hessian updates.
Revised: In order to reduce variance in ˆ H − 1 without incurring such expensive additional computation as in (Moritz et al., 2016; Chang et al., 2019), we propose a novel quasi-Newton method, SLIM-QN which introduces momentum and damping to the Hessian updates.

Original: In SLIM-QN, we ﬁrst obtain the momentum of θ t and g t , from which we subsequently derive parameter and gradient changes s k and y k .
Revised: In SLIM-QN, we ﬁrst obtain the momentum of θ t and g t , from which we subsequently derive parameter and gradient changes s k and y k .

Original: To guarantee positive deﬁniteness of ˆ H − 1 , we further introduce adaptive damping into the update of y k .
Revised: To guarantee positive deﬁniteness of ˆ H − 1 , we further introduce adaptive damping into the update of y k .

Original: With the momentumized and damped parameter and gradient changes, we can construct a consistent ˆ H −throughout the whole optimization process.
Revised: With the momentumized and damped parameter and gradient changes, we can construct a consistent ˆ H − 1 throughout the whole optimization process.

Original: 3. I NTRODUCE M OMENTUM INTO THE H ESSIAN U PDATE
Revised: 3. I NTRODUCE M OMENTUM INTO THE H ESSIAN U PDATE

Original: Inspired by the success of momentum in ﬁrst-order methods, we demonstrate that the Hessian update in L-BFGS can be stabilized via momentum, without requiring a large batch size or even full gradients to reduce stochastic noise.
Revised: Inspired by the success of momentum in ﬁrst-order methods, we demonstrate that the Hessian update in L-BFGS can be stabilized via momentum, without requiring a large batch size or even full gradients to reduce stochastic noise.

Original: In particular, in this paper we apply momentum M θ t , M g t to past parameters θ t and gradients g t during mini-batch training as follows:
Revised: In particular, in this paper we apply momentum M θ t , M g t to past parameters θ t and gradients g t during mini-batch training as follows:

Original: [Equation]
Revised: [Equation]

Original: where β 1 and β 2 are the momentum coefﬁcients for θ t and g t respectively.
Revised: where β 1 and β 2 are the momentum coefﬁcients for θ t and g t respectively.

Original: Assuming that ˆ H − 1 is updated for every L mini-batch iterations, s k and g k are obtained as
Revised: Assuming that ˆ H − 1 is updated for every L mini-batch iterations, s k and g k are obtained as

Original: [Equation]
Revised: [Equation]

Original: This simple but effective technique works surprisingly well when gradients are noisy.
Revised: This simple but effective technique works surprisingly well when gradients are noisy.

Original: To intuitively show improvements of BFGS with momentum over the naive version, we compare the stochastic optimization of a simple quadratic loss function, L = 12 (cid:107) θ (cid:107) 2 .
Revised: To intuitively show improvements of BFGS with momentum over the naive version, we compare the stochastic optimization of a simple quadratic loss function, L = 12 (cid:107) θ (cid:107) 2 .

Original: In stochastic training, we write gradients in each mini-batch iteration as g t =
Revised: In stochastic training, we write gradients in each mini-batch iteration as g t =

Original: θ t + n t , where n t denotes the stochastic noise .
Revised: θ t + n t , where n t denotes the stochastic noise .

Original: For the sake of simplicity, we model the noise as i.i.d.
Revised: For the sake of simplicity, we model the noise as i.i.d.

Original: Gaussian, that is n t ∼ N (0 , σ 2 ) .
Revised: Gaussian, that is n t ∼ N (0 , σ 2 ) .

Original: Table 1 lists the expression for s k and y k for naive BFGS and for BFGS with momentum in terms of θ and the stochastic noise.
Revised: Table 1 lists the expression for s k and y k for naive BFGS and for BFGS with momentum in terms of θ and the stochastic noise.

Original: We compare the stability of y k in the two algorithms via their element-wise relative variance RV j = Var ( y k ( j ))
Revised: We compare the stability of y k in the two algorithms via their element-wise relative variance RV j = Var ( y k ( j ))

Original: E ( y k
Revised: E ( y k

Original: ( j )) , where y k ( j ) denotes the j th element of y k .
Revised: ( j )) , where y k ( j ) denotes the j th element of y k .

Original: As shown in Table 1, it is easy to observe that, RV mom.j using momentum is much lower than RV naivej in naive BFGS.
Revised: As shown in Table 1, it is easy to observe that, RV mom.j using momentum is much lower than RV naivej in naive BFGS.

Original: Hence, given a large momentum, SLIM-QN can signiﬁcantly suppress noise and obtain a more consistent y k .
Revised: Hence, given a large momentum, SLIM-QN can signiﬁcantly suppress noise and obtain a more consistent y k .

Original: With less noise in y k , BFGS with momentum leads to a better approximation of the real Hessian.
Revised: With less noise in y k , BFGS with momentum leads to a better approximation of the real Hessian.

Original: Figure 1 visualizes the optimization trajectory for SGD, naive BFGS, BFGS with momentum and the exact Newton method.
Revised: Figure 1 visualizes the optimization trajectory for SGD, naive BFGS, BFGS with momentum and the exact Newton method.

Original: With the same initialization, BFGS with momentum is as fast as the exact Newton method, and almost twice faster than SGD.
Revised: With the same initialization, BFGS with momentum is as fast as the exact Newton method, and almost twice faster than SGD.

Original: On the other hand, naive BFGS has difﬁculties ﬁnding the right optimization path due to the noise in y k .
Revised: On the other hand, naive BFGS has difﬁculties ﬁnding the right optimization path due to the noise in y k .

Original: 3. G UARANTEE P OSITIVE D EFINITENESS OF THE H ESSIAN VIA D AMPING
Revised: 3. G UARANTEE P OSITIVE D EFINITENESS OF THE H ESSIAN VIA D AMPING

Original: Even though adding momentum stabilizes the Hessian inverse, it cannot guarantee that the approximated Hessian is always positive deﬁnite, especially in practical stochastic and non-convex optimization.
Revised: Even though adding momentum stabilizes the Hessian inverse, it cannot guarantee that the approximated Hessian is always positive deﬁnite, especially in practical stochastic and non-convex optimization.

Original: As analysed in Sec 5.2, negative or very small positive values in the Hessian spectrum are harmful, and they immediately derail the optimization.
Revised: As analysed in Sec 5.2, negative or very small positive values in the Hessian spectrum are harmful, and they immediately derail the optimization.

Original: To effectively prevent radical changes in the Hessian, SLIM-QN further introduces an adaptive damping mechanism to the Hessian update.
Revised: To effectively prevent radical changes in the Hessian, SLIM-QN further introduces an adaptive damping mechanism to the Hessian update.

Original: Positive deﬁniteness of the Hessian hinges on stable and smooth gradient change vectors { y i } , especially in non-convex and stochastic settings.
Revised: Positive deﬁniteness of the Hessian hinges on stable and smooth gradient change vectors { y i } , especially in non-convex and stochastic settings.

Original: Hence, we choose to dampen y i by
Revised: Hence, we choose to dampen y i by

Original: [Equation]
Revised: [Equation]

Original: where ˆ y i is the damped version of y i , and τ is calculated as
Revised: where ˆ y i is the damped version of y i , and τ is calculated as

Original: [Equation]
Revised: [Equation]

Original: where µ = s Ti · y i s Ti · s i , σ L and σ H are the lower and upper thresholds for restraining eigenvalues in ˆ H −and 0 < τ 0 < 1 is a constant coefﬁcient.
Revised: where µ = s Ti · y i s Ti · s i , σ L and σ H are the lower and upper thresholds for restraining eigenvalues in ˆ H −and 0 < τ 0 < 1 is a constant coefﬁcient.

Original: The above damping scheme prevents sudden changes of s T
Revised: The above damping scheme prevents sudden changes of s T

Original: i y i , and guarantees the smoothness of ˆ H − 1 .
Revised: i y i , and guarantees the smoothness of ˆ H − 1 .

Original: Equivalently, equation 3 can be considered as scaling the undamped ˆ H and adaptively shifting its spectrum by a positive constant: τ · ˆ H + (1 − τ ) · I .
Revised: Equivalently, equation 3 can be considered as scaling the undamped ˆ H and adaptively shifting its spectrum by a positive constant: τ · ˆ H + (1 − τ ) · I .

Original: As a result, the eigenvalues of ˆ H are well controlled.
Revised: As a result, the eigenvalues of ˆ H are well controlled.

Original: 3. T HE O VERALL D ESCRIPTION
Revised: 3. T HE O VERALL D ESCRIPTION

Original: With momentum and damping introduced above, in this section we present the complete SLIMQN algorithm.
Revised: With momentum and damping introduced above, in this section we present the complete SLIMQN algorithm.

Original: As shown in Algorithm 2, for each mini-batch iteration, we use M θ t and M g t to accumulate momentum of parameters and gradients.
Revised: As shown in Algorithm 2, for each mini-batch iteration, we use M θ t and M g t to accumulate momentum of parameters and gradients.

Original: For every L iterations, we compute s k and y k , apply damping, and then update the Hessian approximation.
Revised: For every L iterations, we compute s k and y k , apply damping, and then update the Hessian approximation.

Original: Note that in UpdateHessian the inverse Hessian is never explicitly computed, instead we only compute and store the history vectors s k and ˆ y k that are necessary to perform the pre-conditioning step.
Revised: Note that in UpdateHessian the inverse Hessian is never explicitly computed, instead we only compute and store the history vectors s k and ˆ y k that are necessary to perform the pre-conditioning step.

Original: Since at the ﬁrst 2 L iterations, ˆ
Revised: Since at the ﬁrst 2 L iterations, ˆ

Original: H − 1 is not ready yet, we use SGD to conduct a warmup training.
Revised: H − 1 is not ready yet, we use SGD to conduct a warmup training.

Original: After 2 L iterations, we use ˆ H − 1 to ﬁrst pre-condition gradients g t and then apply updates to θ .
Revised: After 2 L iterations, we use ˆ H − 1 to ﬁrst pre-condition gradients g t and then apply updates to θ .

Original: Unlike KFAC, SLIM-QN is compatible to various regularizers such as L2 and gradient regularization(Smith et al., 2021), as long as we can derive gradients from them.
Revised: Unlike KFAC, SLIM-QN is compatible to various regularizers such as L2 and gradient regularization(Smith et al., 2021), as long as we can derive gradients from them.

Original: Finally, we add momentum to the parameter update ∆ θ t after pre-conditioning following the same procedure as in SGD with momentum (which is omitted in
Revised: Finally, we add momentum to the parameter update ∆ θ t after pre-conditioning following the same procedure as in SGD with momentum (which is omitted in

Original: Algorithm 2).
Revised: Algorithm 2).

Original: 1: for t = 1 , · · · , max _ iter do 2:
Revised: 1: for t = 1 , · · · , max _ iter do 2:

Original: Randomly choose mini-batch input S t ∈ X 3:
Revised: Randomly choose mini-batch input S t ∈ X 3:

Original: Compute gradients g t given S t (Forward/Backward) 4:
Revised: Compute gradients g t given S t (Forward/Backward) 4:

Original: Add weight decay: g t =
Revised: Add weight decay: g t =

Original: g t + wd · θ t 5:
Revised: g t + wd · θ t 5:

Original: Compute momentum on θ : M θ t = β 1 · M θ t − 1 + (1 − β 1 ) · θ t (cid:46) M θ 0 = θ6:
Revised: Compute momentum on θ : M θ t = β 1 · M θ t − 1 + (1 − β 1 ) · θ t (cid:46) M θ 0 = θ6:

Original: Compute momentum on g : M g t = β 2 · M g t − 1 + (1 − β 2 ) · g t (cid:46) M g 0 = g7: if t ≤ 2 L then 8:
Revised: Compute momentum on g : M g t = β 2 · M g t − 1 + (1 − β 2 ) · g t (cid:46) M g 0 = g7: if t ≤ 2 L then 8:

Original: Warmup: θ t +1 =
Revised: Warmup: θ t +1 =

Original: θ t − η t · g t 9: else 10:
Revised: θ t − η t · g t 9: else 10:

Original: Pre-condition: ∆ θ t
Revised: Pre-condition: ∆ θ t

Original: = ˆ H − 1 k · g t (cid:46) Algorithm11:
Revised: = ˆ H − 1 k · g t (cid:46) Algorithm11:

Original: Update: θ t +1 =
Revised: Update: θ t +1 =

Original: θ t − η t · ∆ θ t 12: if t % L == 0 and t > L then 13: k
Revised: θ t − η t · ∆ θ t 12: if t % L == 0 and t > L then 13: k

Original: = k
Revised: = k

Original: +14: s : s k
Revised: +14: s : s k

Original: = M θ t − M θ t − L 15: y : y k
Revised: = M θ t − M θ t − L 15: y : y k

Original: = M g t − M g t − L 16:
Revised: = M g t − M g t − L 16:

Original: Damping: ˆ y k
Revised: Damping: ˆ y k

Original: = τ · y k + (1 − τ ) · s k (cid:46) τ from equation17: ˆ H − 1 : ˆ H − 1 k
Revised: = τ · y k + (1 − τ ) · s k (cid:46) τ from equation17: ˆ H − 1 : ˆ H − 1 k

Original: =
Revised: =

Original: UpdateHessian ( ˆ H − 1 k − 1 , s k , ˆ y k )
Revised: UpdateHessian ( ˆ H − 1 k − 1 , s k , ˆ y k )

Original: (cid:46) equation 2, not explicitly computed
Revised: (cid:46) equation 2, not explicitly computed

Original:  T HEORETICAL G UARANTEES FOR SLIM-QN
Revised:  T HEORETICAL G UARANTEES FOR SLIM-QN

Original: In this section, we ﬁrst present convergence guarantees for SLIM-QN, and then discuss the compute and memory costs for SGD, KFAC, and SLIM-QN.
Revised: In this section, we ﬁrst present convergence guarantees for SLIM-QN, and then discuss the compute and memory costs for SGD, KFAC, and SLIM-QN.

Original: 4. C ONVERGENCE G UARANTEES
Revised: 4. C ONVERGENCE G UARANTEES

Original: Following the framework of quasi-Newton method in Wang et al.
Revised: Following the framework of quasi-Newton method in Wang et al.

Original: (2017) in stochastic optimization, we prove that SLIM-QN converges to the optimum at a linear rate, under proper assumptions as follows:
Revised: (2017) in stochastic optimization, we prove that SLIM-QN converges to the optimum at a linear rate, under proper assumptions as follows:

Original: AS 1.
Revised: AS 1.

Original: L ( θ ) is λ -PL in that it satisﬁes Polyak-Lojasiewicz (PL) condition for a constant λ > 0 :
Revised: L ( θ ) is λ -PL in that it satisﬁes Polyak-Lojasiewicz (PL) condition for a constant λ > 0 :

Original: (cid:107)∇L ( θ ) (cid:107) 2 ≥ λ L ( θ ) .
Revised: (cid:107)∇L ( θ ) (cid:107) 2 ≥ λ L ( θ ) .

Original: AS 2.
Revised: AS 2.

Original: (cid:96) i ( θ ) is Λ -smooth for 1 ≤ i ≤ N , Λ > 0 : ∀ θ 1 , θ 2 , (cid:96) i ( θ 2 ) ≤ (cid:96) i ( θ 1 ) + (cid:104)∇ (cid:96) i ( θ 1 ) , θ 2 − θ 1 (cid:105) + Λ2 (cid:107) θ 2 − θ 1 (cid:107) 2 .
Revised: (cid:96) i ( θ ) is Λ -smooth for 1 ≤ i ≤ N , Λ > 0 : ∀ θ 1 , θ 2 , (cid:96) i ( θ 2 ) ≤ (cid:96) i ( θ 1 ) + (cid:104)∇ (cid:96) i ( θ 1 ) , θ 2 − θ 1 (cid:105) + Λ2 (cid:107) θ 2 − θ 1 (cid:107) 2 .

Original: AS 3.
Revised: AS 3.

Original: For every sequence θ 1 , θ 2 , · · · such that lim t →∞ L ( θ t ) = 0 , then for all 1 ≤ i ≤ N , lim t →∞ (cid:96) i ( θ t ) =
Revised: For every sequence θ 1 , θ 2 , · · · such that lim t →∞ L ( θ t ) = 0 , then for all 1 ≤ i ≤ N , lim t →∞ (cid:96) i ( θ t ) =

Original: We note that compared to typical strong convexity assumptions the PL condition in AS 1 (Polyak, 1963) applies to much broader settings including when the loss is nonconvex as it only requires lower bounded variance in gradients, rather than strict positive deﬁniteness required for the Hessian with strong convexity.
Revised: We note that compared to typical strong convexity assumptions the PL condition in AS 1 (Polyak, 1963) applies to much broader settings including when the loss is nonconvex as it only requires lower bounded variance in gradients, rather than strict positive deﬁniteness required for the Hessian with strong convexity.

Original: AS 3 assumes that the global minima of the summands (cid:96) i are the same as the global minima of their sum.
Revised: AS 3 assumes that the global minima of the summands (cid:96) i are the same as the global minima of their sum.

Original: This is in line with what has been observed in over-parameterized deep learning models (Ma et al., 2018).
Revised: This is in line with what has been observed in over-parameterized deep learning models (Ma et al., 2018).

Original: Before we state our main result we require several auxiliary lemmas.
Revised: Before we state our main result we require several auxiliary lemmas.

Original: First Lemma 1 states that s Ti · ˆ y i always lies in [ σ L , σ H ] · s Ti s i .
Revised: First Lemma 1 states that s Ti · ˆ y i always lies in [ σ L , σ H ] · s Ti s i .

Original: With Lemma 1 in place, in Lemma 2 we bound the approximation of the Hessian ˆ H − 1 k at k th Hessian update.
Revised: With Lemma 1 in place, in Lemma 2 we bound the approximation of the Hessian ˆ H − 1 k at k th Hessian update.

Original: Lemma 3 further establishes smoothness on L ( θ ) given each (cid:96) i is smooth.
Revised: Lemma 3 further establishes smoothness on L ( θ ) given each (cid:96) i is smooth.

Original: While Lemma 4 further bounds gradient variance in mini-batch training.
Revised: While Lemma 4 further bounds gradient variance in mini-batch training.

Original: Lemma 1.
Revised: Lemma 1.

Original: Given damping scheme in equation 3, if we choose τ according to equation 4, then σ L ≤ s Ti · ˆ y i s Ti · s i ≤ σ H .
Revised: Given damping scheme in equation 3, if we choose τ according to equation 4, then σ L ≤ s Ti · ˆ y i s Ti · s i ≤ σ H .

Original: Lemma 2.
Revised: Lemma 2.

Original: Given Lemma 1, at the k -th Hessian update, ˆ H − 1 k during the optimization is bounded by ξI (cid:22) ˆ H − 1 k (cid:22)
Revised: Given Lemma 1, at the k -th Hessian update, ˆ H − 1 k during the optimization is bounded by ξI (cid:22) ˆ H − 1 k (cid:22)

Original: Ξ I , where Ξ = ( M + 1) 1 σ L , and ξ = 1 σ H .
Revised: Ξ I , where Ξ = ( M + 1) 1 σ L , and ξ = 1 σ H .

Original: σ L and σ H is the lower and upper bound for s Ti · ˆ y i s Ti · s i in Lemma 1.
Revised: σ L and σ H is the lower and upper bound for s Ti · ˆ y i s Ti · s i in Lemma 1.

Original: Lemma 3.
Revised: Lemma 3.

Original: Assume AS 2 holds, then loss function
Revised: Assume AS 2 holds, then loss function

Original: L ( θ ) is at least Λ -smooth.
Revised: L ( θ ) is at least Λ -smooth.

Original: Lemma 4.
Revised: Lemma 4.

Original: Assume AS 2-3 holds, with Lemma 3, at iteration t with mini-batch input S t , gradient ∇L ( θ t ; S t ) satisﬁes
Revised: Assume AS 2-3 holds, with Lemma 3, at iteration t with mini-batch input S t , gradient ∇L ( θ t ; S t ) satisﬁes

Original: [Equation]
Revised: [Equation]

Original: With ˆ H − 1 k and gradient variance bounded, we can derive the following convergence theorem.
Revised: With ˆ H − 1 k and gradient variance bounded, we can derive the following convergence theorem.

Original: Theorem 1.
Revised: Theorem 1.

Original: Assume AS 1-3 hold at each iteration t of SLIM-QN with mini-batch input S t , then the expectation of L ( θ t ) satisﬁes
Revised: Assume AS 1-3 hold at each iteration t of SLIM-QN with mini-batch input S t , then the expectation of L ( θ t ) satisﬁes

Original: [Equation]
Revised: [Equation]

Original: where α t − 1 = 1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 .
Revised: where α t − 1 = 1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 .

Original: Proofs for the lemmas and theorem are provided in the appendix.
Revised: Proofs for the lemmas and theorem are provided in the appendix.

Original: Remark 1.
Revised: Remark 1.

Original: By choosing η t − 1 such that α t − 1 < 1 , SLIM-QN converges at a linear rate.
Revised: By choosing η t − 1 such that α t − 1 < 1 , SLIM-QN converges at a linear rate.

Original: Remark 2.
Revised: Remark 2.

Original: We note that Theorem 1 also applies to convex settings, as strong convexity
Revised: We note that Theorem 1 also applies to convex settings, as strong convexity

Original: [Table]
Revised: [Table]

Original: layer i .

Original: 4. C OMPUTE AND M EMORY C OST
Revised: 4. C OMPUTE AND M EMORY C OST

Original: As mentioned before, SLIM-QN aims to reduce the complexity of approximating the Hessian.
Revised: As mentioned before, SLIM-QN aims to reduce the complexity of approximating the Hessian.

Original: In this section, we summarize the compute and memory cost of SGD, KFAC, stochastic L-BFGS (sL-BFGS)
Revised: In this section, we summarize the compute and memory cost of SGD, KFAC, stochastic L-BFGS (sL-BFGS)

Original: variants(Chang et al., 2019) and SLIM-QN, and demonstrate the cost advantage of SLIM-QN.
Revised: variants(Chang et al., 2019) and SLIM-QN, and demonstrate the cost advantage of SLIM-QN.

Original: Given a model with parameter θ , we use C fb and M fb to represent the compute and memory cost of a forward/backward (Fwd/Bwd) pass with a batch size of b = 1 .
Revised: Given a model with parameter θ , we use C fb and M fb to represent the compute and memory cost of a forward/backward (Fwd/Bwd) pass with a batch size of b = 1 .

Original: Furthermore, C opt denotes the compute cost of model updates (Opt) which consists of gradient reduction, computing the update ∆ θ , and applying it to θ .
Revised: Furthermore, C opt denotes the compute cost of model updates (Opt) which consists of gradient reduction, computing the update ∆ θ , and applying it to θ .

Original: Table 2 summarizes the compute and memory cost of SGD, KFAC, sL-BFGS and SLIM-QN.
Revised: Table 2 summarizes the compute and memory cost of SGD, KFAC, sL-BFGS and SLIM-QN.

Original: Compared to SGD, during the forward and backward passes, SLIM-QN needs to additionally compute M θ , M g , for which the complexity increases linearly with model size ( α 2 (cid:107) θ (cid:107) ).
Revised: Compared to SGD, during the forward and backward passes, SLIM-QN needs to additionally compute M θ , M g , for which the complexity increases linearly with model size ( α 2 (cid:107) θ (cid:107) ).

Original: The main extra compute SLIM-QN introduces is the Hessian-vector product, in which we need to iterate over { s i } Mi =1 and { y i } Mi =1 , as shown in Algorithm 1.
Revised: The main extra compute SLIM-QN introduces is the Hessian-vector product, in which we need to iterate over { s i } Mi =1 and { y i } Mi =1 , as shown in Algorithm 1.

Original: The complexity increases linearly with the number of stored history vectors and model size ( 2 M (cid:107) θ (cid:107) ).
Revised: The complexity increases linearly with the number of stored history vectors and model size ( 2 M (cid:107) θ (cid:107) ).

Original: While, compared to O ( bC fb ) complexity in forward and backward pass, such operations add relatively marginal cost.
Revised: While, compared to O ( bC fb ) complexity in forward and backward pass, such operations add relatively marginal cost.

Original: As a comparison, KFAC though only approximates diagonal blocks of the Fisher matrix, it still adds signiﬁcant additional computations through 1) multiple backward passes to update factors ( γbC fb with γ ≥ 1 ), 2) matrix inversion ( (cid:80) ( d 3 i
Revised: As a comparison, KFAC though only approximates diagonal blocks of the Fisher matrix, it still adds signiﬁcant additional computations through 1) multiple backward passes to update factors ( γbC fb with γ ≥ 1 ), 2) matrix inversion ( (cid:80) ( d 3 i

Original: + ( (cid:107) θ
Revised: + ( (cid:107) θ

Original: i (cid:107) d i ) 3 ) ) for every L iterations, and 3) Matrix-vector products ( 2 (cid:80) ( d i + (cid:107) θ i (cid:107) d i ) (cid:107) θ i (cid:107) ).
Revised: i (cid:107) d i ) 3 ) ) for every L iterations, and 3) Matrix-vector products ( 2 (cid:80) ( d i + (cid:107) θ i (cid:107) d i ) (cid:107) θ i (cid:107) ).

Original: If the Fisher matrix is updated more frequently (that is for small L ), then the amortized cost for matrix inversion is even more striking.
Revised: If the Fisher matrix is updated more frequently (that is for small L ), then the amortized cost for matrix inversion is even more striking.

Original: On the other hand, sL-BFGS also resorts to computation-intensive operations including full-batch gradients and a separate large batch to estimate the Hessian, which respectively adds amortized costs of O ( bC fb ) and 1 L b H C fb .
Revised: On the other hand, sL-BFGS also resorts to computation-intensive operations including full-batch gradients and a separate large batch to estimate the Hessian, which respectively adds amortized costs of O ( bC fb ) and 1 L b H C fb .

Original: As for memory usage, compared to SGD, SLIM-QN mainly needs O (2 M (cid:107) θ (cid:107) ) to store history vectors { s i } M i =1 and { y i } M i =1  , similar to sL-BFGS.
Revised: As for memory usage, compared to SGD, SLIM-QN mainly needs O (2 M (cid:107) θ (cid:107) ) to store history vectors { s i } M i =1 and { y i } M i =1 . sL-BFGS needs the same storage for s i , y i , and amortized cost of

Original:  While KFAC needs O (2 (cid:80) ( d 2 i + ( (cid:107) θ i (cid:107) d i ) 2 )) to store sub-matricesof Hessian approximation and their inverse, where the actual memory footprint hinges on model architectures.
Revised: O ( 1 L b H M fb ) for additional backward passes. While KFAC needs O (2 (cid:80) ( d 2 i + ( (cid:107) [Table] d i ) 2 )) to store sub-matrices and their inverse, where the actual memory footprint hinges on model architectures.

Original: In practice,typically M is set to be 10 ∼ 20 , which ensures that memory usage is manageable in SLIM-QN.
Revised: In practice, M is set to be 10 ∼ 20 , which ensures that memory usage is manageable in SLIM-QN.

Original:  E MPIRICAL A NALYSES
Revised:  E MPIRICAL A NALYSES

Original: We conduct various experiments on computer vision (CV) problems, where SGD has been widely used.
Revised: We conduct various experiments on computer vision (CV) problems, where SGD has been widely used.

Original: Methods like Adam (Kingma & Ba, 2014) and AdaGrad (Duchi et al., 2011) greatly under-perform
Revised: Methods like Adam (Kingma & Ba, 2014) and AdaGrad (Duchi et al., 2011) greatly under-perform

Original: SGD (Defazio & Jelassi, 2021).
Revised: SGD (Defazio & Jelassi, 2021).

Original: Two metrics are used to evaluate the performance: iteration-wise convergence and wall-clock convergence.
Revised: Two metrics are used to evaluate the performance: iteration-wise convergence and wall-clock convergence.

Original: The iteration-wise metric shows the pure convergence promises of the optimizer; while the wall-clock convergence further captures the impacts of computation complexity on the run-time.
Revised: The iteration-wise metric shows the pure convergence promises of the optimizer; while the wall-clock convergence further captures the impacts of computation complexity on the run-time.

Original: Furthermore, we also conduct an ablation study that investigates how the components in SLIM-QN (momentum and damping) affect the optimization process.
Revised: Furthermore, we also conduct an ablation study that investigates how the components in SLIM-QN (momentum and damping) affect the optimization process.

Original: The current implementation in PyTorch (Paszke et al., 2019) supports various DNN models in a multi-GPU system.
Revised: The current implementation in PyTorch (Paszke et al., 2019) supports various DNN models in a multi-GPU system.

Original: During training, after gradients are synchronized across GPUs, we keep updating the momentum, M θ t and M g t on each GPU.
Revised: During training, after gradients are synchronized across GPUs, we keep updating the momentum, M θ t and M g t on each GPU.

Original: As a result, each GPU stores a copy of the Hessian inverse and locally performs gradient conditioning without communicating across GPUs.
Revised: As a result inverse and locally performs gradient conditioning without communicating across GPUs.

Original: We test models: ResNet18, ResNet50, Vision Transformer on datasets: CIFAR-10(Krizhevsky et al., 2009) and ImageNet(Deng et al., 2009).
Revised: We test models: ResNet18, ResNet50, Vision Transformer on datasets: CIFAR-10(Krizhevsky et al., 2009) and ImageNet(Deng et al., 2009).

Original: We also run SGD and KFAC as two baselines for comparison.
Revised: We also run SGD and KFAC as two baselines for comparison.

Original: 5. E XPERIMENTS ON I MAGE N ET C LASSIFICATION
Revised: 5. E XPERIMENTS ON I MAGE N ET C LASSIFICATION

Original: ImageNet classiﬁcation has been the gold standard for evaluating performance of different optimization algorithms for CV models.
Revised: ImageNet classiﬁcation has been the gold standard for evaluating performance of different optimization algorithms for CV models.

Original: Compared to CIFAR-10, ImageNet consists of much more training and test images ( ∼ 1.2M training and ∼ 50K test images), categorized into 1000 classes.
Revised: Compared to CIFAR-10, ImageNet consists of much more training and test images ( ∼ 1.2M training and ∼ 50K test images), categorized into 1000 classes.

Original: Therefore, convergence on ImageNet can better reveal the optimizer’s promises in practical problems.
Revised: Therefore, convergence on ImageNet can better reveal the optimizer’s promises in practical problems.

Original: During data pre-processing, we resize images to 256 × 256 , and randomly crop to 224 × 224 , and then randomly ﬂip each image.
Revised: During data pre-processing, we resize images to 256 × 256 , and randomly crop to 224 × 224 , and then randomly ﬂip each image.

Original: Each image is normalized using pre-computed mean and variance.
Revised: Each image is normalized using pre-computed mean and variance.

Original: 5.1. R ES N ET - 
Revised: 5.1. R ES N ET -

Original: Figure 2 shows iteration-wise (Top) and wall-clock (Bottom) convergence on ResNet-50 using SGD,
Revised: Figure 2 shows iteration-wise (Top) and wall-clock (Bottom) convergence on ResNet-50 using SGD,

Original: KFAC and SLIM-QN.
Revised: KFAC and SLIM-QN.

Original: Detailed hyper-parameter settings are provided in the appendix.
Revised: Detailed hyper-parameter settings are provided in the appendix.

Original: SLIM-QN enjoys very fast per-iteration convergence, and reaches near optimal accuracy 1 . 5 × faster than SGD, and even 2 × faster in the early-stages.
Revised: SLIM-QN enjoys very fast per-iteration convergence, and reaches near optimal accuracy 1 . 5 × faster than SGD, and even 2 × faster in the early-stages.

Original: Furthermore, it also generalizes well on the validation set, and ﬁnally reaches comparable validation accuracy to SGD.
Revised: Furthermore, it also generalizes well on the validation set, and ﬁnally reaches comparable validation accuracy to SGD.

Original: The beneﬁt of SLIM-QN is even more striking in terms of wall-clock time.
Revised: The beneﬁt of SLIM-QN is even more striking in terms of wall-clock time.

Original: Due to light compute costs, it is 1 .
Revised: Due to light compute costs, it is 1 .

Original: 75 × / 1 .
Revised: 75 × / 1 .

Original: 36 × faster in the early and late stages compared to SGD.
Revised: 36 × faster in the early and late stages compared to SGD.

Original: Whereas in KFAC, the wall-clock performance is signiﬁcantly neutralized by its additional compute costs.
Revised: Whereas in KFAC, the wall-clock performance is signiﬁcantly neutralized by its additional compute costs.

Original: 5.1. V ISION T RANSFORMER
Revised: 5.1. V ISION T RANSFORMER

Original: As a step towards understanding the efﬁcacy of second-order optimizers on contemporary Transformerbased CV models, we perform experiments on a small ( 10 M parameters)
Revised: As a step towards understanding the efﬁcacy of second-order optimizers on contemporary Transformerbased CV models, we perform experiments on a small ( 10 M parameters)

Original: Vision Transformer (ViT)
Revised: Vision Transformer (ViT)

Original: (Dosovitskiy et al., 2020b) using SLIM-QN on ImageNet.
Revised: (Dosovitskiy et al., 2020b) using SLIM-QN on ImageNet.

Original: Details on hyperparameters, model architecture and experiments on further datasets are deferred to the Appendix.
Revised: Details on hyperparameters, model architecture and experiments on further datasets are deferred to the Appendix.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 3, SLIM-QN beneﬁts from faster early-stage convergence compared to the SGD, which is consistent with our ﬁndings for ResNet.
Revised: 3, SLIM-QN beneﬁts from faster early-stage convergence compared to the SGD, which is consistent with our ﬁndings for ResNet.

Original: Furthermore, we observe that SLIM-QN ﬁnds a solution with good generalization achieving a ﬁnal validation accuracy slightly higher than SGD.
Revised: Furthermore, we observe that SLIM-QN ﬁnds a solution with good generalization achieving a ﬁnal validation accuracy slightly higher than SGD.

Original: 5. A BLATION S TUDY : T HE E FFECTS OF M OMENTUM AND D AMPING
Revised: 5. A BLATION S TUDY : T HE E FFECTS OF M OMENTUM AND D AMPING

Original: In this section, we give more insight into the effects of momentum and damping used in SLIM-QN.
Revised: In this section, we give more insight into the effects of momentum and damping used in SLIM-QN.

Original: To this end, we ablate two critical components in SLIM-QN: momentum and damping in the Hessian approximation, and then use the ablated version to train ResNet-18 on CIFAR-10.
Revised: To this end, we ablate two critical components in SLIM-QN: momentum and damping in the Hessian approximation, and then use the ablated version to train ResNet-18 on CIFAR-10.

Original: We focus on
Revised: We focus on

Original: CIFAR-10 since we observed more convergence instability on this dataset compared to others.
Revised: CIFAR-10 since we observed more convergence instability on this dataset compared to others.

Original: Figure 4 shows convergence using the ablated SLIM-QN with only momentum (black), with only damping (purple), and with no momentum or damping (red).
Revised: Figure 4 shows convergence using the ablated SLIM-QN with only momentum (black), with only damping (purple), and with no momentum or damping (red).

Original: Due to stochastic noise, the ablated version of SLIM-QN without momentum or damping diverges easily in the early stages.
Revised: Due to stochastic noise, the ablated version of SLIM-QN without momentum or damping diverges easily in the early stages.

Original: With momentum (black), the whole optimization is signiﬁcantly stabilized.
Revised: With momentum (black), the whole optimization is signiﬁcantly stabilized.

Original: However, it still fails to converge when there is an abrupt change in the loss landscape (for example, when learning rate decays).
Revised: However, it still fails to converge when there is an abrupt change in the loss landscape (for example, when learning rate decays).

Original: With damping (purple), the Hessian approximation is effectively restrained, especially when such sudden changes in the loss landscape happen.
Revised: With damping (purple), the Hessian approximation is effectively restrained, especially when such sudden changes in the loss landscape happen.

Original: It is interesting to observe that while damping prevents divergence, the whole training is still largely affected by stochastic noise.
Revised: It is interesting to observe that while damping prevents divergence, the whole training is still largely affected by stochastic noise.

Original: Notable ﬂuctuation in the loss and accuracy is commonly observed during training.
Revised: Notable ﬂuctuation in the loss and accuracy is commonly observed during training.

Original: As a comparison, the complete
Revised: As a comparison, the complete

Original: SLIM-QN (blue) effectively addresses these issues achieving much more stable convergence.
Revised: SLIM-QN (blue) effectively addresses these issues achieving much more stable convergence.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we propose SLIM-QN, a quasi-Newton method that simultaneously mitigates computation and convergence instability barriers in second-order methods.
Revised: In this paper, we propose SLIM-QN, a quasi-Newton method that simultaneously mitigates computation and convergence instability barriers in second-order methods.

Original: SLIM-QN introduces momentum and damping into the Hessian update, which obviates the need for estimating the Hessian with high costs.
Revised: SLIM-QN introduces momentum and damping into the Hessian update, which obviates the need for estimating the Hessian with high costs.

Original: Empirical analyses on CV models, such as ResNet-50 and Vision Transformer show that
Revised: Empirical analyses on CV models, such as ResNet-50 and Vision Transformer show that

Original: SLIM-QN achieves faster convergence in the early stages, and reaches comparable accuracy to SGD.
Revised: SLIM-QN achieves faster convergence in the early stages, and reaches comparable accuracy to SGD.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: Implementation of SLIM-QN in a multi-GPU platform is described at the beginning of Sec 5.
Revised: Implementation of SLIM-QN in a multi-GPU platform is described at the beginning of Sec 5.

Original: A copy of code is included in supplementary materials.
Revised: A copy of code is included in supplementary materials.

Original: Appendix B lists all hyperparameters for obtaining results on ImageNet using ResNet-50 and ViT models.
Revised: Appendix B lists all hyperparameters for obtaining results on ImageNet using ResNet-50 and ViT models.

Original: Moreover, Appendix C presents more results on CIFAR-10 using ResNet-18 and ViT models.
Revised: Moreover, Appendix C presents more results on CIFAR-10 using ResNet-18 and ViT models.

Original: We also describe some tips of tuning hyperparameters in Appendix F.
Revised: We also describe some tips of tuning hyperparameters in Appendix F.

Original: For theoretical results, proofs of all lemmas and theorems are provided in Appendix A
Revised: For theoretical results, proofs of all lemmas and theorems are provided in Appendix A

Original: A P ROOFS OF L EMMAS AND T HEOREMS
Revised: A P ROOFS OF L EMMAS AND T HEOREMS

Original: This appendix are organized as follows:
Revised: This appendix are organized as follows:

Original: Section A.1-A.4 presents the proof of Lemma 1-4.
Revised: Section A.1-A.4 presents the proof of Lemma 1-4.

Original: Section A.5 presents the proof of Theorem 1.
Revised: Section A.5 presents the proof of Theorem 1.

Original: A. P ROOF OF L EMMA
Revised: A. P ROOF OF L EMMA

Original: Proof.
Revised: Proof.

Original: According to equation 3, s
Revised: According to equation 3, s

Original: Ti ˆ y i =
Revised: Ti ˆ y i =

Original: s
Revised: s

Original: Ti ( τ y i +(1 − τ ) s i )
Revised: Ti ( τ y i +(1 − τ ) s i )

Original: = ( µτ +1 − τ ) s Ti s i , where µ =
Revised: = ( µτ +1 − τ ) s Ti s i , where µ =

Original: s Ti y
Revised: s Ti y

Original: i s Ti s i .
Revised: i s Ti s i .

Original: For µ ≤ σ L , two cases need to be considered: 1) τ
Revised: For µ ≤ σ L , two cases need to be considered: 1) τ

Original: = τ 0 ; 2) τ
Revised: = τ 0 ; 2) τ

Original: = 1 − σ L 1 − µ .
Revised: = 1 − σ L 1 − µ .

Original: If τ = τ 0 , then 1 − σ L 1 − µ ≥ τ 0 , and µτ + 1 − τ ≥ σ L
Revised: If τ = τ 0 , then 1 − σ L 1 − µ ≥ τ 0 , and µτ + 1 − τ ≥ σ L

Original: If τ = 1 − σ L 1 − µ , then µτ + 1 − τ = σ L
Revised: If τ = 1 − σ L 1 − µ , then µτ + 1 − τ = σ L

Original: Therefore, when µ ≤ σ L , s Ti ˆ y i ≥ σ L s Ti s i
Revised: Therefore, when µ ≤ σ L , s Ti ˆ y i ≥ σ L s Ti s i

Original: For σ L < µ < σ H :
Revised: For σ L < µ < σ H :

Original: We can write µτ +1 − τ = µτ 0 +1 − τ 0 .
Revised: We can write µτ +1 − τ = µτ 0 +1 − τ 0 .

Original: It is easy to show that µτ 0 +1 − τ 0 − σ L ≥ (1 − σ L )(1 − τ 0 ) >and µτ 0 + 1 − τ 0 − σ H ≤ (1 − σ H )(1 − τ o ) < 0 .
Revised: It is easy to show that µτ 0 +1 − τ 0 − σ L ≥ (1 − σ L )(1 − τ 0 ) >and µτ 0 + 1 − τ 0 − σ H ≤ (1 − σ H )(1 − τ o ) < 0 .

Original: Therefore, when σ L < µ < σ H , σ L s Ti s i < s
Revised: Therefore, when σ L < µ < σ H , σ L s Ti s i < s

Original: Ti ˆ y i <
Revised: Ti ˆ y i <

Original: σ H s Ti s i .
Revised: σ H s Ti s i .

Original: For µ ≥ σ H , similarly two cases might arise: 1) τ = τ 0 ; 2) τ
Revised: For µ ≥ σ H , similarly two cases might arise: 1) τ = τ 0 ; 2) τ

Original: = σ H − 1 µ − 1 .
Revised: = σ H − 1 µ − 1 .

Original: If τ = τ 0 , then σ H − 1 µ − 1 ≥ τ 0 , and µτ + 1 − τ ≤ σ H .
Revised: If τ = τ 0 , then σ H − 1 µ − 1 ≥ τ 0 , and µτ + 1 − τ ≤ σ H .

Original: If τ = σ H − 1 µ − 1 , then µτ + 1 − τ = σ H .
Revised: If τ = σ H − 1 µ − 1 , then µτ + 1 − τ = σ H .

Original: Therefore, when µ ≥ σ H , s Ti ˆ y i ≤ σ H s Ti s i
Revised: Therefore, when µ ≥ σ H , s Ti ˆ y i ≤ σ H s Ti s i

Original: In summary, σ L ≤ s Ti · ˆ y i s Ti · s i ≤ σ H .
Revised: In summary, σ L ≤ s Ti · ˆ y i s Ti · s i ≤ σ H .

Original: A. P ROOF OF L EMMA
Revised: A. P ROOF OF L EMMA

Original: Proof.
Revised: Proof.

Original: Lower bound : ˆ H − 1 is initialized as ˆ H − 1= s T 

Revised: Lower bound : ˆ H − 1 is initialized as ˆ H − 1= s T

Original: ˆ y 

Revised: ˆ y

Original: ˆ y T 

Revised: ˆ y T

Original: ˆ y 0 · I .
Revised: ˆ y 0 · I .

Original: According to Lemma 1, there exists
Revised: According to Lemma 1, there exists

Original: H 0 ( θ ) (cid:22) σ H I such that ˆ y 

Revised: H 0 ( θ ) (cid:22) σ H I such that ˆ y

Original: = H 0 · s 0 .
Revised: = H 0 · s 0 .

Original: Therefore, s T 

Revised: Therefore, s T

Original: ˆ y
Revised: ˆ y

Original: [Equation]
Revised: [Equation]

Original: T
Revised: T

Original: [Equation]
Revised: [Equation]

Original: y 0 · I =
Revised: y 0 · I =

Original: [Equation]
Revised: [Equation]

Original: Then for k ≥ 1 , assuming ˆ H − 1 k − 1 (cid:23) 1 σ H I hold, based on equation 2, ˆ H − 1 k
Revised: Then for k ≥ 1 , assuming ˆ H − 1 k − 1 (cid:23) 1 σ H I hold, based on equation 2, ˆ H − 1 k

Original: [Equation]
Revised: [Equation]

Original: Because ( I − ρ k ˆ y k s Tk ) T ˆ H − 1 k − 1 ( I − ρ k
Revised: Because ( I − ρ k ˆ y k s Tk ) T ˆ H − 1 k − 1 ( I − ρ k

Original: ˆ y k s Tk ) is positive deﬁnite, we can bound ˆ H − 1 k as: ˆ H − 1 k (cid:23)
Revised: ˆ y k s Tk ) is positive deﬁnite, we can bound ˆ H − 1 k as: ˆ H − 1 k (cid:23)

Original: [Equation]
Revised: [Equation]

Original: Therefore, lower bound of ˆ H − 1 k , ξ = 1 σ H .
Revised: Therefore, lower bound of ˆ H − 1 k , ξ = 1 σ H .

Original: Upper bound : Since H 0 ( θ ) (cid:23) σ L , we can get s T 

Revised: Upper bound : Since H 0 ( θ ) (cid:23) σ L , we can get s T

Original: ˆ y 

Revised: ˆ y

Original: ˆ y T 

Revised: ˆ y T

Original: ˆ y 0 · I =
Revised: ˆ y 0 · I =

Original: [Equation]
Revised: [Equation]

Original: L
Revised: L

Original: [Equation]
Revised: [Equation]

Original: Similarly, for k ≥ 1 , we assume ˆ H − 1 k − 1 (cid:22) k 1 σ L hold.
Revised: Similarly, for k ≥ 1 , we assume ˆ H − 1 k − 1 (cid:22) k 1 σ L hold.

Original: For the ﬁrst part in equation 2, let Q
Revised: For the ﬁrst part in equation 2, let Q

Original: =
Revised: =

Original: ( I − ρ k ˆ y k
Revised: ( I − ρ k ˆ y k

Original: s Tk ) .
Revised: s Tk ) .

Original: Then for ∀ x (cid:54)
Revised: Then for ∀ x (cid:54)

Original: [Equation]
Revised: [Equation]

Original: Let P = s
Revised: Let P = s

Original: k ˆ y Tk
Revised: k ˆ y Tk

Original: ˆ y k
Revised: ˆ y k

Original: s
Revised: s

Original: Tk s T
Revised: Tk s T

Original: k
Revised: k

Original: ˆ y k
Revised: ˆ y k

Original: s T
Revised: s T

Original: k ˆ y k − ˆ
Revised: k ˆ y k − ˆ

Original: y k
Revised: y k

Original: s
Revised: s

Original: Tk + s
Revised: Tk + s

Original: k ˆ y Tk s T
Revised: k ˆ y Tk s T

Original: k ˆ y k , then Q T Q
Revised: k ˆ y k , then Q T Q

Original: = I
Revised: = I

Original: + P .
Revised: + P .

Original: Because P is rank-1 matrix with eigenvalue − 1 and eigenvector ˆ y k , we have x T · Q T
Revised: Because P is rank-1 matrix with eigenvalue − 1 and eigenvector ˆ y k , we have x T · Q T

Original: ˆ H − 1 k − 1 Q · x ≤ kσ L x T · ( I + P ) · x ≤ kσ L x T x
Revised: ˆ H − 1 k − 1 Q · x ≤ kσ L x T · ( I + P ) · x ≤ kσ L x T x

Original: For the second part in equation 2, we can directly get s k s Tk s T
Revised: For the second part in equation 2, we can directly get s k s Tk s T

Original: y k
Revised: y k

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: Therefore, x T · ˆ H − 1 k · x ≤ kσ L x T x + 1 σ L x T x = k +1 σ L x T x
Revised: Therefore, x T · ˆ H − 1 k · x ≤ kσ L x T x + 1 σ L x T x = k +1 σ L x T x

Original: In SLIM-QN, k is at most M , which is the length of history vector, therefore Ξ =
Revised: In SLIM-QN, k is at most M , which is the length of history vector, therefore Ξ =

Original: ( M + 1) 1 σ L .
Revised: ( M + 1) 1 σ L .

Original: In summary, 1 σ H I (cid:22)
Revised: In summary, 1 σ H I (cid:22)

Original: [Equation]
Revised: [Equation]

Original: 1 k (cid:22)
Revised: 1 k (cid:22)

Original: [Equation]
Revised: [Equation]

Original: A. P ROOF OF L EMMA
Revised: A. P ROOF OF L EMMA

Original: Proof.
Revised: Proof.

Original: Given AS 2 hold, we have
Revised: Given AS 2 hold, we have

Original: [Equation]
Revised: [Equation]

Original: (cid:107) .
Revised: (cid:107) .

Original: For L , we have
Revised: For L , we have

Original: [Equation]
Revised: [Equation]

Original: Then, (cid:107)∇L
Revised: Then, (cid:107)∇L

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ) (cid:107)
Revised: ) (cid:107)

Original: [Equation]
Revised: [Equation]

Original: N (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Revised: N (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Original: N (cid:88) i = ∇ (cid:96) i ( θ 1 ) −
Revised: N (cid:88) i = ∇ (cid:96) i ( θ 1 ) −

Original: [Equation]
Revised: [Equation]

Original: (cid:96) i ( θ 2 ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Revised: (cid:96) i ( θ 2 ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Original: TI ≤
Revised: TI ≤

Original: N
Revised: N

Original: N (cid:88) i =
Revised: N (cid:88) i =

Original: (cid:107)∇ (cid:96) i ( θ 1 ) − ∇ (cid:96)
Revised: (cid:107)∇ (cid:96) i ( θ 1 ) − ∇ (cid:96)

Original: [Equation]
Revised: [Equation]

Original: (cid:107)
Revised: (cid:107)

Original: [Equation]
Revised: [Equation]

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Original: “TI" indicates triangle inequality.
Revised: “TI" indicates triangle inequality.

Original: Therefore, L is at least Λ -smooth.
Revised: Therefore, L is at least Λ -smooth.

Original: A. P ROOF OF L EMMA
Revised: A. P ROOF OF L EMMA

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: Expand summation and regroup:
Revised: Expand summation and regroup:

Original: [Equation]
Revised: [Equation]

Original: Take expectation on each sample:
Revised: Take expectation on each sample:

Original: [Equation]
Revised: [Equation]

Original: Because x i and x j are independent
Revised: Because x i and x j are independent

Original: [Equation]
Revised: [Equation]

Original: the second part can be simpliﬁed as:
Revised: the second part can be simpliﬁed as:

Original: [Equation]
Revised: [Equation]

Original: i ( θ t ) (cid:107) 2 ] + b − 1 b (cid:107)∇L ( θ t ) (cid:107)
Revised: i ( θ t ) (cid:107) 2 ] + b − 1 b (cid:107)∇L ( θ t ) (cid:107)

Original: Given AS 2 and Lemma 3, we have
Revised: Given AS 2 and Lemma 3, we have

Original: (cid:107)∇ (cid:96) i ( θ t ) (cid:107) 2 ≤ 2Λ · (cid:96) i ( θ t ) and (cid:107)∇L ( θ t ) (cid:107) 2 ≤ 2Λ · L ( θ
Revised: (cid:107)∇ (cid:96) i ( θ t ) (cid:107) 2 ≤ 2Λ · (cid:96) i ( θ t ) and (cid:107)∇L ( θ t ) (cid:107) 2 ≤ 2Λ · L ( θ

Original: t
Revised: t

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: A. P ROOF OF T HEOREM
Revised: A. P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: Given L ( θ ) is convex, and H ( θ ) (cid:22) Λ I , L ( θ t ) can be bounded by L ( θ t )
Revised: Given L ( θ ) is convex, and H ( θ ) (cid:22) Λ I , L ( θ t ) can be bounded by L ( θ t )

Original: ≤ L ( θ t − 1 ) + ∇L ( θ t − 1 ) T ( θ t − θ t − 1 ) + Λ2 (cid:107) θ t − θ t − 1 (cid:107) 2 for ∀ θ t − 1 , θ t
Revised: ≤ L ( θ t − 1 ) + ∇L ( θ t − 1 ) T ( θ t − θ t − 1 ) + Λ2 (cid:107) θ t − θ t − 1 (cid:107) 2 for ∀ θ t − 1 , θ t

Original: In SLIM-QN, θ t = θ t − 1 − η t − 1 ˆ H − 1 k ∇L ( θ t − 1 ; S t − 1 ) .
Revised: In SLIM-QN, θ t = θ t − 1 − η t − 1 ˆ H − 1 k ∇L ( θ t − 1 ; S t − 1 ) .

Original: Therefore, we can upper bound L ( θ t ) as:
Revised: Therefore, we can upper bound L ( θ t ) as:

Original: [Equation]
Revised: [Equation]

Original: Since ˆ H − 1 k is independent with S t − 1 , we take expectation w.r.t S t − 1 and S t ,
Revised: Since ˆ H − 1 k is independent with S t − 1 , we take expectation w.r.t S t − 1 and S t ,

Original: [Equation]
Revised: [Equation]

Original: According AS 1, we have
Revised: According AS 1, we have

Original: [Equation]
Revised: [Equation]

Original: According to Lemma 4, we have
Revised: According to Lemma 4, we have

Original: [Equation]
Revised: [Equation]

Original: ≤ 2Λ L
Revised: ≤ 2Λ L

Original: [Equation]
Revised: [Equation]

Original: Therefore, E S t [ L ( θ t )
Revised: Therefore, E S t [ L ( θ t )

Original: |S t − 1 ] ≤ L ( θ t − 1 ) − η t − 1 λξ L ( θ t − 1 ) + η 2 t − 1 Λ 2 Ξ 2 L ( θ t − 1 ) .
Revised: |S t − 1 ] ≤ L ( θ t − 1 ) − η t − 1 λξ L ( θ t − 1 ) + η 2 t − 1 Λ 2 Ξ 2 L ( θ t − 1 ) .

Original: After simply regrouping, we can get
Revised: After simply regrouping, we can get

Original: E S t [ L ( θ t )
Revised: E S t [ L ( θ t )

Original: |S t − 1 ] ≤ (1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 )
Revised: |S t − 1 ] ≤ (1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 )

Original: [ L ( θ t − 1 ]
Revised: [ L ( θ t − 1 ]

Original: Apply total expectation rule w.r.t S t , we have
Revised: Apply total expectation rule w.r.t S t , we have

Original: E S t [ L ( θ t )] ≤ (1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 ) E S t − 1 [ L ( θ t − 1 )]
Revised: E S t [ L ( θ t )] ≤ (1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 ) E S t − 1 [ L ( θ t − 1 )]

Original: B S ETTINGS FOR M ODEL T RAINING ON I MAGE N ET B. R ES N ET - 
Revised: B S ETTINGS FOR M ODEL T RAINING ON I MAGE N ET B. R ES N ET -

Original: We train ResNet-50 on ImageNet in a multi-GPU platform, which has 8 Nvidia Quadro RTX
Revised: We train ResNet-50 on ImageNet in a multi-GPU platform, which has 8 Nvidia Quadro RTX

Original: GPUs.
Revised: GPUs.

Original: PyTorch ( ≥ 1.8) and the Distributed Data Parallel (DDP) communication package is used during training.
Revised: PyTorch ( ≥ 1.8) and the Distributed Data Parallel (DDP) communication package is used during training.

Original: Table 3 lists hyperparameters used in SGD, KFAC, and SLIM-QN.
Revised: Table 3 lists hyperparameters used in SGD, KFAC, and SLIM-QN.

Original: Initial learning rate is 0 . 1 , decaying by a factor of 10 at 30 , 60 , 90 th epoch.
Revised: Initial learning rate is 0 . 1 , decaying by a factor of 10 at 30 , 60 , 90 th epoch.

Original: We use a learning warmup for in the ﬁrstepochs.
Revised: We use a learning warmup for in the ﬁrstepochs.

Original: Batch size is 256.
Revised: Batch size is 256.

Original: In SLIM-QN, we use a smaller weight decay ( wd ) as SLIM-QN considers weight decay when conditioning gradients.
Revised: In SLIM-QN, we use a smaller weight decay ( wd ) as SLIM-QN considers weight decay when conditioning gradients.

Original: Therefore a smaller weight decay can achieve as strong regularization as SGD with wd = 0 . 0005 .
Revised: Therefore a smaller weight decay can achieve as strong regularization as SGD with wd = 0 . 0005 .

Original: For the Hessian update frequency ( L ), in SLIM-QN we update the Hessian for every 30 mini-batch iterations; while for KFAC, considering the compute cost, we update the Hessian for every 50 iterations.
Revised: For the Hessian update frequency ( L ), in SLIM-QN we update the Hessian for every 30 mini-batch iterations; while for KFAC, considering the compute cost, we update the Hessian for every 50 iterations.

Original: Lower and upper threshold for restraining eigenvalues in the Hessian ( σ L , σ H ) is set to be (0 . 01 , 1) in all experiments.
Revised: Lower and upper threshold for restraining eigenvalues in the Hessian ( σ L , σ H ) is set to be (0 . 01 , 1) in all experiments.

Original: Initial damping ( 1 − τ 0 ) is 0 .
Revised: Initial damping ( 1 − τ 0 ) is 0 .

Original: 05 , then adapted during training according to equation 4.
Revised: 05 , then adapted during training according to equation 4.

Original: We run 5 random seeds for each optimizer.
Revised: We run 5 random seeds for each optimizer.

Original: [Table]
Revised: [Table]

Original: B. V I T
Revised: B. V I T

Original: We use a small Vision Transformer model with 6 layers, 8 attention heads, a patch size of 16 , and both hidden and MLP dimension of 512 for a total of about 10 M parameters.
Revised: We use a small Vision Transformer model with 6 layers, 8 attention heads, a patch size of 16 , and both hidden and MLP dimension of 512 for a total of about 10 M parameters.

Original: We train for 90 epochs with linear learning rate warmup in the ﬁrst 5 epochs, and decay the learning rate at 80 epochs for
Revised: We train for 90 epochs with linear learning rate warmup in the ﬁrst 5 epochs, and decay the learning rate at 80 epochs for

Original: SLIM-QN.
Revised: SLIM-QN.

Original: For SGD, we train for 100 epochs, decaying the learning rate at 30 , 60 , 90 epochs.
Revised: For SGD, we train for 100 epochs, decaying the learning rate at 30 , 60 , 90 epochs.

Original: A batch size of 1024 is used for both algorithms.
Revised: A batch size of 1024 is used for both algorithms.

Original: We perform 3 runs with different random seeds.
Revised: We perform 3 runs with different random seeds.

Original: Table 4 shows the selected hyperparameters
Revised: Table 4 shows the selected hyperparameters

Original: [Table] [Table]
Revised: [Table] [Table]

Original: C E XPERIMENTS ON CIFARC. R ES N ET - 
Revised: C E XPERIMENTS ON CIFARC. R ES N ET -

Original: Hyperparameter is shown as in Table 5.
Revised: Hyperparameter is shown as in Table 5.

Original: Initial learning rate is 0 . 1 , decaying by a factor of 10 at 150 th epoch.
Revised: Initial learning rate is 0 . 1 , decaying by a factor of 10 at 150 th epoch.

Original: For both SGD and SLIM-QN we used a linear learning rate warmup for 5 epochs.
Revised: For both SGD and SLIM-QN we used a linear learning rate warmup for 5 epochs.

Original: Batch size is 256 for both SGD and SLIM-QN.
Revised: Batch size is 256 for both SGD and SLIM-QN.

Original: Figure 5 shows convergence on ResNet-18 using SGD and
Revised: Figure 5 shows convergence on ResNet-18 using SGD and

Original: SLIM-QN.
Revised: SLIM-QN.

Original: On small dataset CIFAR-10, SGD and SLIM-QN both deliver fast convergence at early stages, while SLIM-QN is slightly better.
Revised: On small dataset CIFAR-10, SGD and SLIM-QN both deliver fast convergence at early stages, while SLIM-QN is slightly better.

Original: Moreover, SLIM-QN achieves more faster convergence and high validation accuracy at later stages.
Revised: Moreover, SLIM-QN achieves more faster convergence and high validation accuracy at later stages.

Original: C. V I T
Revised: C. V I T

Original: We use the same ViT model as in the ImageNet experiment, but with a patch size of 16 .
Revised: We use the same ViT model as in the ImageNet experiment, but with a patch size of 16 .

Original: Tableshows the hyperparameters used in the CIFAR-10 experiments.
Revised: Tableshows the hyperparameters used in the CIFAR-10 experiments.

Original: For both SGD and SLIM-QN we used a linear learning rate warmup for 5 epochs.
Revised: For both SGD and SLIM-QN we used a linear learning rate warmup for 5 epochs.

Original: We perform 3 runs with different random seeds.
Revised: We perform 3 runs with different random seeds.

Original: [Table]
Revised: [Table]

Original: We use a batch size of 1024 and train for 90 epochs and decay the learning rate by a factor of 10 at 100 epochs for SLIM-QN.
Revised: We use a batch size of 1024 and train for 90 epochs and decay the learning rate by a factor of 10 at 100 epochs for SLIM-QN.

Original: We train for 150 epochs and decay the learning rate at 140 epochs for SGD.
Revised: We train for 150 epochs and decay the learning rate at 140 epochs for SGD.

Original: Experimental results on CIFAR-10 using Vision Transformer are depicted in Fig.
Revised: Experimental results on CIFAR-10 using Vision Transformer are depicted in Fig.

Original: We observe that
Revised: We observe that

Original: SLIM-QN converges to a solution with better generalization than SGD and in overall less iterations.
Revised: SLIM-QN converges to a solution with better generalization than SGD and in overall less iterations.

Original: On small datasets such as CIFAR-10 we do not observe signiﬁcant early-stage speedup on ViT
Revised: On small datasets such as CIFAR-10 we do not observe signiﬁcant early-stage speedup on ViT

Original: [Table]
Revised: [Table]

Original: which is consistent with ResNet-18/CIFAR-10 experiments.
Revised: which is consistent with ResNet-18/CIFAR-10 experiments.

Original: Moreover, ViT is better suited for large-scale vision datasets due to the weaker implicit bias resulting from the non-convolutional architecture.
Revised: Moreover, ViT is better suited for large-scale vision datasets due to the weaker implicit bias resulting from the non-convolutional architecture.

Original: [Table]
Revised: [Table]

Original: D S TOCHASTIC T RAINING USING THE C LASSICAL L-BFGS
Revised: D S TOCHASTIC T RAINING USING THE C LASSICAL L-BFGS

Original: In this section, we demonstrate that the classical L-BFGS suffers convergence instability in stochastic training, even using large batch sizes.
Revised: In this section, we demonstrate that the classical L-BFGS suffers convergence instability in stochastic training, even using large batch sizes.

Original: We train ResNet-18 on CIFAR-10, and vary batch size fromto 2048.
Revised: We train ResNet-18 on CIFAR-10, and vary batch size fromto 2048.

Original: Learning rate decays from 0.1 to 0.001 in 100 epochs.
Revised: Learning rate decays from 0.1 to 0.001 in 100 epochs.

Original: Weight decay is 0.0005.
Revised: Weight decay is 0.0005.

Original: The Hessian approximation is updated using the L-BFGS formula for every 50 iterations, with at most 10 history vectors.
Revised: The Hessian approximation is updated using the L-BFGS formula for every 50 iterations, with at most 10 history vectors.

Original: Fig.
Revised: Fig.

Original: 7 shows training accuracy w.r.t iterations (log scale).
Revised: 7 shows training accuracy w.r.t iterations (log scale).

Original: We can observe that L-BFGS fails to converge under various batch sizes.
Revised: We can observe that L-BFGS fails to converge under various batch sizes.

Original: Training using large batches though is a little more stable than small ones at the beginning, still diverges due to stochastic noises or sudden changes in loss landscape.
Revised: Training using large batches though is a little more stable than small ones at the beginning, still diverges due to stochastic noises or sudden changes in loss landscape.

Original: Therefore, large batch only is not a sufﬁcient solution to address convergence instability in the classical L-BFGS.
Revised: Therefore, large batch only is not a sufﬁcient solution to address convergence instability in the classical L-BFGS.

Original: [Figure]
Revised: [Figure]

Original: E SLIM-QN IN D ISTRIBUTED S YSTEMS
Revised: E SLIM-QN IN D ISTRIBUTED S YSTEMS

Original: As mentioned in Sec. 4.2, SLIM-QN requires O (2 M (cid:107) θ (cid:107) ) storage to store required history statistics s k and y k .
Revised: As mentioned in Sec. 4.2, SLIM-QN requires O (2 M (cid:107) θ (cid:107) ) storage to store required history statistics s k and y k .

Original: When training models using data parallelism, each node further needs to store the whole copy of these history vectors.
Revised: When training models using data parallelism, each node further needs to store the whole copy of these history vectors.

Original: Such high memory footprints make it difﬁcult to be used to train very large models such as  (Devlin et al., 2018;
Revised: Such high memory footprints make it difﬁcult to be used to train very large models such as BERT and GPT (Devlin et al., 2018;

Original: Radford et al., 2018).
Revised: Radford et al., 2018).

Original: To address such a limitation, we propose a block-wise SLIM-QN which is much more memory-efﬁcient in distributed systems.
Revised: To address such a limitation, we propose a block-wise SLIM-QN which is much more memory-efﬁcient in distributed systems.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 8, for a neural network, we divide model parameters into multiple blocks, where each block might consist of one or more layers.
Revised: 8, for a neural network, we divide model parameters into multiple blocks, where each block might consist of one or more layers.

Original: During training, these blocks are optimized in parallel using independent SLIM-QN optimizers.
Revised: During training, these blocks are optimized in parallel using independent SLIM-QN optimizers.

Original: In a distributed system, each compute node can conduct optimization on one or more blocks, depending on its capabilities.
Revised: In a distributed system, each compute node can conduct optimization on one or more blocks, depending on its capabilities.

Original: Furthermore, the model can be divided in a way such that each node can store the required statistics s k and y k for at least one block.
Revised: Furthermore, the model can be divided in a way such that each node can store the required statistics s k and y k for at least one block.

Original: Combining with ZeRO data parallelism design (Rajbhandari et al., 2020), block-wise SLIM-QN are capable of training very large models as long as there are enough compute nodes.
Revised: Combining with ZeRO data parallelism design (Rajbhandari et al., 2020), block-wise SLIM-QN are capable of training very large models as long as there are enough compute nodes.

Original: E. C ONVERGENCE G UARANTEES
Revised: [Figure]

Original: In this section, we prove that block-wise SLIM-QN also converges in a linear rate given assumption in Sec. 4.1.
Revised: In this section, we prove that block-wise SLIM-QN also converges in a linear rate given assumption in Sec. 4.1.

Original: Furthermore, the convergence property is guaranteed in any arbitrary way of dividing models.
Revised: Furthermore, the convergence property is guaranteed in any arbitrary way of dividing models.

Original: Theorem 2.
Revised: Theorem 2.

Original: Assume AS 1-3 hold at each iteration t of block-wise SLIM-QN with mini-batch input
Revised: Assume AS 1-3 hold at each iteration t of block-wise SLIM-QN with mini-batch input

Original: S t , and the k -th Hessian update for block i ( i = 1 , · · · , p ) , ˆ B ik ( ˆ B =
Revised: S t , and the k -th Hessian update for block i ( i = 1 , · · · , p ) , ˆ B i k

Original:  ˆ H − 1 ) during the optimization is bounded by ξ i I (cid:22)
Revised: ( ˆ B = ˆ H − 1 ) during the optimization is bounded by ξ i I (cid:22)

Original: ˆ B ik (cid:22)
Revised: ˆ B ik (cid:22)

Original: Ξ i I , then the expectation of L ( θ t ) satisﬁes
Revised: Ξ i I , then the expectation of L ( θ t ) satisﬁes

Original: [Equation]
Revised: [Equation]

Original: where α t − 1 = 1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 , ξ = min ξ i , Ξ = max Ξ i .
Revised: where α t − 1 = 1 − η t − 1 λξ + η 2 t − 1 Λ 2 Ξ 2 , ξ = min ξ i , Ξ = max Ξ i .

Original: Proof.
Revised: Proof.

Original: The bound of ˆ B ik for block i is guaranteed by SLIM-QN.
Revised: The bound of ˆ B ik for block i is guaranteed by SLIM-QN.

Original: Then according to the proof in
Revised: Then according to the proof in

Original: Theorem 1, if we can prove ˆ H − 1 k = diag

Original: [Equation] ˆ [Equation] 1 , ˆ B 2 , · · · , ˆ B p ) is also bounded, then we can show block-wise SLIM-QN converges in a linear rate.
Revised: Theorem 1, if we can prove ˆ H − 1 k = diag ( ˆ B 1 , ˆ B 2 , · · · , ˆ B p ) is also bounded, then we can show block-wise SLIM-QN converges in a linear rate.

Original: For any arbitrary vector x (cid:54) = 0 , x T · ˆ H − 1 k · x can be written as
Revised: For any arbitrary vector x (cid:54) = 0 , x T · ˆ H − 1 k · x can be written as

Original: [Equation]
Revised: [Equation]

Original: where x i is a sub-vector corresponding to block i .
Revised: where x i is a sub-vector corresponding to block i .

Original: Let ξ = min ξ i and Ξ = max Ξ i , therefore, ξ ≤ x T · ˆ H − 1 k · x ≤ Ξ .
Revised: Let ξ = min ξ i and Ξ = max Ξ i , therefore, ξ ≤ x T · ˆ H − 1 k · x ≤ Ξ .

Original: Following the same proof as in Theorem 1, block-wise SLIM-QN also converges in a linear rate.
Revised: Following the same proof as in Theorem 1, block-wise SLIM-QN also converges in a linear rate.

Original: E. E MPIRICAL A NALYSES
Revised: E. E MPIRICAL A NALYSES

Original: We evaluate block-wise SLIM-QN on ResNet-18/CIFAR-10 using 5 random seeds.
Revised: We evaluate block-wise SLIM-QN on ResNet-18/CIFAR-10 using 5 random seeds.

Original: Layers in a
Revised: Layers in a

Original: ResBlock are grouped into one block.
Revised: ResBlock are grouped into one block.

Original: Initial learning rate is 0.1, decaying by a factor of 10 at 150th epoch.
Revised: Initial learning rate is 0.1, decaying by a factor of 10 at 150th epoch.

Original: We used a linear learning rate warmup for 5 epochs.
Revised: We used a linear learning rate warmup for 5 epochs.

Original: Batch size is 256 for all runs.
Revised: Batch size is 256 for all runs.

Original: Table 7 list detailed hyper parameters settings.
Revised: Table 7 list detailed hyper parameters settings.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 9, block-wise SLIM-QN achieves even slightly faster convergence performance as SLIM-QN.
Revised: 9, block-wise SLIM-QN achieves even slightly faster convergence performance as SLIM-QN.

Original: Both SLIM-QN and the block-wise version achieve higher validation accuracy compared to SGD.
Revised: Both SLIM-QN and the block-wise version achieve higher validation accuracy compared to SGD.

Original: Due to the limited time
Revised: Due to the limited time

Original: [Table]
Revised: [Table]

Original: we have not completed experiments on large models/datasets.
Revised: we have not completed experiments on large models/datasets.

Original: We will add that in the future.
Revised: We will add that in the future

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: F H YPERPARAMETER T UNING
Revised: F H YPERPARAMETER T UNING

Original: SLIM-QN involves additional hyperparameters besides common parameters in SGD: learning rate, momentum [Table] and weight decay.
Revised: SLIM-QN involves additional hyperparameters besides common parameters in SGD: learning rate, momentum, and weight decay.

Original: While tuning the hyperparameters is not a huge burden since some parameters are ﬁxed in all the experiments, such as lower and upper threshold for restraining eigenvalues in the Hessian: σ L , σ H , and length of history vectors: M .
Revised: While tuning the hyperparameters is not a huge burden since some parameters are ﬁxed in all the experiments, such as lower and upper threshold for restraining eigenvalues in the Hessian: σ L , σ H , and length of history vectors: M .

Original: For the rest of parameters: damping ( τ 0 ), momentum ( β 1 , β 2 ) and update frequency ( L ) for the Hessian, it is easy to conduct a grid search on a small models and datasets, explore how these parameters affect optimization, then apply them to large-scale model training.
Revised: For the rest of parameters: damping ( τ 0 ), momentum ( β 1 , β 2 ) and update frequency ( L ) for the Hessian, it is easy to conduct a grid search on a small models and datasets, explore how these parameters affect optimization, then apply them to large-scale model training.

Original: For example, after training on CIFAR-10, we found that small L improves generalization performance, and large damping stabilizes optimization but causes optimizer to behave like SGD.
Revised: For example, after training on CIFAR-10, we found that small L improves generalization performance, and large damping stabilizes optimization but causes optimizer to behave like SGD.

Original: With these notions, it is easy to tune these parameters, and then achieve optimal convergence performance and accuracy.
Revised: With these notions, it is easy to tune these parameters, and then achieve optimal convergence performance and accuracy.