Title: Accelerating Large Batch Training via Gradient Signal to Noise Ratio (GSNR)
Venue: NeurIPS 2023
Accepted: no
Keywords: large batch

optimizer
Abstract: As models for nature language processing (NLP), computer vision (CV) and recommendation systems (RS) require surging computation, a large number of GPUs/TPUs are paralleled with a large batch (LB) to improve training throughput. Training such LB tasks often converges to sharp minimum and downgrades final precision. Adversarial learning (ConAdv) and LANS method scales ImageNet and BERT pretraining up to 96k batch size. In this work, we develop the variance reduced gradient descent technique (VRGD) based on the gradient signal to noise ratio (GSNR) and apply it onto popular optimizers such as SGD/Adam/LARS/LAMB. We carry out a theoretical analysis of VR-SGD's convergence rate to explain its fast training dynamics, and a generalization analysis to demonstrate its smaller generalization gap on LB training. Comprehensive experiments demonstrate that VRGD can remarkably accelerate training ($1.7\sim 4 \times$), narrow the generalization gap and improve final accuracy. We push the batch size limit of BERT pretraining up to 128k/64k and DLRM to 512k without noticeable accuracy loss. We improve ImageNet Top-1 accuracy at 96k by $0.52pp$ than LARS and significantly reduce generalization gap by $68.3$%.

Abstract

As models for nature language processing (NLP), computer vision (CV) and recommendation systems (RS) require surging computation, a large number of GPUs/TPUs are paralleled with a large batch (LB) to improve training throughput. Training such LB tasks often converges to sharp minimum and downgrades final precision. Adversarial learning (ConAdv) and LANS method scales ImageNet and BERT pretraining up to 96k batch size. In this work, we develop the variance reduced gradient descent technique (VRGD) based on the gradient signal to noise ratio (GSNR) and apply it onto popular optimizers such as SGD/Adam/LARS/LAMB. We carry out a theoretical analysis of VR-SGD's convergence rate to explain its fast training dynamics, and a generalization analysis to demonstrate its smaller generalization gap on LB training. Comprehensive experiments demonstrate that VRGD can remarkably accelerate training ( 1 . 7 ∼ 4 × ), narrow the generalization gap and improve final accuracy. We push the batch size limit of BERT pretraining up to 128k/64k and DLRM to 512k without noticeable accuracy loss. We improve ImageNet Top-1 accuracy at 96k by 0 . 52 pp than LARS and significantly reduce generalization gap by 68 . 3% .

1 Introduction

Recent machine learning models have grown wider and deeper in their architectures (e.g., GPT-3 [Floridi and Chiriatti, 2020], M6 [Lin et al. , 2021], Switch Transformer [Fedus et al. , 2021]). Training complex models may consume more training data to converge, which needs a surge in computing capacity and efficiency. However, hardware improvement can not keep pace with the expansion of model calculations [Bommasani et al. , 2021].

Several techniques to speed up training are proposed. The aggregation and scattering of gradients among massive workers requires an efficient synchronization algorithm. Since the communication bandwidth between GPUs/TPUs is much higher than CPU-GPU (e.g., NVLink, Foley and Danskin [2017]), several efficient synchronization strategies such as Ring-All-Reduce [Gibiansky, 2017] and software toolkits like Horovod [Sergeev and Del Balso, 2018] are proposed to replace the traditional PS-Worker framework [Li et al. , 2014b,a]. In addition, training with LB can notably improve throughput [You et al. , 2017b; Hoffer et al. , 2019]. You et al. [2020] successfully train BERT using 1024 TPUs and a LB (64k) within 76 minutes. It demonstrates the efficiency of GPUs/TPUs in large scale parallel tasks. Small batch (SB) is not able to fully utilize those powerful GPUs/TPUs.

However, Keskar et al. [2017] theoretically analyze the LB training and finds that it can be easily 32 trapped into sharp local minimum, leading to strong generalization gap. Hoffer et al. [2017] indicate 33 that the generalization gap can be attributed to the fewer update steps in LB training compared with 34

SB when using identical epochs. Dai and Zhu [2018] theoretically demonstrate that training with 35

more steps or expanding the learning rate to batch size ratio helps to converge to a flatter local 36

minimum. Although these issues can be partly resolved by layer-wise adaptive rate scaling (LARS, 37 You et al. [2017a]) and layer-wise adaptive large batch (LAMB, You et al. [2020]), the batch size 38 limit still exists. 39

40

41

42

To push the batch size limit and reduce generalization gap, we propose the element-wise adaptive techniques called variance reduced gradient descent technique (VRGD) based on GSNR of parameters. Our contributions are listed below:

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

Wecarry out theoretical derivations of convergence rate and generalization analysis to explain why VRGD can accelerate LB training and achieve dramatically smaller generalization gap.

We perform comprehensive LB experiments and find that VRGD can remarkably accelerate training ( 1 . 7 ∼ 4 × ), narrow the generalization gap and improve final precision than previous SOTA (e.g., LAMB, LARS).

VR-LAMB pushes the batch size limit of BERT pretraining up to 128k/64k without any accuracy loss, while LAMB stops scaling at 64k/32k. VR-LARS improves the ImageNet Top-1 accuracy to 74 . 82% at 96k, 0.52pp higher than LARS. The generalization gap of ImageNet trained with VR-LARS is dramatically reduced by 68.3% comparing with LARS at 96k. VR-SGD pushes the batch size limit of DLRM from 64k to 512k without noticeable accuracy loss.

2 Related Work

2.1 Large Batch Training

