Title: Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks
Venue: NeurIPS 2025
Accepted: yes
Keywords: feature learning

saddle-to-saddle dynamics

incremental learning

vanishing initialization

emergence

modular addition

Fourier features
Abstract: What features neural networks learn, and how, remains an open question. In this paper, we introduce Alternating Gradient Flows (AGF), an algorithmic framework that describes the dynamics of feature learning in two-layer networks trained from small initialization. Prior works have shown that gradient flow in this regime exhibits a staircase-like loss curve, alternating between plateaus where neurons slowly align to useful directions and sharp drops where neurons rapidly grow in norm. AGF approximates this behavior as an alternating two-step process: maximizing a utility function over dormant neurons and minimizing a cost function over active ones. AGF begins with all neurons dormant, corresponding to an initialization at the origin. At each iteration, a dormant neuron activates, triggering the acquisition of a feature and a drop in the loss. AGF quantifies the order, timing, and magnitude of these drops, matching experiments across several commonly studied architectures. We show that AGF unifies and extends existing saddle-to-saddle analyses in fully connected linear networks and attention-only linear transformers, where the learned features are singular modes and principal components, respectively. In diagonal linear networks, we prove AGF converges to gradient flow in the limit of vanishing initialization. Applying AGF to quadratic networks trained to perform modular addition, we give the first complete characterization of the training dynamics, revealing that networks learn Fourier features in decreasing order of coefficient magnitude. Altogether, AGF offers a promising step towards understanding feature learning in neural networks.


Peer reviews:
comments: SUMMARY:
This paper introduces the concept of alternating gradient flows to explain saddle-to-saddle or stepwise dynamics during training, providing a unified framework across four distinct settings. The central insight is that, despite differing representations (e.g., neurons, SVD modes, Fourier modes), the dynamics exhibit a common mechanism: an initial utility maximization phase over dormant neurons, followed by a cost minimization phase over active neurons.

STRENGTHS AND WEAKNESSES:
**Strengths**

The proposed unification offers novel insights into the training dynamics of two-layer neural networks with small initialization. The paper is clearly written, and the key ideas are well-articulated and supported by analysis.


**Weaknesses**

1. Small initialization leads to slow dynamics in the early stages of training and may result in neuron alignment along only a few directions. In such cases, the network effectively utilizes only a small subset of neurons. This behavior may diverge from typical training regimes used in practice, raising questions about the framework's applicability to real-world scenarios and data.

2. Extending the proposed framework to deeper architectures is non-trivial, potentially limiting its generalizability and broader impact.

