Othmane Marfoq
email: othmane.marfoq@inria.fr

Giovanni Neglia
email: giovanni.neglia@inria.fr

Aurélien Bellet
email: aurelien.bellet@inria.fr

Laetitia Kameni
email: laetitia.kameni@accenture.com

Richard Vidal
email: richard.vidal@accenture.com

Federated Multi-Task Learning under a Mixture of Distributions

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.

Introduction

Federated Learning (FL) [START_REF] Kairouz | Advances and Open Problems in Federated Learning[END_REF] allows a set of clients to collaboratively train models without sharing their local data. Standard FL approaches train a unique model for all clients [START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF][START_REF] Konečný | Federated learning: Strategies for improving communication efficiency[END_REF][START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF][START_REF] Sai | SCAFFOLD: Stochastic controlled averaging for federated learning[END_REF][START_REF] Mohri | Agnostic Federated Learning[END_REF]. However, as discussed in [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF], the existence of such a global model suited for all clients is at odds with the statistical heterogeneity observed across different clients [START_REF] Li | Federated learning: Challenges, methods, and future directions[END_REF][START_REF] Kairouz | Advances and Open Problems in Federated Learning[END_REF]. 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 [START_REF] Chai | An Investigation Into On-device Personalization of End-to-end Automatic Speech Recognition Models[END_REF]. 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) [START_REF] Jiang | Improving federated learning personalization via model agnostic meta learning[END_REF] has been studied in [START_REF] Fallah | Personalized federated learning: A meta-learning approach[END_REF][START_REF] Khodak | Adaptive gradient-based meta-learning methods[END_REF][START_REF] Alp | Debiasing Model Updates for Improving Personalized Federated Training[END_REF] 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 [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF]App. 1]). An alternative 35th Conference on Neural Information Processing Systems (NeurIPS 2021). arXiv:2108.10252v3 [cs.LG] 18 Feb 2022 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 [START_REF] Deng | Adaptive Personalized Federated Learning[END_REF][START_REF] Corinzia | Variational Federated Multi-Task Learning[END_REF][START_REF] Mansour | Three approaches for personalization with applications to federated learning[END_REF]. 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 [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF][START_REF] Ghosh | An Efficient Framework for Clustered Federated Learning[END_REF][START_REF] Mansour | Three approaches for personalization with applications to federated learning[END_REF] 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 [START_REF] Smith | Federated Multi-Task Learning[END_REF][START_REF] Vanhaesebrouck | Decentralized Collaborative Learning of Personalized Models over Networks[END_REF][START_REF] Zantedeschi | Fully Decentralized Joint Learning of Personalized Models and Collaboration Graphs[END_REF][START_REF] Hanzely | Federated Learning of a Mixture of Global and Local Models[END_REF][START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF]. The authors of [START_REF] Smith | Federated Multi-Task Learning[END_REF][START_REF] Vanhaesebrouck | Decentralized Collaborative Learning of Personalized Models over Networks[END_REF] 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 [START_REF] Smith | Federated Multi-Task Learning[END_REF] proposed the MOCHA algorithm for the client-server scenario, while [START_REF] Vanhaesebrouck | Decentralized Collaborative Learning of Personalized Models over Networks[END_REF][START_REF] Zantedeschi | Fully Decentralized Joint Learning of Personalized Models and Collaboration Graphs[END_REF] 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 [START_REF] Hanzely | Federated Learning of a Mixture of Global and Local Models[END_REF][START_REF] Hanzely | Lower bounds and optimal algorithms for personalized federated learning[END_REF][START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF][START_REF] Huang | Personalized cross-silo federated learning on non-iid data[END_REF][START_REF] Li | Ditto: Fair and robust federated learning through personalization[END_REF] 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, [START_REF] Shamsian | Personalized Federated Learning using Hypernetworks[END_REF] 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 [START_REF] Shamsian | Personalized Federated Learning using Hypernetworks[END_REF] 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 [START_REF] Zhang | Personalized Federated Learning with First Order Model Optimization[END_REF] 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., 10 5 -10 7 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 [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF]. 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.

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 D t over X × Y. Local data distributions {D t } t∈T are in general different, thus it is natural to fit a separate model (hypothesis) h t ∈ H to each data distribution D t . The goal is then to solve (in parallel) the following optimization problems

∀t ∈ T , minimize ht∈H L Dt (h t ), (1)
where h t : X → ∆ |Y| (∆ D denoting the unitary simplex of dimension D), l : ∆ |Y| × Y → R + is a loss function, 1 and L Dt (h t ) = E (x,y)∼Dt [l(h t (x), y)] is the true risk of a model h t under data distribution D t . For (x, y) ∈ X × Y, we will denote the joint distribution density associated to D t by p t (x, y), and the marginal densities by p t (x) and p t (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 S t = {s

(i) t = (x (i) t , y (i)
t )} nt i=1 the dataset at client t ∈ [T ] drawn i.i.d. from D t , and by n = T t=1 n t 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 L Dt , and therefore get a model with a better generalization ability to unseen examples.

An Impossibility Result

We start by showing that some assumptions on the local distributions p t (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 [START_REF] Ben-David | Does Unlabeled Data Provably Help? Worst-case Analysis of the Sample Complexity of Semi-Supervised Learning[END_REF] 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 [START_REF] Ben-David | Does Unlabeled Data Provably Help? Worst-case Analysis of the Sample Complexity of Semi-Supervised Learning[END_REF]Conjecture 4]. This conjecture was later proved for the realizable case and hypothesis classes of finite VC dimension [START_REF] Malte Darnstädt | Unlabeled Data Does Provably Help[END_REF]Theorem 1], even when the marginal distribution over the domain set X is known [21, Theorem 2]. 2In the context of FL, if the marginal distributions p t (x) are identical, but the conditional distributions p t (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 p t (x, y) , t ∈ [T ], any FL algorithm can reduce the sample complexity of client-t's problem in [START_REF] Alp | Debiasing Model Updates for Improving Personalized Federated Training[END_REF] 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.

Learning under a Mixture Model

Motivated by the above impossibility result, in this work we propose to consider that each local data distribution D t is a mixture of M underlying distributions Dm , 1 ≤ m ≤ M , as formalized below. Assumption 1. There exist M underlying (independent

) distributions Dm , 1 ≤ m ≤ M , such that for t ∈ T , D t is mixture of the distributions { Dm } M m=1 with weights π * t = [π * t1 , . . . , π * tM ] ∈ ∆ M , i.e. z t ∼ M(π * t ), ((x t , y t ) |z t = m) ∼ Dm , ∀t ∈ T , (2)
where M(π) is a multinomial (categorical) distribution with parameters π.

Similarly to what was done above, we use p m (x, y), p m (x), and p m (y) to denote the probability distribution densities associated to Dm . We further assume that marginals over X are identical. Assumption 2. For all m ∈ [M ], we have p m (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 our attention to discriminative models (e.g., neural networks). 3 More specifically, we consider a parameterized set of models H with the following properties.

where c ∈ R is a normalization constant. The function l(•, •) is then the log-loss associated to p m (y|x).

We refer to the hypotheses in H as component models or simply components. We denote by Θ * ∈ R M ×d the matrix whose m-th row is θ * m , and by Π * ∈ ∆ T ×M the matrix whose t-th row is π * t ∈ ∆ M . Similarly, we will use Θ and Π to denote arbitrary parameters. Remark 1. Assumptions 2-3 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, p t (x, y) depends on Θ * and π * t . Moreover, we can prove (see App. A) that the optimal local model h * t ∈ H for client t is a weighted average of models in H. Proposition 2.1. Let l(•, •) be the mean squared error loss, the logistic loss or the cross-entropy loss, and Θ and Π be a solution of the following optimization problem: minimize

Θ,Π E t∼D T E (x,y)∼Dt [-log p t (x, y|Θ, π t )] , (4)
where D T is any distribution with support T . Under Assumptions 1, 2, and 3, the predictors

h * t = M m=1 πtm h θm (x) , ∀t ∈ T (5)
minimize E (x,y)∼Dt [l(h t (x), y)] and thus solve Problem (1).

Proposition 2.1 suggests the following approach to solve Problem [START_REF] Alp | Debiasing Model Updates for Improving Personalized Federated Training[END_REF]. First, we estimate the parameters Θ and πt , 1 ≤ t ≤ T , by minimizing the empirical version of Problem (4) on the training data, i.e., minimizing:

f (Θ, Π) - log p(S 1:T |Θ, Π) n - 1 n T t=1 nt i=1 log p(s (i) t |Θ, π t ), (6)
which is the (negative) likelihood of the probabilistic model (2). 4 Second, we use [START_REF] Boyd | Fastest Mixing Markov Chain on A Graph[END_REF] to get the client predictor for the T clients present at training time. Finally, to deal with a client t new / ∈ [T ] not seen during training, we keep the mixture component models fixed and simply choose the weights π tnew that maximize the likelihood of the client data and get the client predictor via (5).

Generalizing Existing Frameworks

Before presenting our federated learning algorithms in Section 3, we show that the generative model in Assumption 1 extends some popular multi-task/personalized FL formulations in the literature.

Clustered Federated Learning [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF][START_REF] Ghosh | An Efficient Framework for Clustered Federated Learning[END_REF] 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 π * tc = 1 if task (client) t is in cluster c and π * tc = 0 otherwise. Personalization via model interpolation [START_REF] Mansour | Three approaches for personalization with applications to federated learning[END_REF][START_REF] Deng | Adaptive Personalized Federated Learning[END_REF] relies on learning a global model h glob and T local models h loc,t , and then using at each client the linear interpolation h t = α t h loc,t + (1 -α t )h glob . Each client model can thus be seen as a linear combination of M = T + 1 models h m = h loc,m for m ∈ [T ] and h 0 = h glob with specific weights π * tt = α t , π * t0 = 1 -α t , and π * tt = 0 for t ∈ [T ] \ {t}. Federated MTL via task relationships. The authors of [START_REF] Smith | Federated Multi-Task Learning[END_REF] proposed to learn personalized models by solving the following optimization problem inspired from classic MTL formulations:

min W,Ω T t=1 nt i=1 l(h wt (x (i) t ), y (i) t ) + λ tr (W ΩW ) , (7)
where h wt 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) [START_REF] Kubota | A Framework for Learning Predictive Structures from Multiple Tasks and Unlabeled Data[END_REF][START_REF] Zhou | Clustered Multi-Task Learning Via Alternating Structure Optimization[END_REF]. In App. B, we show that, when predictors h θ * m are linear and have bounded norm, our framework leads to the same ASO formulation that motivated Problem [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF]. Problem [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF] can also be justified by probabilistic priors [START_REF] Zhang | A Convex Formulation for Learning Task Relationships in Multi-task Learning[END_REF] or graphical models [START_REF] Steffen | Graphical models. English[END_REF] (see [START_REF] Smith | Federated Multi-Task Learning[END_REF]App. B.1]). Similar considerations hold for our framework (see again App. B). Reference [START_REF] Zantedeschi | Fully Decentralized Joint Learning of Personalized Models and Collaboration Graphs[END_REF] extends the approach in [START_REF] Smith | Federated Multi-Task Learning[END_REF] 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 FL algorithms, like pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF], FedU [START_REF] Canh T Dinh | FedU: A Unified Framework for Federated Multi-Task Learning with Laplacian Regularization[END_REF], and those studied in [START_REF] Hanzely | Federated Learning of a Mixture of Global and Local Models[END_REF] and in [START_REF] Hanzely | Lower bounds and optimal algorithms for personalized federated learning[END_REF], can be framed as special cases of formulation [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF]. Their assumptions can thus also be seen as a particular case of our framework.

3 Federated Expectation-Maximization

Centralized Expectation-Maximization

Our goal is to estimate the optimal components' parameters Θ * = (θ * m ) 1≤m≤M and mixture weights Π * = (π * t ) 1≤t≤T by minimizing the negative log-likelihood f (Θ, Π) in [START_REF] Bubeck | Convex Optimization: Algorithms and Complexity[END_REF]. A natural approach to solve such non-convex problems is the Expectation-Maximization algorithm (EM), which alternates between two steps. Expectation steps update the distribution (denoted by q t ) over the latent variables z (i) t for every data point s

(i) t = (x (i) t , y (i) t )
given the current estimates of the parameters {Θ, Π}. Maximization steps update the parameters {Θ, Π} 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). Proposition 3.1. Under Assumptions 1 and 2, at the k-th iteration the EM algorithm updates parameter estimates through the following steps:

E-step: q k+1 t (z (i) t = m) ∝ π k tm • exp -l(h θ k m (x (i) t ), y (i) t ) , t ∈ [T ], m ∈ [M ], i ∈ [n t ] (8)
M-step:

π k+1 tm = nt i=1 q k+1 t (z (i) t = m) n t , t ∈ [T ], m ∈ [M ] (9) θ k+1 m ∈ arg min θ∈R d T t=1 nt i=1 q k+1 t (z (i) t = m)l h θ (x (i) t ), y (i) t , m ∈ [M ] (10)
The EM updates in Proposition 3.1 have a natural interpretation. In the E-step, given current component models Θ k and mixture weights Π k , (8) updates the a-posteriori probability q k+1 t (z

(i) t = m) that point s (i)
t of client t was drawn from the m-th distribution based on the current mixture weight π k tm and on how well the corresponding component θ k m classifies s (i)

t . The M-step consists of two updates under fixed probabilities q k+1 t . First, [START_REF] Corinzia | Variational Federated Multi-Task Learning[END_REF] updates the mixture weights π k+1 t to reflect the prominence of each distribution Dm in S t as given by q k+1 t . Finally, [START_REF] Cortes | Learning Bounds for Importance Weighting[END_REF] updates the components' parameters Θ k+1 by solving M independent, weighted empirical risk minimization problems with weights given by q k+1 t . These weights aim to construct an unbiased estimate of the true risk over each underlying distribution Dm using only points sampled from the client mixtures, similarly to importance sampling strategies used to learn from data with sample selection bias [START_REF] Sugiyama | Direct Importance Estimation with Model Selection and Its Application to Covariate Shift Adaptation[END_REF][START_REF] Cortes | Sample Selection Bias Correction Theory[END_REF][START_REF] Cortes | Learning Bounds for Importance Weighting[END_REF][START_REF] Vogel | Weighted Emprirical Risk Minimization: Transfer Learning based on Importance Sampling[END_REF].

Client-Server Algorithm

Federated learning aims to train machine learning models directly on the clients, without exchanging raw data, and thus we should run EM while assuming that only client t has access to dataset S t . The E-step [START_REF] Cohen | EMNIST: Extending MNIST to handwritten letters[END_REF] and the Π update ( 9) in the M-step operate separately on each local dataset S t and can thus be performed locally at each client t. On the contrary, the Θ update (10) requires interaction with other clients, since the computation spans all data samples S 1:T .

In this section, we consider a client-server setting, in which each client t can communicate only with a centralized server (the orchestrator) and wants to learn components' parameters Θ * = (θ * m ) 1≤m≤M and its own mixture weights π * t . We propose the algorithm FedEM for Federated Expectation-Maximization (Alg. 1). FedEM proceeds through communication rounds similarly to most FL algorithms including FedAvg [START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF], FedProx [START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF], SCAFFOLD [START_REF] Sai | SCAFFOLD: Stochastic controlled averaging for federated learning[END_REF], and pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF]. At each round, 1) the central server broadcasts the (shared) component models to 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 ( 9) and updating the local estimates of θ m through a solver which approximates the exact minimization in (10) using only the local dataset S t (see line 7). FedEM can operate with different local solvers-even different across clients-as far as they satisfy some local improvement guarantees (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 [START_REF] Sebastian | Local SGD Converges Fast and Communicates Little[END_REF]). Under the following standard assumptions (see e.g., [START_REF] Wang | Tackling the objective inconsistency problem in heterogeneous federated optimization[END_REF]), FedEM converges to a stationary point of f . Below, we use the more compact notation l(θ; s

(i) t ) l(h θ (x (i) t ), y (i) t )
1 K K k=1 E ∇ Θ f Θ k , Π k 2 F ≤ O 1 √ K , 1 K K k=1 ∆ Π f (Θ k , Π k ) ≤ O 1 K 3/4 , ( 11
)
where the expectation is over the random batches samples, and

∆ Π f (Θ k , Π k ) f Θ k , Π k - f Θ k , Π k+1 ≥ 0.
Theorem 3.2 (proof in App. G.1) expresses the convergence of both sets of parameters (Θ and Π) to a stationary point of f . Indeed, the gradient of f with respect to Θ becomes arbitrarily small (left inequality in [START_REF] Cortes | Sample Selection Bias Correction Theory[END_REF]) and the update in Eq. ( 9) leads to arbitrarily small improvements of f (right inequality in [START_REF] Cortes | Sample Selection Bias Correction Theory[END_REF]).

We conclude this section observing that FedEM allows an unseen client, i.e., a client t new / ∈ [T ] arriving after the distributed training procedure, to learn its personalized model. The client simply retrieves the learned components' parameters Θ K and computes its personalized weights π tnew (starting for example from a uniform initialization) through one E-step [START_REF] Cohen | EMNIST: Extending MNIST to handwritten letters[END_REF] and the first update in the M-step (9).

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 [START_REF] Cyffers | Privacy Amplification by Decentralization[END_REF] and speed-up training as they better use communication resources [START_REF] Lian | Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent[END_REF][START_REF] Marfoq | Throughput-Optimal Topology Design for Cross-Silo Federated Learning[END_REF] and reduce the effect of stragglers [START_REF] Neglia | The Role of Network Topology for Distributed Machine Learning[END_REF]. For these reasons, they have attracted significant interest recently in the machine learning community [START_REF] Lian | Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent[END_REF][START_REF] Vanhaesebrouck | Decentralized Collaborative Learning of Personalized Models over Networks[END_REF][START_REF] Lian | Asynchronous Decentralized Parallel Stochastic Gradient Descent[END_REF][START_REF] Tang | D 2 : Decentralized Training over Decentralized Data[END_REF][START_REF] Bellet | Personalized and Private Peer-to-Peer Machine Learning[END_REF][START_REF] Neglia | Decentralized gradient methods: does topology matter?[END_REF][START_REF] Marfoq | Throughput-Optimal Topology Design for Cross-Silo Federated Learning[END_REF][START_REF] Koloskova | A Unified Theory of Decentralized SGD with Changing Topology and Local Updates[END_REF]. We refer to [START_REF] Nedić | Network Topology and Communication-Computation Tradeoffs in Decentralized Optimization[END_REF] for a comprehensive survey of fully decentralized optimization (also known as consensus-based optimization), and to [START_REF] Koloskova | A Unified Theory of Decentralized SGD with Changing Topology and Local Updates[END_REF] 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. 1, line 10) is replaced by a partial aggregation step, where each client computes a weighted average of its current components and those of a 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, we consider the general assumption in [START_REF] Koloskova | A Unified Theory of Decentralized SGD with Changing Topology and Local Updates[END_REF]Assumption 4] (restated as Assumption 8 in App. E for completeness). For instance, this assumption is satisfied if the graph of clients' communications is strongly connected every τ rounds. D-FedEM converges to a stationary point of f (formal statement in App. E and proof in App. G.2). Theorem 3.3 (Informal). In the same setting of Theorem 3.2 and under the additional Assumption 8, D-FedEM's individual estimates (Θ k t ) 1≤t≤T converge to a common value Θk . Moreover, Θk and Π k converge to a stationary point of f .

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-minimization principle [START_REF] Lange | Optimization Transfer Using Surrogate Objective Functions[END_REF] iteratively minimizes, at each iteration k, a surrogate function g k majorizing the objective function f . The work [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF] studied this approach when each g k is a first-order surrogate of f (the formal definition from [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF] is given in App. F.1).

Our novel federated surrogate optimization framework considers that the objective function f is a weighted sum f = T t=1 ω t f t of T functions and iteratively minimizes f in a distributed fashion using partial first-order surrogates g k t for each function f t . "Partial" refers to the fact that g k t is not required to be a first order surrogate wrt the whole set of parameters, as defined formally below. Definition 1 (Partial first-order surrogate). A function g(u, v) : R du × V → R is a partial-first-order surrogate of f (u, v) wrt u near (u 0 , v 0 ) ∈ R du × V when the following conditions are satisfied:

1. g(u, v) ≥ f (u, v) for all u ∈ R du and v ∈ V; 2. r(u, v) g(u, v) -f (u, v
) is differentiable and L-smooth with respect to u. Moreover, we have r(u 0 , v 0 ) = 0 and

∇ u r(u 0 , v 0 ) = 0. 3. g(u, v 0 ) -g(u, v) = d V (v 0 , v) for all u ∈ R du and v ∈ arg min v ∈V g(u, v ), where d V is non-negative and d V (v, v ) = 0 ⇐⇒ v = v .
Under the assumption that each client t can compute a partial first-order surrogate of f t , we propose algorithms for federated surrogate optimization in both the client-server setting (Alg. 3) and the fully decentralized one (Alg. 5) and prove their convergence under mild conditions (App. G.1 and G.2).

FedEM and D-FedEM can be seen as 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 [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF], as we illustrate in App. F.3.

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 [START_REF] Krizhevsky | Learning multiple layers of features from tiny images[END_REF]), handwritten character recognition (EMNIST [START_REF] Cohen | EMNIST: Extending MNIST to handwritten letters[END_REF] and FEMNIST [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF]), 5 and language modeling (Shakespeare [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF][START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF]). 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 CIFAR10 and EMNIST by distributing samples with the same label across the clients according to a symmetric Dirichlet distribution with parameter 0.4, as in [START_REF] Wang | Federated Learning with Matched Averaging[END_REF]. For CI-FAR100, we exploited the availability of "coarse" and "fine" labels, using a two-stage Pachinko allocation method [START_REF] Li | Pachinko Allocation: DAG-Structured Mixture Models of Topic Correlations[END_REF] to assign 600 sample to each of the 100 clients, as in [START_REF] Sashank | Adaptive Federated Optimization[END_REF]. We also evaluated our method on a synthetic dataset verifying Assumptions 1-3. For all tasks, we randomly split each local dataset into training (60%), validation (20%) and test (20%) sets. Table 1 summarizes datasets, models, and number of clients (more details can be found in App. I.1). Code is available at https://github.com/omarfoq/FedEM.

Other FL approaches. We compared our algorithms with global models trained with FedAvg [START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF] and FedProx [START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF] as well as different personalization approaches: a personalized model trained only on the local dataset, FedAvg with local tuning (FedAvg+) [START_REF] Jiang | Improving federated learning personalization via model agnostic meta learning[END_REF], Clustered FL [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF] and pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF].

For each method and each task, the learning rate and the other hyperparameters were tuned via grid search (details in App. I. 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 on 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 in 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 2 shows the bottom decile of the accuracy of local models, i.e., the (T /10)-th 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 is used for training.

Clients sampling. In cross-device federated learning, only a subset of clients may be available at each round. We ran CIFAR10 experiments with different levels of participation: at each round a given fraction of all clients were sampled uniformly without replacement. We restrict the comparison to FedEM and FedAvg+, as 1) FedAvg+ performed better than FedProx and FedAvg in the previous CIFAR10 experiments, 2) it is not clear how to extend pFedMe and Clustered FL to handle client sampling. Results in Fig. 1 (left) show that FedEM is more robust to low clients' participation levels. We provide additional results on client sampling, including a comparison with APFL [START_REF] Deng | Adaptive Personalized Federated Learning[END_REF], in App. J.6.

Generalization to unseen clients. As discussed in Section 3.2, 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 a 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 one. Unexpectedly, new clients achieve sometimes a significantly higher test accuracy than old clients (e.g., 47.5% against 44.1% on CIFAR100). Our investigation in App. J.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.

Effect of M .

A limitation of FedEM is that each client needs to 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 in practice is quite limited. We used M = 3 in our previous experiments, and Fig. 1 (right) shows that larger values do not yield much improvement and M = 2 already provides a significant level of personalization. In all experiments above, the number of communication rounds allowed all approaches to converge. As a consequence, even if other methods Figure 1: Effect of client sampling rate (left) and FedEM number of mixture components M (right) on the test accuracy for CIFAR10 [START_REF] Krizhevsky | Learning multiple layers of features from tiny images[END_REF].

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 a time-constrained setting, where FedEM is limited to run one third (= 1/M ) of the rounds (Table 7 in App. J.5). Even if FedEM does not reach its maximum accuracy, it still outperforms the other methods on 3 datasets.

Conclusion

In this paper, we proposed a novel federated MTL approach based on the flexible assumption that local data distributions are mixtures of 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 a general federated surrogate optimization framework which can be used to analyze other FL formulations.

Extensive empirical evaluation shows that our approach learns models with higher accuracy and fairness than state-of-the-art FL algorithms, even for clients not present at training time.

In future work, we aim to reduce the local computation and communication of our algorithms. Aside from standard compression schemes [START_REF] Haddadpour | Federated Learning with Compression: Unified Analysis and Sharp Guarantees[END_REF], a promising direction is to limit the number of component models that a client updates/transmits at each step. This could be done in an adaptive manner based on the client's current mixture weights. A simultaneously published work [START_REF] Dieuleveut | Federated Expectation Maximization with heterogeneity mitigation and variance reduction[END_REF] proposes a federated EM algorithm (also called FedEM), which does not address personalization but reduces communication requirements by compressing appropriately defined complete data sufficient statistics.

A second interesting research direction is to study personalized FL approaches under privacy constraints (quite unexplored until now with the notable exception of [START_REF] Bellet | Personalized and Private Peer-to-Peer Machine Learning[END_REF]). 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.

Acknowledgements

This work has been supported by the French government, through the 3IA Côte d'Azur Investments in the Future project managed by the National Research Agency (ANR) with the reference number ANR-19-P3IA-0002, and through grants ANR-16-CE23-0016 (Project PAMELA) and ANR-20-CE23-0015 (Project PRIDE). The authors are grateful to the OPAL infrastructure from Université Côte d'Azur for providing computational resources and technical support.

Appendix

A Proof of Proposition 2.1

For h ∈ H and (x, y) ∈ X × Y, let p h (y|x) denote the conditional probability distribution of y given x under model h, i.e.,

p h (y|x) e c h (x) × exp -l (h (x) , y) , (12)
where

c h (x) -log y∈Y exp -l (h (x) , y) d y . (13)
We also remind that the entropy of a probability distribution q over Y is given by

H (q) - y∈Y q (y) • log q (y) d y, (14)
and that the Kullback-Leibler divergence between two probability distributions q 1 and q 2 over Y is given by

KL (q 1 ||q 2 ) y∈Y q 1 (y) • log q 1 (y) q 2 (y) d y. ( 15
)
Proposition 2.1. Let l(•, •) be the mean squared error loss, the logistic loss or the cross-entropy loss, and Θ and Π be a solution of the following optimization problem:

minimize Θ,Π E t∼D T E (x,y)∼Dt [-log p t (x, y|Θ, π t )] , (4)
where D T is any distribution with support T . Under Assumptions 1, 2, and 3, the predictors

h * t = M m=1 πtm h θm , ∀t ∈ T (5)
minimize E (x,y)∼Dt [l(h t (x), y)] and thus solve Problem (1).

Proof. We prove the result for each of the three possible cases of the loss function. We verify that c h does not depend on h in each of the three cases, then we use Lemma A.3 to conclude.

Mean Squared Error Loss This is the case of a regression problem where Y = R d for some d > 0. For x, y ∈ X × Y and h ∈ H, we have

p h (y|x) = 1 (2π) d • exp - h (x) -y 2 2 , (16)
and

c h (x) = -log (2π) d ( 17
)
Logistic Loss This is the case of a binary classification problem where Y = {0, 1}. For x, y ∈ X × Y and h ∈ H, we have

p h (y|x) = (h (x)) y • (1 -h (x)) 1-y , (18)
and

c h (x) = 0 (19)
Cross-entropy loss This is the case of a classification problem where Y = [L] for some L > 1.

For x, y ∈ X × Y and h ∈ H, we have

p h (y|x) = L l=1 (h (x)) 1 {y=l} , (20)
and

c h (x) = 0 (21) Conclusion For t ∈ T , consider a predictor h * t minimizing E (x,y)∼Dt [l(h t (x), y)]. Using Lemma A.3, for (x, y) ∈ X × Y, we have p h * t (y|x) = M m=1 πtm • p m y|x, θm . (22)
We multiply both sides of this equality by y and we integrate over y ∈ Y. Note that in all three cases we have ∀x ∈ X ,

y∈Y y • p h (•|x) d y = h(x). (23)
It follows that

h * t = M m=1 πtm h θm , ∀t ∈ T . ( 24
)
Supporting Lemmas Lemma A.1. Suppose that Assumptions 1 and 3 hold, and consider Θ and Π to be a solution of Problem (4). Then

p t (x, y| Θ, πt ) = p t (x, y|Θ * , π * t ), ∀t ∈ T . (25)
Proof. For t ∈ T ,

E (x,y)∼Dt -log p t (x, y| Θ, πt ) (26)
= -

(x,y)∈X ×Y p t (x, y|Θ * , π * t ) • log p t (x, y| Θ, πt ) d x d y (27)
= -

(x,y)∈X ×Y p t (x, y|Θ * , π * t ) • log p t (x, y| Θ, πt ) p t (x, y|Θ * , π * t ) d x d y - (x,y)∈X ×Y p t (x, y|Θ * , π * t ) • log p t (x, y|Θ * , π * t ) d x d y (28) = KL p t (•|Θ * , π * t ) p t • | Θ, πt + H [p t (•|Θ * , π * t )] , (29)
Since the KL divergence is non-negative, we have

E (x,y)∼Dt -log p t (x, y| Θ, πt ) ≥ H [p t (•|Θ * , π * t )] = E (x,y)∼Dt [-log p t (x, y|Θ * , π * t )] . (30)
Taking the expectation over t ∼ D T , we write

E t∼D T E (x,y)∼Dt -log p t (x, y| Θ, πt ) ≥ E t∼D T E (x,y)∼Dt [-log p t (x, y|Θ * , π * t )] . (31)
Since Θ and Π is a solution of Problem (4), we also have

E t∼D T E (x,y)∼Dt -log p t (x, y| Θ, πt ) ≤ E t∼D T E (x,y)∼Dt [-log p t (x, y|Θ * , π * t )] . (32)
Combining ( 31), [START_REF] Konečný | Federated learning: Strategies for improving communication efficiency[END_REF], and (29), we have

E t∼D T KL p t (•|Θ * , π * t ) p t • | Θ, πt = 0. ( 33
)
Since KL divergence is non-negative, and the support of D T is the countable set T , it follows that

∀t ∈ T , KL p t (•|Θ * , π * t ) p t • | Θ, πt = 0. (34)
Thus,

p t (x, y| Θ, πt ) = p t (x, y|Θ * , π * t ), ∀t ∈ T . ( 35
)
Lemma A.2. Consider M probability distributions on Y, that we denote q m , m ∈ [M ], and α = (α 1 , . . . , α m ) ∈ ∆ M . For any probability distribution q over Y, we have

M m=1 α m • KL q m M m =1 α m • q m ≤ M m=1 α m • KL (q m q) , (36)
with equality if and only if,

q = M m=1 α m • q m . (37)
Proof.

M m=1 α m • KL (q m q) - M m=1 α m • KL q m M m =1 α m • q m = M m=1 α m • KL (q m q) -KL q m M m =1 α m • q m (38) = - M m=1 α m y∈Y q m (y) • log q (y) M m =1 α m • q m (y) (39) = - y∈Y M m=1 α m • q m (y) • log q (y) M m =1 α m • q m (y) d y (40) = KL M m=1 α m • q m q ≥ 0. (41)
The equality holds, if and only if,

q = M m=1 α m • q m . ( 42
)
Lemma A.3. Consider Θ and Π to be a solution of Problem (4). Under Assumptions 1, 2, and 3, if c h does not depend on h ∈ H, then the predictors h * t , t ∈ T , minimizing

E (x,y)∼Dt [l(h t (x), y)], verify for (x, y) ∈ X × Y p h * t (y|x) = M m=1 πtm • p m y|x, θm . (43)
Proof. For t ∈ T and h t ∈ H, under Assumptions 1, 2, and 3, we have

E (x,y)∼Dt [l(h t (x), y)] = x,y∈X ×Y l(h t (x), y) • p t (x, y|Θ * , π * t ) d x d y. (44)
Using Lemma A.1, it follows that

E (x,y)∼Dt [l(h t (x), y)] = x,y∈X ×Y l(h t (x), y) • p t x, y| Θ, πt d x d y. (45)
Thus, using Assumptions 1 and 2 we have,

E (x,y)∼Dt [l(h t (x), y)] (46) = x,y∈X ×Y l(h t (x), y) • p t x, y| Θ, πt d x d y (47) = x,y∈X ×Y l(h t (x), y) • M m=1 πtm • p m y|x, θm p (x) d x d y (48)
= x∈X c ht (x) + M m=1 πtm • H p m •|x, θm p (x) d x + x∈X M m=1 πtm • KL p m • |x, θm p ht (•|x) p (x) d x. (52)
Let h • t be a predictor satisfying the following equality:

p h • t (y|x) = M m=1 πtm • p m y|x, θm .
Using Lemma A.2, we have

M m=1 πtm • KL p m • |x, θm p ht (•|x) ≥ M m=1 πtm • KL p m • |x, θm p h • t (•|x) (53)
with equality if and only if

p ht (•|x) = p h • t (•|x) . (54)
Since c h does not depend on h, replacing (53) in ( 52), it follows that

E (x,y)∼Dt [l(h t (x), y)] ≥ E (x,y)∼Dt [l(h • t (x), y)] . (55)
This inequality holds for any predictor h t and in particular for h * t ∈ arg min E (x,y)∼Dt [l(h t (x), y)], for which it also holds the opposite inequality, then:

E (x,y)∼Dt [l(h * t (x), y)] = E (x,y)∼Dt [l(h • t (x), y)] , (56)
and the equality implies that

p h * t (•|x) = p h • t (•|x) = M m=1 πtm • p m •|x, θm . (57)

B Relation with Other Multi-Task Learning Frameworks

In this appendix, we give more details about the relation of our formulation with existing frameworks for (federated) MTL sketched in Section 2.3. We suppose that Assumptions 1-3 hold and that each client learns a predictor of the form [START_REF] Boyd | Fastest Mixing Markov Chain on A Graph[END_REF]. Note that this is more general than [START_REF] Zantedeschi | Fully Decentralized Joint Learning of Personalized Models and Collaboration Graphs[END_REF], where each client learns a personal hypothesis as a weighted combination of a set of M base known hypothesis, since the base hypothesis and not only the weights are learned in our case.

Alternating Structure Optimization [START_REF] Zhou | Clustered Multi-Task Learning Via Alternating Structure Optimization[END_REF]. Alternating structure optimization (ASO) is a popular MTL approach that learns a shared low-dimensional predictive structure on hypothesis spaces from multiple related tasks, i.e., all tasks are assumed to share a common feature space P ∈ R d ×d , where d ≤ min(T, d) is the dimensionality of the shared feature space and P has orthonormal columns (P P = I d ), i.e., P is semi-orthogonal matrix. ASO leads to the following formulation: minimize

W,P :P P =I d T t=1 nt i=1 l h wt x (i) t
, y

(i) t + α (tr (W W ) -tr (W P P W )) + β tr (W W ) , (58)
where α ≥ 0 is the regularization parameter for task relatedness and β ≥ 0 is an additional L2 regularization parameter.

When the hypothesis (h θ ) θ are assumed to be linear, Eq. ( 5) can be written as W = ΠΘ. Writing the LQ decomposition6 of matrix Θ, i.e., Θ = LQ, where L ∈ R M ×M is a lower triangular matrix and Q ∈ R M ×d is a semi-orthogonal matrix (QQ = I M ), ( 5) becomes

W = ΠLQ ∈ R T ×d , thus, W = W Q Q, leading to the constraint W -W Q Q 2 F = tr (W W ) -tr (W Q QW ) = 0. If we assume θ m 2 2 to be bounded by a constant B > 0 for all m ∈ [M ], we get the constraint tr (W W ) ≤ T B. It means that minimizing T t=1 nt i=1 l h wt x (i) t , y (i) t
under our Assumption 1 can be formulated as the following constrained optimization problem minimize

W,Q:QQ =I M T t=1 nt i=1 l h wt x (i) t , y (i) t , subject to tr {W W } -tr {W Q QW } = 0, tr (W W ) ≤ T B. (59)
Thus, there exists Lagrange multipliers α ∈ R and β > 0, for which Problem ( 59) is equivalent to the following regularized optimization problem minimize

W,Q:QQ =I M T t=1 nt i=1 l h wt x (i) t , y (i) t + α (tr {W W } -tr {W Q QW }) + β tr {W W } , (60)
which is exactly Problem [START_REF] Chai | An Investigation Into On-device Personalization of End-to-end Automatic Speech Recognition Models[END_REF].

Federated MTL via task relationships. The ASO formulation above motivated the authors of [START_REF] Smith | Federated Multi-Task Learning[END_REF] to learn personalized models by solving the following problem

min W,Ω T t=1 nt i=1 l h wt x (i) t , y (i) t + λ tr (W ΩW ) , (61)
Two alternative MTL formulations are presented in [START_REF] Smith | Federated Multi-Task Learning[END_REF] to justify Problem (61): MTL with probabilistic priors [START_REF] Zhang | A Convex Formulation for Learning Task Relationships in Multi-task Learning[END_REF] and MTL with graphical models [START_REF] Steffen | Graphical models. English[END_REF]. Both of them can be covered using our Assumption 1 as follows:

• Considering T = M and Π = I M in Assumption 1 and introducing a prior on Θ of the form

Θ ∼ N 0, σ 2 I d MN (I d ⊗ Ω) (62)
lead to a formulation similar to MTL with probabilistic priors [START_REF] Zhang | A Convex Formulation for Learning Task Relationships in Multi-task Learning[END_REF].

• Two tasks t and t are independent if π t , π t = 0, thus using Ω t,t = π t , π t leads to the same graphical model as in [START_REF] Steffen | Graphical models. English[END_REF].

Several personalized FL formulations, e.g., pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF], FedU [START_REF] Canh T Dinh | FedU: A Unified Framework for Federated Multi-Task Learning with Laplacian Regularization[END_REF] and the formulation studied in [START_REF] Hanzely | Federated Learning of a Mixture of Global and Local Models[END_REF] and in [START_REF] Hanzely | Lower bounds and optimal algorithms for personalized federated learning[END_REF], are special cases of formulation [START_REF] Tang | D 2 : Decentralized Training over Decentralized Data[END_REF].

C Centralized Expectation Maximization

Proposition 3.1. Under Assumptions 1 and 2, at the k-th iteration the EM algorithm updates parameter estimates through the following steps:

E-step: q k+1 t (z (i) t = m) ∝ π k tm • exp -l(h θ k m (x (i) t ), y (i) t ) , t ∈ [T ], m ∈ [M ], i ∈ [n t ] (8)
M-step:

π k+1 tm = nt i=1 q k+1 t (z (i) t = m) n t , t ∈ [T ], m ∈ [M ] (9) θ k+1 m ∈ arg min θ∈R d T t=1 nt i=1 q k+1 t (z (i) t = m)l h θ (x (i) t ), y (i) t , m ∈ [M ] (10)
Proof. The objective is to learn parameters { Θ, Π} from the data S 1:T by maximizing the likelihood p (S 1:T |Θ, Π). We introduce functions q t (z), t ∈ [T ] such that q t ≥ 0 and M z=1 q t (z) = 1 in the expression of the likelihood. For Θ ∈ R M ×d and Π ∈ ∆ T ×M , we have

log p(S 1:T |Θ, Π) = T t=1 nt i=1 log p t s (i) t |Θ, π t (63) = T t=1 nt i=1 log   M m=1   p t s (i) t , z (i) t = m|Θ, π t q t z (i) t = m   q t z (i) t = m   (64)
≥ T t=1 nt i=1 M m=1 q t z (i) t = m log p t s (i) t , z (i) t = m|Θ, π t q t z (i) t = m (65)
= T t=1 nt i=1 M m=1 q t z (i) t = m log p t s (i) t , z (i)
t = m|Θ, π t - T t=1 nt i=1 M m=1 q t z (i) t = m log q t z (i) t = m (66) L(Θ, Π, Q 1:T ), (67)
where we used Jensen's inequality because log is concave. L(Θ, Π, Q 1:T ) is an evidence lower bound.

