# Large-Scale Deep Learning Optimizations: A Comprehensive Survey

Xiaoxin He<EMAIL_ADDRESS>National University of
<EMAIL_ADDRESS>,  Fuzhao Xue<EMAIL_ADDRESS>National
University of<EMAIL_ADDRESS>,  Xiaozhe Ren<EMAIL_ADDRESS>Huawei Noah’s Ark<EMAIL_ADDRESS>and  Yang You<EMAIL_ADDRESS>National University of<EMAIL_ADDRESS>
###### Abstract.

Deep learning have achieved promising results on a wide spectrum of AI
applications. Larger datasets and models consistently yield better
performance. However, we generally spend longer training time on more
computation and communication. In this survey, we aim to provide a clear
sketch about the optimizations for large-scale deep learning with regard to
the model accuracy and model efficiency. We investigate algorithms that are
most commonly used for optimizing, elaborate the debatable topic of
generalization gap arises in large-batch training, and review the SOTA
strategies in addressing the communication overhead and reducing the memory
footprints.

Deep Learning, Deep Neural Networks, Optimization, Distributed Learning, Large
Batch Training, Communication-Efficient, Memory-Efficient, Survey

††copyright: none

##  1\. Introduction

Nowadays, deep learning (DL) have achieved promising results on a wide
spectrum of AI application domains ranging from computer vision (_e.g.,_ image
classification (He et al., 2016; Huang et al., 2017; Lou et al., 2021), object
detection and segmentation (Girshick, 2015; Ren et al., 2015; He et al., 2020;
Long et al., 2015)), natural language processing (_e.g.,_ language modeling
(Devlin et al., 2019; Xue et al., 2021) and machine translation (Vaswani et
al., 2017; Wu et al., 2016)), information retrieval (_e.g.,_ recommendation
system (He et al., 2017)) and many others. The scale is the main driver behind
the rise of DL (Krizhevsky et al., 2012; He et al., 2016; Simonyan and
Zisserman, 2014; Krizhevsky et al., 2017; Szegedy et al., 2015; Devlin et al.,
2019). Larger datasets and neural networks consistently yield better
performance across all tasks that generally require more computation and
longer training time. Therefore, recent years have witnessed a surge of
interests from both academia and industry in scaling up DL with distributed
training on a large cluster of devices such as TPUs and GPUs with higher
computation capability and memory limit. Data parallelism has become a
dominant practice for distributed training. It distributes a large batch to
multiple devices, where each device holds an identical model replica, computes
the gradient of a local batch and finally gathers the gradients at each
iteration for synchronous parameter update. With recent optimization
techniques, it is now able to train very large batches on thousands of GPU
devices. However, training at such scales requires overcoming both algorithmic
and systems-related challenges. One of the main challenges is the degradation
of model accuracy with large batch size beyond a certain point (e.g., 32k).
Naively increasing the batch size typically results in degradation of
generalization performance and reduces computational benefits. Additionally,
we can not always improve the training speed by just using more processors as
the communication cost is a non-negligible overhead. Intuitively multiple
processors collaboratively training one task can reduce the overall training
time, but the corresponding communication cost between processors is heavy and
limits the model scalibility. Worse still, models with tens of billions to
trillions of parameters clearly do not fit into memory of a single device, and
simply adding more devices will not help scale the training. This limitation
prevents DL researchers from exploring more advanced model architectures.
Existing works investigate and develop optimization techniques to overcome
these problems so as to accelerate training large-scale deep neural networks
(DNNs). We categorise these works into two categories, one endeavors to
maintain/improve the model accuracy in the large-scale setting and the other
emphasises on the model efficiency, designing algorithms that are less hungry
for communication and memory. Importantly, they are not mutually exclusive but
can be used collaboratively to further speed up the training.

###  1.1. Related Surveys

Pouyanfar et al. (2019) give an overview of DL from different perspectives,
including history, challenges, opportunities, algorithms, frameworks,
applications, and parallel and distributed computing techniques. Wang et al.
(2020) provide a quick survey on large-scale distributed deep learning
systems, which concisely introduces parallelisms, parameter server
architectures, synchronization schemes, related applications, and platforms.
While some other surveys focus on a certain scope in deep learning:
communication-efficiency in large-scale parallelism systems (Betzel et al.,
2018; Tang et al., 2020), parallelization strategies (Ben-Nun and Hoefler,
2019) and numerical optimization algorithms (Sun, 2019; Bottou et al., 2018;
Battiti, 1992; Ruder, 2016). Large-scale DL represents a distinctive setting
in which the accuracy, computation, communication and memory are closely
connected and mutually restricted. However, existing surveys either merely
concern part of them or do not address optimizations in the context of large-
scale DL. Different from the two surveys (Wang et al., 2020; Pouyanfar et al.,
2019) that are mostly related to ours, we focus more on the design of the
algorithm rather than the system architecture. The novelty of this paper is
its emphasises on both model accuracy and model efficiency, which captures
critical aspects of large-scale deep learning training by presenting a review
of the state-of-the-art (SOTA) optimization techniques and illustrating the
trade-off in between.

###  1.2. Structure of the Survey

Figure 1. The Overall Structure of this Survey

The overall structure of this survey is presented in Figure 1. Section 2
presents the formulation of a typical neural network optimization problem for
supervised learning. We roughly divide the large-scale DL optimization into
two components: model accuracy and model efficiency. Section 3 introduces the
gradient descent optimization family, including gradient descent variants,
momentum SGD and adaptive gradient algorithms. As large batch training with
data parallelism has increasing popularity in DL and meanwhile introduces
challenges, Section 4 discusses problems in this setting, and reviews main
SOTA optimization strategies to improve the situation. Section 5 dives into
the generalization gap — a debating topic in large batch training. Section 6
introduces second-order optimizations. Then we turn our attention to model
efficiency. Section 7 investigates the communication bottleneck and Section 8
focuses on memory-efficient techniques. Finally, Section 9 concludes this
article.

##  2\. Preliminaries

Following many machine learning applications (Bottou et al., 2018; Xu et al.,
2020; Bottou and Bousquet, 2007), we consider a space of input-output pairs
$(x,y)\in X\times Y$ has a probability distribution $P(x,y)$. The conditional
distribution $P(y|x)$ represents the true relationship between inputs and
outputs. The discrepancy between the predicted output $\hat{y}$ and the real
output $y$ is measured by a smooth but possibly non-convex loss function
$\mathscr{l}(\hat{y},y)$. The objective is to minimize the expected risk

(1) $E(f)=\int\mathscr{l}(f(x),y)P(x,y)=\mathbb{E}[\mathscr{l}(f(x),y)],$

that is,

(2) $f^{*}(x)=\mathop{argmin}_{\hat{y}}\mathbb{E}[\mathscr{l}(\hat{y},y)|x].$

Since $\mathbb{P}$ is an unknown distribution, in practice, one seeks the
solution of a problem that involves an estimate of the empirical risk (Shalev-
Shwartz and Ben-David, 2014)

(3)
$E_{n}(f)=\frac{1}{n}\sum_{i=1}^{n}\mathscr{l}(f(x_{i}),y_{i})=\mathbb{E}_{n}[\mathscr{l}(f(x),y].$

The goal of solving Eq.[3] is to obtain a solution with small generalization
error, i.e., high predictive accuracy on unseen data.

##  3\. Gradient Descent Optimization Algorithms

Training a DNN is an optimization process, i.e., finding the parameters in the
network that minimize the loss function. Gradient descent and its variant
algorithms are the most popular algorithms to optimize neural networks (NNs)
(Ruder, 2016). In order to control the oscillation of gradient descent
methods, the idea of using momentum is introduced. Moreover, adapting the
learning rate w.r.t. the gradient of the previous stages is found beneficial
to avoid the fluctuation. In this section, we briefly sort out the mainstream
optimization algorithms, consisting of gradient descent variants (Section
3.1), momentum (Section 3.2) and adaptive gradient algorithms (Section 3.3).

###  3.1. Gradient Descent Variants

Gradient descent and its variants dominate the optimization algorithms of DL.
The gradient descent (GD) methods aim to minimize the empirical risk of a
model by repeatedly computing the gradient of a loss function on a single
training sample, or a (full) batch of samples, and continuously updating the
model parameters accordingly by following the gradient of the objective
function in the opposite direction. There are three variants in gradient
descent which differ in the number of samples used for each step (updating
model parameters), resulting in different accuracy and learning time.

####  3.1.1. Batch Gradient Descent

Batch gradient descent, a.k.a. vanilla gradient descent, minimizes the loss
function $L(x)$ with the following form:

(4) $L(x)=\frac{1}{|S|}\sum_{s\in S}l(x,s).$

Here $x$ is the weights of a network, $S$ is a labeled training set, $|S|$ is
the number of samples in the training set, and $l(x,s)$ is the loss computed
from sample $s\in S$ and and its label $y$. Typically $l$ is the sum of a
classification loss (e.g., cross-entropy) and a regularization loss on $x$.
And then update the weights

(5) $x=x-\eta\nabla L(x),$

where $\eta$ is a learning rate (LR) which controls how large of a step to
take in the opposite direction of the gradient. As we need to go through the
whole training set to calculate the gradient for one update of weights, batch
gradient descent can be very slow, especially for large datasets (which is
very common in DL tasks). Batch gradient descent also does not allow updating
model online, i.e., with new examples on-the-fly.

####  3.1.2. Stochastic Gradient Descent

Unlike batch gradient descent (GD), which calculates the gradients using the
all training samples, Stochastic Gradient Descent (SGD) performs one weights
update for each training sample

(6) $x=x-\eta\nabla l(x,s).$

Therefore, SGD addresses the computational bottleneck of batch gradient
descent. It is significantly faster than batch gradient descent and can be
used online. The drawback is that the gradient computed from just one sample
is not representative enough for the whole training set. Consequently, the
variance of gradients leads to a fierce fluctuation in the loss function.

####  3.1.3. Mini-batch Stochastic Gradient Descent

Mini-batch SGD takes both advantages of batch GD and SGD by performing weights
update for each mini-batch $B$

(7) $x=x-\frac{\eta}{|B|}\nabla\sum_{s\in B}l(x,s).$

In batch GD, the gradients are computed over the entire dataset, providing an
accurate estimation of the true gradient. It takes lots of time and memory to
do that. But the real handicap is the batch gradient trajectory lands in a bad
spot. While in SGD, parameters are updated by adding the gradient computed on
a single sample of the dataset, which is very noisy and may go off in a
direction far from the batch gradient. However, the noisiness is exactly what
we want in non-convex optimization, because it helps to escape from saddle
points or local minima (Ge et al., 2015). The disadvantage is its terribly
inefficiency of looping over the entire dataset many times to find a good
solution. The mini-batch methodology is a compromise that injects enough noise
to each gradient update, while achieving a relative speedy convergence. Mini-
batch SGD is found to be very effective in the case of large-scale learning
(Bottou and Bousquet, 2007).

###  3.2. Momentum

SGD has trouble navigating a long and narrow valley in the loss function
surface, i.e., the direction of the gradient is almost perpendicular to the
long axis of the valley. In such a situation, the system oscillates back and
forth in the direction of the short axis, and only moves very slowly along the
long axis of the valley.

####  3.2.1. Momentum SGD

Momentum (Qian, 1999) is a strategy that helps to counteract the oscillation
along the short axis meanwhile accumulate contributions along the long axis.
In other words, the momentum strengthens for dimensions whose gradients point
in the same directions and dampens updates for dimensions whose gradients
change directions. This allows Momentum to minimize the training loss in fewer
steps than full batch gradient descent (Park et al., 2019). Specifically,
momentum SGD adds update in previous step to the current update, and
determines the next update $v_{t}$ as a linear combination of the gradient and
the previous update $v_{t-1}$:

(8) $\begin{split}v_{t}&=\beta v_{t-1}+\eta\nabla L(x)\\\
x&=x-v_{t}.\end{split}$

####  3.2.2. Nesterov Accelerated Gradient

In Eq.[8], we know that we are going to move by at least $\beta v_{t-1}$ and a
bit more by $\eta\nabla L(x)$. And in Nesterov Accelerated Gradient (NAG)
(Nesterov, 1983), it looks ahead by calculating the gradient at the partially
updated value of $(x-\beta v_{t-1})$ instead of using the current value:

(9) $\begin{split}v_{t}&=\beta v_{t-1}+\eta\nabla L(x-\beta v_{t-1})\\\
x&=x-v_{t}.\end{split}$

Such ”look before you leap” prevents us from going too fast and results in
increased responsiveness. While the optimization path taken by classical
momentum SGD exhibits large oscillations along the high-curvature vertical
direction, NAG is able to avoid these oscillations almost entirely (Sutskever
et al., 2013a).

###  3.3. Adaptive Gradient Algorithms

The methods mentioned above apply a same LR to all dimensions of the
parameters. Since each dimension of parameters relates to the loss function in
different ways, a per-dimension LR is more advantageous due to the more
accurate and precise control on the step size. Therefore, a variety of
adaptive gradient-based methods have been proposed where gradients are divided
by the component-wise accumulation of previous gradients. For example, AdaGrad
(Duchi et al., 2011) uses the sum of the squares of all past gradients,
whereas Adadelta (Zeiler, 2012), RMSProp (Tieleman and Hinton, 2012) and Adam
(Kingma and Ba, 2017) use an exponentially decaying average.

####  3.3.1. AdaGrad

Previously, we perform updates for all parameters using a same LR, regardless
of their frequency and magnitude. This may lead to a failure in capturing the
knowledge of infrequently occurring updates (which are highly informative and
discriminative). AdaGrad (Duchi et al., 2011) alleviates this problem by
performing larger updates for infrequent parameters and smaller updates for
frequent parameters, which enables it to do well with sparse gradients

(10) $x_{t,i}=x_{t-1,i}-\frac{\eta}{\sqrt{G_{t,ii}+\epsilon}}\nabla
L(x_{t,i}).$

$G_{t}=\sum_{\tau=1}^{T}g_{\tau}g_{\tau}^{T}$ is a diagonal matrix where each
diagonal element $G_{t,ii}$ is the sum of the squares of all past gradient
w.r.t. $x_{i}$ up to time step $t$. And $\epsilon$ is a smoothing term to
avoid division by zero. A vectorized implementation has the following form:

(11) $x_{t}=x_{t-1}-\frac{\eta}{\sqrt{G_{t}+\epsilon}}\odot\nabla L(x_{t}),$

where $\odot$ is an element-wise matrix-vector multiplication. In AdaGrad,
each dimension has its own dynamic LR rate which is inversely dependent to the
gradient magnitude, thus larger gradients have smaller LRs and small gradients
have larger LRs. This is very beneficial for training DNNs since the scale of
gradients in each layer is often different by several orders of magnitude. In
addition, this accumulation of gradients can be regarded as a kind of
simulated annealing which reduces the LRs along the course of training. Most
implementation set the LR $\eta$ to a default value of 0.01, eliminating the
need of manual tuning. However, AdaGrad holds a main drawback with the
accumulation of squares of all past gradients, which keeps growing during the
course of training. As the LRs radically shrink and vanish, the algorithm no
longer gain additional knowledge.

####  3.3.2. Adadelta

Adadelta (Zeiler, 2012) is an extension of AdaGrad (Duchi et al., 2011) that
seeks to tackle its monotonically decreasing LRs. Instead of accumulating the
sum of all previous squared gradients, Adadelta uses an exponentially decaying
average instead. The running average $E[g^{2}]_{t}$ at time step $t$ depends
on the previous average and the current gradient

(12) $E[g^{2}]_{t}=\rho E[g^{2}]_{t-1}+(1-\rho)g_{t}^{2},$

where $\rho$ is a decay constant similar to that used in the momentum method.
As the denominator is the root mean squared (RMS) error criterion of the
gradient, we can replace it with the criterion short-hand:

(13) $\begin{split}RMS[g]_{t}&=\sqrt{E[g^{2}]_{t}+\epsilon}\\\ \Delta
x_{t}&=-\frac{\eta}{RMS[g]_{t}}g_{t}.\end{split}$

Noticing the mismatch of units in Eq.[13], i.e., the units of the update
$\Delta x$ do not match the units of the parameters $x$ which it applies to

(14) $units\,of\,\Delta x\propto units\,of\,g\propto\frac{\partial f}{\partial
x}\propto\frac{1}{units\,of\,x},$

Zeiler (2012) rearranges second order method (i.e., Newton’s method)

(15) $\Delta x=\frac{\frac{\partial f}{\partial
x}}{\frac{\partial^{2}f}{\partial
x^{2}}}\Rightarrow\frac{1}{\frac{\partial^{2}f}{\partial x^{2}}}=\frac{\Delta
x}{\frac{\partial f}{\partial x}}.$

Since $\Delta x_{t}$ for the current time step in unknown, assuming the
curvature is locally smooth, $\Delta x_{t}$ can be approximated by computing
the exponentially decaying RMS of previous $\Delta x$

(16) $\Delta x_{t}=-\frac{RMS[\Delta x]_{t-1}}{RMS[g]_{t}}g_{t}.$

####  3.3.3. RMSProp

RMSprop (Tieleman and Hinton, 2012) was developed independently around the
same time with Adadelta to solve the problem of AdaGrad’s drastically
decreasing gradients. AdaGrad treats all past gradients equally, which is
counter to our intuition that fresh gradient is more informative than the
elder one. RMSProp redefines $v_{t}$ by decaying the past gradients at an
exponential rate

(17) $\begin{split}v_{t}=0.9v_{t-1}+0.1g_{t}^{2}\\\
x_{t}=x_{t-1}-\frac{\eta}{\sqrt{v_{t}+\epsilon}}g_{t}.\end{split}$

####  3.3.4. Adam

Adam (Kingma and Ba, 2017) is one of the most popular optimizers for training
DNNs nowadays. It computes individual LRs for different parameters base on the
estimates of first and second moments of the gradients. In particular, Adam
stores an exponentially moving average of past gradients ($m_{t}$) and squared
gradients ($v_{t}$). The former is an estimate of the first momentum (the
mean) and the latter is an estimate of the second momentum (the uncentered
variance) of the gradients

(18) $\begin{split}m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t}\\\
v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})g_{t}^{2}.\end{split}$

