Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks

Thank you for your review. We will address the two weaknesses you brought up individually below.

Studying small initializations. The study of learning dynamics from small initialization—and its connection to sparsity and simplicity biases—has a rich history in the theory of deep learning literature. As we discuss and cite (e.g., Woodworth et al., 2020; Li et al., 2020; Jacot et al., 2021), this regime has been instrumental in developing theoretical insights into feature learning. You are correct that saddle-to-saddle learning dynamics of neural networks from small initialization differs from the maximal update parameterization used in most large-scale models, where weights move $\mathcal{O}(1)$ in time. We explicitly acknowledge this distinction in lines 72-75. However, recent work (e.g., Atanasov et al., 2025) has shown that while the scale of updates differs, the qualitative behavior of feature learning—such as alignment and progressive recruitment of neurons—remains similar across regimes. We study small initialization precisely because the slower dynamics and separation of timescales make the underlying feature learning process more analyzable. This structure is what enables the development of AGF, which decomposes training into a sequence of simpler optimization problems. In this sense, the very property that makes small initialization less practical for realistic training is what makes it valuable for theory. In our updated manuscript, we will expand Section 6 to include a more explicit discussion of the limitations of the small initialization regime, clarifying how it serves as a tractable proxy for studying feature learning in more realistic settings.

Extending AGF to deeper networks. We agree, as we state on line 386, that the main limitation of AGF is its restriction to two-layer networks. We believe that an extension to deeper networks is possible. In our revised manuscript, we have added a longer discussion in Section 6 on how future work could build off our framework to extend to deeper networks. One promising direction is to extend our numerical implementation of AGF for two-layer networks (discussed in our response to Reviewer D2un) to ResNet-style architectures by treating each residual block as a two-layer network and the encoder and decoder layers as another two-layer network. Other possibilities would be to consider generalizing the definition of a “neuron” to an input-output path as in the path-framework proposed in Saxe et al., 2022 analysis of deep gated ReLU networks.

We understand that your main concern is in the scope of our work's applicability to real-world scenarios and data. We would like to emphasize that the analysis of feature learning from small-scale initialization and in shallow neural networks is an active area of research, with several papers published in the space every year. While we agree that our work will not directly impact researchers training SOTA models, we believe that it can have a significant impact on this community and our understanding of feature learning in neural networks from first principles. We hope you will consider raising your score to reflect this.

We thank the reviewer for the thoughtful and constructive feedback. We’ll respond in two parts: first addressing your specific questions, then the three main weaknesses. While we appreciate your engagement, we were surprised by the low score. We hope our clarifications help and welcome further discussion to better understand your concerns.

Questions

1. Confusing notation on L64-65.

In our notation, $a_i$ is a vector in $\mathbb{R}^c$, without any assumptions on $c$ in general. In the diagonal linear network and linear transformer settings, $c$ is equal to $1$, while in the fully-connected linear network and modular addition settings, we have $c > 1$. In line 64-65, $a_i$ appears in a vector-times-scalar product. In line 72, the Gaussian is an isotropic one, with diagonal covariance. We agree that the expression in line 72 could be confusing, because $i$ indexes the neuron, and not the element. We will rewrite the expressions to be $a_i \sim \mathcal{N}(0, \alpha / \sqrt{2c} \mathbf{I}_c), w_i \sim \mathcal{N}(0, \alpha / \sqrt{2d} \mathbf{I}_d)$.

2. What is a smooth dynamical system in L77-78?

In lines 77-78, we are simply paraphrasing the main result by Bakhtin. In their work, a “smooth” dynamical system is an ODE (or SDE) where the vector field (and diffusion covariance) is $C^2$. Under this definition, GF is smooth when the activation function $\sigma$ is of class $C^3$, which holds for all the settings we study.

3. Tuning down statement on L43-45

