FedPop: A Bayesian Approach for Personalised Federated Learning
Abstract
Personalised federated learning (FL) aims at collaboratively learning a machine learning model taylored for each client.
Albeit promising advances have been made in this direction, most of existing approaches do not allow for uncertainty quantiﬁcation which is crucial in many applications.
In addition, personalisation in the cross-silo and cross-device setting still involves important issues, especially for new clients or those having small number of observations.
This paper aims at ﬁlling these gaps.
To this end, we propose a novel methodology coined FedPop by recasting personalised FL into the population modeling paradigm where clients’ models involve ﬁxed common population parameters and random effects, aiming at explaining data heterogeneity.
To derive convergence guarantees for our scheme, we introduce a new class of federated stochastic optimisation algorithms which relies on Markov chain Monte Carlo methods.
Compared to existing personalised
FL methods, the proposed methodology has important beneﬁts: it is robust to client drift, practical for inference on new clients, and above all, enables uncertainty quantiﬁcation under mild computational and memory overheads.
We provide nonasymptotic convergence guarantees for the proposed algorithms and illustrate their performances on various personalised federated learning tasks.
Introduction
Federated learning (FL) is a recent machine learning paradigm in which distributed clients holding siloed data collaborate in solving a learning problem, usually under the coordination of a central server (Wang et al., 2021; Kairouz et al., 2021).
One of the main focus of FL is on so-called cross-device applications where a large number of personal electronic devices such as mobile phones, wearable devices or home assistants collect and store data at the edges of a decentralised network (McMahan et al., 2017).
While standard FL methods (McMahan et al., 2017; Alistarh et al., 2017; Karimireddy et al., 2020;
Horváth et al., 2019; Li et al., 2020) have focused on training a global model that can be applied to individual agents, the relevance of such inferences has recently been questioned due to statistical heterogeneity between clients.
Indeed, the considered common model may not generalise well on a client with a local data distribution that differs signiﬁcantly from the global data distribution, especially if that client has not participated in the training process.
In fact, it might even be better for such clients to derive a local model from their own data set.
To circumvent these issues, a number of personalised federated learning approaches have recently been proposed, that use local models to ﬁt client-speciﬁc data distribution while capturing some shared knowledge via a federated scheme (Tan et al., 2022).
Personalisation has previously been approached using multi-task learning (Smith et al., 2017), meta-learning (Jiang et al., 2019; Khodak et al., 2019), client clustering (Briggs et al., 2020), data interpolation (Mansour et al., 2020), model interpolation (Hanzely and Richtárik, 2020;
Hanzely et al., 2020) or partially local models (Singhal et al., 2021; Collins et al., 2021).
While these methodologies are promising, they only partially solve the personalisation problem in highly heterogeneous federated settings and have no means of quantifying uncertainty.
In addition, crossdevice FL also faces other important challenges such as (extreme) partial device participation, small local data sets, limited upload bandwidth and device capabilities (Kairouz et al., 2021).
Addressing these problems for personalised FL requires new paradigms regarding how model knowledge is shared and personalisation is performed locally.
Proposed Approach.
In this paper, we adopt a novel perspective to model the problem of personalised
FL.
This paradigm, called mixed-effects modeling (also known as multi-level or population approach) is widely used to analyse data that have a clustered or nested structure, as in medical or biological research where multiple observations per patient are available (Gelman and Hill, 2007;
Long, 2011;
Lavielle, 2014).
Although the hierarchical structure of FL has already been noted (Plassier et al., 2021; Grant et al., 2018; Hong et al., 2022), the mixed-effects paradigm has interestingly never been considered.
Leveraging this framework, we develop a new model for personalised FL called
FedPop and propose an efﬁcient computational solution to perform inference under this model.
More precisely, we introduce a novel class of federated stochastic approximation algorithms based on parallel Markov Chain Monte Carlo (MCMC) methods.
In the proposed framework, we also pay special attention to the cross-device setting by taking into account partial client participation, and by addressing the communication bottleneck with both multiple local updates and the use of lossy compression operators.
Beneﬁts.
Up to the authors’ knowledge, FedPop is the ﬁrst personalised FL approach that allows cheap uncertainty quantiﬁcation with a theoretically-grounded methodology.
The proposed framework also comes with other interesting properties.
First, in contrast to most of personalised FL methods that only consider “ﬁxed-effects” models (Collins et al., 2021; Hanzely et al., 2021; Smith et al., 2017), FedPop provides credibility information (via credibility regions) and allows more accurate inference for clients with small and heterogeneous local data via partial pooling (Gelman and Hill, 2007).
In addition, inference for new clients which did not participate in the training phase can be easily performed by sampling from the prior over the local random effects.
Second, contrary to existing Bayesian FL approaches that aim to provide credibility information by sampling from a target posterior distribution (Hong et al., 2022; Yoon et al., 2018; Vono et al., 2022; El Mekkaoui et al., 2021), FedPop allows to perform personalisation and cheaper on-device uncertainty quantiﬁcation taking an empirical Bayes prediction approach.
Finally, an important beneﬁt of FedPop is its ability to allow for multiple local updates without suffering from the client-drift phenomenon (Karimireddy et al., 2020).
Outline and Contributions.
Our contributions are fourfold.
First, in Section 2, we propose a novel probabilistic methodology, which we call FedPop , to address personalisation under the crossdevice FL paradigm.
To perform efﬁcient inference under this model, we introduce a general class of stochastic approximation algorithms based on MCMC.
Second, we provide in Section 3 nonasymptotic convergence guarantees for the proposed methodology.
Then, we perform in Section 4 a comparison between the proposed approach and exisiting works.
Finally, we illustrate in Sectionthe beneﬁts of our methodology on several federated learning benchmarks involving both synthetic and real data.
Proposed Approach
In this section, we present the statistical estimation problem we are considering and the proposed methodology called FedPop to address it.
Problem Formulation.
We are interested in the cross-device FL setting involving a large number b ∈ N ∗ of clients, potentially unreliable i.e. not necessarily available at each communication round.
These clients are assumed to own sensitive local data sets { D i } i ∈ [ b ] .
In this framework, we aim to make both uncertainty quantiﬁcation and personalised statistical inference by learning a local model taylored to each client.
To this end, and inspired by the population approach used in the biological and physical sciences (Lavielle, 2014), we consider mixed-effects modeling for each client leading to the local marginal likelihood function deﬁned, for any i ∈ [ b ] , by
[Equation]
where φ ∈ Φ ⊆
R d Φ stands for a ﬁxed effect and { z ( i ) } i ∈
[ b ] ∈ Z , Z = (cid:81) bi =1 R d , represent random effects aimed at explaining statistical heterogeneity between local data sets { D i } i ∈
[ b ] .
The objective of the ﬁxed ( i.e. constant across all clients) part is to capture a common representation ( e.g. same features across different classes of images) while the random part, which is typically low-dimensional, performs personalisation and is assumed to be drawn from a population prior whose variance aims at modeling data heterogeneity.
Figure 1 illustrates this statistical framework, referred to as FedPop , by showing its directed acyclic graph (DAG) where grey-ﬁlled shapes indicate observed variables, white-ﬁlled shapes unknown variables and squared shapes variables to be estimated.
When the size of the local data set D i is small, this common prior leverages information from other clients to limit the risk of overﬁtting and is often called partial pooling in the multi-level statistical literature (Gelman and Hill, 2007, Section 12).
Examples of model architectures involving φ and { z ( i ) } i ∈ [ b ] include for instance composition -based architectures p (D i | φ, z ( i ) )
=
p (D i | h φ ◦ h z ( i ) ) where h φ and h z ( i ) are two neural networks (Collins et al., 2021; Arivazhagan et al., 2019).
For the sake of generality, we propose to adopt a ﬂexible energy-based prior distribution of the form for each i ∈ [ b ] ,
[Equation]
Here, Z ( β ) is a normalising constant and E ( · ; β ) represents an energy function, typically a neural network, parameterised by a set of parameters β ∈ B ⊆ R d B (LeCun et al., 2006).
This framework is particularly interesting in the cross-device setting where the number of clients b is large as it allows for efﬁcient enrichment of the model.
However, in the case where b is small, the inference of the parameter β is difﬁcult.
In this situation, a more pragmatic solution is to consider a common prior for the local random effects { z ( i ) } i ∈ [ b ] which is held ﬁxed, i.e. p ( z ( i ) | β ) ∝ exp {− E ( z ( i ) ) } for any β ∈ B .
Finally, for completeness, we allow the use of a prior model p ( φ, β ) for the hyperparameters { φ, β } .
Using Bayes’ rule (Robert, 2001) and by denoting D
= (cid:116) bi =1 D
i the global data set, the posterior distribution associated with these hyperparameters admits a probability density function which can be written as
[Equation]
Set θ = { φ, β } ∈ Θ with Θ =
Φ × B .
In the sequel, we will be interested in solving the maximum a posteriori problem given by
[Equation]
where C ∈ R is a constant independent of θ .
Once we have estimated θ (cid:63) , using an empirical Bayesian approach, we can perform “for free” on-device uncertainty quantiﬁcation for each client i ∈
[ b ] by sampling from the local posterior distribution p ( z ( i ) | D i , φ (cid:63) , β (cid:63) ) , which is typically designed to be low-dimensional.
Algorithm.
To solve the optimisation problem (2), we can either use an alternating maximisation algorithm or perform global maximisation over Θ .
Since the former approach requires more upload bandwidth, in this work we consider the second alternative which is more suitable for FL.
The gradient of the objective function (3) being intractable, we propose to resort to the stochastic approximation framework (Robbins and Monro, 1951) which iteratively deﬁnes ( φ k , β k ) k ∈ N , starting from any ( φ 0 , β 0 )
∈ Θ , via the recursions for any k ∈ N ,
[Equation]
where Π C denotes the projection onto C ∈ { Φ , B } , ( η (1) k , η (2) k ) k ∈ N ∗ are sequences of step-sizes, and { g ( i ) k :
i ∈
[ b ] , k ∈ N ∗ } and { h ( i ) k : i ∈ [ b ] , k ∈ N ∗ } are estimators of the intractable gradients ( φ, β ) (cid:55)→ ∇ β log p
(D i | φ, β ) and ( φ, β ) (cid:55)→ ∇ φ log p
(D i | φ, β ) at ( φ k , β k ) , where p (D i | φ, β ) is deﬁned in (1) for any i ∈ [ b ] .
The choices of the estimators { g ( i ) k : i ∈ [ b ] , k ∈ N ∗ } and { h ( i ) k : i ∈ [ b ] , k ∈ N ∗ } are motivated by the Fisher identity.
More precisely, under mild regularity assumptions, and using the Lebesgue dominated convergence theorem, we have for any, (
[Equation]
which suggests to consider
[Equation]
where M ∈ N ∗ and Z
( i, 1: M ) k
= ( Z ( i,m ) k ) m ∈
[ M ] are approximate samples from p ( z ( i )
| D i , φ, β ) .
More precisely, we consider a family { Q ( i ) γ,θ : γ ∈ (0 , ¯ γ ] , θ ∈ Θ } where for any step-size γ , Q ( i ) γ,θ is a Markov kernel which targets a close approximation of p ( z ( i ) |
D i , θ ) with θ = { φ, β } .
As an example, we can use overdamped Langevin dynamics (Roberts and Tweedie, 1996;
Welling and Teh, 2011) to generate these samples.
In this case, Q ( i ) γ,θ is associated with a Gaussian probability density function q ( i ) γ,θ ( z ( i ) , · ) with mean z ( i ) − γ ∇ z log p
z ( i )
| D i , θ ) and variance 2 γ I d .
Note that the number of Monte Carlo draws per iteration k is considered constant here
but we can easily generalise our scheme to the non-constant setting.
In addition, our scheme can also be generalised by taking into account stochastic gradient estimators of (4) and (5).
For the sake of simplicity, we present our approach with standard gradients.
In this framework, we present the main steps of the corresponding stochastic approximation algorithm, called FedSOUK , in Algorithm 1.
Since we consider the cross-device federated setting, note that only a random subset A k +1 of active ( i.e. available) clients communicates with the central server at each iteration k ∈ N .
In addition, due to limited upload bandwidth, the potentially high-dimensional gradient estimator (5) is compressed locally via an unbiased stochastic compression operator C k + before being sent to the central server (Alistarh et al., 2017; Philippenko and Dieuleveut, 2020).
Stateful and Stateless Versions.
Depending on local memory constraints and the participation rate, we allow for a possible warm-start strategy across communication rounds to improve the convergence properties of the proposed algorithm so that the proposed algorithm becomes stateful , see steps 4 −in Algorithm 1.
In cases the participation rate is very low ( e.g. each client might only participates once to the training process), we replace this warm-start strategy by the initialisation
Z ( i, 0)
k ∼ p ( · | β k ) for any i ∈ [ b ] and k ∈ { 0 , . . .
, K − 1 } .
This yields a stateless version of our algorithm more suitable to the cross-device setting.
Obviously, compared to the previously proposed warm-start strategy, the performances of Algorithm 1 will be affected negatively if we are using the same number of local iterations M .
We end up with an interesting trade-off between local computations and communication: if client-server communication is a bottleneck, the stateless version of the algorithm allows to reduce the communication overhead at the price of longer sampling procedures on each client.
Such analyses will be illustrated empirically in Section 5.
Computation Complexity.
Compared to standard FL methods, our approach has an additional computational cost on the client side associated with Monte Carlo approximations { I ( i ) k } and { J ( i ) k } .
In practice, this cost is negligible.
Indeed, in our experiments, we found that using a small value of
M ∈ [1 , 10] was sufﬁcient to obtain state-of-the-art results in terms of accuracy on the test dataset.
We would like to emphasize that this additional computational cost has also two side advantages compared to existing FL approaches: (1) it allows us to communicate less frequently with the central server and (2) it allows us to converge faster when the number of local iterations increases since
Monte Carlo approximation becomes better.
Communication Overhead.
As pointed out in Table 1, our methodology FedPop improves upon existing FL approaches regarding the communication overhead.
Indeed, FedPop offers the ﬂexibility to use both compression for sending updates to the server, and multiple local steps to reduce the communication frequency.
As such, depending on the bandwidth and local computational power, the practitioner can adapt the number of local iterations and the parameter of the compression operator.
Up to our knowledge, this work is the ﬁrst one combining compression and multiple local steps for personalised FL.
Robustness to client drift.
For simplicity, we will take the example of FedSOUL (see Algorithm 1) which uses the Markov kernel associated with Langevin Monte Carlo to compute gradient estimates of the local marginal likelihood.
However, our answer holds for general Markov kernels (adjusted or unadjusted).
In this scenario, M steps of Langevin Monte Carlo are performed on each device to draw samples { Z ( i,m ) k } used to compute Monte Carlo estimates { I ( i ) k } and { J ( i ) k } .
Increasing the number of local steps M does not slow down convergence but instead allows for more accurate Monte Carlo integration and hence better convergence properties.
In contrast, the client drift phenomenon for classical FL approaches (e.g. FedAvg proposed in McMahan et al.
(2017)) slows down convergence as the number of local iterations M increases.
Simple inference on new clients.
Typical personalized FL approaches such as DITTO or FedRep require additional local training for inference on new clients.
In contrast, the proposed methodology
FedPop allows for a cheaper two-step approach once we have estimated θ (cid:63) =
( φ (cid:63) , β (cid:63) ) , as detailed below for a new client b + 1 with feature vector x :
Sample { Z ( l ) b +1 } l ∈ [ L ] in a i.i.d. manner.
Estimate the posterior predictive function by p ( · | x ) ≈ L − 1 (cid:80) Ll =1 p ( · | φ (cid:63) , Z ( l ) b +1 , x ) .
The prior p ( z b +1 | β (cid:63) ) is typically chosen so that sampling is computationally cheap, e.g. a Gaussian with diagonal covariance matrix as in our experiments, see Section 5.
Theoretical Guarantees
In this section, we present non-asymptotic convergence guarantees for Algorithm 1 when the family of Markov kernels { Q ( i ) γ,θ : γ ∈ (0 , ¯ γ ] , θ ∈ Θ ,
i ∈
[ b ] } is associated to unadjusted, i.e. without
Metropolis acceptance step, overdamped Langevin dynamics (Durmus and Moulines, 2017; Dalalyan, 2017).
The bounds we derive allow to showcase explicitly the impact of FL constraints, namely partial participation and compression.
Results for general unadjusted Markov kernels are postponed to the supplement.
To show our theoretical results and resort to standard assumptions made in the stochastic approximation literature, we consider a minimisation problem and rewrite the opposite of the objective function (3) for any θ ∈ Θ as
[Equation]
where f i ( θ )
= − log p (
[Equation]
[Equation]
1: Input: nb.
outer iterations K , nb.
local iterations M , Markov kernels { Q ( i ) γ,θ } γ,θ,i , step-sizes { η (1) k , η (2) k } k ∈
[ K ] ,i ∈ [ b ] and initial points Z (0) ∈ R d , β 0 ∈ B and φ 0 ∈ Φ .
2: for k = 0 to K − 1 do 3:
Server sends { β k +1 , φ k +1 } to clients belonging to A k +1 .
4: for i ∈
A k +1 //
On active clients A k
+1 do 5: //
Warm-start of the SA scheme if possible 6: if k ≥ 1 then 7:
Set Z ( i, 0) k
= Z ( i,M ) k − 1 .
8: end if 9: //
Computation of key quantities using MCMC 10: for m = 0 to M − 1 do 11:
Draw Z ( i,m +1)
k ∼ Q ( i ) γ,θ k (cid:16) Z
( i,m ) k , · (cid:17) .
For Langevin dynamics 13: // Draw ξ ( i,m +1)
k ∼ N(0 d , I d ) .
14: // Set Z
( i,m +1) k
= Z ( i,m ) k
+ γ ∇ z log p
( Z ( i,m ) k
| D i , φ k , β k ) + √ 2 γξ ( i,m +1) k .
15: end for 16: //
Communication with the server 17:
Set I ( i ) k
= 1 M (cid:80) Mm =1 ∇ β log p (cid:16) Z ( i,m ) k
| β k (cid:17) .
Set J ( i ) k
= 1 M (cid:80) Mm =1 ∇ φ log p (cid:16) D i | Z
( i,m ) k , φ k (cid:17) .
Send I ( i ) k and C k +1 (cid:16) J ( i ) k (cid:17) to the central server.
20: end for 21:
Set β k
+
=
Π B (cid:16) β k + η (1) k
+1 (cid:104) ∇ β log p ( β k ) + b | A k +1 | (cid:80) i ∈
A k
+1 I
( i ) k
(cid:105)(cid:17) .
Set φ k +
=
Π Φ (cid:16) φ k + η (2) k
+1 (cid:104) ∇ φ log p ( φ k )
+ b | A k +1 | (cid:80) i ∈
A k +1 C k + (cid:16) J
( i ) k (cid:17)(cid:105)(cid:17) .
23: end for 24: Output: { φ K , β K } and samples { Z (1: b,m ) K − 1 } Mm =1 .
Non-Asymptotic Convergence Bounds.
For the sake of better readability, we only detail in the main paper assumptions regarding the objective function, compression operators and the partial participation scenario.
Technical assumptions related to the Markov kernels { Q ( i ) γ,θ } are postponed to the supplement.
In spirit, we require, for any i ∈ [ b ] , θ ∈ Θ and γ , that Q ( i ) γ,θ satisﬁes some ergodic condition and can provide samples sufﬁciently close to the local posterior distribution p ( z ( i ) | D i , θ ) .
For the sake of simplicity, we also assume that for any k ∈ N ∗ , η (1) k
= η (2) k
= η k , see Algorithm 1.
We make the following assumptions on Θ and the family of functions { f i : i ∈ [ b ] } .
H1.
Θ is convex, closed subset of R d Θ and Θ ⊂ B(0 , R Θ ) for R Θ > 0 .
H2.
For any i ∈ [ b ] , the following conditions hold.
(i)
The function f i deﬁned in (6) is convex.
(ii) There exist
an open set U ∈ R d Θ and L f > 0 such that Θ ⊂ U , f i ∈ C
( U , R ) and for any θ 1 , θ 2 ∈ Θ ,
[Equation]
The assumption below requires compression operators { C k } k ∈ N ∗ to be unbiased and to have a bounded variance.
Such an assumption is for instance veriﬁed by stochastic quantisation operators, see Alistarh et al.
H3.
The compression operators { C k } k ∈ N ∗ are independent and satisfy the following conditions.
(i) For any k ∈ N ∗ , v ∈ R d , E
[ C k ( v )]
= v .
(ii)
There exists ω ≥ 1 , such that for any k ∈ N ∗ , v ∈ R d , E [ (cid:107) C k ( v ) − v (cid:107) 2 ] ≤ ω (cid:107) v (cid:107) 2 .
We ﬁnally assume that each client has probability p ∈ (0 , 1] to be active at each communication round.
We would like to point out that this partial participation assumption can be associated to a speciﬁc compression operator satisfying H 3.
H4.
For any k ∈ N ∗ , A k = { i ∈ [ b ] : B i,k = 1 } where for any i ∈ [ b ] , { B i,k : k ∈ N ∗ } is a family of i.i.d.
Bernouilli random variables with success probability p ∈ (0 , 1] .
Under these assumptions, the next result establishes that (¯ θ k ) k ∈ N deﬁned by ¯ θ k = (cid:80) kj =1 η j θ j / ( (cid:80) kj =1 η j ) converges towards an element of arg min Θ f .
Theorem 1.
Assume H
1- H 4 along with A 8 detailed in the supplement and let for any k ∈ [ K ] , η k ∈ (0 , 1 /L f ] .
Then, for any K ∈ N ∗ , we have
[Equation]
where E K depends linearly on ( ω/p ) (cid:80)
Kk
=
η 2 k ; and A ( γ ) =
Cγ α with α > 0 and C is independent of ω, p and ( η k ) .
Closed-form formulas for these constants are provided in the supplement.
An interesting feature of Algorithm 1 is that convergence towards a minimum of f is possible and the impact of partial participation and compression vanishes when lim k →∞ η k = 0 .
More precisely, lim sup k →∞ E K / ( (cid:80) Kk
=1 η k ) =
and lim γ → 0 + A ( γ ) = 0 which shows that we can tend towards a minimum of f with arbitrary precision (cid:15) > 0 by setting the step-size γ to a small enough value.
Related Works
As pointed out in Section 1, many different approaches have been proposed to address personalisation and uncertainty quantiﬁcation under the federated learning paradigm.
This section reviews the main related existing lines of research and shows that the proposed methodology provides many beneﬁts; see Table 1.
Interestingly, we also show that FedPop encompasses some of the existing FL models.
Bayesian FL.
One of our main motivations is the possibility to perform grounded uncertainty quantiﬁcation in FL by resorting to the Bayesian paradigm.
In the recent years, many works have suggested to adapt serial workhorses stochastic simulation approaches such as MCMC or variational inference to the FL setting (Chen and Chao, 2020;
Liu and Simeone, 2021b,a; Vono et al., 2022; El Mekkaoui et al., 2021; Corinzia et al., 2019; Bui et al., 2018; Plassier et al., 2021;
Deng et al., 2021a).
Although some of these approaches address important FL challenges such as the communication bottleneck, partial participation or limited computational device resources, they are not suitable for uncertainty quantiﬁcation in the cross-device FL scenario.
Indeed, all these approaches assume that the posterior distribution targeted by each client is parametrised by a single potentially high-dimensional parameter of size
d Φ + d , see (1).
This prevents a sufﬁcient number of samples from being stored locally to perform uncertainty quantiﬁcation and Bayesian model averaging, especially when the model is a large neural network.
In contrast, our approach decouples this unique high-dimensional parameter into a ﬁxed part φ and a low-dimensional random part z
( i ) , signiﬁcantly reducing the memory footprint of local sample storage.
In addition, Bayesian FL methods aim at sampling a random parameter from a target probability distribution where π ( θ )
∝ e − f ( θ ) where f = (1 /b ) (cid:80) bi =1 f
i denotes the negative log-likelihood associated to the i -th client.
On the other hand the proposed methodology considers a mixed-effects modeling approach where parameters are divided into two categories: a ﬁxed component and a random one for each client.
As such, the mixed-effects approach is in essence an empirical Bayesian/marginal likelihood approach (Casella, 1985; Robbins, 1992).
It corresponds to a hierarchical model that aims to combine the modeling ﬂexibility and uncertainty assessment of Bayesian inference with computational pragmatism.
More precisely, a part of the parameters (ﬁxed-effects ) are estimated via marginal likelihood maximisation and the rest (random effects) using common Bayesian techniques, which are in most cases low dimensional.
As a result, up to our knowledge, the model and approach that we propose is novel in FL and comes with many beneﬁts as shown in Table 1.
Personalised FL.
Beside uncertainty quantiﬁcation, we also aim at providing each client with a dedicated personalised model.
Among the numerous existing personalised FL approaches, those related to FedPop can be broadly classiﬁed into two groups: meta-learning and partially local
[Table]
M ( d + d Φ )
[Table]
methods .
Meta-learning based FL methods aim at training a global model conducive to fast training of personalised models.
Such a goal can be achieved, for example, by local ﬁne-tuning (Fallah et al., 2020), regularisation of local models towards their average (Hanzely and Richtárik, 2020; Hanzely et al., 2021) – or the opposite (Li et al., 2021), and model interpolation (Liang et al., 2019).
On the other hand, FL methods based on partial decoupling take an approach similar to ours by splitting the initial model into a backbone component and a local one aimed at personalisation (Collins et al., 2021;
Arivazhagan et al., 2019; Pillutla et al., 2022).
This partial decoupling could also enhance privacy as discussed in Singhal et al.
The main difference with FedPop is that such approaches based on empirical risk minimisation cannot provide credibility information.
FedPop : A Compromise between Standard and Personalised FL.
Interestingly, we show here that the FedPop framework allows existing FL approaches to be retrieved in certain regimes.
To this end, we assume that the prior p ( z ( i ) | β ) is Gaussian with mean µ and covariance matrix σ 2 I d so that β
= { µ, σ } .
If σ → 0 + , then this Gaussian prior tends towards the Dirac distribution centered at µ and the local likelihood becomes p (D i | φ, µ ) , which corresponds to the local objective of standard FL approaches such as FedAvg (McMahan et al., 2017).
On the other hand, when σ → ∞ , no common information β is used to locally regress
z ( i ) and we end up with the FedRep algorithm (Collins et al., 2021).
This shows that FedPop stands for a subtle compromise between standard and personalised FL which should beneﬁt clients with small data sets by pooling information via a common prior.
Finally, in the extreme scenario where φ is the null vector, our approach amounts to the Bayesian FL approach DG-LMC proposed in Plassier et al.
Numerical Experiments
In this section, we illustrate the beneﬁts of our methodology on several FL benchmarks associated to both synthetic and real data.
Since existing Bayesian FL approaches are not suited for personalisation (see Table 1), we only compare the performances of Algorithm 1 with personalised FL methods.
In all our experiments, we use overdamped Langevin dynamics to sample locally and call this speciﬁc instance of Algorithm 1, FedSOUL .
In addition, we set p ( z ( i ) | β ) =
N( µ, σ 2 I d ) with β = { µ, σ } for simplicity.
To be comparable with existing personalised FL approaches that only consider periodic communication via multiple local steps, we do not resort to the proposed compression mechanism although the latter could be of interest for real-world applications.
Additional experiments and details about experimental design are provided in the supplement.
Synthetic Data.
We start by showcasing the beneﬁts of FedSOUL for clients having small and highly heterogeneous data sets as pointed out in Section 1 and Section 2.
To this end, we consider a similar experimental setting as in Collins et al.
(2021) where synthetic observations { y ( i ) j } j ∈ [ N i ] ∈ D i are generated via the following procedure: x ( i ) j ∼ N(0 k , I k ) and y ( i ) j ∼ N( z ( i ) true φ (cid:62) true x
( i ) j , 0 . 1) .
The ground-truth parameters z ( i ) true ∈ R d and φ true ∈ R k × d have been randomly generated beforehand with ( d, k ) =
Compared to Collins et al.
(2021), we use heterogeneous data partitions across clients so that 90% of the b = 100 clients have small data sets of size 5 and the remaining 10% have data sets of size 10.
We compare our results with FedRep (Collins et al., 2021) and FedAvg (McMahan et al., 2017) since they stand for two limiting instances of the proposed methodology, see Section 4 and Gelman and Hill (2007, Section 12).
[Figure]
Figure 2 compares the different approaches by computing the principle angle distance 1 (respectively the (cid:96) 2 norm) between φ true (respectively z ( i ) true ) and its estimated value; the lesser the better.
In contrast to its main competitors and based on both metrics, FedSOUL provides an impressive improvement.
This illustrates the beneﬁts of the introduction of a common prior p ( z ( i ) | β ) which allows to prevent from overﬁtting on clients with small data sets while performing personalisation.
Additional results with other choices for ( b, d, k ) and data partitioning strategies are available in the supplement.
Moreover, to compare our algorithm with a non-FL setting, we perform a non-distributed and nonfederated stochastic approximation algorithm to ﬁnd θ ∗ using a large number of iterations to get an accurate approximation of the optimal parameter θ ∗ .
Then, we use FedPop to obtain an estimate ˜ θ ∗ and measure the relative error in l 2 - distance between θ ∗ and ˜ θ ∗ .
For some outer iterations T = 100 , the relative error was less than 10 − 3 , which illustrates the relevance of our theoretical results.
We also test the performances of the proposed approach when the warm-start strategy is not used.
In this case, we have to set M = 50 to achieve the same performances as in the stateful variant of FedSOUL .
Real Data.
We consider now real image data sets, namely CIFAR-10 and CIFAR-100 (Krizhevsky, 2009).
For our likelihood model deﬁned by p (D i | φ, z ( i ) ) , we use 5-layer convolutional neural networks and perform personalisation for the last layer.
We set b = 100 for convenience and control data heterogeneity by assigning to each client N i images belonging to only S different classes.
[Figure]
Small data sets.
Under this setting, we ﬁrst consider (10%, 50%, 90%) of clients having small data sets of size either N i = 5 or N i = 10 ; while remaining clients have larger data sets of size N i = 25 .
We compare our approach with FedRep since it stands for the state-of-the-art personalised FL approach.
The algorithms are trained fulﬁlling the same computational budget.
Figure 3 shows the average accuracy across clients for the two approaches on both CIFAR10 and CIFAR-100.
We can see that FedSOUL is consistently better than FedRep over different conﬁgurations.
Full data sets.
In addition to show that the proposed approach achieves state-of-the-art performances on small data sets (which is common in the cross-device scenario), we now illustrate that FedSOUL is also competitive on larger data sets.
To this end, we use all training images in CIFAR-10 and
CIFAR-100 image data sets and consider the same data partitioning as in Collins et al.
More precisely, in this case the number of observations and the number of classes per client are uniformly shared over the clients.
Table 2 shows our results in comparison with state-of-the-art personalised
FL approaches.
We can see that that our model outperforms other methods on both CIFAR-
[Table] [Table]
[Table]
FedSOUL (this paper)
[Table]
and CIFAR-100 by a large margin.
Additional results with other personalised FL algorithms are postponed to the supplement.
Uncertainty Quantiﬁcation on Real Data.
As highlighted in Table 1, one advantage of the proposed approach compared to existing personalised FL methods is the ability to perform uncertainty quantiﬁcation by sampling locally from the posterior p ( z ( i )
| D i , φ K , β K ) , see Algorithm 1.
We illustrate this feature by computing on CIFAR-10 calibration curves and scores ( e.g. expected calibration error aka ECE) on a speciﬁc client; and by performing an out-of-distribution analysis on MNIST/FashionMNIST data sets.
Figure 4 shows that the proposed approach provides relevant uncertainty diagnosis.
Additional results on uncertainty quantiﬁcation can be found in the supplement.
E x p ec t e d A cc u r a c y
ECE=2.
ReliabilityDiagram
GapAccuracy 0 .0 .0 .0 .0 .1 .
Conﬁdence
[Figure] [Figure]
Conclusion
In this paper, we proposed a general
Bayesian methodology based on a natural mixed-effects modeling approach to model personalisation in federated learning.
Our FL method is the ﬁrst that allows for both personalisation and cheap uncertainty quantiﬁcation for (cross-device) federated learning.
By introducing a common prior on the local parameters, we tackle the local overﬁtting problem in the scenario where clients have highly heterogeneous and small data sets.
In addition, we have shown that the proposed approach has favorable convergence properties.
Some limitations of
FedPop pave the way for more advanced personalised FL approaches.
As an example, our model does not allow for training heterogeneous architectures across clients because of the introduced common prior, and only satisfy ﬁrst-order privacy guarantees.
Regarding the latter, further works include for instance deriving differentially private versions of our framework.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] We describe these limitations in the conclusion.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] Some assumptions are explicitly stated in Section 3.
For the sake of readability and due to the space limit, we postponed other technical assumptions in the supplement.
(b) Did you include complete proofs of all theoretical results?
[Yes] All our proofs are postponed to the supplement.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See the supplement.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Some details have been stated in Section 5 in the main paper.
Because of the limited number of pages, others have been postponed to the supplement.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] We did it for some experiments we conducted, see e.g. Figure 2.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] This is postponed to the supplement.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See for example the citation for CIFAR image data sets.
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]