# Bayesian Federated Neural Matching that Completes Full Information

Peng Xiao,1 Samuel Cheng2 is the corresponding author.

###### Abstract

Federated learning is a contemporary machine learning paradigm where locally
trained models are distilled into a global model. Due to the intrinsic
permutation invariance of neural networks, Probabilistic Federated Neural
Matching (PFNM) employs a Bayesian nonparametric framework in the generation
process of local neurons, and then creates a linear sum assignment formulation
in each alternative optimization iteration. But according to our theoretical
analysis, the optimization iteration in PFNM omits global information from
existing. In this study, we propose a novel approach that overcomes this flaw
by introducing a Kullback-Leibler divergence penalty at each iteration. The
effectiveness of our approach is demonstrated by experiments on both image
classification and semantic segmentation tasks.

## Introduction

The recent decade has seen a rapid advancement in artificial intelligence,
particularly in the subfield of deep learning. This is large because there is
an abundance of data available. However, the increasing privacy concerns
present a barrier to some datasets’ accessibility in a number of applications,
particularly in the medical and financial areas. Therefore, federated learning
(FL), which involves learning a model from disparate sets of data, is
suggested to address this issue.

Federated learning is a learning paradigm where locally learned models are
combined into a shared global model. Neural networks are naturally used in
federated learning since they play a vital role in deep learning. Considering
the inherent permutation invariance of a neural network, it is reasonable to
match neurons of distinct local models before aggregating them. To do this,
Probabilistic Federated Neural Matching (PFNM) (Yurochkin et al. 2019b) builds
a Bayesian nonparametric framework to match and merge these neurons. In the
process of modeling, PFNM formally characterizes the generative process of
local neurons through the Beta-Bernoulli process (BBP) (Thibaux and Jordan
2007a), and treats the local neurons as noisy realizations of latent global
neurons. In the process of optimizing, PFNM iteratively maximizes the
posterior estimation of latent global neurons, through solving a linear sum
assignment formulation of global neurons and current local neurons. PFNM
subsequently extends in modern architectures such as convolutional neural
networks (CNNs) and long short-term memory (LSTMs) (Wang et al. 2020), and its
variants are also utilized in aggregating various statistical models such as
Gaussian topic models, hierarchical Dirichlet process based hidden Markov
models (Yurochkin et al. 2019a, 2018) etc.

In this paper, we theoretically prove the drawback of the optimizing process
in PFNM under its probabilistic framework and address it by introducing a
Kullback-Leibler (KL) divergence penalty. Our theory can also be generalized
in all PFNM variants. Specifically, the contributions in this work include: 1)
We theoretically prove that the linear sum formulation in the optimizing
process omits the information of global neurons under the Bayesian framework.
2) We fix the missing global information by introducing a KL divergence
penalty to complete the full information. 3) In experiment, we not only
demonstrate the effectiveness of our approach on three image datasets, but
also extend the matching type federated learning algorithm to neural network
with batch normalization layer.

## Related Works

#### Federated Learning

The initial aggregation method in FL is FedAvg (McMahan et al. 2017) in which
parameters of local models are averaged coordinate-wisely. But the performance
of FedAvg deteriorates significantly in non-i.i.d (Independent and Identically
Distributed) data (Karimireddy et al. 2020; Deng, Kamani, and Mahdavi 2021;
Xiao et al. 2020). Subsequent improved methods starts from different
perspectives. FedProx (Li et al. 2018) adds a proximal term in local training
cost to keep dissimilarity between local models in lower bound. SCAFFOLD
(Karimireddy et al. 2020) uses control variates to correct the drift in local
updates. FedPD (Zhang et al. 2020) proposes a primal-dual optimization
strategy to alleviate deterioration in non-i.i.d data. Agnostic federated
learning (Mohri, Sivek, and Suresh 2019) optimizes a centralized distribution
by maximize the worst local model. Federated multi-task learning (Smith et al.
2017) applies a multi-task learning mechanism to customize local models.
Several studies further extend knowledge distillation (Hinton, Vinyals, and
Dean 2015; Buciluǎ, Caruana, and Niculescu-Mizil 2006; Schmidhuber 1992) in
federated learning. And the key idea is to employ the knowledge of pre-trained
teacher neural networks (local models) to learn a student neural network
(global model) (Lin et al. 2020; Chen and Chao 2020; Zhu, Hong, and Zhou
2021). However, all those methods above don’t considering the permutation
invariance in neural network.

#### Parameter Matching

There are also other researches about match the parameters. In (Singh and
Jaggi 2020), the authors align neurons across different NNs by minimizing an
optimal transportation cost matrix. However, it fixes the number of global
neurons, making it impractical when data from different local models is
extremely heterogeneous. Some work (Claici et al. 2020) uses variational
inference to optimize the assignments between global and local components
under a KL divergence. However, the optimization process is complex, and
calculating variational inference is difficult.

## Problem Formulation

Given $S$ fully connected (FC) Neural Networks (NNs) with a hidden layer
trained through different datasets:
$f_{s}(\bm{x})=\bm{W}_{s}^{(1)}\sigma(\bm{W}_{s}^{(0)}\bm{x}),\text{for
}s=1,\cdots,S$ (for simplifying notation, biases are omitted to ), where
$\sigma(\cdot)$ is the nonlinear activation function,
$\bm{W}_{s}^{(1)}\in\mathbb{R}^{K\times J_{s}}$ and
$\bm{W}_{s}^{(0)}\in\mathbb{R}^{J_{s}\times D}$ are the weights; with $D$
being the input dimension, $K$ being the output dimension (i.e., number of
classes), and $J_{s}$ being the number of neurons on the hidden layer of
$s$-th NN. Neuron indexed by $j$ is viewed as a concatenated vector
$\bm{w}_{sj}^{T}=[\bm{W}_{s,j\cdot}^{(0)},{\bm{W}_{s,\cdot j}^{(1)}}^{T}]$,
where $j\cdot$ and $\cdot j$ denote the $j$th row and column correspondingly.
And $s$-th FCNN can also be viewed as a collection of neurons
$\\{\bm{w}_{sj}\in\mathbb{R}^{D+K}\\}_{j=1}^{J_{s}}$. In federated learning,
we want to learn a global neural network with weights
$\Theta^{(0)}\in\mathbb{R}^{J\times D},\Theta^{(1)}\in\mathbb{R}^{K\times J}$,
where $J\ll\sum_{s=1}^{S}J_{s}$ is an inferred variable denoting the number of
global neurons.

#### Permutation Invariance

Expanding the preceding expression of a FCNN:
$f_{s}(\bm{x})=\sum_{j=1}^{J_{s}}\sigma(\langle\bm{W}_{s,j\cdot}^{(0)},\bm{x}^{T}\rangle)\bm{W}_{s,j\cdot}^{(1)}$.
Summation is a permutation invariant operation, thus any permutation
$\tau(1,\cdots,Js)$ of rows of $\bm{W}_{s}^{(0)}$ and columns of
$\bm{W}_{s}^{(1)}$, i.e. the neurons, will not affect the output for any input
$\bm{x}$. Due to the permutation invariance, a neuron indexed by $j$ from one
FCNN is unlikely to correspond to a neuron with the same index from another
FCNN. Thus, we should match neurons from different FCNNs before aggregate them
into a collection of global neurons
$\\{\bm{\theta}_{i}\in\mathbb{R}^{D+K}\\}_{i=1}^{J}$.

## Probabilistic Modeling

PFNM models the generative process of observed local neurons using a Beta
Bernoulli process, which is described in the Appendix A, to infer the global
model while accounting for the inherent permutation invariance. First,
consider the collection of global neurons as prior which are sampled from a
Beta process with a base measure:
$M=\sum_{i}m_{i}\delta_{\bm{\theta}_{i}}\sim\text{BP}(1,\gamma_{0}H)$ and
$\bm{\theta}_{i}\sim H$, where $\gamma_{0}$ is the mass parameter, $m_{i}$ are
the stick-breaking weights, $H$ is the base measure and chosen as a
multivariate Gaussian distribution
$H=\mathcal{N}(\bm{\mu}_{0},\bm{\Sigma_{0}})$ with
$\bm{\mu}_{0}\in\mathbb{R}^{D+K}$ and diagonal
$\bm{\Sigma_{0}}=\bm{I}\sigma_{0}^{2}$.

The Bernoulli process is then used by each local model to select a subset of
global
neurons:$\mathcal{T}_{s}=\sum_{i}a_{si}\delta_{\bm{\theta}_{i}}|M\sim\text{BeP}(M)\text{
for }s=1,\ldots,S$, where $a_{si}|m_{i}\sim\operatorname{Bern}(m_{i})$ is a
random measure representing a subset of global neurons contained in local
model $s$. Considering the inherent permutation invariance of local model, it
denotes $\bm{A}=\\{A_{ij}^{s}\\}_{s,i,j}$ as the assignment variables,
$\sum_{j}A^{s}_{ij}=a_{si}\in\\{0,1\\}$, $\sum_{i}A^{s}_{ij}=1$. Finally,
observed local neurons in model $s$ are treated as noisy measurements of
global neurons under permutation invariance:

$\begin{split}\bm{w}_{sj}\mid\bm{A}^{s},\bm{\theta}\sim\mathcal{N}(\sum_{i}A^{s}_{ij}\bm{\theta}_{i},\bm{\Sigma_{s}})\text{
for }j=1,\ldots,J_{s},\end{split}$ (1)

