# One-Shot Federated Learning with Bayesian Pseudocoresets

Tim d’Hondt
Eindhoven University of Technology
<EMAIL_ADDRESS>
&Mykola Pechenizkiy
Eindhoven University of Technology
<EMAIL_ADDRESS>
&Robert Peharz
Graz University of Technology
<EMAIL_ADDRESS>
###### Abstract

Optimization-based techniques for federated learning (FL) often come with
prohibitive communication cost, as high dimensional model parameters need to
be communicated repeatedly between server and clients. In this paper, we
follow a Bayesian approach allowing to perform FL with one-shot communication,
by solving the global inference problem as a product of local client
posteriors. For models with multi-modal likelihoods, such as neural networks,
a naive application of this scheme is hampered, since clients will capture
different posterior modes, causing a destructive collapse of the posterior on
the server side. Consequently, we explore approximate inference in the
function-space representation of client posteriors, hence suffering less or
not at all from multi-modality. We show that distributed function-space
inference is tightly related to learning Bayesian pseudocoresets and develop a
tractable Bayesian FL algorithm on this insight. We show that this approach
achieves prediction performance competitive to state-of-the-art while showing
a striking reduction in communication cost of up to two orders of magnitude.
Moreover, due to its Bayesian nature, our method also delivers well-calibrated
uncertainty estimates.

##  1 Introduction

Federated learning (FL) is a machine learning setting where multiple data-
owning entities (clients) collaborate on solving a global learning problem
(usually coordinated by a server) without sharing their raw data [24]. For
large neural networks, gradient-based federated optimization (FO) methods are
the de-facto standard class of algorithms [52]. These methods minimize a loss
function over partitioned data using a variation of distributed gradient
descent, where in each step (communication round) the server shares the
current global model parameters and clients return the gradient of their local
likelihood function. For models requiring thousands of server gradient steps
until convergence, the _communication cost_ of FO can become prohibitive, as
it scales proportionally to the number of communication rounds, model size and
number of clients. A popular heuristic to reduce the amount of communication
is federated averaging (FedAvg) [34], i.e. to perform several local gradient
updates per communication round before returning a _pseudo-gradient_ , i.e.
the difference between their updated parameters and the current server
parameters [34, 42]. While empirically effective, this heuristic is known to
cause client drift, a phenomenon where the server model converges to a non-
optimal solution due to the biased pseudo-gradients [27, 10].

Alternatives to optimization-based techniques are Bayesian approaches to FL,
formulating learning as a distributed inference problem, where the global
server posterior over parameters is proportional to the model prior and the
product of local client likelihoods [2, 3]. Besides delivering model
uncertainties, Bayesian inference also allows—in principle—a _one-shot
approach_ to FL, i.e. using only a single communication round. Specifically,
each client approximates its local model posterior, sends it to the server,
which subsequently aggregates these ’local factors’ into a global approximate
posterior. If the client posteriors are exact, this approach will also deliver
the exact global posterior, i.e. the same posterior as if all client data had
been combined into a single dataset. Hence, Bayesian FL does in principle not
suffer from systematic biases such as client drift.

However, when dealing with modern neural networks, a central problem of
Bayesian FL is that likelihood functions are highly non-linear and hard to
approximate. Moreover, it is well known that the parameter space of neural
nets has many symmetries, such as the ones generated by exchanging two neurons
in the same layer. These symmetries result in non-identifiability of the model
and a huge number of equivalent local minima (in an optimization point of
view) or posterior modes (in a Bayesian view). Meanwhile, existing Bayesian FL
methods typically restrict local factors to simple, unimodal functions, often
from the exponential family (e.g [3, 18, 2]). Attempting one-shot Bayesian FL
with these techniques will almost certainly fail, since each local factor will
capture a mode in a different region of the parameter space, causing the
global posterior density to collapse to zero where support is not shared over
all clients [3, 12]. Therefore, these methods still require multi-round
algorithms combined with heavy damping to periodically align the supports of
local factors in the parameter space, again increasing the communication
burden. Also switching to more expressive posterior representation, such as
generative models, seems futile due to the sheer number of equivalent modes in
the parameter posterior.

Function Space FL. The core obstacle to one-shot FL illustrated above is non-
identifiability of model parameters, or in other words, that the relationship
between model parameters and represented functions is ’many-to-one’. The
problem disappears when we adopt a _function-space_ view, that is, when we
treat the Bayesian neural network (BNN) as a stochastic process [6, 44]. In
function space, the posterior will concentrate on a single function in the
infinite data limit (when restricting to the support of the data), i.e., the
model is identifiable in function-space, except for pathological cases. Note
that formally the function-space posterior is given by the push-forward of the
weight-space posterior, reducing all equivalent “symmetric” modes into one
mode in function-space. Thus, working in function space is an amenable way to
perform communication-efficient one-shot Bayesian FL.

Bayesian Pseudocoresets for FL. While the function-space view illustrated
above is elegant, representing and communicating the local function-space
posteriors is non-trivial, especially for neural networks. We propose that
each client learns and shares a small set of key function values that serve as
an approximate sufficient statistic of their data, similar to inducing points
in sparse Gaussian processes [46, 6]. When combined with the functional prior,
these function values induce a posterior that resembles the true local
posterior in function space and automatically captures all equivalent
symmetrical modes in weight space.

We show that learning inducing points in the function space of neural networks
is tightly related to learning a _Bayesian pseudocoreset_ (BPC) [32, 28]. In a
nutshell, a BPC is a collection of synthetic data points whose _weight_
posterior approximates the weight of the entire dataset. Their interpretation
as (approximate) inducing points in function space inference is perhaps not
surprising, but, to the best of our knowledge, has not been deeply studied
yet. This connection motivates to leverage recent advances in BPC learning for
a tractable Bayesian FL algorithm in function space, using one-shot
communication between clients and server. Specifically, each client learns a
BPC on their private data, which is then sent to the server. The server then
trains a model on the combined coresets of all clients, either in a Bayesian
manner (e.g. using Monte Carlo) or via standard gradient-based optimization.

The result of this simple idea is striking, as illustrated in Figure 1,
showing performance as a function of communication for our method and various
competitors, all run on the same neural architecture and the EMNIST-62
benchmark [7]. After just a single communication round, needing less than $7M$
single floats communicated, our method achieves $72\%$ classification
accuracy. This is competitive to the best comparison method, which achieves
ultimately slightly less than $80\%$, but _requires hundreds of communication
rounds._ In total, the most communication-efficient competitor required at
least $119M$ single floats communicated to match our performance of $72\%$.
_Thus, our method is between one and two orders of magnitude more
communication efficient than state-of-the art._

Figure 1: Results on EMNIST-62. Left column shows accuracy (y-axis) vs.
communication in float32s (x-axis). Right column shows negative log likelihood
(y-axis) vs communication in float32s (x-axis) Upper row shows baselines with
10 local steps, bottom row with 100 local steps. Note our one-shot approach
(black $\bullet$) in the upper left corner and our one-shot approach as FedAvg
initialization in red.

Moreover, when using our solution as initialization for FedAvg, the _resulting
method dominates all other methods for any amount of communication spent._
This advantage is in particular evident in the low communication regime (i.e.
one- or few-shot FL). Besides being orders of magnitude more communication-
efficient than weight-space FL approaches, our approach also has the benefits
of being completely asynchronous, producing epistemic uncertainty estimates
and not suffering from client drift.

In summary, our main contributions are:

  * •

We propose a novel approach for ultra communication-efficient one-shot
learning of BNNs in the FL setting by locally fitting and communicating
inducing points in the function-space representation.

  * •

We show how this function-space inference approach is well approximated by
distributed learning of BPCs.

  * •

We propose BPC-FL, a first tractable algorithm for performing function-space
inference in the FL setting.

  * •

In empirical evaluations we demonstrate that our method is competitive to
state-or-the-art FL approaches, while using _up to two orders of magnitude
less communication,_ and produces models with favorable uncertainty estimates.

##  2 Preliminaries

We consider conditional probabilistic models for a random vector
${\mathbf{y}}\in\mathcal{Y}$, conditioned on
${\mathbf{x}}\in\mathcal{X}\subseteq\mathbb{R}^{D}$, where
$\mathcal{Y}\subseteq\\{0,1\\}^{C}$ for classification and
$\mathcal{Y}\subseteq\mathbb{R}^{C}$ for regression tasks. Let
${\mathbf{\Theta}}\subseteq\mathbb{R}^{P}$ be a random vector containing the
model parameters of a Bayesian Neural Network (BNN) with an architecture
defined by $f:\mathcal{X}\times\mathbb{R}^{P}\rightarrow\mathbb{R}^{C}$.
Hence, $f$ is a random function that for a specific parameter realization
${\bm{\theta}}$ of ${\mathbf{\Theta}}$ encodes a mapping from $\mathcal{X}$ to
a set of $C$ logits. Given a dataset
$\mathcal{D}=\\{y_{n},x_{n}\\}^{N}_{n=1}=({\bm{X}},{\bm{y}})$ of $N$
observations of ${\mathbf{y}}$ and ${\mathbf{x}}$, we assume a likelihood
$p({\bm{y}}|f({\bm{X}},{\bm{\theta}}))=\prod^{N}_{n=1}p(y_{n}|f(x_{n},{\bm{\theta}}))$.
Letting $p_{0}$ be a prior density over ${\mathbf{\Theta}}$, the posterior
density for ${\mathbf{\Theta}}={\bm{\theta}}$ is given as:

$p({\bm{\theta}}|\mathcal{D})\propto
p_{0}({\bm{\theta}})p({\bm{y}}|f({\bm{X}},{\bm{\theta}}))$ (1)

which has an intractable normalization constant
$Z(\mathcal{D})=\int_{\mathbb{R}^{P}}p_{0}({\bm{\theta}})p({\bm{y}}|f({\bm{X}},{\bm{\theta}}))d{\bm{\theta}}$
due to $f$ being a nonlinear function in ${\bm{\theta}}$.

In the federated learning (FL) setting, we assume the dataset $\mathcal{D}$ is
partitioned over $M$ data-generating entities (clients):
$\mathcal{D}=\mathcal{D}_{1}\cup...\cup\mathcal{D}_{M}$ that collaborate in
learning a model under the supervision of a server, while any function of a
client’s dataset $\mathcal{D}_{m}$ of size $n_{m}$ can only be evaluated
locally. These local datasets $\mathcal{D}_{m}$ may be non-iid in the sense
that each $\mathcal{D}_{m}$ is sampled from a (slightly) different
distribution.

###  2.1 Federated Posterior Inference

As described in the Section 1, the default FL algorithm is federated
optimization (FO), which is a version of distributed gradient descent and
communication hungry. Federated averaging (FedAvg) [34] aims to save on
communication, but the introduction of pseudo-gradients also introduces a
systematic bias denoted as client drift [27, 10].

Alternatively, one might formulate FL as Bayesian inference problem. The
federated posterior inference problem (FPI) is the task of computing the
following posterior within the restrictions of the FL setting:

$p({\bm{\theta}}|\mathcal{D})\propto
p({\bm{\theta}})\prod^{M}_{m=1}p({\bm{y}}_{m}|f({\bm{X}}_{m},{\bm{\theta}})),$
(2)

where $p({\bm{y}}_{m}|f({\bm{X}}_{m},{\bm{\theta}}))$ is the $m^{\text{th}}$
client likelihood. If the client likelihoods (or equivalently the local
posteriors divided by the prior density) could be communicated exactly, it is
evident that (2) is identical to (1), hence, in principle FPI does not
introduce a systematic bias such as client drift. Even when using approximate
likelihoods, a certain robustness of FPI regarding client drift has been
empirically observed [2].

Since the likelihoods can only be evaluated locally, the usual approach is to
compute $p({\bm{\theta}}|\mathcal{D})$ with a message passing algorithm, which
infers a target density using a collection of localized inferences with an
approximate density $q({\bm{\theta}})\propto
p({\bm{\theta}})\prod^{M}_{m=1}t_{m}({\bm{\theta}}),$ where
$t_{m}({\bm{\theta}})$ are called the local factors, or site parameters [39,
50]. The local factors are then typically fitted to minimize the reverse KL
divergence [5] or its forward substitute [20, 50, 18]. Both objectives have
been applied to the FPI problem [18, 3], with the factors
$t_{m}({\bm{\theta}})$ restricted to be in the exponential family.

Since these client factors $t_{m}({\bm{\theta}})$ are implemented with a
unimodal approximation, they might been drawn to different modes of the
likelihood in local inference, especially in the case of non-identifiable
models such as neural networks. With high probability, this will cause
catastropic failures when aggregating the factors at the server as
$q({\bm{\theta}})\propto p({\bm{\theta}})\prod^{M}_{m=1}t_{m}({\bm{\theta}})$.
Thus, $q({\bm{\theta}})\approx 0$ at any ${\bm{\theta}}$ where even just one
factor has close to zero probability [12]. Therefore, these methods still
require many communication rounds and heavy damping of the server updates to
converge, especially when performing local updates at clients in parallel [3,
18].

###  2.2 Function space representation of BNNs

The posterior collapse problem described above occurs in non-identifiable
models, such as BNNs. A natural remedy is to switch to a ’function space’
view, by regarding the function $f(.,{\mathbf{\Theta}})$ as a stochastic
process on the domain $\mathcal{X}$. Intuitively, $f(.,{\mathbf{\Theta}})$ is
an infinite dimensional random vector indexed by $\mathcal{X}$ that for a
particular parameter instantiation ${\bm{\theta}}$ encodes to what logit
values a NN architecture maps all possible elements of $\mathcal{X}$. Let
${\mathbf{f}}$ denote the random vector $f(.,{\mathbf{\Theta}})$ and
${\mathbf{f}}_{X}$ denote that vector on a subset of the domain
$X\subseteq\mathcal{X}$. For brevity, we will also denote ${\mathbf{f}}$ its
value for a specific parameter ${\bm{\theta}}$ as
${\mathbf{f}}^{\bm{\theta}}$, such that
$f({\bm{X}},{\bm{\theta}})={\mathbf{f}}^{\bm{\theta}}_{\bm{X}}$. Given a prior
probability density $p_{0}$ over ${\mathbf{\Theta}}$, we can define a
function-space prior and posterior distribution as follows [54, 45]:

$\displaystyle p_{0}({\mathbf{f}})$
$\displaystyle=\int_{\mathbb{R}^{P}}p_{0}({\bm{\theta}})\delta({\mathbf{f}}-{\mathbf{f}}^{{\bm{\theta}}})d{\bm{\theta}}$
(3) $\displaystyle p({\mathbf{f}}|\mathcal{D})$
$\displaystyle=\int_{\mathbb{R}^{P}}p({\bm{\theta}}|\mathcal{D})\delta({\mathbf{f}}-{\mathbf{f}}^{{\bm{\theta}}})d{\bm{\theta}}$
(4)

where $\delta$ is the Dirac delta function. These densities can also be
defined more generally as pushforward measures through the tools of measure
theory [6], but we will stick to the definition above, without loss of
generality. It is noticeworthy that the KL divergence between two
distributions over functions generated from different distributions over
parameters applied to the same neural network architecture is well-defined if
the KL divergence between the distributions over parameters is finite, by the
strong data processing inequality [40]:

$\mathbb{D}_{KL}(p({\mathbf{f}}^{\bm{\theta}})||q({\mathbf{f}}^{{\bm{\theta}}}))\leq\mathbb{D}_{KL}(p({\bm{\theta}})||q({\bm{\theta}}))$
(5)

As a consequence, it is possible to perform variational inference in the
function space of neural networks instead of their parameter space [47, 6,
45].

###  2.3 Bayesian Pseudocoresets

Computing or even estimating the BNN posterior (1) is in general
computationally expensive, in particular for large $N$. _Bayesian
pseudocoresets_ (BPCs) introduce the idea to replace the real data with a
small synthetic dataset
$\mathcal{C}=\\{{\bm{z}}_{k},{\bm{u}}_{k}\\}^{K}_{k=1}=({\bm{Z}},{\bm{u}})$ of
pseudolabels ${\bm{u}}_{k}\in\mathcal{Y}$ and pseudoinputs
${\bm{z}}_{k}\in\mathcal{X}$ where $K\ll N$, which produces the coreset
posterior [31]:

$\displaystyle q({\bm{\theta}}|\mathcal{C})$ $\displaystyle\propto
p_{0}({\bm{\theta}})p({\bm{u}}|f({\bm{Z}},{\bm{\theta}}))$ (6)
$\displaystyle=p_{0}({\bm{\theta}})\exp\left(\sum^{K}_{k=1}-\log
p({\bm{u}}_{k}|f({\bm{z}}_{k},{\bm{\theta}}))\right)$ (7)

Typically, $\mathcal{C}$ is trained to minimize some divergence measure
$\mathbb{D}$ between the coreset posterior $q({\bm{\theta}}|\mathcal{C})$ and
the target posterior $p({\bm{\theta}}|\mathcal{D})$:

$\mathcal{C}^{*}=\textrm{arg
}\underset{\mathcal{C}}{\min}\,\mathbb{D}\left[p({\bm{\theta}}|\mathcal{D}),q({\bm{\theta}}|\mathcal{C})\right]$
(8)

Intuitively, a BPC is a summary of the information about ${\bm{\theta}}$ that
is contained in $\mathcal{D}$, i.e. an approximate sufficient statistic. Thus,
$\mathcal{C}$ compresses $\mathcal{D}$ by exploiting redundancy in the data
with respect to $f$ [53].

BPC approximations of posteriors come with some favorable properties over
other approximate posterior families, such as the fact that they preserve
important model structure that is exhibited in the true posterior (symmetries,
subspace structures, heavy tails) through $f$, which makes it an interesting
choice for modeling posteriors involving complex likelihood functions.
Moreover, BPC posteriors are trivially composable through set operations on
the coresets and are inference algorithm agnostic at downstream tasks [53].
Different choices for the divergence measure $\mathbb{D}$ and approaches to
minimize that measure give rise to different BPC algorithms [28, 29, 32, 31].
BPC learning is also closely related to dataset distillation and dataset
condensation [9, 58] as shown in [28].

Intuitively, since BPCs have a particular reminiscence to inducing points in
Gaussian processes (GPs) [46, 6], one might wonder whether they can be
interpreted as a sparse function-space inference method. In the following
section we show that this indeed the case and exploit them for a one-shot
Bayesian FL algorithm.

##  3 Bayesian One-Shot Federated Learning

###  3.1 Federated Learning in Function-space

