Year: 2020

Abstract:
Stochastic gradient descent (SGD) and its variants have been the dominating optimization methods in machine learning. Compared with small batch training, SGD with large batch training can better utilize the computational power of current multi-core systems like GPUs and can reduce the number of communication rounds in distributed training. Hence, SGD with large batch training has attracted more and more attention. However, existing empirical results show that large batch training typically leads to a drop of generalization accuracy. As a result, large batch training has also become a challenging topic. In this paper, we propose a novel method, called stochastic normalized gradient descent with momentum (SNGM), for large batch training. We theoretically prove that compared to momentum SGD (MSGD) which is one of the most widely used variants of SGD, SNGM can adopt a larger batch size to converge to the -stationary point with the same computation complexity (total number of gradient computation). Empirical results on deep learning also show that SNGM can achieve the state-of-the-art accuracy with a large batch size.

heading: 1 INTRODUCTION
text: In machine learning, we often need to solve the following empirical risk minimization problem:
min w∈Rd
F (w) = 1
n n∑ i=1 fi(w), (1)
where w ∈ Rd denotes the model parameter, n denotes the number of training samples, fi(w) denotes the loss on the ith training sample. The problem in (1) can be used to formulate a broad family of machine learning models, such as logistic regression and deep learning models.
Stochastic gradient descent (SGD) Robbins & Monro (1951) and its variants have been the dominating optimization methods for solving (1). SGD and its variants are iterative methods. In the tth iteration, these methods randomly choose a subset (also called mini-batch) It ⊂ {1, 2, . . . , n} and compute the stochastic mini-batch gradient 1/B ∑ i∈It ∇fi(wt) for updating the model parameter, where B = |It| is the batch size. Existing works Li et al. (2014b); Yu et al. (2019a) have proved that with the batch size of B, SGD and its momentum variant, called momentum SGD (MSGD), achieve a O(1/ √ TB) convergence rate for smooth non-convex problems, where T is total number of model parameter updates.
With the population of multi-core systems and the easy implementation for data parallelism, many distributed variants of SGD have been proposed, including parallel SGD Li et al. (2014a), decentralized SGD Lian et al. (2017), local SGD Yu et al. (2019b); Lin et al. (2020), local momentum SGD Yu et al. (2019a) and so on. Theoretical results show that all these methods can achieve aO(1/ √ TKb) convergence rate for smooth non-convex problems. Here, b is the batch size on each worker and K is the number of workers. By setting Kb = B, we can observe that the convergence rate of these distributed methods is consistent with that of sequential methods. In distributed settings, a small number of model parameter updates T implies a small synchronize cost and communication cost. Hence, a small T can further speed up the training process. Based on theO(1/ √ TKb) convergence rate, we can find that if we adopt a larger b, the T will be smaller. Hence, large batch training can reduce the number of communication rounds in distributed training. Another benefit of adopting
large batch training is to better utilize the computational power of current multi-core systems like GPUs You et al. (2017). Hence, large batch training has recently attracted more and more attention in machine learning.
Unfortunately, empirical results LeCun et al. (2012); Keskar et al. (2017) show that existing SGD methods with a large batch size will lead to a drop of generalization accuracy on deep learning models. Figure 1 shows the comparison of training loss and test accuracy between MSGD with a small batch size and MSGD with a large batch size. We can find that large batch training does degrade both training loss and test accuracy. Many works try to explain this phenomenon Keskar et al. (2017); Hoffer et al. (2017). They observe that SGD with a small batch size typically makes the model parameter converge to a flatten minimum while SGD with a large batch size typically makes the model parameter fall into the region of a sharp minimum. And usually, a flatten minimum can achieve better generalization ability than a sharp minimum. Hence, large batch training has also become a challenging topic.
Recently, many methods have been proposed for improving the performance of SGD with a large batch size. The work in Goyal et al. (2017); You et al. (2020) proposes many tricks like warmup, momentum correction and linearly scaling the learning rate, for large batch training. The work in You et al. (2017) observes that the norms of gradient at different layers of deep neural networks are widely different and the authors propose the layer-wise adaptive rate scaling method (LARS). The work in Ginsburg et al. (2019) also proposes a similar method that updates the model parameter in a layer-wise way. Most of these methods lack theoretical evidence to explain why they can adopt a large batch size. Although the work in You et al. (2020) proposes some theoretical explanations for LARS, the implementation is still not consistent with its theorems in which both of the momentum coefficient and weight decay are set as zeros.
In this paper, we propose a novel method, called stochastic normalized gradient descent with momentum (SNGM), for large batch training. SNGM combines normalized gradient Nesterov (2004); Hazan et al. (2015); Wilson et al. (2019) and Polyak’s momentum technique Polyak (1964) together. The main contributions of this paper are outlined as follows:
• We theoretically prove that compared to MSGD which is one of the most widely used variants of SGD, SNGM can adopt a larger batch size to converge to the -stationary point with the same computation complexity (total number of gradient computation). That is to say, SNGM needs a smaller number of parameter update, and hence has faster training speed than MSGD.
• For a relaxed smooth objective function (see Definition 2), we theoretically show that SNGM can achieve an -stationary point with a computation complexity of O(1/ 4). To the best of our knowledge, this is the first work that analyzes the computation complexity of stochastic optimization methods for a relaxed smooth objective function.
• Empirical results on deep learning also show that SNGM can achieve the state-of-the-art accuracy with a large batch size.

