Title: Personalized Federated Learning via Feature Distribution Adaptation
Venue: NeurIPS 2024
Accepted: yes
Keywords: Federated Learning

Data Heterogeneity

Personalization
Abstract: Federated learning (FL) is a distributed learning framework that leverages commonalities between distributed client datasets to train a global model. Under heterogeneous clients, however, FL can fail to produce stable training results. Personalized federated learning (PFL) seeks to address this by learning individual models tailored to each client. One approach is to decompose model training into shared representation learning and personalized classifier training. Nonetheless, previous works struggle to navigate the bias-variance trade-off in classifier learning, relying solely on limited local datasets or introducing costly techniques to improve generalization.
In this work, we frame representation learning as a generative modeling task, where representations are trained with a classifier based on the global feature distribution. We then propose an algorithm, pFedFDA, that efficiently generates personalized models by adapting global generative classifiers to their local feature distributions. Through extensive computer vision benchmarks, we demonstrate that our method can adjust to complex distribution shifts with significant improvements over current state-of-the-art in data-scarce settings.

Personalized Federated Learning via Feature Distribution Adaptation

Abstract

Federated learning (FL) is a distributed learning framework that leverages commonalities between distributed client datasets to train a global model. Under heterogeneous clients, however, FL can fail to produce stable training results. Personalized federated learning (PFL) seeks to address this by learning individual models tailored to each client. One approach is to decompose model training into shared representation learning and personalized classifier training. Nonetheless, previous works struggle to navigate the bias-variance trade-off in classifier learning, relying solely on limited local datasets or introducing costly techniques to improve generalization. In this work, we frame representation learning as a generative modeling task, where representations are trained with a classifier based on the global feature distribution. We then propose an algorithm, pFedFDA, that efficiently generates personalized models by adapting global generative classifiers to their local feature distributions. Through extensive computer vision benchmarks, we demonstrate that our method can adjust to complex distribution shifts with significant improvements over current state-of-the-art in data-scarce settings. Our source code is available on GitHub 1 .

1 Introduction

The success of deep learning models relies heavily on access to large, diverse, and comprehensive training data. However, communication constraints, user privacy concerns, and government regulations on centralized data collection often pose significant challenges to this requirement [31, 37, 18]. To address these issues, Federated Learning (FL) [34] has gained considerable attention as a distributed learning framework, especially for its privacy-preserving properties and efficiency in training deep networks.

The FedAvg algorithm, introduced in the seminal work [34], remains one of the most widely adopted algorithms in FL applications [32, 45, 38, 49, 40, 7]. It utilizes a parameter server to maintain a global model, trained through iterative rounds of distributed client local updates and server aggregation of client models. While effective under independent and identically distributed (i.i.d.) client data, its performance deteriorates as client datasets become more heterogeneous (non-i.i.d.). Data heterogeneity leads to the well-documented phenomenon of client drift [19], where distinct local objectives cause the model to diverge from the global optimum, resulting in slow convergence [20, 28] and suboptimal local client performance [42]. Despite extensive efforts [27, 19, 46, 6] to enhance FedAvg for non-i.i.d. clients, the use of a single global model remains too restrictive for many FL applications.

Personalized federated learning (PFL) has emerged as an alternative framework that produces separate models tailored to each client. The success of personalization techniques depends on

balancing the bias introduced by using global knowledge that may not generalize to individual clients, and the variance inherent in learning from limited local datasets. Popular PFL techniques include regularized local objectives [41, 26], local-global parameter interpolation [6], meta-learning [9, 16], and representation learning [35, 5, 48, 29]. While these techniques have shown significant improvements for clients under limited types of synthetic data heterogeneity (e.g., imbalanced partitioning of an otherwise i.i.d. dataset), we find that current methods still struggle to navigate the bias-variance trade-off with the additional challenge of feature distribution shift and data scarcity, conditions commonly encountered in cross-device FL.

As such, we look to design a method capable of handling real-world distribution shifts, e.g., covariate shift caused by weather conditions or poor camera calibration, (see clients 1 and 2 in Fig. 1) with limited local datasets. To this end, we approach PFL through shared representation learning guided by a global, low-variance generative classifier. Specifically, we select a probability density p with desirable statistical properties (e.g., one that admits an efficient Bayesian classifier) and iteratively estimate the global parameters of this distribution and representation layers to produce features from the estimated distribution (Fig. 1a).

To further navigate the bias-variance trade-off, we introduce a local-global interpolation method to adapt the global estimate to the distribution of each client. At inference time, clients use their adaptive local distribution estimate in a personalized Bayesian classifier (Fig. 1b).

Contributions. We propose a novel P ersonalized Fed erated Learning method based on F eature D istribution A daptation ( pFedFDA ). We contextualize our algorithm using a class-conditional multivariate Gaussian model of the feature space in a variety of computer vision benchmarks. Our empirical evaluation demonstrates that our proposed method consistently improves average model accuracy in benchmarks with covariate shift or client data scarcity, obtaining over 6% in multiple settings. At the same time, our method remains competitive with current state-of-the-art (often within 1%) on more general benchmarks with more moderate data heterogeneity. To summarize, our contributions are three-fold:

Anovel generative modeling perspective for federated representation learning is proposed to enable a new bias-variance trade-off for client classifier learning.

We propose a personalized federated learning method, pFedFDA, which leverages awareness of latent data distributions to guide representation learning and client personalization.

Extensive experiments on image classification datasets with varying levels of natural data heterogeneity and data availability demonstrate the advantages of pFedFDA in challenging settings.

2 Related Work

Federated Learning with Non-i.i.d. Data. Various studies have worked to understand and improve the ability of FL to serve heterogeneous clients. In non-i.i.d. scenarios, the traditional FedAvg method [34] is susceptible to client drift [19], resulting in slow convergence and poor local client accuracy [28, 27]. To tackle this challenge, [27, 1, 21] proposed the use of regularized local objectives to reduce the bias on the global model after local training. Another approach focuses on rectifying the bias of local updates [19, 10] through techniques such as control variates. Other strategies include loss-balancing [15, 47, 3], knowledge distillation [30, 54], prototype learning [43], and contrastive learning [25]. Despite promising results on non-i.i.d. data, their reliance on a single global model poses limitations for highly heterogeneous clients [17].

Personalized Federated Learning. In response to the limitations of a single global model, PFL seeks to overcome heterogeneity by learning models tailored to each client. In this framework, methods attempt to strike a balance between being flexible enough to fit the local distribution and relying on global knowledge to prevent over-fitting on small local datasets. Popular strategies include meta-learning an initialization for client adaptation [16, 9], multi-task learning with local model regularization [41, 26], local and global model interpolation [6], personalized model aggregation [51, 50], client clustering [39, 8], and decoupled representation and classifier learning [5, 29, 35, 48, 3]. Our work focuses on this latter approach, in which the neural network is typically decomposed into the first L -1 layers used for feature extraction, and the final classification layer.

Existing works in this category share feature extraction parameters between clients and rely on client classifiers for personalization. These approaches differ primarily in the acquisition of client classifiers

