Non-Gaussianity of Stochastic Gradient Noise

Abhishek Panigrahi Raghav Somani Microsoft Research, Bangalore, India University of Washington, Seattle, WA, USA [email protected] [email protected]

Navin Goyal Praneeth Netrapalli Microsoft Research, Bangalore, India Microsoft Research, Bangalore, India [email protected] [email protected]

Abstract

What enables Stochastic Gradient Descent (SGD) to achieve better generalization than Gradient Descent (GD) in Neural Network training? This question has at- tracted much attention. In this paper, we study the distribution of the Stochastic Gradient Noise (SGN) vectors during the training. We observe that for batch sizes 256 and above, the distribution is best described as Gaussian at-least in the early phases of training. This holds across data-sets, architectures, and other choices.

1 Introduction

Stochastic Gradient Descent (SGD) algorithm [17] and its variants are workhorses of modern deep learning, e.g. [2, 4]. Not only does SGD allow efficient training of Neural Networks (a non-convex optimization problem) achieving small training loss, it also achieves good generalization, often better compared to Gradient Descent (GD). While striking progress has been made in theoretical understanding of highly over-parameterized networks in recent years, for realistic settings we lack satisfactory understanding. In particular, what property of the distribution of the small-batch gradients arising in SGD is responsible for the efficacy of SGD? Considerable effort has been expended on understanding SGD for Neural Networks. In particular, there have been attempts to relate SGD to a discretization of a continuous time process via Stochastic Differential Equations (SDEs), e.g., [12, 15, 23, 6, 7, 24]. The Stochastic Gradient (SG) can be written as Gradient + Stochastic Gradient Noise (SGN). [12, 6, 7, 24] assume that SGN is approximately Gaussian when the batch size b is large enough, assuming the Central Limit Theorem (CLT) conditions, arXiv:1910.09626v2 [cs.LG] 25 Oct 2019 which we discuss in the next section. [15, 23] add explicit gaussian noise to stochastic gradients in each iteration. One can analyze these algorithms by approximating them with the continuous time process in the form of Langevin diffusion, and one can show that the Markov process exhibited by the continuous time iterates is ergodic with its unique invariant measure whose log-density is proportional to the negative of the objective function [18]. Therefore, SGD can be seen as a first-order Euler-Maruyama discretization of the Langevin dynamics [7, 11]. Since the gradient vectors are very high-dimensional, the number of samples required in the ap- plication of the CLT, so that a multivariate Gaussian distribution provides a sufficiently accurate approximation, may be very high, raising doubts about the applicability of the CLT for SGN com- puted using small batch size. For SGD training of neural networks, the assumption of SGD noise being Gaussian-like was contested recently in [21]. They argue that the SGN instead follows a stable distribution with infinite variance obtained by the application of the Generalized Central Limit Theorem (GCLT). They provided experimental evidence for their claim. Furthermore, they provided

Science meets Engineering of Deep Learning (SEDL) workshop, 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada. arguments based on the properties of the associated continuous time process which can then be viewed as an SDE with the stochastic term being a Levy motion.

2 Preliminaries

The problem of training a Neural Network is essentially an unconstrained optimization problem that usually constitutes a finite sum non-convex objective function. Function f : Rp × Rd+1 → R denotes the composition of the Neural Network and the loss function. On learnable parameters w ∈ Rp, and input (x, y) ∈ Rd × R, function value f(w;(x, y)) measures the accuracy of the underlying n d
n Neural Network on (x, y). For a set of n training data samples {xi, yi}i=1 ∈ R × R , define 1 P
1 P f (w; B) := |B| i∈B f w; {xi, yi} , and ∇wf (w; B) := |B| i∈B ∇wf(w; {xi, yi}). Then the optimization problem is given by w∗ ∈ arg min f w;[n]
. (1) p w∈R (t+1) (t) (t) (t) SGD update rule is given by w ← w − η∇wf(w ,Bt), where w denotes the iterate at time t ≥ 0, η > 0 denotes the learning rate, and Bt ⊆ [n] denotes a batch of training examples of size (t) (t) b ≤ n picked at time t. A random SGN vector at time t is then given by f(w ; Bt) − f(w ;[n]).

Gaussian random variables. Perhaps not as well-known as the scalar CLT, there is a CLT for vector-valued random variables showing that the appropriately scaled sums of i.i.d. random variables with finite covariance matrix tends to multivariate Gaussian in distribution. A useful fact about multivariate Gaussians is that if for a random vector X, all its one-dimensional marginals are univariate Gaussians then X must be multivariate Gaussian (see, e.g., Exercise 3.3.4 in [22]).