Several techniques are proposed to improve the optimization and generalization ability in LB training. Goyal et al. [2017] propose a linear scaling rule on learning rate (LR) to achieve the same accuracy as SB and push the batch size limit of ImageNet to 8k. EXTRAP-SGD uses the extra-gradient to stabilize the optimization trajectory and smooth training [Lin et al. , 2020]. SWAP quickly trains the model with LB in the first stage and refines it by averaging the weights of multiple SB models in the second stage [Gupta et al. , 2020]. Batch Augmentation replicates multiple instances with the same batch size to improve generalization [Hoffer et al. , 2019]. The batch size of the experiments in EXTRAP-SGD/SWAP/Batch-Augmentation are less than 8k and are not compared in our experiments.

DecentLaM removes the growing momentum-incurred bias observed in DmSGD and pushes ImageNet to 32k [Yuan et al. , 2021]. Layer-wise LRs adjustment optimizers such as LARS [You et al. , 2017a], complete layer-wise adaptive rate scaling (CLARS, Huo et al. [2021]), LAMB [You et al. , 2020] successfully improve the batch size up to 64k both for ImageNet and BERT pretraining without accuracy loss. Recently, the concurrent adversarial learning (ConAdv) method pushes the batch size limit of ImageNet training up to 96k [Liu et al. , 2021]. LANS replaces the layer-wise LR adjustment in LAMB with block-wise style [Zheng et al. , 2020] and also pushes BERT training up to 96k. Adasum adds those gradients after scaling with proper scalars and even pushes the batch size limit of BERT up to 128k/32k [Maleki et al. , 2021].

2.2 Gradient Variance and GSNR

Unlike gradient mean, which is widely used in optimizers, gradient variance and its successor GSNR 74 are less used. But gradient variance is frequently discussed in generalization gap. Johnson and Zhang 75 [2013a] propose the stochastic variance reduced gradient (SVRG) with the explicit gradient variance 76 reduction method. Other variants of SVRG like SRVR-NPG, SVRPG and Control Variate methods 77 are also proposed to reduce the gradient variance during training [Liu et al. , 2020b; Wang et al. , 78 2013; Papini et al. , 2018; Miller et al. , 2017]. Rainforth et al. [2018] use GSNR to analyze the 79 variational bounds in variational auto-encoder (V AE). McCandlish et al. [2018] use GSNR to predict 80 the useful upper bound of batch size. Smith et al. [2018]; Devarakonda et al. [2017] adaptively 81 increase the batch size during training to achieve acceleration without accuracy loss. Liu et al. [2020a] 82 theoretically derive a quantitative relationship between GSNR and generalization gap and prove that 83 larger GSNR leads to better generalization performance. Therefore, gradient variance and GSNR are 84 potentially useful to train deep neural networks. 85

3 Preliminaries 86

3.1 GSNR 87

Given a data distribution Z = X × Y , a model 88 ˆ y = f ( x, θ ) parameterized by θ and the loss 89 function L . The parameters' gradient w.r.t. sam90 ple ( x i , y i ) can be written as (Refer to all "nota91 tions" in the Appendix.C): 92

Then j -th parameter' ( θ j ) gradient computed using ( x i , y i ) is g i ( θ j ) . Here we use i to index the 93 data samples and j to index the parameters of θ . We denote the sample-wise gradient mean as 94 g ˜ ( θ ) = E ( x,y ) ∼Z ( g ( x, y, θ )) and variance of g i ( θ ) as ρ 2 ( θ ) = Var ( x,y ) ∼Z ( g ( x, y, θ )) . The GSNR 95 for each model parameter θ j is defined as: 96

Intuitively, GSNR measures the consistency of the gradient direction of each parameter across a batch of data samples. The gradient space of the parameters tends to converge in the same direction when the GSNR is large, but diverge if the GSNR is small (Figure.1).

3.2 GSNR and Generalization Gap

Consider a training set D = { ( x 1 , y 1 ) , ..., ( x n , y n ) } ∼ Z ( n ) , where n samples come from Z , and a test set of dataset size ( n ′ ) from Z ′ ( n ′ ) denoted by D ′ = { ( x ′ 1 , y ′ 1 ) , ..., ( x ′ n ′ , y ′ n ′ ) } ∼ Z ′ ( n ′ ) . The empirical training and test losses can be denoted as:

respectively. Then the empirical generalization gap is given by L [ D ′ ] -L [ D ] . Both the training loss L [ D ] and the test loss L [ D ′ ] would decrease after one training step and can be denoted as ∆ L [ D ] and ∆ L [ D ′ ] respectively. The ratio between the expectations of ∆ L [ D ] and ∆ L [ D ′ ] for one training step can be denoted as:

Assumption 1 (Non-overfitting limit approximation of Liu et al. [2020a]) . The parameters' gradient over the training set and test set i.e., g D ( θ ) and g D ′ ( θ ) obey the same distribution.

Based on Assumption 1 and using a small learning rate λ → 0 , Liu et al. [2020a] derive the relationship between the one-step generalization ratio (eq.4) and GSNR:

where ∆ L j [ D ] is the training loss reduction caused by updating θ j . A more detailed mathematical derivation can be found in Liu et al. [2020a]. This relationship (eq.5) demonstrates that GSNR ( r j ) plays a crucial role in determining the generalization performance of the model. Updating the model parameters with smaller GSNR leads to generalization gap growth. Also note that we have R ( Z , n ) → 1 when n →∞ , which means that training with a larger dataset helps generalization.

4 Proposed Algorithms

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

In this section, we propose VRGD with their general updating rules (taking VR-SGD as an example 118 in Algorithm.1). The SGD is shown in Appendix.D for comparison. 119

120

121

122

123

124

125

126

127

128

129

130

Algorithm 1: V R -SGD

