Work

Implicit Regularization of Large Neural Networks via Mean-Field Formulation

Year: 2026

Type: preprint

Abstract: We propose a mathematical framework to explain implicit regularization from early stopping during the training of overparametrized neural networks. In the mean-field limit, the parameter distribution evolves according to a gradient flow on the space of probability measures. We show that these dynamics admit an equivalent McKean-Vlasov stochastic control formulation through the corresponding Hamilton-Jacobi-Bellman (HJB) equation. The control viewpoint yields a Dynamic Programming Principle (DPP), which we use to define a new metric on probability measures. This metric can be viewed as a mean-field generalization of the control representation of the Wasserstein-2 distance, and it naturally appears as a regularization term selected by early stopping. We further obtain non-asymptotic bounds describing how the induced regularization depends on the stopping time.

Source: arXiv (Cornell University)

DOI: 10.48550/arxiv.2603.20892

Authors: Acciaio, Beatrice, Heiss, Jakob, Pammer, Gudmund, Yan, Qinxin

Institutions: ETH Zurich, University of California, Berkeley, Graz University of Technology, Princeton University

Language: English

Topic: Stochastic Gradient Optimization Techniques

Subfield: Artificial Intelligence

Field: Computer Science

Domain: Physical Sciences

Open Access status: green

OpenAlex ID: W7140216258

Cites: 26

Cited by:

Related to: 10