GradientMix: A Simple yet Effective Regularization for Large Batch Training
Authors: Jihun Yun
Jung Hyun Lee
Eunho Yang
Venue: Submitted to ICLR 2023 (2023)
Keywords: Large Batch Training
Deep Learning Optimization

Abstract: Stochastic gradient descent (SGD) is the core tool for training deep neural networks. As modern deep learning tasks become more complex and state-of-the-art architectures grow as well, network training with SGD takes a huge amount of time; for example, training ResNet on the ImageNet dataset or BERT pre-training can take days to dozens of days. To reduce the network training time, distributed learning using a large batch size for SGD has been one of the main active research areas in recent years, but this approach entails a significant degradation in generalization. To address this issue, in this paper, we propose a simple yet effective regularization technique, GradientMix, for large-scale distributed learning. GradientMix can enhance the generalization in large batch regimes by giving appropriate noise through a mixup of local gradients computed at multiple devices, which is contrary to the conventions that simply average local gradients. Furthermore, GradientMix is optimizer-agnostic, hence can be applied to any popular optimization algorithm as long as the overall loss is expressed as the sum of the subgroup losses. Our extensive experiments show the effectiveness in both small and large-scale problems, and especially we consistently achieve state-of-the-art performance for various optimizers on training ResNet-50 on the ImageNet dataset with 32K batch size.

Under review as a conference paper at ICLR 2023
GRADIENTMIX: A SIMPLE YET EFFECTIVE REGULAR- IZATION FOR LARGE BATCH TRAINING
Anonymous authors Paper under double-blind review
ABSTRACT
Stochastic gradient descent (SGD) is the core tool for training deep neural networks. As modern deep learning tasks become more complex and of-teof-the-art architec- tures grow as well, network training with SGD takes huge amount of time; for example, training ResNet on the ImageNet dataset or BERT pre-training can take days to dozens of days. To reduce the network training time, distributed learning using a large batch size for SGD has been one of the main active research areas recent years, but this approach entails significant degradation in generalization. To address this issue, in this paper, we propose a simple yet effective regularization technique, GradientMix, for large-scale distributed learning. GradientMix enhance the generalization in large batch regimes by giving appropriate noise through mixup of local gradients computed at multiple devices, which is contrary to the conventions that simply average local gradients. Furthermore, GradientMix is optimizer-agnosticc. hence can be applied any popular optimization algorithm as long as the overall loss is expressed as the sum of the subgroup losses. Our ex- tensive experiments show the effectiveness both small and large-scale problems, and especially we consistently achieve ofate-ofttheardt performance for various optimizers on training ResNet-50 on the ImageNet dataset with 32K batch size.
INTRODUCTION
Stochastic gradient descent (SGD) a classical training approach commonly used for training most deep learning models, in which model parameters are updated in the opposite direction of the gradient of the loss function computed on a mini-batch. The recent developments in hardware such as GPU and TPU enables several complex and state-of state-of-the-art models to be trained using large batch whose gradient is computed efficiently via data parallelism. Unfortunately, however, it is known that large batch training typically suffers from severe degradation in generalization compared to small batch training. The reason for poor generalization of large batch training still has not been fully uncovered, which leads many researchers to investigate this phenomenon actively in several directions.
For that reason, many efforts have been made to shed light on whether generalization has to do with sharp minima in the perspective of loss landscape. As one of pioneering work, Keskar et al. (2016) argued that the generalization gap is primarily due to the sharp local minima obtained from large batch training. Another study He et al. (2019a) further showed that local minima as well as sharp minima lie on an asymmetric valley and that a local minimum biased towards the flat side generalizes better than the exact empirical minimizer. Despite some evidence of the relationship between sharp minima and generalization, Dinh et al. (2017) substantiated that the performance degradation in large batch training has nothing to do with sharp minima by showing that all minima of the neural networks can be arbitrarily sharp according to the reparametrization that does not change the network output.
Given that there are limitations to enhance generalization using the curvature information of the loss function, designing an optimization algorithm specific to large batch regime arises as another dimension. The first representative this line of work is LARS optimizer You et al. (2017), which exploits the ratio of the parameter size to gradient size. Yet, as pointed out in Zhang et al. (2020), LARS performs poorly for attention- based model such as Transformer on the ground that LARS is an algorithm based on vanilla SGD. In order to address the issue, the second representative optimizer LAMB is proposed in You et al. (2019), which inherits the spirit of LARS and one of the adaptive gradient methods, Adam Kingma & Ba (2015). The aforementioned optimizers make it possible

---