Input: require device number k ≥ 2 Input: B = GlobalBatchSize/k Input: γ = 0 . 1 1 while θ t not converged do for device d = 1 to k do ˜ g d ( θ t ) ← 1 B ∑ B i =1 ∇ θ L ( y i , f ( x i , θ t -1 )) (Get gradient on each GPU/TPU) ˜ g 2 d ( θ t ) ← ˜ g d ( θ t ) ⊗ ˜ g d ( θ t ) (Element-wise multiply, so as square terms below) ˜ g ( θ t ) ← 1 k ∑ k d =1 ˜ g d ( θ t ) (Reduce gradient over all devices) σ 2 t ← 1 k ∑ k d =1 ˜ g 2 d ( θ t ) -˜ g 2 ( θ t ) (Compute gradient variance) r ( θ t ) ← ˜ g 2 ( θ t ) σ 2 t (Compute GSNR) for layer l = 0 to h do r ( θ ( l ) t ) ← r ( θ ( l ) t ) 1 J ∑ J j =1 r ( θ ( l ) t,j ) (Normalize GSNR so that r ( θ ( l ) t ) = 1 ) r ( θ ( l ) t ) ← { γ, if r ( θ ( l ) t ) < γ 1 , if r ( θ ( l ) t ) > 1 (Confine the max/min ratio within 1 γ ) θ t ← θ t -1 -λ · r ( θ t ) · ˜ g ( θ t ) (Update weights)

4.1 VR-SGD's Updating Rules

Consider the simple updating rule for SGD as follows:

where λ is the learning rate. Previous section demonstrates that updating the weights with larger GSNR confines the model's generalization gap growth during training. Therefore, GSNR can be used in the optimizer for better generalization. In the mathematical derivation of GSNR's role on the generalization gap, all sample-wise gradients for the entire dataset are used to compute the gradient variance, which is less efficient. However, in the LB training training, where each batch is large enough to accurately estimate the gradient variance, we replace the entire dataset with a LB and the sample-wise with device-wise gradient computation. Gradients on each GPU/TPU device can be synchronized using Ring-AllReduce, thus perfectly avoiding the inefficiency of gradient variance computation. The simplified gradient variance computation is as follows:

where k devices are used, each of which computes 1 /k part of the gradient ˜ g d ( θ t ) , the same as what 131 data parallel does. The GSNR can then be easily calculated based on eq.2 ( ρ 2 ( θ j ) is replaced by σ 2 j ). 132 The mean values of GSNR are removed at each layer before applying gradient to the parameters. 133 This normalization of GSNR ensures that the global averaged GSNR remains at 1 . 0 : 134

where l th layer contains J parameters. We constrain the max/min of GSNR within 1 /γ so that 135 those neurons with very small GSNR remain active: 136

where γ is a hyper-parameter used here. For simplicity, we don't tune γ but set it to 0 . 1 in all of our 137 experiments by default. Finally, we element-wisely adapt λ according to GSNR of each parameter 138 and get the updating rule for VR-SGD: 139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

Figure.1 shows the mechanism of VRGD. As for a good estimation of gradient mean (left panel), optimizer should be confident to move along the direction of gradient mean or even further. However, when gradients on the devices are scatteredly distributed (right panel), updating weights with gradient mean may bring noises and slow down convergence, which should be avoided.

Differences compared with existing LB methods:

The linear scaling rule uses the same large LR for all parameters, which tends to diverge when some gradients are too large; LARS/LAMB/LANS use large LRs for some layers but layer-wisely or block-wisely limit LRs when || θ t || is compatible with its updating quantity, i.e., || θ t || ∼ || λ · ˜ g ( θ t ) || ; VRGD that we propose here element-wisely limit the updating quantity for those parameters without confident gradient estimation (Fig.1b, large gradient variance or small GSNR).

GSNR and its relationship with generalization gap is discussed in Liu et al. [2020a], but further work to embed such GSNR into the optimizers is missing. In our work, we apply GSNR in the SGD/LARS/LAMB and demonstrate that GSNR helps the model maintain a small generalization gap in LB training based on the derivations of the generalization gap and ImageNet experiments.

VRGD does not need extra-gradient used in EXTRAP-SGD or the two-stage training like SWAP. Sub gradients used in Batch Augmentation have different transforms each while VRGD uses the same transforms. Adasum adaptively sums two gradients scaled by a constant while VRGD still uses the mean gradient.

4.2 VR-Adam, VR-LAMB and other VRGD optimizers

̸

GSNR can be easily applied on any optimizer using the general updating rules shown above. Here we discuss those popular optimizers frequently used in the research community, e.g., SGD, Adam, LARS and LAMB. As for VR-Adam, GSNR is calculated directly based on ˜ g ( θ t ) and then used to adapt the gradient mean before gradients' momentum estimation. Similar with the gradients' momentum, we apply the momentum mechanism on GSNR ( ˆ p t ) for faster convergence. If we adapt the final update term, i.e. θ t ← θ t -1 -λ · r ( θ t ) · ˆ m t / ( √ ˆ v t + ε ) , the 1 st and 2 nd order momentum estimation ( m t and v t ) for the next training step would be biased (meaning that the update term cannot be inferred merely on ˆ m t and ˆ v t since r ( θ t ) = 1 ).

VR-LAMB is similar to VR-Adam, except that VR-LAMB layer-wisely adapt the LRs for stable convergence when using very large LRs. VR-Adam and VR-LAMB are shown in Appendix.D. VR-LARS and VR-Momentum, which are based on LARS and Momentum, are similar to VR-SGD that it uses GSNR to adapt the gradient means before applying them to the model weights (algorithms omitted).

5 Theoritical Analysis

5.1 Convergence Analysis

Assumption 2 (bounded gradient) . ||∇ L ( θ ) || ≤ G

Assumption 3 ( l -smooth) . ∃ l > 0 satisfies ||∇ L ( x ) -∇ L ( y ) || ≤ l || x -y ||

We mathematically derive the convergence rate of VR-SGD under nonconvex settings and assume the training process satisfies Assumption.2 and Assumption.3, which are widely used in convergence analysis [Shamir and Zhang, 2013; Ghadimi and Lan, 2013; Allen-Zhu and Hazan, 2016; Allen-Zhu et al. , 2019; You et al. , 2020]. Table.1 of Appendix compares the assumptions of ours and those popular optimizers. It shows that our assumptions are weaker than LARS/LAMB/DecentLaM and similar with SGD. Detailed derivations can be found in Appendix.A. Then we have Theorem.1.

where r l and r u are the lower and upper bound of GSNR. 185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

Convergence rates discussion : 1) The convergence rate O ( 1 √ ˆ T ) of VR-SGD is the same as SGD [Johnson and Zhang, 2013b]; 2) VR-SGD's bound depends on the lower ( r l ) and upper bound ( r u ) of GSNR. Larger batch size brings smaller gradient variance (eq.43 of Appendix.B) and larger GSNR (both bigger r l and r u ), then may result in a tighter bound with quicker convergence ( verified by experiments shown in Figure.2 ).

