PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Multi-Task Learning under a Mixture of Distributions

Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, Richard Vidal

arXiv:2108.10252v4cs.LGcs.AImath.OCstat.ML

TL;DR

The paper studies how to personalize federated models when client data distributions are heterogeneous and existing methods lack clear statistical assumptions. It models each client as a mixture of shared components, develops EM-like centralized and decentralized algorithms with convergence guarantees, and reports higher accuracy and fairness, including for unseen clients. The approach also supports principled personalization of clients absent during training.

  • Problem

    A single global model can be arbitrarily bad for some heterogeneous clients, while existing federated MTL personalization lacks clear statistical assumptions about local data distributions.

  • Method

    Each client uses a personalized mixture of M shared component models, trained with EM-like FedEM or D-FedEM algorithms in client-server or fully decentralized settings.

  • Results

    Experiments report higher accuracy and fairness than state-of-the-art FL algorithms, including for clients not present during training.

  • Takeaways & Limitations

    The mixture assumption provides a principled way to transfer knowledge among clients and infer personalized models for unseen clients.

  • Takeaways & Limitations

    The framework's technical assumptions do not generally hold on standard FL benchmark datasets, although the algorithms perform well there.

Abstract

from arXiv · show

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.

1 Introduction

Federated learning personalization is motivated by heterogeneous client distributions that can make a single global model unsuitable. The paper proposes a mixture-based federated MTL framework with convergent algorithms and improved accuracy, fairness, and unseen-client generalization.

  • Motivation: Heterogeneous, non-iid client data can make a single global model arbitrarily poor for individual clients, motivating personalized models.The paper illustrates this with language modeling, where clients may prefer different next tokens after the same prefix.
  • Prior approaches: Existing personalization methods include fine-tuning, global-local interpolation, clustered FL, federated MTL, hypernetworks, and client-model interpolation.These approaches differ in their assumptions, communication patterns, and theoretical support.
  • Open problem: A remaining gap is understanding whether and when clients are guaranteed to benefit from collaboration.The paper addresses this gap by imposing an explicit assumption on local data distributions.
  • Proposed framework: The proposed framework models each client distribution as a mixture of M underlying distributions, encompassing most previously proposed personalized FL approaches.This formulation allows knowledge transfer across clients even when pairs of clients may be arbitrarily different.
  • Algorithms: FedEM and D-FedEM use federated EM-like training for client-server and fully decentralized settings, with convergence guarantees and inference for unseen clients.For unseen clients, shared component models remain fixed while mixture weights are fitted to the new client's data.
  • Results: Experiments on FL benchmarks report higher average accuracy, greater fairness across clients, and better generalization to unseen clients than state-of-the-art personalized and non-personalized methods.The paper presents these as broad empirical outcomes across benchmark datasets.

2 Problem Formulation

The paper formulates federated learning as personalized risk minimization across heterogeneous client distributions, then motivates a mixture model that enables shared components and client-specific predictors. It establishes the mixture formulation, its assumptions, training procedure, and connections to existing personalized FL frameworks.

  • 2 Problem Formulation: Each client has its own data distribution and seeks a model minimizing expected loss under that distribution, while federated learning aims to improve estimation using other clients’ samples.Clients participating initially provide i.i.d. datasets, and additional clients may join later.
  • 2.1 An Impossibility Result: Without assumptions on local distributions, federated learning can reduce a client’s sample complexity by at most a constant factor relative to local learning, regardless of other clients’ data.The impossibility argument reduces federated learning to parallel semi-supervised learning problems when clients share marginals over X but have arbitrarily different conditional label distributions.
  • 2.2 Learning under a Mixture Model: The proposed assumption models each client distribution as a mixture of M underlying distributions, allowing clients to share information even when their distributions differ arbitrarily.The underlying distributions are independent, and client-specific mixture weights determine how components combine.
  • 2.2 Learning under a Mixture Model: Under identical component marginals over X, the analysis can use discriminative models rather than requiring generative models, while the component hypothesis class has a convex hull contained in the client-model class.The paper notes that these assumptions are mainly technical and need not hold for the approach to work in practice.
  • 2.2 Learning under a Mixture Model: Training estimates shared component parameters and client mixture weights by minimizing empirical negative likelihood, then combines the components to obtain each training client’s predictor.The resulting predictors minimize expected loss under the corresponding client distributions under the stated assumptions.
  • 2.2 Learning under a Mixture Model: For an unseen client, the method keeps component models fixed, estimates mixture weights by maximizing that client’s data likelihood, and forms its personalized predictor.This provides an inference procedure for clients that were absent from the initial training phase.
  • 2.3 Generalizing Existing Frameworks: The mixture framework recovers clustered FL by using one-hot mixture weights, and model interpolation by representing global and local models as components with client-specific weights.It also connects federated multi-task learning with task-relationship formulations through linear predictors and suitable component choices.