The centralized EM-algorithm corresponds to iteratively maximizing this bound with respect to Q 1:T (E-step) and with respect to {Θ, Π} (M-step).

E-step.

The difference between the log-likelihood and the evidence lower bound L(Θ, Π, Q 1:T ) can be expressed in terms of a sum of KL divergences:

logp(S 1:T |Θ, Π) -L(Θ, Π, Q 1:T ) = = T t=1 nt i=1    log p t s (i) t |Θ, π t - M m=1 q t z (i) t = m log p t s (i) t , z (i) t = m|Θ, π t q t z (i) t = m    (68) = T t=1 nt t=1 M m=1 q t z (i) t = m   log p t s (i) t |Θ, π t -log p t s (i) t , z (i) t = m|Θ, π t q t z (i) t = m   (69) = T t=1 nt t=1 M m=1 q t z (i) t = m log p t s (i) t |Θ, π t • q t z (i) t = m p t s (i) t , z (i) t = m|Θ, π t (70) = T t=1 nt t=1 M m=1 q t z (i) t = m log q t z (i) t = m p t z (i) t = m|s (i) t , Θ, π t (71) = T t=1 nt i=1 KL q t z (i) t ||p t z (i) t |s (i) t , Θ, π t ≥ 0. (72)
For fixed parameters {Θ, Π}, the maximum of L(Θ, Π, Q 1:T ) is reached when

T t=1 nt i=1 KL q t z (i) t ||p t z (i) t |s (i) t , Θ, π t = 0.
Thus for t ∈ [T ] and i ∈ [n t ], we have:

q t (z (i) t = m) = p t (z (i) t = m|s (i) t , Θ, π t ) (73) = p t (s (i) t |z (i) t = m, Θ, π t ) × p t (z (i) t = m|Θ, π t ) p t s (i) t |Θ, π t (74) = p m (s (i) t |θ m ) × π tm M m =1 p m (s (i) t ) × π tm (75) = p m y (i) t |x (i) t , θ m × p m x (i) t × π tm M m =1 p m y (i) t |x (i) t , θ m × p m x (i) t × π tm (76) = p m y (i) t |x (i) t , θ m × p x (i) t × π tm M m =1 p m y (i) t |x (i) t , , θ m × p x (i) t × π tm , (77)
where (77) relies on Assumption 2. It follows that q t (z

(i) t = m) = p t (z (i) t = m|s (i) t , Θ, π t ) = p m y (i) t |x (i) t , θ m × π tm M m =1 p m y (i) t |x (i) t , θ m × π tm . (78)
M-step. We maximize now L(Θ, Π, Q 1:T ) with respect to {Θ, Π}. By dropping the terms not depending on {Θ, Π} in the expression of L(Θ, Π, Q 1:T ) we write:

L(Θ, Π, Q 1:T ) = T t=1 nt i=1 M m=1 q t z (i) t = m log p t s (i) t , z (i) t = m|Θ, π t + c (79) = T t=1 nt i=1 M m=1 q t z (i) t = m log p t s (i) t |z (i) t = m, Θ, π t + log p t z (i) t = m|Θ, π t + c (80) = T t=1 nt i=1 M m=1 q t z (i) t = m log p θm s (i) t + log π tm + c (81) = T t=1 nt i=1 M m=1 q t z (i) t = m log p θm y (i) t |x (i) t + log p m x (i) t + log π tm + c (82) = T t=1 nt i=1 M m=1 q t z (i) t = m log p θm y (i) t |x (i) t + log π tm + c , (83) (84)
where c and c are constant not depending on {Θ, Π}.

Thus, for t ∈ [T ] and m ∈ [M ], by solving a simple optimization problem we update π tm as follows:

π tm = nt i=1 q t (z (i) t = m) n t . ( 85
)
On the other hand, for m ∈ [M ], we update θ m by solving:

θ m ∈ arg min θ∈R d T t=1 nt i=1 q t (z (i) t = m) × l h θ (x (i) t ), y (i) t . ( 86
q k t z (i) t = m ← π k tm •exp -l(h θ k m (x (i) t ),y (i) t ) M m =1 π k tm •exp -l(h θ k m (x (i) t ),y (i) t )
;

// M-step 10 π k tm ← n t i=1 q k t (z (i) t =m) nt ; θ k m,t ← LocalSolver(J, m, θ k-1 m , q k t , S t ) ; client t sends θ k m,t , 1 ≤ m ≤ M to the server; for component m = 1, . . . , M do θ k m ← T t=1 nt n • θ k m,t ;
Function LocalSolver(J, m, θ, q, S):

for j = 0, . . . , J -1 do Sample indexes I uniformly from 1, . . . , |S|;

θ ← θ -η k-1,j i∈I q(z (i) = m) • ∇ θ l h θ x (i) , y (i) ; return θ; Algorithm 3: Federated Surrogate Optimization Input : u 0 ∈ R du ; V 0 = v 0 t 1≤t≤T ∈ V T ; number of iterations K; number of local steps J Output : u K ; v K t
for iterations k = 1, . . . , K do server broadcasts u k-1 to the T clients; for tasks t = 1, . . . , T in parallel over T clients do Compute partial first-order surrogate function

g k t of f t near u k-1 , v k-1 t ; v k t ← arg min v∈V g k t u k-1 , v ; u k t ← LocalSolver(J, u k-1 t , v k-1 t , g k t , S t );
client t sends u k t to the server;

u k ← T t=1 ω t • u k t ;
Function LocalSolver(J, u, v, g, S):

for j = 0, . . . , J -1 do sample ξ k-1,j from S; u ← u -η k-1,j • ∇ u g(u, v; ξ k-
q k t z (i) t = m ← π k tm •exp -l(h θ k m (x (i) t ),y (i) t ) M m =1 π k tm •exp -l(h θ k m (x (i) t ),y (i) t )
;

// M-step 10 π k tm ← n t i=1 q k t (z (i) t =m) nt ; 11 θ k-1 2 m,t ← LocalSolver(J, m, θ k-1 m,t , q k t , S t , t); Send θ k-1 2 m,t , 1 ≤ m ≤ M to neighbors; Receive θ k-1 2 m,s , 1 ≤ m ≤ M from neighbors; for component m = 1, . . . , M do θ k m,t ← T s=1 w k-1 s,t • θ k-1 2 m,s ;
Function LocalSolver(J, m, θ, q, S, t):

for j = 0, . . . , J -1 do Sample indexes I uniformly from 1, . . . , |S|;

θ ← θ -nt n • η k-1,j i∈I q(z (i) = m) • ∇ θ l h θ x (i) , y (i) ;
return θ;

Algorithm 5: Fully-Decentralized Federated Surrogate Optimization

Input : u 0 ∈ R du ; V 0 = v 0 t 1≤t≤T ∈ V T ; number of iterations K; number of local step J; mixing matrix distributions W k for k ∈ [K] Output : u K t for t ∈ [T ]; v K t for t ∈ [T ] for iterations k = 1, . . . , K do
// Select the communication topology and the aggregation weights

Sample W k-1 ∼ W k-1 ;
for tasks t = 1, . . . , T in parallel over T clients do compute partial first-order surrogate function

g k t of f t near u k-1 t , v k-1 t ; v k t ← arg min v∈V g k t u k-1 t , v ; u k-1 2 t ← LocalSolver(J, u k-1 t , v k-1 t , g k t , t); Send u k-1 2 t to neighbors; Receive u k-1 2 s
from neighbors;

u k t ← T s=1 w k-1 ts × u k-1 2 s ;
Function LocalSolver(J, u, v, g, S, t):

for j = 0, . . . , J -1 do sample ξ k-1,j from S ; u ← u -ω t • η k-1,j ∇ u g(u, v, ξ k-1,j );
return u;

E Details on the Fully Decentralized Setting

As mentioned in Section 3.3, the convergence of decentralized optimization schemes requires certain assumptions on the sequence of mixing matrices (W k ) k>0 , to guarantee that each client can influence the estimates of other clients over time. In our paper, we consider the following general assumption. Assumption 8 ([31, Assumption 4]). Symmetric doubly stochastic mixing matrices are drawn at each round k from (potentially different) distributions W k ∼ W k and there exists two constants p ∈ (0, 1], and integer τ ≥ 1 such that for all Ξ ∈ R M ×d×T and all integers l ∈ {0, . . . , K/τ }:

E ΞW l,τ -Ξ 2 F ≤ (1 -p) Ξ -Ξ 2 F , (87)
where W l,τ W (l+1)τ -1 . . . W lτ , Ξ Ξ 11 T , and the expectation is taken over the random distributions W k ∼ W k . Assumption 8 expresses the fact that the sequence of mixing matrices, on average and every τ communication rounds, brings the values in the columns of Ξ closer to their row-wise average (thereby mixing the clients' updates over time). For instance, the assumption is satisfied if the communication graph is strongly connected every τ rounds, i.e., the graph ([T ], E), where the edge (i, j) belongs to the graph if w h i,j > 0 for some h ∈ {k + 1, . . . , k + τ } is connected. We provide below the rigorous statement of Theorem 3.3, which was informally presented in Section 3.3. It shows that D-FedEM converges to a consensus stationary point of f (proof in App. G.2). Theorem 3.3. Under Assumptions 1-8, when clients use SGD as local solver with learning rate η = a0 √ K , D-FedEM's iterates satisfy the following inequalities after a large enough number of communication rounds K:

1 K K k=1 E ∇ Θ f Θk , Π k 2 F ≤ O 1 √ K , 1 K K k=1 T t=1 n t n KL π k t , π k-1 t ≤ O 1 K , (88)
where Θk = Θ k 1 , . . . Θ k T • 11 T . Moreover, individual estimates Θ k t 1≤t≤T converge to consensus, i.e., to Θk :

min k∈[K] E T t=1 Θ k t -Θk 2 F ≤ O 1 √ K .

F Federated Surrogate Optimization

In this appendix, we give more details on the federated surrogate optimization framework introduced in Section 3.4. In particular, we provide the assumptions under which Alg. 3 and Alg. 5 converge.

We also illustrate how our framework can be used to study existing algorithms.

F.1 Reminder on Basic (Centralized) Surrogate Optimization

In this appendix, we recall the (centralized) first-order surrogate optimization framework introduced in [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF]. In this framework, given a continuous function f : R d → R, we are interested in solving

min θ∈R d f (θ)
using the majoration-minimization scheme presented in Alg. 6.

Algorithm 6: Basic Surrogate Optimization

Input : θ 0 ∈ R d ; number of iterations K; Output : θ K for iterations k = 1, . . . , K do Compute g k , a surrogate function of f near θ k-1 ;
Update solution:

θ k ∈ arg min θ g k (θ);
This procedure relies on surrogate functions, that approximate well the objective function in a neighborhood of a point. Reference [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF] focuses on first-order surrogate functions defined below. Definition F.1 (First-Order Surrogate [START_REF] Mairal | Optimization with first-order surrogate functions[END_REF]). A function g : R d → R is a first order surrogate of f near θ k ∈ R d when the following is satisfied:

• Majorization: we have g(θ ) ≥ f (θ ) for all θ ∈ arg min θ∈R d g(θ). When the more general condition g ≥ f holds, we say that g is a majorant function.

• Smoothness: the approximation error r g -f is differentiable, and its gradient is L-Lipschitz. Moreover, we have r(θ k ) = 0 and ∇r(θ k ) = 0.

F.2 Novel Federated Version

As discussed in Section 3.4, our novel federated surrogate optimization framework minimizes an objective function (u, v 1:T ) → f (u, v 1:T ) that can be written as a weighted sum f (u, v 1:T ) = T t=1 ω t f t (u, v t ) of T functions. We suppose that each client t ∈ [T ] can compute a partial first order surrogate of f t , defined as follows. Definition 1 (Partial first-order surrogate). A function g(u, v) : R du × V → R is a partial first-order surrogate of f (u, v) wrt u near (u 0 , v 0 ) ∈ R du × V when the following conditions are satisfied:

1. g(u, v) ≥ f (u, v) for all u ∈ R du and v ∈ V; 2. r(u, v) g(u, v) -f (u, v) is differentiable and L-smooth with respect to u. Moreover, we have r(u 0 , v 0 ) = 0 and ∇ u r(u 0 , v 0 ) = 0. 3. g(u, v 0 ) -g(u, v) = d V (v 0 , v) for all u ∈ R du and v ∈ arg min v ∈V g(u, v ), where d V is non-negative and d V (v, v ) = 0 ⇐⇒ v = v .
Under the assumption that each client t can compute a partial first order surrogate of f t , we propose algorithms for federated surrogate optimization in both the client-server setting (Alg. 3) and the fully decentralized one (Alg. 5). Both algorithms are iterative and distributed: at each iteration k > 0, client t ∈ [T ] computes a partial first-order surrogate

g k t of f t near u k-1 , v k-1 t (resp. u k-1 t , v k-1 t
) for federated surrogate optimization in Alg. 3 (resp. for fully decentralized surrogate optimization in Alg 5).

The convergence of those two algorithms requires the following standard assumptions. Each of them generalizes one of the Assumptions 4-7 for our EM algorithms.

[∇ u g k t (u, v; ξ)] = ∇ u g k t (u, v) and E ξ ∇ u g k t (u, v; ξ)-∇ u g k t (u, v) 2 ≤ σ 2 .
Assumption 7 . (Bounded dissimilarity) There exist β and G such that

T t=1 ω t • ∇ u g k t (u, v) 2 ≤ G 2 + β 2 T t=1 ω t • ∇ u g k t (u, v) 2 .
Under these assumptions a parallel result to Theorem. 3.2 holds for the client-server setting. Theorem 3.2 . Under Assumptions 4 -7 , when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, the iterates of federated surrogate optimization (Alg. 3) satisfy:

1 K K k=1 E ∇ u f u k , v k 1:T 2 F ≤ O 1 √ K , 1 K K k=1 ∆ v f (u k , v k 1:T ) ≤ O 1 K 3/4 , ( 89
)
where the expectation is over the random batches samples, and

∆ v f (u k , v k 1:T ) f u k , v k 1:T - f u k , v k+1 1:T ≥ 0.
In the fully decentralized setting, if in addition to Assumptions 4 -7 , we suppose that Assumption 8 holds, a parallel result to Theorem. 3.3 holds. Theorem 3.3 . Under Assumptions 4 -7 and Assumption 8, when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, the iterates of fully decentralized federated surrogate optimization (Alg. 5) satisfy:

1 K K k=1 E ∇ u f ūk , v k 1:T 2 ≤ O 1 √ K , 1 K K k=1 T t=1 ω t • d V v k t , v k+1 t ≤ O 1 K , (90
) where ūk = 1 T T t=1 u k t . Moreover, local estimates u k t 1≤t≤T converge to consensus, i.e., to ūk :

1 K K k=1 T t=1 u k t -ūk 2 ≤ O 1 √ K .
The proofs of Theorem 3.2 and Theorem 3.3 are in Section G.1 and Section G.2, respectively.

F.3 Illustration: Analyzing pFedMe with Federated Surrogate Optimization

In this section, we show that pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF] can be studied through our federated surrogate optimization framework. With reference to the general formulation of pFedMe in [16, Eq. ( 2) and (3)], consider

g k t (w) = f t θ k-1 + λ 2 • θ k-1 -ω 2 , (91)
where

θ k-1 = prox f t λ ω k-1 arg min θ f t (θ) + λ 2 • θ -ω k-1 2
. We can verify that g k t is a first-order surrogate of f t near θ k-1 :

1. It is clear that g k t θ k-1 = f t θ k-1 . 2. Since θ k-1 = prox f t λ ω k-1
, using the envelope theorem (assuming that f t is proper, convex and lower semi-continuous), it follows that ∇f t ω k-1 = λ θ k-1 -ω k-1 = ∇g k k ω k-1 . Therefore, pFedMe can be seen as a particular case of the federated surrogate optimization algorithm (Alg. 3), to which our convergence results apply.

G Convergence Proofs

We study the client-server setting and the fully decentralized setting in Section G.1 and Section G.2, respectively. In both cases, we first prove the more general result for the federated surrogate optimization introduced in App. F, and then derive the specific result for FedEM and D-FedEM.

G.1 Client-Server Setting G.1.1 Additional Notations Remark 2. For convenience and without loss of generality, we suppose in this section that ω ∈ ∆ T , i.e., ∀t ∈ [T ], ω t ≥ 0 and

T t =1 ω t = 1.
At iteration k > 0, we use u k-1,j t to denote the j-th iterate of the local solver at client t ∈ [T ], thus

u k-1,0 t = u k-1 , (92)
and

u k = T t=1 ω t • u k-1,J t . ( 93
)
At iteration k > 0, the local solver's updates at client t ∈ [T ] can be written as (for 0 ≤ j ≤ J -1):

u k-1,j+1 t = u k-1,j t -η k-1,j ∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t , (94)
where ξ k-1,j t is the batch drawn at the j-th local update of u k-1 t .

We introduce η k-1 = J-1 j=0 η k-1,j , and we define the normalized update of the local solver at client t ∈ [T ] as,

δk-1 t - u k-1,J t -u k-1,0 t η k-1 = J-1 j=0 η k-1,j • ∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t J-1 j=0 η k-1,j , (95)
and also define

δ k-1 t J-1 j=0 η k-1,j • ∇ u g k t u k-1,j t , v k-1 t η k-1 . ( 96
)
With this notation,

u k -u k-1 = -η k-1 • T t=1 ω t • δk-1 t . (97)
Finally, we define g k , k > 0 as

g k (u, v 1:T ) T t=1 ω t • g k t (u, v t ) . (98)
Note that g k is a convex combination of functions g k t , t ∈ [T ].

G.1.2 Proof of Theorem 3.2

Lemma G.1. Suppose that Assumptions 5 -7 hold. Then, for k > 0, and (η k,j ) 0≤j≤J-1 such that

η k J-1 j=0 η k,j ≤ min 1 2 √ 2L , 1
4Lβ , the updates of federated surrogate optimization (Alg 3) verify

E f (u k , v k 1:T ) -f (u k-1 , v k-1 1:T ) η k-1 ≤ - 1 4 E ∇ u f u k-1 , v k-1 1:T 2 - 1 η k-1 T t=1 ω t • d V v k-1 t , v k t + 2η k-1 L   J-1 j=0 η 2 k-1,j η k-1 L + 1   σ 2 + 4η 2 k-1 L 2 G 2 . ( 99
)
Proof. This proof uses standard techniques from distributed stochastic optimization. It is inspired by [START_REF] Wang | Tackling the objective inconsistency problem in heterogeneous federated optimization[END_REF]Theorem 1].

For k > 0, g k is L-smooth wrt u, because it is a convex combination of L-smooth functions g k t , t ∈ [T ]. Thus, we write

g k u k , v k-1 1:T -g k u k-1 , v k-1 1:T ≤ u k -u k-1 , ∇ u g k (u k-1 , v k-1 1:T ) + L 2 u k -u k-1 2 ,
(100) where < u, u > denotes the scalar product of vectors u and u . Using Eq. ( 97), and taking the expectation over random batches ξ k-1,j t 0≤j≤J-1 1≤t≤T

, we have

E g k u k , v k-1 1:T -g k u k-1 , v k-1 1:T ≤ -η k-1 E T t=1 ω t • δk-1 t , ∇ u g k (u k-1 , v k-1 1:T ) T1 + Lη 2 k-1 2 • E T t=1 ω t • δk-1 t 2 T2 . (101)
We bound each of those terms separately. For T 1 we have

T 1 = E T t=1 ω t • δk-1 t , ∇ u g k u k-1 , v k-1 1:T (102) = E T t=1 ω t • δk-1 t -δ k-1 t , ∇ u g k u k-1 , v k-1 1:T + E T t=1 ω t • δ k-1 t , ∇ u g k u k-1 , v k-1 1:T . ( 103
)
Because stochastic gradients are unbiased (Assumption 6 ), we have

E δk-1 t -δ k-1 t = 0, (104) thus,
T 1 = E T t=1 ω t • δ k-1 t , ∇ u g k u k-1 , v k-1 1:T (105) = 1 2   ∇ u g k u k-1 , v k-1 1:T 2 + E T t=1 ω t • δ k-1 t 2   - 1 2 E ∇ u g k u k-1 , v k-1 1:T - T t=1 ω t • δ k-1 t 2 . ( 106
)
For T 2 we have for k > 0,

T 2 = E T t=1 ω t • δk-1 t 2 (107) = E T t=1 ω t • δk-1 t -δ k-1 t + T t=1 ω t • δ k-1 t 2 (108) ≤ 2 E T t=1 ω t • δk-1 t -δ k-1 t 2 + 2 E T t=1 ω t • δ k-1 t 2 (109) = 2 T t=1 ω 2 t • E δk-1 t -δ k-1 t 2 + 2 1≤s =t≤T ω t ω s E δk-1 t -δ k-1 t , δk-1 s -δ k-1 s + 2 E T t=1 ω t δ k-1 t 2 . ( 110
)
Since clients sample batches independently, and stochastic gradients are unbiased (Assumption 6 ), we have

E δk-1 t -δ k-1 t , δk-1 s -δ k-1 s = 0, (111) thus,
T 2 ≤ 2 T t=1 ω 2 t • E δk-1 t -δ k-1 t 2 + 2 E T t=1 ω t δ k-1 t 2 (112) = 2 T t=1 ω 2 t E J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1,j t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t 2 + 2 E T t=1 ω t δ k-1 t 2 . (113)
Using Jensen inequality, we have

J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1,j t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t 2 ≤ J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1,j t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t 2 , (114)
and since the variance of stochastic gradients is bounded by σ 2 (Assumption 6 ), it follows that

E J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1,j t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t 2 ≤ J-1 j=0 η k-1,j η k-1 σ 2 = σ 2 . ( 115
)
Replacing back in the expression of T 2 , we have

T 2 ≤ 2 T t=1 ω 2 t σ 2 + 2 E T t=1 ω t • δ k-1 t 2 . ( 116
)
Finally, since 0 ≤ ω t ≤ 1, t ∈ [T ] and T t=1 ω t = 1, we have

T 2 ≤ 2σ 2 + 2 E T t=1 ω t • δ k-1 t 2 . ( 117
)
Having bounded T 1 and T 2 , we can replace Eq. ( 106) and Eq. (117) in Eq. ( 101), and we get

E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) ≤ - η k-1 2 ∇ u g k u k-1 , v k-1 1:T 2 + η 2 k-1 Lσ 2 34 - η k-1 2 (1 -2Lη k-1 ) • E T t=1 ω t • δ k-1 t 2 + η k-1 2 E ∇ u g k u k-1 , v k-1 1:T - T t=1 ω t • δ k-1 t 2 . ( 118
) As η k-1 ≤ 1 2 √ 2L ≤ 1 2L , we have E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) ≤ - η k-1 2 ∇ u g k u k-1 , v k-1 1:T 2 + η 2 k-1 Lσ 2 + η k-1 2 E ∇ u g k u k-1 , v k-1 1:T - T t=1 ω t δ k-1 t 2 . ( 119
) Replacing ∇ u g k u k-1 , v k-1 1:T = T t=1 ω t • ∇ u g k t u k-1 , v k-1 t
, and using Jensen inequality to bound the last term in the RHS of Eq. ( 119), we have