5.2 Generalization Gap

This section derives the generalization gap of SGD and VR-SGD during SB and LB scenarios. Detailed derivations can be found in Appendix.B. Citing eq.14 of Liu et al. [2020a] below, i.e., when training satisfies Assumption.1 and λ → 0 , after one training step the expectation of empirical generalization gap at t th step is:

where we use σ 2 t,j and r t,j to denote σ 2 ( θ t,j ) and r ( θ t,j ) for simplicity. Next, we assume that the batch size of LB is k times than that of SB. λ 0 ( λ ) represents the learning rate of SB (LB). The accumulated generalization gap after training T steps for SB using SGD and T/k steps for LB can be derived as follows:

If we assume " σ t,j is t -independent", eq.13 are simplified as E ( GAP SB,SGD ) ≈ λ 0 T ∑ j σ 2 j and E ( GAP LB,SGD ) ≈ λT k 2 ∑ j σ 2 j respectively. Taking λ = k 2 λ 0 , E ( GAP LB,SGD ) will have the same accumulated generalization gap as SB. This is known as the linear/square scaling rules. However, the assumption that " σ t,j is t -independent" is unrealistic. Similarly, the accumulated generalization gap of VR-SGD in LB training can be written as:

The generalization gap of SGD and VR-SGD in LB training:

When training converges ( g t,j → 0 ), we have g 2 t,j < σ 2 t,j because r t,j = g 2 t,j /σ 2 t,j → 0 (verified experimentally by Figure.4 of Liu et al. [2020a]). Therefore, we have λ k ∑ T/k t =1 ∑ j g 2 t,j < λ k ∑ T/k t =1 ∑ j σ 2 t,j , i.e., E ( GAP LB,V R -SGD ) < E ( GAP LB,SGD ) . This inequality demonstrates that VR-SGD has a much smaller generalization gap than SGD in LB training ( verified by our ImageNet experiments shown in Table.3 ).

6 Experiments

In this section, we show comprehensive experiments on commonly used LB benchmarks such as BERT Pretraining [Devlin et al. , 2019], ImageNet-2012 [Russakovsky et al. , 2015] and DLRM [Naumov and Mudigere, 2020]. We mainly adopt the square root rules to scale LRs. We set the hyper-parameters of VRGD as γ = 0 . 1 and k to the minimum GPU devices that can hold the LB without out of memory for resource efficiency (but satisfy k ≥ 8 ) in all experiments. Similar with other optimizers, VRGD can generate a generally good training curve using default sets. The 1 st and 2 nd order decay rates are set to β 1 = β 3 = 0 . 9 , β 2 = 0 . 999 by default. Experiments are performed with TensorFlow on 96 DGX-A100 nodes (768-GPUs).

6.1 BERT Pretraining

BERT pretraining is a common NLP task needs speeding up with LB training. For a fair comparison, we use the same settings as LAMB [You et al. , 2020] except optimizer and learning rate: (1) BERT large pretrains using Wikipedia and BooksCorpus and then finetunes on SQuAD(v1.1) to evaluate its

precision with F1 score; (2) A two-phase training strategy is used. First 90% steps use a sequence 224 length of 128 (phase-1) and last 10% use a sequence length of 512 (phase-2). Mixed-Batch Training 225 is used when batch size is set to 64k/32k, 96k/32k and 128k/64k. 226

∗ means the F1 scores are cited from their work. Using median of repeated experiments is the same as Nado et al. [2021].

We use NVIDIA's best practise 1 to carry out VR-LAMB experiments and tune nothing of the downstream SQuAD(v1.1) tasks (same as LAMB). Detailed hyper-parameters are listed in Appendix.D. Results shown in Table.1 indicate that:

VR-LAMB outperforms LAMB (widely used in BERT LB pretraining) in all batch sizes from 16k to 64k/32k. F1 score is improved up to 91 . 49 at 64k/32k, 0.91pp higher than LAMB.

VR-LAMB also outperforms Adam (with standard bias correction and LR discontinuity removal) and LANS by an improvement of 0.84pp at 64k and 0.63pp at 96k/32k respectively.

VR-LAMB pushes the batch size limit up to 128k/64k using just 4301 steps and maintains a F1 score of 90 . 85 . Although Adasum achieves a F1 score of 90 . 50 at 128k/32k, but it needs 6137 steps to converge ( 30% extra steps than VR-LAMB). VR-LAMB achieves 50% less steps than LAMB at 64k/32k and even 0.45pp higher of F1 score than baseline.

6.2 ImageNet with ResNet50