Class 2

Features

Class 1

Features

Class 2

Features

Class 2

Features

Backbone

Client 3

Client 4

Client 2

Client 1

Backbone

Backbone

Backbone

Class 1

Features

Class 2

Features

Class 1

Features

Class 1

Features

Global Feature

Distributions

Global

Classifier

(a) Global Representation Learning

Class 1

Features

Class 2

Features

Class 1

Features

Class 2

Features

Class 2

Features

Backbone

Client 3

Client 4

Client 2

Client 1

Backbone

Backbone

Backbone

Class 1

Features

Class 2

Features

Class 1

Features

Local Feature

Distributions

Personalized

Classifiers

(b) Local Distribution Adaptation

during training, which influences representation learning. For example, FedRep [5] sequentially trains a strong local classifier while holding the representation fixed, then updates the representation under the fixed classifier. FedBABU [35] proposes to use fixed dummy classifiers to align client objectives, only fine-tuning the classifier layers after the representation parameters have converged. Similarly, FedRoD [3] aims to train a generic representation model and classifier in tandem via balanced softmax loss, later obtaining personalized classifiers through fine-tuning or hypernetworks. FedPAC [48] adopts the learning algorithm of FedRep, but additionally regularizes the feature space to be similar across clients, before learning a personalized combination of classifiers across clients to improve generalization. However, this collaboration comes with an additional computational overhead that scales with the number of active clients. pFedGP [2] leverages a shared feature extractor as a kernel for client Gaussian processes. Although this approach offers improved sample efficiency, it comes at the cost of increased computational complexity and reliance on an inducing points set.

In a similar spirit to our method, FedEM [33] estimates the latent data distribution of clients in parallel to training classification models. FedEM estimates each client data distribution as a mixture of latent distributions, where personalized models are a weighted average of mixture-specific models. Notably, this introduces a significant overhead in both communication and computation as separate models are trained for each mixture. In contrast, our work estimates the distribution of client features in parallel to training a global representation model.

3 Problem Formulation

FL System and Objective. We consider an FL system where a parameter server coordinates with M clients to train personalized models θ i , i = 1 , 2 , · · · , M . Each client i has a local training dataset D i = { ( x j i , y j i ) } n i j =1 , where x ∈ R m and y ∈ { 1 , · · · , C } . The model training objective in PFL is:

where Q is feasible set of model parameters, F i ( θ i ) = E ( x,y ) ∼D i [ L ( θ i ( x ) , y )] is the empirical risk of dataset D i , and L is a loss function of the prediction errors (e.g., cross-entropy). The client population M in FL can be large, resulting in partial client participation [17]. Let q denote the participation rate, meaning that in each round, a client participates in model training with probability q .

Following [5, 48], we approach this as a problem of global representation learning and local classification, in which each θ i consists of a shared backbone, ϕ , responsible for extracting low-level features ( z ∈ R d = ϕ ( x ) ), and a local classifier, h i , for learning a client-specific mapping between features and labels. Considering this decomposition of parameters θ i = ( h i ◦ ϕ ) , we can rewrite the original PFL objective as the following:

where Φ and H are the feasible sets of neural network and classifier parameters, respectively.

In our generative modeling framework, we consider H to be the probability simplex over { 1 , · · · , C } , and our algorithm uses approximations of the posterior distributions as classifiers h i . However, for fair comparison with existing work (as well as other nice properties, discussed in Section 4.1), we select a generative model of the feature space such that h i can be represented with an equivalent linear layer.

̸

Data Heterogeneity. The data distribution of each client i is a joint distribution on X × Y , which can be written as p i ( x, y ) , p i ( y ) p i ( x | y ) , or p i ( x ) p i ( y | x ) . Using the terminology of [17], we refer to each case of data heterogeneity as follows: prior probability shift ( p i ( y ) = p i ′ ( y ) ), concept drift ( p i ( x | y ) = p i ′ ( x | y ) ), covariate shift ( p i ( x ) = p i ′ ( x ) ), and concept shift ( p i ( y | x ) = p i ′ ( y | x ) ). Furthermore, the local dataset volumes D i may have quantity skew , i.e., n i = n i ′ .

̸

4 Methodology

In this section, we introduce pFedFDA, a personalized federated learning method that utilizes a generative modeling approach to guide global representation learning and adapt to local client distributions. We present our method using a classconditional Gaussian model of the feature space, with additional discussion of the selected probability density in Section 4.1.

Algorithm 1 describes the workflow of pFedFDA.

Our algorithm begins with a careful initialization of parameters for the feature extractor ϕ , Gaussian means µ = { µ c } C c =1 , and covariance Σ (Lines 1-2). We initialize ϕ with established techniques (e.g., [12]) such that the output features follow a Gaussian distribution with controlled variance. We similarly use a spherical Gaussian to ensure a stable initialization of the corresponding generative classifier (see Section 4.1).

̸

Algorithm 1: pFedFDA

1 Server initializes feature extractor ϕ 0 g with random Gaussian weights.

2 Server and clients initialize feature distribution estimates ( µ 0 g , Σ 0 g ), ( µ 0 i , Σ 0 i ) as random spherical Gaussians.

3

for

each round

r

= 0

,

· · ·

, R

-

1

do

4

Send

ϕ

r

g

,

µ

r

g

,

Σ

r

g

to participating clients

5

for

each active client

i

do

6

ϕ

r

i

ϕ

r

g

7

Train

ϕ

r

i

for

E

epochs using loss (5)

8

Estimate local

̂

µ

i

,

̂

Σ

i

via (6) and (7)

9

Estimate interpolation parameter

β

i

via (9)

10

µ

r

i

β

i

̂

µ

+(1

-

β

i

)

µ

g

11

Σ

r

i

β

i

̂

Σ+(1

-

β

i

)Σ

g

12

Send

ϕ

r

i

,

µ

r

i

,

Σ

r

i

to the server

13

end

14

Update

ϕ

r

g

,

µ

r

g

,

Σ

r

g

with a weighted average of

each client's

ϕ

r

i

,

µ

r

i

, and

Σ

r

i

15

end

At the start of each FL round r , the server broadcasts the current ϕ r g , µ r g , Σ r g to each participating client. The local training of each client consists of two key components: (1) global representation learning , in which clients train ϕ to maximize the likelihood of local features under the global feature distribution µ r g , Σ r g (Line 7); (2) local distribution adaptation , in which clients obtain robust estimates of their local feature distribution µ r i , Σ r i , using techniques for efficient low-sample Gaussian estimation (Line 8) and local-global parameter interpolation (Lines 9-11). After local training, clients send their ϕ r i , µ r i , Σ r i to the parameter server for aggregation (Line 14).

̸

̸

In the following sections, we provide detailed explanations of each algorithmic component. In Section 4.1 we discuss the benefits of a generative modeling framework and provide the justification for our selected class-conditional Gaussian model. We outline how the resulting generative classifier can be used to guide representation learning in Section 4.2 and describe how we obtain personalized generative classifiers in Section 4.3.

