View source for Shallow neural networks as Wasserstein gradient flows




==Motivation==
 
[https://en.wikipedia.org/wiki/Artificial_neural_network#:~:text=Artificial%20neural%20networks%20(ANNs)%2C,neurons%20in%20a%20biological%20brain.  Artificial neural networks] (ANNs) consist of layers of artificial "neurons" which take in information from the previous layer and output information to neurons in the next layer. Gradient descent is a common method for updating the weights of each neuron based on training data. While in practice every layer of a neural network has only finitely many neurons, it is beneficial to consider a neural network layer with infinitely many neurons, for the sake of developing a theory that explains how ANNs work. In particular, from this viewpoint the process of updating the neuron weights for a shallow neural network can be described by a Wasserstein gradient flow. 


==Single Layer Neural Networks==

{{See also [https://en.wikipedia.org/wiki/Mathematics_of_artificial_neural_networks Mathematics of Artificial Neural Networks]}}

explain, add figure

===Discrete Formulation===
Let us introduce the mathematical framework and notation for a neural network with a single hidden layer.<ref name="Figalli" />  Let <math> D \subset \mathbb{R}^d </math> be open . The set <math> D </math> represents the space of inputs into the network. There is some unknown function <math> f:D \rightarrow \mathbb{R} </math> which we would like to approximate. Let <math> N  \in \mathbb{N} </math> be the number of neurons in the hidden layer. Define

: <math> F_N : D \times \Omega^N \rightarrow \mathbb{R}^k </math>

be given by 
: <math> F_N(x, \omega_1, \dots, \omega_N,\theta_1, \dots, \theta_N) = \frac{1}{N} \sum_{i=1}^N \omega_i h(\theta_i,x) </math> 

where <math> h </math> is a fixed [https://en.wikipedia.org/wiki/Activation_function activation function] and <math> \Omega^N </math> is a space of possible parameters <math> (\omega, \theta)= (\omega_1, \dots, \omega_N,\theta_1, \dots, \theta_N) </math>. The goal is to use training data to repeatedly update the weights <math> \omega_i </math> and <math>\theta_i </math> based on how close <math> f_{N, \omega, \theta} := F_N( \cdot,  \omega_1, \dots, \omega_N,\theta_1, \dots, \theta_N) </math> is to the function <math> f </math>. More concretely, we want to find  <math> \omega, \theta </math> that minimizes the loss function:

: <math> l(f,f_{N, \omega, \theta}) := \frac{1}{2} \int_{D} |f(x)-f_{N,\omega,\theta}(x)|^2dx </math>

A standard way to choose and update the weights is to start with a random choice of weights <math> \bar{\omega}, \bar{\theta} </math> and perform gradient descent on these parameters. Unfortunately, this problem is non-convex, so the minimizer may not be achieved. It turns out in practice, neural networks are surprisingly good at finding the minimizer. A nicer minimization problem that may provide insight into how neural networks work is a neural network model with infinitely many neurons.


===Continuous Formulation===

For the continuous formulation (i.e. when <math> N = \infty </math>), we rephrase the above mathematical framework. In this case, it no longer makes sense to look for weights <math> \omega, \theta </math> that minimize the loss function. We instead look for a probability measure <math> \mu  \in \mathcal{P}(\Omega) </math> such that

: <math> f_{\mu}(x) := \int_{\Omega} \Phi(\xi, x) d \mu (\xi) </math>

minimizes the loss function:

:<math> F(\mu) : = \frac{1}{2} \int_{D} (f - f_\mu)^2 dx </math>.

Here <math>\Phi(\xi, x)</math> is an activation function with parameter <math> \xi = (\omega, \theta) \in \Omega </math>. 

Note that by restricting choices of <math> \mu </math> to probability measures of the form <math> \mu_N = \frac{1}{N} \sum_{i=1}^{N} \delta_{\xi_i} </math>, the above minimization problem generalizes to case with finitely many neurons as well. 

To avoid overfitting the network to the training data, a potential term is added the loss function. For the remainder of this article, we define the loss function <math> F</math> to be:

:<math> F(\mu) : = \frac{1}{2} \int_{D} (f - (\int_{\Omega} \Phi(\xi, x) d \mu (\xi)))^2 dx  +  \int_{\Omega} V(\xi)d\mu(\xi)</math>

for a convex potential function <math> V: \Omega \rightarrow \mathbb{R} </math>. Often we choose <math> V(\xi) = \frac{\lambda}{2} |\xi|^2 </math>. In fact, <math> F </math> is convex (along linear interpolations), in contrast to the minimization function in the finite neuron case. 

==Gradient Flow==

When <math> X \subseteq \mathbb{R}^n </math>, the gradient flow of a differentiable function <math> f: X \rightarrow \mathbb{R} </math> starting at a point <math> x_0 </math> is a curve <math>x(t): [0, T) \rightarrow X </math> satisfying the differential equation 

:<math> \frac{d}{dt}x(t) = - \nabla f (x(t)), x(0)=x_0 </math>.

where <math> \nabla f</math> is the gradient of f.<ref name="Schiebinger"/> 

Crucially, the gradient flow heads in the direction that decreases the value of <math> f </math> the fastest. We would like to use this nice property of gradient flow in our setting with the functional  <math> F </math>. However, it is not immediately straightforward how to do this, since <math> F </math> is defined on the space of probability measures, rather than on <math> \mathbb{R}^n </math>, so the usual gradient is not defined. Before we generalize the notion of gradient flow, note that

:<math> \nabla f (x(t)) = \nabla f \cdot \dot{x}(t) </math>.

Recall that for a vector field <math> V, \rm{div}V = - \nabla \cdot V </math>. 

Motivated by this and using the notion of a subdifferential we can define the [http://34.106.105.83/wiki/Gradient_flows_in_Hilbert_spaces gradient flow in a Hilbert space] <math> \mathcal{H} </math> of a convex and lower semi-continuous functional <math> \Phi </math>. An absolutely continuous curve <math>x: [0, \infty) \rightarrow \mathcal{H} </math> is a gradient flow for <math> \Phi </math> starting at <math> x_0 \in \mathcal{H} </math> if 

: <math> x(0)= x_0, </math> and <math> \dot{x}(t) \in -\partial \Phi(x(t)) </math> for almost every <math> t>0 </math>

where <math> \partial \Phi </math> is the subdifferential of <math> \Phi </math>. <ref name="Glaudo"/> 


===First Variation===

===Wasserstein Subdifferential===




===Wasserstein Distance===
Let us first define the (pth) Wasserstein distance between two probability measures. <ref name="Ambrosio"/> This can be defined for probability measures in any separable metric space. Let <math> \mu, \nu\in \mathcal{P}(\Omega) </math> and let <math> \Gamma(\mu, \nu) </math> denote the space of transport plans from <math> \mu </math> to <math> \nu </math>. Define the pth Wasserstein distance from <math> \mu </math> to <math> \nu </math> to be:

:<math> W_p(\mu,\nu) := \min \{ \int_{\Omega \times \Omega} d(\xi_1,\xi_2)^p d\gamma(\xi_1,\xi_2): \gamma \in \Gamma(\mu,\nu)\} </math>

where <math> d(\xi_1,\xi_2) </math> is the distance between  <math> \xi_1 </math> and <math> \xi_2 </math> in the metric space <math> \Omega \times \Omega </math>. In this context, <math> \Omega \times \Omega </math> is just a subset of <math> \mathbb{R}^{d} </math> for some <math> d </math> with the euclidean metric. Often <math> p = 2 </math>.


===Wasserstein Gradient Flow===

:<math> \partial_t \mu_t = \rm{div}( \mu_t \nabla \frac{\delta F}{\delta \mu} (\mu_t)) </math>

==Main Results==

===Consistency Between Infinite and Finite Cases===


==References==
<references>
<ref name="Figalli">[https://people.math.ethz.ch/~afigalli/lecture-notes-pdf/The-continuous-formulation-of-shallow-neural-networks-as-wasserstein-type-gradient-flows.pdf X. Fernandez-Real and A. Figalli, ''The Continuous Formulation of Shallow Neural Networks as Wasserstein-Type Gradient Flows'']</ref>

<ref name="Ambrosio">[http://www2.stat.duke.edu/~sayan/ambrosio.pdf L. Ambrosio, N. Gigli, G. Savaré, ''Gradient Flows in Metric Spaces and in the Space of Probability Measures'']</ref>

<ref name="Schiebinger"> [https://personal.math.ubc.ca/~geoff/courses/W2019T1/Lecture16.pdf G. Schiebinger, ''Gradient Flow in Wasserstein Space''] </ref>

<ref name="Glaudo"> [https://www.maa.org/press/maa-reviews/an-invitation-to-optimal-transport-wasserstein-distances-and-gradient-flows A. Figalli, F. Glaudo, ''An Invitation to Optimal Transport, Wasserstein Distances, and Gradient Flows''] </ref>

</references>