Federated Multi-Task Learning under a Mixture of Distributions
Authors: Othmane MARFOQ
Giovanni Neglia
Aurélien Bellet
Laetitia Kameni
Richard Vidal
Venue: NeurIPS 2021 Poster (2021)
Keywords: Federated learning
multi-task learning
personalized models
expectation-maximization
personalized federated learning

Abstract: The increasing size of data generated by smartphones and IoT devices motivated the development of Federated Learning (FL), a framework for on-device collaborative training of machine learning models. First efforts in FL focused on learning a single global model with good average performance across clients, but the global model may be arbitrarily bad for a given client, due to the inherent heterogeneity of local data distributions. Federated multi-task learning (MTL) approaches can learn personalized models by formulating an opportune penalized optimization problem. The penalization term can capture complex relations among personalized models, but eschews clear statistical assumptions about local data distributions. In this work, we propose to study federated MTL under the flexible assumption that each local data distribution is a mixture of unknown underlying distributions. This assumption encompasses most of the existing personalized FL approaches and leads to federated EM-like algorithms for both client-server and fully decentralized settings. Moreover, it  provides a principled way to serve personalized models to clients not seen at training time. The algorithms' convergence is analyzed through a novel federated surrogate optimization framework, which can be of general interest. Experimental results on FL benchmarks show that our approach provides models with higher accuracy and fairness than state-of-the-art methods.


Federated Multi-Task Learning under a Mixture of Distributions
Othmane Marfog 1.3 Giovanni Neglia', Aurélien Bellet?, Laetitia Kameni and Richard Vidal
Inria, Université Côte d' Azur, France, (othmane. marfoq. giovanni neglia) @ @nnria.fr Inria, Université de Lille, France, aurelien.belle @o inria.fr Accenture Labs, France, (richard. vidal, anetitia.kamen acccenture.com
Abstract
The increasing size of data generated by smartphones and loT devices motivated the development of Federated Learning (FL), framework for on-device collaborative training of machine learning models. First efforts in FL focused on learning a single global model with good average performance across clients, but the global model may be arbitrarily bad for given client, due to the inherent heterogeneity of local data distributions. Federated multi-task learning (MTL) approaches can learn personalized models by formulating an opportune penalized optimization problem. The penalization term can capture complex relations among personalized models, but eschews clear statistical assumptions about local data distributions.
In this work, propose to study federated MTL under the flexible assumption that each local data distribution a mixture of unknown underlying distributions. This assumption encompasses most of the existing personalized FL approaches and leads to federated EM-like algorithms for both client-server and fully decentralized settings. Moreover, it provides principled way serve personalized models clients not seen at training time. The algorithms'' convergence is analyzed through a novel federated surrogate optimization framework, which can be of general interest. Experimental results on FL benchmarks show that our approach provides models with higher accuracy and fairness than state-of-the-an methods.
Introduction
Federated Learning (FI.) 28
allows a set of clients to collaboratively train models without sharing
their local data. Standard FI. approaches train unique model for all clients 44
32
38
29
48
However, as discussed in [56]. the existence of such a global model suited for all clients is at odds
with the statistical heterogeneity observed across different clients 37
28 Indeed, clients can have non-iid data and varying preferences. Consider for example a language modeling task: given the sequence of tokens "I love eating," the next word can be arbitrarily different from one client to another. Thus, having personalized models for each client is necessity in many FL applications.
Previous work personalized FL. A naive approach for FL personalization consists in learning first a global model and then fine-tuning its parameters at each client via a few iterations of stochastic gradient descent (88. In this case, the global model plays the role of meta-model be used as initialization for few-shot adaptation at each client. particular, the connection between FL and Model Agnostic Meta Learning (MAML) (27) has been studied in [19, 30 1 in order build more suitable meta-model for local personalization. Unfortunately, these methods can fail to build a model with low generalization error (as exemplified by LEAF synthetic dataset 7 App. IJ). An alternative
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

approach is to jointly train a global model and one local model per client and then let each client build a personalized model by interpolating them [14 44). However, if local distributions are far from the average distribution, relevant global model does not exist and this approach boils down to every client learning only on its own local data. This issue formally captured by the generalization bound in 14 Theorem l].
Clustered FL (56 20 44 addresses the potential lack of global model by assuming that clients can be partitioned into several clusters. Clients belonging to the same cluster share the same optimal model, but those models can be arbitrarily different across clusters (see [56] Assumption 2] for a rigorous formulation). During training, clients learn the cluster to which they belong as well as the cluster model. The Clustered FL assumption also quite limiting, as no knowledge transfer is possible across clusters. In the extreme case where each client has its own optimal local model (recall the example on language modeling), the number of clusters coincides with the number of clients and no federated learning is possible.
Multi-Task Learning (MTL) has recently emerged as an alternative approach learn personalized models in the federated setting and allows for more nuanced relations among clients' models (59 63
67
16 The authors of [59 [63] were the first to frame FL personalization as MTL problem. In particulat ar, they defined federated MTL as a penalized optimization problem, where the penalization term models relationships among tasks (clients). The work (59) proposed the MOCHA algorithm for
the client-serve scenario, while 66 67] presented decentralized algorithms for the same problem. Unfortunately, these algorithms can only learn simple models (linear models or linear combination of pre-trained models), because of the complex penalization term. Other MTLL-based approaches (24) 23 16| [6]26]]5 (6) are able to train more general models at the cost of considering simpler penalization terms (e.g., the distance to the average model), thereby losing the capability to capture complex relations among tasks. Moreover, a general limitation of this line of work is that the penalization term is justified qualitatively and not on the basis of clear statistical assumptions on local data distributions.
More recently, [57] proposed pFedHN. pFedHN feeds local clients' representations a global (across clients) hypernetwork, which can output personalized heterogeneous models. Unfortunately, the hypernetwork has a large memory footprint already for small clients' models (e.g., the hypernetwork in the experiments in (57) has 100 more parameters than the output model). Hence, it is not clear if pFedHN can scale to more complex models. Moreover, pFedHN requires each client to communicate multiple times for the server to learn meaningful representations. Therefore, its performance is likely to deteriorate when clients participate only once (or few times) to training, as it is the case for large-scale cross-device FL training. Furthermore, even once the hypernetwork parameters have been learned, training personalized models for new clients still requires multiple client-server communication rounds. More similar to our approach, FedFOMO (68) lets each client interpolate other clients' local models with opportune weights learned during training. However, this method lacks both theoretical justifications for such linear combinations and convergence guarantees. Moreover, FedFOMO requires the presence of a powerful server able to 1) store all individual local models and iitractions-hiinee
2) learn for each client- chient-through repeated useful. Therefore, FedFOMO is not suited for cross-device FL where the number of clients may be very large (e.g., 10--10" participating clients (28) Table 2)) and given client may only participate in a single training round.
other clients' local models may be
Overall, although current personalization approaches can lead to superior empirical performance in comparison to a shared global model or individually trained local models, it is still not well understood whether and under which conditions clients are guaranteed to benefit from collaboration.
Our contributions. In this work, we first show that federated learning impossible without assumptions on local data distributions. Motivated by this negative result, we formulate a general and flexible assumption: the data distribution of each client is a mixture of M underlying distributions. The proposed formulation has the advantage that each client can benefit from knowledge distilled from all other clients' datasets (even if any two clients can be arbitrarily different from each other). We also show that this assumption encompasses most of the personalized FL approaches previously proposed in the literature.
In our framework, a personalized model is a linear combination of M shared component models. All clients jointly learn the M components, while each client learns its personalized mixture weights. We show that federated EM-like algorithms can be used for training. In particular, we propose FedEM and D-FedEM for the client-server and the fully decentralized settings, respectively, and we
2

---

prove convergence guarantees. Our approach also provides principled and efficient way to infer personalized models for clients unseen at training time. Our algorithms can easily be adapted to solve more general problems in a novel framework, which can be seen as federated extension of the centralized surrogate optimization approach [43]. To the best of our knowledge, our paper the first work to propose federated surrogate optimization a algorithms with convergence guarantees.
Through extensive experiments on FL benchmark datasets, we show that our approach generally yields models that 1) are on average more accurate, 2) are fairer across clients, and 3) generalize better to unseen clients than state-of-the of-tth--rt personalized and non-personalized FL approaches.
Paper outline. The rest of the paper is organized as follows. In Section 2 we provide our impossibilitty result, introduce our main assumptions, and show that several popular personalization approaches can be obtained as special cases of our framework. Section [ describes our algorithms, states their convergence results, and presents our general federated surrogate optimization framework. Finally, we provide experimental results in Section 4 before concluding in Section 5
2 Problem Formulation
We consider a (countable) set T of classification (or regression) tasks which represent the set of possible clients. We will use the terms task and client interchangeably. Data at client t E T is generated according to a local distribution D, over X' x y. Local data distributions (Dther are in general different, thus it is natural to fit separate model (hypothesis) h2 H to each data distribution Dr. The goal is then to solve (in parallel) the following optimization problems
Vt E T,
minimize LD, (ht),
(I)
hiEH
where he X + A 31 (AD denoting the unitary simplex of dimension D), L: AD' x y - R+ is a loss function and LD. (he) Exx.))>,, [[(he(x (x), y)] is the true risk of model füp under data distribution D,. For (x, y) E X' Vr we will denote the joint distribution density associated to DD by (x, v). and the marginal densities by (x(x) and pe (v).
A set of T clients [T] 2 (1, T} E T participate to the initial training phase; other clients may join the system in later stage. We denote by S, {s (8) (9) (2) ) the dataset at client t E [T] drawn i.i.d. from D., and by n = Σr= ny the total dataset size.
The idea of federated learning is to enable each client to benefit from data samples available at other clients in order to get better estimation of LD, and therefore get model with better generalization ability to unseen examples.
2.1 An Impossibility Result
We start by showing that some assumptions on the local distributions Pr(x, v), E T are needed for federated learning to be possible, i.e., for each client to be able to take advantage of the data at other clients. This holds even if all clients participate to the initial training phase (i.e.. T = [T').
We consider the classic PAC learning framework where we fix a class of models H. and seek a learning algorithm which is guaranteed, for all possible data distributions over X' x V, to return with high probability a model with expected error e-close to the best possible error in the class H. The worst-case sample complexity then refers to the minimum amount of labeled data required by any algorithm to reach a given c-approximation.
Our impossibility result for FL is based on a reduction to an impossibility result for Semi-Supervised Learning (SSL), which the problem of learning from a training set with only small amount of labeled data. The authors of 4 conjectured that, when the quantity of unlabeled data goes to infinity, the worst-case sample complexity of SSL improves over supervised learning at most by a constant factor that only depends on the hypothesis class [4] Conjecture 4]. This conjecture was later proved for the realizable case and hypothesis classes of finite VC dimension [13] Theorem 1], even when the marginal distribution over the domain set X' is known (21) Theorem 2
In the case of (multi-output) regression, we have h : X' + Rd for some d 2 and 7: Rd Rd + Rt. We note that whether the conjecture in holds in the agnostic case still open problem.
3

---

In the context of FL, if the marginal distributions Pt (x) are identical, but the conditional distributions Pt (3/x) can be arbitrarily different, then each client can learn using: 1) its own local labeled dataset, and 2) the other clients' datasets, but only as unlabeled ones (because their labels have no relevance for t). The FL. problem, with clients, then reduces to 7 parallel SSL. problems, or more precisely, it is at least as difficult T parallel SSL problems (because client has no direct access to the other local datasets but can only learn through the communication exchanges allowed by the FL algorithm), The SSL impossibility result implies that, without any additional assumption on the local distributions pt. (x, y) t E [T], any FL algorithm can reduce the sample complexity of client-t's problem in (I) only by a constant in comparison to local learning, independently of how many other clients participate to training and how large their datasets' sizes are.
2.2 Learning under a Mixture Model
Motivated by the above impossibility result, in this work we propose to consider that each local data distribution D, is a mixture of M underlying distributions Dm, 1 S m S M, as formalized below. Assumption 1. There exist M underlying (independent) distributions Dm S m S M, such that fort E T, De is mixture of the distributions {Dm J with weights = = m+ F E A M L.e.
24 ~ [( ((x,, yc) = m) ~ Dmn Vt T, where M(T) is a multinomial (categorical) distribution with parameters π.
(2)
Similarly to what was done above, we use Pmo (x, v). Pm (x), and Pm (v) to denote the probability distribution densities associated to Dm- We further assume that marginals over X are identical. Assumption 2. For all m E [M], we have pm (x) = p(x).
Assumption 2 is not strictly required for our analysis to hold, but, in the most general case, solving Problem 1 requires to learn generative models. Instead, under Assumption 2 we can restrict attention to discriminative models (e.g., neural networks). More specifically, we consider a parameterized set of models H with the following properties.
Assumption 3. H = [ is a set of hypotheses parameterized by Rd,, whose convex hull is in H. For each distribution Dm with m E [M], there exists a hypothesis hor,- such that
(ho=, (x), v) log Pm Pmmy)lx(1 (y)xx) +c,
(3)
where E IR is normalization constant. The function U(;5)) is then the log-loss associated to Pm (vlx).
We refer to the hypotheses H. as component models or simply components. We denote by O" E M the matrix whose m--th row is 0r and by II E A the matrix whose 1-th row is Ti E AM, Similarly, we will use θ and II to denote arbitrary parameters.
Remark 1. Assumptions 2 are mainly technical and are not required for our approach to work in practice. Experiments in Section 4 show that our algorithms perform well on standard FL benchmark datasets, for which these assumptions do not hold in general.
Note that, under the above assumptions, Pe (xx y) depends O" and कद. Moreover, we can prove (see App. A) that the optimal local model hi E H for client is weighted average of models in H. Proposition 2.1. Let I(..)) be the mean squared error loss, the logistic loss or the cross-entropy loss, and é and I' be a solution of the following optimization problem:
minimize 6,11
E
m(]]
(4)
t~oT (x.v)~D.
where DJ is any distribution with support Under Assumptions and 3 the predictors
M
h
= Σ
#tmlig (x), VLET
(5)
m=1
minimize E(x.a))D, (x), y) and thus solve Problem T
'A possible way to ensure that Assumption 2 holds is to use the batch normalization technique from account for feature shift.
40 to