E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) ≤ - η k-1 2 ∇ u g k u k-1 , v k-1 1:T 2 + η 2 k-1 Lσ 2 + η k-1 2 T t=1 ω t • E ∇ u g k t u k-1 , v k-1 t -δ k-1 t 2 T3 . ( 120
)
We now bound the term T 3 :

T 3 = E ∇ u g k t u k-1 , v k-1 t -δ k-1 t 2 (121) = E ∇ u g k t u k-1 , v k-1 t - J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1,j t , v k-1 t 2 (122) = E J-1 j=0 η k-1,j η k-1 ∇ u g k t u k-1 , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 (123) ≤ J-1 j=0 η k-1,j η k-1 E ∇ u g k t u k-1 , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 (124) ≤ J-1 j=0 η k-1,j η k-1 L 2 E u k-1 -u k-1,j t 2 , ( 125
)
where the first inequality follows from Jensen inequality and the second one follow from the L-smoothness of g k t (Assumption 5 ). We bound now the term

E u k-1 -u k-1,j t for j ∈ {0, . . . , J -1} and t ∈ [T ], E u k-1 -u k-1,j t 2 = E u k-1,j t -u k-1,0 t 2 (126) = E j-1 l=0 u k-1,l+1 t -u k-1,l t 2 (127) = E j-1 l=0 η k-1,l ∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,l t 2 (128) ≤ 2E j-1 l=0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t ; ξ k-1,l t -∇ u g k t u k-1,l t , v k-1 t 2 + 2E j-1 l=0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t 2 (129) = 2 j-1 l=0 η 2 k-1,l E ∇ u g k t u k-1,l t , v k-1 t ; ξ k-1,l t -∇ u g k t u k-1,l t , v k-1 t 2 + 2E j-1 l=0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t 2 (130) ≤ 2σ 2 j-1 l=0 η 2 k-1,l + 2E j-1 l=0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t 2 , ( 131
)
where, in the last two steps, we used the fact that stochastic gradients are unbiased and have bounded variance (Assumption 6 ). We bound now the last term in the RHS of Eq. ( 131),

E j-1 l=0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t 2 = E j-1 l =0 η k-1,l • j-1 l=0 η k-1,l j-1 l =0 η k-1,l ∇ u g k t u k-1,l t , v k-1 t 2 (132) ≤ j-1 l =0 η k-1,l 2 • j-1 l=0 η k-1,l j-1 l =0 η k-1,l E ∇ u g k t u k-1,l t , v k-1 t 2 (133) = j-1 l=0 η k-1,l • j-1 l=0 η k-1,l E ∇ u g k t u k-1,l t , v k-1 t 2 (134) = j-1 l=0 η k-1,l • j-1 l=0 η k-1,l E ∇ u g k t u k-1,0 t , v k-1 t -∇ u g k t u k-1,0 t , v k-1 t + ∇ u g k t u k-1,l t , v k-1 t 2 (135) ≤2 j-1 l=0 η k-1,l • j-1 l=0 η k-1,l • E ∇ u g k t u k-1,0 t , v k-1 t 2 + E ∇ u g k t u k-1,l t , v k-1 t -∇ u g k t u k-1,0 t , v k-1 t 2 (136) =2 j-1 l=0 η k-1,l • j-1 l=0 η k-1,l • E ∇ u g k t u k-1 , v k-1 t 2 + E ∇ u g k t u k-1,l t , v k-1 t -∇ u g k t u k-1 , v k-1 t 2 (137) ≤2 j-1 l=0 η k-1,l j-1 l=0 η k-1,l E ∇ u g k t u k-1 , v k-1 t 2 + L 2 E u k-1,l t -u k-1 2 (138) =2L 2 j-1 l=0 η k-1,l j-1 l=0 η k-1,l • E u k-1,l t -u k-1 2 + 2 j-1 l=0 η k-1,l 2 E ∇ u g k t u k-1 , v k-1 t 2 , ( 139
)
where the first inequality is obtained using Jensen inequality, and the last one is a result of the L-smoothness of g t (Assumption 5 ). Replacing Eq. ( 139) in Eq. ( 131), we have

J-1 j=0 η k-1,j η k-1 • E u k-1 -u k-1,j t 2 ≤ 2σ 2   J-1 j=0 η k-1,j η k-1 • j-1 l=0 η 2 k-1,l   + 4L 2 J-1 j=0 η k-1,j η k-1 j-1 l=0 η k-1,l • j-1 l=0 η k-1,l • E u k-1,l t -u k-1 t 2 + 4   J-1 j=0 η k-1,j η k-1 j-1 l=0 η k-1,l 2   • E ∇ u g k t u k-1 t , v k-1 t 2 . ( 140
) Since j-1 l=0 η k-1,l • E u k-1,l t -u k-1 t 2 ≤ J-1 j=0 η k-1,j • E u k-1,j t -u k-1 t 2
, we have

J-1 j=0 η k-1,j η k-1 • E u k-1 -u k-1,j t 2 ≤ 2σ 2   J-1 j=0 η k-1,j η k-1 • j-1 l=0 η 2 k-1,l   + 4L 2   J-1 j=0 η k-1,j η k-1 j-1 l=0 η k-1,l   •   J-1 j=0 η k-1,j • E u k-1,j t -u k-1 2   + 4   J-1 j=0 η k-1,j η k-1 j-1 l=0 η k-1,l 2   • E ∇ u g k t u k-1 , v k-1 t 2 . ( 141
)
We use Lemma G.11 to simplify the last expression, obtaining

J-1 j=0 η k-1,j η k-1 • E u k-1 -u k-1,j t 2 ≤ 2σ 2 •    J-1 j=0 η 2 k-1,j    + 4η 2 k-1 E ∇ u g k t u k-1 , v k-1 t 2 + 4η k-1 L 2 • J-1 j=0 η k-1,j E u k-1,j t -u k-1 2 . ( 142
)
Rearranging the terms, we have

1 -4η 2 k-1 L 2 • J-1 j=0 η k-1,j η k-1 • E u k-1 -u k-1,j t 2 ≤ 2σ 2 •    J-1 j=0 η 2 k-1,j    + 4η 2 k-1 • E ∇ u g k t u k-1 , v k-1 t 2 . ( 143
)
Finally, replacing Eq. ( 143) into Eq. ( 125), we have

1 -4η 2 k-1 L 2 • T 3 ≤ 2σ 2 L 2 •   J-1 j=0 η 2 k-1,j   + 4η 2 k-1 L 2 • E ∇ u g k t u k-1 , v k-1 t 2 . (144) For η k-1 small enough, in particular if η k-1 ≤ 1 2 √ 2L , then 1 2 ≤ 1 -4η 2 k-1 L 2 , thus T 3 2 ≤ 2σ 2 L 2 •   J-1 j=0 η 2 k-1,j   + 4η 2 k-1 L 2 • E ∇ u g k t u k-1 , v k-1 t 2 . ( 145
)
Replacing the bound of T 3 from Eq. (145) into Eq. ( 120), we have obtained

E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) ≤ - η k-1 2 E ∇ u g k u k-1 , v k-1 1:T 2 + 4η 3 k-1 L 2 T t=1 ω t • E ∇ u g k t u k-1 , v k-1 t 2 + 2η k-1 L   J-1 j=0 η 2 k-1,j L + η k-1   • σ 2 . ( 146
)
Using Assumption 7 , we have

E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) ≤ - η k-1 2 E ∇ u g k u k-1 , v k-1 1:T 2 + 4η 3 k-1 L 2 β 2 • E T t=1 ω t • ∇ u g k t u k-1 , v k-1 t 2 + 2η k-1 L   J-1 j=0 η 2 k-1,j L + η k-1   • σ 2 + 4η 3 k-1 L 2 G 2 . ( 147
)
Dividing by η k-1 , we get

E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) η k-1 ≤ 8η 2 k-1 L 2 β 2 -1 2 E ∇ u g k u k-1 , v k-1 1:T 2 + 2η k-1 L   J-1 j=0 η 2 k-1,j η k-1 L + 1   • σ 2 + 4η 2 k-1 L 2 G 2 . ( 148
) For η k-1 small enough, if η k-1 ≤ 1 4Lβ , then 8η 2 k-1 L 2 β 2 -1 ≤ 1 2 . Thus, E g k (u k , v k-1 1:T ) -g k (u k-1 , v k-1 1:T ) η k-1 ≤ - 1 4 E ∇ u g k u k-1 , v k-1 1:T 2 + 2η k-1 L   J-1 j=0 η 2 k-1,j η k-1 L + 1   • σ 2 + 4η 2 k-1 L 2 G 2 . ( 149
) Since for t ∈ [T ], g k t is a partial first-order surrogate of f t near u k-1 , v k-1 t
, we have (see Def. 1)

g k t u k-1 , v k-1 t = f t u k-1 , v k-1 t , (150)
∇ u g k t u k-1 , v k-1 t = ∇ u f t u k-1 , v k-1 t , (151)
g k t u k , v k-1 t = g k t u k , v k t + d V v k-1 t , v k t . (152)
Multiplying by ω t and summing over t ∈ [T ], we have

g k u k-1 , v k-1 1:T = f u k-1 , v k-1 1:T , ( 153
)
∇ u g k u k-1 , v k-1 1:T = ∇ u f u k-1 , v k-1 1:T , (154)
g k u k , v k-1 1:T = g k u k , v k 1:T + T t=1 ω t • d V v k-1 t , v k t . (155)
Replacing Eq. ( 153), Eq. ( 154) and Eq. (155) in Eq. ( 149), we have

E g k (u k , v k 1:T ) -f (u k-1 , v k-1 1:T ) η k-1 ≤ - 1 4 E ∇ u f u k-1 , v k-1 1:T 2 - 1 η k-1 T t=1 ω t • d V v k-1 t , v k t + 2η k-1 L      J-1 j=0 η 2 k-1,j η k-1    L + 1   • σ 2 + 4η 2 k-1 L 2 G 2 . ( 156
)
Using again Definition 1, we have

g k (u k , v k 1:T ) ≥ f (u k , v k 1:T ), (157)
thus,

E f (u k , v k 1:T ) -f (u k-1 , v k-1 1:T ) η k-1 ≤ - 1 4 E ∇ u f u k-1 , v k-1 1:T 2 - 1 η k-1 T t=1 ω t • d V v k-1 t , v k t + 2η k-1 L   J-1 j=0 η 2 k-1,j η k-1 L + 1   • σ 2 + 4η 2 k-1 L 2 G 2 . ( 158
)
Lemma G.2. For k ≥ 0 and t ∈ [T ], the iterates of Alg. 3 verify

0 ≤ d V v k+1 t , v k t ≤ f t u k , v k t -f t (u k , v k+1 t ) (159) Proof. Since v k+1 t ∈ arg min v∈V g k t u k-1 , v , and g k t is a partial first-order surrogate of f t near {u k-1 , v k-1 t }, we have g k t u k-1 , v k-1 t -g k t u k-1 , v k t = d V v k-1 t , v k t , (160)
thus,

f t u k-1 , v k-1 t -f t u k-1 , v k t ≥ d V v k-1 t , v k t , (161
) where we used the fact that

g k t u k-1 , v k-1 t = f t u k-1 , v k-1 t , (162)
and,

g k t u k-1 , v k t ≥ f t u k-1 , v k t . (163)
Theorem 3.2 . Under Assumptions 4 -7 , when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, the iterates of federated surrogate optimization (Alg. 3) satisfy:

1 K K k=1 E ∇ u f u k , v k 1:T 2 F ≤ O 1 √ K , 1 K K k=1 E ∆ v f (u k , v k 1:T ) ≤ O 1 K 3/4 , (89
) where the expectation is over the random batches samples, and

∆ v f (u k , v k 1:T ) f u k , v k 1:T - f u k , v k+1 1:T ≥ 0. Proof. For K large enough, η = a0 √ K ≤ 1 J min 1 2 √ 2L , 1
4Lβ , thus the assumptions of Lemma G.1 are satisfied. Lemma G.1 and non-negativity of d V lead to

E f (u k , v k 1:T ) -f (u k-1 , v k-1 1:T ) Jη ≤ - 1 4 E ∇ u f u k-1 , v k-1 1:T 2 + 2ηL (ηL + 1) • σ 2 + 4J 2 η 2 L 2 G 2 . ( 164
)
Rearranging the terms and summing for k ∈ [K], we have

1 K K k=1 E ∇ u f u k-1 , v k-1 1:T 2 ≤ 4E f (u 0 , v 0 1:T ) -f (u K , v K 1:T ) JηK + 8 ηL (ηL + 1) • σ 2 + 2J 2 η 2 L 2 G 2 K (165) ≤ 4E f (u 0 , v 0 1:T ) -f * JηK + 8 ηL (ηL + 1) • σ 2 + 2J 2 η 2 L 2 G 2 K , ( 166
)
where we use Assumption 4 to obtain (166). Thus,

1 K K k=1 E ∇ u f u k-1 , v k-1 1:T 2 = O 1 √ K . ( 167
)
To prove the second part of Eq. ( 89), we first decompose

∆ v f u k , v k 1:T -f u k , v k+1 1:T ≥ 0 as follow, ∆ v = f u k , v k 1:T -f u k+1 , v k+1 1:T T k 1 + f u k+1 , v k+1 1:T -f u k , v k+1 1:T T k 2 . ( 168
)
Using again Lemma G.1 and Eq. ( 167), it follows that

1 K K k=1 E T k 1 ≤ O 1 K . ( 169
)
For T k 2 , we use the fact that f is 2L-smooth (Lemma G.12) w.r.t. u and Cauchy-Schwartz inequality. Thus, for k > 0, we write

T k 2 = f u k+1 , v k+1 1:T -f u k , v k+1 1:T (170) ≤ ∇ u f u k+1 , v k+1 1:T • u k+1 -u k + 2L 2 u k+1 -u k 2 . ( 171
)
Summing over k and taking expectation:

1 K K k=1 E T k 2 ≤ 1 K K k=1 E ∇ u f u k+1 , v k+1 1:T • u k+1 -u k + 1 K K k=1 2L 2 E u k+1 -u k 2 (172) ≤ 1 K K k=1 E ∇ u f u k+1 , v k+1 1:T 2 K k=1 E u k+1 -u k 2 + 1 K K k=1 2L 2 E u k+1 -u k 2 , ( 173
)
where the second inequality follows from Cauchy-Schwarz inequality. From Eq. ( 143), with η k-1 = Jη, we have for t ∈ [T ]

E u k -u k-1,J t 2 ≤ 4σ 2 Jη 2 + 8J 3 η 2 • E ∇ u g k t u k-1 , v k-1 t 2 . ( 174
)
Multiplying the previous by ω t and summing for t ∈ [T ], we have

T t=1 ω t • E u k-1 -u k-1,J t 2 ≤ 4J 2 σ 2 η 2 + 8J 3 η 2 • T t=1 ω t E ∇ u g k t u k-1 , v k-1 t 2 . ( 175
)
Using Assumption 7 , it follows that

T t=1 ω t E u k-1 -u k-1,J t 2 ≤ 4J 2 η 2 2JG 2 + σ 2 + 8J 3 η 2 β 2 E T t=1 ω t ∇ u g k t u k-1 , v k-1 t 2 .
(176) Finally using Jensen inequality and the fact that

g k t is a partial first-order of f t near u k-1 , v k-1 t , we have E u k-1 -u k 2 ≤ 4J 2 η 2 2JG 2 + σ 2 + 8J 3 η 2 β 2 E ∇ u f u k-1 , v k-1 1:T 2 . ( 177
)
From Eq. (167) and η ≤ O(1/ √ K), we obtain

1 K K k=1 E u k-1 -u k 2 ≤ O (1) , (178)
Replacing the last inequality in Eq. ( 173) and using again Eq. (167), we obtain

1 K K k=1 E T k 2 ≤ O 1 K 3/4 . (179)
Combining Eq. ( 169) and Eq. ( 179), it follows that

1 K K k=1 E ∆ v f (u k , v k 1:T ) ≤ O 1 K 3/4 . (180)
G.1.3 Proof of Theorem 3.2

In this section, f denotes the negative log-likelihood function defined in Eq. ( 6). Moreover, we introduce the negative log-likelihood at client t as follows

f t (Θ, Π) - log p(S t |Θ, Π) n - 1 n t nt i=1 log p(s (i) t |Θ, π t ). (181)
Theorem 3.2. Under Assumptions 1-7, when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, FedEM's iterates satisfy:

1 K K k=1 E ∇ Θ f Θ k , Π k 2 F ≤ O 1 √ K , 1 K K k=1 ∆ Π f (Θ k , Π k ) ≤ O 1 K 3/4 , (11)
where the expectation is over the random batches samples, and

∆ Π f (Θ k , Π k ) f Θ k , Π k - f Θ k , Π k+1 ≥ 0.
Proof. We prove this result as a particular case of Theorem 3.2 . To this purpose, in this section, we consider that V ∆ M , u = Θ ∈ R dM , v t = π t , and ω t = n t /n for t ∈ [T ]. For k > 0, we define g k t as follows:

g k t Θ, π t = 1 n t nt i=1 M m=1 q k t z (i) t = m • l h θm (x (i) t ), y (i) t
-log p m (x

(i) t ) -log π t + log q k t z (i) t = m -c , (182)
where c is the same constant appearing in Assumption 3, Eq. ( 3). With this definition, it is easy to check that the federated surrogate optimization algorithm (Alg. 3) reduces to FedEM (Alg. 2). Theorem 3.2 then follows immediately from Theorem 3.2 , once we verify that g k t 1≤t≤T satisfy the assumptions of Theorem 3.2 . Assumption 4 , Assumption 6 , and Assumption 7 follow directly from Assumption 4, Assumption 6, and Assumption 7, respectively. Lemma G.3 shows that for k > 0, g k is smooth w.r.t. Θ and then Assumption 5 is satisfied. Finally, Lemmas G.4-G.6 show that for t ∈

[T ] g k t is a partial first-order surrogate of f t w.r.t. Θ near Θ k-1 , π t with d V (•, •) = KL(• •). Lemma G.3. Under Assumption 5, for t ∈ [T ] and k > 0, g k t is L-smooth w.r.t Θ. Proof. g k t is a convex combination of L-smooth function θ → l(θ; s (i) t ), i ∈ [n t ]. Thus it is also L-smooth. Lemma G.4. Suppose that Assumptions 1-3, hold. Then, for t ∈ [T ], Θ ∈ R M ×d and π t ∈ ∆ M r k t (Θ, π t ) g k t (Θ, π t ) -f t (Θ, π t ) = 1 n t nt i=1 KL q k t z (t) i p t z (t) i |s (t) i , Θ, π t ,
where KL is Kullback-Leibler divergence.

Proof. Let k > 0 and t ∈ [T ], and consider Θ ∈ R M ×d and π t ∈ ∆ M , then

g k t Θ, π t = 1 n t nt i=1 M m=1 q k t z (i) t = m • l h θm (x (i) t ), y (i) t -log p m (x (i) t ) -log π t + log q k t z (i) t = m -c , (183)
= 1 n t nt i=1 M m=1 q k t z (i) t = m • -log p m y (i) t |x (i) t , θ m -log p m (x (i) t ) -log π t + log q k t z (i) t = m (184) = 1 n t nt i=1 M m=1 q k t z (i) t = m • -log p m y (i) t |x (i) t , θ m • p m (x (i) t ) • p t z (i) t = m + log q k t z (i) t = m (185) = 1 n t nt i=1 M m=1 q k t z (i) t = m • log q k t z (i) t = m -log p t s (i) t , z (i)
t = m Θ, π t ) (186) = 1 n t nt t=1 M m=1 q k t z (i) t = m log q k t z (i) t = m p t s (i) t , z (i)
t = m|Θ, π t . (187)
Thus,

r k t Θ, π t g k t (Θ, π t ) -f t (Θ, π t ) (188)
= -

1 n t nt t=1 M m=1   q k t z (i) t = m • log p t s (i) t , z (i)
t = m|Θ, π t q k t z (i) t = m   + 1 n t nt i=1 log p t s (i) t |Θ, π t (189) = 1 n t nt t=1 M m=1 q k t z (i) t = m log p t s (i) t |Θ, π t -log p t s (i) t , z (i)
t = m|Θ, π t q k t z (i) t = m (190) = 1 n t nt t=1 M m=1 q k t z (i) t = m log p t s (i) t |Θ, π t • q k t z (i) t = m p t s (i) t , z (i)
t = m|Θ, π t (191) = 1 n t nt t=1 M m=1 q k t z (i) t = m • log q k t z (i) t = m p t z (i) t = m|s (i) t , Θ, π t . (192)
Thus,

r k t (Θ, π t ) = 1 n t nt i=1 KL q k t (•) p t (•|s (t) i , Θ, π t ) ≥ 0. (193)
The following lemma shows that g k t and g k (as defined in Eq. 98) satisfy the first two properties in Definition 1.

Lemma G.5. Suppose that Assumptions 1-3 and Assumption 5 hold. For all k ≥ 0 and t ∈ [T ], g k t is a majorant of f t and r k t

g k t -f t is L-smooth in Θ. Moreover r k t Θ k-1 , π k-1 t = 0 and ∇ Θ r k t Θ k-1 , π k-1 t = 0.
The same holds for g k , i.e., g k is a majorant of f ,

r k g k -f is L-smooth in Θ, r k Θ k-1 , Π k-1 = 0 and ∇ Θ r k Θ k-1 , Π k-1 = 0 Proof. For t ∈ [T ], consider Θ ∈ R M ×d and π t ∈ ∆ M , we have (Lemma G.4)
r k t (Θ, π t ) g k t (Θ, π t ) -f t (Θ, π t ) = 1 n t nt i=1 KL q k t z (t) i p t z (i) t |s (i) t , Θ, π t . ( 194
)
Since KL divergence is non-negative, it follows that g k t is a majorant of f t , i.e.,

∀ Θ ∈ R M ×d , π t ∈ ∆ M : g k t (Θ, π) ≥ f t (Θ, π t ) . (195)
Moreover since, q k t z

(i) t = p t z (i) t |s (i) t , Θ k-1 , π k-1 t for k > 0, it follows that r k t Θ k-1 , π k-1 t = 0. (196)
For i ∈ [n t ] and m ∈ [M ], from Eq. 78, we have

p t z (i) t = m|s (i) t , Θ, π t = p m y (i) t |x (i) t , θ m × π tm M m =1 p m y (i) t |x (i) t , θ m × π tm (197) = exp -l h θm (x (i) t ), y (i) t
× π tm M m =1 exp -l h θ m (x (i) t ), y (i) t × π tm (198)
= exp -l h θm (x

(i) t ), y (i) t + log π tm M m =1 exp -l h θ m (x (i) t ), y (i) t + log π tm . ( 199
)
For ease of notation, we introduce

l i (θ) l h θ (x (i) t ), y (i) t , θ ∈ R d , m ∈ [M ], i ∈ [n t ], (200)
γ m (Θ) p t z (i) t = m|s (i) t , Θ, π t , m ∈ [M ], (201)
and,

ϕ i (Θ) KL q k t z (t) i p t z (i) t |s (i) t , Θ, π t . ( 202
) For i ∈ [n t ], function l i is differentiable because smooth (Assum 5), thus γ m , m ∈ [M ]
is differentiable as the composition of the softmax function and the function {Θ → -l i (Θ) + log π tm }. Its gradient is given by

∇ θm γ m (Θ) = -γ m (Θ) • (1 -γ m (Θ)) • ∇l i (θ m ) , ∇ θ m γ m (Θ) = γ m (Θ) • γ m (Θ) • ∇l i (θ m ) , m = m. (203)
Thus for m ∈ [M ], we have

∇ θm ϕ i (Θ) = M m =1 q k t z (t) i = m • ∇ θm γ m (Θ) γ m (Θ) (204) = m =1 m =m q k t z (t) i = m • γ m (Θ) • γ m (Θ) γ m (Θ) • • ∇l i (θ m ) -q k t z (t) i = m • γ m (Θ) • (1 -γ m (Θ)) γ m (Θ) • ∇l i (θ m ) . (205)
Using the fact that

M m =1 q k t z (t) i = m = 1, it follows that ∇ θm ϕ i (Θ) = γ m (Θ) -q k t z (t) i = m • ∇l i (θ m ) . (206)
Since l i , i ∈ [n t ] is twice continuously differentiable (Assumption 5), and

γ m , m ∈ [M ] is differentiable, then φ i , i ∈ [n t ] is twice continuously differentiable. We use H (ϕ i (Θ)) ∈ R dM ×dM (resp. H (l i (θ)) ∈ R d×d )
to denote the Hessian of ϕ (resp. l i ) at Θ (resp. θ). The Hessian of ϕ i is a block matrix given by

           H (ϕ i (Θ)) m,m = -γ m (Θ) • (1 -γ m (Θ)) • ∇l i (θ m ) • ∇l i (θ m ) + γ m (Θ) -q k t z (t) i = m • H (l i (θ m )) H (ϕ i (Θ)) m,m = γ m (Θ) • γ m (Θ) • ∇l i (θ m ) • ∇l i (θ m ) , m = m.
(207) We introduce the block matrix H ∈ R dM ×dM , defined by

   Hm,m = -γ m (Θ) • 1 -γ m (Θ) • ∇l i (θ m ) • (∇l i (θ m )) Hm,m = γ m (Θ) • γ m (Θ) • ∇ θ l i (θ m ) • ∇l i (θ m ) , m = m, (208)
Eq. ( 207) can be written as

     H (ϕ i (Θ)) m,m -Hm,m = γ m (Θ) -q k t z (t) i = m • H (l i (θ m )) H (ϕ i (Θ)) m,m -Hm,m = 0, m = m.
(209) We recall that a twice differentiable function is L smooth if and only if the eigenvalues of its Hessian are smaller then L, see e.g., [52, Lemma 1.2.2] or [6, Section 3.2]. Since l i and also -l i are L-smooth (Assumption 5), we have for θ ∈ R d ,

-L • I d H (l i (θ)) L • I d . (210)
Using Lemma G.15, we can conclude that matrix H is semi-definite negative. Since

-1 ≤ γ m (Θ) -q k t z (t) i = m ≤ 1, (211)
it follows that

H (ϕ i (Θ)) L • I dM . (212)
The last equation proves that ϕ i is L-smooth. Thus r k t is L-smooth with respect to Θ as the average of L-smooth function.

Moreover, since r k t (Θ k-1 , π k-1 t ) = 0 and ∀Θ, Π; r k t (Θ, π t ) ≥ 0, it follows that Θ k-1 is a minimizer of Θ → r k t Θ, π k-1 t . Thus, ∇ Θ r k t (Θ k-1 , π k-1 t ) = 0.
For Θ ∈ R M ×d and Π ∈ ∆ T ×M , we have

r k (Θ, Π) g k (Θ, Π) -f (Θ, Π) (213)
T t=1 n t n • g k t (Θ, π t ) -f t (Θ, π t ) (214) = T t=1 n t n r k t (Θ, π t ) . (215)
We see that r k is a weighted average of r k t 1≤t≤T . Thus,

r k t is L-smooth in Θ, r k (Θ, Π) ≥ 0, moreover r k t Θ k-1 , Π k-1 = 0 and ∇ Θ r k t Θ k-1 , Π k-1 = 0.
The following lemma shows that g k t and g k satisfy the third property in Definition 1. Lemma G.6. Suppose that Assumption 1 holds and consider Θ ∈ R M ×d and Π ∈ ∆ T ×M , for k > 0, the iterates of Alg. 3 verify

g k (Θ, Π) = g k Θ, Π k + T t=1 n t n KL π k t , π t .
Proof. For t ∈ [T ] and k > 0, consider Θ ∈ R M ×d and π t ∈ ∆ M such that ∀m ∈ [M ]; π tm = 0, we have

g k t (Θ, π t ) -g k t Θ, π k t = M m=1 1 n t nt i=1 q k t z (i) t = m =π k tm (Proposition 3.1) × log π k tm -log π tm (216) = M m=1 π k tm log π k tm π tm (217) = KL π k t , π t . (218)
We multiply by nt n and some for t ∈ [T ]. It follows that

g k Θ, Π k + T t=1 n t n KL π k t , π t = g k (Θ, Π) . (219)
G.2 Fully Decentralized Setting G.2.1 Additional Notations Remark 3. For convenience and without loss of generality, we suppose in this section that

ω t = 1, t ∈ [T ].
We introduce the following matrix notation:

U k u k 1 , . . . , u k T ∈ R du×T (220) Ūk ūk , . . . , ūk ∈ R du×T (221) ∂g k U k , v k 1:T ; ξ k ∇ u g k 1 u k 1 , v k 1 ; ξ k 1 , . . . , ∇ u g k T u k T , v k T ; ξ k T ∈ R du×T (222)
where ūk = 1

T T t=1 u k t and v k 1:T = v k t 1≤t≤T ∈ V T .
We denote by u k-1,j t the j-th iterate of the local solver at global iteration k at client t ∈ [T ], and by U k-1,j the matrix whose column t is u k-1,j t , thus,

u k-1,0 t = u k-1 t ; U k-1,0 = U k-1 , (223) and,
u k t = T s=1 w k-1 st u k-1,J s ; U k = U k-1,J W k-1 . (224)
Using this notation, the updates of Alg. 5 can be summarized as

U k =   U k-1 - J-1 j=0 η k-1,j ∂g k U k-1,j , v 1:T ; ξ k-1,j   W k-1 . (225)
Similarly to the client-server setting, we define the normalized update of local solver at client t ∈ [T ]:

δk-1 t - u k-1,J t -u k-1,0 t η k-1 = J-1 j=0 η k-1,j ∇ u g k t u k-1,j t , v k t ; ξ k-1,j t J-1 j=0 η k-1,j , (226)
and

δ k-1 t J-1 j=0 η k-1,j ∇ u g k t u k-1,j t , v k t η k-1 . (227)
Because clients updates are independent, and stochastic gradient are unbiased, it is clear that

E δ k-1 t -δk-1 t = 0, (228)
and that

∀ t, s ∈ [T ] s.t. s = t, E δ k-1 t -δk-1 t , δ k-1 s -δk-1 s = 0.
(229) We introduce the matrix notation,

Υk-1 δk-1 1 , . . . , δk-1 T ∈ R du×T ; Υ k-1 δ k-1 1 , . . . , δ k-1 T ∈ R du×T . ( 230
)
Using this notation, Eq. ( 225) becomes

U k = U k-1 -η k-1 Υk-1 W k-1 . (231)
G.2.2 Proof of Theorem 3.3

In fully decentralized optimization, proving the convergence usually consists in deriving a recurrence on a term measuring the optimality of the average iterate (in our case this term is

E ∇ u f ūk , v k 1:T
2 ) and a term measuring the distance to consensus, i.e., E T t=1 u k t -ūk 2 . In what follows we obtain those two recurrences, and then prove the convergence. Lemma G.7 (Average iterate term recursion). Suppose that Assumptions 5 -7 and Assumption 8 hold. Then, for k > 0, and

(η k,j ) 1≤j≤J-1 such that η k J-1 j=0 η k,j ≤ min 1 2 √ 2L , 1
8Lβ , the updates of fully decentralized federated surrogate optimization (Alg. 5) verify

E f (ū k , v k 1:T ) -f (ū k-1 , v k-1 1:T ) ≤ - 1 T T t=1 E d V v k t , v k-1 t - η k-1 8 E ∇ u f ūk-1 , v k-1 1:T 2 + (12 + T ) η k-1 L 2 4T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 L • η 2 k-1,j η k-1 + 1   σ 2 + 16η 3 k-1 L 2 T G 2 . ( 232
)
Proof. We multiply both sides of Eq. ( 231) by 11 T , thus for k > 0 we have,

U k • 11 T = U k-1 -η k-1 Υk-1 W k-1 11 T , (233)
since W k-1 is doubly stochastic (Assumption 8), i.e., W k-1 11 T = 11 T , is follows that,

Ūk = Ūk-1 -η k-1 Υk-1 • 11 T , (234)
thus,

ūk = ūk-1 - η k-1 T • T t=1 δk-1 t . ( 235
)
Using the fact that g k is L-smooth with respect to u (Assumption 5 ), we write

E g k ūk , v k-1 1:T = E g k ūk-1 - η k-1 T T t=1 δk-1 t , v k-1 1:T (236) ≤ g k (ū k-1 , v k-1 1:T ) -E ∇ u g k (ū k-1 , v k-1 1:T ), η k-1 T T t=1 δk-1 t + L 2 E η k-1 T T t=1 δk-1 t 2 (237) = g k (ū k-1 , v k-1 1:T ) -η k-1 E ∇ u g k (ū k-1 , v k-1 1:T ), 1 T T t=1 δk-1 t T1 + η 2 k-1 • L 2T 2 E T t=1 δk-1 t 2 T2 , (238)
where the expectation is taken over local random batches. As in the client-server case, we bound the terms T 1 and T 2 . First, we bound T 1 , for k > 0, we have

T 1 = E ∇ u g k (ū k-1 , v k-1 1:T ), 1 T T t=1 δk-1 t (239) = E ∇ u g k ūk-1 , v k-1 1:T , 1 T T t=1 δk-1 t -δ k-1 t =0, because E[δ k-1 t -δk-1 t ]=0 + E ∇ u g k ūk-1 , v k-1 1:T , 1 T T t=1 δ k-1 t (240) = E ∇ u g k ūk-1 , v k-1 1:T , 1 T T t=1 δ k-1 t (241) = 1 2 E ∇ u g k ūk-1 , v k-1 1:T 2 + 1 2 E 1 T T t=1 δ k-1 t 2 - 1 2 E ∇ u g k ūk-1 , v k-1 1:T - 1 T T t=1 δ k-1 t 2 . ( 242
)
We bound now T 2 . For k > 0, we have,

T 2 = E T t=1 δk-1 t 2 (243) = E T t=1 δk-1 t -δ k-1 t + T t=1 δ k-1 t 2 (244) ≤ 2E T t=1 δk-1 t -δ k-1 t 2 + 2 • E T t=1 δ k-1 t 2 (245) = 2 • T t=1 E δk-1 t -δ k-1 t 2 + 2 1≤t =s≤T E δk-1 t -δ k-1 t , δk-1 s -δ k-1
s =0; because of Eq. ( 229)

+ 2E T t=1 δ k-1 t 2 (246) = 2 • T t=1 E δk-1 t -δ k-1 t 2 + 2 • E T t=1 δ k-1 t 2 (247) = 2 • E T t=1 δ k-1 t 2 + 2 • T t=1 1 η 2 k-1 E J-1 j=0 η k-1,j • ∇ u g k t u k-1,j t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t ; ξ k-1,j t 2 . ( 248
)
Since batches are sampled independently, and stochastic gradients are unbiased with finite variance (Assumption 6 ), the last term in the RHS of the previous equation can be bounded using σ 2 , leading to

T 2 ≤ 2 • T t=1 J-1 j=0 η 2 k-1,j η 2 k-1 σ 2 + 2 • E T t=1 δ k-1 t 2 (249) = 2T • σ 2 • T t=1 • J-1 j=0 η 2 k-1,j η 2 k-1 + 2E T t=1 δ k-1 t 2 (250)
≤ 2T • σ 2 + 2 • E T t=1 δ k-1 t 2 . ( 251
)
Replacing Eq. ( 242) and Eq. ( 251) in Eq. ( 238), we have

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ - η k-1 2 E ∇ u g k ūk-1 , v k-1 1:T 2 - η k-1 2 (1 -2Lη k-1 ) E 1 T T t=1 δ k-1 t 2 + L T η 2 k-1 σ 2 + η k-1 2 E ∇ u g k ūk-1 , v k-1 1:T - 1 T T t=1 δ k-1 t 2 . (252)
For η k-1 small enough, in particular for η k-1 ≤ 1 2L , we have

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ - η k-1 2 E ∇ u g k ūk-1 , v k-1 1:T 2 + L T η 2 k-1 σ 2 + η k-1 2 E 1 T T t=1 ∇ u g k t ūk-1 , v k-1 t -δ k-1 t 2 . ( 253
)
We use Jensen inequality to bound the last term in the RHS of the previous equation, leading to

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ - η k-1 2 E ∇ u g k ūk-1 , v k-1 1:T 2 + L T η 2 k-1 σ 2 + η k-1 2T • T t=1 E ∇ u g k t ūk-1 , v k-1 t -δ k-1 t 2 T3 . ( 254
)
We bound now the term T 3 :

T 3 = E ∇ u g k t ūk-1 , v k-1 t -δ k-1 t 2 (255) = E ∇ u g k t ūk-1 , v k-1 t - J-1 j=0 η k-1,j • ∇ u g k t u k-1,j t , v k-1 t η k-1 2 (256) = E J-1 j=0 η k-1,j η k-1 • ∇ u g k t ūk-1 , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 . ( 257
)
Using Jensen inequality, it follows that