For a partitioned dataset over clients
$\mathcal{D}=\mathcal{D}_{1}\cup...\cup\mathcal{D}_{M}$, let us denote
$\mathcal{D}_{m}=({\bm{X}}_{m},{\bm{y}}_{m})$ for each client $m$. Also, let
us split up ${\mathbf{f}}_{\bm{X}}$ into
${\mathbf{f}}_{\bm{X}}=[{\mathbf{f}}_{{\bm{X}}_{1}},...,{\mathbf{f}}_{{\bm{X}}_{M}}]$,
denoting the _latent function values_ indexed by the input values in each
client’s dataset. Further, we denote ${\mathbf{f}}_{*}$ to be the latent
function values at all other inputs in the domain $\mathcal{X}$. We can then
write the federated posterior from equation (2) in its function space
representation as follows:

$\displaystyle p({\mathbf{f}}_{*},{\mathbf{f}}_{\bm{X}}|\mathcal{D})$
$\displaystyle=\frac{p_{0}({\mathbf{f}}_{*},{\mathbf{f}}_{\bm{X}})}{p(\mathcal{D})}\prod^{M}_{m=1}p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
(9)
$\displaystyle=\int_{\mathbb{R}^{P}}p({\bm{\theta}}|\mathcal{D})\delta({\mathbf{f}}_{\bm{X}}^{\theta}-{\mathbf{f}}_{\bm{X}})\delta({\mathbf{f}}_{*}^{\theta}-{\mathbf{f}}_{*})d{\bm{\theta}}$

Instead of estimating the weight space posterior
$p({\bm{\theta}}|\mathcal{D})$ within the restrictions of the FL setting, we
will take an approach inspired by inducing points in sparse GPs [41, 46, 48],
which is—under particular modeling assumptions depicted below—equivalent to
learning BPCs.

Let us assume a set of function values
${\mathbf{f}}_{\bm{Z}}=[{\mathbf{f}}_{{\bm{Z}}_{1}},...,{\mathbf{f}}_{{\bm{Z}}_{M}}]$
where ${\bm{Z}}_{m}$ are _pseudo inputs_ belonging to client $m$. In our
model, each variable ${\mathbf{f}}_{{\bm{Z}}_{m}}$ is treated as a _sufficient
statistic_ of the client’s local data function values
${\mathbf{f}}_{{\bm{X}}_{m}}$, that is, we make the prior assumption that
${\mathbf{f}}_{\bm{X}}$ and ${\mathbf{f}}_{*}$ are conditionally independent
given ${\mathbf{f}}_{{\bm{Z}}_{m}}$ and all ${\mathbf{f}}_{{\bm{X}}_{m}}$ are
conditionally independent given their local ${\mathbf{f}}_{{\bm{Z}}_{m}}$:

$p({\mathbf{f}})=p_{0}({\mathbf{f}}_{\bm{Z}})p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{f}}_{{\bm{Z}}_{m}})$
(10)

where
${\mathbf{f}}=[{\mathbf{f}}_{\bm{X}},{\mathbf{f}}_{\bm{Z}},{\mathbf{f}}_{*}]$
for brevity of notation. If the dimensionality $K$ of each
${\mathbf{f}}_{{\bm{Z}}_{m}}$ is significantly smaller than the client’s
dataset size, i.e. $K\ll n_{m}$, then ${\mathbf{f}}_{{\bm{Z}}_{m}}$ is a
compressed representation of the information about
${\mathbf{f}}_{{\bm{X}}_{m}}$ in local dataset $\mathcal{D}_{m}$. In contrast
to sparse GPs, where pseudo inputs ${\bm{Z}}_{m}$ are often fixed, we let
clients optimize these, additionally to ${\mathbf{f}}_{\bm{Z}}$. Including the
likelihood terms, the full joint model over latent ${\mathbf{f}}$ and observed
${\bm{y}}$ is given as:

$p({\mathbf{f}},{\bm{y}})=p({\mathbf{f}})\prod^{M}_{m=1}p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
(11)

Note that in the FL setting, it is difficult for a client $m$ to perform
inference on ${\mathbf{f}}_{{\bm{Z}}_{m}}$ given $\mathcal{D}_{m}$, because
${\mathbf{f}}_{{\bm{Z}}_{m}}$ is dependent on all other
${\mathbf{f}}_{{\bm{Z}}_{m}}$ and thus other client data through prior
$p_{0}({\mathbf{f}}_{Z})$. One solution would be to iteratively refine the
posterior density over each ${\mathbf{f}}_{{\bm{Z}}_{m}}$ in an EP-type
inference scheme. In order to pursue the goal of one-shot FL, however, we make
the assumption that all ${\mathbf{f}}_{{\bm{Z}}_{m}}$ are independent:
$p_{0}({\mathbf{f}}_{\bm{Z}})=\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})$,
resulting in the following posterior (see Appendix A.1 for a derivation):

$\displaystyle
p({\mathbf{f}}|\mathcal{D})=p({\mathbf{f}}_{*}|{\mathbf{f}}_{{\bm{Z}}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\bm{y}}_{m})$
(12)

Through our assumptions, we have now divided the problem of computing the
global posterior $p({\mathbf{f}}|\mathcal{D})$ into a set of $M$ independent
local posterior computations of
$p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{y}}_{m})$,
which can then be communicated to a server and used to compute the predictive
distribution $p({\mathbf{f}}_{*}|{\mathbf{f}}_{{\bm{Z}}})$ in downstream
tasks.

###  3.2 Approximate Inference

The next question is: how do we locally compute and communicate the
intractable posteriors
$p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{y}}_{m})$?
We choose to fit an approximate posterior
$q({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}})=q({\mathbf{f}}_{{\bm{Z}}_{m}})p({\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{f}}_{{\bm{Z}}_{m}})$
at each client in parallel by minimizing the forward KL divergence
$\mathbb{D}_{KL}\left[p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\bm{y}}_{m})||q({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}})\right]$
with respect to $q$. To be able to apply gradient-based solvers, $q$ is
required to be a differentiable, parametric density function.

Another important insight is that we want to incorporate the functional prior
$p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})$ into our definition of $q$, as this
allows us to model the complex dependencies between all function values in
${\mathbf{f}}_{{\bm{Z}}_{m}}$ through the model architecture. Following this
insight, we choose $q({\mathbf{f}}_{{\bm{Z}}_{m}})$ to be a posterior given
some learnable pseudo-labels, such that
$q({\mathbf{f}}_{{\bm{Z}}_{m}})=p({\mathbf{f}}_{{\bm{Z}}_{m}}|\hat{{\bm{y}}}_{m})\propto
p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})p(\hat{{\bm{y}}}_{m}|{\mathbf{f}}_{{\bm{Z}}_{m}})$.
The client local objective then looks as follows by following the proof from
Proposition 3.1 in [29], which is based on [33] and ignoring terms that are
independent of ${\bm{Z}}_{m}$ and $\hat{{\bm{y}}}_{m}$:

$\displaystyle\mathcal{L}_{m}({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$
$\displaystyle=\mathbb{D}_{KL}\left[p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\bm{y}}_{m})||q({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}})\right]$
$\displaystyle=\mathbb{E}_{p({\mathbf{f}}_{{\bm{Z}}_{m}}|{\bm{y}}_{m})}\left[\log\frac{p(\hat{{\bm{y}}}_{m}|{\mathbf{f}}_{{\bm{Z}}_{m}})}{p(\hat{{\bm{y}}}_{m}|{\bm{Z}}_{m})}\right]$
(13)

Where $p(\hat{{\bm{y}}}_{m}|{\bm{Z}}_{m})$ is the marginal likelihood of
$(\hat{{\bm{y}}},{\bm{Z}}_{m})$. Notice from equation 2 that we can sample
from $p({\mathbf{f}}_{{\bm{Z}}_{m}}|{\bm{y}}_{m})$ using the Law of the
Unconscious Statistician, i.e monte carlo sampling $\tilde{{\bm{\theta}}}\sim
p({\bm{\theta}}|\mathcal{D})$ and mapping it into the function
${\mathbf{f}}^{\tilde{{\bm{\theta}}}}_{{\bm{Z}}_{m}}=f({\bm{Z}},\tilde{{\bm{\theta}}})$
through the network architecture. The objective then simplifies to
$\mathcal{L}_{m}({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})=\mathbb{E}_{p({\bm{\theta}}|\mathcal{D}_{m})}\left[\log\frac{p(\hat{{\bm{y}}}_{m}|{\mathbf{f}}^{\bm{\theta}}_{{\bm{Z}}_{m}})}{p(\hat{{\bm{y}}}_{m}|{\bm{Z}}_{m})}\right]$,
which is the BPC objective when choosing $\mathbb{D}$ to be the forward KL
divergence in equation 8.

Thus, our parameterization of $q({\mathbf{f}}_{{\bm{Z}}_{m}})$ via pseudo-
labels $\hat{{\bm{y}}}_{m}$, can be seen as approximate function space
inference for learning a BPC $({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$. This allows
us to leverage recent advances in BPC learning to tractably minimize (3.2) for
large neural networks. This ’pseudo-label’ parameterization of an approximate
posterior over function space inducing points
$q({\mathbf{f}}_{{\bm{Z}}_{m}})\propto
p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})p(\hat{{\bm{y}}}_{m}|{\mathbf{f}}_{{\bm{Z}}_{m}})$
has also been used in the Sparse Variational Gaussian Process (SVGP)
literature [1].