heading: 2 PRELIMINARIES
text: In this paper, we use ‖·‖ to denote the Euclidean norm, use w∗ to denote one of the optimal solutions of (1), i.e., w∗ ∈ argminw F (w). We call w an -stationary point of F (w) if ‖∇F (w)‖ ≤ . The computation complexity of an algorithm is the total number of its gradient computation. We also give the following assumption and definitions:
Assumption 1 (σ-bounded variance) For any w, E‖∇fi(w)−∇F (w)‖2 ≤ σ2 (σ > 0).
Definition 1 (Smoothness) A function φ(·) is L-smooth (L > 0) if for any u,w,
φ(u) ≤ φ(w) +∇φ(w)>(u−w) + L 2 ‖u−w‖2.
L is called smoothness constant in this paper.
Definition 2 (Relaxed smoothness Zhang et al. (2020)) A function φ(·) is (L, λ)-smooth (L ≥ 0, λ ≥ 0) if φ(·) is twice differentiable and for any w,
‖∇2φ(w)‖ ≤ L+ λ‖∇φ(w)‖,
where∇2φ(w) denotes the Hessian matrix of φ(w).
From the above definition, we can observe that if a function φ(w) is (L, 0)-smooth, then it is a classical L-smooth function Nesterov (2004). For a (L, λ)-smooth function, we have the following property Zhang et al. (2020):
Lemma 1 If φ(·) is (L, λ)-smooth, then for any u,w, α such that ‖u−w‖ ≤ α, we have
‖∇φ(u)‖ ≤ (Lα+ ‖∇φ(w)‖)eλα.
All the proofs of lemmas and corollaries of this paper are put in the supplementary.

