Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks
Authors: Daniel Kunin
Giovanni Luca Marchetti
Feng Chen
Dhruva Karkada
James B Simon
Michael R DeWeese
Surya Ganguli
Nina Miolane
Venue: NeurIPS 2025 poster ()
Keywords: feature learning
saddle-to-saddle dynamics
incremental learning
vanishing initialization
emergence
modular addition
Fourier features
Primary area: theory
TL;DR: We introduce Alternating Gradient Flows, a framework modeling feature learning in two-layer networks with small initialization as utility maximization and cost minimization—unifying saddle-to-saddle analyses and explaining Fourier feature emergence

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.

Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks
Daniel Kunin' Stanford University
Giovanni Luea Marchettit
Feng Chen Stanford University
Dhruva Karkada UC Berkeley
KTH
James B. Simon UC Berkeley and Imbue
Michael R. De
Surya Ganguli
Nina Miolane UC Santa Barbara
UC Berkeley
Stanford University
Abstract
What features neural networks learn, and how, remains an open question. In this paper, 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 utility function over dormant neurons and minimizing a cost function over active ones. AGF begins with all neurons dormant. At each iteration, 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-oo-saddly 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, 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.
Introduction
The impressive performance of artificial neural networks is often attributed to their capacity to learn features from data. Yet, a precise understanding of what features they learn and how remains unclear.
A large body of recent work has sought to understand what features neural networks learn by reverse engineering the computational mechanisms implemented by trained neural networks [1+]]. Known as mechanistic interpretabilliyy (MI), this approach involves decomposing trained networks into interpretable components to uncover their internal representations and algorithmic strategies 7 MI has achieved notable successes in understanding the emergent abilities of large language models [8) including identifying induction heads that enable in-context learning [ and revealing that small transformers trained on algebraic tasks use Fourier features 66 [9. Despite these discoveries,
mechanistic interpretability remains limited by its empirical nature, lacking theoretical framework
to formally define features and predict when and how they emerge [10
Correspondence kunin@berkeley.cmm edu and glma@kth.aa
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Two-layer Nonlinear Networks
Diagonal Linear Networks
Fully-Counected Linear Networks
Atteution-Quly Linear Transformers
Quadratic Networks for Modular Addition
16
Wit
Wit
WHT
0.
W
K
K
K=10
B
X
a
b
Loss
GD AGE
Utility
Maximization
Learus
Learns
Learns principal
Learns
Cusl
zegression
singular
Fourier
Minimization
Steps
coellicient
vectors
component
frequency
Figure I: A unified theory of feature learning in two-layer networks. Left: Alternating Gradient
Flows (AGF) models feature learning as a two-step process alternating between utility maximization
(blue plateaus) and cost minimization (red drops), where each drop reflects learning a new feature
(see Section 2). Middle: AGF unifies prior analyses of saddle-to -o-eddde dynamics (see Sections and 4)
Right: AGF enables new analysis of empirical phenomena (see Section
different line of work, rooted in deep learning theory, has sought to understand how features are learned by directly studying the dynamics of the network during training. Empirical studies suggest that neural networks learn simple functions first, progressively capturing more complex features during training [11444 Termed incremental, stepwise, or saddle-to tolle-to-sddle, this learning process is marked by long plateaus minimal loss change followed by rapid drops. It is conjectured to arise from networks, initialized at small-scale, jumping between saddle points of the loss landscape (15), with each drop corresponding to the acquisition a new feature [6]. This saddle-to-saddle process has been explored across range of simple settings, including diagonal linear networks [17] [8].
fully connected linear networks [19 15), tensor decomposition 220 21). self-supervised learning 22 23). shallow ReLU networks 34 attention-only transformers Bil [32, and multi-index
models 33-39] 39 However, these analyses often rely on different simplifying assumptions about the architecture (e.g., linear activations), data (e.g., orthogonal inputs), optimizer (e.g., layer-wise learning), making it difficult to establish a unified, scalable understanding of feature learning.
In this work, introduce a theoretical framework that unifies existing analyses of saddle-to-saddle dynamics in two-layer networks under the vanishing initialization limit, precisely predicting the order, timing, and magniudee of loss drops, while extending beyond classical settings to explain empirically observed patterns of feature emergence. See Figure for a visual overview of the paper. Our main contributions are:
1. We introduce Alternating Gradient Flows (AGF), a two-step algorithm that approximates gra- dient flow two-layer networks with small initialization by alternating between maximizing a utility over dormant neurons and minimizing cost over active ones (Section 2
2. We prove that AGF converges to the dynamics of gradient flow for diagonal linear networks in the vanishing initialization limit (Section 3
3. We show how AGF unifics and generalizes existing theories of saddle-to-sadee saddle dynamics across fully connected linear networks and attention-only linear transformers (Section 4
4. We use AGF to predict the emergence of Fourier features in modular arithmetic tasks, providing the first theoretical account of both the order in which frequencies appear and the dynamics that drive it (Section 5).
Our work takes step toward unifying deep learning theory and mechanistic interpretability, suggest- ing that what features networks learn, and how, can be understood through optimization dynamics. See Appendix A for further discussion of related theoretical approaches studying feature learning and saddle-to-sade saddee dynamics in two-layer networks, including mean-field analysis 43 and teacher-student frameworks (4H51]).
2 Deriving Two-step Algorithm that Approximates Gradient Flow
We consider a two-layer neural network with H/ hidden neurons of the form [a; +) = Σ (w) gi (x)), where each hidden neuron i has learnable input weights w; Rd and output weights ag E R., and processes potentially neuron-specifie input representation g())).
Setup and notations.

---

The activation d: R R is origin--asssing i.e., (0) = 0. as satisfied by common functions (e.g., linear, ReLU, square, tanh). Let [1 0,) = a,0 ! 9: (2)) denote the output of the 5th neuron, where 0. = (w; a.). The parameters θ = (01..... 02) E I (a+c) for the entire network evolve under gradient flow (GF), é = (()), to minimize the MSF loss [(O)- = Ez [ (lv(a)) [(x; }) where y: R - I is the ground truth function and the expectation is taken over the input distribution. The learning rate n > O rescales time without affecting the trajectory. The parameter vectors for each neuron are initialized a; (0, a² L.). w.; ~ (O, a² Ia). In the limit o - oc, the gradient flow 2d 53 When a = 1, the dynamics correspond to a mean-ffeld or
2n
dynamics enter the kernel regime [52]5]]
maximal update parameterization 40 54 We study the small-scale initialization regime Or < 1,
where the dynamics are conjectured to follow a saddle-to-saddle trajectory as a - [15
2.1 The behavior of gradient flow small-scale initialization
Bakhtin (55) showed that, under suitable time rescaling, the trajectory of smooth dynamical system initialized at a critical point and perturbed by vanishing noise converges to a piecewise constant process that jumps instantaneously between critical points. These results suggest that, in the vanishing initialization limit a - 0, gradient flow converges to similar piecewise constant process, jumping between saddle points of the loss before reaching a minimum. To approximate this behavior, we first examine the structure of critical points in the loss landscape of the two-layer networks we study, and then analyze the dynamics of gradient flow near them. These dynamics reveals two distinct phases, separated by timescales, which in turn motivates the construction of AGF.
Critical points are structured by partitions of neurons into dormant and active sets. As a ++ the initialization converges to the origin θ O, which is a critical point of the loss. The origin is one of potentially an exponential number of distinct (up to symmetry) critical points, that can be structured according to partitions of the H hidden neurons. Formally, split neurons into two disjoint sets: dormant neurons D S [H] and active neurons = D,, with parameters Op = (0.)iED and θ = (0z)iE,A respectively. Due to the two-layer structure of our model and the origin-passing activation function, setting the parameters of the dormant neurons to zero OD = 0 restricts the loss to function x)) of only the active neurons. Now any critical point O of the restricted loss yields critical point θ = (0, ²*a of the original loss. Thus the 2 partitions of dormant and active neurons structure critical ponts of the loss landscape, with the origin corresponding to the special critical point with all neurons dormant.
Dynamics near a saddle point. Consider the dynamics near a saddle point defined by a dormant and active set, where we assume the active neurons are at local minimum of their restricted loss A)) By construction, neurons in our two-layer network only interact during training through the residual (a;; θ)= y(x)= fz; θ) E IS which quantifies the difference between the predicted and target values at each data point x. Near critical point, the residual primarily depends on the active neurons, as the contribution from the dormant neurons is effectively zero. Consequently, dormant neurons evolve independently of each other, while active neurons remain collectively equilibrated. To characterize the dynamics of the dormant neurons, which will determine how the dynamics escape the saddle point, we define the utility function, 24; IRdte + R for each i E D as:
(1(0r)) = R., ())))) where r(x) = v(x) - O.).
Intuitively, the utility U4 quantifies how correlated the 2h dormant neuron is with the residual r, which itself is solely determined by the active neurons. Using this function we can approximate the directional and radial dynamies for the parameters of each dormant neuron as
d 0. n||0,11 2 V (o,l() 0e
11 = my< | (0(1=~((((00))))
(2)
where 0, = 0./110.1 are the normalized parameters, Pd. = (1. 0,0!) is an orthogonal projection matrix, and κ 2 2 is the leading order of the Taylor expansion of the utility from the origin. Note that the directional dynamics scale as ||0]|| X- while the radial dynamics scale as [|@.]|--1 Because ||0:]| < 1 for dormant neurons, the directional dynamics evolves significantly faster than the radial dynamics. Consequently, the contribution of dormant neuron to the residual is suppressed, keeping the residual effectively constant, and this approximation to the dynamics stays valid. This separation of timescales between directional and radial components has been utilized in several previous studies of gradient descent dynamics near the origin [56661].
3

---

Algorithm 1: AGF
Conceptual Illustration: AGF
Initialize: D [H].A A+ 0.8 0ER
while V vL(8p) d0
Neuron
Neuron
Neuron
Activattes
Ep
while Vi E S. S c, do
Active neurons
For: i E D da
sock collectivdyy at lange stidss to miniiiiie the coo
de
na
||:
2p v)(
Vo ())
Activate beuron: A } J. (1.
while ( x) # U do
For do
Dormant indivsiually mustimite their utilry
((O
Remove collapsed neurons: D. + DUC,ANC
T2
Steps
Figure 2: Alternating Gradient Flows (AGF). AGF alternates between utility maximization (blue) over dormant neurons and cost minimization (red) over active neurons, predicting saddle-to-sadee dynamies in training. Utility maximization: dormant neurons evolve independently driven by projected gradient flow to maximize their utility. We keep track of the accumulated utility for each dormant neuron to determine the first neuron to transition to active. Cost minimization: Active neurons work collectively driven by gradient flow to minimize the loss until convergence. After convergence, active neurons compute a new residual that determines the utility at the next iteration, is possible in this step for active neurons to become dormant again (see Appendix
Estimating the jump time from dormant to active. The dynamical approximation in Equation breaks down if a dormant neuron crosses the threshold ||@:|| O1), becoming active by influencing the residual. We denote by & E 'D the first dormant neuron to reach this threshold and by T., the jump time at which this dormant to active transition occurs. To determine T.., we compute for each dormant neuron the earliest time when [[0:(T;)]] > 1, such that i, = arg min Tú This time is defined implicitly in terms of the accumulated utility, the path integral (t) Jo x24, k2l, (s): ds:
T. inf {t {t> 0 Sc(t)
log (0)1) where c; 2ース (0) 1) if k: > 2.
if K = 2,
(3)
In the vanishing initialization limit as - O, the jump time defined in Equation as gradient flow exhibits extreme slowing near the critical point at the origin. Thus, to capture meaningful dynamics in the limit, we accelerate time by setting the learning rate 17 Ee. [c]] where is the threshold defined in Equation 33 and the expectation is taken over the initialization.
diverges
2.2 Alternating Gradient Flow (AGF) approximates GF at small-scale initialization
Based on the behaviors of gradient flow (GF) at small-scale initialization, we introduce Alternating Gradient Flows (AGF) (Algorithm D as an approximation of the dynamics. AGF characterizes phases of constant loss as periods of utility maximization and sudden loss drops as cost minimization steps, accurately predicting both the loss levels and the timing of loss drops during training. At initialization, all neurons are dormant. At each iteration, a neuron leaves the dormant set and enters the active set. The dormant neurons work individually to maximize their utility, and the active neurons collectively work to minimize a cost (Figure 2 Specifically, each iteration in AGF is divided into two steps.
Step 1: Utility maximization over dormant neurons. In this step, only dormant neurons update their parameters. They independently maximize (locally) their utility function over the unit sphere by following the projected gradient flow (left of Equation (2). We keep track of the accumulated utility for each neuron to determine the first neuron to transition and the jump time using Equation (3). At the jump time, we record the norms and directions of the dormant neurons that have not activated, as they will serve the initialization for the utility maximization phase of the next iteration.
Step 2: Cost minimization over active neurons. In this step, all active neurons interact to minimize (locally) the loss, by following the negative gradient flow of L(O.)) For previously active neurons, the initialization is determined by the previous cost minimization step. For the newly activated neuron, the initialization comes from the utility maximization step. It is possible in this step for active neurons to become dormant again if the optimization trajectory brings an active neuron near the origin (see Appendix After convergence, we compute the new residual rax) = y(x) - [x; A)), which defines a new utility for the remaining dormant neurons at the next iteration.

---

Termination. We repeat both steps, recording the corresponding sequence of jump times and loss levels. until there are either no remaining dormant neurons or we are at local minimum of the loss. Through this process, we have generated precise sequence of saddle points and jump times to describe how gradient flow leaves the origin through saddle-to-sadee process.
While AGF and GF exhibit similar behaviors at small-scale initialization (a & 1), a natural question is whether their trajectories converge each other -> O. While general proof remains open (see Section 6). in the next section we present an illustrative setting where convergence can be established.
A Setting Where AGF and GF Provably Converge to Each Other
Diagonal linear networks are simple yet insightful models for analyzing
Loss
Coefficieut A
a 0.99
learning 662 17 66). Central to their analysis is an interpretation of gradient flow in function space as mirror descent with an initialization- dependent potential that promotes sparsity when a is small. Using this perspective, Pesme and Flammarion [8] characterized the sequence of saddles and jump times for gradient flow in the limit & + O. We show that, in this limit, AGF converges to the ex- act same sequence, establishing a set- ting where AGF provably converges to gradient flow (see Figure 3).
10
0
0.50
91
10
1e-02
0
a= 1e-04
-10
3g
10
a 1e-08
3,
a= 1e-16
10-5
AGF
20
10
10
10
Steps
10-1
10
107 Steps
Figure 3: AGF = GF a -+ in diagonal linear networks. Training loss curves for diagonal linear network under the setup described in Section B for various initialization values a. As + O, the trajectory predicted by AGF and the empirics of gradient flow converge. To ensure a meaningful comparison between experiments we set v) - log(a))
We consider two-layer diagonal linear network trained via gradient flow to minimize the MSE 1 lly - xT', where the regression coefficients 3 E Rd are parameterized as the loss L(3) = 2n
element-wise product B = u @ o with u, E Rd, and (X,y) E RRdxN x R are the input, output data. This setup fits naturally within AGF, where the 7th neuron corresponds to 3, with parameters 0, (u,, 04), data representation ga 2;, and
2. Following Pesme and Flammarion (18)
we initialize (0) = (v2a, 0) such that (() = 0, and assume the inputs in general position, standard technical condition in Lasso literature 67 18 to rule out unexpected linear dependencies.
Utility maximization. The utility function for the 7th neuron is U4 uu V B, (( which is maximized on the unit sphere u² + v² when sgn(u,,;) = sgn( (V v. ((aa)) and |u, loal. 3, What makes this setting special is that not
yielding a maximal utility of Uū only can the maximal utility be computed in closed form, but every quantity involved in utility maximization naanelnnaalyyeelleieeeennennaaaa the accumulated utility, jump times, and directional dynamics- admits an exact analytical expression. The key insight is that the normalized utility 27 for each dormant neuron evolves according to a separable Riccati ODE, interpolating from its initial value (0) to its maximum U4," As a result, we can derive an explicit formula for the normalized utility Ui (t), whose integral yields the accumulated utility (t), evolving as:
(T() +1= =
1. log cosh (2na (2 + 1 cosh exp (2 (7 (+())))) 27p
2n
where 7a log( v 22) is the learning rate and the unspecified sign is chosen based on sgn( u,r;) and sgn V A)) as explained in Appendix C
This expression allows us to determine the next neuron to activate, as the first i E D for which S 1, from which the jump time can be computed.
Cost minimization. Active neurons represent non-zero regression coefficients of 3 (see Figure right). During the cost minimization step, the active neurons work to collectively minimize the loss. sgn( V v, ((3x)) If, during
When a neuron activates, it does so with a certain sign sgn(u;e;) the cost minimization phase, a neuron changes sign, then it can do so only by returning to dormancy first. This is due to the fact that throughout the gradient flow dynamics, the quantity u² - oz = 20 conserved and thus, in order to flip the sign of the product u,u; the parameters must pass through their initialization. As result, the critical point reached during the cost minimization step is the
5

---

unique solution to the constrained optimization problem,
{ -
if A, \ 3; (.n-gg((u))) 2 0 if if A.
A = arg min ((3) subject to SERd
(5)
All coordinates where (Bx); = are dormant at the next step of AGF.
AGF in action: sparse regression. We now connect AGF to the algorithm proposed by Pesme and Flammarion (18, which captures the limiting behavior of gradient flow under vanishing initialization. Their algorithm tracks, for each coefficient of B, an integral of the gradient, (t) = L 3,,4a (()()) ds, with a time rescaling (s) = log(a)s. They show that in the limit a + 0, this quantity is piecewise linear and remains bounded in [-1, 1.] Using these properties. each step of their algorithm determines a new coordinate to activate (the first for which (t) =1), adds it to an active set, then solves a constrained optimization over this set, iterating until convergence. Despite differing formulations, this process is identical to AGF in the limit - 0:
Theorem 3.1. Let (BAGr, LAGF) and (Brr, (PF) be the sequences produced by AGF and Algorithm / of Pesme and Flammarion 178), respectively. Then, (BaGr, LAGF) - (Brr, (pr) pointwise as a - O.
The key connection lies in the asymptotic identity log cosh(x) - |x|| as |x] + 0o, which implies that the accumulated utility (t) in AGF converges to the absolute value of the integral S; (t). The unspecified sign in AGF corresponds to the sign of the boundary conditions in their algorithm. Thus, AGF converges to the same saddle-to-saddle trajectory as the algorithm of Pesme and Flammarion [8], and therefore gradient flow. See Appendix for full derivation.
4 AGF Unifies Existing Analysis of Saddle-to-Saddle Dynamics
Fully connected linear network. Linear networks have long served as an analytically tractable setting for studying neural network learning dy- namics (3). Such linear networks exhibit highly nonlinear learning dy- namics. Saxc et al. [70] demonstrated that gradient flow from task- aligned initialization learns a sequential singu- lar value decomposition of the input- output cross- cross-covariance. This behav- ior persists in the vanishing initial- ization limit without task alignment [57
(a) Commuting
(b) Non- conncommuting
Singular Value
4
2-
0
10
10c
10
Steps
Figure 4: Stepwise singular value decomposition. Training two-layer fully connected linear network on Gaussian in- puts with power-law covariance Ege and labels y(æ) Bx generated from a random B. We show the dynamies of the singular values of the network's 's map AW when Exe com- mutes with ET Eyz (a) and when it does not (b). Conjec- ture 4.1 (black dashed lines) predicts the dynamics well.
20 15), and is amplified by depth 75 62). Here, we show how AGF naturally recovers such greedy low-rank learning.
74
T9
We consider fully connected two-layer linear network, fx; 0) AW x, with parameters W E #f and A E Rex# The network is trained to minimize the MSE loss with data generated from the linear map y(x) = Ba, where E Rexd is an unknown matrix. The inputs are drawn independently from a Gaussian
NO, Eex), where Eez E Rdxd is the input covariance, and Σyz BE E Rexd is the onput-output cross-covariance, which we assume are full rank with distinct singular values. A subtlety in applying AGF to this setting is that the standard notion of neuron used in Section 2 is misaligned with the geometry of the loss landscape. Due to the network's invariance under (W, A) + (GW, AWWAAG-1 for any G E GL (R), critical points form manifolds entangling hidden neurons, and conserved quantities under gradient flow couple their dynamics (76] (77). To resolve this, we can reinterpret AGF in terms of evolving dormant and active orthogonal basis vectors instead of neurons. Each basis vector (ā,, wh)) ge+d forms a rank-1 map T Red such that the function computed by the network is the sum over these rank-1 maps, fz: Q) = Eie(in ā,WāJ, x, where H min(c. H, d). From this perspective, at each iteration of AGF the dormant set loses one basis vector, while the active set gains one. See Appendix D for details.
Utility maximization. Let A - ā.we be the function computed by the active basis vectors and m [D]. The total utility over the dormant basis vectors is U Em V.V ((

---

Maximizing this sum while maintaining orthonormality between the basis vectors yields Rayleigh quotient problem, whose solution aligns the dormant basis (ā,, ) with the top m singular modes (u,, o4) of V (BA( Each aligned pair attains maximum utility of U, /2, where o; is the corresponding singular value. The basis vector aligned with the top singular mode activates first, exiting the dormant set and joining the active one.
Cost minimization. The cost minimization step of AGF can be recast as a reduced rank regression problem, minimizing ((3) over Rexd subject to rank(B) = kt where A = |A]]. As first shown in Izenman (78), the global minimum for this problem is an orthogonal projection of the OLS solution 89 = Ug P Eyz E where Pu,- U is the projection onto Us E Re* the top cigenvectors of Egg Pay 2yz- Using this solution, we can show that the next step of utility maximization will be computed with the matrix (gf(31) = Pt. Eye where P = I. Pu.
greedy low-rank learning. In the vanishing initialization limit, AGF reduces to an iterative procedure that selects the top singular mode of VaL(8²), ). transfers this vector from the dormant to the active basis, then minimizes the loss with the active basis to update 8 This procedure is identical to the Greedy Low-Ratd wow-Rank Learning (GLRL) algorithm by Li et al. [20] that characterizes the gradient flow dynamics of two-layer matrix factorization problems with infinitesimal initialization. Encouraged by this connection, we make the following conjecture:
AGF in action:
Conjecture 4.1. In the initialization limit - two-layer fully connected linear network trained by gradient flow, with 17 = log(a), learns one rank at time leading to the sequence
(x) Eish Pu, Eyz
e(R)
Eisk Mss
7(k)
Eish Ar(),
(6)
where 0 OSK. k S min(d, H, c), (us.Pe) are the eigenvectors and eigenvalues of Eyzz End ET o,
the 7th singular value P Uv Eyas and Ar(") = (1-2) 11- Ar(+1) /o (f-1) with Ar(0) 0.
When Enz commutes with Eye E Eym, Conjecture 4.1 recovers the sequence originally proposed by
Gidel et al. 19). Figure eempiricaly supports this conjecture. See Appendix D for details.
Attention-only linear transformer. Pre- trained large language models can learn new tasks from only few examples in context,
GD
.10-
AGF
Singular Value
without explicit fine-tuning [79]. To understand the emergence of this in-context learning ability,
previous empirical [80) and theoretical works 83 88] have examined how transformers
.05-
learn to perform linear regression in context.
Notably, Zhang et al. [3] showed that an attention-only linear transformer learns to implement principal component regression sequentially, with each learned component corresponding to a drop in the loss. We show that AGF recovers their analysis (see Figure 5)
.00
10
10
Steps
Figure 5: Stepwise principal component regres- sion. Training linear transformer to learn linear regression in context. We show the evolution of singular values of Σ V,K,Q,. Horizontal lines show theoretical A and vertical dashed lines show lower bounds for the jump time from Equation 7 with = k - 1. Dashed black lines are numerical AGF predictions.
We consider a attention-one only linear transformer f(X; = X X! Wy,XXXTW..I WK.WW. where X is the input sequence, Wy E (dd+))))(dd11) WK, Wo E #=x(d+1)xR represent the value, key, and query matrices respectively, and H denotes the number of at- tention heads. While this network uses linear activations, it is cubic in its input and parameters. Also, when the rank R = 1, each attention head behaves like homogeneous neuron with 3. making it compatible with the AGF framework.
We consider the linear regression task with input X; = (%) for i < N, where x; ~ AO, Exz). and y; = .Tz, with 3 ~ (0, La). The input covariance Σyn E B has eigenvalues \ 2 ... da > 0 with corresponding eigenvectors uI,..., Uu. The final input token is XN+1 = (#N+1) W+1 ). The network trained by MSE loss to predict yN+1 given the entire sequence X, where the model prediction is taken to be ØN+1
(
f(X; θ) 0)411.441 Following Zhang et al. [32]. we initialize all
7

---

parameters to zero except for some slices, which we denote by V E ER" QQ K E R which are initialized from (V(0, a). The parameters initialized at zero will remain at zero throughout training [2]. and the model prediction reduces to ÚN+1 = > h-1 V1.
>n Un Q We defer B and briefly outline how AGF predicts the saddle-to-saddle dynamics.
derivations Appendix
Utility maximization. At the hthh iteration of AGF, the utility of a dormant attention head is which is maximized over normalized parameters when V. =1//3 and Q..K, = = /3, where the sign is chosen such that the maximum utility is 27" NNA/3v3. In other words, utility maximization encourages dormant attention heads to align their key and query vectors with the dominant principal component of the input covariance not yet captured by any active head. This creates a race condition among dormant heads, where the first to reach the activation threshold, measured by their accumulated utility, becomes active and learns the corresponding component. Assuming instantaneous alignment of the key and query vectors, we can lower bound the jump time for the next head to activate, shown in Equation
4; NV.QT (E), k-1 X
Cost minimization. During cost minimization, because wi,..., Ca form an orthonormal basis, the updates for each attention head are decoupled. Thus, we only need focus on how the magnitude of the newly active head changes. Specifically, we determine the magnitude A& of the newly learned function component, given by fx(X;; 01)a-1,N++1
A En= Un 2x N Solving BC(Aa) = O, we find that the optimal magnitude is A i. +(N +1JAr" from which we can derive 0.d,
the expression for the prediction and loss level after the :tth iteration of AGF, as shown in Equation 77).
AGF in action: principal component regression. At each iteration, the network projects the input onto a newly selected principal component of Enr and fits a linear regressor along that direction. Let a- This process yields the sequence,
(1()) m [[0(()()]] for I < i and Ma = a
(k) UN+1 = D
-i +
K.N
((k) ur E
NA./2 +N+1
((k) 2 (1)
NX.
which recovers the results derived in Zhang et al. [32] and provides excellent approximation to the gradient flow dynamics (see Figure 5)
5 AGF Predicts the Emergence of Fourier Features in Modular Addition
In previous sections, we showed how AGF unifies prior analyses of feature learning in linear networks. We now consider a novel nonlinear setting: a two-layer quadratic network trained on modular addition. Originally proposed as minimal setting to explore emergent behavior (89), modular addition has since become a foundational setup for mechanistic interpretability. Prior work has shown that networks trained on this task develop internal Fourier representations and use trigonometric identities to implement addition as rotations on the circle 6. 90 91). Similar Fourier features have been observed in networks trained on group composition tasks 9. and in large pre-trained language models performing arithmetic [92, 93]. Despite extensive empirical evidence for the universality of Fourier features in deep learning, precise theoretical explanation of their emergence remains open. Recent work has linked this phenomenon to the average gradient outer product framework 194. the relationship between symmetry and irreducible representations (55), implicit maximum margin biases of gradient descent [96). and the algebraic structure of the solution space coupled with simplicity bias [97]. Here, we leverage AGF to unveil saddle corresponds to the emergence of a Fourier feature in the network (see Figure 6).
dynamics, where each
We consider setting similar 990 96 97], but with more general input encodings. Given p E N. the ground-truth function is y: Z/px2/p + Z.pp. (a, b) + a a+ mod p. where Z/p the (additive) Abelian group of integers modulo p. Given a vector x IR,,, we consider the encoding E/p + RP, a + a x. where denotes the action of the cyclic permutation group of order on (which cyclically permutes the components of 2). The modular addition task with this encoding consists of mapping (a xr b x) to (a + b) x. For
is
cp. our encoding coincides with the ooe-hot encoding a - Ca studied prior work. We consider a two-layer neural network with quadratic activation function (2) = 2², where each neuron is parameterized by 6, (u,, Us w;) E 3 and computes the function fa xr. b- x; 0,) = ((u,, a aa-20 x) + (v;, b x)) ) Pu.. The network is composed of the sum of H neurons and is trained over the entire dataset of p² pairs (a, b). We make some technical assumptions in order to simplify the analysis. First, since the network contains no bias term, we assume that the data is centered, i.c., we subtract ((z, 1)/p)1 from Second, in order to encourage
8

---

(a]
Power
(b) Predictions fa 6.
(c) Weights
(d)
Frequency-+ Phases
30 or
Frequeacy
f(u 1x b. -.;
642
004
-03-
5.
10 [10=)
0.89
0:30
9.20
60-
0-
-055
Steps
102
Figure 6: Stepwise Fourier decomposition. We train a two-layer quadratic network on a modular addition task with p 20, using a template vector E I' composed of three cosine waves: [[1] 10, a[3] = and 2 [[]] 2.5. (a) Output power spectrum over time. The network learns the task by sequentially decomposing x: into its Fourier components, acquiring dominant frequencies first. Colored solid lines are gradient descent, black dashed line is AGF run numerically from the same initialization. (b) Model outputs on selected inputs at four training steps, showing progressively accurate reconstructions of the template. (c) Output weight vector w; for all H 18 neurons and (d) their frequency spectra and dominant phase. Neurons are color-coded by dominant frequency. As predicted by the theory, the neurons group by frequency, while distributing their phase shifts.
sequential learning of frequencies, we assume that all the non-conjugate Fourier coefficients of x have distinet magnitude: [=[|]| * [a[k']| for k * = (mod p). Third, in order to avoid dealing with the Nyquist frequency k = p/22, we assume that p is odd. We now describe how AGF applies to this setting (see Appendix F for proofs). Here, denotes the Discrete Fourier Transform (DFT).
Utility maximization. First, we show how the initial utility can be expressed entirely in the frequency domain of the parameters Ô. For a dormant neuron parameterized by 0 = (u, v, w). the Re k We then argue that the unit vectors
initial utility is Ll = (2/r²) Eke [=] 2 a[A:)
maximizing 44 align with the dominant harmonic of .::
Theorem 5.1. Let & be the frequency that maximizes [[(]], k L..... p and denote by 82 the phase of Z(E). Then the unit vectors 0, = (u., V., w.) maximizing the utility function U are:
cos (web + 8v), w.[c] = Ap cos (weC + sw)) (8) where a, b, e E (p) are indices, Su, Su, 8m E are phuse shifts satisfying su + 8y = 8 + 85 (mod 2T)), Ap = (3p) is the amplitude, and we = =2mE/p p is the frequency. Moreover, UI has no other local maxima and achieves maximal value of'r" = <2/(27>>)|<()))
ua[a] = Ap cos (wza + su): [b] = cos
Therefore, after utility minimization, neurons specialize unique frequencies. Now that know the maximal utility, we can estimate the jump time by assuming instantancous alignment as done in Section 4 resulting in the lower bound shown in Equation
Cost minimization. To study cost minimization, we consider regime in which a group A of N S H neurons activates simultaneously, each aligned to the harmonic of frequency &. While this is a technical simplification of AGF, which activates single neuron per iteration, it allows us to analyze the collective behavior more directly. We additionally assume that once aligned, the neurons in A remain aligned under the gradient flow (see Appendix F for a discussion of possible "resonant" escape directions). We then analyze cost minimization for configuration @ = (u', v², w ) of aligned neurons with arbitrary amplitudes and phase shifts, and prove the following result.
Theorem 5.2. The loss function satisfies the lower bound ((OA) Z || (x, 1)2
[ /p. which is tight for N 2 6. When the bound is achieved, the network learns the function f(a- x, b-r;;; A) = (2)=(])/pp) (a+b) Xx,, where X [c] = cos (2wEc/p + Sw),, which leads the new util- ity function for the remaining dormant neurons: lt = (2/p²) Enep) |.e [a[{||][[C] [[ [[]T[[]
Put simply, the updated utility function after the first iteration of AGF has the same form as the old one, but with the dominant frequency Ç of x removed (together with its conjugate frequency - Therefore, at the second iteration of AGF, another group of neurons aligns with the harmonic of the second dominant frequency of ær. Lastly, we argue via orthogonality that the groups of neurons aligned with different harmonics optimize their loss functions independently, implying that each iteration of AGF, another group of neurons learns new Fourier feature of z.

---

AGF in action: greedy Fourier decomposition. Taken together, we have shown that two-layer quadratic network with hidden dimension H 2 3p trained to perform modular addition with a centered encoding vector x E RD sequentially learns frequencies Ç1, {2, ..... ordered by decreasing [+[]] After learning h 2 0 frequencies, the function, level, and next jump time are:
( (a-x, bxr E 2|2(3 (a+b) Xe, ((K) = =i>* [ (()) 2 +1) Ma 100 V
(9)
where p() = ma []6((+())| and learning rate 17 = a Simulating AGF numerically, we find this sequence closely approximates the gradient flow dynamics (see Figure
Extensions to other algebraic tasks. Our analysis of modular addition can naturally extend to a broader class of algebraic problems. First, one can consider modular addition over multiple summands, defined by the map y: (2/p)* + Z/p, (a:)== + dj + +... + ag (mod p). Using a higher-degree activation function (x)) x² the arguments for utility maximization should earry over, while the cost minimization step might be more subtle. Second, one can replace modular integers with an arbitrary finite group and study a network trained to learn the group multiplication map. Non-commutative groups introduce technical challenges due to the involvement of higher- dimensional unitary representations in their Fourier analysis. We leave a detailed analysis of these extensions to future work.
Discussion, Limitations, and Future Work
Loss
In this work introduced Alternttig Gradi- ent Flows (AGF), a framework modeling feature learning in two-layer neural networks as an al- ternating two-step process: maximizing a utility function over dormant neurons and minimizing a cost function over active ones. We showed how AGF converges to gradient flow in diag- onal linear networks, recovers prior saddle-to- saddle analyses in linear networks, and extends to quadratic networks trained on modular ad- dition. While these findings highlight AGF's utility as an ansatz for feature learning, it re- mains open whether its correspondence to gra- dient flow always holds in the vanishing ini- tialization limit. Proving such conjecture is theoretically challenging. Empirical validation also difficult because it requires taking both the initialization scale and learning rate to zero Moreover, this conjecture may simply fail to hold in more general settings. On natural data tasks, loss curves are often not visibly stepwise even at very small initialization scales (see Fig- ure 7), suggesting there may be limitations to AGF. That said, if many dormant neurons reach their activation thresholds in close succession, their nostmminiizztiion phases could interleave, causing the aggregate loss to appear smooth. Re- cent works reconciling the emergent capabilities of large language models with their neural scal-
0.08
Optinaal Linew: Predictor
0.06
0.04
Width 1 Width
Width -
0.02
Width
Width-
Width- -22
103
10
Steps
Figure 7: From a staircase to a slide in a two- layer ReLU network. Training loss for a two- layer ReLU network on a subset of CIFAR-10 under varying hidden widths from an extremely small initialization. Solid lines show gradient de- scent dynamics and dotted lines show the sequence produced by a numerical implementation of AGF from the same initialization. At small widths, gra- dient descent loss curves are stepwise and AGF tracks the jumps closely. As width increases, cost- minimization phases overlap and multiple dormant neurons activate in close succession, producing smoother slide-like loss curve; accordingly, the correspondence with AGF weakens. All experi- mental details are available at a Github code base.
ing laws have made similar suggestions [02) However, there are feature learning regimes in two-layer networks-such as those studied in Saad and Solla [46), Goldt et al. [47). Arnaboldi et al. [T display saddle-to-saddle behavior due to population-level transitions not captured by AGF, Connecting AGF to these multi-indes models and teacher stuuent settings (see Appendix [A for review) remains a key direction for future work. Lastly, the central limitation of the framework is its focus on two-layer networks, leaving open how it might generalize to deeper and more realistic architectures. Possible extensions include leveraging recent analyses of modularity in deep networks [103 105) and adapting insights from the early alignment dynamics of deep networks near the origin (59 106). All together, our results suggest that AGF offers promising step towards a deeper understanding of what features neural networks learn and how.
10

---

Acknowledgments and Disclosure of Funding
We thank Clémentine Dominé, Jim Halverson, Boris Hanin, Christopher Hillar, Alex Infanger, Arthur Jacot, Mason Kamb, David Klindt, Florent Krzakala, Zhiyuan Li, Sophia Sanborn, Nati Srebro, and Yedi Zhang for helpful discussions. Daniel thanks the Open Philanthropy AI Fellowship for support. Giovanni is partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. Surya and Feng are partially supported by NSF grant 1845166. Surya thanks the Simons Foundation, NTT Research, an NSF CAREER Award, and a Schmidt Science Polymath award for support. Nina is partially supported by NSF grant 2313150 and the NSF grant 240158. This work supported in part by the U.S. Army Research Laboratory and the U.S. Army Research Office under Contract No. W91
Author Contributions
Daniel, Nina, Giovanni, and James are primarily responsible for developing the AGF framework in Section 2 Daniel is primarily responsible for the analysis of diagonal and fully-connected linear networks Sections and 4 Feng is primarily responsible for the analysis of the attention-only linear transformer in Section Daniel and Giovanni are primarily responsible for the analysis of the modular addition task in Scction Dhruva is primarily responsible for an implementation of AGF used in the empirics. All authors contributed to the writing of the manuscript.
References
[1] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 5()):e00024-000 2020.
[2] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021.
[3] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, al. In-context learning and induction heads. arXiv preprint arXiv:2209. / 1895, 2022.
[4] Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey, Sparse autoencoders find highly interpretable features in language models. arXiv preprint arXiv:2309.0883 00000, 2023.
[5] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- erly. Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-oodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https:/trrassormer- circuits.
[6] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
Leonard Bereska and Efstratios Gavves. Mechanistic interpretability for ai safety-a review. arXiv preprint arXiv:2404.1408. 2024.
[8] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models, arXiv preprint XrXiv:2206. 07662. 2022.
[9] Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. In International Conference on Machine Learning, pages 6243 6267. PMLR, 2023.
[10] Lee Sharkey, Bilal Chughtai, Joshua Batson, Jack Lindsey, Jeff Wu, Lucius Bushnaq, Nicholas Goldowsky-Dill, Stefan Heimersheim, Alejandro Ortega, Joseph Bloom, et al. Open problems in mechanistic interpretability. arXiv preprint arXiv:2501.16496 2025.
11

---

[11] Devansh Arpit, Stanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, al. A closer look at memorization in deep networks. In International conference on machine learning, pages 233-242. PMLR, 2017.
[12] Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
[13] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21764, 2022.
[14] Alexander Atanasov, Alexandru Meterez, James B Simon, and Cengiz Pehlevan. The optimiza- tion landscape of sgd across the feature learning strength. arXiv preprint 2xiv:22410.06642, 2024.
[15] Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to- saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2 /06. 1,5933, 2021.
[16] Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11546 2019.
[17] Raphaël Berthier. Incremental learning in diagonal linear networks. Journal of Machine Learning Research, 24(171): 2023.
[18] Scott Pesme and Nicolas Flammarion. Saddle-to-saddl dynamics in diagonal linear networks. Advances in Neural Information Processing Systems, 36:7475-7505, 2023.
[19] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[20] Zhiyuan Li. Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, arXiv preprint a7Xiv:2012.09889, 2020.
[21] Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou. Understanding deflation process in over-parametrized tensor decomposition. Advances in Neural Information Processing Systems, 34:1299-1311, 2021.
[22] James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. In International Conference on Machine Learning, pages 11852-31874 PMLR, 2023.
[23] Dhruva Karkada, James B Simon, Yasaman Bahri, and Michael R DeWeese. Solvable dynamics of self-supervised word embeddings and the emergence of analogical reasoning. arXiv preprint arXiv:2502..m 09863, 2025.
[24] Mary Phuong and Christoph H Lampert. The inductive bias of relu networks on orthogonally separable data. In International Conference on Learning Representations, 2020.
[25] Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021.
[26] Mingze Wang and Chao Ma. Early stage convergence and global convergence of training mildly parameterized neural networks. Advances in Neural Information Processing Systems, 35:743 2022.
[27] Etienne Boursier, Loucas Pillaud-Viviss and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances Neural Information Processing Systems, 35:20105-2011 2022.
12

---

[28] Hancheng Min, René Vidal, and Enrique Mallada. Early neuron alignment in two-layer relu networks with small initialization. arXiv preprint arXiv:2307.12851, 2023.
[29] Margalit Glasgow. Sgd finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the xor problem. arXiv preprint arXiv:2309. 15111,, 2023.
[30] Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of relu networks. Advances in Neural Information Processing Systems, 36, 2024.
[31] Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transformers learn through gradual rank increase. Advances in Neural Information Processing Systems, 36:24519-24511, 2023.
[32] Yedi Zhang, Aaditya K Singh, Peter E Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention. arXiv preprint arXiv:2501.1666. 2025.
[33] Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems, 34:26989 2-229999-27002, 2021.
[34] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircess property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference Learning Theory, pages 4782-4881 PMLR. 2022.
[35] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: laap complexity and saddle-to-saddl dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pages 2552-2623, PMLR, 2023.
[36] Alberto Bietti, Joan Bruna, and Loucas Pillaud- Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprini urXiv:2310.1977) 2023.
[37] Berfin Simsek, Amire Bendjeddou, and Daniel Hsu. Learning gaussian multi-index mod- els with gradient flow: Time complexity and directional convergence. arXiv preprint DxXiv:24/1.08700 2024.
[38] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220 2024.
[39] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 2024.
[40] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018
[41] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[42] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal Applied Mathematics, 80(2):725-752, 2020.
[43] Grant Rotskoff and Eric Vanden- Trainability and accuracy of artificial neural net- works: An interacting particle system approach. Communications on Pure and Applied Mathematics, 1899:1889-1935 2022.
[44] David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52(4):4225, 1995.
[45] David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters, 74(21):4337, 1995.
13

---

[46] David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 1995.
[47] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
[48] Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in nigh-dimensional two-layer neural networks. Advances in Neural Information Processing Systems,
2022.
[49] Stefano Sarao Mannelli, Eric Vanden-Eijnden, and Lenka Zdeborová. Optimization and generalization of shallow neural networks with quadratic activation functions. Advances in Neural Information Processing Systems, 33:13445-13455, 2020.
[50] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with a teacher-student model. Advances in Neural Information Processing Systems, 34:18137-18151, 2021.
[51] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high- dimensional & mean-field dynamics to dimensionless odes: A unifying approach sgd in two-layers networks. In The Thirty Sixth Anmual Conference on Learning Theory, pages 1 119--1227, PMLR, 2023.
[52] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[53] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[54] Greg Yang and Edward Hu. Feature learning in infinite-width neural networks. arXiv preprint 1XXi:2001.14522, 2020.
[55] Yuri Bakhtin. Noisy heteroclinic networks. Probability theory and related fields, 150:1-42 42, 2011.
[56] Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint arXiv:1803.0838. 2018
[57] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2021.
[58] Akshay Kumar and Jarvis Haupt. Directional convergence saddles in two-homogeneous neural networks. arXiv preprint arXiv:2402.09222. 2024.
small initializations and
[59] Akshay Kumar and Jarvis Haupt. Early directional convergence in deep homogeneous neural networks for small initializations. arXiv preprint arXiv:2403.08122, 2024.
[60] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations of Computational Mathematics, pages 1-84, 2024.
[61] Daniel Kunin, Allan Raventós, Clémentine Dominé, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning, arXiv preprint 240.0240066 2024.
[62] Daniel Gissin, Shai Shalev-Shwarte and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv: /909. 12051, 2019.
[63] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673, PMLR, 2020.
14

---

[64] Scott Pesme, Loucas Pillaud- Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances Neural Information Processing Systems, 34:29218 8-29230, 2021.
[65] Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability, arXiv preprini arXiv:2302. 08982, 2023.
[66] Hristo Papazov, Scott Pesme, and Nicolas Flammarion. Leveraging continuous time to understand momentum when training diagonal linear networks. In International Conference on Artificial Intelligence and Statistics, pages 3556-3564. PMIR, 2024.
[67] Ryan J Tibshirani. The lasso problem and uniqueness. Electronic Journal of Statistics, 2013.
[68] Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58 --5358..11 1989.
[69] Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 1(04):1E-03, 1998.
[70] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics learning in deep linear neural networks. arXiv preprint arXiv:1312.6122 2013.
[71] Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv: /0374, 2018.
[72] Clémentine CJ Dominé, Nicolas Anguita, Alexandra M Proca, Lukas Braun, Daniel Kunin, Pedro AM Mediano, and Andrew M Saxe. From lazy rich: Exact learning dynamics in deep linear networks. arXiv preprint urXiv:240914400 2024.
[73] Javan Tahir, Surya Ganguli, and Grant M Rotskoff. learning in high-dimensional regression. arXiv preprint arXiv:2410.08194, 2024.
Features are fate: theory of transfer
[74] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34:23831-23845 2021.
[75] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[76] Simon S Du, Wei Hu, and Jason D Lec. Algorithmic regularization in learning deep homoge- neous models: Layers are automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
[77] Daniel Kunin, Javier Sagastuy- Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.047788, 2020.
[78] Alan Julian Izemman. Reduced-rank regression for the multivariate linear model. Journal of multivariate analysis, 5(2):248-25) 264. 1975.
[79] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[80] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? study of simple function classes. Advances in Neural Information Processing Systems, 15:30583-30599 2022.
[81] Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International conference machine learning, pages 19565-19594, PMLR, 2023.
15

---

[82] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in neural information processing systems, 1es14222-12246 2023.
[83] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
[84] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391, 2023.
[85] Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Zhang. Learning threshold neurons via the" edge of stability". arXiv preprint arXiv:2212.07469, 2022.
[86] Johannes Von Oswald. Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pages 35151-35174.
PMLR, 2023.
[87] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614-45650, 2023.
[88] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
[89] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
[90] Andrey Gromov.
Grokking modular arithmetic. arXiv preprint arXiv:2301.02677 2023.
[91]
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Processing Systems, 36, 2024.
[92] Tianyi Zhou, Deqing Fu, Vatsal Sharan, and Robin Jia. Pre-trained large language models use fourier features to compute addition. arXiv preprint arXiv:2406. 034454 2024.
[93] Subhash Kantammeni and Max Tegmark. Language models use trigonometry to do addition. arXiv preprint 2rXiv:2502.0087: 2025.
[94] Neil Mallinar, Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Emergence in non-neural models: grokking modular arithmetic via average gradient outer product. arXiv preprint arXiv:2407.20199, 2024.
[95] Giovanni Luca Marchetti, Christopher J Hillar, Danica Kragic, and Sophia Sanborn. Harmonies of learning: Universal fourier features emerge in invariant networks. In The Thirty Seventh Annual Conference on Learning Theory, pages 3775-3797. PMLR, 2024.
[96] Depen Morwani, Benjamin L Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham M Kakade. Feature emergence via margin maximization: case studies in algebraic tasks. In The Twelfth International Conference on Learning Representations, 2023.
[97] Yuandong Tian. Composing global optimizers to reasoning tasks via algebraic objects in neural nets. arXiy preprint arXiv: 2024.
[98] Yuki Yoshida and Masato Okada. Data-dependence of plateau phenomenon in learning with neural Systems, 32, 2019.
mechanical analysis. Advances in Neural Information Processing
[99] Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems, (6-28699-28722 2023.
16

---

[100] Yoonsoo Nam, Nayara Fonseca, Seok Hyeong Lee, Chris Mingard, and Ard A Louis. An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem. arXiv preprint arXiv:2404. 17563, 2024.
[101] Yunwei Ren, Eshaan Nichani, Denny Wu, and Jason D Lee. Emergence and scaling laws in sgd learning of shallow neural networks. arXiv preprint arXiv:2504. 190::22444.993,,, 2025.
[102] Gérard Ben Arous, Murat Erdogdu, N Mert Vural, and Denny Wu. Learning quadratic neural networks in high dimensions: Sgd dynamics and scaling laws. arXiv preprint arXiv:2508.. S.Ki::2508.03688. 2025.
[103] Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The neural race reduction: Dynamics of abstraction in gated networks. In International Conference Machine Learning, pages 19287-19309, PMLR, 2022.
[104) Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew M Saxe. Make haste slowly: A theory of emergent structured mixed selectivity in feature learning relu networks. arXiv preprint WxXi:255330618/ 2025.
[105] Emmanuel Ameisen, Jack I Lindsey, Adam Pearce, Wes Gurnee, Nicholas L. Turner, Brian Chen, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Mareus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, Andrew Persic, Zhenyi Qi, T. Ben Thompson, Sam Zimmerman, Kelley Rivoire, Thomas Conerly, Chris Olah, and Joshua Batson. Circuit tracing: Revealing computational graphs language models. Transformer Circuits Thread, 2025. URL https://transf ormer- circuits.......c.. pubppppppsssae aattiittiiinnnnno graphs heaphssmethhdd methods html
[106] Akshay Kumar and Jarvis Haupt. Towards understanding gradient flow dynamics of homoge- neous neural networks beyond the origin. arXiv preprint arXiv:2502. 2025.
[107] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203. 20t.:2003.03666, 2022.
[108] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evo- lution in wide neural networks. Advances in Neural Information Processing Systems, 35: 32240-32256, 2022.
[109] Benjamin Aubin, Antoine Maillard, Florent Krzakala, Nicolas Macris, Lenka Zdeborová, et al. The committee machine: Computational to statistical gaps in learning two-layers neural network. Advances in Neural Information Processing Systems, 31, 2018.
[110] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116(12)55451-446 2019.
[111] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
[112] Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Fundamental computational limits of weak learnability in digensiimnii multi-index models. arXiv preprint arXiv:2405.15448 2024.
[113] Leonardo Defilippis, Yatin Dandi, Pierre Mergny, Florent Krzakala, and Bruno Loureiro. Optimal spectral transitions in high-dimensional multi-index models. arXiv preprint arXiv:2502.02545.0 2025.
[114] Yatin Dandi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The computational advan- tage depth: Learning high-dimensional hierarchical functions with gradient descent. arXiv preprint arxiv:2502. 2025.
17

---

[115] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/8 2020.
[116] Libin Zhu, Chaoyue Liu, Adiyannaraynann Radhakrishnan, and Mikhail Belkin. Catapults in sgd: spikes in the training loss and their impact on generalization through feature learning. arXiv preprint arXiv:2306. 23..02500..4815, 2023.
[117] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
[118] Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402. 04980, 2024.
[119] Yatin Dandi, Luca Pesce, Hugo Cui, Florent Krzakala, Yue M Lu, and Bruno Loureiro. A random matrix theory perspective the spectrum of learned features and asymptotic generalization capabilities. arXiv preprint arXiv:2410.184. 2024.
[120] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, (giant) step time. Journal of Machine Learning Research, 2-(349):1-65, 2024.
[121] Dmitry Chistikov, Matthias Englert, and Ranko Lazic. Learning a neuron by a shallow relu Dynamics and implicit bias for correlated inputs. Advances in Neural Information Processing Systems, 36:23748-2341 23760, 2023.
network:
[122] Ioannis Bantzis, James B Simon, and Arthur Jacot. Saddle-to-sadde dynamics in deep relu networks: Low-rank bias in the first saddle escape. arXiv preprint arXiv:2505.21722, 2025.
[123] Uwe Helmke and John Moore. Optimization and dynamical systems. Springer Science Business Media, 2012.
[124] Chandler Davis and William Morton Kahan. The rotation of eigenvectors by perturbation. iii. SIAM Journal on Numerical Analysis, 7(1):1466, 1970.
18

---

A Additional Related Work
A.1 What is feature?
Here, we briefly review definitions of a feature used in the mechanistic interpretability and deep learning theory literature:
In mechanistic interpretability, terms such as feature, circuit, and motif are used to describe network components-suche as directions in activation space, subnetworks of weights, or recurring activation motifs that correspond to eeme-entterpretable concepts or functions within the model's s computation [TO]. While this literature has led to many insights into interpretability, these definitions generally lack precise mathematical formalization.
In deep learning theory, features and feature learning are defined in terms of the neural tangent kernel (NTK). For a network f(z 0). the NTK is given by the kernel K (z, x1:0) = (Vof 0), oo (x'; (a,;))); 0)), formed by the Jacobian feature map [of 0). Feature learn- ing-also referred to as rich learning- occurs when this kernel evolves during training. While this definition is mathematically precise, it offers limited interpretability.
In this work, we implicitly adopt the notion of features from the deep learning theory perspective, but consider learning settings often studied in mechanistic interpretability.
A.2 Analysis of feature learning in two-layer networks
Here, we discuss related thcoretical approaches to studying feature learning in two-layer networks. These approaches operate in distinct regimes-ment -ieissmiiiiiiiiiiiiiiii-------- field/infinite-widthh teacher-student at high- dimensional (thermodynamic) limits, and single-step analyses analyses-whereas whereas AGF focuses on the vanish- ing initialization limit for two-layer networks. Our aim is to position AGF complementary to, not a replacement for, these approaches.
Mean-field analysis. An early line of work in the study of feature learning analyzes the population dynamics of two-layer neural networks under the frean-field parameterization, yielding analytic characterizations of training dynamics in the infinite- -iith limit [40+43]. These ideas have been extended to deep networks via the tensor program framework, culminating in the maximal update parametrization (μPP) [54] [07). Related analyses have also been obtained through self-consistent dynamical mean-field theory [108]. Although feature learning occurs in this regime, these analyses primarily address how to preserve feature learning when changing network depth and width, rather than elucidating what specific features are learned or how they emerge.
Teacher-student setup. The two-layer student framework provides a precise setting for studying feature learning by analyzing how a student network trained on supervised examples generated from fixed teacher recovers the structure of the teacher. Foundational work by Saad and Solla 44 45 [66 provided closed-form dynamics for online gradient descent in soft committee machines trained on i.i.d. Gaussian inputs in the high-dimensional limit, deriving deterministic differential equations governing a set of order parameters. Goldt et al. 47), Veiga et al. 48 extended this dynamical theory to study the evolution of generalization error of one-pass SGD under an
arbitrary learning rate, and range of hidden layer widths. Further developments have analyzed conditions for the absence of spurious minima (49), extended the framework to capture more realistic data settings [50], and provided a unified perspective connecting to the infinite hidden width limit of the mean-field regime [S Much of this analysis use techniques from statistical physics-including the replica method and approximate message passing- to characterize the generalization error and models,
uncover sharp phase transitions and
gaps in
discussed further in Aubin et al. [109], Barbier et al. [110].
Multi-index models. Similar to the teacher-student setting, multi-index models provide a structured framework in which a neural network is trained on data whose labels depend on nonlinear function of low-dimensional projections of nigh-dimensiona inputs. Feature learning in this setting is governed by the network's ability to align with the relevant low-dimensional subspace. This setup often
gives rise to characteristic staircase-like training dynamics [111 36, 37 60 33] Prior analyses
have primarily focused on characterizing generalization and sample complexity, and recent work has established precise theoretical and computational thresholds for weak learnability in high-dimensional multi-index problems [12 [3]. Other studies have demonstrated that reusing batches in gradient
descent can help networks overcome information bottlenecks, accelerating the learning of relevant
19

---

Additionally, incorporating depth into these models has been shown to provide computational advantages for learning hierarchical multi-indes functions [114
projections 38 39].
Single gradient step. A recent line of work has focused on understanding feature learning in neural networks after just single gradient step. Empirically, it has been observed that with large initial learning rates, models can undergo "catapult" dynamics, where sudden spikes in the training loss accelerate feature learning and improve generalization [TS (16). Ba al. [117] initiated the precise analysis of feature learning after a single gradient step by showing, in high-dimensional single- index setting, how the first update to the first-layer weights can dramatically improve prediction performance beyond random features, depending on the learning rate scaling. This single-step analysis has been extended to characterize the asymptotic behavior of the generalization error [T18], establish an equivalence between the updated features and an isotropic spiked random feature model 119], and investigate the effects of multiple gradient steps [20] Collectively, these works highlight that even single gradient step can induce meaningful feature learning in two-layer networks.
Distinet phases of alignment and growth in ReLU networks. Maennel al. [56) first observed striking phenomenon in two-layer ReLU networks trained from small initializations: the first-layer weights concentrate along fixed directions determined by the training data, regardless of network width. Subsequent studies of piecewise linear networks have refined this effect by imposing structural
constraints on the training data, such as orthogonally separable [24 26 28 linearly separable and
symmetric (25,, pair-wise orthonormal [27]. correlated [12]], small angle [30], binary and sparse
29 Across these analyses, consistent observation is that the learning dynamics involve distinct
learning phases separated by timescales: fast alignment phase driven by directional movement and
slow fitting phase driven by radial movement [60] 61. Concurrently, Bantzis al. [22 use closely related directional-ratte
decomposition to identify the optimal escape directions from the saddle the origin in deep ReLU networks.
20

---

B Derivations in Alternating Gradient Flows Framework
In this section, we proivde details and motivations for the general framework of AGE. All code and ex- perimental details are available at githhu.. com/danielkune h http:///ww..aaaaaaoaaaaaaaaaaaaa -la.
B.1 Deriving dormant dynamics near saddle point in terms of the utility function At any point in parameter space the gradient of the MSE loss with respect to the parameters of a neuron can be expressed as
Vo, Ve, (ffi(a; {))]]
(10)
where r(z; 0) 0); neurons into two sets, an active A and dormant D set, then the gradient for a dormant neuron E D can be decomposed into two terms:
v(x) Et 0:) is residual that depends on all neurons. Given partition of
Ve, = Vo Bu (ACEE))))
- -
-Vo, Vo, 10 K (.) (11)
SEA
JED
The first term is the gradient of the utility Vold as defined in Equation (T). The second term only depends the dormant neurons and is thus negligible when all the dormant neurons are small in norm. Taken together, in the vicinity of a saddle point the loss defined by an active and dormant set, the dormant neurons are driven approximately by gradient flow maximizing their utility:
d A2 Vols
(12)
dt
This observation motivates the utility and our AGF ansatz for gradient flow. To formally prove that AGF converges to gradient flow in the limit of vanishing initialization one might consider using tools like Hartman-Grobmen theorem make this step rigorous.
Directional and radial expansion of the dynamics. Using this approximation. we can further decompose the dynamics of a dormant neuron near saddle point into a directional and radial component:
d 0,
P "ollll
d 10.11 17 "110.1"
1 (0,, Voll..
dt |0,11
!7
Voll
dt
(13)
where Py. = -I....
0:07 is projection matrix. When (-)) is a homogeneous function of degree |0:11
k. then the utility a homogeneous function of degree x = k + 1. Using Eulcr's homogencous function theorem, Equation (13) then simplifies to Equation 2). When o(-) is not homogencous, we
can Taylor expand the utility around the origin, such that the utility coincides approximately with a homogeneous function of degree K, where is the leading order of the Taylor expansion.
Dynamics of normalized utility. An interesting observation, that we will use in Appendix C is that the normalized utility function follows Riccati-like differential equation
L i = Tn m 110111
| 12
(14)
x²
In general, this ODE is coupled with the dynamics for both the direction and norm, except when for which the dependency the norm disappears.
B.2 Deriving the jump time To compute Ti we use Equation to obtain the time evolution of the norms of the dormant neurons. For i E DD we obtain:
(xx (0,(1) (110:00))2=* (t)
if = 2,
(15)
where we have defined the accumulated utility as the path integral S,) normalized utility. We find T; the earliest time at which neuron i satisfies [[0((T:)] |I > 1:
J0 (s), of the
log [[0(())]
if = 2,
= inf {t 01 Si(t) 0
a/n), where c=
(16)
[0)]
if 2.
2-A
Note that the expression for c is continuous at 2.
21

---

Lower bound on jump time. When applying this framework to analyze dynamics, computing the exact jump jump time can be challenging due the complexity of integrating U. In such cases, we resort on the following lower bound as useful analytical approximation. Let Li," be the maximal value of the utility, or at least an upper bound on it. Since S((1) I nil, we deduce:
log [0:(0)]
if 2,
20
(17)
if K > 2.
B.3 Active neurons can become dormant again
In the cost minimization phase of AGE, active neurons can become dormant. Here, we motivate and discuss this phenomenon. Intuitively, this happens when the GF trajectory brings an active neuron close the origin. Due the preserved quantities of GF for homogeneous activation functions, the trajectory of active neurons is constrained. Thus, this phenomenon can occur only in specific scenarios, as quantified by the following result.
Lemma B.1. Suppose that & evolves via the GF of C, and let e = (x- ())1, [|,;(0)]* (0)l.. The norm [0,(t)] satisfies the lower bound,
[10,(t)) max (-c, k ====-)
(18)
which holds with equality when either [[w,(t]] f c 2 or [|a;(t]] = 0 for < O.
Proof. For homogeneous activation functions of degree k. it is well known that the quantity c = (N [))la, wy (t)ll is preserved along trajectories of gradient flow. Since [t)
[aa((t]1 + (t)l we have:
10,(1)) -c K (()]
C
(19)
10,(1))2
(1)
K
The claim follows by combining the above inequalities.
Since vanishing initialization (k 1)]|a;112 ~ ||w;|1 we have e 0. Therefore, an active neuron can approach the origin during the cost minimization phase. If this happens, the neuron becomes dormant. When neuron becomes dormant again it does so with an accumulated utility S, & See Figure an example of a diagonal linear network where this behavior is observable.
B.4 Instantaneous alignment in the vanishing initialization limit
In this section, we revisit the separation of the dynamics dictated by revisit Equation 2 In particular, wish to discuss the conditions under which the directional dynamics dominates the radial one, resulting, at the vanishing initialization limit, in instantancous alignment during the utility maximization phase. To this end, we establish the following scaling symmetry with respect to the initialization scale factor a.
Theorem B.2. Suppose
= f((t) solves the initial value problem defined by Equation 2 with initial condition (0) 0,, Then for all a > 0, the scaled solution (,(t) = af R- 2t) solves the initial value problem with initial condition (,(0) ad0oi-
Proof. First, we consider the angular dynamics:
d 0,
d fi(aa-20
dl di [| 20]|
x-2 d fi(s)
ds [[((s]]
1=00
= =1laf((a*~21)) 2p -
22

---

Therefore, (t( afi(a"" (a00(aa-2 2t) satisfies the first identity in Equation ( above. Next, consider the norm dynamics: d 18.11 lafe(a"-20)
d
20)1
dt
dt
=a =aa-1_ aa-111laa [[fi(s)]]
d
ds
= =xllafi(a* uu (0,aaaa--0))
Hence, (t) afi(ta" also satisfies the second equation. Finally, verifying the initial conditions gives, 0,(0) - aa)(()) - alo.;- Therefore, 0,(1) = afa satisfies the given initial condition.
Theorem B.2 establishes transformation among the angular dynamics of different initialization fi(s)
scales. Specifically, for a given point (x, [[.(t in the solution trajectory the initial value problem with initial condition (0) o, the corresponding point in the initial value problem with scaled initial condition (0) ado,s is 2-a fi(e) This correspondence reveals that the angular
[
alignment process is effectively slowed down by factor of a= with initialization scale
Next, examine the alignment speed in accelerated time in the limit of - 0. In this regime, the asymptotics of Equation (16) is given by,
cx(a) log a
if K = 2, if K > 2,
(20)
Therefore, in the accelerated time, the alignment speed is effectively scaled by scaling factor of:
7a) = ca(a)
- log if 2, if
(21)
ax2
2,
And,
+o0
if 2, if
lim 7(a)
(22)
1
The limiting behavior implies that the alignment is instantaneous in accelerated time for almost all initialization f if and only if = 2.
B.5 A neuron-specific adaptive learning rate yields instantaneous alignment
As discussed above, instantaneous alignment of dormant neurons with their local dility-maximiizing
directions occurs only when = 2. For higher-order activations (K 2), the directional dynamics acquire a norm-dependent factor ||0,]| x|= that slows their angular evolution. Consequently, directional and radial dynamics no longer decouple, even in the vanishing initialization limit, and dormant neurons rotate gradually rather than aligning instantaneously.
This dependence can be removed by introducing neuron-specific adaptive learning rate
= ||0 2-*
(23)
where 17 is a global base rate. When x = 2, this scaling has no effect and all neurons evolve at the same rate. For x 2, however, neurons with small norm (1),, || < 1) are accelerated, while those with large norm ([0,1) > 1) are slowed down. This rescaling effectively reparametrizes time that the directional dynamics become norm-independent while the radial dynamics remain norm-dependent. Substituting this adaptive rate into Equation (2) yields an evolution equivalent to that of the 2 case, resulting in the decoupling between directional and radial dynamics in the wanishing-imitializatiioo initialization limit for all K. In practice, this scaling acts analogously to form of RMSProp or Adam-but derived directly from the analytical structure of the x-homogencous gradient flow.
neuron-wise adaptive polimization-rsentiie
23

---

C Complete Proofs for Diagonal Linear Networks
In this section, we derive the AGF dynamics for the two-layer diagonal linear network; see Section for the problem setup and notation. Throughout this section, assume the initialization (,) (V2a, 0). following the convention in Pesme and Flammarion [18].
C.1 Utility maximization
Lemma C.I. After the hthh iteration of AGF, the utility for the 8th neuron is
Ui (0:)r(t)) -1ys V3 (B((:)
(24)
which is maximized on the unit sphere ||0,1| = when sgn(u;u3) = -gn (v,,(((())) and
Jud = -lo] resulting in a maximal utility value of = 2 P (3(())
Proof. Substituting the residual (1)) = y ((K)) into the definition of the utility function:
u((0;;r(h)) = ueu;Ki5Ty = - u. V L (B(K)).
(25)
Observe that the expression for Us is linear in u,t; Therefore, under the normalization constraint ||0:|| = 1, the utility is maximized when |u,] = |vs] and sgn(u,z;) = sgn (V (B(+))) yielding
the maximal value U, 1/2 p (3(
Lemma C.2. At any time t during AGF, the parameters of the 7th neuron satisfy:
ux(t)² - (t))) 20,
2u,(t)eu(t) = + / (t) 4a
(26)
Proof. Both the utility and loss invariant under the transformation (u,, U/)) - (gu,, 9 g0)) for
any * 0. This continuous symmetry, together with the fact that each AGF step consists of gradient flow in of these functions, implies that the quantity ut)+ (t)² is conserved throughout AGE. Plugging the initialization into this expression gives the first identity (,))² - (22)) = 2o. See Kunin et al. (77) for general connection between continuous symmetries and conserved quantities in gradient flow. For the second identity, we solve for the intersection of the hyperbola
(t(2) (t) 2a and the circle [t)(( (t)² (1(1) which gives
(t) 2 2a
00 (+) 2a
u (t) +V
(t(t) V
(27)
2
2
Multiplying these expressions together gives the identity for the product.
Lemma C.3. After the tth iteration of AGF, the normalized utility for the ith neuron is driven by Riccati ODE L = 41 ((z))) -(?) with the unique solution,
[ (T(11)
(k) t) = 27 tanh (5(h) + 4mallit). where g5 tanh- -1
(28)
u
Proof. We begin by observing the gradient of the utility function takes the form
Voll, (0i:((t) - (3(*)
(29)
Evaluating this quantity the normalized parameters o,, and recalling the expression for the maximal utility from Lemma C.T we obtain,
(+))) (2t:)²
(30)
Substituting this into the normalized utility dynamics derived previously (see Equation (14))
we
obtain the Riccati equation. This is a standard ODE with known solution, where the constant is determined by the initial condition ú. (7())..
24

---

Theorem C.4. After the kth iteration of AGF, the accumulated utility for the 8th neuron is
Si (8) +t)
1
log cosh 2M 2 +
Si cosh cxsp-'exp (2naS (+(+)))
exp
(31)
27aa
21/a
where f/ (2)) the learning rate and Ç: = sgn (-V (3(k))) ps (T(1)). The quantity Ps (T(+)) = sgn (us (T())) us (T(K)) is determined by the recursive formula
Pi (7()) + 1) = sgn
Pi (T(+))
cosh eex (2 (x(t)) - (3(+)))
(32)
27a
where Sz(0) 0 and px) 0.
Proof. Using the expression for the normalized utility derived in Lemma C.3 expression for the integral of the normalized utility, i.e,, the accumulated utility:
derive exact
cosh Anali 146,
S (7(8) + = S. (T + So (s) S (7(+)) +
log
(33)
21aa
cosh
Using the hyperbolic identity tanh - (x) = sgn(x) cosh -1 ( (v-=a)) we can express the constant introduced in Lemma
C.3
(1)) )es())()(
6 (k)
V L (sthi) lanh-
sgn
110 []0)(1(8)]11 [10(
Lemma CA
(34)
( L (sthi)) P: (7())
= sgn
cosh
Lemma
(35)
202
(-V v s. L. (s(k))
P: (7(()) cosh -1
cxp (2 Si ()(
sgn
(36)
where in the last equality we used the simplification [+0() )| = 2a exp (2n.. S; (7())). Substi-
tuting this expression into Equation (33). notice that the denominator inside the logarithm simplifies
substantially, as the (x)) term cancels out, yielding Equation (1p. Finally, by Lemma C.2
P
changes sign only when S = 0, yielding Equation 32
Corollary C.5. After the thth iteration of AGF, the next dormant neuron activate is i* arg miggee At, where
A cosh - exp (21)aa - Si cosh -t exp (r(
(37)
27a -2
G = sgm (- (-V V, (B(+)) Pe (t)),,, and the next jump time y (k++)) = (k) Arj..
Proof. From Theorem C.4 we solve for the time I such that Si(T(1) + () = 1 by inverting the expression for accumulated utility. This time is Ar;-
C.2 Cost minimization
Active neurons represents non-zero regression coefficients of During the cost minimization step, the active neurons work to collectively minimize the loss. When neuron activates it does so with certain sign (g,),
sgn( ((3x). If during the cost minimization phase neuron changes sign, then it can do so only by returning to dormancy first. This is due to the fact that throughout the gradient flow dynamics the quantity u² - o² = 2a is conserved and thus order to flip the sign of the product u,ti the parameters must return to their initialization. As result, the critical point reached during the cost minimization step the unique solution the constrained optimization problem,
33 =0
i
8 arg min [()) subject BERd
(38)
3, sgn((cus) 2 0 if if
where uniqueness follows from the general position assumption. All coordinates where (8%); = dormant at the next step of AGE.
25

---

62 log %g)
68 68
68
00
207
(a) Loss ((t)
(b) Function 3(t)
(c) Parameters ((t)
(d) Utility S(t) and 0(t)
Figure 8: AGF = GF as a - O in diagonal linear networks. We consider a diagonal linear network = u O u E R2, trained by gradient flow from initialization u = 22a1, v = O. with varying a, inspired by Figure of Pesme and Flammarion [8]. This example is special as it demonstrates how an active neuron, in this case 3(1,, can return to dormancy during the cost minimization phase, as discussed theoretically in Appendix B In (a) we plot the loss over accelerated time and (b) the loss landscape over 3. For small a, trajectories evolve from g(0) to 3(3), passing near intermediate saddle points g(1) and g(2), Each saddle point is associated with plateau the loss. As - 0. gradient flow spends all its time at these saddles, jumping instantaneously between them 7(), 7(2) ,7(3), matching the stepwise loss drops. In (c) we plot trajectories in parameter space and in (d) the same dynamics visualized in terms of their accumulated utility and angle. The jump times between critical points correspond to when the accumulated utility satisfies Si(r) 1. When the first coordinate returns to dormancy, the accumulated utility first touches zero, causing the angle to flip sign, before reactivating with the opposite sign.
C.3 AGF in action
Combining the results of utility maximization (Appendix C.I and cost minimization (Appendix C.2 yields an explicit recursive expression for the sequence generated by AGF. For each neuron only need to track the accumulated utility S;(() and the sign (4), which at initialization are both zero. We can now consider the relationship between the sequence produced by AGF in the vanishing initialization limit a - 0 and the sequence for gradient flow in the same initialization limit introduced by Pesme and Flammarion [8]
In Pesme and Flammarion 18 they prove that in the vanishing initializa- tion limit a + O, the trajectory of gradient flow, under accelerated time (1) = log(a) 1, converges to- wards a piecewise constant limiting process corresponding to the sequence of saddles produced by Algorithm 2 This algorithm can be split into two steps, which match exactly with the two steps of AGF in the vanishing ini- tialization limit.
Algorithm 2: Pesme and Flammarion
18
Initialize: t 4- 0, B + 0 E R + 0 Rd while VZ( # do
D j E d
VL(3), # 0}
T -' inf (T. 2 S. V =1}
L L
S S - (VL((3)
A20 18. =1,
8 = arg min L(3) where 3 E
BEK"
A-0 18. (-1.1)
return Sequence of (8, t);
We adapt the original notation to fit our framework, highlighting utility maximization (blue) and cost minimization (red) steps.
As in AGF, at each iteration of their algorithm a coefficient that zero becomes non-zero. The new coefficient comes from the set (i E [d] (B); * 0}. which is equivalent to the set of dormant neurons with non-zero utility, as all active neurons will be in equilibrium from the previous cost minimization step. The index for the new active coefficient is determined by the following expression,
i, arg min 73 0: 3i D, T.VL(()) =1}.
(39)
Because T; > 0 and 8; E (-1, for all coefficients E DD then this expression only makes sense if we choose the boundary associated with sgn (3)::. Taking this into account, we get the following simplified expression for the jump times proposed by Pesme and Flammarion [8]
+ sgn (L(())))
(40)
V C(3))
This expression is equivalent to the expression given in Corollary C.5 in the limit 7a oc. To this, we use the asymptotic identity - - + log as - cc. Combined with the
26

---

simplification 2 V
(B)). this allows us to simplify the expression for the jump time and show that P:d; 8; where s; is the integral from Algorithm
The second step of Algorithm 2 is that the new 8 determined by the following constrained minimization problem,
8; =1 1820 0,
8* = arg min (B)) subject to BERA
8: S 0,
if = -1 8: = 0, if (-1,1)
(41)
Using the correspondence A; sgn(s;), and noting that neurons with 8; (-1, correspond to the dormant set, this constrained optimization problem is equivalent to the cost minimization step of AGF, presented in Appendix C.2
27

---

D Complete Proofs for Fully Connected Linear Networks
In this section, we provide the details behind the results around linear networks in Section
We begin by clarifying the notion of 'neuron' in this context. As briefly explained in Section due to the symmetry the parametrization of fully-connected linear networks, the usual notion of a neuron does not define a canonical decomposition of f into rank-1 maps. Instead, the singular value decomposition of the linear map AW computed by the network, defines such a canonical decomposition, being the unique orthogonal one. Therefore, we will think of neurons as basis vectors (ā;, w)s E IRo+d under orthogonality constraints. These vectors align with the singular vectors and partitioned into dormant and active sets, D and A, based on whether their corresponding singular value is (1).
D.I Utility maximization
Differently from the usual setting of AGF, the orthogonality constraint on the dormant basis vectors implies that the utility function is not decoupled. Instead, it is maximized over the space of ['] orthonormal basis vectors, i.e. the Stiefel manifold.
Lemma D.I. After the tth iteration of AGF, the utility function of the 8th dormant basis vector with parameters 0g (āz; 102) is:
Ui (0 = -a.VV - (8%) Wi
(42)
The total utility EiED Ui (0g: is maximized on the Stiefel manifold when {āz, w.) coincides with the set the top [D] = H-kleft and right singular vectors V ())) resulting maximal utility value for the ith dormant basis vector Uf = =o? /2, where o; (k) is the corresponding singular value. Moreover, U4 has no other local maxima.
Proof. After substituting (kk yy- 3(4b thing else follows from the standard theory of Reyleigh quotients over Stiefel manifolds 123
x. the computation of the utility is straightforward. Every-
This means that the gradient flow of the utility function aligns the dormant basis vectors with the singular vectors of (3)
D.2
Cost minimization
As discussed in Section 4 the cost minimization phase is governed by the well-known Eckart- Young theory of reduced-rank regression 78). According to the latter, during cost minimization, the active basis vectors converge
Eye
(43)
where Pu, is the projection onto the top k cigenvectors of Ege - ET Plugging this solution into the expression for the gradient, we find that the utility at the next iteration will be computed with the projection Ug Myyi
V (L(())) = Eyz -B(h) (Ie (L-P Pyx
(44)
D.3 AGF in action
Putting together the utility maximiza- tion and cost minimization steps, AGF progressively learns the projected OLS solution (Equation (43)). coin-
Algorithm 3: Greedy Low-Rank Learning Li et al. [20]
Initialize: r + 0, Wa + 0 E dgdd Uu(oo) E Rdxo while X1 (VW.)) do
ciding with the GLRL algorithm by Li et al. 200 shown in Algorithm 3 with notation from the original work. At each stage, the GLRL algorithm selects top eigenvector (blue: utility maximization) and performs gradient descent in the span of selected direc- tions (red: cost minimization).
+ p-r- + Ur + unit top eigenvector of - L(W
U-(0) + (x>) VEur] E Rax=;
0, 1,. do - U,(t) ()L((.(()):
L 1,
W + U,(OOU,7 (Do)
return Sequence W,
28

---

We now discuss jump times, and motivate Conjecture 4.1 Recall from Section th4 in this setting, because K 2, in the limit of vanishing initialization, the alignment in the utility maximization phase occurs instantaneously. As consequence, the individual utility functions remain, essentially, constantly equal to their corresponding maximal value throughout each utility maximization phase. From the definition of accumulated utility with 7a = -log(a), we immediately deduce the recursive relation:
S (T()) +1) = = (-( + oi t.
(1)
(45)
Pú.., Eye have different singular values and vectors, the relation between the values of Equation (45) as E 'D varies unclear, and is hard determine the first dormant basis vector to accumulate utility value of 1. Yet, suppose that the dormant basis vectors align in such a way that the ordering of the corresponding singular values is preserved across the utility maximization phases. In this case, the ordering of the cumulated utilities in Equation (45) is also preserved for all iterations k and all times l. In particular, it follows
However, since - V = Pú Eye and - VVL((
A
by induction that the basis vector with index i, corresponding to the largest eigenvalue di jumps after a time of:
(k)
1 (1- S.. (T(h=1))
(46)
Once unrolled, the above recursion is equivalent the statement on jump times in Conjecture 4.1
When Egæ commutes with 2T Nyp the left singular vectors of Eye simultancously diagonalize V and Dy ST.. In this case, the singular values of P/, Us Eg correspond to the bottom ones of Eyz. with the same associated singular vectors. Therefore, the dormant basis vectors are correctly aligned already from the first iteration of AGF. By the discussion above, this confirms Conjecture
4.1 in this specific scenario. Moreover, the recursive expression from Equation (46) reduces to
(r 1/or- 1/oR-1, where on is the gth largest singular value of Eyz. Therefore, the jump times are simply:
1
(47)
Ok
These values, together with the loss values in Conjecture coincide with the ones derived by Gidel et al. [19] for GF at vanishing initialization. This means that under the commutativity assumption, AGF and GF converge to the same limit, similarly to the setting of diagonal linear networks.
When Egx and ET Ege do not commute, in order to prove Conjecture is necessary to understand the relation between the SVD of P Σyz and of Eyn. as k varies. The Poincaré separation theorem describes this relation for the singular values, stating that they are interlaced, i.e., they alternate between each other when ordered. This is not sufficient, since the geometry of the singular vectors plays a crucial role in the alignment phase. Even though there exist classical results from linear algebra in this direction [124], we believe that Conjecture 4.1 remains open.
29

---

E Complete Proofs for Attention-only Linear Transformer
Here, we provide additional details and derivations for the results around transformers (see Section for the problem setup and notation). This analysis is based on the setting presented by Zhang et al. 32. to which we refer the reader for further details.
Note that this scenario does not exactly the formalism of Section 2 since we are considering a neural architecture different from fully-connected network. Yet, due our assumption on the rank of the attention heads, each head behaves as a bottleneck, resembling a 'cubical version' of a neuron. Therefore, will interpret attention heads as 'neurons', and apply AGF this context (with 3).
We will show in the following sections that at the end of the hthh iteration of AGF, the 7t attention head, h = 1,.... k, learns the function:
Jn [a(X;
Un o
(48)
trE (N 1]
where Å is the hth largest eigenvalue of Ega, and Up is the corresponding eigenvector. We will proceed by induction on k.
E.1 Utility maximization
We prove the following lemma, establishing, inductively, the utility function for the h iteration of AGF.
Lemma E.1. At the Eth iteration of AGF, assume that the (k - 1) active attention heads have the function form as Equation 448 Then, the utility function of a dormant head with parameters
8, (Vi, Qi, K2) is given by:
Uc (D)r(()) NV. Q
ExRonh Ki
(49)
Moreover, the utility maximized over normalized parameters by (V;" (T,Qi,K) Lt1//3,=14 /3///33 /3), where the sign determined such that U, = N3 Moreover, Uc has other local maxima.
Proof. By assumption, we have:
Us (0gr(t)) ee
UN+1 Σ fx
(50)
The first term, which coincides with the initial utility, can be computed as:
Ue ():: = E (X, (X,,)))4L.NH1
N
E 2
TTN+ V
Σ
Un xn K,
(51)
Vi Q E 11N N+1
B3 B3
Ey [ [T] K;
NV, Q.E_2KK. K,
Denote AG = (trE, (N + 1)Ax) Then, the summand can computed as:
E EN.1. [X, 0h) (1. (X,, 0.)
Ai V.l.a...........
33T
Fn J,KKKK x
Σ xn
(52)
A,V E.Q.E u J
Σ Inat
V,r (KQ.A (NE + N(N + 1) 1)E22))
-VVQQ X3. Un K,,
30

---

where in the second equality, we have used the fact that
N
27
N (En 2 +2
E
Σ
In
= E
Σ
(53)
n=
1SijjSN
-Egg + N(N + 1)
This provides the desired expression for the utility, which corresponds to a Rayleigh quotient. Since the largest eigenvalue of (E Eh=1 ck-1 (fono%) is X the rest of the statement follows from the standard theory of Rayleigh quotients.
E.2
Cost minimization
After utility maximization, new head aligns to the corresponding eigenvector and becomes active. This implies that at the beginning of the cost minimization phase of the 1th iteration of AGF, the active attention heads have parameters On (Vn, Qn, Ks), such that Qs and K₂ coincide up to multiplicative scalar to Us for all S h S Moreover, Equation (48) holds for S < k.
We wish to show that during the cost minimization phase, the newly-activated head with parameters 0g learns function in the same form. First, we prove that the loss function is decoupled for each active attention head.
Lemma E.2. Assume that the k active attention heads have the same function form, up to multiplica- rive scalar, as in Equation (48). Then, the loss over active heads decoupled into sum of individual losses, that is
((OA) = c(
c(.)).
(54)
Moreover, for h = k,...,k; k; the directional derivatives JQ, tional to Uh throughout the cost minimization phase.
ok, BC ((OA) and 8C ((0 are propor-
Proof. The first statement follows from a straightforward calculation using the orthogonality of eigenvectors. As result,
ac (0 iQn
Oc
(0x) ENI. aQn h
Vn
Un yn
Limo Qn
(55)
From this expression, that Qn starts eigenvector of Exz. then subsequent gradient updates will be the same direction. As result, the direction of Q₁ remains unchanged in cost minimization. similar argument holds for K₁, as well.
Since the active heads follow the GF of C, the decoupling the losses implies that during the cost minimization phase, each head actually follows the gradient of the loss of its own parameters, disregarding the other heads. In particular, 6, remains at equilibrium for h = 1....,k its loss has been optimized during the previous iterations of AGF. By Lemmas E. and E.2 the newly-activated head remains aligned to its eigenvector U). learning a function of the form
1, since
N
fa (X; D-IN41 = A Σ T
(56)
where Aş E IR is parameter that is optimized in cost minimization. L(Ox) is convex with respect to Ag, and the corresponding minimization problem easily seen to solved by A₁ 1/(tr (N + 1)..). This concludes our inductive argument.
31

---

E.3 AGF in action
As we have seen above, AGF describes a recursive procedure, where each head sequentially learns principal component. We can easily compute the loss value at the end of each iteration:
g(k) Egg N41 2
UN+1
Σ
A
On Th Toe
h=1
(57)
1 2
NX
A
tr + (N 1)Aa
Furthermore, as an immediate consequence of Lemma accumulated utility the kt iteration of AGF is upper-hounded by S,( S 312, = INXx//V The jump time given by the first head that reaches S,(7(() --(0) 1/ (1]])( which yields
can lower bound the jump times. The
v3
(k) 7(1) 2
(58)
N
where μi max; [r0(] and (())
32

---

Complete Proofs for Generalized Modular Addition
In this section, we provide the proofs for the results summarized in Section Instead of reasoning inductively in Section E for simplicity of explanation will derive in detail the steps of the first iteration of AGF. As we will explain in Section F.3] all the derivations can be straightforwardly extended to the successive iterations, completing the inductive argument.
F.1 Utility maximization
In order to compute the utility for a single neuron, we will heavily rely on the Discrete Fourier Transform (DFT) and its properties. To this end, recall that for a E R., its DFT à R2 is defined as:
a[k] =
ulale 2xtka/p
(59)
a=0
where = - is the imaginary unit. We start by proving a technical result.
Lemma F.1. Let x, u, RP. Then:
p-1
(u, a. (v. s)w, (a + bb b).
a [ 2 [ [[]
(60)
a,b=0
Proof. Notice that the inner product (u, z) can be expressed as
p (u, a Σ
xc(ku[kk + a] (x u)(a],
(61)
where x u the cross-correlation between x and u. Thus, the left-hand side of Equation (60) can be rewritten as:
u)(a] (x D) (xr * w=) + bj (
a,b=0
By using Plancharel's theorem (xr, v) (0, ú) and the cross-correlation property F x [k]] the above expression equals
Σ
(62)
which corresponds to the desired result via the simplification [[A]&(k] [+])
Lemma F.2. Under the modular addition task with embedding vector x and mean centered labels, the initial utility function for single neuron parameterized by (u, w) can be expressed as
u(0:y)
663)
Proof. By definition, the utility is:
u(::v)
1
Σ
(uu,a- + (v,b. x))x (a + b) (x, 12,11/
2 a,b=0)
p
(w, 1) (x 1)
(fu, a x)²² 2 (0.b. x)² + 2(u. a x) ((.bb.2)) 2)))
(w,, (a + x))
2 a.b=0
P
(64)
33

---

Due to the cyclic structure of (w, (a + b)-x), the contributions from the terms (u, a-x)² and (v, b x) terms vanish. This reduces the utility to
p
(u, a (w, b- a) (we, (a + b) x)
Σ (u, a x)
Σ (v b x))
Σ Ww c x))
p
pd
o,b=0
(65)
The first summand in the above expression is provided by Lemma Moreover, since the mean of vector corresponds to the zero-free frequency component &t = 0 of its DFT, the second summand reduces
to:
(u, x))
Σ (v, xx) 1=0
Σ (w, c..)) x)
2
x.
I *
0
w [0)
p
p
(66)
2
1 |0 a
p
Since the latter coincides with the term & 0 in Equation (61 we obtain the desired expression.
We now solve the constrained optimization problem involved in the utility maximization step of our framework.
Theorem F.3. Let & be a frequency that maximizes [a[k]], & = I,..., - L, and denote by Sy the phase of @(). Then the unit vectors O. (u,, Uo W.) that maximize the initial utility functiio ((;; v) take the form
.
E 2 a +
U. [a]
cos
Su
3p
p
2
C
[b
COS
2m b + 80
(67)
3p
p E
w. [c]
2
cos
2n
-
8
3p
p
where a, b, (0.....p-1 1} are indices and Sus Sus Sw are phase shifts satisfying s. + 8, = sw + Sg (mod 2T). They achieve a maximal value of = (22 Moreover, the utility function has local maxima other than the ones described above.
Proof. By Plancharel's theorem, the constraint ||012 ||u|| | + |/ is equivalent to constraint on the squared norm of its DFT, up to scaling. Together with Lemma F2 this allows us to reformulate the original optimization problem in the frequency domain as:
2 p-1
maximize
[
[
(68)
subject to ||0||
|| 2|
To solve this optimization problem, consider the magnitudes Mas Mus Hes Pur R2 and phases Sn, Su, S,, S., E [0, 2x)) of the DFT coefficient æ, u, v, w. respectively. This means that [(k]
Ie [k] exp(iiS (k]) for every k and similarly for u, v,
Since v, and p is odd, the periodicity of the DFT (i.e., a[k] = alk + p)) allows us to simplify the objective by only considering the first half of the frequencies. Substituting the magnitude and phase expressions obtain the equivalent optimization,
real-valued, their DFTs satisfy conjugate symmetry (i.e., a[- = a(k)). Since
4
maximize
Me [c] Mm (k)) cos + S. S.(Mk (k))
(69)
subject
Mu [k]²
Σ k=U
Mulaja
p
2
k=
34

---

The phase terms can be chosen independent of the constraints to maximize the cosine term. Since the only local maximum of the cosine is (mod 2T)), the only local optimum of that term is achieved by setting [k + S.[k] Sw. S. k] = (mod 2x).. But then the optimization problem reduces to maximizing (pp 12/2| Mrx [X:])/ you [R)
subject to the constraints. The only local maximum of this problem easily seen to be achieved by concentrating all the magnitude at the dominant frequency of (and its conjugate), i.e.:
{ IVE 0
i = (mod p)
[u[k] = Pulki [k] Mwl
(70)
otherwise.
This gives a maximal objective value of lt" = (4/p" () 22 By applying the inverse DFT with these choices for magnitudes and phases, while accounting for conjugate symmetry, yields the desired result. Note that in the statement, Sz. Sus Su, Sw denote (4] [4), Sw respectively.
We highlight that Theorem F3 and its proof very similar to Theorem and its proof in Morwani et al. 196]. They derived this optimization problem by looking for a maximum margin solution, suggesting there may an interesting connection between utility maximization and maximum margin biases of gradient descent.
F.2 Cost minimization
We discuss cost minimization. As demonstrated in the previous section, during the utility maximization step, the parameters of each neuron align with the harmonic of the dominant frequency & of x, albeit with phase shifts s s'g. The goal of this section two-fold. First, in Appendix F22. we discuss our assumption that during the cost minimization phase, the neurons remain aligned with the same harmonic, possibly varying the amplitudes and phase shifts of their parameters. Second, in Appendix F.2.2 we solve the cost minimization problem over such aligned neurons. Therefore, throughout the section, we consider N neurons parametrized by θ (u², v², w w') where:
2
C
w' ua[a]
A
cos
2m
s
3p
p
2
2n E
of [b]
1
cos
+
(71)
3p
E
wf [c] = A
COs
2n - c +
3p
p
for some amplitudes A A A E R2o and some phase shifts To begin with, note that loss splits as:
8 syg s E IR.
1
Σ
Σ
fila x, b. b,b-,;0000 x;);) (a + b) x (z. 1)
(())
2T²
P
0.h=0
i=1
(72)
1 |100|2 (x,1)2
= C(0) (()))) +
where
N
D-1
((θ)
= 2p
w
Σ (1 (w))))))) (uw2 a x)) + (...b.)))
(73)
++=1
a.b=0
and u(:: v)
ENN ((:: y) is the cumulated utility function of the N neurons. From Lemma F22
we know that:
2
u(O;;)) = =\ 27p
[ cos(s + ss - 8z).
(74)
35

---

Throughout this section, we repeatedly use the following identity: for any N, & E Z, and s E R.
2nka if 0 (mod p),
p
ccssss),
Σ
cos
(75)
p
0,
otherwise.
a=0
This follows by writing cos Re(e and noting that the geometric sum Ex=0 e 2 vanishes unless k = 0 (mod p).
F.2.1 Preservation of harmonic alignment
To analyze cost minimization, we restrict attention to a regime in which the newly activated neurons do not change their aligned harmonic during this phase.
Assumption F.4. During cost minimization, the N newly activated neurons remain aligned to the harmonic ç.
This restriction allows to solve cost minimization within the subspace spanned by these N neurons, in close analogy with prior sections. However, first we characterize the conditions under which this assumption is valid and clarify in what sense it provides a faithful description of the dynamics.
Theorem F.5. Let h E IR be a vector in the form:
Ma] A cos
2x
(76)
P
for some A,8 E ERR and + 0, +6. for all a, [p)
N
N
Σ
w (u' x)²= Σ
w bu ==)=)=00 =0,
(77)
t=1
1=1
then for all i = I,...,N:: N:
0C h, Ou (0)
(
0c Ou
Oc Ow (0) 0.
h,
(0)
h,
(78)
Proof. A direct calculation leads to the following expressions for the derivatives:
Oc
p
h
(0)
4
(h, a 2) (1u', a 2) + (v', b- x)) (w¹, f(a x, x;) (a b) x).
Ou
h Oc
(0)
(h. b x) ((u' a a) + (wt, b.),b.x)) x)) (w', sla xr x; (a b) .x),
a,b=0
Oc h, Ow (0)
2
2
(lu' a x) + (0* b.),)=))) (h, fla x, z; 2; (a b))
a,b=0
(79)
Via a tedious goniometric computation, the three expressions above can be expanded into a sum of cosine terms. Since E & 5', each of these terms depends on or b, and therefore averages out. We do not report the details of the computation here, but refer to the proof of Lemma F.6 for more details.
The condition Equation 777 rules out square terms that could otherwise generate gradient components at resonant frequencies 5 * =G. If the squared activations (u', =z)² or (v', b 2)² have nonzero aggregate contribution, their trigonometric expansion produces cosine terms at doubled or mixed frequencies that may survive the averaging over b via the cosine-sum identity Equation (75). In this case, the gradient can acquire components outside the span of the aligned harmonic, allowing neurons to escape alignment.
However, pairs of aligned neurons can easily satisfy Equation 777) by coordinating their amplitudes and phases. For example, using the identity 22 = [ + 22 [ - (2))-22) neurons can arrange for square terms to cancel in aggregate while preserving their linear contributions. Such cancellations correspond to directional adjustments that can occur rapidly during cost minimization.
36

---

F.2.2 Cost minimization over aligned neurons Next, in order solve the cost minimization problem aligned neurons, explicitly compute the term (()) in the loss function.
Lemma F.6. We have:
| (())
Σ
A Al, cos(s - sid) -l)) ((AA)++(A)))
+ (A)))) ((A2)²
(AZ)2)
54p²
i,j=1
AM))
cos(2(st, - xa)) + (Ai,
(80)
2 cos(2(s)) -s())
2
+ +2A AtAAAA (cos(s's + 2/ 5²) + cos(s's st t/e + +()))
Proof. This follows from rather tedious computation, which we summarize here. First, note that
2A A
p-1
cos s/4 2TT 20) cos 8/ -l-+255c)
2x A cos(s sw) (81)
(w', w²)
3p
p
p
Next, the shift-equivariane property of the Fourier transform and Plancharel's theorem imply that for all and we have:
(u' aa,a.a0 x) (ie, (2.
2 V 3p
E[E] AM cos s 2T
(82)
and similarly for vi We can plug the above expression into the quadratic term ((u', a- x) (v', b x))² from Equation 773 and expand it via the goniometric identity:
(A cos(a)) (3)) 7 (A2+B3 + B² + A2 cos((2a) + B² A AB(cco(( + B))
(83)
We similarly expand the term ((u), ax) + (2 b 2;) By leveraging on the goniometric identity 2 co) c3s(3) cas(a 3) + cos(a + B), the product of these quadratic terms for the 7th and jth neurons expands into 16 unique terms. Since is odd, 25, 4 * O (mod p), and thus, by Equation (75), only the terms independent from both and do not average out. This leaves four terms, providing the desired result.
We now provide lower bound for the (meaningful terms of the) loss function.
Theorem F.7. We have the following lower bound:
[[S]
C(θ) -((;;))))
(84)
P
Moreover, equality holds if, and only if, we have that EN C2 cos((;)= = - Cssi((a)) = 0 f.ror any choice of (Ci, ax) among
+ (/))
(At, + 2s)))
(AL() 28.).
(85)
(A' 5 (s) {.)), (A m + s + )
and, moreover, Σi= ASA,, sin s + 8x 5 5 and En AgAAA, cos(s + sw s/ -t)) /54pp/
Proof. Given amplitudes Ciy...,
Iap...,C,, ON and angles a,,..., aN. consider the goniometric identity
2 + Σ C Cissinaa)
C.C,cossf a,)
Σ
C cos((a))
(86)
37

---

By expanding each product of two cosines in Equation (80) into a sum of two cosines, we can apply the above identity to each summand by choosing a; { S 2s ; 2s, s (s s ) s sy + sy. s 8x sy s Since both the summands in the right-hand side of Equation (86) are positive, we conclude that:
C( 2
Σ
A', cos(s", + Sx 8 st))
(87)
54p²
with equality holding only if Equation 885 is satisfied. Therefore, using Equation 74
we deduce:
N
(()) - (u())) 2
Σ
A 4 cos(s + - 5//)
54p²
2 V
|
Σ
A A cos(s + -8" - =%)
27ps
i=1
(88)
|
N
|+[[]|
2
L [
A A cos(s", + - 5.
-8)
v54p
vp
i=1
|
p
with equality only when (I=tE))/(V54p)) which concludes the proof.
N ALA cos(s + 8z - 5 -%) = - Vp.
0
For N 2 6, the lower bound from Equation
(84 is tight. For even N, minimizers can be constructed
by setting sny + 8.r
s
N 8y 2 (0, alternating (depending on N), and
A /5pp / (N] For odd, a similar construction holds.
By combining Theorem F.7 and Equation 72
we deduce the following tight lower bound on the
loss function:
(x,1)² 1)2
2|
(()) 2
2 ||
(89)
D
F3 Utility update
Next, we consider the utility for a new neuron after some group of N 2 neurons has undergone the previous steps AGF. Suppose that θ, (ut, vt, wt ) are parameters minimizing the loss as in Section F2 We start by describing the function computed by the network with parameters e..
Lemma F88. For all 0 a, < the network satisfies:
f(a x. b. x,bb,;;,)) = 2|1i[E]
(a + b) Xe:
(90)
where is defined as xelc] = cos (2
Proof. The N neurons compute the function
N Σ
f(a x, b;; O.]cc]
((u a. x)) + (vi, b. 2))² [c]cc.
(91)
The equation above can be expanded via a computation analogous to the one the proof of Lemma
F6 The conditions minimizers from Theorem F.7 imply that the only non- vanishngg term is of the
38

---

form:
N
2
Σ
A cos
28p + 5 - Sy
2T 5 (c 1 b)
2Tp'
p
N
27p
2
cOs
8n
2n _ (c - a -ab)
p
Σ
A A A cos
+ 82 8/ -s)
(92)
/54p/ [6[e]l
2 2|1[
8x
2T (c
a (c-
E
b)).
p
p
where in the first identity we used the fact that Σi= A A Ar sin(s", + st st)) = 0. This concludes the proof.
Finally, we compute the utility function with the updated residual r(a xr b x) (a + b)
(1/p) (z, 1/1 f(a x, x; O.). which is maximized by new neurons the second iteration of AGE.
Theorem F9. The utility function with the updated residual for neuron with parameters 0 (u v, w) is:
u(0:r)=
2 pd (E.t))
Σ
(93)
Proof. The new utility function is
= p² a.b=0
(u,,a + (u,b (,))) w. f(a 2, b,bb2z. x; 0.))
(94)
(
We wish to express A in the frequency domain. By plugging in Equation (90). and due to the cyclic structure of Xe. the squared terms average out:
A
(wu,a x/2 + 2(u, av b. x) + (v, x))) (w, sla b. a; O.))
(95)
a.b=0
p-1
2
(u,a a x) (u, x) (w, f(a x, b,b-r;; O.)).
(96)
0,b=0
Now, by reasoning analogously to the proof of Lemma F.I we obtain:
p-1
A
(x x [)(]
(a * () (xe w)(a + 6]
P
a,b=0
[|]
p-1
(97)
* u
[k]
(x * [ (Xe * w)) [k]
p2
2 ||
+
p
where in the last equality we used the fact that [[(k) -e Lis if k: = = and otherwise. By subtracting the above expression to the expression for the initial utility from Lemma F22 obtain the desired result.
In summary, after an iteration of AGF, the utility for a new neuron has the same form as the initial one, with the summand corresponding to the dominant (conjugate) frequencies + of removed. Consequently, by the same reasoning as in Section F. conclude that during the utility
39

---

maximization phase of the second iteration of AGF, some new group of dormant neurons aligns with the harmonic whose frequency has the second largest magnitude in å.
subtlety arises during the cost minimization phase of the second iteration, as the neurons that aligned during the first phase still involved in the optimization process. However, note that in the loss component (()) (Equation (73), the terms of the form (w,, wy) vanish when i and j are indices of neurons aligned with harmonies of different frequencies. Therefore, during the second cost minimization phase, the loss (()) splits into the sum of two losses, corresponding to the neurons aligned during the first and second iteration of AGF, respectively. The neurons from the first phase already at critical point of their respective loss term, thus the second group of neurons is optimized independently, via the same arguments as Section F.I This scenario analogous to the one discussed in Section E.2 for linear transformers (cf. Lemma E.2). In conclusion, after the second iteration of AGF, the new utility will again have the same form as the initial one, but with the two (conjugate pairs of) frequencies removed. This argument iterates recursively, until either all the frequencies are exhausted, or all the neurons have become active. The quantities in Equation (9) can be derived for the successive iterations of AGF analogously to the previous sections. Lastly, the estimate on jump times is obtained by the same argument as in Section E.3
40

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The abstract describes the contributions the high level, while the introduction contains precise statement of contributions, with references to the sections where such contributions are discussed.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The main limitation of work is that our theory developed for two-layer neural networks. We highlight this limitation in the title, abstract, and introduction. When introducing our framework in Section 2 we clearly explain the assumptions used develop our algorithm. At the end of the paper in Section 6 we discuss additional limitations of our work and how future work could address them.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well- specsfficii,,,nn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
41

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All theorems and formal statements stated either in the main body, in the appendix, with the full set of assumptions and clearly defined notation. All the formal proof' are provided in the appendix with intuition discussed in the main.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
We provide detailed description of our algorithm AGF in Section 2.2 and provide pseudocode in Figure 2 In each setting where we applied AGF we fully describe the experimental setup, including necessary assumptions on initialization and hyperparameters, such that reader could faithfully reconstruct experiments. Every figure clearly discusses what quantities from the experiments are being visualized.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (c.g., in the case
42

---

of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: In supplementary material we enclose set of notebooks that can be run to generate all figures and experiments discussed the paper. In these notebooks we have also implemented version of our main algorithm AGF such that others could adapt our code to run in settings not discussed this work.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/ honpll/e//uuddde//ddedeeeessssss
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips cc/ http://llcggggideeeeeeeeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
43

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: The primary purpose our experiment is to describe the feature learning dynamics of gradient flow. Thus, our experiments do not discuss train test splits, rather they show faithful reconstruction gradient flow dynamics from our proposed theory of AGF. As mentioned above, we provide all the necessary hyperparameters (learning rate and initialization scheme) to reconstruct the figures introduced in this work.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
The experiments we introduce are for gradient flow with complete access to the training data. Thus, in the figures present there should be variability between runs with the same initialization as each run is essentially a solution to an ODE up to numerical precision. Our provided code sets random seed for each experiment such that reader can exactly reproduce each figure introduced our work.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
44

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
All experiments this work involve shallow networks with small constructed datasets that can be run locally. Our provided code describe the necessary packages needed to run our code, but specialized computer resources are not needed.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: We have reviewed the Code Ethics, and believe that the research in the paper conforms to it in every aspect (e.g., complete anonymity).
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper discusses general theoretical properties of neural networks. As such, we do not see any potential societal impact of the research presented here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out
45

---

that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not involve any model dataset with high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
The paper does not use any existing asset.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasete has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
46

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not introduce any new asset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (TRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
47

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: No LLM was used the development of the core results of the paper. LLMs were only used for writing and editing purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
48


Reviews:
[{"review_id": "4rMpBx4Zng", "reviewer": "Reviewer_wsYE", "summary": "This 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.", "questions": "No", "limitations": "No", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths : \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", "was_revised": true, "final_justification": "I thank the authors for the answer that they give. It confirms my overall impression on the paper. I am not very familiar with the modular addition framework, its intrinsic difficulty, etc, but the authors made a point that their algorithm can illustrate well (and not explain!) some mechanisms that appear during the learning of neural networks.\n\nI am willing to keep my score.", "rebuttal": "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. \n\n**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:\n\n- 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.\n\n- 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.\n\nWe 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.\n\n**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.\n\n**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?\n\n**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.\n\nThank 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."}, {"review_id": "Fsuhj4RqiZ", "reviewer": "Reviewer_hRND", "summary": "This 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.", "questions": "- 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?", "limitations": "N/A", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "My questions have been clearly addressed by the authors. I stand by my original review and recommend acceptance of the paper, with the hope that the fuzzy definitions will be clarified in the revised version.", "rebuttal": "We are grateful for the thorough review, the constructive suggestions, and the positive feedback. We will address the weaknesses and questions individually. \n\n**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). \n\n**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.\n\n- 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. \n\n- 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.\n\n- 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.\n\n- 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. \n\n**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.\n\n**Questions** \n- 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. \n\n- 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. \n\n- 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$. \n\nWe 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!"}, {"review_id": "Pk9CzSTEmd", "reviewer": "Reviewer_D2un", "summary": "In 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.", "questions": "1.\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$?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The proposed AGF framework unifies several prior analyses of saddle-to-saddle dynamics and offers intuitive explanations for the behavior observed in these settings. The response from authors addressed my concerns. I believe this is a valuable contribution and therefore recommend acceptance.", "rebuttal": "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. \n\n**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.\n\n**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. \n\n**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.\n\n**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.\n\nPlease 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."}, {"review_id": "km9z2w0z2t", "reviewer": "Reviewer_qgue", "summary": "This 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.", "questions": "- 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", "limitations": "Limitations are discussed in the last section.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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).", "was_revised": true, "final_justification": "The concerns I have raised have been properly addressed by the authors during the rebuttal period. \n\nI remain on the opinion that the paper has a tendency to oversell the centrality of AGF for understanding feature learning, specially given the evidence provided, which is limited to retrieving known phenomenology from a new perspective. Nevertheless, I think the promised changes will help to better place the paper within other approaches to study feature learning on 2LNN, as well as to make some of the claims more sober.\n\nConsidering this, as well as a careful reading of the other reviews, I am happy to reconsider my score towards a weak accept.", "rebuttal": "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.\n\n## Questions\n\n*1. Confusing notation on L64-65.*\n\nIn 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)$.\n\n*2. What is a smooth dynamical system in L77-78?*\n\nIn 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.\n\n*3. Tuning down statement on L43-45*\n\nOur 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. \n\n*4. What is the role of the parametrization on L337-338?*\n\nWe 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. \n\n*5. & 6. What is the distinction with prior analysis, and why do we see different feature learning behavior?*\n\nWe 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. \n\nMoreover, 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: \n\n\n||**Prior Analyses (e.g., Saad, Goldt, Arnaboldi)**|**AGF Framework (our work)**|\n|---|---|---|\n|**Output dimension**|Scalar output only ($c=1$)|Handles both scalar ($c=1$) and vector-valued outputs ($c>1$)|\n|**Optimization**|Often train only the first layer with online SGD and keep the second layer fixed|Always train both first and second layers jointly under GF|\n|**Initialization regime**|Finite initialization variance ($\\alpha \\in O(1)$)|Vanishing initialization ($\\alpha \\to 0$)|\n|**Input dimension**|Thermodynamic limit ($d \\to \\infty$)|Finite input dimension ($d \\in O(1)$)|\n|**Analytical approach**|Track a small number of macroscopic order parameters with deterministic ODEs that can be solved numerically|Track individual neuron dynamics via a timescale separation between direction and norm|\n|**Task and theoretical focus**|Teacher–student setups; generalization under online SGD; each work is specific to a particular architecture or activation|Feature acquisition under gradient flow; general framework that unifies saddle-to-saddle dynamics across multiple settings, can extend to new settings (modular arithmetic)|\n\n## Weaknesses\n*1. The relevance of AGF beyond simple cases*\n\nWe 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. \n\nRegarding 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.\n\n*2. The “broad picture” offered by AGF*\n\nWe 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.\n\n*3. Missing references*\n\nYou 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.\n\nPlease 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."}, {"review_id": "A9URlQFL4H", "reviewer": "Reviewer_vhww", "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.", "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", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe 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.\n\n\n**Weaknesses**\n\n1. 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.\n\n2. Extending the proposed framework to deeper architectures is non-trivial, potentially limiting its generalizability and broader impact.", "was_revised": true, "final_justification": "The rebuttal has adequately addressed my concerns. Considering the paper's novelty and insight, I am raising my overall evaluation to a 5.", "rebuttal": "Thank you for your review. We will address the two weaknesses you brought up individually below.\n\n**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.\n\n**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.\n\nWe 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."}]

Author rebuttal:

Meta-review: This work uses an Alternating Gradient Flow (AGF) algorithm for approximating the gradient flow dynamics of two-layer layer neural networks. 
The motivation is from the fact that convergence of AGF towards GF is established for diagonal linear networks, but I'm still questionable about such a general conjecture (also mentioned by Reviewer wsYE). More discussion are needed, also committed by the authors.

The spirit of AGF follows the saddle-to-saddle dynamics: i) at a saddle-point, a subset of dormant neurons (with small weights) evolve independently to maximize their correlation with the rest, active neurons; ii) the alignment for a neuron leads to a drop of loss until they reach a new plateau. The results show that i) AGF is exact when the initialization scale for diagonal neural networks. ii) linear two-layer networks make AGF recovers the certain subspace. iii) AGF implements stepwise principal component regression for linear transformers. iv) two-layer quadratic networks trained on AGF learn Fourier features a modular addition task.

The results are comprehensive and all of the reviewers vote for acceptance after the rebuttal. Hence I recommend to accept this paper and suggest the authors to incoporate the reviewers' feedback in the final version.
Decision: Accept (poster)