where $J_{s}=card(\mathcal{T}_{s})$, $\bm{\Sigma_{s}}=\bm{I}\sigma_{s}^{2}$ is
also diagonal and represents the noise. The noise is usually caused by
estimation error due to finite sample sizes or variations in the distribution
of each local dataset. $A_{ij}^{s}=1$ indicates that $\bm{w}_{sj}$ is matched
to $\bm{\theta}_{i}$, i.e. $\bm{w}_{sj}$ is the local neuron realization of
the global neuron $\bm{\theta}_{i}$; $A_{ij}^{s}=0$ indicates the inverse.

After the modeling, it is natural to infer the global neurons by maximizing
the posterior probability of the global neurons given local neurons
(likelihoods) under the permutation invariance:

$\begin{split}\max_{\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}}&P(\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}|\\{\bm{w}_{sj}\\})\propto\\\
&P(\\{\bm{w}_{sj}\\}|\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\})P(\\{\bm{A}^{s}\\})P(\\{\bm{\theta}_{i}\\}),\end{split}$
(2)

define $\bm{Z}_{i}=\\{(s,j)|A_{ij}^{s}=1\\}$ be the index set of local neurons
assigned to $i$-th global neuron, and taking negative logarithm it can obtain:

$\begin{split}\min_{\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}}&-\log(P(\\{\bm{A}^{s}\\}))\\\
&-\sum_{i}\bigg{(}\sum_{z\in\bm{Z}_{i}}\log(p(\bm{w}_{z}|\bm{\theta}_{i}))+\log(p(\bm{\theta}_{i}))\bigg{)}\end{split}$
(3)

where $P(\\{\bm{A}^{s}\\})$ is interpreted by Indian Buffet Process (IBP) and
demonstrated in Appendix B. Given $\\{\bm{A}^{s}\\}_{s=1}^{S}$, the closed
form of $\\{\bm{\theta}_{i}\\}$ can be estimated through the Gaussian-Gaussian
conjugacy:

$\bm{\theta}_{i}=\frac{\bm{\mu}_{0}/\sigma_{0}^{2}+\sum_{s,j}\bm{A}_{i,j}^{s}\bm{w}_{sj}/\sigma^{2}_{s}}{1/\sigma_{0}^{2}+\sum_{s,j}\bm{A}^{s}_{i,j}/\sigma^{2}_{s}}\,\text{
for }i=1,\cdots,J.$ (4)

Taking equation (4) into objective (3), it can cast optimization only with
respect to $\\{\bm{A}^{s}\\}_{s=1}^{S}$. Unfortunately, solving all local
assignments together leading an NP-hard combinatorial optimization problem.
Thus, PFNM applies an alternative optimization process to solve the problem.

#### Alternative Optimization

PFNM iteratively optimizes one local assignment variable $\bm{A}^{s^{\prime}}$
at a time by fixing all other assignment variables
$\\{A^{s}_{i,j}\\}_{i,j,s\in-s^{\prime}}$, where $-s^{\prime}$ denotes ”all
but $s^{\prime}$”. It aims to formulate a linear sum assignment problem with
current local assignment $\sum_{i,j}C^{s^{\prime}}_{i,j}A^{s^{\prime}}_{i,j}$
in each iteration, where $\\{C^{s^{\prime}}_{i,j}\\}_{i,j}$ denotes the cost
specification. The Hungarian algorithm is then used to solve the linear sum
assignment problem. In each iteration, it divides terms in objective (3) into
two parts: one is $i=1,\cdots,J_{-s^{\prime}}$, where
$J_{-s^{\prime}}=\max\\{i:\bm{A}^{s}_{ij}=1$, for
$s\in-s^{\prime},j=1,\dots,J_{s}\\}$, this part denotes active global neurons
estimated from $-s^{\prime}$; another is
$i=J_{-s^{\prime}}+1,\cdots,J_{-s^{\prime}}+J_{s^{\prime}}$, it denotes new
global neuron from current local model. Then it has the following proposition
describes the assignment cost $\\{C^{s^{\prime}}_{i,j}\\}_{i,j}$:

######  Proposition 1.

The assignment cost specification $C^{s^{\prime}}_{i,j}$ for finding
$\\{\bm{A}^{s^{\prime}}\\}$ is $C_{i,j}^{s^{\prime}}=$

$\begin{cases}\begin{split}2\log{\frac{S-n_{i}^{-s^{\prime}}}{n_{i}^{-s^{\prime}}}}&-\frac{\big{\lVert}\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\frac{\bm{w}_{s^{\prime}j}}{\sigma^{2}_{s^{\prime}}}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}}\big{\rVert}^{2}}{\frac{1}{\sigma_{0}^{2}}+\frac{1}{\sigma^{2}_{s^{\prime}}}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}\frac{1}{\sigma_{s}^{2}}}\\\
&+\frac{\big{\lVert}\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}}\big{\rVert}^{2}}{\frac{1}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}\frac{1}{\sigma_{s}^{2}}},i\leq
J_{-s^{\prime}}\end{split}\\\
\begin{split}2\log{\frac{i-J_{-s^{\prime}}}{\gamma_{0}/S}}&-\frac{\big{\lVert}\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\frac{\bm{w}_{s^{\prime}j}}{\sigma^{2}_{s^{\prime}}}\big{\rVert}^{2}}{\frac{1}{\sigma_{0}^{2}}+\frac{1}{\sigma^{2}_{s^{\prime}}}}\\\
&+\frac{\big{\lVert}\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}\big{\rVert}^{2}}{\frac{1}{\sigma_{0}^{2}}},\qquad\quad
J_{-s^{\prime}}<i\leq J_{-s^{\prime}}+J_{s^{\prime}},\end{split}\end{cases}$
(5)

where $n_{i}^{-s^{\prime}}=\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}$ denotes
the number of local neurons were assigned to global neuron $i$ outside of
$s^{\prime}$.

The proof can be found in Appendix B and (Yurochkin et al. 2019b).

## Analysis of PFNM

Before we discuss, we have the following proposition:

######  Proposition 2.

For any prior of a global neuron
$\bm{\theta}_{i}\sim\mathcal{N}(\bm{\mu}_{0},\bm{I}\sigma_{0}^{2})$, when it
is assigned with a likelihood observed local neuron
$\bm{w}_{s^{\prime}j}|\bm{\theta}_{i}\sim\mathcal{N}(\bm{\theta}_{i},\bm{I}\sigma_{s^{\prime}}^{2})$,
the posterior distribution of this global neruon is
$\bm{\theta}_{i}|\bm{w}_{s^{\prime}j}\sim\mathcal{N}\Bigl{(}\frac{\bm{\mu}_{0}/\sigma_{0}^{2}+\bm{w}_{s^{\prime}j}/\sigma^{2}_{s^{\prime}}}{1/\sigma_{0}^{2}+1/\sigma^{2}_{s^{\prime}}},\bm{I}\bigl{(}\frac{1}{1/\sigma_{0}^{2}+1/\sigma^{2}_{s^{\prime}}}\bigr{)}\Bigr{)}$;
when it has been assigned local neurons from other assignments
$\bm{Z}_{i}^{-s^{\prime}}=\\{(s,j)|A_{ij}^{s}=1,s\in-s^{\prime}\\}$, the
posterior distribution is
$\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}}\sim\mathcal{N}(\frac{\bm{\mu}_{0}/\sigma_{0}^{2}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\bm{w}_{sj}/\sigma_{s}^{2}}{1/\sigma_{0}^{2}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}/\sigma_{s}^{2}},\bm{I}(\frac{1}{1/\sigma_{0}^{2}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}/\sigma_{s}^{2}}))$.

The proof can be found in Appendix C.

Terms of equation (5) on the left are due to
$P(\bm{A}^{s^{\prime}}|\\{\bm{A}^{-s^{\prime}}\\})$. At initial stage of
iteration all $n_{i}^{-s^{\prime}}$ are identical. So the main differences
compared by cost $C^{s^{\prime}}_{i,j}$ induced from right term.

######  Definition 3.

(Standardized mean square) For any multivariate Gaussian distribution
$\mathcal{N}(\bm{\mu},\bm{I}\sigma^{2})$, we call
$\left\lVert\bm{\mu}\right\rVert^{2}/\sigma^{2}$ as standardized mean square
of this distribution, and denote it as
$\mathbb{S}(\mathcal{N})=\left\lVert\bm{\mu}\right\rVert^{2}/\sigma^{2}$.

The proposition 2 shows that when $i\leq J_{-s^{\prime}}$, the right term is
the difference in standardized mean square between distribution
$p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}})$ and this distribution assigned
with local likelihood
$p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}},\bm{w}_{s^{\prime}j})$; when
$J_{-s^{\prime}}<i\leq J_{-s^{\prime}}+J_{s^{\prime}}$, the right term is the
difference of standardized mean square between prior distribution
$p(\bm{\theta}_{i})$ and this distribution assigned with local likelihood
$p(\bm{\theta}_{i}|\bm{w}_{s^{\prime}j})$. So, by comparing each cost term
$C^{s^{\prime}}_{i,j}$, it essentially compares the standardized mean square
difference between the $i$-th global neuron distribution (either estimated
from $-s^{\prime}$ or a new one from the current local model) and this
distribution assigned with local neuron $w_{s^{\prime}j}$. However, we have
the following proposition shows the issue caused by the cost term:

######  Proposition 4.