heading: 3 RELATIONSHIP BETWEEN SMOOTHNESS CONSTANT AND BATCH SIZE
text: In this section, we deeply analyze the convergence property of MSGD to find the relationship between smoothness constant and batch size, which provides insightful hint for designing our new method SNGM.
MSGD can be written as follows:
vt+1 = βvt + gt, (2) wt+1 = wt − ηvt+1, (3)
where gt = 1/B ∑ i∈It ∇fi(wt) is a stochastic mini-batch gradient with a batch size of B, and vt+1 is the Polyak’s momentum Polyak (1964).
We aim to find how large the batch size can be without loss of performance. The convergence rate of MSGD with the batch size B for L-smooth functions can be derived from the work in Yu et al. (2019a). That is to say, when η ≤ (1− β)2/((1 + β)L), we obtain
1
T T−1∑ t=0 E‖∇F (wt)‖2 ≤ 2(1− β)[F (w0)− F (w∗)] ηT + Lησ2 (1− β)2B + 4L2η2σ2 (1− β)2 ,
=O( B ηC ) +O( η B ) +O(η2), (4)
where C = TB denotes the computation complexity (total number of gradient computation). According to Corollary 1 in Yu et al. (2019a), we set η = √ B/ √ T = B/ √ C and obtain that
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤
√ O( 1√
C ) +O(B
2
C ). (5)
Algorithm 1 SNGM Initialization: η > 0, β ∈ [0, 1), B > 0, T > 0,u0 = 0,w0; for t = 0, 1, . . . , T − 1 do
Randomly choose B function indices, denoted as It; Compute a mini-batch gradient gt = 1B ∑ i∈It ∇fi(wt); ut+1 = βut + gt ‖gt‖ ;
wt+1 = wt − ηut+1; end for
Since η ≤ (1− β)2/((1+ β)L) is necessary for (4), we firstly obtain that B ≤ O( √ C/L). Furthermore, according to the right term of (5), we have to set B such that B2/C ≤ 1/ √ C, i.e., B ≤ C1/4, for O(1/ 4) computation complexity guarantee. Hence in MSGD, we have to set the batch size satisfying
B ≤ O(min{ √ C L , C1/4}). (6)
We can observe that a larger L leads to a smaller batch size in MSGD. If B does not satisfy (6), MSGD will get higher computation complexity.
In fact, to the best of our knowledge, among all the existing convergence analysis of SGD and its variants on both convex and non-convex problems, we can observe three necessary conditions for the O(1/ 4) computation complexity guarantee Li et al. (2014b;a); Lian et al. (2017); Yu et al. (2019b;a): (a) the objective function is L-smooth; (b) the learning rate η is less than O(1/L); (c) the batch size B is proportional to the learning rate η. One direct corollary is that the batch size is limited by the smooth constant L, i.e., B ≤ O(1/L). Hence, we can not increase the batch size casually in these SGD based methods. Otherwise, it may slow down the convergence rate and we need to compute more gradients, which is consistent with the observations in Hoffer et al. (2017).

heading: 4 STOCHASTIC NORMALIZED GRADIENT DESCENT WITH MOMENTUM
text: In this section, we propose our novel methods, called stochastic normalized gradient descent with momentum (SNGM), which is presented in Algorithm 1. In the t-th iteration, SNGM runs the following update:
ut+1 = βut + gt ‖gt‖ , (7)
wt+1 = wt − ηut+1, (8) where gt = 1/B ∑ i∈It ∇fi(wt) is a stochastic mini-batch gradient with a batch size of B. When β = 0, SNGM will degenerate to stochastic normalized gradient descent (SNGD) Hazan et al. (2015). The ut is a variant of Polyak’s momentum. But different from Polyak’s MSGD which adopts gt directly for updating ut+1, SNGM adopts the normalized gradient gt/‖gt‖ for updating ut+1. In MSGD, we can observe that if gt is large, then ut may be large as well and this may lead to a bad model parameter. Hence, we have to control the learning rate in MSGD, i.e., η ≤ (1/L), for a L-smooth objective function. The following lemma shows that ‖ut‖ in SNGM can be well controlled whatever gt is large or small.
Lemma 2 Let {ut} be the sequence produced by (7), then we have ∀t ≥ 0,
‖ut‖ ≤ 1
1− β .

