Gaussian mixture layers for neural networks
Authors: Sinho Chewi
Philippe Rigollet
Yuling Yan
Venue: Accepted by TMLR ()

Abstract: The mean-field theory for two-layer neural networks considers infinitely wide networks that are linearly parameterized by a probability measure over the parameter space. This nonparametric perspective has significantly advanced both the theoretical and conceptual understanding of neural networks, with substantial efforts made to validate its applicability to networks of moderate width. In this work, we explore the opposite direction, investigating whether dynamics can be directly implemented over probability measures. Specifically, we employ Gaussian mixture models as a flexible and expressive parametric family of distributions together with the theory of Wasserstein gradient flows to derive training dynamics for such measures. Our approach introduces a new type of layer—the Gaussian mixture (GM) layer—that can be integrated into neural network architectures. As a proof of concept, we validate our proposal through experiments on simple classification tasks, where a GM layer achieves test performance comparable to that of a two-layer fully connected network. Furthermore, we examine the behavior of these dynamics and demonstrate numerically that GM layers exhibit markedly different behavior compared to classical fully connected layers, even when the latter are large enough to be considered in the mean-field regime.

Published in Transactions on Machine Learning Research (12/2025)
Gaussian mixture layers for neural networks
Sinho Chewi
sinho. Sinho.chewi@yyke..o edu
Deportment of Statistics and Data Science
Yale University
Philippe Rigollet
rigollet@mit.cm edu
Department of Mathematics
Massachusetts Institute of Technology
Yuling Yan
yuling
Bolking-yonneene you Gurisc. edu
Deportment of Statistics
University of Wisconsin- Madison
Reviewed OpenReview: https: https: opwnrevviee. net/ forum? id=sAptI
Abstract
The mean-field theory for two-layer neural networks considers infinitely wide networks that are linearly parameterized by a probability measure over the parameter space. This nonparametric perspective has significantly advanced both the theoretical and conceptual understanding of neural networks, with substantial efforts made to validate its applicability lo networks of moderate width. In this work, explore the opposite direction, investigating whether dynamics can be directly implemented over probability measures. Specifically, we employ Gaussian mixture models as a ficxible and expressive parametric family of distributions logether with the theory of Wasserstein gradient flows to derive training dynamics for such measures. Our approach introduces a new type of layer- the Gaussian mixture (GM) layer that can be integrated into neural network architectures. As proof of concept, we validate proposal through experiments simple classification tasks, where GM layer achieves test performance comparable to that of two-layer fully connected network. Furthermore, we examine the behavior of these dynamics and demonstrate numerically that GM layers exhibit markedly different behavior compared to classical fully connected layers, even when the latter are large enough to be considered in the mean- field regime.
1 Introduction
Deep learning architectures compositions of basic trainable layers, and many major milestones of deep learning models can be traced back to innovations for these fundamental components. For example, convolutional neural networks are based on convolution layers and pooling layers; residual neural networks are based on layers with skip connections (He et al., 2016); and more recently, the striking feats of large language models stem from the use of attention layers ( (Vaswani et al., 2017). Yet without guiding principles map out the potential design space, the development of successful new layers has been largely elusive.
Whereas the vast majority of works on the theory of deep learning focus on understanding and explaining the behavior of existing architectures, in this paper take the approach of applying the theory to propose new ones. In particular, we incorporate insights from the recent literature on the mean-field theory for neural networks (Chizat & Bach, 2018; Mei et al., 2018; 2019; Sirignano & Spiliopoulos, 2020; Chizat, 2022; Nitanda et al., 2022; Rotskoff & Vanden-Eijnden, 2022) in order to propose new type of layer, which we call the Gaussian mixture (GM) layer.
According to fieldfieed theory, which we briefly review in Section the training dynamics of fully connected 2-ayyer neural network converges, under scaling in which divide the output by the number of neurons and

---

Published in Transactions on Machine Learning Research (12/2025)
in the limit of infinite width, to Wasserstein gradient flow in the space of probability measures. Here, the probability measure represents the distribulion of weights in the fully connected layer and thereby shifts our perspective from the evolution of individual neurons to the evolution of the collection thereof. This framework has been successful at making muanced and verifiable prodictions about the training of certain simple ncural networks (Mei el al., 2018; Abbe al., 2023; Berthier al., 2025), but thus far the mean-field limit has mostly been used as a tool for analysis. Indoed, it is not practical to actually implement the Wasserstein gradient flow, since doing so would require prohibitively large widths.
In this work, we explore the consequences of prescriptive mean-field rather than a descriptive one. This proposal results in manipulating infinite dimensional objects- objects-distribtioons over Rd_ and to implement it, we restrict the distribution over weights to a parametric Gaussian misture whose parameters are trained via standard optimization routines. The use of mixture modelling can also be motivated on grounds of clustering phenomena for neurons but we do not explore this question here.
Although our GM layer is proposed as a replacement. for a wide fully connected layer, the parametrization and restriction of the distribution of neurons to the class of finite Gaussian mixtures leads to markedly different training dynamics. We demonstrate this behavior through numerical experiments in Section 5 the MNIST database, which also serve as a proof of concept for the incorporation of GM layers into neural network architecture design. Our initial results are promising and show that a GM layer attains comparable ble test performance to a 2-layer fully connected network. However, we stress that our goal is not to demonstrate superiority of the GM layer over existing architectures which would necessitate going beyond the simple two-layer architecture for which the mean-field theory applies. As a result, while we demonstrate that the GM layer is modular and can be integrated into deep architectures, leave detailed investigation to future research as it would require substantial engineering developments that are beyond the scope of this proposal.
Contributions.
Our primary contribution is our Gaussian layer proposal, which we detail in Section 3.
We also show that whereas the Wasserstein gradient flow over empirical measures is implemented via the Euclidean gradient flow over the locations, the gradient flow over Gaussian mixtures equipped with a certain natural geometry-is is implemented via the Euclidean gradient flow over the means and the square roots of the covariances (Theorem 1).
In Section 4, we suggest an efficient parametrization scheme to speed up implementation.
We conduct numerical experiments in Section 5 on the MNIST and Fashion-MNIST datasets, which serve as a proof of concept and provide some insights into the training dynamics of GM layers. First, we show that GM layers can achieve comparable performance as 2-layer fully connected network. This is despite the fact that the two exhibit quite different training dynamics (see Figure 4). We also check that GM layers exhibit "feature learning" as expected in the mean-field regime- the sense that the distribution over first layer weights moves substantially away from initialization. Finally, we exhibit performance gains obtained by going deeper (i.e., composing multiple GM layers).
Other related work. In recent years, Wasserstein gradient flows have been applied to numerous probabilistic problems, such as sampling and variational inference. A common bottleneck for these applications is implementation of the flow, which can be achieved via stochastic dynamics in the case of sampling (Jordan et al., 1998) but remains challenging in general. The mean-field theory reviewed in Section 2 shows that the empirical distributions of the weights of a neural network along training indced follow a Wassorstein gradient flow, but our goal in this work lo maintain an evolution of a continuous distribution over the weights. The strategy we take here is to restrict the gradient flow to a parametric family, which was introduced in the context of variational inference for the family of Gaussians or mixtures of Gaussians (Lambert al., 2022; Diao et al., 2023) and later applied to filtering (Lambert et al., 2023) and mean-fiell variational inference (Tacker, 2023; Yao & Yang, 2023; Ghosh et al., 2025; Jiang et al., 2025).
2

---

Published in Transactions on Machine Learning Research (12/2025)
2 Review of field theory
We briefly review the mean-ficld theory for 2-layer fully connected neural networks. A width-m neural network computes a function h I2 + R of the form
m Σ Wj ((B j=1
hw,
m
where J:R-Ri is the activation function taken to be the ReLU function ((2) = max(0, z) for the rest of this paper and 3,) E R x Rd are trainable weights, for each neuron j E [m]] Here and throughout, we use boldface denote a collection of parameters. The 1/m scaling above characteristic of the mean-field regime and enable the following perspective. Let p pmm) denote the empirical distribution of the weights, p(m) m- دソ= -n (wyy,))); then we can write ( = f (((3,x)) (l (dw;, (dw,,d) In this formulation, however, we can make sense of this expression even when p(m) is no longer an empirical measure, and we can view h as being parameterized by a probability measure p:
hp(x) = f 118 x)) p(dw,
Consider a loss? objective L and minimize the objective L(ww,s) )) by following the gradient flow for the weights ( This gives rise a curve of parameters (w). ((t)) and corresponding empirical measures (t) m- "J (1)) The insight of mean-field theory is that the evolution of the empirical measures be described the (time-rescaled) gradient flow of the loss function p (hp) over the space of probability measures, equipped with the Wasserstein metric from optimal transport, initialized at p(m)(() (0). We refer to Villani (2003); Ambrosio et al. (2008); Villani (2009); Santambrogio (2015) for background on optimal transport and Wasserstein gradient flows.
The advantage of this reformulation is that it admits a well-defined limit as m + oc: if p(m)) (0) - p(0), then the curve of measures (p(m)(t) (t)) converges to the Wasserstein gradicnt flow of p (hp), initializod p(0). This mean-fiedd limit is, in some cases, easier to study than the original dynamics over the weights, and leads to predictions about the behavior of wide neural networks.
To summarize: the training dynamics of a finito-width neural network correspond to a Wasserstein gradient flow, initialized at (and remaining through its trajectory) an empirical measure, but the Wasserstein gradient flow picture more general because it allows for flows of continuous measures. Unfortunately, in the latter case, the Wasserstein gradient flow does not readily lend itself to tractable implementation. The mean-field theory described above shows that is well-approzimated by a gradient flow started at an empirical measure, but this approximation often requires prohibitively large width. In the next section, we take the familiar approach from statistics of restricting the measures to parametric family, namely, the set of finite Gaussian mixtures.
3 A mean-field theory over the space of Gaussian mixtures
We now introduce the Gaussian Mixture (GM) layer, beginning with the case of single GM layer (corre- sponding with the mean-field theory described in Section 2). Here, we restrict the measure in (1) to be a Gaussian mixture with components:
p= P=.
Σ N Ex), K
Pk E Rd+1
Σg E ((d-1)x(d+1).
(2)
Thus, the measure P is now parameterized by set of means and covariances for the components of the Gaussian mixture. We use the short-hand notation hy == hpn.z With this restriction, we can now train the
Here, &e refers a Dirac measure located 2The loss typically depends training set, but this is irrelevant our discussion here so it omitted.
3

---

Published in Transactions on Machine Learning Research (12/2025)
(J,E). GM layer by minimizing L (hp,z) with respect to the parameters (μ, Σ). For example, in regression task with a labeled dataset (zi, U/ NE[n), we might take the squared loss defined by [h) = E h(x;))²
The use of mixture modelling to model the distribution over neurons is motivated by the empirical observation that for many problems, the neurons tend to cluster (e.g., Papyan et al., 2020; Chen et al., 2023). Given this ansatz, the use of Gaussian mixtures emerges as a natural model for P, although other alternatives could be explored.
As discussed in Section 2, it is well-known from mean-field theory that the Wasserstein gradient flow restricted to empirical measures is implemented, up to time rescaling, by the Euclidean gradient flow with respect to the locations of the particles (Chewi et al., 2025, Proposition 6.16). When we move to Gaussian mixtures, we elfectively replace the particles 03 = (4) with "Gaussian particles" ((u; Σx), which are themselves distributions over 0 but can be viewed simply as (HR, Ex) pairs. In the case K = 1 of a single Gaussian particle, it turns out that the Wasserstcin gradient flow restricted to Gaussian mcasures implemented simply by evolving the parameters (μ,C) C) via the (Euclidean) gradient flow, where Σ = CCT. This is usually known the Bures- Wasserstein gradient flow, discussed next.
3.1 Interpretation as Wasserstein gradient flow
The gradient flows that we consider in this paper are closely related to Wasserstein gradient flows, we describe in detail here.
Sares-Wassersteiie gradient flows. We first consider the case K = 1, so that = (i, Σ) is simply a Gaussian measure. The space of non-degenerate Gaussian measures over IR,, which is naturally identified with Rd x $4++ can be equipped with the Wasserstein metric and is then known as the Bures Wasserstein space (Bhatia et al., 2019). We denote this space by BW(R?). The Riemannian structure of the Wasserstein space endows PW with Riemannian metric, called the Bures- Wasserstein metric.
Given a functional & over the Wasserstein space, we can restrict it to functional L BW(RP) + R via L(H,)) E) = (M, Σ)). If V Vy denote the usual Euclidean gradients of L w.r.t. and E respectively, it is known that the gradient flow of L over (W(PP) is given by
j VuL( E),
and
タ= 2
Vr, E).
See, e.g., Altschuler et al. (2021, Appendix A) or Lambert et al. (2022, Appendix B.3). On the other hand, if we parameterize Σ UUT where U E DDD and we follow the Euclidean gradient flow of (p, U) H L(u, UUT), it is straightforward to see that it coincides with (3). This parametrization also has the advantage of maintaining the positive semidefiniteness of Σ along the optimization without the need for projections, which is convenient for implementation and can also be used to enforce low-rank factorizations (Burer & Monteiro, 2003). In Appendix A, we derive the in mind that they can also be implemented as Fuclidean gradient. flows via the parametrization Σ UUT.
gradient flows for our problems of interest, keeping
Gaussian mixture gradient flows. In the case K 2 2, no longer possible to follow the Wasserstein gradient flow restricted to Gaussian mixtures. Indeed, the latter not explicit, since the space of Gaussian mixtures is not geodesically convex subset of the Wasserstein space. However, we can instead follow a Wasserstein gradient flow for the méring measure == R Σx ((H, over the Bures Wasserstein space. This geometric structure was previously introduced in Chen et al. (2019); Delon & Desolneux (2020) and can be interpreted as the Wasserstein geometry over the curved manifold (BW(RP), W2), see Lambert et al. (2022, Appendix F) for details.
3.2 Gaussian mixture flow for neural networks
We prove that the Gaussian mixture flow admits a simple implementation (sce Appendix A).
Thcorem 1. Let & be a loss function over the space of probability measures. Then, the Gaussian mixture gradient flow for & is equivalent, up to a time rescaling, to the Euclidean gradient flow of the objective (h),,) with respect to the parameters (1,C),, C) where Σn C2CI for each k E [K].
4

