[Paper Review] Function approximation by neural nets in the mean-field regime: Entropic regularization and controlled McKean-Vlasov dynamics

Belinda Tzen, Maxim Raginsky|arXiv (Cornell University)|Feb 5, 2020

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper establishes a mean-field theory for lazy training in two-layer neural networks by formulating function approximation as a free-energy minimization problem that balances $L^2$ approximation risk and KL divergence from a Gaussian prior. It introduces a controlled McKean–Vlasov dynamics via entropic regularization, showing that SGD approximates the optimal distributional dynamics with quantifiable $L^2$ error bounds.

ABSTRACT

We consider the problem of function approximation by two-layer neural nets with random weights that are "nearly Gaussian" in the sense of Kullback-Leibler divergence. Our setting is the mean-field limit, where the finite population of neurons in the hidden layer is replaced by a continuous ensemble. We show that the problem can be phrased as global minimization of a free energy functional on the space of (finite-length) paths over probability measures on the weights. This functional trades off the $L^2$ approximation risk of the terminal measure against the KL divergence of the path with respect to an isotropic Brownian motion prior. We characterize the unique global minimizer and examine the dynamics in the space of probability measures over weights that can achieve it. In particular, we show that the optimal path-space measure corresponds to the Föllmer drift, the solution to a McKean-Vlasov optimal control problem closely related to the classic Schrödinger bridge problem. While the Föllmer drift cannot in general be obtained in closed form, thus limiting its potential algorithmic utility, we illustrate the viability of the mean-field Langevin diffusion as a finite-time approximation under various conditions on entropic regularization. Specifically, we show that it closely tracks the Föllmer drift when the regularization is such that the minimizing density is log-concave.

Motivation & Objective

• To develop a mean-field theory for function approximation in two-layer neural networks under the lazy training regime.
• To characterize the optimal weight distribution as the minimizer of a free-energy functional combining $L^2$ risk and KL divergence from a Gaussian prior.
• To construct a controlled McKean–Vlasov stochastic process that realizes the optimal distributional dynamics through entropic regularization.
• To establish quantitative guarantees on the $L^2$ approximation error of SGD as a greedy approximation to the optimal mean-field dynamics.
• To link the optimal control problem to the Schrödinger bridge problem and the Föllmer drift, enabling a dynamic interpretation of the mean-field limit.

Proposed method

• Formulate the mean-field approximation problem as global minimization of a free-energy functional: $\mathcal{F}_{\tau,\beta}(\mu) = \mathbb{E}_{w \sim \mu}[\|f(w,\cdot) - f^*\|_{L^2}^2] + \beta \cdot D_{\mathrm{KL}}(\mu \| \gamma_\tau)$, where $\gamma_\tau$ is a centered Gaussian prior.
• Prove existence and uniqueness of the global minimizer of the free-energy functional under regularity conditions, with explicit bounds on $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ and $\mathcal{W}_2(\mu^*, \gamma_\tau)$.
• Construct a controlled McKean–Vlasov SDE with drift $\varphi(w,t)$ that solves a stochastic optimal control problem, where the terminal cost encodes both $\beta$ and $\tau$.
• Identify the optimal control law as the Föllmer drift, which corresponds to the solution of the entropic optimal transport (Schrödinger bridge) problem.
• Derive coupled forward (Fokker–Planck) and backward (Hamilton–Jacobi–Bellman) PDEs governing the evolution of the weight distribution and value function.
• Show that SGD with appropriately tuned learning rate $\eta$, network width $N$, and regularization $\beta$ approximates the optimal McKean–Vlasov dynamics with error $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$ with high probability.

Experimental results

Research questions

• RQ1 How can the lazy training regime in two-layer neural networks be formalized via a mean-field free-energy minimization framework?
• RQ2 What is the structure of the optimal weight distribution that minimizes the trade-off between approximation error and deviation from a Gaussian prior?
• RQ3 Can the optimal distributional dynamics be realized as a controlled McKean–Vlasov process, and how is it related to entropic optimal transport?
• RQ4 To what extent does stochastic gradient descent (SGD) approximate the optimal mean-field dynamics in the lazy training regime?
• RQ5 What are the finite-sample and high-probability error bounds for SGD in approximating the optimal distributional dynamics?

Key findings

• The free-energy functional has a unique global minimizer $\mu^*$, and the KL divergence $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ and $L^2$ Wasserstein distance $\mathcal{W}_2(\mu^*, \gamma_\tau)$ are bounded by $O(\beta)$ and $O(\sqrt{\tau D})$, respectively.
• The optimal control law for the McKean–Vlasov dynamics is the Föllmer drift, which solves the Schrödinger bridge problem and ensures the terminal distribution matches the minimizer $\mu^*$.
• The controlled McKean–Vlasov dynamics is governed by a coupled system of Fokker–Planck and Hamilton–Jacobi–Bellman PDEs, with the value function encoding the minimum free energy.
• SGD with $N$ neurons approximates the optimal mean-field dynamics over time $t \in [0,1]$ with high probability $1-\delta$, achieving an $L^2$ error bound of $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$.
• The approximation error scales favorably with $N$ and $\delta$, and can be made small by tuning the variance $\tau$ of the Gaussian prior and the regularization parameter $\beta$.
• The analysis provides a rigorous connection between lazy training, entropic regularization, and McKean–Vlasov dynamics, extending prior work on neural tangent kernels and mean-field limits.