Proving Linear Mode Connectivity of Neural Networks
via Optimal Transport

Damien Ferbach13† Baptiste Goujaud2 Gauthier Gidel1‡ Aymeric Dieuleveut2

Mila, Université de Montréal1 CMAP, Ecole Polytechnique, IPP2 ENS Paris, PSL3

###### Abstract

The energy landscape of high-dimensional non-convex optimization problems is
crucial to understanding the effectiveness of modern deep neural network
architectures. Recent works have experimentally shown that two different
solutions found after two runs of a stochastic training are often connected by
very simple continuous paths (e.g., linear) modulo a permutation of the
weights. In this paper, we provide a framework theoretically explaining this
empirical observation. Based on convergence rates in Wasserstein distance of
empirical measures, we show that, with high probability, two wide enough two-
layer neural networks trained with stochastic gradient descent are linearly
connected. Additionally, we express upper and lower bounds on the width of
each layer of two deep neural networks with independent neuron weights to be
linearly connected. Finally, we empirically demonstrate the validity of our
approach by showing how the dimension of the support of the weight
distribution of neurons, which dictates Wasserstein convergence rates is
correlated with linear mode connectivity.

##  1 Introduction and Related Work

Training deep neural networks on complex tasks is a high-dimensional, non-
convex optimization problem. While stochastic gradient-based methods (i.e.,
SGD and its derivatives) have proven highly efficient in finding a local
minimum with low test error, the loss landscape of deep neural networks (DNNs)
still contains numerous open questions. In particular, Goodfellow et al.
(2014) try to find ways to connect two local minima reached by two independent
runs of the same stochastic algorithm with different initialization and data
orders. This problem has applications in diverse domains such as model
averaging (Izmailov et al., 2018; Rame et al., 2022; Wortsman et al., 2022),
loss landscape study (Gotmare et al., 2018; Vlaar and Frankle, 2022; Lucas et
al., 2021), adversarial robustness (Zhao et al., 2020) or generalization
theory (Pittorino et al., 2022; Juneja et al., 2022; Lubana et al., 2023).

An answer to this question is the _mode connectivity phenomenon_. It suggests
the existence of a continuous low-loss path connecting all the local minima
found by a given optimization procedure. The mode connectivity phenomenon has
extensively been studied in the literature (Goodfellow et al., 2014; Keskar et
al., 2016; Sagun et al., 2017; Venturi et al., 2019; Neyshabur et al., 2020;
Tatro et al., 2020) and _non-linear connecting paths_ have been evidenced for
DNNs trained on MNIST and CIFAR10 by Freeman and Bruna (2016); Garipov et al.
(2018); Draxler et al. (2018).

##### (Linear) mode connectivity.

Formally, let $A:=\hat{f}(.,\theta_{A})$ and $B:=\hat{f}(.,\theta_{B})$ two
neural networks sharing a common architecture $\hat{f}$. They are parametrized
by $\theta_{A}$ and $\theta_{B}$ after training those networks on a data
distribution $P$ with loss $\mathcal{L}$, i.e. by minimizing
$\mathcal{E}(\theta):=\mathbb{E}_{(x,y)\sim
P}[\mathcal{L}(\hat{f}(x,\theta),y)]$ over $\theta$. Let $p$ be a continuous
path connecting $\theta_{A}$ and $\theta_{B}$, i.e. a continuous function
defined on $[0,1]$ with $p(0)=\theta_{A}$ and $p(1)=\theta_{B}$. Frankle et
al. (2020); Entezari et al. (2021) defined the _error barrier height_ of $p$
as
$\sup_{t\in[0,1]}\mathcal{E}\left(p(t)\right)-\left((1-t)\mathcal{E}\left(\theta_{A}\right)+t\mathcal{E}\left(\theta_{B}\right)\right)$.
The two found solutions $\theta_{A}$ and $\theta_{B}$ are said to be _mode
connected_ if there is a continuous path with zero error barrier height
connecting them. Furthermore, if $p$ is linear, that is
$p(t)=(1-t)\theta_{A}+t\theta_{B},\,$ $\theta_{A}$ and $\theta_{B}$ are said
to be _linearly mode connected (LMC)_.

Figure 1: Permuting the neurons in the hidden layer of network $B$ to align
them on network $A$

##### Permutation invariance.

Recently, Singh and Jaggi (2020); Ainsworth et al. (2022) highlighted the fact
that the units in a hidden layer of a given model can be permuted while
preserving the network’s functionality. Figure 1 shows how one can permute the
hidden layer of a two-layer network to match a different target network
without changing the source function. From now on, we will understand LMC
modulo permutation invariance, i.e. two networks $A,B$ are said to be linear
mode connected whenever there exists a permutation of neurons in each hidden
layer of network $B$ such that the linear path in parameter space between
network $A$ and $B$ permuted has low loss.

##### Linear mode connectivity up to permutation.

Singh and Jaggi (2020) proposed to use optimal transport (OT) theory to find
soft alignment providing a “good match” (in a certain sense) between the
neurons of two trained DNNs. Furthermore, the authors propose ways to fusion
the aligned networks together in a federated learning context with local-
steps. Ainsworth et al. (2022) further experimentally studied linear mode
connectivity between two pre-aligned networks. The authors first align network
B’s weights on the weights of network A before connecting both of them by a
linear path in the parameter space. They notably achieved zero-loss barrier
for two trained Resnets with SGD on CIFAR10. Moreover, their experiments
strongly suggest that the error barrier on a linear path gets smaller for
wider networks, with a detrimental effect of big depth.

##### Prior theoretical explanations.

A recent work by Kuditipudi et al. (2019) shows that dropout stable networks
(i.e. networks that are functionally stable to the action of randomly setting
a fraction of their weights and normalizing the others) exhibit mode
connectivity. Shevchenko and Mondelli (2020) use a mean field viewpoint to
show that wide two-layer neural networks trained with SGD are dropout stable
and hence show (non-linear) mode connectivity for two-layer neural networks in
the mean field regime (i.e. one single wide hidden layer). Finally Entezari et
al. (2021) show that two-layer neural networks exhibit linear mode
connectivity up to permutation at initialization for parameters initialized
following uniform independent distribution properly scaled. They highlight
that this result could be extended to networks trained in the Neural Tangent
Kernel regime where parameters stay close to initialization (Jacot et al.,
2018).

##### Contributions.

This paper aims at building theoretical foundations on the phenomenon of
linear mode connectivity up to permutation. More precisely, we theoretically
prove this phenomenon arises naturally on multi-layer perceptrons (MLPs),
which goes beyond two-layer networks on which theoretical works focused so
far. We also provide a new efficient way to find the right permutation to
apply on the units of a neural network’s layer. The paper is organized as
follow:

  * •

In Section 3, we focus on two-layer neural networks in the mean field regime.
While Shevchenko and Mondelli (2020) proved _non-linear_ mode connectivity in
this setting; we go further by proving _linear mode connectivity up to
permutation_. Moreover, we provide an upper bound on the minimal width of the
hidden layer to guarantee linear mode connectivity.

  * •

In Section 4, we use general OT theory to exhibit tight asymptotics on the
minimal width of a multilayer perceptron (MLP) to ensure LMC.

  * •

In Section 5, we apply our general results to networks with parameters
following sub-Gaussian distribution. Our result holds for deep networks,
generalizing the result of Entezari et al. (2021) with better bounds. We shed
light on the dependence in the dimension of the underlying distributions of
the weights in each layer and explain how it connects with previous empirical
observations (Ainsworth et al., 2022). Using a model of approximately low
dimensional weight distribution as a proxy of sparse feature learning, we
yield more realistic bounds on the architectures of DNNs to ensure linear mode
connectivity. We therefore, show why LMC is possible after training and how it
depends on the complexity of the task. Finally we unify our framework with
dropout stability.

  * •

In Section 6, we validate our theoretical framework by showing how the
implicit dimension of the weight distribution is correlated with linear mode
connectivity for MLPs trained on MNIST with SGD and propose a new weight
matching method.

##  2 Preliminaries and notations

Notations. Let two multilayer perceptrons (MLP) $A$ and $B$ with the same
depth $L+1$ ($L$ hidden layers), an input dimension $m_{0}$, intermediate
widths $m_{1},...,m_{L}$ and an output dimension $m_{L+1}$. Given $2(L+1)$
weights matrices $W_{A,B}^{1,...,L+1}$, and a non-linearity $\sigma$, we
define the neural network function of network A by $\hat{f}_{A}$ (respectively
$\hat{f}_{B}$): $\forall x\in{\mathbb{R}}^{m_{0}}$,

$\hat{f}_{A}(x):=\hat{f}(x;\theta_{A}):=W_{A}^{L+1}\sigma\left(W_{A}^{L}\ldots\sigma(W_{A}^{1}x)\right)$
(1)

To $W_{A}^{\ell}\in{\mathcal{M}}_{m_{\ell},m_{\ell-1}}({\mathbb{R}})$ we
associate $\hat{\mu}_{A,\ell}$ the empirical measure of its rows
$[W^{\ell}_{A}]_{i:}\in{\mathbb{R}}^{m_{\ell-1}}$:
$\frac{1}{m_{\ell}}\sum_{i=1}^{m_{\ell}}\delta_{[W^{\ell}_{A}]_{i:}}$ which
belongs to the space of probability measures
${\mathcal{P}}_{1}({\mathbb{R}}^{m_{\ell-1}})$, where $[W^{\ell}_{A}]_{i:}$ is
the $i$-th row of the matrix and $\delta$ denotes the Dirac measure. Note that
$[W^{\ell}_{A}]_{i:}$ is also the weights vector of the $i$-th neuron of the
layer $\ell$ of network $A$. Given an equi-partition111All subsets have the
same number of elements
${\mathcal{I}}^{\ell-1}=\\{I^{\ell-1}_{1},...I^{\ell-1}_{\tilde{m}_{\ell-1}}\\}$
of $[m_{\ell-1}]$ we denote
$W_{A}^{{\mathcal{I}}^{\ell-1}}\in{\mathcal{M}}_{m_{\ell},\tilde{m}_{\ell-1}}({\mathbb{R}})$
the matrix issued from $W_{A}^{\ell}$ where we have summed the columns being
in the same set of the partition ${\mathcal{I}}^{\ell-1}$. In that case
$\hat{\mu}_{A}^{{\mathcal{I}}^{\ell-1}}\in{\mathcal{P}}_{1}\left({\mathbb{R}}^{\tilde{m}_{\ell}}\right)$
denotes the associate empirical measure of its rows.

Denote
$\phi_{A}^{\ell}(x):=\sigma\left(W_{A}^{\ell}\ldots\sigma(W_{A}^{1}x)\right)$
(respectively $\phi^{\ell}_{B}$) the activations of neurons at layer $\ell$ of
network $A$ on input $x$. The data $x$ follows a distribution $P$ in
${\mathbb{R}}^{m_{0}}$.

Given permutations matrices $\Pi_{\ell}\in{\mathcal{S}}_{m_{\ell}},$222We use
interchangeably ${\mathcal{S}}_{m}$ to denote the space of permutations of
$\\{1,\ldots,m\\}$ and the corresponding space of permutations matrices. Given
$\pi\in{\mathcal{S}}_{m}$ its corresponding permutation matrix $\Pi$ is
defined as $\Pi_{ij}=1\iff\pi(i)=j$. $\ell=1,\ldots,L$ of each hidden layer of
network $B$, the weight matrix at layer $\ell$ of the permuted network $B$ is
$\tilde{W}_{B}^{\ell}:=\Pi_{\ell}W_{B}^{\ell}\Pi_{\ell-1}^{T}$ and its new
activation vector is
$\tilde{\phi}_{B}^{\ell}(x):=\Pi_{\ell}\phi^{\ell}_{B}(x)$. Finally, $\forall
t\in[0,1]$ we define $M_{t}$ the convex combination of network $A$ and $B$
permuted, with weights matrices $tW^{\ell}_{A}+(1-t)\tilde{W}^{\ell}_{B}$ and
$\phi_{M_{t}}^{\ell}$ its activations at layer $\ell$.

Preliminaries. We consider networks $A$ and $B$ to be independently chosen
from the same distribution $Q$ on parameters. This is coherent with
considering two networks initialized independently or trained independently
with the same optimization procedure (§3). We additionally suppose the choice
of $A$ and $B$ to be independent of the choice of $x\sim P$, which is valid
when evaluating models on test data not seen during training. We denote
$\mathbb{E}_{Q},\mathbb{E}_{P},\mathbb{E}_{P,Q}$ expectations with respect to
the choice of the networks, the data, or both.

To show linear mode connectivity of networks $A$ and $B$ we will show the
existence of permutations $\Pi_{1},...,\Pi_{L}$ of layers $1,...,L$ that align
the neurons of network $B$ on the closest neurons weights of network $A$ at
the same layer as shown in Figure 1. In other words, we want to find
permutations that minimize for each layer $\ell\in[L]$ the norm
$\|W_{A}^{\ell}-\Pi_{\ell}W_{B}^{\ell}\Pi_{\ell-1}^{T}\|_{2}$. Recursively on
$\ell$, we solve the following optimization problem:

$\begin{split}\Pi_{\ell}&=\operatorname*{arg\,min}_{\Pi\in{\mathcal{S}}_{m_{\ell}}}\|W_{A}^{\ell}-\Pi
W_{B}^{\ell}\Pi_{\ell-1}^{T}\|_{2}^{2}\\\
&=\operatorname*{arg\,min}_{\pi\in{\mathcal{S}}_{m_{\ell}}}\frac{1}{m_{\ell}}\sum_{i=1}^{m_{\ell}}\|[W_{A}^{\ell}]_{i:}-[W_{B}^{\ell}\Pi^{T}_{\ell-1}]_{\pi_{i}:}\|_{2}^{2}\end{split}$
(2)

For each layer, the problem can be cast as finding a pairing of weights
neurons $[W_{A}^{\ell}]_{i:}$ and $[W_{B}^{\ell}\Pi_{\ell-1}^{T}]_{\pi_{i}:}$
to minimize the sum of their Euclidean distances. It is known as the Monge
problem is the optimal transport literature Peyré et al. (2019). More
precisely Equation 2 can be formulated as finding an optimal transport plan
corresponding to the Wasserstein distance between the empirical measures of
the rows of $W_{A}^{\ell}$ and $W_{B}^{\ell}\Pi_{\ell-1}^{T}$. We provide more
details about this connection between Equation 2 and optimal transport in
Section A.2. In the following, the $p-$Wasserstein distance will be denoted
${\mathcal{W}}_{p}(\cdot,\cdot)$ and defined with the underlying distance
$\|\cdot\|_{2}$ unless expressed otherwise.

By controlling the cost in Equation 2 at every layer, we show that the
permuted s of networks $A$ and $B$ are approximately equal. Linearly
interpolating both networks will therefore keep activations of all hidden
layers unchanged except the last layer which acts as a linear function of the
interpolation parameter $t\in[0,1]$.

##  3 LMC for Two-Layer NNs in the Mean Field Regime

We will first study linear mode connectivity between two two-layer neural
networks independently trained with SGD for the same number of steps.

###  3.1 Background on the Mean Field Regime

We will use some notations from Mei et al. (2019) and consider a two-layer
neural network,

$\hat{f}_{N}(x;\theta)=\frac{1}{N}\sum_{i=1}^{N}\sigma_{*}(x;\theta_{i})$ (3)

parametrized by
$\theta_{i}=(a_{i},w_{i})\in{\mathbb{R}}\times{\mathbb{R}}^{d}$ and where
$\sigma_{*}(x;\theta_{i})=a_{i}\sigma(w_{i}x)$. The parameters evolve as to
minimize the following regularized cost $R_{N}(\theta)=\mathbb{E}_{(x,y)\sim
P}[(y-\hat{f}_{N}(x_{;}\theta))^{2}]+\lambda\|\theta\|_{2}^{2}$. Define noisy
regularized stochastic gradient descent (or noiseless regularization-free when
$\lambda=0,\tau=0$) with step size $s_{k}$, and i.i.d. Gaussian noise
$g^{k}\sim{\mathcal{N}}(0,I_{d})$:

$\displaystyle\theta_{i}^{k+1}=(1-2\lambda s_{k})\theta_{i}^{k}$ (SGD)
$\displaystyle+2s_{k}(y_{k}-\hat{f}_{N}(x_{k};\theta^{k}))\nabla_{\theta}\sigma_{*}(x_{k};\theta_{i}^{k})+\sqrt{\tfrac{2s_{k}\tau}{d}}g_{i}^{k}$

It will be useful to consider
$\rho_{N}^{k}:=\tfrac{1}{N}\sum_{i=1}^{N}\delta_{\theta^{k}_{i}}$ the
empirical distribution of the weights after $k$ SGD steps. Indeed some recent
works (Chizat and Bach, 2018; Mei et al., 2018, 2019) have shown that when
setting the width $N$ to be large and the step size $s_{k}$ to be small, the
empirical distribution of weights during training remains close to an
empirical measure drawn from the solution of a partial differential equation
(PDE) we explicit in Section B.1. Especially, the parameters
$\left\\{\theta_{i}^{k},i\in[N]\right\\}$ evolve approximately independently.

###  3.2 Proving LMC in the mean field setting

Define respectively the alignment of a neuron function on the data and the
correlation between two neurons:

$\displaystyle
V(\theta_{1}):=av(w):=-{\mathbb{E}}_{P}[y\sigma_{*}(x;\theta_{1})]$
$\displaystyle
U(\theta_{1},\theta_{2}):=a_{1}a_{2}u(w_{1},w_{2}):={\mathbb{E}}_{P}[\sigma_{*}(x;\theta_{1})\sigma_{*}(x;\theta_{2})]\,.$

and we for $\varepsilon>0$ fixed we note the step size

$s_{k}=\varepsilon\xi(k\varepsilon)$ (4)

where $\xi$ is a positive scaling function. The underlying training time up to
step $k_{T}$ is defined as $T:=\sum_{k=1}^{k_{T}}s_{k}$. We now state the
standard assumptions to work on the mean field regime (Mei et al., 2019),

######  Assumption 1.

The function $t\mapsto\xi(t)$ is bounded Lipschitz. The non-linearity $\sigma$
is bounded Lipschitz and the data distribution has a bounded support. The
functions $w\mapsto v(w)$ and $(w_{1},w_{2})\mapsto u(w_{1},w_{2})$ are
differentiable, with bounded and Lipschitz continuous gradient. The weights at
initialization $\theta_{i}^{0}$ are i.i.d. with distribution $\rho_{0}$ which
has bounded support.

Assumption 1 imposes that the step size is of order
${\mathcal{O}}(\varepsilon)$ and its variations are of order
${\mathcal{O}}(\varepsilon^{2})$. Constant step size $\varepsilon$ will work.
Bounded non-linearity include arctan and sigmoid but excludes ReLU. While it
is a standard assumption in mean field theory ((Mei et al., 2018, 2019)), we
mention in §B that this assumption can be relaxed by the weaker assumption
that the non-linearity stays small on some big enough compact set.

The second assumption is technical and only used for studying noisy
regularized SGD.

######  Assumption 2.

$V,U$ are four times continuously differentiable and
$\nabla_{1}^{k}u(\theta_{1},\theta_{2})$ is uniformly bounded for $0\leq k\leq
4$.

The following theorem states that two wide enough two-layer neural networks
trained independently with SGD exhibit, with high probability, a linear
connection of the prediction modulo permutations for all data.

######  Theorem 3.1.

Consider two two-layer neural networks as in Equation 3 trained with equation
SGD with the same initialization over the weights independently and for the
same underlying time $T$. Suppose Assumptions 1 and 2 to hold. Then
$\forall\delta,{\text{err}},\exists N_{min}$ such that if $N\geq
N_{min},\exists\varepsilon_{max}(N)$ such that if
$\varepsilon\leq\varepsilon_{max}(N)$ in Equation 4, then with probability at
least $1-\delta$ over the training process, there exists a permutation of the
second network’s hidden layer such that for almost every $x\sim P$:
$\begin{split}&|t\hat{f}_{N}(x;\theta_{A})+(1-t)\hat{f}_{N}(x;\theta_{B})\\\
&-\hat{f}_{N}(x;t\theta_{A}+(1-t)\tilde{\theta}_{B})|\leq{\text{err}}\,,\quad\forall
t\in[0,1]\,.\end{split}$

Remark. Assumption 2 is not used when studying noiseless regularization-free
SGD ($\lambda=0,\,\tau=0$).

