Type-II Saddles and Probabilistic Stability of Stochastic Gradient Descent

## Abstract

Abstract Characterizing and understanding the dynamics of stochastic gradient descent (SGD) around saddle points remains an open problem. We first show that saddle points in neural networks can be divided into two types, among which the Type-II saddles are especially difficult to escape from because the gradient noise vanishes at the saddle. The dynamics of SGD around these saddles are thus to leading order described by a random matrix product process, and it is thus natural to study the dynamics of SGD around these saddles using the notion of probabilistic stability and the related Lyapunov exponent. Theoretically, we link the study of SGD dynamics to well-known concepts in ergodic theory, which we leverage to show that saddle points can be either attractive or repulsive for SGD, and its dynamics can be classified into four different phases, depending on the signal-to-noise ratio in the gradient close to the saddle.

## 1 Introduction

Saddle points exist extensively in the loss function of neural networks. An important problem in deep learning theory is to characterize and understand the dynamics of SGD around saddle points. In this work, we divide the saddle point problem into two types: (1) Type-I, where the gradient noise is nonvanishing at the saddle, and (2) Type-II, where there is no gradient noise at the saddle. The main focus of our work is the stability of the Type-II saddles, which are more difficult to escape from using SGD. We show that these saddle points can be either attractive or repulsive for SGD, even if the Hessian matrix contains at least one negative eigenvalue. Recent progress in the learning dynamics of neural networks suggests that the trajectory of SGD is predominantly saddle-to-saddle [, ]. This means that SGD will stay close to many saddle points during training and makes it possible for SGD to actually “stop” at one of these saddle points at the end of training. It is thus important and urgent to understand the stability(^1^11Here, “stability” is a synonym of “attractivity” and should not be confused with the algorithmic stability.) of these saddle points.

Our main contributions are:

- 1.
we propose to classify saddle points in neural networks into two types, among which the Type-II saddles are shown to be difficult to escape;
- 2.
we propose using the probabilistic stability and Lyapunov exponents to study the attractivity of Type-II saddle points of neural networks; this bridges the conventional study of dynamical systems in control theory and ergodic theory to study the dynamics of SGD;
- 3.
we use the probabilistic stability to show that close to a Type-II saddle point, SGD has at least four different phases of learning, which we find to be relevant for understanding the initialization of neural nets.

This work is organized as follows. Section 2 discusses the most relevant works. Section 3 discusses the difference between the two types of saddle points encountered in training neural networks. Section 4 studies the probabilistic stability and Lyapunov exponent of SGD around these fixed points and connects to the well-known results in ergodic theory. Section 5 presents the experiments. All proofs and experimental details are presented in the Appendix.

## 2 Related Works

Dynamical stability. Dynamical stability centers around studying the attractivity and repulsiveness of fixed points. From an optimization perspective, stability can be seen as a worst-case guarantee for optimization because the training algorithm cannot converge to a fixed point that is unstable. Conceptually, the work closest to ours is Ref. [], which uses the variance of SGD to characterize the set of solutions that are preferred by SGD, namely the solutions that are flat and have a rather weak noise. However, as our result shows, moment-based notions of stability cannot be applied to saddle points. Instead, we propose to study this problem with the tools of probabilistic stability, which has been studied extensively in control theory [, , ]. The connection between Lyapunov exponents and probabilistic stability for the random matrix product problem is well known [], but its relevance for studying saddle points in deep learning is unclear because one rarely defines a “saddle point” for matrix product problems. Also, its relevance to deep learning has yet to be studied. Prior works have used similar notions to study the convergence to a local minimum [, ]. In contrast, our focus is on the stability and attractivity of saddle points.

Figure: Figure 1: Escaping from two types of saddles in a ReLU network under SGD. We see that for the escapes from type-I saddles, the escaping process starts immediately for every trajectory. For type-II saddles, the escape only starts significantly after the training starts, despite the gradient noise. The black dashed line shows an exponential fit to the type-II, implying the connection of the problem to Lyapunov exponents. See Appendix A.1 for details on the construction of these saddles.
Refer to caption: extracted/5702367/plots/two-types_of_saddle_inset.png

Saddle points. Due to the nonlinear nature of the loss function of neural networks, escaping from saddle points has been an important problem in deep learning [, , , , ]. As a primary example, Ref. [] mainly focuses on studying saddles from the perspective of gradient descent, where there is no difference between Type-I and Type-II saddles. Also, our focus is not on how to escape from saddles, but on the scientific question of whether a saddle point is attractive or not. See the next section for more discussion about prior literature on saddles.

## 3 Two Types of Saddle Points

Escaping from saddle points is a well-known and fundamental problem in the optimization of neural networks. Conventional wisdom believes that as long as the saddle points are strict (having at least one negative eigenvalue in the Hessian), escaping from it is easy []. This holds true for gradient descent training. However, it is unclear whether strictness is sufficient to guarantee an escape when the training proceeds with SGD, where the gradient noise confounds the training process.

Here, we propose to classify the saddle points into two types and show that the Type-II saddles are much more difficult to escape than the Type-I saddles. A point $\theta$ is said to be a saddle point when $\mathbb{E}[\nabla_{\theta}\ell]=0$ and $\theta$ is not a local minimum. This implies that

$$ $\mathbb{E}[\theta_{t+1}-\theta_{t}]=O(\|\Delta\theta\|),$ (1) $$

where $\|\Delta\theta\|$ is the Euclidean distance from the saddle and $t$ is the parameter at the $t$-th iteration of SGD. However, when stochasticity is taken into account, there are two types of saddles that satisfy this condition:

$$ $\theta_{t+1}-\theta_{t}=\lambda r_{t}+O(\|\Delta\theta\|)\quad\text{(Type-I saddle)};$ (2) $$

$$ $\theta_{t+1}-\theta_{t}=\lambda\hat{H}\Delta\theta_{t}+O(\|\Delta\theta\|^{2}) \quad\text{(Type-II saddle)},$ (3) $$

where $\lambda$ is the learning rate, $r_{t}$ is a random vector with a non-vanishing variance that is independent of $\theta$, and $\hat{H}$ is a random matrix, also independent of theta. The source of the randomness is due to the minibatch sampling over the data, and $\hat{H}$ can be easily identified as the sample-wise Hessian matrix. There is a crucial difference between the two types of saddles: Type-I saddles have a nonvanishing noise at the saddle, and this noise makes it possible for SGD to escape from the saddle. Meanwhile, the Type-II saddles have a vanishing noise at the saddle, making it more difficult to escape from. See Figure 1 for a comparison of escaping from the two types of saddles during neural network training. Here, the Type-II saddle is induced by the permutation symmetry between neurons in the hidden layer. This discussion motivates the following definition for Type-II saddles. Let $H(\theta)$ denote the Hessian of the loss at $\theta$, and $P$ denote the projection to the subspaces where $H$ has non-positive eigenvalues.

###### Definition 1 .

$\theta$ is a Type-II saddle of $\ell(\theta,x)$ if (1) $P\nabla_{\theta}\ell(\theta,x)=0$ for all $x$ and (2) $\theta$ is not a local minimum of $\mathbb{E}[\ell(\theta,x)].$

Setting $P=I$ recovers Eq. (3) and so this definition generalizes Eq. (3). Prior works suggest that Type-II saddle points exist abundantly in the loss function of neural networks and the matrix $P$ depends on the architecture of the network. The main cause of the saddle points in neural networks is the vast number of parameter symmetries, such as permutation symmetry [, , , ], rescaling symmetry [, ], and rotation symmetry [].
More recently, Ref. [] proves that close to these symmetry-induced saddle points, the SGD dynamics is always described by the type of dynamics in Eq. (4). Convergence to these symmetry-induced saddles has been identified as a main cause of collapsing of models into low-capacity solutions []. Literature regarding SGD escaping saddle points (e.g., [, , , ]) often takes advantage of the assumption of non-vanishing noise thus studies escaping type-I saddle. Thus, Type-II saddles are understudied even despite their relevance in deep learning.

A main result in [] shows that as long as the saddle is symmetry-induced, all escaping directions are contained within the subspace given by a projection matrix $P$, and the dynamics (3) take a low-rank form:

$$ $P\theta_{t+1}=P\theta_{t}-\lambda P^{T}\hat{H}(x_{t})P(\theta_{t}-\theta^{*}),$ (4) $$

where $\lambda$ is the learning rate, $\theta^{*}$ is the critical point under consideration, $\hat{H}(x_{t})$ is a random symmetric matrix that is a function of the random variables $x_{t}$, which can stand for both a single data point or a minibatch of data points.
Essentially, it is the data distribution of $\hat{H}$ that matters. When we have the same data distribution but a different batch size $S$, the distribution of $\hat{H}$ is different. For our purpose, it suffices to say that it is interesting and important to study the case when $\hat{H}$ has a certain low-rank structure with probability $1$, which is frequently the case in neural networks.

Lastly, while our focus is on the saddles, our theory also applies to interpolation minima in neural networks. Here, because the loss function for every data point vanishes, and the dynamics around it also obeys an identical form to Eq. (3). The dynamics we consider is more general than that around an interpolation minimum because the Hessians $\hat{H}$ in Eq. (4) are allowed to have both positive and negative eigenvalues, whereas an interpolation minimum only allows for positive semidefinite $\hat{H}$. Therefore, the general solution of Eq. (4) also helps us understand this type of minima better once we restrict the study to PSD Hessians.