For a global neuron distribution
$p(\bm{\theta})=\mathcal{N}(\bm{\mu}_{0},\bm{I}\sigma_{0}^{2})$, after it is
assigned with a likelihood
$p(\bm{w}|\bm{\theta})=\mathcal{N}(\bm{\theta},\bm{I}\sigma^{2})$, the
standardized mean square difference between $p(\bm{\theta})$ and
$p(\bm{\theta}|\bm{w})$ is proportion to
$\left\lVert\bm{w}\right\rVert^{2}/\sigma^{2}$.

###### Proof.

Denote the mean and covariance matrix of $p(\bm{\theta}|\bm{w})$ as
$\tilde{\bm{\mu}}$ and $\bm{I}\tilde{\sigma}^{2}$. From definition 3, we can
write the the standardized mean square difference between $p(\bm{\theta})$ and
$p(\bm{\theta}|\bm{w})$ as

$\displaystyle\mathbb{S}(p(\bm{\theta}))-\mathbb{S}(p(\bm{\theta}|\bm{w}))$
$\displaystyle=$
$\displaystyle\frac{\left\lVert\bm{\mu}_{0}\right\rVert^{2}}{\sigma_{0}^{2}}-\frac{\left\lVert\tilde{\bm{\mu}}\right\rVert^{2}}{\tilde{\sigma}^{2}}$
$\displaystyle=$
$\displaystyle\log\frac{\exp{(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}})}}{(2\pi\sigma_{0}^{2})^{\frac{D+K}{2}}p(\bm{\theta})}-\log\frac{\exp{(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\tilde{\bm{\mu}}}{2\tilde{\sigma}^{2}})}}{(2\pi\tilde{\sigma}^{2})^{\frac{D+K}{2}}p(\bm{\theta}|\bm{w})}$
$\displaystyle=$
$\displaystyle\log\frac{\exp{(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}})}}{(2\pi\sigma_{0}^{2})^{\frac{D+K}{2}}p(\bm{\theta})}\cdot\frac{(2\pi\tilde{\sigma}^{2})^{\frac{D+K}{2}}p(\bm{\theta}|\bm{w})}{\exp{(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\tilde{\bm{\mu}}}{2\tilde{\sigma}^{2}})}}.$
(6)

Because

$\displaystyle(\frac{\tilde{\sigma}^{2}}{\sigma_{0}^{2}})^{\frac{D+K}{2}}\cdot\exp(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}}-\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\tilde{\bm{\mu}}}{2\tilde{\sigma}^{2}})$
$\displaystyle=$
$\displaystyle(\frac{\tilde{\sigma}^{2}}{\sigma_{0}^{2}})^{\frac{D+K}{2}}\cdot\exp(\bm{\theta}^{T}\bm{\theta}(-\frac{1}{2\sigma_{0}^{2}}+\frac{1}{2\tilde{\sigma}^{2}})+\bm{\theta}^{T}(\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}-\frac{\tilde{\bm{\mu}}}{\tilde{\sigma}^{2}})),$
(7)

from proposition 2, we can obtain

$\displaystyle\frac{1}{2}(-\frac{1}{\sigma_{0}^{2}}+\frac{1}{\tilde{\sigma}^{2}})=\frac{1}{2\sigma^{2}},\qquad\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}-\frac{\tilde{\bm{\mu}}}{\tilde{\sigma}^{2}}=-\frac{\bm{w}}{\sigma^{2}},$
(8)
$\displaystyle(\frac{\tilde{\sigma}^{2}}{\sigma_{0}^{2}})^{\frac{D+K}{2}}=(1+\frac{\sigma_{0}^{2}}{\sigma^{2}}))^{-\frac{D+K}{2}}\text{
is a constant in PFNM setting,}$

and from Bayesian theory

$p(\bm{\theta}|\bm{w})\propto p(\bm{\theta})p(\bm{w}|\bm{\theta}),$ (9)

thus we have

$\displaystyle\mathbb{S}(p(\bm{\theta}))-\mathbb{S}(p(\bm{\theta}|\bm{w}))$
$\displaystyle\propto\log\frac{p(\bm{w})}{\exp(\frac{-\bm{\theta}^{T}\bm{\theta}+2\bm{\theta}^{T}\bm{w}}{2\sigma^{2}})}$
$\displaystyle\propto-\frac{\bm{w}^{T}\bm{w}}{2\sigma^{2}}.$ (10)

∎

From proposition 4, We can know that the differences compared by each
$C^{s’}_{ij}$ are only related to local neuron. As a consequence, for one
fixed local neuron $\bm{w}_{s^{\prime}j}$ and two different global neurons
$\bm{\theta}_{i}$ and $\bm{\theta}_{i^{*}}$, the cost does not discriminate
global neurons $\bm{\theta}_{i}$ and $\bm{\theta}_{i^{*}}$. Hence, the cost
specifications induced from original PFNM can’t find optimal solutions for
$\\{\bm{A}^{s^{\prime}}\\}$.

### Kullback-Leibler Divergence

From the above analysis, we know that the cost specification
$C^{s^{\prime}}_{i,j}$ induced from PFNM omits information of global neuron,
i.e. the $i$ index in the cost. So, it is natural to fix the issue by adding a
regularized term that contains both information of $i$ and $j$. As we point
out before, the cost term $C^{s^{\prime}}_{i,j}$ essentially is the
standardized mean square difference between $i$-th global distribution and
this distribution assigned with local neuron $w_{s^{\prime}j}$. From the
probabilistic perspective, minimizing the cost is meant to hope these two
distributions are as close as possible. So how to measure the distance between
two distributions? Kullback-Leibler becomes the first answer that gets into
mind. Can KL-divergence fix the issues caused by PFNM? To answer this
question, we have the following proposition:

######  Proposition 5.

For a global neuron distribution $p(\bm{\theta})$, after it is assigned with a
likelihood $p(\bm{w}|\bm{\theta})$, the Kullback–Leibler Divergence Penalty
between $p(\bm{\theta})$ and $p(\bm{\theta}|\bm{w})$ is equal up to
$\mathbb{E}_{p(\bm{\theta})}\bigl{[}\log p(\bm{w}\mid\bm{\theta})\bigr{]}.$

###### Proof.

$\displaystyle\operatorname{KL}\bigl{(}p(\bm{\theta})\|p(\bm{\theta}|\bm{w})\bigr{)}$
$\displaystyle=$ $\displaystyle\int
p(\bm{\theta})\log\frac{p(\bm{\theta})}{p(\bm{\theta}|\bm{w})}\,\rm{d}\bm{\theta}$

and from Bayesian theory

$p(\bm{\theta}|\bm{w})\propto p(\bm{\theta})p(\bm{w}|\bm{\theta}),$ (12)

so we have

$\displaystyle\operatorname{KL}\bigl{(}p(\bm{\theta})\|p(\bm{\theta}|\bm{w})\bigr{)}\cong-\mathbb{E}_{p(\bm{\theta})}\bigl{[}\log
p(\bm{w}\mid\bm{\theta})\bigr{]}$ (13)

∎

From above proposition, KL divergence between two distribution is essentially
equal up to the expectation of the assigned local likelihood distribution, and
the expectation is taken over the global distribution $p(\bm{\theta})$. That
means the KL divergence contains both information of local and global neuron.

Thus, to fix the drawback induced from PFNM, we can formulate a new cost
specifications that regularize original cost specification with the KL
penalty, i.e., $\widetilde{\bm{C}}^{s^{\prime}}_{i,j}=$

$\begin{cases}\begin{split}\bm{C}^{s^{\prime}}_{i,j}+\lambda\operatorname{KL}\Bigl{(}p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}})\big{\|}p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}},\bm{w}_{s^{\prime}j})\Bigr{)},i\leq
J_{-s^{\prime}},\quad\end{split}\\\
\begin{split}\bm{C}^{s^{\prime}}_{i,j}+\lambda\operatorname{KL}\Bigl{(}p(\bm{\theta}_{i})\big{\|}p(\bm{\theta}_{i}|\bm{w}_{s^{\prime}j})\Bigr{)},&J_{-s^{\prime}}<i\leq
J_{-s^{\prime}}+J_{s^{\prime}}.\end{split}\end{cases}$ (14)

where the coefficient $\lambda$ is the adjusting ratio. And the KL divergence
between two multivariate normal distributions can be calculated by the
following Lemma:

######  Lemma 5.1.

The Kullback-Leibler divergence between
$\mathcal{N}_{x}(\bm{\mu}_{x},\bm{\Sigma}_{x})$ and
$\mathcal{N}_{y}(\bm{\mu}_{y},\bm{\Sigma}_{y})$, is:

$\frac{1}{2}[\operatorname{tr}(\bm{\Sigma}_{y}^{-1}\bm{\Sigma}_{x})+(\bm{\mu}_{y}-\bm{\mu}_{x})^{T}\bm{\Sigma}_{y}^{-1}(\bm{\mu}_{y}-\bm{\mu}_{x})-(D+K)+\ln{\frac{|\bm{\Sigma}_{y}|}{|\bm{\Sigma}_{x}|}}],$
(15)

where $D+K$ is the dimension of $\bm{\mu}_{x}$ and $\bm{\mu}_{y}$.

The proof can be found in (Duchi 2007).

We call our new method as neural aggregation with full information (NAFI). And
the process of NAFI can be summarized in algorithm 1.

As demonstrated by the study in (Wang et al. 2020), directly applying the
matching algorithms fails on deep architectures designed for more complex
tasks. To address this issue, we extend NAFI to a layer-wise matching scheme,
which can be found in Appendix D.

Algorithm 1 The process of NAFI

0: Local weights $\bm{w}_{sj}$ from $S$ local models;

0: Matching assignments $\\{\bm{A}^{s}\\}_{s=1}^{S}$, global weights
$\\{\bm{\theta}_{i}$ };