heading: 4.1 SMOOTH OBJECTIVE FUNCTION
text: For a smooth objective function, we have the following convergence result of SNGM:
Table 1: Comparison between MSGD and SNGM for a L-smooth objective function. C denotes the computation complexity (total number of gradient computation).
1 T
∑T−1
t=0 E‖∇F (wt)‖ learning rate batch size
MSGD √ O( 1√C ) +O( B2 C ) B√ C min{ √ C L , C 1/4} SNGM O( 1C1/4 ) √ B√ C √ C
Theorem 1 Let F (w) be a L-smooth function (L > 0). The sequence {wt} is produced by Algorithm 1. Then for any η > 0, B > 0, we have
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] ηT + Lκη + 2σ√ B , (9)
where κ = 1+β(1−β)2 .
Proof 1 See the supplementary.
We can observe that different from (4) which needs η ≤ O(1/L), (9) is true for any positive learning rate. According to Theorem 1, we obtain the following computation complexity of SNGM:
Corollary 1 Let F (w) be a L-smooth function (L > 0). The sequence {wt} is produced by Algorithm 1. Given any total number of gradient computation C > 0, let T = dC/Be,
B =
√ C(1− β)σ2
2L(1 + β)(F (w0)− F (w∗)) ,
and
η =
√ 2(1− β)3(F (w0)− F (w∗))B
(1 + β)LC .
Then we have
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2 √ 2 4
√ 8L(1 + β)[F (w0)− F (w∗)]σ2
(1− β)C = O( 1 C1/4 ).
Hence, the computation complexity for achieving an -stationary point is O(1/ 4).
It is easy to verify that the η and B in Corollary 1 make the right term of (9) minimal. However, the η and B rely on the L and F (w∗) which are usually unknown in practice. The following corollary shows the computation complexity of SNGM with simple settings about learning rate and batch size.
Corollary 2 Let F (w) be a L-smooth function (L > 0). The sequence {wt} is produced by Algorithm 1. Given any total number of gradient computation C > 0, let T = dC/Be, B = √ C and
η = √ B/C. Then we have
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] C1/4 + L(1 + β) (1− β)2C1/4 + 2σ C1/4 = O( 1 C1/4 ).
Hence, the computation complexity for achieving an -stationary point is O(1/ 4).
According to Corollary 2, the batch size of SNGM can be set as O( √ C), which does not rely on the smooth constant L, and the O(1/ 4) computation complexity is still guaranteed (see Table 1). Hence, SNGM can adopt a larger batch size than MSGD, especially when L is large.

heading: 4.2 RELAXED SMOOTH OBJECTIVE FUNCTION
text: Recently, the authors in Zhang et al. (2020) observe the relaxed smooth property in deep neural networks. According to Definition 2, the relaxed smooth property is more general than L-smooth property. For a relaxed smooth objective function, we have the following convergence result of SNGM:
Theorem 2 Let F (w) be a (L, λ)-smooth function (L ≥ 0, λ > 0). The sequence {wt} is produced by Algorithm 1 with the learning rate η and batch size B. Then we have
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] ηT + 8Lκη + 4σ√ B , (10)
where κ = 1+β(1−β)2 and η ≤ 1/(8κλ).
Proof 2 The proof is similar to that of Theorem 1. See the supplementary.
According to Theorem 2, we obtain the computation complexity of SNGM:
Corollary 3 LetF (w) be a (L, λ)-smooth function (L ≥ 0, λ ≥ 0). The sequence {wt} is produced by Algorithm 1. Given any total number of gradient computation C > 0, let T = dC/Be, B = √ C
and η = 4 √ 1/C ≤ 1/(8κλ). Then we have
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] C1/4 + 8L(1 + β) (1− β)2C1/4 + 4σ C1/4 = O( 1 C1/4 ).
Hence, the computation complexity for achieving an -stationary point is O(1/ 4).
According to Corollary 3, SNGM with a batch size of B = √ C can still guarantee a O(1/ 4) computation complexity for a relaxed smooth objective function.

heading: 5 EXPERIMENTS
text: All experiments are conducted with the platform of PyTorch, on a server with eight NVIDIA Tesla V100 (32G) GPU cards. The datasets for evaluation include CIFAR10 and ImageNet.