Under review as a conference paper at ICLR 2023
to train ResNet-50 on the ImageNet dataset and BERT around an hour by scaling the batch size to more than 32K. However, Nado et al. (2021) empirically confirmed that traditional optimizers such as Nesterov can achieve as high performance as large batch optimizers (e.g. LARS, LAMB, AdaScale SGD, and etc.) if one puts the same effort into the hyperparameter tuning for each optimizer. Owing to such an observation, it is debatable whether the superiority of large-batch optimizers like LARS and LAMB is due to their well-designed concepts or just merely due to dense hyperparameter tuning.
Last but not least, there have been several studies elucidate how the noise can enhance generaliza- tion. Bottou (1991); Neelakantan et al. (2015) showed that noise can play a significant role in training neural networks although they do not focus on large batch training. In addition, Zhou et al. (2019)
theoretically demonstrated that adding noise to the gradient is helpful in circumventing a spurious local optimum, thus allowing for converging to global optimum, and Liu et al. (2021) showed that noisy gradient descent finds flat minimum in a non-convex matrix factorization. Smith et al. (2020) empirically corroborated that noise in stochastic gradients can enhance generalization and Wen et al. (2018) dipped into how injecting curvature noise can be useful in large batch training.
In the light of previous studies regarding gradient noise, one can say "certain noise can regularize the gradient descent procedure well". However, most of the aforementioned studies gradient noise are restricted to the Gaussian-type noise. Therefore, is still questionable whether other types of noise could be beneficial in the aspect of generalization especially for a large batch regime.
In this paper, we focus on the approach to improve generalization for large batch training with the use of the mixup Zhang et al. (2018) technique that has been recently highlighted in the data augmentation context. Toward this, we propose the novel way of adding noise for the optimizer by mixing up the gradients and show that this technique in fact encourages convergence flat minimum.
Our main contributions with some details are as follows:
We propose simple yet effective regularization GradientMx for large batch training. Gradient- Mix performs a linear combination of local gradients computed at each device (or each sample) with arbitrary mixing rates. Our GradienMix is optimizer-agnostic,, hence it can be applied to any optimization algorithms including LARS or LAMB for large batch training.
We mathematically investigate that the optimization with GradientMix could reduce the trace of generalized Beuss-Newton matrix of the objective. Then, we provide the convergence analysis of the optimization with GradientMix in a non-convex regime under the standard assumptions.
We validate the GradientMix popular problems in deep learning communities. Our extensive experiments show that GradientMix could result in better generalization especially for the large batch settings. Specifically, various optimizers with GradientMix consistently achieve state-of- the-art performance on the task of training ResNet-50 on the ImageNet with 32K batch size.
METHOD
In this section, we introduce our main regularization technique, GradientMix. Then, we show that GradientMix encourages reducing the trace of generalized Gauss-Newton (GGN) of the loss function.
2.1 GRADIENTMIX: MIXING GRADIENTS WITH RANDOM MIXING RATE
As a motivating example, we start from the distributed training with multiple GPU machines. Let D be the training dataset with inputs X' fo (I, 3) D) and the corresponding labels y = (y= (x. E D). We study solving the optimization problem under the distributed environment:
1 Σ (1)))))))))
minimize L(() = BERd
where 0 E Rd is the model parameter and E(J. y) denotes the instance wise loss function between the prediction ÿ and the true label y. Given K devices, the vanilla SGD updates the model parameter 0 as
B
Ve, (1 Oc+1 = 02 -/n.iinall
9tinal
K
B
local gradient
2

---

Under review as a conference paper at ICLR 2023
Algorithm A general optimization framework with GradienMMix
1: Input: Training dataset D, stepsize {ıe Ii=1= variance k > O, noise distribution P(T)) satisfying Eq. 12, and optimization algorithm A (e.g. SGD, Adam,...), momentum parameter [0, 1).
2: Initialize: Model parameter 01 E R (Nesterov) momentum v1 E IR..
3: for = .,,,, T do
4: 0 +- 02 p0
1 = 1, ... BB do 5: for
6: Randomly sample b-th datapoint x7) from D
This can be done in parallel way Sample-wise gradient
7:
9t - VoLx :0+4)
8:
end for
9:
Sampling random noise from noise distribution, (T1, "2, (T) P(T))
10:
Tb9t
(b)
GradientMix in sample-wise manner
B
11:
1t+1 - por - 1/t
Σ
V
V6(0,+))
D Nesterov momentum construction
b=1
12:
0.+1 - 01 - A(91.99 92, ... 96 T1)
Parameter update
13: end for
14:
Output: 0T+1
where is the local batch size for each device and k2, represents the b-th datapoint at device & (hence, b E [B). For convenience, we abbreviate the gradient (e[[(x;; 0), Mi) evaluated at the datapoint (z,, 34) as (f()). In Eq. 2. the local gradient computed at device k are averaged across all devices to yield the final gradient. Inspired by several studies (Srivastava al., 2014; Zhu al., 2019; Simsekli et al., 2019; Smith et al., 2020) on the importance of noise generalization, we propose mixing up local gradients at each device k. Toward this, we first generate the random noise Th from some distribution P(T)) at each device k every iteration and perform linear combination of local gradients using TT as a mixing coefficient. summary, the final gradient computed by
K
B
k ~ (TT) with E [[T 1/K
Ak E [K], Ihinal
Σ
T
Ven (f(0))
(3)
B
A=1
b=1
Preserving the unbiasedness of the stochastic gradient, we further assume that each noise satisfies E[TR] = KK in Eq. 3, for example, Gaussian distribution with mean 1/K.
In the one hand, the gradient mixing in Eq. 3 could also be applied in sample-wise manner. Suppose that the gradient is computed with the batch size B, then the sample-wise gradient mixing would be
T P(T)) with B[TT] = 1/B 4b E [B], Gtinal
m (l))(
where To is a sample-wiss noise. The mixing gradients in sample-wise manner Eq. is a generalized extension of the batch-wise version in Eq. We call the procedures in Eq. 4 as GradientMix.
Note that Gradien:Mix is optimizer-agnostic since it does not require any specific update rule; therefore, it can be applicable any popular optimizers such Adam, LARS, LAMB, and Thus, we consider a unified framework for GradientMix. We let A be a randomized optimization algorithm (e.g. SGD) and A(g1 92, ........ 9t)) produce the update vector time where 9, means the gradient computed at time т. Algorithm summarizes an optimization framework with GradientMix.
One of the important properties of GradientMix is that the variance of mixed gradient is always higher than the traditional averaged gradient since the variance of mixed gradient in Eq. 4 is smallest when the noise na in Eq. satisfies ग. = 1/B for b [B]. We will investigate how this noisy property affect the loss landscape and the convergence of optimization in the following section.
2.2 GRADIENTMIX CAN REDUCE THE TRACE OF GENERALIZED GAUSS-NEWTON MATRIX
In this section, we provide more intuition of GradientMix. To this end, we consider an L-layer deep neural network following the notations in (Lee al., 2019) defined by
(1+1) [0w(+)1) (1+1) (I+1) (1(11)
3

---

Under review as a conference paper at ICLR 2023
for x(0) = and [ = 0, 0, L. Here, <(-) is an element- wise (non-linear) activation function such as ReLU, W(O and b() are the weight and bias parameter of appropriate shape at I-th layer, respectively. For simplicity, let E I be the vector of all network parameters and f(z; 0) = (x) E IR* denote the network output (or logit). Under this setup, we first introduce the generalized Gauss Newton (GGN) matrix. Toward this, revisit an optimization problem in Eq. 1:
minimize L(0) = _ n Σ ((((x))) - Σ (((())) ni
0
Thanks (Schraudolph, 2002; Kunstner et al., 2019), the Hessian of the objective (with respect to 0) can be written as
H = V2c(O)
n Σ Jo(ss)
(136 + Remaining term
(5)
G(0)
where (1;) represents the sample-wise Jacobian 0f evaluated single datapoint (2i, y:) and 80
(f()) is the Hessian of E, with respect to the network output. The matrix G(Ø) defined in Eq. known be the generalized Gauss-Newton (GGN) matrix. In case of convex loss l, the GGN is always positive (semi)definite and well-justified approximation of the Hessian (0) Kunstner et al. (2019), so it is popular curvature approximation the non-convex problems.
We study the relationship between the GradientMix and the GGN matrix. The gradient of objective function can be written
vol(((()) - 1 Σ
VoC(O)
A
(6)
where the matrix of all sample-wise gradients A 1 |... Voln E Rdx= and = = (1/n,... 1/n) = 11 E R" is a vector of L's. Under this formulation, GradientMix in effect replaces the deterministic vector with the noisy one 2 as
VoC(O) A
where the vector 3 a random variable satisfying [| 1 for the unbiasedness of the gradient estimator in Eq. 3. We call Eq. 7 mixed gradient.
In the optimization analysis, the convergence is generally measured by how fast the gradient norm [|Vo((0)]l approaches to zero. GradientMix would change this convergence criterion to
|Vo(())] ->
(8)
Plugging the mixed gradient Eq. 7 into Eq. 8, the convergence measure Eq. 8 is nothing but
E 2 A +
Provided that 2 satisfies Cov [3] 22 I for some positive constant 0, owing Hutchinson's
trace estimator Hutchinson (1989), Eq. 9 is equivalent to
Tr A) +
(10)
Rearranging the matrix in the trace operator, we finally arrive the following:
Tr Σ (o (xc10))))(()))) Je(x))) + 0 1=1
(11)
In case of negative log-likelihood loss (as is many cases in deep learning problems), the matrix in the trace in equation Eq. 11 is the GGN matrix G(0) (refer to Proposition 1 in Kunstner et al. (2019)). Consequently, the convergence with GradientMix roughly boils down to reducing the trace of GGN matrix, which could represents the shaprness of minima to some extent in several previous work (Zhu et al., 2019; Simsekli et al., 2019; Lin et al., 2020). Hence, it might be possible that GradientMix could help to find the flatter solution than the obtained via usual averaged gradient.
4

---

Under review as a conference paper at ICLR 2023
Note that this equivalence does not rely the batch size, but expect that GradientMix would be most helpful large batch regime, which is crucially vulnerable to the sharpness of the curvature.
The remaining question is that which noise class reduce the trace of GGN matrix and whether the gradient norm really goes to zero when the optimization is equipped with GradientMix. For the first question, already assume the mild conditions the noise given the batch size B:
1 B In and Cov[TT] = In
E π
(12)
where 1g E _D is just vector of ones. There are several distributions that satisfy Eq. 12, and we adopt two practical distributions: (i) Gaussian and (ii) Rademacher.
Based on our mathematical intuition, we give the empirical evidence on the trace of GGN ma- trix. We train ResNet- 18 on CIFAR-10 dataset with 8K batch size for large batch setting. The table report the trace of GGN matrix of the final models trained with each method. As ex- pected, GradientMix indeed reduces the trace of GGN to a meaningful extent, which means that GradientMix could find flatter solution than the usual average approach with almost no additional computation overhead.
Table I: The trace comparison for trained ResNet- 18 on CIFAR-I0 dataset.
Method
log((T((((O))
SGD Baseline
5.87
SGD GradientMix (Gaussian)
4.21
SGD GradientMi (Rademacher)
4.58
We show the effectiveness of GradientMix with these noise structures in large batch training in the experimental section. Now, we answer the second question in the next section.
3 CONVERGENCE ANALYSIS
In this section, we provide the convergence analysis of optimization with GradientMix in Algorithm 1. Our goal is find e-stationary point for the problem Eq. 1. For this purpose, need the standard assumptions for the analysis in non-convex optimization:
(C-1) (L-smoothness) The loss function f differentiable, L-smooth, and lower-bounded: Vx, y, [[v ff(v) < L||= - yll and f(x") > -co for the optimal solution x=.
(C-2) (Bounded variance) The stochastic gradient 9r time t is unbiased and has the bounded variance: Eg [9c] = f(,)), [ - f..))) < where & represents the randomness
from data distribution.
The conditions (C-1) and (C-2) are standard in the analysis of non-convex optimization Ghadimi & Lan (2013); Zaheer al. (2018); Chen et (2019); Yun al. (2022). Also, GradientMix preserves the unbiasedness of the stochastic gradient, thereby satisfying the condition (C-2). Since Algorithm presents the general framework with GradientMix, we specifically provide the convergence analysis for vanilla SGD with Nesterov momentum, which is the state-of-the-art optimizer used in deep learning communities. The update rule for SGD Nesterov momentum with GradientMix would be
Σ (b) Vb((++)) 0,+1 = 0, 4 U1+1
0++3 = 0, + 400, Vt+1 = fult
(13)
(b) where j denotes the momentum parameter, n. is the sampled noise of b-th datapoint at time satisfying Eq. 12. We further define the quantities n and
Tt B Σ
7
Σ t=
where T' is the total iterations. Note that the quantity roughly measures the total amount of noise injected the gradient time and 7 means the average amount noise over all time. Now, we are ready to state our main theorem.
Theorem (Convergence for SGD + Nesterov momentum with GradientMix). Let 02 denote an iterate uniformly randomly chosen from (0..... ... ,0T44) Under the conditions (C-1) and (C-2) with
5

---

Under review as a conference paper at ICLR 2023
2(1 2(1-p)² the stepsize v 1 11r+1yy the Algorithm I with SGD + Nesterov momentum yields
En [[V SO
LA(HE +1) T(1 11)
2LAo BT( 11)
(14)
where A - f(06) - (00) with optimal point 0".
Remarks. Our Theorem is generalized version of Theorem 4.2 in Lin et al. (2020). Importantly the second term in the upper bound in Eq. 14 is asymptotically dominant under the batch size condition B O
T²T1
In this regime, we can obtain a linear speedup as increasing
(F | for e-stationary point. Note that the second term in the batch size condition involves the average noise 7 attaining the smallest value 1 (so, 2 1) when the sampled probability is = 1/B for all and t, which corresponds to the traditional averaged gradients. In other words, the mixed gradient with random noise x. has always higher value of 7 than the averaged gradient, which in effect allows larger critical batch size but with requiring more iterations. For this reason, GradientMix might show performance degradation in tight budget of epochs, but tuning the variance K in Eq. 12, which affects 7 in result, can play a role to balance the batch size and the number of total iterations.
the batch size with the total iteration order o(
Bea
Since GradientMix do not hurt the unbiasedness of the stochastic gradient but with little higher variance, one can easily show the convergence of each optimizer with GradientMix upon the previous results. We defer the convergence analysis on LARS and LAMB with Gradientix in the Appendix.
4 RELATED WORK
Large batch training. Many researchers have paid attention to large-scale training to speed up training and achieve as high accuracy as small minibatch training. To reduce the generalization gap, Goyal et al. (2017) propose a linear scaling rule to increase the learning rate proportionally to the minibatch size. As exploiting large learning rate can cause neural network to diverge, You et al. (2017) suggest Layer-wise Adaptive Rate Scaling (LARS) to scale updates by multiplying the ratio of the f2-norm of weights to that of gradients for each layer. Furthermore, You et al. (2019) apply this strategy to the Adam optimizer (Kingma & Ba, 2015), which is called LAMB, so as to accelerate training attention-based models, especially BERT, by using large batch sizes. In contrast to (You et al., 2017; 2019), without any layer-wise normalization, Nado et al. (2021) show that standard optimizers are enough for large batch training given sufficient hyperparmeter tuning. Apart from them, in order to narrow the generalization gap, Lin et al. (2020) employ an extragradient technique for smoothing and stabilizing optimization dynamics, and Johnson et al. (2020) propose AdaScale SGD that reliably and automatically adjusts the learning rate depending on the variance of gradients. However, all these methods mentioned above are based on the average of gradients, which is likely to preclude the possibility of exploration during training.
Mixup. Mixup Zhang et al. (2018) is one of the popular methods for enhancing the generalization and robustness of a neural network. Taking advantage of the beta distribution, Zhang et al. (2018) interpolates two inputs and their corresponding labels linearly to generate new data, and Verma et al. (2019) interpolates hidden representations linearly for a smooth decision boundary. Following these works, several extensions Yun et al. (2019); Kim et al. (2020) have been proposed to substitute some area of input for another input's patch, but all the aforementioned studies are discussed only in the context of data augmentation.
Benefit of noise for network training. In recent years, there have been several attempts to figure out connections between noise and generalization. Zhu et al. (2019) showed that noise in stochastic gradients helps escape from sharp minima and converge to flat minima which generalize well. Simsekli et al. (2019) analyze the dynamics of stochastic gradient descent driven by noise and showed that stochatic gradient descent prefers wide minima that performs better than narrow minima on the test set. Smith et al. (2020) verify that the noise in the stochastic gradient can improve generalization. Wu et al. (2020) showed that injecting multiplicative noise to the full-batch gradient generalizes as the usual stochastic gradient evaluated on mini-batch but the noise structure they consider is not applicable to reducing the trace of curvature as we discussed in Section 2.2. Also, the effect of noise on the generalization is guaranteed in theory only in terms of the linear regression.

---

Under review as a conference paper at ICLR 2023
Resket-18 18 on CIFAN-10 dataset
Resket-18 18 on CIFAR-10 datasse
95.8
Performance
95.6
Performance
95.6
95.4
95.4-
95.2
5G3
95.2
LAKS
GradientMix (Goussian)
Gradientii (Gaussian)
95.0
0.2K 1K
2K
5K
0.2K 1K
2K
5K
Batch Size
Batch Size
(a) SGD with
(b) LARS
Reshvt-118 on CIFAR-11 dataset
Resivet-18 CIFAR-10 dataset
Performance
95.4
Performance
95.4
95.2
95.2
95.0
Adam
LAMB
Adam Gradienthix (Saussian)
LAMB GradientMix (Gaussian)
94.8 0.2K 1K
Adam wrnderitiii (Redemacher)
95.0
LAMB GradientMic (Rademacher
2K
5K
0.2K 1K
2K
5K
Batch
Batch Size
(c) Adam
(d) LAMB
Figure 1: Results on training ResNet-18 on the CIFAR-10 10 dataset with GradientMix.
ResNet-18 on CIFAR-100 100 dataset
79.5
79.25
Performance
Performance
79.0
79.00
78.5
78.75
78.0
78.50
LARS
GrodlontMix (Saussion)
GradiantMix (Saussian) (Rademacher)
77.5
Gradientiis (lademacchrr)
78.25
0.2K 1K
0.2K 1K
2K
5K
Batch Size
Batch Size
(a) SGD with momentum
(b) LARS
ResNet:18 18 CIFAR-100 dataset
Restiet 18 on CIFAR-100 dataset
78.5
Performance
78.0
Performance
78.0
77.5
77.5
77.0
76.5
vitiix (Pademacher)
77.0
GradientMi (Rademachor)
0.2K 1K
2K
5K
0.2K 1K
2K
5K
Batch Size
Batch Size
(c) Adam
(d) LAMB
Figure 2: Results training ResNet-18 on the CIFAR-100 dataset with GradientMix.
For more exploration in large batch training, instead of the averaged gradients, our work the first in-depth study applying mixup to the local gradients using suitable distribution and we empirically justify that injecting appropriate noise into the gradient can bridge the generalization gap with large batch well in the next section.
7

---

Under review as a conference paper at ICLR 2023
Transformer on Multi20k dataset
Transformer on Multi3 dataset
70.5
LAMB
70
Acom Gradienteex (Gaussion)
GradientMiz GradientMix [Rademacher]
Performance
Acom Gradientiey (Rademachen
Performance
70.0
LAMB
69
69.5
69.0
68
68.5
0.6K 1K
2K
5K
0.6K 1K
2K
5K
Batch Size
Batch Size
(a) Adam
(b) LAMB
Figure 3: Results on training Transformer on Multi30k dataset with GradientMi.
EXPERIMENTS
We consider two sets of experiments. The first set aims to purely see the effectiveness of Gradient- Mix relatively small-scale problems and the second set is to evaluate GradientMix in a large scale. The details on experimental settings are provided at each section.
5.1 CIFAR CLASSIFICATION
We train ResNet- 18 He et al. (2016) on the CIFAR datasets using two sets of optimizers, which is one of benchmark tasks in large batch training. The first set compares SGD with momentum and LARS, and the second set compares Adam and LAMB since LARS and LAMB are extensions of SGD and Adam respectively. For large batch training. the polynomial LR scheduling with gradual warmup Nado al. (2021) is recommended, so follow the linear LR scaling Goyal et al. (2017) with the base LR Masse = 0.1 for the batch size 200 for SGD and LARS. Similarly, we choose Thase = 10 for the batch size 200 for Adam and LAMB. Throughout this experiment, we consider the total 200 epochs with 20 warmup epochs. In particular, the trust coefficient of LARS is fixed as 10-3, While main goal is how much GradienlMix can improve the performance in a large batch regime, also include the results on the small batch size for better understanding.
Figure I illustrates the results for the CIFAR-I0 dataset. GradientMix shows the consistent im- provement in generalization for all optimizers considered acrosss all batch sizes except the smallest batch size 0.2K. The interesting point in Figure/ 1 is that for all optimizers, the performance with GradientMix tends to increase slightly and then decrease (see 0.2K ~ 2K at z-axis) as the batch size increases while the generalization of all baselines only gets worse as the batch size gets larger. At large butch size such as 2K or 5K, which is our main focus, Gradien(Mix show consistent superiority to the baselines without overlapping error bars. This might be due to the fact that the small batch size has already enough noise, in which GradientMix can interfere with model generalization.
The similar dynamics can be seen in Figure 2 for the CIFAR-100 dataset. Note that all the optimizers with GradientMix achieve great improvement over the baseline for the largest batch size 5K as experiments for the CIFAR-10 dataset. Also, we can see the increase-ther -heeeerraaeee behavior in Figure (see 0.2K 2K at æ-axis) for all optimizers considered in the CIFAR- 10 experiment.
5.2 TRANSFORMER ON MULTI30K
In order to evaluate GradientMix on various problems in deep learning, we consider language modeling task. To this end, we train Transformer base model (Vaswani et al., 2017) on Multi30k dataset (Elliott et al., 2016). Following the experimental settings in Lin et al. (2020), we employ the linear LR scaling and gradual warmup scheme (Goyal et al., 2017) with inverse square root scheduling (Vaswani et al., 2017). The warmup step is set to 4000 for the base batch size 64 and linearly decayed by the batch size. It is known that the adaptive gradient methods work well for attention models, so we compare the Adam and LAMB optimizers for this experiment.
Figure 3 demonstrate the comparison of validation accuracy. Similar to the results in Section 5.1, GradientMix tends to be more effective as the batch size increases for both optimizers Adam and LAMB. Importantly, Gradien(Mix achieve great improvements especially large batch size such as
8

---

Under review as a conference paper at ICLR 2023
Table 2: Top-1 accuracy with GradientMix using 32K batch size for training ResNet-50 on ImageNet.
Top-1 Accuracy (%) for 32K Batch Size
Optimization Algorithm
Fine-tuned hyperparameters
Standard hyperparameters with epochs
with 64 epochs
Noise Type
Nesterov
LARS
Adam
LAMB
No Noise (Baseline)
||
75.93
76. 2004
76.13 + 0.00
76.35
Gaussian (Ours)
76.17 0.06
76.30 4: 0.05
76.33 +: 0.06
76.50 +: 0.08
Rademacher (Ours)
76.25 + 0.44
76.36 + 0.0
76.31 1
76.47 = 0.07
2K and 5K. Combined with the results in Section 5.1, we believe that GradienrMix has the potential to generalize better for other several tasks and optimizers in deep learning.
5.3 RESNET-50 ON IMAGENET WITH 32K BATCH SIZE
Training ResNet-50 on ImageNet has been considered as standard benchmark in large batch opti- mization, so we evaluate our GradientMix on this task using various optimization algorithms. In large batch training, several regularization techniques and hyperparameters should be considered carefully in order to achieve the comparable performance of small batch training. Toward this, we first introduce the experiment settings for each optimizer.
The recent study Nado et al. (2021) empirically shows that the traditional optimizers such as Nesterov can achieve the competitive performance to the optimizer specifically designed for large batch training, such as LARS or LAMB, with the same effort in the hyperparameter tuning. Inspired by this work, we investigate whether GradientMix can further enhance the generalization even under highly fine-tuned hyperparameters. Recommended in Nado et al. (2021), the weight decay should be applied only to the network weight parameters except the bias and batch normalization parameters (You et al., 2017; Goyal et al., 2017). One of the most essential tricks to improve the generalization (He et al., 2019b; Nado et al., 2021) is to tune the initial scale parameter "/o of the final batch normalization layer of each residual block (Goyal al., 2017). We use the polynomial LR scheduling (Nado et al., 2021) with tuning each hyperparameter. We summarize the hyperparameter details in the Appendix.
For Adam and LAMB, unfortunately, there is no previous reference using highly fine-tuned hyper- paramaters for ResNet-50 training. so we follow the experimental pipelines in (You et al., 2019). We (0.9, 0.999) with 10 for both Adam and LAMB optimizers as in (You et al., 2019) and the polynomial LR scheduling with gradual warmup (Nado et al., 2021) employed, The number of warmup epochs twarmup set to 20 among total 90 epochs.
use the momentum parameters (31, (2)
Table 2 illustrates the Top-1 accuracy of each method for 32K batch size. For all baseline optimizers, GradientMix could successfully achieve the state-of-te of-the-art generalization even under highly fine- tuned hyperparameters (see Nesterov and LARS) as well as under standard hyperparameter settings (see Adam and LAMB). As seen Table 2, the performance gain for GradientMix with Nesterov and Adam is slightly larger than that for GradientMix with LARS and LAMB respectively. This might be due to the fact that LARS and LAMB algorithms are already optimized designs to some extent specifically for large batch training, but results prove that GradientMix is indeed optimizer- agnostic. We emphasize that our improvement (without overlapping confidence intervals) is NOT marginal since the performance of baselines is already highly optimized for ResNet-50 architecture. Also, our improvements are bigger than the ones achieved in recent studies related to regularization in large batch regime such Yuan al. (2020); Huo et al. (2021).
6 CONCLUSION
We proposed an optimizer-agnostic and effective regularization for large-batch training, GradientMix, which is the mixup of local gradients with arbitrarily sampled noise. We also show that Gradient- Mix could reduce the sharpness of the loss landscape. Finally, we empirically verify the effectiveness various datasets/models and achieve state-of-the-a the performance mark task, training ResNet-50 on ImageNet dataset. As future work, we plan to investigate how the various type of noise affects the performance for large-batch training both in theory and practice.
of GradientMix
the bench-

---

Under review as a conference paper at ICLR 2023
REFERENCES
Léon Bottou. Stochastic gradient learning in neural networks. Proceedings of Newro-Numes, 91( 12, 1991.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam- type algorithms for non-convex optimization. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Desmond Elliott, Stella Frank, Khalil Sima' an, and Lucia Specia. Multi30K: Multilingual English- German image descriptions. In Proceedings of the Sth Workshop on Vision and Language, pp. 70-74, 2016.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth- methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23344)2341123666 2013.
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: training imagenet in hour. CoRR, abs/ 70s.00066.,777, 2017. URL http: /tt:://arxivvi wwtpsiivoorrrooooooooooo alss 701.0/00000000 02677.
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019a. URL https: htttt::://proceeediii www.eeiiiiinnnnnnnneeeeeeeeeeree neurips wttpapee//019///e/eeeee fflee/ Paper. pdf.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 558-567, 2019b.
Zhouyuan Huo, Bin Gu, and Heng Huang. Large batch optimization for deep learning using new complete layer-wise adaptive rate scaling. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI, pp. 2-9, 2021.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statisties-Simulation and Computation, 18(3):1059-1076 1989.
Tyler B. Johnson, Pulkit Agrawal, Haijie Gu, and Carlos Guestrin. Adascale sgd: A user-friendly algorithm for distributed training. In Proceedings of the 37th International Conference on Machine Learning, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 16Xiv:1609.0448 2016.
Hyun Kim, Wonho Choo, and Hyun Oh Song. Puzzle mix: Exploiting saliency and local statistics for optimal mixup. Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learn- ing Research, pp. 5275-5285 PMLR, 13-18 Jul 2020. URL press/ /v/vvv1111 v1v9/l11vk/.... 1tt119/mm2m00 0b. html.
mlr.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. URL http:// //twxivvv.v//////////ee///aaaav/ org/ abs 1112. 6980.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approx- imation for natural gradient descent. Advances in neural information processing systems, 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32:8572- 2019.
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Extrapolation for large-batch training in deep learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6004-6104. PMLR, 13-18 Jul 2020. URL http: http:://occceeddeeeeeoo wwagss.mr..peerr wwl.s.pres//ee /v119/ 1in20b. html.
Tianyi Liu, Yan Li, Song Wei, Enlu Zhou, and Tuo Zhao. Noisy gradient descent converges to flat minima for nonconvex matrix factorization. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Staristics, volume 130 of Proceedings of Machine Learning Research, pp. 1891-1899 PMLR, 13-15 Apr 2021. URL http:// ppwppeoocediiooomc ngs ml r. press / v130 wuleelllllleee e. html.
Zachary Nado, Justin Gilmer, Christopher J. Shallue, Rohan Anil, and George E. Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. CoRR, abs/2102.06356, 2021. URL https: arxiv .rg/ http//b///// abs//b/b....// 06356.
Arvind Neelakantan, Luke Vilnis, Quoe V Le, Ilya Sutskever, Lukasz Kaiser, Karol Kurach, and James Martens, Adding gradient noise improves learning for very deep networks. arXiv preprint arXiv:1511.06807, 2015.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):1723-1733) 1738, 2002.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Reseurch, pp. 5827-5837. PMLR, 09-15 Jun 2019. URL http:/ http://ppcccennngeeee www. www.pppess//vpe /v2 httsssisssskiiiaee html.
Samuel L Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. arXiv preprini arXiv:2006. 15081, 2020.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-199 1958, 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi, Ioannis Mitliagkas, David Lopez-Paz. and Yoshua Bengio. Manifold mixup: Better representations by interpolating hidden states. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pp. 1238--447..2 PMLR, 09-15 Jun 2019. URL http: /tttp://proee ppoceed ing. ngs ml r. www.r..re///... /
vermal 1 9a. html.
Yeming Wen, Kevin Luk. Maxime Gazeau, Guodong Zhang, Harris Chan, and Jimmy Ba. Exploring curvature noise in large-batch stochastic optimization. 2018.
Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu On the noisy gradient descent that generalizes as sgd. In International Conference on Machine Learning, pp. 10367-10376. PMLR, 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint Xrxiv 008888 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.0092, 2019.
11

---

Under review as a conference paper at ICLR 2023
Qiwei Yuan, Weizhe Hua, Yi Zhou, and Cunxi Yu. Contrastive weight regularization for large minibatch sgd. arXiv preprint arXiv:2011.08966 2020.
Jihun Yun, Aurelie Lozano, and Eunho Yang. Adablock: Sgd with practical block diagonal matrix adaptation for deep learning. In International Conference on Artificial Intelligence and Statistics, pp. 2574-2606. PMLR, 2022.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), October 2019.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances in neural information processing systems, pp. 9793-9803,
2018.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018. URL https: / //oppnreview.... net. forum? id=r1Dép1-RR.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 15383-15393, Curran Asso- ciates, Inc., 2020. URL /tttpp::/prroeeedngg neurips cc/ pctppppe//02//////eeeccp///p file/ Paper..
Mo Zhou, Tianyi Liu, Yan Li, Dachao Lin, Enlu Zhou, and Tuo Zhao. Toward understanding the importance of noise in training neural networks. In International Conference on Machine Learning. pp. 7594-7603 2019.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 1664-7663. PMLR, 09-15 Jun 2019. URL http: http://pooccedddee mlr, www..ppressvve.. zhul 9e. html.
12

---

Under review as a conference paper at ICLR 2023
SUPPLEMENTARY MATERIALS
A HYPERPARAMETER DETAILS FOR NESTEROV AND LARS IN SECTION 5.3
As described in Section 5.3, we follow the same hyperparameter settings in Nado et al. (2021) for Nesterov and LARS optimizers. First, we employ the polynomial learning rate scheduling designed for large-batch training:
Minit +
(Mpeak
M/mit
teamuup
Caameup
(15)
Pircry
Thoal +
(Mpaak
T
tuanuup
We train ResNet-550 on ImageNet dataset with the batch size B 32768 (32K). We summarize the detailed values for each hyperparameter introduced in Section 5.3 in the following table.
Table 3:
The highly fine-tuned hyperparameter configuration for training ResNet-50 on ImageNet dataset for Nesterov and LARS optimizers.
Nesterov
LARS
Total steps (Corresponding to 64 epochs)
2512
2512
Ewarmup
706
706
Pwarmup
2.0
1.0
Pdecay
2.0
2.0
BatchNorm momentum
0.9
0.9
BatchNorm
10-5
10-5 5
1/paak
7.05
29.0
Thinll
10
10
Optimizer momentum
0.97603
0.929
Weight decay
5.8 x 10-5
10
Label smoothing
0.15
0.1
Initial scale parameter of the final BN layer of each residual block
0.4138
0.0
Trust coefficient (only for LARS)
N/A
10-3
B PROOF OF THEOREM
Our analysis based the convergence proof of SGD Nesterov momentum Lin al. (2020). We first derive the convergence batch-wise GradientMix and present analysis sample-wise GradientMix, Recall that the iterate of SGD + Nesterov momentum with GradientMix is
K
B
0 0. + pos, U141
(0,
0.t1 + 1441
b=1
where πj. means that the sampled probability of k-th device at time and f is the momentum parameter. For simplicity, we define the as
B
9i
fug (0z)
B
k=
b=
Then, the momentum update would become
Ue+1 μud 199.+
We also define the quantities and Č
G =
(k)
13

---

Under review as a conference paper at ICLR 2023
The main changes in our analysis how to bound the term
Σ
110
11
Toward this, we have
==
K
I3
g1+ 3 11]
(R)
V fRn (0
(0.44)
(16)
Var
(R)
fhe (0 (0c+)))
B
(K)
fky ++
(17)
B
B
1
bm1
Am
Le B
f + ||
+
(18)
Henec, we obtain
B
1
92+ 1 17] S 5 oo 1 T
B
Σ
vf ++ ||
(19)
T
Following the proofs in Lin et al. (2020), we define an auxiliary sequence yr as
if 0
y1
I-p 0144
1 01 0.-4 + -
if 2 1
(20)
Lemma (Lemma A.I Lin et al. (2020)). The sequence {Nc} in equation satisfiex
"7
y44 y 1- 11 9c+
Lemma (Lemma A.2 Lin et al. (2020)). For sequence ( for 2 O, the following holds
T-1
12
T
0++3
944
Now, we derive the convergence bound. From smoothness condition, we have
[ff(+++1) - f(ye)] s B (VI())). ut+1 -+) - - |lye+1 - y/
= E
7 -
[5(0++))] -1 (V((y)) v.
2 22+ 1 11?
Ti1| "
1-
by Lemma 1. Further we have
p (v V.(00 v((0++))) [/00)))))
-1- (vI())) - v)(O.4))
n 3/12
3L
[I 4J3/2
Vf(v)) Vf
75 1[00+3]
V
2
2(1 n)3
by the inequality (I. v) S [1 [lvll3). Then, we get
[[[()+1]] [(v:)]
<
E 1 -
71
V V (5((0-4) +
20
1-
V [(()) v[(0
2(1 4)3
L
v5(0 (.. +)
2(1 12)
17² L
Me+ $
71 L4x
I|
(1 [-)=
1/ - l|e
gct 11]
B
1- 1 2(1 11)3
2/13
2(1 11)²
By telescoping over t 0 ~~ T 1, we obtain
T 1
1
[f(()))))
-((r))
E(s(c)) - -[v0)]
T
T
1
Ln²
T-1
vf(++ $11 +B
(1 L)
11a
T
p
+
2(1 n)3
T
2/13
p1) T
lct 1
1-
1/² L
T
+ E
9
2(1 11)2
T
14

---

Under review as a conference paper at ICLR 2023
Here, apply derivations equation 18 then have
[((())
=(())
7 Ln
n²L nn²
Lun 21)
T
r²L
Luni
50²
I- 2(1 14
2(1
2(1
T ==
2 nz)
2(1 11)3 B
By rearranging all the items, we have
T-1 T +0
nL Z0z
A
1 2(-a)²
T
2(1 μn² B
In order get the final results, need the following lemmas from Lin et al. (2020)
Lemma (Lemma A.4 in Lin et al. (2020)). For every non uon-negative sequence {re h1z0 and any parameters d 2 0, c 2 0, and T 2 0, there exists constant S 1/d such that for any constant stepsizes ne it holds
VT
r++1 + cnn)
S
dA (T +
+
(21)
Then, using VT =
4L 201 ュリートン p)² B in Lemma 3, we arrive at the following results by case study
T
where
AR < or
L similar to Lin et al. (2020).
[T
L
Σ v ( ;) [1] =o
LA( +1) T(1 p)
2L BT p))
+
(22)
Now, if we replace the = Bo and B = 1 for some batch size Bo. we obtain the convergence of sample-wise GradientMix. Going further, for batch size Bo. the noise quantity Çe in equation 16 should be at least of order Ba, in other words
G
(m 10)2
2 1/Bo = Q(1/Bb)
since Çe is smallest when all the noise satisfy ( order remove dependency on the batch size Bo. we newly define the following quantity
1/Bo (by Cauchy-Sehward inequality). For clear analysis,
Bo
Tt
Bo
7
7
T
Lmm
Then, the final convergence bound for Nesterov with sample-wise GradienrMix would be
E
T-1 |
Σ
LAL + +1 -(1-p)
2LAo²g
(23)
BoT(1 11)
for the batch size Bo (in effect, replaces ? with 7).
B. CONVERGENCE ANALYSIS OF LARS WITH GradientMix
From this section, analyze the convergence of LARS/LAMB with gradient mixup based upon the previous analysis You et al. (2019). Here, we make conditions specific LARS LAMB as follows.
We assume that the loss function f(-) is Ly-smooth with respect the paramter of the 1-th layer g()) for forl= E [h], which means
V a) - V.f s)ll s [ - ") Var y E Rd
where represents a random datapoint from data distribution. We use I. (I1,.... I.).. We also assume the bounded variance of stochastic gradient
Vif of
15

---

Under review as a conference paper at ICLR 2023
for all layer i [k). Furthermore, in You et al. (2019), assume that B|ll s)}, [V(())) S ह,, va R for j E [d]. Under this condition, we use the following notation simplicity as
(dr, oh)
(ō1, ,.,)
Lastly, we assume the bounded gradient condition [vf(x,)) s)l, I G for all j E [d] and 184.
Recall that the iterate of LARS You et al. (2019) is
8
(110
II)
gi l1g.
(24)
for all layer [L] and o represents the parameter of I-th layer time t. The scaling function @(-) satisfies au S ( I Ou and practice we use just a clipping function d(z) min(max(2 " 7u). Here, stochastic gradient Jr with gradient mixup is computed as
91
( ((0r..))
(2.5
( where Sh denotes the b-th sample in mini- mini-batch of size B and {=( Jb-1 is unit simplex time Note that this more general setting of gradient mixup and datapoint-wise gradient mixup can surely cover the case minihatch-wise gradient mixup. Here, we define the following quantitics Çe and similarly as Nesterov momentum
1 T
Č= T
Σ VC.
The main changes in our analysis is how to bound the term
B
lla
of
(26)
where (0.) means the gradient function computed at time with respect to parameter 0,00 I-th layer. The above term can be easily bound under the bounded varaince condition as
E|1
( (.f((0,, v.f(
(27)
F
( (V. 86) v.f(o.))
(28)
6=1
B
Var
(b) (V.f(0e: m6) v.S(0.)) -..((.)) + mt
Vif(le, 86) V, v(f())))
(29)
?
(30)
Henec, we have
R - v.f((e)]] s VGa
(31)
We replace the inequality (6) in the Appendix A in You et al. (2019) with our equation 31. Then, the following holds
e) S f(0.) vaa
[vf(0.)]|
+ 2na.. /Glo
||
2
By telescoping over t = ~~ T,, we have
T ΣΣ 1=1
B[f(Or+1)] f(0)) you
2n
2
Now, we arrive at the final convergence rate as
1
ΣΣΣ BIlV. S Tno
A
20,
[|
T
201
where A =((0)) f(0") with the optimal point 0"
16

---

Under review as a conference paper at ICLR 2023
B.2 CONVERGENCE ANALYSIS OF LAMB WITH GradientMix
this section, we provide the convergence analysis LAMB, which has the following update
= Σ
Ht (b)
(fr,,
my 81 m 1 (1 B1)9. = 8 (1 32)91
me = my (1 81)
2 = 12 / (1 82)
vec +
$(100
0 = 0
(r2 + )
1Js ||
As in You et al. (2019), we deal with the case of B1 = O and À = O, but B2 > O. Then, the update vector would be
9t
Similar to the convergence of LARS, the major changes in our analysis is in how to bound the following term from inequality (7) in the Appendix You et al. (2019)
T S
32 (0(100"))) x [..((.)], x gi.) G
(
:110{ x v.](0.]] x
r
(sigg((V * sigu(r()))
||r:(0
22 @((10_1)) x (v.f(.))]: .
4
(1.(( P(sign([V # sign sign( (
Here, the mixed gradient has a different bound with the averaged gradient as
L(1 - 82)
d
B[Z+] < /
((()))) + nan V
0
G2d
j=1
Therefore, we have
L(1 82) G2d
17 ?²
B[S(0:+1) S f(0.) - V
2
Telescoping over t = ~ T yicids
L(1 82) G
T
+o. + 1²a2T 7
R[f(O++1] f(1)) nou
Finally, we have the following convergence guarantce LAMB optimizer with gradient mixup as
L(1 82) 1
Σ
A
naa
[T
Gdd T
Thos
201
17


Reviews:
[{"review_id": "saAm840cvZ", "reviewer": "Reviewer_LFbe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tries to use large-batch training to accelerate neural network training. However, traditional large-batch training algorithms usually suffer from generalization issues. To solve this problem, this paper proposes an algorithm GradientMix, which tries to mixup of local gradients computed at multiple devices. The experimental results illustrate that GradientMix can further improve the accuracy of LARS/LAMB for large-batch training.", "summary_of_the_review": "NA", "strength_and_weaknesses": "Strength:\nThis paper is very easy to follow. The authors propose GeadientMix to solve the generalization issue in large-batch training. \nThe proposed method is easy to reproduce. As shown in Algorithm 1, the main difference is about sampling random noise and weighted gradient. \n\n\nWeakness:\n\nI still think the proposed method is not novel. The proposed method (mixup of local gradients computed at multiple devices) is similar to the weighted gradient. \n\nThe main purpose of Large-batch training is to accelerate the training process. Therefore, I think you also need to report the training time for these methods. Since GradientMix introduces more computation cost.\n\nAs you mentioned, generalization is an important issue for large-batch training. So I think you should also report the generalization gap (training accuracy - test accuracy) in your experiments to verify GradientMix can narrow the gap. \n\nYou need to provide more experiments about sensitivity analysis. For example, the results with different number of GPUs and different noise. That is because the number of GPUs and the selection of noise (with different mean and variance) is an important hyper-parameters for your proposed method.  \n\nMore results about recent regularization methods for large-batch training should be reported. The main reason is that LARS  is proposed in 2017 and there are many algorithms that try to further improve the performance of large-batch training.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "UM4u6nDuthu", "reviewer": "Reviewer_yVcz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed to incorporate mix-up to the gradient synchronization. Traditionally, the gradients are aggregated by a simple averaging with uniform weight to ensure unbiasedness and low variance. However, this paper argues that random weights following certain distribution pushes the optimizer to minimize the trace of the the generalized Gauss-Newton (GGN) matrix, helping to find a flatter solution than the one obtained via usual average. Since the random weights are reduced to uniform weights in expectation, the convergence of the proposed GradientMix is still guaranteed. The proposed method can be used in conjunction with any optimizers and achieves consistently better performance in training ResNet on CIFAR-10, 100, ImageNet datasets for different batch sizes. The authors also validate the performance on training a small Transformer model on the Multi30k dataset.", "summary_of_the_review": "This paper introduces a simple and interesting approach for improving the generalization performance. The intuition is backed up by a sketch on how GradientMix tends to minimize the trace of GGN matrix. However, the empirical improvements are limited and slight. Given that ResNet training has been well studies and the proposed achieves very marginal improvement, it is more interesting to see if the proposed method could bring any benefit on training a multi-billion parameter GPT model, which is known be more challenging for large batch training.", "strength_and_weaknesses": "Strength:\n1. The paper introduces an interesting method for improving the accuracy of large batch training. The intuition of reducing trace of GGN matrix makes sense to me and flat minima generalizing better is backed up by several existing works.\n\nWeakness:\n1. The empirical improvement looks marginal. For the majority of the experiments, the improvement over the baselines are within 0.4%. It is possible that the gap could be closed with some hyper-parameter tuning.\n2. Large batch training for ResNet model has been well studied. (Sun, P., et al., 2019) shows that we can train ResNet-50 on ImageNet in 7.3 mins with a large batch size of 64K, which is already larger than the largest batch size considered in the paper. I think it is more interesting if GradientMix could improve the large batch training of multi-billion parameter GPT model, which is known to suffer from accuracy degradation of large batch and hasn't been studied in the literature.\n\nSun, P., Feng, W., Han, R., Yan, S., & Wen, Y. (2019). Optimizing network performance for distributed dnn training on gpu clusters: Imagenet/alexnet training in 1.5 minutes. arXiv preprint arXiv:1902.06855", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "UEqyZRh1Wqw", "reviewer": "Reviewer_HmFS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes GradientMix, a regularization method where minibatch gradients are computed as a linear combination of per-example gradients with randomized rather than uniform example weights. The primary application it explores is improving generalization at large batch sizes.\n\nThe paper demonstrates GradientMix in a few standard deep learning settings (ConvNets on CIFAR and ImageNet and a sequence-to-sequence Transformer) and shows better generalization at large batch sizes. The authors also provide a theoretical explanation based on GradientMix reducing the trace of the generalized Gauss-Newton matrix, a common approximation of the sharpness of the minimum.", "summary_of_the_review": "I really appreciate straightforward papers presenting simple but novel improvements to deep learning methods, and that's exactly what I see here. But I would also like to see both more theory (especially discussion of gradient noise scales, as in McCandlish 2018) and more experimental results (at least the train and test losses to show generalization gap).", "strength_and_weaknesses": "Strengths:\n- The main goal of the work (better optimization at large batch sizes) is very important to the community.\n- The core method is simple and seems to work well, representing meaningful progress on this goal.\n- CIFAR-10, CIFAR-100, and ImageNet are all standard image classification benchmarks, and while CIFAR datasets and models are small by contemporary standards they're still good testbeds for rapid iteration.\n- The paper is fairly convincing: I think this method is worth trying in my own work.\n- The paper relates and compares GradientMix to similar ideas in data augmentation.\n\nWeaknesses:\n- The experimental settings for Transformers are quite limited. Multi30k contains only 30k examples, so the entire dataset is around the same size as a single batch in contemporary large-batch Transformer training.\n- More generally, contemporary use cases driving the desire for large-batch training are frequently in the regime where each example is visited once (LLMs, recommenders), but none of the experiments cover this case.\n- Experimental results are reported only for test accuracy (I think?) but showing both train and test curves would be valuable for understanding the differential impact of GradientMix on optimization and generalization. \n- The ImageNet experiment uses a single batch size, but it would be ideal to demonstrate that the critical batch size (the point at which linear convergence scaling wrt batch size ends) is higher for GradientMix than for the baseline; this would require experimental results at multiple batch sizes.\n- The paper observes that the advantage of GradientMix is low or reversed at small batch sizes, and makes a suggestion as to the cause without really digging into the question.\n- This and other aspects of the picture here could be made clearer by use of the gradient noise scale concept/quantity introduced in McCandlish 2018.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6AW0cGhuVM", "reviewer": "Reviewer_SMM4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper applies the mix-up approach proposed in [1] to the parallel gradients in the data-parallel large-batch training regime and shows that it can provide a minor improvement on the results on validation dataset son two tasks (image classification and machine translation) and three small NN architectures (ResNet-15, ResNet-50, and a transformer model).", "summary_of_the_review": "The paper applied the mixed-up technique to the distributed gradients in data-parallel large-batch training. According to the experiments presented in the paper, I feel the improvement is rather minor, and am not sure how it can be applied in real large-batching training settings.", "strength_and_weaknesses": "Strength:\n- The technical part of the writing is clear and easy to follow\n- The idea of using mixed-up gradients to reduce GCN seems sound to me\n\nWeakness:\n- The introduction is not well-written. It seems to just repeat the related work section\n- The effectiveness of the proposed methods is questionable. From the experiments presented in the paper I can tell that the improvement of the proposed technique provides very minor advantages; and not sure how it can actually help distributed training in practice.\n\nSee detailed comments in the next section", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject