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

AI summary

Overview

Research area: Deep learning theory — loss landscape geometry, optimization dynamics, and inductive bias in neural networks. The work sits at the intersection of optimization theory and architecture-specific analysis (fully-connected, convolutional, and attention-based models), from the Gatsby Computational Neuroscience Unit and Sainsbury Wellcome Centre at University College London.

Technical level: Advanced. The paper is written for readers comfortable with gradient flow dynamics, fixed points and saddles, singular value decompositions, invariant manifolds, and timescale separation arguments. The intuition sections are accessible, but the core results are stated as theorems and propositions with proofs in appendices.

Scope: A unified theoretical account of why gradient descent learns increasingly simple-to-complex solutions — the "simplicity bias" — arguing that saddle-to-saddle dynamics is a common mechanism across linear, ReLU, convolutional, quadratic, and linear self-attention networks, with the notion of simplicity defined as the number of effective units (hidden neurons, convolutional kernels, or attention heads).

What This Paper Is About

Neural networks trained with gradient descent often learn simple solutions first and more complex ones later, frequently in stage-like fashion with long plateaus in the loss interrupted by bursts of rapid improvement. This "dynamical simplicity bias" has been observed across many architectures and tasks, but the paper argues that existing theoretical treatments are architecture-specific and lack a unifying framework.

The paper's goal is to supply that framework: to show that a single mechanism — trajectories passing near a sequence of saddle points — explains the simplicity bias across a broad class of architectures, and to identify what "simple" means in each case (few hidden neurons, kernels, or attention heads).

Key Contributions

1. A general theory of embedded fixed points. Theorem 1 shows that a fixed point of a network with (H−1) units can be embedded into a network with H units, yielding a fixed point there. The paper gives four constructions (Equations 4–7). Equations 4 and 5 were first discovered by Fukumizu and Amari (2000); Equations 6 and 7 are new extensions. Corollary 2 iterates the result across multiple units and multiple layers, so a narrow deep network's fixed points remain fixed points of any wider network of the same depth.

2. A general theory of invariant manifolds. Theorem 3 identifies four conditions under which a relationship between weights is preserved forever under gradient flow: equal weights between two units (any activation function), zero weights in a unit, proportional weights between two units under homogeneity, and linear dependence among any number of units under linearity. On these manifolds the network behaves as if it had fewer units than its actual width, so simplicity is preserved along the connecting trajectories between embedded fixed points.

3. Mechanistic analyses of saddle-to-saddle dynamics in two concrete families. For two-layer networks where the layer function is a homogeneous polynomial in the weights, the paper analyzes the linear case (Theorem 4) and the quadratic case (Proposition 5), showing that both produce saddle-to-saddle dynamics — but for different reasons.

4. A separation of data-induced and initialization-induced saddle-to-saddle dynamics, with testable predictions. The paper shows that timescale separation from the data distribution leads to low-rank weights, while timescale separation from the initialization leads to sparse weights, and uses this to predict how network width, data distribution, and initialization structure change the number and duration of plateaus.

Main Findings

• Simplicity means effective unit count. Across the architectures studied, the operative notion of simplicity is the number of units that suffice to express the network's input-output map. 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.

• Fixed points of narrow networks become saddles in wider networks. Because the global minimum of a narrow network (which may incur nonzero loss) is embedded as a fixed point of a wider network, the wider network's landscape contains a nested hierarchy of saddles associated with increasing effective width. The paper notes these are guaranteed to be saddles in deep linear networks with rank-r (r ≥ 1) target maps, and are saddles under mild conditions in general architectures.

• Fixed points visited during learning come from three of the four constructions. In the simulations shown in Figure 1, the saddles visited during learning fall under Equations 5, 6, and 7 — not Equation 4. Panels (B, C) correspond to rank-one and rank-two weights (Equation 7); panels (D, E) to one and two rays of proportional weights (Equation 6); panels (F, G) to one or two units with large weights and the rest near zero (Equation 5).

• One saddle-to-saddle transition is: escape, approach, repeat. During a plateau the network escapes a saddle associated with effective width h onto an invariant manifold with effective width (h+1); during the rapid transition it approaches a fixed point on that manifold, which is itself a saddle for a width-(h+1) network. Repeating this yields the stage-like loss curve.

• Linear case: timescale separation is between directions, driven by the data. Theorem 4 analyzes the linear dynamical system that governs the early phase under small initialization. Weights grow exponentially along the singular vectors q_k and r_k at rate e^{s_k t}. Relative to the dominant rate e^{s_1 t}, components along other singular vectors decay as e^{(s_k − s_1)t}. When the projection onto the top r singular vectors reaches O(1), the projection onto the remaining subspace is O(ε^{1 − s_{r+1}/s_1}) almost surely. With a unique largest singular value (r = 1), growth is approximately rank-one, placing the network near an invariant manifold with one more effective unit.

