Why Do We Need Residuals? A Perspective from DeepNet

为什么需要残差？一个来自DeepNet的视角 苏剑林 · 2022-03-19 · Machine Learning model optimization deep learning gradient

In What's So Hard About Training a 1000-Layer Transformer? we introduced DeepNet, a technique proposed by Microsoft that manages to train 1000-layer Transformers. When it comes to DeepNet, readers generally react in one of two ways: some are impressed and give it a thumbs-up, while others find it an unremarkable rehash of old ideas. Those in the latter camp usually feel this way because DeepNet's two proposed improvements — enlarging the weight on the identity path and shrinking the initialization of the residual branch — are simply too mundane, and similar conclusions have appeared in other work before, so there's not much novelty to be found there.

Admittedly, judging purely by its conclusions, DeepNet isn't all that interesting. But in the author's view, the process behind DeepNet is far more important than its conclusions. What makes it interesting is that it offers a simple and effective way of analyzing gradient magnitudes, one that can be applied to many related problems — including the question this post wants to discuss, "why do we need residuals," for which it can give an answer that gets fairly close to the essence of the matter.

Increment Explosion

Why do we need residuals? The answer is that residuals make it easier to train deep models — and here "deep" might mean a hundred layers, a thousand layers, or even ten thousand layers. So the question becomes: why is it hard to train deep models without residuals?

Many readers would probably answer: vanishing or exploding gradients. These are indeed two important issues. However, with appropriate initialization methods and normalization techniques, we can already make the gradients of an ordinary feedforward neural network quite stable — and yet, even so, training a deep feedforward network remains difficult. This tells us that the underlying cause isn't just vanishing/exploding gradients, but something else as well — namely the "increment explosion" problem, which we already discussed in What's So Hard About Training a 1000-Layer Transformer?.

Understanding increment explosion isn't difficult. Suppose the loss function is $\mathcal{L}(\boldsymbol{\theta})$, with $\boldsymbol{\theta}$ being its parameters. When the parameters change from $\boldsymbol{\theta}$ to $\boldsymbol{\theta}+\Delta\boldsymbol{\theta}$:

\begin{equation}\Delta\mathcal{L} = \mathcal{L}(\boldsymbol{\theta}+\Delta\boldsymbol{\theta}) - \mathcal{L}(\boldsymbol{\theta}) \approx \langle\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta}),\Delta\boldsymbol{\theta}\rangle\end{equation}

For SGD we have $\Delta\boldsymbol{\theta}=-\eta \nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$, so $\Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert^2$. Suppose the model has $N$ layers, with an average parameter count of $K$ per layer. If the vanishing/exploding gradient problem has been solved, then we can assume each parameter's gradient is of order $\mathcal{O}(1)$, so we get $\Delta\mathcal{L}=\mathcal{O}(\eta NK)$. Thus, the size of the update at each step of training is proportional to the model depth $N$ (width isn't within the scope of this post). The deeper the model, the larger the update — which means that in the early stage of training, a deeper model is more likely to fall into a poor local optimum, causing training to stall or even collapse. This is the "increment explosion" problem.

Treating the Symptoms

In short, "increment explosion" means that as the number of layers grows, a tiny change in the parameters can trigger a huge change in the loss function — which is especially harmful during the early stages of training. A straightforward workaround here is warmup: start with an extremely small learning rate at first, then gradually increase it, so as to avoid learning too fast at the outset. Once the model safely gets through this initial "danger zone," training can proceed normally.

However, although warmup does help to some extent, it really only treats the symptoms rather than the underlying cause. The fact that "a tiny change in the parameters causes a huge change in the loss" means that the model itself is inherently unstable — in more technical terms, its loss landscape is extremely non-smooth, which is not a property a good model should have. Therefore, we should fix this by modifying the model itself, rather than resorting to a "surface-level" fix like lowering the learning rate.