Our motivation for using the forward KL divergence is that this divergence
forces $q({\mathbf{f}}_{{\bm{Z}}_{m}})$ to have a broad support over all
functions that are probable under
$p({\mathbf{f}}_{{\bm{Z}}_{m}}|\mathcal{D}_{m})$, which is important for
server aggregation and offsets some of the bias that we introduced by assuming
an independent prior $p_{0}({\mathbf{f}}_{\bm{Z}})$. Moreover, contrary to the
reverse KL divergence, choosing the forward KL divergence prevents us from
having to backpropagate through a Monte Carlo sampling process that estimates
the expectation in equation 3.2, which makes the procedure more scalable [28].

###  3.3 BPC-FL

In this section, we propose a first FL algorithm that solves the defined
inference problem, which we will call BPC-FL. The general idea is simple and
summarized in Algorithm 1: In parallel, each client $m$ learns his BPC
$({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$ by minimizing equation (3.2) and finally
communicating $q({\mathbf{f}}_{{\bm{Z}}_{m}})$ represented as
$({\bm{Z}}_{m},\hat{{\mathbf{y}}}_{m})$ to the server. The server then simply
concatenates the BPCs as ${\bm{Z}}=[{\bm{Z}}_{1},...,{\bm{Z}}_{M}]$ and
$\hat{{\mathbf{y}}}=[\hat{{\bm{y}}}_{1},...,\hat{{\bm{y}}}_{M}]$ to obtain
$q({\mathbf{f}}_{\bm{Z}})=p({\mathbf{f}}_{\bm{Z}}|\hat{{\bm{y}}})$, which can
then be used to compute $p({\mathbf{f}}_{*}|{\bm{f}}_{\bm{Z}})$ in downstream
tasks.

Learning the local pseudodata Having shown the connection between BPC learning
and equation 3.2, we apply the BPC-fKL algorithm from [28] to minimize
equation 3.2 locally at each client $m$ with respect to his pseudodata, which
we denote in short as $\mathcal{C}_{m}=({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$.
This choice of BPC inference algorithm is orthogonal to our general one-shot
FL approach, so performance of our proposed method is expected to increase if
new, better BPC inference algorithms are proposed over time.

The gradient of the local objective $\mathcal{L}_{m}$ with respect to
$\mathcal{C}_{m}=({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$ is given as follows [28]:

$\nabla_{\mathcal{C}}\mathcal{L}_{m}=\mathbb{E}_{p({\bm{\theta}}|\mathcal{D}_{m})}\left[\nabla_{\mathcal{C}}\log
p(\hat{{\bm{y}}}_{m}|f({\bm{Z}}_{m},{\bm{\theta}}))\right]+\mathbb{E}_{p({\bm{\theta}}|\mathcal{C}_{m})}\left[-\nabla_{\mathcal{C}}\log
p(\hat{{\bm{y}}}_{m}|f({\bm{Z}}_{m},{\bm{\theta}}))\right]$ (14)

and can then be plugged into any gradient-based solver of choice. The
expectations over $p({\bm{\theta}}|\mathcal{D}_{m})$ and
$p({\bm{\theta}}|\mathcal{C}_{m})$ are estimated with short-chain MCMC
sampling, to make the gradient estimator similar to the one proposed in the
Contrastive Divergence algorithm [21]. The BPC-fKL algorithm applies a number
of tricks to the estimation of (14) to improve performance and reduce
computation and memory cost. Most importantly, it pretrains MAP optimization
trajectories on $p({\bm{\theta}}|\mathcal{D})$ and saves checkpoints of
weights to disk, which can be loaded during training as initialization points
for the MCMC sampler [9]. We share the initialization seeds for these
trajectories among clients to ensure similar spread of sampling from
$p({\bm{\theta}}|\mathcal{D})$ throughout the parameter space. For
completeness, we provide the pseudocode for BPC-fKL in Appendix B.

Initialization of ${\bm{Z}}$ and $\hat{{\bm{y}}}_{m}$ We initialize each
${\bm{z}}\in{\bm{Z}}_{m}$ randomly from a Gaussian centered at the element-
wise mean $\bar{{\bm{X}}}_{m}$ of a client $m$’s data:
${\bm{z}}\sim\mathcal{N}(\bar{{\bm{X}}}_{m},\sigma_{{\bm{z}}})\hskip
5.69054pt\forall{\bm{z}}\in{\bm{Z}}_{m}$, with $\sigma_{{\bm{z}}}$ being a
hyperparameter to determine the initialization noise. For classification
tasks, we initialize pseudolabels $\hat{{\bm{y}}}_{m}$ in order by cycling
through the unique labels in a client’s dataset, sorted by their occurence.
For regression tasks, we initialize $\hat{{\bm{y}}}_{m}$ from the dataset
average output ${\bm{y}}_{m}$.

Server aggregation Constructing the server posterior
$q({\mathbf{f}}_{\bm{Z}})=\prod^{M}_{m=1}q({\mathbf{f}}_{{\bm{Z}}_{m}})$ is as
simple as concatenating all received ${\bm{Z}}_{m}$ and ${\bm{y}}_{m}$
vectors. In experiments, we found that it is helpful to weigh each
$q({\mathbf{f}}_{{\bm{Z}}_{m}})$ proportional to $\frac{n_{m}}{N}$, to put
more emphasis on $q({\mathbf{f}}_{{\bm{Z}}_{m}})$ that carry information from
more data.

Algorithm 1 BPC-FL

1: Input: local data $\mathcal{D}_{1},...,\mathcal{D}_{M}$

2: for $m\in\\{1,...,M\\}$ in parallel do

3: $({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})\leftarrow$ BPC-fKL$(\mathcal{D}_{m})$
[28]

4: Send $({\bm{Z}}_{m},\hat{{\bm{y}}}_{m})$ to server

5: end for

6:
${\bm{Z}}\leftarrow[{\bm{Z}}_{1},...,{\bm{Z}}_{M}],\hat{{\bm{y}}}\leftarrow[\hat{{\bm{y}}}_{1},...,\hat{{\bm{y}}}_{M}]$

7: Return $q({\mathbf{f}}_{\bm{Z}})$

###  3.4 Downstream inference tasks

After completing Algorithm 1, the server has access to $q({\bm{f}}_{\bm{Z}})$,
which provides end users with more flexibility in what downstream tasks we can
tackle than standard FL approaches.

One of the most important downstream tasks will be to predict a set of unseen
test datapoints ${\bm{X}}_{*}$. This can be done by estimating the predictive
distribution
$\mathbb{E}_{q({\bm{f}}_{\bm{Z}})}({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})=\int_{\mathbb{R}^{P}}p({\bm{\theta}}|\mathcal{C})\delta({\mathbf{f}}_{*}^{\bm{\theta}}-{\mathbf{f}}_{*})d{\bm{\theta}}$.
We have total freedom in whether we want to obtain a MAP estimate from
$p({\bm{\theta}}|\mathcal{C})$ with for example SGD or whether we want to
obtain samples from it using Hamiltonian Monte Carlo.

Thinking a little more creative, BPC-FL could also be used to enhance existing
FL methods. For example, we could find a mode of
$q({\bm{\theta}}|\mathcal{C})$ and use it as an initializer for an FO
algorithm as a ’jump start’ to save costly communication rounds, something
which we show in our experimental results in Figure 1.

##  4 Related work

Formulating the goal of federated learning as a Bayesian inference problem as
in equation (2) was first proposed in [2], where the authors estimated
Gaussian local factor approximations to un-bias client pseudo-gradients used
in an FO algorithm, but did not learn a posterior on the server side. Solving
a distributed posterior inference problem with a product of local factors has
a long history in Expectation Propagation (EP) [50, 20], Bayesian Committee
Machines (BCM) [49] and embarrassingly parallel VI [37, 38], but was first
mentioned in the context of FL for Partitioned Variational Inference (PVI)
[3]. [18] then adapted PVI to instead minimize the forward KL divergence. All
of these works use a unimodal approximation for the local factors, which fails
for one-shot inference as mentioned in [3, 12]. An alternative approach is to
estimate the global posterior by collecting samples using distributed MCMC
algorithms fitted to the FL setting [15, 51, 25, 11], though these methods
require many communication rounds, similar to FO algorithms.

Learning a model in the FL setting in one communication round (one shot) has
been proposed through model distillation at the server in [11, 17, 30, 19],
though these works assume a publicly available dataset at the server. As a
response, Zhang et al. [57] replaced the public dataset with a generative
model.

Another approach to one-shot FL is assume a publicly available feature
extractor and learn a classification head by having each client send a
function-space distribution for each class [4].

A number of works exist that achieve one-shot FL with dataset distillation;
also through sharing synthetic datasets with the server [59, 16, 22, 56].
Dataset distillation techniques are actually closely related to Bayesian
Pseudocoresets [28], but typically use optimization instead of MCMC sampling.

Formulating FL as a function space inference problem has been proposed in
[23], though in the context of performing federated variational inference on
the parameter prior in a hierarchical Bayesian model, and with only
preliminary empirical results. In [13], it is proposed to perform the server
aggregation step of an FO algorithm in function space using estimations of
local Fisher information matrices, in order to reduce client drift.

##  5 Experimental Results

In this section we compare the predictive performance and uncertainty
estimates of models produced by our method and models trained by other
federated learning methods. Because communication efficiency is the main
motivation for this paper, we pay special attention to the amount of
communication required to obtain a certain result for each method.

###  5.1 Experimental Setup

Datasets and model architecture Empirical results are obtained from learning
tasks on two synthetic datasets and a large FL benchmark dataset from the LEAF
repository [7]. The synthetic datasets; one regression problem and one 2d
classification problem, were chosen mainly to evaluate the uncertainty
estimates of models and to get an intuitive view on what the pseudodata at the
server looks like. To evaluate predictive performance on relatively large and
complex FL learning tasks, we use the EMNIST-62 dataset from LEAF that is
split across clients as to carry natural non-iid-ness [7, 42]. MLP
architectures are applied to the synthetic dataset problems and a ConvNet
architecture is selected for EMNIST-62 based on the one used in [42]. Details
can be found in Appendix C.1.

Baselines We compare against FedAvg [34] and two FO algorithms that are
designed to prevent client drift: MIME [26] and FedPA [2]. We test these
methods using different amounts of local updates, in order to improve their
communication efficiency, but also to evaluate the expected reduction in
performance due to client drift. For the synthetic datasets, we also compare
against Distilled One Shot Federated Learning (DOSFL); a method that also
promises one-shot federated learning, but using dataset distillation [58].
Because DOSFL requires backpropagation through an entire gradient descent
trajectory, we were not able to compute a baseline for this algorithm on
EMNIST-62 due to memory limitations. Hyperparameters of baseline methods were
set according to their best performing configuration reported in their
respected paper, if available. These and other details can be found in
Appendix C.3.

Evaluation metrics To evaluate the predictive performance of our method, we
apply different algorithms on $p({\bm{\theta}}|\mathcal{C})$ to produce neural
network weights. We use SGD and Adam to compute MAP estimates of
$p({\bm{\theta}}|\mathcal{C})$ and use HMC [36] to produce samples that we use
in a Monte Carlo estimation of the Bayesian predictive
$\mathbb{E}_{p({\bm{\theta}}|\mathcal{C})}[p(y^{\prime}|f(x^{\prime},{\bm{\theta}}))]$
of a test example $(x^{\prime},y^{\prime})$. We also run SGD on
$p({\bm{\theta}}|\mathcal{C})$ and use the resulting MAP estimate as an
initializer for FedAvg, with the goal of having FedAvg converge in less
communication rounds by giving it a ’head start’ through sending $\mathcal{C}$
first. Details and hyperparameters for these downstream samplers and
optimizers are given in Appendix C.4.

For EMNIST-62, we record negative log likelihoods and accuracies on a test set
that is produced by withholding 20 percent of each client’s data from
training. We also record intermediary estimates of these 2 metrics after each
communication round on a random sample of clients, as to visualize convergence
of FO methods. To evaluate the uncertainty calibration of models, we report
the expected calibration error (ECE) [35] for each method. Finally, we keep
track of the communication cost of each method as the amount of single floats
that are communicated from and to a (fictive) server, to evaluate the
communication cost - performance trade off. All metrics are reported are an
average over 5 different seeds.

Implementation details All experiments on the benchmark tasks were conducted
using an adapted version of FedJAX [43]. Author’s implementations were used
for FedAvg and MIME, whereas FedPA and DOSFL were implemented by us.

###  5.2 Main results

Synthetic regression Figure 2 shows the predictive distributions of a model
using weights that are learned by FedAvg, DOSFL and our method BPC-FL on a
synthetically generated regression dataset that was partitioned over 5
clients. Each client has 100 datapoints and learns a pseudodataset of size
$K=6$ in BPC-FL, which are then sent to the server. Concluding from Figure 2,
BPC-FL seems to produce predictive distributions on the inducing points that
generalize well to the entire dataset. This is especially noticeworthy from a
communication cost perspective when considering that FedAvg requires around
100 communication rounds to converge to a good result with 10 local update
steps, whereas BPC-FL only requires 6 one dimensional pseudodatapoints to be
sent to the server in one shot. DOSFL was not tested on a regression task in
the original paper and therefore does not produce good results on this task
for any hyperparameter setting. More results on smaller client datasets can be
found in Appendix C.5, where it is shown that FedAvg starts to overfit, but
BPC-FL keeps producing good results.

Figure 2: Synthetic regression: client data is visualized in a different color
for each client. Upper row presents the produced server models by FedAvg
(left) and DOSFL (right). Lower row shows ${\bm{Z}}$ as black crosses. The
purple line left shows a model learned by running Adam on
$q({\bm{\theta}}|\mathcal{C})$, whereas the purple line and area show the mean
and one standard deviation values for the Bayesian predictive distribution
estimated by HMC samples from $q({\bm{\theta}}|\mathcal{C})$.

2d classification Figure 3 shows the predictive distributions of models
produced by the methods. With a pseudodataset of size only 5, BPC-FL with SGD
produces a decent predictive distribution and BPC-FL with HMC produces a
predictive distribution that estimates the epistemic uncertainty on the full
dataset very well. For points that were intitialized randomly from a Gaussian
distribution, ${\bm{Z}}$ converges to a logical solution. After an extensive
hyperparameter search, the performance for DOSFL is again disappointing.

Figure 3: 2d classification: red-blue filling shows each method’s predictive distribution. On the bottom row, the crosses represent ${\bm{Z}}$ and their color represents $\hat{{\bm{y}}}_{m}$. | Negative log likelihood | Accuracy | ECE
---|---|---|---
Method | Final | Communication | Final | Communication |
FedAvg (10 local steps) | $1.038\pm 0.035$ | $264e^{6}$ | $0.791\pm 0.006$ | $250e^{6}$ | $0.246\pm 0.035$
FedPA (10 local steps) | $1.538\pm 0.083$ | $244e^{6}$ | $0.790\pm 0.006$ | $216e^{6}$ | $0.252\pm 0.031$
MIME (10 local steps) | $2.254\pm 0.113$ | $435e^{6}$ | $0.775\pm 0.010$ | $398e^{6}$ | $0.261\pm 0.033$
FedAvg (100 local steps) | $1.703\pm 0.058$ | $138e^{6}$ | $0.775\pm 0.007$ | $119e^{6}$ | $0.303\pm 0.04$
FedPA (100 local steps) | $4.960\pm 0.259$ | - | $0.768\pm 0.005$ | $150e^{6}$ | $0.264\pm 0.076$
MIME (100 local steps) | $1.471\pm 0.461$ | $1545e^{6}$ | $0.722\pm 0.014$ | $3043e^{6}$ | $0.167\pm 0.083$
BPC-FL + SGD (ours) | $1.203\pm 0.004$ | $6.6e^{6}$ | $0.723\pm 0.003$ | $6.6e^{6}$ | $0.086\pm 0.020$
BPC-FL + Adam (ours) | $1.750\pm 0.001$ | - | $0.696\pm 0.010$ | - | $0.028\pm 0.003$
BPC-FL + HMC (ours) | $1.349\pm 0.011$ | - | $0.675\pm 0.001$ | - | $0.139\pm 0.010$
Table 1: Negative log likelihood (NLL), accuracy and ECE results for
EMNIST-62. For NLL and accuracy, we report both the final metrics and the
average amount of communication (in number of floats) it takes to match the
performance of BPC-FL + SGD. Standard deviation of a metric over different
seeds is also given.

EMNIST-62 For this larger FL benchmark, we mostly focus on evaluating the
trade-off between performance and communication cost. We consider here a
scenario with a pool of 100 training clients and a pseudodataset size of 62.
Metrics are reported over test clients that were not in the training pool.
Considering the final negative log likelihood and accuracy in Table 1, we see
that BPC-FL cannot achieve the final performance of FedAvg with 10 local
steps, but does actually beat all baselines on negative log likelihood when we
increase the amount of local steps, likely because the baseline methods start
to suffer from client drift. Also, we observed in experiments that the test
performance of FO methods started to decrease over time, an indication of
overfitting. While BPC-FL does not achieve state-of-the-art final performance,
we can always run SGD on $p({\bm{\theta}}|\mathcal{C})$, use this as an
initialization for any FO algorithm and still obtain FO algorithm final
performance, though at a fraction of the communication cost, as is shown in
Figure 1 by the red line.

In terms of communication cost though, BPC-FL obtains its competitive
predictive results with orders of magnitude less communication, concluding
from Table 1. FedAvg with 10 local steps requires on average 40 times more
communication to achieve the same performance as BPC-FL with SGD as a
downstream algorithm. This clearly demonstrates that this novel approach of
communicating function-space inducing points has a lot of potential in
reducing the communication cost of FL. Finally, we also see that BPC-FL
produces models that have better calibrated uncertainty estimates, as shown by
the Expected Calibration Error. We provide additional results showing the
influence of the amount of participating clients and the amount of inducing
points on predictive performance in Appendix C.5.

##  6 Conclusions

In this paper, we explored doing FL in a single communication round by
performing local approximate inference on a set of inducing points in the
function space representation of neural networks. We showed empirically that
this approach has potential to decrease communication cost of FL by orders of
magnitude, while still being able to produce server models with good
predictive performance and favorable uncertainty estimates over existing FL
algorithms.

Limitations A clear limitation at this point is the computation and memory
cost of BPC-fKL for learning the client BPCs. It requires a very large amount
of disk space to store the pretrained optimization trajectories and has in our
experience a 10-100 fold longer runtime over e.g FedAvg. Also, with BPC-fKL
being based on Contrastive Divergence [21], it is notoriously instable and
difficult to tune. We expect our method can be significantly improved once the
scientific field of BPC inference (which is still in its infancy) matures and
produces better and more efficient algorithms.

Another limitation is the lack of privacy guarantees or clear compatibility
with contemporary privacy-guaranteeing techniques. Differential privacy
guarantees have been shown to be easily implemented in BPC learning [32]
through using DP-SGD and leveraging the post-processing property of
differential privacy. However, since the fKL-BPC algorithm saves many
intermediate training states to disk and uses these states throughout the
optimization of the BPC, the sensitivity of a BPC to training data is non-
trivial, as is the sensitivity of the final server BPC to an individual or an
entire client dataset. In general, adding privacy guarantees to dataset
distillation (a specific instance of BPC [28]) is an ongoing and controversial
topic [14, 8], and this paper merely serves as a motivating use-case for more
work to be done in this direction. Nonetheless, we consider the proposed
method as an interesting change of perspective on how to approach distributed
inference or federated learning of neural networks with extreme communication
efficiency.

##  7 Acknowledgements

This project was partially supported by KPN SmartTwo and H2020 SmartChange. We
also thank SURF (www.surf.nl) for the support in using the National
Supercomputer Snellius.

## References

  * [1] Vincent Adam, Paul Chang, Mohammad Emtiyaz E Khan, and Arno Solin.  Dual parameterization of sparse variational gaussian processes.  Advances in Neural Information Processing Systems, 34:11474–11486, 2021.
  * [2] Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh.  Federated learning via posterior averaging: A new perspective and practical algorithms.  In International Conference on Learning Representations, 2021.
  * [3] Matthew Ashman, Thang D. Bui, Cuong V. Nguyen, Stratis Markou, Adrian Weller, Siddharth Swaroop, and Richard E. Turner.  Partitioned variational inference: A framework for probabilistic federated learning.  arXiv preprint arXiv:2202.12275, 2022.
  * [4] Mahdi Beitollahi, Alex Bie, Sobhan Hemati, Leo Maxime Brunswic, Xu Li, Xi Chen, and Guojun Zhang.  Parametric feature transfer: One-shot federated learning with foundation models.  arXiv preprint arXiv:2402.01862, 2024.
  * [5] Thang D Bui, Cuong V Nguyen, Siddharth Swaroop, and Richard E Turner.  Partitioned variational inference: A unified framework encompassing federated and continual learning.  arXiv preprint arXiv:1811.11206, 2018.
  * [6] David R. Burt, Sebastian W. Ober, Adrià Garriga-Alonso, and Mark van der Wilk.  Understanding variational inference in function-space.  In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
  * [7] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečnỳ, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar.  Leaf: A benchmark for federated settings.  arXiv preprint arXiv:1812.01097, 2018.
  * [8] Nicholas Carlini, Vitaly Feldman, and Milad Nasr.  No free lunch in” privacy for free: How does dataset condensation help privacy”.  arXiv preprint arXiv:2209.14987, 2022.
  * [9] George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A. Efros, and Jun-Yan Zhu.  Dataset distillation by matching training trajectories.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pages 4750–4759, June 2022.
  * [10] Zachary Charles and Jakub Konečný.  Convergence and accuracy trade-offs in federated learning and meta-learning.  In Arindam Banerjee and Kenji Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 2575–2583. PMLR, 13–15 Apr 2021.
  * [11] Hong-You Chen and Wei-Lun Chao.  Fed{be}: Making bayesian model ensemble applicable to federated learning.  In International Conference on Learning Representations, 2021.
  * [12] Daniel A De Souza, Diego Mesquita, Samuel Kaski, and Luigi Acerbi.  Parallel mcmc without embarrassing failures.  In International Conference on Artificial Intelligence and Statistics, pages 1786–1804. PMLR, 2022.
  * [13] Nikita Dhawan, Nicole Mitchell, Zachary Charles, Zachary Garrett, and Gintare Karolina Dziugaite.  Leveraging function space aggregation for federated learning at scale.  arXiv preprint arXiv:2311.10291, 2023.
  * [14] Tian Dong, Bo Zhao, and Lingjuan Lyu.  Privacy for free: How does dataset condensation help privacy?  In International Conference on Machine Learning, pages 5378–5396. PMLR, 2022.
  * [15] Khaoula El Mekkaoui, Diego Mesquita, Paul Blomstedt, and Samuel Kaski.  Distributed stochastic gradient mcmc for federated learning.  In Conference on Uncertainty in Artificial Intelligence, 2021.
  * [16] Jack Goetz and Ambuj Tewari.  Federated learning via synthetic data, 2020.
  * [17] Neel Guha, Ameet Talwalkar, and Virginia Smith.  One-shot federated learning.  arXiv preprint arXiv:1902.11175, 2019.
  * [18] Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric Xing.  Federated learning as variational inference: A scalable expectation propagation approach.  In The Eleventh International Conference on Learning Representations, 2023.
  * [19] Mohsin Hasan, Guojun Zhang, Kaiyang Guo, Xi Chen, and Pascal Poupart.  Calibrated one round federated learning with bayesian inference in the predictive space.  In Proceedings of the AAAI conference on artificial intelligence, volume 38, pages 12313–12321, 2024.
  * [20] Leonard Hasenclever, Stefan Webb, Thibaut Lienart, Sebastian Vollmer, Balaji Lakshminarayanan, Charles Blundell, and Yee Whye Teh.  Distributed bayesian learning with stochastic natural gradient expectation propagation and the posterior server.  The Journal of Machine Learning Research, 18(1):3744–3780, 2017.
  * [21] Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  Neural computation, 14(8):1771–1800, 2002.
  * [22] Shengyuan Hu, Jack Goetz, Kshitiz Malik, Hongyuan Zhan, Zhe Liu, and Yue Liu.  Fedsynth: Gradient compression via synthetic data in federated learning.  In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022), 2022.
  * [23] Michael Hutchinson, Matthias Reisser, and Christos Louizos.  Federated functional variational inference.  Proceedings of the NeurIPS 2021 worksop on Bayesian Deep Learning, 2021.
  * [24] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D’Oliveira, Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao.  Advances and open problems in federated learning.  Foundations and Trends® in Machine Learning, 14(1–2):1–210, 2021.
  * [25] Avetik Karagulyan and Peter Richtárik.  Elf: Federated langevin algorithms with primal, dual and bidirectional compression.  arXiv preprint arXiv:2303.04622, 2023.
  * [26] Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh.  Breaking the centralized barrier for cross-device federated learning.  Advances in Neural Information Processing Systems, 34:28663–28676, 2021.
  * [27] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh.  SCAFFOLD: Stochastic controlled averaging for federated learning.  In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5132–5143. PMLR, 13–18 Jul 2020.
  * [28] Balhae Kim, Jungwon Choi, Seanie Lee, Yoonho Lee, Jung-Woo Ha, and Juho Lee.  On divergence measures for bayesian pseudocoresets.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 757–767. Curran Associates, Inc., 2022.
  * [29] Balhae Kim, Hyungi Lee, and Juho Lee.  Function space bayesian pseudocoreset for bayesian neural networks.  In ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
  * [30] Qinbin Li, Bingsheng He, and Dawn Song.  Practical one-shot federated learning for cross-silo setting.  arXiv preprint arXiv:2010.01017, 2020.
  * [31] Dionysis Manousakas, Hippolyt Ritter, and Theofanis Karaletsos.  Black-box coreset variational inference.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 34175–34187. Curran Associates, Inc., 2022.
  * [32] Dionysis Manousakas, Zuheng Xu, Cecilia Mascolo, and Trevor Campbell.  Bayesian pseudocoresets.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 14950–14960. Curran Associates, Inc., 2020.
  * [33] Alexander G de G Matthews, James Hensman, Richard Turner, and Zoubin Ghahramani.  On sparse variational methods and the kullback-leibler divergence between stochastic processes.  In Artificial Intelligence and Statistics, pages 231–239. PMLR, 2016.
  * [34] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In Artificial intelligence and statistics, pages 1273–1282. PMLR, 2017.
  * [35] Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht.  Obtaining well calibrated probabilities using bayesian binning.  In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
  * [36] Radford M Neal et al.  Mcmc using hamiltonian dynamics.  Handbook of markov chain monte carlo, 2(11):2, 2011.
  * [37] Willie Neiswanger, Chong Wang, and Eric Xing.  Embarrassingly parallel variational inference in nonconjugate models.  arXiv preprint arXiv:1510.04163, 2015.
  * [38] Willie Neiswanger, Chong Wang, and Eric P Xing.  Asymptotically exact, embarrassingly parallel mcmc.  In Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, pages 623–632, 2014.
  * [39] Judea Pearl.  Fusion, propagation, and structuring in belief networks.  Artificial Intelligence, 29:241–288, 1986.
  * [40] Yury Polyanskiy and Yihong Wu.  Strong data-processing inequalities for channels and bayesian networks.  In Convexity and Concentration, pages 211–249. Springer, 2017.
  * [41] Carl Edward Rasmussen.  Gaussian processes in machine learning.  In Summer school on machine learning, pages 63–71. Springer, 2003.
  * [42] Sashank Reddi, Zachary Burr Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan, editors.  Adaptive Federated Optimization, 2021.
  * [43] Jae Hun Ro, Ananda Theertha Suresh, and Ke Wu.  Fedjax: Federated learning simulation with jax.  arXiv preprint arXiv:2108.02117, 2021.
  * [44] Tim GJ Rudner, Zonghao Chen, and Yarin Gal.  Rethinking function-space variational inference in bayesian neural networks.  In Third Symposium on Advances in Approximate Bayesian Inference, 2020.
  * [45] Tim GJ Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal.  Tractable function-space variational inference in bayesian neural networks.  Advances in Neural Information Processing Systems, 35:22686–22698, 2022.
  * [46] Edward Snelson and Zoubin Ghahramani.  Sparse gaussian processes using pseudo-inputs.  Advances in neural information processing systems, 18, 2005.
  * [47] Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse.  Functional variational bayesian neural networks.  In International Conference on Learning Representations, 2019.
  * [48] Michalis Titsias.  Variational learning of inducing variables in sparse gaussian processes.  In Artificial intelligence and statistics, pages 567–574. PMLR, 2009.
  * [49] Volker Tresp.  A bayesian committee machine.  Neural computation, 12(11):2719–2741, 2000.
  * [50] Aki Vehtari, Andrew Gelman, Tuomas Sivula, Pasi Jylänki, Dustin Tran, Swupnil Sahai, Paul Blomstedt, John P Cunningham, David Schiminovich, and Christian P Robert.  Expectation propagation as a way of life: A framework for bayesian inference on partitioned data.  The Journal of Machine Learning Research, 21(1):577–629, 2020.
  * [51] Maxime Vono, Vincent Plassier, Alain Durmus, Aymeric Dieuleveut, and Eric Moulines.  Qlsd: Quantised langevin stochastic dynamics for bayesian federated learning.  In International Conference on Artificial Intelligence and Statistics, pages 6459–6500. PMLR, 2022.
  * [52] Jianyu Wang, Zachary Burr Charles, Zheng Xu, Gauri Joshi, Brendan McMahan, Blaise Hilary Aguera-Arcas, Maruan Al-Shedivat, Galen Andrew, A. Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang He, Samuel Horvath, Zhouyuan Huo, Alex Ingerman, Martin Jaggi, Tara Javidi, Peter Kairouz, Satyen Chandrakant Kale, Sai Praneeth Karimireddy, Jakub Konečný, Sanmi Koyejo, Tian Li, Luyang Liu, Mehryar Mohri, Hang Qi, Sashank Reddi, Peter Richtarik, Karan Singhal, Virginia Smith, Mahdi Soltanolkotabi, Weikang Song, Ananda Theertha Suresh, Sebastian Stich, Ameet Talwalkar, Hongyi Wang, Blake Woodworth, Shanshan Wu, Felix Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang, Chunxiang (Jake) Zheng, Chen Zhu, and Wennan Zhu.  A field guide to federated optimization.  Technical report, arxiv, 2021.
  * [53] Steven Winter, Trevor Campbell, Lizhen Lin, Sanvesh Srivastava, and David B. Dunson.  Machine learning and the future of bayesian computation.  arXiv preprint arXiv:2304.11251, 2023.
  * [54] David H Wolpert.  Bayesian backpropagation over io functions rather than weights.  Advances in neural information processing systems, 6, 1993.
  * [55] Yuxin Wu and Kaiming He.  Group normalization.  In Proceedings of the European conference on computer vision (ECCV), pages 3–19, 2018.
  * [56] Yuanhao Xiong, Ruochen Wang, Minhao Cheng, Felix Yu, and Cho-Jui Hsieh.  Feddm: Iterative distribution matching for communication-efficient federated learning.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16323–16332, 2023.
  * [57] Jie Zhang, Chen Chen, Bo Li, Lingjuan Lyu, Shuang Wu, Shouhong Ding, Chunhua Shen, and Chao Wu.  Dense: Data-free one-shot federated learning.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 21414–21428. Curran Associates, Inc., 2022.
  * [58] Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen.  Dataset condensation with gradient matching.  In International Conference on Learning Representations, 2021.
  * [59] Yanlin Zhou, George Pu, Xiyao Ma, Xiaolin Li, and Dapeng Wu.  Distilled one-shot federated learning.  arXiv preprint arXiv:2009.07999, 2020.

##  Appendix A Derivations

###  A.1 Derivation of independent posterior

$\displaystyle p({\mathbf{f}}|\mathcal{D})$
$\displaystyle=\frac{1}{p(\mathcal{D})}p({\mathbf{f}})\prod^{M}_{m=1}p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
$\displaystyle=\frac{1}{p(\mathcal{D})}p_{0}({\mathbf{f}}_{\bm{Z}})p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{f}}_{{\bm{Z}}_{m}})p({\mathbf{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
$\displaystyle=\frac{1}{p(\mathcal{D})}p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{f}}_{{\bm{Z}}_{m}})p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
$\displaystyle=\frac{1}{p(\mathcal{D})}p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}})p({\mathbf{f}}_{{\bm{X}}_{m}}|{\mathbf{f}}_{{\bm{Z}}_{m}})p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
$\displaystyle=\frac{1}{p(\mathcal{D})}p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}})p({\bm{y}}_{m}|{\mathbf{f}}_{{\bm{X}}_{m}})$
$\displaystyle=\frac{p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}},{\bm{y}}_{m})}{\int
p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p_{0}({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}},{\bm{y}}_{m})d{\mathbf{f}}}$
$\displaystyle=\frac{p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}},{\bm{y}}_{m})}{\prod^{M}_{m=1}p(\mathcal{D}_{m})}$
$\displaystyle=p({\mathbf{f}}_{*}|{\mathbf{f}}_{\bm{Z}})\prod^{M}_{m=1}p({\mathbf{f}}_{{\bm{Z}}_{m}},{\mathbf{f}}_{{\bm{X}}_{m}}|{\bm{y}}_{m})$