T 3 ≤ J-1 j=0 η k-1,j η k-1 • E ∇ u g k t ūk-1 , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 (258) = J-1 j=0 η k-1,j η k-1 • E ∇ u g k t ūk-1 , v k-1 t -∇ u g k t u k-1 t , v k-1 t + ∇ u g k t u k-1 t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 (259) ≤ 2 • E ∇ u g k t ūk-1 , v k-1 t -∇ u g k t u k-1 t , v k-1 t 2 + 2 • J-1 j=0 η k-1,j η k-1 • E ∇ u g k t u k-1 t , v k-1 t -∇ u g k t u k-1,j t , v k-1 t 2 (260) ≤ 2L 2 • E ūk-1 -u k-1 t 2 + 2L 2 • J-1 j=0 η k-1,j η k-1 • E u k-1,j t -u k-1,0 t 2 , (261)
where we used the L-smoothness of g k t (Assumption 5 ) to obtain the last inequality. As in the centralized case (Lemma G.1), we bound terms u k-1,j t -u k-1,0 t 2 , j ∈ {0, . . . , J -1}. Using exactly the same steps as in the proof of Lemma G.1, Eq. ( 143) holds with u k-1,0 t instead of u k-1 t , i.e.,

1 -4η 2 k-1 L 2 • J-1 j=0 η k-1,j η k-1 • E u k-1,0 t -u k-1,j t 2 ≤ 2σ 2 •    J-1 j=0 η 2 k-1,j    + 4η 2 k-1 • E ∇ u g k t u k-1,0 t , v k-1 t 2 . ( 262
)
For η k-1 small enough, in particular for η k-1 ≤ 1 2 √ 2L , we have

J-1 j=0 η k-1,j η k-1 • E u k-1,0 t -u k-1,j t 2 ≤ 8η 2 k-1 • E ∇ u g k t u k-1,0 t , v k-1 t 2 + 4σ 2 •    J-1 j=0 η 2 k-1,j    (263) ≤ 8η 2 k-1 • E ∇ u g k t u k-1,0 t , v k-1 t -∇ u g k t ūk-1 , v k-1 t + ∇ u g k t ūk-1 , v k-1 t 2 + 4σ 2 •    J-1 j=0 η 2 k-1,j    (264) ≤ 16η 2 k-1 • E ∇ u g k t u k-1,0 t , v k-1 t -∇ u g k t ūk-1 , v k-1 t 2 + 16η 2 k-1 • ∇ u g k t ūk-1 , v k-1 t 2 + 4σ 2 •    J-1 j=0 η 2 k-1,j    (265) ≤ 16η 2 k-1 L 2 • E u k-1 t -ūk-1 2 + 16η 2 k-1 • ∇ u g k t ūk-1 , v k-1 t 2 + 4σ 2 •    J-1 j=0 η 2 k-1,j    , (266)
where the last inequality follows from the L-smoothness of g k t . Replacing Eq. (266) in Eq. ( 261), we have

T 3 ≤ 32η 2 k-1 L 4 • E u k-1 t -ūk-1 2 + 8L 2 σ 2 •    J-1 j=0 η 2 k-1,j    + 32η 2 k-1 L 2 • E ∇ u g k t ūk-1 , v k-1 t 2 + 2L 2 • E ūk-1 -u k-1 t 2 . ( 267
)
For η k small enough, in particular if

η k ≤ 1 2 √
2L we have,

T 3 ≤ 6L 2 E u k-1 t -ūk-1 2 + 8L 2 σ 2 J-1 j=0 η 2 k-1,j + 32η 2 k-1 L 2 ∇ u g k t ūk-1 , v k-1 t 2 . ( 268
)
Replacing Eq. ( 268) in Eq. ( 254), we have

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ 3η k-1 L 2 T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 T L • η 2 k-1,j η k-1 + 1   σ 2 - η k-1 2 E ∇ u g k ūk-1 , v k-1 1:T 2 + 16η 3 k-1 L 2 T T t=1 ∇ u g k t ūk-1 , v k-1 t 2 . (269)
We use now Assumption 7 to bound the last term in the RHS of the previous equation, leading to

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ 3η k-1 L 2 T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 T L • η 2 k-1,j η k-1 + 1   σ 2 - η k-1 • 1 -32η 2 k-1 L 2 β 2 2 E ∇ u g k ūk-1 , v k-1 1:T 2 + 16η 3 k-1 L 2 T G 2 . ( 270
)
For η k-1 small enough, in particular, if η k-1 ≤ 1 8Lβ , we have

E g k (ū k , v k-1 1:T ) -g k (ū k-1 , v k-1 1:T ) ≤ - η k-1 4 E ∇ u g k ūk-1 , v k-1 1:T 2 + 3η k-1 L 2 T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 T L • η 2 k-1,j η k-1 + 1   σ 2 + 16η 3 k-1 L 2 T G 2 . ( 271
)
We use Lemma G.14 to get

E g k (ū k , v k-1 1:T ) -f (ū k-1 , v k-1 1:T ) ≤ - η k-1 8 E ∇ u f ūk-1 , v k-1 1:T 2 + (12 + T ) η k-1 L 2 4T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 L • η 2 k-1,j η k-1 + 1   σ 2 + 16η 3 k-1 L 2 T G 2 . (272)
Finally, since g k t is a partial first-order surrogate of f t near u k-1 , v k-1 t , we have

E f (ū k , v k 1:T ) -f (ū k-1 , v k-1 1:T ) ≤ - 1 T T t=1 E d V v k t , v k-1 t - η k-1 8 E ∇ u f ūk-1 , v k-1 1:T 2 + (12 + T ) η k-1 L 2 4T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 L • η 2 k-1,j η k-1 + 1   σ 2 + 16η 3 k-1 L 2 T G 2 . ( 273
)
Lemma G.8 (Recursion for consensus distance, part 1). Suppose that Assumptions 5 -7 and Assumption 8 hold. For k ≥ τ , consider m = k τ -1 and

(η k,j ) 1≤j≤J-1 such that η k J-1 j=0 η k,j ≤ min 1 4L , 1 4Lβ
then, the updates of fully decentralized federated surrogate optimization (Alg 5) verify

E T t=1 u k t -ūk 2 F ≤ (1 - p 2 )E U mτ -Ūmτ 2 F + 44τ 1 + 2 p L 2 k-1 l=mτ η 2 l E U l -Ūl 2 F + T • σ 2 • k-1 l=mτ    η 2 l + 16τ L 2 1 + 2 p •    J-1 j=0 η 2 l,j       + 16τ 1 + 2 p G 2 k-1 l=mτ η 2 l + 16τ 1 + 2 p β 2 k-1 l=mτ η 2 l E ∇ u f ūl,j , v l 1:T 2 .
Proof. For k ≥ τ , and m = k τ -1, we have

E T t=1 u k t -ūk 2 F = E U k -Ūk 2 F (274) = E U k -Ūmτ -Ūk -Ūmτ 2 F (275) ≤ E U k -Ūmτ 2 F , (276)
where we used the fact that A -

Ā 2 F = A • I -11 T F ≤ I -11 T 2 • A 2 F = A 2
F to obtain the last inequality. Using Eq. ( 231) recursively, we have

U k = U mτ k-1 l =mτ W l - k-1 l=mτ η l Υl k-1 l =l W l . (277)
Thus,

E T t=1 u k t -ūk 2 F ≤ E U mτ k-1 l =mτ W l -Ūmτ - k-1 l=mτ η l Υl k-1 l =l W l 2 F (278) = E U mτ k-1 l =mτ W l -Ūmτ - k-1 l=mτ η l Υ l k-1 l =l W l + k-1 l=mτ η l Υ l -Υl k-1 l =l W l 2 F (279) = E U mτ k-1 l =mτ W l -Ūmτ - k-1 l=mτ η l Υ l k-1 l =l W l 2 F + E k-1 l=mτ η l Υ l -Υl k-1 l =l W l 2 F + 2E U mτ k-1 l =mτ W l -Ūmτ - k-1 l=mτ η l Υ l k-1 l =l W l , k-1 l=mτ η l Υ l -Υl k-1 l =l W l F . (280)
Since stochastic gradients are unbiased, the last term in the RHS of the previous equation is equal to zero. Using the following standard inequality for Euclidean norm with α > 0,

a + b 2 ≤ (1 + α) a 2 + 1 + α -1 b 2 , (281)
we have

E T t=1 u k t -ūk 2 F ≤ (282) (1 + α) E U mτ k-1 l =mτ W l -Ūmτ 2 F + 1 + α -1 E k-1 l=mτ η l Υ l k-1 l =l W l 2 F + k-1 l=mτ η 2 l E Υ l -Υl k-1 l =l W l 2 F . (283)
Since k ≥ (m + 1)τ and matrices W l l≥0 are doubly stochastic, we have

E T t=1 u k t -ūk 2 F ≤ (1 + α) E U mτ    (m+1)τ -1 l =mτ W l    -Ūmτ 2 F + 1 + α -1 E k-1 l=mτ η l Υ l 2 F + k-1 l=mτ η 2 l E Υ l -Υl 2 F (284) ≤ (1 + α) E U mτ    (m+1)τ -1 l =mτ W l    -Ūmτ 2 F + 1 + α -1 • (k -mτ ) k-1 l=mτ η 2 l E Υ l 2 F + k-1 l=mτ η 2 l E Υ l -Υl 2 F , (285)
where we use the fact that AB F ≤ A 2 B F and that A = 1 when A is a doubly stochastic matrix to obtain the first inequality, and Cauchy-Schwarz inequality to obtain the second one. Using Assumption 8 to bound the first term of the RHS of the previous equation and the fact that that k ≤ (m + 2)τ , it follows that

E T t=1 u k t -ūk 2 F ≤ (1 + α)(1 -p)E U mτ -Ūmτ 2 F + 2τ 1 + α -1 k-1 l=mτ η 2 l E Υ l 2 F + k-1 l=mτ η 2 l E Υ l -Υl 2 F . (286)
We use the fact that stochastic gradients have bounded variance (Assumption 6 ) to bound

E Υ l -Υl 2 F as follows, E Υ l -Υl 2 F = T t=1 E δ l t -δl t 2 (287) = T t=1 E J-1 j=0 η l,j η l • ∇ u g l+1 t u l,j t , v k-1 t -∇ u g l+1 t u l,j t , v l t ; ξ l,j t 2 (288)
≤ T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v k-1 t -∇ u g l+1 t u l,j t , v l t ; ξ l,j t 2 (289)
≤ T t=1 J-1 j=0 η l,j η l σ 2 (290) = T • σ 2 , (291)
where we used Jensen inequality to obtain the first inequality and Assumption 6 to obtain the second inequality. Replacing back in Eq. ( 286), we have

E T t=1 u k t -ūk 2 F ≤ (1 + α)(1 -p)E U mτ -Ūmτ 2 F + 2τ 1 + α -1 k-1 l=mτ η 2 l E Υ l 2 F + T • σ 2 • k-1 l=mτ η 2 l . (292)
The last step of the proof consists in bounding E Υ l 2 F for l ∈ {mτ, . . . , k -1},

E Υ l 2 F = T t=1 E δ l t 2 (293) = T t=1 E J-1 j=0 η l,j η l • ∇ u g l+1 t u l,j t , v l t 2 (294)
≤ T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v l t 2 (295)
≤ T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v l t -∇ u f t u l t , v l t + ∇ u f t u l t , v l t 2 (296) ≤ 2 T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v l t -∇ u f t u l t , v l t 2 + 2 T t=1 E ∇ u f t u l t , v l t 2 . (297) 53
Since g l+1 t is a first order surrogate of f near u l t , v l t , we have

E Υ l 2 F ≤ 2 T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v l t -∇ u g l+1 t u l,0 t , v l t 2 + 2 T t=1 E ∇ u f t u l t , v l t -∇ u f t ūl , v l t + ∇ u f t ūl , v l t 2 (298) ≤ 2 T t=1 J-1 j=0 η l,j η l • E ∇ u g l+1 t u l,j t , v l t -∇ u g l+1 t u l,0 t , v l t 2 + 4 T t=1 E ∇ u f t u l t , v l t -∇ u f t ūl , v l t 2 + 4 T t=1 E ∇ u f t ūl , v l t 2 . ( 299
)
Since f is 2L-smooth w.r.t u (Lemma G.12) and g is L-smooth w.r.t u (Assumption 5 ), we have

E Υ l 2 F ≤ 2 T t=1 J-1 j=0 η l,j η l • L 2 E u l,j t -u l,0 t 2 + 16L 2 • T t=1 E u l t -ūl 2 + 4 T t=1 E ∇ u f t ūl , v l t 2 . (300)
We use Eq. ( 266) to bound the first term in the RHS of the previous equation, leading to

E Υ l 2 F ≤ 32η 2 l L 2 T t=1 E ∇ u g l+1 t ūl,j , v l t 2 + 16L 2 1 + 2η 2 l L 2 • T t=1 E u l t -ūl 2 + 4 T t=1 E ∇ u f t ūl , v l t 2 + 8T L 2 σ 2 •    J-1 j=0 η 2 l,j    . (301)
Using Lemma G.14, we have

E Υ l 2 F ≤ 4 1 + 16η 2 l L 2 • T t=1 E ∇ u f t ūl,j , v l t 2 + 16L 2 1 + 6η 2 l L 2 • T t=1 E u l t -ūl 2 + 8L 2 σ 2 T •    J-1 j=0 η 2 l,j    . (302)
For η l small enough, in particular, for η l ≤ 1 4L , we have

E Υ l 2 F ≤ 8 T t=1 E ∇ u f t ūl,j , v l t 2 + 22L 2 E U l -Ūl 2 F + 8L 2 σ 2 T    J-1 j=0 η 2 l,j    . (303)
Replacing Eq. (303) in Eq. ( 292), we have

E T t=1 u k t -ūk 2 F ≤ (1 + α)(1 -p)E U mτ -Ūmτ 2 F + 44τ 1 + α -1 L 2 k-1 l=mτ η 2 l E U l -Ūl 2 F + 16τ 1 + α -1 k-1 l=mτ η 2 l T t=1 E ∇ u f t ūl,j , v l t 2 + T • σ 2 • k-1 l=mτ    η 2 l + 16τ L 2 1 + α -1 •    J-1 j=0 η 2 l,j       . (304)
Using Lemma G.13 and considering α = p 2 , we have

E T t=1 u k t -ūk 2 F ≤ (1 - p 2 )E U mτ -Ūmτ 2 F + 44τ 1 + 2 p L 2 k-1 l=mτ η 2 l E U l -Ūl 2 F + T • σ 2 • k-1 l=mτ    η 2 l + 16τ L 2 1 + 2 p •    J-1 j=0 η 2 l,j       + 16τ 1 + 2 p G 2 k-1 l=mτ η 2 l + 16τ 1 + 2 p β 2 k-1 l=mτ η 2 l E ∇ u f ūl,j , v l 1:T 2 . ( 305
)
Lemma G.9 (Recursion for consensus distance, part 2). Suppose that Assumptions 5 -7 and Assumption 8 hold. Consider m = k τ , then, for

(η k,j ) 1≤j≤J-1 such that η k J-1 j=0 η k,j ≤ min 1 4L , 1 4Lβ
, the updates of fully decentralized federated surrogate optimization (Alg 5) verify

E T t=1 u k t -ūk 2 F ≤ (1 + p 2 )E U mτ -Ūmτ 2 F + 44τ 1 + 2 p L 2 k-1 l=mτ η 2 l E U l -Ūl 2 F + T • σ 2 • k-1 l=mτ    η 2 l + 16τ L 2 1 + 2 p •    J-1 j=0 η 2 l,j       + 16τ 1 + 2 p G 2 k-1 l=mτ η 2 l + 16τ 1 + 2 p β 2 k-1 l=mτ η 2 l E ∇ u f ūl,j , v l 1:T 2 . ( 306
)
Proof. We use exactly the same proof as in Lemma G.8, with the only difference that Eq. (284)-Eq. ( 286) is replaced by

E T t=1 u k t -ūk 2 F ≤ (1 + α)E U mτ -Ūmτ 2 F + 2τ 1 + α -1 k-1 l=mτ η 2 l E Υ l 2 F + k-1 l=mτ η 2 l E Υ l -Υl 2 F , (307)
resulting from the fact that

(m+1)τ -1 l =mτ
W l is a doubly stochastic matrix.

Lemma G.10. Under Assum. 5 -7 and Assum 8. For η k,j = η J with

η ≤ min 1 4L , p 92τ L , 1 4βL , 1 32 √ 2 • p τ β ,
the iterates of Alg. 5 verifies

(12 + T )L 2 4T K k=0 E U k -Ūk 2 F ≤ 1 16 K k=0 E ∇ u f ūk , v k 1:T 2 +16A• 12 + T T • τ L 2 p (K +1)η 2 ,
(308) for some constant A > 0 and K > 0.

Proof. Note that for k > 0, η k = J-1 j=0 η kj = η, and that

k-1 l=mτ η 2 l = k-1 l=mτ η 2 ≤ 2τ • η 2
Using Lemma G.8 and Lemma G.9, and the fact that p ≤ 1, we have for m

= k τ -1 E U k -Ūk 2 F ≤ (1 - p 2 )E U mτ -Ūmτ 2 F + 132τ p L 2 η 2 k-1 l=mτ E U l -Ūl 2 F +η 2 2τ T σ 2 1 + 16τ L 2 J 1 + 2 p + 16τ 1 + 2 p G 2 A + 16τ p β 2 η 2 k-1 l=mτ E ∇ u f ūl , v l 1:T 2 . ( 309
)
and for m = k τ ,

E U k -Ūk 2 F ≤ (1 + p 2 )E U mτ -Ūmτ 2 F + 132τ p L 2 η 2 k-1 l=mτ E U l -Ūl 2 F +η 2 2τ T σ 2 1 + 16τ L 2 J 1 + 2 p + 16τ 1 + 2 p G 2 A + 16τ p β 2 D η 2 k-1 l=mτ E ∇ u f ūl , v l 1:T 2 . (310)
Using the fact that η ≤ p 92τ L , it follows that for m = k τ -1

E U k -Ūk 2 F ≤ (1 - p 2 )E U mτ -Ūmτ 2 F + p 64τ k-1 l=mτ E U l -Ūl 2 +η 2 A + Dη 2 k-1 l=mτ E ∇ u f ūl , v l 1:T 2 , (311)
and for m = k τ ,

E U k -Ūk 2 F ≤ (1 + p 2 )E U mτ -Ūmτ 2 F + p 64τ k-1 l=mτ E U l -Ūl 2 F +η 2 A + Dη 2 k-1 l=mτ E ∇ u f ūl , v l 1:T 2 . ( 312
)
The rest of the proof follows using [31, Lemma 14]

with B = (12+T )L 2 4T , b = 1 8 , constant (thus 8τ p -slow 7 ) steps-size η ≤ 1 32 √ 2 p τ β = 1 16 p/8
Dτ and constant weights ω k = 1.

Theorem 3.3 . Under Assumptions 4 -7 and Assumption 8, when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, the iterates of fully decentralized federated surrogate optimization (Alg. 5) satisfy:

1 K K k=1 E ∇ u f ūk , v k 1:T 2 ≤ O 1 √ K , (313)
and,

1 K K k=1 T t=1 ω t • E d V v k t , v k+1 t ≤ O 1 K , (314)
where ūk = 1 T T t=1 u k t . Moreover, local estimates u k t 1≤t≤T converge to consensus, i.e., to ūk :

1 K K k=1 T t=1 E u k t -ūk 2 ≤ O 1 √ K . (315)
Proof. We prove first the convergence to a stationary point in u, i.e. Eq. (313), using [START_REF] Koloskova | A Unified Theory of Decentralized SGD with Changing Topology and Local Updates[END_REF]Lemma 17], then we prove Eq. (314) and Eq. (315).

Note that for

K large enough, η ≤ min 1 4L , p 92τ L , 1 4βL , 1 32 √ 2 • p τ β .
Proof of Eq. 313. Rearranging the terms in the result of Lemma G.7 and dividing it by η we have

1 η • E f (ū k , v k 1:T ) -f (ū k-1 , v k-1 1:T ) ≤ - 1 8 E ∇ u f ūk-1 , v k-1 1:T 2 + (12 + T ) L 2 4T • E U k-1 -Ūk-1 2 + ηL T 4L J + 1 σ 2 + 16η 2 L 2 T G 2 . ( 316
)
Summing over k ∈ [K + 1], we have

1 η • E f (ū K+1 , v K+1 1:T ) -f (ū 0 , v 0 1:T ) ≤ - 1 8 K k=0 E ∇ u f ūk , v k 1:T 2 + (12 + T ) L 2 4T • K k=0 E U k -Ūk 2 + (K + 1)ηL T 4L J + 1 σ 2 + 16(K + 1) • η 2 L 2 T G 2 . (317)
Using Lemma G.10, we have

