FedAPA: Server-side Gradient-Based Adaptive Personalized Aggregation for Federated Learning on Heterogeneous Data

1 Introduction

As a prominent distributed machine learning paradigm, Federated Learning (FL) allows clients to collaboratively train models while keeping data local [Yang et al., 2019]. In crosssilo FL, client data is often non-IID (non-Independent and Identically Distributed), exhibiting statistical heterogeneity that leads to client drift and reduces the effectiveness of methods like FedAvg [McMahan et al., 2017].

To address the challenges posed by heterogeneous data distributions, Personalized Federated Learning (PFL) has been developed to tailor models to the unique data characteristics of individual clients[Wu et al., 2020]. PFL can be broadly categorized into three main approaches: (1) Personalized model-based approaches, which directly create a personalized local model for each client, typically without relying on a global model [Li et al., 2021b; Huang et al., 2021; Luo and Wu, 2022; Shamsian et al., 2021]; (2) Global modelbased approaches, where a global model is first constructed and then personalized through techniques like client-specific fine-tuning [Xu et al., 2023; Zhang et al., 2024]; and (3) Other PFL approaches include those that leverage information beyond model parameters, such as prototypes, to optimize local models [Tan et al., 2022]. Among these, personalized model-based approaches, the primary focus of this paper, have gained significant attention for their ability to directly tailor local models to each client's unique data distribution.

∗Corresponding Author

We categorize existing personalized model-based PFL approaches into three groups based on their aggregation strategies for creating each personalized model: (1) Static metricbased methods: Approaches like FedAMP[Huang et al., 2021] and FedPHP[Li et al., 2021b] define aggregation weights using predefined metrics derived from parameter similarity or training progress. However, these predefined metrics do not directly reflect client-specific optimization objectives (e.g., local losses or gradients). (2) Local gradientbased methods: Approaches such as APPLE[Luo and Wu,

2022] and FedALA[Zhang et al., 2023] calculate aggregation weights locally via gradient descent. While these methods share model parameters between clients or use elementwise model combination, they incur high computational complexity and potential communication overhead. Additionally, their focus on individual client objectives limits effective cross-client coordination. (3) Auxiliary network-based methods: Approaches like pFedHN[Shamsian et al., 2021] utilize Hypernetworks[Ha et al., 2016] to derive personalized models or aggregation weights. However, the need to train auxiliary models introduces significant computational overhead. In summary, these limitations highlight the need for new model-aggregation techniques that not only enable adaptive personalization in non-IID scenarios but also reduce computational and communication costs.

To address these limitations, we propose a novel personaliezed model based PFL approach, FedAPA (Adaptive Personalized Aggregation for Federated Learning), as illustrated in Fig. 1. FedAPA introduces a server-side aggregation strategy that adaptively learns aggregation weights using gradients of client parameter changes with respect to the aggregation weights, generating aggregated personalized models for each client. Specifically, the server: (1) maintains the client parameters and an aggregation-weight vector, where each element quantifies the relevance of other clients' model knowledge during aggregation; (2) constructs personalized client parameters through weighted aggregation of all collaborative clients' parameters using the learned weights; and (3) updates the weight vector based on the client-parameter delta-the change in the client's local model parameters after multiple local update steps. As an additional feature, FedAPA adopts a partial model sharing strategy, transferring only the client's feature extractor to reduce irrelevant information during aggregation.

The key innovation of FedAPA lies in its aggregation strategy, which updates aggregation weights on the server side using the gradient of the client-parameter delta with respect to the aggregation weights, optimizing the personalized aggregation process. Half of the squared norm of the clientparameter delta serves as a proxy for the local loss, effectively capturing each client's optimization state. By dynamically adjusting aggregation weights to reduce this proxy loss, FedAPA generates an aggregated personalized model for each client, better aligning the model with the client's local objectives. Compared to existing model-aggregation approaches, FedAPA's strategy offers the following advantages: (1) Adaptability: FedAPA dynamically learns aggregation weights based on gradients derived from the proxy loss, enabling the aggregation weights to adapt effectively to each client's unique optimization trajectory. (2) Centralized Efficiency: By centralizing weight updates on the server, FedAPA minimizes communication and computational overhead while enhancing cross-client coordination. (3) Simplicity: FedAPA avoids the need for auxiliary networks, eliminating the computational complexity associated with training additional models. In summary, FedAPA's aggregation strategy effectively balances personalization and collaboration in non-IID scenarios, enhancing computational and communication efficiency.

We summarize our contributions as follows:

• We propose FedAPA, a novel model-aggregation PFL framework for heterogeneous data scenarios, featuring a server-side gradient-based model-aggregation strategy that enhances FL performance while maintaining low computational and communication overhead. while maintaining low computational and communication overhead.
• To the best of our knowledge, FedAPA is the first to adaptively update aggregation weights on the server side to generate aggregated personalized models, using gradients derived from client-parameter changes.
• We provide a theoretical convergence guarantee for the FedAPA algorithm.
• We experimentally demonstrate that our approach outperforms 10 PFL competitors across three benchmarks under two non-IID settings, achieving top accuracy, optimal computational efficiency, and low communication overhead. Ablation studies highlight the critical contribution of the aggregation strategy and the complementary benefits of the partial model sharing strategy.

2 RelatedWork

2.1 PersonalizedModel-based PFL Approaches

Personalized model-based PFL approaches, also known as PFL approaches via learning personalized local models, directly create personalized models for individual clients, typically without the need for a global model. These approaches employ various aggregation strategies to construct each personalized model, which can be broadly categorized as follows: (1) Static metric-based methods: These methods define aggregation weights using predefined, static metrics. For example, FedAMP [Huang et al., 2021] calculates similarity as weights based on the Euclidean distance of model parameters, while FedPHP [Li et al., 2021b] measures relevance using training progress. Although straightforward, these metrics do not directly align with client-specific optimization objectives, limiting their adaptability to non-IID data distributions. (2) Local gradient-based methods: They learn aggregation weights via gradient descent on the client side. For example, APPLE [Luo and Wu, 2022] learns interclient relations by sharing model parameters between clients, leading to high communication costs and privacy concerns. FedALA [Zhang et al., 2023] computes element-wise combination weights between client and global models, resulting in high computational complexity. While these methods utilize non-local information, their focus on local objectives limits cross-client coordination and adaptability to data heterogeneity. (3) Hypernetwork-based methods: Hypernetworks [Ha et al., 2016] have been utilized as the auxilliary networks to enhance personalization. For example, pFedHN [Shamsian et al., 2021] employs Hypernetworks to generate personalized models However, these methods incur significant computational overhead due to the complexity of training Hypernetworks

2.2 Global Model-based PFL Approaches

Global model-based PFL approaches, also referred to as PFL approaches via global model personalization, involve training a global model first and then personalizing it to derive local models for individual clients. For example, FedPAC [Xu et al., 2023] leverages a probabilistic framework for personalization using Bayesian neural networks and a global model, allowing personalized models to adaptively incorporate uncertainty during client-specific training. DBE [Zhang et al., 2024] enhances the generalization of the global feature extractor through bi-directional knowledge transfer, reducing representation bias and fine-tuning client models for improved personalization.

2.3 Other PFL Approaches

PFL can also leverage non-model information, such as prototypes, to achieve personalization. For example, FedProto [Tan et al., 2022] leverages local prototypes to construct global prototypes, which are used to improve feature extraction and classification. As a combined method, FedGH[Yi et al., 2023] uses client-uploaded prototypes (i.e., classaveraged representations) to train a shared global head on the server, which is then downloaded to replace local heads.

3 Method

3.1 ProblemFormulation

Personalized Federated Learning (PFL) aims to collaboratively train personalized models for a set of clients, each with its own private data. PFL addresses the challenge of data heterogeneity across clients, ensuring that each client's model is tailored to its specific data while maintaining data privacy and improving individual model performance.

Let M be the total number of clients participating in the federated learning process. Each client i ∈ {1, · · · ,M} has a local dataset Di; with its own data distribution Pi on X × Y . In PFL, each client i seeks to learn a personalized model parameterized by ωi. The objective function for PFL can be formulated as:

where |Di| is the size of the local dataset Di; N is the total number of data points across all clients; F (ωi; x) is the model parameterized by ωi evaluated on input x; y is the true label for input x; and Li is the local loss function for client i, such as the following cross-entropy loss function:

where C is the number of classes; yc is the binary indicator (0 or 1) if class label c is the correct classification for input x; and F (ωi; x)c is the predicted probability of input x belonging to class c under the model parameterized by ωi.

3.2 FedAPAAlgorithm

The framework of FedAPA is demonstrated in Fig1. Each FedAPA client uploads its feature extractor, receives a personalized feature extractor after server-side aggregation, and then updates the entire client model locally, including both the feature extractor and the classifier. On the server side, the model parameters of each client's feature extractor and the corresponding aggregation-weight vector are maintained. The core of FedAPA lies in the server-side aggregation, where the server constructs a personalized feature extractor for each client by adaptively aggregating all clients' model parameters using weights updated via gradient descent.

FedAPA adjusts each aggregation weight vector on the server using the gradient calculated from the client parameter delta, which represents the change in local model parameters after multiple local updates within a communication round. At the end of each round, the updated client parameters can be considered as an approximation of the local task's optimal solution, with the client-parameter delta approximately reflecting the local optimization direction (as detailed in subsubsection 3.2). For each client, the server iteratively updates the client's aggregation-weight vector based on the gradient of the client-parameter delta with respect to the vector. The learned weights represent adaptively acquired knowledge contributions from other clients, which may potentially reflect the data distribution correlation among clients. By aggregating beneficial model knowledge from collaborative clients, FedAPA addresses non-IID data challenges for individual clients effectively.

Local Training

As shown in Fig1, for each client i, the local model F(ωi) is partitioned into two parts: the feature extractor f(θi), which takes the input sample x and produces the representation z = f(θi; x), and the decision layer h(ϕi), which processes the representation z and outputs the predicted label ˆ y = h(ϕi; z). The shared parameter θi is uploaded to the server, while the private parameter ϕi is kept locally. Thus, the local objective function in PFL can be expressed as

Upon downloading the parameters ¯ θi as the current feature extractor, client i obtains its current entire model ωi := [θi; ϕi]. Each client conducts several local update steps on the local data using the following update rule:

where αi denotes the learning rate of the local model parameters ωi. After the local updates, client i separates θi from the updated ωi, and uploads the updated θi as the shared client parameters to the server.

Server-Side Aggregation of Client-Parameters

For each client i, FedAPA server manages a trainable aggregation-weight vector Ai = (ai,1, . . . , ai,j, . . . , ai,M)T where i, j ∈ {1, · · · ,M}, and M indicates the total number of clients in the federated learning system. Here, ai,j denotes the weight of client j during the weighted aggregation used to form the target parameters ¯ θi for client i. During the initialization of the learning process, we set the self-weight ai,i to 1, while the other weights ai,j to 0 for i ̸= j. For the t-th training round, the server computes the personalized parameters ¯ θi for client i by performing a weighted aggregation of all clients' shared parameters according to the following formula:

Server-Side Updating of Weight-Vectors

Upon receiving the client-parameters θi from each client i, the FedAPA server calculates the difference between the current θi and the previous parameter versions ¯ θi stored on the server, obtaining the client-parameter delta as follows:

Subsequently, the server leverages this client-parameter delta to update the aggregation-weight vector Ai for client i according to the following rule:

where η is the learning rate for updating the aggregationweight vector on the server.

Equation (7) is pivotal to the FedAPA algorithm. To provide intuition behind this formula, consider the optimal solution for the local task's model parameters, θ∗ i = arg min θi Li. After multiple local updates within a communication round, θi can be regarded as an approximation of θ∗ i . Consequently, 1 2 ∥θ∗ i - ¯ θi∥22 approximates 12 ∥θi - ¯ θi∥22 , which is equivalent to half of the squared norm of ∆θi. This measure, 12 ∥θ∗ i - ¯ θi∥22 = 12 ∥¯ θi - θ∗ i ∥22 , can be interpreted as a surrogate for the local loss. Thus, half of the squared norm of ∆θi serves as a proxy for the local loss Li, where ∆θi can be interpreted as -∇¯ θi Li. According to the chain rule, (∇Ai ¯ θi)T ∆θi represents -∇Ai Li. Consequently, Equation (7) updates the aggregation weights Ai on the server in a way that minimizes the local loss Li for each client, ensuring that the personalized aggregation aligns with the client-specific optimization trajectory.