1: for $t$ = 1, 2, 3, $\cdots$ do

2: for $s^{\prime}$ = 1, 2, $\cdots$, $S$ do

3: Giving assignments $\\{\bm{A}^{s}_{i,j}\\}_{i,j,s\in-s^{\prime}}$,
acquiring corresponding global model via equations (4):
$\tilde{\bm{\theta}}_{i}=\frac{\bm{\mu}_{0}/\sigma_{0}^{2}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}\bm{w}_{sj}/\sigma_{s}^{2}}{1/\sigma_{0}^{2}+\sum_{s\in{-s^{\prime}},j}\bm{A}_{i,j}^{s}/\sigma_{s}^{2}}$;

4: Acquire the assignment cost
$\\{\widetilde{\bm{C}}^{s^{\prime}}_{i,j}\\}_{i,j}$ via equation (14);

5: Obtain $\\{\bm{A}^{s^{\prime}}_{i,j}\\}_{i,j}$ by solving the linear
assignment problem via Hungarian algorithm.

6: end for

7: Apply $\\{\bm{A}^{s}\\}_{s=1}^{S}$ to update the global model.

8: end for

## Experiments

To evaluate the efficiency of the proposed NAFI, we presents an empirical
study of NAFI and compares it with PFNM, FedAvg (McMahan et al. 2017), FedProx
(Li et al. 2018) and optimal transport fusion (Singh and Jaggi 2020) (which we
refer as OT fusion in remaining part). The experiments are carried out over
three datasets with three types of neural networks: FCNN, shallow CNN and
U-net. The experiments below indicate that our algorithm can aggregate
multiple neural networks into a more efficient global neural
network111https://github.com/XiaoPeng24/NAFI.

#### Datasets, models and metrics

Our algorithm is evaluated on three datasets: MINST, CIFAR 10 and Carvana
Image Masking Challenge (CIMC). MINST and CIFAR-10 are image classification
datasets and each contains ten classes on handwriting digits and objects in
real life respectively. CIMC is a binary semantic segmentation dataset
consisting of photos of cars, and the task is to split out the car and the
background. For MNIST, we use an FCNN model and evaluate it with accuracy; for
CIFAR 10, we use a ConvNet with 3 convolutional and 2 fully-connected layers
and evaluate it with accuracy; and for CIMC, we use the U-net (Ronneberger,
Fischer, and Brox 2015) and evaluate it with dice coefficient.

#### Partition strategies of local data

To simulate a federated learning scenario, we partition each dataset in a
heterogeneous strategy in which the number of data points and class
proportions in each local model is unbalanced. We follow prior works
(Yurochkin et al. 2018) in the heterogeneous partition of local data for three
datasets, which apply $K$-dimensional Dirichlet distribution $Dir(\alpha)$ to
create non-identical independent distribution (iid) data, with a smaller
$\alpha$ indicating higher data heterogeneity. Specifically, for dataset with
class number $K$, we sample the proportion of the instances of class $k$ to
local model $s$, $p_{k,s}$, via $p_{k,s}\sim Dir_{k}(\alpha)$. For MNIST and
CIFAR10, $K=10$. For CIMC, $K=1$ as it is a binary semantic segmentation
dataset. In each dataset, we execute 5 trials to obtain the mean and standard
deviation of the performances.

#### Baselines

Our method is compared to the original PFNM, FedAvg, OT fusion and FedProx.
FedAvg and FedProx are executed in local neural networks that have been
trained using the same random initialization as proposed by (McMahan et al.
2017). We note that while a federated averaging variant without shared
initialization is likely to be more realistic when attempting to aggregate
pre-trained models, it performs significantly worse than all other baselines.
In CNN architectures, the original PFNM we compared is actually its deep model
extensive version—FedMA (Wang et al. 2020). Considering the high heterogeneity
in each trial, we set $\lambda$ in a wide grid (from $10^{-8}$ to 1) and
choose the best result for NAFI. We subsequently check the sensitivity of
$\lambda$ and find it only have tiny fluctuation.

Table 1: Hyperparameter settings for training neural networks | MNIST | CIFAR-10 | CIMC
---|---|---|---
Model | FCNN | ConvNet | U-net
Optimizer | Adam | SGD | RMSprop
Learning rate | 0.01 | 0.01 | 0.01
Size of minibatch | 32 | 32 | 3
Epochs | 10 | 10 | 3
Table 2: Test accuracy of FCNN in MNIST | S | N | Local NN | FedAvg | FedProx | OT fusion | PFNM | NAFI
---|---|---|---|---|---|---|---|---
nets | 15 | 1 | 71.9 $\pm$ 2.57 | 75.47 $\pm$ 5.90 | 75.65 $\pm$ 5.93 | 81.44 $\pm$ 3.21 | 83.93 $\pm$ 0.14 | 87.34 $\pm$ 0.57
20 | 1 | 69.44 $\pm$ 2.50 | 75.02 $\pm$ 4.03 | 75.17 $\pm$ 3.99 | 82.71 $\pm$ 3.57 | 83.23 $\pm$ 3.31 | 86.73 $\pm$ 2.15
25 | 1 | 67.99 $\pm$ 2.00 | 75.46 $\pm$ 3.33 | 75.24 $\pm$ 3.30 | 82.53 $\pm$ 3.05 | 84.87 $\pm$ 1.66 | 86.46 $\pm$ 2.56
30 | 1 | 65.90 $\pm$ 2.66 | 73.43 $\pm$ 4.49 | 73.11 $\pm$ 4.44 | 82.74 $\pm$ 2.41 | 83.39 $\pm$ 2.23 | 86.23 $\pm$ 2.57
layers | 10 | 2 | 73.21 $\pm$ 3.05 | 66.56 $\pm$ 6.82 | 66.34 $\pm$ 6.75 | 83.77 $\pm$ 4.23 | 79.09 $\pm$ 4.73 | 84.58 $\pm$ 4.54
10 | 3 | 70.28 $\pm$ 2.97 | 52.01 $\pm$ 6.52 | 51.79 $\pm$ 6.49 | 67.56 $\pm$ 5.64 | 60.71 $\pm$ 4.59 | 72.04 $\pm$ 4.75
Table 3: Test accuracy of Convnet in CIFAR10 | S | Local NN | FedAvg | FedProx | OT fusion | PFNM (FedMA) | NAFI
---|---|---|---|---|---|---|---
nets | 5 | 25.21 $\pm$ 2.30 | 51.26 $\pm$ 2.97 | 51.43 $\pm$ 3.01 | 52.12 $\pm$ 2.78 | 50.39 $\pm$ 0.94 | 52.56 $\pm$ 0.81
10 | 18.92 $\pm$ 1.41 | 46.78 $\pm$ 2.36 | 46.94 $\pm$ 2.32 | 45.43 $\pm$ 3.45 | 46.27 $\pm$ 1.73 | 47.46 $\pm$ 1.42
15 | 15.85 $\pm$ 0.51 | 42.40 $\pm$ 2.25 | 42.06 $\pm$ 2.20 | 40.79 $\pm$ 4.29 | 42.82 $\pm$ 2.46 | 45.72 $\pm$ 1.35
20 | 14.11 $\pm$ 0.53 | 34.20 $\pm$ 2.54 | 34.02 $\pm$ 2.52 | 37.56 $\pm$ 4.13 | 42.61 $\pm$ 2.07 | 44.96 $\pm$ 1.93
Table 4: Test dice coefficient of U-net in CIMC | S | Local NN | FedAvg | FedProx | OT fusion | PFNM (FedMA) | NAFI
---|---|---|---|---|---|---|---
nets | 8 | 67.60 $\pm$ 9.38 | 53.28 $\pm$ 10.63 | 41.53 $\pm$ 0.82 | 63.57 $\pm$ 5.48 | 90.31 $\pm$ 2.90 | 96.47 $\pm$ 0.76
16 | 27.84 $\pm$ 12.16 | 42.62 $\pm$ 0.87 | 48.50 $\pm$ 9.52 | 54.49 $\pm$ 5.62 | 75.47 $\pm$ 3.54 | 83.46 $\pm$ 3.41

#### Training setup

We use PyTorch (Paszke et al. 2017) to implement these networks and train them
by the Adam (Kingma and Ba 2014), SGD (Bottou 2010) and RMSprop (Hinton,
Srivastava, and Swersky 2012) with hyperparameter settings which are
summarized in Supplement table 1.

#### Performance Overview

In real world applications of Federated Learning, the discrepancy of data
distribution among local models and communication cost will inevitably
increase as the number of local models increases. This is mainly due to the
variability of data generation paradigms in the system (Li et al. 2020). This
suggests that testing fusing algorithms for various numbers of local models
$S$ is important. For MNIST, we firstly apply various baselines with 15, 20,
25 and 30 local models with one hidden layer FCNN. Local NN in Table 2 reports
the average of separately tested network accuracies. The lower extremes of
aggregating are determined by the performance of local NNs. As shown in Table
2, accuracy of NAFI on MNIST is 4% higher than PFNM on average. We also
examine how the performance of all methods is impacted by the number of hidden
layers $N$. We train 10 neural networks with 2 and 3 hidden layers
respectively and then fuse them using various baselines. The performance of
PFNM decreases with the increasing number of hidden layers—even worse than OT
fusion. This phenomenon has already been proposed in (Wang et al. 2020), and
it can be fixed by a layer-wise training manner which we apply in CNN
architectures. It is worth mentioning that NAFI seems have the potential to
address the drawback of PFNM in deep architecture, as NAFI still maintain at a
moderate level in deep architecture. This may be explained by an accumulating
error effect. From above analysis, we discovered that PFNM is unable to
discriminate those global neurons because the cost specification does not
include global neuron information, while NAFI equipped with KL divergence can
correct this. Thus, with the fusing process of neural networks is going layer-
by-layer, the drawback of PFNM described above takes effect, and global
neurons are not generated correctly. As a result, the incorrectness of PFNM in
picking global neurons is superimposed, and the performance disparity between
PFNM and NAFI is amplified as presented in Table 2.