$\beta_{1}$ and $\beta_{2}$ are hyper-parameters controlling the decaying
rates of theses moving averages. Since the moving averages are initialed as
0’s, the estimates of first and second moments are biased towards zero,
especially in the beginning of training. Adam utilizes correction terms to
counteract the initialization bias

(19) $\begin{split}\hat{m}_{t}=\frac{m_{t}}{1-\beta_{1}^{t}}\\\
\hat{v}_{t}=\frac{v_{t}}{1-\beta_{2}^{t}}.\end{split}$

Then Adam applies the update rule

(20) $x_{t}=x_{t-1}-\frac{\eta}{\sqrt{\hat{v}_{t}+\epsilon}}\hat{m}_{t}.$

Adam is found to be robust and well-suited to a wide range of non-convex
optimization problems in the field of DL. There are several variants of Adam.
AdaMax (Kingma and Ba, 2017) is an extension to Adam that generalizes the
approach to the infinite norm (max) and may result in a more effective
optimization on some problems. Nesterov-accelerated Adaptive Moment Estimation
(NAdam) (Dozat, 2016) incorporates NAG into Adam. It shows better convergence
speed in some cases. While these algorithms have been successfully employed in
several practical applications, they may fail to converge to optimal solution
even in convex setting, or even diverge in DL training. Reddi et al. (2019)
pinpoint the exponential moving average of past squared gradients as a reason
for such failures. Recall that the introduction of the exponential average was
well-motivated to tackle the key flaw of the Adagrad algorithm: it should
prevent the LRs to become infinitesimally small as training progresses by
limiting the reliance of the update on essentially only the past few
gradients. However, this short-term memory of the gradients can indeed cause
significant convergence issues in other scenarios. To resolve this issue, the
authors propose new variants of Adam — AMSGrad, which relies on long-term
memory of past gradients. AMSGrad uses the maximum of past squared gradients
rather than the exponential average to update the parameters. Liu et al.
(2019) argue that the root cause of the bad convergence problem suffered by
Adam is that the adaptive LR has undesirably large variance in the early stage
of model training, due to the limited amount of training samples being used.
Thus, to reduce such variance, it is better to use smaller LRs in the first
few epochs of training. The authors propose Rectified Adam (RAdam) to rectify
the variance of the adaptive LR.

Choosing an optimizer is a crucial step when training DNNs since it is woven
with the training speed and the final predictive performance. Despite the fact
that adaptive optimization methods, including AdaGrad, RMSProp, AdaDelat and
Adam, are becoming increasingly popular, to date, how to choose an optimal one
is still theoretically elusive and intractable. Instead practitioners rely on
empirical studies (Wilson et al., 2017) and bench-marking (Schneider et al.,
2019). Wilson et al. (2017) observed that the solutions found by adaptive
methods generalize worse (often significantly worse) than SGD, even when these
solutions have better training performance. However, Choi et al. (2019)
suggest that popular adaptive gradient methods never under-perform momentum or
gradient descent. They point out the comparisons among optimizers are
sensitive to the hyper-parameter tuning protocols.

##  4\. Large Batch Training

Input: Dataset $X$,minibatch size $b$ per node, the number of nodes $N$,
optimization function SGD, init parameters $w={w[0],\cdots,w[M]}$

for  _$t=0,1,\cdots$_ do

$G_{t}^{k}\leftarrow 0$;

for  _$i=1,\cdots,B$_ do

Sample data $x$ from $X$;

$G_{t}^{k}\leftarrow G_{t}^{k}+\frac{1}{Nb}\nabla f(x;w_{t})$

end for

All-Reduce $G_{t}^{k}:G_{t}\leftarrow\sum_{k=1}^{N}G_{t}^{k}$;

$w_{t+1}\leftarrow\textit{SGD}(w_{t},G_{t})$

end for

Algorithm 1 Distributed Synchronous SGD on Node k.

Large DNNs and large datasets have fueled the development of deep learning
(Krizhevsky et al., 2012; He et al., 2016; Simonyan and Zisserman, 2014;
Krizhevsky et al., 2017; Szegedy et al., 2015; Devlin et al., 2019). However,
training large models on massive datasets is compute-intensive. For instance,
training the SOTA DL models like BERT and ResNet-50 takes 3 days on 16 TPUv3
chips and 29 hours on 8 Tesla P100 gpus respectively (Devlin et al., 2019; He
et al., 2016). An intuitive way to accelerate training is to add more
computational power (e.g., more GPU nodes) and use data parallel (see Alg.1).
Considering communication (i.e., synchronizing the updates at each iteration)
is an issue, each GPU must be utilized as much as possible to amortize the
communication cost. Therefore, large batch should be used to distribute more
data to each GPU. The nontrivial growth of batch size often results in test
performance degradation, as observed in (Krizhevsky, 2014; Keskar et al.,
2017; Li et al., 2014b; Hoffer et al., 2017). We describe the training
difficulties introduced by large batch in Section 4.1, a recipe for large
batch training (i.e., linear LR scaling with a warmup strategy) in Section
4.2, other supplementary strategies such as adaptive layer-wise learning in
Section 4.3 and adaptive batch size in Section 4.4, and finally discuss the
extent to which we can scale up the batch size in Section 4.5.

###  4.1. Large Batch Training Difficulties

Although large batches are preferable to increase the parallelism by
distributing the workload to multiple nodes, they may slow down convergence
rate in practice (Byrd et al., 2012). Empirically, an increase in mini-batch
size after a certain point (e.g. 1024) without a careful optimization scheme
typically decreases the rate of convergence. The test accuracy of the
converged solution becomes significantly lower than the baseline (Goyal et
al., 2018; Keskar et al., 2017; Hoffer et al., 2017; Li et al., 2014b). In
addition to a degradation of the test performance, Masters and Luschi (2018)
provide evidence that increasing the batch size also results in a
progressively smaller range of LRs that allows stable training.

Keskar et al. (2017) find a drop in generalization (often denoted as
generalization gap) to be as high as 5% even for smaller networks, and
correlate the generalization gap with the sharpness of the loss landscape.
They argue that large-batch methods tend to converge to sharp minimizers of
the training and testing functions, whereas small-batch methods consistently
converge to flat minimizers. Hoffer et al. (2017) deny the existence of
inherent generalization gap and suggest that training longer will help the
algorithm to generalize better and keep the accuracy higher. Goyal et al.
(2018) admit that large batches cause optimization difficulties, but when
these are addressed the trained networks exhibit good generalization. They
tried to bridge the generalization gap with heuristics of LR scaling (Goyal et
al., 2018) with a warpup strategy. However, empirical study (Shallue et al.,
2019) shows that LR scaling heuristics with the batch size do not hold across
all problems or across all batch sizes. Later You et al. (2017) proposed
Layer-wise Adaptive Rate Scaling (LARS) to solve the large batch optimization
difficulties. Several recent works successfully scaled the batch size to large
values using adaptive learning rates without degrading the performance.

###  4.2. Learning Rate Scaling for Large Batch

A nice property of large batch is its lower variance of the gradient. This is
because when we take the gradient over more examples, the variance is
obviously lower. Consequently, large batch allows us to take a larger step per
iteration. Followings are two commonly used LR heuristics: linear scaling and
sqrt scaling, to guide us to adapt the LR for large batches.

####  4.2.1. Linear Scaling

(Krizhevsky, 2014; Goyal et al., 2018; Bottou et al., 2018) suggest linearly
scaling up LR with batch size, i.e., when the mini-batch size is multiplied by
$k$, multiply the LR by $k$. Intuitively, after $k$ iterations of mini-batch
SGD , we have

(21) $x_{t+k}=x_{t}-\eta\frac{1}{|S|}\sum_{i<k}\sum_{s\in S}\nabla
l(x_{t+i},s),$

while after one iteration of large mini-batch $\bigcup_{j}B_{j}$ of size
$|S|=k|B|$ we have

(22) $\hat{x}_{t+1}=x_{t}-\hat{\eta}\frac{1}{k|B|}\sum_{j<k}\sum_{s\in
B_{j}}\nabla l(x_{t},s).$

If we assume $\nabla l(x_{t+i})\approx\nabla l(x_{t})$ for $i<k$, then the
adjustment $\hat{\eta}=k\eta$ would yield $\hat{x}_{t+1}\approx x_{t+k}$.
Noted that this assumption holds with the premises: (1) $k$ cannot be
infinite. That is, we cannot scale up the batch size without limits; (2) $t$
cannot be too small. Because at the beginning of training, the gradients
change rapidly, and thus the difference between $\nabla l(x_{t})$ and $\nabla
l(x_{t+i})$ is no longer negligible. Using LR warmup and linear scaling, Goyal
et al. (2018) trained Resnet-50 with batch B=8K without loss in accuracy.

####  4.2.2. Sqrt Scaling

Another scaling strategy is sqrt scaling, i.e., when the mini-batch size is
multiplied by $k$, multiply the LR by $\sqrt{k}$. In SGD, the co-variance
matrix of the parameters update $\Delta x$ is (Hoffer et al., 2017)

(23) $cov(\Delta x,\Delta
x)\thickapprox\frac{\eta^{2}}{|B|}(\frac{1}{N}\sum_{n=1}^{N}g_{n}g_{n}^{T}).$

A simple way to keep this co-variance constant when we change the batch size
is to choose $\eta\propto\sqrt{|B|}$. Hoffer et al. (2017) find that by using
”Ghost Batch Normalization” and sqrt scaling, the generalization gap can be
significantly decreased. However, the largest batch size used was 4,096, which
does not rule out an effect appearing at still larger batch sizes, as
suggested by the work of Goyal et al. (2018). Moreover, establishing this
invariant co-variance remains poorly justified, and often sqrt scaling is
found to degrade model quality in practice, see (Krizhevsky, 2014; Goyal et
al., 2018; Jastrzebski et al., 2017).

####  4.2.3. Warmup

After adjusting the LR with these strategies, the main obstacle for scaling up
batch size is the instability of training with high LR, especially in the
initial epochs when the gradients change dramatically. This issue can be
alleviated by a properly designed warmup strategy by using less aggressive LRs
in the initial epochs.

Constant warmup. Verma et al. (2017) use a low ”safe” constant LR for the
first few epochs of training and after that return to the target LR
$\hat{\eta}=k\eta$. Goyal et al. (2018) find constant warmup particularly
helpful for prototyping object detection and segmentation methods (Girshick,
2015; Ren et al., 2015; He et al., 2020), but not sufficient enough to solve
the large batch optimization problem. In particular, a transition out of the
low LR warmup phase can cause the training error to spike. This motivates them
to use a more moderate warmup stragegy — gradual warmup.

Gradual warmup. Unlike constant warmup, gradual warmup avoid a sudden increase
of LR by gradually arising the LR from a small to a large value. We denote the
LR of the $t$-th iteration as $lr(t)$ and the maximum LR during training as
$lr_{max}$. Given a predefined time frame $T_{warmup}$, the LR scheduler for
the $t$-th iterations is defined as

(24) $lr(t)=\frac{t}{T_{warmup}}lr_{max},\quad t\leq T_{warmup}.$

After this warmup stage, the LR will be set by classical LR schedulers (e.g.,
cosine decay). A LR warmup stage is proved to be beneficial when training NNs
with extremely large batch size (You et al., 2020; Goyal et al., 2018). Liu et
al. (2019) claim that the benefit of the warmup stage comes from reducing the
variance for the adaptive LR in the Adam optimizer. They further propose
Rectified Adam (RAdam) by introducing a term to rectify the variance of the
adaptive LR. Additionally, Xiong et al. (2020) find the LR warm-up stage also
helps quite a lot for other optimizers.

###  4.3. Adaptive Layerwise Learning