###### Definition 1 .

## 4 Probabilistic Stability of SGD

In this section, we introduce the main theoretical framework. We first define the probabilistic stability, and then apply it to Type-II saddles.

### 4.1 Probabilistic Stability

We use $\|\cdot\|_{{p}}$ to denote the ${p}-$norm of a matrix or vector and $\|\cdot\|$ to denote the case where ${p}=2$. The notation $\to_{p}$ indicates convergence in probability.

###### Definition 2 .

A sequence (of random variables) $\{\theta_{t}\}_{t}^{\infty}$ generated by the SGD algorithm is probabilistically stable at a constant vector $\theta^{*}$ if $\theta_{t}\to_{p}\theta^{*}$.

Here, the notation $\to_{p}$ denotes convergence in probability. A sequence ${\theta}_{t}$ converges in probability to $\theta^{*}$ if $\lim_{t\to\infty}\mathbb{P}(\|\theta_{t}-\theta^{*}\|{>}\epsilon)=0$ for any $\epsilon>0$. While this notion of stability is a simple use of standard probability theory, it is a common tool in control theory to study stability. We will see that this notion of stability is especially suitable for studying the stability of the saddle points in SGD. The standard alternative is the notion of attractivity based on statistical moments.

###### Definition 3 .

A sequence $\{\theta_{t}\}_{t}^{\infty}$ is $L_{p}$-norm stable at $\theta^{*}$ if $\lim_{t\to\infty}\mathbb{E}\|\theta_{t}-\theta^{*}\|_{p}^{p}\to 0$.

For deep learning, the sequence of $\theta_{t}$ is the model parameters obtained by the iterations of the SGD algorithm for a neural network. For dynamical systems in general and deep learning specifically, it is impossible to analyze the convergence behavior of the dynamics starting from an arbitrary initial condition. Therefore, we have to restrict ourselves to the neighborhood of a given stationary point and consider the linearized dynamics around it.

Figure: Figure 2: SGD exhibits a complex phase diagram through the lens of probabilistic stability. Left: $a$ denotes the parameter in the data distribution, as discussed in detail in section 5.3. For a matrix factorization saddle point, the dynamics of SGD can be categorized into at least five different phases. Phase I, II, and IV correspond to a successful escape from the saddle. Phase III is where the model converges to a low-rank saddle point. Phase I corresponds to the case $w_{t}\to_{p}u_{t}$, which signals correct learning. In phase Ia, the model also converges in variance. Phase II corresponds to stable but incorrect learning, where $w_{t}\to_{p}-u_{t}$. Phase IV corresponds to complete instability. Right: the phases of SGD can quantified by the sign of the Lyapunov exponent $\Lambda$. Where $\Lambda<0$, SGD collapses to a saddle point; when $\Lambda>0$, SGD escapes the saddle and enters an escaping phase. The two escaping phases are qualitatively different. For a small learning rate, the model is in a learning phase due to the repulsiveness of the saddle point at a small learning rate, and the model is likely to converge to local minima close to the saddle. For a very large learning rate, SGD escapes the saddle due to the dynamical instability of SGD, and the model will move far away from the saddle. Besides, the magnitude of the Lyapunov exponent can also quantity the speed of the learning dynamics. See Appendix A.2 for numerical details of this example.
Refer to caption: extracted/5702367/plots/phase_diagram_thesis.png

###### Definition 2 .

###### Definition 3 .

### 4.2 Rank-1 Dynamics

An analytically solvable case is when the dynamics lie in a 1d manifold. Let $\hat{H}(x)=h(x)nn^{T}$ is rank-$1$ for a random scalar function $h(x)$, and a fixed unit vector $n$ for all data points $x$. Thus, the dynamics simplifies to a 1d dynamics, where $h(x)\in\mathbb{R}$ is the corresponding eigenvalue of $\hat{H}(x)$:

$$ $\theta_{t+1}=\theta_{t}-\lambda h(x)(\theta_{t}-\theta^{*}).$ (5) $$

###### Theorem 1 .

Let $\theta_{t}$ follow Eq. (5). Then, for any distribution of $h(x)$, $n^{T}(\theta_{t}-\theta^{*})\to_{p}0$ if and only if

$$ $\mathbb{E}_{x}[\log|1-\lambda h(x)|]<0.$ (6) $$

The l.h.s. of the inequality can be identified as the Lyapunov exponent for the process. The condition (6) is a sharp characterization of when a critical point becomes attractive. It also works with weight decay. When weight decay is present, the diagonal terms of $\hat{H}$ are shifted by $\gamma$, and so $h=h^{\prime}+\gamma$.

At what learning rate is the condition violated? To leading orders in $\lambda$, this can be identified by expanding the logarithmic term up to the second order in $\lambda$:

$$ $\mathbb{E}_{x}[\log|1-\lambda h(x)|]=-\lambda\mathbb{E}[h(x)]-\frac{1}{2} \lambda^{2}\mathbb{E}_{x}[h(x)^{2}]+{O}(\lambda^{3}).$ $$

Ignoring the second-order term, we see that the dynamics always follow the sign of $\mathbb{E}[h(x)]$, in agreement with the belief that the GD algorithm always escapes a strict saddle. When the second-order term is taken into consideration, the fluctuation of $h(x)$ now decides the attractivity of the critical point, which is attractive if

$$ $\lambda>2\frac{-\mathbb{E}[h(x)]}{\mathbb{E}[h(x)^{2}]}.$ (7) $$

This condition directly points to the reason why Type-II saddle points are difficult (or impossible) to escape: the critical point can be attractive even if $h<0$, when the point has become a saddle point.
The r.h.s. of the condition also has a natural interpretation as a signal-to-noise ratio (SNR) in the gradient. The numerator is the Hessian of the original loss function, which determines the signal in the gradient. The denominator is the strength of the gradient noise in the minibatch []. An illustration of this solution is given in Figure 2. We show the probabilistic stability conditions for a rank-$1$ saddle point with a rescaling symmetry (see Section 5.3). The loss function is $\ell(u,w)=-xyuw+o(u^{2}+w^{2})$. Here, the data points $xy=1$, and $xy=a$ are sampled with equal probability. These saddles appear naturally in matrix factorization problems and also in recent sparse learning algorithms [, , , ].

###### Theorem 1 .

### 4.3 Insufficiency of Norm-Stability

Theorem 1 provides a perfect example to compare the probabilistic stability with the norm-stability. First of all, it is easy to see that if SGD converges to a point in $L_{p}$-norm, it must converge in probability. Therefore, moment stability always implies probabilistic stability, but not vice versa. Thus, norm stability is a more restricted notion than probabilistic stability.

The following result shows that saddle points are always unstable for moment stability, which is constructively established by the following proposition.

###### Theorem 2 .

Let $\theta_{t}$ follow Eq. (4) around a critical point $\theta^{*}$ of an arbitrary loss function. Then, for any fixed $\lambda$,

- 1.
there exists a data distribution such that $\theta_{t}$ is probabilistically stable but not $L_{p}$-stable;
- 2.
if $\theta^{*}$ is a saddle point and $p\geq 1$, the set of $\theta_{0}$ that is $L_{p}$-stable has Lebesgue measure zero.

Therefore, the $L_{p}$-stability is not useful in understanding the stability of SGD close to saddle points. One reason is that the outliers strongly influence the $L_{p}$ norm in the data, whereas the probabilistic stability is robust against such outliers, a point we will discuss in the next section. This theorem points to a very special but interesting property of Type-II saddles. Namely, that studying Type-II saddles really requires some rather special notion of attractivity. While probabilistic stability is likely not the only way to study Type-II saddles, it is the most well-understood one given the prior literature on the Lyapunov exponents of random matrix products.

###### Theorem 2 .

### 4.4 Lyapunov Exponent and Probabilistic Stability

In reality, the notion of probabilistic stability is rather vague and cannot be directly experimentally explored. This problem can be avoided by noticing that the dynamics in Eq. (3) is the same as a random matrix product:

$$ $\theta_{t}-\theta^{*}=\prod_{i}^{d}Z_{t}(\theta_{0}-\theta^{*}),$ (8) $$

where $Z_{t}=I-\lambda\hat{H}_{t}$. Thanks to the celebrated Furstenberg-Kesten theorem [], one can prove that the condition for the probabilistic attractivity of the saddle point under SGD is the same as the Lyapunov exponent of the process being negative. The maximal Lyapunov exponent of a point $\theta^{*}$ is defined as

$$ $\Lambda=\max_{\theta_{0}}\lim_{t\to\infty}{\frac{1}{t}}\mathbb{E}\left[\log \frac{\|\theta_{t}-\theta^{*}\|}{\|\theta_{0}\|}\right].$ (9) $$

Here, the expectation is taken over the random samplings of the SGD algorithm. In general, $\Lambda$ does not vanish. The following theorem shows that SGD is probabilistically stable at a point if and only if its Lyapunov exponent is negative. The maximum Lyapunov exponent is initialization-independent. If $\bar{H}$ is a $d$-by-$d$ matrix, a well-known fact is that the initialization-dependent Lyapunov exponent takes at most $d$ distinctive values. It is also interesting to study the Lyapunov exponent in each of these $d$ subspaces.