######  Corollary 3.2.

Under assumptions of Theorem 3.1, $\forall\delta,{\text{err}}>0,\,\exists
N^{\prime}_{\min},\,\forall N\geq
N^{\prime}_{\min},\,\exists\varepsilon^{\prime}_{\max}(N),\,\forall\varepsilon\leq\varepsilon^{\prime}_{\max}(N)$
in Equation 4, then with probability at least $1-\delta$ over the training
process, there exists a permutation of the second network’s hidden layer such
that $\forall t\in[0,1]$:
$\begin{split}&{\mathbb{E}}_{P}\big{[}\big{(}\hat{f}_{N}(x;t\theta_{A}+(1-t)\tilde{\theta}_{B})-y\big{)}^{2}\big{]}\leq{\text{err}}\\\
&+{\mathbb{E}}_{P}\big{[}t(\hat{f}_{N}(x;\theta_{A})-y)^{2}+(1-t)(\hat{f}_{N}(x;\theta_{B})-y\big{)}^{2}\big{]}\end{split}$

Discussion. Two wide enough two-layer neural networks wide enough trained with
SGD are therefore Linear Mode Connected with an upper bound on the error
tolerance we explicit in Appendix B. We have extensively used the independence
between weights in the mean field regime to apply OT bounds on convergence
rates of empirical measures. To go beyond the two-layer case, we will need to
make such an assumption on the distribution of weights. Note that this is true
at initialization and after training for two-layer networks. Studying the
independence of weights in the multi-layer case is a natural avenue for future
work, already studied in Nguyen and Pham (2020).

##  4 General Strategy for proving LMC of multi-layer networks

We now build the foundations to study the case of multi-layer neural networks
(see Equation 1).

We first write one formal property expressing the existence of permutations of
neurons of network $B$ up to layer $\ell$ such that the activations of network
$A$, network $B$ permuted and the mean network $M_{t}$ are close up to layer
$\ell$. This property is trivially satisfied at the input layer. We then show
that under two formal assumptions on the weights matrices of networks $A$ and
$B$, this property still hold at layer $\ell+1$.

###  4.1 Formal Property at layer $\ell$

Let $\varepsilon>0,m_{\ell}\geq\tilde{m}_{l}$ and
$m_{\ell+1}\geq\tilde{m}_{\ell+1}$. Assume
$\frac{m_{\ell}}{\tilde{m}_{\ell}},\frac{m_{\ell+1}}{\tilde{m}_{\ell+1}}\in{\mathbb{N}}$
to simplify technical details but this hypothesis can easily be removed.

######  Property 1.

There exists two constants $\underaccent{\bar}{E}_{\ell},E_{\ell}$ such that
given weight matrices up to layer $\ell$,
$W_{A,B}^{1,\cdots,\ell},W_{B}^{1,\cdots,\ell}$ one can find $\ell$
permutations $\Pi_{1},\cdots,\Pi_{\ell}$ of the neurons in the hidden layers
$1$ to $\ell$ of network $B$, an equi-partition
${\mathcal{I}}^{\ell}=\\{I_{1}^{\ell},\ldots,I^{\ell}_{\tilde{m}_{\ell}}\\}$,
and a map $\underaccent{\bar}{\phi}^{\ell}(x)\in{\mathbb{R}}^{n}$ such that
$\forall k\in[\tilde{m}_{\ell}]\,,\,\forall i,j\in
I_{k}^{\ell},\underaccent{\bar}{\phi}^{\ell}_{i}(x)=\underaccent{\bar}{\phi}^{\ell}_{j}(x)$
such that:

$\displaystyle\mathbb{E}_{P,Q}\|\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\leq\underaccent{\bar}{E}_{\ell}m_{\ell}$
$\displaystyle\mathbb{E}_{P,Q}\|\phi^{\ell}_{A}(x)-\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\leq
E_{\ell}m_{\ell}$
$\displaystyle\mathbb{E}_{P,Q}\|\tilde{\phi}^{\ell}_{B}(x)-\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\leq
E_{\ell}m_{\ell}$
$\displaystyle\mathbb{E}_{P,Q}\|\phi^{\ell}_{M_{t}}(x)-\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\leq
E_{\ell}m_{\ell}\,,\quad\forall t\in[0,1],$

This property not only requires proximity between activations
$\phi_{A}^{\ell}(x),\tilde{\phi}_{B}^{\ell}(x)$ at layer $\ell$ but requires
the existence of a vector $\underaccent{\bar}{\phi}^{\ell}(x)$ whose
coefficients in the same groups of the partition ${\mathcal{I}}^{\ell}$ are
equal, and therefore lives in a $\tilde{m}_{\ell}$. It bounds the size of the
function space available at layer $\ell$ and hence allows to use an effective
width $\tilde{m}_{\ell}$ independent of the real width $m_{\ell}$, which can
be much larger. It is crucial in order to show LMC for neural networks of
constant width across layers. The introduction of such a map
$\underaccent{\bar}{\phi}^{\ell}(x)$ is non trivial and is an important
contribution since it allows to extend results of Entezari et al. (2021)
beyond two layers.

###  4.2 Assumptions on the weight distribution

We now make an assumption on the empirical distribution of the weights
$\hat{\mu}_{A,\ell+1}$ at layer $\ell+1$ of $W_{A}^{\ell+1}$.

######  Assumption 3.

There exists an integer $\tilde{m}_{\ell+1}$ such that for all equi-partiton
${\mathcal{I}}^{\ell}$ of $[m_{\ell}]$ with $\tilde{m}_{\ell}$ sub-sets, there
exists a random empirical measure $\hat{\mu}_{\tilde{m}_{\ell+1}}$ independent
of $A$ and $B$ composed of $\tilde{m}_{\ell+1}$ vectors in
${\mathbb{R}}^{m_{\ell}}$, such that
$\mathbb{E}_{Q}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell}}_{A,\ell+1},\hat{\mu}^{{\mathcal{I}}^{\ell}}_{\tilde{m}_{\ell+1}})]\leq
C_{1}$.

This assumption requires that the empirical distribution with $m_{\ell+1}$
points of the neurons’ weights of network $A$ at layer $\ell+1$ can be
approximated by an empirical measure with a smaller $\tilde{m}_{\ell+1}$
number of points. Note that it implies proximity in Wasserstein distance
between $\hat{\mu}^{{\mathcal{I}}^{\ell}}_{A}$ and
$\hat{\mu}^{{\mathcal{I}}^{\ell}}_{B}$ by a triangular inequality.

We finally assume some central limit behavior when summing the errors made for
each neuron of layer $\ell$.

######  Assumption 4.

There exists a constant $C_{2}$ such that $\forall X\in{\mathbb{R}}^{m_{l}}$
we have:

$\max\left\\{\mathbb{E}_{Q}[\|W_{A}^{\ell+1}X\|_{2}^{2}],\mathbb{E}_{Q}[\|W_{\tilde{m}_{l+1}}X\|_{2}^{2}]\right\\}\leq
C_{2}\frac{m_{\ell+1}}{m_{\ell}}\|X\|_{2}^{2},$

Finally, we consider the following assumption on the non-linearity, verified
for example by pointwise ReLU.

######  Assumption 5.

$\sigma$ is pointwise, $1$-Lipschitz, $\sigma(0)=0$.

###  4.3 Propagating Property 1 to layer $\ell+1$

We state now how Property 1 propagates throughout the layers using Assumptions
5, 3 and 4 with new parameters $\underaccent{\bar}{E}_{\ell+1},E_{\ell+1}$. We
give a proof in Section A.6.

######  Lemma 4.1.

Let $\ell\in\\{0,\cdots,L-1\\}$ and suppose Property 1 to hold at layer $\ell$
and Assumptions 5, 3 and 4 to hold, then Property 1 still holds at the next
layer with $\tilde{m}_{\ell+1}$ given in Assumption 3 and
$\begin{split}&\underaccent{\bar}{E}_{\ell+1}=C_{2}\underaccent{\bar}{E}_{\ell}\\\
&E_{\ell+1}=2C_{2}E_{\ell}+2C_{1}\tilde{m}_{\ell}\underaccent{\bar}{E}_{\ell}\end{split}$

##  5 LMC for Randomly Initialized NNs with sub-Gaussian Distributions

We will make the following assumption on the empirical distribution of neurons
weights $\hat{\mu}_{A,\ell},\hat{\mu}_{B,\ell}$ of $W_{A}^{\ell},W_{B}^{\ell}$
at layer $\ell$.

######  Assumption 6 (Independence of neurons weights).

$\hat{\mu}_{A,\ell},\hat{\mu}_{B,\ell}$ correspond to two i.i.d drawings of
vectors with distribution $\mu_{\ell}$ i.e.,
$\hat{\mu}_{A,\ell},\hat{\mu}_{B,\ell}$ have the law of
$\frac{1}{m_{\ell}}\sum_{i=1}^{m_{\ell}}\delta_{x_{i}}$ where
$x_{i}\sim\mu_{\ell}$ i.i.d.

Assumption 6 is verified for example at initialization but more generally when
weights do not depend too much one of each other. This case still holds for
wide two-layer neural networks trained with SGD and is at the heart of the
proof of Theorem 3.1.

###  5.1 Showing LMC for multilayer MLPs under Gaussian distribution

We first examine the case
$\mu_{\ell}={\mathcal{N}}\left(0,\tfrac{I_{m_{\ell-1}}}{m_{\ell-1}}\right)$.
We moreover assume that the input data distribution has bounded second moment:
$\mathbb{E}_{P}[\|x\|_{2}^{2}]\leq m_{0}$.

Our strategy detailed in Section A.7 consists in showing that wide enough such
networks will satisfy Assumptions 3 and 4 with well controlled constants
$C_{1},C_{2}$. We can then apply Lemma 4.1 successively $L$ times to get the
following lemma:

######  Lemma 5.1.

Under normal initialization of the weights, given $\varepsilon>0$, if
$m_{0}\geq 5$, there exists minimal widths
$\tilde{m}_{1},\ldots,\tilde{m}_{L}$ such that if
$m_{1}\geq\tilde{m}_{1},\ldots,m_{L}\geq\tilde{m}_{L}$, Property 1 is verified
at the last hidden layer $L$ for
$\underaccent{\bar}{E}_{L}=1,E_{L}=\varepsilon^{2}$. Moreover,
$\forall\ell\in[L],\exists T_{\ell}$ which does only depend on $L,\ell$ such
that one can define recursively $\tilde{m}_{\ell}$ as $\tilde{m}_{0}=m_{0}$
and
$\tilde{m}_{\ell}=\tilde{{\mathcal{O}}}\left(\frac{T_{\ell}}{\varepsilon}\right)^{\tilde{m}_{\ell-1}}$

Discussion. The hypothesis $m_{0}\geq 5$ is technical and could be relaxed at
the price of slightly changing the bound on $\tilde{m}_{1}$. Lemma 5.1 shows
that given two random networks whose widths $m_{\ell}$ is larger than
$\tilde{m}_{\ell}$, we can permute neurons of the second one such that their
activations at layer $\ell$ are both close to the one of the networks on a
linear path in parameter’s space.

As $\varepsilon$ goes to $0$, the width of the layer $\ell+1$ must scale at
least as $\left(\frac{1}{\varepsilon}\right)^{\tilde{m}_{\ell-1}}$. This is a
fundamental bound due to the convergence rate in Wasserstein distance of
empirical measures. It imposes a recursive exponential growth in the width
needed with respect to depth. This condition appears excessive as compared to
the typical width of neural networks used in practice. We highlight here that
Ainsworth et al. (2022) empirically demonstrates that networks at
initialization do not exhibit LMC and that the loss barrier is erased only
after a sufficient number of SGD steps.

###  5.2 Showing Linear Mode Connectivity

We make the following assumption on the loss function to show LMC from Lemma
5.1.

######  Assumption 7.

$\forall y\in{\mathbb{R}}^{m_{L+1}}$, the loss ${\mathcal{L}}(\cdot,y)$ is
convex and $1$-Lipschitz.

We finally prove the following bound on the loss of the mean network $M_{t}$
in Section A.8:

######  Theorem 5.2.

Under normal initialization of the weights, for
$m_{1}\geq\tilde{m}_{1},\cdots,m_{L}\geq\tilde{m}_{L}$ as defined in Lemma
5.1, $m_{0}\geq 5$, and under Assumption 7 we know that $\forall t\in[0,1]$,
with $Q$-probability at least $1-\delta_{Q}$, there exists permutations of
hidden layers $1,\ldots,L$ of network $B$ that are independent of $t$, such
that:
$\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{M_{t}}(x),y\right)\right]\leq
t\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{A}(x),y\right)\right]+\\\
(1-t)\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{B}(x),y\right)\right]+\tfrac{4\sqrt{m_{L+1}}}{\delta_{Q}^{2}}\varepsilon$

Discussion. The minimal width at layer $\ell$ needed for Theorem 5.2 is
recursively $\tilde{m}_{l}\sim\varepsilon^{-\tilde{m}_{l-1}}$. Applied to
randomly initialized two-layer networks, we need a hidden layer’s dimension of
$\varepsilon^{-m_{0}}$ as opposed to Entezari et al. (2021) which prove a
bound of $\varepsilon^{-(2m_{0}+4)}$.

###  5.3 Tightness of the bound dependency with respect to the error tolerance

We discuss here the tightness of the minimal width $\tilde{m}_{\ell}$ we
require in Lemma 5.1 with respect to the error tolerance $\varepsilon$. The
recursive exponential growth of the width in the form
$\tilde{m}_{\ell}\sim\left(\frac{1}{\varepsilon}\right)^{\tilde{m}_{\ell-1}}$
is a consequence of the convergence rate of Wasserstein distance of empirical
measures in dimension $\tilde{m}_{\ell-1}$ at the rate
$\nicefrac{{1}}{{\tilde{m}_{\ell-1}}}$. Theorem 5.3 provides a corresponding
lower bound which shows that this recursive exponential growth is tight at the
precise rate $\left(\tfrac{1}{\varepsilon}\right)^{\tilde{m}_{\ell-1}}$ (just
take
$n=\tilde{m}_{\ell-1},\,m=\tilde{m}_{\ell},\,\mu=\mu_{\ell},\,x=\phi_{A}^{\ell-1}(x),W_{A,B}=W_{A,B}^{\ell}$).
A proof is given in Section A.11.

######  Theorem 5.3.

Let $n\geq 1,x\sim P\in{\mathcal{P}}_{1}({\mathbb{R}}^{n})$ and
$\mu\in{\mathcal{P}}({\mathbb{R}}^{n})$ such that
$\frac{d\mu}{d{\text{Leb}}}\leq F_{1}$. Suppose $\Sigma=\mathbb{E}[xx^{T}]$ is
full rank $n$. Let $m\geq 1$ and
$W_{A},W_{B}\in{\mathcal{M}}_{m,n}({\mathbb{R}})$ whose rows are drawn i.i.d.
from $\mu$. Then, there exists $F_{0}$ such that
$\mathbb{E}_{W_{A},W_{B}}[\min_{\Pi\in{\mathcal{S}}_{m}}\mathbb{E}_{P}\|(W_{A}-\Pi
W_{B})x\|_{2}^{2}]\geq F_{0}\left(\frac{1}{m}\right)^{2/n}$

Remark 1. Using an effective width $\tilde{m}_{\ell-1}$ smaller and
independent of the real width $m_{\ell-1}$ allows to show LMC for networks of
constant hidden width $m_{1}=m_{2}=\ldots=m_{L}$ as soon as they verify
$m_{1}\geq\tilde{m}_{1},\ldots,m_{L}\geq\tilde{m}_{L}$ where
$\tilde{m}_{1},\ldots,\tilde{m}_{L}$ are defined in Lemma 5.1. Without this
trick, we need a recursive exponential growth of the real width
$m_{\ell}\sim\left(\tfrac{1}{\varepsilon}\right)^{m_{\ell-1}}$.

Remark 2. Motivated by the fact that feature learning may concentrate the
weight distribution on low dimensional sub-space, we could extend our proofs
to the case where the underlying weight distribution has a support with
smaller dimension to get recursive bounds no longer at rate
$\tilde{m}_{\ell-1}$ but at a smaller one. Note this is unlikely to happen as
we expect the matrix of weight vectors of a given layer to be full rank.
Therefore, we study in the next section the case when this matrix is
approximately low rank, or equivalently when the weight distribution is
concentrated around a low dimensional approximated support.

###  5.4 Approximately low dimensional supported measures

For the sake of clarity, assume from now on that the layer $\ell-1$ of network
$A$ has been permuted such that for
${\mathcal{I}}^{\ell-1}=\\{I_{1}^{\ell-1},\ldots,I_{\tilde{m}_{\ell-1}}^{\ell-1}\\}$
(given in Property 1) we have
$I_{1}^{\ell-1}=\\{1,\ldots,p_{\ell-1}\\},\,\ldots,\,I_{\tilde{m}_{\ell-1}}^{\ell-1}=\\{m_{\ell-1}-p_{\ell-1}+1,\ldots,m_{\ell-1}\\}$
with $p_{\ell-1}=\nicefrac{{m_{\ell-1}}}{{\tilde{m}_{\ell-1}}}\\}$. This
assumption is mild since we can always consider a permuted version of network
$A$ without changing the problem.

Motivated by the discussion in Section A.9.1 we consider the model where the
weights at layer $\ell$ are initialized i.i.d. multivariate Gaussian
$\mu_{\ell}={\mathcal{N}}(0,\Sigma^{\ell-1})$ with

$\Sigma^{\ell-1}:=\operatorname{Diag}\left(\lambda^{\ell}_{1}I_{p_{\ell-1}},\lambda_{2}^{\ell}I_{p_{\ell-1}},\ldots,\lambda^{\ell}_{\tilde{m}_{\ell-1}}I_{p_{\ell-1}}\right)$

with
$\tfrac{1}{m_{\ell-1}}\tfrac{\tilde{m}_{\ell-1}}{k_{\ell-1}}\geq\lambda_{1}^{\ell}\geq\lambda_{2}^{\ell}\geq\ldots\geq\lambda^{\ell}_{\tilde{m}_{\ell-1}}$
with $k_{\ell-1}\leq\tilde{m}_{\ell-1}$ an approximate dimension of the
support of the underlying weights distribution. Note that to balance the low
dimensionality of the weights distribution, we have replaced the upper-bound
on the eigenvalues $\tfrac{1}{\tilde{m}_{\ell-1}}$ by the greater value
$\tfrac{1}{m_{\ell-1}}\tfrac{\tilde{m}_{\ell-1}}{k_{\ell-1}}$ to avoid
vanishing activations when $\ell$ grows which would have made our result
vacuous.

The following assumption states that the weights distribution
$\mu_{\ell}^{{\mathcal{I}}^{\ell-1}}$ at layer $\ell$ considered in
${\mathcal{P}}_{1}({\mathbb{R}}^{\tilde{m}_{\ell-1}})$ (with the operation
explicited in Section 2) is approximately of dimension
$k_{\ell-1}=e\tilde{m}_{\ell-1}$. The approximation becomes more correct as
$\eta\rightarrow 0$.

######  Assumption 8 (Approximately low-dimensionality).

$\exists\eta,e\in(0,1),\forall\ell\in[L],\,\frac{\sqrt{\sum_{j=k_{\ell-1}}^{\tilde{m}_{\ell-1}}\lambda^{\ell}_{j}}}{4\sqrt{\sum_{j=1}^{k_{\ell-1}}\lambda^{\ell}_{j}}}\leq\eta,\,\tfrac{k_{\ell-1}}{\tilde{m}_{\ell-1}}=e$

######  Theorem 5.4.

Under Assumptions 8 and 7, given $\varepsilon>0$, if $em_{0}\geq 5$ there
exists minimal widths $\tilde{m}_{1},\ldots,\tilde{m}_{L}$ such that if
$\eta^{-k_{0}}\geq m_{1}\geq\tilde{m}_{1},\ldots,\eta^{-k_{L-1}}\geq
m_{L}\geq\tilde{m}_{L}$, Property 1 is verified at the last hidden layer $L$
for $\underaccent{\bar}{E}_{L}=1,E_{L}=\varepsilon^{2}$. Moreover,
$\forall\ell\in[L],\,\exists T^{\prime}_{\ell}$ which does only depend on
$L,e,\ell,$ such that one can define recursively $\tilde{m}_{\ell}$ as
$\tilde{m}_{\ell}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime}_{\ell}}{\varepsilon}\right)^{k_{\ell-1}}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime}_{\ell}}{\varepsilon}\right)^{e\tilde{m}_{\ell-1}}$
where $\tilde{m}_{0}=m_{0}$. Moreover there exists permutations of hidden
layers $1,\ldots,L$ of network $B$ s.t. $\forall t\in[0,1]$, with
$Q$-probability at least $1-\delta_{Q}$:
$\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{M_{t}}(x),y\right)\right]\leq
t\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{A}(x),y\right)\right]\\\
+(1-t)\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{B}(x),y\right)\right]+\tfrac{4\sqrt{m_{L+1}}}{\sqrt{e}\delta_{Q}^{2}}\varepsilon$