For CNN application, we apply ConvNet (2 convolutional layers and 3 fully
connected layers) on 5, 10, 15, and 20 local models and use various methods to
fuse the models trained in CIFAR 10’s heterogeneous dataset partition. We
apply PFNM and NAFI in an iteratively layer-wise way mentioned above, so we
also set the communications rounds of FedAvg, FedProx and OT fusion to 5 equal
to the number of layers in ConvNet for equality. As shown in Table 3, NAFI
outperforms the other methods on fusing convolutional neural networks.
Concretely, the accuracy of NAFI on CIFAR 10 is 3% higher than PFNM on average
for a various number of local models.

U-Net (Ronneberger, Fischer, and Brox 2015) is widely used on medical image
segmentation (Du et al. 2020), and medical images such as tumor scan images
generated by multiple institutions are generally forbidden to be exchanged for
privacy issues. Therefore, it is critical to federate well-known image
segmentation algorithms. To this end, we compare our method and baselines on
U-Net. The U-Net architecture is described in (Ronneberger, Fischer, and Brox
2015), and it employs the batch normalization layer, which is not extended by
the matching type federated learning algorithm yet. Here we apply a technique
which fuses batch normalization layer into convolutional layer, and extend the
matching type algorithm to batch normalization layer. The details can be
referred in Appendix E. As shown in Table 4, FedAvg, FedProx and OT fusion
have very subpar performances. And FedAvg and FedProx even cannot converge to
a stationary result in such complex network architecture (the results go up
and down during communication rounds). The result shown in the table is taken
from the best result during all 19 communication rounds which are equal to the
number of layers in U-net. As shown in Table 4, the priority of NAFI is
significantly more important among popular Federated Learning methods than it
is in the case of 2- or 3-layer neural networks. With 8 local models, NAFI on
CIMC has accuracy that is at least 7% higher than PFNM, and accuracy that is
almost 8% higher with 16 local models.

#### Amount of fused global model

In real federated application, the amount of fused global model determines the
burden of communication. Hence we also measure the size of global model fused
by PFNM and NAFI respectively. In figure 1, we plot the logarithm ratio
between the amount of fused global and the total amount of all local models
$\log{\frac{J}{\sum_{s}{J_{s}}}}$ in FCNN multiple nets scenario. As shown in
figure 1, the logarithm ratio decrease with the increase in number of local
models and is kept lower than -0.5, which means the matching type method
produces a relative small model while can capture the most advantages of local
models. It is worth noting that the ratio of NAFI is always small than PFNM,
while the performance of fused global model in NAFI outperforms than PFNM
accordint to above experiment results. This indicates the availability of NAFI
in some model compression tasks.

Figure 1: The logarithm ratio between the amount of fused global and the total
amount of all local models in FCNN multiple nets scenario.

### Sensitivity Analysis

The weight of the KL-divergence term, $\lambda$, is a hyper-parameter that we
introduce into NAFI. Therefore, it is essential to demonstrate the sensitivity
of $\lambda$. Here, we set $\lambda$ to various positive values for NAFI
applied on fusing FCNNs and CNNs trained in MNIST and CIFAR10 respectively. We
also take the impact of the number of local models into account. As shown in
figure 2, the heat map indicates the accuracy on the test data, and for
various number of clients, there is only a tiny fluctuation of prediction
accuracy for a fused global model when $10^{-8}\leq\lambda\leq 1$. Although
NAFI with $\lambda=10^{-8}$ performs significantly worse than NAFI with
$10^{-3}\leq\lambda\leq 0.5$, the fused global model maintains a high level of
performance. When $\lambda=1$, the performance of the fused global model drops
sharply, this indicates that $\lambda$ should not set too high. To summarize,
NAFI is robust on the hyperparameter $\lambda$ under a variety of conditions.

Figure 2: KL regularization coefficients sensitivity analysis

## Conclusion

In this study, we theoretically analyze the flaw of PFNM and propose a new
federated neural matching method to fix this flaw. It is empirically shown
that the new method outperforms other state-of-the-art algorithms for
federated learning of neural networks, and the comparison of fused global
model size also indicates the potential availability of our method in model
compression and pruning tasks. In future work, we will extend our algorithm
into more advanced architectures, like transformers etc. Additionally, the
success of KL-divergence suggests it is also likely to try other metrics which
can measure the similarity between two probability distributions.

## Acknowledgements

During the process of preparing this paper, the FIFA World Cup Qatar 2022 is
underwaying. The Argentina football player Lionel Messi and his Argentina
men’s football team have also successfully qualified for the group stage. For
a long time, Messi has maintained a brilliant and lasting career through his
own efforts and self-discipline. This kind of hard work and self-discipline
also deeply moved the author, and Messi also become the spiritual mentor of
the author during his career path. Lack of a World Cup trophy is the only
barrier for Messi being the world’s number one football player. In 2014, Messi
regretted passing it by. Here, the author wishes Messi to win the World Cup
trophy in this year through his effort. And all of us should believe that God
rewards those who work hard. Let’s work hard for our goals.

## References

  * Bottou (2010) Bottou, L. 2010.  Large-scale machine learning with stochastic gradient descent.  In _Proceedings of COMPSTAT’2010_ , 177–186. Springer.
  * Buciluǎ, Caruana, and Niculescu-Mizil (2006) Buciluǎ, C.; Caruana, R.; and Niculescu-Mizil, A. 2006.  Model compression.  In _Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining_ , 535–541.
  * Chen and Chao (2020) Chen, H.-Y.; and Chao, W.-L. 2020.  Fedbe: Making Bayesian model ensemble applicable to federated learning.  _arXiv preprint arXiv:2009.01974_.
  * Claici et al. (2020) Claici, S.; Yurochkin, M.; Ghosh, S.; and Solomon, J. 2020.  Model Fusion with Kullback–Leibler Divergence.  _arXiv preprint arXiv:2007.06168_.
  * Deng, Kamani, and Mahdavi (2021) Deng, Y.; Kamani, M. M.; and Mahdavi, M. 2021.  Distributionally robust federated averaging.  _arXiv preprint arXiv:2102.12660_.
  * Du et al. (2020) Du, G.; Cao, X.; Liang, J.; Chen, X.; and Zhan, Y. 2020.  Medical image segmentation based on U-net: A review.  _Journal of Imaging Science and Technology_ , 64(2): 20508–1.
  * Duchi (2007) Duchi, J. 2007.  Derivations for linear algebra and optimization.  _Berkeley, California_ , 3(1): 2325–5870.
  * Ghahramani and Griffiths (2006) Ghahramani, Z.; and Griffiths, T. L. 2006.  Infinite latent feature models and the Indian buffet process.  In _Advances in neural information processing systems_ , 475–482.
  * Hinton, Srivastava, and Swersky (2012) Hinton, G.; Srivastava, N.; and Swersky, K. 2012.  Neural networks for machine learning lecture 6a overview of mini-batch gradient descent.  _Cited on_ , 14(8): 2.
  * Hinton, Vinyals, and Dean (2015) Hinton, G.; Vinyals, O.; and Dean, J. 2015.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_.
  * Karimireddy et al. (2020) Karimireddy, S. P.; Kale, S.; Mohri, M.; Reddi, S.; Stich, S.; and Suresh, A. T. 2020.  Scaffold: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , 5132–5143. PMLR.
  * Kingma and Ba (2014) Kingma, D. P.; and Ba, J. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Li et al. (2020) Li, T.; Sahu, A. K.; Talwalkar, A.; and Smith, V. 2020.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Processing Magazine_ , 37(3): 50–60.
  * Li et al. (2018) Li, T.; Sahu, A. K.; Zaheer, M.; Sanjabi, M.; Talwalkar, A.; and Smith, V. 2018\.  Federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_.
  * Lin et al. (2020) Lin, T.; Kong, L.; Stich, S. U.; and Jaggi, M. 2020.  Ensemble distillation for robust model fusion in federated learning.  _arXiv preprint arXiv:2006.07242_.
  * McMahan et al. (2017) McMahan, B.; Moore, E.; Ramage, D.; Hampson, S.; and y Arcas, B. A. 2017.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial Intelligence and Statistics_ , 1273–1282.
  * Mohri, Sivek, and Suresh (2019) Mohri, M.; Sivek, G.; and Suresh, A. T. 2019.  Agnostic Federated Learning.  In _International Conference on Machine Learning_ , 4615–4625.
  * Paszke et al. (2017) Paszke, A.; Gross, S.; Chintala, S.; Chanan, G.; Yang, E.; DeVito, Z.; Lin, Z.; Desmaison, A.; Antiga, L.; and Lerer, A. 2017.  Automatic differentiation in pytorch.
  * Ronneberger, Fischer, and Brox (2015) Ronneberger, O.; Fischer, P.; and Brox, T. 2015.  U-net: Convolutional networks for biomedical image segmentation.  In _International Conference on Medical image computing and computer-assisted intervention_ , 234–241. Springer.
  * Schmidhuber (1992) Schmidhuber, J. 1992.  Learning complex, extended sequences using the principle of history compression.  _Neural Computation_ , 4(2): 234–242.
  * Singh and Jaggi (2020) Singh, S. P.; and Jaggi, M. 2020.  Model fusion via optimal transport.  _Advances in Neural Information Processing Systems_ , 33.
  * Smith et al. (2017) Smith, V.; Chiang, C.-K.; Sanjabi, M.; and Talwalkar, A. S. 2017.  Federated multi-task learning.  In _Advances in Neural Information Processing Systems_ , 4424–4434.
  * Teh, Grür, and Ghahramani (2007) Teh, Y. W.; Grür, D.; and Ghahramani, Z. 2007.  Stick-breaking construction for the Indian buffet process.  In _Artificial Intelligence and Statistics_ , 556–563.
  * Thibaux and Jordan (2007a) Thibaux, R.; and Jordan, M. I. 2007a.  Hierarchical Beta Processes and the Indian Buffet Process.  _Journal of Machine Learning Research_ , 2(3): 564–571.
  * Thibaux and Jordan (2007b) Thibaux, R.; and Jordan, M. I. 2007b.  Hierarchical Beta processes and the Indian buffet process.  In _Artificial Intelligence and Statistics_ , 564–571.
  * Wang et al. (2020) Wang, H.; Yurochkin, M.; Sun, Y.; Papailiopoulos, D.; and Khazaeni, Y. 2020.  Federated learning with matched averaging.  _arXiv preprint arXiv:2002.06440_.
  * Xiao et al. (2020) Xiao, P.; Cheng, S.; Stankovic, V.; and Vukobratovic, D. 2020.  Averaging is probably not the optimum way of aggregating parameters in federated learning.  _Entropy_ , 22(3): 314.
  * Yurochkin et al. (2019a) Yurochkin, M.; Agarwal, M.; Ghosh, S.; Greenewald, K.; and Hoang, N. 2019a.  Statistical model aggregation via parameter matching.  _Advances in Neural Information Processing Systems_ , 32: 10956–10966.
  * Yurochkin et al. (2019b) Yurochkin, M.; Agarwal, M.; Ghosh, S.; Greenewald, K.; Hoang, N.; and Khazaeni, Y. 2019b.  Bayesian Nonparametric Federated Learning of Neural Networks.  In _International Conference on Machine Learning_ , 7252–7261.
  * Yurochkin et al. (2018) Yurochkin, M.; Fan, Z.; Guha, A.; Koutris, P.; and Nguyen, X. 2018.  Scalable inference of topic evolution via models for latent geometric structures.  _arXiv preprint arXiv:1809.08738_.
  * Zhang et al. (2020) Zhang, X.; Hong, M.; Dhople, S.; Yin, W.; and Liu, Y. 2020.  Fedpd: A federated learning framework with optimal rates and adaptivity to non-iid data.  _arXiv preprint arXiv:2005.11418_.
  * Zhu, Hong, and Zhou (2021) Zhu, Z.; Hong, J.; and Zhou, J. 2021.  Data-Free Knowledge Distillation for Heterogeneous Federated Learning.  _arXiv preprint arXiv:2105.10056_.