4.1 Generative Model of Feature Distributions

Motivation for Generative Classifiers. A central theme in FL is exploiting inter-client knowledge to train more generalizable models than any client could attain using only their local dataset. This presents an important bias-variance trade-off, as incorporating global knowledge naively can introduce significant bias. Fortunately, under a generative modeling approach, this bias can be naturally handled, enabling efficient inter-client collaboration.

First note that local class priors p i ( y ) can be approximated with local counts: p i ( y = c ) ≈ n c i ∑ c ′ ∈ C n c ′ i := π c i , where n c i is the number of local samples whose labels are c . This leaves the primary source of bias to the mismatch between local and global feature distributions p g ( z | y ) and p i ( z | y ) . Crucially, it turns out that this bias is controllable due to the dependence of z on global representation parameters ϕ . Consequentially, we propose to minimize this bias through our classification objective, which we discuss further in Section 4.2.

Class-Conditional Gaussian Model. In this work we approximate the distribution of latent representations using a class-conditional Gaussian with tied covariance, i.e., p i ( z | y = c ) = N ( z | µ c i , Σ i ) . We show the resulting generative classifier under this model in Eq. 4. Note that it has a closed form and results in a decision boundary that is linear in z . I.e., if we know the underlying local feature distribution mean and covariance, we can efficiently compute the optimal header parameters h i for the inner objective in Eq. 2.

In addition to the convenient form of the Bayes classifier, we select this distribution as the Gaussianity of latent representations is likely to hold in practice. Notably, by adopting the common technique of Gaussian weight initialization (e.g., [12]), the resulting feature space is highly Gaussian at the start of training. It has also been observed that the standard supervised training of neural networks with cross-entropy objectives results in a feature space that is well approximated by a class-conditional Gaussian distribution [24], i.e., the corresponding generative classifier Eq. 4 has equal accuracy to the learned discriminative classifier. We provide a further discussion of this modeling assumption in Appendix A.

4.2 Global Representation Learning

Next, we describe our process for training the shared feature extractor ϕ . Similar to existing works [5, 48], our local training consists of training ϕ via gradient descent to minimize the cross-entropy loss of predictions from fixed client classifiers. We obtain our client classifiers through Eq. 4, using global estimates of µ g , Σ g and local estimated priors π i . For computational efficiency, we avoid inverting the covariance matrix by estimating Σ -1 µ c with the least-squares solution w = min w ′ ∥ Σ w ′ -µ c ∥ .

The loss of client i for an individual training sample ( x, y ) is provided in Eq. 5.

Note that for a spherical Gaussian Σ = I and uniform prior π , we recover a nearest-mean classifier under Euclidean distance. This resembles the established global prototype regularization [43], which minimizes the Euclidean distance of features from their corresponding global class prototypes. Notably, FedPAC [48] uses this prototype loss to align client features. However, this implicitly assumes that all feature dimensions have equal variance, and additionally requires a hyperparameter

λ to balance the amount of regularization with the primary objective. In contrast, our generative classifier naturally aligns the distribution of client features by training ϕ with our global generative classifier.

4.3 Local Distribution Adaptation

Local Estimation. Akey component of pFedFDA is the estimation of local feature distribution parameters, used both for model personalization and for updating the global distribution for representation learning.

Given a set of n extracted features Z with n c examples per class c , a maximum likelihood estimate of the class means and an unbiased estimator of the covariance, respectively, are given by:

where, with slight abuse of notation, ¯ Z ∈ R n × d denotes the matrix of centered features with rows corresponding to each original feature z j centered by their respective means, i.e.,

Estimators Eq. 6 and Eq. 7 may be noisy on clients with limited local data. To illustrate this, consider the common practical scenario where n i ≪ d . The feature covariance matrix Σ i at client i will be degenerate; in fact, it will have a multitude of zero eigenvalues. In these cases, we can add a small diagonal ϵ I to Σ , and replace the non-positive-definite matrices with the nearest positive definite matrix with identical variance. This can be efficiently computed by clipping eigenvalues in the corresponding correlation matrix and followed by converting it back to a covariance matrix with normalization to maintain the initial variance. We refer readers to [11] for a review of low-sample covariance estimation.

Local-Global Interpolation. We introduce this fusion because even with the aforementioned correction to ill-defined covariances, the variance of the local estimates remains highly noisy, indicating the necessity of leveraging global knowledge. It is essential to consider that in the presence of data heterogeneity, clients with differing local data distributions and dataset sizes have varying requirements for global knowledge.

For our Gaussian parameters µ , Σ , we consider the introduction of global knowledge through a personalized interpolation between local and global estimates, which can be viewed as a form of prior. We provide an analysis of the high-probability bound on estimation error for an interpolated mean estimate in simple settings in Theorem 1. The full derivation is deferred to Appendix E.

Theorem 1 (Bias-Variance Trade-Off) . Let C = 1 . Define µ i as the sample mean of client i 's local features µ i := 1 n i ∑ n i j =1 z j i , and µ g as the global sample mean using all N samples across M clients: µ g := 1 N ∑ M i =1 ∑ n i j =1 z j i . Assume client features are independent and distributed as z i ∼ N ( θ i , Σ i ) , with true global feature distribution N ( θ g , Σ g ) . We consider the use of global knowledge at client i through an interpolated estimate: ̂ µ i := βµ i +(1 -β ) µ g , where β ∈ [0 , 1] . For any δ ∈ (0 , 1) , with probability at least 1 -δ , it holds that

where c > 0 is an absolute constant.

Intuitively, the estimation error and optimal β depend on the bias introduced by using global knowledge ∥ θ g -θ i ∥ 2 2 , the variance of local and global features, and the respective data volumes.

We formulate this as an optimization problem, in which clients estimate interpolation coefficients β i to combine local and global estimates of ( µ , Σ ) with minimal k -fold validation loss:

where D k is the dataset consisting of the validation samples for the k -th fold, and ( ̂ µ k , ̂ Σ k ) are the local distribution estimates Eq. 6 and Eq. 7 estimated using the training samples from the k -th fold. In our experiments, we avoid additional forward passes on the local dataset by preemptively storing the feature-label pairs obtained over the latest round of training.

We solve Eq. 9 using off-the-shelf quasi-newton methods (e.g., L-BFGS-B). We additionally explore using separate β terms for the means and covariance (Section 5.3) and recommend the use of a single β term for most applications.

After obtaining β , we set our local estimates of µ i , Σ i to their interpolated versions. These estimates are then sent to the server for aggregation. Notably, the server update rule can be viewed as a moving average [52] between the previous round estimate and the client average scaled by β , reducing the influence of local noise in the global distribution estimate. At test time, clients use their local distribution estimates for inference through the classification rule in Eq. 4.

5 Experiments

5.1 Experimental Setup

