Technical Summary: Born Discrete, Made Smooth: Variational Formulation of Shallow Neural Networks

Problem Statement

The theoretical understanding of gradient-based optimization in highly over-parameterized, non-convex neural network objectives remains elusive. Existing frameworks face distinct limitations:

• Mean-Field/Wasserstein Approaches: While offering a particle interpretation of gradient descent and connections to interacting particle systems, these methods often suffer from limited regularity of minimizers and require non-trivial approximations to transition from infinite-width continuum models to finite-width networks.
• Neural Tangent Kernel (NTK) Framework: This approach achieves global convexity but only by linearizing the network around its initialization, restricting theoretical guarantees to the "lazy-training" regime where features do not significantly change.

The central challenge addressed is whether a nonlinear continuum model for shallow neural networks can admit a unique, stable minimizer with high regularity without linearization, and whether this continuum formulation can serve as an exact relaxation of the discrete training problem.

Methodology

The authors propose a paradigm shift by replacing the discrete training problem of shallow neural networks with a well-posed continuum variational surrogate defined over parameter densities.

1. Variational Formulation

The training of a shallow network approximating a target function $f$ is reformulated as minimizing a risk functional over a space of parameter densities $u(\theta)$ in a weighted Sobolev space $W = W^{1,2}(\Omega) \cap L^2_\omega(\Omega)$. The parameter space $\Omega$ represents the hidden units, and the activation function $\sigma$ is assumed continuous (e.g., ReLU, sigmoid).

The regularized objective functional $F^{(f)}_{\alpha, \beta}(u)$ is defined as:
$$F^{(f)}_{\alpha, \beta}(u) = R(f, u d\theta) + \alpha \|u\|^2_{L^2_\omega} + \beta \|\nabla u\|^2_{L^2}$$
where:

• $R(f, u d\theta)$ is the squared risk (data fidelity).
• $\alpha \|u\|^2_{L^2_\omega}$ provides coercivity and $\lambda$-convexity.
• $\beta \|\nabla u\|^2_{L^2}$ acts as a Sobolev regularization term, promoting smoothness and inducing a diffusive effect analogous to entropy regularization in Wasserstein formulations but within a Hilbert-space setting.

2. Analytical Framework

The analysis relies on classical calculus of variations and elliptic PDE theory within a linear Hilbert-space setting. Key theoretical tools include:

• Convex Analysis: Establishing $\lambda$-convexity of the functional.
• Elliptic Regularity: Analyzing the Euler-Lagrange equation associated with the minimizer.
• Gradient Flow Theory: Utilizing the Hille-Yosida and Brézis-Komura theorems to study the continuous-time limit of gradient descent (the gradient flow in $L^2_\omega$).

Key Contributions and Results

1. Exact Variational Characterization (No Lavrentiev Gap)

The paper proves that the infimum of the risk is identical whether optimization is performed over:

• Atomic measures (finite-width networks).
• Sobolev densities.
• Smooth, compactly supported functions.

Theorem 1 establishes that there is no "Lavrentiev gap," meaning the continuum formulation is an exact relaxation of the discrete problem. Furthermore, Proposition 2 quantifies the convergence of finite-width networks to the continuum optimum, showing that networks of width $N$ achieve the continuum solution with an error rate of $O(1/N)$.

2. Global Well-Posedness and Convexity

Unlike the NTK framework, which requires linearization, this formulation achieves global convexity for the nonlinear problem:

• The functional is $2\min(\alpha, \beta)$-convex on the space $W$, ensuring the existence and uniqueness of a minimizer $u^*$.
• The functional is $2\alpha$-convex on $L^2_\omega(\Omega)$, which implies that the associated gradient flow converges exponentially fast to the equilibrium $u^*$ at a rate of $e^{-2\alpha t}$ (Corollary 9).

3. Near-$C^3$ Regularity of the Minimizer

A significant finding is the regularity of the optimal parameter density $u^*$.

• Theorem 3 demonstrates that $u^*$ is a classical solution to an explicit Euler-Lagrange equation:
  $$-\beta \Delta u^*(\theta) + \alpha \omega(\theta) u^*(\theta) - Q(\theta) + \int_\Omega K(\theta, \vartheta) u^*(\vartheta) d\vartheta = 0$$
• The minimizer $u^*$ belongs to $C^{2,s}_{loc}(\Omega)$ for every $s < 1$. This "near-$C^3$" smoothness prevents pathological mass concentrations in the parameter space and suggests that the network predictions concentrate around graphs of smooth functions. This regularity holds even for locally Lipschitz activation functions, contrasting with the regularity constraints often found in Wasserstein-based approaches.

4. Stability and Generalization

The framework provides explicit stability controls:

• Proposition 4 and Proposition 6 show that the minimizer depends linearly and Lipschitz-continuously on the target data $f$.
• Corollary 7 establishes that the generalization error scales linearly with the noise level $\varsigma$ and is controlled by $1/\alpha$. This provides a concrete form of implicit bias toward smooth, well-generalizing solutions.

5. Computational Implications

Due to the quadratic nature of the functional and the convexity of the problem, the minimizer can be obtained by solving a single linear system of equations (reducing to ridge regression in a chosen basis). This bypasses the need for iterative optimization or stochastic gradient descent entirely in the continuum limit.

Significance and Claims

The paper claims to bridge the gap between the Neural Tangent Kernel (NTK) and feature-learning regimes by providing a principled framework for understanding over-parameterization through variational calculus.

• Beyond Linearization: The authors assert that their $\lambda$-convexity holds for the same nonlinear problem arising in the infinite-width limit, without requiring the linearization assumptions necessary for NTK analysis.
• Mechanism for Non-Overfitting: The near-$C^3$ smoothness of the optimal density offers a concrete explanation for the non-overfitting phenomenon in over-parameterized networks, suggesting an implicit attraction to low-dimensional smooth manifolds in parameter space.
• Complementarity: The approach is presented as complementary to existing mean-field and Sobolev-penalty literature. While previous works focused on the conditioning of optimization or particle approximations, this work focuses on the variational structure of the objective and the regularity of its minimizers.
• Limitations: The authors explicitly note that the current formulation is restricted to shallow (single-hidden-layer) networks. Extending this to deep architectures introduces strong nonlinearity, making the Euler-Lagrange system a coupled nonlinear PDE where existence of smooth minimizers requires additional structural assumptions (e.g., small initial conditions).

In summary, the paper posits that casting shallow network training as a variational problem over parameter densities unlocks global convexity, elliptic regularity, and gradient flow theory simultaneously, providing an analytically controlled and dynamically stable model that is exactly tied to finite-width networks.