3 Federated Expectation-Maximization

FedEM models each client distribution as a mixture of shared components with personalized weights, yielding EM-like algorithms for centralized and decentralized federation. The framework supports unseen-client personalization and provides convergence guarantees through federated surrogate optimization.

  • Client-Server Algorithm: FedEM alternates local latent-weight and mixture-weight updates with component-model updates aggregated by a central server.Clients perform local updates, transmit updated components, and the server aggregates them each communication round.
  • Convergence Guarantees: Under the paper’s assumptions, FedEM’s parameters converge toward a stationary point of the federated objective.The theorem controls both the component-model gradient and the improvement from the mixture-weight update.
  • Client Personalization: FedEM enables an unseen client to personalize a model by retrieving learned components and updating its mixture weights locally.The new client can initialize its weights, perform one E-step and the first M-step update, and obtain a personalized model without retraining the federation.
  • Fully Decentralized Algorithm: D-FedEM replaces centralized aggregation with neighborhood-based weighted averaging, while requiring communication connectivity sufficient for every client to influence others.Under the stated connectivity assumption, individual component estimates converge to a common value and the parameters converge to a stationary point.
  • Federated Surrogate Optimization: The federated surrogate optimization framework minimizes a weighted sum of client objectives using partial first-order surrogates and covers both client-server and decentralized algorithms.FedEM and D-FedEM are particular instances, and convergence follows under mild conditions when clients can compute suitable partial surrogates.

4 Experiments

Experiments cover five federated benchmark datasets and compare FedEM with global, local, clustered, and other personalized approaches. FedEM achieves the best average performance across datasets, improves the worst-client accuracy, is robust to low participation, and generalizes to unseen clients.

  • Average Performance: FedEM obtains the best average weighted accuracy across all evaluated datasets.The evaluation spans image classification, handwritten character recognition, and language modeling tasks.
  • Fairness Across Clients: FedEM also improves fairness: clients in the bottom accuracy decile perform better than with the compared training methods.The bottom decile is reported instead of minimum accuracy because minimum values are noisy when local test sets are small.
  • Client Sampling: FedEM is more robust than FedAvg+ when fewer clients participate in each communication round.The comparison varies the uniformly sampled client fraction in CIFAR10 experiments.
  • Generalization to Unseen Clients: New clients can learn personalized models at least as good as FedAvg’s global model and always better than FedAvg+’s model.In one CIFAR100 result, unseen clients reached 47.5% accuracy versus 44.1% for old clients.
  • Computational Cost: FedEM requires roughly M times more computation and M times larger messages because each client updates and transmits all M components.This is the principal stated systems cost of increasing the number of mixture components.

5 Conclusion

