The Impact of the Mini-batch Size on the Dynamics of SGD: Variance and Beyond
Authors: Xin Qian
Diego Klabjan
Venue:  (2021)

Abstract: We study mini-batch stochastic gradient descent (SGD) dynamics under linear regression and deep linear networks by focusing on the variance of the gradients only given the initial weights and mini-batch size, which is the first study of this nature. In the linear regression case, we show that in each iteration the norm of the gradient is a decreasing function of the mini-batch size $b$ and thus the variance of the stochastic gradient estimator is a decreasing function of $b$. For deep neural networks with $L_2$ loss we show that the variance of the gradient is a polynomial in $1/b$. The results theoretically back the important intuition that smaller batch sizes yield larger variance of the stochastic gradients and lower loss function values which is a common believe among the researchers. The proof techniques exhibit a relationship between stochastic gradient estimators and initial weights, which is useful for further research on the dynamics of SGD. We empirically provide insights to our results on various datasets and commonly used deep network structures. We further discuss possible extensions of the approaches we build in studying the generalization ability of the deep learning models.

Under review as a conference paper at ICLR 2021
THE IMPACT OF THE MINI-BATCH SIZE ON THE DYNAMICS OF SGD: VARIANCE AND BEYOND
Anonymous authors Paper under double-hlind review
ABSTRACT
We study mini-batch stochastic gradient descent (SGD) dynamics under linear re- gression and deep linear networks by focusing on the variance of the gradients only given the initial weights and mini-batch size, which is the first study of this nature. In the linear regression case, we show that each iteration the norm of the gradient is decreasing function of the mini-batch size b and thus the vari- ance of the stochastic gradient estimator is a decreasing function of b. For deep neural networks with Lg loss show that the variance of the gradient is poly- nomial in 1/b. The results theoretically back the important intuition that smaller batch sizes yield larger variance of the stochastic gradients and lower loss function values which is a common believe among the researchers. The proof techniques exhibit relationship between stochastic gradient estimators and initial weights, which is useful for further research on the dynamics of SGD. We empirically pro- vide insights our results on various datasets and commonly used deep network structures. We further discuss possible extensions of the approaches we build in studying the generalization ability of the deep learning models.
INTRODUCTION
Deep learning models have achieved great success in a variety of tasks including natural language processing, computer vision, and reinforcement learning (Goodfellow et al.] (0166). Despite their practical success, there are only limited studies of the theoretical properties of deep learning; see survey papers (Sun) 2019 Fan al. 2019) and references therein. The general problem underlying deep learning models is to optimize (minimize) a loss function, defined by the deviation of model predictions on data samples from the corresponding true labels. The prevailing method to train deep learning models is the mini-batch stochastic gradient descent algorithm and its variants (Bottou, 1998 Bottou et al. (2018). SGD updates model parameters by calculating a stochastic approximation of the full gradient of the loss function, based random selected subset of the training samples called a mini-batch.
It is well-accepted that selecting large mini-batch size reduces the training time of deep learning models, as computation on large mini-batches can be better parallelized on processing units. For example.
al. (2017) scale ResNet-50 (He et al. (0166) from mini-batch size of 256 images and training time of 29 hours, to larger mini-batch size of 8, 192 images. Their training achieves the same level of accuracy while reducing the training time to one hour. However, noted by many researchers, larger mini-batch sizes suffer from worse generalization ability (LeCun et al. 2012. Keskar al., 2017). Therefore, many efforts have been made develop specialized training pro- cedures that achieve good generalization using large mini-batch sizes (Hoffer et al. 2017 Goyal jet al., 2017). Smaller batch sizes have the advantage of allegedly offering better generalization (at the expense of higher training time).
The focus of this study is on the behavior of SGD subject to the conditions on the initial point. This is different from previous results which analyze SGD via stringing one-step recursions together. The dynamics of SGD are not comparable if we merely consider the one-step behavior, as the model pa- rameters change iteration by iteration. Therefore, fixing the initial weights and the learning rate give us a fair view of the impact of different mini- batch sizes on the dynamics of SGD. We hypoth- esize that, given the same initial point, smaller sizes lead to lower training loss and, unfortunately,

---

Under review as a conference paper at ICLR 2021
decrease stability of the algorithm average. The latter follows from the fact that the smaller is the batch size, more stochasticity and volatility is introduced. After all, if the batch size equals to the number of samples, there is stochasticity in the algorithm. To this end, we conjecture that the variance of the gradient in each iteration is a decreasing function of the mini-batch size. The conjecture is the focus of the work herein.
Variance correlates to many other important properties of SGD dynamics. For example, there is substantial work on variance reduction methods (Johnson & Zhang 2013] Allen-Zhu & Hazan 2016 Wang et al. 2013) which show great success on improving the convergence rate by controlling the variance of the stochastic gradients. Mini- batch size is also a key factor deciding the performance of SGD. Some research focuses on how to choose an optimal mini-batch size based on different criteria (Smith & Le 2017 7, Gower et al. 2019). However, these works make strong assumptions on the loss function properties (strong or point or quasi convexity, or constant variance near stationary points) or about the formulation of the SGD algorithm (continuous time interpretation by means of differential equations). The statements are approximate in nature and thus not mathematical claims. The theoretical results regarding the relationship between the mini-batch size and the variance (and other performances, like loss and generalization ability) of the SGD algorithm applied to general machine learning models are still missing. The work herein partially addresses this gap by showing the impact of the mini-batch size on the variance of gradients in SGD. We further discuss possible extensions of the approaches we build in studying the generalization ability.
We able prove the hypothesis about variance in the convex linear regression case and to show significant progress in a deep linear neural network setting with samples based on a normal distribu- tion. In this case we show that the variance a polynomial in the reciprocal of the mini-batch size and that it is decreasing if the mini-batch size is larger than a threshold (further experiments reveal that this threshold can be as small as 2). The increased variance as the mini-batch size decreases should also intuitively imply convergence to lower training loss values and in turn better prediction and generalization ability (these relationships are yet to be confirmed analytically; but we provide empirical evidence to their validity).
The major contributions of this paper are as follows.
For linear regression, we show that in each iteration the norm of any linear combination of sample-wise gradients is a decreasing function of the mini-batch size (Theorem [ As special case, the variance of the stochastic gradient estimator and the full gradient at the iterate in step t are also decreasing functions of b at any iteration step & (Theorem 2 In addition, the proof provides a recursive relationship between the norm of gradients and the model parameters at each iteration (Lemma 2). This recursive relationship can be used to calculate any quantity related to the full/stochastic gradient or loss any iteration with respect to the initial weights.
For the deep linear neural network with L2-loss and samples drawn from a normal distribu- tion, we take two-layer linear network as an example and show that in each iteration step t the trace of any product of the stochastic gradient estimators and weight matrices is a poly- nomial in 1/b with coefficients a sum of products of the initial weights (Theorem 3). As a special case, the variance of the stochastic gradient estimator is polynomial in 1/6 without the constant term (Theorem 4 and therefore it is a decreasing function of b when is large enough (Thcorem 5). The results and proof techniques can be easily extended to general deep linear networks. As a comparison, other papers that study theoretical propertics of two-layer networks either fix one layer of the network, or assume the wee-pparameterized property of the model and they study convergence, while our paper makes no such assump- tions on the model capacity. The proof also reveals the structure of the coefficients of the polynomial, and thus serving as a tool for future work on proving other properties of the stochastic gradient estimators.
The proofs are involved and require several key ideas. The main one is to show a more general result than it is necessary in order to carry out the induction. The induction is on time step t. The key idea is to show a much more general result that lets us carry out induction. New concepts and definitions are introduced in order to handle the more general case. Along the way we show a result of general interest establishing expectation of several rank one matrices sampled from a normal distribution intertwined with constant matrices.
2

---

Under review as a conference paper at ICLR 2021
We verify the theoretical results various datasets and provide further understanding. We further empirically show that the results extend to other widely used network structures and hold for all choices of the maini-bachh sizes. We also empirically verify that, average, in each iteration the loss function value and the generalization ability (measured by the gap between accuracy on the training and test sets) are all decreasing functions of the mini- batch size.
In conclusion, we study the dynamics of SGD under linear regression and a two-layer linear net- work setting by focusing the decreasing property of the variance of stochastic gradient estimators with respect to the mini-batch size. The proof techniques can also be used to derive other properties of the SGD dynamics in regard to the mini-batch size and initial weights. To the best of authors' knowledge, the work is the first one to theoretically study the impact of the mini-batch size on the variance of the gradient subject to the conditions on the initial weights, under mild assumptions on the network and the loss function. We support our theoretical results by experiments. We further experiment on other state-of-the- arr deep learning models and datasets to empirically show the va- lidity of the conjectures about the impact of mini-batch size on average loss, average accuracy and the generalization ability of the model.
The rest of the manuscript is structured follows. In Section 2 we review the literature while in Se present the theoretical results on how mini-batch sizes impact the variance of stochastic gradient estimators, under different models including linear regression and deep linear networks. Section 4 introduces (part of) the experiments that verify our theorems and provide further insights into the impact of the mini-batch sizes on SGD performance. We defer the complete experimental details to Appendix A and the proofs of the theorems and other technical details to to Appendix B
2 LITERATURE REVIEW
Stochastic gradient descent type methods are broadly used in machine learning (Bottou) 1991 Le- Cun et al. 1998 Bottou et al. 2018). The performance of SGD highly relies on the choice of the mini-batch size. has been widely observed that choosing a large mini-bater size to train deep neural networks appears to deteriorate generalization (LeCun et al. 2012). This phenomenon exists even the models are trained without any budget or limits, until the loss function value ceases to improve (Keskar et al., 2017). One explanation for this phenomenon is that large mini-batch SGD produces "sharp" minima that generalize worse (Hochreiter & Schmidhuber, 1997| Keskar et al.] 2017]. Specialized training procedures to achieve good performance with large mini-batch sizes have also been proposed (Hoffer et al. 2017 Goyal et al. 2017).
It is well-known that SGD has slow asymptotic rate of convergence due to its inherent variance (Johnson & Zhang 2013). Variants of SGD that can reduce the variance of the stochastic gradient estimator, which yield faster convergence, have also been suggested. The use of the information of full gradients to provide variance control for stochastic gradients is addressed in Johnson & Zhang (2013) Roux al. (2012); Shhlev-Shwartz & Zhang] (2013). The works in Lei al.) (2017); et al. (2014); Schmidd et al. (2017) further improve the efficiency and complexity of the algorithm by carefully controling the variance.
There is prior work focusing on studying the dynamics of SGD. Neelakantan et al. (2015) propose to add isotropic white noise the full gradient study the "structured" variance. The works in Li et al. (2017); Mandt et al. (2017): Dastrebbsk et al. (2017) connect SGD with stochastic differential equations to explain the property of converged minima and generalization ability of the model. (Smihh & Le (017) propose an "optimal" mini-batch size which maximizes the test set accuracy by Bayesian approach. The Stochastie Gradient Langevin Dynamics (SGLD, a variant of SGD)
algorithm for non-convex optimization is studied in Zhang et al. (2017): Mou et al. (2018).
In most of the prior work about the convergence of SGD. is assumed that the variance of stochastic gradient estimators is upper-bounded by a linear function of the norm of the full gradient, e.g. Assumption 4.3 in Bottou et al. (2018). Gower et al. (2019) give more precise bounds of the variance under different sampling methods and Khaled & Richtárik (2020) extend them smooth non-convex regime. These bounds are still dependent on the model parameters at the corresponding iteration. To the best of the authors' knowledge, there is no existing result which represents the variance of
3

---

Under review as a conference paper at ICLR 2021
stochastic gradient estimators only using the initial weights and the mini-batch size. This paper partially solves this problem.
ANALYSIS
Mini-batch SGD is a lighter-weight version of gradient descent. Suppose that we are given loss function L(w) where w is the collection (vector, matrix, or tensor) of all model parameters. At each iteration t, instead of computing the full gradient (ww)) SGD randomly samples a mini-batch set B, that consists of b = [B,] training instances and sets W+++ + wA a, V w.,., (w;), where the positive scalar ar is the learning rate (or step size) and V. (ws) denotes the stochastic gradient estimator based on mini-batch B,.
An important property of the stochastic gradient estimator V.. (w²) is that it is an unbiased estimator, i.e. EV (wx) (wx), where the expectation is taken over all possible choices of mini-hatch B,. However, it is unclear what is the value of var (V.. T.B, (w²) an E V (mrx)
[IE ( (w,) >|1 Intuitively, we should have var (Vu LB, (w)) x a² b var (V (wx))) where n is the number of training samples and stochasticity the right-hand side comes from mini-batch samples behind wy (Smith & Le 2017 Gower et al. 2019). However, even the quantities V (()) still challenging compute we do not have direct formulas of their precise values. Besides, as we choose different b's, their values are not comparable as we end up with different wt's.
and var (-
A plausible idea to address these issues is to represent EV LB, (wa) and var Lv (wx) using the fixed and known quantities Wo, b, l, and ag. In this way, can further discover the properties, like decreasing with respect to b, of EV u LR, (wy) and var LV (.) The biggest challenge is how to connect the quantities in iteration with those of iteration 0. This is similar to discovering the properties of a stochastic differential equation at time given only the dynamics of the stochastic differential equation and the initial point.
In this section, we address these questions under two settings: linear regression and deep linear network. In Section 3.1 with a linear regression setting, we provide explicit formulas for calcu- lating any norm of the linear combination of sample-wise gradients. We therefore show that the var LB, (w)) is a decreasing function of the mini-batch size b. In Section 3.2 with a deep linear network setting and samples drawn from normal distribution, we show that any trace of the product of weight matrices and stochastic gradient estimators is a polynomial in 1/b with finite degree. We further prove that var (V.. mn T.s, (w)) is decreasing function of the mini-batch size b > by for some constant by-
For a random matrix M, we define var ( (M) E [[]vec( [[vee(((M)) ([vee where vec( denotes the vectorization of matrix M. We denote [m : n²] + {m, m 1..... n) if m & n, and Ø otherwise. We use [n] + [1: : n] as an abbreviation. For clarity, we use the superscript b to distinguish the variables with different choices of the mini-batch size b. In each iteration the batch of samples (or sample indices) to calculate the stochastic gradient. We denote by Fb the filtration of information before caleulating the stochastic gradient in the t-th iteration, i.e. F!. +
B& denote
{ (wo, , 1
3.1 LINEAR REGRESSION
In this subsection, we discuss the dynamics of SGD applied in linear regression. Given data points (I1, 11),.).... ... (Zn, Vn). where I; E Ip' and yi R. we define the loss function to be L(w) "E (w) -m= (w )² where w e IR" are the model parameters. We consider minimizing L(w) by mini-batch SGD. Note that the bias term in the general linear regression models is omitted, however, adding the bias term does not change the result of this section. Formally, we first choose a mini-batch size b and initial weights wi- In each iteration I, we sample By a subset of [n] with cardinality b, and update the parameters by wh., w/-a where of 1 Li VL, (wf).
We first show the relationship between the variance of stochastic gradient of and the full gradient VL (w?) and sample-wise wise gradient V L. (w+²) í E [n]. derived by considering all possible choices
4

---

Under review as a conference paper at ICLR 2021
of the mini-batch B Readers should note that Lemma actually holds for all models with Lz-loss, not merely linear regression (since in the proof we do not need to know the explicit form of Lz(w)).
Lemma 1.
Let cb
0. For any matrix A E Rexp we have var (Ag) (7))
E
|
FP] - IAVL (w?) 2 cb | AVL, - (w?) (²)
Lemma 1 provides a bridge to connect the norm and variance of oe with sample-wise gradients VI, (wx) ; [n]] Therefore, if we can further discover the properties of Vi, (w²) i [r]]. we
are able calculate the variance of g.- Lemma addressses this problem by showing the relationship between any linear combination of V L; (wt)* 's and VI; (w(_)'ss.
Lemma 2. For any set of square matrices (A1,... An) E TRXX, if we denote A = | A (wi+1) Fo 1
Ei_1 A then we have E
E
Σ
L, (w:))
|| Fo +
Ek= E [ E
0f00 a²
B (w²)
where B; = A at A: B A if k 11
B.d A iff==liii k, and B equals the zero matrix, otherwise.
Lemma 2 provides the tool to reduce the iteration t by one. Therefore, we can easily use it to recur- sively calculate the norm of any linear combinations of the sample-wise gradients, for all iterations t. Combining the fact that Op is decreasing function of b, we are able to show Theorem
Theorem 1. For any N and any matrices A, RPxP, i [n]] EL A.V. (w Fo is decreasing function of b for b [ni]
Theorem D states that the norm of any linear combinations of the sample-wise gradients de- creasing function of b. Combining Lemma which connects the variance of 9; with the linear combination of LL, (uf) "s, and the fact that VL (w?) - VL, (w?). we have Theorem
Fixing initial weights wo. both var (Bgy Fo) and var (BVL (w!) (Fo) are decreasing functions of mini-batch size for all b E [n]] t e N, and all square matrices gexp.
Theorem 2.
As a special case, Corollary [ guaranttes that the variance of the stochastic gradient estimator is a decreasing function of b.
Corollary 1. Fixing initial weights wp, both var (of Fo) and var (VL (w?) (Fo) are decreasing functions of mini--btth size for all b E [n] and t E N.
In conclusion, we provide a framework for calculating the explicit value of variance of the stochastic gradient estimators and the norm of any linear combination of sample-wise gradients. We further show that the variance of both the full gradient and the stochastic gradient estimator are decreasing function of the mini-batch size Readers should note that the framework here is not limited to showing the decreasing property of the variance, but can also be used in many other cireumstance. For example, we can use Lemma 2 to induct on f and easily show that E Σ> A,V, (w²) ( is polynomial of with degree most t and estimate the coefficients therein.
3.2 DEEP LINEAR NETWORKS WITH ONLINE SETTING
In this section, we study the dynamics of SGD on deep linear networks. We take the two-layer linear network
example while the results and proofs can be easily extended to deep lin- ear network with any depth (see Appendix B.3 for more details). We consider the population
loss L(w) Bz=N(0.F,))
WWW Wtt"#=?
21 under the teacher-student learning frame-
(W1, W2) tuple of two matrices. Here W1 1 xp and W2 E P are parameter matrices of the student network and W and W are the fixed ground-truth parameters of the teacher network. We use online SGD to minimize the population loss L(w). Formally, we first choose a mini-batch size b and initial weight matrices (Wo,1, Wo,2)- In each iteration I, draw b independent and identically distributed samples 2ı,ú, i E [b] from A(O0, I2) to form the mini-batch By and update the weight matrices by W 1+,11 Wh: augh, and
work (Hinton et al. 2015) with
5

---

Under review as a conference paper at ICLR 2021
W 4112 Wh,2 agg22. where
t.1 1 b Σ w (2 W -e me,s (³)
_
w (wiaw"t w w/) ae
(1)
2/22
w
W Wh We. - me.
P))
( W2 w.fWF)) on w.
(2)
The derivation follows from the formulas in Petersen Pedersen (001)) In the following, we use W. W W7, 1an-WW3 weights.
W2 WT to denote the gap between the product of model weights and ground-truth
For case of developing our proofs, we first introduce the definition of a multiplicative term in Def- inition ! Intuitively, a multiplicative term is a matrix which equals to the product of its parameter matrices and constant matrices (and their transpose). The degree of matrix A in a multiplicative term M is the number of appearance of A and A in M. The degree of M is exactly the number of appearances of all weight matrices in M.
Definition For any of matrices S, we denote S = (M2 -M S}. Given a set of parameter matrices X =
, } and constant matrices C = (C,,, Cn we say that a matrix M is a multiplicative term of parameter matrices X' and constant matrices C can be written in the form of M- = M) = 11 A.. where A; E X C.C. We write deg( X5; M) = E (1 A, + 1 1 - Af)) E [m.]] as the degree of parameter matrix X, in M, de = Ei (1{C, +
j [ne] as the degree of constant matrix C3 in M, and deg(M) Σ 1 {) ード deg(X;; M) as the total degree of the parameter matrices of M.
As pointed out in the Section T the difficulty of studying the dynamics of SGD is how to connect the quantities in iteration with fixed variables, like initial weights Wo,1, Wo, and mini-batch size b. We overcome this challenge by carefully calculating the relationship between g and gi... i= 1, so 9 in Appendix [B.2 we can represent S.11. = 2,2 using multiplicative terms of oh,, i = 1, 2 and some other constant matrices. Theorem 3 precisely gives the representation in the form of polynomial of 1/ and the coefficients the sum of multiplicative terms of parameter matrices { (Wa1. W and constant matrices (WT, W.
that we can reduce the iteration & step by step. With the help of Lemmas 8 and
Theorem 3. Given t O. for any multiplicatiie terms M,, E [0: [0:m] of parameter matrices (of.1 00,119922 and constant matrices {W W 1,2, W W with degree di respectively, we denote M I tr (M.) Mo. d = Σ de and d' -i= (deg (W M.) + deg(W?; M.)).. There exists a set of multiplicative terms {M i [mm1.je E [o: med, k E [0: q1} of parameter matrices (Wo, 3 and constant matrices {W Wx) such that E [M[Fo] = Ng + N1 + ..... + N Ng 69
where N₂ 1 Ir (MZ) Mb. k E [0: q]. Here my, mai and y 1/ (3 - 1)d + 1
1)d are constants independent b, and Ln deg (MK) < 3'd d').
As special case of Theorem Theorem shows that the variance of the stochastic gradient esti- mators is also a polynomial the variance is approximately inversely proportional to the mini-batch size b. Besides, note that if we consider - oo, intuitively we should have var (oF - 0, 1,2. This observation aligns with the statement of Theorem
but with no constant term. This backs the important intuition that
Theorem 4. Given > O, value var (st..)o), 1,2 can be written a polynomial of 1 with degree at most 2 3º with no constant term. Formally, we have var (FF 31 + B 1.
where r x 2 3 and each B3 is a constant independent of b.
One should note that the polynomial representation of var Fo). 1,2 does not have the constant term. Therefore, to show the that the variance is decreasing function of b, we only need to show that the leading coefficient 3 is nen-negative. This is guarantced by the fact that variance is always nen-tegative. We therefore have Theorem 5

---

Under review as a conference paper at ICLR 2021
Theorem 5. Given t E N. there exists constant bu such that for all b > by function var (o.,(Fo) i= 1,2 2,226 is a decreasing function of b.
The constant by is the largest root of the equation 3,8"-1 + 32b"~2 +....
........ +.......... B, = 0. See the proof
of Theorem [5 in Appendix B.2 for more details. Although we cannot calculate the precise value of bo. we verify that by is smaller than in many experiments. From the proofs we conclude that the scale of each 3 is of the order O (I) where M is multiplicative term of parameter matrices [Wo,1. 2o,2, WT, W and constant matrix Ø with degree 2 3
Unlike the linear regression setting where we can iteratively calculate the variance by Lemma 2 the closed form expressions for the variance of the stochastic gradients in the deep linear network setting are much harder to calculate. However, we are able to iteratively deducing t one by one and provide a polynomial representation for any multiplicative terms of parameter matrices { and constant matrices (WT. W } using only the initial weights Wo,1, Wo.2 and the mini-batch size b. As we further study the polynomial representation of var (o)) = 1,2, we are also able to show the decreasing property of the variance of stochastic gradient estimators with respect to b.
h W t,3: i 1,
4 EXPERIMENTS
In this section, we present numerical results to support the theorems in Section 3aand provide further insights into the impact of the mini-batch size on the dynamics of SGD. The experiments are con- ducted on four datasets and models that are relatively small due to the computational cost of using large models and datasets. We only report the results on the MNIST dataset here due the limited space. A complete empirical study is deferred in Appendix A
For all experiments, we perform mini-batol SGD multiple times starting from the same initial weights and following the same choice of the learning rates and other hyper-parameee if applica- ble. This enables us to calculate the variance of the gradient estimators and other statistics in each iteration, where the randomness comes only from different samples of SGD.
4.1 RESULTS ON MNIST DATASET
The MNIST dataset is to recognize digits in handwritten images of digits. We use all 60,000 training samples and 10,000 validation samples of MNIST. We build three-layer fully connected neural network with 1024, 512 and 10 neurons in each layer. For the two hidden layers, we use the ReLU activation function. The last layer is the softmax layer which gives the prediction probabilities for the 10 digits. We use miii-batch SGD to optimize the cross- loss of the model. The model deviates from our analytical setting since has non-linear activations. it has the cross-entropy loss function (instead of L2), and empirical loss (as opposed to population). MNIST is selected due to its fast training and popularity in deep learning experiments. The goal to verify the results in this different setting and to back up our hypotheses.
Different initial weights
(b) Log
(c) Log of
Gap accuracy (zoomed-in)
Figure I: Experimental results for the MNIST dataset. (a) The median, min, and max of the log of variance of the stochastic gradient estimators for two different mini-batch sizes (distinguished by colors) and five different initial weights. The solid lines show the median of five initial weights while the highlighted regions show the min and max of the log of variance. (b) The log of the training and validation loss vs epochs. (c) The log of training and validation crror vs epochs. Here crror is defined as one minus predicting accuracy. The plot does not show the epochs if error equals to zero. (d) The gap of accuracy on training and test sets vs epochs starting from epoch 100.
7

---

Under review as a conference paper at ICLR 2021
As shown in Figure aa), run SGD with two batch sizes 64 and 128 on five different initial weights with 50 runs for each initial point. This plot shows that, even the smallest value of the variance among the five different initial weights with mini-batch size of 64, is still larger than the largest variance of mini-hatch size 128. We observe that the sensitivity to the initial weights is not large. This plot also empirically verifies our conjecture in the introduction that the variance of the stochastic gradient estimators is a decreasing function of the mini-batch size, for all iterations of SGD in general deep learning model.
In addition, we also conjecture that there exists the decreasing property for the expected loss, error and the generalization ability with respect the mini-batch size. Figure shows that the expected loss (again, randomness comes from different runs of SGD through the different mini mini-batches with the same initial weights and learning rates) on the training set is decreasing function of b. However, this decreasing property does not hold on the validation set when the loss tends to be stable or increasing, in other words, the model starts to be over-fitting We hypothesize that this is because the learned weights start to bounce around a local minimum when the model is over-fitting. As the larger mini-batch size brings smaller variance, the weights are closer to the local minimum found by SGD, and therefore yield smaller loss function value. Figure [e) shows that both the expected error on training and validation sets are decreasing functions b.
Figure dd) exhibits relationship between the model's generalization ability and the mini-bate size. As suggested by [Simard et al. (2013), we build a test set by distorting the 10,000 images of the validation set. The prediction accuracy is obtained the gap between these two accuracies every 100 epochs. We use this gap to measure the model generalization ability (the smaller the better). di shows that the gap an increasing function of b starting at epoch 500, which partially aligns with our conjecture regarding the relationship between the generalization ability and the mini-batch size. We test this on multiple choices of the soper-parameter which control the degree of distortion in the test set and this pattern remains clear.
both training and test sets and we calculate
SUMMARY AND FUTURE WORK
We examine the impact of the mini-batch size on the dynamics of SGD. Our focus is on the variance of stochastic gradient estimators. For linear regression and two-layer linear network, we are able to theoretically prove that the variance conjecture holds. We further experiment on multiple models and datasets verify our claims and their applicability to practical settings. Besides, also empirically address the conjectures about the expected loss and the generalization ability.
A challenging research direction is to theoretically investigate the impact of the mini-batei size on the generalization ability. There are existing works studying the relationship between the variance of the stochastic gradients and the generalization ability (Gorbunov et al. 2020 Meng al. (016). Together with the tools developed herein, it would be possible to bridge the mini batch size with the generalization ability of a neural network. We can further choose an optimal mini-batch size which minimizes the generalization ability by solving the polynomial equation if we have more precise estimations of the coefficients.
Another appealing direction is using our variance estimations to develop better variance reduction methods. As a results, the upper-bound of the variance decides the convergent rate of these al- gorithms. Researchers usually assume a much larger upper-bound at each iteration, like a linear function of the norm of the full gradient. With the help of our techniques, we should calculate the variance more precisely and further improve the algorithms.
Further interesting work is to extend our techniques to more complicated and sophisticated networks. Although the underlying model of this paper corresponds to deep linear network networks, we are able to show a deeper relationship between the variance and the mini- bani-bacch size, the polynomial in 1/b, while the common knowledge is simply that the variance is proportional to 1/b. The extension to other optimization algorithms, like Adam and Gradient Boosting Machines, are also very attractive. We hope our theoretical framework can serve as tool for future research of this kind.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Mohan S Acharya, Asfia Armaan, and Aneeta S Antony. A comparison of regression models for pre- diction of graduate admissions. In 2019 International Conference on Computational Intelligence in Data Science, pp. 1-5, 2019.
Zeyuan Allen-Zhu and Elad Hazan. Variance reduction for faster non-convex optimization. In International conference on machine learning, pp. 699-707, 2016.
Léon Bottou. Stochastic gradient learning in neural networks. Proceedings of Neuro-Nimes, 91(8):
12, 1991.
Léon Bottou. Online learning and stochastic approximations. On-line Learning in Neural Networks, 17(9):12,, 1998.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018.
Jianqing Fan, Cong Ma, and Yiqiao Zhong. A selective overview of deep learning. arXiv preprint arXiv:1904.05526, 2019.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. A unified theory of sgd: Variance reduc- tion, sampling, quantization and coordinate descent. In International Conference Artificial Intelligence and Statistics, pp. 680-690, 2020.
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. In International Conference on Machine Learning, pp. 5200-5209, 2019.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training Imagenet in hour. arXiv preprint arXiv: 1706.0677 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. Proceedings of the IEEE conference on Computer Vision and Puttern Recognition, pp- 770-778, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint
2015.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. In Advances in Neural Information Processing Systems, pp. 1731-1741, 2017.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint arXiv:1711.04623, 2017.
Rie Johnson and Tong Zhang, Accelerating stochastic gradient descent using predictive variance reduction. In Advances in Neural Information Processing Systems, pp. 315-323, 2013.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang. Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, 2017, 2017.
Ahmed Khaled and Peter Richtárik. Better theory for sgd in the nonconvex world. arXiv preprint arXiv:2002.0333 033399, 2020.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- baadie learning applied to document recognition. Proceedings of the Institute of Electrical and Electronics Engineers, 86 (11):2278-2324 1998.

---

Under review as a conference paper at ICLR 2021
Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Rober Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Lihua Lei, Cheng Ju, Jianbo Chen, and Michael Jordan. Non-convex finite-sum optimization via
SCSG methods. In Advances in Neural Information Processing Systems, pp. 2348-2358, 2017.
Mu Li, Tong Zhang. Yuqiang Chen, and Alexander J Smola. Efficient mini-batch training for stochastic optimization. Proceedings of the 20th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pp. 661-670, 2014.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning, pp. 2101-211000 10. PMLR, 2017.
Jan R Magnus. The moments of products of quadratic forms in normal variables. Instituut Actuariaat en Econometrie, 1978.
Stephan Mandt, Matthew Hoffman, and David Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1):4873 4907, 2017.
Qi Meng, Yue Wang, Wei Chen, Taifeng Wang, Zhi-Ming Ma, and Tic- Yan Liu. Generalization error bounds for optimization algorithms via stability. arXiv preprint 1609::1609.08800 2016.
Wenlong Mou, Liwei Wang. Xiyu Zhai, and Kai Zheng. Generalization bounds of SGLD for non- convex learning: Two theoretical viewpoints. In Conference On Learning Theory, pp. 605-635 638,, 2018.
Arvind Neelakantan, Luke Vilnis, Quoc 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.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems, pp- 8024-8035, 2019.
K. B. Petersen and M. S. Pedersen. The matrix cookbook, 2012. Version 20121115.
Nicolas L Roux, Mark Schmidt, and Francis R Bach. A stochastic gradient method with expo- nential convergence rate for finite training sets. In Advances in Neural Information Processing Systems, pp. 2663-2671, 2012.
Mark Schmidt, Nicolas Le Roux, and Francis Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming, 162(1-2):83-112, 2017.
Shai Shalev-Shwartz and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Research, 14(Feb):567-59,, 2013.
P. Y. Simard, D. Steinkraus, and J. C. Platt. Best practices for convolutional neural networks applied to visual document analysis. In Seventh International Conference on Document Analysis and Recognition, pp. 958-963, 2013.
Samuel L Smith and Quoe V Le. A bayesian perspective on generalization and stochastic gradient descent. arXiv preprint 1710:077,. 2017.
Ruoyu Sun. Optimization for deep learning: theory and algorithms. arXiv preprint arXiv:1912.089577 2019.
Chong Wang. Xi Chen, Alexander J Smola, and Eric P Xing. Variance reduction for stochastic gradient optimization. Advances in Neural Information Processing Systems, pp. 181-189, 2013.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances Neural Information Processing Systems, pp. 5754-5764, 2019.
10

---

Under review as a conference paper at ICLR 2021
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. In Advances in Neural Information Processing Systems, pp. 649-657, 2015.
Yuchen Zhang, Percy Liang, and Moses Charikar. A hitting time analysis of stochastic gradient langevin dynamics. In Conference on Learning Theory, pp. 1980-2022, 2017.

---

Under review as a conference paper at ICLR 2021
A EXPERIMENTS
In this section, we present numerical results to support the theorems in Section to backup the hypotheses discussed in the introduction, and provide further insights into the impact of the mini- batch size on the dynamics of SGD. The experiments are conducted on four datasets and models that are relatively small due to the computational cost using large models and datasets.
We cannot present the complete numerical results in the main paper due to the space limit. Therefore, we move the whole experimental section to Appendix. In order to keep a smooth reading, some of the content is overlapping with Section 4
Remark:
A.1 DATASETS AND SETTINGS
For all experiments, we perform mini-batch SGD multiple times starting from the same initial weights and following the same choice of the learning rates and other eyper-parameters if applica- ble. This enables us to calculate the variance of the gradient estimators and other statistics in each iteration, where the randomness comes only from different samples of SGD. The learning rate o, is selected to be inversely proportional to iteration I, or fixed, depending on the task at hand.
All models are implemented using Py PyTorch version 1.4 (Paszke al.) 2019) and trained NVIDIA 2080T1/1080 GPUs. We have also tested several other random initial weights and ground-truth weights, and learning rates, and the results and conclusions are similar and not presented.
A.1.1 GRADUATE ADMISSION DATASET
The Graduate Admission datase (Acharya et al. 2019) is to predict the chance of graduate ad- mission using linear regression. The dataset contains 500 samples with 6 features and is normalized by mean and variance of each feature. This is a popular regression dataset with clean data. We build a linear regression model to predict the chance of acceptance (we include the intercept term in the model) and minimize the empirical L2 loss using
SGD, stated in Section 3.1
For the experiment in Figure 2 a), we randomly select an initial weight vectors wi and run SGD for 2,000 iterations where it appears to converge. We record all statistics at every iteration. There are in total 1,000 runs behind each observation which yields a p p-value lower than 0.05. As for Figure select 20 different O's and run SGD from the same initial point for 40 iterations. There are in total of 200,000 runs to make sure the p-value of all statistics are lower than 0.05. In all experiments, the learning rate is chosen be a ž E [2000] because this rate yields theoretical convergence guaranteed (factor 1/2 has been fine tuned). The purpose of this experiment is to empirically study the rate of decrease of the variance. The theoretical study exhibited Section 3.1 establishes the non-increasing property but it does not state anything about the rate of decrease.
A.1.2 SYNTHETIC DATASET
We build a synthetic dataset of standard normal samples to study the setting in Section 3.2 We fix the teacher network with 64 input neurons, 256 hidden neurons and 128 output neurons. We optimize the population T22 loss by updating the two parameter matrices of the student network using online SGD, as stated in Section 3.2 In this variance as a function of b and show the decreasing property of the variance of the stochastic gradient estimators for large mini-batch sizes. However, we do not show the decreasing property for every b. With this experiment we confirm that the conjecture likely holds. In the experiment, we randomly select two initial weight matrices Wo.1, Wo,2 and the ground-truth weight matrices W W². We run SGD for 1,000 iterations which appears to be a good number for convergence while there are 1,000 runs of SGD in total to again give p-value below 0.05. We record all statistics at every iteration. The learning rate is chosen to be 0t Ioe: E [1000] for the same reason as in the regression experiment.
have proved the functional form of the
https:: htw.:/gwg.....ooooaaaaaaaaayrgga admi ons
12

---

Under review as a conference paper at ICLR 2021
A.1.3 MNIST DATASET
The MNIST dataset is recognize digits in handwritten images of digits. We use all 60,000 training samples and 10,000 validation samples of MNIST. The images are normalized by mapping each entry 1-1, 1]. We build a three-layer fully connected neural network with 1024, 512 and 10 neurons in each layer. For the two hidden layers, we use the ReLU activation function. The last layer is the softmax layer which gives the prediction probabilities for the 10 digits. We use mini- batch SGD to optimize the cross-entropy loss of the model. The model deviates from our analytical setting since it has non-linear activations, has the cross-entropy loss function (instead of L2), and empirical loss (as opposed to population). MNIST is selected due to its fast training and popularity in deep learning experiments. The goal is to verify the results in this different setting and to back up our hypotheses.
We run SGD for 1,000 epochs on the training set which is enough for convergence. The learning rate is a constant set to 3 10-3 (which has been tuned). For the experiment in Figure there are in total 100 runs to give us the p-value below 0.05. For the experiment in Figure 4a), [). we randomly select five different initial points and we have 50 runs for each initial point. For the experiment corresponding to Figure 4b) we choose a = 8 and a = 2 as in Simard et al. (2133). The initial weights and other param-earseeters chosen to be the same in Figure 5
A.1.4 YELP REVIEW DATASET
The Yelp Review dataset from the Yelp Dataset Challenge (Zhang al. 2013) contains 1,569,266 264 samples of customer reviews with positive/negative sentiment labels. We use 10,000 samples as our training set and 1,000 samples as the validation set. We use XLNet (Yang et al. 2019) to perform sentiment classification on this dataset. Our XLNet has 6 layers, the hidden size of 384, and 12 attention heads. There are in total 35,493, 122 parameters. We intentionally reduce the number of layers and hidden size of XLNet and select relatively small size of the training and validation sets since training of XLNet is very consuminn (Yang et al. (2019) train on 512 TPU v3 chips for 5.5 days) and need to train the model for multiple runs. This setting allows us train our model in several hours on a single GPU card. We train the model using the Adam weight decay optimizer, and some other techniques, as suggested in Table 8 of Yang et al. 2019)). This dataset represents sequential data where we further consider the hypotheses.
We randomly select set of initial parameters and run Adam with two different mini-batch sizes of 32 and 64. For computational tractability reasons, for each mini-batch size there are in total of 100 runs and each run corresponds to 20 epochs. We record the variance of the stochastic gradient, loss and accuracy in every step of Adam. The statistics reported in Figure 6 are averaged through each epoch. In all experiments, the learning rate is set be 10-5 and the parameter of Adam is set to be 10-8 (these two have been tuned). The stochastic gradients of all parameter matrices are clipped with threshold I in each iteration. We use the same setup for the learning rate warm-up strategy as suggested in Yang et al. (2019). The maximum sequence length is set to be 128 and we pad the sequences with length smaller than 128 with zeros.
A.2 DISCUSSION
As observed in Figure a). under the linear regression setting with the Graduate Admission dataset, the variance of the stochastic gradient estimators and full gradients are all strictly decreasing func- tions of for all iterations. This result verifies the theorems in Section 3.1 Figure 2bb) further studies the rate of decrease of the variance. From the proofs in Section 3.1 we that var (g)l Fo)
a polynomial of } with degree t 1. Therefore, for every f,, we can approximate this polynomial by sampling many different b's and calculate the corresponding variances. We pick b to cover all numbers that are either a power of 2 or multiple of 40 in [2, 500] (there are a total of 21 such values) and fit a polynomial with degree 6 (an estimate from the analyses) at f 10, 20, 30, 40. Figure b shows the fitted polynomials. As we observe, the value var (o)(Fo) (approximated by the value of the polynomial) is both decreasing with respect to the mini-batch size b and iteration t. Further, the rate of decrease in b is slower as the increasing. This provides a further insight into the dynamics of training linear regression problem with SGD.
13

---

Under review as a conference paper at ICLR 2021
balth
batch atschaalic
-0.5
itoration
Log of varience of w2ln and o7
0-
balth 16,
balth B4,
batch 86, stachedlt gred
batch 16, sincheele gred
iteration 20
-2-
batth gred
batch stochasti
Log polemariial valuee
-1.
iteraion 40
-6-
2.0
-0-
-2.5%
-10-
-3.0
1250 1500
1750
2000
500
Rerelion
Mini batch size
(a) Variance stochastic gradients and full gradients
(b) Fitting polynomials of mini-bator size b
Figure 2:
Experimental results for the Graduate Admission dataset.
Left:
log (var (s((Fa) and
log (var (VL(w)) (Fo)) vs iteration for 4 different mini-hatch sizes. Right: The log of polynomial val- when futting polynomials on selected mini-batch sizes at certain iterations.
batch eite gnnd grod
batch stortaticc grad
Log of variance of PayEtorf) and pi
Log of variance of (ouc(oof) and of::
grad
grod
eetk grad
-0-
batt1 Size gmd
Stechactic goyd stretwetic grac grac good
15; grad
-10
grdd
grad
-10-
200
400
800
800
1000
200
400
600
800
1000
Iteration
Iteration
(a) Variance of gradients with respect to W
(b) Variance of gradients with respect to W2
Figure
Experimental results for the Synthetic dataset.
Left:
log (var (Fo)) log (var (g)) and
log (var (Vw, W22) vs iteration
Right:
log (var (Vw,C(WE,, W1.2) Fo)) iteration t.
Under the two-layer linear network setting with the synthetic dataset, Figure verifies that the variance of the stochastic gradient estimators and full gradients are all strictly decreasing functions of b for all iterations. This figure also empirically shows that the constant bo in Theorem 5 could be small decreasing property remains to hold. We also test this on multiple choices of initial weights and learning rates and this pattern remains clear.
4. In fact, also experiment with the mini-batch size of and 2, and the
In aforementioned two experiments we use SGD in its original form by randomly sampling mini- batches. In deep learning with large-scale training data such a strategy is computationally prohibitive and thus samples are scanned in a cyclic order which implies fixed mini mini-batches are processed many times. Therefore, the next two datasets we perform standard "epoch" based training to empirically study the remaining two hypotheses discussed in the introduction (decreasing loss and emoor as a function of and sensitivity with respect the initial weights. Note that we are using cross-entropy loss in the MNIST datasct and the Adam optimizer in the Yelp dataset and thus these experiments do not meet all of the assumptions of the analysis Scction
As shown in Figure 4 a), we SGD with two batch sizes 64 and 128 on five different initial weights. This plot shows that, even the smallest value the variance among the five different initial weights with mini-batch size of 64, is still larger than the largest variance of 128. We observe that the sensitivity to the initial weights is not large. This plot also empirically verifies our conjecture in the introduction that the variance of the stochastic gradient estimators is
size
14

---

Under review as a conference paper at ICLR 2021
batch
0.049
botch size 128, median
0.048
-6-
Ciap accurncy
0.047
0.046
Leg
-12
0.044
batch di: trai
-14
0.043
batch slee 125: train bater dize 250: train acc test acx
100
200
400
800
900
Epochs
Epachs
(8) Different initial weights
Gap accuracy (zoomed-in)
Figure 4: Experimental results for the MNIST datasct. Left: of the stochastic gradient estimators ferent initial weights. The solid lines show the median of all five initial weights while the highlighted regions show the min and max of the log of variance. Right: The gap of accuracy on training and test sets vs epochs starting from epoch 100.
The median, min, and max of the log of variance two different mini-bate sizes (distinguished by colors) and five dif-
botch 922 128: bai butch aie Do6
butch size 64: train
batch sue 128, vol dotion
batch size 64: vullcatioon
cain
vellcation
batch size 256, volidatioon
Log of
-6
Log
-6-
200
400
600
800
1000
200
a00
600
800
1000
Epochs
(a) Log of loss for training and validation sets
(b) Log of error for training and validation sets
Figure 5: Experimental results for the MNIST dataset. Left: The log of the training and validation loss vs epochs. Right: The log of training and validation error epochs. Here error is defined as one minus predicting accuracy. The plot does show the epochs if error equals to
bunch
butih trom beeth vildbrion
butec tran
harr valtaroon
18840
18850
10000
0.20
0.55
3800
0.3
05
(a) Variance of stochastic gradients
(b) Training and validation loss
(c) Training minus validation error
Figure 6: Experimental results for the XLNNt model on the Yelp dataset. Left: The variance of stochastic gradient estimators vs epochs. Middle: The training and validation loss epochs. Right: The training and validation error epochs.
a decreasing function of the mini-batch size, for all iterations of SGD in a general deep learning model.
In addition, we also conjecture that there exists the decreasing property for the expected loss, error and the generalization ability with respect the mini-batch size. Figure 5 shows that the expected
15

---

Under review as a conference paper at ICLR 2021
loss (again, randomness comes from different runs of SGD through the different mini-batches with the same initial weights and learning rates) on the training set is decreasing function of b. However, this decreasing property does not hold the validation set when the loss tends to be stable or increasing, in other words, the model starts to be over-fitting. We hypothesize that this is because the learned weights start to bounce around a local minimum when the model is over-fitting As the larger mini-batch size brings smaller variance, the weights are closer to the local minimum found by SGD, and therefore yield smaller loss function value. Figure (5)b shows that both the expected error on training and validation sets are decreasing functions of b.
Figure 4b)) exhibits relationship between the model's generalization ability and the mini- batch size. As suggested by (Simard et al. (2013), we build a test set by distorting the 10,000 images of the validation set. The prediction accuracy is obtained on both training and test sets and we calculate the gap between these two accuracies every 100 epochs. We use this gap to measure the model generalization ability (the smaller the better). (b) shows that the gap is an increasing function of 6 starting at epoch 500, which partially aligns with our conjecture regarding the relationship between the generalization ability and the mini-batch size. We also test this on multiple choices of the hyper-parament which control the degree of distortion in the test set and this pattern remains clear.
Figure shows the similar phenomenon that the variance of stochastic estimators and the expected loss and error on both training and validation sets are decreasing functions of even we train XLNet using Adam. This example gives confidence that the decreasing properties are not merely restricted on shallow neural networks or vanilla SGD algorithms. They actually appear in many advanced models and optimization methods.
B LEMMAS AND PROOFS
B.1 LEMMAS AND PROOFS OF RESULTS IN SECTION 3.1
For two matrices A, B with the same dimension, we define the inner product (A, B) + tr (AT
Lemma 3. Suppose that f(x) and g(æ) are both smooth, non-negative and decreasing functions of x E R. Then h(r)- = f()) is also non-negative and decreasing function of x.
Proof. It is obvious that h(z) is non-negative for all x. The first-order derivative of h is
h'(z) f(())(()) f()) E0,
and thus h(x) is also decreasing function of x.
Proof of Lemmal 7 ber. Note that
Throughout the paper, We use Ch [[] to denote the combinatorial num-
96
b
(g))
J
Σ VL
(w))
Σ
VL (wi)
1E8
h-1
h-2
Σ
Li; (w LL (2x))²
V L; (w?) VL (w0z)² T
1
6
Ch itj
b Σ
Li (w)) VL (w?) T + b(b 1)
Σ
VL, (w VL, (wi)) 7
B
n(n 1) itj
1
n2nn 1) Σ
b(n b)
VL, (w)) VL, (wx)² T + b(b - 1)
n(n 1) 2 VL (w Σ VL, (2)))T
n n-6 b
(b- 1)m VL (w?) 7L (w?) T
L4 (w)) VL, (wi)
bn - 1)
11
1
16

