Proximal Mean Field Learning in Shallow Neural Networks

TL;DR

This paper introduces a novel Sinkhorn regularized proximal algorithm for mean field learning in shallow neural networks, enabling meshless computation of parameter distributions.

Contribution

It develops a computational mean field learning algorithm using Wasserstein gradient flows, bridging theoretical insights with practical meshless implementation.

Findings

Effective in binary and multi-class classification tasks

Performs gradient descent on the risk functional's free energy

Enables meshless, particle-based computation of neural network dynamics

Abstract

We propose a custom learning algorithm for shallow over-parameterized neural networks, i.e., networks with single hidden layer having infinite width. The infinite width of the hidden layer serves as an abstraction for the over-parameterization. Building on the recent mean field interpretations of learning dynamics in shallow neural networks, we realize mean field learning as a computational algorithm, rather than as an analytical tool. Specifically, we design a Sinkhorn regularized proximal algorithm to approximate the distributional flow for the learning dynamics over weighted point clouds. In this setting, a contractive fixed point recursion computes the time-varying weights, numerically realizing the interacting Wasserstein gradient flow of the parameter distribution supported over the neuronal ensemble. An appealing aspect of the proposed algorithm is that the measure-valued recursions allow meshless computation. We demonstrate the proposed computational framework of interacting weighted particle evolution on binary and multi-class classification. Our algorithm performs gradient descent of the free energy associated with the risk functional.