The proposed federated MTL approach uses mixture-based local-distribution assumptions and EM-like algorithms to learn shared components with personalized weights. It provides convergence guarantees and achieves higher accuracy and fairness, including for clients absent during training.

  • 5 Conclusion: FedEM and D-FedEM jointly learn shared component models and personalized mixture weights in client-server and fully decentralized settings.The approach is based on modeling each local distribution as a mixture of underlying distributions.
  • 5 Conclusion: The federated surrogate optimization framework provides convergence guarantees for the proposed algorithms and can analyze other federated formulations.
  • 5 Conclusion: Higher accuracy and fairness than state-of-the-art FL algorithms are reported even for clients not present during training.
  • 5 Conclusion: Reducing local computation and communication is identified as future work, including selectively updating or transmitting component models.
  • 5 Conclusion: Privacy-constrained personalized FL remains an open research direction, including differentially private versions and privacy-utility analysis.

A Proof of Proposition 2.1

The proof establishes that, under the stated assumptions, solutions of the mixture-model optimization problem yield predictors minimizing expected loss for each client. It handles mean squared error, logistic, and cross-entropy losses using entropy and KL-divergence arguments.

  • A Proof of Proposition 2.1: Proposition 2.1 covers mean squared error, logistic, and cross-entropy losses under Assumptions 1, 2, and 3.
  • A Proof of Proposition 2.1: The final result shows that the predictors minimize E_(x,y)∼D_t[l(h_t(x), y)] and therefore solve Problem (1).
  • A Proof of Proposition 2.1: For each loss case, the proof verifies that the term c_h is independent of h before applying Lemma A.3.
  • A Proof of Proposition 2.1: Non-negativity of KL divergence supplies the key inequality used to compare arbitrary predictors with client-wise loss minimizers.
  • A Proof of Proposition 2.1: Lemma A.2 relates a mixture representation to KL divergence, with equality characterized by the corresponding equality condition.

B Relation with Other Multi-Task Learning Frameworks

The appendix relates the mixture-based formulation to existing multi-task learning models and derives federated EM and surrogate-optimization algorithms. It also states convergence results for client-server and decentralized settings under explicit assumptions.

  • Relation with Other Multi-Task Learning Frameworks: The formulation is more general than approaches where clients combine fixed base hypotheses because it learns both the component models and the mixture weights.
  • Relation with Other Multi-Task Learning Frameworks: Under linear hypotheses, the mixture formulation connects to alternating structure optimization through a low-dimensional shared structure and corresponding constrained or regularized objectives.
  • Relation with Other Multi-Task Learning Frameworks: Several personalized FL formulations, including pFedMe, FedU, and formulations in [24] and [23], arise as special cases of the proposed formulation.
  • Relation with Other Multi-Task Learning Frameworks: Centralized EM alternates an E-step maximizing an evidence lower bound over auxiliary distributions and an M-step maximizing that bound over component parameters and mixture weights.
  • Relation with Other Multi-Task Learning Frameworks: D-FedEM converges to a consensus stationary point under Assumptions 1–9, while decentralized mixing requires communication matrices that connect clients over time.
  • Relation with Other Multi-Task Learning Frameworks: The general federated surrogate-optimization framework yields parallel convergence results for client-server and fully decentralized algorithms under their respective assumptions.

F.2 Novel Federated Version

The novel federated surrogate framework represents the objective as a weighted sum of client functions and lets clients iteratively construct partial first-order surrogates. Its convergence results cover both client-server and fully decentralized optimization and include pFedMe as a special case.

  • F.2 Novel Federated Version: The objective is a weighted sum of client functions, with each client computing a partial first-order surrogate with respect to shared variables.
  • F.2 Novel Federated Version: A valid partial surrogate upper-bounds the objective, has a smooth residual touching the objective at the current point, and induces a nonnegative distance condition for local variables.
  • F.2 Novel Federated Version: The resulting client-server and fully decentralized algorithms are iterative and distributed, with clients computing partial surrogates at each iteration.
  • F.2 Novel Federated Version: Under the stated smoothness, gradient, variance, and bounded-dissimilarity assumptions, the client-server method satisfies a convergence result after sufficiently many communication rounds.
  • F.2 Novel Federated Version: The fully decentralized method has a corresponding convergence result when the additional mixing-matrix assumption holds.
  • F.2 Novel Federated Version: pFedMe can be viewed as a particular case of federated surrogate optimization, so the framework's convergence results apply to it.