Datasets, Tasks, and Models: We consider image classification tasks and evaluate our method on four popular datasets. The EMNIST [4] dataset is for 62-class handwriting image classification. The CIFAR10/CIFAR100 [22] datasets are for 10 and 10-class color image classification. The TinyImageNet [23] dataset is for 200-class natural image classification. For EMNIST and CIFAR10/100 datasets, we adopt the 4-layer and 5-layer CNNs used in [48]. On the larger TinyImageNet dataset, we use the ResNet18 [13] architecture. Notably, the feature dimension d for EMNIST/CIFAR CNNs is 128, and 512 for ResNet. We provide additional details in Appendix C.1.

Clients and Dataset Partitions: The EMNIST dataset has inherent covariate shifts due to the individual styles of each writer. We partition the dataset by writer following [6], and train with M = 1000 total clients (writers), participating with rate q = 0 . 03 . On CIFAR and TinyImageNet datasets, we simulate prior probability shift and quantity skew by partitioning the dataset according to a Dirichlet distribution with parameters α ∈ (0 . 1 , 0 . 5) , where lower α indicates higher levels of heterogeneity. On these datasets, we use M = 100 clients with participation rate q = 0 . 3 . Additional details of the partitioning strategy are provided in Appendix C.1.2.

We split each client's data partition 80-20% between training and testing.

Covariate Shift and Data Scarcity: We introduce two modifications to client partitions to simulate the challenges of real-world cross-device FL. We first consider common sources of input noise for natural images, which may result from the qualities of the measuring devices (e.g., camera calibration, lens blur) or environmental factors (e.g., weather, lighting). To simulate this, we select ten image corruptions at five levels of severity defined in [14], and corrupt the training and testing samples of the first 50 clients in CIFAR10/100 with unique corruption-severity pairs. We leave the remaining 50 client datasets unchanged. We refer to these datasets with natural covariate shifts as CIFAR10-S/CIFAR100-S and detail the specific corruptions in Appendix C.1.1.

Second, we perform uniform subsampling of client training sets, leaving them with (75%, 50%, or 25%) of their original samples. These low-sample settings are more realistic for cross-device FL, where clients rely more on knowledge sharing.

Baselines and Metrics: We compare pFedFDA to the following baselines: Local, in which each client trains its model in isolation; FedAvg [34] and FedAvg with fine-tuning (FedAvgFT); APFL [6]; Ditto [26]; pFedMe [41]; FedRoD [3]; FedBABU [35]; FedPAC [48]; FedRep [5]; and LG-FedAvg [29]. We report the average and standard deviation of client test accuracies.

Model Training: We train all algorithms with mini-batch SGD for E = 5 local epochs and R = 200 global rounds. We apply no data augmentation besides normalization into the range [ -1 , 1] . For pFedFDA, we use k = 2 cross-validation folds to estimate a single β i term for each client. Additional training details and hyperparameters for each baseline method are provided in Appendix C.2.

5.2 Numerical Results

Performance under covariate shift and data scarcity. We first present our evaluation under natural client covariate shift with varying data scarcity in Table 1. In all experiments, pFedFDA outperforms

the other methods in test accuracy, demonstrating the effectiveness of our method in adapting to heterogeneous client distributions. Additionally, pFedFDA has an increasing benefit relative to other methods in data-scarce settings: on CIFAR10, we improve 4.2% over the second-best method with 100% of training samples and 6.9% with 25%. On CIFAR100, the same improvements range from 0.1% to 6.5%. This indicates the success of our method in navigating the bias-variance trade-off.

Evaluation in more moderate scenarios. Our evaluation of all four datasets in the traditional setting (no added covariate shift, full training data) is presented in Table 2. We note that: (1) our method is still competitive, always ranking within the top 3 methods, and (2) the gap between top methods is smaller than in the previous experimental setting. For example, on EMNIST/CIFAR10, we see that FedAvgFT, FedPAC, and pFedFDA are within ∼ 1% accuracy. We observe larger performance gaps for CIFAR100, with FedPAC and pFedFDA having the best results.

Results under extreme data scarcity. We present additional results at the limits of data scarcity on CIFAR10/100 datasets in Table 3, where we assign a single mini-batch (50) of training examples to each client. Notably, even as n i ≪ d , which poses a challenge to local covariance estimation, pFedFDA clients obtain the best test accuracy, indicating the robustness of our local-global adaptation.

Generalization to new clients. We further analyze the ability of our generative classifiers to generalize on clients unseen at training time. To simulate this setting, we first train the server model model using half of the client population. We then evaluate each method on the set of clients not encountered throughout training, using their original input data, as well as their dataset

transformed using each corruption from CIFAR-S. Further benchmark details, including fine-tuning (personalization) procedures, are provided in Appendix C.3. As demonstrated in Table 4, our method generalizes well even on clients with covariate shifts not encountered at training time. Moreover, observe that pFedFDA has the highest accuracy on the original clients, highlighting the efficacy of structured generative classifiers when less training data is available (i.e., having 50 rather than 100 clients).

5.3 Ablation of Method Components

We conduct two studies to verify the efficacy of our local-global interpolation method. In Table 5, we see that our interpolated estimates always perform better than using only local data, indicating the benefits of harnessing global knowledge. Learning separate β terms for the means and covariance may be beneficial in low-sample or covariate-shift settings when the local distribution estimate may fluctuate further from the global estimate. However, using a single scalar β appears sufficient and comes with the lowest computational cost (associated with the time to solve Eq. 9).

Weadditionally visualize the spread of learned β across clients as a function of their dataset corruption in Fig. 2. As expected, clients with clean datasets rely more on global knowledge (smaller β values) than corrupted clients. Moreover, corruptions with higher β values (e.g., contrast) often align with the more difficult corruptions encountered in Table 4.

5.4 Communication and Computation

The parameter count and relative communication load of our generative classifiers compared to a simple linear classifier varies depending on class count C and feature dimension d . In our experimental configurations (datasets, architectures), the overhead in total parameter count ranges from 1.1% to 6.8%. See Appendix D.3 for additional details.

In Table 6, we compare the local training time (client-side computation) and total runtime of pFedFDA to baseline methods on CIFAR10. We observe a slight increase in training time relative to FedAvg, which can be attributed primarily to cost of learning our parameter interpolation coefficient β . However, this increase is comparable to the existing methods and is lower than representation-

Motion

Blur

Defocus

Blur

Noise

Shot

Noise

Impulse

Noise

Frost

Fog

Brightness JPEG

Contrast

None

Client Corruption

0.0

0.2

0.4

0.6

0.8

Learned

learning methods FedRep and FedPAC. This demonstrates the relative efficiency of our generative classifier formulation in comparison to classifiers obtained through local fine-tuning.

6 Conclusion

Balancing local model flexibility and generalization remains a central challenge in personalized federated learning (PFL). This paper introduces pFedFDA, a novel approach that addresses the biasvariance trade-off in client personalization through representation learning with generative classifiers. Our extensive evaluation on computer vision tasks demonstrates that pFedFDA significantly outperforms current state-of-the-art methods in challenging settings characterized by covariate shift and data scarcity. Furthermore, our approach remains competitive in more general settings, showcasing its robustness and adaptability. The promising results underline the potential of our method to improve personalized model performance in real-world federated learning applications. Future work will focus on exploring the scalability of pFedFDA and its application to other domains.