By "modifying the model," we mean adjusting the model's architecture or its initialization scheme so as to naturally cancel out the effect that the number of layers $N$ has on the update magnitude. Based on the earlier result $\Delta\mathcal{L} \approx -\eta\Vert\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})\Vert^2$ and $\Delta\mathcal{L}=\mathcal{O}(\eta NK)$, to cancel out the effect of depth we need the gradient $\nabla_{\boldsymbol{\theta}}\mathcal{L}(\boldsymbol{\theta})$ to become order $\mathcal{O}(1/\sqrt{N})$. In other words, the gradient of each parameter needs to shrink as the number of layers grows.

Stable Propagation

If all we wanted to do were shrink the gradient, that would be simple — just lower the initialization variance as much as possible. But in practice, while shrinking the gradient we also need to preserve the stability of forward propagation, because forward-propagation stability encodes a kind of prior knowledge about the task we're solving — it corresponds to a better starting point for the model. As discussed in A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers, forward-propagation stability can be measured via second moments. For a simple linear layer,

\begin{equation}\boldsymbol{y} = \boldsymbol{x}\boldsymbol{W}, \quad \boldsymbol{x}\in\mathbb{R}^n, \boldsymbol{W}\in\mathbb{R}^{n\times m}\end{equation}

we already know that in order for the second moment of $\boldsymbol{y}$ to match that of $\boldsymbol{x}$, we need an initialization scheme with zero mean and variance $1/n$; if we account for the activation function, there's an extra constant-level scale factor — for instance, for the $\text{relu}$ activation, the variance becomes $2/n$. As for backpropagation, we have

\begin{equation}\frac{\partial\mathcal{L}}{\partial \boldsymbol{x}} = \frac{\partial\mathcal{L}}{\partial \boldsymbol{y}}\frac{\partial\boldsymbol{y}}{\partial \boldsymbol{x}} = \frac{\partial\mathcal{L}}{\partial \boldsymbol{y}}\boldsymbol{W}^{\top}\end{equation}

As we can see, backpropagation goes exactly the other way: if we want to stabilize the second moment of the backward pass, we need an initialization scheme with zero mean and variance $1/m$. Xavier initialization takes the average of the two, $2/(n+m)$; for more details, see Thoughts on the Dimension-Averaging Strategy for Non-Square Matrices in Initialization.

In other words, if we want to stabilize forward propagation, the initialization variance must be $1/n$, and the second moment of the backward pass then becomes $m/n$ times the original. Since $m,n$ are hyperparameters fixed in advance, with no necessary relationship to the number of layers, there's no way to use them to make the gradient shrink to $1/\sqrt{N}$ times its original size. This means that for a residual-free deep feedforward network,

\begin{equation}\phi_l(\phi_{l-1}(\phi_{l-2}(\cdots\phi_1(\boldsymbol{x}\boldsymbol{W}_1 + \boldsymbol{b}_1)\cdots)\boldsymbol{W}_{l-1} + \boldsymbol{b}_{l-1})\boldsymbol{W}_l + \boldsymbol{b}_l)\end{equation}

as long as its forward propagation is stable, its backward propagation is then also fixed, and we cannot make the gradient depend on the number of layers. So at best we can solve the vanishing/exploding gradient problem for a deep feedforward network, but we cannot solve the "increment explosion" problem mentioned at the start of this post — which is why deep feedforward networks are inevitably hard to train.

Enter the Residual

This is where residuals come to the rescue! Without loss of generality, assume the input and output dimensions are equal, and consider

\begin{equation}\boldsymbol{y} = \boldsymbol{x} + \varepsilon \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\theta})\end{equation}

Clearly, as long as $\varepsilon$ is small enough, forward propagation will necessarily be stable; and

\begin{equation}\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}} = \boldsymbol{I} + \varepsilon\frac{\partial \boldsymbol{f(\boldsymbol{x};\boldsymbol{\theta})}}{\partial \boldsymbol{x}}\label{eq:bp}\end{equation}

so we can likewise see that as long as $\varepsilon$ is small enough, backward propagation is also stable. As for the parameter gradient,