---

Published in Transactions on Machine Learning Research (12/2025)
In other words, compared to ordinary neural networks, training GM layers is accomplished by simply incorporating gradient steps for the square rools of the covariances.
4
Implementation
In this ion, we enhance the flexibility and tractability of GM layers by allowing for vector-valued outputs (needed for multi-class classification), reduced parametrization, and composition.
4.1 Multi-class classification and vector-valued outputs
Consider a multi-class classification problem with 1 labels denoted {O..... ..... L}. Suppose we are given a
dataset {a,, Vi ke(m)) where cach 2; Rd and Yi E (0, 1,..., L}. It suffices to describe how to paramcterize vector-valued function h Rd + R2,, since we can then apply the logistic loss
2(h)
log(1 logg
Σ exp(h(a.)e))
The most straightforward way parameterize the function h Rd + R is via (1) and (2), where now (w, B) R² Rd, ie., the Gaussian mixture Pu. distribution over Rd In other words,
[,x) = K-1
(w. (13, x13,
(4)
kE|K]
However, the number of parameters becomes ((dd + 2)2 K), which prohibitively large.
4.2 Reducing the number of parameters
To reduce the number of parameters, we propose to incorporate sparsity into the model by considering only diagonal covariance matrices for 8: ~ NNu, diag(o")) where diag(o?) is diagonal matrix whose entries are given by the vector o² O E Rd Moreover, for jointly Gaussian (w, B), the conditional mean of w given 3 is affine: E[w []] UB u. Since (4) only requires to know the conditional expectation E[w []] we model each component of the Gaussian mixture as
(p, diag(o))
E[w | UB++0 v.
The trainable parameters for such a component are Rd, E Rd, U BLxd and E R4. Hence, the parameters for the full Gaussian mixture are 0 = (H, = {[ ok, Uk, ux) uk) This leads to
ha(z)= = K1 Σ Eg- 1)BBB ((3,x)))]
55)
kE(K]
This use of (5) cuts down the number of parameters to (ddLL), substantial savings when L << d (as is typical). An additional benefit is that parametrization by automatically maintains the positive semidefiniteness of the covariances during training as they are of the form diag(o?), without recourse to costly projection steps.
To summarize, for multi-class classification, we train the parameters 0 via an optimization algorithm (in our experiments, vanilla stochastic gradient descent) the objective L Lho)), where L is the multi-class logistic loss and he is given in (5).
Comparison with fully connected networks. In summary, a fully connected layer with input and output dimensions and respectively, and width W, has ((d++L) parameters. A GM layer with components and full covariance parametrization has @(dd + T.)2 K) parameters. A GM layer with K components and reduced parametrization has ((dKL) parameters. Typically, K < W, e.g., in our experiments we often take K 5 5 or K 10, and for classification problem with not too many classes, L & d as well. The parameter count directly translates into memory and computation time per gradient update.
We may also write μ = 10² avoid confusion with (2).
5

---

Published in Transactions on Machine Learning Research (12/2025)
4.3 Stacking GM layers
In the previous subsection, we showed how to parameterize function h Rd + R4 as a GM layer. Since the input and output dimensions are arbitrary, this construction can be readily composed with other types of layers including GM layers themselves in order to build deep neural network architectures. Indeed, as depicted in Figure 1, the GM layer can be dropped in as a replacement for an (infinitely wide) fully connected layer and, so long as square root parameterization used for the covariance, can be trained standard backpropagation by Theorem
We leave the question of designing and optimizing deep architectures that integrate the GM layer to future research. Instead, we focus on a single GM layer, and call it GM network.
x
I
Bm
AB8 x)
1
w1,1
WL
hy (x:)
(x)
J [wi ABT x)) pldw, d3)
(1Ses)) L)
Figure 1: A GM layer (right) can act replacement for fully connected layer (left).
5
Experiments
The source code for all experiments in this section is available at httpp/////gtttto....//////////// laayeerr.
Datasct. We test the performance of neural networks with GM layers on multi-class classification on two widely used datasets: MNIST (LeCun & Cortes, 2010) and Pashion-MNIST (Xiao et al., 2017). Both datasets consist of 60,000 training examples and 10,000 test examples, where each example is a 28 28 grayscale image, associated with a label from one of 10 classes. Each image is vectorized and normalized to have zero mean and unit standard deviation. Throughout this section, test error refers to the misclassification error evaluated over the test set.
Setup. The number of components K hyperparameter: larger K enables more expressive G.M layers, while smaller K speeds up computation. We consider K E (5, 10, 20) for different experiments. For GM layer with parameters nr a, U, and u, we initialize the entries of n², U and i.i.d. from (O,,,²), and the entries of a all equal to 7, for some 7 > 0. For most of the experiments we fix 1/2, but by adjusting the value of y we can investigate the role the initialization scale, as discussed below. We train the work using SGD with batch size 64 and fixed learning rate 1 for the parameter o and 0.1 for all other parameters. All of these choices are used for simplicity and were found via minimal tuning.
Test error. For both datasets, we test the performance of network with one GM layer, where the number of components K takes values in (5, 10, 20). The results are presented in Figure 2. As we can see, a single GM layer with 20 components achieves test error of x 2.77% for MNIST, and z 12.13% for Pashion-MNISTT Increasing K leads to better performance, but the marginal improvement is very small after K exceeds 10, suggesting that K = 10 already leads to sufficiently expressive Gaussian mixtures. In practice one may choose K between 10 and 20 to strike a balance between expressive power and computational efficiency.
We also compare the test error of the GM network with the 2-layer fully connected network. We set the width of the latter network be 1000 (in fact the test error curve almost the same for m (100, 500, 1000), we just present the result. for 1000 for simplicity). We use two different methods to initialize the