###### Theorem 3 .

Assuming that $\Lambda\neq 0$, the linearized dynamics of SGD is probabilistically stable at $\theta^{*}$ for any $\theta_{0}$ if and only if $\Lambda<0$.

This result can also be proved by adapting the classical results in Ref. []. Note that for a finite dataset size, one can always choose a learning rate such that $I-\lambda\hat{H}$ is all positive with probability $1$, this further makes it possible to apply the central limit theorem of the classical Furstenberg-Kesten theory.

For a finite-size dataset, it is straightforward to find an upper and lower bound for $\Lambda$. we can define $r_{\max}$ to be larger than the absolute value of the eigenvalues of $I-\lambda\hat{H}(x)$ for all $x$ (which exists because there is only finitely many $x$). Similarly, we can define $r_{\min}>0$ to be smaller than all the absolute values of all the eigenvalues of $I-\lambda\hat{H}(x)$ for all $x$. Therefore,

$$ $\log r_{\min}<\Lambda<\log r_{\max}.$ (10) $$

While it is in general challenging but worthwhile to theoretically estimate the exponent [, , ], it is a quantity that can be easily experimentally studied, and studying how Lyapunov exponents change with respect to hyperparameters of training algorithms.

Now, we give two quantitative estimates about when the Lyapunov exponent will be negative. This discussion also implies a sufficient but weak condition for a general type of multidimensional dynamics to converge in probability. Let $h^{*}(x)$ be the largest eigenvalue of $\hat{H}(x)$ and assume that $1-h^{*}(x)>0$ for all $x$. Then, the following condition implies that ${\theta}\to_{p}0$: $\mathbb{E}_{x}[\log|1-\lambda h^{*}(x)|]<0$, which mimics the condition we found for rank-$1$ systems. An alternative estimate can be made by approximating every $\hat{H}(x)$ as a diagonal matrix, which is a common approximation in the literature that requires computing the Hessian or Fisher information of the neural networks. If $\hat{H}$ has rank $d$, this reduces the problem to $d$ separated rank-1 dynamics, and Theorem 1 again gives the exact solutions in each subspace. Numerical evidence shows that the diagonal approximation quite accurately predicts the onset of low-rank behavior the actual rank (see Section 5.3 and Appendix A.6).

Another potential concern is whether this theorem is trivial for SGD at a high dimension in the sense that it could be the case that $\Lambda$ could be identically zero independent of the dataset. One can show that the Lyapunov exponent is generally nonzero for all datasets that satisfy a mild condition. Let $\mathbb{E}[\hat{H}]$ be full rank. By definition,

$$ $\displaystyle\Lambda$ $\displaystyle=\lim_{t\to\infty}\frac{1}{t}\mathbb{E}\left[\log\theta_{0}^{T}FF ^{T}\theta_{0}\right]=-\frac{2\lambda\theta_{0}^{T}H\theta_{0}}{\|\theta_{0}\| ^{2}}+O(\lambda^{2}),$ $$

where $F=\prod_{j}^{t}(I-\lambda\hat{H}_{i_{j}})$.

Therefore, as long as $\lambda$ is sufficiently small, the sign of the Lyapunov exponent is opposite to the sign of the eigenvalues of $H$. This proves something quite general for SGD at an interpolation minimum: with a small learning rate, the model converges to the minimum exponentially fast, in agreement with common analysis in the optimization literature. See Figure 2-right for the numerical computation of Lyapunov exponents of a matrix factorization problem and the corresponding phases.

###### Theorem 3 .

## 5 Application and Experiment

### 5.1 Solving lasso with SGD

As a simple illustrative example, we first apply our theory to spred, a recent algorithm for solving the lasso problem with gradient descent []. Here, the target function is $\ell=((u\odot w)x-y)^{2}$, where $\odot$ is the element-wise product, $u,\ w,\ ,x\in\mathbb{R}^{200}$, and $y\in\mathbb{R}$. It has been proved that all local minima of this loss function is identitcal to the unique solution of lasso, and thus have the same level of sparsity. We choose the dataset such that the lasso solution has $0.5$ sparsity. The solutions sparser than $50\%$ are the saddle points where $u_{i}=w_{i}=0$ for some $i$. Apparently, each of these points is a Type-II saddle because the gradient of $w_{i}$ is proportional to $u_{i}$ and vice versa. The theory suggests that SGD is attracted to these saddle points at a large learning rate. See Figure 3.

Figure: Figure 3: At a small learningrate, the spred algorithm [] converges to the ground truth solution of lasso. At a large learning rate, however, it is biased towards sparser solutions than the ground truth because the sparser solutions are Type-II saddles and thus attractive for SGD at a large learning rate.
Refer to caption: extracted/5702367/plots/SGD_spred.png

### 5.2 Comparison with moment stability

Conventionally, it is thought that SGD can only select a local or global minimum as a solution. Our result suggests that saddle points can also be selected as solutions. In this section, we present one such case study that directly compares the attractivity of fixed points predicted by the probabilistic stability and the $L_{p}$ stability. We consider a two-layer network with a single hidden neuron with the swish activation function: $f(w,u,x)=u\times{\rm swish}(wx)$, where ${\rm swish}(x)=x\times{\rm sigmoid}(x)$. We generate $100$ data points $(x,y)$ as $y=0.1{\rm swish}(x)+0.9\epsilon$, where both $x$ and $\epsilon$ are sampled from normal distributions. See Figure 4 for an illustration of the training loss landscape. There are two local minima: solution A at roughly $(-0.7,-0.2)$ and solution B at $(1.1,-0.3)$. Here, the solution with better generalization is A because it captures the correct correlation between $x$ and $y$ when $x$ is small. Solution A is also the sharper one; its largest Hessian eigenvalue is roughly $h_{a}=7.7$. Solution B is the worse solution; it is also the flatter one, with the largest Hessian value being $h_{b}=3.0$. There is also a saddle point C at $(0,0)$, which performs significantly better than B and slightly worse than A in terms of generalization.

If we initialize the model at A, $L_{2}$ stability theory would predict that as we increase the learning rate, the solution moves from the sharper minimum A to the flatter minimum B when SGD loses $L_{2}$ stability in A; the model would then lose total stability once SGD becomes $L_{2}$-unstable at B. As shown by the red arrows in Figure 4. In contrast, probabilistic stability predicts that SGD will move from A to C as C becomes attractive and then lose stability, as the black arrows indicate. See the right panel of the figure for the comparison with the experiment for the model’s generalization performance. The dashed lines in the middle and right panels show the predictions of the $L_{2}$ stability and probabilistic theories, respectively. We see that the probabilistic theory predicts both the error and the place of transition right, whereas $L_{2}$ stability neither predicts the right transition nor the correct level of performance. If we initialize at B, the flatter minimum, $L_{2}$ stability theory would predict that the solution will only have one jump from B to divergence as we increase the learning rate. Thus, from $L_{2}$ stability, SGD would have roughly the performance of B until it diverges, and having a large learning rate will not help increase the performance. In sharp contrast, the probabilistic stability predicts that the solution will have two jumps: it stays at B for a small $\lambda$ and jumps to C as it becomes attractive at an intermediate learning rate. The model will ultimately diverge if C loses stability. Thus, our theory predicts that the model will first have a bad performance, then show a better performance at an intermediate learning rate, and finally diverge. See the middle panel of Figure 4. We see that the prediction of the probabilistic stability agrees with the experiment and correctly explains why SGD leads to better performance.

Figure: Figure 4: How SGD selects a solution. Left: The landscape of a two-layer network with the swish activation function []. The black arrow corresponds to the experimental trajectory and the prediction of probabilistic stability, while the red arrow corresponds to the (false) prediction of the $L_{2}$ stability. Middle, Right: the generalization performance of the model for different learning rates. Middle: Initialized at solution B, SGD first jumps to C and then diverges. Right: Initialized at A, SGD also jumps to C and diverges. In both cases, the behavior of SGD agrees with the prediction of the probabilistic stability instead of the $L_{2}$ stability. Instead of jumping between local minima, SGD, at a large learning rate, transitions from minima to saddles.
Refer to caption: extracted/5702367/plots/swish_example.png

### 5.3 Phases of Learning

A good observation is that almost all neural networks are initialized close to a Type-II saddle. Let $f(x)=W^{(D)}\sigma(W^{(D-1)}\sigma(...\sigma(W^{(1)}x+b^{(1)}))+b^{(D-1)})$
be a generic neural network with depth $D$ and activation $\sigma$, and $W^{(D)}\in\mathbb{R}^{d_{D-1}\times d_{y}}$, $W^{(i)}\in\mathbb{R}^{d_{i-1}\times d_{i}}$, and $W^{(1)}\in\mathbb{R}^{d_{x}\times d_{1}}$ are of arbitrary dimensions that match the input and output dimensions. Here, $\sigma(x)=c_{0}x+O(x^{2})$ is any nonlinearity that is locally linear at $x=0$. Let $y=y(x)$ be the label of $x$. Assuming that the per-sample loss is differentiable, one can prove the following result, which is relevant for the standard small initialization of neural networks.

###### Theorem 4 .

For any data distribution, loss function $\ell$, and any $\sigma$, the point $(u_{i},w_{i})=(0,0)$ is a Type-II saddle of $\ell(f(x),y(x))$ if it is not a local minimum.