Discussion. We give a proof in Section A.10. For $\eta$ small enough, the
distribution of weights is approximately lower dimensional. It yields faster
convergence rates until $m$ becomes exponentially big in $\eta$. This prevents
the previous recursive exponential growth of width with respect to depth,
though asymptotically, we recover the same rates as in Theorem 5.2. The
smaller $e$, the lower dimensional are the distributions, and the less the
width needs to grow when $\varepsilon\rightarrow 0$. The problem in that model
is that the constant $T^{\prime}_{i}$ explodes if $e\rightarrow 0$, which
prevents using a model with fixed $k_{\ell}$ across the layers for the weight
distribution. We want to highlight here that the proof can be extended to such
a case, but we need to assume that the constant $C_{2}$ is bounded and not
depending on $e$ across the layers in Lemma 4.1 (recall that with our proof,
we had $C_{2}=\tfrac{1}{e}$). This assumption seems coherent because the
average activations don’t explode across layers in the model. Assuming this,
the bound we obtain for $\tilde{m}_{\ell}$ in Theorem 5.4 is completely
independent on $\tilde{m}_{\ell-1}$, and there is no recursive exponential
growth in the width needed across the layers. We give a more explicit
discussion in Section A.12.

###  5.5 LMC for sub-Gaussian distributions

Still under the setting of Assumption 6 assume that the underlying
distribution $\mu_{\ell}$ verifies for each layer $\ell\in[L+1]$: if
$X\sim\mu_{\ell}$ then, $\forall j\neq k\in[m_{l-1}],X_{j}\amalg X_{k}$.
Moreover $\forall i\in[\tilde{m}_{\ell-1}],\forall j,k\in I^{\ell-1}_{i}$,

$\mathbb{E}[X_{j}^{2}]=\mathbb{E}[X_{k}^{2}]=\lambda_{i}^{\ell}$

Finally suppose the variables are sub-Gaussian i.e., $\exists K>0,\forall
i\in[\tilde{m}_{\ell-1}],\forall j\in I^{\ell-1}_{i},\,\forall c>0$,

${\mathbb{P}}(|X_{j}|\geq c)\leq 2\exp(-\frac{c^{2}}{K\lambda_{i}^{\ell}})$

We explain in Section A.13 why both Theorem 5.2 (in the case
$\lambda_{1}^{\ell}=\ldots=\lambda^{\ell}_{\tilde{m}_{l-1}}=\nicefrac{{1}}{{m_{\ell-1}}}$)
and Theorem 5.4 hold with mild modifications in the constants.

It extends our previous result considerably to LMC for any large enough
networks whose weights are i.i.d. and whose underlying distribution has a sub-
Gaussian tail (for example uniform distribution).

###  5.6 Link with dropout stability

In Section A.14, we build a first step towards unifying our study with the
dropout stability viewpoint (Kuditipudi et al., 2019; Shevchenko and Mondelli,
2020) by showing in a simplified setting how networks become dropout stable in
the same asymptotics on the width as the one needed in our Theorem 5.2.

(a) Mean test loss of the trained networks $A$ and $B$ and error barrier on
the linear path $M_{t},\,t\in[0,1]$ across different learning rate values for
each matching problem.

(b) Approximate dimension
${\text{Dim}}(S):=\nicefrac{{\operatorname{tr}(S)^{2}}}{{\operatorname{tr}(S^{2})}}$
of the matrices considered in each matching problem.

Figure 2: Statistics of the average network $M$ over the linear path between
networks $A$ and $B$ using respectively weight matching (blue), weight
matching using covariance of activations and activations (green), and
activation matching (red)

##  6 Experimental validation and new method to find permutations

Our previous study shows the influence of the dimension of the underlying
weight distribution on LMC effectiveness. Based on this insight we develop a
new weight matching method at the crossroads between previous naive weight
matching (WM) and activation matching (AM) methods (Ainsworth et al., 2022).
Given $n$ training points $x_{i},\,i\in[n]$, denote
$Z_{A}^{\ell}\in{\mathcal{M}}_{m_{\ell},n}({\mathbb{R}})$ (respectively
$Z_{B}^{\ell}$) the activations $\phi^{\ell}_{A}(x_{i})$ for the $n$ data
points $x_{i}$. Further denote
$\Sigma_{A}^{\ell}:=\tfrac{1}{n}Z_{A}^{\ell}[Z_{A}^{\ell}]^{T}\approx\mathbb{E}_{P}\left[\phi_{A}^{\ell}(x)[\phi_{A}^{\ell}(x)]^{T}\right]$.
We aim at finding for each layer $\ell$ the optimal permutation $\Pi$
minimizing the cost (respectively for naive WM, our new WM method and AM):

$\displaystyle\min_{\Pi\in S_{m_{\ell}}}\left\|W_{A}^{\ell}-\Pi
W_{B}^{\ell}\Pi_{\ell-1}^{T}\right\|_{2}^{2}\,,$ (Naive WM)
$\displaystyle\min_{\Pi\in S_{m_{\ell}}}\left\|W_{A}^{\ell}-\Pi
W_{B}^{\ell}\Pi_{\ell-1}^{T}\right\|^{2}_{2,\Sigma_{A}^{\ell-1}}\,,$ (WM
(ours)) $\displaystyle\min_{\Pi\in S_{m_{\ell}}}\left\|Z_{A}^{\ell}-\Pi
Z_{B}^{\ell}\right\|_{2}^{2}\,,$ (AM)

where $\|\cdot\|_{2,\Sigma_{A}^{\ell-1}}$ is the norm333Semi-norm in full
generality (if $\Sigma_{A}^{\ell-1}$ is not full rank) induced by the scalar
product $(X,Y)\mapsto\operatorname{tr}(X\Sigma_{A}^{\ell-1}Y^{T})$. We both
theoretically support the gain of our method in Theorem C.2 and empirically
verify that this method constantly and substantially outperforms naive Weight
Matching across different learning rates when training with SGD.

We train a three hidden layer MLP of width $512$ on MNIST with learning rates
varying between $10^{-4}$ and $10^{-1}$ across $4$ runs. We plot on Figure
2(b) the approximate dimension of the considered covariance matrix for each
matching method: $W_{A}^{\ell}[W_{A}^{\ell}]^{T}$ for WM (naive),
$W_{A}^{\ell}\Sigma_{A}^{\ell-1}[W_{A}^{\ell}]^{T}$ for WM (ours) and
$\Sigma_{A}^{\ell}$ for AM (see §C.1). Our code is available at
https://github.com/9aze/OT_LMC/tree/main. We see the detrimental effect of
high approximate dimension on LMC effectiveness, therefore validating our
theoretical approach. Note that for a learning rate of $10^{-1}$ the
correlation is less clear but a trend is visible on decreasing dimension for
naive WM as it performs better (and increasing dimension for AM and our WM
method as it performs comparatively less well). An alternative would be to use
a proxy taking the diameter of the distributions into account (and not only
the dimension of their support). Finally, experiments on Adam lead to less
clear results that we did not report as more experimental investigation is
needed. In particular, understanding the impact of the optimizer on the
independence of weights during training is crucial, as it is a central
assumption in our study.

##  7 Discussion

Optimal transport serves as a good framework to study linear mode connectivity
of neural networks. This paper uses convergence rates of empirical measures in
Wasserstein distance to upper bound the test error of the linear combination
of two networks in weight space modulo permutation symmetries. Our main
assumption is the independence of all neuron’s weight vectors inside a given
layer. This assumption is trivially true at initialization but remains valid
for wide two-layers networks trained with SGD. We experimentally demonstrate
the correlation between the dimension of the underlying weight distribution
with LMC effectiveness and design a new weight matching method that
significantly outperforms existing ones. A natural direction for future work
is to focus on the behaviour of the weights distribution inside each layer of
DNNs and their independence. Moreover, extending our results to only assuming
approximate independence of weights is a natural direction as it seems a more
realistic setting.

##  8 Acknowledgements

The work of B. Goujaud and A. Dieuleveut is partially supported by
ANR-19-CHIA-0002-01/chaire SCAI, and Hi!Paris. This work was partly funded by
the French government under management of Agence Nationale de la Recherche as
part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001
(PRAIRIE 3IA Institute).

## References

  * Ainsworth et al. (2022) S. K. Ainsworth, J. Hayase, and S. Srinivasa.  Git re-basin: Merging models modulo permutation symmetries.  _arXiv preprint arXiv:2209.04836_ , 2022.
  * Chizat and Bach (2018) L. Chizat and F. Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  _Advances in neural information processing systems_ , 31, 2018.
  * Draxler et al. (2018) F. Draxler, K. Veschgini, M. Salmhofer, and F. Hamprecht.  Essentially no barriers in neural network energy landscape.  In J. Dy and A. Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 1309–1318. PMLR, 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/draxler18a.html.
  * Entezari et al. (2021) R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur.  The role of permutation invariance in linear mode connectivity of neural networks.  _arXiv preprint arXiv:2110.06296_ , 2021.
  * Frankle et al. (2020) J. Frankle, G. K. Dziugaite, D. Roy, and M. Carbin.  Linear mode connectivity and the lottery ticket hypothesis.  In _International Conference on Machine Learning_ , pages 3259–3269. PMLR, 2020.
  * Freeman and Bruna (2016) C. D. Freeman and J. Bruna.  Topology and geometry of half-rectified network optimization.  _arXiv preprint arXiv:1611.01540_ , 2016.
  * Garipov et al. (2018) T. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson.  Loss surfaces, mode connectivity, and fast ensembling of dnns.  _Advances in neural information processing systems_ , 31, 2018.
  * Goodfellow et al. (2014) I. J. Goodfellow, O. Vinyals, and A. M. Saxe.  Qualitatively characterizing neural network optimization problems.  _arXiv preprint arXiv:1412.6544_ , 2014.
  * Gotmare et al. (2018) A. Gotmare, N. S. Keskar, C. Xiong, and R. Socher.  Using mode connectivity for loss landscape analysis.  _arXiv preprint arXiv:1806.06977_ , 2018.
  * Izmailov et al. (2018) P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson.  Averaging weights leads to wider optima and better generalization.  _arXiv preprint arXiv:1803.05407_ , 2018.
  * Jacot et al. (2018) A. Jacot, F. Gabriel, and C. Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Juneja et al. (2022) J. Juneja, R. Bansal, K. Cho, J. Sedoc, and N. Saphra.  Linear connectivity reveals generalization strategies.  _arXiv preprint arXiv:2205.12411_ , 2022.
  * Keskar et al. (2016) N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_ , 2016.
  * Kuditipudi et al. (2019) R. Kuditipudi, X. Wang, H. Lee, Y. Zhang, Z. Li, W. Hu, R. Ge, and S. Arora.  Explaining landscape connectivity of low-cost solutions for multilayer nets.  _Advances in neural information processing systems_ , 32, 2019.
  * Lubana et al. (2023) E. S. Lubana, E. J. Bigelow, R. P. Dick, D. Krueger, and H. Tanaka.  Mechanistic mode connectivity.  In _International Conference on Machine Learning_ , pages 22965–23004. PMLR, 2023.
  * Lucas et al. (2021) J. Lucas, J. Bae, M. R. Zhang, S. Fort, R. Zemel, and R. Grosse.  Analyzing monotonic linear interpolation in neural network loss landscapes.  _arXiv preprint arXiv:2104.11044_ , 2021.
  * Mei et al. (2018) S. Mei, A. Montanari, and P.-M. Nguyen.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 115(33):E7665–E7671, 2018.
  * Mei et al. (2019) S. Mei, T. Misiakiewicz, and A. Montanari.  Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.  In _Conference on Learning Theory_ , pages 2388–2464. PMLR, 2019.
  * Neyshabur et al. (2020) B. Neyshabur, H. Sedghi, and C. Zhang.  What is being transferred in transfer learning?  _Advances in neural information processing systems_ , 33:512–523, 2020.
  * Nguyen and Pham (2020) P.-M. Nguyen and H. T. Pham.  A rigorous framework for the mean field limit of multilayer neural networks.  _arXiv preprint arXiv:2001.11443_ , 2020.
  * Peyré et al. (2019) G. Peyré, M. Cuturi, et al.  Computational optimal transport: With applications to data science.  _Foundations and Trends® in Machine Learning_ , 11(5-6):355–607, 2019.
  * Phillips (2012) J. M. Phillips.  Chernoff-hoeffding inequality and applications.  _arXiv preprint arXiv:1209.6396_ , 2012.
  * Pittorino et al. (2022) F. Pittorino, A. Ferraro, G. Perugini, C. Feinauer, C. Baldassi, and R. Zecchina.  Deep networks on toroids: removing symmetries reveals the structure of flat regions in the landscape geometry.  In _International Conference on Machine Learning_ , pages 17759–17781. PMLR, 2022.
  * Rame et al. (2022) A. Rame, M. Kirchmeyer, T. Rahier, A. Rakotomamonjy, P. Gallinari, and M. Cord.  Diverse weight averaging for out-of-distribution generalization.  _Advances in Neural Information Processing Systems_ , 35:10821–10836, 2022.
  * Sagun et al. (2017) L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou.  Empirical analysis of the hessian of over-parametrized neural networks.  _arXiv preprint arXiv:1706.04454_ , 2017.
  * Shevchenko and Mondelli (2020) A. Shevchenko and M. Mondelli.  Landscape connectivity and dropout stability of sgd solutions for over-parameterized neural networks.  In _International Conference on Machine Learning_ , pages 8773–8784. PMLR, 2020.
  * Singh and Jaggi (2020) S. P. Singh and M. Jaggi.  Model fusion via optimal transport.  _Advances in Neural Information Processing Systems_ , 33:22045–22055, 2020.
  * Tatro et al. (2020) N. Tatro, P.-Y. Chen, P. Das, I. Melnyk, P. Sattigeri, and R. Lai.  Optimizing mode connectivity via neuron alignment.  _Advances in Neural Information Processing Systems_ , 33:15300–15311, 2020.
  * Venturi et al. (2019) L. Venturi, A. S. Bandeira, and J. Bruna.  Spurious valleys in one-hidden-layer neural network optimization landscapes.  _Journal of Machine Learning Research_ , 20:133, 2019.
  * Villani et al. (2009) C. Villani et al.  _Optimal transport: old and new_ , volume 338.  Springer, 2009.
  * Vlaar and Frankle (2022) T. J. Vlaar and J. Frankle.  What can linear interpolation of neural network loss landscapes tell us?  In _International Conference on Machine Learning_ , pages 22325–22341. PMLR, 2022.
  * Weed and Bach (2019) J. Weed and F. Bach.  Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance.  _Advances in Neural Information Processing Systems_ , 2019.
  * Wortsman et al. (2022) M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, et al.  Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.  In _International Conference on Machine Learning_ , pages 23965–23998. PMLR, 2022.
  * Wu (2016) Y. Wu.  Packing, covering, and consequences on minimax risk.  http://www.stat.yale.edu/$\sim$yw562/ teaching/598/lec14.pdf, 2016.  [Online; accessed October-10-2023].
  * Zhao et al. (2020) P. Zhao, P.-Y. Chen, P. Das, K. N. Ramamurthy, and X. Lin.  Bridging mode connectivity in loss landscapes and adversarial robustness.  _arXiv preprint arXiv:2005.00060_ , 2020.

##  Appendix A Proofs and details about Optimal Transport Theory for LMC

###  A.1 Background on optimal transport and convexity lemma

Optimal transport is a mathematical framework that aims at quantifying
distances between distributions. We refer to the books Villani et al. [2009]
and Peyré et al. [2019] (focused on computational aspects and applications to
data science) for an extensive overview of this topic.

######  Definition A.1 (Wasserstein distance [Villani et al., 2009]).

Let $(\mathcal{X},d)$ a Polish metric space, $p\in[1,\infty)$.
$\forall\mu,\nu\in{\mathcal{P}}_{1}(\mathcal{X})$, define the $p$-Wasserstein
distance between $\mu$ and $\nu$ by:

${\mathcal{W}}_{p}(\mu,\nu)=\left(\inf_{\pi\in\Pi(\mu,\nu)}\int_{\mathcal{X}^{2}}d^{p}(x,y)d\pi(x,y)\right)^{\nicefrac{{1}}{{p}}}$

Recall that $\Pi(\mu,\nu)$ denotes the set of coupling between $\mu$ and
$\nu$, i.e.,

$\pi\in\Pi(\mu,\nu)\Leftrightarrow\pi\in{\mathcal{P}}_{1}(\mathcal{X}^{2})\text{
with marginals }\mu,\nu$

This defines a distance and especially it satisfies the triangular inequality.
Moreover we state a Jensen type inequality proved in Villani et al. [2009]:

######  Lemma A.2 (Convexity of the optimal cost (Theorem 4.8 in Villani et
al. [2009])).

Let $d:{\mathbb{R}}^{n}\times{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}_{+}$ a
distance and for $p\in[1,\infty)$,
${\mathcal{W}}_{p}:{\mathcal{P}}_{1}({\mathbb{R}}^{n})^{2}\rightarrow{\mathbb{R}}_{+}$
its associate Wasserstein distance. Let $(\Theta,\lambda)$ be a probability
space and $\mu_{\theta},\nu_{\theta}$ two measurable functions on that space
taking values in ${\mathcal{P}}_{1}({\mathbb{R}}^{n})$. Then,

${\mathcal{W}}_{p}^{p}\left(\int_{\Theta}\mu_{\theta}\lambda(d\theta),\int_{\Theta}\nu_{\theta}\lambda(d\theta)\right)\leq\int_{\Theta}{\mathcal{W}}_{p}^{p}(\mu_{\theta},\nu_{\theta})\lambda(d\theta)$

###### Proof.

To apply Theorem $4.8$ in Villani et al. [2009], we just need to notice that
$d^{p}(\cdot,\cdot)$ is continuous, $d^{p}(\cdot,\cdot)\geq 0$ and the
associated optimal cost functional is ${\mathcal{W}}_{p}^{p}(\cdot,\cdot)$. ∎

###  A.2 Validity of the OT viewpoint: Birkhoff’s theorem

We have motivated in Section 2 the following minimization problem (Equation
2):

$\begin{split}\Pi_{\ell}&=\operatorname*{arg\,min}_{\Pi\in{\mathcal{S}}_{m_{\ell}}}\|W_{A}^{\ell}-\Pi
W_{B}^{\ell}\Pi_{\ell-1}^{T}\|_{2}^{2}\\\
&=\operatorname*{arg\,min}_{\pi\in{\mathcal{S}}_{m_{\ell}}}\frac{1}{m_{\ell}}\sum_{i=1}^{m_{\ell}}\|[W_{A}^{\ell}]_{i:}-[W_{B}^{\ell}\Pi^{T}_{\ell-1}]_{\pi_{i}:}\|_{2}^{2}\end{split}$

Since LMC effectiveness will be related to the effectiveness of this
optimization problem, we want to quantify the minimization error:

$\min_{\Pi\in{\mathcal{S}}_{m_{\ell}}}\|W_{A}^{\ell}-\Pi
W_{B}^{\ell}\Pi_{\ell-1}^{T}\|_{2}^{2}$

We highlight the similarity with previous Definition A.1. The main difference
being that the latter minimizes the cost among all couplings
$\pi\in\Pi(\mu,\nu)$ between the two distributions, especially transport plans
that can split mass. Permutation must be on the other hand seen as Monge maps,
i.e. deterministic maps. This ambiguity is solved with the following lemma:

######  Lemma A.3 (Proposition 2.1 in Peyré et al. [2019]).

Let $m,n$ integers and $x_{1},...,x_{m},y_{1},...,y_{m}$, $2m$ points of
${\mathbb{R}}^{n}$. Let
$\hat{\mu}_{m}=\frac{1}{m}\sum_{i=1}^{m}\delta_{x_{i}},\hat{\nu}_{m}=\frac{1}{m}\sum_{i=1}^{m}\delta_{y_{i}}$
their associated empirical measures. Consider
$d:{\mathbb{R}}^{n}\times{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}_{+}$ a
distance function and for $p\in[1,\infty),{\mathcal{W}}_{p}$ its associated
Wasserstein distance. Then one has:

${\mathcal{W}}_{p}:=\inf_{\pi\in\Pi(\hat{\mu}_{m},\hat{\nu}_{m})}\left(\int_{{\mathbb{R}}^{n}\times{\mathbb{R}}^{n}}d(x,y)^{p}d\pi(x,y)\right)^{\frac{1}{p}}=\min_{\pi\in{\mathcal{S}}_{m}}\left(\frac{1}{m}\sum_{i=1}^{m}d(x_{i},y_{\sigma_{i}})^{p}\right)^{\frac{1}{p}}$

###### Proof.

$\Pi(\hat{\mu}_{m},\hat{\nu}_{m})$ is the convex envelope of its extremal
points which are described by permutations by Birkhoff’s theorem. Moreover,

$\pi\in{\mathcal{P}}_{1}({\mathbb{R}}^{n}\times{\mathbb{R}}^{n})\rightarrow\int
d(x,y)^{p}d\pi(x,y)$

is linear and therefore its infimum is attained on one extremal point of
$\Pi(\hat{\mu}_{m},\hat{\nu}_{m})$ ∎

This lemma implies equality between the Wasserstein distance between two
empirical measures and the minimum cost over the set of permutations. We can
therefore restrict our study to permutations while still using tools from
general optimal transport theory and convergence rates of empirical measures
in Wasserstein distance.

###  A.3 Technical lemmas

The following lemma will be very useful in the following and shows that
Binomial variables are concentrated around their expectation.

######  Lemma A.4 (Hoeffding inequality for Binomial variables).

Let ${\mathcal{B}}_{\frac{p}{2},m}\sim{\mathcal{B}}(\frac{p}{2},m)$ a binomial
variable with $p\in[0,1]$. Then,

${\mathbb{P}}(\frac{{\mathcal{B}}_{\frac{p}{2},m}}{m}\geq
p)\leq\exp(-\frac{p^{2}m}{2})$

Proof: This is a simple application of Hoeffding concentration inequality
[Phillips, 2012] since a Binomial variable is a sum of independent Bernoulli
variables.

######  Lemma A.5.

Let $a,b>0$ such that $a+b=1$ and let
$\mu_{1},\mu_{2},\nu_{1},\nu_{2}\in{\mathcal{P}}_{1}(\mathcal{X})$ with
$({\mathcal{X}},d)$ a Polish space. Then, $\forall p\in[1,\infty)$:

${\mathcal{W}}_{p}^{p}(a\mu_{1}+b\mu_{2},a\nu_{1}+b\nu_{2})\leq
a{\mathcal{W}}_{p}^{p}(\mu_{1},\nu_{1})+b{\mathcal{W}}_{p}^{p}(\mu_{2},\nu_{2})$
(5)

###### Proof.

Just apply Lemma A.2 with $\mu_{\Theta},\nu_{\Theta}$ such that
${\mathbb{P}}\left((\mu_{\Theta},\nu_{\Theta})=(\mu_{1},\nu_{1})\right)=a$ and
${\mathbb{P}}\left((\mu_{\Theta},\nu_{\Theta})=(\mu_{2},\nu_{2})\right)=b$. ∎

######  Lemma A.6 (Hölder, Remark $6.6$ in Villani et al. [2009]).

Let $(\mathcal{X},d)$ a Polish space, let $p,q\in[1,\infty]$ such that $p\leq
q$:

$\forall\mu,\nu\in{\mathcal{P}}_{1}(\mathcal{X}),\,{\mathcal{W}}_{p}\leq{\mathcal{W}}_{q}$
(6)

###### Proof.

Just apply Hölder’s inequality. ∎

###  A.4 Definitions and technical lemma on packing numbers

Let $n\in{\mathbb{N}}^{*}$, $S\subset{\mathbb{R}}^{n}$ and $d(\cdot,\cdot)$ a
distance of ${\mathbb{R}}^{n}$. We recall the definitions of packing numbers
and covering numbers as well as two lemmas stated and proved in Wu [2016]:

######  Definition A.7 ($\varepsilon$-covering number [Weed and Bach, 2019]).

The $\varepsilon$-covering of a set $S$ denoted
${\mathcal{N}}_{\varepsilon}(S)$ is the minimum number $m$ of closed balls
$B_{1},\ldots,B_{m}$ of radius $\varepsilon$ such that
$S\subset\bigcup_{i=1}^{m}B_{i}$.

######  Definition A.8 ($\varepsilon$-packing number).

The $\varepsilon$-packing number of a set $S$ denoted
${\mathcal{P}}_{\varepsilon}(S)$ is the maximum number $m$ of distinct points
$\theta_{1},...,\theta_{m}\in S$ such that $\forall i\neq j$,
$\|\theta_{i}-\theta_{j}\|>\varepsilon$.

######  Lemma A.9.

If the distance $d(\cdot,\cdot)$ comes from a norm, ($d(x,y)=\|x-y\|)$,
denoting ${\text{Leb}}_{n}$ the Lebesgue measure in ${\mathbb{R}}^{n}$ we
have: $\forall S\subset{\mathbb{R}}^{n},\forall\varepsilon>0$,

${\mathcal{N}}_{\varepsilon}(S)\leq\frac{{\text{Leb}}_{n}(S+{\mathcal{B}}(0,\varepsilon/2))}{{\text{Leb}}_{n}({\mathcal{B}}(0,\varepsilon/2))}$

###### Proof.

We prove first
${\mathcal{N}}_{\varepsilon}(S)\leq{\mathcal{P}}_{\varepsilon}(S)$. Indeed
considering $m={\mathcal{P}}_{\varepsilon}(S)$ and
$\theta_{1},\ldots,\theta_{m}$ associated, we know by definition of the
$\varepsilon$-packing number that $\forall\theta\in S,\exists i\in[m]$ such
that $\|\theta_{i}-\theta\|\leq\varepsilon$. This shows that
$S\subset\bigcup_{i=1}^{m}{\mathcal{B}}(\theta_{i},\varepsilon)$

Now, on the other hand, we know that all balls
${\mathcal{B}}(\theta_{i},\frac{\varepsilon}{2})$ are disjoint and
$\bigcup_{i=1}^{m}{\mathcal{B}}(\theta_{i},\frac{\varepsilon}{2})\subset
S+{\mathcal{B}}(0,\varepsilon/2)$. This yields the result by a volume (i.e.,
Lebesgue measure) argument. ∎

######  Lemma A.10.

If the distance $d(\cdot,\cdot)$ comes from a norm, ($d(x,y)=\|x-y\|)$ we
have: $\forall S\subset{\mathbb{R}}^{n},\forall\varepsilon>0$,

${\mathcal{N}}_{\varepsilon}(S)\geq\left(\frac{1}{\varepsilon}\right)^{n}\frac{{\text{Leb}}_{n}(S)}{{\text{Leb}}_{n}({\mathcal{B}}(0,1))}$

###### Proof.

Notice that given a covering
$\bigcup_{i=1}^{{\mathcal{N}}_{\varepsilon}(S)}{\mathcal{B}}(x_{i},\varepsilon)\supset
S$, by a volume argument we get

${\text{Leb}}_{n}(S)\leq\varepsilon^{n}{\mathcal{N}}_{\varepsilon}(S){\text{Leb}}_{n}({\mathcal{B}}(0,1))$

∎

where we have used homogeneity of the norm.

###  A.5 Lemmas on convergence rates of empirical measures

This section is devoted to proving convergence rates in Wasserstein distance
of empirical measure towards the underlying distribution whose points are
drawn. More precisely, given $\mu\in{\mathcal{P}}_{1}({\mathbb{R}}^{n})$ a
probability measure on an euclidean space and $p\in[1,\infty)$, we will focus
on bounding the quantity
$\mathbb{E}\left[{\mathcal{W}}_{p}^{p}(\hat{\mu}_{m},\mu)\right]$ as $m$
grows, where $\hat{\mu}_{m}$ is a random empirical measure i.e.,
$\hat{\mu}_{m}=\frac{1}{m}\sum_{i=1}^{m}\delta_{x_{i}}$ where
$x_{i}\overset{\text{i.i.d.}}{\sim}\mu$.

We will first prove the following lemma:

######  Lemma A.11.

Consider $X_{n}\sim{\mathcal{N}}\left(0,\tfrac{I_{n}}{n}\right)$ a random
variable whose law is parametrized by $n\in{\mathbb{N}}^{*}$. There exists a
universal constant $D_{0}$ such that

$\forall n\in{\mathbb{N}}^{*},\forall
c>1,\mathbb{E}[\|X_{n}\|_{2}^{2}\;|\;\|X_{n}\|_{2}>c]\leq D_{0}c^{2}$

Jensen inequality implies:

$\forall n\in{\mathbb{N}}^{*},\forall
c>1,\mathbb{E}[\|X_{n}\|_{2};|\;\|X_{n}\|_{2}>c]\leq\sqrt{D_{0}}c$

###### Proof.

We write

$\begin{split}\mathbb{E}[\|X_{n}\|_{2}^{2}|\|X_{n}\|_{2}>c]&=\frac{\int_{c}^{+\infty}r^{2}r^{n-1}e^{-\tfrac{r^{2}n}{2}}dr}{\int_{c}^{+\infty}r^{n-1}e^{-\tfrac{r^{2}n}{2}}dr}\\\
&\leq
4c^{2}+\sum_{m=2}^{\infty}\frac{\int_{mc}^{(m+1)c}r^{2}r^{n-1}e^{-\tfrac{r^{2}n}{2}}dr}{\int_{c}^{2c}r^{n-1}e^{-\tfrac{r^{2}n}{2}}dr}\\\
&\leq
4c^{2}+\sum_{m=2}^{\infty}\frac{\int_{0}^{1}((m+1)c)^{n+1}e^{-\tfrac{((m+t)c)^{2}n}{2}}cdt}{\int_{0}^{1}c^{n-1}e^{-\tfrac{((1+t)c)^{2}n}{2}}cdt}\end{split}$

But, $\forall t\in[0,1],\forall m\geq 2$,

$\frac{((m+1)c)^{n+1}e^{-\tfrac{((m+t)c)^{2}n}{2}}}{c^{n-1}e^{-\tfrac{((1+t)c)^{2}n}{2}}}\leq
c^{2}(m+1)^{n+1}e^{-\tfrac{c^{2}n}{2}(m^{2}-1)}\leq
c^{2}(2m)^{2n}e^{-\tfrac{c^{2}n}{2}(m^{2}-1)}\leq
c^{2}e^{-n\left(\tfrac{c^{2}}{2}(m^{2}-1)-2\log(2m)\right)}$

It is clear that uniformly in $n$,
$\sum_{m=2}^{\infty}c^{2}e^{-n\left(\tfrac{c^{2}}{2}(m^{2}-1)-2\log(2m)\right)}\xrightarrow[c\rightarrow\infty]{}0$
which proves the lemma. ∎

We will now prove a bound on the rate of convergence in Wasserstein distance
of an empirical measure to the underlying distribution when this one has a
bounded support. Denote ${\mathcal{B}}_{2}^{k}(0,r)$ the euclidean ball
centered around $0$ of radius $r$ in dimension $k$.

######  Lemma A.12.

Let $\mu\in{\mathcal{P}}_{1}({\mathbb{R}}^{n})$ be a measure whose support is
included in
${\mathcal{B}}_{2}^{n}\left(0,\frac{1}{12}\right)\subset{\mathbb{R}}^{n}$ with
$n\geq 5$ Then, $\forall m\geq 1$ we have

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq
D_{1}\left(\frac{1}{m}\right)^{2/n}$

Where $D_{1}=27^{2}\left(2+\frac{1}{\sqrt{3}-1}\right)$

###### Proof.

We know from Lemma A.9 that when considering $\|\cdot\|_{2}$ the distance for
defining covering number, $\forall\varepsilon^{\prime}\leq\tfrac{1}{6}$:

${\mathcal{N}}_{\varepsilon^{\prime}}({\mathcal{B}}_{2}^{n}(0,\nicefrac{{1}}{{12}}))=\left(\frac{\tfrac{1}{12}+\tfrac{\varepsilon^{\prime}}{2}}{\tfrac{\varepsilon^{\prime}}{2}}\right)^{n}\leq\left(\frac{\tfrac{1}{6}+\varepsilon^{\prime}}{\varepsilon^{\prime}}\right)^{n}\leq(3\varepsilon^{\prime})^{-n}$

and therefore also when $\varepsilon^{\prime}\leq\tfrac{1}{27}$ Applying
Proposition 15 from Weed and Bach [2019] we get that that since
${\text{Supp}}(\mu)\subset{\mathcal{B}}_{2}^{n}(0,\nicefrac{{1}}{{12}})\subset{\mathcal{B}}_{2}^{n}(0,\nicefrac{{1}}{{12}})+{\mathcal{B}}_{2}^{n}(0,\varepsilon)$
for any $\varepsilon>0$, if $n\geq 5$,

$\forall m\geq 1,\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq
27^{2}\left(2+\frac{1}{\sqrt{3}-1}\right)\left(\frac{1}{m}\right)^{2/n}$

as well as if $n\geq 3$,

$\forall m\geq 1,\mathbb{E}[{\mathcal{W}}_{1}(\hat{\mu}_{m},\mu)]\leq
27\left(2+\frac{1}{\sqrt{3}-1}\right)\left(\frac{1}{m}\right)^{1/n}$

∎

We can prove the same kind of inequality when $\mu$ concentrates mass around
an approximately low dimensional set.

######  Lemma A.13.

Let $\mu\in{\mathcal{P}}_{1}({\mathbb{R}}^{n})$ be a measure whose support is
included in
${\mathcal{B}}_{2}^{k}\left(0,\nicefrac{{1}}{{12}}\right)\times{\mathcal{B}}^{n-k}_{2}(0,r)$
with $k\geq 5$. Then, $\forall m\leq(3r)^{-k}$ we have

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq
D_{1}\left(\frac{1}{m}\right)^{2/k}$

where $D_{1}=27^{2}\left(2+\frac{1}{\sqrt{3}-1}\right)$

###### Proof.

This is the same proof as before, just notice that
${\text{Supp}}(\mu)\subset{\mathcal{B}}_{2}^{k}(0,\nicefrac{{1}}{{12}})\times\\{0\\}^{n-k}+{\mathcal{B}}^{n}_{2}(0,r)$
and as before:

${\mathcal{N}}_{\varepsilon^{\prime}}({\mathcal{B}}_{2}^{k}(0,\nicefrac{{1}}{{12}})\times\\{0\\}^{n-k})\leq{\mathcal{N}}_{\varepsilon^{\prime}}({\mathcal{B}}_{2}^{k}(0,\nicefrac{{1}}{{12}}))\leq(3\varepsilon^{\prime})^{-k}$
(7)

if $\varepsilon^{\prime}\leq\tfrac{1}{27}$. ∎

We will now extend our results to unbounded variables very concentrated around
bounded sets, beginning with multivariate normal random variable.

######  Lemma A.14.

Consider a centered multivariate normal distribution $\mu$ on
${\mathbb{R}}^{n}$ with covariance matrix
$\operatorname{Diag}(\lambda_{1},\ldots,\lambda_{n})$ where
$\frac{1}{n}\geq\lambda_{1}\geq\ldots\geq\lambda_{n}\geq 0$. There exists two
universal constants $D_{2},E_{2}$ such that $\forall n\geq 5,\forall
m\in{\mathbb{N}}^{*}$, if $m\geq E_{2}^{n}$ then,

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq\frac{D_{2}}{n}\log(m)\left(\frac{1}{m}\right)^{2/n}$

In that case,

$\mathbb{E}[{\mathcal{W}}_{1}(\hat{\mu}_{m},\mu)]\leq\frac{\sqrt{D_{2}}}{\sqrt{n}}\sqrt{\log(m)}\left(\frac{1}{m}\right)^{1/n}$

###### Proof.

We will use previous Lemma A.12. The problem is that it applies only for a
bounded distribution. Therefore we will have to bound the mass of a
multivariate Gaussian outside of a euclidean ball. We will prove the lemma for
$\lambda_{1}=...=\lambda_{n}=\frac{1}{n}$ by noticing that it extends for
smaller eigenvalues by rescaling the axis.

Let $f\in(0,1),c>0,X\sim\mu$.

Lemma 1 from Weed and Bach [2019] tells us:

${\mathbb{P}}\left(\|X\|_{2}^{2}\geq c^{2}\sum_{i=1}^{n}\lambda_{i}\right)\leq
e^{-\tfrac{c^{2}}{4}}$

Noticing $\sum_{i=1}^{n}\lambda_{i}\leq 1$ and taking
$c=2\sqrt{\log\left(\tfrac{1}{f}\right)}$ we get that:

${\mathbb{P}}(\|X\|_{2}^{2}\geq c^{2})\leq f$

Using Lemma A.4 with $p=2f$ we get that with probability at least
$1-\exp(-2f^{2}m)$, a fraction at least $1-2f$ of vectors $x_{i}$ lies in
${\mathcal{B}}_{2}(0,c)$. We denote $H_{f}$ this event such that
${\mathbb{P}}(H_{f})\geq 1-\exp(-2f^{2}m)$.

Further denote $I$ the corresponding set of indices for $x_{i}$,
$\hat{\mu}_{m,I}=\frac{\sum_{i\in I}\delta_{x_{i}}}{|I|}$ and
$\hat{\mu}_{m,I^{c}}=\frac{\sum_{i\notin I}\delta_{x_{i}}}{|I^{c}|}$. Finally
for a Borel set $U\subset{\mathbb{R}}^{n}$ denote $\mu_{|U}$ the renormalized
restricted measure $\mu$ on $U$: $\mu_{|U}=\frac{1}{\mu(U)}\mu\mathds{1}_{U}$

We will now consider two cases:

##### 1st case

We consider the case $\frac{|I|}{m}\leq\mu({\mathcal{B}}(0,c))$ and denote
${\text{Case}}_{1}$ this set.

In that case, we can write using Lemma A.5:

$\begin{cases}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)&\leq\frac{|I|}{m}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I},\mu_{|{\mathcal{B}}_{2}(0,c)})+\frac{|I^{c}|}{m}{\mathcal{W}}_{2}^{2}\left(\hat{\mu}_{m,I^{c}},\frac{\mu_{|{\mathcal{B}}_{2}(0,c)}-\frac{|I|}{m}\mu_{|{\mathcal{B}}_{2}(0,c)}}{\frac{|I^{c}|}{m}}\right)\\\
{\mathcal{W}}_{1}(\hat{\mu}_{m},\mu)&\leq\frac{|I|}{m}{\mathcal{W}}_{1}(\hat{\mu}_{m,I},\mu_{|{\mathcal{B}}_{2}(0,c)})+\frac{|I^{c}|}{m}{\mathcal{W}}_{1}\left(\hat{\mu}_{m,I^{c}},\frac{\mu_{|{\mathcal{B}}_{2}(0,c)}-\frac{|I|}{m}\mu_{|{\mathcal{B}}_{2}(0,c)}}{\frac{|I^{c}|}{m}}\right)\par\end{cases}$

By previous Lemma A.12, we know the existence of $D_{1}$ such that if $n\geq
5$:

$\mathbb{E}\left[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I},\mu_{|{\mathcal{B}}_{2}(0,c)})\right]\leq
D_{1}(12c)^{2}\left(\frac{1}{|I|}\right)^{2/n}\leq(144D_{1})c^{2}\left(\frac{1}{\frac{|I|}{m}m}\right)^{2/n}$

Therefore we get since $\frac{2}{n}\leq 1$ and $\frac{|I|}{n}\leq 1$:

$\mathbb{E}\left[\frac{|I|}{m}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I},\mu_{|{\mathcal{B}}_{2}(0,c)})\right]\leq(144D_{1})c^{2}\left(\frac{1}{m}\right)^{2/n}$

We know from Lemma A.11 the existence of a universal constant $D_{0}$ such
that: $\forall c\geq 1,\mathbb{E}[\|X\|_{2}^{2}|\|X\|_{2}\geq c]\leq
D_{0}c^{2}$. Using a triangular inequality

$\begin{split}\mathbb{E}\left[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I^{c}},\mu_{|{\mathcal{B}}_{2}(0,c)^{c}})\right]&\leq
4D_{0}c^{2}\end{split}$

