A Theory of Saddle Escape in Deep Nonlinear Networks

[NOTE] Still a work in progress; some parts are rough.

Generally, when a network is trained from small initialization (i.e. all its weights are initially very small), we tend to see a pretty characteristic pattern: long plateaus where loss barely drops, with sharp drops as new features are acquired. This "saddle-to-saddle" structure is pretty well understood for deep linear networks (where the activation function is $\operatorname{Id}: x \mapsto x$). For nonlinear networks, why this happens remains mysterious; here we try to provide a result that shows that deep nonlinear networks are not that different than their linear counterparts.

Background#

Deep Linear Networks#

Much of the existing theory for optimization dynamics is based in the deep linear network model: an $L$-layer network without any activation function, just a long chain of matrix multiplications $\hat y = W_L \cdots W_1 x$. Although the model is simple, deep linear networks exhibit surprisingly nontrivial dynamics during gradient flow ($\dot W = - \nabla_W \mathcal L$) and many features of deep learning (e.g. saddle-to-saddle structure, low-rank bias, progressive feature learning) appear here in an analytically tractable form [SMG14].

A lot of the reason deep linear networks are so nice to analyze is due to the conservation law of the imbalance of the weight matrix norms between adjacent layers: $$\Delta_\ell \doteq \|W_{\ell+1}\|_F^2 - \|W_\ell\|_F^2.$$ For linear networks, one can show $\dot \Delta_\ell = 0$ for all $\ell$ [DHL18]. Intuitively, the network only ever sees the composition $W_{\ell+1}\sigma(W_\ell h)$, so if $\sigma$ is linear we can shrink $W_\ell$ by any factor $c$ and grow $W_{\ell+1}$ by $1/c$ without changing the output function. The loss then has a flat direction, and gradient flow cannot move along a direction in which the loss doesn't change, so whatever that rescaling moves is frozen.

In a broader sense, this follows for any activation function that satisfies Euler's homogeneity identity: $z \cdot \sigma'(z) = \sigma(z)$, e.g. for $\sigma(z) = z$ or ReLU.¹ So, if we define the functional $$\varphi_\sigma(z) \doteq z\sigma'(z) - \sigma(z),$$ for a 1-homogeneous function, it vanishes exactly. This is actually a much bigger deal than it seems; this conserved quantity is what makes linear network analysis work!

Nonlinear Activations#

For a general nonlinear activation $\sigma$, $\varphi_\sigma \not\equiv 0$ and our neat conservation law is not guaranteed.

To remedy this, let us study the Taylor expansion of $\varphi_\sigma$ near zero. Writing $\sigma(u) = \alpha u + a_q u^q + O(u^{q+1})$ with $\alpha = \sigma'(0) > 0$, the linear part cancels in $\varphi_\sigma$ (that is what Euler's identity says), and the first surviving term has order $q$: $$\varphi_\sigma(z) = (q-1)a_q z^q + O(z^{q+1}).$$ For $\operatorname{tanh}$, $\sigma(z) \approx z - z^3/3$, so $\varphi_{\tanh}(z) = z\mathrm{sech}^2(z) - \tanh(z) \approx -\tfrac{2}{3}z^3$, giving $q = 3$. For a quadratic activation $\sigma(z) = z + a_2 z^2$ one gets $\varphi_\sigma(z) = a_2 z^2$, so $q = 2$.

That's kinda the main idea underlying the whole paper. Near the origin (where our network is initialized, since we take initialization tending to zero), the pre-activations (i.e. the values passed into the activation) are small, so $\varphi_\sigma(z) \approx c_q z^q$ is high order small. Thus, the conservation law is only "slightly broken" and we can track exactly its drift to order $q$, which gives us enough control to do what we want to with our analysis. A similar argument in principle is the physics small angle approximation; $\sin\theta \neq \theta$, but the error is cubic, so for small swings we may use the linear theory and carry a correction to be bounded later.

Imbalance and Drift#

Consider an $L$-layer network $\hat y = W_L \sigma(W_{L-1} \cdots \sigma(W_1 x) \cdots)$ with pre-activations $z_\ell \doteq W_\ell \sigma(z_{\ell-1})$ and population loss $\mathcal{L}$.

The proof is quite natural; it follows by considering scaling before and after activation is applied; we outline it below.