##  Appendix B BPC-fKL pseudocode (from [28])

Algorithm 2 BPC-fKL

1: Require: Set of expert trajectories $\mathcal{T}$, trained on
$\mathcal{D}$, each parameter trajectory saves parameters to disk at a
specified interval

2: Require: Number of SGD sampler steps with the BPC $L_{\mathcal{C}}$ and the
data $L_{\mathcal{D}}$, total training steps $K$, number of Gaussian samples
$S$, variances $\Sigma_{\mathcal{C}}$, $\Sigma_{\mathcal{D}}$, sampling
learning rate $\eta$ and BPC learning rate $\psi$.

3: Input: local data $\mathcal{D}=({\bm{X}},{\bm{y}})$

4: Initialize $\mathcal{C}=({\bm{Z}},\hat{{\bm{y}}})$ from $\mathcal{D}$

5: for $k\in\\{1,...,K\\}$ do

6: Sample expert trajectory $\tau=\\{{\bm{\theta}}^{(r)}_{*}\\}^{T}_{r=0}$
from trajectory set $\mathcal{T}$

7: Randomly choose a checkpoint to start $r\leq L_{\mathcal{D}}$ and
initialize
${\bm{\theta}}^{(0)}_{\mathcal{C}}={\bm{\theta}}^{(0)}_{\mathcal{D}}={\bm{\theta}}^{(r)}_{*}$