1 η • E f (ū K+1 , v K+1 1:T ) -f (ū 0 , v 0 1:T ) ≤ - 1 16 K k=0 E ∇ u f ūk , v k 1:T 2 + 16A • 12 + T T • τ L 2 p (K + 1)η 2 + (K + 1)ηL T 4L J + 1 σ 2 + 16(K + 1)η 2 L 2 T G 2 . ( 318
)
Using Assumption 4 , it follows that

1 16 K k=0 E ∇ u f ūk , v k 1:T 2 ≤ f (ū 0 , v 0 1:T ) -f * η + 16A • 12 + T T • τ L 2 p (K + 1)η 2 + (K + 1)ηL T 4L J + 1 σ 2 + 16(K + 1)η 2 L 2 T G 2 .
(319)

We divide by K + 1 and we have

1 16(K + 1) K k=0 E ∇ u f ūk , v k 1:T 2 ≤ f (ū 0 , v 0 1:T ) -f * η(K + 1) + 16A • 12 + T T • τ L 2 p η 2 + ηL T 4L J + 1 σ 2 + 16η 2 L 2 T G 2 . ( 320
)
The final result follows from [START_REF] Koloskova | A Unified Theory of Decentralized SGD with Changing Topology and Local Updates[END_REF]Lemma 17].

Proof of Eq. 315. We multiply Eq. (308) (Lemma G.10) by 1 K+1 , and we have

1 K + 1 K k=0 E U k -Ūk 2 F ≤ 1 16(K + 1) K k=0 E ∇ u f ūk , v k 1:T 2 F + 64Aτ p(K + 1)
Kη 2 , (321) since η ≤ O 1 √ K , using Eq. ( 313), it follows that

1 K K k=1 E U k -Ūk 2 F ≤ O 1 √ K . (322)
Thus,

1 K K k=1 T t=1 E u k t -ūk 2 F ≤ O 1 √ K . ( 323
)
Proof of Eq. 314. Using the result of Lemma G.7 we have

1 T T t=1 E d V v k t , v k-1 t ≤ E f (ū k-1 , v k-1 1:T ) -f (ū k , v k 1:T ) + (12 + T ) η k-1 L 2 4T • T t=1 E u k-1 t -ūk-1 2 + η 2 k-1 L T   4 J-1 j=0 L • η 2 k-1,j η k-1 + 1   σ 2 + 16η 3 k-1 L 2 T G 2 . ( 324
)
The final result follows from the fact that η = O 1 √ K and Eq. (315).

G.2.3 Proof of Theorem 3.3

We state the formal version of Theorem 3.3, for which only an informal version was given in the main text. Theorem 3.3. Under Assumptions 1-8, when clients use SGD as local solver with learning rate η = a0 √ K , D-FedEM's iterates satisfy the following inequalities after a large enough number of communication rounds K:

1 K K k=1 E ∇ Θ f Θk , Π k 2 F ≤ O 1 √ K , 1 K K k=1 T t=1 n t n KL π k t , π k-1 t ≤ O 1 K , (325)
where Θk = Θ k 1 , . . . Θ k T • 11 T . Moreover, individual estimates Θ k t 1≤t≤T converge to consensus, i.e., to Θk :

min k∈[K] E T t=1 Θ k t -Θk 2 F ≤ O 1 √ K .
Proof. We prove this result as a particular case of Theorem 3.3 . To this purpose, we consider that V ∆ M , u = Θ ∈ R dM , v t = π t , and ω t = n t /n for t ∈ [T ]. For k > 0, we define g k t as follow,

g k t Θ, π t = 1 n t nt i=1 M m=1 q k t z (i) t = m • l h θm (x (i) t ), y (i) t
-log p m (x

(i) t ) -log π t + log q k t z (i) t = m -c , ( 326
)
where c is the same constant appearing in Assumption 3, Eq. ( 3). With this definition, it is easy to check that the federated surrogate optimization algorithm (Alg. 5) reduces to D-FedEM (Alg. 4).

Proof. Consider arbitrary u, v ∈ R du × V, and for t ∈ [T ], consider g {u,v} to be a partial first-order surrogate of f t near {u, v}. We write Assumption 7 for g

ω t • ∇ u f t (u, v) 2 ≤ G 2 + β 2 T t=1 ω t • ∇ u f t (u, v) 2 . ( 333
)
Remark 4. Note that the assumption of Lemma G.13 is implicitly verified in Alg. 3 and Alg. 5, where we assume that every client t ∈ T canfunction compute a partial first-order surrogate of its local objective f t near any iterate (u, v) ∈ R du × V. Lemma G.14. For k > 0, the iterates of Alg. 5, verify the following inequalities:

g k ūk-1 , v k-1 1:T ≤ f ūk-1 , v k-1 1:T + L 2 T t=1 ω t ūk-1 -u k-1 t 2 , ∇ u f ūk-1 , v k-1 1:T 2 ≤ 2 ∇ u g k ūk-1 , v k-1 1:T 2 + 2L 2 T t=1 ω t ūk-1 + u k-1 t 2 ,
and,

∇ u g k ūk-1 , v k-1 1:T 2 ≤ 2 ∇ u f ūk-1 , v k-1 1:T 2 + 2L 2 T t=1 ω t ūk-1 -u k-1 t 2 ,
Proof. For k > 0 and t ∈ [T ], we have

g k t ūk-1 , v k-1 t = g k t ūk-1 , v k-1 t + f t ūk-1 , v k-1 t -f t ūk-1 , v k-1 t (334) = f t ūk-1 , v k-1 t + r k t ūk-1 , v k-1 t (335) = f t ūk-1 , v k-1 t + r k t ūk-1 , v k-1 t -r k t u k-1 t , v k-1 t + r k t u k-1 t , v k-1 t . (336)
Since

g k t u k t , v k-1 t = f t u k t , v k-1 t
(Definition 1), it follows that

g k t ūk-1 , v k-1 t = f t ūk-1 , v k-1 t + r k t ūk-1 , v k-1 t -r k t u k-1 t , v k-1 t . (337)
Because r k t is L-smooth in u (Definition 1), we have

r k t ūk-1 , v k-1 t -r k t u k-1 t , v k-1 t ≤ ∇ u r k t u k-1 t , v k-1 t , ūk-1 -u k-1 t + L 2 ūk-1 -u k-1 t 2 . ( 338
)
Since g k t is a partial first order surrogate of We have ∇ u r k t u k-1 t , v k-1 t = 0, thus

g k t ūk-1 , v k-1 t ≤ f t ūk-1 , v k-1 t + L 2 ūk-1 -u k-1 t 2 . ( 339
)
Multiplying by ω t and summing for t ∈ [T ], we have

g k ūk-1 , v k-1 1:T ≤ f ūk-1 , v k-1 1:T + L 2 T t=1 ω t ūk-1 -u k-1 t 2 , (340)
and the first inequality is proved.

+ M m=1     M m =1 m =m (α m • α m • x m • u m • u m • x m )     (353) = M m=1     -α m • (1 -α m ) • x m , u m 2 + α m • x m , u m M m =1 m =m α m • x m , u m     . ( 354
) Since α ∈ ∆ M , ∀m ∈ [M ], M m =1 m =m α m = (1 -α m ) , (355) thus,
x

• H • x = M m=1 α m • x m , u m • M m =1 m =m α m x m , u m -x m , u m (356) = M m=1 α m • x m , u m • M m =1 α m x m , u m -x m , u m (357) = M m=1 α m • x m , u m 2 - M m=1 α m • x m , u m 2 . ( 358
)
Using Jensen inequality, we have x • H • x ≤ 0.

I Details on Experimental Setup I.1 Datasets and Models

In this section we provide detailed description of the datasets and models used in our experiments. We used a synthetic dataset, verifying Assumptions 1-3, and five "real" datasets (CIFAR-10/CIFAR-100 [START_REF] Krizhevsky | Learning multiple layers of features from tiny images[END_REF], sub part of EMNIST [START_REF] Cohen | EMNIST: Extending MNIST to handwritten letters[END_REF], sub part of FEMNIST [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF][START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF] and Shakespeare [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF][START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF]) from which, two (FEMNIST and Shakespeare) has natural client partitioning. Below, we give a detailed description of the datasets and the models / tasks considered for each of them.

I.1.1 CIFAR-10 / CIFAR-100 CIFAR-10 and CIFAR-100 are labeled subsets of the 80 million tiny images dataset. They both share the same 60, 000 input images. CIFAR-100 has a finer labeling, with 100 unique labels, in comparison to CIFAR-10, having 10 unique label. We used Dirichlet allocation [START_REF] Wang | Federated Learning with Matched Averaging[END_REF], with parameter α = 0.4 to partition CIFAR-10 among 80 clients. We used Pachinko allocation [START_REF] Sashank | Adaptive Federated Optimization[END_REF] with parameters α = 0.4 and β = 10 to partition CIFAR-100 on 100 clients. For both of them we train MobileNet-v2 [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF] architecture with an additional linear layer. We used TorchVision [START_REF] Marcel | Torchvision the Machine-Vision Package of Torch[END_REF] implementation of MobileNet-v2.

I.1.2 EMNIST

EMNIST (Extended MNIST) is a 62-class image classification dataset, extending the classic MNIST dataset. In our experiments, we consider 10% of the EMNIST dataset, that we partition using Dirichlet allocation of parameter α = 0.4 over 100 clients. We train the same convolutional network as in [START_REF] Sashank | Adaptive Federated Optimization[END_REF]. The network has two convolutional layers (with 3 × 3 kernels), max pooling, and dropout, followed by a 128 unit dense layer.

I.1.3 FEMNIST

FEMNIST (Federated Extended MNIST) is a 62-class image classification dataset built by partitioning the data of Extended MNIST based on the writer of the digits/characters. In our experiments, we used a subset with 15% of the total number of writers in FEMNIST. We train the same convolutional network as in [START_REF] Sashank | Adaptive Federated Optimization[END_REF]. The network has two convolutional layers (with 3 × 3 kernels), max pooling, and dropout, followed by a 128 unit dense layer.

I.1.4 Shakespeare

This dataset is built from The Complete Works of William Shakespeare and is partitioned by the speaking roles [START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF]. In our experiments, we discarded roles with less than two sentences. We consider character-level based language modeling on this dataset. The model takes as input a sequence of 200 English characters and predicts the next character. The model embeds the 80 characters into a learnable 8-dimensional embedding space, and uses two stacked-LSTM layers with 256 hidden units, followed by a densely-connected layer. We also normalized each character by its frequency of appearance.

I.1.5 Synthetic dataset

Our synthetic dataset has been generated according to Assumptions 1-3 as follows:

J.6 Additional Results under Client Sampling

In our experiments, except for Figure 1, we considered that all clients participate at each round. We run extra experiments with client sampling, by allowing only 20% of the clients to participate at each round. We also incorporate APFL [START_REF] Deng | Adaptive Personalized Federated Learning[END_REF] into the comparison. Table 8 summarizes our findings, giving the average and standard deviation of the test accuracy across 3 independent runs.

J.7 Convergence Plots

Figures 3 to 8 show the evolution of average train loss, train accuracy, test loss, and test accuracy over time for each experiment shown in Table 2.

Assumption 3 .

 3 H = {h θ } θ∈R d is a set of hypotheses parameterized by θ ∈ R d , whose convex hull is in H. For each distribution Dm with m ∈ [M ], there exists a hypothesis h θ * m , such that l h θ * m (x) , y = -log p m (y|x) + c,

  t (x), y) • p m y|x, θm d y p (x) d x (49) = x∈X M m=1 πtm c ht (x) -y∈Y p m y|x, θm log p ht (y|x) d y p (x) d x (50) = x∈X c ht (x) -M m=1 πtm y∈Y p m y|x, θm log p ht (y|x) d y p (x) d x (51)

Figure 3 :

 3 Figure 3: Train loss, train accuracy, test loss, and test accuracy for CIFAR10 [33]. .

Figure 4 :

 4 Figure 4: Train loss, train accuracy, test loss, and test accuracy for CIFAR100 [33].

Figure 5 :

 5 Figure 5: Train loss, train accuracy, test loss, and test accuracy for EMNIST [8].

Figure 6 :

 6 Figure 6: Train loss, train accuracy, test loss, and test accuracy for FEMNIST [7, 47].

Figure 7 :

 7 Figure 7: Train loss, train accuracy, test loss, and test accuracy for Shakespeare [7, 47].

.

  Assumption 4. The negative log-likelihood f is bounded below by f * ∈ R. ) 2 ≤ σ 2 . Assumption 7. (Bounded dissimilarity) There exist β and G such that for any set of weights α ∈ ∆ M :

	Algorithm 1: FedEM (see also the more detailed Alg. 2 in App. D.1)
	Input : Data S 1:T ; number of mixture distributions M ; number of communication rounds K Output : θ K m , m ∈ [M ]
	for iterations k = 1, . . . , K do
	server broadcasts θ k-1 m , 1 ≤ m ≤ M , to the T clients;
	for tasks t = 1, . . . , T in parallel over T clients do
	for component m = 1, . . . , M do
	update q k t (z t = m) as in (8), ∀i ∈ {1, . . . , n t }; (i)
	update π k tm as in (9);
	θ k m,t ← LocalSolver(m, θ k-1 m , q k t , S t );
	client t sends θ k m,t , 1 ≤ m ≤ M , to the server;
	for component m = 1, . . . , M do
	θ k m ←	T t=1	nt n × θ k m,t ;
	Assumption 7 limits the level of dissimilarity of the different tasks, similarly to what is done in [66].
	Theorem 3.2. Under Assumptions 1-7, when clients use SGD as local solver with learning rate η = a0 √ K , after a large enough number of communication rounds K, FedEM's iterates satisfy:
	Assumption 5. (Smoothness) For all t ∈ [T ] and i ∈ [n t ], the function θ → l(θ; s (i) t ) is L-smooth
	and twice continuously differentiable.
	Assumption 6. (Unbiased gradients and bounded variance) Each client t ∈ [T ] can sample a random
	batch ξ from S t and compute an unbiased estimator g t (θ, ξ) of the local gradient with bounded
	variance, i.e., E ξ [g t (θ, ξ)] = 1 nt	nt i=1 ∇ θ l(θ; s (i) t ) and E ξ g t (θ, ξ) -1 nt	nt i=1 ∇ θ l(θ; s (i)
	T t=1	n t n	1 n t	nt i=1	M m=1	α m • l(θ; s (i) t )	2	≤ G 2 + β 2 1 n	t=1 T	i=1 nt	m=1 M	α m • l(θ; s (i) t )	2	.

t

Table 1 :

 1 2). 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. Datasets and models (details in App. I.1).

	Dataset	Task	Clients	Total samples	Model
	FEMNIST [7]	Handwritten character recognition	539	120, 772	2-layer CNN + 2-layer FFN
	EMNIST [8]	Handwritten character recognition	100	81, 425	2-layer CNN + 2-layer FFN
	CIFAR10 [33]	Image classification	80	60, 000	MobileNet-v2 [55]
	CIFAR100 [33]	Image classification	100	60, 000	MobileNet-v2 [55]
	Shakespeare [7, 47]	Next-Character Prediction	778	4, 226, 158	Stacked-LSTM [25]
	Synthetic	Binary Classification	300	1, 570, 507	Linear model

Table 2 :

 2 Test accuracy: average across clients / bottom decile. At each round, clients train for one epoch. Results for D-FedEM are in App. J.1. A comparison with MOCHA [59], which can only train linear models, is presented in App. J.2.

	Dataset	Local	FedAvg [47] FedProx [38] FedAvg+ [27] Clustered FL [56] pFedMe [16] FedEM (Ours)
	FEMNIST	71.0 / 57.5	78.6 / 63.9	78.9 / 64.0	75.3 / 53.0	73.5 / 55.1 74.9 / 57.6 79.9 / 64.8
	EMNIST	71.9 / 64.3	82.6 / 75.0	83.0 / 75.4	83.1 / 75.8	82.7 / 75.0 83.3 / 76.4 83.5 / 76.6
	CIFAR10	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
	CIFAR100	31.5 / 19.9	40.9 / 33.2	41.0 / 33.2	39.0 / 28.3	41.5 / 34.1 41.8 / 32.5 44.1 / 35.0
	Shakespeare 32.0 / 16.6 46.7 / 42.8	45.7 / 41.9	40.0 / 25.5	46.6 / 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

Table 3 :

 3 Average test accuracy across clients unseen at training (train accuracy in parenthesis).

Dataset

FedAvg

[START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF]

FedAvg+ [27] FedEM (Ours) FEMNIST 78.3 (80.9) 74.2 (84.2) 79.1 (81.5) EMNIST 83.4 (82.7) 83.7 (92.9) 84.0 (83.3) CIFAR10 77.3 (77.5) 80.4 (80.5) 85.9 (90.7) CIFAR100 41.1 (42.1) 36.5 (55.3) 47.5 (46.6) Shakespeare 46.7 (47.1) 40.2 (93.0) 46.7 (46.6) Synthetic 68.6 (70.0) 69.1 (72.1) 73.0 (74.1)

  Client-Server Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Fully Decentralized Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . .

	Table of Contents
	A Proof of Proposition 2.1
	B Relation with Other Multi-Task Learning Frameworks
	C Centralized Expectation Maximization
	D Detailed Algorithms
	D.1

E Details on the Fully Decentralized Setting F Federated Surrogate Optimization F.1 Reminder on Basic (Centralized) Surrogate Optimization . . . . . . . . . . . . . F.2 Novel Federated Version . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.3 Illustration: Analyzing pFedMe with Federated Surrogate Optimization . . . . . G Convergence Proofs G.1 Client-Server Setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . G.2 Fully Decentralized Setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . G.3 Supporting Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . H Distributed Surrogate Optimization with Black-Box Solver H.1 Supporting Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . H.2 Proof of Theorem H.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . H.3 Proof of Theorem H.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . I Details on Experimental Setup I.1 Datasets and Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . I.2 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . J Additional Experimental Results J.1 Fully Decentralized Federated Expectation-Maximization . . . . . . . . . . . . J.2 Comparison with MOCHA . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . J.3 Generalization to Unseen Clients . . . . . . . . . . . . . . . . . . . . . . . . . J.4 FedEM and Clustering . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . J.5 Effect of M in Time-Constrained Setting . . . . . . . . . . . . . . . . . . . . . J.6 Additional Results under Client Sampling . . . . . . . . . . . . . . . . . . . . . J.7 Convergence Plots . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

  )Alg. 2 is a detailed version of Alg. 1 (FedEM), with local SGD used as local solver.Alg. 3 gives our general algorithm for federated surrogate optimization, from which Alg. 2 is derived. R d for 1 ≤ m ≤ M ; for tasks t = 1, . . . , T in parallel over T clients do

	D Detailed Algorithms
	D.1 Client-Server Algorithm
	Algorithm 2: FedEM: Federated Expectation-Maximization
	Input : Data S 1:T ; number of mixture components M ; number of communication rounds K;
	number of local steps J
	Output :θ K m for 1 ∈ [M ]; π K t for t ∈ [T ]
	// Initialization
	server randomly initialize θ 0 m ∈ Randomly initialize π 0 t ∈ ∆ M ;
	// Main loop
	for iterations k = 1, . . . , K do
	server broadcasts θ k-1

m , 1 ≤ m ≤ M to the T clients; for tasks t = 1, . . . , T in parallel over T clients do for component m = 1, . . . , M do // E-step 8 for sample i = 1, . . . , n t do 9

  Alg. 4 shows D-FedEM, the fully decentralization version of our federated expectation maximization algorithm.Alg. 5 gives our general fully decentralized algorithm for federated surrogate optimization, from which Alg. 4 is derived. Algorithm 4: D-FedEM: Fully Decentralized Federated Expectation-Maximization Input : Data S 1:T ; number of mixture components M ; number of iterations K; number of local steps J; mixing matrix distributions W k for k ∈ [K] Output : θ K m,t for m ∈ [M ] and t ∈ [T ]; π t for t ∈ [T ] // Initialization for tasks t = 1, . . . , T in parallel over T clients do Randomly initialize Θ t = (θ m,t ) 1≤m≤M ∈ R M ×d ; Randomly initialize π 0 t ∈ ∆ M ; // Main loop for iterations k = 1, . . . , K do // Select the communication topology and the aggregation weights

	D.2 Fully Decentralized Algorithm
	return Θ;