ImageNet training with ResNet50 v1.5 [He et al. , 2016a] is a standard CV benchmark for LB training. We use the default sets of official best practise of Google Tensorflow 2 with linear LR warm-up, label smoothing and cosine LR decay (to 0 ). It is the same setup as LARS [Liu et al. , 2021]. We merely adjust the optimizers and learning rate for a fair comparison. We find some successful LB applications using Momentum, LAMB and LARS, but not for Adam, AdaGrad or AdamW optimizers [Goyal et al. , 2017; You et al. , 2020; Liu et al. , 2021]. LARS based on Momentum is more fitful on CV tasks. Therefore, we merely apply VR-LARS on ImageNet. Detailed hyper-parameters are listed in the appendix.D.

The results shown in Table.2 indicate that:

VR-LARS outperforms Momentum, DecentLaM, LAMB and LARS (previous SOTA) in all batch sizes (from 0.03pp to 0.56pp ). The improvements are higher for larger batch size.

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

VR-LARS achieves 75 . 86% accuracy at 64k batch size, 0.56pp higher than LARS. When batch size reaches up to 96k , VR-LARS maintains 74 . 82% accuracy, close to the MLPerfv0.5 standard ( 74 . 9% ).

Generalization Gap : Table.3 demonstrates that VR-LARS can dramatically narrow the generalization gap in LB training. The generalization gap is only 1.46 for VR-LARS at 96k ( 68.3% smaller than LARS), even smaller than ConAdv+AA ( 2 . 2 ; Liu et al. [2021]). Note that VR-LARS can be used together with ConAdv+AA and other techniques for further improvement.

6.3 DLRMTraining

Criteo Terabyte click logs dataset ( 4 billion records) trained with DLRM is a standard CTR prediction benchmark newly added in MLPerf-v0.7. DLRM is used following NVIDIA's best practise 1 . For a fair comparison, we merely modify LRs and optimizers (hyper-parameters are listed in Appendix.D). Settings of Linear LR warm up, polynomial decay and training with 1 epoch are used by their default set up. Results in Table.4 indicates that:

VR-SGD outperforms SGD in all batch size settings. Similar with experiments shown above, the improvement of VR-SGD w.r.t SGD increases along with larger batch sizes (from 0.12pp to 2.26pp ).

VR-SGD pushes the batch size limit up to 512k and maintains a high AUC of 0.8013 , close to the baseline of 0.8014. Note that Google's submission of MLPerf v0.7 merely uses a maximum batch size of 64k [Kumar et al. , 2021].

7 Ablation Studies

7.1 Orthogonal Experiments

In this section, we demonstrate that GSNR is important in optimization and VRGD can be applicable 274 to most popular optimizers using CIFAR10. During CIFAR10 training with ResNet56 [He et al. , 275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

2016a,b], we use the default sets of the official best practice for Google Tensorflow 2 and mainly add square-root LR scaling rules to perform the 216 composite experiments shown in Figure.2. Additional linear LR warm-up, label smoothing and cosine LR decay (to 0 ) techniques are used to stabilize LB training experiments shown in Table.5, the same as ImageNet training. Detailed hyper-parameters are listed in Appendix.D. As for the test accuracy curves, Figure.2 shows the averaged composite test accuracy curve of all 216 experiments for the LR-batch size pairs. Training with VR-Momentum/VR-Adam/VR-LAMB converge much faster ( 1 . 7 ∼ 4 × ). As for the final precision, Table.5 demonstrate that VR-Momentum/VR-Adam/VR-LAMB/VR-LARS dramatically outperform Momentum/Adam/LAMB/LARS when batch size is larger than 4k, which demonstrates that VRGD is applicable to most popular optimizers in LB training. The improvements of VRGD comparing with their base optimizers grow with the increase of batch size. VRGD optimizers remains convergent when batch size reaches 8k.

7.2 GSNR's Behaviour

To understand GSNR's behaviour in VRGD optimizers, we perform the linear regression experiments. The true weights are set to W i = i, i ∈ [1 , 10] and the corresponding parameters w i are initialized to zero. Given randomly generated inputs X , we have the true labels as Y = WX and the MSE loss as L = || Y -wX || 2 . Finally, optimize w with 100 steps.

Training about 50 (half) steps, VR-SGD is able to converge to the test loss where SGD requires 100 steps (Figure.1a of Appendix.D). The weights of VR-SGD (dashed lines of Figure.1b of Appendix.D) converge faster to their ground truth. We find that w 5 , w 6 converge firstly, then w 3 , w 8 and finally w 1 , w 10 . Consistently, the GSNR of w 5 , w 6 arise firstly (updating w 5 , w 6 with larger LRs), then w 3 , w 8 while the GSNR of w 5 , w 6 decrease slowly (no need to update the converged weights using large LRs). Finally after step 60 , the GSNR of w 1 , w 10 begin to arise. Intuitively, GSNR helps element-wisely fine-tune the LRs for different weights.

7.3 Hyper-parameters Sensitivity

There are two main hyper-parameters in VRGD, i.e., normalization strength factor γ and the equivalent GPU device number k . We take use of linear regression trained with VR-SGD using batchsize = 2048 shown above to examine the hyperparameter sensitivity.

Figure.3 shows that the optimal γ is around (0 . 04 , 0 . 2) for linear regression. Test loss would be larger if γ → 1 , which means VR-SGD is reduced to SGD . It again demonstrates that GSNR is valuable to improve final precision. On the other hand, the optimal k is around [32 , 256] . This means that each gradient mean calculated using [8 , 64] samples on each GPU/TPU device, and gradient variance calculated using [32 , 256] values of the gradient mean will return a good evaluation of GSNR. In fact, we do not use the optimal hyper-parameters. Instead, above experiments use γ = 0 . 1 and set k to the minimum GPU devices that can hold the LB without out of memory (but satisfy k ≥ 8 , refer all of the hyper-parameters in Appendix.D). Fine-tuning γ and k may further improve the results.

8 Summary