Finally, conditioned on the event $H_{f}$ and that we are in
${\text{Case}}_{1}$, we get that if $c\geq 1$:

$\begin{cases}\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)|H_{f}\bigcap{\text{Case}}_{1}]\leq(144D_{1})c^{2}\left(\frac{1}{m}\right)^{2/n}+8fD_{0}c^{2}\\\
c=2\sqrt{\log(\tfrac{1}{f})}\end{cases}$

##### 2nd case

We consider the case $\frac{|I|}{m}>\mu({\mathcal{B}}_{2}(0,c))$ and denote
${\text{Case}}_{2}$ this set.

In that case, we denote $I^{\prime}\subset I$ taken randomly uniformly, such
that $|I^{\prime}|=\max\left\\{k\geq
1,\frac{k}{m}\leq\mu({\mathcal{B}}(0,c))\right\\}$ and denote
$\hat{\mu}_{m,I^{\prime}}$ the renormalized empirical measure with points in
$I^{\prime}$ and $\hat{\mu}_{m,I\bigcap I^{{}^{\prime}c}}$ the renormalized
empirical measure with points in $I\bigcap I^{\prime c}$.

We can write:

$\begin{split}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)&\leq\frac{|I^{\prime}|}{m}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I^{\prime}},\mu_{|{\mathcal{B}}_{2}(0,c)})+\frac{|I|-|I^{\prime}|}{m}{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I\bigcap
I^{\prime c}},\mu_{|{\mathcal{B}}_{2}(0,c)^{c}})\\\
&+\left(\mu\big{(}{\mathcal{B}}_{2}(0,c)\big{)}-\frac{|I^{\prime}|}{m}\right){\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I\bigcap
I^{\prime
c}},\mu_{|{\mathcal{B}}_{2}(0,c)})+\frac{|I^{c}|}{m}{\mathcal{W}}_{2}^{2}\left(\hat{\mu}_{m,I^{c}},\mu_{|{\mathcal{B}}_{2}(0,c)^{c}}\right)\end{split}$

Provided $fm>1$, we know that $\frac{|I^{\prime}|}{m}\geq 1-3f$.

We can repeat all the previous arguments and get that if $c\geq 1$ and $m$:

$\begin{cases}\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)|H_{f}\bigcap{\text{Case}}_{2}]\leq(144D_{1})c^{2}\left(\frac{1}{m}\right)^{2/n}+16fD_{0}c^{2}\\\
c=2\sqrt{\log(\tfrac{1}{f})}\end{cases}$

Finally,

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq{\mathbb{P}}(H_{f})\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)|H_{f}]+(1-{\mathbb{P}}(H_{f}))\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)|H_{f}^{c}]$

We can easily bound as before
$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)|H_{f}^{c}]\leq
4D_{0}c^{2}$

which yields finally:

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq(144D_{1})4\log(\frac{1}{f})\left(\frac{1}{m}\right)^{2/n}+16fD_{0}4\log(\frac{1}{f})+2\exp(-f^{2}m)4D_{0}4\log(\frac{1}{f})$

Taking $f=\left(\frac{1}{m}\right)^{2/n}$ we get

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq\frac{1152D_{1}}{n}\log(m)\left(\frac{1}{m}\right)^{2/n}+D_{0}\frac{128}{n}\log(m)\left(\frac{1}{m}\right)^{2/n}+\frac{64}{n}\exp(-m^{1-4/n})\log(m)$

Note now that there exists a universal constant $C>0$ such that $\forall n\geq
5,\forall m\geq 1$:

$\exp(-m^{1-4/n})\leq\exp(-m^{1/5})\leq C\left(\frac{1}{m}\right)^{2/5}\leq
C\left(\frac{1}{m}\right)^{2/n}$

Finally we get the existence of universal constants $D_{2},E_{2}$ such that if
$\left(\frac{1}{m}\right)^{2/n}\leq\frac{1}{E_{2}^{2}}$ (to ensure $c\geq 1$
take for example $E_{2}=\exp(\frac{1}{4})$),

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq\frac{D_{2}}{n}\log(m)\left(\frac{1}{m}\right)^{2/n}$

To prove the second part of the lemma just apply Lemma A.6 and Jensen
inequality.

∎

We can finally extend Lemma A.13 to unbounded distributions as we just
extended Lemma A.12 to unbounded distributions.

######  Lemma A.15.

Let $\lambda_{1}\geq\ldots\geq\lambda_{n}$ and
$\mu={\mathcal{N}}(0,\operatorname{Diag}((\lambda_{i})_{i=1}^{n}))$, with
$k\geq 5$. Suppose $1\geq\sum_{i=1}^{k}\lambda_{i}$. Denote
$\eta=\frac{\sqrt{\sum_{i=k+1}^{n}\lambda_{i}}}{4\sqrt{\sum_{i=1}^{k}}\lambda_{i}}$.
We know the existence of two universal constants
$D_{2}^{\prime},E^{\prime}_{2}$ such that if $\eta^{-k}\geq m\geq
E_{2}^{\prime k}$, then:

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq\frac{D^{\prime}_{2}}{k}\log(m)\left(\frac{1}{m}\right)^{2/k}]$

In that case,

$\mathbb{E}[{\mathcal{W}}_{1}(\hat{\mu}_{m},\mu)]\leq\frac{\sqrt{D^{\prime}_{2}}}{\sqrt{k}}\sqrt{\log(m)}\left(\frac{1}{m}\right)^{1/k}]$

###### Proof.

We will follow the same steps as previously. Let $X\sim\mu$ and denote
$\underaccent{\bar}{X}=(X_{1},\ldots,X_{k})\in{\mathbb{R}}^{k}$,
$\bar{X}=(X_{k+1},\ldots,X_{n})\in{\mathbb{R}}^{n-k}$.

$\begin{split}&{\mathbb{P}}\left(\|\underaccent{\bar}{X}\|_{2}^{2}\geq
c^{2}\sum_{i=1}^{k}\lambda_{i}\right)\leq e^{-\frac{c^{2}}{4}}\\\
&{\mathbb{P}}\left(\|\bar{X}\|_{2}^{2}\geq
c^{2}\sum_{i=k+1}^{n}\lambda_{i}\right)\leq e^{-\frac{c^{2}}{4}}\end{split}$

Take $c=2\sqrt{\log(\tfrac{2}{f})}$. Then, by the same arguments as before,
using Lemma A.4 and union bounds, with probability at least
$1-2\exp(-f^{2}m/2)$ a fraction at least $1-2f$ of points $x_{i}$ are in
${\mathcal{B}}_{c}:={\mathcal{B}}_{k}(0,c\sqrt{\sum_{i=1}^{k}\lambda_{i}})\times{\mathcal{B}}_{n-k}(0,c\sqrt{\sum_{i=k+1}^{n}\lambda_{i}})$.
We denote $H_{f}$ such an event and $I$ such a set of indices.

By using Lemma A.13, we know that in that case we can bound

$\mathbb{E}\left[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m,I},\mu_{|{\mathcal{B}}_{c}})|H_{f}\right]\leq\left(12c\sqrt{\sum_{i=1}^{k}\lambda_{i}}\right)^{2}C_{1}\left(\frac{1}{m}\right)^{2/k}\leq(12c)^{2}C_{1}\left(\frac{1}{m}\right)^{2/k}$

if
$m\leq\left(\frac{3\sqrt{\sum_{i=k+1}^{m}\lambda_{i}}}{12\sqrt{\sum_{i=1}^{k}\lambda_{i}}}\right)^{-k}$
where we have used $\sum_{i=1}^{k}\lambda_{i}\leq 1$.

Moreover, we know that

$\begin{split}\mathbb{E}\left[\|X\|_{2}^{2}|X\notin{\mathcal{B}}^{k}_{2}(0,c\sqrt{\sum_{i=1}^{k}\lambda_{i}})\times{\mathcal{B}}_{2}^{n-k}(0,c\sqrt{\sum_{i=k+1}^{n}\lambda_{i}})\right]&\leq\mathbb{E}\left[\|\underaccent{\bar}{X}\|_{2}^{2}|\underaccent{\bar}{X}\notin{\mathcal{B}}^{k}_{2}(0,c\sqrt{\sum_{i=1}^{k}\lambda_{i}})\right]\\\
&+\mathbb{E}\left[\|\bar{X}\|_{2}^{2}|\bar{X}\notin{\mathcal{B}}^{n-k}_{2}(0,c\sqrt{\sum_{i=k+1}^{n}\lambda_{i}})\right]\\\
&\leq
D_{0}c^{2}\sum_{i=1}^{k}\lambda_{i}+D_{0}c^{2}\sum_{i=k+1}^{n}\lambda_{i}\\\
&\leq 2D_{0}c^{2}\sum_{i=1}^{k}\lambda_{i}\\\ &\leq 2D_{0}c^{2}\end{split}$

We just need to differentiate the same two cases as in the proof of Lemma A.14
to get that finally, if
$m\leq\left(\frac{\sqrt{\sum_{i=k+1}^{m}\lambda_{i}}}{4\sqrt{\sum_{i=1}^{k}\lambda_{i}}}\right)^{-k}$
we can bound as before, with $c=2\sqrt{\log(\tfrac{2}{f})}$:

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq(12c)^{2}C_{1}\left(\frac{1}{m}\right)^{2/k}+16f*4C_{0}c^{2}+2\exp(-\frac{f^{2}m}{2})*4*2C_{0}c^{2}$

Taking $f=\left(\frac{1}{m}\right)^{2/k}$, we see as before the existence of
$E_{2}^{\prime},D_{2}^{\prime}$ such that if $E_{2}^{\prime k}\leq
m\leq\eta^{-k}$ then,

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{m},\mu)]\leq\frac{D^{\prime}_{2}}{k}\log(m)\left(\frac{1}{m}\right)^{2/k}$

We choose $E^{\prime}_{2}=\sqrt{2e^{\nicefrac{{-1}}{{4}}}}$ such that
$c>1\Leftrightarrow 2\sqrt{\log(\tfrac{2}{f})}>1\Leftrightarrow
2\sqrt{\log(\tfrac{2}{\left(\frac{1}{m}\right)^{2/k}})}>1\Leftrightarrow
m>\sqrt{2e^{\nicefrac{{-1}}{{4}}}}^{k}$

For the second part of the lemma, apply Lemma A.6 and Jensen inequality.

∎

###  A.6 Proof of Lemma 4.1

See 4.1

###### Proof.

We know from Assumption 3 the existence of a random empirical measure with
$\tilde{m}_{l+1}$ points $\hat{\mu}_{\tilde{m}_{l+1}}$ such that
$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{A}^{{\mathcal{I}}^{\ell}},\hat{\mu}^{{\mathcal{I}}^{\ell}}_{\tilde{m}_{l+1}})]\leq
C_{1}$. Therefore by using Lemma A.3 and a carefully chosen permutation, we
can consider the (random) matrix $W_{\tilde{m}_{l+1}}$ associated to
$W_{A}^{\ell+1}$ such that
$\mathbb{E}[\|W_{\tilde{m}_{l+1}}^{{\mathcal{I}}^{\ell}}-W_{A}^{{\mathcal{I}}^{\ell}}\|_{2}^{2}]\leq
C_{1}m_{\ell+1}$. Note that since $W_{\tilde{m}_{\ell+1}}$ comes from an
empirical measure with only $\tilde{m}_{\ell+1}$ points one can denote
${\mathcal{I}}^{\ell+1}=\\{I^{\ell+1}_{1},\ldots,I^{\ell+1}_{\tilde{m}_{\ell+1}}\\}$
the (random) equi-partition of $[m_{\ell+1}]$ delimiting its equal rows. In
the same way, since $A$ and $B$ have the same weights distribution, we can
find a permutation $\Pi_{\ell+1}$ of the layer $\ell+1$ of network $B$ such
that denoting
$\tilde{W}^{\ell+1}_{B}=\Pi_{\ell+1}W_{B}^{\ell+1}\Pi_{\ell}^{T}$ and taking
expectations over the choice of the weights matrices, we have
$\mathbb{E}[\|\tilde{W}^{\ell+1}_{B}-W_{\tilde{m}_{\ell+1}}\|_{2}^{2}]\leq
m_{\ell+1}C_{1}$. Consider
$W^{\ell+1}_{M_{t}}=tW_{A}^{\ell+1}+(1-t)\tilde{W}_{B}^{\ell+1}$ and denote
$\underaccent{\bar}{\phi}^{\ell+1}(x)=\sigma(W_{\tilde{m}_{\ell+1}}\underaccent{\bar}{\phi}^{\ell}(x)),\phi_{A}^{\ell+1}(x)=\sigma(W^{\ell+1}_{A}\phi^{\ell}_{A}(x)),\phi^{\ell+1}_{B}(x)=\sigma(W_{B}^{\ell+1}\phi^{\ell}_{B}(x)),\phi^{\ell+1}_{M_{t}}(x)=\sigma(W^{\ell+1}_{M_{t}}\phi^{\ell}_{M_{t}}(x))$.
It is clear that $\forall k\in[\tilde{m}_{\ell+1}],\forall i,j\in
I^{\ell+1}_{k},\underaccent{\bar}{\phi}^{\ell+1}_{i}(x)=\underaccent{\bar}{\phi}^{\ell+1}_{j}(x)$
by definition of the choice of the equi-partition ${\mathcal{I}}^{\ell+1}$.

We will finally denote
$\underaccent{\bar}{\phi}^{\prime\ell+1}(x)\in{\mathbb{R}}^{\tilde{m}_{\ell+1}},\underaccent{\bar}{\phi}^{\prime\ell}(x)\in{\mathbb{R}}^{\tilde{m}_{\ell}}$
the vectors
$\underaccent{\bar}{\phi}^{\ell+1}(x),\underaccent{\bar}{\phi}^{\ell}(x)$
where we have kept only one index in each of the elements of the partitions
respectively ${\mathcal{I}}^{\ell+1},{\mathcal{I}}^{\ell}$.

Moreover using that the non-linearity is pointwise $1$-Lipschitz and
$\sigma(0)=0$,

$\mathbb{E}[\|\underaccent{\bar}{\phi}^{\ell+1}(x)\|_{2}^{2}]\leq\mathbb{E}_{P,Q}[\mathbb{E}[\|W_{\tilde{m}_{\ell+1}}\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}|\underaccent{\bar}{\phi}^{\ell}(x)]]\leq\mathbb{E}_{P,Q}[C_{2}\frac{m_{\ell+1}}{m_{\ell}}\|\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}]\leq
m_{\ell+1}C_{2}\underaccent{\bar}{E}_{\ell}$

which yields
$\underaccent{\bar}{E}_{\ell+1}=C_{2}\underaccent{\bar}{E}_{\ell}$ where we
have used Assumption 4.

Then

$\begin{split}\mathbb{E}_{P,Q}[\|\phi^{\ell+1}_{A}(x)-\underaccent{\bar}{\phi}^{\ell+1}(x)\|_{2}^{2}]&\leq\mathbb{E}_{P,Q}[\|W^{\ell+1}_{A}\phi^{\ell}_{A}(x)-W_{\tilde{m}_{\ell+1}}\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\\\
&\leq
2\mathbb{E}_{P,Q}[\|W_{A}^{\ell+1}(\phi^{\ell}_{A}(x)-\underaccent{\bar}{\phi}^{\ell}(x))\|_{2}^{2}+\|(W_{A}^{\ell+1}-W_{\tilde{m}_{\ell+1}})\underaccent{\bar}{\phi}^{\ell}\|_{2}^{2}]\\\
&\leq
2m_{\ell+1}C_{2}E_{\ell}+2\mathbb{E}_{P,Q}[\|(W_{A}^{{\mathcal{I}}^{\ell}}-W^{{\mathcal{I}}^{\ell}}_{\tilde{m}})\underaccent{\bar}{\phi}^{\prime\ell}\|_{2}^{2}]\\\
&\leq
2m_{\ell+1}C_{2}E_{\ell}+2\mathbb{E}_{P,Q}[\|(W_{A}^{{\mathcal{I}}_{\ell}}-W^{{\mathcal{I}}_{\ell}}_{\tilde{m}_{\ell+1}})\|_{2}^{2}]\mathbb{E}_{P,Q}\|\underaccent{\bar}{\phi}^{\prime\ell}\|_{2}^{2}]\\\
&\leq
2m_{\ell+1}C_{2}E_{\ell}+2m_{\ell+1}C_{1}\frac{\tilde{m}_{\ell}}{m_{\ell}}\mathbb{E}_{P,Q}\|\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}\\\
&\leq
2m_{\ell+1}C_{2}E_{\ell}+2m_{\ell+1}C_{1}\tilde{m}_{\ell}\underaccent{\bar}{E}_{\ell}\end{split}$

where we have used
$(W^{\ell+1}_{A}-W_{\tilde{m}_{\ell+1}})\underaccent{\bar}{\phi}^{\ell}(x)=(W_{A}^{{\mathcal{I}}^{\ell}}-W^{{\mathcal{I}}^{\ell}}_{\tilde{m}_{\ell+1}})\underaccent{\bar}{\phi}^{\prime\ell}$
and
$\|\underaccent{\bar}{\phi}^{\prime\ell}(x)\|_{2}^{2}=\frac{\tilde{m}_{\ell}}{m_{\ell}}\|\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}$.

We do the same computations for
$\mathbb{E}_{P,Q}[\|\phi^{\ell+1}_{B}(x)-\underaccent{\bar}{\phi}^{\ell+1}(x)\|_{2}^{2}]$.

Finally,

$\begin{split}\mathbb{E}_{P,Q}[\|\phi^{\ell+1}_{M_{t}}(x)-\underaccent{\bar}{\phi}^{\ell+1}(x)\|_{2}^{2}]&\leq
t\mathbb{E}_{P,Q}[\|W^{\ell+1}_{A}\phi^{\ell}_{M_{t}}(x)-W_{\tilde{m}_{\ell+1}}\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}]+(1-t)\mathbb{E}_{P,Q}[\|W^{\ell+1}_{B}\phi_{M_{t}}(x)-W_{\tilde{m}_{\ell+1}}\underaccent{\bar}{\phi}^{\ell}(x)\|_{2}^{2}]\\\
&\leq
2C_{2}E_{\ell}m_{\ell+1}+2C_{1}\tilde{m}_{\ell}\underaccent{\bar}{E}_{\ell}m_{\ell+1}\end{split}$

where we have used convexity for the first inequality and then the same proof
as above for both terms. This yields
$E_{\ell+1}=2C_{2}E_{\ell}+2C_{1}\tilde{m}_{\ell}\underaccent{\bar}{E}_{\ell}$.

∎

###  A.7 Proof of Lemma 5.1

######  Lemma A.16 (Version of Assumption 3 for normal distribution).

Consider $\mu_{\ell}\in{\mathcal{P}}_{1}({\mathbb{R}}^{m_{\ell-1}})$ a
multivariate Gaussian distribution with covariance matrix
$\Sigma^{\ell}=\operatorname{Diag}(\lambda^{\ell}_{1}I_{p_{\ell-1}},\ldots,\lambda^{\ell}_{\tilde{m}_{\ell}}I_{p_{\ell-1}})$
where $p_{\ell-1}=\frac{m_{\ell-1}}{\tilde{m}_{\ell-1}}$. Suppose that
$\frac{1}{m_{\ell-1}}\geq\lambda^{\ell}_{1},\geq\ldots\geq\lambda_{\tilde{m}_{\ell-1}}^{\ell}$
with $\tilde{m}_{\ell-1}\geq 5$. Then, there exists two universal constants
$D_{3},E_{3}$ such that $\forall\tilde{m}_{\ell}\geq
E_{3}^{\tilde{m}_{\ell-1}}$ there exists a random empirical measure
$\hat{\mu}_{\tilde{m}}$ with only $\tilde{m}_{\ell}$ points such that $\forall
m_{\ell}\geq\tilde{m}_{\ell}$

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{\tilde{m}_{\ell}})]\leq\frac{D_{3}}{\tilde{m}_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}$

In that case:

$\mathbb{E}[{\mathcal{W}}_{1}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{\tilde{m}_{\ell}})]\leq\frac{\sqrt{D_{3}}}{\sqrt{\tilde{m}_{\ell-1}}}\sqrt{\log(\tilde{m}_{\ell})}\left(\frac{1}{\tilde{m}_{\ell}}\right)^{1/\tilde{m}_{\ell-1}}$

###### Proof.