Linear/Sqrt LR scaling with warmup mitigates the vulnerability to the
fluctuation of gradients in the initial epoch by taking less aggressive steps,
starting from a small LR which is safe enough for all layers and gradually
increasing it to the target value. Hoffer et al. (2017) use less aggressive
sqrt scaling with ”Ghost Batch Normalization” to train Alexnet with $B=8K$,
but still the accuracy ($53.93\%$) was much worse than baseline ($57.10\%$).
Goyal et al. (2018) use LR warmup and linear scaling to train Resnet-50 with
batch B=8K without loss in accuracy. While these works demonstrate the
feasibility of these strategies for reducing the wall time for training large
DNNs, they are not general enough if we want further enlarge the batch size.
For instance, You et al. (2017) applied linear scaling and warmup scheme to
train Alexnet with batch normalization on Imagenet, and observed a $2.2\%$
drop when $B=8K$ in the test accuracy. You et al. (2017) explain their method
to solve this problem: To analyze the training stability with large LRs we
measured the ratio between the norm of the layer weights and norm of gradients
update. We observed that if this ratio is too high, the training may become
unstable. On the other hand, if the ratio is too small, then weights don’t
change fast enough. This ratio works like a hint about how to adapt the LR for
each layer. In this section, we will first introduce a general adaptive
layerwise strategy motivated by this ratio, followed by two specific
algorithms, LARS (You et al., 2017) and LAMB (You et al., 2020).

####  4.3.1. General Layerwise Strategy

Suppose we use an iterative base algorithm $\mathscr{A}$ (e.g., SGD or Adam)
in the small batch setting with the following layerwise update rule

(25) $x_{t+1}=x_{t}+\eta_{t}u_{t},$

where $u_{t}$ is the update made by $\mathscr{A}$ at time step $t$. You et al.
(2020) propose the following two changes to the update for large batch
settings:

  1. (1)

The update is normalized to unit $l_{2}$-norm. This is ensured by modifying
the update to the form $u_{t}/\lVert u_{t}\rVert$. Such a normalization is
done layer-wise, i.e., the update for each layer is ensured to be unit
$l_{2}$-norm.

  2. (2)

The LR is scaled by $\phi(\lVert x_{t}\rVert)$ for some function
$\phi:\mathbb{R}^{+}\rightarrow\mathbb{R}^{+}$. Similar to the normalization,
such a scaling is done layer-wise.

Suppose the base algorithm $\mathscr{A}$ is SGD, then the modification results
in the following update rule

(26)
$x_{t+1}^{(i)}=x_{t}^{(i)}-\eta_{t}\frac{\lVert\phi(x_{t}^{(i)})\rVert}{\lVert
g_{t}^{(i)}\rVert}g_{t}^{(i)}$

for all layers $i\in[h]$. The normalization modification $g_{t}^{(i)}/\lVert
g_{t}^{(i)}\rVert$ is similar to one typically used in normalized gradient
descent except that it is done layer-wise. Normalization of this form provides
robustness to exploding/vanishing gradients (where the gradient can be
arbitrarily large/small) by essentially ignoring the size of the gradient but
preserving the direction. As for the scaling step, the scaling term involving
$\phi$ ensures that the norm of the update is of the same order as that of the
parameter. When the parameters are small, we take a small step and vice versa.

There are two notable differences between this general strategy and other
adaptive algorithms such as Adam or RMSProp: (1) it uses a separate LR for
each layer and not for each weight. (2) the magnitude of the update is
controlled w.r.t the weight norm for better control of training speed. Both
LARS (You et al., 2017) and LAMB (You et al., 2020) are based on this general
strategy, using momentum and Adam optimizer as the base algorithm
respectively.

####  4.3.2. LARS

The first instantiation of the general strategy is the LARS algorithm (You et
al., 2017), which is obtained by using momentum optimizer as the base
algorithm $\mathscr{A}$ in the framework. LARS stands for Layer-wise Adaptive
Rate Scaling, which was proposed for large batch learning for ResNet on
ImageNet. Specifically, a local LR $\lambda^{l}$ is defined for each layer $l$

(27) $\lambda^{l}=\eta\frac{{\lVert x\rVert}_{2}^{2}}{{\lVert\nabla
L(x)\rVert}_{2}^{2}}.$

The hyper-parameter $\eta<1$ describes the extent to which we can trust the
layer to update its weights during each epoch. At the beginning of training,
the numerator ${\lVert x\rVert}_{2}^{2}$ above is relatively small. In
contrast, the denominator ${\lVert\nabla L(x)\rVert}_{2}^{2}$ is probably
large since when everything is wrong, the loss and gradients are large. Any
steps we take are likely to be small. In this way we naturally warm up as the
weights increase. As we approach 0 loss, the gradients become smaller and the
local LR increases again, encouraging jumping out of the local minima to
prevent over-fitting. The parameter update is

(28) $\Delta x_{t}^{l}=\gamma*\lambda^{l}*\nabla L(x_{t}^{l}),$

where $\lambda$ is the global LR. In this way, each layer can learn at its own
pace accurately. The training for SGD with LARS are summarized in the
Algorithm 2.

Input: base LR $\gamma_{0}$, momentum $m1$, weight decay $\beta$, ”trust”
coefficient $\eta$, number of steps $T$

Init: $t=0$; $v=0$. Init weight $w_{0}^{l}$ for each layer $l$

while  _$t <T$ for each layer $l$_ do

$g_{t}^{l}\leftarrow\nabla L(w_{t}^{l})$ ;

//obtain a stochastic gradient for the current mini-batch

$\gamma_{t}\leftarrow\gamma_{0}*(1-\frac{t}{T})^{2}$;

//compute the global LR

$\lambda^{l}\leftarrow\frac{\lVert w_{t}^{l}\rVert}{\lVert
g_{t}^{l}\rVert+\beta\lVert w_{t}^{l}\rVert}$;

//compute the local LR

$v_{t+1}^{l}\leftarrow mv_{t}^{l}+\gamma_{t}*\lambda^{l}*(g_{t}^{l}+\beta
w_{t}^{l})$;

//update the momentum

$w_{t+1}^{l}\leftarrow w_{t}^{l}-v_{t+1}^{l}$ ;

//update the weights

end while

Example with weight decay, momentum and polynomial LR decay.

Algorithm 2 SGD with LARS.

Example with weight decay, momentum and polynomial LR decay.

Several works successfully scaled the batch size to large values using LARS
without degrading the performance, thereby, finishing ResNet-50 training on
ImageNet in a few minutes (You et al., 2017; Ying et al., 2018; Yamazaki et
al., 2019). LARS also applies to tasks such as self-supervised image
representation learning and contrastive learning of visual representations
(Grill et al., 2020; Chen et al., 2020).

####  4.3.3. LAMB

LAMB is the second instantiation of the general strategy, which is obtained by
using Adam as the base algorithm $\mathscr{A}$. The pseudo-code is provided in
Algorithm 3. The adaptivity of LAMB is two-fold: (1) per dimension
normalization w.r.t the square root of the second moment used in Adam and (2)
layer-wise normalization obtained due to layer-wise adaptivity. By using LAMB,
You et al. (2020) scale the batch size of BERT pre-training to 64K without
losing accuracy, thereby, reducing the BERT training time from 3 days to
around 76 minutes. LAMB is also the first large batch adaptive solver that can
achieve the SOTA accuracy on ImageNet training with RESNET-50. LAMB has also
been adopted by many other work (Lan et al., 2019).

Despite of the popularity of LARS and LAMB, their utility as a ”large batch
optimizer” is challenged by (Nado et al., 2021), which argues that they are
more indirect regularizers than optimizers. By sophisticated tuning,
traditional, generic algorithms (e.g., Momentum or Adam) achieve strong
results across batch size. They appeal to researchers that the superiority of
one particular optimizer over others should be claimed with extreme caution
since the fair comparisons between optimizers crucially depend on the effort
spent tuning hyperparameters for each optimizer.

Input: $x_{1}\in\mathbb{R}^{d}$, LR $\\{\eta_{t}\\}_{t=1}^{T}$, parameters
$0<\eta_{1},\eta_{2}<1$, scaling function $\phi$, $\epsilon>0$

Init: Set $m_{0}=0,\,v_{0}=0$

for  _$t=1$ to T_ do

$g_{t}=\nabla L(x_{t})$;

//obtain a stochastic gradient for the current mini-batch

$m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t}$;

$v_{t}=\beta_{2}v_{t-1}+(1-\beta_{2})g_{t}^{2}$;

$m_{t}=m_{t}/(1-\beta_{1}^{t})$;

$v_{t}=v_{t}/(1-\beta_{2}^{t})$;

$r_{t}=\frac{m_{t}}{\sqrt{v_{t}}+\epsilon}$;

$x_{t+1}^{(i)}=x_{t}^{(i)}-\eta_{t}\frac{\phi(\lVert
x_{t}^{(i)}\rVert)}{\lVert r_{t}^{(i)}+{\lambda
x_{t}^{(i)}}\rVert}(r_{t}^{(i)}+{\lambda x_{t}^{(i)}})$

end for

Algorithm 3 LAMB

###  4.4. Adaptive Batch Size

It is a common practice to decay the LR during training. When one decays the
LR, one simultaneously decays the ”noise scale”, i.e., the scale of random
fluctuations in the SGD dynamics (Smith and Le, 2018)

(29) $\begin{split}g&=\frac{\epsilon}{1-m}(\frac{N}{B}-1)\\\
&\thickapprox\frac{\epsilon N}{(1-m)B}.\end{split}$

When we decay the LR, the ”noise scale” falls, enabling us to converge to the
minima of the loss function. We can achieve the same reduction in noise scale
at constant LR by increasing the batch size. Smith et al. (2018) and
Devarakonda et al. (2017) empirically demonstrated the equivalence between
decaying LR and increasing the batch size. Instead of decaying the LR by a
factor of $\alpha$, they increase the batch size by $\alpha$ during training.
This strategy reaches equivalent test accuracy after the same number of
training epochs, but with fewer parameter updates, leading to greater
parallelism and shorter training times. Crucially, such strategy is
complementary to existing training schedules requiring no hyper-parameter
tuning.

###  4.5. Efficient Scaling

Increasing the batch size is one of the most appealing ways to accelerate NN
training on data parallel hardware. Ideally, parallel mini-batch SGD can
achieve a linear speed-up of the training time w.r.t. the number of workers
compared with SGD over a single worker. However, such linear scalability in
practice is significantly limited by the growing demand for gradient
communication as more workers are involved. Moreover, when batch very large,
the stochastic gradients become very close to true gradients, so increasing
the batch does not give much additional gradient information comparing to
smaller batches.

A series of work has conducted comprehensive experiments on the relationship
between batch size and training time for NNs (Shallue et al., 2019; Smith et
al., 2020; Zhang et al., 2019). Shallue et al. (2019) experimentally measure
the effects of data parallelism training across different families of NNs,
training algorithms and data sets, finding no evidence that larger batch sizes
degrade out-of-sample performance. They observed three distinct scaling
regimes in the relationship between batch size and training time: a ”perfect
scaling” regime where doubling the batch size halves the number of training
steps required to reach a target out-of-sample error, followed by a regime of
”diminishing returns”, and finally a ”maximal data parallelism” regime where
further increasing the batch size does not reduce training time, even assuming
idealized hardware. They also provide experimental evidence that the critical
batch size depends on the model architecture, the dataset and regulation
technology.

##  5\. Generalization Gap

Optimization in general is an extremely difficult task, especially for
training NNs. With non-convex and high-dimensional functions, it is possible
to have many local minima and saddle points. Optimization methods, such as
SGD, generally converge to different regions of parameter space, highly
dependent on the design of network architecture, the choice of optimizer,
variable initialization, and a variety of other considerations (Shallue et
al., 2019). The term generalization refers to how well a hypothesis applies
even to new examples that it hasn’t seen in the training set. As mentioned in
Section 4.1, it is observed that while yielding similar values of training
functions, models trained with large-batch methods perform worse on test data
compared to small-batch methods (Keskar et al., 2017; Hoffer et al., 2017;
Shallue et al., 2019; Masters and Luschi, 2018). Such persistent degradation
in generalization performance is referred to as the generalization gap.
Identifying the origin of this gap and finding ways to close it is of
significant practical importance whereas remains an open problem.

This section is structured as follows. Section 5.1 introduces the concept of
sharp and flat (wide) minima; Section 5.2 addresses the relationship between
sharpness/flatness of local minima and their generalization ability; Section
5.3 provides explanation for the so-called generalization gap and Section 5.4
provides a somewhat opposing account.

###  5.1. Sharp and Flat (Wide) Minima

When training a DL model, we are seeking for a solution that minimizes a loss
function on a given training set. This solution lies in a very high
dimensional space (thousands, millions or even billions of parameters to
learn) called parameter space. The landscape of parameter space is showed
empirically crucial to generalize well. That being said, the wider the
solution’s local geometry, the better the generalization (Chaudhari et al.,
2017; Keskar et al., 2017; Li et al., 2018). Figure 2 provides an intuitive
explanation. There is generally a shift of the loss function in the parameter
space, flat minima is more robust to the perturbation of parameter than the
sharp one and thus generalizes better.

There are various definitions for ”sharpness/flatness” of the landscape.
Hochreiter and Schmidhuber (1997) define ”flatness” as a large connected
region in weight space where the error remains approximately constant. Keskar
et al. (2017) characterize ”flatness” by the magnitude of the eigenvalues of
Hessian, and propose a computational feasible $\epsilon$-sharpness measure.
Dinh et al. (2017) show that flat minima in practical DL hypothesis spaces can
be turned into sharp minima via re-parameterization without affecting the
generalization gap. Chaudhari et al. (2017) exploit the local geometric
properties of the objective function and use ”local entropy” as a measure of
”flatness”, which is invariant to the simple re-parametrization in (Dinh et
al., 2017). Foret et al. (2021) capture the ”sharpness” at parameter $w$ by
measuring how quickly the training loss can be increased by moving from $w$ to
a nearby parameter value.

Empirically, optimizers like SGD, Adam, etc. implicitly converge towards wide
valleys solutions. But there is no guarantee that this will always be the
case. This has motivated the creation of algorithms that will actively look
for flat minima such as Entropy SGD (Chaudhari et al., 2017), Sharpness-Aware
Minimization (SAM) (Foret et al., 2021) and many others.

Figure 2. A Conceptual Sketch of Flat and Sharp Minima, src: (Keskar et al.,
2017)

###  5.2. Generalization Gap and Sharp Minima

With regards to large batch training, Keskar et al. (Keskar et al., 2017)
observed that naively increasing the batch size typically results in
degradation of generalization performance and reduces computational benefit.
They speculate that “the lack of generalization ability is due to the fact
that large-batch methods tend to converge to sharp minima of the training
functions”. Specifically, large-batch methods are more vulnerable to sharp
minima whose training function increases rapidly in a relatively small
neighborhood (see Fig. 2). Such kind of high susceptibility to training
functions hampers the trained model from fitting on the test data. By
contrast, small-batch methods can resist the attraction of these minima and
converge to a flat minima whose training function varies slowly in a
relatively large neighborhood. They contribute the success of small batch
methods to their noisy gradients in the computation step. On one hand, the
noise expels the iterations from the trap of sharp minima. On the other hand,
the noise encourages them to move towards and retain in a flatter loss
landscape. However, as a larger batch size depicts a more accurate gradient,
their noise is no longer sufficient enough to eject the iterations from the
basin of sharp minima.