---

Published in Transactions on Machine Learning Research (12/2025)
MNIST
Fashion-MNIST
0.30
0.40
K=10
K=10
0.25
K-20
0.35
K-20
0.20
0.30
test error
test error
0.15
0.25
D.10
0.20
0.05
0.15
0.00
0.10
0
20
40
60
80
100
0
20
40
60
80
100
number of epochs
number of epochs
Figure 2: Test error (with error bars) for a GM Network with K = 5, 10, 20 number of components vs. the number of epochs. The left (resp. right) panel shows the result. for MNIST (resp. Fashion-MNIST) dataset. The error bars are computed over 5 independent trials.
MNIST
Fashion- Fashion-MNIST
0.30
0.40
PyTorch default init
PyTorch default init
random init from
random init from
0.25
0.35
GM Net: random init (y= 1/2)
GMM Net: random init (y= 1/2)
0.20
0.30
test testerroo error
test error
0.15
0.25
0.10
0,20
0.05
0.00
0.10
20
100
20
100
number epochs
number epochs
Figure 3: Test error (with error bars) for a 2-layer fully connected network using two different initialization schomes vs. the number epochs. The left (resp. right) panel shows the results for the MNIST (resp. Fashion- MNIST) dataset. The error bars are computed over 5 independent trials.
parameters of fully-connected layer: Torch'ss default method (Kaiming uniform (He et al., 2015)), and random initialization drawn i.i.d from the initial distribution we used in training the GM network with K = 20 components (for a fair comparison with the GM network). The results are presented in Figure 3. We can see that although in the end all curves converges to comparable test errors, the number of epochs required to achieve low test are different: the GM network performs better than fully connected network with matching initialization, i.e., with weights drawn randomly from the initial Gaussian mixture distribution of the GM network, but worse than that with PyTorch's default initialization. This observation raises the question of whether the slight mnder-performance of the GM network compared to the fully connected network with default initialization could be overcome via better initialization schemes for GM layers.
Evolution of Gaussian components. In order to visualize the evolution of the Gaussian mixture (p:)220 during the training phase, we Irain the network for T = 200 epochs, compute the top 2 principal subspace the final Gaussian mixture distribution pr (marginalized over 3), and project the entire training trajectory (Pt)oSesT (marginalized over 3) onto this 2-dimensional subspace. Figure 4 depicts the evolution of this 2-dimensional distribution across 200 epochs for 5, using two types of initialization schemes y E (1/2, 1/256). As we see, the means of the five Gaussian components move far away from their

---

