Alternating Gradient Flows (AGF) in Neural Networks

Last updated: June 11, 2025

Alternating Gradient Flows ° (AGF) constitute a rigorous theoretical framework for describing feature learning dynamics ° in two-layer neural networks ° trained from small initialization. AGF captures observed training trajectories—characterized by loss plateaus ° punctuated by abrupt drops—through an alternating, two-step process that has been analyzed mathematically and validated empirically in recent work (Kunin et al., 6 Jun 2025 ° ).

Significance and Background

A pivotal question in deep learning theory ° concerns what features neural networks acquire and the mechanisms by which they do so. Empirical and analytical studies have demonstrated that training two-layer neural networks with very small random weights ° produces staircase-like loss curves: periods of near-constant loss (“plateaus °”) separated by sudden drops in loss, each corresponding to the network’s acquisition of a new informative feature (Kunin et al., 6 Jun 2025 ° ).

Prior research described aspects of these phenomena in special cases such as diagonal or fully connected linear networks ° by analyzing "saddle-to-saddle" paths in the loss landscape. AGF extends and unifies these analyses, providing explicit predictions for the timing, sequence, and magnitude of loss drops in a broad class of two-layer architectures, including diagonal linear, fully-connected linear, and quadratic networks (Kunin et al., 6 Jun 2025 ° ).

Foundational Concepts: The AGF Mechanism

AGF models feature learning ° as a sequence of discrete phases, each decomposed into two alternating steps:

1. Utility Maximization ° (Dormant Neurons): During a plateau, each dormant (inactive) neuron maximizes its individual "utility" function—quantifying the alignment between its representational direction and the current residual error. For neuron $i$, the utility is given by

$\mathcal{U}_i(\theta; r) = \mathbb{E}_x[\langle f_i(x; \theta), r(x) \rangle],$

where $r(x) = y(x) - f(x; \Theta_\mathcal{A})$ is the residual and $\mathcal{A}$ indexes the set of currently active neurons [(Kunin et al., 6 Jun 2025 ° ), Sec. 3.2].

1. Cost Minimization ° (Active Neurons): When a dormant neuron’s weight norm reaches a threshold and it becomes active, all active neurons jointly minimize the overall cost (the residual). During this phase, the dormant neurons’ weights change negligibly.

This alternation arises directly from the analysis of the time-scale separation ° in the continuous-time gradient flow dynamics ° when initialization is small: the active directions evolve rapidly (minimizing loss), while dormant directions change slowly (maximizing alignment) [(Kunin et al., 6 Jun 2025 ° ), Sec. 3.2].

The two-layer network is given by

$f(x; \Theta) = \sum_{i=1}^{H} a_i \sigma(w_i^\intercal g_i(x))$

where $\sigma$ is an origin-passing activation and $g_i(x)$ are potentially neuron-specific input mappings [(Kunin et al., 6 Jun 2025 ° ), Eq. 2].

For each dormant neuron, the evolution of the norm and direction (with timescale separation) is governed by:

$\frac{d}{dt}\|\theta_i\| = \eta\kappa \|\theta_i\|^{\kappa - 1} \mathcal{U}_i(\bar{\theta}_i; r), \qquad \frac{d}{dt} \bar{\theta}_i = \eta\|\theta_i\|^{\kappa - 2} \mathbf{P}^\perp_{\theta_i}\nabla_{\theta}\mathcal{U}_i(\bar{\theta}_i; r),$

where $\kappa$ is the degree of the leading term in the Taylor expansion ° of $\mathcal{U}_i$ at the origin ($\kappa=2$ for linear, $3$ for quadratic activations) [(Kunin et al., 6 Jun 2025 ° ), Eq. 5].

The "jump" or activation time for neuron $i$ is given by

$\tau_i = \inf\left\{ t > 0 : \mathcal{S}_i(t) = \frac{c_i}{\eta} \right\}$