The details of the proof can be found in section B.4. Now, we solve an analytical model that exhibits complex and dramatic phase transition-like behaviors as we change the learning rate of SGD. We will see that the solutions we obtained from this model are relevant for understanding the initialization of real neural networks.

Figure: Figure 5: Phase diagram at $N=100$. The colors indicate the phases as in Figure 2. At a finite-size, the phase boundaries have a fractal-like structure. The bottom-left of the phase diagram has a smooth boundary and is shared across almost all phase diagrams we plotted.
Refer to caption: extracted/5702367/plots/phase_diagram_100.png

As an analytically tractable model, let the loss function be $\mathbb{E}_{x}[(\sum_{i}u^{(i)}\sigma(w^{(i)}x)-y)^{2}/2]$. We let $c_{0}=1$ and both $x$ and $y$ be one-dimensional. At $u,\ w\approx 0$, the model $u^{T}w$ is either rank-$1$ or rank-0 0. The rank-0 0 point where $u^{(i)}=w^{(i)}=0$ for all $i$ is a saddle point as long as $\mathbb{E}[xy]\neq 0$. Consider the linearized dynamics around the saddle at $w^{(i)}=u^{(i)}=0$. The expanded loss function takes the following form: $\ell(u,w{;x,y})=-{xy}\sum^{d}_{i}u^{(i)}w^{(i)}+const.$ For learning to happen, SGD needs to escape from the saddle point. Let us consider a simple data distribution where $xy=1$ and $xy=a$ with equal probability. When $a>-1$, correct learning happens when ${\rm sign}(w)={\rm sign}(u)$. We thus focus on the case of $a>-1$. The case of $a<-1$ is symmetric to this case up to a rescaling. This example is already presented in Figure 2. There are five phases of learning in this simple example

- •
Ia. correct learning with prob. and norm stability ($w_{t}-u_{t}\to_{L_{2}}0$, $w_{t}+u_{t}$ diverges);
- •
Ib. correct learning with prob. but not norm stability ($w_{t}-u_{t}\to_{p}0$, $w_{t}-u_{t}\not\to_{L_{2}}0$, $w_{t}+u_{t}$ diverges);
- •
II. incorrect learning under probabilistic stability ($w_{t}-u_{t}$ diverges, $w_{t}+u_{t}\to_{p}0$);
- •
III. convergence to low-rank saddle point ($w_{t}-u_{t}\to_{p}0$, $w_{t}+u_{t}\to_{p}0$);
- •
IV. completely unstable ($w_{t}+u_{t}$, $w_{t}-u_{t}$ diverges in p.).

Two aspects agree well with the theory: (1) SGD can indeed converge to low-rank saddle points; however, this happens only when the gradient noise is sufficiently strong and when the learning rate is large (but not too large); (2) the region for convergence to saddles (region III) is exclusive with the region for convergence in mean square (Ia), and thus one can only understand the saddle-seeking behavior of SGD within the proposed probabilistic framework. Let $B$ denote a mini-batch and $S$ be its cardinality. The following proposition precisely characterizes the phase boundaries.

###### Proposition 1 .

For any $w_{0},u_{0}\in\mathbb{R}/\{0\}$. $w_{t}-u_{t}\to_{p}0$ if and only if $\mathbb{E}_{B}[\log|1-\lambda\sum_{(x,y)\in B}xy{/S}|]<0$. $w_{t}+u_{t}$ converges to 0 0 in probability if and only if $\mathbb{E}_{B}[\log|1+\lambda{\sum_{(x,y)\in B}}xy{/S}|]<0$.

This shows that the phase diagram of SGD strongly depends on the data distribution, and it is interesting to explore and compare a few different settings. Now, we consider a size-$N$ Gaussian dataset. Let $x_{i}\sim\mathcal{N}(0,1)$ and noise $\epsilon_{i}\sim\mathcal{N}(0,4)$, and generate a noisy label $y_{i}=\mu x_{i}+(1-\mu)\epsilon_{i}$. See the phase diagram for this dataset in Figure 5. We see that the phase diagram has a rich structure at a finite size. There are three rather surprising observations about the phase diagrams: (1) as $N\to\infty$, the phase diagram becomes smoother and smoother and each phase takes a connected region (cf. experiments in Appendix A.3); (2) phase II seems to disappear as $N$ becomes large; (3) the lower part of the phase diagram seems universal, taking the same shape for all samplings of the datasets and across different sizes of the dataset. This suggests that the convergence to low-rank structures can be a universal aspect of SGD dynamics, which corroborates the widely observed phenomenon of collapse in deep learning [, , ]. The theory also shows that if we fix the learning rate and noise level, increasing the batch size makes it more and more difficult to converge to the low-rank solution (see Figure 11, for example). This is expected because the larger the batch size, the smaller the effective noise in the gradient.

###### Theorem 4 .

###### Proposition 1 .

### 5.4 Initialization of Deep Neural Networks

As Theorem 4 shows, almost all neural networks are close to Type-II saddles at initialization. This means that they will also have significant problems in initiating training if the SNR in the gradient is small. To verify this, we construct and modify datasets to exhibit different SNR during training for deep nonlinear networks.

We start with a controlled experiment where, at every training step, we sample input $x\sim\mathcal{N}(0,I_{200})$ and noise $\epsilon\sim\mathcal{N}(0,4I_{200})$, and generate a noisy label $y=\mu x+(1-\mu)\epsilon$. Note that $1-\mu$ controls the level of the noise. Training proceeds with SGD on the MSE loss. We train a two-layer model with the architecture: $200\to 200\to 200$. See Figure 7 for the theoretical phase diagram, which is estimated under the diagonal approximation. Under SGD, the model escapes from the saddle with a finite variance to the right of the dashed line and has an infinite variance to its left. In the region $\lambda\in(0,0.2)$, this loss of the $L_{2}$ stability condition coincides with the condition for the convergence to the saddle. The experiment shows that the theoretical boundary agrees well with the numerical results. The Adam optimizer [] also have a similar phase diagram. See Appendix A. This suggests that the effects we studied are rather universal, not just a special feature of SGD.

Then, we train independently initialized ResNets on CIFAR-10 with SGD. The training proceeds with SGD without momentum at a fixed learning rate and batch size $S=32$ (unless specified otherwise) for $10^{5}$ iterations. Our implementation of Resnet18 contains $11$M parameters and achieves $94\%$ test accuracy under the standard training protocol, consistent with the established values. To probe the effect of noise, we artificially inject a dynamical label noise during every training step, where, at every step, a correct label is flipped to a random label with probability $noise$, and we note that the phase diagrams are similar regardless of whether the noise is dynamical or static (where the mislabelling is fixed). See Figure 10 for the phase diagram of static label noise. Interestingly, the best generalization performance is achieved close to the phase boundary when the noise is strong. This is direct evidence that SGD noise has a strong regularization effect on the trained model. We also study the sparsity of the ResNets in different layers in Figure 7, and we observe that the phase diagrams are all qualitatively similar. Also, see Appendix A for the experiment with a varying batch size. These experiments show that Type-II saddles are indeed a major obstacle in the initial phase of training. This result may also explain why the warmup is needed for training transformers. Training transformers is only successful if one starts with a small learning rate and increases it slowly toward a maximal value. If the training starts at a large learning rate, one often observes that the model becomes trapped in a low-capacity state soon after initialization. This agrees with the theoretical expectation that a small learning rate has a larger SNR for training.

Figure: Figure 6: Convergence to low-rank solutions in nonlinear neural networks. At every training step, we sample input $x\sim\mathcal{N}(0,I_{200})$ and noise $\epsilon\sim\mathcal{N}(0,\sqrt{2}I_{200})$, and generate a noisy label $y=\mu x+(1-\mu)\epsilon$, where $1-\mu$ controls the level of the noise. We compute the rank of the second layer of the weight matrix after training. Left: Linear network. Right: tanh network. The white dashed line shows the theoretical prediction of the appearance of low-rank structure computed by numerically integrating the condition in Proposition 1.
Refer to caption: extracted/5702367/plots/saddle_linear.png

## 6 Discussion

We have shown that special attention should be paid to the study of Type-II saddle points in neural networks. To this purpose, we demonstrate that the probabilistic stability serves as an essential notion for understanding the attractivity of Type-II saddle points as no moment-based notions of stability can be used to study when a saddle point becomes attractive. These effects are only present for SGD and not for GD, demonstrating that the algorithmic regularization due to SGD is qualitatively different from that of GD. We also link the saddle point problem to the conventional study of Lyapunov exponents in ergodic theory and dynamical systems, both of which can now be leveraged for us to design better algorithms of training. Our result also sheds light on how a solution in the loss landscape is chosen by SGD. In the probabilistic-stability perspective, SGD performs a selection between converging to saddles and to local minima, not between sharp minima and flat ones. A main limitation of our work is that we only considered simple analytical models on the theory side, but we believe they serve as good examples to connect contemporary deep learning research to the conventional fields of dynamical systems.

We have also suggested a new approach to the saddle point problem in neural networks. Depending on how one views the problem of Type-II saddles, there are two interesting future directions. If we regard the Type-II saddles as major obstacles to optimization, an important future problem is to design algorithms to better escape from them or to remove them from the loss landscape. Alternatively, the type-II saddles can function as effective capacity constraints on the models, and one might also be interested in leveraging them to better regularize neural networks.

