Proximal Mean Field Learning in Shallow Neural Networks
Authors: Alexis Teter
Iman Nodozi
Abhishek Halder
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (12/2023)
Proximal Mean Field Learning in Shallow Neural Networks
Alexis M.H. Teter Department of Applied Mathematics University of California, Santa Cruz
amicler@uuccc.ddv
Iman Nodozi
inodozi@uccc..o edu
Department Electrical and Computer Engineering University of California, Santa Crus
Abhishek Halder
ahalder edu
Department of Aerospace Engineering lowa State University
Reviewed OpenReview: pttps::/openreviite net/ forun? id= (d
Abstract
We propose a custom learning algorithm for shallow over-parameterized neural networks, ie., networks with single hidden layer having infinite width. The infinite width of the hidden layer serves as an abstraction for the over-parameterivation.. Building on the recent mean field interpretations of learning dynamics in shallow neural networks, we realize mean field learning computational algorithm, rather than design a Sinkhorn regularized proximal algorithm lo 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 Wasser- stein 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 inter- acting weighted particle evolution binary and multi-class classification. Our algorithm performs gradient cescent of the free energy associated with the risk functional.
analytical tool. Specifically,
1
Introduction
While universal function approximation theorems for neural networks have long been known (Cybenko 1989
Barron 1993 Hornik et al. 1989), such guarantees do not account for the dynamics of the learning algorit
Starting in 2018, several works
Mei et al. 2018 Chizat & Bach 2018 Rotskoff & Vanden- Eijnden 201 g:
Sirignnn & Spiliopoulos
2020 Rotskoff & Vanden-Eijnden 2022: Boursier et al. 2022) pointed out that
the first order learning dynamics for shallow (i.e., single hidden layer) neural networks in the infinite width
(i.e., over- parametteeeiiiiooon) limit leads to a nonlinear partial differential equation (PDE) that depends on a
pair of advection and interaction potentials.
The Cauchy initial value problem associated with the PDE describes the evolution of neuronal parameter ensemble induced by the learning dynamics. This result can be interpreted dynamical version of the universal approximation theorem. In particular, the potentials depend on both the loss function as well as the activation functions of the neural network.
The advection potential in this nonlinear PDE induces a drift, while the interaction potential induces a nonlocal force. Remarkably, this PDE can be interpreted as an infinite dimensional gradient flow of the population risk w.r.t. the Wasscrstein metric arising from the theory of optimal mass transport (Villani. 2009 2021

---

Published in Transactions on Machine Learning Research (12/2023)
The nonlocal nonlinear PDE interpretation makes connection with the so-called "propagation of chaos"-a term due to Kac (Kac 1956) that has grown into a substantial literature in statistical physics (McKean Jr. 1966 Sznitman 1991 Carmona & Delarue 2018). From this viewpoint, the first order algorithmic dynamics makes the individual ncurons in the hidden layer behave as interacting particles. These particle-love or microscopic interactions manifest as a population-level or macroscopic gradient flow.
As an analytic tool, the mean field Wasserstein gradient flow interpretation helps shed light on the conver- gence of first order learning dynamics in wet-parameteriied networks. In this work, we propose Wasserstein proximal recursions to realize the mean field learning dynamics meshless computational algorithm.
1.1 Computational challenges
Transcribing the mean field Wasserstein gradient flow PDE from an analytical tool to computational algorithm is particularly challenging in the neural network context. This is because the derivation of the PDE in (Mci ct al. 2018 Chizat & Bach 2018 Rotskoff & Handen-Fijnden 2018 Sirignano & Spiliopoulos 2020). and the corresponding infinite dimensional gradient descent interpretation, is an asymptolic consistency result. Specifically, the PDE describes the time evolution of the joint population density (or population measure in general) for the hidden layer neuronal ensemble. By definition, this interpretation is valid in the mean field (infinite width) limit of the single hidden layer. In other words, to leverage the gradient flow PDE perspective in computation, the number of neurons in the hidden layer must be large.
However, from numerical perspective, explicitly evolving the joint population in the large width regime is problematic. This is because the large width implies that the time-varying joint neuronal population distributions have high dimensional supports. Even though existing software tools routinely deploy stochas- tic gradient descent (SGD) algorithms at the microscopic (i.e., particle) level, is practically infeasible to estimate the time-narying population distributions using Monte Carlo or other posteriori function approx- imation algorithms near this limit, One also cannot resort to standard finite difference-type discretization approach for solving this gradient flow PDE because the large width limit brings the curse of dimensionality. Therefore, it is not obvious whether the mean field dynamics can lead to a learning algorithm in practice.
1.2 Related works
Beyond mcan field learning, Wasserstein gradient flows appear in many other scientific (Ambrosio et al. 2005
Santambrogio 2017) and engineering (Halder & Georgiou 2017
Caluya & Halder
el al. 2022
contexts. Thus, there exists a substantial literature on numerically implementing the Wasserstein gradient
flows - both with grid (Peyré] 2015 Benamou ct al. 2016 Carlier et al.) 2017 Carrillo et al. 2022 and
without grid (Liu et al., 2019 Caluya Halder 2019a Halder et al. 2020). The latter class of algorithums
are more relevant for the mean field learning context since the underlying parameter space (i.e., the support)
is high dimensional. The proximal recursion we consider is closely related to the forward or backward discretization (Salim et al. 2020 Frogner & Poggio 2020 of the
(JKO) scheme
(Jordan et al. 1998).
To bypass numerical optimization over the manifold of measures or densities, recent works (Mokrov et al. Alvarez-Melis et al.) 2021 Bunne et al.) 2022 propose using input convex neural networks {Amos et al. 2017) to perform the Wasserstein proximal time-stepping by learning the convex potentials (Brenier) 1991 associated with the respective pushforward maps. To alleviate the computational difficulties in high dimensions, Bonet et al. (2021) proposes replacing the Wasserstein distance with the Waseeddessssstninn distance [Rabin et al. 2011) scaled by the ambient dimension.
2021
Wc note here that there exists extensive literature on the mean field limit of learning in neural networks from other perspectives, including the Neural Tangent Kernel (NTK) and the Gaussian process viewpoints, see e.g., (Jacot et al. 2018 Lee et al. 2019 Novak et al. 2019 Xiao et al. 2018 Li & Nguyen 2018 Matthews et al. 2018). Roughly speaking, the key obscrvation is that the infinite width limit, the learning evolves as suitably defined Gaussian process with network architecture-depeneen kernel. We mention this in the passing since in this work, we will only focus on the Wasserstein gradient flow viewpoint. We point out that
2

---

Published in Transactions on Machine Learning Research (12/2023)
unlike the NTK, the mean field limit in Wasserstein gradient flow viewpoint does not approximate dynamical nonlinearity. In particular, the associated initial value problem involves nonlinear PDE (see (14)
1.3 Contributions
With respect to the related works referenced above, the main contribution of the present work is that we propose a meshless Wasserstein proximal algorithm that directly implements the macroscopic learning dynamics in a fully connected shallow network. We do so by evolving population densities as weighted scattered particles.
Different from Monte Carlo-type algorithms, the weight updates in our setting are done explicitly by solving regularized dual ascent. This computation occurs within the dashed box highlighted in Fig. [ The particles' location updates are done via nonlocal Fuler-Maruyama. These two updates interact with each other (sce Fig. 1 and together set up a discrete time-stepping scheme.
The discrete time-stepping procedure we propose, is a novel interacting particle system in the form of a meshless algorithm. Our contribution advances the state-of-the-art as it allows for evolving the neuronal population distribution in an online manner, as needed in mcan field lcarning. This in contrast posteriori function approximation in existing Monte Carlo methods (ef. Sec. 1.1]. Explicit proximal weight updates allows to bypass offline high dimensional function approximation, thereby realizing mean field learning at an algorithm level.
With respect to the computational challenges mentioned in Sec. 1.1 it is perhaps surprising that we are able to design an algorithm for explicitly evolving the population densities without directly discretizing the spatial domain of the underlying PDE. Our main idea to circunvent the computational difficulty is to solve the gradient flow PDE not as a PDE, but instead direct the algorithmic development for a proximal recursion associated with the gradient flow PDE. This allows to implement the associated proximal recursion over suitably discrete time without directly discretizing the parameter space (the latter is what makes the compulation otherwise problematic in the mean field regime).
For specificity, we illustrate the proposed framework on two numerical experiments involving binary and multi-class classification with quadratic risk. The proposed methodology should be of broad interest to other problems such as the policy optimization (Zhang al. 2018: Chu al. 2019 Zhang al., 2020 and
the adversarial learning (Domingo-Enrien Enrich al. 2020 Mroueh & Nguyen 2021 Lu 2023).
We emphasize that the perspective taken in this work is somewhat non-standard w.r.t. the existing literature in that our main intent is to explore the possibility of designing new class of algorithms by leveraging the connection between the mean field PDE and the Wasserstein proximal operator. This is a new line of idea for lcarning algorithm design that we show is feasible. As such, we do not aim to immediately surpass the carefully engineered existing state-of art in experiments. Instead, this study demonstrates a proof-of- concept which should inspire follow up works.
1.4 Notations and preliminaries
We use the standard convention where the boldfaced lowercase letters denote vectors, boldfaced uppercase letters denote matrices, and non-boldfaced letters denote scalars. We use the symbols V and A to denote the Euclidean gradient and Laplacian operators, respectively. In case of potential confusion, attach subscripts to these symbols to clarify the differential operator is taken w.r.t. which variable. The symbols @ Ø exp, and tanh denote elementwise multiplication, division, exponential, and hyperbolic tangent, respectively. Furthermore, rand and randn denote draw of the uniform and standard normal distributed random vector of appropriate dimension. In addition, represents a vector ones of appropriate dimension.
Let Z,, Ze S Rd. The squared 2- Wasserstein metric W2 (with standard Euclidean ground cost) between two probability measures mi (dzı) and #2(12)) (or between the corresponding densities when the measures are
3

---

Published in Transactions on Machine Learning Research (12/2023)
absolutely continuous), where 21 21, 22 Z2, is defined as
inf W3 (T) xEll(m). 22) x,x
| 22l12 dT (21
(1)
In 11, the symbol II (2) denotes the collection of joint measures (couplings) with finite second moments, whose first marginal is Ti, and the second marginal is T2.
1.5 Organization
The remainder of this paper is organized as follows. In Sec. provide the necessary background for the empirical risk minimization and for the corresponding mean field limit. The proposed proximal algorithm (including its derivation, convergence guarantee and implementation) is detailed in Sec. 3 We then report numerical case studies in Sec. and Sec. for binary and multi-class classifications, respectively. Sec. G concludes the paper. Supporting proofs and derivations are provided in Appendices results for a synthetic one dimensional case study of learning sinusoid using the proposed proximal algorithm is provided in Appendix D
and Numerical
2 From empirical risk minimization to proximal mean field learning
To motivate the mean field learning formulation, we start by discussing the more familiar empirical risk minimization set up. We then explain the infinite width limit for the same.
2.1 Empirical risk minimization
We consider a supervised learning problem where the dataset comprises of the features X C R"*,, and the labels y E y IR, i.e. the samples of the dataset tuples of the form
(x,v) X VCR" x R.
The objective of the supervised learning problem to find the parameter vector I2 such that y f(z 0) where is some function class parameterized by 0. In other words, f maps from the feature space X to the label space D. To this end, we consider shallow neural network with single hidden layer having nn neurons. Then, the parameterized function admits representation
f(x,0) = YAH Σ (x. 0:),
(2)
where & (x. == ((w, ()) bb)) for all E [mn] == (1,2,.... nH} and o()) is a smooth activation function. The parameters (g, w; and b, are the scaling, weights, and bias of the ęth hidden neuron, respectively, and together comprise the specific parameter realization 0 R2 i [лн].
We stack the parameter vectors of all hidden neurons
0 = (0,,0, 0e..0.....,0
0mn) E R
and consider minimizing the following quadratic loss:
1,.,z,0 = 1ly. f(x,
(1))((r)))
(3)
quadratic loss
We suppose that the training data follows the joint probability distribution 7, i.e. (). 7. Define the population risk R as the expected loss given by
R()) E(e,N)y My, 2v,
(4)
4

---

Published in Transactions on Machine Learning Research (12/2023)
In practice, y is unknown, approximate the population risk with the empirical risk
Bdata
R(f)
Σ
(5)
datt
where Idata is the number of data samples. Then, the supervised learning problem reduces to the empirical risk minimization problem
min R()-
(6)
ERR
Problem (6 is large but finite dimensional optimization problem that is nonconvex in the decision variable 0. The standard approach is to employ first or second order search algorithms such as the variants of SGD or ADAM (Kingma & Ba 2014
2.2 Mean field limit
The mean field limit concerns with continuum of hidden layer neuronal population by letting nh + Then, we view (2) as the empirical average associated with the ensemble average
MiManFielld := J2r
(
du(00
hidden neuronal population mass
where denotes the joint population measure supported on the hidden neuronal parameter space in Rr. Assuming the absolute continuity of for all times, we write du(0) = p(0)) where denotes the joint population density function (PDF).
Thus, the risk functional R, now viewed as a function of the joint PDF P. takes the form
y ( / (1 ())(
F(p) == R( R(fMeanFiedd (I. p)) = ((
Fu v(0)p(0)dd
( (().0)p()pp(0)ppp))
(8)
where
Fo = ((r.y) [v]]]
V(0): = E (x.v) 2 0)], U 0) = [a (0 (x,
(9)
Therefore, the supervised learning problem, in this mean field limit, becomes an infinite dimensional varia- tional problem:
min F(P)
(10)
where is a sum of (hree functionals. The first summand Fo is independent of p. The second summand is a potential energy given by expected value "drift potential" V and is linear in p. The last summand is a bilinear interaction energy involving an "intcraction potential" U and is nonlincar p.
The main result. in Mei al. (2018) flow of the functional F w.r.t. the 2- Wasserstein metric We, i.e., the mean field learning dynamics results in joint PDF trajectory p(t, 0). Then, the minimizer in 10 can be obtained from the large limit of the following nonlinear PDE:
that using first order SGD learning dynamics induces gradient
3p
(W2F(())
(11)
at
where the 2- 2-Wasserstein gradient
(Villani 2021 Ch. 9.1)
Ambrosio et al. 2005 Ch. 8) of F is
(pv
W2 F(p) : -V
0p
and A denotes the functional derivative w.r.t. P.
5

---

Published in Transactions on Machine Learning Research (12/2023)
In particular, Mei et al. (2018) considered the regularized risk functional
J plog p d0,
+
B 0,
(12)
by adding a strictly convex regularizer (scaled negative entropy) to the unregularized risk F. In that case, the sample path dynamics corresponding the macroscopic dynamics (11b preciscly becomes the noisy SGD:
dl =- ( + J U(0.0)p(0)dB ) dt + /23-1 dn, O(t = 0) Po
(13)
where 7 is the standard Wiener process in R2, and the random initial condition A(t = 0) follows the law of suitable PDF Po supported over IRP.
In this regularized case,
I1 results in the following nonlinear PDE initial value problem (IVP):
Pp 06
Vo (o (vi0)+ J. ((0.0)p(0)d)) ) BB--DOP. p(t 0,0) Po-
(14)
In other words, the noisy SGD induces evolution of a PDF-valued trajectory governed by the advection, nonlocal interaction, and diffusion the latter originating from regularization. Notice that a large value of B 0 implies small entropic regularization (12), hence a small additive process noise in 13 and consequently, a small diffusion term in the PDE 14p.
The regularized risk functional F3 in 12 can be interpreted of the advection potential energy and interaction energy. The term 3 JRx plog pdl contributes internal energy due to the noisy fluctuations induced by the additive Wiener process noise d2 13
free energy wherein F contributes sum
In Mei et al. (2018). asymptotic guarantees were obtained for the solution of (14) to converge the minimizer of Fg. Our idea, outlined next, is to solve the minimization of F3 using measure-valued proximal recursions.
2.3 Proximal mean field learning
For numerically computing the solution of the PDE IVP (14), we propose proximal recursions over P2 (RR), defined as the manifold of joint PDFss supported over IR having finite second moments. Symbolically,
P2 (R") = Lebesgue integrable over I2 p 0, / pdo / 0 0 d0 < ox)
Proximal updates generalize the concept. of gradient steps, and are of significant interest in both finite and infinite dimensional optimization (Rockafellar 1976a b Bauschke et al. 2011 Teboulle 1992
Bertsekas
et al. 2011 Parikh et al. 2014). For given input, these updates take the form of a structured optimization problem:
proximal update
dist (decision variable, input) time step functional (decision variable)
arg inf decision variaille
(15)
for some suitable notion of distance dist (--+) on the space of decision variables, and some associated func- tional. It is usual view 15) as an operator mapping input + proximal update, thus motivating the term prozimal operator.
The connection between 15 and the gradient flow comes from recursively evaluating 115 with some initial choice for the input. For suitably designed pair (dist (,,-), functional), in the small time step limit, the sequence of proximal updates generated by (15) converge to the infimmm of the functional. In other words, the gradient descent of the functional w.r.t. dist may be computed as the fixed point of the proximal operator
The parameter > is referred inverse temperature.

---

Published in Transactions on Machine Learning Research (12/2023)
For a parallel between gradient descent and proximal recursions in finite and infinite dimensional gradient descent, see e.g., (Halder & Georgiou 2017 Sec. 1). Infinite dimensional proximal recursions over the manifold of PDI's have recently appeared in uncertainty propagation (Caluya & Halder 2019b Halder
15
ct al., 2022), stochastic filtering (Halder & Georgiou 2017
2018
2019 and stochastic optimal control
(Caluya & Halder a))
In our context, the decision variable p P2 (IP)) and the distance metric dist = W2. Specifically, we propose recursions over discrete time f4-1 (A: 1)h where the index i N, and h >> 0 is constant time step-
size. Leveraging that (14) describes gradient flow of the functional Fg w.r.t. the W2 distance metric, the associated proximal recursion is of the form
(W2 (e. 2h 1 h }
2k = -prxx
Wa (0k-1) = arg inf (EP2(
(16)
W2 where (x-1(-) ((- and @0 Po. The notation prox (ek-1) can be parsed as "the proximal operator of the scaled functional hFg w.r.t. the distance W2, acting on the input Øk E Pa ".") Our idea is evaluate the recursion in the small limit, Le., for h 0.
To account for the nonconvex bilinear term appearing in 12 following (Benamou et al. 2016 Sec. 4), we employ the approximation:
J J (10,0)g())e()ddddd = 010.0)g(0een Ak E N.
gg
We refer to the resulting approximation of F3 as Fsx i.e.,
Pe (e. (x.ex=1) = / Fo (Fo+V V + (f (x0.0)en--10)l + + loo e(0) e(0)do.
Notice in particular that Fg depends on both e @k-1, e N. Consequently, this approximation results in a semi-implee semi-implicit variant of [66.. given by
(W₂ (e. (k-1)) (e. }
2k = arg inf ((P)
(17)
We have the following consistency guarantee, stated informally, amoug the solution of the PDE IVP 14 and that of the variational recursions (17). The rigorous statement and proof are provided in Appendix A
Theorem 1. (Informal) Consider the regularized risk functional the small time step (h 4 0) limit, the prozimal updates
12 wherein is given by 17 with Po Po converge to the solution for the
PDE IVP (4).
We next detail the proposed algorithmic approach to numerically solve
17
3 ProxLearn: proposed proximal algorithm
The overall workflow of our proposed proximal mean field learning framework is shown in Fig, [ We generate N samples from the known initial joint PDF @o and store them weighted point cloud {0) - i- Here, el : do (0)) for all i E [N]. In other words, the weights of the samples are the joint PDF values evaluated at those samples.
For each & E N, the weighted point clouds (f N are updated through the two-step process outlined in our proposed Algorithm [ referred to as PROXLEARN. At a high level, lines 9-18 in Algorithm 1 perform nonlincar block co-ordinate recursion on internally defined vectors 2. q whose converged values yield the proximal update (line 19). We next explain where these recursions come from detailing both the derivation of PROXLEARN and its convergence guarantee.
7

---

Published in Transactions on Machine Learning Research (12/2023)
3.1 Derivation of ProxLearn
The main idea behind our derivation that follows, is to regularize and dualize the discrete version of the optimization problem in [17). This allows us to leverage certain structure of the optimal solution that emerges from the first order conditions for optimality, which in turn helps design a custom numerical recursion.
Specifically, to derive the recursion given PROXLEARN, we first write the discrete version of 17 as
2k arg min min 1 () h (UR-1 + 1-- 8- log loge.)) }. KEN, (18)
where
II (@x-1, = MM E M 2 0 (clementwise). M1 2k-1, M 1 = e)
(19)
Uk-1 = (0x-1),
(20)
1 U 0 00--110x-1)-
(21)
and Ck E RNNN denotes the squared Euclidean distance matrix, i.e.
Cx(i.)) = 10 0 1 v(ij) E [N] [N].
We next follow a "rogulariza-thendnnalll approach. In particular, we regularize regularization H (M) = (M, log M), and write
18 by adding the entropic
1 (C M) + (M + (UR-1 Un 10k-1 B log logeee)
Q: arg min
min
KEN (22)
Metl(e 1:e) 2
where 0 is regularization parameter.
Following Karlsson & Ringh
2017 Lemma 3.5), Caluya & Halder 2019aa Sec. III,, the Lagrange dual
problem associated with 22 is
pop (qpt)
arg max ARRNE
(A. -f
(->1) - (exp h/ee) exp ( (-Cx/2e) exp (Aihffe)
(23)
where
= sup(((()p(())) -))())
(24)
is the Legendre-Fenchel conjugate of the free energy Fg in 17 and the optimal coupling matrix Mopt
poopt (1, (2 has the Sinkhorn form
mpt (i, = cxp (An(1))/c)))))) exp (- exp (A1 (t)/c)..
(25)
To solve 23
considering
12 we write the "discrete free energy" as
F((e) = (Ux-1 RK-1R--1 + B log 0.
(26)
Its Legendre-Fencee Fenchel conjugate, by 24 is
FG == sup (X - 1 _ 10- e T U
14- g3- log
(27)
Setting the gradient of the objective in 27 w.r.t. @ zero, and solving for e gives the maximizer
2max exp (3 (A Uz-1 8-1 U -x-
(28)
8

---

Published in Transactions on Machine Learning Research (12/2023)
Substituting 28
back into 27
obtain
(A) = 3 exp (3 (A 1 vk-1 UR U 1).
(29)
Fixing Ao and laking the gradient of the objective in 23 w.r.t. X1 gives
exp (Aih/c) O cxp Cx)) exp (Aph/e)) = exp(-3ure
3U 3-K-1 1 (exp (A,n/c))
(30)
Likewise, lixing AI, and taking the gradient of the objective in [23 w.r.t. Ao gives
exp (Aoh/c) o (exp (k/2c)cog) exp (Ath/c)) =
(31)
Next, letting I& exp (-Ck/2e) q exp Anoh/)) 2 exp (Anh//)), and Ex-1 exp(-Buk-1
BU k-10k-1 1), express (30)
2 2 gT - 11000 x
(32)
and 31 as
q 3 (Tx2) -
(33)
Finally using 19
we obtain
2: (M
Σ m opt (1,1) z(1) Σ (1
(34)
In summary, 334 allows us to numerically perform the proximal update.
Remark Note that in Algorithm (t.e., PROXLEARN) presented in Sec. 3 the lines 11, 12, and 19 respectively.
correspond to 32 33 and 34
3.2
Convergence of ProxLearn
Our next result provides the convergence guarantee for our proposed PROXLEARN algorithm derived in Sec.
3
Proposition 1. The recursions given in lines 7-18 in Algorithm | PROXLEERN, converge to a unique faced point pt 2pp) E R x R Consequently, the prozimal update 34 (t.e., the evaluation at line 19 in Algorithm 7 is unigue.
Proof. Notice that the mappings 2(,,,))
2(, given lines and 12 Algorithm
are cone preserving since these mappings preserve the product orthant RS By This is direct consequence of the definition of q. in terms of Ag A
Now the idea is to show that the recursions in lincs 11 and 12 in Algorithm L as composite nonlincar maps, are in fact ractive w.r.t. a suitable metric on this cone. Following Caluya & Halder (2019a Theorem 3); the z iteration given in line 11 in Algorithm PROXLEARN, for L = 1,2, is strictly contractive the Thompson's part metric (Thompson, 1963) and thanks to the Banach contraction mapping theorem, converges a unique fixed point np E P.50
We note that our definition of &x-1 is slightly different compared to the same in Caluya & Halder 2019a Theorem 3), but this does not affect the proof. From definition of Ck, have Cs E [0, oo) which implies ( E (0, 1]. Therefore, Г₁ is a positive linear map for each & E N. Thus, by (linear) Cerron-Frobeniug theorem, the linear maps Tş are contractive. Consequently the iterates also converge to unique fixed point got E E 50
Since converged pair (g opt opt) E N
N
unique, so the proximal update 34 ie., the evaluation
line 19 in Algorithm D
We next discuss the implementation details for the proposed PROXLEARN algorithm.

---

Published in Transactions on Machine Learning Research (12/2023)
1 step delay
k-1
Euler-Maruyama scheme
W2 prox FFF
Ck-1
Qk
step delay
Figure Schematic of the proposed proximal algorithm for mean field learning, updating scattered point cloud {0k-1) e N for k E N. The location of the points { updated via the Euler-Maruyama scheme; the corresponding probability weights are computed via proximal updates highlighted within the dashed box. Explicitly performing the proximal updates via the proposed algorithm, and thereby enabling mean filed learning as an interacting weighted particle system, is our novel contribution
3.3 Implementation of ProxLearn
We start by emphasizing that PROXLEARN updates both the parameter sample locations 0f and the joint PDF values of evaluated those locations, without gridding the parameter space. In particular, the PDF values updated online, not as offline posteriorí function approximation as in traditional Monte Carlo algorithms.
We will apply PROXLEARN, as outlined here, in Sec. In Sec. 5 PROXLEARN to showcase its flexibility.
we will detail additional modifications of
Required inputs of PROXLEARN are the inverse temperature 8, the time step-size h, a regularization parameter 8, and the number of samples N. Additional required inputs are the training feature data [a and the corresponding training labels well the weighted point cloud (0 1 for each ho E N. Furthermore, PROXLEARN requires two internal parameters as user input: the numerical tolerance 6, and the maximum number iterations L.
X = Enaua
[y Brasal E 18
For N, let
0 of
0
NRP
Q =
E
RRN
In line 2 of Algorithm 1 PROXLBARN updates the locations of the parameter vector samples 0f in TR via Algorithm 2 EULERMARUYAMA. This location update takes the form:
0 =0=00- AV (V (0-z) (0)-) 1)) + 2 (ni m
(35)
10

---

Published in Transactions on Machine Learning Research (12/2023)
Algorithm Proximal Algorithm
procedure PROXLEARN 0 h, e, N, N.
2:
1k-1, 1- 0 +- RULPRMARRIAMMA(M e, 1, U. @x-1) Update the location of the samples
Ck(i,))
0.' 0
TA exp( C2/2e)
b Lines 4-8 define the terms needed in re-expressing as 32
fk- t exp(-Bur.-
3U K- -
20 randNx1
2 < 2o, ON)
8:
F - 2 (I ( ONX(L-1)
9:
L ==
10:
while & do
11:
z(;, 11) + (fx-1 2 (T ((: Fहे.
4 Following
32
12:
q(: ( 1) - Q @ (Taz(; (1)
D. Following
33
13:
if (g (1,,++10 1) (i. < and [ l 1) - -((;0) < then
14:
Break
15:
else
16:
++1+1
17:
end if
18:
end while
19:
return @s < +x;; () @ (TZ(()
b. Use 34 to map 0k- to &
20: end procedure
A Algorithm Euler-Maruyama Algorithm
procedure BULERMARUYAMA(H. 0,--,, 1k-1)
Pk-1 <- ((0
Lines 2-4 construct the argument of the gradient in 35
Uş-1 t- 1/niasa k- PR-I k-1
U;-1 +- -k-10k-1
5: Uk-1 +- -2/ndata PE: PR-19
6: D + BACKWARD (u Vk-1)
> Approximate the gradient of 35 using PyTorch library
BACKWARD Passke et al. 2017
7:
G+ 2h/3 x randn
O 6 0 D G
D. Complete the location update via.
end procedure
where w(.) = ((0 and vi == ny (k: 1)4), vk E N.
To perform this update, EULERMARUYAMA constructs a matrix Pk-1 whose (i, (ttt element P (2x, 0 1)))) From Pk-1; we construct U:-1 and U as in lines and 5. In line 6 of Algorithm EULERMARUYAMA uses the automatic differentiation module of PyTorch Library, BACKWARD Paszke 2017,, to calculate the gradients needed to update θ O₁ Wk N.
Once θk, U4-1, and - have been constructed via BULERMARUYAMA, PROXLEARN maps the x 1 vector 0k-1 to the proximal update @k-
We next illustrate the implementation of PROXLEARN for binary and multi-class classification case studies. A GitHub repository containing our code for the implementation of these applications can be found -aa-immm-aaa----l--eeeaaaiiiiiie therein for outline of the structure of our code.
httpp://githul:
git Please refer the Readme file
11

---

Published in Transactions on Machine Learning Research (12/2023)
4 Case study: binary classification
In this Section, we report numerical results for our first case study, where we apply the proposed PROXLEARN algorithm for binary classification.
For this case study, we perform two implementations on different computing platforms. Our first implemen- tation on PC with 3.4 GHz 6-Core Intel Core i5 processor, and GB RAM. For runtimo improvement, we then use Jetson TX2 with a NVIDIA Pascal GPU with 256 CUDA cores, 64-bit NVIDIA Denver and ARM Cortex-A57 CPUs.
4.1 WDBC data set
We apply the proposed algorithm perform binary classification on the Wisconsin Diagnostic Breast Cancer (henceforth, WDBC) data available at the UC Irvine machine learning repository (Dua & Graff 2017 This data set consists of the data of scans from 569 patients. There n 30 features from each scan. Scans are classified as "benign" (which we label as -1) or "malignant" (labeled as +1).
In 2 we define & 0-1) == - lanh((wf x) + b))1) Vi E [N] after (k (k-1) 1) updates. The parameters @k -k-1: w/ and b are the scaling, weight and bias of the 7th sample after (k 1)
1) updates, respectively. Letting p dr the parameter vector of the 8th sample after (I- - 1) updates is
d
0 0k ==
bi
k-1 ERP, ViE
w
We set en = ([.,9 x 1-0.1,0.1] x [-1,1]"*), a uniform joint. PDF supported over np = fl + 2 32 dimensional mean field parameter space.
We use 70% of the entire data set training data. As discussed in Soc. 3 distribution via weighted scaltered point cloud evolution using PROXLEARN. We then the confusion matrix method (Visa et al. 2011) to evaluate the accuracy of the obtained model over the test data, which is the remaining 30% of the full data set, containing Thess points.
we learn the mean field paramcter
For each test point Feest E R we construct
+ Frest 0111
T [ Frest O
1
(p(ttcst):
=
E CRN
+ (Erest ON
where 0i-1 1 obtained from the training process. We estimate fMeanField in (7) in two ways. First, we estimate fMeanField as a sample average the elements of sp Second, we estimate fMeanField by numerically approximating the integral in ( using the propagated samples (0) 1 for k E N. We refer these as the "unweighted estimate" and "weighted estimate," respectively. While the first estimate is empirical average, the second uses the weights { obtained from the proposed proximal algorithm. The MMeanField "unweighted estimate" and "weighted estimate" are then passed through the SOFTMAX and ARGMAX functions respectively, to produce the predicted labels.
4.2 Numerical experiments
We set the number of samples N = 1000, numerical tolerance 5 = 10 10-3, the maximum number of iterations L = 300, and the regularizing parameter e = 1. Additionally, we set the time step to h: = 10-3 We run the simulation for different values the inverse temperature 3, and list the corresponding classification accuracy Table The "weighted estimato," " the ensemble average using proximal updates, produces more accurate results, whereas the "unweighted estimate," " the empirical average, is found be more sensitive the inverse temperature B.
12

---

Published in Transactions on Machine Learning Research (12/2023)
Unwwighned
Wighhed
Unweighted
Weighted
3xx 1x10
Tbsk functionnl Fz
10'0
Rish funetooal F
10-7
160
10
102
10P
10
107
10
10F
107
x 10"
10
10'2
10 107
107
10"
100
107
105
100
102
100
Irerations
Iierations
(a) Regularized risk calculated CPU.
(b) Regularized risk calculated via GPU.
Figure 2: The solid line shows the regularized risk functional Fa versus the number of proximal recursions shown for the WDBC dataset with B 0.05. The shadow shows the Fg variation range for different values of B E (0.03, 0.05, 0.07).
Table 1: Classification accuracy of the proposed computational framework for the WDBC Dataset
B
Unweighted
Weighted
0.03
91.17%
92.35%
0.05
92.94%
92.94%
0.07
78.23%
92.94%
For each fixed 8, we perform 106 proximal recursions incurring approx. hours of computational time. Fig. 2a shows the risk functional, computed the averaged loss over the test data using each of the two estimates described above.
To improve the runtime of our algorithm, we run our code Jetson TX2 module, converting data and variables to PyTorch variables.
We begin calculations in Float32, switching to Float64 only when needed to avoid not-a-number (NaN) errors. This switch typically occurs after 2 x 10 to 10 iterations. As shown Table 2 we train the neural network to comparable accuracy in only 2.5 x 10 itcrations. The new runtime around 6% of the original runtime for the CPU-based computation. Fig. 2b shows the risk functional calculated via this updated code. We parallelize these calculations, taking advantage of the GPU capacity of the Jetson TX2.
We utilize these improvements in runtime to additionally experiment with our choice of E. Table B reports the final values of the regularized risk functional Fg and corresponding runtimes for varying e. As expected, larger E entails more smoothing and lowers runtime. The corresponding final regularized risk values show no significant variations, suggesting numerical stability.
4.3 Computational complexity
In this case study, we determine the computational complexity of PROXLEARN as follows. Letting a:- := (ak-1 ... N b = (b/_1 b) WK-1 == (wh_ ... w create the N x Aldat
matrix
Pk- :: (otnn1")
bb
(36)
which has complexity O(ndaaaN The subsequent creation of matrix UK-1 in line 3 of Algorithm has complexity O(ddara N2), and creating Uk-1 and u takes complexity O O(Nrdata) and O(N2) respectively.
13

---

Published in Transactions on Machine Learning Research (12/2023)
Table 2: Classification accuracy on Jetson Tx2, after 2.5 x 10 iterations
B
Unweighted
Weighted
Runtime (hr)
0.03
91.18%
91.18%
1.415
0.05
91.18%
92.94%
1.533
0.07
90.59%
91.76%
1.704
Table 3: Comparing final Fg and runtimes for various E
Unweighted Final F3
Runtime (s)
0.1
1.4348 x 10-2
32863
0.5
1.3740 x 10-2
11026
1.0412 x 10-2
5022
1.1293 x 10
4731
10
9.8849 10
4766
The complexity in caleulating the relevant derivatives of Ux- and uu is Mritallt)) (these derivatives are calculated in Appendices B and C). Updating O using these results has complexity of O(Np) = O(Nn.). Therefore, the process of updating k via EULERMARUYAMA is D(N
The significant complexity the remainder of PROXLEARN arises in the construction of matrix Cs line and the matrix-vector multiplications within the while loop in lines 11, 12.
The creation of the N N matrix Ck. in which each element is the vector norm of ng vector, is N²). In worst-case scenario, the while loop runs L times. The operations of leading complexity within the while loop are the multiplications of the D& matrix of size N x N with the N 1 vectors, which have a complexity of O(N2). Therefore, the while loop has complexity of O(LN2).
Updating o thus has a complexity of O(n L) In practice, the while loop typically ends far before reaching the maximum number of iterations
From this analysis, find that the overall complexity of PROXLEARN is O (Mdataaa L)). In comparison, the per iteration complexity for JKO-ICNN is O (Nimne (N + 1) Nbatch + Nº) where Niner denotes the number of inner optimization steps, and Nbathh denotes the batch sizc. The per iteration complexity for SWGF O (Ninner-Npro, Nbatch log Nbatch) where Nproj denotes the number of projections to approximate the sliced Wasserstein distance.
4.4 Comparisons to existing results
From Fig. we observe that there is a significant burn in period for the risk functional curves. These trends in learning curves agree with those reported in (Mei et al. 2018). In particular, (Mei et al. 2018 Fig. 3) and Fig. 7.3 in that reference's Supporting Information, show convergence trends very similar to our Fig. 2 slow decay until approx. 105 iterations and then a significant speed up. The unusual convergence trend was explicitly noted in (Mei et al.] 2018): "We observe that SGD converges to network with very small risk, but this convergence has nontrivial structure and presents long flat regions". It is interesting note that (Mei et al.) 2018) considered experiment that allowed rotational symmetry and simulated the radial (i.e., with one spatial dimension) discretized PDE, while we used the proposed proximal recursion directly in the neuron population ensemble to solve the PDE IVP, i.c., similar convergence trends were observed using different numerical methods applied to the same mean field PDE IVP. This makes us speculate that the convergence trend is specific to the mean field PDE dynamics itself, and is less about the particular numerical algorithm. This observation is consistent with recent works such as (Wojtowyssch & Weinan 2020) which investigate the mean field learning dynamics in two layer ReLU networks and in that setting, show that the learning can indeed be slow depending on the target. function.
14

---

Published in Transactions on Machine Learning Research (12/2023)
Table 4: Comparison of average classification accuracy from Bonet et al. (0022 Table 1) to our algorithm,
ProxLearn
Dataset
JKO-ICNN
SWGF + RealNVP
ProxLearn, Weighted
ProxLearn, Unweighted
Banana
0.550 + 10-2
0.559 + 10-
0.551 + 10-2
0.535 = 5- 10-2
Diabetes
0.777 7 10
0.778 = 2- 10-3
0.736 + 2 10
0.731 = 10
Twonorm
0.981 22
0.981 6- 10
0.972 + 2 10-3
0.972 2- 10
As a first study, our numerical results achieve reasonable classification accuracy compared to the state-of-art, even though proposed meshless proximal algorithm is very different from the existing implementations. We next compare the numerical performance with existing methods as in Mokrov et al. (2021) and Bonet
et al. 2022)) We apply our proximal algorithm for binary classification to three datasets also considered in Mokrov et al. (2021) and Bonet et al. (0022): ; the banana, diabetes, and twonorm datasets.
The banana dataset consists of 5300 data points, each with ng = features, which we rescale to lie between and 8. We set 3 = 0.05, draw our initial weights w from Unif (1-2 and bias b from (n-O( 0.3)), and set Po Unif (0, 1000). We run our code for 3500 iterations, splitting the data evenly between test and training data.
The diabetes dataset consists of n, 8 features from each of 768 patients. Based on our experimental results, we make the following adjustments to our algorithm: we redefine B = 0.65, Po = Unif (0, 1000), and draw our initial weights w from Unif ([-2, 2)" We rescale the data to lie between 0 and 1, and use half of the dataset for training purposes, and the remainder as test data. In this case, we run our code for 4.99 x 10 iterations.
The twonorm datasel consists of 7,400 samples drawn from two different normal distributions, with n, 20 features. We again consider 50% of the same as training data and used the remaining 50% as test data, and rescale the given data by factor of 8. Based on our empirical observations, we redefine 3 = 1.95, and once more draw our initial weights w from Unif Uni 2)² and set @u = Unif (0, 1000). In this case, we perform 10a proximal recursions in each separate run.
We run our code five times for each of the three datasots under considcration and compute "unweighted estimates" and "weighted estimates" in each case, as described above. These estimates assign each data point a value: negative values predict the label as 0, while positive values predict the label as 1. From these results, calculate the weighted and unweighted accuracy by finding the percentage of predicted test labels that match the actual test labels. The average accuracy over all five runs is reported in Table [ alongside the results reported Bonet et al. 2022
Table 1). We achieve comparable accuracy to these recent results.
5 Case study: multi-class classification
We next apply the proposed proximal algorithm to ten-class classification problem using the Semeion Handwritten Digit (hereafter SHD) Data Set (Dua & Graff 2017 This numerical experiment is performed on the aforementioned Jetson TX2.
5.1 SHD data set
The SHD Data Set consists of 1593 handwritten digits. By viewing each digit as 16 16 pixel image, cach image represented by ny 162 256 features. Each feature is boolean value indicating whether a particular pixel is filled. We subsequently re-scale these features such that 2; 1-
15

---

Published in Transactions on Machine Learning Research (12/2023)
5.2 Adaptations to ProxLearn for multi-class case
To apply PROXLEARN for multi-class case, make several adaplations. For instance, rather than al- tempting to determine f(x) ~ y, we redefine f(x) to represent a mapping of imput data to the predicted likelihood of the correct label. We therefore redefinc the variables, paramcters, and risk function as follows.
Each label is represented by a x 10 vector of booleans, stored a Adat x 10 matrix Y where Y,, = if the ith data point z, has label j. and Yt, = 0 otherwise.
We construct the N X Adtaa matrix Pk-1 by defining the (J. element of Pk-1 as
P Pk-1(j.1) = ((,,),,
(37)
where (: denotes the standard Euclidean inner product. The SOFTMAX function in 37 produces 10 1 vector of non-negative entries that sum to 1. By taking the inner product of this vector with the Boolean vector :i,;), we define Px-1(j, Φ ((i,)) Yi,:)) as the perceived likelihood that the data point is labeled correctly by sample j. As our model improves, this value approaches 1, which causes the probability of an incorrect label to drop.
As this newly defined Pk-1 does not call for bias or scaling, the weights alone are stored in the N x p matrix O--1- In this case, p= 10 as each of the ny features requires a distinct weight for each of the ten labels. For convenience, we reshape θR- = (0) - ..... 0 1), where each 0 is 10 x ny matrix. Therefore, 0x-1 is 10 X n2 N tensor.
We redefine the unregularized risk to reflect our new & as follows:
F(p) = E( w)
(000).)))00000 (000))
F y.
(38)
Expanding the above, arrive at form that resembles
g now with the following adjusted definitions:
Fu == 1, V(0) == E( - 20 x. v)).
(39)
(40)
4 U (0,0) == E(z.y) [0 2b x. 8)].
(41)
We use the regularized risk functional F3 as in (12) where F' now is given by (38 Due the described changes in the structure of O the creation of Cx in line 3 of PROXLEARN results in 10 N x N tensor, which we then sum along the ten element axis, returning N N matrix.
Finally, we add a scaling in line 7 of EULERMARUYAMA, scaling the noise by a factor of 1/100.
5.3 Numerical experiments
With the adaptations mentioned above, we set the inverse temperature B = 0.5, e = 10, the step size h = 10-3, and N = 100. We draw the initial weights from Unif ([-1,1] 10===) ). We take the first Pliata 1000 images as training data, reserving the remaining Nttst = 593 images as test data, and execute 30 independent runs of our code, cach for 10 proximal recursions.
To evaluate the training process, we create matrix - Nxmea using test. data rather than training data, but otherwise defined as in (37). We then calculate weighted approximation of Fg:
Fg 2
1
test
(42)
Ptest
and an unweighted approximation:
1
F8 ~
1
Prest
(43)
Moest
N
16

---

Published in Transactions on Machine Learning Research (12/2023)
Unweighted
Weighted
1x
x 10
Risk functional
6x
x 10
3 10
10F
10'
10F
102
10"
10'
10P
Iverations
Figure 3: The solid line shows the average regularized risk functional Fg versus the number of proximal recursions shown for the Semeion dataset with 3 0.5. The narrow shadow shows the Fg variation range with the same 3 using the results of 30 independent runs, each starting from the same initial point cloud (0), eb
which we use to produce the risk and weighted risk log-log plots shown in Fig.
Notably, despite the new activation function and the adaptations described above, our algorithm produces similar risk plots in the binary and multi-class cases. The run time for 10t iterations is approximately 5.3 hours.
To evaluate our multi-class model, calculate the percentage of accurately labeled test data by first laking ARGMAX (X 0x along the ten dimensional axis, to determine the predicted labels for each test data point using each sample of - We then compare these predicted labels with the actual labels. We achieved 61.079% accuracy for the test data, and 75.330% accuracy for the training data.
5.4 Updated computational complexity of ProxLearn
In the binary case, the creation of matrix C₁ requires O(n,NV) flops. In the case of multi-class classification concerning m classes, C is redefined as the sum of m such matrices. Therefore, crcating the updated matrix Cs takes (m),,N2)) flops. Thus, updating @k-1 is of complexity O(mm, + L)N2). The complexity of EULERMARUYAMA can be generalized from the discussion in Sec.
6
Conclusions
6.1
Summary
This work presents proximal mean field learning algorithm to train shallow neural network in the over- parameterized regime. The proposed algorithm is meshless, non-parametric and implements the Wasserstein proximal recursions realizing the gradient descent of entropic-regularize risk. Numerical case studies in binary and multi-class classification demonstrate that the idcas of mean field learning can be attractive computational framework, beyond purely theoretical interests. Our contribution should be of interest other learning and variational inference tasks such the policy optimization and adversarial learning.
17

---

Published in Transactions on Machine Learning Research (12/2023)
We clarify that the proposed algorithm is specifically designed for neural network with single hidden layer in large width regime. For multi-hidden layer neural networks, the mean field limit in the sense of width as pursued here, is relatively less explored for the training dynamies. In the multiple hidden layer setting, thooretical undorstanding of the limits a fronticr of current research; sec c.g., Fang et al. (2021); Sirignanoo & Spiliopoulos (2022). Extending these ideas to design variants of proximal algorithms requires new lines of thought, and as such, is out-of-scope of this paper. In the following, we outline the scope for such future work.
6.2 Scope for future work
Existing efforts to generalize the tcal results for the mean field limit in this work, from single to multi-hidden layer networks, have been pursued in two different limiting sense. One line of investigations [Sirignano & Spiliopoulos) 2020 2022 Yu & Spiliopoulos 2023 take the infinite width limit one hidden layer at a time while holding the (variable) widths of other hidden layers fixed. More precisely, if the ith hidden layer has N; neurons, then the limit is taken by first normalizing that layer's output by N for some fixed T [1/2, 1] and then letting N; + oo for an index while holding other Nj's fixed and finite, +
A different line of investigations (Araújo et al. 2019 Nguyen 2019) consider the limit where the widths of all hidden layers simultancously go to infinity. In this setting, the population distribution over the joint (across hidden layers) parameter space is shown to evolve under SGD as per a McKean- Vlasov type nonlinear IVP; see (Araújo et al. 2019 Def. 4.4 and Sec. 5). We anticipate that the proximal recursions proposed herein can be extended to this setting by effectively lifting the Wasserstein gradient flow to the space of measure- valued paths. Though not quite the same, but this is similar in spirit to how classical bi-mariginal Schrödinger bridge problems (Léonard 2012 Chen et al. 2021) have been generalized to multi-marginal settings over the path space and have led to significant algorithmic advances in recent years (Haasler et al. 2021 Carlier 2022 Chen et al.) 2023). Pursuing such ideas for designing proximal algorithms in the multiple hidden layer case will comprise our future work.
Acknowledgment
This work was supported by NSF award 2112755. The authors acknowledge the reviewers' perceptive feed- back to help improve this paper.
References
David Alvarez-Melis, Yair Schiff, and Youssef Mrouch. Optimizing functionals on the space of probabilities with input convex neural network. In Annual Conference on Neurol Information Processing Systems, 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155. PMLR, 2017.
Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint Xrxi::1906 00190, 2019.
Andrew R Barron. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information theory, 39(3):930-8 1993.
Heinz Н Bauschke, Patrick Combettes, et al. Conver analysis and monotone operator theory in Hilbert spaces, volume 408. Springer, 2011.
Jean-David Benamou, Guillaume Carlier, and Maxime Laborde. An augmented Lagrangian approach Wasserstein gradient. flows and applications. ESAIM: Proceedings and surveys, 54:1-17, 2016.
18

---

Published in Transactions on Machine Learning Research (12/2023)
Dimitri P Bertsekas et al. Incremental gradient, subgradient, and proximal methods for convex optimization: A survey. Oplimization for Machine Learning, 2010(1-38):3, 2011.
Clément Bonel, Nicolas Courly, Frangois Septier, and Lucas Drumetz. SlicedWasserssle gradient flows. arXiv preprint arXit:2110.108 2021.
Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Efficient gradient flows in sliced- wasserstein space. arXiv preprint 0rzi0:2110.10972, 2022.
Etienne Boursier, Loucas Pillaud- Fillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances in Neural Information Processing Systems, 35: 20105 20118, 2022.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375 417, 1991.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal transport modeling population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511-6528. PMLRR, 2022.
Kenneth F Caluya and Abhishek Halder. Gradient flow algorithms for density propagation in stochastic systems. TEEE Transactions on Automatic Control, 65(10):391144044, 2019a.
Kenneth F Caluya and Abhishek Halder. Proximal recursion for solving the Fokker-Planck equation. In 2019 American Control Conference (ACC), pp. (098-4108 LEEE, 2019b.
Kenneth F Caluya and Abhishek Halder. Reflected Schrödinger bridge: Density control with path constraints. In 2021 American Control Conference (ACC), pp. 1137-1142. IEEE, 2021a.
Kenneth F Caluya and Abhishek Halder. Wasserstein proximal algorithms for the Schrödinger bridge prob- lem: Density control with nonlinear drift. IEEE Transactions on Automatic Control, 6363):1133-1178, 2021b.
Guillaume Carlier. On the linear convergence of the multimarginal sinkhorn algorithm. SIAM Journal on Optimization, 32(2):786-794, 2022.
Guillaume Carlier, Vincent Duval, Gabriel Peyré, and BeSchruhard SchmitzeSchr. Convergence of entropic schemes for optimal transport and gradient flows. SIAM Journal on Mathematical Analysis, 19(2):1385- 1418, 2017.
René Carmona and François Delarue. Probabilistic theory of mean field games with applications I-II. Springer, 2018.
José A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for Wasserstein gradient flows. Foundations of Computational Mathematics, 22(2):389 443, 2022.
Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos A Theodorou. Deep momentum malti-marginal schr\" odinger bridge. arXiv preprint arXiv:2303.0r 011.0011111111 2023.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on schrodinger bridge. Siam Review, 63(2):249-313, 2021.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
Casey Chu, Jose Blanchet, and Peter Glynn. Probability functional descent: A unifying perspective on GANS, variational inference, and reinforcement learning. In International Conference on Machine Learn- ing, pp. 1212- PMLR, 2019.
George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303 314, 1989.
19

---

Published in Transactions on Machine Learning Research (12/2023)
Carles Domingo-Enrich, Samy Jelassi, Arthur Mensch, Grant Rotskoff, and Joan Bruna. A mean-field analysis of two-player zero-sum games. Advances in neural information processing systems, 33:20215- 20226, 2020.
Dheeru Dua and Casey Graff. UCI machine learning machine learning repository. 2017. URL "http:
7/archive.ich ics uus. edu/ml"
Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: mean-field framework for over- paraparraeteri deep neural networks. In Conference on learning theory, pp. 1887-1936. PMLR, 2021.
Charlie Frogner and Tomaso Poggio. Approximate inference with Wasserstein gradient flows. In International PMLR, 2020.
Conference on Artificial Intelligence and Statistics, pp.
Isabel Haasler, Axel Ringh, Yongxin Chen, and Johan Karlsson. Multimarginal optimal transport with a troe-structured cost and the schrodinger bridge problem. SIAM Journal on Control and Optimization, 59 (4):2428 2453, 2021.
Abhishek Halder and Tryphon T Georgion. Gradient flows in uncertainty propagation and filtering of linear gaussian systems. In 2017 IEEE 56th Annual Conference Decision and Control (CDC), pp. 3081-3088. TEFE, 2017.
Abhishek Halder and Tryphon T Georgion. Gradient flows in filtering and Fisher-Rao geometry. In 2018 Annual American Control Conference (ACC), pp. 4281-4286. IEEE, 2018.
Abhishek Halder and Tryphon T Georgiou. Proximal recursion for the Wonham filter. In 2019 TEFE 58th Conference on Decision and Control (CDC), pp. 660 665. IEEE, 2019.
Abhishek Halder, Kenneth F Caluya Bertrand Travacca, and Scott J Moura. Hopfield neural network flow: A geometric viewpoint. IEEE Transactions on Newral Networks and Learning Systems, 31(11):88994488. 2020.
Abhishek Halder, Kenneth F Caluya, Pegah Ojaghi, and Xinbo Geng. Stochastic uncertainty propagation in power system dynamics using measure-valued proximal recursions. IEFF Transactions on Power Systems, 2022.
Kurt Hornik, Maxwell Stinchoombe, and Halbert White. Multilayer feedforward networks are universal approximators. Newral networks, 2(5):359-360 1989.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convorgence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Richard Jordan, David Kinderlchrer, and Felix Otto. The variational formulation of the Fokker-Planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Mark Kac. Foundations of kinetic theory. In Proceedings of The third Berkeley symposium on mathemalical statistics and probability, volume 3, pp. 171-197, 1956.
Johan Karlsson and Axel Ringh. Generalized Sinkhorn iterations for regularizing inverse problems using optimal mass transport. STAM Journal on Imaging Sciences, (0(4):1935-196) 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXio preprint arXiv:1412.600 2014.
flows. Topological Optimization and Optimal Transport: In the Applied Sciences, 17:304 332, 2017. Maxime Laborde. On some nonlinear evolution systems which are perturbations of Wasserstein gradient
Jachoon Lee, Lechao Xiao, Samuel Schocnholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
20

---

Published in Transactions on Machine Learning Research (12/2023)
Christian Léonard. From the schrödinger problem to the monge- kange-kantooovich problem. Journal of Functional Analysis, 262(4))1879-199) 1920, 2012.
Ping Li and Phan-Minh Nguyen. On random deep weight- autoencoders: Exact asymptotic analysis, phase transitions, and implications to training. In International Conference on Learning Represenlations, 2018.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning, pp. 4082-4092. PMLR, 2019.
Yulong Lu. Two-scale gradient descent ascent dynamics finds mixed nash equilibria of continuous games: A mean-field perspective. In International Conference on Machine Learning, pp. 12790-22811 PMLR, 2023.
Alexander G de G Matthows, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahramani. Gaus- sian process behaviour in wide deep neural networks. In International Conference on Learaing Represen- tations, 2018.
Henry P McKean Jr. A class of Markov processes associated with nonlinear parabolic equations. Proceedings of the National Academy of Sciences, 56(6):1907-191) 1966.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, (1(333))7665--7777 2018.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243- 15256, 2021.
Youssef Mrouch and Truyen Nguyen. On the convergence of gradient descent in GANs: MMD GAN gradient flow. In International Conference on Artificial Intelfigence and Statistics, pp. 1720 1728. PMLR, 2021.
Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv: 10xx:190..02800 2019.
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A Alemi, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In Interna- tional Conference on Learning Representations, 2019.
Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and trends@ in Optimization, 1(3):
127-239, 2014.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4):2323 225122335,,, 2015.
Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot. Wasserstein barycenter and its application texture mixing. In International Conference Scale Space and Variational Methods in Computer Vision pp. 435-446. Springer, 2011.
R. Tyrrell Rockafellar. Augmented Lagrangians and applications of the proximal point algorithm in convex programming. Mathematics of operations research, 1(2):97-116, 1976a.
R Tyrrell Rockafellar. Monotone operators and the proximal point algorithm. SIAM journal on control and oplimization, 14(5):877 898, 1976b.
Grant Rotskoff and Eric Vanden Eijnden Trainability and accuracy of artificial neural networks: An inter- acting particle system approach. Communications on Pure and Applied Mathematics, 75(9):1889-199) 2022.
21

---

Published in Transactions on Machine Learning Research (12/2023)
Grant M Rotskoff and Eric Vanden-Eijnden. Neural networks interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stal, 1050:22, 2018.
Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein proximal gradient algorithm. Advances in
Newral Information Processing Systems, 33:12356 126666,2666666 2020.
Filippo Santambrogio. [Euclidean, metric, and Wasserstein| } gradient flows: an overview. Bulletin of Math- ematical Sciences, 7(1):87-154, 2017.
Justín Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3):1820 1852, 2020.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis deep neural networks. Mathematics of Operations Research, 47(1):120-152, 2022.
Alain-Sol Sznitman. Topics in propagation of chaos. In Ecole d'été de probabilités de Saint-Flowt XIX-1989 pp. 165 251. Springer, 1991.
Mare Teboulle. Entropic proximal mappings with applications to nonlinear programming. Mathemalics of Operations Research, 17(3):670-66) 6900, 1992.
Anthony C Thompson. On certain contraction mappings in a partially ordered vector space. Proceedings of the American Mathematical Society, 14(3):438-445 443, 1963.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Cédric Villani. Topics in optimal transportation, volume 58, American Mathematical Soc., 2021.
Sofia Visa, Brian Ramsay, Anca L Ralescu, and Esther Van Der Knaap. Confusion matrix-based feature selection. MAICS, 110:120-127, 2011.
Stephan Wojtowytsch and E. Weinan. Can shallow neural networks beat the curse of dimensionality? mean field training perspective. IEEE Transactions Artificial Intelligence, 1(2):121 129, 2020.
LecHao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of CNNs: How train 10, 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pp. 5393-544 PMLR, 2018
Jiahui Yu and Konstantinos Spiliopoulos. Normalization effects on deep neural networks. Foundations of
Data Science, 5(3):389-44)) 465, 2023.
Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gra- dient method for reinforcement learning with general utilities. Advances in Neural Information Processing Systems, 33:4572-4583, 2020.
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as Wasserstein gradient flows. In International Conference on Machine Leurning, pp. 5737-5746. PMLR, 2018.
A Proof of Theorem 1
We provide the formal statement followed by the proof.
Theorem 1. Consider the regularized risk functional
12 wherein F is given by
- D. Let p(t, solve
the IVP [14,, and let (ee- be the sequence generated by 17 with Do = Pa. Defone the interpolation
[0, 50) x Re + (0) as
(1, = e (h. E [(k kh), & E.N.
Then 0n (t,00) 440 p(t,0) )) in (.PP).
22

---

Published in Transactions on Machine Learning Research (12/2023)
Proof. Our proof follows the general development in Laborde (2017 Sec. 12.3). In the following, sketch the main ideas.
We have the semi-implicit free energy
Fale, (x-)) = Eg Fo - v + / J ex- +3 Egg [log EEN,
+
=Vudvoc(e)
wherein the summand
Vadvec = Eg Fo v(0) / ((0,0)em-
is linear in and contributes E- [log e) results in from diffusion regularization and contributes
effective advection potential energy. The remaining summand internal energy term.
We note from equation g that the functional Vadvec(@) is lower bounded for all E P2 (RRP). Furthermore, (edvec(O) and vod (0) are uniformly Lipschitz continuous, i.e., there exists C₁ > 0 such that
V ( (e) (RP)) + V V Vadvee (e)lL C1
for all P2 (IP) where the constant C1 > 0 is independent of and 7 denotes the Euclidean Hessian operator.
Moreover, there exists C2 > such that for all o, 0 P2(R"), we have
Vadvec(e) - V Vudvee
L (RP) S C2W2 (e,
Thus, (e)voc(e) satisfy the hypotheses in Laborde 2017
Sec. 12.2).
For E [0, we say that p(t, E smooth compactly supported test function yp E C ([0, x RR), we have
P2 [RP) is a weak solution of the IVP equation 14 if for any
(
10p 01 (Ve. (adveo p + 3 dedi = -- -/- -(1=00,)0p00)
(44)
RP
Following Laborde (2017 Sec. 12.2), under the stated hypotheses on Vndver (e), there exists weak solution of the IVP equation 14 that continuous w.r.t. the W2 metric.
The remaining of the proof follows the outline below.
Using the Deltord-Pettis' theorem, establish that the sequence of functions (c,(hh))eee solving equation 17 is unique.
Define the interpolation en (t) (h, 0) ift ((k-1)4 kh), kh] for all e N. Then establish that en (t)
solves a discrete approximation of equation 44
Finally combine the gradient estimates and pass to the limit h 4. O, to conclude that (x(() in this limit solves converges to the weak solution of equation 44 in strong L (RP) sense.
For the detailed calculations on the passage to the limit, we refer the readers Laborde 20117 Sec. 12.5).
B Expressions involving the derivatives of v
We define matrices
T = tanh [ T + b1 )
S sech? 2 (WX T
+ b1
23

---

Published in Transactions on Machine Learning Research (12/2023)
where 1 is vector of all ones of size Ndata 1, and the functions tanh(-) and sech²()) are elementwise. Notice that T,S E R.N R.Naaata
Then
2
a 2 (tanh(WXX + b1 }g)
3 (Ty).
Adata
Bddata
Proposition 2. With the above notations in place, have
Ovk
2 Ty. idata
1
(45)
0a
da
and
2v
N
00k
a @ Sy.
(46)
Öb
Öb
idata
Furthermore,
N
OuR
[(al]) o (S cyl })]
(47)
OW
Adaaa
Proof. The 8th element of v is Us
2 ax
> siil
[ T.
Thus,
00
for j.
0a;
2
Vldaaa
T TR. for h: j.
ndata
1=
So the matrix de da
is diagonal, and
Ov da
2
Rddat
1
T
Ty)s y)
nddata
Adata
Hence, obtain
Ov
2
1
Ty
Da
Tidata
which is 45
On the other hand,
OvR
Ränta
ak
T
Adata
2
Adata
0
Rdatt
bb
Note that
a
θ
D2 ( Wis) + bb
tanh
yi
Ob
bb.
The
sech?
Σ
W
kJj
( + +
81 =Sk.yi
24

---

Published in Transactions on Machine Learning Research (12/2023)
Therefore,
Ov
for
k + J.
ak Σ
Tddata
(
for k j.
Ddaaaa
1=
As the matrix de 36
is diagonal, we get
1
Σ
ag [yy] ks
0b
k
dadta
2m
Adata
and so
1
④ Sy,
Öb
Iddtt
which is indeed 46
Likewise, we take an element-wise approach to the derivatives with respect to weights W W.m- Note that such a weight Wa km will only appear in the kth element of v, and so we only need to compute
dok
Phanto
as
( 1.).
ÖW k.m
didata
O k.m
Since
0 oW kmm
TR
tanh
W X + bb yi
OW k.m
= sech²
Σ
W ki. Xin) + ba
X Xim = S imm
we have
Ovk
for + m, S for k m.
ow m.j
2 ak E
Vddata
Thus,
N
Ovr
2
Thiaza
am
S Sm,iXisn
X iimyil
OW m.j
Dedata
2
am
[ (X
3 (y1 ()]
Rdata
Therefore, considering 1 Ra= we write
O0R
o (S [(11])
C [))]
aw
Ndata
thus arriving at 47 This completes the proof.
C Expressions involving the derivatives of u
Expressions involving the derivatives of u, are summarized in the Proposition next. These results find use in Sec. 4 We start by noting that
u ndata
a p.
Adata
25

---

Published in Transactions on Machine Learning Research (12/2023)
Proposition 3. With the above notations in place, have
du
1
[((en) 2
(T 1
TT
(48)
da
Adata
NxN Nx1
and
iu
[(a11) ② (ST ((a
1
0b
Pdata
NxN Nx1
((1a)) o (ST o ((a 2 1))
(49)
Furthermore,
N
N
Adata
agax(0i + (b) BB) ii
(50)
OW.,
Tdaata
Proof. Letting t denote the ith row of T, we rewrite as follows:
atti
u
(piatt +... + pNaNtn)
Thdata
aNt
ant (pranti + +... pNaNtN)
Addta
(NNaiiiaati +... + pNaNtn)
For + k, we thus have
Ou: act, (pata)
a:pkt th.
da;
ndata
Ndata
Likewise, for i= k, have
Ouk
an patt bu + 1o (piatti+ +... + pvaxtn).
Da
Idata
Pldaaa
Combining the above, we obtain and hence follows. On the other hand, for # k. we have
8a
Öt: att phansk. ou: at, Pkak ab. Ob ndata
ndaca
and for = k, we obtain
0u: Ot. ab
ai
(pratti +... + pNaNtn)
Adaaa
1
Ol
untt piack bb.
Odata
(Sw) (platti t... pNaNtn)
Rdata
1
agti PkanSk.
Rdata
Combining the above, obtain 8u 86
and hence 49 follows.
26

---

Published in Transactions on Machine Learning Research (12/2023)
Iteration# 1
Iteration## 1000
#reration50000
Regularized risk functional F3
10-2
X 10
0
10
10
k
10a
103
Figure 4: The evolution of the regularized risk f3 versus iteration index k for the proposed PROXLEARN. Inset plots compare the ground truth (simusoid) with the output from the network at three specific iterations.
Finally, noting that W j. is in the ith row of T, for + k, we obtain
Ou, OW ر.
Ol,
asty p:0; OW i.j Ndata
arth praz(s; ① zz),
Adata
where By the jth column of X. Likewise, for = we get.
Ouk ak
Öyg aw i.j
(pranti +... +..... pNaNtN)
aW s.j
Adata
Ot, antt pkacs OWi,
Tdtata
ax(s: O ( (prant +.... pNaNtx)
Ndata
pxap Oz₂)..
Adaata
au Combining the above, we obtain aw,,'
thereby arriving at 50
Learning sinusoid
To better visualize the functionality of the proposed algorithm, perform synthetic study of learning sinusoid following the set up as in (Novak et al. 2019 Sec. 2.1). We performed 5000 itera- tions of PROXLEARN with N 1000 samples (no mini-batch) from the initial PDF (o = (a, b, w)) = Unif ([-1, x [-1,1] x [-1.5, 1.5)), and used algorithm parameters 8 = 0.3, h = 10-4, = 10 E 10
L 10. The evolution of the associated regularized risk functional and the Iearnt functions are shown Fig.
Fig. shows the function approximations learnt by our proposed algorithm the end of 3200 iterations for 20 randomized with the same initial PDF and parameters as reported here.
27

---

Published in Transactions on Machine Learning Research (12/2023)
sin(a)
1.00
0.75
0.50
0.25
0.00
0.25
0.50
-0.75
-1.00
22
2
E
Figure Comparison of the ground truth (here sin(x)) with the learnt approximants obtained from the proposed PROXLEARN after 3200 iterations for 20 randomized runs. All randomized runs the same initial PDF and parameters as reported here.
Table 5: Comparing final Fg for varying N
N
Final F3
500
0.01241931
700
0.01075817
1000
0.00806645
2000
0.00762518
To illustrate the effect of finite N on the algorithm's performance, we report the effect of varying N on the final regularized risk value Fg for a specific synthetic experiment. We observe that increasing N improves the final regularized risk, as expected intuitively.
28


Reviews:
[{"review_id": "fIk43yWfKd", "reviewer": "Reviewer_kope", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThere is quite a lot of seemingly novel math here.  The overall logic of the algorithm is easy to follow despite the complexity caused by the many mathematical steps and operations.\n\nThis provides a different viewpoint on learning in shallow neural networks.  This viewpoint is potentially valuable to really understand the behavior of how learning happens in greater detail.\n\n## Weaknesses\n\nIn my opinion, the greatest weakness of this work is that it is not clear how the developments in this paper can be used and built upon.  The author's highlight at the end of Section 1.3\n> This is a new line of idea for learning algorithm design that we show is feasible. As such, we do not aim to immediately surpass the carefully engineered existing state-of-the-art in experiments. Instead, this study demonstrates a proof-of- concept which should inspire follow up works.\n\nI generally think that it's good to explore new ideas, but the claim here is rather vague.  What are the follow-up works?  How could this be used?  I don't see how a similar framework would be created for deeper models.  This manuscript would be greatly enhanced by discussing at much greater length how the framework could be used in future work.\n\nThe overall landscape of analysis of shallow networks should be enhanced.  For example, there are the viewpoints where we can view and learn with kernels or gaussian processes.\n\nI would really like to understand what this algorithm is learning functionally and how it compares to alternative approaches.  I would suggest looking at synthetic data with 1-dimensional inputs so that you can visualize the input-output space, and then looking at the functions learned by this approach versus others.  How does it compare to gradient descent or kernel approximations functionally.  For example, consider Figure 1 in [1].\n\n[1] Novak, Roman, et al. \"Neural tangents: Fast and easy infinite neural networks in python.\" ICLR 2020.\n\n## Questions\n\nThe learning curves in Figure 2 are unusual.  Why is there such a significant burn-in period where the risk functional does not get better?  Why is the discrepancy so large between the GPU and CPU solutions? e.g., the GPU has big jumps where the risk functional gets worse.  I understand that they have different numerical precisions, but I wouldn't expect them to be so different.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript proposes a novel optimization algorithm for shallow neural networks.  In particular, this looks at a PDE approach to neural network learning and derives a new algorithm to approximate the neural network solution by weighted scattered particles.  The derivation of this algorithm requires significant mathematical developments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have three required changes:\n\n1. A bigger discussion about the potential long-term utility of this work, including more specific examples of how this could be used either theoretically or applied going forward.\n2. A comparison of how different functions learned by this approach versus other approaches are.\n3. A discussion of the learning dynamics, and why the algorithm is so slow in early iterations.\n\nSome additional suggested changes:\n1. There's not much evidence on what the different parameters of the algorithm do.  While the primary goal is a proof-of-concept, it would be beneficial to understand what the tuning parameters do.  I'm particularly interested in $N$.\n2. Algorithm 1 and 2 are challenging to read.  I'd suggest adding comments throughout the algorithms to explicitly describe the steps and link them back to the paper.  Furthermore, it may be useful to further break up Algorithm 1 for clarity.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rJnTUJf84y", "reviewer": "Reviewer_QGHP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The logic of this paper is clear.\n2. The derivation seems solid and correct.\n\nWeaknesses: \n1. Writing needs some improvement.\n2. Some parts of the experiments are confusing.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a proximal algorithm to approximate the distributional flow of the learning dynamics of shallow neural networks. Such method tolerates evolving the neuronal population distribution in an online manner. Also, the motivation and derivation of the algorithm is supported by solid mathematical works.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Sections 2 and 3 are too detailed with many equations. I don't think they are needed in the main body of the paper. It makes it difficult to follow the paper. Some major insights and conclusions should be emphasized to make the contribution clearer. Some discussions and explanations, e.g., for Theorem 1, sections 3.1 and 3.2,  can be added to help the understanding. \n\n2. For experiments, (a) Section 4.3 seems irrelevant to your algorithm. If so, it is better to remove it. (b) I think Table 3 shows that the results are worse than existing methods. (c) it is better to add the comparison of computational complexities between different methods.", "broader_impact_concerns": "No such concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "44K1sGUoqJ", "reviewer": "Reviewer_7qYb", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is very clearly written and does an admirable job of discussing related work. The motivation for and construction of the algorithm is done very precisely and approximations are mentioned throughout. Comparing the limiting dynamics to particle-based methods remains a significant topic of interest, so I believe this numerical approach is also timely and important. \n\nFor me, the main weaknesses involve the numerical results: first, the role of regularization parameter in the Sinkhorn is not really described in terms of either the stability or accuracy of the results. It would be good to include some numerical results that show how performance varies with the degree of regularization. \n\nSecond, in both the weighted and unweighted cases, the evolution of the loss is not particular reassuring (Fig. 2). The risk functional does not decay smoothly (as opposed to Fig. 3) and the numerical stability looks rather questionable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a proximal algorithm for a Wasserstein gradient flow in the space of probability measures to carry out shallow neural network training in the mean-field regime. The main contributions are the construction of an algorithm to carry out this scheme in practice (Eq. 34), a convergence result ensuring asymptotic convergence to a unique fixed point, and numerical experiments assessing the performance of the method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the authors should study the role of $\\epsilon$ and give better assurances of stability in their implementation. I would also like some analysis of the time evolution of the risk---why does it seem to suddenly start decreasing only after $10^5$ iterations?", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Given the above 2 points, I am happy to recommend this paper for acceptance into TMLR. The authors have already made several updates to their paper in light of the reviews, and the paper can be accepted as-is (after changing blue changes back to black).
Decision: Accept (poster)