Published in Transactions on Machine Learning Research (12/2025)
MNIST (y= 1/2)
Fashion-MNIST (y= 1/2)
7.5
5.0
5
0.0
0-
-2.5
-5.0
-5
7.5
-10.0
-10
-5
5
10
-10
-5
5
10
MNIST (y= 1/256)
Fashion-MNIST (y= 1/256)
6
6
2
0
0
-2
-4
66
-6-
-8
-6 2
-10.0
-7.5
-5.0
-2.5
0.0
2.5
5.0
7.5
Figure 4: The evolution of the Gaussian components (marginalized over 3) and weights Bg of the neurons, projected onto the top 2 PCs the final GM distribution, across 200 epochs of training. The number of components for the GM net (resp. number of neurons the 2-layer neural network) K = 5. The projected Gaussian components are represented by their covariance ellipses centered at their means, while the projected weights of the neurons are depicted dots. We use the same color for the evolution of the same Gaussian component and neuron, with increasing opacity the number of epochs increases. The left (resp. right) plots show results for MNIST (resp. Mushion-MNIST)) while the top (resp. bottom) plots use initialization scale = 1/2 (resp. = 1/256).
initializations (even when they are initialized near zero, which the case when y 1/256), and the covariance matrices also become non-isotropic quickly. This also shows that the training dynamics of networks with GM layers are not overly sensitive to initialization.
We also train a fully-connected 2-layer neural network with width m = K to compare the training dynamics of networks with a GM layer and a fully-connected layer. To make the comparison fair, we initialize the neurons at 3x- = 44 and W/ = U.e + e for 1 < k < K where {ux WihgKK.. U and v are the parameters of the initial Gaussian mixture distribution. We train the neural netw using the same SGD algorithm with learning rate 0.1. The evolution of these neurons (projected onto the same 2-dimensional subspace) across 200 epochs is also shown in Figure4 4. We can that they exhibit drastically different training dynamics from that of the GM network. For example, when initialized at scale 7 1/2, the Gaussian components of the GM layer tend to move far away from zero, while the neurons of the fully-connected layer do not exhibit this trend.
Mean field vs. "NTK" regime. We design a simple experiment to inspect whether the training of a network with single GM layer is in the neural tangent kernel (NTK) a.k.a "lazy training" regime (Jacot et al., 2018; Chizat al., 2019; Du et al., 2019; Bartlett et al., 2021), or the mean-field regime. Indoed, since the presence of feature learning is powerful motivation for the mean-field regime, it is important to check
8

---

Published in Transactions on Machine Learning Research (12/2025)
MNIST
Fashion-MNIST
Random beta Trained beta
0.9
Random beta
Trained beta
0.8
test error
test error
0.6
0.4
0.2
0
0
10
number of epochs
number of epochs
Figure 5: Test error (with error bars) for training with or without updating the marginal distribution over B the number of epochs. The left (resp. right)) panel shows the result for the MNIST (resp. Fashion-MNIST) dataset. The error bars are computed over 5 independent trials.
MNIST
Fashion-MNIST
0.40
0.50
Fully-Connected layer
Fully-Connected layer
0.35
GM layer benchmark
0.45
GM layer benchmark
test error
0.25
error
0.35
0.20
0.30
0.10
0.20
0.05
0.15
0.10
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
number neurons
number neurons
Figure 6: Test error (with error bars) for fully connected 2-layer networks constructed via subsampling vs. width. The neurons are sampled from trained GM layer, whose test error is also plotted benchmark. The left (resp. right) panel shows the result for the MNIST (resp. Fashion-MNIST) dataset. The error bars are computed over 5 independent trials.
that this intuition also carries over to GM layers. Although lazy training is not formally defined for GM layers, can loosely take it to be the case when the distribution over the "first layer weights" B does not significantly move away from its initialization.
We fix the marginal distribution over 3 at its initialization (which can be achieved by setting the learning rates for 1 and a to be 0) and only update U and u for each Gaussian component. Figure compares the performance of training the network with fixed 3 vs. the network with trained 3. If we fix the marginal distribution over 8, the trained network can only achieve a test error of 40% for both the MNIST and Fashion-MNIST datasets. This shows that good test performance can only be achieved if moves significantly from its initialization, and thus "feature learning" is indeed crucial for the performance of GM layers.
Monte Carlo reduction to fully connected layers. After training network with GM layer, which gives a distribution jointly over (w, one natural question is whether it possible construct fully connected 2-layer network with similar performance by sampling a reasonable number of neurons from this

---

Published in Transactions on Machine Learning Research (12/2025)
MNIST
Fashion-MNIST
one GM layer two GM layers
0.9
one GM layer
two GM layers
0.8
0.7
test error
test error
0.6
0.5
0.4
0.2
0.3
0
2
10
0
2
8
10
number of epochs
number of epochs
Figure 7: Test error (with error bars) for networks with one or two GM layers vs. the number of epochs. The left (resp. right) panel shows the result for the MNIST (resp. Fashion-MNIST) dataset. The error bars are computed over 5 independent trials.
Gaussian mixture distribution. To answer this question, we first train a network with a GM layer with K = 20 components and 20 epochs on MNIST. Then we construct fully connected 2-layer networks by sampling m neurons (i.e., (w, 3) pairs) from the trained Gaussian mixture distribution, and evaluate their test errors without training. The results reported in in Figure 6 are indicating of the classical 1/vmm convergence rate of Monte Carlo approximation. Unfortunately, this convergence is to slow from the perspective of test error: even for m 10 there is still a gap between the performance of the GM network and its Monte Carlo approximation. This is a consequence of the high-dimensional nature of the space of parameter: we sample, average 10* /20 500 points per component, and this is not sufficient to estimate accurately Gaussian integral in dimension 28 x 28 = 784.
Stacking multiple GM layers. In the previous experiments, we used networks with a single GM layer. In this experiment, we demonstrate that stacking GM layers can improve classification performance. Specifically, we stack a GM layer with an input dimension of 784 and an output dimension of L' = 100 with second GM layer with an input dimension of L' 100 and an output dimension of 9, both layers utilizing K 10 Gaussian components. This configuration yields a lower test error compared to a single GM layer with the same number of Gaussian components, as shown in Figure 7. Our primary objective is to provide straightforward example illustrating the benefits of stacking multiple GM layers. We have not optimized hyperparameters such as the number of components K and the width of the middle layer m for optimal performance. Additionally, constructing a deeper architecture to achieve competitive performance on classic benchmark tasks like image classification would necessitate the integration and optimization of convolutional layers. This would shift the focus to secondary aspects that are beyond the scope of this paper.
Experiments on CIFAR datasets. Finally, we conduct an experiment based on the CIFAR-10 and CIFAR-100 datasets (Krizhevsky et al., 2009), which are larger in scale. We start with an effective VGG-like network architecture, which consists of four convolutional layers with batch normalization and max pooling, followed by three fully connected layers with dropout (see Appendix B for details). This network structure achieves a test error 12.62% for CIFAR-10, and test error 61.48% for CIFAR-100 over 50 training epochs. To incorporate the GM layer into this structure, we replace the last two fully connected layers with GM layer (with K = 10 components), which leads to an improved test errors 11.04% for CIFAR-10 and 43.55% for CIFAR-100, as shown Figure In addition to demonstrating the applicability of GM layers on more challenging datasots, this experiment also shows that GM layers can be composed with other types of layers and thus incorporated into complex architectures.
10

---

Published in Transactions on Machine Learning Research (12/2025)
CIFAR-10
CIFAR-100
VGG network VGG with GM layer
VGG network
0.9
VGG with GM layer
0.5
0.8
test error
test error
0.4
0.7
0.3
0.6
0.5
0.1
10
20
30
40
50
10
20
30
40
50
number of epochs
number of epochs
Figure Test error (with error bars) for classifying CIFAR-10 and CIFAR-100 datascts using VGG-type network (with and without GM layers) vs. the nunber of epochs. The error bars computed over independent trials.
6 Conclusion
In this paper, we have introduced GM layers as novel layer type for neural network architectures, offering a fresh perspective that bridges concepts from mean-field theory and variational inference. This approach opens up a wealth of unexplored possibilities for layer design, suggesting new aventes for incorporating diverse "variational families" to model p beyond Gaussian mixtures, and for optimizing with respect to various goometries within the variational family. For instance, while our current model restricts to Gaussian mixtures with equal weights, this limitation can be addressed by optimizing over Gaussian mixtures with unequal weights via the Wasserstein Fisher- Rao geometry (Liero et al., 2016; Chizat et al., 2018; Liero et al., 2018; Lu et al., 2019; Yan et al., 2024), as demonstrated in Lambert et al. (2022, Appendix H). Such advancements may also inspire alternatives for other types of layers, including convolutional and attention layers.
A notable limitation of our work the preliminary nature of our experiments. Further empirical investigation is required to refine our design choices and compare them with existing architectures. For example, the sparse parametrization in Subsoction 4.2 quite drastic, suggesting the need explore alternatives like low-rank factorization of the covariance matrix. Moreover, developing effective initialization and training strategies for GM layers remains open question. We anticipate that addressing these challenges will pave the way for more robust and versatile neural network archilectures in the future.
References
Emmanuel Abbe, Enric Boix-Adserà, and Theodor Misiakiewiez. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. In Gergely Neu and Lorenzo Rosasco (eds.), Proceedings of Thirty Sixth Conference Learning Theory, volume 195 of Proceedings of Machine Learning Research, pp- 2552-2623. PMLR, 7 2023.
Jason M. Altschuler, Sinho Chewi, Patrik Gerber, and Austin J. Stromme. Averaging on the Bures Wasserstein manifold: dimension-fres convergence of gradient descent. Advonces in Neural Information Processing Systems, 34:22132-222 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows in metric spaces and in the space of probability measures. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition, 2008.
Peter L. Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: statistical viewpoint. Acta Numer, 30:87-201, 2021.
11

---

Published in Transactions on Machine Learning Research (12/2025)
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Found. Compul. Math., 25(5):1627-1710, 2025.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Bures- Mures-Wassersteiin distance between positive definite matrices. Expo. Math., 37(2):165-191, 2019.
Samuel Burer and Renato D. C. Monteiro. A nonlinear programming algorithm for solving semidefinite
programs via low-rank factorization. Math. Program., 95(2):329-357, 2003. Computational semidefinite and second order cone programming: the state of the art.
Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: how gradient noise attracts SGD dynamics lowards simpler subnetworks. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp- 35027-35063. Curran Associates, Inc., 2023.
Yongxin Chen, Tryphon T. Georgiou, and Allen Tannenbaum. Optimal transport for Gaussian mixture models. IEEE Access, 7:6269-6277 6278, 2019.
Sinho Chewi, Jonathan Niles-Weed, and Philippe Rigollet. Statistical optimal transport, volume 2364 of Lecture Notes in Mathematics. Springer, Cham, 2025. Ecole d'Été de Probabilités de Saint-Flour XLIX 2019.
Lénaïc Chizat. Mcan-field Langevin dynamics: exponential convergence and annealing. Transactions on Machine Learning Research, 2022.
Lénaic Chizat and Francis Bach. On the global convergence of gradient descent for over- par-ammttrrtzzedd models using optimal transport. Advances in Newral Information Processing Systems, 31, 2018.
Lénaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xave Xavier Vialard. An interpolating distance between optimal transport and Fisher-Rao metries. Found. Comput. Math., 18(1):1-44, 2018.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training differentiable programming. H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (ods.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Julie Delon and Agnès Desolncux. Wasscrstein-type distance in the space of Gaussian mixture models. SIAM J. Imaging Sci., 13(2):936 970, 2020.
Michael Z. Diao, Krishnakumar Balasubramanian, Sinho Chewi, and Adil Salim. Forward-backwaa Gaussian variational inference via JKO in the Bures Wasserstein space. In Andreas Krause, Emna Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 1960-7991. PMLR, 2023.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. In International Conference on Learning Representations, 2019.
Soumyadip Ghosh, Yingdong Lu, Tomasz Nowicki, and Edith Zhang. On representations of mean-field variational inference. Journal of Applied Analysis, 2025.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers surpassing human- level performance on ImageNet classification. In Proceedings of the IEEE International Conference Computer Vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
12

---

Published in Transactions on Machine Learning Research (12/2025)
Yiheng Jiang, Sinho Chewi, and Aram-Alexandre Pooladian. Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space. Found. Compul. Math., 2025.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker-Plann equation. STAM J. Math. Anal., 29(1):1-17, 1998.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Danicl Lacker. Independent projections of diffusions: gradient flows for variational inference and optimal mean field approximations. arXiv preprint 13392, 2023.
Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Mare Lambert, Silvère Bonnabel, and Francis Bach. Variational Gaussian approximation of the Kushner optimal filter. In Geometric science of information. Part volume 14071 of Lecture Notes in Comput, Sci., pp. 395-404. Springer, Cham, 2023.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL http:://tannnn .ewp:///.aanneeeu.... lecun. com/exdb/mnist/.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal transport in competition with reaction: the Helinger-Kantorovich distance and geodesic curves. SIAM Math. Anal., 48(4))2869-2911, 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transpor problems and a new Kantorovich distance between positive measures. Invent. Math., 211(3):969-111, 2018.
Yulong Lu, Jianfeng Lu, and James Nolen. Accelerating Langevin sampling with birth-death. arXiv preprint 1905.09863, 2019.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proc. Natl. Acad. Sci. USA, (16633)) 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Alina Beygelzimer and Daniel Hsu (eds.), Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. PMLR, 2019.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field Langevin dynamics. In Gustan Camps-Vall Francisco J. R. Ruiz, and Isabel Valera (eds.), Proceedings of the 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, PMLR, 3 2022.
pp.
Vardan Papyan, Xiaoyan Han, and David L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training, Proc. Natl. Acad, Sci. USA, 117(40):24652 -246633, 2020.
Grant M. Rotskoff and Eric Vanden-Bijnden. Trainability and accuracy of artificial neural networks: an interacting particle system approach. Comm. Pure Appl. Math., 75(9):1889-13) 1935, 2022.
Filippo Santambrogio. Optimal transport for applied mathematicians, volume 87 of Progress in Nonlinear Differential Equations and their Applications. Birkhäuser/Spprineer Cham, 2015. Calculus of variations, PDEs, and modeling.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: a central limit thoorem. Stochastic Process. Appl., 13800)112--- 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Ilaia Polosukhin. Attention is all you neod. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Cédric Villani. Topics in optimal transportation, volume 58 of Graduate Studies in Mathematics. American Mathematical Society, Providence, RI, 2003.
Cédric Villani. Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften [Funda- mental Principles of Mathematical Sciences). Springer-Verlag, Berlin, 2009, Old and new.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: a novel image datasct for benchmarking machine learning algorithms. arXiv preprint arXiv:1 2017.
Yuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning Gaussian mixtures using the Wasserstein- Fisher-Rao gradient flow. Ann. Stalist., 52(4):1774-177) 2024.
Rentian Yao and Yun Yang. Mean-field variational inference via Wasserstein gradient flow. arXiv preprint 2207.08074, 2023.
A Exact derivations
In this section, we record exact expressions for the Gaussian mixture gradient flows. These expressions were used validate the correctness of our GM layer implementations in PyTorch and could be useful for future investigations.
We consider the following two problems.
Regression. We are given a dataset {(0, Ve)lve()) where cach Ia E IR and y/s E R. We consider the squared loss
(h)
[ - h(xx))²²
Multi-class classification. We are given a dataset f(x yi hegi where cach x; Rd and E 1,......, L. and denotes the number of classes. For this problem, we consider the multi-class logistic loss
hxxi)e - -log + Σ Exp((Mxs)))
L(h)
where h: Rd R4
For regression, we derive exact expressions for the gradient flows for the case of K = 1 (i.e., the Gaussian mixture simply Gaussian) and for the case K 1 for the full paramctrization (4). For multi-class classification, we focus on the case K with the diagonal parametrization (5).
Notation. For 0 = (w, (), the density of ((, We also use & to denote loss over the space of probability measures (whereas L denotes a loss over the space of functions).
the shorthand notation f(x, 0) w x(),z))). We write o(- p, 2) for
A.1 Regression
A.1.1 Gaussians
We first restrict, p to be single Gaussian, namely
RD 0 w
B p N(H,)) -N((Ur
Dw
Dus E ))
where D d 1. Then write the objective function as
min C(p) = Σ yi - ( 0(cx))) (=
(6)
14

---

Published in Transactions on Machine Learning Research (12/2025)
The Bures- Bures-Wassee Wasserstein gradient flow for minimizing is characterized by the following theorem.
Theorem 2. The Bures Wasserslein gradient flow (Pt)s20 for minimizing & in (6) is given by pr = N Ex) that evolves according to the ODE system
yi
E, (222, Ep. Vof(z,,
n Σ
(yi p (ri, 0)) E. Vol(a; 0) 0 pb) (0 14)) (of 0)
n
Proof. The loss function & for the regression problem (6) can be written
-f ((z),))
{(p) =
Σ
y/s
i=1
Σ yi
2
2
ΣVn 3/ f [EE +
+ n Σ
Σ
(f (f O(=,))
i=1
11
= Lo + 2 / v(0) v(0) p(d0)
U/(00,000 p(d))) p(d0").
where to = n -' i= v is some constant that does not depend on P, and the two functions V RRR - RR and
U D x RD + are given by
((0) = 1 Σ (x; and U(0,0') = Σ [ai [(;;, 0')
(7)
By (3), it suffices to compute the Euclidean gradients w.r.t. the variables p and E. We first compute V (p((pp) as follows:
v // /1110000000000000
(xb(p)=2) 2V
d0
@ 2 ((0) ((0)V 2) 11 p. E) do do
(
(( ((0)0~,(010
p. E) do' U(0,0) vn0(0 d0 6(0')
E) / U(8 Vod(O 14 E) do (00' |r [), E) de
2
(v)
VoV 018 de +2 /|/ VoU 6(0 (4 E) de] ((8' 14, Σ) de
2 [oV / VoU(0,00 60'' n, 1) 6(0 22
(8)
Here, step (i) uses the Leibniz integral rule (since U and are continuous and Φ is sufficiently smooth); step (ii) follows from the chain rule and the fact that U@0 0¹) U(' 9); step (iii) follows from (15) A Appendix A.3; and step (iv) follows from integration by parts. Following similar steps, can compute
15

---

Published in Transactions on Machine Learning Research (12/2025)
(p)(p) follows:
Vxf(p)= = / V(0) VZ0(0 14 d) do + / U(0,0) 83) Vro(0 n (4,
do 6(0'
14, E) do
(1)
V(0) V30 u. 2. do + / U(0,0) v8 010 n. 6(0' p4 23 d0'
/ VOV V00 nr d)
JU VOU(6,6) @ Voo(O) 14. Σ) do @(8' μ. Σ) d0'
(oV(O) 2 ()
(0;,, 6(0' n, I,, do' E-1
(6)
Voo 14, 2) VoV(0)
u4 E) a. 6(0' E)
1 Eー / p)) & [o)(0)((0))
- (0-
de a(0') p. E,E)dd do
Here, step (a) follows from (16) in Appendix A.3; steps (b) and (b') both follow from applying integration by
parts to the formula following step (a); whereas steps (c) and (c') both follow from (14) in Appendix A.3. According to (3), the Wassersessnn gradient flow for minimizing & is given by Pt NWr, Et), t 22 0, where
Ae vx((p))
Σ, -2 Vs(((p) Ex 25, Vz((pp)
In order to derive closed-form expression, we compute
T2 Σ Yi Vof(x,, 0) Σ yi
n
12
113 xc.)
VoV(0) =
and
(9)
w ( xz)) Iz
n
[u w (BT xT)) I.
x))
VoU(0,0')
Vef(z 0) f(E,,))
f(a,00))
(10)
where ((x) = 1xx 2 0} is the derivative of ReLU at any x * 0. Hence, we arrive at
14 = - (p((p)) = -2 -2 / / + (0U(0,0) (00 Im, E,) a8 6(0 pr. Σ) do
2 n2
Σ n
y1 Ep. (x; 0)) Ep. ()
'In general, the Wasserstein gradient measure z an element (p) and therefore only defined almost everywhere; Ambrosia al. (2008, Chapter for details.
16

---

Published in Transactions on Machine Learning Research (12/2025)
and
Σ = - -2 Vz((pp)
= =2 VoV(0) (4- (r: E.)
2 /1 V00(0,000(00-1)) 0( pt, 24. do) 6(0') pr, 21) do
/ (0 mm) 8 (2V(0) 60 Pr. 2.) do
2 /1/ mm) @ VoU 6(0 Mr, Ex) d0 610' ue, 2x) do
2 Σ y f 0) (Ep. Vof (0- ue) + Ep [(O uo) @ Vof(z..)))
This completes the derivation.
Although Theorem 2 derives equations for the Bures Fures-Wassersted gradient flow, the expressions involve expectations which must also be computed. We now proceed to show that these expectations can be computed in closed form for ReLU activations, which allows for exact implementation software. The following derivations are tedious, but the resulting equations are readily programmed..
We need to compute
1 x)) 1 xx) Fi
(xi, Ep
1 x;)],
Vo 0) E,
and
w
w (B
13Tx;) TT
V) Ep
w² l xz) 0z
W 13 x;)x, 8
Basically, we need to compute, for each 1 < i n,
A, Ep 133 xz):
B. Ep[w 1 T x;));
C, E, 1/3 x;)],
D. EPlw2 / x;)]]
and for each 1 i
Prj = Epl11B1 x.) B,),
Qij Ep|w (33 x;)
Then we can express
A B,,;
Ep Vol(,,
Exs(x;, = =C,
and
C, [Payhsssa
]
Ai
60 (0 () 1n]] =
V
2 n.
Dia;
I, 8 Qualsssa
Therefore, the update rule looks like
M4+1 = + 19t
and
Σ Eg + n/G,
where
2 n
A
9t
(Bi C₁)
B
71
17

---

Published in Transactions on Machine Learning Research (12/2025)
and
Σ (yi Ci) {
C. Pighsssa
Ai
G.
Bra: 8 u
Dag &
Qishisisa
1=1
2 n
M/ ct) {
[
C PPrllisssa
DaaT
A; B
}
Quilisssa @ 24
- 14 Ø
n
In addition, the objective function
[(pe)
1
Σ
1 -C,2
72
Let us first compute universal rule. Let X
Y
BT x4 and Z By. We have
x
μ1
of
P1,20102
) p2.00203 og
P1,30103
Y
N
P2
P1,20102
oz
(11)
Z
A3
11,30103
22,30203
where
A1
1120
of 01.
P1,20100 of
P1,30103
(at))
Eud
Eur
P2
Dy e
P2,30203
x DBA e, 3
x. T
n.. D.T
[4
ee a
01.301033
02,30203
03
oj 2
e, e,
We know that
cov(X - aY, Y) P1,20102 -o. =
when
a= P1,2
01
and
02
(Z 3 al) - ,Y, Y) P230203 702 =0
when
7 = P2,3
03
and
02
COV % - (x - aY) - 7Y, x a8) = cov x aY) 3 var X - -YY)= = 0
when
-1.30103-0122302000 032233203
of -20p1,201 +
Therefore X aY, Y, and Z - (x - at) yY mutually independent. We first compute A₁ and B,.
By direct computation,
A, max (Y,0) 122 02 E max
Y
P2
P2
02
02
000
I
12 dr 122 40
P2
P2 +
122/02 V2T
02
P2 02
de - xx2 3/12 12
p1 1
12
P1
+
02
06
122
P2 4
1/2
12/02 2
02
02
02
and
B. = - + +EaYl
-
= E| - P(Y > 0) + a
max max
= E,F, + aA,,
where we further define
E - aY| = 11-1000 aup and F= P(Y > 0) = + 2
18

---

Published in Transactions on Machine Learning Research (12/2025)
We then compute
C = E[x - E/w. + Ely max (Y,0)) A,E; + aG;
X = B max [Y E(aY max maxx(Y00) 0]]
where we define, for Y (N2,
x²
G = E max Y. 0}} B|y
e (x--2a)) /22)//(2222 1(2 dx
/2m02
(12 029)²
e- -2 dy
by change of variable y d2
I H1
-w2/02
V2T
= 12
1
12 12 dy + 2
00
1212 dy + c
000
o2y² - 5/ dy
E
122002 22T
M22002 22T
-D2/02
22T
$
U2/ 022]
+ 22222 (1- 122022)
y 12 00
+ + 1
12 dy
02 2x
-2//02
-psl02 2tT
(12 + o [1
क 112 02)]
+ 2222 Φ0
1 122
We also need to compute
D: = _| 1 o = [X Ty +
2a - -at)
F,11; - a² G; + 2aA,E;
where
= at = - 1.--201 -2 + a²0
(1i a2222
Then, we compute
P.. E max [Y,0)] = -Y) max (Y,0J) + B[TY max (Y,0))
=E >Y B|ax + ElY max (Y.0)) 03]
+
A,,M, GG,
where we let
Mej = B[Z - -Y) 133 714-
Finally, we compute
Qu -IX
+
E((X a1) - (2
- - - (Z 3x a(X-aY) -XX)
+ +I((X
+ +[(X aY)ly>0777 "\=
= E,NN, + a..N, + 3F,H + + AAE G.
where we define
Nr = E|z 3 - 7Y = P3 8pg + a.31/ - 7112
We have provided explicit expressions for all of the terms.
A.1.2 Gaussian mixtures
Consider a K-component Gaussian mixture distribution Pv parameterized by v:
K
1 K
Pu = (1((0) I()),
1
where
V' = K
( (20))
(200))))000)
K
A=1
Here v is a discrete probability measure over Fd Sd We start by deriving the Gaussian mixture gradient flow for general loss L.
19

---

Published in Transactions on Machine Learning Research (12/2025)
Theorem 3. Let & be functional over the Wasserstein space, The Gaussian mirture gradient flow (vi)120 for minimizing & inilialized at distribulion K E (nb"), 206)) with K aloms is given by
V/ = == Ek= 8 (n9 (1a0i,s s t 2 O, where for each &o E [K], the dynamics of (14 Jezo and (E) )220 are governed by the ODE system
( -Bxcuf" 201) Vol(pm))
5(k) -x(()) BN s (P - B s(tt) (PS((p)) ((k)
Here, 8L SpP refers to the first variation of L (see Santambrogio, 2015, Chapter 7).
Proof. We refer to Lambert et al. (2022, Appendix F) for background. We calculate the first variation of v + L(p)) in terms of the first variation of L: note that for any 8 >>> 0 and any measure X' satisfying JawrRa) dλ' = 0 and such that P + 8X for sufficiently small 8 > 0, if & = (n, 23 and Pe denotes N (n, 2), we have
[(vu4ox) -(pp = (L(()) d(pu+ox - Pv) p+)+ + = /1/ 82 pp((p) dps ox((E)) + 008)
which, by the definition of the first variation, shows that
8xl(p)): : I (b((p)) dps.
Since the Gaussian mixture gradient flow is, by definition, a Wasserstein gradient flow over the Bures Wasserstein space, the particle interpretation of Wasserstein gradient flows shows that each (ML x evolves by the Bures Bures-Wassersen Wasserstein gradient of (p,,) It follows from Lambert et al. (2022, Appendix C) that the Gaussian mixture flow takes the claimed form.
Alternatively, we can derive the gradient flow more explicitly. Noting that & = (1,, E), we have
V (o = / (l(p))))))))) Vo0(O n4 E) do -0 (C(pu)Voole n. E) do
/ VoDC(p)) 6(0 | d.
and
VrobC(p))(E)) / 68(p)) Vzr(( u1 d0 -
n, a0
(go)(()))dl 600 1a, E)
where we used the expressions in Appendix A.3. Recalling (3), it completes the derivation.
With Theorem 3 hand, we can now prove Theorem 1.
Proof of Theorem 1. Consider the loss function
Elu, C)
Ck. L(pr)
where
K
K Σ Tm)
with Σx = C.C
D =
Note that for any A Rdd and any & > O, we have
((m =A, .2...., MK, Ci...., Cx) E( ........ pk, C Cu....,CA (V C), 0
lim
8
20

---

Published in Transactions on Machine Learning Research (12/2025)
The left-hand side of the above equation can also be expressed
[(pv.) ((pw)
1
lim
where
V22 K
+od,5)) +
Fa)
2 2-
By the definition of first variation, we know that
/ ol(pe) d(pu (o) d(s)
[(pve) [(pv)
/ (p dps) ex((d)) + o(e) where
x= U/ - v
1 K [/ (p((v) dP(n
(k((w) dp(n) dpw d(E))
Taking the above three relations collectively yields
Em, 1 (V I dp(p dp K
Since the above equation holds for any A E Rd, we know that
Vi (), = K VM / (LB(p)) dp(p) = E.) VdL(pv))
where we used the expressions in Appendix A.3 and integration by parts. Therefore the Euclidean gradient flow w.r.t. ((*) is given by
4 = K E ( Vol((p))
(12)
Similarly, we can show that
1 Vy, / [00(0)) 2K Vr, -)= K
1 EN(H En) v28L(pp.)
which then leads to
C) C(C 1 Ey(nE (o5L(p)))
K
Hence the Euclidean gradient flow w.r.l. ((k) is given by
((k)
K E N(r x(1) (
therefore s)(k) c(())((())) satisfies
(*))_ (())) ((A)T + C.
K [BNoa 2it) VOLC(p.)) ) x1
1 (K)
[E x(t)) (oDL(p)))
(13)
K
By comparing (12) and (13) with the ODE system in Theorem 3, we can see that they are equivalent up to a scaling factor of K.
Then, we specialize to the objective function
min C(pr) = "57
(v: - f(x;,)) p())))
1=1
21

---

Published in Transactions on Machine Learning Research (12/2025)
Theorem 4. The Gaussian mirture gradient flow (vi)120 for minimizing & initialized at distribution 2 0, where for each (8) )120 and (z(")))20 governed by the ODE system t [K], the dynamics of (14
Vo K- K--EE 8 ( with atoms given by ut K- -K--E Eimi (144) sk (14"))))))
(2)
2
yb Eo,, f(z, 0)) (x((((,,())) v Vof(z,,0),
1=1
((()
(yi pp,, (zi,
m1
E N( 1 )(0) (V)((;; (0 me (x)) + (0 tee (x))) @ Vof(z; 0).
Proof- The first. variation is given by
: V(0)
where U and are (7). Then, we can apply Theorem 3, which seen yield the desired equations after substituting in the definitions of U and V and performing integration by parts.
As before, write out more explicit expressions for the gradient flow. Thankfully, calculations. Compared to Appendix A.1.1, we only need to compute
reuse previous
K
K
= [er, Epu -W 1 x;)l
(BTx:))
1 K
BN Mon) (B x;))
1
Σ
cl
K
k=1
Then, the update becomes
M'+1 Pi
(k) (k)
199t
(8)
and
(b)
241
->((A)
m)(((h)
where
2
A Bid
(y)
C
1=1
and
2 n
Σ (y: C)) {
C. [Padlessdd
A Bidis
Dees @ @Qishisisl
@ m
i=1
2
y c {
G. [
DaoT Quhhsssa 8
A; B
n
A.2 Multi-class classification
While our implementation relies on PyTorch's Automatic Differentiation engine, present here exact computations of gradients for multi-class classification. Their complexity indicates that it is largely preferable to employ automatic differentiation and that a study of first-order optimality conditions appears challenging.
For this setting, we focus on the parametrization as described Subsection 4.2, that is, 3 N(w E) and E|w | [3]
UB + v, where we apply the Euclidean gradient flow the parameters 1²,, U, v, and the square root of Σ. To simplify the notation, we write p = p (The sparse parametrization in Subsection 4.2 corresponds to further restricting Σ to be diagonal.)
We first compule the gradients w.r.l. U and v. Write
U1
U1
U
and
u
UL
22

---

Published in Transactions on Machine Learning Research (12/2025)
The following expectations are understood to be taken over the Gaussian mixture. We have
EV a (- x) V A3Tx) T B+ 1B x) = -E/u T 3 (B x) B + E (3 x)] u.
The loss function is
1 1
1
/ fe(o,,0) - log 1 + Σ exp / [((os)))
[()
Σ
Σ 1 wit
-lgg
71 i=1
t=1
where
(B
B] + up E 1B
We can compute
1
exp fe( 0) p(d))
(ip =
E (3 x) 1)
π1 i=1
1+ 1 exp f fer (x 8) p(d0)
1
exp f fe(x,, 0) p(d))
V vpp
1+ 1+E+= exp f fe (z,, 0) p( E 133 x)
n1
Next, we can compute
n
L
n= VH /
n
El exp( fe( 0) p(d0)) v fe(xs, 0) p(d0)
1 + exp fe 0) p(d0)
1+ 1 exp f fe 9) p(d()) } V / ()(()) p(d0)
exp J fe(z,, pdd0)
4 =
0)
n
and similarly.
1+ Σ -== exp s fe(x,, 0) p(d)) } v / f((x..0pp(00)
exp f fe 0) p(d0)
Vx((p) = n1
n=
t=1
We have
V/ / Je(os, 0) pldo) = V, 3 + V 2r E /3 x)
= =VV> x) u7
EV l x))
x)] ue + Ue E (BTT x))=.
We also have
V / fe(x.,0) = x3 B (
where
1
②)
u
3
E
-- I-1
2
Σ
un [1]
2
2)
ue) (B -4]] x-1
2
E-
1) & 3x 1B3 xx) :))s]]
2
23