heading: 5.1 ON CIFAR10
text: First, we evaluate SNGM by training ResNet20 and ResNet56 on CIFAR10. CIFAR10 contains 50k training samples and 10k test samples. We compare SNGM with MSGD and an existing large batch training method LARS You et al. (2017). We implement LARS by using the open source code 1. The standard strategy He et al. (2016) for training the two models on CIFAR10 is using MSGD with a weight decay of 0.0001, a batch size of 128, an initial learning rate of 0.1, and dividing the learning rate at the 80th and 120th epochs. We also adopt this strategy for MSGD in this experiment. For SNGM and LARS, we set a large batch size of 4096 and also a weight decay of 0.0001. Following You et al. (2017), we adopt the poly power learning rate strategy and adopt the gradient accumulation Ott et al. (2018) with a batch size of 128 for the two large batch training methods. The momentum coefficient is 0.9 for all methods. Different from existing heuristic methods for large batch training, we do not adopt the warm-up strategy for SNGM.
The results are presented in Figure 2. As can be seen, SNGM achieves better convergence rate on training loss than LARS. The detailed information about the final convergence results is presented in Table 2. We can observe that MSGD with a batch size of 4096 leads to a significant drop of test accuracy. SNGM with a batch size of 4096 achieves almost the same test accuracy as MSGD with a batch size of 128. But the other large batch training method LARS achieves worse test accuracy than MSGD with a batch size of 128. These results successfully verify the effectiveness of SNGM.
1https://github.com/noahgolmant/pytorch-lars

heading: 5.2 ON IMAGENET
text: We also compare SNGM with MSGD by training ResNet18 and ResNet50 on ImageNet. The standard strategy He et al. (2016) for training the two models on ImageNet is using MSGD with a weight decay of 0.0001, a batch size of 256, an initial learning rate of 0.1, and dividing the learning rate at the 30th and 60th epochs. We also adopt this strategy for MSGD in this experiment. For SNGM, we set a larger batch size of 8192 and a weight decay of 0.0001. We still adopt the poly power learning rate and the gradient accumulation with a batch size of 128 for SNGM. We do not adopt the warmup strategy for SNGM either. The momentum coefficient is 0.9 in the two methods. The results are
presented in Figure 3 and Table 3. As can be seen, SNGM with a larger batch size achieves almost the same test accuracy as MSGD with a small batch size.

heading: 6 CONCLUSION
text: In this paper, we propose a novel method called stochastic normalized gradient descent with momentum (SNGM), for large batch training. We theoretically prove that compared to MSGD which is one of the most widely used variants of SGD, SNGM can adopt a larger batch size to converge to the -stationary point with the same computation complexity. Empirical results on deep learning also show that SNGM can achieve the state-of-the-art accuracy with a large batch size.

heading: A APPENDIX

heading: A.1 PROOF OF LEMMA 1
text: The proof follows Zhang et al. (2020). We put it here for completeness. For any u,w, let r(x) = x(u−w) +w, p(x) = ‖∇φ(r(x))‖, x ∈ [0, 1]. Then we have
p(x) =‖∇φ(r(x))‖ = ‖ ∫ x 0 Hφ(r(y))r ′(y)dy +∇φ(r(0))‖
=‖ ∫ x 0 Hφ(r(y))(u−w)dy +∇φ(w)‖
≤‖u−w‖ ∫ x 0 ‖Hφ(r(y))‖dy + ‖∇φ(w)‖
≤α ∫ x 0 (L+ λ‖∇φ(r(y))‖)dy + ‖∇φ(w)‖
=Lα+ ‖∇φ(w)‖+ λα ∫ x 0 p(y)dy.
According to Gronwall’s Inequality, we obtain
p(x) ≤ (Lα+ ‖∇φ(w)‖)eλα.

heading: A.2 PROOF OF LEMMA 2
text: According to (7), we have
‖ut+1‖ ≤β‖ut‖+ 1 ≤β2‖ut−1‖+ β + 1 ≤βt+1‖u0‖+ βt + βt−1 + · · ·+ 1
≤ 1 1− β .