Acknowledgments and Disclosure of Funding

We gratefully acknowledge the support from the National Science Foundation CAREER award under Grant No. 2340482, the Army Research Laboratory under Cooperative Agreement Number W911NF-23-2-0014, the Sony Faculty Innovation Award, and the National Defense & Engineering Graduate (NDSEG) Fellowship Program. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Laboratory, the National Science Foundation, or the U.S. government. The U.S. government is authorized to reproduce and distribute reprints for government purposes notwithstanding any copyright notation herein. We also thank Ming Xiang for valuable discussions and feedback on this work.

References

Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations , 2021.

Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems , 34:8392-8406, 2021.

Hong-You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In International Conference on Learning Representations , 2022.

Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. Emnist: an extension of mnist to handwritten letters, 2017.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference on Machine Learning , pages 2089-2099. PMLR, 2021.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 , 2020.

Zhaoyang Du, Celimuge Wu, Tsutomu Yoshinaga, Kok-Lim Alvin Yau, Yusheng Ji, and Jie Li. Federated learning for vehicular internet of things: Recent advances and open issues. IEEE Open Journal of the Computer Society , 1:45-61, 2020.

Moming Duan, Duo Liu, Xinyuan Ji, Yu Wu, Liang Liang, Xianzhang Chen, Yujuan Tan, and Ao Ren. Flexible clustered federated learning for client-level data distribution shift. IEEE Transactions on Parallel & Distributed Systems , 33(11):2661-2674, November 2022.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems , 33:3557-3568, 2020.

Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng-Zhong Xu. Feddc: Federated learning with non-iid data via local drift decoupling and correction. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10112-10121, 2022.

Yaqian Guo, Trevor Hastie, and Robert Tibshirani. Regularized linear discriminant analysis and its application in microarrays. Biostatistics (Oxford, England) , 8:86-100, 02 2007.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision , pages 1026-1034, 2015.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 770-778, 2016.

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Learning Representations , 2019.

Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real-world data distribution. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part X 16 , pages 76-92. Springer, 2020.