---

Published in Transactions on Machine Learning Research (12/2025)
and
1 Σ E-1 1-13 μ)) & V 13 x)]
V& (B"
E / x)
0 (3 00-
2 E
= - p)8
Hence, we have
L
1 Σ exp J Se(r,, 0) p(d0) V. / ()
exp J fe(x,, 0) p(d0)
# = (. =
exp S Se( 0) p(d0)
1 1=1 exp S Se (Ti, 0) p(d))
(3 x) u 3 x + E8 (B x)] us + ue (B x]]=]
exp f (11, 0) p(d0)
1+ Σド= exp f Se(a;, 0) p(d0)
| / x;) 2| + B 21)) ue 2s 1B x]
and
= -
Σ
exp f Je(#;, pld0)
1+ E exp f fe p(do)
E
(3 x2) 4 3x, (3 x uu 1 1 xz) 2 (B 18-m)) [)]
+ E
x,) uu B x2 x3 uet 1 (3 (3 @ (B - -+]] }
Let
A; E 1(3 x.).
F= (33Tz))
Ra BB B,1.
and
Pag = E 13 (3) 211111
== x.) 1u 33,.
S.t Elwi 1 x2,]
We have
exp S..
[1g Ps 20 Asue + us P.t:]
y m
71
+ /'
exp Si,
and
1
cxp S
n
1 - exp S..
L-1
x [1x 2 Qie + us @ P.+ Pr
un R3 I, @ p - A 8 H Fat; @ p)+
+ (x4 8 Q1. + 2 P. + 3 R. i
T R2 x4 @ μ
A 2 n- Fxx,
We also have
1 1
S 1+ - exp S
V C())T
1 n
Ú-
71
VUL T
exp 1-1
1==--1
1+ Σ r1 exp S.
24