G.1.2 Proof of Theorem 3.2′

The proof establishes Theorem 3.2′ by bounding successive terms using Jensen’s inequality, smoothness, unbiased stochastic gradients, and bounded variance. These bounds yield the stated behavior of federated surrogate optimization after sufficiently many communication rounds.

  • Term-wise bounds: The proof decomposes the relevant expressions into terms such as T1, T2, and T3 and bounds them separately.The derivation repeatedly substitutes intermediate bounds into earlier inequalities.
  • Inequality tools: Jensen’s inequality is used to control several terms in the recursive bounds.The proof applies Jensen’s inequality both during individual term estimates and when bounding a final right-hand-side term.
  • Stochastic assumptions: Unbiased stochastic gradients with variance bounded by σ2 provide the stochastic-error control required in the estimates.These properties are invoked under Assumption 6′.
  • Smoothness and step sizes: The proof also uses L-smoothness of gk and sufficiently small learning rates to complete the recursive bound.The small-step conditions include explicit upper bounds on ηk−1 involving L and β.
  • Theorem statement: Theorem 3.2′ assumes Assumptions 4′–8′ and SGD with learning rate η = a0 √K, then states an iterate bound after enough rounds.The expectation is taken over random batch samples.

G.1.3 Proof of Theorem 3.2

The proof of Theorem 3.2 derives FedEM’s convergence result as a special case of the federated surrogate optimization theorem. It verifies smoothness, majorization, and surrogate properties before applying the general result.

  • Objective definition: The client-level negative log-likelihood is introduced as the local objective used in the analysis.The global function f is identified as the negative log-likelihood, with a corresponding client-level definition.
  • Theorem statement: Theorem 3.2 assumes Assumptions 1–8 and SGD with learning rate η = a0 √K, and states inequalities for FedEM after enough communication rounds.The theorem’s formal statement is connected to the displayed result for FedEM’s iterates.
  • Reduction to FedEM: FedEM is obtained by specializing the surrogate framework to V = ∆M, u = Θ, vt = πt, and client weights ωt = nt/n.With the corresponding definition of gk, federated surrogate optimization reduces to FedEM.
  • Surrogate properties: The proof verifies that the specialized surrogate gk is smooth with respect to Θ and is a majorant of the local and global objectives.The majorization argument uses the non-negativity of KL divergence.
  • Smoothness verification: The proof establishes differentiability and smoothness of the mixture responsibilities and the resulting surrogate components.It uses differentiability of the loss and the softmax-based responsibilities to analyze gradients and Hessians.

G.2.2 Proof of Theorem 3.3′

The proof of Theorem 3.3′ analyzes fully decentralized federated surrogate optimization through recursions for average-iterate optimality and consensus distance. Smoothness, stochastic-gradient assumptions, and network mixing properties combine to yield convergence bounds.

  • Average iterate: The proof derives a recursion for the average-iterate term measuring optimality of the decentralized average iterate.Lemma G.7 provides the corresponding recursion under the stated assumptions.
  • Consensus distance: It separately derives recursions for the distance to consensus among local iterates.Lemmas G.8 and G.9 handle two parts of the consensus-distance analysis.
  • Stochastic control: The stochastic terms are controlled using independent batches, unbiased gradients, and bounded variance under Assumption 6′.These assumptions bound the final stochastic-gradient terms in the recursions.
  • Combining bounds: The analysis uses L-smoothness, Jensen’s inequality, Cauchy–Schwarz, and doubly stochastic mixing matrices to combine the recursions.The proof also imposes sufficiently small learning-rate conditions involving L and β.
  • Theorem statement: Theorem 3.3′ assumes Assumptions 4′–8′ and 9, uses SGD with η = a0 √K, and gives bounds after sufficiently many communication rounds.The result concerns the iterates of fully decentralized federated surrogate optimization.
  • Reduction to D-FedEM: The decentralized theorem specializes to D-FedEM by setting V = ∆M, u = Θ, vt = πt, and ωt = nt/n.Under this specialization, Algorithm 5 reduces to D-FedEM.