## Acknowledgment

Liu Ziyin is supported by the JSPS fellowship during the writing of this paper. The research of B. Li is funded by the Chair Capital Markets Tomorrow: Modeling and Computational Issues under the aegis of the Institut Europlace de Finance, a joint initiative of Laboratoire de Probabilités, Statistique et Modélisation (LPSM) / Université Paris Cité and Crédit Agricole CIB, with the support of Labex FCD (ANR-11-LABX-0019-01). This work was supported by KAKENHI Grant No. JP22H01152 from the Japan Society for the Promotion of Science. We gratefully acknowledge the support from the CREST program “Quantum Frontiers” (Grant No. JPMJCR23I1) by the Japan Science and Technology Agency.

## Appendix A Experimental Concerns

### A.1 Setting for Figure 1

The following derivation gives the construction of the Type-I and Type-II saddle points in Figure 1. Let us first study the case when there is a single neuron. The network is

$$ $f(x,\theta)=\sum_{i}u_{i}\sigma(w_{i}^{T}x),$ (11) $$

where $\sigma={\rm ReLU}$ is the activation, and $x$ contains an element of $1$ to include the effect of having a bias.

The loss function is $L=\frac{1}{N}\sum_{x}\ell(x)$, where

$$ $\ell(x)=(f(x)-y(x))^{2}$ (12) $$

Without loss of generality, we can write $\mathbb{E}_{x}$ in the place of $\frac{1}{N}\sum_{x}$, where the expectation is taken over the training set. Let $\sigma={\rm ReLU}$. We have that $\sigma^{\prime\prime}=0$. The gradients are

$$ $\nabla_{u}\ell(x)=\mathbbm{1}_{>}w^{T}x(uw^{T}x-y),$ (13) $$

$$ $\nabla_{w}\ell(x)=\mathbbm{1}_{>}u(uw^{T}x-y)x.$ (14) $$

where we have used $\mathbbm{1}_{>}$ as an indicator function of the event $w^{T}x\geq 0$. The Hessian is

$$ $H(x)=\begin{bmatrix}\mathbbm{1}_{>}(w^{T}x)^{2}&\mathbbm{1}_{>}w^{T}xux+ \mathbbm{1}_{>}(u(w^{T}x)-y)x^{T}\\ \mathbbm{1}_{>}w^{T}xux+\mathbbm{1}_{>}(u(w^{T}x)-y)x&\mathbbm{1}_{>}u^{2}xx^{ T}\end{bmatrix}$ (15) $$

There are two cases: (1) $u=0$ and (2) $u\neq 0$.

For the case $u=0$, we have that $\nabla_{w}L=0$. For $\nabla_{u}L=0$, we must have

$$ $\nabla_{u}L=\mathbb{E}[\mathbbm{1}_{>}w^{T}x(-y)]=0$ (16) $$

which is solved by any $w$ such that $w^{T}\mathbb{E}[\mathbbm{1}_{>}xy]=0$. These solutions are all saddle points. To see this, note that the sample Hessian is:

$$ $H(x)=\begin{bmatrix}\mathbbm{1}_{>}(w^{T}x)^{2}&-\mathbbm{1}_{>}yx^{T}\\ -\mathbbm{1}_{>}yx&0\end{bmatrix},$ (17) $$

The expected Hessian is:

$$ $\mathbb{E}[H(x)]=\begin{bmatrix}\mathbb{E}[\mathbbm{1}_{>}(w^{T}x)^{2}]&- \mathbb{E}[\mathbbm{1}_{>}yx^{T}]\\ -\mathbb{E}[\mathbbm{1}_{>}yx]&0\end{bmatrix},$ (18) $$

which is a rank-2 matrix whose eigenproblem is solved by eigenvalues

$$ $\lambda_{\pm}=\frac{1}{2}\left(\mathbb{E}[\mathbbm{1}_{>}(w^{T}x)^{2}]\pm\sqrt {\mathbb{E}[\mathbbm{1}_{>}(w^{T}x)^{2}]^{2}+\|\mathbb{E}[\mathbbm{1}_{>}yx^{T }]\|^{2}}\right),$ (19) $$

and eigenvectors:

$$ $v_{\pm}=(\lambda_{\pm},-\mathbb{E}[\mathbbm{1}_{>}yx^{T}]).$ (20) $$

Obviously, $\lambda_{-}<0$, and these solutions are strict saddle points, except for the directions of $w$ that $\mathbb{E}[\mathbbm{1}_{>}yx^{T}]=0$.

Now, a crucial observation is that the saddle points given by this condition can be divided into two classes with dramatically different properties under a stochastic optimization algorithm: (I) $w\neq 0$ and (II) $w=0$. For the type I saddle, we have that the sample-wise gradient is

$$ $\nabla_{u}\ell=\mathbbm{1}_{>}w^{T}x(uw^{T}x-y),$ (21) $$

which is a random variable that does not vanish in general. This random noise facilitates the escape from the saddle.

For the type II, we have that with probability $1$,

$$ $\nabla_{u}\ell=0,$ (22) $$

which signals a vanishing noise, the lack of which means that this type of saddle is difficult to escape.
Alternatively, let $u=u_{0}+\Delta u$ and $w=w_{0}+\Delta w$ denote a small deviation from the saddle $(u_{0},w_{0})$. Then,

$$ $\nabla_{u}\ell(x)=\begin{cases}\Theta(1)&\text{if $w_{0}\neq 0$};\\ O(\|\Delta w\|)&\text{if $w_{0}=0$}.\end{cases}$ (23) $$

$$ $\nabla_{w}\ell(x)=O(\|\Delta u\|).$ (24) $$

### A.2 Experimental Detail for Figure 2 -right

The experiment is performed for a two-dimensional system whose dynamics is specified in (4). The expectation of the Hessian $\mathbb{E}[\hat{H}]$ is chosen to be $\text{diag}(0.1,-0.1)$, while the noise is generated via a normal $2\times 2$ random matrix $M_{\text{noise}}$. The noisy Hessian is obtained as

$$ $\hat{H}=\mathbb{E}[\hat{H}]+M_{\text{noise}}+M_{\text{noise}}^{T},$ (25) $$

and one can verify that such $\hat{H}$ is symmetric and consistent with our choice of $\mathbb{E}[\hat{H}]$. The initial state is sampled from a unit circle. The dynamics stops at time step $t$, and the Lyapunov exponent is calculated as $\frac{1}{t}\log||\theta_{t}||$, if one of the three following conditions is satisfied: $||\theta||$ reaches the upper cutoff of $10^{100}$; $||\theta||$ reaches the lower cutoff of $10^{-140}$; the preset maximal number of steps of $5000$ is reached. For each learning rate, the Lyapunov is obtained as the average of the results collected in $800$ independent runs.

### A.3 Phases of Finite-Size Datasets

See Figure 8.

Figure: Figure 8: Phase diagrams of SGD stability for finite-size dataset. The data sampling is the same as in Figure 9. From upper left to lower right: $N=3,\ 4,\ 8,\ 10,\ 24,\ 100$. As the dataset size tends to infinity, the phase diagram converges to that in Figure 9. The lower parts of all the phase diagrams look similar, suggesting a universal structure.
Refer to caption: extracted/5702367/plots/phase_diagram_3.png

Figure: Figure 9: Phase diagrams of SGD stability. The definitions of the phases are the same as those in Figure 2. We sample a dataset of size $N$ such that $x\sim\mathcal{N}(0,1)$ and noise $\epsilon\sim\mathcal{N}(0,4)$, and generate a noisy label $y=\mu x+(1-\mu)\epsilon$. Left: the $\lambda-\mu$ (noise level) phase diagram for $S=1$ and $N=\infty$. Right: The $\lambda-S$ (batch size) phase diagram for $\mu=0.06$ and $N=\infty$.
Refer to caption: extracted/5702367/plots/phase_diagram_300000_edited.png

### A.4 Resnet18 with Static Label noise

See Figure 10, where the labels of the training set is corrupted at a fixed probability to a different fixed label during training. We see that the problem of initialization remains and features a similar phase boundary.

Figure: Figure 10: Rank (left) and test accuracy (right) of the ResNet18 trained in a data set with static noise. The transition of rank has a clear boundary. The model has a full rank but random-guess level performance for large noise and small learning rates. Here, noise refers to the probability that the data point is mislabeled.
Refer to caption: extracted/5702367/plots/cifar10_static_noise.png

### A.5 Effect of Changing Batch size on Resnet18

See Figure 11.

Figure: Figure 11: The rank of the penultimate-layer representation of Resnet18 trained with different levels of batch sizes. In agreement with the phase diagram, the model escapes from the low-rank saddle as one increases the batch size.
Refer to caption: extracted/5702367/plots/resnet_bs.png

### A.6 Diagonal Approximation

Here, we compare the empirical rank of the solution with the diagonally approximated critical learning rates obtained in (7). See Figure 12. The experiment is run on a two-layer fully connected linear network: $50\to 50\to 50$, which is equivalent to a matrix factorization problem. The model is initialized with the standard Kaiming init. The dataset we consider is one with a sparse but full-rank signal.