Before sending the updated aggregation-weight vector Ai to client i , the server further optimizes the weights using following three post-processing steps: (1) Clipping: To improve training stability and convergence, we control the update range of each weight element in the vector Ai by clipping each element value ai,j to be within [0, 1], where j ∈ {1, · · · ,M}. (2) Self-weight adjusting: Each client should prioritize its own local knowledge over that of other clients. Therefore, for the clipped aggregation-weight vector of each client i, we set the weight of the element ai,i (also known as the self-weight µ) to 0.5, ensuring it is not less than the weights of the other elements in the vector. (3) Normalization: To balance client contributions and minimize variance during aggregation, thus improving training generalization, we normalize each aggregation-weight vector from the previous step to ensure that the total sum of the weights equals 1.

Algorithm 1 FedAPA

• Input: dataset {D1, . . . ,DM}, learning rates η and αi, total communication rounds T, local rounds E, clientparameter set {θ1, . . . , θM}, aggregation-weight vector set {A1, . . . , AM}. Output: Trained personalized model set {¯ θ1, ¯ θ2, . . . , ¯ θM }. 1: ∀i ∈ {1, · · · ,M}, initialize θi on server 2: ∀i ∈ {1, · · · ,M}, initialize Ai = (ai,1, . . . , ai,M)T on server 3: Function Server Executes 4: for each communication round t = 1, . . . , T do 5: Sample a subset of clients St 6: for each client i ∈ St in parallel do 7: ¯ θi ← ∑Mj=1 ai,jθj {Weighted aggregation} 8: θi ← ClientUpdate(¯ θi) 9: ∆θi ← θi - ¯ θi {Computing ∆θi} 10: Update Ai via Eq 7 {Updating based on ∆θi} 11: ai,j ← min(max(ai,j, 0), 1) , j = 1, . . . ,M {Clipping} 12: ai,i ← 0.5 13: Normalize Ai 14: end for 15: end for 16: End Function 17: Function Clientupdate(¯ θi) 18: Download ¯ θi 19: θi ← ¯ θi . 20: ωi ← [θi; ϕi] 21: for each local epoch e = 1, . . . ,E do 22: for mini-batch B ⊆ Di do 23: ωi ← ωi - αi∇ωi Li(B) {Local training} 24: end for 25: end for 26: [θi; ϕi] ←ωi 27: return θi {Uploading θi} 28: End Function

In summary, during a communication round for each client i, the FedAPA server first constructs the personalized client parameters ¯ θi by performing a weighted aggregation of the collaborative clients' parameters. Next, the client downloads ¯ θi and uploads the optimized client parameters θi after local training. Subsequently, the server computes the clientparameters change ∆θi, and uses it to update the aggregationweight vector Ai, adaptively learning beneficial knowledge from collaborative clients. Algorithm 1 outlines the overall learning process in FedAPA.

4 Convergence Analysis

In this section, we analyze the convergence of FedAPA under suitable conditions. Additional notations used in this section are detailed in Appendix.

Assumption 1. (Lipschitz Smooth). Each local objective function is L1-Lipschitz smooth, which also means the gradient of local objective function L is L1-Lipschitz continuous. This assumption is reasonable as it encapsulates the differentiability and smoothness properties inherent to many common loss functions in deep learning, aligning with the mathematical prerequisites for the convergence of gradientbased optimization methods fundamental to neural network training,

which implies the following quadratic bound,

Assumption 2. (Unbiased Gradient and Bounded Variance). The stochastic gradient gr(t) i = ∇L(ω(t) i , ξi) is an unbiased estimator of the local gradient for each client, where the random variable ξi follows the distribution Di. This foundational assumption is pivotal, as it ensures that the expectation of the stochastic gradient, when averaged over the local dataset, corresponds to the true gradient of the local loss function. Despite the inherent variability in data distribution across different clients, this property remains intact, underscoring its relevance in Federated Learning where data heterogeneity is often the rule rather than the exception,

and its variance is bounded by the constant σ2 :