where $\mathcal{S}_i(t) = \int_0^t \kappa \mathcal{U}_i(\bar{\theta}_i(s); r)\,ds$ and $c_i$ is determined by the initialization norm [(Kunin et al., 6 Jun 2025 ° ), Eq. 6].

A staircase trajectory then results: during each plateau, dormant neurons align with the residual; at each jump time, the first to reach threshold activates, triggering a sudden drop in loss.

Principal Findings

AGF provides a detailed quantitative description and successful predictions across multiple network architectures [(Kunin et al., 6 Jun 2025 ° ), Secs. 3–6, Figs. 1–3]:

• Order of Feature Learning: Neurons become active in the order of greatest available utility, aligning with principal components °, singular modes, or dominant Fourier coefficients—depending on architecture.
• Timing of Loss Drops: The explicit activation times predicted by AGF closely match the durations of plateaus observed in experiments, and provide lower bounds in theoretical analysis.
• Magnitude of Loss Drops: Each loss drop is quantitatively tied to the portion of the target signal explained by the new feature (e.g., the variance of a principal component or the coefficient of a Fourier harmonic).
• Unification and Extension of Previous Analyses: AGF encompasses and extends "saddle-to-saddle" analyses for diagonal linear networks (where features correspond to individual coordinates) and for fully connected linear networks (where features correspond to dominant singular modes), and provides convergence guarantees in the small-initialization limit [(Kunin et al., 6 Jun 2025 ° ), Theorem 1].
• Novel Results for Quadratic Networks and Modular Arithmetic: In two-layer quadratic networks trained on modular addition, AGF characterizes the stepwise acquisition of Fourier components ° in decreasing order of coefficient magnitude—a previously unexplained phenomenon [(Kunin et al., 6 Jun 2025 ° ), Sec. 6].

Illustrative Table: AGF-Predicted Feature Learning

Empirical validations throughout the paper confirm the accuracy and generality of these predictions [(Kunin et al., 6 Jun 2025 ° ), Figs. 1–2, 6].

Applications and Scope

AGF serves as an analytic tool for understanding and predicting:

• Feature emergence and ordering: Especially in networks where the task or data structure admits interpretable components (principal components, SVD, Fourier modes).
• Grokking and generalization in algorithmic tasks: In modular addition, AGF explains the sequential appearance of Fourier features that underlies grokking behavior ° [(Kunin et al., 6 Jun 2025 ° ), Sec. 6].
• Plateau phenomena and training dynamics: AGF mathematically characterizes training plateaus as periods of dormant neuron alignment, interrupted by rapid capacity increases as new features are learned [(Kunin et al., 6 Jun 2025 ° ), Sec. 3].

While the framework is currently established for two-layer networks, the underlying mechanisms—timescale separation and alternation between utility maximization and cost minimization—suggest broader applicability [(Kunin et al., 6 Jun 2025 ° ), Discussion].

Limitations and Open Directions

AGF is primarily developed for two-layer networks trained under continuous-time gradient flow from small initialization. The current formalism:

• Assumes vanishing initialization and exact gradient flow: The role of stochasticity and the effects of moderate or large initial weights remain open questions [(Kunin et al., 6 Jun 2025 ° ), Discussion].
• Titles its extension to deeper, modular, or convolutional architectures ° as future work: Although the alternation and timescale principles could plausibly extend, full theoretical development is outstanding.
• Does not directly account for stochastic optimization: The interaction with SGD noise, common in practical training, is an important area for further research.

Table: Strengths and Caveats

Speculative Note

While AGF demonstrates remarkable alignment between theoretical predictions and experimental outcomes, its extension to more complex network classes, to stochastic settings, and to more realistic initializations is an area of active research [(Kunin et al., 6 Jun 2025 ° ), Discussion].

References

• Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks (Kunin et al., 6 Jun 2025 ° ). (See main text: Sections 2–6, Figures 1–2, Theorem 1, and Discussion for theoretical and empirical details.)