In this paper, we propose the VRGD for large batch training using GSNR. We carry out theoretical 320 derivations of convergence rate and generalization analysis to explain why VRGD can accelerate 321 large batch training and reduce generalization gap. Comprehensive experiments on BERT-pretraining, 322 ImageNet and DLRM verify that VRGD can push the batch size limit than previous SOTA optimizers 323 in LB training and perform better. Codes will be released when published. 324

1.6

k=2

Best test loss searched in LRs

1.4

k=4

1.2

k=16

k=32

1.0

optimalarea

0.8

k=256

0.6

0.4

0.2

0.0

0.01

.02

O

0.6

S200

Y

10

y= 0.0001

Y=0.006

y=0.08

Best test loss searched in LRs

8-

y= 0.0002

000°0=A

Y= 0.0006

一

y=0.01

V=0.0125

Y=0.008

y=0.2

y=0.4

y=0.1

Y = 0.0008

y=0.02

y=0.6

y= 0.001

y=0.04

y=0.8

6-

y = 0.002

y= 0.004

y=0.06

y=1.0

optimalarea

4

0

16

32

6

128

256

Equivalent GPU Device Number (k)

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

References

Zeyuan Allen-Zhu and Elad Hazan. Variance reduction for faster non-convex optimization. In International conference on machine learning , pages 699-707. PMLR, 2016.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258 , 2021.

Xiaowu Dai and Yuhua Zhu. Towards theoretical understanding of large batch training in stochastic gradient descent. CoRR , abs/1812.00542, 2018.

Aditya Devarakonda, Maxim Naumov, and Michael Garland. Adabatch: Adaptive batch sizes for training deep neural networks. arXiv preprint arXiv:1712.02029 , 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT , pages 4171-4186. Association for Computational Linguistics, 2019.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961 , 2021.

Luciano Floridi and Massimo Chiriatti. Gpt-3: Its nature, scope, limits, and consequences. Minds and Machines , 30(4):681-694, 2020.

Denis Foley and John Danskin. Ultra-performance pascal gpu and nvlink interconnect. IEEE Micro , 37(2):7-17, 2017.

Saeed Ghadimi and Guanghui Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming. SIAM J. Optim. , 23(4):2341-2368, 2013.

Andrew Gibiansky. Bringing hpc techniques to deep learning. Baidu Research, Tech. Rep. , 2017.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Large-batch training that generalizes well. In 8th International Conference on Learning Representations, ICLR . OpenReview.net, 2020.

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 , pages 770-778, 2016.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision , pages 630-645. Springer, 2016.

Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems , pages 1731-1741, 2017.

Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: better training with larger batches. arXiv preprint arXiv:1901.09335 , 2019.

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 , pages 7883-7890. AAAI Press, 2021.

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems , 26, 2013.

Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Christopher J. C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger, editors, Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of a meeting held December 5-8, 2013, Lake Tahoe, Nevada, United States , pages 315-323, 2013.

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. In 5th International Conference on Learning Representations, ICLR . OpenReview.net, 2017.

Sameer Kumar, Yu Emma Wang, Cliff Young, James Bradbury, Naveen Kumar, Dehao Chen, and Andy Swing. Exploring the limits of concurrency in ML training on google TPUS. In Alex Smola, Alex Dimakis, and Ion Stoica, editors, Proceedings of Machine Learning and Systems 2021, MLSys 2021, virtual, April 5-9, 2021 . mlsys.org, 2021.

Mu Li, David G Andersen, Jun Woo Park, Alexander J Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J Shekita, and Bor-Yiing Su. Scaling distributed machine learning with the parameter server. In 11th USENIX Symposium on Operating Systems Design and Implementation , pages 583-598, 2014.

Mu Li, David G Andersen, Alexander J Smola, and Kai Yu. Communication efficient distributed machine learning with the parameter server. Advances in Neural Information Processing Systems , 27, 2014.

Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Extrapolation for large-batch training in deep learning. In International Conference on Machine Learning , pages 6094-6104. PMLR, 2020.

Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, et al. M6: A chinese multimodal pretrainer. arXiv preprint arXiv:2103.00823 , 2021.

Jinlong Liu, Guoqing Jiang, Yunzhi Bai, Ting Chen, and Huayan Wang. Understanding why neural networks generalize well through GSNR of parameters. In 8th International Conference on Learning Representations, ICLR . OpenReview.net, 2020.

Yanli Liu, Kaiqing Zhang, Tamer Basar, and Wotao Yin. An improved analysis of (variance-reduced) policy gradient and natural policy gradient methods. Advances in Neural Information Processing Systems , 33:7624-7636, 2020.

Yong Liu, Xiangning Chen, Minhao Cheng, Cho-Jui Hsieh, and Yang You. Concurrent adversarial learning for large-batch training. arXiv preprint arXiv:2106.00221 , 2021.

Saeed Maleki, Madan Musuvathi, Todd Mytkowicz, Olli Saarikivi, Tianju Xu, Vadim Eksarevskiy, Jaliya Ekanayake, and Emad Barsoum. Scaling distributed training with adaptive summation. In Alex Smola, Alex Dimakis, and Ion Stoica, editors, Proceedings of Machine Learning and Systems 2021, MLSys 2021, virtual, April 5-9, 2021 . mlsys.org, 2021.

Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. An empirical model of large-batch training. CoRR , abs/1812.06162, 2018.

Andrew Miller, Nick Foti, Alexander D'Amour, and Ryan P Adams. Reducing reparameterization gradient variance. Advances in Neural Information Processing Systems , 30, 2017.

Zachary Nado, Justin M Gilmer, Christopher J Shallue, Rohan Anil, and George E Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. arXiv preprint arXiv:2102.06356 , 2021.

Maxim Naumov and Dheevatsa Mudigere. Dlrm: An advanced, open source deep learning recommendation model, 2020.

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