Our intention with this sentence was to convey that prior analyses each rely on different simplifying assumptions—be it about architecture, data, or optimizer—which makes it difficult to extract general principles. We are not trying to imply that our work does not also consider simplifying assumptions. We agree that this sentence could sound too dismissive; we will rephrase it in order to tune it down.

4. What is the role of the parametrization on L337-338?

We believe that the parametrization you described is identical to the one we use, but with a different notation: the $\beta_i$ in your parameterization corresponds to the pair $(u_i, v_i)$ in ours. Under this change of notation, $\langle \beta_i, z \rangle = \langle u_i, a \cdot x \rangle + \langle v_i, b \cdot x \rangle$. We use this notation for two reasons. First, it matches the notation used in prior works such as Gromov 2023 or Morwani et al. 2023, and second, the vector $u_i$ and $v_i$ will have different behaviors in the forthcoming analysis. While they will both be aligned to the same frequency, they can have different phase shifts. Therefore, it is notationally convenient to split the vector $\beta_i \in \mathbb{R}^{2p}$ in these two components.

5. & 6. What is the distinction with prior analysis, and why do we see different feature learning behavior?

We wish to highlight that the key distinction between our analysis and the prior works you cite lies in the initialization regime. In those prior analyses, the network is not initialized near the saddle point at the origin. In contrast, our analysis focuses on the vanishing initialization limit ($\alpha \rightarrow 0$), placing the network infinitesimally close to this critical point. This regime is central to AGF (and to related frameworks such as those in Jacot et al., Li et al., and Pesme & Flammarion); it is the primary mechanism inducing a separation of timescales, allowing the gradient flow dynamics to decouple. This motivates AGF, where such a separation is enforced by design. Instead, the line of research you mention focuses on the large architecture (or large input dimension) limit. Crucially, this might cause a population-distributed behavior, rather than a neuron-specialized one. As a result, the dynamics captured by AGF and the “specialization transition” studied in the prior works appear to arise from fundamentally different mechanisms—though both may reflect transitions between saddle points in the loss landscape. The difference between the picture that emerges from the quadratic analyses you cite and our results with AGF is rooted in this distinction as well.

Moreover, there are other differences in assumptions, analysis techniques, and theoretical goals. We summarize them in the table below, and will discuss them in our updated related work section:

Weaknesses

1. The relevance of AGF beyond simple cases

We agree that there are valid reasons to question the broader applicability of AGF, particularly given its current formulation is limited to two-layer networks. However, we believe that this criticism applies to many analysis in theoretical deep learning. While proving a general equivalence between AGF and GF in the vanishing initialization limit is likely to be extremely challenging, we view our convergence result in the diagonal linear case as strong support that such an equivalence may hold under the right conditions. Moreover, the precise match between theory and experiment across the four very different settings we study, along with the recovery of earlier analyses as special cases, gives us reason to believe that AGF captures a fundamental phenomenon. It is possible that the four settings are the only ones where AGF applies, but given their diversity, we think it is far more plausible that AGF provides a broadly useful lens for analyzing feature learning in two-layer networks.

Regarding the quadratic network example, we are, unfortunately, not sure which additional assumption the reviewer is referring to. The analysis in Section 5 closely follows the setup used in Gromov (2023) and Morwani et al. (2023), and in fact generalizes those works by allowing an arbitrary mean-centered input vector $x$. We would appreciate clarification so we can address this concern more precisely.

2. The “broad picture” offered by AGF

We believe the greatest value of AGF lies in mathematically formalizing an intuition that has been present in the saddle-to-saddle literature for a long time: that learning, at small initialization scales, unfolds in distinct phases. Our goal was to show that a wide variety of prior analyses—previously treated as separate—can be unified under a single framework, and that this framework can explain empirical observations that had lacked a precise theoretical account. We agree that precisely predicting the order, timing, and magnitude of loss drops in simple examples is not valuable on its own. But showing that such predictions can be made from a unified mechanism that captures these transitions analytically is, we believe, a meaningful contribution. As you note, the formulation of the utility function may also have broader implications, particularly for estimating the timescale of escaping from saddles. We see AGF as a starting point, and we can imagine future work applying it to more general two-layer settings, using it to analyze the dynamics of generalization, or extending its principles to deeper networks.

