Federated Multi-Task Learning under a Mixture of Distributions

1 Introduction . Federated Learning ( FL ) [ 28 ] allows a set of clients to collaboratively train models without sharing their local data . Standard FL approaches train a unique model for all clients [ 47 , 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 a necessity in many FL applications . Previous work on 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 [ 58 ] . In this case , the global model plays the role of a meta-model to be used as initialization for few-shot adaptation at each client . In particular , the connection between FL and Model Agnostic Meta Learning ( MAML ) [ 27 ] has been studied in [ 19 , 30 , 1 ] in order to build a 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 . 1 ] ) . An alternative 35th Conference on Neural Information Processing Systems ( NeurIPS 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 , 9 , 44 ] . However , if local distributions are far from the average distribution , a relevant global model does not exist and this approach boils down to every client learning only on its own local data . This issue is formally captured by the generalization bound in [ 14 , Theorem 1 ] . Clustered FL [ 56 , 20 , 44 ] addresses the potential lack of a 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 is 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 to learn personalized models in the federated setting and allows for more nuanced relations among clients ’ models [ 59 , 63 , 67 , 24 , 16 ] . The authors of [ 59 , 63 ] were the first to frame FL personalization as a MTL problem . In particular , 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-server scenario , while [ 63 , 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 MTL-based approaches [ 24 , 23 , 16 , 26 , 36 ] 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 to 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 2 ) learn for each client—through repeated interactions—which other clients ’ local models may be useful . Therefore , FedFOMO is not suited for cross-device FL where the number of clients may be very large ( e.g. , 105–107 participating clients [ 28 , Table 2 ] ) and a given client may only participate in a single training round . 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 is 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 prove convergence guarantees . Our approach also provides a 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 a federated extension of the centralized surrogate optimization approach in [ 43 ] . To the best of our knowledge , our paper is the first work to propose federated surrogate optimization 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-art personalized and non-personalized FL approaches . Paper outline . The rest of the paper is organized as follows . In Section 2 we provide our impossibility result , introduce our main assumptions , and show that several popular personalization approaches can be obtained as special cases of our framework . Section 3 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 ∈ T is generated according to a local distribution Dt over X × Y . Local data distributions { Dt } t∈T are in general different , thus it is natural to fit a separate model ( hypothesis ) ht ∈ H to each data distribution Dt . The goal is then to solve ( in parallel ) the following optimization problems ∀t ∈ T , minimize ht∈H LDt ( ht ) , ( 1 ) where ht : X 7→ ∆|Y| ( ∆D denoting the unitary simplex of dimension D ) , l : ∆|Y| × Y 7→ R+ is a loss function,1 and LDt ( ht ) = E ( x , y ) ∼Dt [ l ( ht ( x ) , y ) ] is the true risk of a model ht under data distribution Dt . For ( x , y ) ∈ X ×Y , we will denote the joint distribution density associated to Dt by pt ( x , y ) , and the marginal densities by pt ( x ) and pt ( y ) . A set of T clients [ T ] , { 1 , 2 , . . . T } ⊆ T participate to the initial training phase ; other clients may join the system in a later stage . We denote by St = { s ( i ) t = ( x ( i ) t , y ( i ) t ) } nt i=1 the dataset at client t ∈ [ T ] drawn i.i.d . from Dt , and by n = ∑T t=1 nt 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 a better estimation of LDt , and therefore get a model with a better generalization ability to unseen examples . 2.1 An Impossibility Result . We start by showing that some assumptions on the local distributions pt ( x , y ) , t ∈ 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 × Y , to return with high probability a model with expected error -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 -approximation . Our impossibility result for FL is based on a reduction to an impossibility result for Semi-Supervised Learning ( SSL ) , which is the problem of learning from a training set with only a 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 ] . 2 1In the case of ( multi-output ) regression , we have ht : X 7→ Rd for some d ≥ 1 and l : Rd × Rd 7→ R+ . 2We note that whether the conjecture in [ 4 ] holds in the agnostic case is still an open problem . In the context of FL , if the marginal distributions pt ( x ) are identical , but the conditional distributions pt ( y|x ) can be arbitrarily different , then each client t 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 T clients , then reduces to T parallel SSL problems , or more precisely , it is at least as difficult as T parallel SSL problems ( because client t 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 ∈ [ T ] , any FL algorithm can reduce the sample complexity of client-t ’ s problem in ( 1 ) 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 .

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$.  The 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.  The authors show convergence proofs under appropriate assumptions and show that FedEM and D-FedEM are special cases of  more general federated surrogate optimization.  == after rebuttal ==  I have read the rebuttal and raised my score accordingly