Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network Architectures

TL;DR

The paper addresses why gradient descent training often yields progressively more complex solutions across architectures, proposing a universal saddle-to-saddle learning framework. It centers on embedded fixed points and invariant manifolds, predicting incremental recruitment of effective units (hidden neurons, kernels, or heads) via plateaus and rapid transitions, with complexity measured by the minimal effective width. Two concrete mechanisms drive the dynamics: (i) a linear-case timescale separation determined by the input-output covariance spectrum, and (ii) a quadratic-case timescale separation driven by initialization, plus a generalization to nonlinear activations. The framework yields testable predictions about how width, data structure, and initialization shape learning dynamics and provides a principled, architecture-aware notion of simplicity with broad implications for understanding and pruning neural networks.

Abstract

Neural networks trained with gradient descent often learn solutions of increasing complexity over time, a phenomenon known as simplicity bias. Despite being widely observed across architectures, existing theoretical treatments lack a unifying framework. We present a theoretical framework that explains a simplicity bias arising from saddle-to-saddle learning dynamics for a general class of neural networks, incorporating fully-connected, convolutional, and attention-based architectures. Here, simple means expressible with few hidden units, i.e., hidden neurons, convolutional kernels, or attention heads. Specifically, we show that linear networks learn solutions of increasing rank, ReLU networks learn solutions with an increasing number of kinks, convolutional networks learn solutions with an increasing number of convolutional kernels, and self-attention models learn solutions with an increasing number of attention heads. By analyzing fixed points, invariant manifolds, and dynamics of gradient descent learning, we show that saddle-to-saddle dynamics operates by iteratively evolving near an invariant manifold, approaching a saddle, and switching to another invariant manifold. Our analysis also illuminates the effects of data distribution and weight initialization on the duration and number of plateaus in learning, dissociating previously confounding factors. Overall, our theory offers a framework for understanding when and why gradient descent progressively learns increasingly complex solutions.

Figures (6)

