Batch Normalization (BN), introduced in Batch Normalization, is a significantly important technique in current deep learning models (especially those related to computer vision). It accelerates training, provides a certain degree of regularization against overfitting, and allows us to use larger learning rates. Overall, it offers many benefits, provided you can afford a relatively large batch size.
But how exactly does BN work? Early explanations were primarily based on probability distributions, suggesting that normalizing the input distribution of each layer to \mathcal{N}(0,1) reduces the so-called Internal Covariate Shift, thereby stabilizing and accelerating training. While this explanation seems plausible, it reveals issues upon closer inspection: first, the input to any layer cannot strictly satisfy a normal distribution, so simply standardizing the mean and variance cannot achieve a standard \mathcal{N}(0,1) distribution; second, even if it could achieve \mathcal{N}(0,1), this interpretation fails to explain why other normalization methods (such as Instance Normalization or Layer Normalization) also work.
In the 2018 paper "How Does Batch Normalization Help Optimization?", the authors explicitly raised these doubts, refuted some of the original views, and proposed their own new understanding of BN: They believe the primary role of BN is to make the landscape of the entire loss function smoother, thereby allowing for more stable training.
This blog post shares the conclusions of that paper, but the method of argumentation is my own "behind-closed-doors" (independent) construction. I personally find the arguments in the original paper somewhat obscure, especially the mathematical parts. Therefore, this article attempts to express the same viewpoint as intuitively as possible.
(Note: Before reading this article, please ensure you already know what BN is; this article will not repeat the basic concepts and procedures of BN.)
Basic Conclusions
In this section, we first present a core inequality, then derive gradient descent, and obtain some basic conclusions about model training to lay the groundwork for the subsequent analysis of BN.
Core Inequality
Assume that the gradient of a function f(\theta) satisfies the Lipschitz constraint (L-constraint), i.e., there exists a constant L such that the following holds: \Vert \nabla_{\theta} f(\theta + \Delta \theta) - \nabla_{\theta} f(\theta)\Vert_2 \leq L\Vert \Delta\theta\Vert_2 Then we have the following inequality: f(\theta+\Delta\theta) \leq f(\theta) + \left\langle \nabla_{\theta}f(\theta), \Delta\theta\right\rangle + \frac{1}{2}L \Vert \Delta\theta\Vert_2^2 \label{eq:core-eq}
Proof: The proof is not difficult. Define an auxiliary function f(\theta + t\Delta\theta), t \in [0, 1], and then directly obtain: \begin{aligned} f(\theta + \Delta\theta) - f(\theta) &= \int_0^1 \frac{\partial f(\theta + t\Delta\theta)}{\partial t} dt \\ &= \int_0^1 \left\langle \nabla_{\theta} f(\theta + t\Delta\theta), \Delta\theta \right\rangle dt \\ &= \left\langle \nabla_{\theta} f(\theta), \Delta\theta \right\rangle + \int_0^1 \left\langle \nabla_{\theta} f(\theta + t\Delta\theta) - \nabla_{\theta} f(\theta), \Delta\theta \right\rangle dt \\ &\leq \left\langle \nabla_{\theta} f(\theta), \Delta\theta \right\rangle + \int_0^1 \Vert \nabla_{\theta} f(\theta + t\Delta\theta) - \nabla_{\theta} f(\theta) \Vert_2 \cdot \Vert \Delta\theta \Vert_2 dt \\ &\leq \left\langle \nabla_{\theta} f(\theta), \Delta\theta \right\rangle + \int_0^1 L \Vert \Delta\theta \Vert_2^2 t dt \\ &= \left\langle \nabla_{\theta} f(\theta), \Delta\theta \right\rangle + \frac{1}{2} L \Vert \Delta\theta \Vert_2^2 \end{aligned}
Gradient Descent
Assume f(\theta) is the loss function, and our goal is to minimize f(\theta). This inequality tells us a lot. First, since we are minimizing, we naturally hope that each step results in a decrease, i.e., f(\theta+\Delta\theta) < f(\theta). Since \frac{1}{2}L \Vert \Delta\theta\Vert_2^2 is necessarily non-negative, the only choice to ensure a decrease is to have \left\langle \nabla_{\theta}f(\theta), \Delta\theta\right\rangle < 0. A natural choice is: \Delta\theta = -\eta \nabla_{\theta}f(\theta) \label{eq:gd} where \eta > 0 is a scalar, namely the learning rate.
We can see that Equation [eq:gd] is the update formula for gradient descent. Thus, this is a derivation of gradient descent, and the information contained in this derivation is richer because it is a strict inequality, allowing us to draw further conclusions about training.
Lipschitz Constraint
Substituting the gradient descent formula into inequality [eq:core-eq], we get: f(\theta+\Delta\theta) \leq f(\theta) + \left(\frac{1}{2}L\eta^2 - \eta\right) \Vert \nabla_{\theta}f(\theta)\Vert_2^2 Note that a sufficient condition to ensure the loss function decreases is \frac{1}{2}L\eta^2 - \eta < 0. To achieve this, either \eta must be sufficiently small, or L must be sufficiently small. However, a very small \eta means the learning speed will be quite slow. Therefore, the more ideal situation is for L to be small enough. Reducing L allows for a larger learning rate, which speeds up learning—this is one of its benefits.
However, L is an intrinsic property of f(\theta), so L can only be reduced by adjusting f itself.
How BN is Forged
This section will show that BN can be naturally derived with the goal of reducing the L-constant of the neural network’s gradient. That is to say, BN reduces the L-constant of the neural network’s gradient, making the network easier to learn (e.g., allowing for larger learning rates). Intuitively, reducing the L-constant of the gradient means making the loss function less "turbulent," which is what it means to make the landscape smoother.
Note: We have discussed L-constraints before, but previously we discussed the neural network satisfying an L-constraint with respect to the input, which leads to spectral regularization and spectral normalization of weights (please refer to "Lipschitz Constraints in Deep Learning: Generalization and Generative Models"). This article discusses the neural network (its gradient) satisfying an L-constraint with respect to the parameters, which leads to various normalization methods for inputs, with BN being the most natural one.
Gradient Analysis
Taking supervised learning as an example, assume the neural network is represented as \hat{y}=h(x;\theta) and the loss function is l(y,\hat{y}). Our task is: \theta = \mathop{\text{argmin}}_{\theta}\, \mathbb{E}_{(x,y)\sim p(x,y)}\left[l(y, h(x;\theta))\right] That is, f(\theta)=\mathbb{E}_{(x,y)\sim p(x,y)}\left[l(y, h(x;\theta))\right], so: \begin{aligned} \nabla_{\theta}f(\theta) &= \mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{\theta}l(y, h(x;\theta))\right] \\ &= \mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{h}l(y, h(x;\theta))\nabla_{\theta}h(x;\theta)\right] \end{aligned}
As a side note, none of the symbols in this article are bolded, but depending on the context, they may represent either scalars or vectors.
Non-linearity Assumption
Clearly, f(\theta) is a non-linear function, and its non-linearity comes from two sources:
The loss function l(y,\hat{y}) is generally non-linear;
The activation functions in the neural network h(x;\theta) are non-linear.
Regarding activation functions, most current mainstream activation functions satisfy a property: the absolute value of the derivative does not exceed a certain constant. We now consider whether this property can be extended to the loss function: (During the entire training process) is the gradient of the loss function \nabla_{h}l(y, h(x;\theta)) confined within a certain range?
At first glance, this assumption usually does not hold. For example, cross-entropy is -\log p, and its derivative is -1/p, which clearly cannot be constrained within a finite range. However, when the loss function is considered together with the activation function of the last layer, this constraint is usually satisfied. For instance, in binary classification, the last layer usually uses a sigmoid activation; combined with cross-entropy, this becomes: -\log \text{sigmoid}(h(x;\theta)) = \log \left(1 + e^{-h(x;\theta)}\right) In this case, its gradient with respect to h is between -1 and 1. Of course, there are cases where this does not hold, such as regression problems using MSE as the loss function where the last layer usually has no activation function. In such cases, its gradient is a linear function and will not be confined to a finite range. In these situations, we can only hope that the model has good initialization and a good optimizer so that \nabla_{h}l(y, h(x;\theta)) remains relatively stable throughout the training process. This "hope" might seem strong, but neural networks that are successfully trained basically satisfy this "hope."
Cauchy-Schwarz Inequality
Our goal is to explore the extent to which \nabla_{\theta}f(\theta) satisfies the L-constraint and to investigate methods for reducing this L. To this end, let’s first consider the simplest single-layer neural network (vector input, scalar output) h(x;w,b)=g\left(\left\langle x, w\right\rangle + b\right), where g is the activation function. In this case: \begin{aligned} \mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{b}f(w,b)\right] &= \mathbb{E}_{(x,y)\sim p(x,y)}\left[\frac{\partial l_{w,b}}{\partial g}\dot{g}\left(\left\langle x, w\right\rangle + b\right)\right] \\ \mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{w}f(w,b)\right] &= \mathbb{E}_{(x,y)\sim p(x,y)}\left[\frac{\partial l_{w,b}}{\partial g}\dot{g}\left(\left\langle x, w\right\rangle + b\right)x\right] \end{aligned} \label{eq:grads}
Based on our assumption, \frac{\partial l_{w,b}}{\partial g} and \dot{g}\left(\left\langle x, w\right\rangle + b\right) are both restricted within a certain range, so we can see that the gradient of the bias term b is very stable, and its updates should also be very stable. However, the gradient of w is different; it is directly related to the input x.
Regarding the gradient difference for w, we have: \begin{aligned} &\big\Vert\mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{w}f(w+\Delta w,b)\right] - \mathbb{E}_{(x,y)\sim p(x,y)}\left[\nabla_{w}f(w,b)\right]\big\Vert_2 \\ =& \Bigg\Vert\mathbb{E}_{(x,y)\sim p(x,y)}\left[\left(\frac{\partial l_{w+\Delta w,b}}{\partial g}\dot{g}\left(\left\langle x, w+\Delta w\right\rangle + b\right) - \frac{\partial l_{w,b}}{\partial g}\dot{g}\left(\left\langle x, w\right\rangle + b\right)\right)x\right]\Bigg\Vert_2 \end{aligned}
Let the term in the parentheses be \lambda(x, y; w,b,\Delta w). According to the previous discussion, it is constrained within a certain range and remains a stable term. As such, we might as well assume it naturally satisfies the L-constraint, i.e.: \Vert\lambda(x, y; w,b,\Delta w)\Vert_2 = \mathcal{O}\left(\Vert\Delta w\Vert_2\right) Now we only need to care about the extra x. According to the Cauchy-Schwarz inequality, we have: \begin{aligned} &\Big\Vert\mathbb{E}_{(x,y)\sim p(x,y)}\left[\lambda(x, y; w,b,\Delta w) x\right]\Big\Vert_2 \\ \leq & \sqrt{\mathbb{E}_{(x,y)\sim p(x,y)}\left[\lambda(x, y; w,b,\Delta w)^2\right]} \times \sqrt{\big|\mathbb{E}_{x\sim p(x)}\left[x\otimes x\right]\big|_1} \end{aligned} \label{eq:kexi}
In this way, we obtain the term \big|\mathbb{E}_{x\sim p(x)}\left[x\otimes x\right]\big|_1, which is independent of the (current layer’s) parameters. If we wish to reduce the L-constant, the most direct method is to reduce this term.
Subtracting Mean and Dividing by Standard Deviation
It is important to note that while we strongly desire to reduce the L-constant of the gradient, there is a prerequisite—it must be done without significantly reducing the original neural network’s fitting capability. Otherwise, one could simply multiply by 0 to reduce L to 0, but that would be meaningless.
The result of Equation [eq:kexi] tells us that finding a way to reduce \big|\mathbb{E}_{x\sim p(x)}\left[x\otimes x\right]\big|_1 is a direct approach. This means we need to transform the input x. Given the prerequisite of "not reducing fitting capability," the simplest and potentially effective method is a translation transformation. That is, we consider x \to x - \mu. In other words, we consider an appropriate \mu that minimizes: \big|\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]\big|_1 This is merely a quadratic minimization problem, and it is not difficult to solve for the optimal \mu: \mu = \mathbb{E}_{x\sim p(x)}\left[x\right] \label{eq:mu} This is exactly the mean of all samples. Thus, we obtain:
Conclusion 1: Subtracting the mean of all samples from the input can reduce the L-constant of the gradient. It is an operation that facilitates optimization without reducing the neural network’s fitting capability.
Next, we consider a scaling transformation, i.e., x - \mu \to \frac{x - \mu}{\sigma}, where \sigma is a vector of the same size as x, and the division is element-wise. This leads to: \big|\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]\big|_1 \to \left|\frac{\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]}{\sigma\otimes \sigma}\right|_1
\sigma is a direct scaling factor for L. But the question is, what is a good value to scale to? If we only pursued a smaller L, we could simply let \sigma \to \infty, but then the neural network would have no fitting capability at all. Conversely, if \sigma is too small, resulting in an excessively large L, it would be detrimental to optimization. So we need a standard.
What should be the standard? Looking back at the gradient expression [eq:grads], as mentioned before, the gradient of the bias term is not significantly affected by x, so it seems like a reliable standard. If so, it is equivalent to scaling this part of the input x weight directly to 1. That is to say, \frac{\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]}{\sigma\otimes \sigma} becomes an all-ones vector, or in other words: \sigma = \sqrt{\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]} \label{eq:sigma}
Thus, a relatively natural choice is to set \sigma as the standard deviation of the input. In this case, we can sense that dividing by the standard deviation acts more like an adaptive learning rate correction term. To some extent, it eliminates the differences in parameter optimization caused by inputs at different layers, making the optimization of the entire network more "synchronized," or making each layer of the neural network more "equal." This allows the entire neural network to be more fully utilized and reduces the possibility of overfitting in a specific layer. Of course, when the magnitude of the input is too large, dividing by the standard deviation also helps reduce the L-constant of the gradient.
Thus, we have:
Conclusion 2: Dividing the input (after subtracting the mean) by the standard deviation of all samples acts similarly to an adaptive learning rate. it makes the updates of each layer more synchronized and reduces the possibility of overfitting in a specific layer, thereby improving the performance of the neural network.
The Emergence of BN
Although the previous derivation used a single-layer neural network (vector input, scalar output) as an example, the conclusions are sufficiently representative because a multi-layer neural network is essentially a composition of single-layer networks (for more on this, refer to my previous work "From Boosting to Neural Networks: Seeing Mountains as Mountains?").
With these two conclusions, BN can basically be implemented: during training, subtract the mean and divide by the standard deviation for the input of each layer. However, since each batch is only an approximation of the whole, and the expectations [eq:mu] and [eq:sigma] are the mean and standard deviation of all samples, BN inevitably requires a larger batch size to be effective, which places demands on computational power. Furthermore, a conclusion of this analysis is: BN should be placed before the fully connected or convolutional layers.
Additionally, we need to maintain a set of variables to store the mean and variance during the training process for use during prediction; these are the mean and variance variables tracked via moving averages in BN. As for the \beta and \gamma terms added after subtracting the mean and dividing by the standard deviation in the standard BN design, I believe they are "icing on the cake" and not strictly necessary, so I won’t provide further explanation for them.
Simple Summary
This article analyzed the principles of how BN works from an optimization perspective. The viewpoint held is basically consistent with "How Does Batch Normalization Help Optimization?", but the mathematical demonstration and description used here are, in my opinion, simpler and easier to understand. The final conclusion is that subtracting the mean helps reduce the L-constant of the neural network’s gradient, while dividing by the standard deviation acts more like an adaptive learning rate, making the updates of each parameter more synchronized and preventing overfitting to a specific layer or parameter.
Of course, the above interpretation is just a rough guide. Fully explaining BN is a very difficult task. The effect of BN is more like the composite result of multiple factors. For example, for our mainstream activation functions, the interval [-1, 1] is basically where non-linearity is strongest. Therefore, making the input have a mean of 0 and a variance of 1 also allows the non-linear capabilities of the activation functions to be more fully utilized, avoiding a waste of the neural network’s fitting capacity.
In short, the theoretical analysis of neural networks is a difficult undertaking, far beyond what I can master. I can only write these blog posts and tell some perhaps trivial stories for everyone’s amusement.
Original URL: https://kexue.fm/archives/6992
For more details on reprinting, please refer to: "Scientific Space FAQ"