Yihan Jiang, Jakub Koneˇ cn` y, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.12488 , 2019.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Georgios Kaissis, Alexander Ziller, Jonathan Passerat-Palmbach, Théo Ryffel, Dmitrii Usynin, Andrew Trask, Ionésio Lima, Jason Mancuso, Friederike Jungmann, Marc-Matthias Steinborn, Andreas Saleh, Marcus Makowski, Daniel Rueckert, and Rickmer Braren. End-to-end privacy preserving deep learning on multi-institutional medical imaging. Nature Machine Intelligence , 3(6):473-484, Jun 2021.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning , pages 5132-5143. PMLR, 2020.

Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 4519-4529. PMLR, 2020.

Jinkyu Kim, Geeho Kim, and Bohyung Han. Multi-level branched regularization for federated learning. In International Conference on Machine Learning , pages 11058-11073. PMLR, 2022.

Alex Krizhevsky. Learning multiple layers of features from tiny images. University of Toronto , 05 2012.

Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin. A simple unified framework for detecting out-of-distribution samples and adversarial attacks. Advances in Neural Information Processing Systems , 31, 2018.

Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10713-10722, 2021.

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning , pages 6357-6368. PMLR, 2021.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems , 2:429-450, 2020.

Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations , 2020.

Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523 , 2020.

Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems , 33:2351-2363, 2020.

Ximeng Liu, Lehui Xie, Yaopeng Wang, Jian Zou, Jinbo Xiong, Zuobin Ying, and Athanasios V Vasilakos. Privacy and security issues in deep learning: A survey. IEEE Access , 9:4566-4593, 2020.

Nathalie Majcherczyk, Nishan Srishankar, and Carlo Pinciroli. Flow-fl: Data-driven federated learning for spatio-temporal predictions in multi-robot systems. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 8836-8842. IEEE, 2021.

Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. Advances in Neural Information Processing Systems , 34:15434-15447, 2021.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Jaehoon Oh, SangMook Kim, and Se-Young Yun. FedBABU: Toward enhanced representation for federated image classification. In International Conference on Learning Representations , 2022.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems , 32, 2019.

Kilian Pfeiffer, Martin Rapp, Ramin Khalili, and Jörg Henkel. Federated learning for computationally constrained heterogeneous devices: A survey. ACM Computing Surveys , 55(14s):1-27, July 2023.

Swaroop Ramaswamy, Rajiv Mathews, Kanishka Rao, and Françoise Beaufays. Federated learning for emoji prediction in a mobile keyboard. arXiv preprint arXiv:1906.04329 , 2019.

Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning: Modelagnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems , 32(8):3710-3722, 2020.

Micah J Sheller, G Anthony Reina, Brandon Edwards, Jason Martin, and Spyridon Bakas. Multi-institutional deep learning modeling without sharing patient data: A feasibility study on brain tumor segmentation. In Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 4th International Workshop, BrainLes 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 16, 2018, Revised Selected Papers, Part I 4 , pages 92-104. Springer, 2019.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 8432-8440, 2022.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Chunnan Wang, Xiang Chen, Junzhe Wang, and Hongzhi Wang. Atpfl: Automatic trajectory prediction model design under federated learning framework. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 6563-6572, June 2022.

Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in Neural Information Processing Systems , 33:7611-7623, 2020.

Lixu Wang, Shichao Xu, Xiao Wang, and Qi Zhu. Addressing class imbalance in federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 10165-10173, 2021.

Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference on Learning Representations , 2023.

Timothy Yang, Galen Andrew, Hubert Eichner, Haicheng Sun, Wei Li, Nicholas Kong, Daniel Ramage, and Françoise Beaufays. Applied federated learning: Improving google keyboard query suggestions. arXiv preprint arXiv:1812.02903 , 2018.

Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Fedala: Adaptive local aggregation for personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 11237-11244, 2023.

Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez. Personalized federated learning with first order model optimization. In International Conference on Learning Representations , 2021.

Sixin Zhang, Anna Choromanska, and Yann LeCun. Deep learning with elastic averaging sgd, 2015.

Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVI 16 , pages 561-578. Springer, 2020.

Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In International Conference on Machine Learning , pages 12878-12889. PMLR, 2021.

A Limitations

The selected class-conditional Gaussian distribution may not work well for all neural network architectures. For example, if the output features are the result of an activation such as ReLU, a truncated Gaussian distribution may be a better model. Future work can look to exploit knowledge of the neural network architecture to improve the accuracy of the feature distribution estimate.

In this work, we leverage the insights from a fusion of global and local feature space. As in many applications there is often an underlying cluster structure between clients datasets, future works may explore the identification and efficient estimation of feature distributions of client clusters, in order to reduce the degree of bias introduced in client collaboration.

B Broader Impacts

Federated learning has become the main trend for distributed learning in recent years and has deployed in many popular consumer devices such as Apple's Siri, Google's GBoard, and Amazon's Alexa. Our paper addresses the practical limitations of personalization methods in adapting to clients with covariate shifts and/or limited local data, which is a central issue in cross-device FL applications. We are unaware of any potential negative social impacts of our work.

C Details of Experimental Setup

All experiments are implemented in PyTorch 2.1 [36] and were each trained with a single NVIDIA A100 GPU. Compute time per experiment ranges from approximately 2 hours for CIFAR10/100 and 20 hours for TinyImageNet. Code for re-implementing our method is provided at the following GitHub URL: https://github.com/cj-mclaughlin/pFedFDA .

C.1 Dataset Description

The EMNIST [4] dataset contains over 730,000 28 × 28 grayscale images of 62 classes of handwritten characters. The CIFAR10/CIFAR100 [22] datasets contain 60,000 32 × 32 color images in 10 and 100 different classes of natural images, respectively. TinyImageNet [23] contains 120,000 64 × 64 color images of natural images.

C.1.1 CIFAR-S Generation.

We implement the following 10 common image corruptions at 5 levels of severity as described in [14]: Gaussian noise, shot (Poisson) noise, impulse noise, defocus blur, motion blur, fog, brightness, contrast, frost, JPEG compression. We apply a unique corruption-severity pair to all samples of the first 50 clients.

C.1.2 Non-i.i.d. Partitioning.

On CIFAR and TinyImageNet datasets, we distribute the proportion of samples of class C across M clients according to a Dirichlet distribution: q c , m ∼ Dir M ( α ) , where we consider α ∈ (0 . 1 , 0 . 5) as in [30].

We provide a visualization of Dirichlet partitioning strategies on CIFAR10 below. The size of each point represents the number of allocated samples. Notably, as α increases, Dir( α ) becomes less heterogeneous.

C.2 Training Settings

All methods are trained using mini-batch SGD for 200 global rounds with 5 local epochs of training. We use a fixed learning rate of 0.01, momentum of 0.5, and weight decay of 5e-4. The batch size is set to 50 for all experiments, except for EMNIST, where we use a batch size of 16. We sample the set of active clients uniformly with probability q=0.3 for CIFAR and TinyImageNet and q=0.03 for EMNIST. The last global round of training employs full client participation. We split the data of each client 80-20% between training and testing.

0

20

40

60

80

100

Client Index

0

1

2

3

4

5

6

7

8

9

Class Index

(a) CIFAR10 Dirichlet(0.1)

0

20

40

60

80

100

Client Index

0

1

2

3

4

5

6

7

8

9

Class Index

(b) CIFAR10 Dirichlet(0.5)

Hyper-parameters. For APFL, we tune α over [0.25, 0.5, 0.75, 1.0], and set α = 0 . 25 . For pFedMe, we tune λ over [1.0, 5.0, 10.0, 15.0] and set λ = 5 . 0 . For Ditto, we use five local epochs for personalization and tune µ over [0.05, 0.1, 0.5, 1.0, 2.0] and set µ = 1 . 0 . For FedRep and FedBABU, we use five local epochs for training the head parameters. For FedPAC, we tune λ over [0.1, 0.5, 1.0, 5.0, 10.0], and set λ = 1 . 0 . FedPAC uses one local epoch for training head parameters with a higher learning rate of 0.1, following the original implementation.

C.3 Evaluation on New Clients

Our fine-tuning procedure on new clients largely follows the methodology above. For FedAvgFT, we fine-tune the global model for five local epochs. For FedBABU and FedPAC, we personalize the model in 2 different ways and report the best result: (1) fine-tuning only the head for 5 local epochs, and (2) fine-tuning both the body and head for 5 local epochs. For pFedFDA, each new client estimates their local interpolated statistics (i.e., lines 8-11 of Algorithm 1) to obtain a personalized generative classifier.

For our covariate shift evaluation, we apply a medium severity corruption (level 3) to all samples.

D Additional Results

D.1 Multi-Domain FL

In Table 7, we present results on the DIGIT-5 domain generalization benchmark [53]. This presents an alternate form of covariate shift, as the data from each client is drawn from one of 5 datasets (SVHN, USPS, SynthDigits, MNIST-M, and MNIST). In particular, we use 20 clients trained with full participation, and assign 4 clients to each domain. Within each domain, we use the Dirichlet(0.5) partitioning strategy to assign data to each client. We observe that pFedFDA is effective in all settings, but has the most significant benefits over prior work in the low-data regime.

D.2 Effect of Local Epochs

In many FL settings, we would like clients to perform more local training between rounds to reduce communication costs. However, too much local training can cause the model to diverge. In Fig. 4, we compare the effect of the local amount of epochs for CIFAR100 and CIFAR100-S-25% sample datasets. We observe that (1) pFedFDA outperforms FedAvgFT at all equivalent budgets of E , (2)

both methods follow exhibit a general plateau in accuracy after E = 5 , and (3) pFedFDA learns much faster than FedAvgFT, with significantly higher accuracy for E = 1 .

40

35

30

25

20

15

10

5

0

1

5

FedAvgFT

pFedFDA

CIFAR100 Dir(0.5)

CIFAR100-S Dir(0.5) 25%

10

20

Local Epochs

Test Accuracy (%)

D.3 Communication Load Examples

In Table 8, we compare the number of distinct parameters in our Gaussian estimates to that of a typical linear classifier for the models and datasets used in this paper, along with some additional examples. We display the resulting overhead relative to the base parameter count of the shared representation backbone.

D.4 Runtime of Method Components

In Table 9, we evaluate the proportion of each local iteration of pFedFDA associated with each line of our algorithm. Network Passes refers to the time taken to train the base network parameters ϕ (Line 7 of Alg. 1). Mean/Covariance Est. refers to the time taken to estimate the local mean and covariance from features extracted during model training (Line 8 of Alg. 1). Interpolation Optimization refers to the time taken to optimize the local coefficient β (Line 9 of Alg. 1). Overall, we find that the majority of the overhead of our method comes from estimating the interpolation parameter β .

E On the Bias-Variance Tradeoff

This section justifies the bias-variance tradeoff under some simplified technical assumptions. For simplicity, we assume that at any given round, the extracted feature vectors for a class are independent. We illustrate the bias-variance tradeoff in estimating the mean feature of a given class c at round t .

Proof of Theorem 1. For ease of exposition, we drop the time index and class index.

Let i be an arbitrary client with local dataset size n i of class c . Let N be the total data volume of class c over the entire FL system. Assuming that the distribution of client i 's features z follow a multivariate Gaussian distribution N ( θ i , Σ i ) , and the global feature distribution follows N ( θ, Σ) where θ g := ∑ M i =1 n i θ i / ( ∑ i ∈ [ M ] n i ) , Σ g := ∑ M i =1 n 2 i Σ i / ( ∑ i ∈ [ M ] n i ) 2 . Note θ i , θ g are deterministic parameters.

We denote the local and global mean estimates as:

Let ̂ µ i be the local estimate that interpolates between local and global knowledge, defined as

We will focus on bounding the high probability local estimation error ∥ ̂ µ i -E [ µ i ] ∥ 2 .

Note that Eq.(10) can be further expanded as

̸

where γ := β +(1 -β ) n i N , and ¯ µ := 1 N ∑ i ′ = i ∑ n i ′ j =1 µ j i ′ .

Thus µ i ∼ N ( θ i , 1 n i Σ i ) and ¯ µ ∼ N ( Nθ g -n i θ i N , N Σ g -n i Σ i N 2 ) . Since µ i and ¯ µ are independent, we have

Let g ∼ N ( 0 , I ) . We have

̸

̸

where ̂ Σ 1 / 2 is the square root matrix of ̂ Σ := γ 2 1 n i Σ i +(1 -β ) 2 N Σ g -n i Σ i N 2 . It holds that

Taking the expectation with respect to the randomness in the Gaussian random variable g and by the law of total expectation, we have

and

where equality ( a ) holds because ( ̂ Σ 1 / 2 ) ⊤ ( ̂ Σ 1 / 2 ) = ̂ Σ as ̂ Σ 1 / 2 is symmetric.

By Hanson-Wright inequality [44, Theorem 6.2], we conclude that with probability at least 1 -δ (for any given δ ∈ (0 , 1) ),

where c > 0 is some absolute constant, inequality ( b ) holds as a · b ≤ ( a + b 2 ) 2 , and N ≥ n i , inequality ( c ) holds because of triangular inequality ∥ A + B ∥ F ≤ ∥ A ∥ F + ∥ B ∥ F , that ∥ A ∥ F = √ ∥ A ∥ 2 F = √ Tr( A ⊤ A ) = √ Tr( A 2 ) if matrix A is symmetric, and that max { a, b } ≤ a + b , inequality ( d ) holds because Tr( A 2 ) ≤ (Tr( A )) 2 for positive semidefinite matrix A and that Tr(Σ i ) , β 2 ≥ 0 , and Tr(Σ g ) are by definition non-negative.

The first term (1 -β ) 2 ∥ θ g -θ i ∥ 2 2 is the bias introduced when client i uses global knowledge; the smaller the β , the more bias introduced. The last term reveals the interaction of β and the tradeoff between local and global variance. When β approaches 0, we have the global feature variance Tr(Σ) reduced by the average of N global samples. When β approaches 1, we have local feature variance Tr(Σ i ) reduced by the average of only n i local data. Thus the bias-variance tradeoff on client i crucially depends on the degree of local-global distribution shift, ∥ θ g -θ i ∥ 2 2 , the local data volume n i and its quality (i.e., Σ i ), and the volume and quality of the data across clients N, Σ g .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The scope of the paper is on an important topic of client model personalization in federated learning. We faithfully state our contributions in both the abstract and introduction.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations and considerations of our Gaussian modelling of the feature space both in the main text and with additional notes in the appendix. While the focus of the paper is in improving client personalization in the challenging setting of data scarcity and client distribution shift, we additionally benchmark our method in more general settings to demonstrate the widespread applicability of our work. Finally, we provide an assessment of the communication and computation overhead of our method compared to state-of-art approaches.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the key assumptions in the main text. The missing proof is deferred to Appendix E.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide detailed experimental setups and review the hyperparameters in Section 5.4 and Appendix C.2. We additionally provide code and instructions to train our method.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our evaluations are based on open-accessed datasets that are publically available. An official implementation code is provided in the supplementary materials.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide detailed experimental setups and review the hyperparameters in Section 5.4 and Appendix C.2. We additionally provide code and instructions to train our method.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: For our main experiments we include the standard deviation of client accuracies, and include std error bars in our ablation visualization of the method components.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Please find the software/hardware specifications in Appendix C.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The NeurIPS code of ethics is strictly enforced throughout our research.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have discussed the broader impacts of our work in Appendix B. Please find details therein.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The existing assets used in this paper has been adequately cited or credited to.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: Our attached code is well documented and comes with a README file indicating how reviewers may set up our experiments and train the proposed method.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFederated Learning (FL) combines data from multiple clients to train a global model but struggles with heterogeneous data. Personalized Federated Learning (PFL) creates individual models for each client, addressing this issue. Traditional methods face challenges with bias-variance trade-offs and rely on limited local data or costly techniques. This paper proposes pFedFDA, an algorithm that frames global representation learning as a generative modeling task. It adapts global generative classifiers to local feature distributions, improving performance in data-scarce settings.\n\nSTRENGTHS:\n1. To implement personalized federated learing from generative view looks interesting and promising.\n2. This paper is well-structured entirely, but the notations could be improved.\n3. The mathematical proof looks good and sufficient.\n\nWEAKNESSES:\n1. Your method looks similar to the traditional personalzed federated learning. just as the algorithm 1, the shared backbone is just the weighted summation of the corresponding part in all involved clients. \n2. You mentioned that your shared backbone are trained in a generative way, but I can not get the core of of it. Actually, there are some other works[1] use the generative models (e.g., autoregression) as backbone to implement personalized federated learning, what is the difference between your methods and such baselines? \n3. Your global distribution parameters are weighted summation of all clients, so in my opinion, the global distribution may be bias due to data size skewness. There are many ways to overcome this, such as [2], etc. do you think such methods can be adopted by you to fix this issue of data skewness?\n\n[1] Kou, W.B., Lin, Q., Tang, M., Xu, S., Ye, R., Leng, Y., Wang, S., Chen, Z., Zhu, G. and Wu, Y.C., 2024. pFedLVM: A Large Vision Model (LVM)-Driven and Latent Feature-Based Personalized Federated Learning Framework in Autonomous Driving. arXiv preprint arXiv:2405.04146.\n[2] Kou, W.B., Lin, Q., Tang, M., Wang, S., Zhu, G. and Wu, Y.C., 2024. FedRC: A Rapid-Converged Hierarchical Federated Learning Framework in Street Scene Semantic Understanding. arXiv preprint arXiv:2407.01103.\n\nQUESTIONS:\n1. \"where Φ and H are the feasible sets of neural network and classifier parameters, respectively.\", in this context, it looks better to replace \"neural network\" to \"shared backbone\", do you think so?\n2. The notaions sometimes looks confused, which could be improved.\n3. In table 5, it should be pFedFDA not pFedDFA, right?\n4. The used dataset looks simple, could you use more complex dataset, such as cityscpaes, camvid, kitti, imagenet, etc to veriry your contributions?\n\nLIMITATIONS:\nthe proposed method is over-reliance on the local distribution, so it is difficult to handle following cases: 1. a new client is added into the system and its local distribution is far away the original global distribution and the original distributions of all clients; 2. the local distribution of all clients are depondent on time (i.e., dynamic)", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces pFedFDA, a novel approach to personalized federated learning that conceptualizes global representation learning as a generative modeling task. Specifically, the method involves shared representation learning, guided by a generative classifier characterized by a low-variance global probability density. Additionally, it iteratively refines global parameters and employs a local-global interpolation technique to tailor these estimates to individual client distributions. The authors validate the effectiveness of the proposed solution through experiments on benchmark datasets, demonstrating its robustness and applicability.\n\nSTRENGTHS:\nThe paper presents an interesting and meaningful idea, complemented by a theoretical analysis of the bound on high probability estimation errors for the interpolated mean estimate, offering valuable insights and inspiration to the community. \n\nThe writing is clear and straightforward, with well-designed figures that effectively support and clarify the presented concepts.\n\nWEAKNESSES:\nThere are several concerns that need addressing:\n1.In some cases, the performance improvements appear minimal. For instance, Table 2 shows that pFedFDA is outperformed significantly by FedBABU and pFedME.\n\n2.In Table 3, the authors employ a Dirichlet distribution (Dir(0.5)) for experiments under extreme data scarcity. It would be beneficial to lower the Dirichlet value, perhaps to 0.1, to conduct a more extensive evaluation under greater data scarcity.\n\n3.Table 5 indicates that the runtime of pFedFDA is longer than several other methods, including FedAvg, Ditto, and FedRoD, which calls into question the efficiency of the proposed solution.\n\n4.\tIn Equations 3 and 5, should there be a superscript ‘c’ on Σ to clarify the notation further?\n\nQUESTIONS:\nplease check \"Weaknesses\". I like the idea of this work and I am looking forward to the responses from authors.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a personalized Federated Learning (FL) method that adapts global generative classifiers to local feature distributions. The authors show that their method can handle complex distribution shifts for computer vision tasks.\n\nSTRENGTHS:\n- The paper proposes a personalized FL method that uses a generative classifier by considering bias-variance trade-off.\n- The authors conduct extensive experiments.\n- The paper is well-structured and easy to understand.\n\nWEAKNESSES:\n- Sharing the statistics of the generative classifier could increase privacy risks.\n- The algorithm’s additional computational cost (in Algo.1 L8-9) could increase with larger local data. Calculating the local covariance matrix for larger data and the covariance inversion matrix for each beta value could be computationally intensive.\n\nQUESTIONS:\n- In Fig. 1, the colors representing client 1 and client 2 are not easily distinguishable.\n- In Eq. (4), wouldn’t the cost of the covariance matrix inversion be high? Could the server calculate it once and send it to the client to reduce the client’s computational cost, even if it increases communication cost?\n- In Eq. (9), should “\\in min” be replaced with “= \\argmax”?\n- If the client’s generative classifier’s mean and variance are sent to the server, wouldn’t the privacy risk increase compared to sending only the weight of the conventional linear classifier?\n- If the generative classifier is not used and only the weight parameter from the conventional linear classifier is interpolated locally, could the performance be expected to be mid-range between FedAvgFT and pFedFDA?\n- In experiments like Table 2, where there is less covariate shift between clients within the federation, FedPAC performs better. Could the degree of covariate shift within the federation be determined by looking at the difference in the generative classifier’s statistics coming from the client, and then choose the appropriate FL method?\n\nLIMITATIONS:\nPlease refer to the weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper uses Class-Conditional Gaussian Model to formulate the latent representation of the global generative part; for the other part,  a personalized federated learning algorithm pFedFDA is designed via Federated Distribution Adaptation. This paper then proves a bound on the bias-variance trade-off of pFedFDA under the assumption of independent and Gaussian distributed local dataset features.  And the performance of pFedFDA against dataset scarcity and heterogeneity is presented.\n\nSTRENGTHS:\n1.\tThe introduction and related work are well-organized. The explanation on the replacement of generative global model of the embedding module is convincing. \n2.\tThis paper gives a solid proof on the analysis of bias-variance trade-off.\n\nWEAKNESSES:\n1.\tThe experiment is concentrated on the CIFAR-10 and CIFAR-100. However, it is also the dataset where pFedFDA has best performance on (Tab. 2). Also, the robustness of pFedFDA against strong heterogeneity is not good. The weakness of pFedFDA against heterogeneity is still a limitation in this paper. \n2.\tThe assumption of Theorem 1 is very strong, however, this paper does not give an evaluation on the datasets in the experiment of longtailed-ness or normality. Also, the basic result of Theorem 1 is about the upper-bound of local bias, there is no matric in the experiment to compare the performance of bias-variance balance among the algorithms, especially in the scarcity experiment (Tab. 1, 3). \n3.\tThis paper does not explain the connection between bias-variance balance and robustness against data scarcity.\n\nQUESTIONS:\n1.\tCan the Class-Conditional Gaussian Kernel represent the optimal Gaussian distribution?\n2.\tFurther explain the experiment Tab. 1. How can this relate to the advantage of pFedFDA on bias-variance balancing? Also, other models (i.e. APFL, Ditto) have a stable performance enduring the dataset scarcity. How can this be explained?\n3.\tIn Tab. 2, the line of FedRep, please double check the value of the last cell .145(0.4). Is that correct?\n4.\tDo you perform normality test on the training set of agents? If not, can you give a further explanation of the feature distribution of datasets to match the assumption of Theorem 1?\n\nLIMITATIONS:\n1.\tThe assumption of Theorem 1 is strong. It would be better to consider the skewed distribution, and even long-tailed distributions. Based on this analysis, the result in Tab. 3 could be expanded to more general few-shot or one-shot learning cases, with a modified version of pFedFDA. \n2.\tThe design of pFedFDA is still not robust to heterogeneity and is not capable of better personalization.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces pFedFDA, a personalized Federated Learning (pFL) method designed to address the issue of client heterogeneity in federated learning. pFedFDA combines global knowledge through server aggregation with local knowledge through client-specific training and distribution estimation, enhancing the model's personalized performance on each client. The authors propose an algorithm that efficiently generates personalized models, demonstrating significant improvements in data-scarce settings through extensive computer vision benchmarks. The paper is well-structured, and the writing is clear, making it easy to follow the authors' arguments and experimental results.\n\nSTRENGTHS:\n1.The method of decomposing model training into shared representation learning and personalized classifier training, followed by adaptation to local feature distributions, is both innovative and promising for handling non-i.i.d. data in FL environments.\n2.The paper provides strong empirical evidence through comprehensive experiments on various datasets, showcasing the superiority of pFedFDA in challenging distribution shift and data scarcity scenarios.\n3.The paper is well-written, with a clear structure that logically progresses from the introduction of the problem to the presentation of the methodology and results.\n\nWEAKNESSES:\n1. In this paper, the mean value and covariance are utilized as the characteristic distribution of the data. Have you explored other statistical measures to describe this distribution?\n2. In the ablation study, the results obtained by calculating multiple β values were unexpectedly lower compared to using a single β value. Intuitively, multiple β values should provide a more comprehensive understanding of the data distribution and thus yield better results. However, the experimental findings do not support this expectation, and there is no clear explanation for this discrepancy.\n3.Note that pFedFDA uses the same feature extractor for all clients but employs heterogeneous classifiers. A key feature of pFL is the heterogeneity of client models. pFedFDA is somewhat limited in this regard. Are there methods to handle heterogeneous feature extractors in pFL scenarios?\n\nQUESTIONS:\nThe main concerns and questions are listed in the weaknesses. Please provide answers to address these concerns.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-25"}