Theoretical Analysis of Auto Rate-Tuning by Batch Normalization
Sanjeev Arora, Zhiyuan Li, Kaifeng Lyu
Introduction
Batch Normalization (abbreviated as BatchNorm or BN) (Ioffe & Szegedy, 2015) is one of the most important innovation in deep learning, widely used in modern neural network architectures such as ResNet (He et al., 2016), Inception (Szegedy et al., 2017), and DenseNet (Huang et al., 2017). It also inspired a series of other normalization methods (Ulyanov et al., 2016; Ba et al., 2016; Ioffe, 2017; Wu & He, 2018).
BatchNorm consists of standardizing the output of each layer to have zero mean and unit variance. For a single neuron, if is the original outputs in a mini-batch, then it adds a BatchNorm layer which modifies the outputs to
where and are the mean and variance within the mini-batch, and are two learnable parameters. BN appears to stabilize and speed up training, and improve generalization. The inventors suggested (Ioffe & Szegedy, 2015) that these benefits derive from the following:
By stabilizing layer outputs it reduces a phenomenon called Internal Covariate Shift, whereby the training of a higher layer is continuously undermined or undone by changes in the distribution of its inputs due to parameter changes in previous layers.,
Making the weights invariant to scaling, appears to reduce the dependence of training on the scale of parameters and enables us to use a higher learning rate;
By implictly regularizing the model it improves generalization.
But these three benefits are not fully understood in theory. Understanding generalization for deep models remains an open problem (with or without BN). Furthermore, in demonstration that intuition can sometimes mislead, recent experimental results suggest that BN does not reduce internal covariate shift either (Santurkar et al., 2018), and the authors of that study suggest that the true explanation for BN’s effectiveness may lie in a smoothening effect (i.e., lowering of the Hessian norm) on the objective. Another recent paper (Kohler et al., 2018) tries to quantify the benefits of BN for simple machine learning problems such as regression but does not analyze deep models.
Taking derivatives one finds that the gradient at equals to the gradient at multiplied by a factor . Thus, even though the scale of weight parameters of a linear layer proceeding a BatchNorm no longer means anything to the function represented by the neural network, their growth has an effect of reducing the learning rate.
Our paper considers the following question: Can we rigorously capture the above intuitive behavior? Theoretical analyses of speed of gradient descent algorithms in nonconvex settings study the number of iterations required for convergence to a stationary point (i.e., where gradient vanishes). But they need to assume that the learning rate has been set (magically) to a small enough number determined by the smoothness constant of the loss function — which in practice are of course unknown. With this tuned learning rate, the norm of the gradient reduces asymptotically as in iterations. In case of stochastic gradient descent, the reduction is like . Thus a potential way to quantify the rate-tuning behavior of BN would be to show that even when the learning rate is fixed to a suitable constant, say , from the start, after introducing BN the convergence to stationary point is asymptotically just as fast (essentially) as it would be with a hand-tuned learning rate required by earlier analyses. The current paper rigorously establishes such auto-tuning behavior of BN (See below for an important clarification about scale-invariance).
We note that a recent paper (Wu et al., 2018) introduced a new algorithm WNgrad that is motivated by BN and provably has the above auto-tuning behavior as well. That paper did not establish such behavior for BN itself, but it was a clear inspiration for our analysis of BN.
1 Our contributions
In this paper, we show that the scale-invariant parameters do not require rate tuning for lowering the training loss. To illustrate this, we consider the case in which we set learning rates separately for scale-invariant parameters and scale-variant parameters . Under some assumptions on the smoothness of the loss and the boundedness of the noise, we show that
In full-batch gradient descent, if the learning rate for is set optimally, then no matter how the learning rates for is set, converges to a first-order stationary point in the rate , which asymptotically matches with the convergence rate of gradient descent with optimal choice of learning rates for all parameters (Theorem 3.1);
In the usual case where we set a unified learning rate for all parameters, our results imply that we only need to set a learning rate that is suitable for . This means introducing scale-invariance into neural networks potentially reduces the efforts to tune learning rates, since there are less number of parameters we need to concern in order to guarantee an asymptotically fastest convergence.
In our study, the loss function is assumed to be smooth. However, BN introduces non-smoothness in extreme cases due to division by zero when the input variance is zero (see equation 1). Note that the suggested implementation of BN by Ioffe & Szegedy (2015) uses a smoothening constant in the whitening step, but it does not preserve scale-invariance. In order to avoid this issue, we describe a simple modification of the smoothening that maintains scale-invariance. Also, our result cannot be applied to neural networks with ReLU, but it is applicable for its smooth approximation softplus (Dugas et al., 2001).
We include some experiments in Appendix D, showing that it is indeed the auto-tuning behavior we analysed in this paper empowers BN to have such convergence with arbitrary learning rate for scale-invariant parameters. In the generalization aspect, a tuned learning rate is still needed for the best test accuracy, and we showed in the experiments that the auto-tuning behavior of BN also leads to a wider range of suitable learning rate for good generalization.
2 Related works
Previous work for understanding Batch Normalization. Only a few recent works tried to theoretically understand BatchNorm. Santurkar et al. (2018) was described earlier. Kohler et al. (2018) aims to find theoretical setting such that training neural networks with BatchNorm is faster than without BatchNorm. In particular, the authors analyzed three types of shallow neural networks, but rather than consider gradient descent, the authors designed task-specific training methods when discussing neural networks with BatchNorm. Bjorck et al. (2018) observes that the higher learning rates enabled by BatchNorm improves generalization.
Convergence of adaptive algorithms. Our analysis is inspired by the proof for WNGrad (Wu et al., 2018), where the author analyzed an adaptive algorithm, WNGrad, motivated by Weight Normalization (Salimans & Kingma, 2016). Other works analyzing the convergence of adaptive methods are (Ward et al., 2018; Li & Orabona, 2018; Zou & Shen, 2018; Zhou et al., 2018).
General framework
In this section, we introduce our general framework in order to study the benefits of scale-invariance.
Scale-invariance is common in neural networks with BatchNorm. We formally state the definition of scale-invariance below:
(Scale-invariance) Let be a loss function. We say that is a scale-invariant parameter of if for all , ; if is not scale-invariant, then we say is a scale-variant parameter of .
We consider the following -layer “fully-batch-normalized” feedforward network for illustration:
BN has the property that the output is unchanged when the batch inputs are scaled or shifted simultaneously. For being the output of a linear layer, it is easy to see that is scale-invariant, and thus each row vector of weight matrices in are scale-invariant parameters of . In convolutional neural networks with BatchNorm, a similar argument can be done. In particular, each filter of convolutional layer normalized by BN is scale-invariant.
With a general nonlinear activation, other parameters in , the scale and shift parameters and in each BN, are scale-variant. When ReLU or Leaky ReLU (Maas et al., 2013) are used as the activation , the vector of each BN at layer (except the last one) is indeed scale-invariant. This can be deduced by using the the (positive) homogeneity of these two types of activations and noticing that the output of internal activations is processed by a BN in the next layer. Nevertheless, we are not able to analyse either ReLU or Leaky ReLU activations because we need the loss to be smooth in our analysis. We can instead analyse smooth activations, such as sigmoid, tanh, softplus (Dugas et al., 2001), etc.
2 Framework
3 The intrinsic optimization problem
Thanks to the scale-invariant properties, the scale of each weight does not affect loss values. However, the scale does affect the gradients. Let be the set of normalized weights, where . The following simple lemma can be easily shown:
To make to be small, one can just scale the weights by a large factor. Thus there are ways to reduce the norm of the gradient that do not reduce the loss.
For this reason, we define the intrinsic optimization problem for training the neural network. Instead of optimizing and over all possible solutions, we focus on parameters in which for all . This does not change our objective, since the scale of does not affect the loss.
Let be the intrinsic domain. The intrinsic optimization problem is defined as optimizing the original problem in :
In this paper, we aim to show that training neural network for the original optimization problem by gradient descent can be seen as training by adaptive methods for the intrinsic optimization problem, and it converges to a first-order stationary point in the intrinsic optimization problem with no need for tuning learning rates for .
4 Assumptions on the loss
We assume is defined and twice continuously differentiable at any satisfying none of is . Also, we assume that the expected loss is lower-bounded by .
Furthermore, for , where , we assume that the following bounds on the smoothness:
Smoothed version of motivating neural networks. Note that the neural network illustrated in Section 2.1 does not meet the conditions of the smooothness at all since the loss function could be non-smooth. We can make some mild modifications to the motivating example to smoothen it Our results to this network are rather conceptual, since the smoothness upper bound can be as large as , where is the number of layers and is the maximum width of each layer.:
The activation could be non-smooth. A possible solution is to use smooth nonlinearities, e.g., sigmoid, tanh, softplus (Dugas et al., 2001), etc. Note that softplus can be seen as a smooth approximation of the most commonly used activation ReLU.
The formula of BN shown in equation 3 may suffer from the problem of division by zero. To avoid this, the inventors of BN, Ioffe & Szegedy (2015), add a small smoothening parameter to the denominator, i.e.,
Since the variance of inputs is usually large in practice, for small , the effect of the smoothening term is negligible except in extreme cases.
Using the above two modifications, the loss function is already smooth. However, the scale of scale-variant parameters may be unbounded during training, which could cause the smoothness unbounded. To avoid this issue, we can either project scale-variant parameters to a bounded set, or use weight decay for those parameters (see Appendix C for a proof for the latter solution).
5 Key observation: the growth of weights
The following lemma is our key observation. It establishes a connection between the scale-invariant property and the growth of weight scale, which further implies an automatic decay of learning rates:
For any scale-invariant weight in the network , we have:
and are always perpendicular;
Let be all the parameters in other than . Taking derivatives with respect to for the both sides of , we have The right hand side equals , so the first proposition follows by taking . Applying Pythagorean theorem and Lemma 2.2, the second proposition directly follows. ∎
Using Lemma 2.4, we can show that performing gradient descent for the original problem is equivalent to performing an adaptive gradient method for the intrinsic optimization problem:
Let . Then for all ,
where is a projection operator which maps any vector to .
Wu et al. (2018) noticed that Theorem 2.5 is true for Weight Normalization by direct calculation of gradients. Inspiring by this, they proposed a new adaptive method called . Our theorem is more general since it holds for any normalization methods as long as it induces scale-invariant properties to the network. The adaptive update rule derived in our theorem can be seen as with projection to unit sphere after each step.
which implies the first equation. The second equation is by Lemma 2.4. ∎
Training by full-batch gradient descent
In this section, we rigorously analyze the effect related to the scale-invariant properties in training neural network by full-batch gradient descent. We use the framework introduced in Section 2.2 and assumptions from Section 2.4. We focus on the full-batch training, i.e., is always equal to the whole training set and .
This matches the asymptotic convergence rate of GD by Carmon et al. (2018).
2 Proof sketch
The high level idea is to use the decrement of loss function to upper bound the sum of the squared norm of the gradients. Note that . For the first part , we have
Thus the core of the proof is to show that the monotone increasing has an upper bound for all . It is shown that for every , the whole training process can be divided into at most two phases. In the first phase, the effective learning rate is larger than some threshold (defined in Lemma 3.2) and in the second phase it is smaller.
The full proof is postponed to Appendix A.
Training by stochastic gradient descent
In this section, we analyze the effect related to the scale-invariant properties when training a neural network by stochastic gradient descent. We use the framework introduced in Section 2.2 and assumptions from Section 2.4.
Assumptions on learning rates. As usual, we assume that the learning rate for is chosen carefully and the learning rate for is chosen rather arbitrarily. More specifically, we consider the case that the learning rates are chosen as
2 Proof sketch
We delay the full proof into Appendix B and give a proof sketch in a simplified setting where there is no and . We also assume there’s only one , that is, and omit the index .
Taking expectation over equation 14 and summing it up, we have
Plug the above bounds into the above inequality, we complete the proof.
Conclusions and future works
In this paper, we studied how scale-invariance in neural networks with BN helps optimization, and showed that (stochastic) gradient descent can achieve the asymptotic best convergence rate without tuning learning rates for scale-invariant parameters. Our analysis suggests that scale-invariance in nerual networks introduced by BN reduces the efforts for tuning learning rate to fit the training data.
However, our analysis only applies to smooth loss functions. In modern neural networks, ReLU or Leaky ReLU are often used, which makes the loss non-smooth. It would have more implications by showing similar results in non-smooth settings. Also, we only considered gradient descent in this paper. It can be shown that if we perform (stochastic) gradient descent with momentum, the norm of scale-invariant parameters will also be monotone increasing. It would be interesting to use it to show similar convergence results for more gradient methods.
Thanks Yuanzhi Li, Wei Hu and Noah Golowich for helpful discussions. This research was done with support from NSF, ONR, Simons Foundation, Mozilla Research, Schmidt Foundation, DARPA, and SRC.
References
Appendix A Proof for Full-Batch Gradient Descent
By the scale-invariant property of , we know that . Also, the following identities about derivatives can be easily obtained:
Thus, the assumptions on the smoothness imply
Using Taylor expansion, we have , such that for ,
By the inequality of arithmetic and geometric means, we have
Using the assumption on the smoothness, we can show that the gradient with respect to is essentially bounded:
A.1 Fix all the parameters except . Then can be written as a function on the variable . Let . Since is continuous and is compact, there must exist such that for all . Note that is scale-invariant, so is also a minimum in the entire domain and .
For an arbitrary , let . Let be a curve such that , , and goes along the geodesic from to on the unit sphere with constant speed. Let . By Taylor expansion, we have
The following lemma gives an upper bound to the weight scales.
the following inequality on weight scales at time holds:
Taking sum over all and also subtracting on the both sides, we have
where is used at the second line. ∎
Combining the lemmas above together, we can obtain our results.
Thus converges in the rate of
Appendix B Proof for Stochastic Gradient Descent
Let be the filtration, where denotes the sigma field.
For any with ,
Let be the minimum such that . Let be the maximum such that exists. Let . Then we know that
Thus, . ∎
Conditioned on , by Taylor expansion, we have
By the inequality , we have
Taking this into equation 21 and summing up for all , we have
For any , , we have
Fix . First we bound . Recall that
We can get the bound for by combining equation 22 and equation 23.
Combining Lemma B.2 and Lemma B.4, for , we have
Appendix C Proof for the smoothness of the motivating neural network
In this section we prove that the modified version of the motivating neural network does meet the assumptions in Section 2.4. More specifically, we assume:
We use the network structure in Section 2.1 with the smoothed variant of BN as described in Section 2.4;
The objective is twice continuously differentiable, lower bounded by and Lipschitz ();
The activation is twice continuously differentiable and Lipschitz ();
We add an extra weight decay (L2 regularization) term to the loss in equation 2 for some .
First, we show that (containing all scale and shift parameters in BN) is bounded during the training process. Then the smoothness follows compactness using Extreme Value Theorem.
We use the following lemma to calculate back propagation:
Let for some function . If , then
Since ,
If is bounded by a constant, there exists some constant such that .
Fix a time in the training process. Consider the process of back propagation. Define
where is the output of the -th neuron in the -th layer in the -th data sample in the batch. By the Lipschitzness of the objective, can be bounded by a constant. If can be bounded by a constant, then by the Lipschitzness of and Lemma C.1, the gradient of and in layer can also be bounded by a constant. Note that
Thus and in layer can be bounded by a constant since
Also Lemma C.1 and the Lipschitzness of imply that can be bounded if and in the layer can be bounded by a constant. Using a simple induction, we can prove the existence of for bounding the norm of for all time . ∎
If is bounded by a constant, then satisfies the assumptions in Section 2.4.
Appendix D Experiments
In this section, we provide experimental evidence showing that the auto rate-tuning behavior does empower BN in the optimization aspect.
In this network, every kernel is scale-invariant, and for every BN layer except the last one, the concatenation of all and parameters in this BN is also scale-invariant. Only and parameters in the last BN are scale-variant (See Section 2.1). We consider the training in following two settings:
Train the network using the standard SGD(No momentum, learning rate decay ,weight decay and dropout);
Train the network using Projected SGD (PSGD): at each iteration, one first takes a step proportional to the negative of the gradient calculated in a random batch, and then projects each scale-invariant parameter to the sphere with radius equal to its -norm before this iteration, i.e., rescales each scale-invariant parameter so that each maintains its length during training.
Note that the projection in Setting removes the adaptivity of the learning rates in the corresponding intrinsic optimization problem, i.e., in equation 9 remains constant during the training. Thus, by comparing Setting 1 and Setting 2, we can know whether or not the auto-tuning behavior of BN shown in theory is effective in practice.
As in our theoretical analysis, we consider what will happen if we set two learning rates separately for scale-invariant and scale-variant parameters. We train the network in either setting with different learning rates ranging from to for epochs.
First, we fix the learning rate for scale-variant ones to , and try different learning rates for scale-invariant ones. As shown in Figure 1, for small learning rates (such as ), the training processes of networks in Setting 1 and 2 are very similar. But for larger learning rates, networks in Setting 1 can still converge to for all the learning rates we tried, while networks in Setting 2 got stuck with relatively large training loss. This suggests that the auto-tuning behavior of BN does takes effect when the learning rate is large, and it matches with the claimed effect of BN in Ioffe & Szegedy (2015) that BN enables us to use a higher learning rate. Though our theoretical analysis cannot be directly applied to the network we trained due to the non-smoothness of the loss function, the experiment results match with what we expect in our analysis.
D.2 Unified Learning Rate
Next, we consider the case in which we train the network with a unified learning rate for both scale-invariant and scale-variant parameters. We also compare Setting 1 and 2 with the setting in which we train the network with all the BN layers removed using SGD (we call it Setting 3).
As shown in Figure 2, the training loss of networks in Setting 1 converges to . On the contrast, the training loss of networks in Setting 2 and 3 fails to converge to when a large learning rate is used, and in some cases the loss diverges to infinity or NaN. This suggests that the auto-tuning behavior of BN has an effective role in the case that a unified learning rate is set for all parameters.
D.3 Generalization
Despite in Setting 1 the convergence of training loss for different learning rates, the convergence points can be different, which lead to different performances on test data.
In Figure 3, we plot the test accuracy of networks trained in Setting 1 and 2 using different unified learning rates, or separate learning rates with the learning rate for scale-variant parameters fixed to . As shown in the Figure 3, the test accuracy of networks in Setting 2 decreases as the learning rate increases over , while the test accuracy of networks in Setting 1 remains higher than . The main reason that the network in Setting 2 doesn’t perform well is underfitting, i.e. the network in Setting 2 fails to fit the training data well when learning rate is large. This suggests that the auto-tuning behavior of BN also benefits generalization since such behavior allows the algorithm to pick learning rates from a wider range while still converging to small test error.