What Does BN Actually Do? A "Reinventing the Wheel" Analysis
BN, i.e. Batch Normalization, is a fairly important trick in current deep learning models (especially vision-related ones). It can speed up training, offers some resistance to overfitting, and lets us use larger learning rates — overall it has quite a few benefits (provided you can afford a reasonably large batch size).
So how exactly does BN work? Early explanations were mostly based on probability distributions: roughly speaking, the idea was that normalizing the input distribution of each layer to $\mathcal{N}(0,1)$ reduces the so-called Internal Covariate Shift, thereby stabilizing and even accelerating training. This explanation seems reasonable at first glance, but on closer inspection it's actually problematic: no matter which layer's input we're talking about, it can never strictly follow a normal distribution, so simply standardizing the mean and variance cannot achieve a standard distribution $\mathcal{N}(0,1)$; moreover, even if we could achieve $\mathcal{N}(0,1)$, this interpretation still fails to explain why other normalization techniques (such as Instance Normalization and Layer Normalization) work.
In last year's paper How Does Batch Normalization Help Optimization?, the authors explicitly raised the above objections, rejecting some of the original claims, and proposed their own new understanding of BN: they argue that the main effect of BN is to make the loss landscape smoother, allowing training to proceed more stably.
This post is mainly about sharing the conclusions of that paper, but the way I present the argument was worked out "behind closed doors" on my own. I feel the original paper's exposition is overly obscure, especially the math, which is not easy to follow, so this post tries to convey the same viewpoint as intuitively as possible.
(Note: before reading this post, please make sure you already know what BN is — this post will not repeat the concept and procedure of BN.) more
Some Foundational Results
In this section we first give a core inequality, then derive gradient descent from it, and obtain some basic conclusions about model training, as groundwork for the later analysis of BN.
The Core Inequality
Suppose the gradient of function $f(\theta)$ satisfies a Lipschitz constraint (the $L$ constraint), i.e. there exists a constant $L$ such that the following always holds:
\begin{equation}\Vert \nabla_{\theta} f(\theta + \Delta \theta) - \nabla_{\theta} f(\theta)\Vert_2\leq L\Vert \Delta\theta\Vert_2\end{equation}
Then we have the following inequality:
\begin{equation}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}\end{equation}
The proof is not hard: define the auxiliary function $f(\theta + t\Delta\theta),t\in[0, 1]$, and we directly obtain \begin{equation}\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}\end{equation}.
Gradient Descent
Suppose $f(\theta)$ is the loss function, and our goal is to minimize $f(\theta)$. This inequality then tells us a great deal. First, since we're minimizing, naturally we want every step to be a decrease, i.e. $f(\theta+\Delta\theta) < f(\theta) $, and since $\frac{1}{2}L \Vert \Delta\theta\Vert_2^2$ is necessarily non-negative, the only way to guarantee a decrease is $\left\langle \nabla_{\theta}f(\theta), \Delta\theta\right\rangle < 0$. A natural choice satisfying this is
\begin{equation}\Delta\theta = -\eta \nabla_{\theta}f(\theta)\label{eq:gd}\end{equation}
Here $\eta > 0$ is a scalar, i.e. the learning rate.
We can see that equation $\eqref{eq:gd}$ is exactly the update formula for gradient descent, so this is essentially a derivation of gradient descent — and one that carries richer information than usual, since it's a strict inequality and can therefore tell us some further conclusions about training.
The Lipschitz Constraint
Substituting the gradient descent formula into inequality $\eqref{eq:core-eq}$, we get
\begin{equation}f(\theta+\Delta\theta) \leq f(\theta) + \left(\frac{1}{2}L\eta^2 - \eta\right) \Vert \nabla_{\theta}f(\theta)\Vert_2^2\end{equation}
Notice that a sufficient condition for guaranteeing the loss decreases is $\frac{1}{2}L\eta^2 - \eta < 0$. To achieve this, we either need $\eta$ to be small enough, or $L$ to be small enough. But making $\eta$ small enough means learning will be quite slow, so the more desirable situation is for $L$ to be small enough — lowering $L$ lets us use a larger learning rate, which speeds up learning. This is one of its benefits.
However, $L$ is an intrinsic property of $f(\theta)$, so the only way to reduce $L$ is by adjusting $f$ itself.
How BN Comes to Be
This section will show that, with the goal of lowering the $L$ constant of the neural network's gradient, BN emerges quite naturally. In other words, BN lowers the $L$ constant of the neural network's gradient, which makes learning easier — for instance, allowing the use of a larger learning rate. Intuitively, lowering the $L$ constant of the gradient means making the loss function less "bumpy," i.e. making the landscape smoother.
Note:
We have discussed the $L$ constraint before. Previously, we discussed the case where the neural network satisfies the $L$ constraint with respect to the "
input
", which led to spectral regularization and spectral normalization of weights (see
Lipschitz Constraints in Deep Learning: Generalization and Generative Models
). In this post, we instead discuss the case where the neural network (its gradient) satisfies the $L$ constraint with respect to the "
parameters
", which leads to the various normalization techniques applied to inputs, of which BN is the most natural.
Gradient Analysis
Take supervised learning as an example: suppose the neural network is denoted $\hat{y}=h(x;\theta)$, and the loss function is $l(y,\hat{y})$. Then what we want to do is
\begin{equation}\theta = \mathop{\text{argmin}}_{\theta}\, \mathbb{E}_{(x,y)\sim p(x,y)}\left[l(y, h(x;\theta))\right]\end{equation}
that is, $f(\theta)=\mathbb{E}_{(x,y)\sim p(x,y)}\left[l(y, h(x;\theta))\right]$, so
\begin{equation}\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}\end{equation}
By the way, none of the notation in this post is bolded, but depending on context it may represent either a scalar or a vector.
The Nonlinearity Assumption
Clearly, $f(\theta)$ is a nonlinear function, and its nonlinearity comes from two sources:
1. the loss function $l(y,\hat{y})$ is generally nonlinear;
2. the activation functions inside the neural network $h(x;\theta)$ are nonlinear.
Regarding activation functions, essentially all mainstream activation functions today satisfy one property: the absolute value of the derivative is bounded by some constant. Let's now consider whether this property can be extended to the loss function, i.e. whether the gradient of the loss function $\nabla_{h}l(y, h(x;\theta))$ (throughout training) is confined to some bounded range.
At first glance, this assumption often doesn't hold. For example, cross-entropy is $-\log p$, and its derivative is $-1/p$, which clearly cannot be bounded within a finite range. However, if we consider the loss function together with the activation function of the last layer, this constraint is usually satisfied. For instance, in binary classification the last layer typically uses a sigmoid activation, and combined with cross-entropy this gives
\begin{equation}-\log \text{sigmoid}(h(x;\theta)) = \log \left(1 + e^{-h(x;\theta)}\right)\end{equation}
In this case the gradient with respect to $h$ lies between -1 and 1. Of course, there are indeed cases where this doesn't hold — for example, regression problems typically use MSE as the loss, and the last layer usually has no activation function, in which case the gradient is a linear function and won't be confined to a finite range. In such cases, we can only hope that the model has good initialization and a good optimizer, so that $\nabla_{h}l(y, h(x;\theta))$ remains fairly stable throughout training. This "hope" may seem like a strong assumption, but in practice, neural networks that train successfully basically do satisfy it.
The Cauchy Inequality
Our goal is to examine the extent to which $\nabla_{\theta}f(\theta)$ satisfies the $L$ constraint, and to explore ways of lowering 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{equation}\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}\end{equation}
Based on our assumptions, both $\frac{\partial l_{w,b}}{\partial g}$ and $\dot{g}\left(\left\langle x, w\right\rangle + b\right)$ are confined within some range, so we can see that the gradient of the bias term $b$ is quite stable, and its updates should also be quite stable. But the gradient of $w$ is different — it is directly tied to the input $x$.
For the gradient difference with respect to $w$, we have
\begin{equation}\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}\end{equation}
Denote the term in parentheses as $\lambda(x, y; w,b,\Delta w)$. Based on the earlier discussion, it is confined to some range — this is still a stable term. Given this, let's simply assume it naturally satisfies the $L$ constraint, i.e.
\begin{equation}\Vert\lambda(x, y; w,b,\Delta w)\Vert_2=\mathcal{O}\left(\Vert\Delta w\Vert_2\right)\end{equation}
At this point we only need to focus on the extra term $x$. By the Cauchy inequality, we have
\begin{equation}\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}\end{equation}
In this way we obtain a term $\big|\mathbb{E}_{x\sim p(x)}\left[x\otimes x\right]\big|_1$ that is independent of the (current layer's) parameters. If we want to lower the $L$ constant, the most direct approach is to reduce this term.
Subtracting the Mean and Dividing by the Standard Deviation
Note that although we very much want to lower the $L$ constant of the gradient, this comes with a precondition — it must not significantly reduce the original fitting capacity of the neural network. Otherwise, we could simply multiply by 0 and lower $L$ to zero, but that would be meaningless.
The result in equation $\eqref{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, which means we need to transform the input $x$. Then, given the precondition of "not reducing fitting capacity," the simplest — and potentially effective — method is a shift transformation, i.e. we consider $x \to x - \mu$. In other words, we consider choosing an appropriate $\mu$ to minimize
\begin{equation}\big|\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]\big|_1\end{equation}
This is nothing more than a minimization problem for a quadratic function, and it's easy to solve that the optimal $\mu$ is
\begin{equation}\mu = \mathbb{E}_{x\sim p(x)}\left[x\right]\label{eq:mu}\end{equation}
which is exactly the mean over all samples. So we arrive at:
Conclusion 1: Subtracting the mean of all samples from the input lowers the $L$ constant of the gradient, and is an operation that benefits optimization without reducing the neural network's fitting capacity.
Next, let's consider the scaling transformation, i.e. $x - \mu \to \frac{x - \mu}{\sigma}$, where $\sigma$ is a vector the same size as $x$, and the division is element-wise. This gives
\begin{equation}\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\end{equation}
$\sigma$ is the most direct scaling factor for $L$, but the question is where should it be scaled to? If we simply chase a smaller $L$, we could just set $\sigma\to \infty$ directly, but then the neural network would have entirely lost its fitting capacity; on the other hand, if $\sigma$ is too small, causing $L$ to be too large, that's also unfavorable for optimization. So we need a standard to go by.
What should this standard be? Let's go back and look at the gradient expression $\eqref{eq:grads}$ once more. As mentioned earlier, the bias term's gradient is not significantly affected by $x$, so it seems like a reasonable reference point. If so, this is equivalent to scaling the weight on this term of the input $x$ 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:
\begin{equation}\sigma = \sqrt{\mathbb{E}_{x\sim p(x)}\left[(x-\mu)\otimes (x-\mu)\right]}\label{eq:sigma}\end{equation}
Given this, a relatively natural choice is to take $\sigma$ to be the standard deviation of the input. At this point, we can sense that dividing by the standard deviation acts more like an adaptive learning-rate correction term: to some extent it removes the discrepancies across layers in how inputs affect parameter optimization, making the optimization of the whole network more "synchronized," or making each layer of the neural network more "equal," thereby making fuller use of the whole network and reducing the chance of overfitting at any single layer. Of course, if the magnitude of the input is too large, dividing by the standard deviation also helps lower the $L$ constant of the gradient.
So we have the conclusion:
Conclusion 2: Dividing the input (after subtracting the mean of all samples) by the standard deviation of all samples has an effect similar to an adaptive learning rate, making the updates of each layer more synchronized, reducing the chance of overfitting at any single layer, and is an operation that improves neural network performance.
The Derivation Runs Out, and BN Appears
Although the derivation above only used a single-layer neural network (vector input, scalar output) as an example, its conclusions are already sufficiently representative, because a multi-layer neural network is essentially just a composition of single-layer neural networks (on this point, see my earlier post From Boosting to Neural Networks: Is a Mountain Still a Mountain?).
So, with the two conclusions above in hand, BN essentially falls into place: during training, we simply subtract the mean and divide by the standard deviation for the input of each layer. However, since each batch only gives an approximation of the whole, while the expected value $\eqref{eq:mu},\eqref{eq:sigma}$ is the mean and standard deviation over all samples, BN inevitably works better with a larger batch size, which places demands on compute. In addition, one conclusion of this analysis is: BN should be placed before the fully-connected/convolutional layer.
Additionally, we need to maintain a set of variables that store the running mean and variance from training for use at inference time — this is exactly the mean/variance statistics tracked via moving averages in BN. As for the $\beta,\gamma$ term added after subtracting the mean and dividing by the standard deviation in BN's standard design, I think it merely serves as a finishing touch rather than being strictly necessary, so I won't elaborate further on it.
A Brief Summary
This post analyzed the principle behind BN's effectiveness from an optimization perspective. The viewpoint here is basically consistent with that of How Does Batch Normalization Help Optimization?, but I believe the mathematical argument and manner of presentation used here are simpler and easier to understand. The final conclusion is that subtracting the mean helps lower the $L$ constant of the neural network's gradient, while dividing by the standard deviation mainly acts like an adaptive learning rate, making the updates of each parameter more synchronized so as to avoid overfitting on any particular layer or parameter.
Of course, the above interpretation is only a rough guide — fully explaining BN is a very difficult task, and BN's effect is more like the combined result of multiple factors. For instance, for our mainstream activation functions, $[-1, 1]$ is basically the interval where nonlinearity is strongest, so shaping the input to have mean 0 and variance 1 also lets the activation function's nonlinear capacity be exploited more fully, avoiding wasting the neural network's fitting capacity.
In short, theoretical analysis of neural networks is a very difficult undertaking, far beyond what I'm truly capable of — so all I can do here is write a blog post, tell a story that may or may not matter, and offer it up for others' amusement.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.