3. Missing references

You are absolutely correct, this was an oversight on our part. We were primarily focused on works studying saddle-to-saddle dynamics and overlooked the extensive body of works studying feature learning in two-layer networks using different techniques such as mean-field, teacher–student, multi-index models, and single-step analyses. In our updated manuscript, we have included a new appendix section devoted to related works on feature learning in two-layer neural networks. We cite all the works you referenced, and a number of other works using similar techniques.

Please let us know if you have any other questions. We hope you will consider raising your score and look forward to further discussion. Thank you again for your review.

We thank the reviewer for thoroughly reviewing our paper, and for highlighting areas needing further clarification. We appreciate the positive feedback on our work. We will address each of the weaknesses and questions mentioned individually.

Full end-to-end analysis. We agree that the AGF framework, like most theoretical frameworks, requires simplifying assumptions to enable tractable analysis. However, we would like to clarify that in the diagonal linear network setting, AGF does in fact yield a full end-to-end analysis of training dynamics, and we prove convergence of AGF to gradient flow (see Theorem 3.1). In the fully-connected linear network case, we provide a complete analysis under the assumption that the input covariance commutes with the input-output cross-covariance; in the more general non-commuting case, we state a conjecture supported by empirical evidence (see Conjecture 4.1). As correctly pointed out in the review, in the linear transformer and modular addition settings, we assume perfect execution of the utility maximization step, in order to keep the analysis tractable. This is precisely why our the predicted jump times in Figures 5 and 6 are presented as lower bounds. Overall, while we acknowledge that the full generality of AGF remains analytically intractable in some settings, our goal is to isolate and analyze regimes where rigorous statements can be made, and the resulting predictions remain consistent with observed behavior.

Active to Dormant. Yes, it is possible in the cost minimization step for active neurons to become dormant again, if the optimization trajectory brings an active neuron near the origin. This is mentioned on lines 142-143, and discussed more thoroughly in the Appendix, Section A.3 (unfortunately, the link on line 144 incorrectly points to Section B). This describes how reversion to dormancy is theoretically possible. While we rarely observe this empirically, it happens for diagonal linear networks – see figure 7 in the Appendix. Here, the first coordinate $\beta^{(1)}$ activates, followed by $\beta^{(2)}$. During the second step of cost minimization, $\beta^{(1)}$ deactivates as it most change sign and thus the optimization takes the neuron near the origin. At the next step $\beta^{(1)}$ activates with the new sign to reach the global minimum. We will highlight this example in the main text.

Numerical implementation of AGF. To explore the applicability of AGF beyond analytically-tractable settings, we have developed a numerical implementation of AGF for arbitrary two-layer neural networks trained in PyTorch (shared via GitHub, though we are unable to link it here due to NeurIPS rebuttal policies). Using this tool, we have compared loss curves between gradient descent (at small initializations) and AGF in general non-orthogonal settings, where theoretical analysis is currently infeasible (e.g., a two-layer ReLU network trained on CIFAR-10). Our findings suggest that AGF can closely match the behavior of gradient descent in these settings. However, this requires careful tuning of hyperparameters, particularly those controlling the termination of the cost minimization step (e.g., number of iterations or gradient norm threshold). We will include an additional plot in Section 6 showing representative examples from these experiments, along with a discussion of when and why numerical AGF may diverge from standard training dynamics. We agree that identifying the precise conditions under which AGF faithfully approximates full gradient descent is a crucial direction for future work, and are actively exploring this question.