Stable distribution. For a real-valued random variable X, let X1 and X2 denote its i.i.d. copies. X is said to be stable if for any a, b > 0, random variable aX1 +bX2 has the same distribution as cX +d for some constants c > 0 and d. Stable distributions are a four parameter family S(α, β, γ, δ) of univariate probability distributions. There is no closed form formula known for these distributions in general, instead they are defined using their characteristic function. We do not include that definition here since we will not need it. Instead, we will just note the properties of stable distributions that are relevant for us. See [14] for more information. The special case of S(α, 0, 1, 0) is represented as SαS distribution, and can be considered as the symmetric, centered and normalized case of stable distribution. α ∈ (0, 2] is the stability parameter; we omit the discussion of other parameters. This family includes the standard Gaussian distribution (S(2, 0, 1, 0)), and standard Cauchy distribution (S(1, 0, 1, 0)). Interest in stable distributions comes from the Generalized Central Limit Theorem (see [14]) which informally states that the distribution of appropriately scaled sums of i.i.d. random variables with infinite variance tends to a stable distribution. This was the basis of the suggestion in [21] that SGN vectors follow a stable distribution with infinite variance. More specifically, [21] posited that the coordinates of SGN vectors are i.i.d. stable random variables. In parallel with multivariate Gaussians, there is a notion of multivariate stable distributions (see [14]) which satisfies the property that one-dimensional marginals are univariate stable. A more general hypothesis for SGN vectors could be that they are multivariate stable.

3 Method, Experiments and Observations

Gaussianity testing is a well-studied problem in statistics; we use tests due to [19] and [1]. These were found to be the top two tests in a comparison study [16]. Both of the tests assume the null hypothesis that the given set of scalar examples come from a Gaussian distribution. A small p-value implies that the null hypothesis is likely to be false. We follow the following steps to perform statistical tests on stochastic gradient noises (SGN).

1. We train a model using minibatch SGD and consider SGN vectors for 1000 independent mini- batches, which are of the same size as the training batch-size and randomly sampled with replace- ment, at multiples of 100 iterations. 2. We project each SGN vector along 1000 random unit vectors, and perform Gaussianity tests on the projections along each direction. We plot the average confidence value across the dimensions,

2 (a) (b) (c)

Figure 1: (a) Gaussianity tests on SαS variables; (b) and (c) training Resnet18 with 256 batch-sized SGD at l.r. 10−2, (b) Accuracy & loss curves, and (c) Gaussianity tests on projection.

as given by Shapiro–Wilk test [19], and the fraction of directions accepted by the Anderson Wilk test [1] and compare with the case when actual Gaussian data is fed to the tests.

3.1 Experimental Setup

We conduct experiments on different models, namely a 3-layer fully connected network, Alexnet, Resnet18 and VGG16 model, with and without batch-normalization (BN), and different data-sets, namely CIFAR10 and MNIST [8, 10, 9, 5, 20]. All the layers were initialized with Xavier initializa- tion [3]. The models were trained with constant learning rate mini-batch SGD, where mini-batch- size was varied from 32, 256 and 4096, on cross entropy loss. The learning rates were varied in 10{−1,−2,−3}.

3.2 Observations and Conclusion

We first perform sanity checks to see how well the hypothesis tests do on generated SαS samples. It can be seen in Figure 1a that for 0.2 . α . 1.8 the tests very confidently and correctly declare that the distribution is not Gaussian. As α → 2, we expect the distribution to become closer to Gaussian, and for very small α the tails become very heavy. Thus the behavior of the statistical tests on SαS distribution is as expected. We see that SGN looks Gaussian to statistical tests in Fig 2a for Resnet18 throughout the training for batch-size 4096. In contrast, in Fig. 2c the behavior is not Gaussian anywhere for batch-size 32. For intermediate batch-sizes such as 256, behavior was Gaussian at the beginning of the training, and later becomes non-Gaussian. We observe this behavior consistently across different architectures, data-sets, and larger batch-sizes. This is in contrast with the results in [21], which suggests stable distribution right from the start for batch-size 500. When the behavior is Gaussian in 1000 random directions, by the characterization of multivariate Gaussian mentioned above, this suggests—though does not definitively prove—that SGN vectors are multivariate Gaussians early on in the training. For batch-size 32 is it possible that the SGN distribution here is stable? This is unlikely because by GCLT we would expect this behavior at higher batch-sizes, but we observe Gaussian behavior. We briefly explain possible reasons why our results do not agree with those in [21]. They use the estimator of [13] to estimate α. The application of this estimator in [21] suffers from 2 errors: (1) In [13], the estimator is shown to give a good estimate of α under the assumption that the random variable is indeed α-stable; it makes no claims about the estimate when this assumption is not satisfied. [21] seems to tacitly assume stability. (2) The estimator assumes that the coordinates of SGN vector are i.i.d. This assumption is invalid in the typical over-parameterized setting.

3 (a) (b) (c) (d) Figure 2: (a) and (b) training Resnet18 with 4096 batch-sized SGD at l.r. 10−2; (a) Gaussianity tests on projections (b) Accuracy & loss behavior with training; (c) and (d) training Resnet18 with 32 batch-sized SGD at l.r. 10−2; (c) Gaussianity tests on projections (d) Accuracy & loss curves. For (a) and (c), legend is the same as the legend for Fig. 1c.