• Quadratic case: timescale separation is between units, driven by initialization. Proposition 5 analyzes a "rich-get-richer" dynamic in which one unit grows much faster than the others. The paper illustrates the mechanism with the scalar dynamics v̇_i = v_i², whose solution is v_i(t) = (1/v_i(0) − t)^{−1}, showing that when the fastest unit reaches O(1) the remaining units are still O(ε). This places the network near an invariant manifold of effective width one.

• Width affects the two cases differently. Because the linear case separates timescales across directions, adding units has little effect on loss curves in linear networks provided there are enough units to learn all directions. Because the quadratic case separates timescales across units, adding units reduces the gaps between adjacent initial samples and shortens the plateaus. In Figure 2A, increasing the number of attention heads in linear self-attention speeds learning, while increasing the width of fully-connected linear networks does not — described in the paper as a theoretically grounded advantage of scaling linear self-attention over scaling fully-connected linear networks.

• Data distribution controls plateau structure. In Figure 2B the singular values of Σ_yz (linear network) or the positive singular values of Σ_yZ (linear self-attention) follow a power law s_n = n^{−κ} for n = 1, 2, 3, normalized so that the three singular values sum to 1. Decreasing κ narrows the gaps between singular values and shortens plateaus in both models. Setting κ = 0 makes all singular values equal and eliminates plateaus in linear networks (except the initial one corresponding to escaping the saddle at zero), because the largest singular value then has multiplicity r = D in Theorem 4, so the solution jumps directly from effective width 0 to D. In linear self-attention, κ = 0 shortens but does not eliminate plateaus. In these simulations κ = 1 for both models; H is 100 for linear networks and 25 for linear self-attention.

• Initialization structure matters in two ways. Linear networks with small isotropic initialization or with large low-rank initialization both exhibit saddle-to-saddle dynamics. Initializing near an invariant manifold but away from saddles is not a necessary condition for saddle-to-saddle behavior: in Figure 2C the large low-rank initialization is away from saddles, so there is no plateau at the start — the loss first drops exponentially and then shows plateaus followed by rapid transitions. The paper's cartoon illustrates that large rank-r weights place a linear network near an invariant manifold with r effective units, so learning approaches saddles. Increasing the scale of isotropic random initialization shortens the plateaus (κ = 1 for the panels C and D).

• Higher-order and general nonlinear activations. For a layer function that is a homogeneous polynomial of degree p > 2 in the weights, the paper conjectures that a timescale separation between units still exists, possibly stronger than in the quadratic case, by analogy with the scalar dynamics v̇_i = v_i^p; the cubic (p = 3) case is reported to be consistent with this intuition (Figure 4G). For a general nonlinear activation, Taylor expansion around zero means the early dynamics is dominated by the lowest-order non-vanishing term. In a two-layer tanh network that term is linear, so the network develops rank-one weights early, but the subsequent dynamics need not be saddle-to-saddle because rank-one weights do not generally correspond to invariant manifolds for tanh networks (Figure 4D). By contrast, a network with activation φ(x; u) = u^⊤x · tanh(u^⊤x) has a quadratic lowest-order term, giving timescale separation between units and saddle-to-saddle dynamics (Figure 4F).

Methodology in Plain English

The authors study gradient flow — the continuous-time limit of gradient descent with a small learning rate — on a training loss averaged over a supervised dataset, for a general layer template that covers fully-connected layers (φ is an activation of a linear pre-activation), convolutional layers (φ is an activation of a convolution), and self-attention layers (φ produces the softmax-weighted attention output). This template lets a single set of statements cover neurons, kernels, and heads.

The argument proceeds in three layers. First, they prove algebraically that you can take a fixed point of a narrower network and construct weights for a wider network that implement exactly the same function and are themselves a fixed point — the "embedded fixed points" of Theorem 1, extended recursively to deep networks by Corollary 2. Second, they prove that certain simple relationships between weights (two units identical, one unit zeroed, two units proportional, or one unit a linear combination of others) are preserved by gradient flow, defining invariant manifolds on which the network is effectively narrower than its width. Third, they show that these manifolds connect one embedded fixed point to the next, so a trajectory that breaks exactly one constraint moves from a saddle of effective width h to a fixed point of effective width h+1.