Convergence rate. We believe that this is an interesting question, and something we have not spent time thinking about. Likely, the best place to start answering this question is the setting of diagonal linear networks, where we can prove that AGF and GF do converge to each other. Extracting convergence rates would require a careful analysis of the proof by Pesme and Flammarion, who first formally connected gradient flow with a step-wise process. We believe that this might be possible, but it would be subtle and technical. Yet, it is an important point to bring up for future work. We will add to the appendix a short discussion on this.

Please let us know if you have any other questions regarding our work. We hope you will consider raising your score and continue to consider our paper an important contribution to the NeurIPS community.

We are grateful for the thorough review, the constructive suggestions, and the positive feedback. We will address the weaknesses and questions individually.

Validity of the AGF approximation. We acknowledge that Equation 2 is introduced abruptly, and that the derivations in Appendix A.1 are rather compressed. In the updated version of the manuscript, we will expand the derivations in Section A.1, and summarize them in the main body of the paper (Section 2). We wish to point out that Equation 2, together with its related derivations, is meant to approximate the GF dynamics around the origin. Indeed, that equation is used only in the utility maximization step of AGF. This step only involves dormant neurons, which, by definition, lie close to the origin. This justifies treating the utility (of dormant neurons) as a $\kappa$-homogeneous function, and applying Euler’s theorem. In the updated version, we will stress the fact that Equation 2, and the derivations in Section A.1, assume that $\theta_i \sim \mathbf{0}$. Of course, during the utility maximization phase of AGF, the dormant neurons could, technically, leave the region where the approximation is valid. Indeed, AGF is at its core an ansatz; its derivation is not completely formal, and its equivalence with GF at the vanishing initialization limit ($\alpha \rightarrow 0$) is conjectural in the general case (see lines 150-152).

Definition of AGF. We agree that the pseudocode in Algorithm 1 is not completely precise. Our intention was to balance between mathematical rigor and intuition. Below, we address your main concerns raised regarding the formalism of Algorithm 1.

• Our intention with the expression $\frac{d \theta_j}{d t} = - \nabla L (\Theta_{\mathcal{A}})$ was to imply that the gradient is taken w.r.t. $\theta_j$. We will change the notation for the RHS to $\nabla_{\theta_j} L (\Theta_{\mathcal{A}})$, hoping to resolve the confusion.

• It is true that the equilibria during the cost minimization phase can take infinite time to be reached. In fact, the cost minimization step should be interpreted as taking infinite time, akin to how Greedy Low-Rank Learning (GLRL) is expressed in pseudocode in Li et al. 2020 (Algorithm 1 in their work). Like in their work, in practice, we approximate the infinite time limit by running sufficiently many steps. In fact, we find that the amount of time we run cost minimization becomes a crucial hyperparameter when running AGF numerically – see our reply to Reviewer D2un. Additionally, note that in AGF, time can only be accumulated in the utility maximization phase. That is why in Figure 2 (right) – which is meant to represent GF at the limit $\alpha \rightarrow 0$, where the loss becomes piecewise constant – the cost minimization phase is depicted as instantaneous.

• The halting condition $\nabla\mathcal{L}(\Theta_\mathcal{D}) \neq 0$ is a way to express two conditions: (1) the dormant set is empty and (2) if all remaining dormant neurons have zero utility, i.e. their gradient is zero, then the next step of utility maximization will take infinite time and thus the algorithm should be terminated. In this paper we only consider settings where all neurons will eventually activate, however, if for example using ReLU, one could imagine a setting where all remaining dormant neurons are dead and thus no more neurons can activate.

• If the $i$-th neuron becomes dormant again, its accumulated utility $\mathcal{S}_i$ should be initialized at the maximal allowed value, i.e., $c_i$. Such utility will then start decreasing at the next utility maximization phase. In order for this to be consistent with the current halting condition for the utility maximization step, the inequality $\mathcal{S}_i < c_i$ in Algorithm 1 should be corrected into $\mathcal{S}_i \leq c_i$ (otherwise, that neuron will immediately jump and become active again). We thank the reviewer for bringing this up, and will elaborate on this in the updated version, including the correction to the inequality.