##  Appendix A Appendix A: BBP and IBP

Beta-Bernoulli Process and Indian Buffet Process Denote $Q$ as a random
measure drawn from a Beta process:
$Q|\gamma_{0},H\sim\text{BP}(1,\gamma_{0}H)$, where $\gamma_{0}$ is the mass
parameter, $H$ is the base measure over some domain $\Omega$ such that
$H(\Omega)=1$. One can show that $Q$ is a discrete measure with
$Q=\sum_{i}q_{i}\delta_{\bm{\theta}_{i}}$, which can be characterized by an
infinitely countable set of (weight, atom) pairs
$(q_{i},\bm{\theta}_{i})\in[0,1]\times\Omega$. The atoms $\bm{\theta}_{i}$ can
be drawn i.i.d from $H$ and the weights $\\{q_{i}\\}_{i=1}^{\infty}$ can be
generated via a stick-breaking process (Teh, Grür, and Ghahramani 2007):
$q_{1}\sim\operatorname{Beta}(\gamma_{0},1),q_{i}=\prod_{g=1}^{i}q_{g}$. Then
subsets of atoms in the random measure $Q$ can be picked via a Bernoulli
process. That is, each subset $\mathcal{T}_{s}$ for $s=1,\cdots,S$ can be
distributed via a Bernoulli process with base measure $Q$:
$\mathcal{T}_{s}|Q\sim\operatorname{BeP}(Q)$. Hence, subset $\mathcal{T}_{s}$
can also be viewed as a discrete measure
$\mathcal{T}_{s}:=\sum_{i}a_{si}\delta_{\bm{\theta}_{i}}$, which is formed by
pairs $(a_{si},\bm{\theta}_{i})\in\\{0,1\\}\times\Omega$, where
$a_{si}|q_{i}\sim\operatorname{Bernoulli}(q_{i}),\forall i$ is a binary random
variable indicating whether $\bm{\theta}_{i}$ belongs to subset
$\mathcal{T}_{s}$. We call such collection of subsets a Beta-Bernoulli process
(Thibaux and Jordan 2007b).

The Indian buffet process (IBP) specifies distribution on sparse binary
matrices (Ghahramani and Griffiths 2006). IBP involves a metaphor of a
sequence of customers tasting dishes in an infinite buffet: the first customer
tastes $\operatorname{Poisson}(\gamma_{0})$ dishes, every subsequent $s$th
customer tastes each dish that is previously selected with probability
$n_{i}/s$, where $n_{i}=\sum_{s=1}^{S-1}a_{si}$, and then tastes
$\operatorname{Poisson}(\gamma_{0}/s)$ new dishes. Marginalizing over Beta
Process distributed $Q$ above will induce dependencies among subsets and
recover the predictive distribution
$\mathcal{T}_{S}|\mathcal{T}_{1},\cdots,\mathcal{T}_{S-1}\sim\operatorname{BeP}(H\frac{\gamma_{0}}{S}+\sum_{i}\frac{n_{i}}{S}\delta_{\bm{\theta}_{i}})$.
That is equivalent to the IBP.

##  Appendix B Appendix B: Proof of Proposition 1 (Yurochkin et al. 2019b)

PFNM maximizes a posterior probability of the global
atoms$\\{\bm{\theta}_{i}\\}_{i=1}^{\infty}$ and assignments of observed neural
network weight estimates to global atoms$\\{\bm{A}^{s}\\}_{s=1}^{S}$. Given
estimates of the local model weights $\\{\bm{w}_{sj}\text{ for
}j=1,\dots,J_{s}\\}_{s=1}^{S}$, it has:

$\max_{\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}}P(\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}|\\{\bm{w}_{sj}\\})\\\
\propto
P(\\{\bm{w}_{sj}\\}|\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\})P(\\{\bm{A}^{s}\\})P(\\{\bm{\theta}_{i}\\}),$
(16)

by taking negative natural logarithm it can obtain:

$\min_{\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}}-\sum_{i}\bigg{(}\sum_{s,j}\bm{A}^{s}_{i,j}\log(p(\bm{w}_{sj}|\sim\bm{\theta}_{i}))+\log(q(\bm{\theta}_{i}))\bigg{)}\\\
-\log(P(\\{\bm{A}^{s}\\})),$ (17)

expand probability function of multi-dimensional Gaussian distributions in
equation (17), it obtains:

$\min_{\\{\bm{\theta}_{i}\\},\\{\bm{A}^{s}\\}}\frac{1}{2}\sum_{i}\Bigg{(}\frac{||\hat{\bm{\theta}}_{i}-\bm{\mu}_{0}||^{2}}{\sigma_{0}^{2}}+(D+K)\log(2\pi\sigma_{0}^{2})\\\
+\sum_{s,j}A^{s}_{i,j}\frac{||\bm{w}_{sj}-\hat{\bm{\theta}}_{i}||^{2}}{\sigma_{s}^{2}}\Bigg{)}-\log(P(\\{\bm{A}^{s}\\})).$
(18)

We now consider the first part of equation (18). Through the closed-form
expression of $\\{\bm{\theta}_{i}\\}$ estimated according to the Gaussian-
Gaussian conjugacy:

$\hat{\bm{\theta}_{i}}=\frac{\bm{\mu}_{0}/\sigma_{0}^{2}+\sum_{s,j}A_{i,j}^{s}\bm{w}_{sj}/\sigma^{2}_{s}}{1/\sigma_{0}^{2}+\sum_{s,j}A^{s}_{i,j}/\sigma^{2}_{s}}\text{
for }i=1,...,J,$ (19)

where for simplicity we assume $\bm{\Sigma}_{0}=\bm{I}\sigma^{2}_{0}$ and
$\bm{\Sigma}_{s}=\bm{I}\sigma^{2}_{s}$, we can now cast first part of equation
(18) with respect only to $\\{\bm{A}^{s}\\}_{s=1}^{S}$:

$\displaystyle\frac{1}{2}\sum_{i}\Bigg{(}\frac{||\hat{\bm{\theta}}_{i}-\bm{\mu}_{0}||^{2}}{\sigma_{0}^{2}}+(D+K)\log(2\pi\sigma_{0}^{2})+\sum_{s,j}A^{s}_{i,j}\frac{||\bm{w}_{sj}-\hat{\bm{\theta}}_{i}||^{2}}{\sigma_{s}^{2}}\Bigg{)}$
(20) $\displaystyle\cong$
$\displaystyle\frac{1}{2}\sum_{i}\Bigg{(}\langle\hat{\bm{\theta}}_{i},\hat{\bm{\theta}}_{i}\rangle(\frac{1}{\sigma_{0}^{2}}+\sum_{s,j}\frac{A^{s}_{i,j}}{\sigma^{2}_{s}})+(D+K)\log(2\pi\sigma_{0}^{2})-2\langle\hat{\bm{\theta}}_{i},\sum_{s,j}A^{s}_{i,j}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}})\rangle\Bigg{)}$
$\displaystyle=$
$\displaystyle-\frac{1}{2}\sum_{i}\Bigg{(}\frac{||\sum\limits_{s,j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{(1/\sigma_{0}^{2}+\sum\limits_{s,j}\bm{A}^{s}_{i,j}/\sigma_{s}^{2}}-(D+K)\log(2\pi\sigma_{0}^{2})\Bigg{)}.$

Partition equation (20) between $i=1,...,J_{-s^{\prime}}$ and
$i=J_{-s^{\prime}}+1,...,J_{-s^{\prime}}+J_{s^{\prime}}$, and because it is
now solving for $\bm{A}^{s^{\prime}}$, it can subtract terms independent of
$\bm{A}^{s^{\prime}}$:

$\displaystyle\sum_{i}\Bigg{(}\frac{||\sum_{s,j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{(1/\sigma_{0}^{2}+\sum_{s,j}\bm{A}^{s}_{i,j}/\sigma_{s}^{2}}-(D+K)\log(2\pi\sigma_{0}^{2})\Bigg{)}$
(21) $\displaystyle\cong$
$\displaystyle\sum_{i=1}^{J_{-s^{\prime}}}\Bigg{(}\frac{||\sum_{j}\bm{A}^{s^{\prime}}_{i,j}\frac{\bm{w}_{s^{\prime}j}-\bm{\mu}_{0}}{\sigma_{s^{\prime}}^{2}}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{1/\sigma_{0}^{2}+\sum_{j}\bm{A}^{s^{\prime}}_{i,j}/\sigma_{s^{\prime}}^{2}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}/\sigma^{2}_{s}}-\frac{||\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{1/\sigma_{0}^{2}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}/\sigma^{2}_{s}}\Bigg{)}$
$\displaystyle\quad+\sum_{i=J_{-s^{\prime}}+1}^{J_{-s^{\prime}}+J_{s^{\prime}}}\Bigg{(}\frac{||\sum_{j}\bm{A}^{s^{\prime}}_{i,j}\frac{\bm{w}_{s^{\prime}j}-\bm{\mu}_{0}}{\sigma_{s^{\prime}}^{2}}||^{2}}{1/\sigma_{0}^{2}+\sum_{j}\bm{A}^{s^{\prime}}_{i,j}/\sigma^{2}_{s^{\prime}}}\Bigg{)},$

observe that $\sum_{j}\bm{A}^{s^{\prime}}_{i,j}\in\\{0,1\\}$, i.e. it is 1 if
some neuron from dataset $s^{\prime}$ is matched to global neuron $i$ and 0
otherwise. Thus equation (21) can rewritten as a linear sum assignment
problem:

$\displaystyle\sum_{i=1}^{J_{-s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\Bigg{(}\frac{||\frac{\bm{w}_{s^{\prime}j}-\bm{\mu}_{0}}{\sigma_{s^{\prime}}^{2}}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{1/\sigma_{0}^{2}+1/\sigma_{s^{\prime}}^{2}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}/\sigma^{2}_{s}}-\frac{||\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}\frac{\bm{w}_{sj}-\bm{\mu}_{0}}{\sigma_{s}^{2}}||^{2}}{1/\sigma_{0}^{2}+\sum_{s\in-s^{\prime},j}\bm{A}^{s}_{i,j}/\sigma^{2}_{s}}\Bigg{)}$
(22)
$\displaystyle\quad+\sum_{i=J_{-s^{\prime}}+1}^{J_{-s^{\prime}}+J_{s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\Bigg{(}\frac{||\frac{\bm{w}_{s^{\prime}j}-\bm{\mu}_{0}}{\sigma_{s^{\prime}}^{2}}||^{2}}{1/\sigma_{0}^{2}+\sum_{j}\bm{A}^{s^{\prime}}_{i,j}/\sigma^{2}_{s^{\prime}}}\Bigg{)}.$

Then consider the second term of equation (18), by subtracting terms
independent of $\bm{A}^{s^{\prime}}$ it has:

$\log(P(\bm{A}^{s^{\prime}}))=\log(P(\bm{A}^{s^{\prime}}|\bm{A}^{-s^{\prime}}))+\log(P(\bm{A}^{-s^{\prime}})).$
(23)

First, it can ignore $\log(P(\bm{A}^{-s^{\prime}}))$ since now are optimizing
for $\bm{A}^{s^{\prime}}$ . Second, due to exchange ability of datasets (i.e.
customers of the IBP), $\bm{A}^{s^{\prime}}$ can always be treated as the last
customer of the IBP. Denote
$n^{-s^{\prime}}_{i}=\sum_{-s^{\prime},j}\bm{A}^{s^{\prime}}_{i,j}$ as the
number of times local weights were assigned to global atom $i$ outside of
group $s^{\prime}$. Now it can obtain the following:

$\displaystyle\log P(\bm{A}^{s^{\prime}}|\bm{A}^{-s^{\prime}})$
$\displaystyle\cong\sum_{i=1}^{J_{-s^{\prime}}}\Bigg{(}\bigg{(}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\bigg{)}\log\frac{n^{-s^{\prime}}_{i}}{S}+\bigg{(}1-\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\bigg{)}\log\frac{S-n^{-s^{\prime}}_{i}}{S}\Bigg{)}$
(24)
$\displaystyle\qquad-\log\Bigg{(}\sum_{i=J_{-s^{\prime}}+1}^{J_{-s^{\prime}}+J_{s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\Bigg{)}+\Bigg{(}\sum_{i=J_{-s^{\prime}}+1}^{J_{-s^{\prime}}+J_{s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\Bigg{)}\log\frac{\gamma_{0}}{J}.$

equation (24) thus can be rearranged as a linear sum assignment problem:

$\sum_{i=1}^{J_{-s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\log\frac{n^{-s^{\prime}}_{i}}{S-n^{-s^{\prime}}_{i}}+\sum_{i=J_{-s^{\prime}}+1}^{J_{-s^{\prime}}+J_{s^{\prime}}}\sum_{j=1}^{J_{s^{\prime}}}\bm{A}^{s^{\prime}}_{i,j}\Bigg{(}\log\frac{\gamma_{0}}{S}-\log(i-J_{-s^{\prime}})\Bigg{)}.$
(25)

Combining equation (22) and equation (24), we arrive at the cost specification
shown in (6) of the main text. That completes the proof of Proposition 1 in
the main text.

##  Appendix C Appendix C: Proof of Proposition 2

###### Proof.

i) Let’s prove the first part. From Bayesian theory, the posterior is given by

$\displaystyle p(\bm{\theta}_{i}|\bm{w}_{sj})$ $\displaystyle\propto
p(\bm{\theta}_{i})p(\bm{w}_{sj})$ (26)
$\displaystyle\propto\exp{(\frac{-\bm{\theta}_{i}^{T}\bm{\theta}_{i}+2\bm{\theta}_{i}^{T}\bm{\mu}_{0}-\bm{\mu}_{0}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}})}$
(27)
$\displaystyle\exp{(\frac{-\bm{w}_{sj}^{T}\bm{w}_{sj}+2\bm{\theta}_{i}^{T}\bm{w}_{sj}-\bm{\theta}_{i}^{T}\bm{\theta}_{i}}{2\sigma_{s}^{2}})}.$
(28)

From Gaussian conjugate, the product of two Gaussians is still a Gaussian, we
will rewrite this in the form

$\displaystyle p(\bm{\theta}_{i}|\bm{w}_{sj})\propto$
$\displaystyle\exp{\bigl{[}-\frac{\bm{\theta}_{i}^{T}\bm{\theta}_{i}}{2}(\frac{1}{\sigma_{0}^{2}}+\frac{1}{\sigma_{s}^{2}})+\bm{\theta}_{i}^{T}(\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\frac{\bm{w}_{sj}}{\sigma_{s}^{2}})-(\frac{\bm{\mu}_{0}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}}+\frac{\bm{w}_{sj}^{T}\bm{w}_{sj}}{2\sigma_{s}^{2}})\bigr{]}}$
(30)
$\displaystyle\mathop{=}\limits^{def}\exp{[\frac{-\bm{\theta}_{i}^{T}\bm{\theta}_{i}+2\bm{\theta}_{i}^{T}\bm{\hat{}}{\mu}-\bm{\hat{}}{\mu}^{T}\bm{\hat{}}{\mu}}{2\hat{\sigma}^{2}}]}=\exp{[-\frac{1}{2\hat{\sigma}^{2}}(\left\lVert\bm{\theta}_{i}-\bm{\hat{}}{\mu}\right\rVert^{2})]},$
(31)

where $\bm{\hat{}}{\mu}$ and $\hat{\sigma}^{2}$ denote the mean and variance
of the posterior Gaussian. By completing the square, we first match the
coefficients of $\bm{\theta}_{i}^{T}\bm{\theta}_{i}$—second power of
$\bm{\theta}_{i}$, and find $\hat{\sigma}^{2}$ is given by

$\displaystyle\frac{1}{\hat{\sigma}^{2}}$
$\displaystyle=\frac{1}{\sigma_{0}^{2}}+\frac{1}{\sigma_{s}^{2}},$ (32)
$\displaystyle\hat{\sigma}^{2}$
$\displaystyle=\frac{1}{\frac{1}{\sigma_{0}^{2}}+\frac{1}{\sigma^{2}_{s}}},$
(33)

and then match the coefficients of first power $\bm{\theta}_{i}^{T}$ we get

$\displaystyle\frac{\bm{\hat{}}{\mu}}{\hat{\sigma}^{2}}$
$\displaystyle=\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\frac{\bm{w}_{sj}}{\sigma_{s}^{2}},$
(34)

hence

$\displaystyle\bm{\hat{}}{\mu}=\hat{\sigma}^{2}(\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\frac{\bm{w}_{sj}}{\sigma_{s}^{2}}).$
(35)

ii) Now let’s prove the second part. When $\bm{\theta}_{i}$ has been assigned
local neurons from other assignments
$\bm{Z}_{i}^{-s^{\prime}}=\\{(s,j)|A_{ij}^{s}=1,s\in-s^{\prime}\\}$,from
Bayesian theory, the posterior is given by

$\displaystyle p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}})$
$\displaystyle\propto
p(\bm{\theta}_{i})\prod_{z\in\bm{Z}_{i}^{-s^{\prime}}}p(\bm{w}_{z})$ (36)
$\displaystyle\propto\exp{(\frac{-\bm{\theta}_{i}^{T}\bm{\theta}_{i}+2\bm{\theta}_{i}^{T}\bm{\mu}_{0}-\bm{\mu}_{0}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}})}$
(37)
$\displaystyle\exp{(-\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}^{T}\bm{w}_{sj}}{\sigma_{s}^{2}}+\bm{\theta}_{i}^{T}\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}}-\bm{\theta}_{i}^{T}\bm{\theta}_{i}\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{1}{2\sigma_{s}^{2}})}.$
(38)

Since the product of two Gaussians is a Gaussian, we will rewrite this in the
form

$\displaystyle p(\bm{\theta}_{i}|\bm{Z}_{i}^{-s^{\prime}})$
$\displaystyle\propto\exp{\bigl{[}-\frac{\bm{\theta}_{i}^{T}\bm{\theta}_{i}}{2}(\frac{1}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{1}{\sigma_{s}^{2}})+\bm{\theta}_{i}^{T}(\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}})-(\frac{\bm{\mu}_{0}^{T}\bm{\mu}_{0}}{2\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}^{T}\bm{w}_{sj}}{2\sigma_{s}^{2}})\bigr{]}}$
(39)
$\displaystyle\mathop{=}\limits^{def}\exp{(\frac{-\bm{\theta}_{i}^{T}\bm{\theta}_{i}+2\bm{\theta}_{i}^{T}\bm{\hat{}}{\mu}-\bm{\hat{}}{\mu}^{T}\bm{\hat{}}{\mu}}{2\hat{\sigma}^{2}})}=\exp{[-\frac{1}{2\hat{\sigma}^{2}}(\left\lVert\bm{\theta}_{i}-\bm{\hat{}}{\mu}\right\rVert^{2})]},$
(40)