---

Under review as a conference paper at ICLR 2021
For any A E R we have
E
A
J
B
(9t)
T
AT
Agi
= tr (())) A" A0t
E
tr
(AT
Agt (9t)
F
A
AE
ge (9t)
F
(b- 1)m AT AVI. (w?) VL (1x))²
AT AVT.; (w?)) VL. (w?)²
bn(n
b(n -n
i=1
n
AVL; (w 2 + b(n-1) || |AVL (m,²) 1
(b-1)n
bn(n - 1)
AVL. (u:)]² - - (1at)) (ut)() + 1|AVL (
cb
Therefore, have
var F4)
TR.
AAg b1
Fb - EE [A F
=IE
||A0
- |AVL (w 2
AVL: (w,) - (wx))
Lemma 4. For any of square matrices (A,,".. ; And E TRPXP, if we denote A then we have
Σn Awtr
Σ A. (0211)
50 =3 [ B.L (u:)
Fo
of
Σ E
Σ B vL. (wt)
[50]
Here B, = A A: B A k, B = Affi 1, k and Bkd equals the zero
3z A
matrix, otherwise.
17

---

Under review as a conference paper at ICLR 2021
Proof of Lemma 4 Let C2 zt and C
E== C.. For the given A1... An, denote
A
A,C, Then we have
2
n
E
A L. (w(-1)
Fo
E
E
A,
(w2+1)
Ft
Fo
12
E
E
Σ
A (of w/++ y.)
Ft
Fu
E
E
Σ
A (1) (w) (n.9) - wi) It
Ft
Fo
22
E
E
Σ
A A. (w²) - 19
Fr
Fo
E
Σ
A. (w²)
Fo
20
E
Σ A A.VLe (w?)
Art
F
Fo
+
af E
EE
||>T)
E
L. (w))
Fo
20
A. (w²) AVL (wtt)
Fo
+ a
+ |AV -IAVL(
Fo
cb
72
E
Σ
A L, (w²) (
Fo
+ az 23
AV
(VV ? Fo
7
afch n²
E
A VL, (w?) c.AVL(u
Fo
> it
E |||||115 |AVL. (w AVL, (
2
E
Σ Ai 0t A)
Li (w²)
Fo
+
07c0
n
n
E
|AALL (w. - AVLL, (1o [Fo}
2
1=15=
Therefore, if we set B, A₁
A and
A i
A
i= 1,,i k.
otherwise,
have
22
Σ
n
V
(wi+1)
Fo
=E
Σ B. L. (w+)
Fo
o
Cb
n 12
Σ B v L; (utt)
Fo
n²
Proof of Theorem 7 We use induction to show this statement.
When = 0, | (w+) (1 = I_T_, A. (wo) which is invariant There- fore, it is decreasing function of b.
18

---

Under review as a conference paper at ICLR 2021
Suppose the statement holds for t. For any set of matrices {A1,-- ......... And in Rex,, by Lemma know that there exist matrices (B1,... B..) and {B k, [n]} such that
2 we
Σ A V L. (wf+1)
Fo =E
Σ B L. (w²)
Fo
+
Σ B VL (uft)
Fo
n
By induction, we know that E
[-=1 B; (w+) (1)) and all E
E VL (w | |7
afea n²
of (n-
are non-negative and decreasing functions of b. Besides, clearly
non son-negative
bn (n-1)
and decreasing function of b. By Lemma 3 know that
aiee n2
E
| B B. Li (w || Fo is
also non-negative and decreasing function of Finally, E
E
A,V (wh+1) Fo as the
sum of non-negative and decreasing functions in b, is non-negative and decreasing function of b.
In order to prove Theorem 2 we split the task to two separate theorems about the full gradient and the stochastic gradient and prove them one by one.
Theorem 6. Fixing initial weights wo, var (BVL (wf) FF)) size b for all be [n], t N, and all square matrices B IRPxP,
decreasing function of mini-bater
Theorem 7. Fixing initial weights wo, var (Bgy Fo) is a decreasing function of mini-batch size for all b [m]. N, and all square matrices B Rexp.
Proof of Theorem 0
We induct on to show that the statement holds. For 0, we have
(BV L (w?) 0 for any matrix B. Suppose the statement holds for 2 0. Note
that from
1
VL (w))
2i [ wt y1)
n
Σ
24 (x, (w - argi-n) yi)
n
1
Σ
2i (x.² whi y1))
Σ
n
xa 9:
:-1
VL (w}_) CgCggin
19

---

Under review as a conference paper at ICLR 2021
have
var (BVL (w,) (Fo)
var (BVL (w) 1 Fo)
|BVL L (w/_1) - 0r BC 84-11 11" Fi) - [BVL - -
E
E
[BVL (w) 1) 1- 20r (BVL 10
E
BBVL (w 2 Fu + 0
E
BC 11 アタ -
20, (BVL BC (CG-)))))
| [ [VLL (w a.
F
| 2
BCVL, (w) -IBBVL (w( + ||BBVL (wp- 1)
E BVL ( )) Fu + az
20, (BVI. (w)_1) BCVL (w)- Foo
E [VL (w6-1) -o, BCVL (w) [ (3)
E [BB af
BCVL, (w -BCVL (w Fo
-
n 1=1
-18BB11-0003((((,111111 1 Σ
var (bL( (5) + a
E
|BCVL; (u²- -E |B (w (
var a (w) Fo) + afcs n (ナ
Σ E
BCVL, (wi_1) - BCVL, (wo-
(4)
where (3) is by Lemma
By induction, know that the first term of 4 is decreasing function
b. Taking A, = BC, A, = BC, A = 0, k E [n]] in Theorem we know that
E IBCVLL (w BCVL, (wi))) Fo
is also decreasing function of b. Note that ai decreases increases. By Lemma - that (4) is a decreasing function of b and hence we have completed the induction.
we learn
Proof of Theorem 7
We have
var (Bg. Fo) =
Bgg
1
Fo -
E [Bg Fol
E
Bgg
2
Fo
E E Bg F
BVL (wmt) ( ( Fo 1)
c
E
BVL, (wx))
Fo -
BVL (w)) || [o] [B (w Fo11
+ E
Σ
E
BV Lu (wx) BVL, (w²) / [o] + var (BVL (m+) (Fo)
n²
Taking A, B,A, B, = kE (t..j in Theorem I we know that
E
LVL, (w?) BVL, (w?) Fo
is decreasing and non-negative function of for all i.j [n]. By Theorem 6 we know that
var (BVL (w?) (F)) is also a decreasing function of b. Therefore, var (Bgf Fo). as the sum of two decreasing functions of b, is also decreasing function of b.
20

---

Under review as a conference paper at ICLR 2021
Proof of Corollary 7
Simply taking B Ip in Theorem yields the proof.
B.2 PROOFS FOR RESULTS
3.2
Remark. We often rely the trivial facts that 21 2 Inzi and mr aan =raz Intzzq
Lemma 5. Given multiplicative term of parameter matrices {unny uu vs RP,iE [mil]ufAj
A3 E Rexp, j [na]] and constant matrix (Tp) such that deg( un of M) 2 1, we have
tr (M) = = M'u1,
where M' is multiplicative term of parameter matrices {ul'y' 14,25 R,, RP.ie [ni]] U
{A₃ : A, E RPxp, e [m2]] and constant matrix {1p) such that deg(M) = deg(M") + 1, de M) = deg(Aj; M'), E [m2], deg( u ;M) = deg(u,)) M), e [2:n|] and deg(urof of; M) = deg(urrf; M') + 1.
Proof. By the definition of multiplicative terms, we know that there exist two multiplicative terms MI, M2 parameter matrices {u.uf us, 02 E IR,, constant matrix (Tp) such that
18,ic i [m]]] u {A₃ A₁ e RPxx, [me]]] and
M Minnii M2:
where deg(M) deg(Mi) + deg(M2) 1, deg((eg))))) M) deg(Aç; Reg(AA;MM)) M1) deg(A) (Aj; M2).j [nz], deg uevf M) deg(uref Mi) + deg(urf; M2), ; E [2:mı] and deggurvy; M) oeg(urof M1) deg( (uoo M2) + 1. Therefore have
tr (M) = (Muuuy) M2) ir (vT M2 vi 2i" M2
Note that M' = M2 satisfies that deg(M") = deg(Mi) + deg( M2),, deg( M') deg(A M1) + deg(A) (A; M [m2], deg( u M) = deg(u,u? T. Mi) + deg(u; :f: M2), i M2) [2:n]] and deg(uro] M") = deg(us (eg(uv] M1) + (2)) + 1. We have finished the proof.
The following two lemmas focus on the expectation of the product of quadratic forms of the standard normal samples. Lemma O focuses on single sample while 7 focuses on the same form with b i.i.d. samples drawn from the standard normal distribution.
Lemma 6. Given matrices A₁ RPXP, j e [m- 11, we have
Nm
Eg-N(0.T,)) [xx Ayæx" A2 ... Am 1100" 12."
tr (M&) Ma.
where Nm and n4, E [Nm] are constants depending on m and (Mi, k E n0 i E [Nm] are multiplicative terms of parameter matrices {Aj.] E [m - 1]} and constant matrix {Ip}- Fur- thermore, for every i [Nm] we have En deg( Aj; Max) = 1, [m 1] and therefore
>x deg (Mx) = m
Proof. See Magnus 1978).
Lemma 7. We given matrices A₃ RPxP .j [m and random vectors x,, E [b] indepen-
dently and identically drawn from NO, Ip). We assume that the multi-set S { 1/ 1/: : [m]] satisfies that for every S, i is an element of [b] and the number of appearance of in S is even. Then
E2,~N(0,1,))
Ti T AI 2.
T Ag. ...... Am Am-
tr (Mix Ma.
(5)
where No and n; are constants depending on (and independent of and Ma k [0 n.) i [Nm]] are multiplicative terms of parameter matrices {A,r] E Im - 1}} and constant matrix (1p).
Furthermore, for every E [Nm]. we have Σ deg(A Maa) 1.j [m - 1] and therefore Dni deg (Mx) m
21

---

Under review as a conference paper at ICLR 2021
Proof. Let 3 [6] be the number of appearances of in S, which are even by assumption. We induct on the quantity N Σi (3: * 0}.
For the base case of N = 1, all elements in the multi-set S have the same value. generality, we assume 2
Without loss of
1.j [m]. Then
Ex, ((0,1) a A13i, 2 Am 1ii. 2
B Be.-W(O.L)) [ Aaai ... Am Am-10111
which is the statement of Lemma G
Suppose the statement holds for N > and consider the case of N + 1. Note that #7 AgRig-1
_ Ajæi is scalar so that we can move around without changing the value of the expression We distinguish two cases.
Let i * En. Without loss of generality, we assume it = 1. We can always change the order of N Af e [m - 1] (and flip it to be aT, Azzy if necessary) such that all vi i's appear in the form of ar
on o A T. A2... Am Iim x = 21 (x) Arzin 2 A2... Am- Am-11im)) o
x Āizia Ag... Ag - winy Azg Tor
[b]] 3 * σ1 and A are multiplicative terms of parameter matri- ces (auat T u, v En,ne(2:b)
where Fo E {2i, i E
[2: [4, ; [m 1]} and constant matrix {Tp} such that
deg(=wa7) As) = m 0 - - and Σ deg(A, (A;; Ax) = 1,j E [m -
Applying Lemma and the law of iterative expectations, we have
Ex, (.,,)((11,)) Ti x A x1 Am-
= Exy me
x T Aiaia Ag... - 211 Ag tx..
N
Eaggy aat
Σ
Ir (Ma) Mi
E
Exg iss
tr (Ma) Mio
À To
where Nm and n, are constant depending on m (and independent of b) and Mik, h [0: i [Nm] are multiplicative terms of parameter matrices {Annj E 1 1] and constant matrix {4p}. Furthermore, for every i E [Nm]] we have Σx= deg(A, Ay; Max) = 1.j e 2 1] and therefore Σ deg (Mxx)
1.
Combining the definition of Az's, is's we know that Ma, e [0: n.:. de [Nm] are multiplica- tive terms of parameter matrices (2uzī u, v E [2: 6]} u (Ay : j E [m- 1]} and constant
For example, we can rewrite
zi a Al J Arti Ii (I T ALTi2) [ A2tis Agtis
T = Wij I/g
J A2tia (*f AL
Di Th (If Altiz))
T 5/g
1 A2 (2 A
T
xíg
For example, we can rewrite
of A223 = A 22 4 Aς .noe A Azāaz,
21 A A Astras =x (1 A123))
{ Asai x2 31 (xf (e
12 Auars god Asaa) x2
where A₁ = AΙ A2, Ag = As and ã = uy. Besides, m = 4, = 4, thus the degree of a.we in all Ås
sum up m 41 1-1-1
22

---

Under review as a conference paper at ICLR 2021
matrix {Ip) such that for every i [NNm., we have ve(:b] Zk= deg(E Mak
m Pt 1 and _ k- deg( Aj; Ma2) = 1,j [m - 1).
Applying Lemma 5 for every & [0: and every [. there exists Uik Vik (or j [2:0] and multiplicative term M of parameter matrices {aue] : u, 22:b] u [AA:] th:je [m 1]] and constant matrix (Tp) such that
tr (M+x) = To My.
Therefore, we have
14
Ne IIll k=1
I ur (Mas) Mi
Ao J.t.
1l (LuM))))) Mo 2od Mn 1=1
(u Maers) ox = Us.
Note that for every i E [. have
nz
deg(a;; A2)
deg(z,; Ma) + deg(2;; Ma) deg (m::Ag) ) deg(zs; deg (I:: x!7
and for every E Im
11. we have
deg( M1+) + deg(Aj; + deg (Aji 1.
e [N U, has the form of 1. 2 Amd but there no appearance of I1. Here "i, {xg.je [ b]}, and A, E [0:m] are multiplicative terms of parameter matrices {A,.jE 1mm 1]} and constant matrix {4p) Furthermore, for every j E [m 1]. we have Ix= deg(AA; - 1. Note that here we use the liberty of adding identity matrices if more than two consecutive I's appear. Since have reduced N 1 by one, we can use induction on 2;; xTT Ar xT Am-1 o% and finish the proof.
In other words, for every Aoz; of "y Avæ xx Am-
The two constant matrices Åo and Am do not change the result of expectation since E (AoXAm)) Amr) = A
If i'm without loss of generality assume, i and E (note that all oif AfTijmn E [m - 1] are inter-changeable and there is at least one element in S that is not equal to i1). We change the orders of 27 AyDist1:J e [m 1] (and flip it to be " Azæz, if necessary) such that all 's appear in consecutive form of xma::
vi, x A1 A2. ... Am Am- xh Ii, (a A1 A2 A2.. Am Am- Tz
{ Ag Fa Ān ... Ag
As 82
of
where Z1, 82 E (or E [b]] * Ii and A are multiplicative terms of parameter matrices (aaad u. E [2:b]] (Aj: j [m 1]} and constant matrix {4p) such that
31 2
Σ
deg(I xr Au) =
2
81 and Yi deg(Aj Aj; As) 1, Im 1]. The remaining reasoning is the same as the previous case.
23

---

Under review as a conference paper at ICLR 2021
Remark. If one of the 83 numbers of appearance of 2z. [6] is odd, then it is easy to that the result in is the zero matrix.
As pointed out in the Section I the difficulty of studying the dynamics of SGD is how to connect the quantities in iteration with fixed variables, like initial weights W Wo,2 and mini-batch size
b. We overcome this challenge by the following two lemmas. Lemma 8 provides the relationship between gº,, = 1,2 and W = 1,2 by taking expectation over the distribution of random
samples in B3. Lemma shows the relationship between W f,i i = 1,2 and 06-1, = 1.2 using and 2
Lemma 8. For multiplicative terms Mi, E [0:m] of parameter matrices bi.1.hh.) and constant matrices {Wh,, W 1,2: WT. W2) with degree di respectively, we denote M 1T,,, tr (M) Mu and d Ei dz. There exists a set of multiplicative terms {M i [m=]..j E [0: mamm k [0: of parameter matrices (Wh,, W and constant ma- trices {Wr, W such that
1 Nayz
E [ No + N.T""
where N2. -i=1 1ー (M+) M & []: Here my, mak are constants independent of b, & 3d + En (deg (W7,: M,) + deg( W1 M.)).
and 21-0 deg (M
Lemma 9. For multiplicative term MG [0: m] of parameter matrices {Wh, Wh,2) f,1 1,23 and con-
stant matrices {WT, W of degree di let d 2dot..-dm There exists set of multiplica- tive terms {Ma,i [0: m] k [d]] of parameter matrices {ot.1: 98,2 and constant matrices Whh, W Wf, W } such that
H tr (M.) Mo
Σ I tr (Mx) Mok.
where En deg (Mx)= & d.
Proof of Lemma g
By and ( we have
M tr (M)) My
8d Σ
tr (Mai MRo-
(6)
where each MRi k famaz [b]} and constant matrices {W w Wt}. Let Mr = I. tr (Mki Mao. k
[b e [0:m] is multiplicative term of parameter matrices
[bq]. We split set Me h e [b$]] into disjoint and non-empty sets (equivalent classes) Si...., San such that
1. for every i E [nar]] and every M1, .2 Sü, we have E [Mi = [Mg
2. for every i,j E [nar], * j and every Mi S; and M2 E Sj, we have [M] *
E [2]]
Note that UMMS = Mr } Let Ma Sa represent the equivalent class SN (it can any member of Sc). For every E [n we can always write |S|| e;,o + ct, + , such that eij N, < b,j [0: d) (actually e,," are the digits of the base-b representation of (.))). Then
24

---

Under review as a conference paper at ICLR 2021
have
ba
E
M
Σ M
F4
E
Σ
Cio + + ... + Egab") M,
bd
F
k=1
1 nar
(ci, ci,b + ... + ca )
E M F+
bd
Cid
+ +
1 + ... + ci,0 gd d
E
M, F.
It is important to note that n.M the number of different equivalent classes, is independent (and so as Mr F includes a finite number of weight matrices W and W with degree less than or equal to 3d + -i (deg ( (W21; M.) + (e: f,2i Mc)) (see Lemma 7). Thus the number of partition sets is bounded by a quantity independent of b.
of b. This follows from the fact that each E Mr Fb
Note that each Mgi can be represented as
MAi A a0
T T A .......A--110 1 Ftiss Ad: Fii, zin
for some matrices , that are multiplicative term of parameter matrices { {Wi, W W } constant matrix {4p} (we stress again that some A matrices can be identities, based on the definition of multiplicative terms), and w ".... stisa ( Fe,0]- We have
(Mkx) = tr (Ab Alt AKi... - 3i,sa All
ki 22,00
T A 2,,,, " A
AApi - 1
For every A E [bt]] we have
tr (M&) MRo
A4
A
A.,AA 1
A A0_
T
A
A..
1 1,
21,600
A0
T
A
22,5a,
ki
A
4ki
4
2th
Att...
A
1
A
A
2,800
AB
60
Ad
which can be rewritten as
[1 9-1
d
M = I tr (Mai M10
Af
A0 A0 A
A
2ti, It,500
Note that the randomness of each Mr given F3 only comes from the randomness of Te i's, i.e. for all l [dd] we have
d
E
M Fbb
E
A Ab"
Afiee [ st Afarss
2T. Ap'
(8)
ri n
= Σ
tr Mi
Mr
25

---

Under review as a conference paper at ICLR 2021
where the last equation comes from Lemma] 7 Here n n i E |nt k [64] constants inde-
pendent of b, Mt is are multiplicative terms of parameter matrices {W W.,,. Wj) and constant matrix { such that for every [ J. we have
deg W WW: M,;))
(9)
and
Σ deg ( My) + deg (W/, M -=0
d+ Σ
(deg (W + deg( 2,23 M,)) (10)
P=0
These degree relationships can be observed from ① and the fact that each o¹, or lh.1 con-
tributes W and of Wh:, W in Iī tr My
M Note that W, W W
W For every i E [ntr]. if we replace all appearances of W. I MM M and expand all parentheses of (W?, WF)+), we have
n
I
M
(11)
where Mhl's 's are multiplicative terms of parameter matrices { (Wh,, r,1, Wh2 and constant matrices [Wr, W such that
deg Wh My) deg (wh, M7. < 3d + Σ (deg (W M, deg(W 1,23 M,))
where the inequality comes from and 10 and the fact that each sh,1 or وf contributes 2 (12)
degrees in the form W 1,2
Wh or W+ respectively.
Combining
and
we have
E
M
F
=
Σ
ck,d + ekdd
ek.0
bd
IE
M
F
2t
T
Σ
ΣI 1111j-1
(M
Ckid + ckid ... eN.0
tr
yd
= No + No
+ No No
where
n 24 x.
aM
N,
Ck.d-r
ΣΣΠ tr
M
(13)
Note that all constants in 13 are independent of b and combining with (12 have finished the proof.
Proof of Lemma 9 Simply using the fact that Whb = Wh ong 1,6, = 1,2, if we replace
each Whb in the left-hand-side of (13 by W angp-1. and expand all the parentheses, then each M,, m] becomes the sum of 2d, multiplicative terms parameter matrices {9L1, gi and constant matrices [WA Who W W with degree at most d,- As result, IT. tr (M) M becomes the sum of 2d terms in the form of [1. tr (Mx) Mot where deg (Ma) 2d, and
therefore E deg (Mx) I1.0 2d. d.
26

---

Under review as a conference paper at ICLR 2021
Proof of Theorem 3 We use induction to show this result. The base case of t 0 it is the same as the statement in Lemma
Suppose that the statement holds for é > 0, and consider the case of t + 1. By Lemma 8 there exists a set multiplicative terms (M 2. [m++1,*], E [0 m4+1, [0:d]] of parameter matrices (W (+1,1) W W and constant matrices {WI. W2) such that
1 E [M No+1.0 + No+1.15
1
(14)
Σm L I1 L (Ma)) MR M4 E [0:d]. Here m4+1,k, 1,K,, are constants independent of b, and Ex=o deg (M & 3d + dd'
where 1
For each i [m++1.k] and each & [0: d]. by Lemma 9 there exists set of multiplicative terms [Maaj.kad e [m++1ak], e [de,s.a]] of parameter matrices 1of,1.99.2 and constant matrices (W,, Wh 1,2 Wf, W such that
deii m.t1kkki
II tr (MELLad) M Mi+1.0 J-1
Σ
II
tr (MaaARa) Mes.0,
(15)
where dr,i, = 2-5-0 1.
hid (dey( (a wegge:::ddetgtWW..M..ss is a constant independent of b and
Σ
deg (Maig.t)) & 3d + d'
(16)
and
Σ
(deg (W,,;; M.....)) + deg (W,,; Mr M.jgy.))) & 3d + d'
(17)
Combining (14 and 15 we have for every & [0 d]
m.+1. deat
m
N1+1
Σ
Σ
Lr (Mesja.)) M1s.0...-
(18)
j-1
Note that
[M E [
[M
M
E [ Fol E IN
+ E E
mc+1.0 de,io
Σ
Σ
E
I 5-1
tr (Mess.o)) Me.i00.0 Fo +
de,ca
Σ
Σ
E
tr (Mesdad) Mesonn Fo
+
b
Mi+1. de.
Ma+1,a.4
Σ
Σ
I
tr (Measaa)) M Mes, Fo
(19)
bd
and each My..J.k. is a multiplicative term of parameter matrices {gi.1 gi.2 and constant matrices (WA,, 1.1. W2 W., Wx) such that the degree is at most 1. Therefore, by induction, for every i, k. have
E
II tr (M.......) M...o,k. Fo NE.KA NARRAI ....N...... bas
(20)
where de & d' + (3' - 1)(3d + d') and No,,..,..."" N...i, are sum of multiplicative terms of parameter matrices (Wa., Wh,2) and constant matrices (Wf, W3) with degree most d 3'.
27

---

Under review as a conference paper at ICLR 2021
Combining (19) and (20). we can rewrite
E [M| = No + Ni b + + N Na
in the same form in the statement. Here q < 3q. 1(3 1)d 133 1)dd and Σn deg (My) 3 x (3d + d') 3 (3d + d') follow from (T6) and (17.
In conclusion, we have shown that the statement holds for + 1, and therefore finishes the proof.
By changing the role of parameter and constant matrices in Theorem we obtain the following corollary.
Corollary 2. Given > 0, for any multiplicatiie terms M,, [0: [0:m] of parameter matri- ces {Wā1, 21, W and constant matrices {WT. W2) such that Σm=1 deg (W",; M)) = d and deg (We; M) = d', we denote M = - tr (M.) Mo- There exists set of multiplicative terms {M i E [m e m00m & e [0: q]} parameter matrices {Wo, W and constant matri-
ces (W*, W such that
1 + N b E [M = No+ NoN
where N3 E S
mee tr (M+) Mr: k E [0:4]. Here my, mis and q & 3 (d+ 2d') are con- stunts independent b, and 2リ- deg (Mz) < 3 (d )
Proof of Corollary 2 We simply note that M can be written as the sum of at most 2 multiplica- tive terms of parameter matrices (W2,, f,1, W f,2: W W 3 and constant matrix (To). Then we apply
Lemmas 8 and 9
iteratively in the same way as in the proof of Theorem 3 finish the proof.
Proof of Theorem 4
We only show the case for It, since the proof for 9r,2 can be tackled similarly.
Note that
(1
1 b
var (F)) var
Σ Wh "W 22 T Fo
Σ
var (Wh.)
T "Woa Fo)
82 1=1
1
(w:p2" T' " xe, c Fo
var
b
b
1
E
W 2,2
T
W x4, o
Fo -11E
t,2
T "w W 2271 [oo 1
1
E
tr
T o Weewo "a [oo
E
[wote T " a T [Fo] 11"
b
1
E
E
tr
21, " w Who w T "W x
| Fou
| E
Who T "Were | [Fo 111²)
b
1
(p
2)tr (w!" w w "w |5o]
E T I m.m Fol [1()
E
b
1
(p+
2)tr
E w Wh 1, 2, "w Fol)
w I
E wHo I mi! Fo 11²)
b
1 (p +
(++
2)1r
E Ww Wh 1,2 w "w (o(
E
W T w. Fo |1"
b
Here we have used the fact that E (NxT Axx") (p+2)tr (A). By Corollary we know that there exists a set of multiplicative terms { MM i e [m=]]. [0 mes] k E [0: g]} of parameter matrices { {We,, ] and constant matrices (WF, W such that
tr E w W.e T "w]
1
2.2
70
n 0
+
79 T
(21)
28

---

Under review as a conference paper at ICLR 2021
where 7/ Σ >> Tyy=o tr (M k []: Here my, mai and q & 6 3' are constants indepen- dent of b, and : deg (M,) < 6.3". Note that W W are fixed, and we have Nk. k e [0 ql are constants independent of b.
such that Similarly we observe that there exist constants o' & 2- 3 and 2/: A: E [0:
Wh W Fo ||| = % + 1
1/ ba
(22)
By defining 7/2 0, 9 and 7/2 0, g'. and combining
21 and (2)
we have
var (o) = 5 (p + 2)tr [E I W W/: T "w Fo
W T W Fo
11]
p +
Nu + ni 1
t... + 70b9 (1)
1
1 (16
1 + rii + + 1/1 bs
b
b
max n.n')
((o+ 1)m -1) T
Note that "/k 's and 7/6 's are all constants independent of b, and max (9,4) & 2.3++1, This completes the proof.
Proof of Theorem
We first show that in
1
1 B.
var (o) B17,
we have 31 > 0. If r 1, the statement obviously holds. Let us assume that the statement does not
3,br-2+ 2.... + B, 0 yields hold for r > 1, i.c. 31 < O. Taking large enough such that - +
1
var Fo) br (3,0--1 11 + -2b0-22
+ 3.) < 0,
> 0. Therefore, we have $1 0. which contradicts the fact that var (gf., (o))
Let by be large enough such that for all b > bu. we have 3,6r-1 + 2,325 4... 2 We
denote f(b) $1 5/ + B2 अ ... + 3r O. For all b > bo we have
(((b) b+I (B,0 1 2,320
+ r3,,) $
Therefore, for all b > bo we have (var (ot. Foo) (Fo. is decreasing function of for all > bg.
F ((b) + 5 ((b) & 0, and thus
-
B.3 EXTENSION TO DEEP LINEAR NETWORKS
The extension from two-layer linear network to deep linear network straightforward. Here we only provide the ideas on how to translate the proof of two-layer network to d-layer network, but not the strict proof. For simplicity, we remove all superscripts of matrices in this subsection.
Assume that the d-layer linear network is given by s(x; w) WaWa-1' W2W11, where W,, E [d] is the parameter matrix on the i-th layer and w (W1.... Wa). The population loss is defined
as
[(w) = E (O-.
W Wi W w
Similar to and
we have
1 W W2, W ... W7 Wisnel?
Ot.k
V Wi.6
WT 1 ... -hr (Wa... Wi W W ( 2e, WTi --Ff- k ke [d].
29

---

Under review as a conference paper at ICLR 2021
We denote W/ Wi.d.... Wt,1 W W The remaining are all the same as the proofs in Ap- pendix B.2 except we should replace all appearance of (W,,2, Wi,1} to (We,d, Widd-.... W.,1 and all { WT] to { {Wa, gt.k still the sum of multiplicative terms parameter matrices { and constant matrices (W... Wi,,, Wa.
d. -........ W We can do this because the stochastic gradient
d: WF) so the Lemmas in Appendix B.2 still apply.
In conclusion, and without the constant term. By the same approach in the proof of Theorem we can show that the variance is decreasing function of the mini-bateh size b.
again represent var (gt. & (o)), k [d] polynomial of 1 with finite degree
30


Reviews:
[{"review_id": "skA84bWPS3", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper shows that the variance of the gradient has an inverse dependence on the batch size in linear networks, subject to the knowledge of the initial weights. The main novelty of the paper is the computation of an exact dependence between batch size and variance of the gradient in the linear regression setting. In addition to that, the authors conduct a lot of experiments involving non-linear networks and real-world datasets, that show the inverse dependence of the variance of gradient and batch size throughout the training.\n\n\nMy major concern is that the authors don't provide an application of their theorems, i.e. a setting where the exact knowledge of the variance of the gradient is useful. It will be useful to know if their theorems help tighten the convergence rates of mini-batch SGD in convex regression or improve generalization bounds for mini-batch SGD. Or can the authors say that their theorems help to tune the learning rate when going from large batch SGD to small-batch SGD? Without any such application, I believe the theorems are incomplete.\n\nMy other concerns are the following:\n1. The authors prove that the variance is inversely dependent on the batch size. However, their theorems also imply the variance of the gradient keeps on increasing with time. Hence, a convergence theorem with an appropriate learning rate decay will be helpful. I see in the experiments the authors have used the learning rate as $1/t$. A justification will be really helpful.\n2. Most of the experiments in the paper have been conducted on non-linear networks and non-gaussian data (like MNIST and YELP). Hence, a rough idea about how the theorems can be extended to ReLU networks, even in the over parametrized case, will be highly appreciated. \n3. The experiments show that small-batch SGD can reach lower training loss. Can the authors show why the training loss decreases with increasing variance in gradient, at least in the linear regression case?\n\nMy scores are on the lower side mainly because I believe the authors need to show some application of their theorems. I am happy to discuss this with the authors and other reviewers during the discussion period.\n\nMinor concern:\nThe abstract and introduction have writing issues. However, I haven't taken them into account in my score. I would request the authors to improve those sections soon.", "extra_scores": {}, "extra_text": {}}, {"review_id": "XcppNu1oxDU", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studied mini-batch stochastic gradient descent (SGD) for linear regression and linear neural networks. The analysis show that the variance of the stochastic gradient estimator is a decreasing function of b (mini-batch size). The authors also empirically validate their observations on MNIST and YELP review datasets.\n\nI think the main result of this paper is not interesting. There are some comments.\n\n1. The connection between mini-batch size and the variance of gradient estimator almost will be considered in any paper about SGD and I do not think this paper provide any novel perspective.\n2. The analysis in section 3.2 discuss the extension on linear network, however, we are more interested in nonlinear case. The current analysis looks cannot be extended to nonlinear model directly.\n3. The experimental section also should be improved. The structure of the network should be introduced in details. It is prefer to validate the proposed theory on more difficult task. MNIST and YELP cannot reflect the power of deep learning model.", "extra_scores": {}, "extra_text": {}}, {"review_id": "I-xDU_Hn5lW", "reviewer": "AnonReviewer4", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary and contributions\n\nThis paper tackles the problem of the impact of mini-batch size on the variance of the gradients of SGD. Unlike most work on the topic, it does not study the variance of the gradient conditionned over the last iteration, but rather the absolute variance of the gradient conditioned only on the initial point.\n\nThe paper restricts itself to linear models, either with least mean square regression or deep linear model with 2 layers (which would be non convex but still linear).\nThe paper shows two results. For linear models, the variance of the gradients conditionned on the initial point is decreasing with the batch size b.\nFor a deep linear model, the result shows that the variance of gradient is a polynomial in 1/b with no constant term. Therefore, for b >= b0 for some b0, the variance is decreasing with b.\n\n## Review\n\nThis paper tackles the problem of the variance in SGD from a novel angle, namely the total variance conditioning only on the initial weights, and not on the previous iteration as usually done.\n\nThe result on linear regression seems very natural, and the proof is done elegantly (I checked up to the proof of Theorem 1).\n\nI haven't checked the proof for the deep linear network. The result is more interesting than in the linear regression case because the 2 layers linear network is non convex, and therefore, one can imagine having multiple local minima and 2 different trajectory starting from the same point to diverge at some point, which could lead to drastically different gradients. This is not really possible for least mean square with a full rank hessian, as there is a single optimum, and in any case all trajectories will end up at the same place.\n\nHaving a small total variance conditioned only on the initial point means that somehow the trajectories for different samplings cannot diverge too much. Of course, if the batch size goes to infinity, one is doing gradient descent, and all the trajectories are exactly the same which is compatible with theorem 4. I think the proof technique is interesting to be able to bound or study deviations between trajectories. However, in its current form it is non practical as the bound is very complex and in particular can increase as the iteration increases. In essence, the results says that for any iteration, as B increases, trajectories get closers, but it does not say that they will stay close for any iteration number.\n\nI must say that while the theoretical part is doing a good job, the experiment part in Section 4 is quite poor.\nThe author tries to open up to the problem of generalization. \nSpeaking of figure 1 b), the authors say that the validation loss improves with the batch size, i.e. the generalization ability is better with large batch sizes. This however contradicts previous work that have been mentioned by the authors in their very introduction. The authors do not comment on this contradiction. My own conclusion is that deriving general results on generalization from MNIST is a perilous exercice, if not plain wrong.\n\n\n## Conclusion\n\nOverall, I would say that this paper is just above the acceptance bar, because the theory holds up well and could be of interest for finer analysis of the dynamics of SGD, and in particular of different trajectory starting from the same point (how quickly will they diverge?), although doing so would require significant extension to the present work.\nThe authors try to open up to the problem of generalisation but fail to a proper theoretical link with their own work, while their experimental results are obtained only on MNIST and therefor subject to caution (and in fact contradict previous work).\n\n\n## Notes and remarks\n\n- In the abstract, the authors write \"for deep neural network with L2 loss, we show...\". This is not true, the result is only for two layers deep linear network.\n- Definition 1: the definition of the degree is a bit hand wavy, there should be some minimum over all possible decomposition of M. This is especially needed if for some X, both X and X^{-1} belong to \\mathcal{X}.\n- page 18, at the top: \"We denote A = \\sum ...\", A has already been introduced before end of page 17, but with x_i x_i^T instead of C_i (both are equal of course).\n- Section 4, comments on the loss of Figure 1: phrasing \"loss is decreasing with b\" is ambigous or wrong. For the training, the loss increases with the batch size. I guess the authors meant that the loss worsen with b (as low is good), but worsen != decreasing.", "extra_scores": {}, "extra_text": {}}, {"review_id": "3BoKGuFN-lD", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the variance of stochastic gradient in SGD conditioned on the initialization point. It shows that the variance of stochastic gradient is a decreasing function of minibatch size for linear regression and deep linear network. Compared with previous works that show similar the results for one step in SGD, the results in this work only rely on initialization point. \nAlthough the technique in this paper is different from previous works,  the setting in this paper can only handle linear model which is simple and limited. Besides, only showing that the tendency of gradient variance as minibatch size changes is not enough,  and it is better to give the decreasing rate. So the impact of the refinement on the estimation of gradient variance is not clearly stated. If the authors can provide some examples to show that the analyses in previous works are too coarse to give the wrong direction, the impact of this paper can be enhanced. In current version of the paper, I can not get any new insights from the theoretical and experimental results because there have been many existing works that can tell the relation (or tendency) between gradient variance and minibatch.  \n\nDetailed comments: \n1. The notation \"var()\" is not rigorous because the gradient is high-dimensional. It should be replaced by the covariance matrix of a random vector.  Given that the gradient is a high-dimensional vector and the \"variance\" is a matrix, it is not clear that what the meaning of \"the variance is a decreasing function of minibatch size\" is. Does it mean that all the diagonal elements in the covariance matrix is a decreasing function?\n2. It must be clarified that which distribution the expectation and variance are taken over. For example, var(\\nabla L(w_t)) should be written as var_{w_t}(\\nabla L(w_t)).\n3. In the setting of deep linear network, the distribution of data is assumed to be normal distribution, which is strong. Is this condition necessary in the proof? \n4. The experiments are not sufficient to support the claim. Can you provide results for more selections of minibatch sizes, more datasets and more NN architectures?\n5. In future works, the authors point out many interesting research problems, but it is not clear that why the techniques and results in this work can help to solve those problems. Can you provide connections between the main results and the future work? For example, which technique can help design better variance reduction method? Clarifying this will be helpful to evaluate the impact of this paper.\n ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work analyses the impact of mini-batch size on the variance of the gradients during SGD, in the context of linear models. It shows an inverse relationship between the variance of the gradient and the batch size for such models, under certain assumptions. Reviewers generally agree that the work is theoretically sound. However, all reviewers believe that the contributions of this work are limited. This concern was not adequately addressed during the discussion phase and led to the ultimate decison to reject.
Decision: Reject