8: Obtain
${\bm{\theta}}_{\mathcal{D}}^{(L_{\mathcal{D}})}={\bm{\theta}}_{*}^{(r+L_{\mathcal{D}})}$
from memory

9: for $t\in\\{1,...,L_{\mathcal{C}}\\}$ do

10: Update coreset parameter
${\bm{\theta}}^{(t)}_{\mathcal{C}}\leftarrow{\bm{\theta}}_{\mathcal{C}}^{(t-1)}-\eta\nabla_{\bm{\theta}}\log
p(\hat{{\bm{y}}}|f({\bm{Z}},{\bm{\theta}}_{\mathcal{C}}^{(t-1)}))$

11: end for

12: Sample random Gaussian noises
$\\{\epsilon^{(s)}_{\mathcal{D}},\epsilon^{(s)}_{\mathcal{C}}\\}^{S}_{s=1}\overset{i.i.d}{\sim}\mathcal{N}(0,\sigma_{\epsilon})$

13: $g\leftarrow\frac{1}{S}\sum^{S}_{s=1}\left(\nabla_{\mathcal{C}}\log
p(\hat{{\bm{y}}}|f({\bm{Z}},{\bm{\theta}}_{\mathcal{D}}^{(L_{\mathcal{D}})}+\Sigma^{1/2}_{\mathcal{C}}\epsilon^{(s)}_{\mathcal{D}}))-\nabla_{\mathcal{C}}\log
p(\hat{{\bm{y}}}|f({\bm{Z}},{\bm{\theta}}^{(L_{\mathcal{C}})}_{\mathcal{C}}+\Sigma^{1/2}_{\mathcal{C}}\epsilon^{(s)}_{\mathcal{C}}))\right)$