heading: A.3 PROOF OF THEOREM 1
text: Let zt = wt + β1−β (wt −wt−1), then we have wt+1 = wt − η gt ‖gt‖ + β(wt −wt−1) and
zt+1 = 1
1− β wt+1 −
β
1− β wt
=zt − η 1− β gt ‖gt‖ .
Using the smooth property, we obtain
F (zt+1) ≤F (zt)− η
1− β ∇F (zt)T gt ‖gt‖ + Lη2 2(1− β)2
=F (zt)− η
1− β ‖gt‖+
Lη2
2(1− β)2
− η 1− β [(∇F (zt)−∇F (wt))T gt ‖gt‖ + (∇F (wt)− gt)T gt ‖gt‖ ]
≤F (zt)− η
1− β ‖gt‖+
Lη2
2(1− β)2 +
η
1− β [L‖zt −wt‖+ ‖∇F (wt)− gt‖] (11)
Since wt+1 −wt = β(wt −wt−1)− ηgt/‖gt‖, we obtain
‖wt+1 −wt‖ ≤ β‖wt −wt−1‖+ η ≤ η
1− β .
Hence, ‖wt −wt−1‖ ≤ η/(1− β) and
‖zt −wt‖ = β
1− β ‖wt −wt−1‖ ≤
βη
(1− β)2 . (12)
Combining the above equations, we obtain
‖gt‖ ≤ (1− β)[F (zt)− F (zt+1)]
η +
Lη
2(1− β) +
Lβη
(1− β)2 + ‖∇F (wt)− gt‖.
Since ‖∇F (wt)‖ ≤ ‖∇F (wt)− gt‖+ ‖gt‖, we obtain
‖∇F (wt)‖ ≤ (1− β)[F (zt)− F (zt+1)]
η +
Lη
2(1− β) +
Lβη
(1− β)2 + 2‖∇F (wt)− gt‖.
Using the fact that E‖∇F (wt)− gt‖ ≤ σ/ √ B and summing up the above inequality from t = 0 to T − 1, we obtain
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] ηT + Lκη + 2σ√ B .