\begin{equation}\frac{\partial \mathcal{L}}{\partial \boldsymbol{\theta}} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{y}}\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{\theta}} = \varepsilon\frac{\partial \mathcal{L}}{\partial \boldsymbol{y}}\frac{\partial \boldsymbol{f(\boldsymbol{x};\boldsymbol{\theta})}}{\partial \boldsymbol{\theta}}\end{equation}

this shows that we can control $\varepsilon$ to achieve a depth-dependent gradient scaling! For instance, if we want the gradient scaled down to $1/\sqrt{N}$, we simply set $\varepsilon=1/\sqrt{N}$.

With this result in hand, we can now answer why we need residuals:

Because the residual structure is a design that simultaneously stabilizes both forward and backward propagation, while also allowing us to scale the parameter gradient so as to resolve increment explosion — and this is what enables us to train much deeper models.

Small Enough

We just said "$\varepsilon$ small enough" twice — but how small is small enough? Is $\varepsilon=1/\sqrt{N}$ enough?

Suppose we have a one-dimensional model; then $\frac{\partial y}{\partial x} = 1 + \varepsilon\frac{\partial f}{\partial x}$, and it's generally assumed that $\frac{\partial f}{\partial x}$ is $\mathcal{O}(1)$, so we can approximate the order of magnitude using $\frac{\partial y}{\partial x}=1+\varepsilon$. After propagating through $N$ layers, the "expansion factor" is then approximately $(1+\varepsilon)^N$. And we know that

\begin{equation}\left(1 + \frac{1}{N}\right)^N < \lim_{N\to\infty} \left(1 + \frac{1}{N}\right)^N = e\end{equation}

That is to say, for a one-dimensional model, in order for backward propagation not to explode as the number of layers grows, we would need at least $\varepsilon$ to be at least $\mathcal{O}(1/N)$ — and $\varepsilon=1/\sqrt{N}$ really isn't quite large enough.

However, for high-dimensional models the situation improves somewhat. If we multiply both sides of equation $\eqref{eq:bp}$ by an arbitrary vector $\boldsymbol{v}$:

\begin{equation}\boldsymbol{v}\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}} = \boldsymbol{v} + \varepsilon\boldsymbol{v}\frac{\partial \boldsymbol{f(\boldsymbol{x};\boldsymbol{\theta})}}{\partial \boldsymbol{x}}\end{equation}

note that in the early stage of training $\frac{\boldsymbol{f(\boldsymbol{x};\boldsymbol{\theta})}}{\partial \boldsymbol{x}}$ is also essentially a randomly initialized matrix with zero mean, and as we discussed in Understanding Model Parameter Initialization Strategies from a Geometric Perspective, such a matrix is close to (some multiple of) an orthogonal matrix. So in the early stage, $\boldsymbol{v}$ and $\varepsilon\boldsymbol{v}\frac{\partial \boldsymbol{f(\boldsymbol{x};\boldsymbol{\theta})}}{\partial \boldsymbol{x}}$ are nearly orthogonal, and hence

\begin{equation}\left\Vert\boldsymbol{v}\frac{\partial \boldsymbol{y}}{\partial \boldsymbol{x}}\right\Vert^2 = \mathcal{O}\big((1 + \varepsilon^2)\Vert\boldsymbol{v}\Vert^2\big)\end{equation}

In plain terms, this means that in the high-dimensional case, the expansion factor at each layer is closer to $1+\varepsilon^2$ rather than $1+\varepsilon$. Based on the result of the one-dimensional discussion, we only need $\varepsilon^2=\mathcal{O}(1/N)$, so $\varepsilon=1/\sqrt{N}$ is basically enough.

Summary

This post discussed the question of "why we need residuals." Inspired by DeepNet, we arrived at the conclusion that residuals can simultaneously stabilize forward propagation and backward propagation and resolve increment explosion, thereby making deep models much easier to train — whereas an ordinary residual-free feedforward network cannot solve all three of these problems at once, which is why it becomes hard to train once made deep.