Assumption 3. (Lipschitz Continuity). Each local loss function L is L2-Lipschitz continuous with respect to θ. This assumption is deemed reasonable, akin to assumption 1, as it reflects a common property of loss functions, ensuring the boundedness of the gradient's impact on parameter updates, which is a prerequisite for the stability and convergence of optimization algorithms,

Theorem1. (One-roundDeviationBound). Let Assumptions 1 to 3 hold. For an arbitrary client, after every communication round, we have,

This theorem indicates the deviation bound of the local objective function for an arbitrary client after each communication round. Convergence of each client's local objective function can be guaranteed when there is a certain expected one-round decrease, which can be achieved by choosing appropriate η and α. The proof of Theorem 1 is in Appendix.

5 Experiments

5.1 Experimental Setup

Datasets. We employ three public image-classification datasets of varying complexities: FMNIST (simple), CIFAR10 (medium), and CIFAR-100 (complex). Experiments are conducted with 20 and 50 clients, respectively. For each dataset, we construct heterogeneous scenarios following the two widely-used settings: the pathological and the practical settings[McMahan et al., 2017; Li et al., 2021a]. For the pathological setting, we allocate disjoint and unbalanced data of 2/2/10 classes to each client out of a total of 10/10/100 classes on FMNIST/CIFAR-10/CIFAR-100 datasets. For the practical setting, we use a Dirichlet distribution[Yurochkin et al., 2019], denoted as Dir(α), to sample data points from each class in a given dataset to each client, where a smaller value of the concentration parameter α reflects stronger data heterogeneity. We set α to 0.1/0.1/0.01 for FMNIST/CIFAR10/CIFAR-100 datasets.

Competitors. We evaluated FedAPA against FedAvg, the standard non-PFL baseline, and 10 state-of-the-art PFL algorithms. The comparison includes six personalized model-based methods (FedPHP, FedAMP, APPLE, FedALA, pFedHN), two global model-based approaches (FedPAC and DBE), and two other PFL approaches (FedProto, and FedGH). The dashed lines in Table 1 and Table 2 sequentially divide the competitor methods into these three groups.

Model and Training. In all experiments, each client uses the same LeNet-5 architecture, with the first four layers serving as shared feature extraction layers with the server, and the final fully connected (FC) layer acting as a private decision layer. To accommodate the size and color channels of the images in each dataset, we set the input channels/FC layer input dimension/output channels of each client model to 1/256/10 for FMNIST, 3/400/10 for CIFAR-10, and 3/400/100 for CIFAR-100. We utilize the default training-to-test set ratio of each dataset, namely 6:1 for FMNIST and 5:1 for CIFAR-10 and CIFAR-100. Unless otherwise specified, we use the following training settings. we use SGD with a momentum of 0.9 as the client optimizer. We run two local training epochs in each iteration with a batch size of 64 and a learning rate α = 0.01 for 50 communication rounds with the server. Each round involves participation from at least 60% of all users. For our FedAPA algorithm, we initialize the self-weight µ in aggregation-weight vectors to 0.5 and set the learning rate η for updating weight vectors to 0.01.

Implementation. The experiments were conducted on a workstation running Ubuntu 22.04 LTS. The implementation utilized Python 3.8.10, CUDA version 11.3, and the PyTorch 1.11.0 deep learning framework. The workstation features a 12 vCPU Intel Xeon Platinum 8255C @ 2.50GHz CPU and an NVIDIA GeForce RTX 2080 Ti GPU.

5.2 Accuracy Evaluation

To evaluate and compare the performance of FedAPA with the competitors, we ran each FL algorithm and computed the average accuracy from three random data allocations. Accuracy was calculated as the ratio of correct to total predictions. Table 1 shows the test accuracy of all methods across 12 scenarios, covering three datasets, two client numbers, and two non-IID settings. It also includes the number of classes c per client for the pathological settings, and the concentration parameter α of the Dirichlet distribution for the practical settings. As shown in Table 1, FedAPA outperforms all competitors in test accuracy across 10 out of 12 scenarios. The only exceptions are the pathological non-IID settings on the FMNIST and CIFAR-100 datasets with 20 clients, where FedAPA trails the top-performing methods by just 0.06% (vs. FedPAC) and 0.51% (vs. FedALA), respectively. FedAPA demonstrates a more pronounced performance advantage in practical non-IID settings, which better reflect realistic heterogeneous data distributions. Overall, FedAPA consistently exhibits superior performance compared to all competing methods.