Apparent inconsistency with Pesme and Flammarion. We believe that AGF is actually coherent with Pesme and Flammarion’s algorithm – please see Figure 7 in our appendix. This figure actually recreates Figure 2 in Pesme and Flammarion, and visualizes empirical gradient flow dynamics for the parameters, accumulated utility, parameter directions and the predictions for these quantities under AGF. The subtlety is that while the cost minimization phase in AGF continues until convergence, the active neurons are removed immediately when they collapse, i.e., when they approach their initialization again (see Appendix A.3 for a formal discussion on the condition to become dormant again). This is the reason the sentence “Remove collapsed neurons” appears inside the loop of the cost minimization phase in Algorithm 1. We acknowledge that the instantaneous removal of collapsed active neurons is not clear in Algorithm 1; we will elaborate on this in the updated version of the paper. For the example discussed (using the notation in Figure 7 of our work), when the first coordinate $\beta_1$ collapses at time $\tau^{(2)}$, in AGF it would be removed immediately from the active set, and the cost minimization phase would continue for the other active neuron $\beta_2$ until it reaches a value corresponding to the second saddle. Then utility maximization would proceed with neuron $\beta_1$, which is dormant again, starting with an accumulated utility $\mathcal{S}_i = c_i$. It would first lose utility until $\mathcal{S}_i = 0$ at $(\tau^{(2)}+ \tau^{(3)}) / 2$, switching signs, then gaining utility until it activates at $\tau^{(3)}$ with the correct sign. Then during cost minimization $\beta_1$ and $\beta_2$ now both active with the correct sign would converge to the values associated with the minimum.

Questions

• We agree that it is possible for multiple dormant neurons to activate simultaneously. This is an edge case of AGF, where the accumulated utilities of multiple neurons reach their corresponding maximal value $c_i$ at the same exact time. AGF can still be applied in this case; the only nuance is that $i_*$ will not be a single index, but rather a subset of $\mathcal{D}$, which will be subtracted from $\mathcal{D}$ at the end of utility maximization. We will correct this inaccuracy in line 117 by mentioning that, as an edge case, $i_*$ could contain multiple indices, and is not necessarily a singleton.

• If dormant neurons are treated as lying at the origin (which is exact at the limit $\alpha \rightarrow 0$, and only approximate away from it), the intermediate points in AGF are critical points (for the loss) of the type described in lines 85-94. Moreover, the trajectory of the dynamics passes in their proximity, approaching and then departing (excluding the last point, where the dynamics converges). At the limit $\alpha \rightarrow 0$, this implies that such points possess a stable (infinitesimal) direction, and an unstable one, i.e., they are saddles of the loss.

• We agree that one could intuitively expect the dormant neurons to align instantaneously (in the limit $\alpha \rightarrow 0$), and to stay aligned throughout utility maximization. However, as stated in the review, this is not what happens, even for scenarios considered in our work. We discuss this question in Appendix A.4, where we argue that instantaneous alignment is expected for $\kappa = 2$, but not otherwise. This is why the jump times for linear transformers and quadratic networks, for which $\kappa = 3$, are stated as (approximate) lower bounds (Equations 7 and 9). When the alignment is not instantaneous, the direction $\bar{\theta}_i$ of a dormant neuron can change during a single utility maximization phase (i.e., without any neuron becoming active). Additionally, even for $\kappa = 2$ one could see a change in direction for dormant neurons, but this change would be instantaneous and associated with the accumulated utility reaching zero. In fact that is exactly what we see in Figure 7 for the neuron associated with $\beta_1$.

We hope that the above discussions will clarify your concerns, and that this will increase your confidence in the importance of our work. Please let us know if you have any other questions. Thank you again for your very detailed and constructive review!