Matteo Papini, Damiano Binaghi, Giuseppe Canonaco, Matteo Pirotta, and Marcello Restelli. Stochastic variance-reduced policy gradient. In International conference on machine learning , pages 4026-4035. PMLR, 2018.

Tom Rainforth, Adam R. Kosiorek, Tuan Anh Le, Chris J. Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML , volume 80 of Proceedings of Machine Learning Research , pages 4274-4282. PMLR, 2018.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV) , 115(3):211-252, 2015.

Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensorflow. arXiv preprint arXiv:1802.05799 , 2018.

Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning , pages 71-79. PMLR, 2013.

Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V. Le. Don't decay the learning rate, increase the batch size. In 6th International Conference on Learning Representations, ICLR . OpenReview.net, 2018.

Chong Wang, Xi Chen, Alexander J Smola, and Eric P Xing. Variance reduction for stochastic gradient optimization. Advances in Neural Information Processing Systems , 26, 2013.

Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv:1708.03888 , 2017.

Yang You, Igor Gitman, and Boris Ginsburg. Scaling sgd batch size to 32k for imagenet training. arXiv preprint arXiv:1708.03888 , 6(12):6, 2017.

Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. In 8th International Conference on Learning Representations, ICLR . OpenReview.net, 2020.

Kun Yuan, Yiming Chen, Xinmeng Huang, Yingya Zhang, Pan Pan, Yinghui Xu, and Wotao Yin. Decentlam: Decentralized momentum SGD for large-batch deep training. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021 , pages 3009-3019. IEEE, 2021.