QUESTIONS:
There are no major technical questions. The methodology and arguments are well-presented. The main concerns lie in the assumptions and scope of the setup, as noted in the weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces an *Alternating Gradient Flow* (AGF) algorithm for approximating the gradient flow dynamics of two-layer layer neural networks. Inspired by the saddle-to-saddle dynamics observed accross different ML tasks, AGF mimicks by an iterative, two-step procedure: (1) Utility maximization (plateau): at a saddle-point, a subset of *dormant* neurons (with small weights) evolve independently to maximize their correlation with the rest, *active* neurons; (2) Cost minimization (drop): Once a dormant neuron develops meaningful correlation, it becomes active, leading to a drop of loss until they reach a new plateau.\n\nTHe paper then investigates AGF in different settings. First, they show that AGF is exact when the initialization scale $\\alpha\\to 0^{+}$ for diagonal neural networks. Second, they show that in linear two-layer networks AGF recovers the step-wise singular value recovery previously observed in the context of commuting covariate-covariate and label-covariate covariances, and conjecture that this remains approximately true in the non-commuting case. Third, they show that AGF implements stepwise principal component regression for linear transformers. Finally, they show that two-layer quadratic networks trained on AGF learn Fourier features a modular addition task.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**: The paper is well-written and the idea of unifying the dynamics of different settings studied in the literature under a single framework is appealing. \n\n**Weaknesses**: The principal weakness of the paper is that the scope and relevance of AGF beyond the simple cases analyzed remain unclear. Equivalence with GF is proved or conjectured only in the linear setting, and the single non-linear example - a quadratic network - relies on an additional assumption that, while reasonable, falls outside the stated scope of AGF.\n\nOn a bigger picture, it is not clear to what extent the \"broad picture\" offered here is useful or insightful. Because gradient flow is a strictly descent method, once the critical points are characterized, its trajectory is fully determined by the initial conditions: it proceeds from initialization toward a minimum, potentially passing through successive saddles. Hence, the substantive challenge lies in calculating these critical points (and their indices) for concrete problems of interest. In this light, the principal contribution I see in this work is the identification of a utility function that potentially simplifies part of this analysis - estimating the time scale for escaping saddles — under certain conditions whose generality is unclear (see questions below). The framework does not, however, appear to simplify other essential aspects of the dynamics, such as determining the eventual point of convergence or whether that point is a local or global minimum for a given initialization.\n\nOn a minor note, the paper also ignores a large body of literature that has been dedicated to similar questions on quadratic nets and single-/multi-index models (see comment below).\n\nQUESTIONS:\n- The notation is confusing. In lines 64-65, the 2LNN is written as $f(x;\\Theta) = \\sum_{i=1}^{H}a_{i}\\sigma(w_{i}^{\\top}g_{i}(x))$ with $a_{i}\\in\\mathbb{R}^{c}$. Do you actually mean $c=1$ or this is understood as a vector $\\times$ scalar multiplication? In L72 you write $a_{i}\\sim\\mathcal{N}(0,\\alpha/\\sqrt{2c})$ suggesting $a_{i}$ is a scalar. Or should I understand this as i.i.d. component-wise?\n\n- How do you define a *smooth dynamical system* in L77-78? Is it clear that GF is a smooth dynamical system for any activation function $\\sigma$?  \n\n- L43-45: \"*However, these analyses often rely on specific simplifying assumptions about the architecture (e.g., linear activations), data (e.g., orthogonal inputs), and optimizer (e.g., layer-wise learning)*\". This is a strange way of contrasting the paper to the cited literature, since all cases analyzed here are for simplified architectures (2 layer, linear or quadratic) and a simplified optimizer (GF, AGF), and from the results it is not really clear to what extent they generalize beyond it. Moreover, most of the cited papers only use one of these simplifying assumptions. I suggest tuning this down. \n\n- What is the role of the parametrization of the quadratic network in L337-338 in the dynamics? For instance, given the data $z=(a\\dot x,b\\dot x)\\in\\mathbb{R}^{2d}$, one could fit a standard quadratic network $f_{i}(x;\\Theta) = \\langle \\beta_{i}, z\\rangle^{2}w_{i}$ with double number of neurons $\\beta_{i}\\in\\mathbb{R}^{2d}$, how does the AGF dynamics look like? Does it spontaneously learn to split the neurons?  \n\n- The population dynamics of non-linear 2 layer neural networks has been widely studied, specially in the context of multi-index models (Saad, Solla 1995; Goldt et al., 2018; Arnaboldi et al., 2023a). A picture arising from these works is that the saddle-to-saddle dynamics in these models correspond to a *specialization transition* (Saad, Solla 1995) from a regime where the correlation with the target directions is homogeneous (all neurons attempt to fit the linear part of the target independently) to a regime where neurons specialize to directions of the target (see for instance fig. 1 of Goldt et al., 2018 or fig. 2 of Arnaboldi et al., 2023a). In particular, the exact details of the specialized regime depends on the target architecture. For instance, if the target is itself a 2LNN and the model has a number of neurons which is a multiple of the target the dynamics can prioritize groups of neurons specializing at some target weights, while if the model width is not a multiple of the target, this can induce neurons to \"shut down\" instead of grouping (Goldt et al., 2018). Moreover, (Arnaboldi et al., 2023a) has proven that in this context the model neurons are only approximately independent in the large-width (mean-field) limit, with the correlation between neurons being non-negligeable for a finite width network. Note that all these works preceed [36], which confirm and prove part of this analysis under some additional assumptions. Together, this seems to paint a quite different picture from AGF, where individual neurons specialize sequentially and independently, not collectively. How do the authors see these works in the light of AGF?\n\n- On a note similar to the above, the SGD dynamics of quadratic neural networks have been widely studied as a simple non-linear setting in the literature (Sarao Mannelli et al., 2020; Arnaboldi et al., 2023b; Martin et al., 2024). The picture that emerge from these works is that every target direction is learned simultaneously, although at different speeds (depending on the corresponding eigenvalue). This seems to lack the separation of time implied by the AGF where directions are fully learned sequentially. Am I missing something?\n\n**References**:\n\n- (Saad, Solla 1995) David Saad, Sara A. Solla. \"*Dynamics of On-Line Gradient Descent Learning for Multilayer Neural Networks*\". NeurIPS 1995\n\n- (Ben Arous et al., 2021) Gerard Ben Arous, Reza Gheissari, Aukosh Jagannath. \"*Online stochastic gradient descent on non-convex losses from high-dimensional inference*\". JMLR 2021\n\n- (Goldt et al., 2018) Sebastian Goldt, Madhu S. Advani, Andrew M. Saxe, Florent Krzakala, Lenka Zdeborová. \"*Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup*\". NeurIPS 2019. \n\n- (Arnaboldi et al., 2023a) Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro. \"*From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks*\". COLT 2023.\n\n- (Sarao Mannelli et al., 2020) Stefano Sarao Mannelli, Eric Vanden-Eijnden, Lenka Zdeborová. \"*Optimization and Generalization of Shallow Neural Networks with Quadratic Activation Functions*\". NeurIPS 2020\n\n- (Arnaboldi et al., 2023b) Luca Arnaboldi, Florent Krzakala, Bruno Loureiro, Ludovic Stephan. \"*Escaping mediocrity: how two-layer networks learn hard generalized linear models with SGD*\". arXiv:2305.18502\n\n- (Martin et al., 2024) Simon Martin, Francis Bach, Giulio Biroli. \"*On the Impact of Overparameterization on the Training of a Shallow Neural Network in High Dimensions*\". ICML 2024\n\nLIMITATIONS:\nLimitations are discussed in the last section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors study the training dynamics of shallow neural networks, with a focus on saddle-to-saddle dynamics that emerge under small initialization. They propose a framework called Alternating Gradient Flows (AGF) to describe the behavior in this regime. AGF consists of two alternating steps: maximizing a utility function over dormant (i.e., small) neurons, and minimizing a cost function over active (i.e., large) neurons. This framework unifies several existing saddle-to-saddle dynamics, including those found in diagonal linear networks, linear networks, and attention-only linear transformers. It also provides a new perspective on the dynamics involved in learning modular addition. Experimental results are included to support the proposed claims.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1.\tUnderstanding the training dynamics of neural networks, particularly in the feature learning regime and in settings like the saddle-to-saddle dynamics studied in this paper, is an important and active area of research.\n2.\tThe proposed Alternating Gradient Flows (AGF) framework unifies several prior analyses of saddle-to-saddle dynamics and offers intuitive explanations for the behavior observed in these settings.\n3.\tThe framework also sheds light on the training dynamics involved in learning modular addition, which appears to be a novel contribution.\n\nWeaknesses:\n\n1.\tWhile the AGF framework offers useful insights into saddle-to-saddle dynamics, it does not provide a full end-to-end analysis of the training process. Many of the examples rely on heuristics or assume that each step in the alternation (e.g., utility maximization and cost minimization) is perfectly executed. While these assumptions are helpful for building intuition, they limit the framework's predictive or theoretical completeness.\n\n    Additionally, the utility maximization step may not be tractable in general. In some scenarios not covered by the paper, solving this maximization could be NP-hard.\n\nQUESTIONS:\n1.\tIs it possible for an active neuron to become dormant again? In all the plots, it seems that once a neuron becomes active, it remains active. Could the authors comment on whether reversion to dormancy is theoretically possible or observed in practice?\n2.\tDoes the AGF framework accurately capture the dynamics of general (non-orthogonal) two-layer neural networks, such as the one shown on the left side of Figure 1? If not, are there specific conditions under which the agreement breaks down?\n3.\tDo the authors have an estimate for how small the initialization scale $\\alpha$ needs to be in order for the actual dynamics to closely follow the AGF dynamics? Specifically, is there any known rate of convergence with respect to $\\alpha$?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Alternating Gradient Flows (AGF), an algorithm designed to approximate the saddle-to-saddle dynamics observed when training two-layer neural networks with vanishing initialization. AGF alternates between a utility maximization phase, during which dormant neurons evolve independently (corresponding to plateaus in the loss), and a cost minimization phase, where active neurons collectively reduce the loss. The authors prove that AGF converges to the saddle-to-saddle dynamics of gradient flow in the limit of vanishing initialization for diagonal linear networks, and they conjecture that it captures the dynamics more broadly across other architectures.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\nThe paper addresses an important problem: understanding how features emerge during training. The proposed framework, Alternating Gradient Flows (AGF), offers an original perspective on this question. While the derivation involves many heuristic arguments, the approach is conceptually clear and quite compelling. A notable strength is that the framework applies broadly to two-layer architectures, rather than being tailored to a specific model. The introduction of new quantities—such as the utility function guiding dormant neuron dynamics—is (to the best of my knowledge) novel and insightful. Finally, the paper is supported by clear and informative illustrations (e.g., Figures 1 and 2), which help convey the underlying intuition.\n\n**Weaknesses**\nThe paper relies on heuristics and informal arguments, and some of the mathematics lacks rigour and clarity. For example:\n- the key equation in (2) is introduced rather abruptly. Since they underpin the AGF framework, it would be important to explain more carefully in what sense they validly approximate the gradient flow dynamics. \n- The use of Euler’s homogeneous function theorem (e.g., line 643 in the appendix) appears questionable. While the utility function  $U_i$ may be locally $\\kappa$-homogeneous near the origin, the theorem is applied in regions where $\\Vert \\theta_i \\Vert$ is not necessarily small, which undermines its validity. Several steps in the derivations in the appendix are skipped; more detailed computations would be helpful.\n- Another issue lies in the definition of AGF: the stopping conditions for the two while loops are ill-defined. Neither $\\nabla L(\\Theta_D)$ nor $\\nabla L(\\Theta_A)$ will reach exactly zero in finite time, so the algorithm as written is not well-posed. It is unclear how this should be interpreted. Moreover, the second loop (cost minimization) is said to run until convergence, or let's say \"nearly convergence\", which in practice would take some nonzero time. Yet, in the illustrations (e.g., Figure 2, right), the step is depicted as instantaneous. How come?\n- In the definition of AGF, if a neuron collapses during the cost minimization phase and becomes dormant again, it is unclear how its associated accumulated utility $S_i$ should be initialized in the next utility maximization phase. This detail is not specified in the current version of the algorithm, yet it is essential for the consistency of the method.\n- The current definition of AGF does not appear to align with the algorithm in Pesme and Flammarion. Consider their example in Figure 2 of their paper: initially, the active set is empty. At time $t_1$, coordinate 1 is added to the active set, and phase 2 proceeds until the gradient vanishes over this active set. We then return to phase 1, and at time $t_2$, the second coordinate is added to the active set. We then enter phase 2 again, and according to AGF,  this phase is meant to continue until the gradient is zero over both active coordinates. However, this is not what occurs in the dynamics: the iterates pass through an intermediate saddle point $\\beta_{c,2}$ due to the collapse of the first coordinate. Therefore, in its current form, the definition of AGF does not capture this behaviour, and as such, Theorem 3.1—which claims equivalence in the $\\alpha \\to 0$ limit—cannot be correct without modification of the AGF algorithm.\n\n\n**Minor comments**\n-In the definition of AGF, the quantity $\\eta_\\alpha$ is used without prior definition.\n- In the def of AGF: the expression $\\frac{d \\theta_j}{d t} = -\\nabla L(\\Theta_A)$ is problematic: the left-hand side is a vector corresponding to neuron $j$, while the right-hand side is the full gradient over all active neurons.\n- In my opinion, the use of random initializations $a_i \\sim \\mathcal{N}(0, \\frac{\\alpha}{2c})$, $w_i \\sim \\mathcal{N}(0, \\frac{\\alpha}{\\sqrt{2d}})$ makes things unnecessarily complicated. Why not simply consider deterministic initializations of the form $a_i^{\\alpha} = \\alpha a_i$ for fixed vectors $a_i$? This setup strictly generalizes the random case.\n\nQUESTIONS:\n- In the definition of AGF, it is implicitly assumed that a single neuron $i_*$ activates at each iteration. However, in practice—particularly in fully connected networks or ReLU networks—multiple neurons often activate simultaneously in clusters. This phenomenon is not captured by the current formulation.\n- Line 148 : \"Through this process, we have generated a precise sequence of saddle points and jump times to describe how gradient flow leaves the origin through a saddle-to-saddle process.\"Can you clearly justify why the intermediate points in this sequence correspond to actual saddle points of the loss function. \n- Intuitively, during the first phase of AGF (utility maximization by dormant neurons), one would expect the normalized weights $\\bar{\\theta}_i$ to quickly align with the direction that maximizes the utility function and then remain stationary. However, to my knowledge, this is not what happens in practice—for example, in ReLU networks (and even in certain linear settings), the alignment direction can change over time, even without any neuron becoming active. Could you comment on this?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Alternating Gradient Flows (AGF), a framework that models how features emerge in two-layer networks trained from small initialization. AGF captures the staircase-shaped loss dynamics as an alternating process: dormant neurons activate one by one, each time acquiring a new feature and causing a sharp drop in loss. It predicts the timing and magnitude of these drops, aligning with experiments.\n\nThe authors claim that AGF unifies prior analyses across architectures (e.g., linear networks, transformers), and in some cases provably matches gradient flow. Applied to modular addition tasks, it shows that networks learn Fourier features in decreasing order. Overall, AGF offers a structured view of feature learning dynamics.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths : \n\n- The idea of this type of algorithmic/sequential dynamics that mimics the gradient flow has been in the air since 4-5 years and this attempt to unify these works and put some clear mathematical content is a good thing\n- The identification of the independent dynamics of *dormant neurons* and the correlated dynamics of *active ones*\n- The theorem linking AGF and diagonal linear networks\n- The modular addition example\n\nWeaknesses : \n- The other side of the coin is that this article only seems to put words in a phenomenon that people understood (at least since Maennel, 2018). The mathematical content, beyond the equivalence with diagonal linear networks, is not very rich, and the work is mostly a phenomenology.\n- I don't believe that a general conjecture between AGF and gradient flow is plausible : in the case of orthogonal inputs, for two layer relu networks, there might be several neurons that contribute to the loss, while it appears to be that with the AGF only two will be active (one for positive labels one for negative ones). Yet the loss dynamics shape should match (but not parameters dynamics).\n- Sometimes the authors play a bit with the notion of what is *dormant*: one time it corresponds to small in magnitude, another time it comes from the fact that it is orthogonal to the signal etc... I feel that the precise definition of it has to be polished. \n\nMinor typos : \n- line 87 : exponentially many saddle points --> I believe in some settings, the set of saddle points to be connected\n- critical points\n\nQUESTIONS:\nNo\n\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}