Towards Understanding Gradient Flow Dynamics of Homogeneous Neural Networks Beyond the Origin

Technical Summary: Towards Understanding Gradient Flow Dynamics of Homogeneous Neural Networks Beyond the Origin

Problem Statement
Modern deep neural networks trained via gradient-based methods often exhibit strong generalization despite over-parameterization, a phenomenon attributed to implicit regularization induced by the training algorithm. While significant progress has been made in understanding training dynamics in the Neural Tangent Kernel (NTK) regime (large initialization) and the late stages of training, the "small initialization" regime remains less understood. In this regime, networks operate in a feature-learning mode, yet the highly non-linear evolution of weights poses analytical challenges.

Previous works (e.g., Maennel et al., 2018; Kumar and Haupt, 2024, 2025) established that for homogeneous neural networks with small initialization, weights remain near the origin for an extended period. During this phase, the weights converge in direction toward a Karush-Kuhn-Tucker (KKT) point of a Neural Correlation Function (NCF), a phenomenon known as "early directional convergence." Furthermore, these early-stage KKT points often exhibit sparsity. However, the behavior of the gradient flow after the weights escape the origin and the preservation of the emerging sparsity structure during this transition were not fully characterized. This paper addresses the gap by analyzing the gradient flow dynamics of homogeneous neural networks immediately after they escape the origin, up to the first saddle point encountered.

Methodology and Setup
The paper studies homogeneous neural networks $H(x; w)$ that are $L$-positively homogeneous ($L \ge 2$), meaning $H(x; cw) = c^L H(x; w)$ for $c \ge 0$. The analysis assumes the network has locally Lipschitz gradients, which includes deep linear networks and networks with polynomial ReLU activations ($p \ge 2$), but explicitly excludes standard ReLU networks ($p=1$) due to non-differentiability at the origin, though empirical results suggest potential extensions.

The training objective is the minimization of a loss function $\mathcal{L}(w)$ via gradient flow:
$$\dot{w} = -\nabla \mathcal{L}(w)$$
The authors build upon the concept of the Neural Correlation Function (NCF), defined as $N(u) = \bar{e}_y^\top H(X; u)$, where $\bar{e}_y = -\ell'(0, y)$. They analyze the constrained optimization of $N(u)$ on the unit sphere.

The core methodology involves:

1. Limiting Trajectory Construction: Defining a limiting solution $p(t)$ that represents the path of the gradient flow after escaping the origin. This is achieved by time-shifting the solution $\psi(t, \delta w_0)$ (where $\delta$ is the initialization scale) by a term dependent on $\ln(1/\delta)$ and taking the limit as $\delta \to 0$.
2. Stability Analysis: Utilizing the properties of $\Delta$-second-order positive KKT points of the constrained NCF. The authors assume the initial direction of the weights lies in the stable set of such a KKT point.
3. Zero-Preserving Subsets: Introducing the concept of a "zero-preserving subset" of weights—weights initialized at zero that remain zero throughout training under gradient flow. This is used to analyze the preservation of sparsity.

Key Contributions and Results

1. Gradient Flow Dynamics Post-Escape (Theorems 4 and 10)
The paper establishes that for sufficiently small initialization, the trajectory of the gradient flow after escaping the origin closely follows a specific limiting path $p(t)$, regardless of the specific small initialization scale $\delta$, provided the initial direction aligns with a stable set of a second-order positive KKT point of the NCF.

• Two-Homogeneous Networks ($L=2$): Theorem 4 shows that after a time scaling of $O(\ln(1/\delta))$, the trajectory $\psi(t, \delta w_0)$ converges to $p(t)$ with an error bound of $O(\delta^{\frac{\Delta}{\Delta + 8N(w^*)}})$.
• Deep Homogeneous Networks ($L > 2$): Theorem 10 extends this to $L$-homogeneous networks. The escape time scales as $O(\delta^{-(L-2)})$, and the trajectory similarly converges to a limiting path $p(t)$ with a specific error bound.
• Implication: The behavior of the gradient flow immediately after escaping the origin is primarily determined by the KKT point of the NCF that the weights approached while near the origin.

2. Characterization of the First Saddle Point (Corollaries 5, 6, 11, 12)
The authors demonstrate that the limiting trajectory $p(t)$ converges to a stationary point (saddle point) of the training loss. Consequently, for sufficiently small initialization, the actual gradient flow trajectory $\psi(t, \delta w_0)$ will approach this same saddle point after escaping the origin.

• The paper distinguishes between cases where the saddle point is finite (bounded trajectory) and cases where it is at infinity (unbounded trajectory, common with logistic loss).
• The results imply that the gradient flow follows a "saddle-to-saddle" trajectory, specifically characterizing the first saddle point encountered after the origin.

3. Preservation of Sparsity in Feed-Forward Networks (Section 4 and Theorem 18)
A critical contribution is the proof that the sparsity structure emerging during the "near-origin" phase is preserved after escaping the origin, up to the first saddle point.

• Zero-Preserving Subsets: The authors prove that for feed-forward homogeneous networks, the rows and columns of weight matrices corresponding to zero-magnitude entries in a KKT point of the NCF form a "zero-preserving subset." If these weights are initialized at zero, they remain zero throughout training.
• Sparsity Preservation: Since the gradient flow with small initialization converges in direction to a KKT point with specific sparsity (where certain weights are zero), and these zero weights form a zero-preserving subset, the sparsity structure is maintained as the network escapes the origin and approaches the first saddle point.
• Theorem 18: Formalizes that for both $L=2$ and $L>2$ networks, the magnitude of weights in the zero-preserving subset remains negligible ($O(\delta^{\text{exponent}})$) even after the escape time, effectively preserving the sparsity pattern.

Experimental Validation
The paper provides numerical experiments (Figures 1, 4, 5, 6) using square activation functions ($L=2$) and polynomial ReLU activations.

• Square Activation: Experiments confirm that weights become sparse before escaping the origin and that this sparsity is preserved at the next saddle point.
• ReLU Activation: Although the theoretical framework excludes standard ReLU ($L=1$), experiments on ReLU networks show that while the weights do not become perfectly sparse (more non-zero entries than in the polynomial case), the relative sparsity structure (where small weights remain small) is preserved. This suggests the findings may extend to ReLU networks empirically, despite the lack of local Lipschitz gradients at the origin.

Significance and Claims
The paper claims to provide a rigorous characterization of a specific segment of the training dynamics: the transition from the "near-origin" phase to the first saddle point.

• Mechanism of Implicit Regularization: It links the implicit regularization observed in small initialization regimes to the geometry of the NCF and the specific KKT points the network converges to.
• Stability of Sparsity: It offers a theoretical explanation for why sparsity patterns observed early in training persist, attributing this to the "zero-preserving" nature of specific weight subsets in feed-forward architectures.
• Limitations: The authors explicitly state modesty in their scope. The results are limited to:
  1. Homogeneous networks with locally Lipschitz gradients (excluding standard ReLU).
  2. The dynamics only up to the first saddle point encountered after escaping the origin. The behavior beyond this point (e.g., escaping the saddle to a global minimum) is not covered and is identified as a direction for future work.
  3. The assumption that the initial direction lies in the stable set of a second-order positive KKT point.

In summary, the paper bridges the gap between early-stage directional convergence and the subsequent training phase, demonstrating that the structural properties (specifically sparsity) established while near the origin dictate the trajectory and stability of the network as it begins to learn features and approach the first saddle point.