heading: A.4 PROOF OF THEOREM 2
text: Let zt = wt + β1−β (wt −wt−1), then we have wt+1 = wt − η gt ‖gt‖ + β(wt −wt−1) and
zt+1 = 1
1− β wt+1 −
β
1− β wt
= 1
1− β [wt − η gt ‖gt‖ + β(wt −wt−1)]− β 1− β wt
= 1
1− β wt −
β
1− β wt−1 −
η 1− β gt ‖gt‖
=zt − η 1− β gt ‖gt‖ .
Using the Taylor theorem, there exists ξt such that
F (zt+1) ≤F (zt)− η
1− β ∇F (zt)T gt ‖gt‖ + ‖HF (ξt)‖η2 2(1− β)2
=F (zt)− η
1− β ‖gt‖+
‖HF (ξt)‖η2
2(1− β)2
− η 1− β [(∇F (zt)−∇F (wt))T gt ‖gt‖ + (∇F (wt)− gt)T gt ‖gt‖ ]. (13)
Let ψt(w) = (∇F (w)−∇F (wt))T gt‖gt‖ . Using the Taylor theorem, there exists ζt such that
|ψt(zt)| =|ψt(wt) +∇ψt(ζt)(zt −wt)| = |∇ψ(ζt)(zt −wt)| ≤‖HF (ζt)‖‖zt −wt‖. (14)
Combining (13) and (14), we obtain
‖gt‖ ≤ (1− β)[F (zt)− F (zt+1)] η + ‖HF (ξt)‖η 2(1− β)
+ (‖HF (ζt)‖‖zt −wt‖+ ‖∇F (wt)− gt‖). (15)
Since wt+1 −wt = β(wt −wt−1)− ηgt/‖gt‖, we obtain
‖wt+1 −wt‖ ≤ β‖wt −wt−1‖+ η ≤ η
1− β .
Hence, ‖wt −wt−1‖ ≤ η/(1− β) and
‖zt −wt‖ = β
1− β ‖wt −wt−1‖ ≤
βη
(1− β)2 . (16)
Combining (15) and (16), we obtain
‖gt‖ ≤ (1− β)[F (zt)− F (zt+1)] η + ‖HF (ξt)‖η 2(1− β) + ‖HF (ζt)‖βη (1− β)2
+ ‖∇F (wt)− gt‖. Since ‖∇F (wt)‖ ≤ ‖∇F (wt)− gt‖+ ‖gt‖, we obtain
‖∇F (wt)‖ ≤ (1− β)[F (zt)− F (zt+1)]
η +
η
2(1− β) ‖HF (ξt)‖+
βη
(1− β)2 ‖HF (ζt)‖
+ 2‖∇F (wt)− gt‖.
Next, we bound the two Hessian matrices. For convenience, we denote κ = 1+β(1−β)2 . Since ‖zt − wt‖ ≤ βη/(1− β)2 and
‖zt+1 −wt‖ ≤‖zt+1 − zt‖+ ‖zt −wt‖
≤η( 1 1− β + β (1− β)2 )
≤κη
≤ 1 λ ,
we obtain ‖HF (ζt)‖ ≤ L+ (L+ λ‖∇F (wt)‖)e, ‖HF (ξt)‖ ≤ L+ (L+ λ‖∇F (wt)‖)e.
Then we obtain
‖∇F (wt)‖ ≤ (1− β)[F (zt)− F (zt+1)]
η + [
η
2(1− β) +
βη
(1− β)2 ][L+ (L+ λ‖∇F (wt)‖)e]
+ 2‖∇F (wt)− gt‖
≤ (1− β)[F (zt)− F (zt+1)] η + 4κη[L+ λ‖∇F (wt)‖]
+ 2‖∇F (wt)− gt‖. Since 4λκη ≤ 1/2, we obtain
‖∇F (wt)‖ ≤ 2(1− β)[F (zt)− F (zt+1)]
η + 8Lcη + 4‖∇F (wt)− gt‖.
Summing up the above inequality from t = 0 to T − 1, we obtain
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2(1− β)[F (w0)− F (w∗)] ηT + 8Lκη + 4σ√ B .
where η ≤ 18λκ and we use the fact that E‖∇F (wt)− gt‖ ≤ σ/ √ B.

heading: A.5 PROOF OF COROLLARY 1
text: Let x = 2(1− β)[F (w0)− F (w∗)], y = Lκ, z = 2σ. Then we have
xB Cη + yη + z√ B ≥ 2
√ xyB
C + z√ B ≥ 2
√ 2z √ xyB
C = 2 √ 2 4
√ xyz2
C . The equal sign works if and only if η = √ Bx/Cy, B = √ Cz2/(4xy). Then we obtain
1
T T−1∑ t=0 E‖∇F (wt)‖ ≤ 2 √ 2 4
√ 8L(1 + β)[F (w0)− F (w∗)]σ2
(1− β)C .

heading: A.6 PROOF OF COROLLARY 2
text: By plugging T = dC/Be, B = √ C and η = 4 √ 1/C into (9), we obtain the result.

heading: A.7 PROOF OF COROLLARY 3
text: By plugging T = dC/Be, B = √ C and η = 4 √ 1/C into (10), we obtain the result.

Review summaries:
This paper proposes a new stochastic normalized gradient descent method with momentum (SNGM) for large batch training. They prove that unlike mometum SGD (MSGD), SNGM can adopt larger batch size to converge to the epsilon-stationary point with the same computation complexity (total number of gradient computation). The paper shows that SNGM with large batches is comparable to MSGD with small batches for training ResNet on CIFAR10 and ImageNet. The paper also shows that SNGM outperforms LARS on CIFAR10.

Large batch training has been observed to not only significantly improve the training speed but also lead to a worse generalization performance. This paper considers how to improve the performance of MSGD in large batch training. They propose the so called normalized MSGD where instead of the gradient, the algorithm uses the normalized gradient to update the momentum. They also provide theoretical justification of this change by considering smooth and relaxed smooth function. O(1/\epsilon^4)  convergence rate is established.