14: $\mathcal{C}\leftarrow\mathcal{C}-\psi g$

15: end for

16: Return $\mathcal{C}$

##  Appendix C Experiment details

###  C.1 Datasets and models

Synthetic regression We generated data on a domain defined on 3 intervals:
$\\{(-0.8,-0.6),(-0.2,0.0),(0.5,0.8)\\}$ using a ground truth function
$y(x)=1.5\sin(0.4\pi x)+1.5\cos(2\pi x)+\epsilon$ where
$\epsilon\sim\mathcal{N}(0,0.3)$. We generate 5 clients that sample datapoints
from these intervals and compute outputs with $y(x)$. We introduce non-iid-
ness to client datasets by sampling a probability vector over the 3 intervals
from a Dirichlet distribution with $\alpha=[1.0,1.0,1.0]$. Each client thus
samples more from specific intervals than other clients, which introduces the
non-iid-ness in the data. We use a 3 layer MLP with 128 units per linear layer
with Swish activation functions as neural network architecture.

2d classification The dataset is the ’moons’ dataset from the scikit-learn
library111https://scikit-learn.org/. We add Gaussian noise with a standard
deviation of $0.14$ to all samples. We generate 5 clients that then sample a
dataset of only 20 points from the moons dataset. As a network architecture we
use a 3 layer MLP with 50 units and relu activations, combined with Group
Normalization with two groups in the first two layers [55].