5.3 Computation and Communication Overhead

The computation times for each FL algorithm are summarized in the computation column of Table 2, showing total time (in minutes) and time per iteration (in seconds) for 20 clients using the medium CIFAR-10 dataset in the practical non-IID setting with 2 training epochs. As shown, FedAPA requires 17.66 (approximately 18) seconds per iteration, comparable to the baseline FedAvg. Thus, FedAPA incurs only an additional 0.34 seconds per iteration while achieving significant accuracy improvements. It offers state-of-the-art performance with minimal computation overhead among all compared PFL algorithms.

The communication costs per iteration are listed in the communication column of Table 2. Prototype-based methods, such as FedProto and FedGH, incur the lowest communication overhead by transmitting prototypes instead of model parameters. FedAPA, while achieving higher accuracy, maintains communication efficiency by transmitting only partial model parameters, resulting in the least communication overhead among all non-prototype FL algorithms. In summary, FedAPA achieves superior accuracy with the lowest communication cost among non-prototype methods.

5.4 Ablation Study

(1) FedAPA incorporates the Adaptive Personalized Aggregation (APA) strategy and the Partial Model Sharing (PMS) strategy. To assess the contributions of these two strategies, we conduct an ablation study, as shown in the top half of Table 3. Starting with FedAvg as the baseline, which uses average aggregation and entire-model sharing, we achieve an accuracy of 56.68%. Introducing the APA strategy significantly improves accuracy by 28.25%, reaching 84.93%. Adding the PMS strategy further enhances accuracy by 0.92%, resulting in a final accuracy of 85.85%. These results underscore the pivotal contribution of the APA strategy as the primary driver of performance improvement, while the PMS strategy offers a smaller, supplementary boost. (2) Following the server-side gradient update of the weight vectors, FedAPA applies three post-processing steps to the weights: clipping, self-weight adjustment, and normalization. Ablation studies in the lower half of Table 3 show that removing each step individually reduces accuracy by 2.58%, 1.44%, and 5.51%, respectively, with a significant 26.14% drop when all are omitted. These results highlight the importance of each step in improving model accuracy.

5.5 Convergence Curve Analysis

The convergence curves in Fig2 show the performance of the 12 methods over training rounds. FedAPA reaches convergence within approximately 10 rounds, demonstrating its rapid convergence rate compared to the competing methods. The curve for FedAPA exhibits a smooth progress towards convergence, indicating stable training and consistent performance improvement. Upon convergence, FedAPA achieves an accuracy of around 86%, which is higher than all other competitors. These results confirm that FedAPA is one of the fastest-converging algorithms, achieving the highest accuracy with stable convergence behavior.

5.6 Effect of Self-weight µ

FedAPA adjusts the self-weight of each aggregation-weight vector using the hyperparameter µ. We evaluate the effect of µ on accuracy by varying it from 0 to 1, using 20 clients on CIFAR-10 in a Practical Non-IID setting. As shown in Table 4, the best accuracy is achieved when µ = 0.5. Lower values underweight local knowledge, while higher values overly emphasize local tasks, reducing the benefits of shared knowledge. Thus, µ = 0.5 balances local and shared contributions effectively.

6 Conclusion

In this paper, we introduce a new PFL framework, FedAPA, designed to automatically consolidate desired knowledge from collaborative clients to form personalized local models. FedAPA adaptively learns aggregation weights on the server by utilizing gradients derived from changes in client parameters. We provide a theoretical proof for the convergence of our FedAPA algorithm. Experimental results under heterogeneous data scenarios demonstrate that FedAPA achieves higher accuracy and lower computational overhead compared to the competing PFL methods, while maintaining communication efficiency. These contributions highlight the effectiveness and efficiency of FedAPA in improving PFL performance, suggesting its potential suitability for handling large-scale heterogeneous data in practical applications.

Acknowledgements

This research was partially supported by funding from Guangdong Basic and Applied Basic Research Foundation (No. 2021A1515012297), the Hong Kong RGC General Research Fund (No. 152211/23E, 15216424/24E, and 152115/25E), the PolyU Internal Fund (No. P0056171), and the Huawei Gifted Fund.