We know that the distribution on the rows of $W^{{\mathcal{I}}^{\ell-1}}_{A}$
in ${\mathbb{R}}^{\tilde{m}_{\ell-1}}$ is multivariate Gaussian with
covariance matrix $\frac{I_{\tilde{m}_{\ell-1}}}{\tilde{m}_{\ell-1}}$ since
each parameters is obtained by summing the $p_{\ell-1}$ corresponding
parameters of the row of $W_{A}^{\ell}$ which has covariance matrix
$\frac{I_{m_{\ell-1}}}{m_{\ell-1}}$ by hypothesis. Therefore using Lemma A.14,
we know the existence of constants $D_{2},E_{2}$ such that if $m_{\ell}\geq
E_{2}^{\tilde{m}_{\ell-1}}$,
$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\mu_{\ell}^{{\mathcal{I}}^{\ell-1}})]\leq\frac{D_{2}}{\tilde{m}_{\ell-1}}\log(m_{\ell})\left(\frac{1}{m_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}$.

Therefore considering $\hat{\mu}_{\tilde{m}_{\ell}}$ with the same law but
only a fixed number $\tilde{m}_{\ell}$ of elements, we get for
$m_{\ell}\geq\tilde{m}_{\ell}$:

$\begin{split}&\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\mu^{{\mathcal{I}}^{\ell-1}}_{\ell})]\leq\frac{D_{2}}{\tilde{m}_{\ell-1}}\log(m_{\ell})\left(\frac{1}{m_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}\leq\frac{D_{2}}{\tilde{m}_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}\\\
&\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{\tilde{m}_{\ell}}^{{\mathcal{I}}^{\ell-1}},\mu_{\ell}^{{\mathcal{I}}^{\ell-1}})]\leq\frac{D_{2}}{\tilde{m}_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}\end{split}$

Indeed, the first inequality can be obtained by noticing that
$\left(x\mapsto\log(x)\left(\frac{1}{x}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}\right)$
is decreasing for $x\geq\sqrt{e}^{\tilde{m}_{\ell-1}}$ and hence one can just
increase the constant $E_{3}$ considered: taking
$E_{3}=\max\\{\sqrt{e},E_{2}\\}$ and $D_{3}=4D_{2}$, by triangular inequality,

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{\tilde{m}_{\ell}})]\leq
2\left(\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\mu_{\ell}^{{\mathcal{I}}^{\ell-1}})]+\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{\tilde{m}_{\ell}},\mu_{\ell}^{{\mathcal{I}}^{\ell-1}})]\right)\leq\frac{D_{3}}{\tilde{m}_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{\nicefrac{{2}}{{\tilde{m}_{\ell-1}}}}$

For the second part of the lemma, just apply Lemma A.6 and Jensen inequality.

∎

######  Lemma A.17 (Version of Assumption 4 for normal variable).

$\forall X\in{\mathbb{R}}^{m_{\ell-1}}$ we have:

$\begin{split}&\mathbb{E}[\|W^{\ell}_{A}X\|_{2}^{2}]\leq\frac{m_{\ell}}{m_{\ell-1}}\|X\|_{2}^{2}\\\
&\mathbb{E}[\|W_{\tilde{m}_{\ell}}X\|_{2}^{2}]\leq\frac{m_{\ell}}{m_{\ell-1}}\|X\|_{2}^{2}\end{split}$

###### Proof.

First, given $X\in{\mathbb{R}}^{m_{\ell-1}},\forall
i\in[m_{\ell}],(W^{\ell}_{A}X)_{i}=\sum_{j=1}^{m_{\ell-1}}(W^{\ell}_{A})_{i,j}X_{j}$
where $(W^{\ell}_{A})_{i,j}$ are iid following
${\mathcal{N}}(0,\frac{1}{m_{\ell-1}})$. Therefore, $\forall
i\in[m_{\ell}],\mathbb{E}[(W^{\ell}_{A}X)_{i}^{2}]=\frac{1}{m_{\ell-1}}\|X\|_{2}^{2}$.
Finally,

$\mathbb{E}[\|W^{\ell}_{A}X\|_{2}^{2}]=\sum_{i=1}^{m_{\ell}}\mathbb{E}[(W_{A}^{\ell}X)_{i}^{2}]=\frac{m_{\ell}}{m_{\ell-1}}\|X\|_{2}^{2}$

Since the weight matrix associated to $\hat{\mu}_{\tilde{m}_{\ell}}$ denoted
$W^{\prime}_{\tilde{m}_{\ell}}\in{\mathcal{M}}_{\tilde{m}_{\ell},m_{\ell-1}}({\mathbb{R}})$
has only $\tilde{m}_{\ell}$ raws, we have expanded it to a matrix
$W_{\tilde{m}_{\ell}}\in{\mathcal{M}}_{m_{\ell},m_{\ell-1}}({\mathbb{R}})$ by
cloning raws in the same element of the partition ${\mathcal{I}}^{\ell}$ given
by the pairing between
$\hat{\mu}^{{\mathcal{I}}^{\ell}}_{A},\hat{\mu}_{\tilde{m}_{\ell-1}}^{{\mathcal{I}}^{\ell}}$
that minimizes the Wasserstein distance. Since
$W^{\prime}_{\tilde{m}_{\ell}}\in{\mathcal{M}}_{\tilde{m}_{\ell},m_{\ell-1}}({\mathbb{R}})$
built in the proof of Lemma A.16 has the same law on raws $\mu_{\ell}$ as
$W^{\ell}_{A},W^{\ell}_{B}$ but with only $\tilde{m}_{\ell}$ raws, we can use
what preceeds to get:

$\forall
X\in{\mathbb{R}}^{m_{\ell-1}},\mathbb{E}[\|W^{\prime}_{\tilde{m}_{\ell}}X\|_{2}^{2}]\leq\frac{\tilde{m}_{\ell}}{m_{\ell-1}}\|X\|_{2}^{2}$

Noting that
$\|W_{\tilde{m}_{\ell}}X\|_{2}^{2}=\frac{m_{\ell}}{\tilde{m}_{\ell}}\|W^{\prime}_{\tilde{m}_{\ell}}X\|_{2}^{2}$,
we get

$\forall
X\in{\mathbb{R}}^{m_{\ell-1}},\mathbb{E}[\|W_{\tilde{m}_{\ell}}X\|_{2}^{2}]\leq\frac{m_{\ell}}{m_{\ell-1}}\|X\|_{2}^{2}$

which concludes the proof of the lemma. ∎

Having the two assumptions we need, we can prove Lemma 5.1.

We recall it here:

See 5.1

###### Proof.

From $\mathbb{E}_{x\sim P}[\|x\|_{2}^{2}]\leq m_{0}$ we get immediately
Property 1 at the input layer with $\underaccent{\bar}{E}_{0}=1,E_{0}=0$.

By the recursive relation of Lemma 4.1 and using Lemmas A.16 and A.17, we get
Property 1 at each hidden layer $\ell\in[L]$ with $\tilde{m}_{\ell}$ to be
chosen later with
$m_{\ell}\geq\tilde{m}_{\ell}\geq\min\\{5,E_{3}^{\tilde{m}_{l-1}}\\}$ and:

$\begin{cases}\underaccent{\bar}{E}_{\ell}=1\\\
E_{\ell}=\sum_{i=1}^{\ell}2^{\ell+1-i}D_{3}\log(\tilde{m}_{i})\left(\frac{1}{\tilde{m}_{i}}\right)^{\nicefrac{{2}}{{\tilde{m}_{i-1}}}}\underaccent{\bar}{E}_{i-1}\end{cases}$

Therefore, just take $\forall
i\in[L],\log(\tilde{m}_{i})\left(\frac{1}{\tilde{m}_{i}}\right)^{\frac{2}{\tilde{m}_{i-1}}}\leq\varepsilon^{2}\frac{1}{2^{L+1-i}L}$,
i.e.

$\tilde{m}_{i}=\tilde{{\mathcal{O}}}\left(\frac{T_{i}}{\varepsilon}\right)^{\tilde{m}_{i-1}}$

where $T_{i}=\sqrt{2^{L+1-i}L}$ and the notation
$\tilde{{\mathcal{O}}}(\cdot)$ hides logarithmic terms.

In that case,

$\begin{cases}\underaccent{\bar}{E}_{L}=1\\\ E_{L}=\varepsilon^{2}\end{cases}$

∎

###  A.8 Proof of Theorem 5.2

We prove here Theorem 5.2 that we recall:

Under normal initialization of the weights, for
$m_{1}\geq\tilde{m}_{1},\ldots,m_{L}\geq\tilde{m}_{L}$ as defined in Lemma
5.1, $m_{0}\geq 5$, and under Assumption 7 we know that $\forall t\in[0,1]$,
with $Q$-probability at least $1-\delta_{Q}$, there exists permutations of
hidden layers $1,\ldots,L$ of network $B$ that are independent of $t$, such
that:
$\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{M_{t}}(x),y\right)\right]\leq
t\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{A}(x),y\right)\right]+(1-t)\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{B}(x),y\right)\right]+\tfrac{4\sqrt{m_{L+1}}}{\delta_{Q}^{2}}\varepsilon$

###### Proof.

Under assumptions of Lemma 5.1, given $A,B$, we know the existence of (random)
permutations of the hidden layers $\Pi_{1},\ldots,\Pi_{L}$ such that for
$1\leq\ell\leq L$, denoting $M_{t}$ the mean network of weight matrix at layer
$\ell$: $tW_{A}^{\ell}+(1-t)\Pi_{\ell}W_{B}^{\ell}\Pi_{\ell-1}^{T}$ we know
the existence of
$\underaccent{\bar}{\phi}^{L}:{\mathbb{R}}^{m_{0}}\rightarrow{\mathbb{R}}^{m_{L}}$
such that:

$\begin{split}&\mathbb{E}_{P,Q}[\|\phi_{M_{t}}^{L}(x)-\phi_{A}^{L}(x)\|_{2}^{2}]\leq\varepsilon^{2}m_{L}\\\
&\mathbb{E}_{P,Q}[\|\phi_{M_{t}}^{L}(x)-\phi_{B}^{L}(x)\|_{2}^{2}]\leq\varepsilon^{2}m_{L}\end{split}$

Then, by convexity, we get at the last layer:

$\begin{split}\mathbb{E}_{P,Q}[\|\left(tW_{A}^{L+1}+(1-t)W_{B}^{L+1}\Pi_{L}^{T}\right)]\phi_{M_{t}}^{L}(x)&-tW_{A}^{L+1}\phi_{A}^{L}(x)-(1-t)W_{B}^{L+1}\Pi_{L}^{T}\phi_{B}^{L}(x)\|_{2}^{2}]\\\
&\leq
t\mathbb{E}[\|W_{A}^{L+1}(\phi_{M_{t}}^{L}(x)-\phi_{A}^{L}(x))\|_{2}^{2}]+(1-t)\mathbb{E}[\|W_{B}^{L+1}\Pi_{L}^{T}(\phi_{M_{t}}^{L}(x)-\phi_{B}^{L}(x))\|_{2}^{2}]]\\\
&\leq\varepsilon^{2}m_{L+1}\end{split}$

Finally, by Jensen inequality,

$\mathbb{E}[\|\left(tW_{A}^{L+1}+(1-t)W_{B}^{L+1}\Pi_{L}^{T}\right)]\phi_{M_{t}}^{L}(x)-tW_{A}^{L+1}\phi_{A}^{L}(x)-(1-t)W_{B}^{L+1}\Pi_{L}^{T}\phi_{B}^{L}(x)\|_{2}]\leq\sqrt{m_{L+1}}\varepsilon$

Therefore, we get by applying two successive Markov lemma, that with
probability at least $1-\delta_{Q}$ over the choice of the networks $A,B$:

$\mathbb{E}_{x\sim
P}[\|\left(tW_{A}^{L+1}+(1-t)W_{B}^{L+1}\Pi_{L}^{T}\right)]\phi_{M_{t}}^{L}(x)-tW_{A}^{L+1}\phi_{A}^{L}(x)-(1-t)W_{B}^{L+1}\Pi_{L}^{T}\phi_{B}^{L}(x)\|_{2}]\leq\frac{\sqrt{m_{L+1}}\varepsilon}{\left(\tfrac{\delta_{Q}}{2}\right)^{2}}$

Indeed remember that we have introduced an intermediate random measure
$\hat{\mu}_{\tilde{m}_{\ell}}$ which the permutations depend on and that
intervenes in the expectation.

Using convexity of the loss and $1$-Lipschitzness we get for all $t\in[0,1]$
that with probability at least $1-\delta_{Q}$ over the choice of the networks
$A,B$:

$\mathbb{E}_{x\sim P}[{\mathcal{L}}(f_{M_{t}}(x),y)]\leq
t\mathbb{E}[{\mathcal{L}}(f_{A}(x),y)]+(1-t)\mathbb{E}[{\mathcal{L}}(f_{B}(x),y)]+\frac{4\sqrt{m_{L+1}}\varepsilon}{\delta_{Q}^{2}}$

∎

###  A.9 Approximately low dimensional underlying weights distribution

####  A.9.1 Motivation on the structure of the covariance matrix

Remind that we are given a partition of the input layer of the weights
$[m_{\ell-1}]$ in $\tilde{m}_{\ell-1}$ different groups of the same size
${\mathcal{I}}^{\ell-1}=\\{I^{\ell-1}_{1},...,I^{\ell-1}_{\tilde{m}_{\ell-1}}\\}$.
Suppose we have already permuted the first layer of network $A$ and $B$ we can
suppose that $I^{\ell-1}_{1}=\\{1,\ldots,p_{\ell-1}\\},\ldots$ where
$p_{\ell-1}:=\frac{m_{\ell-1}}{\tilde{m}_{\ell-1}}$. We want a covariance
matrix that respects the fact that incoming neurons in a given group behave
the same. Therefore the covariance matrix must be invariant under the
permutations of indices inside a set of the equi-partition. We will write the
Kroenecker product $\otimes$.

######  Lemma A.18.

To respect symmetries of the incoming layer, the covariance matrix of weights
is necessarily of the form

$\Sigma_{\ell}=D_{\ell}\otimes
I_{p_{\ell-1}}+B_{\ell}\otimes\mathds{1}_{p_{\ell-1}}$

where $D_{\ell}\in{\mathcal{M}}_{\tilde{m}_{\ell-1}}({\mathbb{R}})$ is
diagonal, $B_{\ell}\in{\mathcal{S}}_{\tilde{m}_{\ell-1}}({\mathbb{R}})$ is
symmetric.

###### Proof.

Denoting the matrix of covariances by blocks like:

$\Sigma_{\ell}=\begin{pmatrix}A_{11}&A_{12}&...&A_{1\tilde{m}_{\ell-1}}\\\
A_{21}&A_{22}&...&A_{2\tilde{m}_{\ell-1}}\\\ ...&...&...&...\\\
A_{\tilde{m}_{\ell-1}1}&A_{\tilde{m}_{\ell-1}2}&...&A_{\tilde{m}_{\ell-1}\tilde{m}_{\ell-1}}\end{pmatrix}$

we get the relation
$\forall\Pi_{1},...,\Pi_{\tilde{m}_{\ell-1}}\in{\mathcal{S}}_{p_{\ell-1}}$
permutation matrices, denoting
$\Pi=\operatorname{Diag}(\Pi_{1},\ldots,\Pi_{\tilde{m}_{\ell-1}})$

$\Sigma_{\ell}=\mathbb{E}\left[XX^{T}\right]=\mathbb{E}\left[(\Pi X)(\Pi
X)^{T}\right]=\Pi\Sigma\Pi^{T}=\operatorname{Diag}(\Pi_{1},...,\Pi_{\tilde{m}_{\ell-1}})\Sigma\operatorname{Diag}(\Pi_{1}^{T},...,\Pi_{\tilde{m}_{\ell-1}}^{T})$

Evaluating this relation for any
$\Pi_{1},\Pi_{2},...,\Pi_{\tilde{m}_{\ell-1}}$ we get
$\forall\Pi_{1}\in{\mathcal{S}}_{p_{\ell-1}},\Pi_{1}A_{11}\Pi_{1}^{T}=A_{11}$
and therefore $A_{11}$ is of the form
$d_{11}I_{p_{\ell-1}}+b_{11}\mathds{1}_{p_{\ell-1}}$.

We do the same for $A_{ii},i\geq 2$.

Moreover we get for $A_{12}$ that:

$\forall\Pi_{1},\Pi_{2}\in{\mathcal{S}}_{p_{\ell-1}},\Pi_{1}A_{12}\Pi_{2}^{T}=A_{12}$

which brings that $A_{12}$ is of the form $b_{12}\mathds{1}_{p_{\ell-1}}$. We
do the same for all $A_{ij}$ where $i\neq j$. This concludes the proof.

Finally by summing over columns inside the partition we get:

$\Sigma_{\ell}^{{\mathcal{I}}^{\ell-1}}=p_{\ell-1}D_{\ell}+p_{\ell-1}^{2}B_{\ell}$

∎

The model that we chose in Section 5.4 is a particular case that corresponds
to choosing
$D_{\ell}=\operatorname{Diag}(\lambda^{\ell}_{1},...,\lambda^{\ell}_{\tilde{m}_{\ell-1}})$
and $B_{\ell}=0$. This is not the most general since it implies independence
between weights coming from different groups but is sufficient to show the
influence of low feature dimensionality on LMC efficiency.

####  A.9.2 Non diagonal model

A natural direction is to consider the case where the matrix $B_{\ell}$ is non
zero. Since $\Sigma_{\ell}^{{\mathcal{I}}^{\ell-1}}$ is symmetric positive,
one can orthogonally change the basis where it becomes diagonal. The arguments
to prove Assumption 3 remains unchanged since they rely exclusively on the
eigenvalues of $\Sigma_{\ell}^{{\mathcal{I}}^{\ell-1}}$.

However one important point to check, is about Assumption 4. Indeed having
covariance between weights of a given block implies that the errors at all
neurons of a given layer may sum up. Take for example
$X=(1,\ldots,1)^{T},D_{\ell}=0,B_{\ell}=\mathds{1}_{\tilde{m}_{\ell-1}}$. In
the general case, the constant $C_{2}$ in Assumption 4 will be a depending on
the matrix $B_{\ell}$ and therefore potentially on the dimension. We propose a
way to address this issue in Section A.12.

###  A.10 Proof of Theorem 5.4

######  Lemma A.19 (Version of Assumption 3 for approximately low dimensional
distribution).

Denote $\mu_{\ell}$ the law of a multivariate normal distribution of
covariance matrix
$\operatorname{Diag}(\lambda_{1}^{\ell}I_{p_{\ell-1}},\ldots,\lambda^{\ell}_{\tilde{m}_{\ell-1}}I_{p_{\ell-1}})$
where $p_{\ell-1}=\frac{m_{\ell-1}}{\tilde{m}_{\ell-1}}$ and
$\frac{1}{k_{\ell-1}}\geq\lambda_{1}^{\ell}\ldots\geq\lambda^{\ell}_{\tilde{m}_{\ell-1}}$.
Let $k_{\ell-1}\geq 5$ and denote
$\eta:=\frac{\sqrt{\sum_{i=k_{\ell-1}+1}^{m}\lambda_{i}^{\ell}}}{4\sqrt{\sum_{i=1}^{k_{\ell-1}}\lambda_{i}^{\ell}}}$.

There exists two universal constants $D^{\prime}_{3},E^{\prime}_{3}$ such that
$\forall\tilde{m}_{\ell}\geq 1$ such that
$E_{3}^{{}^{\prime}k_{\ell-1}}\leq\tilde{m}_{\ell}\leq\eta^{-k_{\ell-1}}$,
there exists a random empirical measure $\hat{\mu}_{\tilde{m}_{\ell}}$ with
only $\tilde{m}_{\ell}$ points such that $\forall m_{\ell}\geq 1$ such that
$\tilde{m}_{\ell}\leq m_{\ell}\leq\eta^{-k_{\ell-1}}$ we have:

$\mathbb{E}[{\mathcal{W}}_{2}^{2}(\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{A},\hat{\mu}^{{\mathcal{I}}^{\ell-1}}_{\tilde{m}_{\ell}})]\leq\frac{D^{\prime}_{3}}{k_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{\nicefrac{{2}}{{k_{\ell-1}}}}$

###### Proof.

We do exactly the same as for the proof of Lemma A.16, i.e. a triangular
inequality but now we use rate of convergence of empirical measures in
Wasserstein distance with approximately low dimensional support as expressed
in Lemma A.15. ∎

######  Lemma A.20 (Version of Assumption 4 for approximately low dimensional
distribution).