EMNIST-62 The original EMNIST-62 dataset is an extended version of the classic
MNIST dataset to 62 classes, adding both capital and undercase letters. The
authors of LEAF partitioned this dataset into a federated dataset based on the
author identification that came with the original data collection of EMNIST
[7]. This gives the federated EMNIST-62 dataset a natural non-iid-ness between
clients which allows for more realistic simulations, as the difference between
author handwriting styles is a good example of the type of non-iid-ness that
one could expect to see when implementing FL on real world data. As a network
architecture we adopted the ConvNet architecture from [42], but without
Dropout, as we found that negatively affected our method’s performance.

###  C.2 BPC-FL Hyperparameters

We provide all relevant hyperparameters for BPC-FL for each experiment in
Table 2

Hyperparameter | Synthetic Regression | 2d classification | EMNIST-62
---|---|---|---
batch size | 20 | 50 | 25
client update steps | 400 | 700 | 500
prior precision / regularization | 1e-2 | 1e-1 | 1e-2
random init std $\sigma_{z}$ | 1.0 | 6e-1 | 1e-3
step size $x$ | 0.01 | 1e-3 | 1e2
step size $y$ | 1.0 | 0.0 | 0.0
weight mc samples per update | 10 | 20 | 5
Client local hyperparameters |  |  |
number of pretrain seeds | 100 | 100 | 100
number of pretrain trajectories | 100 | 100 | 100
pretrain save interval | 5 | 5 | 5
number of pretrain steps | 300 | 400 | 750
pretrain step size | 1e-2 | 1e-2 | 1e-2
mcmc sampler | Adam | Adam | SGD
mcmc sampler step size | 1e-2 | 1e-2 | 1e-2
pseudodata optimizer | SGD | SGD | SGD
data mcmc chain length | 150 | 150 | 200
pseudodata mcmc chain length | 200 | 200 | 250
number of noise samples $S$ | 10 | 10 | 10
noise std $\sigma_{\epsilon}$ | 1e-2 | 1e-2 | 1e-4
Table 2: Used hyperparameters for experiments

###  C.3 Baselines details

FedAvg We use the adaptive implementation of FedAvg from [26] implemented in
FedJAX. For EMNIST-62 we used the optimal hyperparameters from that paper. For
the synthetic datasets we performed a grid search over
$\\{1e-3,1e-2,1e-1\\},\\{1e-2,1e-1,1e0\\}$ for the client and server step
sizes, respectively. All hyperparameters are summarized in 3.

Hyperparameter | Synthetic regression | 2d classification | EMNIST-62
---|---|---|---
Client optimizer | Adam | Adam | SGD
Client step size | 1e-2 | 1e-2 | 1e-1
Server optimizer | Adam | Adam | SGD
Server step size | 0.1 | 0.1 | 1.0
Table 3: Hyperparameters FedAvg baseline

DOSFL We implemented the DOSFL algorithm ourselves following the details in
[59]. We found that backpropagating through optimization trajectories made the
algorithm too memory demanding for our resources, thus we could not produce
results on the larger EMNIST-62 task. When setting the hyperparameters to
different configurations that match the communication load of our method,
DOSFL did not produce a useful server model. Therefore, we gradually increased
distillation steps and epochs until a reasonable result was produced by the
algorithm. We found that random masking and soft resetting did not improve
performance. We also were not able to produce a good result for DOSFL on the
synthetic regression task for any hyperparameter configuration that was within
our memory constraints. The hyperparameters used for producing the results in
the paper are given in Table 4.

Hyperparameter | Synthetic regression | 2d classification
---|---|---
$x$ step size | 0.1 | 0.01
$x$ optimizer | SGD ($m=0.9$) | SGD ($m=0.9$)
$y$ step size | 0.1 | 0.0
$y$ optimizer | SGD ($m=0.9$) | SGD ($m=0.9$)
$\eta$ step size | 0.0 | 0.0001
$\eta$ optimizer | SGD | Adam ($\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=1e-8$)
$\eta$ initial value | 0.01 | 0.01
distillation steps | 20 | 30
distillation epochs | 1 | 20
distillation batch size | 10 | 30
Table 4: Hyperparameters for DOSFL baseline

FedPA We implemented the Federated Posterior Averaging (FedPA) algorithm
ourselves, following mostly following the author’s implementation
222https://github.com/alshedivat/fedpa. We used the paper’s suggested
hyperparameters for the EMNIST-62 learning task [2]. Interestingly, FedPA
really tended to overfit on training clients when the training client pool was
relatively small.

MIME We used the MIME implementation in FedJAX written by the paper authors.
Interestingly, we did not find MIME to perform according to the reported
results in [26] when using the corresponding hyperparameters. MIME typically
performed worse on EMNIST-62 when we increased the amount of local steps.

Hyperparameter | EMNIST-62
---|---
client optimizer | SGD ($m=0.9$)
client step size | 0.01
server optimizer | SGD ($m=0.9$)
server step size | 0.5
number of samples | 10
burn in rounds | 500
shrinkage | 0.1
Table 5: FedPA hyperparameters Hyperparameter | EMNIST-62
---|---
base optimizer | Adam ($\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=1e-8$)
base step size | 0.001
server optimizer | SGD
server step size | 1.0
gradient batch size | 128
Table 6: MIME hyperparameters

###  C.4 Downstream evaluation

We searched hyperparameters for the downstream optimizers and samplers by
running Bayesian Optimization search over a fixed learned pseudodataset for a
specific task. The final hyperparameters are shown below in table 7.

SGD hyperparameters | Synthetic regression | 2d classification | EMNIST-62
---|---|---|---
step size | - | 2e-2 | 1e-2
number of steps | - | 500 | 2000
Adam hyperparameters |  |  |
step size | 1e-2 | - | 1e-2
$\beta_{1}$ | 0.9 | - | 0.9
$\beta_{2}$ | 0.999 | - | 0.999
$\epsilon$ | 1e-8 | - | 1e-8
number of steps | 300 | - | 2000
HMC hyperparameters |  |  |
step size | 1e-3 | 2e-3 | 2e-4
inverse mass matrix | $\text{diag}(50)$ | $\text{diag}(100)$ | $\text{diag}(10)$
num integration steps | 20 | 30 | 10
number of steps | 200 | 40 | 150
number of samples | 40 | 50 | 10
Table 7: Hyperparameters for downstream inference algorithms on evaluation
coreset

###  C.5 Additional results

####  C.5.1 Synthetic regression, small datasets

Figure 4 shows the synthetic regression problem for a setting where clients
have very little data. We simulated 5 clients that only have 20 datapoints.
BPC-FL inferred a pseudodataset of 6 datapoints as before. It is interesting
to see that FedAvg starts to produce functions that overfit when clients have
smaller datasets, while BPC-FL still produces decent results.

Figure 4: Synthetic regression - small datasets

####  C.5.2 EMNIST62 - influence of amount of pseudodata size

In Table 8 we show the final performance of models produced by SGD after
running BPC-FL on 100 clients for different BPC sizes. Note how performance
improves with a larger amount of pseudodatapoints. We did observe that for
larger pseudodatasets, pseudoinputs started to exactly mimic ${\bm{x}}$-values
of datapoints with only a few classes in the dataset.

Number of pseudo-input | Negative log likelihood | Accuracy | ECE
---|---|---|---
30 | $1.504\pm 0.150$ | $0.703\pm 0.005$ | $0.111\pm 0.040$
62 | $1.118\pm 0.051$ | $0.720\pm 0.003$ | $0.080\pm 0.027$
100 | $1.00\pm 0.03$ | $0.755\pm 0.004$ | $0.064\pm 0.023$
Table 8: Influence of pseudodata size on performance

####  C.5.3 EMNIST62 - influence of number of clients

In Table 9, we show the final performance of models produced by SGD after
running BPC-FL on different amount of clients, with a pseudodataset size of
62. It is shown that for more clients that participate, the quality of the
server model improves, though it seems that after 100 clients the performance
increase starts to stall. We suspect that because we learn a compressed
representation of client information, some fine-grain information in the data
posterior will go lost, which will not be gained back by simply adding more
clients that also carry that approximation error.

Number of clients | Negative log likelihood | Accuracy | ECE
---|---|---|---
30 | $1.822\pm 0.132$ | $0.613\pm 0.009$ | $0.150\pm 0.004$
50 | $1.385\pm 0.006$ | $0.651\pm 0.006$ | $0.111\pm 0.009$
100 | $1.210\pm 0.029$ | $0.712\pm 0.103$ | $0.083\pm 0.005$
200 | $1.152\pm 0.073$ | $0.719\pm 0.003$ | $0.092\pm 0.031$
300 | $1.01\pm 0.002$ | $0.718\pm 0.005$ | $0.071\pm 0.009$
Table 9: Influence of the amount of clients on performance

EMNIST62 ${\bm{Z}}$ examples To give an idea of what the learned pseudodata
${\bm{Z}}$ look like visually for the EMNIST-62 dataset, we provide some
examples below.

Figure 5: Some example x-values from the server set in EMNIST-62