---

Proposition 2.1 suggests the following approach to solve Problem T. First, we estimate the parame-
ters θ and #1 S 1 - T, by minimizing the empirical version of Problem (4) on the training data,
Le., minimizing:
f(e, II) A (ggp(Str/O,11)
log p(
(6) (0, me);
n
t=1
which is the (negative) likelihood of the probabilistic model (2) Second, we use (5) to get the client predictor for the clients present at training time. Finally, deal with a client (new & [7] not seen during training, we keep the mixture component models fixed and simply choose the weights The that maximize the likelihood of the client data and get the client predictor via (5)
2.3 Generalizing Existing Frameworks
Before presenting our federated learning algorithms in Section 3 we show that the generative model extends some popular multi-task sullft-taskpeersonaaizeed FL formulations the literature.
in Assumption T
Clustered Federated Learning [56, 20] assumes that each client belongs to one among C clusters and proposes that all clients in the same cluster learn the same model. Our framework recovers this scenario considering M C and ntc = if task (client) is in cluster and nic = otherwise.
relies on learning a global model hglab and T local models Mis,t, and then using each client the linear interpolation h2 = ag/ioe, + (1 a+) Niglob- Each client model can thus be seen as a linear combination of M T + 1 models htm No.. for Cty, and ,..
Personalization via model interpolation [4 14
m E [T] and ho = highab with specific weights के. = Qs, TTo
for 1' E [7] {4}.
Federated MTL via task relationships. The authors of (59) proposed to learn personalized models by solving the following optimization problem inspired from classic MTL formulations:
T 1he (xx ).
min W.
(2) (2)
+ Atr (WΩWT),
where lw, are linear predictors parameterized by the rows of matrix W and the matrix Ω captures task relationships (similarity). This formulation is motivated by the alternating structure optimization method (ASO) 22 70). In App. B we show that, when predictors he: are linear and have bounded norm, our framework leads to the same ASO formulation that motivated Problem 77.. Problem (7)
can also be justified by probabilistic priors (69) or graphical models (35) (see [59] App. B. B.1]). Similar
considerations hold for our framework (see again App. B Reference [67] extends the approach
[59] by letting each client learn a personalized model as a weighted combination of M known hypotheses. Our approach is more general and flexible as clients learn both the weights and the hypotheses. Finally, other personalized FI. algorithms, like pFedMe [16), FedU [17]. and those studied
in (24] and in [23). be framed as special cases of formulation Their assumptions thus also be seen as a particular case of our framework.
Federated Expectation-Mxxiizzttoo
3.1 Centralized Expectation-Maximizatii
Our goal is to estimate the optimal components' parameters O" = (0%))smm and mixture weights I1" (πf)s by minimizing the negative log-likelihood f(e, II) in (6. A natural approach to solve such non-convex problems is the Dnneectation-axxmizzttinn algorithm (EM), which alternates between two steps. Expectation steps update the distribution (denoted by q:)
over the latent variables 2(() for every data point (s) (x(") (8) Vi (8) given the current estimates of the parameters [0,11}. Maximization steps update the parameters (0, 11) by maximizing the expected log-likelihood, where the expectation is computed according to the current latent variables' distributions.
The following proposition provides the EM updates for our problem (proof in App- C
$ As the distribution D over tasks in Proposition 2.1 is arbitrary, positively weighted sum of clients'
empirical losses could be considered.
5

---

Proposition 3.1. Under Assumptions 7 and 2 at the k-th iteration the EM algorithm updates parameter estimates through the following steps:
E-step:
Ht
k+1
(xL
(1)
m.) n exp
( -1(hog (x .))) E [T], E [M], i E [na] (8)
M-step:
k+1 tm
Zis1 qt
++ (8)
(2+
m)
t E [2], m E [M] (9)
n4
arg min BERd
(2)
(0) (hx v(v())
m E [M] (10)
The EM updates in Proposition 3.1 have a natural interpretation. In the E-step. given current component models O and mixture weights I1 (8) updates the a-posteriori probability qf (2("
m) that point s(") of client was drawn from the m-th distribution based on the current mixture weight rim and on how well the corresponding component 0% classifies 8
(z) The M-step consists of two updates under fixed probabilities of+11 First, (9) updates the mixture weights π-- to reflect the prominence of each distribution Dm in S; as given by g Finally, 10) updates the components' parameters @++1 by solving M independent, weighted empirical risk minimization problems with weights given by d: g These weights aim to construct an unbiased estimate of the true risk over each underlying distribution D.,, using only points sampled from the client mixtures, similarly to importance sampling strategies used to learn from data with sample selection bias 61 T0
3.2 Client-Server Algorithm
Federated learning aims to train machine learning models directly on the clients, without exchanging raw data, and thus should run EM while assuming that only client t has access to dataset S,. The E-step (8) and the II update (9) in the M-step operate separately on each local dataset S, and can thus be performed locally at each client i. On the contrary, the @ update (TO) requires interaction with other clients, since the computation spans all data samples S1.T.
In this section, we consider client-server setting, in which each client i can communicate only with a centralized server (the orchestrator) and wants to learn components' parameters e* =(0m) (fm)) and its own mixture weights n²-
We propose the algorithm FedEM for Federated Stpectation-Maximizaaiii (Alg. FedEM pro- ceeds through communication rounds similarly to most FL algorithms including FedAvg (47). FedProx [38). SCAFFOLD [29]. and PFedMe [16]. At each round, 1) the central server broadcasts the (shared) component models the clients, 2) each client locally updates components and its personalized mixture weights, and 3) sends the updated components back to the server, 4) the server aggregates the updates. The local update performed at client t consists in performing the steps in (8) and [ and updating the local estimates of fm through a solver which approximates the exact minimization in (TO) using only the local dataset S; (see line 77. FedEM can operate with different different across clients as far as they satisty some local improvement guarantees
local solvers-evs (see the discussion in App. H.. In what follows, we restrict our focus on the practically important case where the local solver performs multiple stochastic gradient descent updates (local SGD (60)). Under the following standard assumptions (see e.g., (66)). FedEM converges to a stationary point of s.
Below, we use the more compact notation 1(0) s(.)) & ((ho(xx))) (()), ()))
Assumption 4. The negative -igeeiill f is bounded below by f E R.
Assumption 5. (Smoothness) For all E [T] and E [ni], the function 0 + 1(0; (()) is L-smooth and twice continuously differentiable.
Assumption 6. (Unbiased gradients and bounded variance) Each client E [T] can sample a random
batch & S₁ and compute an unbiased estimator (0,0g) of the local gradient with bounded
variance, i.e. Eg,g, (), 1 Em Vol(0; ((0) and Eg g (0, E)
ni Eie Vol(O; (0)) )112 < o².
1
Assumption 7. (Bounded dissimilarity) There exist 33 and G such that for any set of weights a E M;
M
am 1(0; (.00 10)
S G2 + 52 1
am 1(0 06 (1)
T

---

Algorithm 1: FedEM (see also the more detailed Alg. 2 in App. D.1
Input Data SI:T; number of mixture distributions M; number of communication rounds K Output: 0 m. m E [M]
for iterations k = ..... K do
server broadcasts 0% 1 S I M. to the T' clients; for tasks I,.... T in parallel over T' clients do
for component 1...... .M do
update of (2g = m) as in ViE nc. update Tm as 9
0m. < LocalSollvrr (m, 0k-1, g. S,); client sends 0 m.t 1 S m S M, to the server;
for component = ,,..,, M do
0
2 x Oh
10
2t=1 n
Assumption Ilimits the level of dissimilarity of the different tasks, similarly what done in 66
Theorem 3.2. Under Assumptions 7+7 when clients use SGD local solver with learning rate
4g after a large enough number of communication rounds K, FedEN's iterates satisfy: VK'
1
Σ
Vef
E
1' S
So
1 VR
1
Σ
K
II1I) s So
)
(11)
K
K
K3/4
where the expectation over the random batches samples, and Anf(O*, .Inh) - f (0*, I1*)
111++1)+++ 0.
Theorem 3.2 (proof in App. G. expresses the convergence of both sets of parameters (O and II) to a stationary point of f. Indeed, the gradient of f with respect to θ becomes arbitrarily small (left inequality (11) and the update Eq. 9 leads arbitrarily small improvements of (right
inequality in
We conclude this section observing that FedEM allows unseen client, i.e., client tne & [T] arriving after the distributed training procedure, to learn its personalized model. The client simply retrieves the learned components' parameters O and computes its personalized weights Tr. (starting for example from a uniform initialization) through one E-step (8) and the first update in the M-step (9)
3.3 Fully Decentralized Algorithm
In some cases, clients may want to communicate directly in a peer-to-peer fashion instead of relying on the central server mediation [see 28 Section 2.1]. In fact, fully decentralized schemes may provide stronger privacy guarantees 12 and speed-up training as they better use communication resources [41] 46 and reduce the effect of stragglers [50]. For these reasons, they have attracted
significant interest recently in the machine learning community [41 63 62 3 51
46 31].
We refer to 449 for comprehensive survey of fully decentralized optimization (also known as consensus-based optimization). and to 31 for a unified theoretical analysis of decentralized SGD.
We propose D-FedEM (Alg. 4 in App. D.2 a fully decentralized version of our federated expectation maximization algorithm. As in FedEM, the M-step for θ update is replaced by an approximate maximization step consisting of local updates. The global aggregation step in FedEM (Alg. line I0) replaced by partial aggregation step. where each client computes a weighted average of its current components and those of subset of clients (its neighborhood), which may vary over time. The convergence of decentralized optimization schemes requires certain assumptions to guarantee that each client can influence the estimates of other clients over time. In our paper, consider the general assumption in [1 Assumption 4] (restated as Assumption ii App. [ for completeness). For instance, this assumption is satisfied if the graph of clients' communications is strongly connected every 7 rounds.
D- D-FedEM converges to stationary point of f (formal statement in App. and proof in App. G.2

---

Theorem 3.3 (Informal). In the same setting of Theorem 3.2 and under the additional Assumption g D- D-FedEM's individual estimates 1 converge to common value O". Moreover, O and IIe converge to a stationary point of f.
3.4 Federated Surrogate Optimization
FedEM and D-FedEM can be seen as particular instances of a more general framework-of potential interest for other applications- that we call federated surrogate optimization.
The standard majorization majorization-miniiizzaoo principle (34) iteratively minimizes, at each iteration k, surrogate function وf majorizing the objective function f. The work (43) studied this approach when
each g² is a first-order surrogate of f (the formal definition from [43 is given in App. F.I.
Our novel federated surrogate optimization framework considers that the objective function of is Σi= wiJe of T functions and iteratively minimizes f in distributed fashion using partial first-order surrogates of for each function ft. "Partial" refers to the fact that of is not required to be a first order surrogate wrt the whole set of parameters, as defined formally below.
weighted sum
Definition (Partial first-order surrogate). A function g(u, v): Rd x V - R is a partial-first-one surrogate of f(u, v) wrt u near (uo, vo) E Rd when the following conditions are satisfied:
1. g(,, v) 2 fu, v) for all Rdd and v V:
2. r(u. v) - g(u - f(u, v) is differentiable and L-smooth with respect to u. Moreover, we have r(u,, Vo) = 0 and V Vur(up, Vu) = O.
3. g(u, vo) - g(u, v) dy (vo, for is non-negative and dv(v,
for all Rdd and v E arg vev gu, v'), where dy 0 = v=v.
Under the assumption that each client i can compute a partial first-order surrogate of ft, we propose algorithms for federated surrogate optimization in both the client- setting (Alg. 33 and the fully decentralized one (Alg. 5) and prove their convergence under mild conditions (App. G. and G.2). FedEM and D D-FedMM can be seen particular instances of these algorithms and Theorem. 3.2 and Theorem. 3.3 follow from the more general convergence results for federated surrogate optimization. We can also use our framework to analyze the convergence of other FL algorithms such as pFedMe [16], as we illustrate in App.
4 Experiments
Datasets and models. We evaluated our method on five federated benchmark datasets spanning a wide range of machine learning tasks: image classification (CIFAR10 and CIFAR100 (33 ), handwrit- ten character recognition (EMNIST (8) and FEMNIST (7) and language modeling (Shakespeare
Shakespeare dataset (resp. FEMNIST) was naturally partitioned by assigning all lines from the same characters (resp. all images from the same writer) to the same client. We created federated versions of CIFARIO and EMNIST by distributing samples with the same label across the clients according to a symmetric Dirichlet distribution with parameter 0.4, as in [65]. For CI- FAR100, we exploited the availability of "coarse" and "fine" labels, using a two-stage Pachinko allocation method (39) to assign 600 sample to each of the 100 clients, as in (54). We also evaluated our method on a synthetic dataset verifying Assumptions 3 For all tasks, we randomly split each local dataset into training (60%). validation (20%) and test (20%) sets. Table [ summarizes
47
datasets, models, and number of clients (more details can be found in App. I.1 Code is available at https://gi thub. com/ omarroq/////oooooooooooooooooooo tamarfoq/FedEM FedEM
Other FL approaches. We compared our algorithms with global models trained with FedAvg [47
and FedProx (38) as well as different personalization approaches: a personalized model trained only on the local dataset, FedAvg with local tuning (FedAvg+) (27). Clustered FL. (56) and pFedMe [16]. For each method and each task, the learning rate and the other hyperparameters were tuned via grid search (details in App. I.2p. FedAvg+ updated the local model through a single pass on the
local dataset. Unless otherwise stated, the number of components considered by FedEM was M = 3, training occurred over 80 communication rounds for Shakespeare and 200 rounds for all other datasets.
For training, sub- sub-sampled 10% and 15% from EMNIST and FEMNIST datasets respectively.

---

Table 1: Datasets and models (details in App. LT
Dataset
Task
Clients
Total samples
Model
FEMNIST
2
Handwrillen churacter recognition
539
120, 772
2/aayer CNN + 2-layer FFN
FMNIST
Handwritten character resognition
100
81, 425
2-layer CNN 2. FPN
CIFARIO
Image classification
80
60, 000
MobileNet-v22 55
CIFARIOU
Image classification
100
60,
MobiliNeet-2 P3
Shakespeure
47
Next Character Prediction
778
4, 226, 158
Slacked LSTM
Symheric
Rinary Classification
300
1, 570, 507
Linear
Table 2: Test accuracy: average across clients / /bottom decile.
Dataset
Local
Fedtvg (47
FedProx
38
Fedlvg++
Clustered FL 556
pFedNe 16
FedEx (Ours)
FEMNIST
71.0/
78.6
18.6/63.9 63.9
78.9 64.0
75.3 / 53.0
73.5 55.1 55.1
74.9 (..9/M 57.6
79.9
64.8
EMNIST
71.9/61.3
82.6/
83.0
83.1 75.8
82.7/75.0
83.3
76.4
76.6
CIFARIO
70.2/48.7
78.2 72.4
78.0 70.8
82.3 / 70.6
78.6/71.2
81.7 73.6
84.3
78.1
CIFARIO0
31.5/
10.9/33.2 33.2
33.2
39.0 39.0/28.3 28.3
11.5/31.1
11.8
11.5
44.1/
35.0
Shakespeare
32.0/ 16.6
46.7
42.8
45.7 / 41.9
40.0 25.5
46.6/40 42.7
41.2 / 36.8
46.7
43.0
Synthetic
65.7 58.4
68.2
58.9
68.2 / 59.0
68.9 / 60.2
69.1 59.0
69.2 / 61.2
74.7
66.7
At each round, clients train for one epoch. Results for D-FedEM are in App. J.1 MOCHA (59. which can only train linear models, is presented in App. J.2
A comparison with
Average performance of personalized models. The performance of each personalized model (which is the same for all clients in the case of FedAvg and FedProx) is evaluated the local test dataset (unseen at training). Table 2 shows the average weighted accuracy with weights proportional to local dataset sizes. We observe that FedEM obtains the best performance across all datasets.
Fairness across clients. FedEM's improvement terms of average accuracy could be the result of learning particularly good models for some clients at the expense of bad models for other clients. Table Shows the bottom decile of the accuracy of local models, i.e., the 11th worst accuracy (the minimum accuracy is particularly noisy, notably because some local test datasets are very small). Even clients with the worst personalized models are still better off when FedEM used for training.
Clients sampling, In cross-device federated learning, only subset of clients may be available at each round. We ran CIFARIO experiments with different levels of participation: at each round given fraction of all clients were sampled uniformly without replacement. We restrict the comparison FedEM and FedAvg+, as 1) FedAvg+ performed better than FedProx and FedAvg the previous CIFAR1O experiments, 2) it is not clear how to extend PFedMe and Clustered FL to handle client sampling. Results Fig. (left) show that FedEM more robust to low clients' participation levels. We provide additional results on client sampling, including comparison with APFL 14 in App. [.6
FedEM allows new clients arriving after the distributed training to easily learn their personalized models. With the exception of FedAvg+ it is not clear how the other personalized FL algorithms should be extended to tackle the same goal (see discussion in App. J.3). In order to evaluate the quality of new clients' personalized models, we performed an experiment where only 80% of the clients ("old" clients) participate to the training. The remaining 20% join the system in second phase and use their local training datasets to learn their personalized weights. Table 3 shows that FedEM allows new clients to learn a personalized model at least as good as FedAvg's global one and always better than FedAvg+'s 'ss one. Unexpectedly, new clients achieve sometimes a significantly higher test accuracy than old clients (e.g., 47.5% against 44. 1 5% on CIFAR 100). Our investigation in App. [.3 suggests that, by selecting their mixture weights on local datasets that were not used to train the components, new clients can compensate for potential overfitting in the initial training phase. We also investigate in App. J.3 the effect of the local dataset size on the accuracy achieved by unseen clients, showing that personalization is effective even when unseen clients have small datasets.
Generalization to unseen clients. As discussed in Section 3.2
Effect of M. A limitation FedEM is that each client needs update and transmit M components at each round, requiring roughly M times more computation and M times larger messages. Nevertheless, the number of components to consider practice quite limited. We used M 3 in our previous experiments, and Fig. [ (right) shows that larger values do not yield much improvement and M = already provides a significant level of personalization. In all experiments above, the number of communication rounds allowed all approaches to converge. As consequence, even if other methods

---

Table 3: Average test accuracy across clients unseen at training (train accuracy in parenthesis).
Fedkeg
FedAvg+
27
FedEN (Ours)
FEMNIST
(80.9)
(8..2)
79.1 (81.5)
FMNIST
88.4
83.7 (92.9)
84.0 (83.3)
CIFARIO
(77.
(80.5)
85. (90.7)
CIFARI00
(1.1(12.1)
(53.3)
47.5
Shakespeare
(47.1)
(93.0)
46.7 (46.6)
Synthetic
68.6
69.1 (72.1)
(74.1)
84
Test accuracy
Test accuracy
02
82
75
80
FedEM
76
FedAvg+
78
2
Sampling rate
M
Figure 1: Effect client sampling rate (left) and FedEM number mixture components M (right) on the test accuracy for CIFARIO [3
trained over M = 3 times more rounds in order to have as much computation and communication as FedEM-the conclusions would not change. As a final experiment, we considered time-constrained setting, where FedEM is limited to run one third (= 1/M) of the rounds (Table 7 in App. A.p. Even if FedEM does not reach its maximum accuracy, still outperforms the other methods on datasets.
Conclusion
In this paper, we proposed novel federated MTL approach based on the flexible assumption that local data distributions are mixtures underlying distributions. Our EM-like algorithms allow clients to jointly learn shared component models and personalized mixture weights in client-server and fully decentralized settings. We proved convergence guarantees for our algorithms through general federated surrogate optimization framework which be used to analyze other FL formulations. Extensive empirical evaluation shows that our approach learns models with higher accuracy and fairesss than state-of-the-an FL algorithms,
for clients not present at training time.
In future work, we aim reduce the local computation and communication our algorithms. Aside from standard compression schemes [22]. promising direction is to limit the number of component models that client updates/transmits each step. This could be done an adaptive manner based on the client's 's current mixture weights. A simultaneously published work [15] proposes federated EM algorithm (also called FedEM), which does not address personalization but reduces communication requirements by compressing appropriately defined complete data sufficient statistics.
second interesting research direction straints (quite unexplored until now with the notable exception of (3). Some features of our algorithms may be beneficial for privacy (e.g., the fact that personalized weights are kept locally and that all users contribute to all shared models). We hope to design differentially private versions of our algorithms and characterize their privacy-utility trade-offs.
study personalized FL approaches under privacy con-
Acknowledgements
This work has been supported by the French government, through the 3IA Côte d' Azur Investments the Future project managed by the National Research Agency (ANR) with the reference number and through grants ANR- ANR-16-CE23-0001 (Project PAMELA) and ANR-20- CF23-0015 (Project PRIDE). The authors are grateful to the OPAL infrastructure from Université Côte d' Azur for providing computational resources and technical support.
ANR-19-P3IA-000
10

---

References
[1] Durmus Alp Emre Acar al. "Debiasing Model Updates for Improving Personalized Federated Training". In: Proceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, July 2021, pp. 21-31. URR: https: htatp::/prroceedings mlr, press/ 1339/ acarala......... htmllll..lmm
[2] Ric Kubota Ando and Tong Zhang. "A Framework for Learning Predictive Structures from Multiple Tasks and Unlabeled Data". In: Journal of Machine Learning Research 6.61 (2005),
pp. 1817-1853.
[3] Aurélien Bellet, Rachid Guerraoui, Mahsa Taziki, and Mare Tommasi. "Personalized and Private Peer-to-Peer Machine Learning". In: AISTATS. 2018.
[4] Shai Ben-David, Tyler Lu, and D. Pál. "Does Unlabeled Data Provably Help? Worst-case Analysis of the Sample Complexity of Semi-Supervised Learning". In: COLT. 2008.
[5] Stephen Boyd, Persi Diaconis, and Lin Xiao. "Fastest Mixing Markov Chain on A Graph". In: SIAM REVIEW 46 (2003), pp. 667-689.
[6] Sébastien Bubeck. Convex Optimization: Algorithms and Complexity. 2015. arXiv: 1405, 4980
[caath.CC]
[7] Sebastian Caldas al. "Leaf: A benchmark for federated settings". In: arXiv preprint arXiv:1812.01099 (2018). Presented at the 2nd International Workshop on Federated Learning for Data Privacy and Confidentiality (in conjunction with NeurlPS 2019).
[8] Gregory Cohen. Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. "EMNIST: Extending MNIST to handwritten letters". In: 2017 International Joint Conference on Neural Networks (IJCNV). IEEE. 2017, pp. 2921-2926.
[9] Luca Corinzia and Joachim M. Buhmann. Variational Federated Multi-Tass Learning. 2019. arXiv: 1906. 06268 LG] [10] Corinna Cortes, Yishay Mansour, and Mehryar Mohri. "Learning Bounds for Impor- tance Weighting". In: Advances in Neural Information Processing Systems. Ed. by J. Laf- ferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta. Vol. 23. Curran Asso- ciates, Inc., 2010. URL: lhtpps::... / proceedding nwwrris....... cc caper / 2010 /ile/ 5S933016884462216668 25te//55268822225223e2r pdf
[11] Corinna Cortes, Mehryar Mohri, Michael Riley, and Afshin Rostamizadeh. "Sample Selection Bias Correction Theory". In: ALT. 2008
[12] Edwige Cyffers and Aurélien Bellet. Privacy Amplification by Decentralization. Presented
at the Privacy Preserving Machine Learning workshop (in conjunction with NeurlPS 2020).
2021. arXiv: 2012. 05326 [cs.LG]
[13] Malte Darnstädt, H. U. Simon, and Balázs Szörényi. "Unlabeled Data Does Provably Help".
In: STACS. 2013.
[14] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. "Adaptive Personalized Federated Learning". In: arXiv preprint arXiv:2003. 13461 (2020).
[15] Aymeric Dieuleveut, Gersende Fort, Eric Moulines, and Geneviève Robin. "Federated Expec- tation Maximization with heterogeneity mitigation and variance reduction". In: Advances in Neural Information Processing Systems. Vol. 34. 2021.
[16] Canh T Dinh, Nguyen H Tran, and Tuan Dung Nguyen. "Personalized Federated Learning with Moreau Envelopes". In: 34th Conference on Neural Information Processing Systems (NeurIPS 2020). 2020.
[17] Canh T Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang. "FedU: A Unified Framework for Federated Multi-Task Learning with Laplacian Regularization". In: arXiv preprint arXiv:2102.0m 04848 (2021).
[18] P. Erdös and Rényi. "On Random Graphs I". In: Publicationes Mathematicae Debrecen (1959), p. 290.
[19] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. "Personalized federated learning: A meta-learning approach". In: 34th Conference on Neural Information Processing Systems (NeurIPS 2020). 2020.
[20] Avishek Ghosh, Jichan Chung, Dong Yin. and Kannan Ramchandran, "An Efficient Framework for Clustered Federated Learning". In: NewrIPS. 2020.

---

[21] Christina Göpfert, Shai Ben- Den-David, Olivier Bousquet, Sylvain Gelly, Ilya Tolstikhin, and Ruth Urner. "When can unlabeled data improve the learning rate?" In: Conference on Learning Theory. PMLR. 2019, pp. 1500-1518.
[22] Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. "Federated Learning with Compression: Unified Analysis and Sharp Guarantees". In: ICML.
2021.
[23] Filip Hanzely, Slavomıér Hanzely, Samuel Horváth, and Peter Richtárik. "Lower bounds and optimal algorithms for personalized federated learning". In: 34th Conference on Neural Information Processing Systems (NewrIPS 2020). 2020.
[24] Filip Hanzely and Peter Richtárik. "Federated Learning of a Mixture of Global and Local Models". In: (2020). arXiv: 2002.. 05516 [cs LG]]
[25] Sepp Hochreiter and Jürgen Schmidhuber. "Long Short-Term Memory". In: Neural Computa-
tion 9.8 (1997), pp. 1735-1780.
[26] Yutao Huang et al. "Personalized cross-silo federated learning on non-iid data". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 35. 9. 2021, pp. 7865-7873.
[27] Yihan Jiang. Jakub Konečný, Keith Rush, and Sreeram Kannan. "Improving federated learning personalization via model agnostic meta learning". In: arXiv preprint arXiv: 1909. 12488 (2019). Presented at NeurlPS FL. workshop 2019.
[28] Peter Kairouz et al. "Advances and Open Problems in Federated Learning". In: Foundations and Trends@ in Machine Learning 14.1 2 (2021). ). pp. 1-210. ISSN: 1935-8237. DOI: 10
1561/2200000083 URL: http:://xx: doi org/ 10. 1561/2200000083
[29] Sai Pranceth Karimireddy, Satyen Kale, Mchryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. "SCAFFOLD: Stochastic controlled averaging for federated learning". In: International Conference on Machine Learning. PMLR. 2020, pp. 5132-5143. [30] Mikhail Khodak, Maria-Florina F Balcan, and Ameet Tallaakar. "Adaptive gradient-based meta-lcarning methods". In: Advances in Neural Information Processing Systems. 2019, pp. 5917-5928.
[31] Anastasia Koloskova, N. Loizou, Sadra Boreiri, M. Jaggi, and S. Stich. "A Unified Theory of Decentralized SGD with Changing Topology and Local Updates". In: ICML. 2020.
[32] Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. "Federated learning: Strategies for improving communication efficiency". In: arXiv preprint arXiv: 1610.5549 (2016). Presented at NIPS 2016 Workshop on Private Multi-Party Machine Learning.
[33]
Alex Krizhevsky. "Larnin multiple layers of features from tiny images". MSc thesis. 2009. Kenneth Lange, David R. Hunter, and Ilsoon Yang. "Optimization Transfer Using Surrogate Objective Functions", " In: Journal of Computational and Graphical Statistics 9.1 (2000), pp. 1-
[34]
20. ISSN: 10618600. URL: htt::htp::// //www/ jstor. org/ httg//table//e 1390060
[35] Steffeen L. Lauritzen. Graphical models. English. Oxford Statistical Science Series 17. Claren-
don Press, 1996. ISBN: 0198522193.
[36] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. "Ditto: Fair and robust federated learning through personalization". In: International Conference on Muchine Learning. PMLR. 2021. pp. 6357-6368.
[37] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. "Federated learning: Chal- lenges, methods, and future directions". In: IEEE Signal Processing Magazine 37.3 (2020), pp. 50-60.
[38] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. "Federated Optimization in Heterogeneous Networks". In: Third MLSys Conference. 2020.
[39] Wei Li and Andrew McCallum. "Pachinko Allocation: Struutuuttrrd Mixture Models of Topic Correlations". In: Proceedings of the 23rd International Conference on Machine Learning. ICML, '06. Pittsburgh, Pennsylvania, USA: Association for Computing Machinery, ISBN: 1595933832. DOI: 10. /145/1143844 1143917
2006, pp. 577-584.
URL: https
//doi owoooog//111 1115/ /t////143444... 1143917
[40] Xiaoxiao Li, Meirui JIANG, Xiaofei Zhang, Michael Kamp, and Qi Dou. "FedBN: Federated learning on Non-IIID Features via Local Batch Normalization". In: International Conference on Learning Representations. 2020.
12

---

[41] Xiangru Lian, Ce Zhang, Huan Zhang. Cho-Jui Hsieh, Wei Zhang, and Ji Liu. "Can Decentral- ized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent". In: Proceedings of the 31 International Conference on Neural Information Processing Systems. NIPS' 17. Long Beach, California, USA: Curran Associates Inc., 2017, pp. 5336-5346. ISBN: 97815100 510860964.
[42] Xiangru Lian, Wei Zhang, Ce Zhang, and Ji Liu. "Asynchronous Decentralized Parallel
Stochastic Gradient Descent". In: ICML. 2018.
[43] Julien Mairal. "Optimization with first-order surrogate functions". In: International Conference
on Machine Learning 2013, pp. 783-791.
[44] Yishay Mansour, Mehryar Mohri, Jac Ro, and Ananda Theertha Suresh. "Three approaches for personalization with applications to federated learning". In: arXiv preprint arXiv:2002.1000 (2020).
[45] Sébastien Marcel and Yann Rodriguez. "Torchvision the Machine- Package of Torch". In: Proceedings of the 18th ACM International Conference on Multimedia. MM '10. Firenze, Italy:
Association for Computing Machinery, 2010, pp. 1485- 1488. ISBN: 978 9781605589336. DOI:
10. 11111/117399111111 187 1874 URR: https: org/ 10. 1145/1873951. 1874254
[46]
Othmane Marfoq, Chuan Xu, Giovanni Neglia, and Richard Vidal. Marroughput-Optimma
Topology Design for Cross-Silo Federated Learning" In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin.
Vol. 33. Curran Associates, Inc., 2020, pp. 19478- 19487. URL: https: /tttps:/pprcceedinng neurips.. ctt.app.ccccppeee cttccc/pppppee pwwpp/app///// pttpp/pppppp/////e
pdf
[47] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. CCommunication-effiicii learning of deep networks from decentralized data", In: Artificial Intelligence and Statistics. PMLR. 2017, pp. 1273-1282.
[48] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. "Agnostic Federated Learning". In: International Conference on Machine Learning. 2019, pp. 4615-4625.
[49] A. Nedić, A. Olshevsky, and M. G. Rabbat. "Network Topology and Communication- Computation Tradeoff's in Decentralized Optimization"" In: Proceedings of the IEEE 106.5 (2018), pp. 953-976. DOT: 10. 1 1111/119990C// 2018. 28174611
[50] Giovanni Neglia, Gianmarco Calbi, Don Towsley, and Gayane Vardoyan. "The Role of Network Topology for Distributed Machine Learning". In: JEEE INFOCOM 2019 JEEE Conference on Computer Communications. 2019, pp. 2350-2358. Do1: [00 1001/ INFOCOM. 2019. 8737602
[51] Giovanni Neglia, Chuan Xu, Don Towsley, and Gianmarco Calbi. "Decentralized gradient methods: does topology matter?" In: AISTATS. 2020.
[52] Y. Nesterov. Introductory Lectures on Convex Optimization: A Basic Course. 1st ed. Applied Optimization. Springer, 2003. URL: http: //gen, lib. rus wc/book/////eeeeeeeeeeeec.cceece index poadexx.pppmddee php7md5-
48843c366629960211ff007 1675df 06776
[53] Adam Paszke et al. "PyTorch: An Imperative Style, High-Performance Deep Learning Li- brary". In: Advances in Neural Information Processing Systems 32. Ed. by H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. Fox, and R. Garnett. Curran Associates, Inc., 2019, pp. 8024-8035. imperative- wtw..oo---ggghppeooooooooeeeeee. www.perooooooeeeeeeeeeeee
wwtpprisssccpppppe/eee/ee URL: http: //ttpers neurips cc/ pwwpp 10aperr/9115-yytrchhh an
learning- 1ese--ibbrrrr. pdf
[54] Sashank J. Reddi al. "Adaptive Federated Optimization". " In: International Conference
on Learning Representations. 2021. URL: https: / openreview net / forum? ? id LKFG31B13US
[55] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang- Chen. "Mobilenetv2: Inverted residuals and linear bottlenecks" " In: Proceedings of the IEEE confer- ence on computer vision and pattern recognition. 2018, pp. 4510-4520.
[56] Felix Sattler, Klaus- Müller, and Wojciech Samek. "Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints". In: IEEE Transactions on Neural Networks and Learning Systems (2020).
[57] Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. "Personalized Federated Learning using Hypernetworks". In: Proceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang, Vol. 139. Proceedings of Machine Learning Research. PMLR, July 2021, pp. 9489-9502. URL: https: h/tpoee//pooeeeinggs mlr. press/ #139/shams saams i 2n21a. html
13

---

[58] Khe Chai Sim, Petr Zadrazil, and Françoise Beaufays. "An Investigation Into On-device Personalization of End-to-end Automatic Speech Recognition Models". In: INTERSPEECH. 2019.
[59] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. "Federated Multi- Task Learning". In: Proceedings of the 31st International Conference on Neural Information Processing Systems. NIPS' 17. Long Beach, California, USA: Curran Associates Inc., 2017,
pp. 4427-4437. ISBN: 97815 510860944
[60] Sebastian U Stich. "Local SGD Converges Fast and Communicates Little". In: International Conference on Learning Representations. 2018.
[61] Masashi Sugiyama, Shinichi Nakajima, Hisashi Kashima, Paul Buenau, and Motoaki Kawan- abc. "Direct Importance Estimation with Model Selection and Its Application Covariate Shift Adaptation". In: NIPS. 2008.
[62] Hanlin Tang, Xiangru Lian, Ming Yan, Ce Zhang, and Ji Liu. "D²; Decentralized Training over Decentralized Data". In: ICML. 2018.
[63] Paul Vanhaesebrouck, Aurélien Bellet, and Mare Tommasi. "Decentralized Collaborative Learning of Personalized Models over Networks". In: AISTATS. 2017.
[64] Robin Vogel, Mastane Achab, Stéphan Clémençon, and Charles Tillier. "Weighted Emprirical Risk Minimization: Transfer Learning based on Importance Sampling". In: ESANN. 2020.
[65] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khaza- eni. "Federated Learning with Matched Averaging". In: International Conference on Learning Representationss 2020. URL: /tttpp:://ppereeviiii
net/ f forum?dd BklgqgsFD
[66] Jianyu Wang. Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. "Tackling the objective inconsistency problem in heterogeneous federated optimization". In: 34th Conference Neural Information Processing Systems (NeurIPS 2020). 2020.
[67] Valentina Zantedeschi, Aurélien Bellet, and Marc Tommasi. "Fully Decentralized Joint Learn- ing of Personalized Models and Collaboration Graphs". In: ed. by Silvia Chiappa and Roberto Calandra. Vol. 108. Proceedings of Machine Learning Research. Online: PMLR, Aug. 2020, pp. 864-874. URL: URL: http: http:::/p/oceeednnggge/////////e mlr wwwss/ vess/v100 http//vv0//zatteeee cantedeschi20s httl
[68] Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M Alvarez. "Personalized Federated Learning with First Order Model Optimization". In: International Conference on Learning Representations. 2020.
[69] Yu Zhang and Dit Yan Yeung. "A Convex Formulation for Learning Task Relationships in Multi-task Learning.. : In: Proceedings of the 26th Conference on Uncertainty in Artificial Intelligence, UAI 2010. 2010, p. 733.
[70] Jiayu Zhou, Jianhui Chen, and Jieping Ye. "Clustered Multi-Task Learning Via Alternating Structure Optimization". In: Advances in Neural Information Processing Systems. Ed. by J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K. Q. Weinberger. Vol. 24, Curran Associates, Inc., 201 1. URL: https: http///pocceeneaae///eeeee/ae//a neurips wwtppppp/p/paiiiccceeee 201a//re/01//flee/ 325i887cccccc222233223444424266 Paper.
14


Reviews:
[{"review_id": "t5VwvqlzuRo", "reviewer": "Reviewer_mgAM", "summary": "The paper proposes an approach to learn local models in a federated learning setup where local distributions are assumed to be a mixture of a given set of ground distributions. The idea is to learn a ground model for each ground distribution and the mixture coefficients for the local distributions, and then obtain a local model as the weighted average of ground models (weights are given by the mixture coefficients). ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well-written, presents the contributions clearly, and discusses relevant related work. The proposed method is sound and novel. The idea of learning a set of distributions and their mixture coefficients is an interesting take on the non-iid problem in federated learning, especially given that the paper shows this to be a generalization of other frameworks. The empirical evaluation is solid with a decent number of datasets, a large number of clients (100), and a variety of model types. \n\nThe paper puts a lot into the appendix, which is fair, but I felt that some of the results (especially the fully decentralized algorithm and the federated surrogate optimization) feel a bit displaced in the main text and could be moved to the appendix to rather provide more details on the main contributions in the main text. \n\nOverall, the paper provides a solid contribution that is backed by theory and empirical results. I vote for acceptance.\n\nDetailed comments:\n\n- the negative result is presented too informally and in its formulation only holds for equal feature distributions. At the same time, the result seems quite trivial: if one client holds no relevant information for another client, then no form of collaboration can improve local learning. \n- way too many arxiv citations, including older ones (2016, 2018). For some, this is just due to lazy bibliography writing. E.g., [28] is actually published: Kang, Jiawen, et al. \"Incentive design for efficient federated learning in mobile networks: A contract theory approach.\" 2019 IEEE VTS Asia Pacific Wireless Communications Symposium (APWCS). IEEE, 2019. and [32] was published at the Workshop on Private Multi-Party Machine Learning.\n- In order to overcome the limitation of assumption 2, the authors might consider using the results in [1], which deal with different feature distributions p_m(x), but do not address different conditionals. \n\n[1] Li, Xiaoxiao, et al. \"FedBN: Federated Learning on Non-IID Features via Local Batch Normalization.\" International Conference on Learning Representations. 2020.\n\n----------------------------------------------------------------------------------\n\nI have read my fellow reviewer's reviews and the authors rebuttal. Many valid points of critique have been adequately addressed by the authors. While some points still remain open, e.g., the proper discussion of related work in PFL, I am still convinced that this is a good paper with an interesting take on the problem of non-iid local data. So, while this is surely not a perfect paper, I think it is a valuable contribution to the conference and I keep my vote on acceptance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper addresses the limitations of the approach. The paper discusses societal impact in sufficient detail.", "review_text": "", "extra_scores": {"ethical_concerns": "The potential ethical concerns are related to the societal impact and thus discussed together.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "RD8qd81ir3", "reviewer": "Reviewer_j2i1", "summary": "In this paper, the authors consider model personalisation as an important aspect of federated learning. Based on an impossibility result, the authors are motivated to learn a mixture model $p(y|x)=\\sum_m^M\\pi_m p(y|x, \\theta_m)$ as opposed to a single mapping $p(y|x,\\theta)$. As opposed to standard FedAvg approaches, the mixture weights $\\pi_m$ are local, i.e. specific to each client $t$. \nThe authors propose an EM-type algorithm for learning the parameters $\\theta_m$ and $\\pi_{m,t}$ in a federated environment, coining their algorithm FedEM. Furthermore they propose a decentralised variant of FedEM. \nThe authors show convergence proofs under appropriate assumptions and show that FedEM and D-FedEM are special cases of  more general federated surrogate optimization.\n\n== after rebuttal ==\n\nI have read the rebuttal and raised my score accordingly", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Thx for presenting this interesting paper in a relevant topic in FL research. I enjoyed reading it, it is presented concisely.\n\nGeneral paper outline:\nThe paper consists of several elements, the connection of which I would like to see improved. The main technical motivation for FedEM seems to be 2.1, i.e. the impossibility result. The authors state that \"The main consequence for FL is that, without further assumptions, a client cannot provably benefit from larger amounts of data available at other clients.\". How does this motivate a mixture modelling approach? Further, without having read [3], how is the following not true: Imagine client 1 learning a feature extractor through some unsupervised fashion on a small vs. a huge amount of data. Client 2 will benefit from the shared feature extractor trained on a large data-set more than from the model trained on the small data-set, presumably through its better learned representation. \nGoing further, the authors introduce their mixture model, before showing how it generalizes existing approaches. Comparing against [47], though, FedEM requires sending all M=C models to all clients, whereas [47] only ever communicates 1 model. \nThe authors proceed to describe learning in the mixture model, both in the centralised and in the federated setting. Then the authors introduce a decentralised variant as well as a generalised 'federated surrogate optimization' few, without showing experiments in the later section. It seems to me that the paper would be improved by moving those two extensions to the appendix / a separate paper, as it doesn't add to the main contribution of the paper, which is the mixture model optimised with EM in the FL setting. \n\nThe experimental section is characterised by breath of data-sets considered, as well as competing methods. \nThere are, however, a few trivial baselines that I would like to see included:\nThe authors consider personalised model accuracy, i.e. not the performance of a global model at the server (as in FedAvg). As such, the authors need to include the performance of FedAvg when the last global model is allowed to be fine-tuned on the local data-sets for an appropriate amount of steps. Assuming $D_{train,t} \\sim p(D_t)$ and  $D_{test,t} \\sim p(D_t)$, fine-tuning should improve the performance of the global model. Moreover, due to the limited size of local training-sets, overfitting needs to be taken into account - similarly to how the authors speculate about $\\pi_t$ to overfit due to its localised nature. \n\nAnother trivial baseline to compare FedEM(m) with is a model with m-times the parameter count. It needs to be established that the additional model performance is not due to additional model parameters and instead indeed due to the mixture formulation. Furthermore it would be interesting to see the entropy of $\\pi$ to understand if specialisation is happening. Is there a guarantee or at least an empirical observation that EM can recover a 'true' clustering, i.e. assuming all clients within a cluster share the same data-distribution? \n\nThe authors are talking about a synthetic dataset, but do not include what that is about (also not in the appendix). \n\nI find it peculiar that except for Figure 1, all experiments do not include client sub-sampling. Full participation is a strange assumption to make in FL and I would consider the experimental section much stronger if the authors consider the more established settings. E.g. FEMNIST usually has >3.4k clients, sampling around 100 clients per round. Otherwise the authors should re-focus their claims for the cross-silo FL setting. \nFurthermore the authors make the in my opinion very strong assumption that there will be a large-enough labelled training set available at a new client. What I would like to see is the average performance of new clients with zero labelled examples (and therefore uniform $\\pi$, adding progressively more labelled examples, up until the full local labelled training set is assumed to be available. \n\nI see that the code to reproduce experiments is attached in the supplementary, however I would appreciate a table of hyper parameters in the actual paper (appendix). \n\nI have not read the proofs of various propositions/theorems in the paper. \n\nIf the authors improve their empirical evaluation along the lines that I have suggested, I will consider raising my score.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discuss the 'Fairness' of their method, both in the experimental section as well as in Appendix L. While it is clear to me what the exact definition of fairness is that the authors relate to, I suggest acknowledging the overloaded nature and general unclear definition of that term, also in connection with FL. More specifically it might be the case that \"Even clients with the worst personalized models are still better off when FedEM is used for training\", however it might still not be 'fair' as some clients benefit more (absolutely or relatively) than others. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "QpY8aSPGL4l", "reviewer": "Reviewer_kvAM", "summary": "The authors propose a new approach for personalized federated learning (PFL) by formulating PFL as multi-task learning (MTL). by assuming that each local data distribution is a mixture of unknown underlying distributions the model aims to learn a linear combination of clients' models based on their data distribution.\n\nThe authors claim the following contributions:\n-  Showing that federated learning (FL) is impossible without assumptions on local data distributions\n- Proposing FedEM and D-FedEM an expectation-maximization (EM) based approaches for PFL\n- The proposed methods achieve better accuracy and generalize better than SOTA PFL methods.\n\n\n\n", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Pros/Cons**:\n\n[+] The paper is well written and structured.\n\n[+] Good Theoretical analysis of FedEM's convergence.\n\n[+] The proposed methods achieve higher performance compared to previous methods over different datasets and federation sizes.\n\n[-] The authors did not address more recent approaches in PFL [1,2] and did not compare to them. FedFomo[1] aims to learn a linear combination of clients' weights W/O assuming anything about underlying data distributions and pFedHN[2] performs better than pFedMe under identical data heterogeneity learning settings.\n\n[-] The authors should provide STD analysis on their experimental results (i.e. running on multiple seeds)\n\n[-] \"the learning rate and the other hyper-parameters were tuned via grid search\" - on \"Datasets and models\" subparagraph the authors describe the data splitting as follow \" For all tasks, we randomly split each local dataset into training (80%) and test (20%) sets. **This implies that the HPs were chosen based on the test split!**\n\n[-] Generalization to unseen clients - “it is not clear if and how the other personalized FL algorithms can tackle the same goal.” I totally disagree with this statement. PFL methods can handle situations in which new clients join the federation. Let us take pFedMe for example, after the training process is finished the Hub/Server obtains a global model which is optimized based on the federation. This model can be sent to the new clients to be fine-tuned locally. \n\n[-] All the clients in the federation take part in each communication round - this is a big limitation for large-scale federations, in real life FL systems federations can contain millions of clients such that training this federation will cost lots of computational power and money. Other PFL methods sample a subset of clients at each COM round (most of them sample $5$ clients per round).\nThis limitation raises another question, have the authors run all methods on the same number of communications? each COM round in FedEM performs $T$ communications and $T$ denotes the number of clients while for pFedMe $5$ communications occur per COM round. In FL we care about the number of communications it takes the model to converge.\n\n**References**:\n\n[1] Zhang, Michael, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez. \"Personalized Federated Learning with First Order Model Optimization.\" arXiv preprint arXiv:2012.08565 (2020).\n\n[2] Shamsian, Aviv, Aviv Navon, Ethan Fetaya, and Gal Chechik. \"Personalized Federated Learning using Hypernetworks.\" arXiv preprint arXiv:2103.04628 (2021).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors should discuss the possible societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "at least 5 hours"}, "extra_text": {}}, {"review_id": "KJurSs2DpcQ", "reviewer": "Reviewer_fDTX", "summary": "This paper proposed a federated multi-task learning approach to learn personalized models.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros.\nThis paper has proposed a theretical sound algorithm for federated MTL. The idea is valid with convergence results supprt.\n\nHowever, the work can be improved by addressing the following comments.\n\n1. Section 2.1 is not well-written and self-contained. The impossibility result is not formally established and only some intuitions are given. In addition, the SSL problem is not well defined and no citation is given as well. This makes me hard to justify the motivation of the proposed work.\n\n2. There are more recent published papers in personalized FL. For example,\n\na. FedBN: Federated Learning on Non-IID Features via Local Batch Normalization\n\nb. Personalized Cross-Silo Federated Learning on Non-IID Data\n\nThese two apporaches, APFL and SCAFFOLD are able to handle clients sampling. Thus it is better to compare with them than pFedMe and clustered FL.\n\n3. For the synethic data, the prospoed work should be compared with MOCHA which is the most closest work to the proposed work for linear model.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Not applied", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers are generally in favor of accepting the paper for its algorithmic and theoretical contributions on federated multi-task learning. Based on that, I recommend acceptance. However, please make sure to incorporate the reviewers' comments and the rebuttal into the final version. 
Decision: Accept (Poster)