Technical Summary: Global Convergence of Gradient Descent for Wide Shallow Models with Bounded Nonlinearities

1. Problem Statement

The paper investigates the global convergence of gradient descent (specifically, gradient flow) for training wide shallow neural networks and attention mechanisms. While gradient descent is a first-order algorithm that typically risks converging to non-optimal critical points in non-convex landscapes, empirical evidence suggests it often finds global minimizers in over-parameterized regimes.

Previous theoretical work, notably [CB18], established global convergence for wide networks under specific conditions:

• ReLU activations: Where the mapping is positively 2-homogeneous.
• Sigmoid activations with scalar output weights: Where the mapping is partially 1-homogeneous.

However, these results do not cover:

1. Vector output weights: Common in multi-class classification and attention mechanisms (where output dimension $d_w > 1$).
2. Attention layers: Specifically, multi-head attention layers which rely on bounded nonlinearities (softmax) and vector-valued parameters.
3. General initializations: Previous proofs often required the initial parameter distribution to have compact support, excluding popular Gaussian initializations.

The authors aim to bridge these gaps by proving that, for a large class of models including multi-head attention and two-layer networks with vector output weights, all non-global minimizers are unstable under gradient descent dynamics. Consequently, if the initial distribution has full support (e.g., Gaussian), the continuous-time gradient flow converges to global minimizers in the limit of infinite width.

2. Methodology

2.1 Mean-Field Framework

The analysis relies on the mean-field limit, where the parameters of $m$ neurons are encoded as a probability distribution $\mu_t$ over the parameter space $\Omega$. As $m \to \infty$, the discrete gradient descent dynamics converge to a Wasserstein gradient flow of a functional $F$ defined on the space of probability measures $\mathcal{P}(\Omega)$.

The objective function is $F(\mu) = \mathcal{R}(\int_\Omega \Phi d\mu)$, where $\mathcal{R}$ is the loss and $\Phi(w, \theta) = \phi(\theta)w$ represents the model. Here, $w$ corresponds to output weights (or value matrices in attention), and $\theta$ to hidden parameters (or query/key matrices).

2.2 Well-Posedness and Stability (Section 2)

Before analyzing convergence, the authors establish the well-posedness of the mean-field dynamic.

• Assumption: The initial measure $\mu_0$ is sub-Gaussian (a weaker condition than the compact support required in [CB18]). This includes Gaussian initializations.
• Result: They prove the existence and uniqueness of the Wasserstein gradient flow for all time $t \geq 0$. Furthermore, they establish a stability estimate: if a sequence of empirical measures $\mu_{m,0}$ converges to $\mu_0$ in $W_2$ distance, the corresponding flows converge uniformly on finite time intervals.
• Technique: Unlike [CB18] which used metric space gradient flow theory, this proof utilizes a fixed-point argument based on Dobrushin-type estimates for interacting particle systems, leveraging the sub-Gaussian tail control to handle the linear growth of the velocity field.

2.3 Global Convergence via "Escaping Active Sets" (Section 3)

The core theoretical contribution is the proof that non-global minimizers are unstable.

• Mechanism: The authors introduce the concept of an "escaping active set" ($E(g)$). If a candidate limit point $\mu$ is not a global minimizer, there exists a region in the parameter space where the output weights $w$ must grow uniformly at a positive rate.
• Theorem 2: If the initial distribution has full support, the Wasserstein gradient flow cannot converge to any measure $\mu$ for which the associated vector field $g_\mu$ admits an escaping active set.
• Corollary: If every non-zero $g_\mu$ admits an escaping active set, then the flow converges only to global minimizers (where $g_\mu \equiv 0$).

2.4 Correction and Extension of [CB18]

• Scalar Case ($d_w=1$): The authors provide a corrected proof for the existence of escaping active sets for bounded nonlinearities (sigmoid) with scalar outputs. The original proof in [CB18, Appendix C] contained an erroneous argument regarding the boundedness of the gradient on the boundary of the active set. The authors fix this by analyzing the behavior at infinity using uniform convergence assumptions on the mapping and its derivative.
• Vector Case ($d_w > 1$): The authors extend the analysis to vector output weights. They demonstrate that the natural generalization of the scalar proof fails due to the complexity of the vector dynamics (direction changes). They propose a second construction of the active set using a fixed direction vector $v$. While they prove this construction works for $d_\theta = 1$ (scalar hidden parameters), they conjecture that the result holds for $d_\theta > 1$ under mild regularity assumptions, though a rigorous proof for the general vector case remains open.

3. Key Results

1. Well-Posedness for Sub-Gaussian Initializations: The mean-field training dynamic is well-posed and stable for sub-Gaussian initial distributions, extending previous results that required compact support.
2. Global Convergence for Bounded Nonlinearities: For models with bounded nonlinearities (e.g., sigmoid) and scalar output weights, the Wasserstein gradient flow converges to global minimizers, provided the initial distribution has full support.
3. Application to Attention and Vector Weights:
   • The framework applies to multi-head attention layers (modeled as $\Phi(A, V) = V \psi(A)$ where $A=K^T Q$).
   • The authors prove that the softmax attention mechanism converges uniformly to a "hardmax" mechanism as parameters grow large, a result of independent interest.
   • They establish that the conditions for global convergence are met for these architectures, though the rigorous proof for the vector output case ($d_w > 1$) relies on a conjecture regarding the existence of escaping active sets in higher dimensions.
4. Correction of Prior Work: The paper identifies and corrects a flaw in the proof of [CB18, Proposition C.4] regarding the instability of non-global minimizers for sigmoid networks.

4. Significance and Claims

The paper claims to significantly broaden the theoretical understanding of why gradient descent succeeds in over-parameterized deep learning:

• Bridging the Gap: It moves beyond the restrictive NTK regime (where features do not learn) and the specific case of ReLU activations. It addresses the feature learning regime for models with bounded nonlinearities, which are ubiquitous in modern architectures (Transformers, Sigmoid networks).
• Realistic Architectures: By including multi-head attention layers and vector output weights, the results apply to architectures that are central to current state-of-the-art models, rather than just simplified two-layer networks.
• Robust Initialization: By relaxing the initialization requirement from compact support to sub-Gaussian, the theory aligns better with practical training practices (e.g., Gaussian initialization).
• Modesty on Open Problems: The authors explicitly state that the complete proof for the vector output case ($d_w > 1$) with general hidden dimensions ($d_\theta > 1$) remains an open question. They provide a sufficient condition and a local analysis but do not claim a full proof for the general vector case. They also note that their analysis excludes regularization terms (like $\ell_2$ weight decay) and non-bounded activations like GELU/SiLU, which are common in practice but fall outside their current theoretical scope.

In summary, the paper provides a rigorous mathematical foundation suggesting that for wide shallow models with bounded nonlinearities, the landscape of the training loss is such that gradient flow is repelled from non-optimal critical points, provided the initialization is sufficiently broad. This supports the empirical observation that gradient descent finds global solutions in these settings.