Thank you for your detailed responses. My questions and concerns have been addressed.

I will keep my score as it is. However, I would strongly encourage that the revised version takes these concerns into account. I am fine with the idea that AGF is an ansatz—but it should be crystal clear to the reader which parts of the construction are heuristic or approximate, and which are formal.

Moreover, in my view, there are too many opaque or ill-defined components in the current formulation of the AGF algorithm. I don’t mind having an intuitive or informal version in the main text, but I believe there should also be a fully precise and rigorous definition somewhere in the paper, one that eliminates ambiguities such as those I raised. This is particularly important given that the paper states formal theorems, for which one expects a corresponding level of rigour in both definitions and proofs.

We appreciate the thoughtful review, and we are happy to hear you think our work merits acceptance. We will respond individually to the weaknesses and questions raised regarding our paper.

The mathematical content. We agree that many of the ideas behind AGF (i.e., two phases of optimization separated by timescales) existed in prior literature, albeit not described in a unified framework. However, we disagree that the mathematical content in this work, beyond the equivalence with DLNs, is not rich. To this end, we wish to point out that nearly all the analysis in Section 5 (“AGF Predicts the Emergence of Fourier Features in Modular Addition”) is extremely novel and non-trivial. In particular, we highlight the following two theorems in this section:

• Theorem 5.1 proves that the utility maximizing directions for a quadratic neuron in the modular addition task are cosine waves aligned with the dominant frequency of the encoding vector. Proving this involves recasting a nonlinear constrained optimization problem into a solvable optimization problem in the Fourier domain. Here we make use of the shift-equivariance property of the Fourier transform and Plancharel’s theorem.

• Theorem 5.2 proves that there exists a lower bound on the loss within the subspace spanned by neurons aligned to the same frequency, and provides precise conditions on the relationship between these neurons in order to achieve this lower bound. The proof again involves tools from harmonic analysis and several subtle trigonometric identities.

We do not consider our work as “phenomenological”, but rather as a precise mathematical description of intuition that unifies existing analysis, and that can be used to derive novel analysis leading to predictions matching empirical observations.

General conjecture. Thank you for your thoughts here. We agree that proving the general conjecture between AGF and GF is likely to be extremely challenging. Yet, given that convergence of AGF towards GF is established in the setting of diagonal linear networks, we believe that there is substantial evidence to support a general conjecture, under the right constraints. For example, a constraint requiring that all neurons are needed to achieve zero training error, something that is true in the cases we studied and would prevent the situation you suggested in your example, might be needed. In order to clarify this, we will add a further discussion in Section 6 around when a general conjecture might be plausible, together with steps to be considered for future work.

Definition of dormant. In the context of AGF, a neuron is defined to be dormant if its accumulated utility is less than its initialization-dependent threshold – see Section 2.2 or Figure 2. We are always consistent with this definition throughout the work. When motivating AGF, we introduce the concept of dormant and active neurons to describe the structure of saddle points in the loss landscape. Here, dormant and active are consistent with the definitions used in AGF. However, when describing GF near saddle points, we use dormant and active as adjectives to describe neurons that are either small in norm or not. In a sense, since GF is continuous, at finite scale a neuron is never dormant, but rather active to a degree. We understand that this language could cause confusion, and we will work to improve the writing here. Is there another section in the paper where you think we could be more formal with the definition of dormant?

Minor typos. We agree that, in some settings, the saddle points described in lines 85-94 can be connected, in the sense that some of them can coincide. Note, on line 87 we state “one of potentially an exponential number of critical points”. For example, in diagonal linear networks, with generic data, the number of critical points is exponential. However, it is possible to construct data where the number of saddle points is less. We use the term "critical points" instead of "saddle points" because these points are not necessarily saddles; they could be (local) minima.

Thank you again for your constructive feedback. We hope we addressed some of your concerns regarding our work. If we have, we would appreciate it if you could raise your score to reflect this.