where $\bm{\hat{}}{\mu}$ and $\hat{\sigma}^{2}$ denote the mean and variance
of the posterior Gaussian. By completing the square, we first match the
coefficients of $\bm{\theta}_{i}^{T}\bm{\theta}_{i}$—second power of
$\bm{\theta}_{i}$, and find $\hat{\sigma}^{2}$ is given by

$\displaystyle\frac{1}{\hat{\sigma}^{2}}$
$\displaystyle=\frac{1}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{1}{\sigma_{s}^{2}},$
(41) $\displaystyle\hat{\sigma}^{2}$
$\displaystyle=\frac{1}{\frac{1}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{1}{\sigma_{s}^{2}}},$
(42)

and then match the coefficients of first power $\bm{\theta}_{i}^{T}$ we get

$\displaystyle\frac{\bm{\hat{}}{\mu}}{\hat{\sigma}^{2}}$
$\displaystyle=\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}},$
(43)

hence

$\displaystyle\bm{\hat{}}{\mu}=\hat{\sigma}^{2}(\frac{\bm{\mu}_{0}}{\sigma_{0}^{2}}+\sum_{s\in{-s^{\prime}},j}A_{i,j}^{s}\frac{\bm{w}_{sj}}{\sigma_{s}^{2}}).$
(44)

∎

##  Appendix D Appendix D: Iteratively Layer-wise Matched Aggregation via NAFI

As the empirical study in (Wang et al. 2020) demonstrates, directly applying
the matching algorithms fails on deep architectures which are necessary to
solve more complex tasks. Thus to alleviate this problem , we also extend NAFI
to the following layer-wise matching scheme. Firstly, the server only collects
the first layer weights from the clients and applies NAFI to acquire the first
layer weights of the federated model. Then, the server broadcasts these
weights to the clients, which proceed to train all consecutive layers on their
datasets while keeping the matched layers frozen. Repeat this process until
the last layer, where we make a weighted average based on class proportions of
each client’s data points. We summarize our layer-wise version of NAFI in
Algorithm 2. The layer-wise approach requires communication rounds that equal
to the number of layers in a neural network. Experimental results show that
with layer-wise matching, NAFI performs well on the ConvNets even for U-nets
which has a complex architecture. In the more challenging heterogeneous
setting, NAFI outperforms FedAvg, FedProx trained with same number of
communication rounds (5 for ConvNet and 19 for U-net).

Algorithm 2 Iteratively Layer-wise NAFI

0: Collected local weights of $N$-layer architectures
$\\{\bm{W}_{s}^{(0)},\cdots,\bm{W}_{s}^{(N-1)}\\}_{s=1}^{S}$ from $S$ clients;

0: New constructed global weights $\\{\bm{W}^{(0)},\cdots,\bm{W}^{(N-1)}\\}$.

1: $n=0$;

2: while layers $n\leq N$ do

3: if n ¡ N-1 then

4:
$\\{A_{s}\\}_{s=1}^{S}=\operatorname{NAFI}(\\{\bm{W}_{s}^{(n)}\\}_{s=1}^{S})$;

5: $\bm{W}^{(n)}=\frac{1}{S}\sum_{s}\bm{W}_{s}^{(n)}A^{T}_{s}$;

6: else

7: $\bm{W}^{(n)}=\sum_{s}\bm{p}_{s}\bm{\cdot}\bm{W}^{(n)}_{s}$ where
$\bm{p}_{s}$ is vector of fraction of data points with each label on worker
$s$, and $\bm{\cdot}$ denotes the dot product;

8: end if

9: for $s\in\\{1,\cdots,S\\}$ do

10: $\bm{W}_{s}^{(n+1)}=A_{s}\bm{W}_{s}^{(n+1)}$;

11: Train $\\{\bm{W}_{s}^{(n+1)},\cdots,\bm{W}_{s}^{(N-1)}\\}$ with
$\bm{W}_{s}^{(n)}$ frozen;

12: end for

13: $n=n+1$

14: end while

##  Appendix E Appendix E: Extending federated matching algorithm to Batch
Normalization Layer

Although (Wang et al. 2020) shows how to apply the PFNM to CNNs, it doesn’t
enable additional deep learning building blocks, e.g., batch normalization
layer, to the matching algorithm. However, widely used deep CNNs such as U-net
often contain batch normalization layer in their architectures. In this paper,
we utilize a common setup which merges the batch normalization layer with a
preceding convolution to incorporate GI-FNM with batch normalization layer.

Without loss of generality, we assume that the feature map size is the same as
the filter size. Let $W_{\rm{conv}}\in\mathbb{R}^{C^{out}\times(C^{in}\cdot
w\cdot h)}$ and $b_{\rm{conv}}\in\mathbb{R}^{C^{out}}$ be the parameters of
the convolutional layer that precedes batch normalization. Given a input data
$x\in\mathbb{R}^{(C^{in}\cdot w\cdot h)}$, the convolutional operator can be
simply expressed as:

$f=W_{\rm{conv}}*x+b_{\rm{conv}}.$ (45)

Batch normalization (BN) is a popular method used in modern neural networks as
it often reduces training time and potentially improves generalization. Given
the outputting feature of preceding convolutional layer, it can be normalized
as follows:

$\hat{f}=\gamma\frac{f-\mu}{\sqrt{\sigma^{2}+\epsilon}}+\beta,$ (46)

where $\mu$ and $\sigma^{2}$ are the mean and variance computed over a batch
of feature, $\epsilon$ is a small constant included for numerical stability,
$\gamma$ is the scaling factor and $\beta$ the shift factor. The parameters
$\gamma$ and $\beta$ are slowly learned with gradient descent together with
the other parameters of the network.

If we take $W_{\rm{conv}}$ and $b_{\rm{conv}}$ into the Eq. (46) we can get
the new weights and bias as:

  * •

weights:
$W_{\rm{BN}}=\gamma\cdot\frac{W_{\rm{conv}}}{\sqrt{\sigma^{2}+\epsilon}}$;

  * •

bias:
$b_{\rm{BN}}=\gamma\cdot\frac{(b_{\rm{conv}}-\mu)}{\sqrt{\sigma^{2}+\epsilon}}+\beta$.

Thus the batch normalized feature can be directly obtained by:

$\hat{f}=W_{\rm{BN}}*x+b_{\rm{BN}}.$ (47)

By matching the fused weights $W_{\rm{BN}}$ and bias $b_{\rm{BN}}$, we enable
the batch normalization layer in GI-FNM.