It is widely thought that small-batch SGD produces “flat” minima that
generalize well, while large batches converge to “sharp” minima with poor
generalization (Hochreiter and Schmidhuber, 1997; Keskar et al., 2017;
Chaudhari et al., 2017). However, there are some disputes about the effects of
batch size on model’s generalization ability. Hoffer et al. (2017) deny the
existence of inherent generalization gap and show empirically that the
”generalization gap” stems from the relatively small number of updates rather
than the batch size, and can be completely eliminated by adapting the number
of weight updates. Goyal et al. (2018) hold the view that optimization
difficulty is the main issue with large mini-batches, rather than the poor
generalization (at least on ImageNet). Specifically, using linear scaling and
warmup strategy, they show no loss of accuracy when training with large mini-
batch sizes up to 8,192 images on the ImageNet dataset.

###  5.3. Gradient Noise Ratio

As mentioned above, how batch size affects sharpness and generalization is
controversial. Smith and Le (2018) show that the test accuracy peaks at an
optimal batch size, if one holds the other hyper-parameters constant. They
believe that the arise of peak is not controlled by the batch size itself, but
the underlying scale of random fluctuations in the SGD dynamics.

Consider a simple model of SGD; the estimated gradient step is
$\eta\nabla_{SGD}(x)=\frac{\eta}{|B|}\sum_{i\in B}\nabla l_{i}(x)$, which can
be restated as the true gradient and a gradient noise term

(30) $\eta\nabla_{SGD}(x)=\underbrace{\eta\nabla
l(x)}_{gradient}+\underbrace{\frac{\eta}{|B|}\sum_{i\in B}(\nabla
l_{i}(x)-\nabla l(x))}_{noise\,term}.$

Smith and Le (2018) analogy between SGD and stochastic differential equations
(SDEs) to describe the noise in the SGD dynamics. In particular, they depicted
Eq.[30] as the discrete update of a stochastic differential equation (SDE) and
derive an analytical expression for the stochastic ”noise scale”
$g=\eta(\frac{N}{B}-1)\approx\eta\frac{N}{B}$, which controls the scale of
random fluctuations in the SGD dynamics. Noise drives SGD away from sharp
minima, and therefore there is an optimal batch size which maximizes the test
accuracy. This optimal batch size is proportional to the LR and training set
size $B_{opt}\propto\eta N$. Therefore, they attribute the so-called
”generalization gap” observed in (Dinh et al., 2017) as a consequence of
scaling batch size above this optimal batch size. Similarly, Jastrzebski et
al. (2017) derive a ”stochastic noise” using a different SDE. They verify
experimentally that the ratio of LR to batch size, $\eta/B$, influences the
width of the minima found by SGD, and that higher values of the ratio lead to
wider minima and often better generalization. Despite the slightly difference
in the form of ”stochastic noise”, both (Smith and Le, 2018; Jastrzebski et
al., 2017) indicate that gradient noise can be beneficial, especially in non-
convex optimization. Also they theoretically explain the empirical finding in
(Hoffer et al., 2017; Goyal et al., 2018) that rescaling the LR with the
square root of the batch size and train for more epochs, or linearly with
batch size, can reach the same generalization with a large batch size.

###  5.4. Train longer, Generalize Better

Another observation in (Keskar et al., 2017) is that large batch methods are
more likely to be attracted to minima close to the initial point, whereas
small batch methods are more explorative and always locate minima that are
farther away, with a ratio of $\lVert x_{S}^{*}-x_{0}\rVert_{2}/\lVert
x_{L}^{*}-x_{0}\rVert$ in the range of 3 to 10. Hoffer et al. (2017) further
find that the weight distance from initialization point increases
logarithmically with the number of training iterations (weight updates),
$\lVert w_{t}-w_{0}\rVert\sim log\,t$. They therefore deny the existence of
inherent generalization gap and believe that ”generalization gap” stems from
the relatively small number of updates rather than the batch size.
Specifically, they ”stretched” the time-frame of the optimization process,
where each time period of $e$ epochs in the original regime will be
transformed to $\frac{B_{L}}{B_{S}}e$ epochs according to the mini-batch size
used. However, such modification anneals the speedup effect of large batch
training.

##  6\. Second Order Optimization

Optimizations in DL, both theoretically and empirically, are presently
dominated by first-order gradient methods (Agarwal et al., 2016b, a;
Bollapragada et al., 2016; Carmon et al., 2018; Conn et al., 2000; Xu et al.,
2020). Second-order optimization methods that involve second derivatives
and/or second order statistics of the data, are far less prevalent despite
strong theoretical properties, due to their prohibitive computation, memory
and communication costs. In this section, we setup second-order optimization
basics in Section 6.1, start from the classical Newton’s method in Section
6.2, and turn to some up-to-date algorithms such as the Hessian-Free Method
(in Section 6.3), K-FAC (in Section 6.4) and Shampoo (Gupta et al., 2018) (in
Section 6.5).

###  6.1. Second-Order Optimization Basics

Basically, many problems in machine learning can be simply described as
minimizing the loss function over variables $x\in\mathbb{R}^{d}$

(31) $\mathop{min}\limits_{x\in{\mathbb{R}}^{d}}F(x).$

When training the weights of a NN, we are trying to get as far down the error
surface as possible. In most cases, we often use SGD to update the parameter
vector to solve this optimization problem

(32) $\begin{split}x_{t+1}&=x_{t}-\eta_{t}g_{t}.\end{split}$

Another very popular family of algorithms used in practice are the adaptive
optimization algorithms (e.g., AdaGrad (Duchi et al., 2011), Adadelta (Zeiler,
2012), RMSProp (Tieleman and Hinton, 2012), Adam (Kingma and Ba, 2017), etc.).
These are basically algorithms that update for each individual entry in the
parameter vector. Each entry has its own step size which is an adaptive update
using past gradients.

(33) $x_{t+1}^{(i)}=x_{t}^{(i)}-\eta_{t}^{(i)}g_{t}^{(i)}\quad i=1,...,d$

And there are also momentum variants of these methods that have slightly
different update rules. Potentially more powerful family of algorithms are
known as the preconditioned algorithms which use some matrices called
preconditioners to transform the gradient before taking a step. In general,
the idea of second-order optimization is to model the objective function $f$
by the local approximation

(34) $f(x+\delta)\approx M(x)\equiv f(x)+\nabla
f(x)^{T}\delta+\frac{1}{2}\delta^{T}B(x)\delta.$

Here, $B$ is a symmetric preconditioner and $\delta$ is the change in
parameters. In Newton’s method, $B=H$, or $B=H+\lambda I$. Fully optimizing
$M(x_{t})$ w.r.t. $\delta$ gives

(35) $\delta^{*}=\mathop{argmin}_{\delta}M(x_{t})=-B^{-1}\nabla f,$

then apply the update

(36) $x_{t+1}=x_{t}+\delta^{*}.$

This family includes algorithms mentioned above such as AdaGrad (Duchi et al.,
2011), Adam (Kingma and Ba, 2017), where the preconditioners are diagonal. But
they can be more powerful if the metrics are not diagonal but full
preconditioners, for exapmle, full AdaGrad, Natural Gradient (Amari, 1998) and
also classical algorithms like Newton’s method, Quasi-Newton methods
(Schraudolph et al., 2007; Goldfarb et al., 2020) and so on. It is well-known
in optimization that preconditioning often leads to faster convergence or
better ”condition number” in many different scenarios. But it comes with
obvious caveats: supposing the number of parameter is $n$, we need (1) at
least quadratic space $\Omega(n^{2})$ in the dimension to store the
preconditioner. (2) $n^{3}$ time to invert the preconditioner to apply to the
gradient vector. Generally speaking, these second-order methods are not very
practical where the cost of computation and memory is formidable in the DL
settings. Alternatively practitioners use the diagonal approximation or using
the SGD again.

Recently, there has been considerable advancement in the development of
second-order methods, seeking a balance between between full matrics and the
diagonal case. These methods usually approach preconditioners of the gradient
in a modular way, which is as powerful (or nearly powerful) as the full matrix
case, but can be used in practical like the diagonal case in terms of storage
and run-time. Inspired by the idea of the natural gradient method (Amari et
al., 2000), Martens and Grosse (2015) use a Kronecker-factored approximation
to the Fisher matrix as its preconditioning matrix that can be applied to
multi-layer perceptrons (MLPs), which was subsequently extended to other
architectures, such as convolutional neural networks (CNNs) (Grosse and
Martens, 2016) and recurrent neural networks (RNNs) (Osawa et al., 2019).
Kronecker-factored preconditioners based on the structure of the Hessian and
quasi-Newton methods have also been developed (Goldfarb et al., 2020; Ren and
Goldfarb, 2021).

###  6.2. Newton’s Method

Recall that in GD method, the gradient of a function is defined as the vector
of partial derivatives.

(37) $g_{t}\triangleq\nabla f(x)=<\frac{\partial f}{x_{1}},\frac{\partial
f}{x_{2}},\cdots,\frac{\partial f}{x_{n}}>$

It means we are assuming that the error surface of the NNs locally looks and
behaves like a circle. And we are ignoring all curvatures of the surface,
which may lead our training to progress very slowly. To rectify this, we can
use information from the second derivative of a function. The idea of Newton’s
method is to apply a linear transformation that turns ellipses into circles.
If we apply that transformation to the gradient vector, it will be as if we
were going downhill in a circular error surface. Formally, Newton’s method use
the Hessian matrix as preconditioner

(38) $H=\left[\begin{array}[]{cccc}\frac{\partial^{2}f}{\partial
x_{1}^{2}}&\frac{\partial^{2}f}{\partial x_{1}\partial
x_{2}}&\cdots&\frac{\partial^{2}f}{\partial x_{1}\partial x_{n}}\\\
\frac{\partial^{2}f}{\partial x_{2}\partial
x_{1}}&\frac{\partial^{2}f}{\partial
x_{2}^{2}}&\cdots&\frac{\partial^{2}f}{\partial x_{2}\partial x_{n}}\\\
\vdots&\vdots&\ddots&\vdots\\\ \frac{\partial^{2}f}{\partial x_{n}\partial
x_{1}}&\frac{\partial^{2}f}{\partial x_{n}\partial
x_{2}}&\cdots&\frac{\partial^{2}f}{\partial x_{n}^{2}}\end{array}\right],\quad
H_{ij}=\frac{\partial^{2}f}{\partial x_{i}\partial x_{j}}.$

The Hessian is a function of the parameters and we need to take its inverse
and multiply the gradient by that. Then we need to go some distance in that
direction

(39) $\Delta x=-\eta H(x)^{-1}\nabla f(x).$

If it is a truly quadratic surface and we choose the LR correctly, we will
arrive at the minima of the surface in a single step. However, that single
step involves something complicated which is inverting that Hessian matrix.
Assuming that we only have a million parameters in our NN, the Hessian matrix
will have a trillion terms which is completely infeasible to invert.

Curvature Matrices. Each element in the curvature matrix specifies how the
gradient in one direction changes as we move in some other direction. The off-
diagonal terms in a curvature matrix correspond to ”twists” in the error
surface. A twist means that when you travel in one direction, the gradient in
another direction changes. If we have a nice circular bulb, all those off-
diagonal terms are zero. As we travel in one direction, the gradient in other
directions doesn’t change. But when we have an elliptical error surface, as we
travel in one direction, the gradient in another direction changes. This is
actually what is going wrong with GD. As GD updates one of the weights, at the
same time it is updating all the other weights, causing a change in the
gradient for the first weight. That means when we update it we may actually
make things worse. The gradient may have actually reversed sign due to the
changes in all the other weights. And so the more weights we get, the more
cautious about changing each one of them we need to be, because the
simultaneous changes in all the other weights can change the gradient of a
weight.

How to avoid inverting a huge matrix. The intensive computation of the
curvature has limited the applications of second-order optimization methods in
DL settings. To address this problem, there are various ideas in the
literature. One very popular line of work looks at diagonal approximations,
e.g., Adagrad (Duchi et al., 2011), RMSProp(Tieleman and Hinton, 2012),
Adadelta (Zeiler, 2012) and many others (Botev et al., 2017; Bordes et al.,
2009). But these diagonal terms consist only a tiny fraction of the
interactions, so we are ignoring most of the terms (nearly all of them) in the
curvature matrix. And the experimental evidence indicates that there is
limited or almost no improvement in practice when compared to well-tuned SGD
with or without momentum (see (Zeiler, 2012; Botev et al., 2017)). The
benefits of these diagonal approaches seem to lie mainly in the ease of
choosing the LR, but may not provide any fundamental benefits beyond that.
Another thing we could do is to approximate the curvature matrix with much
lower rank matrix but capturing its main aspects. Limited-memory BFGS (L-BFGS)
(Schraudolph et al., 2007) is the most well-known example. Again there is
limited/non-existent empirical success for NN optimization.

More recently, many focus are put on the trade-off between the full matrix and
the diagonal case, seeking a balance between these two extreme case. Some
researchers find something in between that is as powerful (or nearly powerful)
as the full matrix case, but can be used in practical like the diagonal case
in terms of storage and run-time. Some recent approaches for approximating a
full-matrix preconditioner are K-FAC (Martens and Grosse, 2015) and Shampoo
(Gupta et al., 2018). Others incorporate automatically the Hessian operator,
such as Hessian-Free method (Martens, 2010) and trust-region (Conn et al.,
2000; Xu et al., 2020).

###  6.3. Hessian-Free Method

The Hessian-Free method (Martens, 2010) is a quasi-Newton method that uses no
low-rank approximations. Named ”free” because it never explicitly computes the
preconditioner $B$ but instead does approximate minimization of quadratic
model $M(\delta)$ (see Eq.[35]). The Hessian-Free (HF) method is motivated by
two observations. The first one being that it is relatively easy to compute
the matrix-vector product $Hv$ for an arbitrary vectors $v$, e.g., use finite
differences to approximate the limit.

(40) $Hv=\mathop{lim}\limits_{\epsilon\rightarrow 0}\frac{\nabla f(x+\epsilon
v)-f(x)}{\epsilon}$

The second motivating observation is that linear conjugate gradient (CG)
minimizes positive definite quadratic cost functions using only matrix-vector
products, which is relatively easy to obtained (as shown in Eq.[40]).
Conjugate gradient is a very clever method that instead of trying to go
straight to the minimum like in Newton’s method, it tries to minimize in one
direction at a time. It starts off by taking the direction of GD and goes to
the minimum in that direction that might involve re-evaluating the gradient or
re-evaluating the error a few times to find the minimum in that direction.
Once it is done, CG method now finds another direction and goes to the minimum
in that second direction. The clever thing about the technique is that it
chooses the second direction in such a way that doesn’t mess up the
minimization it already did in the first direction, which is called a
conjugate direction. ”Conjugate” means that as we go in the new direction we
do not change the gradients in the previous directions. What CG achieves is
that it gets to the global minimum of an $n$-dimensional quadratic surface in
only $n$ steps. More importantly, in many less than $n$ steps on a typical
quadratic surface, it will have reduced the error very close to the minimum
value. And that’s why we use it. As doing the full $n$ steps that would be as
expensive as inverting the whole matrix, we are going to do many less than $n$
steps and get quite close to the minimum. Pseudo-code for a simple variant of
damped HF optimization is provided in Algorithm 4.