Let $\odot$ denote the Hadamard product. The input data is generated as $x=m\odot X$, where $X\sim\mathcal{N}(0,I_{50})$ and $m$ is a random mask where a random element is set to be $1$, and the rest is zero. The labels $Y$ are generated as $Y=\mu x+(1-\mu)(m\odot\epsilon)$, where $\epsilon\sim\mathcal{N}(0,2\text{diag}(0.01,0.05,....,2.01))$ is the noise.

Figure: Figure 12: Rank vs learning rate in a vanilla matrix factorization problem for, from upper left to lower right, $\mu=0.05,\ 0.15,\ 0.25,\ 0.35$. The theoretical curve is from the diagonal approximation where each subspace of the model collapses at the critical learning rate $\lambda=-2\frac{\mathbb{E}[h(x)]}{\mathbb{E}[h^{2}(x)]}$.
Refer to caption: extracted/5702367/plots/RankLR/synopsis5.png

Figure: Figure 13: Norms of gradient ($g$) of a matrix factorization problem trained with SGD. $n$ is a low-rank direction after training. During training, it is commonly the case the model does not converge to a stationary point but to a stationary distribution. Our theory is compatible with this case because it is possible and common for the model to converge to a point in some subspace, even if it is not converging to a point overall.
Refer to caption: extracted/5702367/plots/stationary_distribution.png

### A.7 Convergence to Stationary distributions

As we mentioned in the main text, our theory is compatible with the case when the model converges to a stationary distribution but not a stationary point, which is more commonly the case during actual deep learning practice []. The experimental setup is the same as in the previous section.

See Figure 13, where we plot the norm of the gradient $g$ and the norm of $n^{T}g$, where $n$ is a low-rank direction after training. Here, we see that the norm of the gradient does not converge to zero, but to a positive value, signaling a convergence to a stationary distribution. At the same time, the norm of $n^{T}g$ does converge to zero, which means that in some subspace, the parameters do converge to a point. This justifies our argument.

Figure: Figure 14: Rank of the converged solution for two-layer linear (upper left), tanh (upper right), relu (lower left) and swish (lower right) models.
Refer to caption: extracted/5702367/plots/Adam/Adam_linear.png

### A.8 Experiment with Adam

We note that the phenomena we studied is not just a special feature of the SGD, but, empirically, seems to be a universal feature of first-order optimization methods that rely on minibatch sampling. Here, we repeat the experiment in Figure 9. We train with the same data and training procedure, except that we replace SGD with Adam [], the most popular first-order optimization method in deep learning. Figure 14 shows that similarly to SGD, Adam also converges to the low-rank saddles in similar regions of learning rate and $\mu$.

## Appendix B Proofs

### B.1 Proof of Theorem 1

###### Proof.

Consider Eq. (4):

$$ $\theta_{t+1}=\theta_{t}-\lambda\hat{H}(x)(\theta_{t}-\theta^{*}).$ (26) $$

Defining $w_{t}=\theta_{t}-\theta^{*}$, this equation can be written as

$$ $w_{t+1}=w_{t}-\lambda\hat{H}(x)w_{t},$ (27) $$

which is mathematically equivalent to the case when $\theta^{*}=0$. Therefore, without loss of generality, we write the dynamics in the form of Eq. (27) in this proof and the rest of the proofs.

Now, when $\hat{H}\propto nn^{T}$ is rank-$1$, we can multiply $n^{T}$ from the left on both sides of the dynamics to obtain

$$ $n^{T}w_{t+1}=n^{T}w_{t}-\lambda h(x)n^{T}w_{t}.$ (28) $$

The dynamics thus becomes one-dimensional in the direction of $n^{T}$.

Let $h_{t}$ denote the eigenvalue of the Hessian of the randomly sampled batch at time step $t$. The dynamics in Eq. (5) implies the following dynamics

$$ $\|n^{T}w_{t+1}\|/\|n^{T}w_{t}\|=|1-\lambda h_{t}|,$ (29) $$

which implies

$$ $\|n^{T}w_{t+1}\|/\|n^{T}w_{0}\|=\prod_{\tau=1}^{t}|1-\lambda h_{\tau}|.$ (30) $$

We can define auxiliary variables $z_{t}:=\log(\|n^{T}w_{t+1}\|/\|n^{T}w_{0}\|)-m$ and $m:=\mathbb{E}[\log(\|n^{T}w_{t+1}\|/\|n^{T}w_{0}\|)]=t\mathbb{E}_{x}[\log|1-
\lambda h_{t}|]$. Let $\epsilon>0$. We have that

$$ $\displaystyle\mathbb{P}(\|n^{T}w_{t}\|<\epsilon)$ $\displaystyle=\mathbb{P}(\|n^{T}w_{0}\|e^{z_{t}+m}<\epsilon)$ (31) $\displaystyle=\mathbb{P}\left(\frac{1}{{t}}z_{t}<\frac{1}{{t}}(\log\epsilon/\| n^{T}w_{0}\|-m)\right)$ (32) $\displaystyle=\mathbb{P}\left(\frac{z_{t}}{t}<-\mathbb{E}_{x}\log|1-\lambda h_ {t}|+o(1)\right).$ (33) $$

By the law of large numbers, the left-hand side of the inequality converges to 0 0, whereas the right-hand side converges to a constant. Thus, we have, for all $\epsilon>0$,

$$ $\lim_{t\to\infty}\mathbb{P}(\|n^{T}w_{t}\|<\epsilon)=\begin{cases}1&\text{if $ m<0$};\\ 0&\text{if $m>1$}.\end{cases}$ (34) $$

This completes the proof.
∎

###### Proof.

### B.2 Proof of Proposition 2

###### Proof.

Part 2 of the proposition follows immediately from Proposition 2, which we prove below. Here, we prove part 1.

It suffices to consider a dataset with two data points for which $h(x_{1})=1/\lambda$ and $h(x_{2})=c_{0}$, where each data point is sampled with equal probability. Let $c_{0}$ be such that

$$ $|1-\lambda c_{0}|^{p}>\frac{1}{2}.$ (35) $$

Now, we claim that this dynamics converges to zero in probability. To see this, note that

$$ $\|z_{t+1}\|=\begin{cases}\|z_{t}\||1-\lambda/\lambda|=0&\text{with probability $0.5$};\\ \|z_{t}\||1-\lambda c_{0}|&\text{with probability $0.5$}.\end{cases}$ (36) $$

Therefore, at time step $t$, $\mathbb{P}(z_{t}=0)\geq 1-2^{-t}$, which converges to 0 0. This means that $z_{t}$ converges in probability to 0 0.

Meanwhile, the $p$-norm is

$$ $\displaystyle\mathbb{E}[\|z_{t+1}\|^{p}]$ $\displaystyle=\frac{1}{2}\mathbb{E}[\|z_{t}\|^{p}]|1-\lambda c_{0}|^{p}$ (37) $\displaystyle\propto\frac{1}{2^{t}}|1-\lambda c_{0}|^{pt}\to 0.$ (38) $$

The convergence to zero follows from the construction that $|1-\lambda c_{0}|^{p}>\frac{1}{2}$. This completes the proof.
∎

###### Proposition 2 .

(No convergence in $L_{p}$.) For every strict saddle point $\theta^{*}$, there exists an initialization $\theta_{0}$ such that for any $\lambda\in\mathbb{R}_{+}$ and distance function $f(\cdot,\theta^{*})$, $\theta^{*}$ is unstable in $f$.

###### Proof.

This problem is easy to prove when $\theta$ is one-dimensional. For a high-dimensional $\theta$, the dynamics of SGD is

$$ $\theta_{t+1}=(I-\lambda\hat{H}_{t})\theta_{t}.$ (39) $$

Note that the expected value of $\theta_{t}$ is the same as the gradient descent iterations:

$$ $\mathbb{E}[\theta_{t+1}]=(I-\lambda\mathbb{E}[\hat{H}])\mathbb{E}[\theta_{t}]= (I-\lambda\mathbb{E}[\hat{H}])^{t}\theta_{0},$ (40) $$

which diverges if $\theta_{0}$ is in one of the escape directions of $\mathbb{E}[\hat{H}]$, which exist by the definition of strict saddle points.

Taking the $f-$distance of both sides and taking expectation, we obtain

$$ $\displaystyle\mathbb{E}[f(\theta_{t},\theta^{*})]$ $\displaystyle\geq f(\mathbb{E}[\theta_{t}],\theta^{*})$ (41) $\displaystyle=f\left((I-\lambda\mathbb{E}[\hat{H}])^{t}\theta_{0},\theta^{*} \right)\not\to 0.$ (42) $$

The first line follows from the fact that the distance function is convex by definition, and so one can apply Jensen’s inequality.

Therefore, as long as $\theta_{0}$ overlaps with the concave directions of $\mathbb{E}[\hat{H}]$, the argument of $f$ diverges, which implies that the distance function converges to a nonzero value. The expected value of $\theta_{t}$ is just the gradient descent trajectory, which diverges for any strict saddle point.

By definition, $\mathbb{E}[\hat{H}]$ contains at least one negative eigenvalue, and so the directions that do not overlap with this direction are strict linear subspaces with dimensions lower than the the total available dimensions. This is a space with Lesbegue measure zero. The proof is complete.
∎

###### Proof.

###### Proposition 2 .

###### Proof.

### B.3 Proof of Theorem 3