To explain why gradient descent would actually follow such paths, the authors specialize to two-layer networks with homogeneous polynomial layer functions and analyze the dynamics near small initialization. In the linear case they linearize around small weights, relate the resulting dynamics to the singular value decomposition of the input-output correlation matrix Σ_yz, and show that growth rates differ exponentially across singular directions. In the quadratic case they analyze the approximate dynamics of individual units and show that the unit with the largest initialization runs away from the rest. Both analyses are backed by simulations of architectures including linear fully-connected and convolutional networks, quadratic networks, linear self-attention, ReLU networks, and tanh networks. Experimental details are reported in Appendix I; derivations appear in Appendices E through H.

Why This Matters

Impact on research. The paper reframes a large collection of architecture-specific observations as instances of one mechanism. Instead of asking whether linear networks, ReLU networks, convolutional networks, and transformers each have their own story, it provides a shared vocabulary — embedded fixed points, invariant manifolds, and timescale separation — and identifies the properties of a layer function (having a zero, being homogeneous, being linear) that determine which parts of the theory apply. It also offers a principled distinction between two causes of stage-like learning that are easy to conflate: the data distribution and the initialization. The theory additionally predicts when stage-like behavior will not appear (for example, when singular values are equal, or for tanh networks where rank-one weights are not on an invariant manifold), which is as useful as predicting when it will.

Real-world applications (the paper does not report applied benchmarks or deployed systems; these are areas where the theory's predictions are relevant):

• Training-dynamics diagnostics. Because low-rank weights indicate data-induced saddle-to-saddle paths and sparse weights indicate initialization-induced ones, inspecting weight structure during training could help identify which regime a model is in.
• Initialization and scaling decisions. The differing width-scaling predictions for linear networks versus linear self-attention are directly relevant to decisions about where added capacity pays off.
• Data curation and spectral structure. The result that the spectrum of the input-output correlation matrix governs plateau duration in the linear case points to dataset properties that shape training curves.
• Transformer capacity planning. The finding that adding attention heads shortens plateaus, whereas adding width to fully-connected linear layers does not, speaks to the practical scaling of attention-based models.

Industry relevance. Stage-like training curves with long plateaus are a familiar operational reality in large-model training, and questions about whether to add heads, add width, change initialization scale, or alter data spectra are routine. This paper offers a theory that makes those choices non-arbitrary, while being explicit that the sharpest quantitative results are derived for two-layer networks with squared loss and small initialization — so the work is best read as a mechanistic explanation rather than a prescription for a specific production system.

Future Directions

• Extending the dynamic analysis to deeper networks. The fixed-point and invariant-manifold results (Theorem 1, Corollary 2, Theorem 3) already apply to deep networks with the layer template of Equation 1, but the detailed dynamical analyses (Theorem 4 and Proposition 5) are carried out for two-layer networks. Whether the same timescale-separation arguments control saddle-to-saddle learning layer by layer in deep networks remains to be established.

• Higher-degree and general activation functions. The paper states the p > 2 case as a conjecture and reports the cubic case as consistent with its intuition. A general characterization of which lowest-order Taylor term produces saddle-to-saddle behavior — and when, as with tanh, it produces rank-one weights without subsequent stages — is left open.

• Characterizing when stage-like behavior fails to appear. The paper predicts non-stage-like dynamics for equal singular values in linear networks and for tanh networks, and cites prior work showing that smooth exponential dynamics can occur by changing initialization. A complete account of the boundary between stage-like and smooth regimes is an open question.

• The truncated discussion of initialization structure. The analysis of initializing near an invariant manifold but away from saddles is cut off in the available text, and the paper's broader claims about initialization effects on plateau duration and number would benefit from full comparison with simulations beyond the cases reported in Figure 2.

Target Audience

This paper is aimed at machine learning theorists and mathematically inclined researchers — particularly those working on optimization dynamics, loss landscape geometry, and implicit bias — who want a single framework spanning fully-connected, convolutional, and attention-based architectures. It is also relevant to empirically minded deep learning researchers and engineers who want a principled account of plateau-and-jump training curves and of how width, initialization, and data spectra shape them. Readers without a background in dynamical systems, fixed-point analysis, and linear algebra will find the intuition sections readable but the theorem statements demanding.

Authors’ 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 disentangles data-induced and initialization-induced saddle-to-saddle dynamics. In particular, the former leads to low-rank weights while the latter to sparse weights. Equipped with the theory, we predict the effects of data distribution and weight initialization on the duration and number of plateaus in learning. Overall, our theory offers a framework for understanding when and why gradient descent progressively learns increasingly complex solutions.