SGD batch saturation for training wide neural networks

Critical Batchsize: According to [1], the number of training steps is halved by each doubling of the batch size and that diminishing returns exist beyond a critical batch size. Is the definition of "the critical batchsize" used in the paper the same as in [1]? I cannot understand why "the critical batchsizes" are $b^* = \frac{\gamma}{\frac{1}{n} \Vert X \Vert^2}$ (linear least squares case; Page 4) and $b^* = \frac{4 \gamma + \frac{50 \gamma L^2 r^2}{\alpha}}{\Vert F(w_0) \Vert^2 + L^2 r^2}$ (nonlinear least squares case; Page 5). Please give mathematically the definition of "the critical batchsize" and show that "the critical batchsizes" are $b^*$ (Pages 4 and 5).

[1] Christopher J Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv:1811.03600, 2018.

Assumptions and numerical results: This paper gives sufficient conditions to satisfy Assumptions 1 and 2. I understand the authors try to give such sufficient conditions. However, we need to set $r$ (see (4.2)), $m$ (see (4.2)), $\beta$ (see (4.3)), and $\eta$ before implementing SGD using a constant step-size $\eta = \frac{1}{\beta}$. There is no how to set such parameters in Section 5 (Experimental results). Could the authors show evidences to emphasize that the parameters in Section 5 are based on Theorems 4.2-4.4?

Numerical results: I think that the neural networks in Section 5 are small. I would like to check the performance of SGD for training wide-resnets on CIFAR-100 and ImageNet. I also would like to suggest that

• the authors plot the number of iterations to achieve high test/training accuracy versus the batch size and show the existence of critical batchsizes.

Moreover, I would like to suggest that

• the authors compare not only SGD using a constant step-size (e.g., $\eta$ is chosen on the basis of a grid search of $\eta \in (10^{-1}, 10^{-2}, 10^{-3} )$) but also the existing optimizers (e.g., momentum, Adam, and variants of Adam) with SGD using $\eta = \frac{1}{\beta}$. If SGD using $\eta = \frac{1}{\beta}$ performs better than other optimizers, SGD using $\eta = \frac{1}{\beta}$ is a good way to train deep neural networks.

Optimization problem: This paper considers Problem (2.1) that is to minimize MSE: $\mathcal{L}(w) = \frac{1}{2n} \sum_{i=1}^n (f_i(w) - y_i)^2$. This setting seems to be limited to consider "a broader class of nonconvex problem" (see Abstract). Can we consider a more general nonconvex optimization, such as: Minimize $\mathcal{L}(w) = \frac{1}{n} \sum_{i=1}^n f_i (w)$, where $f_i$ is nonconvex and differentiable?

Iteration complexity: The authors use "iteration complexity." Please define mathematically the iteration complexity used in the paper. Is the definition of the iteration complexity used in the paper the same as the definition in [2]?

[2] Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan J. Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1):165– 214, 2023.

Convergence speed: The abstract says that the paper gives a convergence speed of SGD. Please give a rate of convergence of SGD under Assumptions 1 and 2.

Mathematical Preliminaries: The mathematical preliminaries are insufficient. For example,

• Theorem 1.1: $\lambda_{\min}$, $c$, and $\Vert \cdot \Vert_{\mathrm{op}}$ are not defined.
• (2.2): $\mathbb{E}_S$ is not defined.
• Line -8; Page 2: We cannot use $m$, since $m$ has been used to define the width $m$ (see Theorem 1.1).
• Line +2; Page 3: $S$ has been used to define the batch size $|S|$. The same notation is not good.
• Theorem 4.2 (4.1): $f$ is not defined.
• ...

Proofs: Please give the proofs of Theorems 4.1--4.4 in the main body (I think that the proofs of Lemma 3.1, Theorem 3.2, Theorem 3.3 can move to the Appendix section. Then, the authors have sufficient space of the manuscript). I do not believe that Theorems 4.1--4.4 are true, since there is no proof of theorems.

Typos:

• Theorem 1.1: where ... "and $m$ is the width of the network" is deleted.
• Figure 1: in each. --> in each
• Line -1; Page 7: opt
• ...

My recommendation: The paper includes theoretically interesting results. However, I have some concerns (the critical batchsize, the assumptions, the implementation of SGD, and numerical results). Hence, at this point, my score is 1: strong reject. However, I strongly believe that the authors could address my concerns. Hence, I hope to raise my score.

Please see my comments above.

• On the convergence bound: The result provided is a high probability result that needs the existence of a solution in a ball whose radius scales with the failure probability. Are you choosing $m$ after fixing the failure probabilities to ensure the existence of such a solution ? Also I believe $\delta_1 \in (0, 1/5)$ and not of $1/3$ ?
• Why not simply reduce the moment bound assumption (3.1) to simply $\gamma$ bounded $x_i$ ? I have difficulty seeing any other scenario where this assumption is naturally shown to hold without assuming bounded x_i's.
• On the experiments: You choose to fix the same step size for all the batch sizes to illustrate the existence of a critical batch size. Could the authors explain why they make this choice especially since the empirical papers they cite studying critical batch sizes carefully note in the abstract that "how batch size affects model quality can largely be explained by differences in metaparameter tuning"(Shallue et al. (2018)) (and more importantly since Theorem 2.1 has a stepsize depending on $\beta$).
• Can you give more details on what is plotted in figure 2 column 2. Is the predicted rate obtained from computing the first terms of (4.3) on the datasets ? In addition, the range of values plotted is very small (from 3e-4 to 4e-4 for example on MNIST) so is the effect of batch size on the rate actually very small in practice ?
• I believe $w$ should be $w_0$ whenever you give expressions for $b*$ and $\beta$, so if my understanding is correct it is possible to simply compute the optimal batch size afterwhich there is diminished returns, could you tell us what they are for MNIST for example ?
• It doesn't appear that NanoGPT is converging linearly nor does it look like it's approaching zero training error, so I am unsure it fits in your setting of wide neural networks.

Please see above.