Let us first state the Furstenberg-Kesten theorem.

###### Theorem 5 .

(Furstenberg-Kesten theorem)
Let $X_{1},\ X_{2},\ X_{3},...$ be independent random square matrices drawn from a metrically transitive time-independent stochastic process and $\mathbb{E}[\log_{+}\|X^{1}\|<\infty]$, then(^2^22$\log_{+}(x)=\max(\log x,0)$.)

$$ $\lim_{n\to\infty}\frac{1}{n}\log\|X_{1}X_{2}...X_{n}\|=\lim_{n\to\infty} \mathbb{E}\left[\frac{1}{n}\log\|X_{1}X_{2}...X_{n}\|\right]$ (43) $$

with probability $1$, where $\|\cdot\|$ denotes any matrix norm.

Namely, the Lyapunov exponent of every trajectory converges to the expected value almost surely. Essentially, this is a law of large numbers for the Lyapunov exponent.

Now, we present the proof of Theorem 3.

###### Proof.

First of all, we define $m_{t}=\log\|\theta_{t}-\theta^{*}\|$ and $z_{t}=m_{t}-\mathbb{E}[m_{t}]$. By definition, we have

$$ $\displaystyle\mathbb{P}(g_{t}<\epsilon)$ $\displaystyle=\mathbb{P}(e^{z_{t}+m_{t}}<\epsilon)$ (44) $\displaystyle=\mathbb{P}\left(\frac{1}{{t}}(z_{t}+\mathbb{E}[m_{t}])<\frac{1}{ {t}}\log\epsilon\right)$ (45) $\displaystyle=\mathbb{P}\left(\frac{1}{{t}}(z_{t}+\mathbb{E}[m_{t}])<o(1) \right).$ (46) $$

We can lower bound this probability by

$$ $\mathbb{P}\left(\frac{1}{{t}}(z_{t}+\mathbb{E}[m_{t}])<o(1)\right)\geq\mathbb{ P}\left(\frac{1}{{t}}\max_{\theta_{0}}(z_{t}+\mathbb{E}[m_{t}])<o(1)\right).$ (47) $$

By the definition of SGD, we have

$$ $\displaystyle\frac{1}{t}\max_{\theta_{0}}\left(z_{t}(\theta_{0})+\mathbb{E}[m_ {t}]\right)$ $\displaystyle=\frac{1}{t}\max_{\theta_{0}}\log\left\|\prod_{i}^{t}(I-\lambda \hat{H}_{i})(\theta_{t}-\theta_{0})\right\|.$ (48) $$

By the Furstenberg-Kesten theorem [], this quantity converges to the constant $\Lambda=\lim_{t\to\infty}\mathbb{E}[m_{t}]/t\in\mathbb{R}$ almost surely. Namely, $z_{t}/t$ converges to 0 0 for almost every SGD trajectory.

Thus, for every $\epsilon$, if $\Lambda<0$, Eq. (46) can be bounded as

$$ $\lim_{t\to\infty}\mathbb{P}(g_{t}<\epsilon)=\mathbb{P}(\Lambda<0)=1.$ (49) $$

Because $\Lambda$ is a constant, we have that if $\Lambda<0$, all trajectories from all initialization converge to 0 0. This finishes the first part of the proof. For the second part, simply let $z_{t}$ be the trajectory starting from the trajectory that achieves the maximum Lyapunov exponent. Again, this dynamics escapes with probability $1$ by the Furstenberg-Kesten theorem. The proof is complete.
∎

###### Theorem 5 .

###### Proof.

### B.4 Proof of Theorem 4

###### Proof.

With the linear approximation of activation function $\sigma(x)\approx c_{0}x$, the neural network takes the forms of $f(x)=c_{0}^{D}\prod_{i=1}^{D}W^{(i)}x+\sum_{i=1}^{D-1}c_{0}^{(D-i)}\prod_{j=i+
1}^{D}W^{(j)}b^{(i)}$. The gradient of the loss $\ell(f(x),y(x))$ is thus

$$ $\displaystyle\nabla_{\theta}\ell(f(x),y(x))=(\nabla_{f(x)}\ell)^{T}\nabla_{ \theta}f(x).$ (50) $$

For layer $i$, the gradient

$$ $\displaystyle\nabla_{W^{(i)}}f(x)$ $\displaystyle\propto c_{0}^{D}\left(\prod_{j=1}^{i-1}W^{(j)}x\right)\left( \prod_{j=i+1}^{D}W^{(j)}\right)+\sum_{k=1}^{i-1}c_{0}^{(D-k)}\left(\prod_{j=k+ 1}^{i-1}W^{(j)}b^{(k)}\right)\left(\prod_{j=i+1}^{D}W^{(j)}\right);$ (51) $\displaystyle\nabla_{b^{(i)}}f(x)$ $\displaystyle\propto c_{0}^{(D-i)}\prod_{j=i+1}^{D}W^{(j)}.$ (52) $$

As there is no constant term in the gradient, the gradient vanishes at $\theta=0$ for all $x$, and so is the projection of the gradient. Thus, $\theta=0$ is a type-II saddle, under the assumption that $\theta=0$ is not a local minimum.
∎

###### Proof.

### B.5 Proof of Proposition 1

###### Proof.

We consider the dynamics of SGD around a saddle:

$$ $\ell=-\chi\sum_{i}u_{i}w_{i},$ (53) $$

where we have combined ${\frac{1}{S}\sum_{(x,y)\in B}}xy$ into a single variable $\chi$. The dynamics of SGD
is

$$ $\begin{cases}w_{i,t+1}=w_{i,t}+\lambda\chi u_{i,t};\\ u_{i,t+1}=u_{i,t}+\lambda\chi w_{i,t}.\end{cases}$ (54) $$

Namely, we obtain a set of coupled stochastic difference equations. Since the dynamics is the same for all values of the index $i$, we omit $i$ from now on. This dynamics can be decoupled if we consider two transformed parameters: $h_{t}=w_{t}+u_{t}$ and $m_{t}=w_{t}-u_{t}$. The dynamics for these two variables is given by

$$ $\begin{cases}h_{t+1}=h_{t}+\lambda\chi h_{t};\\ m_{t+1}=m_{t}-\lambda\chi m_{t}.\end{cases}$ (55) $$

We have thus obtained two decoupled linear dynamics that take the same form as that in Theorem 1. Therefore, as immediate corollaries, we know that $h$ converges to 0 0 if and only if $\mathbb{E}_{B}[\log|1+\lambda\chi|]<0$, and $m$ converges to 0 0 if and only if $\mathbb{E}_{B}[\log|1-\lambda\chi|]<0$.

When both $h$ and $m$ converge to zero in probability, we have that both $w$ and $u$ converge to zero in probability. For the data distribution under consideration in section 5.3 and for batch size one, we have

$$ $\mathbb{E}[\log|1+\lambda\chi|]=\frac{1}{2}\log\left|(1+\lambda)(1+\lambda a)\right|$ (56) $$

and

$$ $\mathbb{E}[\log|1-\lambda\chi|]=\frac{1}{2}\log\left|({1-\lambda})({1-\lambda a })\right|.$ (57) $$

There are four cases: (1) both conditions are satisfied; (2) one of the two is satisfied; (3) neither is satisfied. These correspond to four different phases of SGD around this saddle.
∎

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: See above.
soundness: 2
strengths: The paper is focused on the interesting question of studying the attractivity of symmetry-induced saddles. Even though such saddles might have a negative eigenvalue in the Hessian, they suffer from degeneracy due to the symmetries. 

The one-dimensional model is illustrative and explains the escape mechanism from the saddle. 

Figures are done with care.
summary: The paper is focused on the important question of the role of saddles in learning dynamics. In particular, they classify saddles into two types: Type 1 and Type 2. Type-II saddles are attractive under SGD dynamics due to the so-called vanishing gradient noise. The theory and numerics are provided to support the claims.
weaknesses: I think the paper needs a significant revision. I cannot read the mathematical notation as it is so I cannot evaluate the validity of the results. For example, three places where the classification of Type I vs Type II saddles are done -- all three unclear to me. 

1. Line 36: what does it mean "GD noise persists at the saddle"
2. Line 43: Type-I where the **noises of the gradient vanishes in the escape directions**, and Type-II saddles where the **gradient noise vanish in the directions of escape**. 

I read this sentence maybe 5 times, but the two definitions are identical up to a permutation of words.

* Can you provide a precise mathematical definition of "GD noise persists at the saddle" ?
* Can you clarify the distinction between "escape directions" and "directions of escape", or use consistent terminology if they mean the same thing ?
* Can you define these concepts explicitly when they are first introduced?


3. Line 102-103: What is $\hat{H}(x_t)$ here?

Figure 3 seems to capture the core result of the paper and hence could be put on the second page.  Can you provide a clear definition of this concept of $L_2$ stability when it is first introduced, as this seems to be an important concept in the paper ?

Why does SGD escape the Type-II saddle in Figure 1 (even though in a longer time scale). Aren't Type-II saddles are attractive under SGD by definition? Can you explicitly address this apparent contradiction and clarify the conditions under which Type-II saddles are escaped versus when they are attractive?

Finally, the attractivity of saddles and even local max is studied in Chen et al 2023 which is a very important but missing citation in this submission. They also study the symmetry-induced saddles and low-rank matrices for linear networks. 