---

Published in Transactions on Machine Learning Research (12/2025)
and
exp St
Vny (p)T
1=
1+ = c'== exp Ste
a n
8 6=
Az-
-
exp.S; 1+ _e'=1 exp St,e
Next, compute each quantity. As before, consider (11), except now define X up B. This time, have
Ug μ
of P1,20102 01,30103
11 of
P1.30
U
P2
J, T
1
P2,30203
=
Σ u8
Zi
Egl
[43
e 1
P2,0200
og
Then, as before, we have
A = m[max(YYO)
F = P(Y > 0),
R1, = EI 1y
and
P,, M[max (Y, 0]
and
Q. = B[
The only new quantity compute is
J E[Zily y 0ol = E[(z - - >ol + E/yY1y>a
E|Z yY P(Y > 0) 7)ax)
-FMMi 7Ai
and
S = Elwe 1 x4)] = B x2) u 3 E 1 x.) = C2+ U.A
A.3 Auxiliary Gaussian calculus
Recall that of- μ, 2 is the density function of N 2) Then we can compute the gradient and Hessian w.r.t. the variable 0:
0- - I-=--1 # 2 1 -
V Ved(0) pr, E) 6( 14, Σ) Vo
2
=6 14, EE Σ (0 -1)) 11),
(14)
Vi0(0) | 12 = -0 μ)) Voo 11, 1لل - 6 11 2 -
=( μ4, E) (0 11) - Φ(0 μ. Σ) Σー
In addition, also compute the gradient w.r.t. the parameters and E:
vn0(0 n. -) 6(0 1 M. 2) V ( 1000000~~~00000 4 p T 1)
= = | 14 2) -- -4) -0 ,, E)
(15)
and
V (000|1E)) E)
V exp - (0 - T -- - m)]
V(2T) (d det Σ
2
1
1
+ exp
(0 p1) T Σ (0 - [1) V
2
V det E
1
1
2 5(0 14, E) E 10 1) 1) T 6(0 14, E)
V det
det E
1
610
11,
Σ-E))-((00 (0 n) Tx- 2 60 14, Σ
2
V30(0)
(16)
2
25

---

Published in Transactions on Machine Learning Research (12/2025)
B Experimental details
The numerical experiments conducted in this paper are implemented with PyTTrrc in Python, using 2023 MacBook Pro with Apple M2 Pro chip and 32GM memory. The fully-connected layers are implemented using Py Torch'ss built-in functions. The G.M layers are implemented using the derivation in Appendix (thanks to PyTorch's Automatic Differentiation engine, we only need to implement the loss function, and there is no need implement the gradients explicitly).
The error bars in Figures 2, 5, 6 and represent standard error computed 5 independent trials. When implementing two-layer GM network, the output of the first GM layer is normalized to have unit norm, before being sent to the second layer as input.
Finally, we include the architecture of the two network structures used in the CIFAR experiments. The VGG-like network structure is as follows:
Conv Layer 1: Input: 32 x 32 x 3 CIFAR-10 image; Convolution: 64 filters, 3 x 3 kernel, padding 1; Output: 32 x 32 x 64; Batch Normalization, ReLU Activation.
Conv Layer 2: Input: 32 x 32 64; Convolution: 128 filters, 3 kernel, padding 1; Output: 32 x 32 128;
Batch Normalization; ReLU Activation; Max Pooling: x 2; Output: 16 16 x 128.
Conv Layer 3: Input: 16 x 16 x 128; Convolution: 256 filters, 3 3 kernel, padding 1; Output: 16 x 16 256;
Batch Normalization; ReLU Activation.
Conv Layer 4: Input: 16 x 16 x 256; Convolution: 256 filters, 3 kernel, padding 1; Output: 16 16 x 256;
Batch Normalization; ReLU Activalion; Max Pooling: x 2; Output: 8 x 8 256.
Fully Connected Layer 1: Input: 16384 (after flattening); Output: 256; ReLU activation; Dropout: 0.5.
Fully Connected Layer 2: Input: 256; Output: 128; ReLU activation.
Fully Connected Layer 3: Input: 128; Output: 10 (for CIFAR-10) or 100 (for CIFAR-100).
To incorporate the GM layer into the above structure, we replace the Fully Connected Layers 2 and 3 with GM layer (10 components). When training these two networks, we use data augmentation techniques including random horizontal flips and random cropping.
26


Reviews:
[{"review_id": "6WP9HjjlkY", "reviewer": "Reviewer_nXho", "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": "", "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": "Summary of Contributions:\nThis paper introduces Gaussian Mixture (GM) layers as a new neural network component inspired by mean-field theory. Instead of representing infinite-width layers through empirical measures, the authors approximate weight distributions with Gaussian mixtures and derive training dynamics via Wasserstein gradient flows. The formulation reduces to Euclidean updates on mixture parameters, making the approach computationally tractable. Proof-of-concept experiments on MNIST and Fashion-MNIST show that GM layers can achieve accuracy comparable to fully connected networks while exhibiting distinct training dynamics and enabling interpretability through visualizations of Gaussian components.\n\nStrengths\n- Novel and conceptually clear idea bridging mean-field theory and practical architectures.\n- Provides interpretability via Gaussian mixture components and ensemble visualizations.\n- Offers a new perspective on layer design that is computationally tractable.\n- Well-written and easy to follow, with clear motivation and exposition.\n\nWeaknesses\n- Empirical validation is limited to small-scale datasets (MNIST, Fashion-MNIST).\n- Reproducibility is uncertain without released code.\n- Competitiveness against stronger baselines or on larger datasets is not established.\n- Performance may be sensitive to the number of mixture components and initialization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper’s claims are supported by accurate and convincing empirical evidence. The proposed Gaussian Mixture layers are derived from mean-field theory and shown to be computationally tractable through their Euclidean parameter updates. Proof-of-concept experiments on MNIST and Fashion-MNIST demonstrate that the method achieves performance comparable to fully connected networks while exhibiting distinct training dynamics and offering interpretability. Although the empirical scope is limited to small-scale datasets, the claims are clearly framed as proof-of-concept, and the evidence presented is consistent with that framing.", "audience_yn": "Yes", "audience_explanation": "This work will interest a portion of the TMLR audience because it introduces a new type of neural network layer that directly links mean-field theory with practical implementation. The approach is conceptually novel, provides interpretability through Gaussian components, and demonstrates promising results on benchmark datasets. Readers focused on neural network theory, architecture design, and interpretable models would find these findings valuable, even if the experiments are limited in scale.", "requested_changes": "- Release code or detailed reproducibility materials to ensure independent verification of results.\n- Expand the discussion of limitations, especially regarding scalability to larger datasets and sensitivity to hyperparameters such as the number of mixture components and initialization.\n- Add experiments or analysis on more challenging datasets (e.g., CIFAR-10 or CIFAR-100) to demonstrate broader applicability.\n- Provide a sensitivity or ablation study showing the effect of different numbers of mixture components and initialization strategies on accuracy and stability.\n- Include runtime or computational cost comparisons with standard fully connected layers to quantify trade-offs.\n- Add further visualizations of Gaussian components and decision boundaries to illustrate interpretability benefits.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "iC7HhVVard", "reviewer": "Reviewer_3QH6", "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": "", "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": "In this article, the authors introduce Gaussian mixture layers (GM layers) as a novel layer type for neural network architectures, combining concepts from mean-field theory and Wasserstein gradient flows. The study demonstrates that traditional fully connected layers can be effectively replaced by GM layers. Experiments indicate that networks utilizing GM layers can achieve similar performance to classical two-layer fully connected networks.\n\nThe primary contribution lies in modeling the distribution of weights in an infinite-width fully connected layer through a Gaussian mixture model, offering a new perspective on layer design.\n\nStrengths:\n\n- The utilization of mean-field theory in a prescriptive manner, rather than purely analytical, presents an interesting approach.\n\n- The concept explored in the paper is academically interesting and may hold future potential in advancing neural network architectures.\n\nWeaknesses:\n\n- Despite the promise of GM layers, additional empirical evidence is necessary to see the specific advantages of Gaussian mixtures over traditional fully connected layers (even though it's not about performance; what are other benefits?)\n\n- Certain sections regarding the implementation and training dynamics of GM layers lack clarity\n\n- A clear and compact one-to-one comparison with fully connected layers is missing\n\n- The authors state that “the training dynamics of networks with GM layers are not sensitive to initialization,” while also suggesting a need for future research into better initialization schemes. This contradictory statement raises questions about the necessity of further investigation if the training dynamics are indeed not sensitive.\n\nMinor Issues:\n\n- The article can occasionally be difficult to follow, and greater elaboration in certain areas would be beneficial.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The claims made in the paper are generally supported by robust theoretical foundations and experimental results. However, it would be better to write Theorem 1 in a more formal way.", "audience_yn": "Yes", "audience_explanation": "The paper is well-suited for readers interested on innovations in neural network architecture and probabilistic training strategies.", "requested_changes": "Address the stated weaknesses and include a table that outlines the major differences between GM layers and fully connected layers, such as the number of parameters, training methods, memory, etc.", "broader_impact_concerns": "At this stage, assessing the broader impact is challenging, as the paper does not clearly highlight any benefits of the new layer.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BL6U6YQM4I", "reviewer": "Reviewer_8U2b", "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": "", "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": "The paper explores the role of theoretical tools, such as mean-field analysis, in discovering new architectures that are true to the expected theoretical dynamics predicted. The authors introduce a new Gaussian Mixture (GM) Layer to completely replace a dense hidden layer in a standard multi-layer perceptron, while remaining true to the Wasserstein gradient flow requirements of mean-field theory. This is shown to be practically feasible for a 2-layer neural network, under certain assumptions of sparsity and compositeness. \n\nThe experiments are limited to comparison with only dense layers (no convolution, attention, or recurrences), and also only up to 2 hidden GM layers. The practicality of a real deep and wide network with GM layers is considered beyond the scope of this work. In this regard, the GM layer represents a successful new architecture design (for the 2-layer theoretical playground) that utilizes the mean-field parameterization, admitting feature learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "*Edited after the author's response to review.*\n\nAlthough empirical experiments can be stronger, the primary objective of the paper appears to be met: proof of concept for GMMs as a practical layer option.\n\n\n~~Please note that the above \"No\" is completely hinged on the points below.~~\n\n1. ~~Firstly, to draw attention to Sections 2-3 being extremely crucial and requiring more explanations, and aggressive referencing of the Appendix and the theorems. Crucial ideas and notations are there, which go a long way in clarifying these sections, especially Section 3.~~\n\n    ~~Some of the figures can be reduced in size, be tighter to gain paper space. The recommendation is to update Sections 2 & 3 with this afforded extra space. Please refer below to *Requested Changes* for some direct suggestions (esp. points 5-11).~~ \n\n    ~~Overall, the flow of Sections 1-3 could be improved.~~\n\n2. ~~The exact motivation for the choice of Gaussian Mixture Models (GMMs) as the probability measure for the neuronal distribution is unclear. The resulting theory certainly validates the choice of these GMMs and shows the empirical ease of optimizing GMM parameters. However, it is unclear if GMMs were chosen *for* the convenience of re-parameterization of GMM parameters enabling the relevant gradient flow in the relevant space, **or** if the theory emerges from a trivial choice of GMMs for $\\rho$ in Equation 1.~~ \n\n~~Having the above clarified suitably will lead to the updating of \"No\" above.~~", "audience_yn": "Yes", "audience_explanation": "Anybody interested in Deep Learning Theory could be interested in the theorems and proofs offered, especially for the connection between the Wasserstein gradient flow in the space of probability measures, and the Euclidean flow in the parameter space, and the optimization of the GM layer's mixing parameters under a suitable reformulation.\n\nThis is also interesting for both manual designers and search-based seekers of new neural architecture designs.\n\nThe empirical analysis for *feature learning* was another contribution of the paper for the application of Neural Tangent Kernel (NTK) and $\\mu$Parameterization.", "requested_changes": "Please note that the questions below imply that the manuscript either requires or could do with extra clarification on those points. The authors are not obligated to answer them here to the reviewer (except for the points above for the *claims*).\n\n\nQuestions from reviewer / Clarifications required in the manuscript:\n1. Pg 1, para 3: `under a natural scaling` requires clarification.\n2. Pg 2: Could the contributions be enumerated, please?\n3. Pg 3, after Eq. 1: please clarify $\\delta$?\n4. Pg 3, 3 paras after Eq. 1, until Sec. 3: in the first sentence, unclear if `Wasserstein gradient flow` and the `Wasserstein gradient flow picture`\n5. Pg 3, Sec. 3:\n  * The motivation for a GMM to represent $\\rho$ is not quite clear here, despite the paras after Eq. 2 attempting to do so\n  * \"As discussed in Section 2, it is well known from mean-field theory that...\": would be good practice to cite/refer for the \"well known\" claim here\n  * For the empirical observation that neuronal distributions tend to cluster after training, it is unclear how this alone justifies the use of GMMs to represent $\\rho$\n  * Clearly, GMMs make intuitive sense, and the GMM math allows for the gradient flow analysis better; if this requirement was a motivator for GMMs, then it needs to be mentioned more clearly\n  * What is the authors' opinion on whether at least one other alternative representation of $\\rho$ should be evaluated empirically (e.g. non-Gaussian kernels)?\n6. Pg 4, Theorem 1: Please use content (or refer) from the Appendix more liberally to support the statements here. Parsing the section helped clarify many parts of this theorem. \n  * This section could do with simpler explanations or notations, since the primary intuitive connection the reader needs to make here is how the $K \\ge 2$ case easily can be brought to the $K = 1$ case with a simple reformulation.\n7. Pg 4, Sec. 4.2: Are there other choices of sparsity not considered here?\n8. Pg 4-5, Sec. 4.2: could do with clarifications, references, and notations:\n  * The affine relation for the conditional means w.r.t Eq. 3 comes from? (e.g., mentioning the law of total expectation would help the reader, or other relevant detail(s)).\n  * Could the notation $\\mu^{\\beta}_k$ be simplified to $\\mu_k$, or, to avoid overloading of $\\mu_k$ with the general Equation 1, can we be consistent and let $\\sigma_k$ be $\\sigma^{\\beta}_k$? \n  * (last para) The claim of positive semi-definiteness is from the role of the trained $\\sigma$ or an outcome of the diagonal sparsity assumption? (here, can suitably refer to Appendix A, Pg 13).\n10. Pg 5, 4.3: given this part comes right after the conclusion of 4.2 describing the training of a GM layer, is it worthwhile to have a comment here on how the backpropagation for deeply stacked GMs layer work out (not detailed proof, but suitable enough for intuition, s.t. the reader to thinks and understands it, right now, there is no mention of it at all).\n  * especially, how does the approximation of 4.2 work out when the $x$ in Eq. 4 depends on intermediate layers, with a deeper compositeness?\n11. Pg 5, Sec 5: \n  * Motivation for different learning rates and their choices?\n  * Motivation for initializing $\\mu^{\\beta}, U, v$ from the same distribution parameterized by a $\\gamma$. \n  * How is the value for $\\gamma$ ($=1/2$) decided upon?\n  * Why or what limits the marginal gains with increasing $K$ (mixture components), mathematically or intuitively?\n12. Pg 6, Fig. 3: How is the `NN: random init from GM` actually constructed? Does one need to initialize a GM Net using K=20 and $\\gamma=1/2$, and then sample every weight parameter for a dense NN from this parameterized GM?\n  * What does one actually conclude here, explaining the worse generalization performance for relatively *simpler* datasets such as MNIST and Fashion-MNIST?\n13. Pg 7, first sentence: How is this conclusion arrived at? \"Training dynamics of networks with GM layers are not sensitive to initialization.\"\n14. Pg 7, para 2, last sentence: unclear from here and Figure 4, how to conclude and verify if fully connected layers do not show \"feature learning\" or movement away from 0.\n15. Pg 8, para 2, last sentence: could there be more clarification here as to why a poorer performance of random $\\beta$ be enough to conclude that feature learning is absent? \n\n\n*suggestion*: quantifying any differences in memory footprint and inference speed, given a net reduction in the number of parameters, would be an interesting addition to the empirical analysis, especially as a *new* architecture design.", "broader_impact_concerns": "There are none.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)