1,j

);

Sample W k-1 ∼ W k-1 ;

for tasks t = 1, . . . , T in parallel over T clients do for component m = 1, . . . , M do // E-step

8

for sample i = 1, . . . , n t do 9

  Assumption 4 . The objective function f is bounded below by f * ∈ R. Assumption 5 . (Smoothness) For all t ∈ [T ] and k > 0, g k t is L-smooth wrt to u. Assumption 6 . (Unbiased gradients and bounded variance) Each client t ∈ [T ] can sample a random batch ξ from S t and compute an unbiased estimator ∇ u g k t (u, v; ξ) of the local gradient with bounded variance, i.e., E ξ

  {u,v} , is a partial first-order surrogate of f t near {u, v}, it follows that

	T	ω t • ∇ u g t {u,v}	(u, v)	2	≤ G 2 + β 2	T	ω t • ∇ u g t {u,v}	(u, v)	2	.	(332)
	t=1					t=1
	{u,v} t Since g
		T
		t=1

  1. Sample weightπ t ∼ Dir (α) , t ∈ [T ] from a symmetric Dirichlet distribution of parameter α ∈ R + 2. Sample θ m ∈ R d ∼ U [-1, 1] d , m ∈ [M ] for uniform distribution over [-1, 1]

d . 3. Sample m t , t ∈ [T ] from a log-normal distribution with mean 4 and sigma 2, then set n t = min (50 + m t , 1000). 4. For t ∈ [T ] and i ∈ [n t ], draw x (i) t ∼ U [-1, 1] d and (i) t ∼ N (0, I d ). 5. For t ∈ [T ] and i ∈ [n t ], draw z (i) t ∼ M (π t ).

Table 7 :

 7 Test and train accuracy comparison across different tasks. For each method, the best test accuracy is reported. For FedEM we run only K M rounds, where K is the total number of rounds for other methods-K = 80 for Shakespeare and K = 200 for all other datasets-and M = 3 is the number of components used in FedEM.

	Dataset	Local	FedAvg [47]	FedProx [38]	FedAvg+ [27]	Clustered FL [56]	pFedMe [16]	FedEM (Ours)

Table 8 :

 8 Test accuracy under 20% client sampling: average across clients with +/-standard deviation over 3 independent runs. All experiments with 1200 communication rounds. ± 0.14 77.7 ± 0.16 77.8 ± 0.07 78.2 ± 0.27 82.1 ± 0.13 CIFAR100 [33] 40.6 ± 0.17 39.7 ± 0.75 39.9 ± 0.08 40.3 ± 0.71 43.2 ± 0.23 Synthetic 68.2 ± 0.02 69.0 ± 0.03 69.1 ± 0.03 69.1 ± 0.04 74.7 ± 0.01

	Dataset	FedAvg [47]	FedAvg+ [27]	pFedMe [16]	APFL [14]	FedEM (Ours)
	CIFAR10 [33]	73.1

In the case of (multi-output) regression, we have ht : X → R d for some d ≥ 1 and l : R d × R d → R + .

We note that whether the conjecture in[START_REF] Ben-David | Does Unlabeled Data Provably Help? Worst-case Analysis of the Sample Complexity of Semi-Supervised Learning[END_REF] holds in the agnostic case is still an open problem.

A possible way to ensure that Assumption 2 holds is to use the batch normalization technique from[START_REF] Li | FedBN: Federated Learning on Non-IID Features via Local Batch Normalization[END_REF] to account for feature shift.

As the distribution DT over tasks in Proposition 2.1 is arbitrary, any positively weighted sum of clients' empirical losses could be considered.

For training, we sub-sampled 10% and 15% from EMNIST and FEMNIST datasets respectively.

Note that when Θ is a full rank matrix, this decomposition is unique.

The notion of τ -slow decreasing sequence is defined in [31,Defintion 2].

As shown by[START_REF] Wang | Tackling the objective inconsistency problem in heterogeneous federated optimization[END_REF] Theorem. 2], the convergence is guaranteed in two scenarios: 1) G 2 = 0, 2) All clients use take the same number of local steps using the same local solver. Note that we allow each client to use an arbitrary approximate local solver.

https://github.com/gingsmith/fmtl

Theorem 3.3 then follows immediately from Theorem 3.3 , once we verify that g k t 1≤t≤T satisfy the assumptions of Theorem 3.3 . Assumption 4 , Assumption 6 , and Assumption 7 follow directly from Assumption 4, Assumption 6, and Assumption 7, respectively. Lemma G. [START_REF] Bellet | Personalized and Private Peer-to-Peer Machine Learning[END_REF] shows that for k > 0, g k is smooth w.r.t. Θ and then Assumption 5 is satisfied. Finally, Lemmas G.4-G.6 show that for t ∈ [T ] g k t is a partial first-order surrogate of f t near Θ k-1 t , π t with d V (•, •) = KL(• •).

G.3 Supporting Lemmas

Lemma G.11. Consider J ≥ 2 and positive real numbers η j , j = 0, . . . , J -1, then:

Proof. For the first inequality,

For the second inequality

For the third inequality,

Lemma G.12. Suppose that g is a partial first-order surrogate of f , and that g is L-smooth, where L is the constant appearing in Definition 1, then f is 2L-smooth.

Proof. The difference between f and g is L-smooth, and g is L-smooth, thus f is 2L-smooth as the sum of two L-smooth functions.

Lemma G. [START_REF] Malte Darnstädt | Unlabeled Data Does Provably Help[END_REF]. Consider f = T t=1 ω t • f t , for weights ω ∈ ∆ T . Suppose that for all (u, v) ∈ R du × V, and t ∈ [T ], f t admits a partial first-order surrogate g {u,v} t near {u, v}, and that g {u,v} = T t=1 ω t • g {u,v} t verifies Assumption 7 for t ∈ [T ]. Then f also verifies Assumption 7 .

Writing the gradient of Eq. (337), we have

Multiplying by ω t and summing for t ∈ [T ], we have

Thus,

where (344) follows from a

The proof of the last inequality is similar, it leverages a + b 2 ≤ 2 a 2 + 2 a 2 to upper bound (343).

then H is a semi-definite negative matrix.

Proof.

We have:

H Distributed Surrogate Optimization with Black-Box Solver

In this section, we cover the scenario where the local SGD solver used in our algorithms (Alg. 3 and Alg. 5) is replaced by a (possibly non-iterative) black-box solver that is guaranteed to provide a local inexact solution of ∀m ∈ [M ], minimize

with the following approximation guarantee. Assumption 9 (Local α-approximate solution). There exists 0 < α < 1 such that for t

where

t ), θ k m,t is the output of the local solver at client t and θ k-1 m is its starting point (see Alg. 2).

We further assume strong convexity.

Assumption 9 is equivalent to the γ-inexact solution used in [START_REF] Li | Federated learning: Challenges, methods, and future directions[END_REF] (Lemma. H.2), when local functions (Φ t ) 1≤t≤T are assumed to be convex. We also need to have G 2 = 0 in Assumption 7 as in [START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF]Definition 3], in order to ensure the convergence of Alg. 2 and Alg. 4 to a stationary point of f , as shown by [START_REF] Wang | Tackling the objective inconsistency problem in heterogeneous federated optimization[END_REF]Theorem. 2]. 8 Theorem H.1. Suppose that Assumptions 1-7, 9 and 10 hold with G 2 = 0 and α < 1 β 2 κ 4 , then the updates of federated surrogate optimization converge to a stationary point of f , i.e.,

and

As in App. G, we provide the analysis for the general case of federated surrogate optimization (Alg. 3) before showing that FedEM (Alg. 2) is a particular case.

We suppose that, at iteration k > 0, the partial first-order surrogate functions g k t , t ∈ [T ] used in Alg. 3 verifies, in addition to Assumptions 4 -7 , the following assumptions that generalize Assumptions 9 and 10, Assumption 9 (Local α-inexact solution). There exists 0 < α < 1 such that for t ∈ [T ] and k > 0,

where u k t, * ∈ arg min u∈R du g k t u, v k t . Assumption 10 . For t ∈ [T ] and k > 0, g k t is µ-strongly convex in u.

Under these assumptions a parallel result to Theorem. H.1 holds.

Theorem H.1 . Suppose that Assumptions 4 -7 , Assumptions 9 and 10 hold with G 2 = 0 and α < 1 β 2 κ 4 , then the updates of federated surrogate optimization converges to a stationary point of f , i.e., lim

and

H.1 Supporting Lemmas

First, we prove the following result. Lemma H.2. Under Assumptions 5 , 9 and 10 , the iterates of Alg. 2 verify for k > 0 and t

where κ = L/µ.

), we have using Assumption 9 ,

thus,

Combining Eq. (367) and Eq. (369), we have

Lemma H.3. Suppose that Assumptions 5 , 7 , 9 and 10 hold with G 2 = 0. Then,

where α = β 2 κ 4 α, and

where the last inequality is a result of Lemma H.2. Using Jensen inequality, we have

Using Assumption 7 and Jensen inequality with the " √ •" function, it follows that

Since g k is L-smooth in u as a convex combination of L-smooth function, we have

Using Polyak-Lojasiewicz (PL), we have

(384) Using the L-smoothness of g k in u, we have

Thus,

Thus,

= 0 and that r k t is non-negative and L-smooth in u. Lemma H.4. Suppose that Assumptions 4 and 5 hold and that

then

If we moreover suppose that Assumption 10 holds and that there exists 0 < α < 1 such that for all k > 0,

where

It follows that the sequence f u k , v k k≥0 is a non-increasing sequence. Since f is bounded below (Assum. 4 ), it follows that f u k , v k k≥0 is convergent. Denote by f ∞ its limit. The sequence

Proof of Eq. 390 Using the fact that

Thus,

By summing over k then passing to the limit when k → +∞, we have

Proof of Eq. 391 Because the L-smoothness of u → r k u, v k 1:T , we have

Thus,

because r k is a non-negative function (Definition 1). Finally, using Eq. (390), it follows that

Proof of Eq. 393 We suppose now that there exists 0 < α < 1 such that

It follows that,

then,

and by using the definition of g k we have,

Since

From Eq. (406) and Eq. ( 407), it follows that,

Finally, since

Since g k is µ-strongly convex in u (Assumption 10), we write

It follows that,

and,

Proof.

(414) Computing the gradient norm, we have,

Since g k is L-smooth in u, we write

Thus by replacing Eq. (418) in Eq. ( 416), we have

Using Lemma H.3, there exists 0 < α < 1, such that

Thus, the conditions of Lemma H.4 hold, and we can use Eq. ( 391) and (393), i.e.

Finally, combining this with Eq. (419), we get the final result

and,

Proof. We prove this result as a particular case of Theorem H.1 . To this purpose, we consider that V ∆ M , u = Θ ∈ R dM , v t = π t , and ω t = n t /n for t ∈ [T ]. For k > 0, we define g k t as follow,

where c is the same constant appearing in Assumption 3, Eq. ( 3). With this definition, it is easy to check that the federated surrogate optimization algorithm (Alg. 3) reduces to FedEM (Alg. 2). Theorem H.1 then follows immediately from Theorem H.1 , once we verify that g k t 1≤t≤T satisfy the assumptions of Theorem H.1 . Assumption 4 , Assumption 6 , Assumption 7 , Assumption 9 and Assumption 10 follow directly from Assumption 4, Assumption 6, Assumption 7, Assumption 9 and Assumption 10, respectively. Lemma G.3 shows that for k > 0, g k is smooth w.r.t. Θ and then Assumption 5 is satisfied. Finally, Lemmas G.4-G.6 show that for t ∈ [T ] g k t is a partial first-order surrogate of

Dataset

GPU Simulation time Shakespeare [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF][START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF] Quadro RTX 8000 4h42min FEMNIST [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF] Quadro RTX 8000 1h14min EMNIST [START_REF] Cohen | EMNIST: Extending MNIST to handwritten letters[END_REF] GeForce GTX 1080 Ti 46min CIFAR10 [START_REF] Krizhevsky | Learning multiple layers of features from tiny images[END_REF] GeForce GTX 1080 Ti 2h37min CIFAR100 [START_REF] Krizhevsky | Learning multiple layers of features from tiny images[END_REF] GeForce GTX 1080 Ti 3h9min Synthetic GeForce GTX 1080 Ti 20min

Table 5: Learning rates η used for the experiments in Table 2. Base-10 logarithms are reported.

Dataset FedAvg [START_REF] Mcmahan | Communication-efficient learning of deep networks from decentralized data[END_REF] FedProx [START_REF] Li | Federated Optimization in Heterogeneous Networks[END_REF] FedAvg+ [START_REF] Jiang | Improving federated learning personalization via model agnostic meta learning[END_REF] Clustered FL [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF] pFedMe [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF] FedEM (Ours)

I.2 Implementation Details

I.2.1 Machines

We ran the experiments on a CPU/GPU cluster, with different GPUs available (e.g., Nvidia Tesla V100, GeForce GTX 1080 Ti, Titan X, Quadro RTX 6000, and Quadro RTX 8000). Most experiments with CIFAR10/CIFAR-100 and EMNIST were run on GeForce GTX 1080 Ti cards, while most experiments with Shakespeare and FEMNIST were run on the Quadro RTX 8000 cards. For each dataset, we ran around 30 experiments (not counting the development/debugging time). Table 4 gives the average amount of time needed to run one simulation for each dataset. The time needed per simulation was extremely long for Shakespeare dataset, because we used a batch size of 128.

We remarked that increasing the batch size beyond 128 caused the model to converge to poor local minima, where the model keeps predicting a white space as next character.

I.2.2 Libraries

We used PyTorch [START_REF] Paszke | PyTorch: An Imperative Style, High-Performance Deep Learning Library[END_REF] to build and train our models. We also used Torchvision [START_REF] Marcel | Torchvision the Machine-Vision Package of Torch[END_REF] implementation of MobileNet-v2 [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF], and for image datasets preprossessing. We used LEAF [START_REF] Caldas | Leaf: A benchmark for federated settings[END_REF] to build FEMNIST dataset and the federated version of Shakespeare dataset.

I.2.3 Hyperparameters

For each method and each task, the learning rate was set via grid search on the set 10 -0.5 , 10 -1 , 10 -1.5 , 10 -2 , 10 -2.5 , 10 -3 . FedProx and pFedMe's penalization parameter µ was tuned via grid search on 10 1 , 10 0 , 10 -1 , 10 -2 , 10 -3 . For Clustered FL, we used the same values of tolerance as the ones used in its official implementation [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF]. We found tuning tol 1 and tol 2 particularly hard: no empirical rule is provided in [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF], and the few random setting we tried did not show any improvement in comparison to the default ones. For each dataset and each method, Table 5 reports the learning rate η that achieved the corresponding result in Table 2.

J Additional Experimental Results

J.1 Fully Decentralized Federated Expectation-Maximization D-FedEM considers the scenario where clients communicate directly in a peer-to-peer fashion instead of relying on the central server mediation. In order to simulate D-FedEM, we consider a binomial Erdős-Rényi graph [START_REF] Erdös | On Random Graphs I[END_REF] with parameter p = 0.5, and we set the mixing weight using Fast Mixing Markov Chain [START_REF] Boyd | Fastest Mixing Markov Chain on A Graph[END_REF] rule. We report the result of this experiment in Table 6, showing the average weighted accuracy with weight proportional to local dataset sizes. We observe that D-FedEM often performs better than other FL approaches and slightly worst than FedEM, except on CIFAR-10 where it has low performances.

J.2 Comparison with MOCHA

In the case of synthetic dataset, for which train a linear model, we compare FedEM with MOCHA [START_REF] Smith | Federated Multi-Task Learning[END_REF]. We implemented MOCHA in Python following the official implementation 9 in MATLAB. We tuned the parameter λ of MOCHA on a holdout validation set via grid search in {10 1 , 10 0 , 10 -1 , 10 -2 , 10 -3 }, and we found that the optimal value of λ is 10 0 . For this value, we ran MOCHA on the synthetic dataset with three different seeds, and we found that the average accuracy is 73.4 ± 0.05 in comparison to 74.7 ± 0.01 achieved by FedEM. Note that MOCHA is the second best method after FedEM on this dataset. Unfortunately, MOCHA only works for linear models.

J.3 Generalization to Unseen Clients

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 one. Unexpectedly, new clients achieve sometimes a significantly higher test accuracy than old clients (e.g., 47.5% against 44.1% on CIFAR100).

In order to better understand this difference, we looked at the distribution of FedEM personalized weights for the old clients and new ones. The average distribution entropy equals 0.27 and 0.92 for old and new clients, respectively. This difference shows that old clients tend to have more skewed distributions, suggesting that some components may be overfitting the local training dataset leading the old clients to give them a high weight.

We also considered a setting where unseen clients progressively collect their own dataset. We investigate the effect of the number of samples on the average test accuracy across unseen clients, starting from no local data (and therefore using uniform weights to mix the M components) and progressively adding more labeled examples until the full local labeled training set is assumed to be available. Figure 2 shows that FedEM achieves a significant level of personalization as soon as clients collect a labeled dataset whose size is about 20% of what the original clients used for training.

As we mentioned in the main text, it is not clear how the other personalized FL algorithms (e.g., pFedMe and Clustered FL) should be extended to handle unseen clients. For example, the global model learned by pFedMe during training can then be used to perform some "fine-tuning" at the new clients, but how exactly? The original pFedMe paper [START_REF] Canh T Dinh | Personalized Federated Learning with Moreau Envelopes[END_REF] does not even mention this issue. For example, the client could use the global model as initial vector for some local SGD steps (similarly to what done in FedAvg+ or the MAML approaches) or it could perform a local pFedMe update (lines 6-9 in [16, Alg. 1]). The problem is even more complex for Clustered FL (and again not discussed in [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF]). The new client should be assigned to one of the clusters identified. One can think to compute the cosine distances of the new client from those who participated in training, but this would require the server to maintain not only the model learned, but also the last-iteration gradients of all clients that participated in the training. Moreover, it is not clear which metric should be considered to assign the new client to a given cluster (perhaps the average cosine similarity from all clients in the cluster?). This is an arbitrary choice as [START_REF] Sattler | Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints[END_REF] does not provide a criterion to assign clients to a cluster, but only to decide if a given cluster should be split in two new ones. It appears that many options are possible and they deserve separate investigation. Despite these considerations, we performed an additional experiment extending pFedMe to unseen clients as described in the second option above on CIFAR-100 dataset with a sampling rate of 20%. pFedMe achieves a test accuracy of 40.5% ± 1.66%, in comparison to 38.9% ± 0.97% for FedAvg and 42.7% ± 0.33% for FedEM. FedEM thus performs better on unseen clients, and pFedMe's accuracy shows a much larger variability.

J.4 FedEM and Clustering

We performed additional experiments with synthetic datasets to check if FedEM recovers clusters in practice. We modified the synthetic dataset generation so that the mixture weight vector π t of each client t has a single entry equal to 1 that is selected uniformly at random. We consider two scenarios both with T = 300 client, the first with M = 2 component and the second with M = 3 components. In both cases FedEM recovered almost the correct Π * and Θ * : we have cosine_distance Θ * , Θ ≤ 10 -2 and cosine_distance Π * , Π ≤ 10 -8 . A simple clustering algorithm that assigns each client to the component with the largest mixture weight achieves 100% accuracy, i.e., it partitions the clients in sets coinciding with the original clusters.

J.5 Effect of M in Time-Constrained Setting

Recall that in FedEM, each client needs to update and transmit M components at each round, requiring roughly M times more computation and M times larger messages than the competitors in our study.

In this experiment, we considered a challenging time-constrained setting, where FedEM is limited to run one third (= 1/M ) of the rounds of the other methods. The results in Table 7 show that even if FedEM does not reach its maximum accuracy, it still outperforms the other methods on 3 datasets.

We additionally compared FedEM with a model having the same number of parameters in order to check if FedEM's advantage comes from the additional model parameters rather than by its specific formulation. To this purpose, we trained Resnet-18 and Resnet-34 on CIFAR10. The first one has about 3 times more parameters than MobileNet-v2 and then roughly as many parameters as FedEM with M = 3. The second one has about 6 times more parameters than FedEM with M = 3. We