Chen, Feng, et al. "Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks." Advances in Neural Information Processing Systems 36 (2024).

How does this paper compare to Chen et al 2023?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Overall, while I think that the authors did a good job explaining the intutions between the observed behavior, I find some of the definitions a bit hard to parse, and some parts unclear:
- Is the type I / II classification a strict classification in the sense that every saddle is either type I or type II?
- Why do you define Lp-norm stability to then discard it saying that it is not useful for type II saddles two sections later? I do find Theorem 2 interesting, but maybe it would make sense to define Lp-norm stability in section 4.3, rather than introducing it earlier which suggests that it is a tool that you are going to use throughout the paper. And maybe explain that Lp-norm stability is a common and useful tool to do this in general but you show that it fails for type II saddles, or something like that.
- I find the definition of Lyapunov exponent a bit unclear, when you take the max over the initialization $\theta_0$ can you take any possible initial $\theta_0$, and when you later say that the Lyapunov exponent is independent of initialization, do you mean that it is independent of $\theta_0$?

There are two papers that are quite relevant to the symmetry induced saddles and SGD and how it can lead to a low-rank/sparsity bias (https://arxiv.org/pdf/2306.04251v3 and https://arxiv.org/abs/2305.16038). Both rely on the fact that the noise of the gradient vanish along directions that would allow the network to increase the rank, thus making these symmetry induced region attractive, which is similar to the phenomenology around type II saddles. In the second paper this effect is further amplified by the presence of weight decay (as you briefly note in the solvable 1D dynamics).
soundness: 3
strengths: I really appreciate the choice to focus on type II saddles, and agree that it is common for such saddles to play an important role. I believe that many previous works on SGD have tried to avoid such saddles (for example by approximating SGD with noisy/Langevin GD or some variant thereof, thus guaranteeing that all saddles are type I and therefore easily avoidable), while type II saddles appear harder to study,  their peculiarity should not be avoided, so I appreciate this paper for going this less common direction.

The experiments are on simple models but the analysis is quite thorough with phase diagrams and comparison with the computed Lyapunov exponents.
summary: The paper proposes a classification of saddles into two types and then focuses on the second type of saddles, which are characterized by the fact that the variance of the gradient goes to zero as one approaches the saddle. The authors study the stability properties of such saddles and show that they can actually be attractive for large learning rates, as determined by a Lyapunov exponent that can be computed empirically.

The paper then studies a few simple neural network models and observe that SGD can in certain indeed converge to such type II saddle. Interestingly, these saddles tend to have low-rank weights (because they arise as symmetry induced saddles, right?).
weaknesses: The theoretical results describing conditions for stability of the saddles are in terms of the Lyapunov exponents, which I am not particularly familiar with and seem hard to work with theoretically. To be honest going into the paper I was expecting that it would be possible to characterize stability/attractivity in terms of statistics of the random gradient/Hessian at the saddle (expectations or variances) like in the solvable 1D dynamics that they present. Do the authors believe that no such general condition could be obtained? And is the section on insufficiency of norm stability trying to explain why such a condition is not possible in general?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: -	For eq (1), I don’t understand how to get it. Also, the definition of $\|\|\Delta\theta\|\|$ is not very clear. It is the distance from saddle point to which point?
-	In eq (3) and (4), the Hassian matrix $\hat{H}(x_t)$ is at which point? I’m assuming it is at saddle point $\theta^*$, but it does not seem to mention this in the paper.
-	For eq (4), I believe it’s worth mentioning that this is an approximate or linearized dynamic instead of the actual dynamic. The current way of writing is not clear about this.
-	In eq (5), are $\theta_t$ and $\theta_*$ scalars or vectors? The description of 1d dynamics makes me think they are scalars, but in Theorem 1 it seems they are vectors.
-	In eq (8), what are $i$ and $d$ here?
-	In Figure 2, which point is the $\Lambda$ computed at?
-	I don’t see how to get the equation in line 288. Is it assuming $\theta_*=0$? Where does the $\|\|\theta_0\|\|^2$ in the denominator come from?
-	For the experiments in section 5.1, I wonder in general how to make the predictions using $L_2$ or probability stability. This seems not discussed in the paper. As a related question, if there are multiple attractive/stable points, how should one predict which point the algorithm should converge to? For example in the description in line 317, I believe both B and C are probability stable, so why SGD converge to C instead of B?
-	I would not agree with the interpretation before Theorem 4 that all neural networks are initialized close to a Type-II saddle. The initialization is not close to the 0 point for example in the sense of Frobenius norm.
-	In Figure 6 and 7, I wonder why use rank and sparsity as the metric to measure whether SGD escapes a saddle point. It is not clear to me why these results show that ‘’Type-II saddle are indeed a major obstacle in the initial phase of training’’ (line 472).

Typo:

-	In line 43, the description of Type-I and Type-II saddle points are the same.
soundness: 2
strengths: -	Understanding the dynamics of SGD near saddle points is an interesting and important research problem
-	The classification of saddle points into Type-I and Type-II seems to be new and interesting.
-	Some of the theoretical insights from simple linearized dynamics can transfer to more complex settings also seem to be interesting.
summary: In this paper, the authors focus on the dynamics of stochastic gradient descent (SGD) around the saddle points. They propose a fine-grained characterization of saddle points: Type-I and Type-II, where gradient noise vanishes in Type-II saddles. The paper focuses on Type-II saddles which are intuitively harder than Type-I. By analyzing the linearized dynamics in simple settings, the authors show that there are different dynamics depending on stepsize and data. Numerical experiments are provided to support the claims.
weaknesses: Many places are not very clear to me as a reader. See questions section below.

---

confidence: 2
contribution: 3
initial_rating: 3
presentation: 1
questions: **Q1:** Sec 3. line 96:  loss $\mathcal l$ not introduced. I have a hard time to understand anything if I do not know the underlying problem and the probability space and SGD scheme is also not defined.
Could you provide: 
- A formal definition of the loss function $\mathcal{l}$
- A clear description of the probability space and problem setting
- A precise definition of the SGD update rule used

**Q2:** Eq. (1): Could you clarify the dimension of $\theta $? From the context, it seems to be a vector in $ \mathbb{R}^d $, but it’s unclear how a $ d $-dimensional vector could be $ \mathcal{O} $ of a real value, and whether this holds for any $ t $. I’m not familiar with this notation and am uncertain about its meaning here. Does this equation apply to any saddle point, and how is it linked to an expected gradient norm of zero? Additionally, when you refer to “$\lVert \Delta \theta \rVert$ as the Euclidean distance from the saddle,” do you mean $\lVert \theta_t - \theta \rVert $, or is it some quantity independent of $t$?


**Q3:** Following Q1 and Q2, I’m having some difficulty understanding the derivation of the different saddle types. Could you provide a more rigorous derivation of the saddle types. Are there other types not covered by this classification? Could you explain why only these two types are considered, if that is indeed the case.

**Q4:** Could you clarify what $P$ denotes the projection of— is it $\theta$ or $H$? Based on Definition 1, it seems likely that $\theta$ is intended, but this is just my assumption. 

**Q5:** How does Def. 1 imply eq. (3)? And does property (3) ($\theta$ is a saddle), not already imply property (2) (that $\theta$ is not a local minimum)? Compare to your definition of a saddle in line 96.

**Q6:** In line 113, you refer to $ H(\theta)$, but in Eq. (4) and throughout the rest of the paper, you use $ H(x) $ instead. Could you clarify the difference between these notations and what has changed?


**Q7:** Can you specify which result in Ziyin (2023) leads to eq. (4)? 

**Q8:** Proof of Thm. 1: you start by assuming eq (4), but why do you not need the projection $P$? And I thought you needed symmetry-induced saddles for this equation to hold.. are all type two saddle symmetry induced?



**Remark**

I made a genuine effort to understand the content of this paper, as I find the topic highly interesting overall. Assuming that Eq. (4) holds (without projection) around the so-called Type 2 saddles, the proofs appear solid, and the experimental results are indeed promising. I would be glad to raise my score if the questions above are addressed satisfactorily.
soundness: 2
strengths: 1. The distinction in type 1 and type 2 saddles is new.

2. The authors not only provide theoretical evidence that type two saddles cause problems for SGD but also provide numerical examples to underline the results.
summary: The authors introduce a new classification of saddle points, Type 1 and Type 2, based on the variance exhibited by the SGD process near each type of saddle point. Using this distinction, they demonstrate that Type 2 saddle points can slow down or even prevent the convergence of SGD. This leads to a better understanding of SGD’s convergence behavior in neural network optimization.
weaknesses: While the general idea behind the classification in these two types of saddle points is intuitive and interesting the following points weaken the contribution of the paper:

1. The notation used is often not formally introduced, details are missing. Even the under lying problem and SGD scheme is not introduced.

2. Given 1, it is hard to verify the correctness of the derived equations and theorems. See questions below.

Meta-review: The paper investigates the dynamics of SGD around saddle points, distinguishing two types of saddles and introducing a novel framework using random matrix theory and Lyapunov exponents. While the theoretical insights into the behavior of SGD near saddle points are interesting, reviewers raised significant concerns about the clarity of some theoretical arguments. Despite the authors' rebuttal, these concerns remained unresolved. As a result, I recommend rejection.

Decision: Reject