H.2 Proof of Theorem H.1′

The proof of Theorem H.1′ establishes convergence of federated surrogate optimization to a stationary point under additional assumptions. The result is then specialized to FedEM through the same surrogate construction.

  • Stationarity result: Theorem H.1′ states that federated surrogate optimization converges to a stationary point of f under Assumptions 4′–8′, 10′, and 11′.The theorem requires G2 = 0 and a bound on α involving β2 and κ4.
  • Proof strategy: The proof combines previous lemmas with bounds on the surrogate-gradient relation and applies the resulting inequalities to obtain the final convergence claim.The final step combines Eq. (422) with the preceding result.
  • FedEM specialization: FedEM’s corresponding theorem is obtained as a particular case by setting V = ∆M, u = Θ, vt = πt, and ωt = nt/n.The specialized federated surrogate algorithm reduces to FedEM.
  • FedEM condition: The theorem applies a condition with G2 = 0 and α bounded by an expression involving β2 and κ5 for the FedEM specialization.The specialized statement gives the corresponding condition for convergence to a stationary point.
  • Assumption verification: The proof verifies the specialized assumptions through the earlier smoothness, majorization, and partial first-order surrogate lemmas.The verification transfers Assumptions 4′, 6′, 8′, 10′, and 11′ from their non-primed counterparts.

I.1.5 Synthetic dataset

The synthetic-data experiments specify a mixture-based client-data generator and compare federated methods using average client test accuracy. FedEM is evaluated against MOCHA on this setting.

  • Clients receive mixture weights sampled from a symmetric Dirichlet distribution with parameter α.
  • Client sample counts are generated from a log-normal distribution and clipped to the range implied by nt = min(50 + mt, 1000).
  • Each client sample draws a latent component assignment z(i)t from its mixture distribution M(πt).
  • FedEM is compared with MOCHA on the synthetic dataset using a linear model.MOCHA’s penalty parameter λ is tuned on a holdout validation set.
  • FedEM achieves 74.7 accuracy versus 73.4 ± 0.05 for MOCHA on the synthetic dataset.

J.3 Generalization to Unseen Clients

FedEM supports personalization for clients absent from training by learning mixture weights from local data. Experiments report strong unseen-client performance, while broader evaluation also examines sampling, convergence, and computational constraints.

  • Generalization to unseen clients: FedEM lets new clients learn personalized models at least as good as FedAvg’s global model and better than FedAvg+’s model.
  • Generalization to unseen clients: 47.5% versus 44.1% test accuracy is observed for new versus old clients on CIFAR100.
  • Generalization to unseen clients: New clients have average mixture-weight entropy 0.92 versus 0.27 for old clients, suggesting less skewed component usage.
  • Generalization to unseen clients: FedEM reaches significant personalization after unseen clients collect about 20% of the original clients’ training samples.
  • Generalization to unseen clients: FedEM obtains 42.7% ± 0.33% on unseen CIFAR100 clients, compared with 40.5% ± 1.66% for pFedMe and 38.9% ± 0.97% for FedAvg.
  • Effect of M in time-constrained setting: FedEM updates and transmits M components each round, requiring roughly M times more computation and message size than the compared methods.
  • Effect of M in time-constrained setting: With one third as many rounds, FedEM still outperforms other methods on three datasets.
  • Scope and limitations: Finding effective model architectures remains challenging because of the large search space, bias-variance trade-off, and FL-specific characteristics.
Loading 2108.10252v4…