for  _$n=1$ to max-epochs_ do

compute gradient $g_{t}=\nabla f(x_{t})$;

choose/adapt $\eta_{t}$ according to some heuristic;

define the function $B_{t}(v)=Hv+\eta_{t}v$;

$p_{t}=CGMinimize(B_{t},-g_{t})$;

$x_{t+1}=x_{t}+p_{t}$

end for

Algorithm 4 Damped Hessian-Free Optimization

However, common variants of HF don’t work particular well for NNs.
Enhancements techniques (e.g., the Gauss-Newton approximation to the Hessian,
early CG stopping, damping, etc.) are provided in (Martens, 2010). More
recently research has revealed that DNN learning is easier than previously
thought using simple methods (Sutskever et al., 2013b). Carefully tuned
momentum methods suffice for dealing with the curvature issues in deep and
recurrent network training objectives without the need for sophisticated
second-order methods. Despite SGD with or without momentum still being the
most widely used and best method in most situations, the fact that HF uses
100-1000x fewer iterations than SGD supports the idea that a second order
method can help a lot in principle, provided that we can make these iterations
cheap enough to compute.

###  6.4. K-FAC

Kronecker-Factored Approximate Curvature (K-FAC) (Martens and Grosse, 2015) is
one of the natural gradient approximation methods where the preconditioner is
a high-quality approximation of Fisher information matrix (FIM). We first give
a brief introduction about natural gradient descent and then explained the
outline of K-FAC.

####  6.4.1. Natural Gradient Descent

Natural Gradient Descent (NGD) (Amari, 1998) is a second order optimization
method based on information geometry. NGD acquires the loss landscape
correctly by using FIM as curvature of loss function and converges faster in
term of iterations than a simple first-order method. The FIM associated with
network’s distribution $P_{y|x}(\theta)$ is

(41) $F=E[\nabla log(p(y|x;\theta))\nabla log(p(y|x;\theta))^{T}].$

Importantly, one property of $F$ is that it can be interpreted as the negative
expected Hessian of our model’s log likelihood (Martens, 2020)

(42) $F=-E_{p(y|x;\theta)}[H_{logp(x|\theta)}].$

Knowing this result, we can see the role of $F$ as a measure of curvature of
the log likelihood function. Thus the immediate application of $F$ is as drop-
in replacement of $H$ in second order optimization methods. Using KL-
divergence to measures how different two models are, the update rule of NGD is

(43) $\theta_{t+1}\leftarrow\theta_{t}-\eta_{t}F^{-1}\nabla f(\theta_{t}).$

Here the inverse of the FIM is applied to the gradient of loss, and the
gradient preconditioned by the FIM is called the natural gradient. For the
parameters of size $N$, the size of FIM is $N\times N$, and NNs used in DL
tend to have a massive number of parameters (e.g., 60 million parameters in
AlexNet for ImageNet classification) so the inverse of the FIM is intractable,
and it limits the number of the applications of NGD to DL. In recent years,
some works have proposed methods that approximate or avoid inversing the FIM.

####  6.4.2. K-FAC Approximation.

K-FAC approximates the FIM so that the inverse matrix is easy to calculate.
Firstly, K-FAC approximates F as $\hat{F}$, a diagonal block matrix where each
block represents one layer in a NN with $L$ layers

(44) $\hat{F}=diag(\hat{F}_{1},...,\hat{F}_{l},...,\hat{F}_{L}).$

Next, each diagonal block matrix $F_{l}$ is approximated as a Kronecker
product

(45) $\hat{F}_{l}\approx A_{l-1}\otimes G_{l}.$

This is called Kronecker factorization and $G_{l}$, $A_{l-1}$ are called
Kronecker factors, representing the gradient of the output of the $l$-th layer
and the activation of the ($l$-1)-th layer respectively. By using the critical
property of the Kronecker product of the matrices $(A\otimes
B)^{-1}=A^{-1}\otimes B^{-1}$, the inverse of $\hat{F}_{l}$ can be computed as

(46) ${\hat{F}_{l}}^{-1}={A_{l-1}}^{-1}\otimes{G_{l}}^{-1}.$

The final update step of parameters $w_{l}$ in the $l$-th layer is as follows:

(47) $w_{l}^{(t+1)}=w_{l}^{(t)}-\eta^{(t)}{G_{l}}^{-1}\nabla
L_{l}(w_{l}^{(t)})A_{l-1}^{-1}.$

In most implementations, Tikhonov regularization is used to avoid ill-
conditioned matrix inverses with K-FAC by adding a damping parameter $\gamma$
to the diagonal of $\hat{F}_{l}$ (Pauloski et al., 2020; Grosse and Martens,
2016)

(48) $(\hat{F}_{l}+\gamma I)^{-1}={({A_{l-1}+\gamma
I})^{-1}}\otimes{({G_{l}+\gamma I})^{-1}}.$

A standard K-FAC update step for one layer requires inverting two matrices
$(A_{l-1}+\gamma I)$ and $(G_{l}+\gamma I)$, which can be computed implicitly
using an alternative method based on the eigendecompostion of $\hat{F}_{l}$
(Grosse and Martens, 2016; Pauloski et al., 2020).

(49) $\begin{split}V_{1}&=Q_{G}^{T}L_{i}(w_{i}^{(k)})Q_{A}\\\
V_{2}&=V1/(v_{G}(v_{A})^{T}+\lambda)\\\ (\hat{F}_{l}+\gamma I)^{-1}\nabla
L_{i}(w_{i}^{(k)})&=Q_{G}V_{2}{Q_{A}}^{T}\end{split}$

In practice, practitioners avoid significant computation and communication by
reducing the frequency of computing these factors and and their
eigendecompositions, at the cost of introducing staled information. For
example, Pauloski et al. (2020) update K-FAC statistics for every 500
iterations for ResNet scaling experiments on 64 GPUs.

####  6.4.3. Distributed K-FAC

Some studies have used K-FAC and implemented the algorithm in a distributed
computing environment (Osawa et al., 2019; Pauloski et al., 2020). With only
35 epochs and a 16K batch size, ResNet50 can be trained to achieve 75% Top1
accuracy in ImageNet (Osawa et al., 2019). More recently, Pauloski et al.
(2020) scales up K-FAC for training CNNs. It mainly refers to the calculation
scheme of preconditioned gradient in (Grosse and Martens, 2016) and uses
feature decomposition to replace matrix inversion.

Algorithm | Preconditioning | Memory | Computation
---|---|---|---
Full Matrix AdaGrad |  $H_{t}=(\sum_{s=1}^{t}g_{s}g_{s}^{T})^{\frac{1}{2}}$ $W_{t+1}=W_{t}-\eta_{t}H_{t}^{-1/2}$ | $O((mn)^{2})$ | $O((mn)^{2})$
Shampoo |  $L_{t}=L_{t-1}+g_{t}g_{t}^{T}$ $R_{t}=R_{t-1}+g_{t}^{T}g_{t}$ $W_{t+1}=W_{t}-{L_{t}^{-1}G_{t}R_{t}^{-1}}$ | $O(m^{2}+n^{2})$ | $O(m^{2}+n^{2})$
K-FAC |  $\hat{F}_{i}=A_{i-1}\otimes G_{i}$ $A_{i-1}=a_{i-1}a_{i-1}^{T}$ $G_{i}=g_{i}g_{i}^{T}$ $W_{l}^{(t+1)}=W_{l}^{(t)}-\eta^{(t)}{G_{l}}^{-1}\nabla L_{l}(W_{l}^{(t)})A_{l-1}^{-1}$ | $O(m^{2}+n^{2})$ | $O(m^{3}+n^{3})$
Diagonal AdaGrad |  $H_{t,ij}=\sum_{s\leq t}g^{2}_{s,ij}$ $W_{t+1}=W_{t}-\eta_{t}H_{t}^{-1/2}$ | $O(mn)$ | $O(mn)$
SM3 |  $\widehat{H_{t,ij}}=min(L_{t-1,i},R_{t-1,j})+g_{t,ij}^{2}$ $L_{t,i}=\mathop{max}\limits_{j}(\widehat{H_{t,ij}})$ $R_{t,j}=\mathop{min}\limits_{i}(\widehat{H_{t,ij}})$ $W_{t+1,i}=W_{t,i}-\eta g_{t,i}{\widehat{H_{t,ij}}}^{-1/2}$ | $O(m+n)$ | $O(mn)$
Table 1. Summary of Preconditioning Methods. Example of a fully connected
layer $[m,n]$.

###  6.5. Shampoo

Shampoo (Gupta et al., 2018) is another preconditioning algorithm, which is an
approximation version of full matrix AdaGrad (Duchi et al., 2011). It first
approximates by treating each layer independently (block diagonal). And it
uses small matrices whose Kronecker product approximates the full AdaGrad
statistics. Such two approximations make Shampoo more expressive than the
diagonal preconditioning and practical to store and compute at large scale.
Mathematically, the preconditioner Shampoo looking for can be written as a
Kronecker product of two smaller matrices $L$ and $R$

(50) $\mathop{argmin}\limits_{H=L\otimes R\atop L,\,R\succ
0}\,\\{H^{-1}\bullet\overline{G_{t}}+Tr(H)\\}.$

Though it cannot solve the exact optimization problem, it has a nice limit
that relaxes the upper bounds in a matrix sense:

(51)
$\frac{1}{\sqrt{r}}\underbrace{(\sum\limits_{t=1}^{T}g_{t}(g_{t})^{T})^{\frac{1}{2}}}_{full\,AdaGrad\,precond.}\preceq\underbrace{(\sum\limits_{t=1}^{T}{G_{t}{G_{t}}^{T}})^{\frac{1}{4}}}_{L_{t}}\otimes\underbrace{{(\sum\limits_{t=1}^{T}{{G_{t}}^{T}G_{t}})^{\frac{1}{4}}}}_{R_{t}}$

The full AdaGrad peconditioner is given on the left, bounded by a Kronecker
product of two smaller matrices. The update statistic of Shampoo is given as
follows:

(52) $\begin{split}L_{t}=L_{t-1}+G_{t}G_{t}^{T},\quad
R_{t}=R_{t-1}+G_{t}^{T}G_{t}\\\
W_{t+1}=W_{t}-{L_{t}^{-1}G_{t}R_{t}^{-1}}\end{split}$

Both Shampoo and K-FAC employ a preconditioning scheme that approximates the
FIM. Despite their similarity in construction, they differ in several
important ways. The differences are based on choices such as the empirical FIM
or FIM, moving average or sum, and the inverse component. Another key
difference is that Shampoo construction is agnostic to layer types. K-FAC
relies heavily on the structure of the back-propagated gradients in a feed-
forward neural network. In contrast, Shampoo is virtually oblivious to the
particular model structures and only depends on standard gradient information.
More recently, Anil et al. (2020) extend Shampoo in a number of ways so as to
make it applicable to a larger range of deep architectures.

Despite the fact that first-order methods have been dominant in the recent
decade, recently second order methods, such as K-FAC and Shampoo, show some
promise. They mitigate the space and run-time costs of full-matrix second-
order algorithms and have been applicable to a larger range of deep
architectures (see Table 1). It is interesting to see whether second order
methods can outperform first order ones in the future.

##  7\. Communication

Large-scale distributed training improves the productivity of training deeper
and larger models, where data parallelism is adopted so as to take full
advantage of the compute capability on multiple workers. SGD is usually
selected as the optimization method because of its high computation efficiency
and well support by the DL tool-kits, such as TensorFlow (Abadi et al., 2016),
PyTorch (Paszke et al., 2019) and DeepSpeed (Rasley et al., 2020). In data-
parallel SGD, each worker processes a random mini-batch of its training data,
and then the local updates are synchronized by making an All-Reduce step or
through a centralized parameter server, which aggregates stochastic gradients
from all workers, and taking a Broadcast step that transmits the updated
parameter vector back to all workers. The process of gradient synchronization
is repeated until an appropriate convergence criterion is met.

Increasing the number of workers and taking advantage of data parallelism help
to reduce the computation time on the same size training data dramatically.
However, as the scale of distributed systems grows up, gradient and parameter
synchronization prolongs the communication time and hinders the perfect
scalability (Li et al., 2014a; Wen et al., 2017). Therefore, the high network
communication cost becomes a significant bottleneck of distributed training.
There have been many attempts to reduce the communication overhead in data-
parallel SGD. One notable method is to let each worker use compressed
gradients rather than raw gradients for communication. For example, quantized
SGD or sparcified SGD allow each worker to use fewer bits to pass gradients by
sacrificing the convergence to a mild extent. Another notable method is to
reduce the frequency of communication (Zinkevich et al., 2010; McDonald et
al., 2010; Zhang et al., 2016; Kamp et al., 2018; Lin et al., 2020). We detail
on these two communication-efficient methods in the subsequent sections.

###  7.1. Gradient Compression

Lin et al. (2018) find that 99.9% of the gradient exchange in distributed SGD
are redundant. One promising solution is gradient compression, e.g., through
gradient quantization (Alistarh et al., 2017; Wen et al., 2017; Ramezani-
Kebrya et al., 2021; Seide et al., 2014; Tang et al., 2021; Li et al., 2021)
and/or gradient sparsification (Aji and Heafield, 2017; Lin et al., 2018;
Dryden et al., 2016; Strom, 2015). Sparcification means transmitting only
those gradients that are important (e.g., gradients with large absolute
values), while quantization refers to using fewer bits to represent the
original gradient. Their difference is described in Fig. 3.

####  7.1.1. Gradient Quantization

Figure 3. Comparison of Quantization and Sparsification, src: (Tang et al.,
2020)

Quantizing the gradients to low-precision values can reduce the communication
bandwidth. In full-precision data-parallel SGD, each processor broadcasts its
locally computed stochastic gradient vector at every iteration, whereas in
quantized data-parallel SGD, each processor quantizes its stochastic gradient
before broadcasting. Gradient quantization is usually done by mapping a
continuous space of gradient values onto a discrete set. Take the classic
Quantized Stochastic Gradient Descent (QSGD) (Alistarh et al., 2017) as an
example, the quantization function is denoted with $Q_{s}(v)$:

(53) $Q_{s}(v_{i})=\lVert v\rVert_{2}\cdot sign(v_{i})\cdot\xi_{i}(v,s)$

where $\xi_{i}(v,s)$ are independent random variables. Let $0\leq l<s$ be an
integer such that $v_{i}/\lVert v\rVert_{2}\in[l/s,(l+1)/s]$. That is,
$[l/s,(l+1)/s]$ is the quantization interval corresponding to $v_{i}/\lVert
v\rVert_{2}$. Then $\xi_{i}(v,s)$ is defined as follows