Denote $\mu_{\ell}$ the law of a multivariate normal distribution of
covariance matrix
$\operatorname{Diag}(\lambda^{\ell}_{1}I_{p_{\ell-1}},\ldots,\lambda^{\ell}_{\tilde{m}_{\ell-1}}I_{p_{\ell-1}})$
where $p_{\ell-1}=\frac{m_{\ell-1}}{\tilde{m}_{\ell-1}}$ and
$\frac{1}{k_{\ell-1}}\geq\lambda_{1}^{\ell}\ldots\geq\lambda^{\ell}_{\tilde{m}_{\ell-1}}$.
$\forall X\in{\mathbb{R}}^{m_{\ell-1}}$ we have:

$\begin{split}&\mathbb{E}[\|W^{\ell}_{A}X\|_{2}^{2}]\leq\frac{\tilde{m}_{\ell-1}}{k_{\ell-1}}\|X\|_{2}^{2}\\\
&\mathbb{E}[\|W^{\ell}_{\tilde{m}_{\ell}}X\|_{2}^{2}]\leq\frac{\tilde{m}_{\ell-1}}{k_{\ell-1}}\|X\|_{2}^{2}\end{split}$

###### Proof.

Just notice that $\lambda_{1}\leq\frac{1}{k_{\ell-1}}$ and repeat the same
steps as in the proof of Lemma A.17. ∎

We will now re-state and prove Theorem 5.4:

######  Theorem A.21.

Under Assumptions 8 and 7, given $\varepsilon>0$, if $em_{0}\geq 5$ there
exists minimal widths $\tilde{m}_{1},\ldots,\tilde{m}_{L}$ such that if
$\eta^{-k_{0}}\geq m_{1}\geq\tilde{m}_{1},\ldots,\eta^{-k_{L-1}}\geq
m_{L}\geq\tilde{m}_{L}$, Property 1 is verified at the last hidden layer $L$
for $\underaccent{\bar}{E}_{L}=1,E_{L}=\varepsilon^{2}$. Moreover,
$\forall\ell\in[L],\,\exists T^{\prime}_{\ell}$ which does only depend on
$L,e,\ell,$ such that one can define recursively $\tilde{m}_{\ell}$ as
$\tilde{m}_{\ell}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime}_{\ell}}{\varepsilon}\right)^{k_{\ell-1}}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime}_{\ell}}{\varepsilon}\right)^{e\tilde{m}_{\ell-1}}$
where $\tilde{m}_{0}=m_{0}$. Moreover $\forall t\in[0,1]$, with
$Q$-probability at least $1-\delta_{Q}$, there exists permutations of hidden
layers $1,\ldots,L$ of network $B$ s.t.,
$\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{M_{t}}(x),y\right)\right]\leq
t\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{A}(x),y\right)\right]+(1-t)\mathbb{E}_{P}\left[{\mathcal{L}}\left(\hat{f}_{B}(x),y\right)\right]+\tfrac{4\sqrt{m_{L+1}}}{\sqrt{e}\delta_{Q}^{2}}\varepsilon$

###### Proof.

We just need to prove the first part of the theorem as proving the similarity
of loss is exactly the same as in the proof of Theorem 5.2 when we have proved
that Property 1 holds at layer $L$ with $E_{L}=\varepsilon^{2}$. The only
change comes from the constant $C_{2}$ in Lemma A.20 which is not $1$ anymore
but $\frac{1}{e}$, hence the additional factor $e$.

To prove Property 1 at layer $L$ with $E_{L}=\varepsilon^{2}$ we just combine
$L$ different times the two previous lemma Lemmas A.19 and A.20 and Lemma 4.1.

Lemma A.20 brings that at layer $\ell$ the constant $C_{2}$ is
$\frac{\tilde{m}_{\ell-1}}{k_{\ell-1}}\leq\frac{1}{e}$ by Assumption 8

Moreover, Lemma A.19 brings that at each layer $\ell$,
$C_{1}=\frac{D^{\prime}_{3}}{k_{\ell-1}}\log(\tilde{m}_{\ell})\left(\frac{1}{\tilde{m}_{\ell}}\right)^{2/k_{\ell-1}}$

It brings:

$\begin{cases}\underaccent{\bar}{E}_{i+1}=\frac{1}{e}\underaccent{\bar}{E}_{i}\\\
E_{i+1}=\frac{2}{e}E_{i}+2\frac{D^{\prime}_{3}}{k_{i}}\log(\tilde{m}_{i+1})\left(\frac{1}{\tilde{m}_{i+1}}\right)^{2/k_{i}}\tilde{m}_{i}\underaccent{\bar}{E}_{i}\end{cases}$

From there we see that if we have chosen at each layer

$D^{\prime}_{3}\log(\tilde{m}_{i+1})\left(\frac{1}{\tilde{m}_{i+1}}\right)^{2/k_{i}}=\frac{\sqrt{\frac{L2^{L-i+1}}{e^{L}}}}{\varepsilon^{2}}$

if

$\tilde{m}_{i}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime}_{i}}{\varepsilon}\right)^{e\tilde{m}_{i-1}}$

where $\tilde{{\mathcal{O}}}(\cdot)$ hides logarithmic terms and we define
$T^{\prime}_{i}=\sqrt{\frac{D^{\prime}_{3}e^{L}}{L2^{L-i+1}}}$

∎

###  A.11 Proof of Theorem 5.3

See 5.3

###### Proof.

First since $\Sigma$ is full rank we see that when writing $\Sigma=ODO^{T}$
where $OO^{T}=I_{\tilde{n}}$, the problem is equivalent to consider the
matrices $W_{A}O,W_{B}O$ and $\Sigma=D$. In that case, the raws of
$W_{A},W_{B}$ are still i.i.d. and follow the same law as $\mu$ modulo a non-
degenerated dilatation. We can then still assume
$\frac{d\mu}{d{\text{Leb}}}\leq F_{1}$ for a certain constant $F_{1}$.

Let $\tau=\frac{1}{2}$. $\forall S\subset{\mathbb{R}}^{n}$ a Borel set such
that $\mu(S)\geq 1-\tau$, we know that
${\text{Leb}}(S)\geq\frac{1-\tau}{F_{1}}$. In that case, applying Lemma A.10
we get
${\mathcal{N}}_{\varepsilon}(S)\geq\left(\frac{1}{\varepsilon}\right)^{n}\frac{\frac{1-\tau}{F_{1}}}{{\text{Leb}}({\mathcal{B}}_{2}(0,1)}$.
Denote
$F_{2}=\left(\frac{\frac{1-\tau}{F_{1}}}{{\text{Leb}}({\mathcal{B}}_{2}(0,\varepsilon)}\right)^{\nicefrac{{1}}{{n}}}$.
Using notations of Weed and Bach [2019] we get
${\mathcal{N}}_{\varepsilon}(\mu,\tau)\geq\left(\frac{F_{2}}{\varepsilon}\right)^{n}$.

Applying Proposition $6$ from Weed and Bach [2019] we get that

${\mathcal{W}}_{2}^{2}(\hat{\mu}_{A},\mu)\geq
F_{3}\left(\frac{1}{m}\right)^{\nicefrac{{2}}{{n}}}$

Finally noticing that $\mathbb{E}_{W_{B}}[\hat{\mu}_{B}]=\mu$ and applying
Lemma A.2, we get that:

$\mathbb{E}_{W_{A},W_{B}}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{A},\hat{\mu}_{B})]\geq\mathbb{E}_{W_{A}}[{\mathcal{W}}_{2}^{2}(\hat{\mu}_{A},\mathbb{E}_{W_{B}}[\hat{\mu}_{B}])]\geq
F_{3}\left(\frac{1}{m}\right)^{\nicefrac{{2}}{{n}}}$

Finally, remark that sinc
$\Sigma=D=\operatorname{Diag}(\lambda_{1},\ldots,\lambda_{\tilde{n}})$ and
noting $\lambda(\Sigma)=\min\\{d_{i},1\leq i\leq n\\}>0$ the smallest
eigenvalue of $\Sigma$,

$\begin{split}\mathbb{E}_{W_{A},W_{B}}[\min_{\Pi\in{\mathcal{S}}_{m}}\mathbb{E}_{x\sim
P}\|(W_{A}-\Pi
W_{B})x\|_{2}^{2}]&\leq\mathbb{E}_{W_{A},W_{B}}[\min_{\Pi\in{\mathcal{S}}_{m}}\mathbb{E}_{x\sim
P}\operatorname{tr}((W_{A}-\Pi W_{B})^{T}(W_{A}-\Pi W_{B})D)]\\\
&\geq\lambda(\Sigma)\mathbb{E}_{W_{A},W_{B}}[\min_{\Pi\in{\mathcal{S}}_{m}}\mathbb{E}_{x\sim
P}\|(W_{A}-\Pi W_{B})\|_{2}^{2}]\\\
&\geq\lambda(\Sigma)F_{3}\left(\frac{1}{m}\right)^{\nicefrac{{2}}{{n}}}\end{split}$

∎

###  A.12 Discussion about a model with no growth in the width needed

For proving Theorem 5.4, we have used constants $C_{1},C_{2}$ in Lemma 4.1
given by Lemmas A.19 and A.20. However we would like to highlight that Lemma
A.20 is very sub-optimal, though it can not really be improved in the general
case. Indeed, $\frac{1}{e^{\ell}}$ grows to infinity while
$\mathbb{E}[\|\phi^{\ell}(x)\|_{2}^{2}]$ is supposed to remain bounded for
$1\leq\ell$. Therefore from now on suppose that we have the following version
of Lemma A.22:

######  Lemma A.22 (Extending Assumption 4 for approximately low dimensional
distribution).

Denote $\mu_{\ell}$ the law of a multivariate normal distribution of
covariance matrix
$\operatorname{Diag}(\lambda_{1}^{\ell}I_{p_{\ell-1}},\ldots,\lambda^{\ell}_{\tilde{m}_{\ell-1}}I_{p_{\ell-1}})$
where $p_{\ell-1}=\frac{m_{\ell-1}}{\tilde{m}_{\ell-1}}$ and
$\frac{1}{k_{\ell-1}}\geq\lambda_{1}^{1}\ldots\geq\lambda_{\tilde{m}_{\ell-1}}^{\ell}$.
Suppose we have:

$\begin{split}\mathbb{E}[\|W_{A}(\phi^{\ell-1}_{B}(x)-\phi_{A}^{\ell-1}(x)\|_{2}^{2}]\leq\frac{m_{\ell}}{m_{\ell-1}}\mathbb{E}[\|\phi^{\ell-1}_{B}(x)-\phi_{A}^{\ell-1}(x)]\|_{2}^{2}\end{split}$

Notice that it is equivalent to making an assumption on the distribution of
$\phi^{\ell-1}_{B}(x)-\phi_{A}^{\ell-1}(x)$ which must not put too much mass
on the worst case coordinates $(\lambda^{\ell}_{i}$ for $i$ small).

In that case, adapting the proof as in Theorem 5.4, we could get an inequality
on $\tilde{m}_{i}$ of the form

$\tilde{m}_{i}=\tilde{{\mathcal{O}}}\left(\frac{T^{\prime\prime}_{i}}{\varepsilon}\right)^{k_{i-1}}$

where $T^{\prime\prime}_{i}$ is independent of $e$ which lead to controlled
bounds as $L\rightarrow\infty$ (without the exponent $e\tilde{m}_{i}$ as in
Theorem 5.4).

###  A.13 Extension of Theorems 5.2 and 5.4 to sub-Gaussian variables

Still under the setting of Assumption 6, suppose now that at a given layer
$\ell$, all the parameters of $W_{A}^{\ell}$ are still drawn independently but
no longer from ${\mathcal{N}}(0,\frac{1}{m_{\ell-1}})$ Instead we assume that
the underlying distribution $\mu_{\ell}$ verifies for each layer
$\ell\in[L+1]$: if $X\sim\mu_{\ell}$ then, $\forall j\neq
k\in[m_{l-1}],X_{j}\amalg X_{k}$. Moreover $\forall
i\in[\tilde{m}_{\ell-1}],\forall j,k\in I^{\ell-1}_{i}$,

$\mathbb{E}[X_{j}^{2}]=\mathbb{E}[X_{k}^{2}]=\lambda_{i}^{\ell}$

Finally suppose the variables are sub-Gaussian i.e., $\exists K>0,\forall
i\in[\tilde{m}_{\ell-1}],\forall j\in I^{\ell-1}_{i},\,\forall c>0$,

${\mathbb{P}}(|X_{j}|\geq c)\leq 2\exp(-\frac{c^{2}}{K\lambda_{i}^{\ell}})$

Further suppose that we are in the setting of Theorem 5.2 (The case of Theorem
5.4 is treated similarly):
$\tfrac{1}{m_{\ell-1}}\geq\lambda_{1}^{\ell}\geq\ldots\lambda_{\tilde{m}_{\ell-1}^{\ell}}$.

It is clear that Lemma A.17 is still valid for a constant $C_{2}=1$, the proof
being exactly the same.

We therefore just need to prove Lemma A.16 for $C_{1}$ to be determined. To
prove Lemma A.16, one just needs an equivalent of Lemma A.14 for sub-Gaussian
variables. To prove Lemma A.14, recall that we have used the fact that a
normal distribution doesn’t put to much mass outside of a ball of radius $c$
when c grows logarithmically. More precisely we have used the property, that
if $X\sim{\mathcal{N}}(0,\frac{I_{m_{l-1}}}{m_{l-1}})$, then:

$\forall c>1,{\mathbb{P}}(\|X\|_{2}^{2}\geq c^{2})\leq e^{-\frac{c^{2}}{4}}$

In our case, for a sub-Gaussian distribution, we know the existence of a
constant $K$ such that $\forall c>0$:

${\mathbb{P}}(|X_{j}|\geq c)\leq 2\exp(-\frac{c^{2}}{K\lambda_{i}^{\ell}})$

Therefore, by plugging this into the proof, and scaling parameter $c$ by
$\sqrt{K}$, we get exactly the same version of Lemma A.16 for sub-Gaussian
variable, with different constants scaled by a factor $\sqrt{K}$.

Propagating Property 1 with the recurrence formula of Lemma 4.1 we get LMC for
networks with sub-Gaussian distributions in the same form as for normal
variables.

##### Remark

Finally notice that results for sub-Gaussian variables can be extended in the
same way to variables whose tail decreses sufficiently fast (exponentially,
polynomially, etc…). The assymptotics of the tail will affect the convergence
rate in Wasserstein of the corresponding empirical measure.

###  A.14 Link with dropout stability

We relate now our previous study to a line of work exploring mode connectivity
through dropout stability.

Kuditipudi et al. [2019] define $\varepsilon$-dropout stable networks, as
networks $\hat{f}(\cdot;\theta)$ as defined in Equation 1 for which there
exists in each layer $\ell\in[L]$, a subset of at most $\frac{m_{\ell}}{2}$ of
neurons (i.e., rows of the weight matrix $W^{\ell}$) such that after
renormalizing each layer, the expected loss of the new network increases by no
more than $\varepsilon$ with respect to the original loss. Kuditipudi et al.
[2019] shows that two $\varepsilon$-dropout stable networks are mode connected
(with error barrier height $\varepsilon$) and Shevchenko and Mondelli [2020]
uses this result to show that two wide enough two-layer neural networks
trained with SGD are mode connected (where the continuous path may be non-
linear).

Recall that we have shown in Section 3 the stronger statements that two such
networks are in the same local minima modulo permutation symmetries. However,
note that Shevchenko and Mondelli [2020] don’t allow permutations of neurons).
We discuss here how to embrace in the same view our framework with dropout
stability results, showing how networks with independent neuron’s weights
become dropout stable in the same asymptotics of large width than the
condition of Lemma 5.1.

Consider the simplified setting of a $1$-hidden layer neural network with
$1$-Lipschitz activation where the weights of the second layer are fixed to
$\frac{1}{N}$: $\hat{f}(x;\theta)=\frac{1}{N}\sum_{i=1}^{N}\sigma(w_{i}x)$
where $w_{i}=W_{i:}\in{\mathbb{R}}^{d}$ is the $i-th$ row of the weight matrix
$W$. Suppose that $w_{i}$ are sampled independently from a sub-Gaussian
distribution and the data follows a distribution $(x,y)\sim P$ with
${\text{Supp}}(P)\subset{\mathcal{B}}_{2}(0,1)$. Denote
${\mathcal{A}}=\left[\frac{N}{2}\right]$. Dropout stability can be quantified
by controlling the error between the correctly renormalized network with
weights in ${\mathcal{A}}$ and the original one,

$\mathbb{E}\left[\left|\frac{2}{N}\sum_{i\in{\mathcal{A}}}\sigma(w_{i}x)-\frac{1}{N}\sum_{i=1}^{N}\sigma(w_{i}x)\right|\right]\leq{\mathcal{W}}_{1}(W^{\mathcal{A}},W^{{\mathcal{A}}^{c}})$
(8)

where we have denoted $W^{\mathcal{A}}$ (respectively $W^{{\mathcal{A}}^{c}}$)
the matrix $W$ where we have kept only the rows in ${\mathcal{A}}$
(respectively (${\mathcal{A}}^{c}$)). The right hand term can be connected to
convergence rates of empirical measure (Lemma A.16 and the extension to sub-
Gaussian distribution discussed in Section A.13):

$\begin{split}{\mathcal{W}}_{1}(W^{\mathcal{A}},W^{{\mathcal{A}}^{c}})\approx\left(\frac{1}{N}\right)^{1/d}\end{split}$

In a nutshell, showing that previous Equation 8 is tight would provide a
formal connection between dropout stability and our results. It is an
interesting direction for future work and note that it has strong connections
with the dual expression of the Wasserstein $1$ distance.

In that case, the bound on the dropout error evolves as
$\left(\frac{1}{N}\right)^{1/d}$, as for the linear mode connectivity error.
Hence networks become dropout stable in the same asymptotics as to exhibit
linear mode connectivity.

This is consistent with the idea that LMC requires the information to be
distributed evenly among neurons without any neuron responsible for the
particular behavior of one layer. This is similar to the intuitive requirement
for dropout stability.

##  Appendix B Proof of LMC for two-layer neural networks in the mean field
regime

###  B.1 Description of the mean field regime

When training a two-layer neural network with fixed input and output
dimensions but with a very wide hidden layer using SGD, the parameters of each
neuron can be seen as particles evolving independently one from each other:
the dynamic of each neuron’s weights depends only on the average distribution
of the weights and itself.

The main object of study is therefore the empirical distribution of the
neurons weights in the intermediate layer after $k$ Stochastic Gradient
Descent (SGD) steps. We denote it
$\rho_{k}^{(N)}=\frac{1}{N}\sum_{i=1}^{N}\delta_{\theta^{k}_{i}}$ where
$\theta^{k}_{i}=(w^{k}_{i},a^{k}_{i})\in{\mathbb{R}}^{d+1}$.

Multiple works [Chizat and Bach, 2018, Mei et al., 2018, 2019] show that if
the hidden layer’s width $N$ is big, the learning rate $s_{k}$ is small and
setting $T=\sum_{i=1}^{k}s_{i}$ a time re-normalization after $k$ steps, then
$\rho_{k}^{(N)}$ can be well approximated by $\rho_{t}$ which follows the
following partial differential equation (PDE):

$\displaystyle\partial_{t}\rho_{t}=2\xi(t)\nabla_{\theta}\cdot(\rho_{t}\nabla_{\theta}\Psi(\theta;\rho_{t})),\text{
}$ $\displaystyle\Psi(\theta;\rho_{t})=V(\theta)+\int
U(\theta,\tilde{\theta})\rho_{t}(d\tilde{\theta})$ $\displaystyle
V(\theta)=-{\mathbb{E}}[y\sigma_{*}(x;\theta)],\text{ }$ $\displaystyle
U(\theta_{1},\theta_{2})={\mathbb{E}}[\sigma_{*}(x;\theta_{1})\sigma_{*}(x;\theta_{2})]$

Here $\xi(t)$ represent a scaling of the learning rate where
$s_{k}=\varepsilon\xi(k\varepsilon)$. $U(\theta_{1},\theta_{2})$ represents a
correlation between neurons. $V(\theta)$ is an energy quantifying the
alignment of a neuron function with the data. In the following, as in Mei et
al. [2019] we will work with $\xi=\frac{1}{2}$ a constant step size function.
As in Mei et al. [2019], we highlight that the proof remains valid under
Assumption 1.