• Figure 1: Saddle-to-saddle dynamics occurs in the gradient descent training of a wide range of architectures and leads to a dynamical simplicity bias. (A) Saddle-to-saddle dynamics on a cartoon loss landscape. The cyan and yellow curves represent invariant manifolds, on which the network implements input-output maps expressible by the architecture with one and two units, respectively. In general, saddle-to-saddle dynamics operates by repeating: i) during the plateau, escaping from a saddle associated with a width-$h$ network onto an invariant manifold with effective width $(h+1)$; ii) during the rapid transition phase, approaching a fixed point on that manifold, which is a saddle associated with a width-$(h+1)$ network. This figure shows two repeats of this process. (B-G) Loss and weight dynamics for various architectures. Each panel shows the loss during training (top), and the first-layer weights during the intermediate plateau (bottom left, phase 3 in panel A) and at the end of learning (bottom right). The first-layer weights to each hidden unit are two-dimensional and plotted as black dots. During the intermediate plateau, all networks visit a saddle, at which the input-output map of the network can be expressed by the architecture with only one unit. The network then converges to a stable fixed point, at which the input-output map is expressible with two units. The weight structures in BC, DE, and FG correspond to three categories of weight configurations of fixed points in \ref{['thm:fixed-point']}; see \ref{['sec:fixed-point']} for details. A video version of this figure is provided at https://yedizhang.github.io/simplicity.html. Dynamics with other two-layer architectures and deep networks are provided in \ref{['fig:mnist', 'fig:s2s-more', 'fig:deep']}. Experimental details are provided in \ref{['supp:implementation']}.
• Figure 2: The effect of network width, data distribution, and initialization on learning dynamics. Singular values of ${\bm{\Sigma}}_{yz}$ (linear network) or positive singular values of ${\bm{\Sigma}}_{yZ}$ (linear self-attention) follow a power law, $s_n = n^{-\kappa}, n=1,2,3$, and are normalized such that $\sum_{n=1}^3 s_n=1$. (A) Increasing the number of units $H$ has little effect on the loss curves of linear networks, but shortens the plateaus in linear self-attention. $\kappa=1$ for both models. (B) Decreasing the power law exponent $\kappa$ shortens the plateaus in both linear networks and linear self-attention. Setting $\kappa=0$ eliminates plateaus in linear networks but does not eliminate plateaus in linear self-attention. $H$ is 100 for linear networks and 25 for linear self-attention. (C) Linear networks with small isotropic initialization or large low-rank initialization exhibit saddle-to-saddle dynamics. The loss landscape cartoon illustrates large rank-$r$ weights places a linear network near an invariant manifold with $r$ effective units and thus approaches saddles during learning. (D) Increasing the scale of isotropic random initialization shortens the plateaus. $\kappa=1$ for panels C,D.
• Figure 3: Saddle-to-saddle dynamics in two-layer fully-connected linear and ReLU networks trained for binary classification of MNIST digits. The input dimension is $28\times28=784$, the hidden layer width is 1000, and the target outputs are two-dimensional one-hot vectors. The intermediate plateau is longer when the two digits are harder to distinguish. For example, digits 3/5 are harder to distinguish than digits 0/1. The colored curves represent the top three singular values of the first-layer weight matrix, ${\bm{U}} \in \mathbb{R}^{1000\times784}$. Consistent with our theory, the growth of the first and second singular values coincides with the first and second abrupt drops in the training loss, respectively, corresponding to the increase in the effective width. The third largest singular value is close to zero, meaning the rank of the first-layer weight matrix is at most two, approximately. Details: The batch size is 64. The learning rate is 0.01. The initial weights are sampled independently from $\mathcal{N}(0,10^{-12})$.
• Figure 4: Learning dynamics in two-layer networks with other activation functions. Each panel shows the loss during training (top), and the first-layer weights right after the first abrupt loss drop (bottom left) and at the end of learning (bottom right). The first-layer weights to each hidden unit are two-dimensional and plotted as black dots. (A) The softmax self-attention model is the same as linear self-attention in \ref{['fig:cover']}F, except for adding the softmax activation function. The training data set is the same as that of \ref{['fig:cover']}F. (B) Two-layer fully-connected sigmoid network, i.e., $\phi({\bm{x}};{\bm{u}}) = \mathsf{sigmoid}\left({\bm{u}}^\top {\bm{x}}\right)$. (C) Two-layer fully-connected sinusoid network, i.e., $\phi({\bm{x}};{\bm{u}}) = \mathsf{sin}\left({\bm{u}}^\top {\bm{x}}\right)$. (D) Two-layer fully-connected tanh network, i.e., $\phi({\bm{x}};{\bm{u}}) = \mathsf{tanh}\left({\bm{u}}^\top {\bm{x}}\right)$. (E,F,H) Two-layer fully-connected networks with the given activation functions. (G) Two-layer fully-connected cubic network, i.e., $\phi({\bm{x}};{\bm{u}}) = \left({\bm{u}}^\top {\bm{x}}\right)^3$. Details: Except for panel A, the training set is generated by a width-2 teacher network with the same activation function, $y=\phi({\bm{x}};{\bm{u}}_1^*)+\phi({\bm{x}};{\bm{u}}_2^*), {\bm{x}}\in\mathbb{R}^2, y\in\mathbb{R}$. The input is sampled from $\mathcal{N}({\bm{0}},{\bm{I}})$. For panels B-F, the teacher network has ${\bm{u}}_1^*=[1,0]^\top, {\bm{u}}_2^*=[0,2]^\top$. For panels G,H, the teacher network has ${\bm{u}}_1^*=[1,0]^\top, {\bm{u}}_2^*=[0,1]^\top$. The number of training samples is 8192. The learning rate is 0.02. The initial weights are sampled independently from $\mathcal{N}(0,10^{-12})$ for panels A-D, and from $\mathcal{N}(0,0.005^2)$ for panels E-H. The width is $H=50$ for panels A-D, and $H=10$ for panels E-H.
• Figure 5: Learning dynamics in deep networks. Each panel shows the loss over training time (top), and the first-layer weights right after the first abrupt loss drop (bottom left) and at the end of learning (bottom right). The first-layer weights to each hidden unit are two-dimensional and plotted as black dots. The training sets in panels A,C,E are the same as those in \ref{['fig:cover']}(B,D,F). The training sets in panels B,D,F split the scalar output in \ref{['fig:cover']}(C,E,G) into a two-dimensional vector output. (A) Three-layer linear fully-connected network. (B) The network has a convolutional linear layer as the first hidden layer and a fully-connected linear layer as the second hidden layer. (C) Three-layer ReLU fully-connected network. (D) The network has a convolutional ReLU layer as the first hidden layer and a fully-connected ReLU layer as the second hidden layer. (E) One-layer linear transformer, consisting of one linear self-attention layer and two fully-connected linear layers. (F) The network has a fully-connected layer with quadratic activation as the first hidden layer and a fully-connected linear layer as the second hidden layer. Details: The number of training samples is 8192. The learning rate is 0.02. The initial weights are sampled independently from $\mathcal{N}(0,0.005^2)$ for panels A-E, and from $\mathcal{N}(0,0.05^2)$ for panel F. The width is $H=50$ for panels A-D, and $H=10$ for panels E,F.

Theorems & Definitions (14)

• Definition 1
• Theorem 1: Embedded fixed points
• Remark 1
• Corollary 2
• Theorem 3: Invariant manifolds
• Theorem 4: Timescale separation between directions
• Proposition 5: Timescale separation between units
• proof
• Definition 2: Invariant manifold
• proof