Shuai Zheng, Haibin Lin, Sheng Zha, and Mu Li. Accelerated large batch optimization of BERT pretraining in 54 minutes. CoRR , abs/2006.13484, 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nthis paper proposes a gradient descent technique for learning deep neural nets with large batch sizes. the authors focus on the setting where the model size is small to medium (10M to 300M parameters) and the dataset size is also small to medium (up to 1M images in vision or ~ 3B words in NLP) but with as large batch size as possible. the goal of this works to accelerating the learning of such models(e.g., ResNet on Imagenet, BERT on the original BERT-used corpus) by using less training steps and less training time. The authors were able to show the proposed method can achieve good results (without big drop in eval metrics) while using larger batch sizes than prior arts. And when the proposed method is compared with prior arts at the same batch size, the proposed method seems to outperform prior arts as well.\n\nSTRENGTHS:\nThe papers covers standard benchmarks like ResNet on ImageNet and BERT pretraining. and therefore can be fairly compared against many prior arts on the same tasks\n\nWEAKNESSES:\nThis paper claims \"training acceleration\" as a key contribution. But throughout the paper, the comparison on speed up is based on number of steps or number of epochs. It is unclear what are the speed advantage in terms of wall-clock time by using the propose technique. I also checked the supp. pdf. In other works (such as LARS and LANs, which are cited by this work), the authors usually report actual wall clock time speedups as they increase the batch size and the compute infra. It was disappointing to not see any mention of that given the authors are using 768 GPUs (therefore I expect very interesting scaling behaviors)\n\nQUESTIONS:\n* how does the proposed method help with the large scale trainings such as CLIP models or DiNOv2 models, which the datasets are much bigger than traditional settings like ImageNet? \n\n* what would be an ideal scenario for using the proposed technique in a computer vision task? from Table 2, there is a clear trade-off between accuracy vs batch size.\n\n* in Table 1 and Table 2, how much time does each training take (under different batch sizes)?\n\nLIMITATIONS:\nthe proposed technique is only validated on small-medium size model on small-medium size dataset. it is not applicable (at least no evidence provided) to large scale model training (either large in dataset size or large in model size.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors strive to propose a heuristic training strategy, called variance reduced gradient descent technique (VRGD), based on the gradient signal to noise ratio, i.e. the ratio between the norm and the variance of gradient. Compared to vanilla training, VRGD scales the learning rate with the gradient signal to noise ratio during each iteration. Then, the authors prove that the proposed method can converge within finite training steps, and claim that it will give better generalization gap compared to the vanilla training. Finally, the authors show the effectiveness of their methods via Bert-training and ImageNet training with up to 96k batch size.\n\nSTRENGTHS:\n***Strengths***\n\n1. The paper is clearly written and easy to follow.\n\n2. With the development of hardwares, training with large batch would gradually become a basic requirement for training large-scale models on gigantic data, such as GPT. In my opinion, the authors are focusing on a very topic worthy to probe.I believe that this paper may have a potential significance not only in academia but also in the industry. However, some modification may be required currently.\n\n3. The authors show some interesting and valuable experiment results, but not comprehensive enough.\n\nWEAKNESSES:\n***Weakness***\n\n\n1. (Major) The authors repeatedly mention that large batch training can lead to sharp minima in Abstract and Introduction, which seems to suggest that the proposed method can avoid such problem. However, I have not found discussions or observations regarding the proposed method can solve this issue. So can the proposed escape these bad minima?\n\n    Considering that many recent works that focuses on guiding training to converge to flat minima i.e. SAM family/gradient norm regularization, I am quite curious what would be like to adopt the proposed method in these algorithms when given large batch training. And from the results, I find that using the proposed method alone would not result better performance than these flat-minima-based methods. I have listed some typical works below.\n\n    [1] Foret, Pierre, et al. \"Sharpness-aware minimization for efficiently improving generalization.\" ICLR2020.\n\n    [2] Kwon, Jungmin, et al. \"Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.\" ICML2021.\n\n    [3] Zhuang, Juntang, et al. \"Surrogate gap minimization improves sharpness-aware training.\" ICLR2021.\n\n    [4] Zhao, Yang, et al. \"Penalizing gradient norm for efficiently improving generalization in deep learning.\" ICML 2022.\n\n2. (Major) Continuing with the previous comment, I find the proposed method is rather heuristic, lacking a clear motivation. I could not find clear rationale and sufficient analysis that the proposed method can benefit training. Essentially, the proposed method is simply to scale the learning rate adaptively based on a specific parameter (Eq. 10). For me, I am not quite convinced that such a learning rate scaling policy will lead to reasonable performance gain. Meanwhile, the author argue that the proposed method can lead to smaller generalization gap. However, the core that supports this claim is based on an empirical observation, which makes the mathematical proof not quite rigorous and the claim much weaker and unhelpful. \n\n3. (Major) I would like to discuss the convergence of the proposed method. Firstly, to my understanding, the convergence analysis focuses on analyzing to what extent can training converge on the given training samples, not testing set. So, it is not quite appropriate to use the convergence curve on the testing set to demonstrate the conclusion regarding the convergence analysis. i.e. Figure 2. Secondly, from Figure 2, the authors state that the proposed method converge 1.7~4 times faster than the conventional optimizers. But, I could not observe such a big gap between them from Figure 2, so could the authors explain how to measure the convergence here. Thirdly, a tighter bound in convergence would not give any information regarding the testing performance. A looser bound and slower convergence rate can give better testing performance in many cases, for example SAM. The author can refer to the paper below. \n\n    [5] Andriushchenko, Maksym, and Nicolas Flammarion. \"Towards understanding sharpness-aware minimization.\" ICML 2022.\n\n    Note that I am not saying a faster convergence is harmful. In my opinion, the core meaning of this convergence section is to prove that the proposed method can converge in finite time. And it is not surprising that The proposed method shares the same convergence rate as that in SGD, i.e. O(1/sqrt{T}) given that this method is to scale the learning rate compared to SGD. In the current version, it appears that the convergence section seems to heavily imply that the proposed method outperforms SGD without a promise of testing performance, which I disagree with.\n\n4. (Minor) Line 32. \"However, Keskar et al. [2017] theoretically analyze the LB training and finds that it can be easily trapped into sharp local minimum, leading to strong generalization gap\". Actually, the cited paper (Keskar et al. 2017) has not provided theoretical analysis.\n\n5. (Minor) Line 95. To my understanding, the variance of random vectors is a matrix, i.e. covariance matrix. Why is a scalar here?\n\n6. (Minor) It is highly encouraged to show the results of training vision transformers with the proposed methods.\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nI have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the learning method in deep learning. Still, it is highly encouraged to add corresponding discussions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper examines the improvement of training throughput in a large batch parallel training setting. By employing the gradient signal to noise ratio (GSNR) as a measurement of the generalization gap during training, the authors introduce a variance-reduced gradient descent method designed for large batch training scenarios. The authors provide theoretical analysis to substantiate that the proposed variance-reduced gradient descent (VRSGD) method exhibits superior generalization compared to stochastic gradient descent (SGD) and potentially achieves faster convergence. Additionally, experimental evaluations on BERT pretraining, ResNet training, and DLRM training are presented to demonstrate the superiority of the proposed method over alternative approaches for large batch training. Furthermore, orthogonal experiments and analyses pertaining to the behavior of GSNR and sensitivity to hyperparameters are conducted to further support the superiority of the proposed method.\n\nSTRENGTHS:\n1. The research topic of large batch training is of considerable interest.\n\n2. The proposed method is both simple and effective for large-batch training.\n\n3. The theoretical analyses of the convergence rate and generalization are persuasive.\n\n4. The proposed method consistently demonstrates improvements over other baseline approaches.\n\n5. The paper is well-written and easy to understand.\n\nWEAKNESSES:\n1. The contribution is limited. It seems that the authors just introduce GSNR into the large batch training, but do not highlight why GSNR is important to large batch training.\n\n2. Although the theoretical analysis and experimental results presented in the study are persuasive, the rationale behind the necessity of the GSNR method for large-batch training remains unclear. It is essential to provide a more comprehensive explanation of why GSNR is specifically relevant and advantageous in the context of large batch training. While GSNR can potentially enhance generalization in various settings, a more explicit justification is required to elucidate its particular suitability and effectiveness for large-batch training scenarios.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper proposes a new method for large batch training. It is based on the insight that the gradient-to-signal-noise-ratio for each parameter should be reflected in its learning rate, and hence modifies gradient descent to reduce the variance of the gradients. The paper then shows convergence rates of VRSGD, which are the same as SGD asymptotically, and states that VR-SGD is particularly suited to large batch training where the GSNR will be high. The paper also shows that the proposed method has a smaller generalization gap than SGD in the large batch setting. Experiments are then performed to support these claims on a variety of tasks and architectures.

STRENGTHS:
1. The method is based on a solid insight.
2. The empirical results indicate decent gains over baselines.

WEAKNESSES:
1. The assumptions of smoothness and bounded gradients seem a bit too strong.

QUESTIONS:
Can the assumptions for the theoretical analysis be relaxed?

LIMITATIONS:
See weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: SUMMARY:
This paper focus on using large-batch training to accelerate the training of neural network. Specifically, the authors try to use variance reduced gradient descent technique to scale up the batch size and therefore to accelerate the training. The experimental results illustrate that the proposed method can scale up to larger batch size and further accelerate the training of ResNet, BERT and DLRM.

STRENGTHS:
Strength:
1. This paper focuses on an important problem, accelerate neural network training. especially for large-batch training. 
2. The proposed method is very easy to follow.
3. The authors provide some results to verify the performance of proposed method.

WEAKNESSES:
Weakness:
1. The authors can provide more visualization and analysis about the proposed method and why the proposed method can further scale the batch size. For example, the proposed method can help the model converge to s flat region?
2. The authors should provide more results about the wall time of each experiment and verify whether the proposed can save the training time.
3. I'm not sure whether you should compare your method with more baselines since LARS and LAMB is not current SOTA.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11