(54)
$\xi_{i}(v,s)=\left\\{\begin{array}[]{cc}l/s&with\,probability\,1-p(\frac{|v_{i}|}{\lVert
v\rVert_{2}},s)\\\ (l+1)/s&otherwise\end{array}\right.$

Here, $p(a,s)=as-l$ for any $a\in[0,1]$. For gradients $v$, quantization is
used to randomly convert gradient values in each dimension $v_{i}$ to some
discrete value in a predetermined discrete set. After normalized by the
Euclidean norm of the gradients ($|v_{i}|/\lVert v\rVert$), the value of each
dimension will fall on a sub-interval $[0,1]$, and we approximate it to one of
the endpoints of the sub-interval with a certain probability each time, so
that the continuous value space of the original gradient value can be replaced
by a set of finite discrete values. Here $\xi_{i}(v,s)$ is a binary random
variable which guarantees that each value is quantized in a way which
preserves the value in expectation,$E[\xi_{i}(v,s)]=|v_{i}|/\lVert v\rVert$,
and introduce minimal variance. An instance of QSGD is provided in Fig. 5.

This gradient quantization method greatly reduces the amount of communication
required by a single node. Instead of passing $n$ 32-bit floating-point
gradients, only one 32-bit floating-point gradient with one bit for gradient
sign and $log(s)$ bits for $\xi_{i}(v,s)$ on each dimension are required. In
addition, there is another method TernGrad (Wen et al., 2017) developed
simultaneously with QSGD. Their underlying idea is essentially similar, where
TernGrad can be viewed as a special case of QSGD when $l=1$. TernGrad randomly
quantizates gradient $g_{t}$ to a ternary value vector with value of
$\\{-1,0,1\\}$. Formally, with a random binary vector $b_{t}$, gradient is
ternarized as

(55) $\begin{split}\tilde{g}_{t}&=ternarize(g_{t})=s_{t}\cdot sign(g_{t})\circ
b_{t}\\\ s_{t}&\triangleq\lVert g_{t}\rVert_{\infty}\triangleq
max(abs(g_{t})).\end{split}$

Here, $\circ$ is the Hadamard product. TernGrad also adopts techniques such as
layer-wise ternarizing and gradient clipping to improve convergence. Ramezani-
Kebrya et al. (2021) propose nonuniform quantization levels (NUQSGD) and
demonstrate superior empirical results compared to QSGD. Horvath et al. (2019)
propose natural compression and natural dithering, where the latter is a
special case of logarithmic quantization.

Figure 4. QSGD Example with $s=4,l=3$

Figure 5. Top-k Example

Unlike QSGD and its variants which use stochastic rounding which has an
unbiased gradient expectation, there are methods adopt biased ones. Methods
performing updates only based on the sign of each coordinate of the gradient
have gained popularity for training DL models (Seide et al., 2014; Bernstein
et al., 2018; Karimireddy et al., 2019; Tang et al., 2021; Li et al., 2021).
Seide et al. (2014) proposed SignSGD, a.k.a. 1-bit SGD, to quantize the
gradients aggressively to one bit per value. In this scheme, gradient updates
greater than or equal to zero are encoded using the value 1, and those less
than zero with the value 0. The reconstruction values are chosen to be the
means of the non-negative and negative updates, respectively, in order to
minimize the square quantization error. This is done column-wise over the
weight matrix. In each data exchange, the two reconstruction values are
transmitted along with their respective quantized column. Bernstein et al.
(2018) later provided convergence guarantees for a variant of SignSGD.
Karimireddy et al. (2019) proposed EF-SignSGD, which is an improved version of
SignSGD. More recently, gradient compression with error compensation has been
successfully applied to adaptive optimizer such as Adam (1-bit Adam (Tang et
al., 2021)) and LAMB (1-bit LAMB (Li et al., 2021)), further scaling up
training algorithms in the distributed setting.

To recap, while the analyses of gradient quatization have largely been
restricted to unbiased compression schemes (Alistarh et al., 2017; Ramezani-
Kebrya et al., 2021; Wen et al., 2017; Horvath et al., 2019), biased schemes
which perform extreme compression practically perform much better often
without any loss in convergence or accuracy (Seide et al., 2014; Karimireddy
et al., 2019; Strom, 2015; Lin et al., 2018; Tang et al., 2021; Li et al.,
2021).

####  7.1.2. Gradient Sparsification

Gradient sparsification is an orthogonal approach to quatization methods,
which reduces the communication bandwidth by sending only the important
gradients. Since zeroing small gradients damages convergence, small gradients
are accumulated over time locally until they become large enough to be
transmitted. Thus, we send the large gradients immediately but eventually send
all of the gradients. Strom (2015) proposed threshold quantization by
considering only gradient elements whose absolute values exceed a threshold. A
fixed threshold $\tau$ is chosen in advance. Gradient updates greater than
$\tau$ are encoded with the value 1, and those less than $-\tau$ with the
value of 0. Updates of magnitude less than $\tau$ are not sent at all,
reducing the volume of data sent. The reconstructed value is $\tau$ and
$-\tau$ respectively, and error feedback is used as normal. However, the
threshold is hard to choose in practice and, moreover, it can change over time
during optimization. As a resolve, Top-k sparsification selects the top-k
gradients in terms of absolute values at each iteration (Stich et al., 2018;
Alistarh et al., 2018) (see Fig. 5 for an example). Dryden et al. (2016)
choose an adaptive threshold so as to keep a constant proportion of gradients
each iteration. Aji and Heafield (2017) sparsify gradient updates by removing
the R% smallest gradients by absolute value, dubbing this Gradient Dropping.
This approach is slightly different from (Dryden et al., 2016) as it uses a
single threshold based on absolute value, instead of dropping the positive and
negative gradients separately. Concurrently, Chen et al. (2018) localize
selection of gradient residues and automatically tunes the compression rate
depending on local activity. Lin et al. (2018) further push the compression
ratio by employing momentum correction, local gradient clipping, momentum
factor masking, warm-up training on top of the gradient sparsification while
maintaining model performance. Table 2 summarises the gradient quantization
and sparsification methods.

Method | Taxonomy | Reference
---|---|---
Quant. | Unbiased | QSGD (Alistarh et al., 2017), NQSGD (Ramezani-Kebrya et al., 2021), TernGrad (Wen et al., 2017), Natural (Horvath et al., 2019)
Biased |  1-bit SGD (Seide et al., 2014) /Adam (Tang et al., 2021)/LAMB (Li et al., 2021) EF-SignSGD (Karimireddy et al., 2019)
Spars. | Random | Random-k (Wangni et al., 2018)
. | Deterministic |  Fixed threshold (Strom, 2015) Top-K (Stich et al., 2018; Alistarh et al., 2018), Adaptive threshold (Dryden et al., 2016; Aji and Heafield, 2017; Chen et al., 2018; Lin et al., 2018)
Table 2. Summary of Gradient Compression Methods

###  7.2. Reducing Communication Frequency

A parallel line of work reduces the communication cost by reducing the
frequency of communication. For instance, local SGD saves the communication
cost by allowing each worker to perform more than one batch update on local
data and exchange the updated weights rather than the gradients among workers.

We consider a distributed SGD framework with $K$ worker nodes where all
workers communicate with others via a central server or via direct inter-
worker communication. In local SGD, each worker $k\in[K]$ performs $H$
sequential mini-batch SGD updates locally, and then the local models are
synchronized by averaging weights among workers. Thus, the overall update rule
at the $k$-th worker is given by

(56)
$\begin{split}w_{(t)+h+1}^{k}&:=w_{(t)+h}^{k}-\eta_{(t)}[\frac{1}{B_{loc}}\sum_{i\in{I^{k}_{(t)+h}}}\nabla
f_{i}(w_{(t)+h}^{k})]\\\
w_{(t+1)}^{k}&:=\frac{1}{K}\sum_{k=1}^{K}w_{(t)+H}^{k}\end{split}$

where $w_{(t)+h}^{k}$ denotes the local model on worker $k$ with batch size
$B_{loc}$ after $t$ global synchronization and $h$ local SGD updates. Mini-
batch SGD is a special case of local SGD, with $H=1$, that is, the local
models are synchronized after every iteration. The convergence results for
convex and non-convex objectives are provided in (Stich, 2019; Zhou and Cong,
2018). However, while local updates reduce the communication frequency by
performing global synchronization periodically instead of at per iteration,
the discrepancies between local models can result in an inferior error-
convergence. A larger value of $H$ (i.e., the number of sequential local SGD
updates), which means less frequent averaging, saves communication delay and
reduces the run-time per iteration. But on the other hand, a larger $H$ leads
to slower convergence w.r.t. the number of iterations. The trade-off in
between still need more exploration.

In addtion to being communication efficient, Lin et al. (2020) find local SGD
also exhibits good generalization behaviour. They argue that local SGD is a
way to inject and control stochastic noise to the whole training procedure,
and thus proposed post-local SGD as large batch training alternative for
better generalization.

##  8\. Memory

Larger models usually require more computation and memory resources to train.
The amount of memory required to train these models can be several orders of
magnitude larger than the amount of memory available on a single GPU. In this
section, we will see how some popular techniques successfully reduce the
memory requirements of training NNs without compromising model performance.
Section 8.1 introduces how mix-precision training (Micikevicius et al., 2018)
lowers the burden on memory using fewer bits to preserve the weights and
gradients during training. Section 8.2 introduces two memory-efficient
adaptive optimizers, Adafactor (Shazeer and Stern, 2018) and SM3 (Anil et al.,
2019). And as orthogonal to the above methods, ZeRO (Rajbhandari et al., 2020)
do not change the model optimization method or affect model convergence, but
instead reduces the memory cost by removing the redundancy in data-parallel
(Section 8.3).

###  8.1. Mix-Precision Training

Modern DL training systems use single-precision (FP32) format, which takes 32
bits of memory. However, lower-precision (FP16) takes 16 bits of memory
instead. Modern accelerators like Google TPUs and NVIDIA GPUs can run
operations faster in the FP16 format, as they have specialized hardware to run
16-bit computations and 16-bit dtypes can be read from memory faster. These
lower-precision provides numerous benefits. First, they require less memory,
enabling the training and deployment of larger NNs. Second, they lowers the
burden on memory since fewer bits are required to preserve the same number of
values than the FP32 format, thereby speeding up data transfer operations.
Third, they speed up the mathematical computation since low-precision
calculation is less time-consuming, especially on GPUs with Tensor Core
support for that precision. However, low precision training also introduces a
trade-off of the number of bits used versus the statistical accuracy: the
fewer bits used, the lower accuracy.

Figure 6. Workflow of Mix Precision Training

Mixed precision training is a very simple and practical technique, jointly
published by Baidu and Google in (Micikevicius et al., 2018), which almost
halves the memory footprint while maintaining the model accuracy. This is
achieved by identifying the steps that require full precision and using FP32
for only those steps while using FP16 everywhere else. We explain the workflow
of mix precision training in Figure 6. Due to the differences in representable
ranges, simply converting the model to FP16 can cause gradient underflow or
overflow problems. We can avoid these challenges with the following four
steps: (1) Conversion into FP16. In step one, we cast the inputs from FP32 to
FP16 for compute intensive operations. (2) Use FP32 to compute the loss.
Because FP16 might cause underflow or overflow issues, we do loss calculation
in FP32 in the backward pass and cast gradients back to FP16, which means that
the weights and the gradients are still in FP16. (3) FP32 master weights. In
the backward pass, gradients are small compared to the parameters. If we try
to update our parameters with the gradients that are much smaller than the
parameters, then we might lose those parameter updates. To compensate we will
maintain the master copy of weights in FP32. This means that, to the end of
the backward pass, FP16 gradients will be cast into FP32 and thereby applied
to FP32 weights. In the forward pass, we will cast the weights into FP16 so
that the gradient computations remain in FP16. So effectively we have a master
copy of all the parameters which are weights and biases stored in FP32 but all
the computational operations will see the casted version which is in FP16. (4)
Loss (Gradient) scaling. The last step is to do loss scaling to avoid a
gradient underflow problem. Before computing gradients from the FP32 loss, we
scale a loss by multiplying it with a loss scale factor. By doing so,
gradients are pushed to larger values and we can safely represent them in
FP16. Later when updating the weights we can re-scale the gradients by
dividing them with the same loss scale factor.

Their article (Micikevicius et al., 2018) is not the first to propose the use
of lower precision for training, but its influence is far-reaching, and many
current programs are designed based on this work. Jia et al. (2018) apply
mixed-precision training to large-batch strategies such as LARS. Using LARS
with mixed-precision training, ResNet-50 with the mini-batch size of 64K,
could maintain the top-1 accuracy as 76.2%.

###  8.2. Memory Efficient Adaptive Optimization

Some stochastic optimization methods (e.g., RMSProp, Adam (Kingma and Ba,
2017), Adadelta (Zeiler, 2012)), keep first and second moment estimates of the
per-parameter gradients to scale the gradients which triples the required
memory. As models continue to grow, the memory overhead will pose more
limitation on the quality of the trained model. Motivated by these challenges,
memory efficient adaptive optimization methods are proposed to retain the
benefits of standard per-parameter adaptivity while significantly reduce
memory overhead. For instance, Adafactor (Shazeer and Stern, 2018) was
proposed as a way to reduce the memory costs of AdaGrad (Duchi et al., 2011),
primarily for training large language models, and SM3 (Anil et al., 2019)
saves memory by sharing moments of similar magnitude.

####  8.2.1. Adafactor

Adafactor (Shazeer and Stern, 2018) is a space-efficient adaptive optimization
which achieves a drastic reduction in auxiliary memory usage without hurting
the performance (compared to that obtained using full accumulators). One of
the key contribution is the use of factored second momentum estimation.

Consider a matrix-shaped parameter subset $X$ with second moment estimate $V$.
They want to identify a low-rank representation of $V$ as a product of two
factors $R$ and $S$, i.e., $V\approx RS$ which is compatible with exponential
moving averaging. This would allow us to store just the low-rank factors
across iteration, cutting down a memory usage. More formally, if factorization
$F:V\mapsto(R,S)$, we want $F(\eta V_{t-1}+(1-\eta)G_{t}^{2})=\eta
F(V_{t-1})+(1-\eta)F(G_{t}^{2})$. In particular, by using techniques from non-
negative matrix factorization using I-divergence, the low-rank approximation
can be converted into following optimization problem:

(57) $\begin{split}\mathop{minimize}_{R\in\mathbb{R}^{n\times
k},S\in\mathbb{R}^{k\times
m}}\sum_{i=1}^{n}\sum_{j=1}^{m}d(V_{ij},[RS]_{ij})\\\ subject\,to\quad
R_{ij}>0,S_{ij}>0.\end{split}$

In particular, for the case of rank one factors, i.e., $k=1$, the solution set
of the optimization problem can be characterized as the set of all pairs
$(R,S)$, whose product is equal to the expression below:

(58) $\\{(R,S):RS=\underbrace{V1_{m}}_{row\atop
sums}\underbrace{1_{n}^{T}V}_{colum\atop
sums}/\underbrace{1_{n}^{T}V1_{m}}_{sum\,of\atop all\,entries}\\}.$

The right hand side can be broken down into the vector of row sums and column
sums, and the denominator is the sum of all entries. In addition to the
factored second moment estimation, other key changes in Adafactor include
$\eta_{2}$ varies with time, update cliping, relative step size and no
momentum.

####  8.2.2. SM3 Algorithm

Adaptive gradient methods, such as AdaGrad (Duchi et al., 2011), have proved
to be particularly useful in training sparse models. Crucially, however,
Adagrad must maintain auxiliary sequence of accumulators (i.e., the diagonal
preconditioner) $H_{t}$ (also in Eq.[10]) :

(59) $H_{t,ii}=\sum_{s\leq t}g_{s,ii}^{2}$

and thus needs $\Omega(n)$ additional space $n$ is the number of parameters.
SM3 (Anil et al., 2019) provides a memory-efficient methods with comparable
convergence characteristics which refrains from maintaining the full vectors.
SM3 is short for save memory by sharing moments of similar magnitude. This is
because they observe that the diagonal preconditioners $H_{t}$ accumulated by
AdaGrad are actually similar in rows and columns, and by sharing moments cross
rows and columns, the memory requirements therefore drop from $\Theta(mn)$ to
merely $\Theta(m+n)$.

(60)
$\begin{split}\widehat{H_{t+1},ij}&=min(R_{t,i},C_{t,j})+g_{t+1,ij}^{2}\\\
R_{t,i}&=\mathop{max}\limits_{j}(\widehat{H_{t+1},ij})\\\
C_{t,j}&=\mathop{max}\limits_{i}(\widehat{H_{t+1},ij})\end{split}$

SM3 can be viewed as a diagonal version of Shampoo (see Table 1). We refer
readers about the implementation details to (Anil et al., 2019).

###  8.3. ZeRO

Figure 7. Comparing the per-device memory consumption of model states, with
three stages of ZeRO-DP optimizations. Src:(Rajbhandari et al., 2020).

The Zero Redundancy Optimizer (Rajbhandari et al., 2020) (abbreviated as ZeRO)
is a novel memory optimization technology for large-scale distributed DL.
Contrary to Adafactor (Shazeer and Stern, 2018) and SM3 (Anil et al., 2019)
which reduce memory consumption of adaptive optimization methods by
maintaining coarser-grained statistics of model parameters and gradients, ZeRO
do not change the model optimization method or affect the model convergence.

We show how ZeRO works in Fig.[7]. The first row shows the memory map while
training a model in data parallel. The first row shown in blue represents the
memory consumed by the parameters; the second row in orange shows the memory
consumed by the gradients; and the big chunk in green shows the memory
consumed by the optimizer states (e.g., this could be momentum and variance
for Adam). So the key thing is that these optimizer states gradients and
parameters (which we collectively call the model states) are replicated across
all the different GPUs in distributed data parallel training. The way ZeRO
works is by removing this redundancy across these GPUs. Since there are three
different types of model states, there are three different phases of ZeRO,
each of them removing the redundancy for one of these states by simply
partitioning these model states across GPUs instead of replicating them.

In addition to reducing the memory from model states, ZeRO also has a bunch of
more optimizations that allows reduction in memory from other components (see
Table 3). For example, just like how model states are replicated across
multiple GPUs in data parallel training, ZeRO removes the redundancy in
activation memory by partitioning the activations across these model parallel
GPUs. We can also offload these activation memories to CPU if we don’t have
enough memory to train extremely large models. The next optimization that ZeRO
can do is to convert fragmented memory to defragmented memory on the fly
during training. During training if the memory is fragmented, we might still
run out of memory even though there might be enough fragmented memory that can
satisfy the request if they were contiguous. In ZeRO the memory
defragmentation will on the fly defragment these memory fragments so that all
the memory is contiguous and you are able to satisfy these larger memory
requests. So with all these different memory optimizations, ZeRO is able to
train models with up to 200 billion parameters up to 10 times faster than the
SOTA.

Memory Consumption | Optimization
---|---
Model State Memory | Partition optimizer state, gradient, and parameters
Activation Memory | Partition activations; Offload to GPU
Fragmented Memory | Proactively manage memory w.r.t tensor lifetime
Table 3. Different Memory Optimizations in ZeRO

##  9\. Conclusion

Given larger datasets and bigger models consistently yielding significant
improvements in accuracy, large-scale deep learning has become an inevitable
trend. As datasets increase in size and DNNs in complexity, the computational
intensity, communication cost and memory demands of deep learning increase
proportionally. Considerable efforts have been devoted to accelerating the
training speed. In this article, we give an overview of large-scale deep
learning optimization. The goal in general is two-fold: model accuracy and
model efficiency. As for the model accuracy, we investigate algorithms that
are most commonly used for optimizing, spanning from the gradient descent
variants to the (large-batch) adaptive methods, and from first-order to
second-order methods. Further, we elaborate the debatable topic of
generalization gap arises in large-batch training. As for the model
efficiency, we summarise the SOTA techniques in addressing the expensive cost
of communication overhead and memory footprint. We hope this article can
provide a clean sketch for those who are interested in training large-scale
training.

## References

  * (1)
  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Z. Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek Gordon Murray, Benoit Steiner, Paul A. Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqian Zhang. 2016\.  TensorFlow: A system for large-scale machine learning. In _OSDI_.
  * Agarwal et al. (2016a) Naman Agarwal, Brian Bullins, and Elad Hazan. 2016a.  Second-order stochastic optimization in linear time.  _stat_ 1050 (2016), 15.
  * Agarwal et al. (2016b) Naman Agarwal, Zeyuan Allen Zhu, Brian Bullins, Elad Hazan, and Tengyu Ma. 2016b.  Finding Approximate Local Minima for Nonconvex Optimization in Linear Time.  _CoRR_ abs/1611.01146 (2016).
  * Aji and Heafield (2017) Alham Fikri Aji and Kenneth Heafield. 2017.  Sparse Communication for Distributed Gradient Descent. In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017_. 440–445.
  * Alistarh et al. (2017) Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. 2017.  QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding. In _NIPS_. 1709–1720.
  * Alistarh et al. (2018) Dan Alistarh, Torsten Hoefler, Mikael Johansson, Nikola Konstantinov, Sarit Khirirat, and Cédric Renggli. 2018.  The Convergence of Sparsified Gradient Methods. In _Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada_. 5977–5987.
  * Amari (1998) Shun-ichi Amari. 1998\.  Natural Gradient Works Efficiently in Learning.  _Neural Comput._ 10, 2 (1998), 251–276.
  * Amari et al. (2000) Shun-ichi Amari, Hyeyoung Park, and Kenji Fukumizu. 2000\.  Adaptive Method of Realizing Natural Gradient Learning for Multilayer Perceptrons.  _Neural Comput._ 12, 6 (2000), 1399–1409.
  * Anil et al. (2020) Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. 2020.  Scalable second order optimization for deep learning.  _arXiv preprint arXiv:2002.09018_ (2020).
  * Anil et al. (2019) Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. 2019\.  Memory Efficient Adaptive Optimization. In _NeurIPS_.
  * Battiti (1992) Roberto Battiti. 1992\.  First- and Second-Order Methods for Learning: Between Steepest Descent and Newton’s Method.  _Neural Computation_ 4, 2 (1992), 141–166.  https://doi.org/10.1162/neco.1992.4.2.141
  * Ben-Nun and Hoefler (2019) Tal Ben-Nun and Torsten Hoefler. 2019.  Demystifying Parallel and Distributed Deep Learning: An In-depth Concurrency Analysis.  _ACM Comput. Surv._ 52, 4 (2019), 65:1–65:43.
  * Bernstein et al. (2018) Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. 2018.  signSGD: Compressed optimisation for non-convex problems. In _International Conference on Machine Learning_. PMLR, 560–569.
  * Betzel et al. (2018) Filipe Betzel, S. Karen Khatamifard, Harini Suresh, David J. Lilja, John Sartori, and Ulya R. Karpuzcu. 2018.  Approximate Communication: Techniques for Reducing Communication Bottlenecks in Large-Scale Parallel Systems.  _ACM Comput. Surv._ 51, 1 (2018), 1:1–1:32.
  * Bollapragada et al. (2016) Raghu Bollapragada, Richard Byrd, and Jorge Nocedal. 2016\.  Exact and Inexact Subsampled Newton Methods for Optimization.  arXiv:1609.08502 [math.OC]
  * Bordes et al. (2009) Antoine Bordes, Léon Bottou, and Patrick Gallinari. 2009\.  SGD-QN: Careful Quasi-Newton Stochastic Gradient Descent.  _J. Mach. Learn. Res._ 10 (2009), 1737–1754.
  * Botev et al. (2017) Aleksandar Botev, Hippolyt Ritter, and David Barber. 2017\.  Practical gauss-newton optimisation for deep learning. In _International Conference on Machine Learning_. PMLR, 557–565.
  * Bottou and Bousquet (2007) Léon Bottou and Olivier Bousquet. 2007.  The Tradeoffs of Large Scale Learning. In _Proceedings of the 20th International Conference on Neural Information Processing Systems_. 161–168.
  * Bottou et al. (2018) Léon Bottou, Frank E. Curtis, and Jorge Nocedal. 2018\.  Optimization Methods for Large-Scale Machine Learning.  _SIAM Rev._ 60, 2 (2018), 223–311.
  * Byrd et al. (2012) Richard H. Byrd, Gillian M. Chin, Jorge Nocedal, and Yuchen Wu. 2012.  Sample size selection in optimization methods for machine learning.  _Math. Program._ 134, 1 (2012), 127–155.
  * Carmon et al. (2018) Yair Carmon, John C. Duchi, Oliver Hinder, and Aaron Sidford. 2018.  Accelerated Methods for NonConvex Optimization.  _SIAM J. Optim._ 28, 2 (2018), 1751–1772.
  * Chaudhari et al. (2017) Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer T. Chayes, Levent Sagun, and Riccardo Zecchina. 2017.  Entropy-SGD: Biasing Gradient Descent Into Wide Valleys. In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_.
  * Chen et al. (2018) Chia-Yu Chen, Jungwook Choi, Daniel Brand, Ankur Agrawal, Wei Zhang, and Kailash Gopalakrishnan. 2018\.  Adacomp: Adaptive residual gradient compression for data-parallel distributed training. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 32.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. 2020.  A Simple Framework for Contrastive Learning of Visual Representations. In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ _(Proceedings of Machine Learning Research, Vol. 119)_. PMLR, 1597–1607.
  * Choi et al. (2019) Dami Choi, Christopher J. Shallue, Zachary Nado, Jaehoon Lee, Chris J. Maddison, and George E. Dahl. 2019\.  On Empirical Comparisons of Optimizers for Deep Learning.  _CoRR_ abs/1910.05446 (2019).
  * Conn et al. (2000) Andrew R. Conn, Nicholas I. M. Gould, and Philippe L. Toint. 2000\.  _Trust Region Methods_.  SIAM.
  * Devarakonda et al. (2017) Aditya Devarakonda, Maxim Naumov, and Michael Garland. 2017\.  AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks.  _ArXiv_ abs/1712.02029 (2017).
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In _NAACL-HLT (1)_.
  * Dinh et al. (2017) Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. 2017.  Sharp Minima Can Generalize For Deep Nets. In _International Conference on Machine Learning_. PMLR, 1019–1028.
  * Dozat (2016) Timothy Dozat. 2016\.  Incorporating nesterov momentum into adam.  (2016).
  * Dryden et al. (2016) Nikoli Dryden, Tim Moon, Sam Ade Jacobs, and Brian Van Essen. 2016\.  Communication quantization for data-parallel training of deep neural networks. In _2016 2nd Workshop on Machine Learning in HPC Environments (MLHPC)_. IEEE, 1–8.
  * Duchi et al. (2011) John Duchi, Elad Hazan, and Yoram Singer. 2011.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of machine learning research_ 12, 7 (2011).
  * Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. 2021.  Sharpness-Aware Minimization for Efficiently Improving Generalization.  arXiv:2010.01412 [cs.LG]
  * Ge et al. (2015) Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. 2015\.  Escaping from saddle points—online stochastic gradient for tensor decomposition. In _Conference on learning theory_. PMLR, 797–842.
  * Girshick (2015) Ross Girshick. 2015\.  Fast r-cnn. In _Proceedings of the IEEE international conference on computer vision_. 1440–1448.
  * Goldfarb et al. (2020) Donald Goldfarb, Yi Ren, and Achraf Bahamou. 2020.  Practical Quasi-Newton Methods for Training Deep Neural Networks. In _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020_.
  * Goyal et al. (2018) Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. 2018.  Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.  arXiv:1706.02677 [cs.CV]
  * Grill et al. (2020) Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Pires, Zhaohan Guo, Mohammad Azar, et al. 2020\.  Bootstrap Your Own Latent: A new approach to self-supervised learning. In _Neural Information Processing Systems_.
  * Grosse and Martens (2016) Roger B Grosse and James Martens. 2016.  A Kronecker-factored approximate Fisher matrix for convolution layers. In _ICML_ , Vol. 48. 573–582.
  * Gupta et al. (2018) Vineet Gupta, Tomer Koren, and Yoram Singer. 2018\.  Shampoo: Preconditioned Stochastic Tensor Optimization. In _ICML_ , Vol. 80. 1837–1845.
  * He et al. (2020) Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross B. Girshick. 2020.  Mask R-CNN.  _IEEE Trans. Pattern Anal. Mach. Intell._ 42, 2 (2020), 386–397.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016\.  Deep residual learning for image recognition. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 770–778.
  * He et al. (2017) Xiangnan He, Lizi Liao, Hanwang Zhang, Liqiang Nie, Xia Hu, and Tat-Seng Chua. 2017\.  Neural Collaborative Filtering. In _WWW_. 173–182.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. 1997.  Flat minima.  _Neural computation_ 9, 1 (1997), 1–42.
  * Hoffer et al. (2017) Elad Hoffer, Itay Hubara, and Daniel Soudry. 2017.  Train longer, generalize better: closing the generalization gap in large batch training of neural networks. In _Proceedings of the 31st International Conference on Neural Information Processing Systems_. 1729–1739.
  * Horvath et al. (2019) Samuel Horvath, Chen-Yu Ho, Ludovit Horvath, Atal Narayan Sahu, Marco Canini, and Peter Richtárik. 2019.  Natural Compression for Distributed Deep Learning.  _CoRR_ abs/1905.10988 (2019).
  * Huang et al. (2017) Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q Weinberger. 2017.  Densely Connected Convolutional Networks. In _CVPR_. 2261–2269.
  * Jastrzebski et al. (2017) Stanisław Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. 2017\.  Three factors influencing minima in sgd.  _arXiv preprint arXiv:1711.04623_ (2017).
  * Jia et al. (2018) Xianyan Jia, Shutao Song, W. He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, Tiegang Chen, Guangxiao Hu, Shaohuai Shi, and Xiaowen Chu. 2018\.  Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes.  _ArXiv_ abs/1807.11205 (2018).
  * Kamp et al. (2018) Michael Kamp, Linara Adilova, Joachim Sicking, Fabian Hüger, Peter Schlicht, Tim Wirtz, and Stefan Wrobel. 2018.  Efficient Decentralized Deep Learning by Dynamic Model Averaging. In _ECML/PKDD (1)_.
  * Karimireddy et al. (2019) Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian Stich, and Martin Jaggi. 2019.  Error feedback fixes signsgd and other gradient compression schemes. In _International Conference on Machine Learning_. PMLR, 3252–3261.
  * Keskar et al. (2017) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. 2017.  On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima.  arXiv:1609.04836 [cs.LG]
  * Kingma and Ba (2017) Diederik P. Kingma and Jimmy Ba. 2017.  Adam: A Method for Stochastic Optimization.  arXiv:1412.6980 [cs.LG]
  * Krizhevsky (2014) Alex Krizhevsky. 2014\.  One weird trick for parallelizing convolutional neural networks.  _arXiv preprint arXiv:1404.5997_ (2014).
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. 2012\.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ 25 (2012), 1097–1105.
  * Krizhevsky et al. (2017) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. 2017\.  ImageNet classification with deep convolutional neural networks.  _Commun. ACM_ 60, 6 (2017), 84–90.
  * Lan et al. (2019) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2019\.  Albert: A lite bert for self-supervised learning of language representations.  _arXiv preprint arXiv:1909.11942_ (2019).
  * Li et al. (2021) Conglong Li, Ammar Ahmad Awan, Hanlin Tang, Samyam Rajbhandari, and Yuxiong He. 2021.  1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB’s Convergence Speed.  _ArXiv_ abs/2104.06069 (2021).
  * Li et al. (2018) Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. 2018.  Visualizing the loss landscape of neural nets. In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_. 6391–6401.
  * Li et al. (2014a) Mu Li, David G Andersen, Alexander J Smola, and Kai Yu. 2014a.  Communication efficient distributed machine learning with the parameter server.  _Advances in Neural Information Processing Systems_ 27 (2014), 19–27.
  * Li et al. (2014b) Mu Li, Tong Zhang, Yuqiang Chen, and Alexander J Smola. 2014b.  Efficient mini-batch training for stochastic optimization. In _Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining_. 661–670.
  * Lin et al. (2020) Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. 2020.  Don’t Use Large Mini-Batches, Use Local SGD.  arXiv:1808.07217 [cs.LG]
  * Lin et al. (2018) Yujun Lin, Song Han, Huizi Mao, Yu Wang, and Bill Dally. 2018.  Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training. In _International Conference on Learning Representations_.
  * Liu et al. (2019) Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. 2019.  On the Variance of the Adaptive Learning Rate and Beyond. In _International Conference on Learning Representations_.
  * Long et al. (2015) Jonathan Long, Evan Shelhamer, and Trevor Darrell. 2015\.  Fully convolutional networks for semantic segmentation. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 3431–3440.
  * Lou et al. (2021) Yuxuan Lou, Fuzhao Xue, Zangwei Zheng, and Yang You. 2021\.  Sparse-MLP: A Fully-MLP Architecture with Conditional Computation.  _arXiv preprint arXiv:2109.02008_ (2021).
  * Martens (2010) James Martens. 2010\.  Deep learning via Hessian-free optimization. In _Proceedings of the 27th International Conference on Machine Learning (ICML-10), June 21-24, 2010, Haifa, Israel_. 735–742.
  * Martens (2020) James Martens. 2020\.  New Insights and Perspectives on the Natural Gradient Method.  _J. Mach. Learn. Res._ 21 (2020), 146:1–146:76.
  * Martens and Grosse (2015) James Martens and Roger Grosse. 2015.  Optimizing neural networks with kronecker-factored approximate curvature. In _International conference on machine learning_. PMLR, 2408–2417.
  * Masters and Luschi (2018) Dominic Masters and Carlo Luschi. 2018.  Revisiting Small Batch Training for Deep Neural Networks.  _CoRR_ abs/1804.07612 (2018).
  * McDonald et al. (2010) Ryan McDonald, Keith Hall, and Gideon Mann. 2010.  Distributed training strategies for the structured perceptron. In _Human language technologies: The 2010 annual conference of the North American chapter of the association for computational linguistics_. 456–464.
  * Micikevicius et al. (2018) Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Frederick Diamos, Erich Elsen, David García, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. 2018\.  Mixed Precision Training.  _ArXiv_ abs/1710.03740 (2018).
  * Nado et al. (2021) Zachary Nado, Justin Gilmer, Christopher J. Shallue, Rohan Anil, and George E. Dahl. 2021\.  A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes.  _ArXiv_ abs/2102.06356 (2021).
  * Nesterov (1983) Yurii E Nesterov. 1983\.  A method for solving the convex programming problem with convergence rate O (1/k^ 2). In _Dokl. akad. nauk Sssr_ , Vol. 269. 543–547.
  * Osawa et al. (2019) Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. 2019\.  Large-scale distributed second-order optimization using kronecker-factored approximate curvature for deep convolutional neural networks. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_. 12359–12367.
  * Park et al. (2019) Daniel S. Park, William Chan, Yu Zhang, Chung-Cheng Chiu, Barret Zoph, Ekin Dogus Cubuk, and Quoc V. Le. 2019.  SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition. In _Interspeech 2019_. 2613–2617.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019\.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ 32 (2019), 8026–8037.
  * Pauloski et al. (2020) J Gregory Pauloski, Zhao Zhang, Lei Huang, Weijia Xu, and Ian T Foster. 2020.  Convolutional neural network training with distributed K-FAC. In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–12.
  * Pouyanfar et al. (2019) Samira Pouyanfar, Saad Sadiq, Yilin Yan, Haiman Tian, Yudong Tao, Maria E. Presa Reyes, Mei-Ling Shyu, Shu-Ching Chen, and S. S. Iyengar. 2019.  A Survey on Deep Learning: Algorithms, Techniques, and Applications.  _ACM Comput. Surv._ 51, 5 (2019), 92:1–92:36.
  * Qian (1999) Ning Qian. 1999\.  On the momentum term in gradient descent learning algorithms.  _Neural networks_ 12, 1 (1999), 145–151.
  * Rajbhandari et al. (2020) Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020.  Zero: Memory optimizations toward training trillion parameter models. In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–16.
  * Ramezani-Kebrya et al. (2021) Ali Ramezani-Kebrya, Fartash Faghri, Ilia Markov, Vitaly Aksenov, Dan Alistarh, and Daniel M. Roy. 2021.  NUQSGD: Provably Communication-efficient Data-parallel SGD via Nonuniform Quantization.  _CoRR_ abs/2104.13818 (2021).
  * Rasley et al. (2020) Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. 2020.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_. 3505–3506.
  * Reddi et al. (2019) Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. 2019\.  On the convergence of adam and beyond.  _arXiv preprint arXiv:1904.09237_ (2019).
  * Ren et al. (2015) Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. 2015\.  Faster r-cnn: Towards real-time object detection with region proposal networks.  _Advances in neural information processing systems_ 28 (2015), 91–99.
  * Ren and Goldfarb (2021) Yi Ren and Donald Goldfarb. 2021.  Kronecker-factored Quasi-Newton Methods for Convolutional Neural Networks.  _CoRR_ abs/2102.06737 (2021).
  * Ruder (2016) Sebastian Ruder. 2016\.  An overview of gradient descent optimization algorithms.  _CoRR_ abs/1609.04747 (2016).
  * Schneider et al. (2019) Frank Schneider, Lukas Balles, and Philipp Hennig. 2019\.  DeepOBS: A Deep Learning Optimizer Benchmark Suite. In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_.
  * Schraudolph et al. (2007) Nicol N Schraudolph, Jin Yu, and Simon Günter. 2007\.  A stochastic quasi-Newton method for online convex optimization. In _Artificial intelligence and statistics_. PMLR, 436–443.
  * Seide et al. (2014) Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 2014.  1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns. In _Fifteenth Annual Conference of the International Speech Communication Association_. Citeseer.
  * Shalev-Shwartz and Ben-David (2014) Shai Shalev-Shwartz and Shai Ben-David. 2014.  _Understanding Machine Learning - From Theory to Algorithms_.  Cambridge University Press.  http://www.cambridge.org/de/academic/subjects/computer-science/pattern-recognition-and-machine-learning/understanding-machine-learning-theory-algorithms
  * Shallue et al. (2019) Christopher J. Shallue, Jaehoon Lee, Joseph M. Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. 2019.  Measuring the Effects of Data Parallelism on Neural Network Training.  _J. Mach. Learn. Res._ 20 (2019), 112:1–112:49.
  * Shazeer and Stern (2018) Noam Shazeer and Mitchell Stern. 2018.  Adafactor: Adaptive learning rates with sublinear memory cost. In _International Conference on Machine Learning_. PMLR, 4596–4604.
  * Simonyan and Zisserman (2014) Karen Simonyan and Andrew Zisserman. 2014.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ (2014).
  * Smith et al. (2020) Samuel Smith, Erich Elsen, and Soham De. 2020.  On the Generalization Benefit of Noise in Stochastic Gradient Descent. In _International Conference on Machine Learning_. PMLR, 9058–9067.
  * Smith et al. (2018) Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V Le. 2018.  Don’t Decay the Learning Rate, Increase the Batch Size. In _International Conference on Learning Representations_.
  * Smith and Le (2018) Samuel L Smith and Quoc V Le. 2018.  A Bayesian Perspective on Generalization and Stochastic Gradient Descent. In _International Conference on Learning Representations_.
  * Stich (2019) Sebastian U. Stich. 2019\.  Local SGD Converges Fast and Communicates Little. In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_.
  * Stich et al. (2018) Sebastian U Stich, Jean-Baptiste Cordonnier, and Martin Jaggi. 2018.  Sparsified SGD with Memory.  _Advances in Neural Information Processing Systems_ 31 (2018), 4447–4458.
  * Strom (2015) Nikko Strom. 2015\.  Scalable distributed DNN training using commodity GPU cloud computing. In _INTERSPEECH 2015, 16th Annual Conference of the International Speech Communication Association, Dresden, Germany, September 6-10, 2015_. 1488–1492.
  * Sun (2019) Ruoyu Sun. 2019\.  Optimization for deep learning: theory and algorithms.  _CoRR_ abs/1912.08957 (2019).
  * Sutskever et al. (2013a) Ilya Sutskever, James Martens, George E. Dahl, and Geoffrey E. Hinton. 2013a.  On the importance of initialization and momentum in deep learning. In _Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013_ _(JMLR Workshop and Conference Proceedings, Vol. 28)_. 1139–1147.
  * Sutskever et al. (2013b) Ilya Sutskever, James Martens, George E. Dahl, and Geoffrey E. Hinton. 2013b.  On the importance of initialization and momentum in deep learning. In _Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013_ _(JMLR Workshop and Conference Proceedings, Vol. 28)_. 1139–1147.
  * Szegedy et al. (2015) Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott E. Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. 2015.  Going deeper with convolutions. In _IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2015, Boston, MA, USA, June 7-12, 2015_. IEEE Computer Society, 1–9.
  * Tang et al. (2021) Hanlin Tang, Shaoduo Gan, Ammar Ahmad Awan, Samyam Rajbhandari, Conglong Li, Xiangru Lian, Ji Liu, Ce Zhang, and Yuxiong He. 2021.  1-bit Adam: Communication Efficient Large-Scale Training with Adam’s Convergence Speed. In _ICML_.
  * Tang et al. (2020) Zhenheng Tang, Shaohuai Shi, Xiaowen Chu, Wei Wang, and Bo Li. 2020.  Communication-Efficient Distributed Deep Learning: A Comprehensive Survey.  _CoRR_ abs/2003.06307 (2020).
  * Tieleman and Hinton (2012) T. Tieleman and G. Hinton. 2012.  Lecture 6.5—RmsProp: Divide the gradient by a running average of its recent magnitude.  COURSERA: Neural Networks for Machine Learning.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017\.  Attention is All you Need. In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (Eds.). 5998–6008.
  * Verma et al. (2017) Abhishek Verma, Hussam Qassim, and David Feinzimer. 2017\.  Residual squeeze CNDS deep learning CNN model for very large scale places image recognition.  _2017 IEEE 8th Annual Ubiquitous Computing, Electronics and Mobile Communication Conference (UEMCON)_ (2017), 463–469.
  * Wang et al. (2020) Meng Wang, Weijie Fu, Xiangnan He, Shijie Hao, and Xindong Wu. 2020.  A survey on large-scale machine learning.  _IEEE Transactions on Knowledge and Data Engineering_ (2020).
  * Wangni et al. (2018) Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. 2018\.  Gradient Sparsification for Communication-Efficient Distributed Optimization. In _Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada_ , Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (Eds.). 1306–1316.
  * Wen et al. (2017) Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. 2017.  TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning. In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_. 1509–1519.
  * Wilson et al. (2017) Ashia C. Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. 2017.  The Marginal Value of Adaptive Gradient Methods in Machine Learning. In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_. 4148–4158.
  * Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Lukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. 2016.  Google’s Neural Machine Translation System: Bridging the Gap between Human and Machine Translation.  _CoRR_ abs/1609.08144 (2016).
  * Xiong et al. (2020) Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. 2020\.  On Layer Normalization in the Transformer Architecture. In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ _(Proceedings of Machine Learning Research, Vol. 119)_. 10524–10533.
  * Xu et al. (2020) Peng Xu, Fred Roosta, and Michael W. Mahoney. 2020.  Second-order Optimization for Non-convex Machine Learning: an Empirical Study. In _Proceedings of the 2020 SIAM International Conference on Data Mining, SDM 2020, Cincinnati, Ohio, USA, May 7-9, 2020_. 199–207.
  * Xue et al. (2021) Fuzhao Xue, Ziji Shi, Yuxuan Lou, Yong Liu, and Yang You. 2021.  Go Wider Instead of Deeper.  _arXiv preprint arXiv:2107.11817_ (2021).
  * Yamazaki et al. (2019) Masafumi Yamazaki, Akihiko Kasagi, Akihiro Tabuchi, Takumi Honda, Masahiro Miwa, Naoto Fukumoto, Tsuguchika Tabaru, Atsushi Ike, and Kohta Nakashima. 2019.  Yet Another Accelerated SGD: ResNet-50 Training on ImageNet in 74.7 seconds.  _ArXiv_ abs/1903.12650 (2019).
  * Ying et al. (2018) Chris Ying, Sameer Kumar, Dehao Chen, Tao Wang, and Youlong Cheng. 2018.  Image Classification at Supercomputer Scale.  _CoRR_ abs/1811.06992 (2018).
  * You et al. (2017) Yang You, Igor Gitman, and Boris Ginsburg. 2017.  Large Batch Training of Convolutional Networks.  arXiv:1708.03888 [cs.CV]
  * You et al. (2020) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. 2020\.  Large Batch Optimization for Deep Learning: Training BERT in 76 minutes.  arXiv:1904.00962 [cs.LG]
  * Zeiler (2012) Matthew D. Zeiler. 2012\.  ADADELTA: An Adaptive Learning Rate Method.  _ArXiv_ abs/1212.5701 (2012).
  * Zhang et al. (2019) Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. 2019\.  Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model.  _Advances in neural information processing systems_ 32 (2019), 8196–8207.