Scaling layer $\ell$ and scaling layer $\ell+1$ both act on the same object $z_{\ell+1} = W_{\ell+1}\sigma(z_\ell)$, but they enter in different places; one multiplies the argument of $\sigma$, the other multiplies its output. Differentiating each at $s=0$, $$\frac{d}{ds}\Big[W_{\ell+1}\sigma\bigl((1+s)z_\ell\bigr)\Big] = W_{\ell+1}\bigl(z_\ell \odot \sigma'(z_\ell)\bigr), \qquad \frac{d}{ds}\Big[(1+s)W_{\ell+1}\sigma(z_\ell)\Big] = W_{\ell+1}\sigma(z_\ell).$$ Subtracting gives that the difference is $\varphi_\sigma$. The imbalance drifts at the rate by which $\sigma$ fails to commute with scaling.

Some important comments on the above: (1) the identity is exact and assumes basically nothing, and (2) the right-hand side is a correlation between the upstream gradient at layer $\ell+1$ and $\varphi_\sigma$ applied to the pre-activations at layer $\ell$. When $\varphi_\sigma \equiv 0$ we recover exact conservation; when it doesn't vanish, the drift is controlled by how large $\varphi_\sigma(z_\ell)$ is, and near the saddle that is $O(\varepsilon^q)$.²

Note that we are using an approximation that is valid near the origin in order to analyze the event which is fundamentally defined as taking the network away from the origin (learning features). So, "escape" is declared at a small fixed threshold rather than at the end of training. Then the time integral we take is dominated by the early part, where the network is deepest inside the regime the approximation covers.³ Most of the actual work in the paper is showing that this defect stays negligible along the entire trajectory, and not just at initialization.

Reducing the ODE#

Approximate conservation now gives us the same thing the linear theory has, plus an error we need to control. If the bottleneck layers (those initialized at scale $\varepsilon$) start at comparable scales, they will stay comparable for the whole escape, so instead of tracking $r$ matrices we can reduce the problem to only tracking their shared gain.

Note that these layer gains have no real reason to agree with each other, and once the escape is over they definitely don't, but through the transition itself they move in unison.

The other neat reduction we will want to make is that near the saddle every layer's gradient collapses to rank one, and every layer points along the same rank-one direction. So, we also want to track the network's overlap with the one direction it's being pushed in; call it $\gamma$. Then, we have the following relation between the layer gain and the overlap: $$\dot Y \asymp \Gamma\gamma, \qquad \dot\gamma \gtrsim \Gamma Y^{r-1},$$ where $Y$ is the squared gain of the first bottleneck layer and $\Gamma$ is the strength of the signal the data presents (we'll pin it down in a moment). This is fundamentally a feedback loop! The network can only grow its channel as fast as it is correlated with the target, and it can only become correlated with the target as fast as the channel lets signal through.⁴ Near the origin both are tiny, which is what causes the plateau. Integrating this loop gives us the main escape time law.

Critical-Depth Law#

Suppose $r$ of the $L$ layers initialize at scale $\varepsilon \to 0^+$ (call these the bottleneck layers) and the remaining $L - r$ initialize at scale $\Theta(1)$. By the reduction above, all $r$ bottleneck layers share a magnitude $y(t)$, with $Y = y^2$.

Now, note that the gradient driving any one bottleneck layer is the product of the signals passing through all of the other bottleneck layers, so it scales as $y^{r-1}$ (a layer does not appear in its own gradient). The $L-r$ full-size layers contribute $O(1)$ factors and only really change the prefactor. So the scalar ODE for the shared bottleneck magnitude is $$\dot y \asymp \Gamma y^{r-1},$$ and the time to get from $y(0) = \varepsilon$ out to a small fixed threshold $\vartheta$ is given by $$\tau \asymp \frac{1}{\Gamma}\int_\varepsilon^\vartheta y^{-(r-1)}dy.$$ Then, everything we need follows immediately from the behavior of that integral!

The $-(r-2)$ exponent has a pretty neat interpretation: one power of $\varepsilon$ is consumed because each layer's gradient is $y^{r-1}$ rather than $y^r$, and a second is absorbed by the integral. The threshold at $r=2$ is pretty interesting since it's the minimal bottleneck for which the escape time diverges at all as $\varepsilon \to 0$.

It is also worth remarking on the absence of width or total depth in the asymptotic theorem. In addition, the escape time law does not depend on the activation or number of features in the target. At first, this is somewhat surprising: width multiplies the number of directions the flow could move in, extra depth multiplies the number of compositions, and changing the activation changes the loss landscape. However, none of these matter because of the rank-one structure from the earlier section; since we can reduce the high-dimensional ODE to a 1D ODE, the only things that matter are the quantities given in the escape law.

Data Contribution#

Near the origin the activations are all evaluated close to zero, where $\sigma(u) \approx \alpha u$, so the network is a linear map to leading order. A linear map can only correlate with the linear part of the target, so every nonlinearity in the target is invisible to the network until it has already moved. So the only property of the data that can possibly set the initial clock is the input-output correlation $$\Gamma = \bigl\|\mathbb{E}[yx]\bigr\|.$$

It is important to note that $\Gamma$ is the norm of a vector sum. If the target is built from several features $v_k$ with amplitudes $\beta_k$, each one contributes a vector $\propto \beta_k u_k \hat a_k$, where $\hat a_k$ is the direction of $\Sigma v_k$, and these add with signs: $$\Gamma = h_\sigma\Bigl\|\textstyle\sum_k \beta_k u_k \hat a_k\Bigr\|.$$ Here $\Sigma = \mathbb{E}[xx^\top]$ is the input covariance, the target is $y = \sum_k \beta_k \sigma(v_k^\top x/s_k)$ for unit vectors $v_k$ with $s_k^2 = v_k^\top \Sigma v_k$, $u_k = \|\Sigma v_k\|/s_k$, and $h_\sigma = \mathbb{E}[g\sigma(g)]$ is the linear part of $\sigma$ for a standardized input $g$. Thus, two features whose covariance images point in opposing directions partially cancel one another. If you rotate them in favor of them cancelling out, $\Gamma$ decreases, and the network sits on the plateau proportionally longer, even though neither feature has become any weaker on its own. This also has the interesting consequence that a teacher direction that is misaligned with $\Sigma$ gets presented to the network more strongly than its own variance would suggest.⁵

Because $\Gamma$ is just a correlation, you never really need to decompose the target into features to compute it, and nothing in the definition dictates whether expectation is over a population or a finite sample. So, on a real dataset it is $\hat\Gamma = \bigl\|n^{-1}\sum_i y_i x_i\bigr\|$, which is computable from real data without any knowledge of the true data-generating function.

Feature Arrival Timing#

Looking at the first figure (the loss staircase), the obvious interpretation is that a single step corresponds to a single feature being learned. Our theory says when the network leaves the first saddle, but the picture above makes it tempting to assume that departure is really several departures, one per feature, that are averaged together. This isn't necessarily the case, however!

This part is a bit difficult to explain why without getting quite technical, but if you're interested, this is covered in Section 5 of the paper.