When considering noisy SGD, the limit PDE becomes:

$\displaystyle\partial_{t}\rho_{t}=2\xi(t)\nabla_{\theta}\cdot(\rho_{t}(\theta)\nabla_{\theta}\Psi_{\lambda}(\theta;\rho_{t}))+2\xi(t)\tau
d^{-1}\Delta_{\theta}\rho_{t}$
$\displaystyle\Psi_{\lambda}(\theta;\rho)=\Psi(\theta;\rho)+\frac{\lambda}{2}\|\theta\|_{2}^{2}$

The crucial point about the mean field view is to show that the empirical
distribution of parameters is well enough approximated by $\rho_{t}$. Then the
study of the neural network can be reduced to the study of the partial
differential equation. For example global convergence of the test loss results
can be deduced as in Chizat and Bach [2018]. This view is also convenient to
get insights of typical behaviours of the dynamics while smoothing the effects
of local minima [Mei et al., 2019]. In our case, the mean field view allows us
to use convergence results in Wasserstein distance of empirical measures
towards the underlying distribution. We can then show Linear Mode Connectivity
for two-layer neural networks independently trained in the mean field regime.

###  B.2 Proving LMC for noiseless regularization-free SGD

To prove our results we have to show that the empirical distribution of
weights can be well approximated by the solution of the mean field PDE. To
achieve this, Mei et al. [2019] introduce four intermediate dynamics that stay
close one of each other.

First note that our Assumption 1 implies Assumptions $1$ to $4$ of Mei et al.
[2019]. Especially the non-linearity being Lipschitz implies its gradient
distribution on the data is bounded and hence sub-Gaussian.

####  B.2.1 Intermediate dynamics

Mei et al. [2019] introduce 4 different intermediate dynamics between the
empirical distribution of weights optimized by SGD and the solution of the PDE
that we recall here:

Nonlinear dynamics

Let consider $\bar{\theta}_{i}^{t}$ with initialization
$\bar{\theta}_{i}^{0}\sim\rho_{0}$ i.i.d. and which follows the dynamics

$\bar{\theta}_{i}^{t}=\bar{\theta}_{i}^{0}+2\int_{0}^{t}\xi(s)G(\bar{\theta}_{i}^{s};\rho_{s})ds$

or equivalently

$\frac{d}{dt}\bar{\theta_{i}^{t}}=-2\xi(t)\left[\nabla
V(\bar{\theta}_{i}^{t})+\int\nabla_{1}U(\bar{\theta}_{i}^{t},\theta)\rho_{t}(d\theta)\right]$

where $G(\theta;\rho)=-\nabla\Psi(\theta;\rho)$. An important fact is that
$\bar{\theta}_{i}^{t}$ is random because of the random initialization.
Moreover its law at time $t$ is $\rho_{t}$. It corresponds to the evolution of
particles under a velocity field $-2\xi(t)\left[\nabla
V(\bar{\theta}_{i}^{t})+\int\nabla_{1}U(\bar{\theta}_{i}^{t},\theta)\rho_{t}(d\theta)\right]$
which depends only on the position of the optimized particle and the overall
distribution of all particles.

Particle Dynamics

Let $\underaccent{\bar}{\theta}_{i}^{t}$ have the same initialization as the
nonlinear dynamics $\underaccent{\bar}{\theta}_{i}^{0}=\bar{\theta}_{i}^{0}$,
and
$\underaccent{\bar}{\rho}_{t}^{(N)}=\frac{1}{N}\sum_{i=1}^{N}\delta_{\underaccent{\bar}{\theta}_{i}^{t}}$
denote the empirical distribution of $\underaccent{\bar}{\theta}_{i}^{t}$.
Then the particle dynamics is given by:

$\underaccent{\bar}{\theta}_{i}^{t}=\underaccent{\bar}{\theta}_{i}^{0}+2\int_{0}^{t}\xi(s)G(\underaccent{\bar}{\theta}_{i}^{s};\underaccent{\bar}{\rho}_{s}^{(N)})ds$

or equivalently

$\frac{d}{dt}\underaccent{\bar}{\theta}_{i}^{t}=-2\xi(t)\left[\nabla
V(\underaccent{\bar}{\theta}_{i}^{t})+\frac{1}{N}\sum_{j=1}^{N}\nabla_{1}U(\underaccent{\bar}{\theta}_{i}^{t},\underaccent{\bar}{\theta}_{j}^{t})\right]$

Gradient descent dynamics

Let $\tilde{\theta}_{i}^{k}$ with initialization
$\tilde{\theta}_{i}^{0}=\bar{\theta}_{i}^{0}$ following the dynamics:

$\tilde{\theta}_{i}^{k}=\tilde{\theta}_{i}^{0}+2\varepsilon\sum_{l=0}^{k-1}\xi(l\varepsilon)G(\tilde{\theta}_{i}^{l};\tilde{\rho}_{l}^{(N)})$

or equivalently:

$\tilde{\theta}_{i}^{k+1}=\tilde{\theta}_{i}^{k}-2s_{k}\left[\nabla
V(\tilde{\theta}_{i}^{k})+\frac{1}{N}\sum_{j=1}^{N}\nabla_{1}U(\tilde{\theta}_{i}^{k},\tilde{\theta}_{j}^{k})\right]$

Stochastic Gradient Descent Dynamics

Consider $\theta_{i}^{k}$ with initialization
$\theta_{i}^{0}=\bar{\theta}_{i}^{0}$ following the dynamics:

$\theta_{i}^{k}=\theta_{i}^{0}+2\varepsilon\sum_{l=0}^{k-1}\xi(l\varepsilon)F_{i}(\theta^{l};z_{l+1})$

or equivalently

$\theta_{i}^{k+1}=\theta_{i}^{k}-2s_{k}F_{i}(\theta^{k};z_{k+1})$

where
$F_{i}(\theta^{k};z_{k+1})=(y_{k+1}-\hat{y}_{k+1})\nabla_{\theta}\sigma_{*}(x_{k+1};\theta_{i}^{k})$,
$z_{k}=(x_{k},y_{k})$ and
$\hat{y}_{k+1}=\frac{1}{N}\sum_{j=1}^{N}\sigma_{*}(x_{k+1};\theta_{j}^{k})$.

One can use Proposition 26,28,29 from Mei et al. [2019] to show the following
lemma:

######  Lemma B.1.

Consider a two-layer neural network trained with noiseless regularization-free
SGD for an underlying time $T$. Then under Assumption 1, there exists
constants $K$ and $K_{0}$ such that, if
$\varepsilon\leq\min\left\\{\frac{1}{K_{0}e^{K_{0}(1+T)^{3}}},\frac{1}{K_{0}(d+\log(N)+z^{2})e^{K_{0}(1+T)^{3}}}\right\\}$,
then with probability at least $1-3e^{-z^{2}}$ we have:

$\begin{split}\max_{k\in[0,T/\varepsilon]\bigcap{\mathbb{N}}}\max_{i\in[N]}\|\theta_{i}^{k}-\bar{\theta}_{i}^{k\varepsilon}\|_{2}&\leq
Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]\\\
&+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}[\sqrt{d+\log(N)}+z]\end{split}$

###### Proof.

This is direct application from Proposition 26,28,29 from Mei et al. [2019] by
doing two union bounds and two triangular inequalities. ∎

We moreover recall that as highlighted in Mei et al. [2019],
$\\{\bar{\theta}_{i}^{t},1\leq i\leq N\\}$ are independent from each other and
each follows the distribution $\rho_{t}$ when initialized i.i.d. as
$\rho_{0}$. Therefore, when considering two two-layer neural networks
initialized randomly as $\rho_{0}$ and trained for the same underlying time
$T$ with noiseless regularization-free SGD, we know from the previous lemma
that the parameters of both networks are close to two samples from $\rho_{t}$.

####  B.2.2 Proof of Theorem 3.1 in the case of noiseless regularization-free
SGD

We now prove Theorem 3.1 in case of noiseless regularization-free SGD.

See 3.1

###### Proof.

We know from Lemma B.1 that with probability at least $1-3e^{-z^{2}}$, if
$\varepsilon\leq\min\left\\{\frac{1}{K_{0}e^{K_{0}(1+T)^{3}}},\frac{1}{K_{0}(d+\log(N)+z^{2})e^{K_{0}(1+T)^{3}}}\right\\}$,

$\begin{split}\max_{k\in[0,T/\varepsilon]\bigcap{\mathbb{N}}}\max_{i\in[N]}\|\theta_{A,i}^{k}-\bar{\theta}_{A,i}^{k\varepsilon}\|_{2}&\leq
Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]\\\
&+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}[\sqrt{d+\log(N)}+z]\end{split}$

which means that $\theta_{A,i}$ is close to the non linear dynamics which are
samples from $\rho_{t}$. By a union bound, with probability $1-6e^{-z^{2}}$
this is true for both networks $A$ and $B$.

Denoting as before $\theta_{A,i}=(w_{A,i},a_{A,i})\in{\mathbb{R}}^{d+1}$
(respectively $\theta_{B,i}$), $A_{A}=(a_{A,1},\ldots,a_{A,N})$ (respectively
$A_{B}$) and $W_{A}\in{\mathcal{M}}_{N,d}({\mathbb{R}})$ the concatenation of
vectors $w_{A,i}\in{\mathbb{R}}^{d}$ (respectively $W_{B}$). Given
$t\in[0,1]$, we aim at finding a permutation $\Pi\in{\mathcal{S}}_{N}$ of the
second network’s hidden layer to get
$\tilde{\theta}_{B}=(\tilde{A}_{B},\tilde{W}_{B})=(A_{B}\Pi^{T},\Pi W_{B})$
bounding

$\displaystyle|t\hat{f}(x;\theta_{A})+(1-t)\hat{f}(x;\theta_{B})-\hat{f}(x;t\theta_{A}+(1-t)\tilde{\theta}_{B})|$
$\displaystyle=\frac{1}{N}|tA_{A}\sigma(W_{A}X)+(1-t)\tilde{A}_{B}\sigma(\tilde{W}_{B}X)-(tA_{A}+(1-t)\tilde{A}_{B})\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)|$
(9) $\displaystyle\leq
t\left|\frac{A_{A}}{N}(\sigma(W_{A}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X))\right|+(1-t)\left|\frac{\tilde{A}_{B}}{N}(\sigma(\tilde{W}_{B}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X))\right|$
$\displaystyle\leq
t\|A_{A}\|_{\infty}\frac{\|\sigma(W_{A}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}+(1-t)\|\tilde{A}_{B}\|_{\infty}\frac{\|\sigma(\tilde{W}_{B}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}$

Both terms
$\frac{\|\sigma(W_{A}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}$ and
$\frac{\|\sigma(\tilde{W}_{B}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}$
can be bounded.

Indeed, first using lemma 22 from Mei et al. [2019] and that from Assumption 1
${\text{Supp}}(\rho_{0})$ is bounded, we get that

${\text{Supp}}(\rho_{t})\subset{\mathcal{B}}_{2}(0,K((1+T^{2})T)+1)$

is bounded with a diameter depending only on the initialization
${\text{Supp}}(\rho_{0})$ and underlying time $T$.

Therefore we can apply Theorem 1 from Weed and Bach [2019] and get for $s>d$
the existence of a constant $C$ such that with probability at least
$1-\frac{\delta}{2}$, there exists a permutation $\gamma\in{\mathcal{S}}_{N}$
such that by considering $\|\cdot\|_{1}$ as a distance for the Wasserstein:

${\mathcal{W}}_{1}(\hat{\mu}_{A},\hat{\mu}_{B})=\frac{1}{N}\sum_{i=1}^{N}\|\bar{\theta}_{A,i}-\bar{\theta}_{B,\gamma_{i}}\|_{1}\leq\frac{C}{\delta}N^{-1/s}$

Note that, while $C$ is independent of $N$, it depends on the distribution
$\rho_{t}$ and therefore on $d$, ${\text{diam}}({\text{Supp}}(\rho_{t}))$
(i.e., $T$) and on the constants from Assumption 1.

Recall that we suppose the data distribution P bounded and denote
${\text{Supp}}(P)\subset[-H_{x},H_{x}]^{d}\times[-H_{y},H_{y}]$.

Therefore, we get that with probability at least
$1-\frac{\delta}{2}-6e^{-z^{2}}$:

$\displaystyle\forall
X\in[-H_{x},H_{x}]^{d},\frac{\|\sigma(W_{A}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}\leq
g_{2}(T,z,\delta,N,\varepsilon)$ (10)
$\displaystyle:=L_{\sigma}\Bigg{(}H_{x}\frac{C}{\delta}N^{-\frac{1}{s}}+2H_{x}\sqrt{d}\bigg{[}\frac{Ke^{K(1+T)^{3}}}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}[\sqrt{d+\log(N)}+z]\bigg{]}\Bigg{)}$

and same for the other term:

$\begin{split}\forall
X\in[-H_{x},H_{x}]^{d},\frac{\|\sigma(\tilde{W}_{B}X)-\sigma((tW_{A}+(1-t)\tilde{W}_{B})X)\|_{1}}{N}\leq
g_{2}(T,z,\delta,N,\varepsilon)\end{split}$

Using lemma 20 from Mei et al. [2019] we know that $\forall
i\in[N],\bar{a}_{i}^{T}\leq K(1+T)$ for a certain constant $K$. Therefore we
can bound $\|A_{A}\|_{\infty},\|A_{B}\|_{\infty}\leq
g_{1}(T,z,N,\varepsilon):=K(1+T)+Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{d+\log(N)}+z)$.

Taking $z=\sqrt{\log\left(\frac{12}{\delta}\right)}$, we have shown the
existence of a permutation $\gamma$ with probability at least $1-\delta$ such
that almost surely on the choice of $x\sim P$ and $\forall t\in[0,1]$, we
have:

$\begin{split}|t\hat{f}(x&;\theta_{A})+(1-t)\hat{f}(x;\theta_{B})-\hat{f}(x;t\theta_{A}+(1-t)\tilde{\theta}_{B})|\leq
g_{1}(T,z,N,\varepsilon)g_{2}(T,z,\delta,N,\varepsilon)\\\
&\leq\left(K(1+T)+Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{d+\log(N)}+z)\right)\\\
&\Big{(}L_{\sigma}(H_{x}\frac{C}{\delta}N^{-\frac{1}{s}}+2H_{x}\sqrt{d}(Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]\\\
&+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{d+\log(N)}+z)))\Big{)}\end{split}$

For fixed $T,\delta$ and $z=\sqrt{\log\left(\frac{12}{\delta}\right)}$, denote
${\text{err}}(N,\varepsilon)$ the right hand term.

It is clear that:

$\begin{split}\forall{\text{err}}>0\exists N_{\min}\forall N\geq
N_{\min}\exists\varepsilon_{\max}(N)\forall\varepsilon\leq\varepsilon_{\max}(N),\,{\text{err}}(N,\varepsilon)\leq{\text{err}}\end{split}$

This brings the first part of the theorem.

Discussion: Let’s look more closely at the term
$Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{d+\log(N)}+z))$
from Mei et al. [2019] which comes from the mean field approximation. When
taken alone, this term yields an error which is independent of the input
dimension $d$, since taking $N$ large leads to small error (provided
$\varepsilon$ is small). However, here the growth of the hidden layer $N$
depends on the input dimension $d$ through the exponent $s>d^{*}_{1}(\mu)$
(with $d^{*}_{1}(\mu)\leq d$) using notations from Weed and Bach [2019]. This
is due to Wasserstein convergence rates of empirical measures in dimension
$d$. Without any further assumption on the weight distribution or precise
study of the PDE we have to consider that ${\text{Supp}}(\mu))$ has dimension
$d$. To remove this dependence, one could study a precise model for the data
and look more closely at the PDE evolution to better understand the support of
the distribution $\rho_{t}$. ∎

To prove the second part of the theorem, first notice that as already
mentioned, $\|A_{A}\|_{\infty},\|A_{B}\|_{\infty}\leq
g_{1}(T,z,N,\varepsilon)$. Moreover, the data distribution is bounded and the
weights of the first layer of the approximating PDE live in a bounded set
thanks to Section B.2.2. It brings the existence of a constant $K(T)$ such
that:

$\forall t\leq T,\forall w\in{\text{Supp}}(\rho(t)),\forall
X\in{\text{Supp}}(P),|w.x|\leq K(T)$

Using we see that if
$\varepsilon\leq\min\\{\frac{1}{K_{0}e^{K_{0}(1+T)^{3}}},\frac{1}{K_{0}(d+\log(N)+z^{2})e^{K_{0}(1+T)^{3}}}\\}$,
with probability at least $1-\frac{\delta}{2}$, setting
$z=\sqrt{\tfrac{12}{\delta}}$

$\begin{split}\max_{k\in[0,T/\varepsilon]\bigcap{\mathbb{N}}}&\max_{i\in[N]}|w_{A,i}^{k}x|\leq
H_{x}\sqrt{d}(Ke^{K(1+T)^{3}}\\\
&\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{d+\log(N)}+z))+K(T)\end{split}$

Up to changing previous $N_{\min},\varepsilon_{\max}$, we can suppose that
$\forall N\geq
N_{\min},\forall\varepsilon\leq\varepsilon_{\max}(N),\,Ke^{K(1+T)^{3}}\frac{1}{\sqrt{N}}[\sqrt{\log(NT)}+z]+Ke^{K(1+T)^{2}T}\varepsilon+Ke^{K(1+T)^{2}T}\sqrt{\varepsilon}(\sqrt{D+\log(N)}+z)\leq
1$

Therefore, since the data distribution is bounded we know that there exists a
constant $K(T)$ such that if $N\geq
N_{\min},\varepsilon\leq\varepsilon_{\max}(N),\,P-\text{almost-surely}$:

$\begin{cases}|y|\leq K(T)\\\ |\hat{f}(x:\theta_{A})|\leq K(T)\\\
|\hat{f}(x:\theta_{B})|\leq K(T)\\\ \end{cases}$

We make a general assumption on the loss of the form: $\forall
y\in{\mathbb{R}}(x\rightarrow{\mathcal{L}}(x,y))$ is convex and $\forall
K>0,\exists C_{K},\forall
x_{1},x_{2},y\in[-K,K],|{\mathcal{L}}(x_{1},y)-{\mathcal{L}}(x_{2},y)|\leq
C_{K}|x_{1}-x_{2}|_{2}$. In particular this is true for the square loss.

Combining convexity of the loss, the first part of the theorem already proved
and Lipschitzness on a compact domain of the loss, we get the second part of
the theorem with a term $C_{K(T)}{\text{err}}$ instead of err. To solve this,
just consider $\min\left\\{\frac{{\text{err}}}{C_{T}},{\text{err}}\right\\}$
in the first part of the theorem.

We have supposed in the beginning that the non-linearity was bounded. But the
previous study shows that with probability at least $1-\delta$, $|w_{A,i}x|$
is upper bounded for all $i$ during the training up to time T by some constant
depending only on $T,\delta$ provided $N\geq
N_{\min},\varepsilon\leq\varepsilon_{\max}(N)$ Therefore assuming that the
non-linearity is bounded on a big enough compact set is enough to get the the
result since it doesn’t change the dynamics of the parameters considered.
However the size of this set is not made explicit here.

###  B.3 Proving LMC for general SGD

We will now study LMC of neural networks trained under general SGD using
Theorem 4 part B from Mei et al. [2019]. The study is very similar to the case
of noiseless SGD and will yield similar results.

More precisely in that case the PDE writes as:

$\displaystyle\partial_{t}\rho_{t}=2\xi(t)\nabla_{\theta}\cdot(\rho_{t}(\theta)\nabla_{\theta}\Psi_{\lambda}(\theta;\rho_{t}))+2\xi(t)\tau
d^{-1}\Delta_{\theta}\rho_{t}$
$\displaystyle\Psi_{\lambda}(\theta;\rho)=\Psi(\theta;\rho)+\frac{\lambda}{2}\|\theta\|_{2}^{2}$

Notice that our Assumptions 1 and 2 imply assumptions 1 to 6 in Mei et al.
[2019].

####  B.3.1 Intermediate dynamics for general SGD

Mei et al. [2019] define as before intermediate dynamics:

Non linear dynamics

Let consider $\bar{\theta}_{i}^{t}$ with initialization
$\bar{\theta}_{i}^{0}\sim\rho_{0}$ i.i.d. which follows the dynamics