References [1] T. W. Anderson and D. A. Darling. A test of goodness of fit. Journal of the American Statistical Association, 49(268):765–769, 1954. [2] L. Bottou, F. Curtis, and J. Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223–311, 2018. [3] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedfor- ward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256, 2010. [4] Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016. [5] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [6] Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of nonconvex stochastic gradient descent. arXiv preprint arXiv:1705.07562, 2017. [7] 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. [8] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009. [9] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In F. Pereira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems 25, pages 1097–1105. Curran Associates, Inc., 2012. [10] Yann LeCun, LD Jackel, Leon Bottou, A Brunot, Corinna Cortes, JS Denker, Harris Drucker, I Guyon, UA Muller, Eduard Sackinger, et al. Comparison of learning algorithms for handwritten digit recognition. In International conference on artificial neural networks, volume 60, pages 53–60. Perth, Australia, 1995. [11] Qianxiao Li, Cheng Tai, et al. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 2101–2110. JMLR. org, 2017. [12] Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. In International conference on machine learning, pages 354–363, 2016. [13] Mohammad Mohammadi, Adel Mohammadpour, and Hiroaki Ogata. On estimating the tail index and the spectral measure of multivariate α-stable distributions. Metrika, 78:549–561, 07 2015.

4 [14] J. P. Nolan. Stable Distributions - Models for Heavy Tailed Data. Birkhauser, Boston, 2018. In progress, Chapter 1 online at http://fs2.american.edu/jpnolan/www/stable/stable.html. [15] Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis. arXiv preprint arXiv:1702.03849, 2017. [16] Nornadiah Mohd Razali and Yap Bee Wah. Power comparisons of shapiro-wilk, kolmogorov- smirnov, lilliefors and anderson-darling tests. J. Stat. Model. Analytics, 2, 01 2011. [17] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400–407, 1951. [18] Gareth O Roberts and Osnat Stramer. Langevin diffusions and metropolis-hastings algorithms. Methodology and computing in applied probability, 4(4):337–357, 2002. [19] S. S. Shapiro and M. B. Wilk. An analysis of variance test for normality (complete samples). Biometrika, 52(3/4):591–611, 1965. [20] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014. [21] Umut Simsekli, Levent Sagun, and Mert Gürbüzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5827–5837, Long Beach, California, USA, 09–15 Jun 2019. PMLR. [22] Roman Vershynin. High-Dimensional Probability. Cambridge University Press, 2018. [23] Yuchen Zhang, Percy Liang, and Moses Charikar. A hitting time analysis of stochastic gradient langevin dynamics. COLT 2017, 2017. [24] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from minima and regularization effects. arXiv preprint arXiv:1803.00195, 2018.

5 A Appendix

Here we present the various plots for the remaining data-sets and models.

(a) (b) (c)

Figure 3: Gaussianity Test Experiments on VGG16, without BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 4096.

(a) (b) (c)

Figure 4: Gaussianity Test Experiments on VGG16, with BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 4096.

6 (a) (b) (c)

Figure 5: Gaussianity Test Experiments on Resnet18, without BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 4096.

(a) (b) (c)

Figure 6: Gaussianity Test Experiments on Resnet18, with BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 4096.

(a) (b) (c)

Figure 7: Gaussianity Test Experiments on VGG16, without BN and trained on CIFAR10 with 256 minibatch-size SGD, (a) learning rate is 10−1 (b) learning rate is 10−2 (c) learning rate is 10−3.

7 (a) (b) (c)

Figure 8: Gaussianity Test Experiments on VGG16, with BN and trained on CIFAR10 with 256 minibatch-size SGD, (a) learning rate is 10−1 (b) learning rate is 10−2 (c) learning rate is 10−3.

(a) (b) (c)

Figure 9: Gaussianity Test Experiments on Resnet18, without BN and trained on CIFAR10 with 256 minibatch-size SGD, (a) learning rate is 10−1 (b) learning rate is 10−2 (c) learning rate is 10−3.

(a) (b) (c)

Figure 10: Gaussianity Test Experiments on Resnet18, with BN and trained on CIFAR10 with 256 minibatch-size SGD, (a) learning rate is 10−1 (b) learning rate is 10−2 (c) learning rate is 10−3.

8 Figure 11: Gaussianity Test Experiments on Alexnet, without BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 1024.

(a) (b) (c)

Figure 12: Gaussianity Test Experiments on Alexnet, with BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 1024.

(a) Gaussianity testing on gradient (b) Gaussianity testing on gradient (c) Gaussianity testing on gradient projection projection projection

Figure 13: Gaussianity Test Experiments on a 3 layer fully connected network, with BN and trained on CIFAR10 at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 1024.

9 Figure 14: Gaussianity Test Experiments on a 3 layer fully connected network, without BN and trained on MNIST at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 1024.

(a) Gaussianity testing on gradient (b) Gaussianity testing on gradient (c) Gaussianity testing on gradient projection projection projection

Figure 15: Gaussianity Test Experiments on a LeNet, without BN and trained on MNIST at l.r. 10−2, (a) Mini batch-size is 32 (b) Mini batch-size is 256 (c) Mini batch-size is 1024.

10