FedGPS: Statistical Rectification Against Data Heterogeneity in Federated Learning

We are grateful to the reviewers for their diligent efforts and for acknowledging both the significance of our research questions and the interpretability of our methodology.

Q1: Lacks novelty and is incremental.

A1: Thank the reviewers for the questions they raised, which gave us the opportunity to further clarify the contributions of FedGPS. Extensive experiments have shown that sharing statistical information has great potential to alleviate data heterogeneity, which also inspires FedGPS. Furthermore, FedGPS takes into account statistical information from both the distribution and gradient perspectives to enhance a global view. Specifically, the shift at the distribution level is corrected through two-stage alignment based on the surrogate dataset. At the gradient level, the gradient information of other clients is considered before local training from a more global perspective. FedGPS has covered many heterogeneous scenarios across various datasets, further demonstrating its robustness.

Q2: Related work is incomplete and lacks some key SOTA works.

A2: We thank the reviewers for pointing out the shortcomings of related work, and we will add a discussion of the latest paper in the next version to improve our related work[R1-R4]. Additionally, we add FedMR[R1] as our new baselines; the results are listed in Tables 1-6. We will also add the results of this heuristic search method in the next version.

Table 1.The results in FedGPS under Heterogeneous scenario 1 with CIFAR-10, the best accuracy is reported.

Table 2. The results in FedGPS under Heterogeneous scenario 1 with CIFAR-10, average, and std accuracy are reported.

Table 3. The results in FedGPS under Heterogeneous scenario 1 with SVHN, the best accuracy is reported.

Table 4. The results in FedGPS under Heterogeneous scenario 1 with SVHN, average, and std accuracy are reported.

Table 5.The results in FedGPS under Heterogeneous scenario 1 with CIFAR-100, the best accuracy is reported.

Table 6. The results in FedGPS under Heterogeneous scenario 1 with CIFAR-100, average, and std accuracy are reported.

Q3: The performance improvement seems limited. I suggest that the author use the mean and standard deviation to show the experimental results.

A3: Thanks to the reviewers for their suggestions, we show part of the mean and standard deviation results in Table 7. Based on the mean and std metrics, it can be shown that FedGPS can achieve higher performance gain on mean compared to the highest accuracy. More results will be updated in the next version. Table 7. The mean and standard deviation results under CIFAR-10 of the original results in our paper .

Q4: The authors seem to have only evaluated extreme heterogeneous scenarios, not normal heterogeneous scenarios, such as $\alpha$=0.3, 0.6, and IID, which are commonly used settings.

A4: We sincerely appreciate the reviewers' constructive feedback, which has significantly advanced FedGPS's applicability in federated scenarios, allowing it to cover settings from homogeneous to extremely heterogeneous. To demonstrate this, we present the experimental results for FedGPS on the CIFAR-10 dataset, with K=100 and E=1, $\lambda_s=10$ in Table 8.

Table 8. The results of different federated scenarios under CIFAR-10.

Q5: Lacking experiments on more clients (e.g., K=500, 1000) and different participation ratios (the author only evaluates 10% of devices participating in local training).

A5: Thanks to the reviewers for motivating a more comprehensive evaluation of FedGPS, we conduct more clients (K=500) in Table 9 to demonstrate that FedGPS still surpasses other baselines.

Table 9. More number of clients experimental results

Furthermore, inspired by reviewers, we also perform experiments with different sampling rates; the results are listed in Table 10.

Table 10. Different sampling rates for K=100.

Due to limited space, in addition to the highest acc reported, we will also include the results of mean and standard deviation in the revised version, and thank the reviewers for their suggestions.

Q6: It is unreasonable for the author to set E=1 for the main experiments, which will cause the FL communication overhead to increase sharply. In addition, according to the experimental results in the appendix, when E=5, the advantage of the method decreases significantly, and the author needs to give an explanation.

A6: Thanks to the reviewers for their suggestions on local epochs. In fact, our setting follows many papers [R4-R7] which set local epochs to 1. In addition, we also experimented with E=5, and because the number of local epochs increased, the overfitting of the client to the local data and the deviation from the global makes the larger local epochs more difficult to improve and optimize, but our method still achieves the SOTA effect in all scenarios, which demonstrates the robust performance of FedGPS.

References:
[R1] Hu M, et al. Is aggregation the only choice? federated learning via layer-wise model recombination. In KDD, 2024.
[R2] Xing X, et al. Flexible Sharpness-Aware Personalized Federated Learning. In AAAI, 2025.
[R3] Hu M, et al. FedCross: Towards accurate federated learning via multi-model cross-aggregation. In ICDE, 2024.
[R4] Caldarola D, et al. Beyond Local Sharpness: Communication-Efficient Global Sharpness-aware Minimization for Federated Learning. In CVPR, 2025.
[R5] Tang Z, et al. Virtual homogeneity learning: Defending against data heterogeneity in federated learning. In ICML, 2022.
[R6] Yang Z, et al. Fedfed: Feature distillation against data heterogeneity in federated learning. In NeurIPS, 2023.
[R7] Kim, Junhyung Lyle, et al. Adaptive Federated Learning with Auto-Tuned Clients. In ICLR, 2024.

We sincerely thank the reviewer for the feedback, noting our paper is 'well-motivated and clearly structured', as well as 'novel and intuitive'. We deeply appreciate the reviewer's valuable efforts in improving the quality of our paper, and we address the concerns as follows:

Q1: How exactly the surrogate distribution is generated? It would be beneficial to provide some justifications for why Gaussian distribution was chosen.

A1: We thank the reviewers for their constructive suggestions, which have significantly enhanced FedGPS's readability. Our method assigns a distinct Gaussian distribution to each class in the original dataset. The surrogate data is then generated by sampling from these class-specific Gaussian distributions, with each client sampling from the same distribution for a given class (Note that, unlike heterogeneous, where clients only have partial classes, for the surrogate dataset, each client has surrogate data for all classes). The underlying insight is that any real data distribution can be effectively transformed from a Gaussian distribution [R1]. This approach also benefits from the ease of sampling Gaussian noise, making it well-suited for resource-constrained environments like FL. Furthermore, we evaluate other data distributions, such as the Laplace distribution, and as shown in Table 1, they yielded comparable performance to the Gaussian distribution. Ultimately, we select Gaussian noise for FedGPS.

Table 1. Performance comparison of Gaussian and Laplace distribution sampling for surrogate dataset generation.

Q2: Lacks sufficient discussion on any potential privacy risks from their proposed two techniques. The authors claim that the surrogate distribution is a privacy-preserving approach since no raw statistical information is shared. However, the local surrogate distributions might still be vulnerable to reconstruction or membership inference attacks. Additionally, the dynamic path-oriented gradient rectification explicitly utilizes gradient information derived from other clients. More comprehensive analysis and clarification regarding the potential privacy risk associated with this gradient-sharing technique are needed.

A2: Thank the reviewers for their comments on the paper, which have helped us improve the quality of the paper. In summary, the two technologies of FedGPS do not transmit any additional client's information to any other client compared to traditional methods. On the contrary, FedGPS replaces raw data prototypes with surrogate prototypes instead, thereby further protecting privacy compared to previous works [R2-R5]. We list the privacy clarification among surrogate distribution and gradient information in the following:

• Regarding the surrogate distribution privacy issue, because the surrogate is sampled from different Gaussian distributions, it does not contain any information related to the local data. Further, we transmit the aggregated class-wise prototypes of surrogate data. After aggregating all the high-dimensional embeddings of each class, the surrogate data information is further protected. Many papers[R2-R5] also transmit using the original data prototypes, which offer a weaker level of information protection than what we have. (Detailed information can be referred to the response to Reviewer 3XPZ Q1).
• Regarding the gradient information privacy issue, FedGPS doesn't transmit the gradient information of a certain client to any other client. Every client only uploads its own information to the server and downloads the aggregated information from the server. This process is the same as most other federated methods in uploading and downloading the aggregated information. (Detailed information can be referred to the response to Reviewer 2KuF Q1).

We also add these discussions to the next version. Thanks for your suggestions.

Q3: "This insight builds on the ingenious concept that a malicious client in FL …, offering extra advantages," seems problematic and may require further clarification.

A3: Thanks to the reviewers for suggesting where misunderstandings may have occurred. Our motivation for incorporating the non-self gradient indeed draws a high-level concept from the model replacement strategy in the backdoor to attack federated models. In this scenario, the malicious client leverages a profound understanding of the aggregation process and the collective behavior of benign clients. By accurately anticipating how other clients' updates contribute to the global model, the attacker strategically crafts and scales their own malicious update. This sophisticated manipulation allows them to effectively "replace" the global model with their backdoored version, despite the presence of numerous benign participants. The core insight here is that knowledge of the aggregate influence of other clients provides a significant advantage in shaping the global model. We thank the reviewer again for pointing out this confusing part. We will clarify these insights in the next version.

Q4: Have the authors considered evaluating other non-iid scenarios beyond the Dirichlet partitioning method, such as assigning each client data from exactly C distinct classes (C=2 or C=3), ensuring that each client only sees a specific subset of all labels?

A4: Thanks to the reviewers for their suggestions, we consider more heterogeneous partitioning methods to better show that FedGPS is applicable to other partitions besides Dirichlet. We conduct experiments where each client's data only has C classes. We test both (C=2 and C=3), the results are listed in Table 2.

Table 2. CIFAR-10 results with $\alpha=0.1, K=100, \lambda_s=10\%, R=500, E=1$.

References:
[R1] Rezende D, Mohamed S. Variational inference with normalizing flows. In ICML, 2015.
[R2] Huang, Wenke, et al. Rethinking Federated Learning with Domain Shift: A Prototype View. In CVPR, 2023.
[R3] Tan, Yue, et al. Fedproto: Federated prototype learning across heterogeneous clients. In AAAI, 2021.
[R4] Dai, Yutong, et al. Tackling data heterogeneity in federated learning with class prototypes. In AAAI, 2023
[R5] Wang, Lei, et al. Taming cross-domain representation variance in federated prototype learning with heterogeneous data domains. In NeurIPS, 2024.

We deeply appreciate the reviewer's efforts, particularly their recognition that the proposed method is intuitive and could be applied in a wide range of cases, as well as their affirmation of our experimental results. We also address the questions raised in the following.

Q1: Formally define the local surrogate distribution and global surrogate distribution.

A1: Thank the reviewers for helping to make our paper clearer and easier to understand. In summary,

• We give the formal definition of local and global surrogate distribution.
• The detailed process of how to use local and global surrogate distribution of FedGPS in practice.

We first recap some notations in Table 1(we will revise some notations in the next version). Then we give the formal define of local surrogate distribution and global surrogate distribution in the following:

Table 1. Some notations.

Definition 1.(Local Surrogate Distribution). For a client $k$ in a federated learning system, the Local Surrogate Distribution ($\mathcal{P}_{k}^{s}$) is conceptually defined as the set of feature embeddings obtained by passing each data point from surrogate dataset through the $k$-th local model's feature extractor($\psi_k$).

Definition 2.(Global Surrogate Distribution). At the server side, the Global Surrogate Distribution ($\mathcal{P}^{s}$) is defined as the set of feature embeddings obtained by passing each data point from surrogate dataset through the global extractor($\psi$).

In the implementation, to ensure privacy and reduce communication overhead, what is transmitted to the server is a compressed, privacy-preserving statistical representation of surrogate distribution. Typically these embeddings are then aggregated to form a set of class-wise prototype vectors (e.g., 512-dimensional), with each prototype representing a specific class. This distribution serves as a compact proxy for local surrogate distribution: \mathcal{E}_{k}^{c}=\frac{1}{|\mathcal{D}_s^c|}\sum\psi_k(\xi_s^c),\tag{1}

where $\mathcal{D}^c_s$ denotes the subset of the surrogate dataset belonging to class $c$. Follow the Eq(1) we can get the $C$ prototypes for each class in the surrogate dataset to present the local surrogate distribution in a privacy-preserving and communication-friendly way. The global surrogate distribution is also replaced with aggregated local surrogate prototypes (only selected clients'). \mathcal{E}^{c}=\sum_{k\in\mathcal{S}_t}\mathcal{E}_k^c\tag{2} This prototype-based method also been applied in many FL works e.g.,[R1-R4]. Moreover, FedGPS only transmit the the aggregated surrogate data's prototype, which is safer than the local data's prototype.

Q2: More information needs to be transmitted between the server and clients such as the global features and gradients of other clients. It should be made clear what information is transmitted in each round of communication.

A2: We sincerely appreciate your kind suggestion. Due to the limited rebuttal space, the detailed communication overheads are listed in Table 1 of the response to Reviwer 2KuF. In summary,

• Our method only introduces additional $C*512$ communication overheads during the upload process, which is much smaller than the original model.
• During the download process, in addition to the global model, an additional aggregated gradient information from the previous round of participating clients is required. This is the same size as the original model, and there is also global aggregated global surrogate prototypes (this can still be ignored). Therefore, the additional communication overhead of FedGPS is limited.
• Inspired by the comments, we implement a communication-friendly version of FedGPS, called FedGPS-CF. Compared to FedAvg, only additional class-wise prototypes that can be ignored are required still show comparable peroformance with FedGPS.

Q3: Notation problems:

A3: We appreciate your careful review and highlighting these notation inconsistencies in Sections 3 and 4. We will revise the text accordingly in the next version.

• In line 111 the set of clients is K but in line 112 it is S.
  • We will fix this problem in the next version to replace $\mathcal{S}$ with $\mathcal{K}$ in Line 111.
• In line 163, sum of distributions.
  • We will replace these notations according to the response to Q1.
• In lines 175-176, Stage 1 and Stage 2 are not further defined or referenced.
  • Stage 1 denotes the alignment at local period in the Fig.3(a) and Stage 2 is the alignment at local global stage in Fig.3(a). We will fix these problems.
• In line 180, \delta in bound.
  • We will fix this problem in the next version to replace $\delta$ with $\kappa$ in the Theorem 4.1.

Q4: Why is it necessary to introduce the surrogate distributions? Based on the assumption that the global surrogate distribution is computed from the global real distribution (which is not defined in the paper), if equation (6) is the objective function for local update, then why not simply use the real global and local statistics, i.e., use the distance between the local feature and global feature (sent from the server) in the loss function? In addition, what is purpose of aligning the local real distribution and the local surrogate distribution if the latter stems from the former one?

A4: The surrogate dataset is sampled from different Gaussian distributions representing distinct classes. Thus, the surrogate dataset is different from local datasets (without any information about raw data). Secondly, extensive experiments also demonstrate the potential of sharing statistical information for alleviating data heterogeneity. Therefore, from a distribution perspective, we use a two-stage alignment method to align the locally biased statistical information with the global data while protecting privacy. Lastly, FedGPS employs a two-stage approach. On one hand, it prevents the directly exposure of the statistical information of the original data. On the other hand, all clients use the same surrogate distribution, which reduces the deviation of both global and local statistical information. (The gap can be bridged by aligning the distribution of different models on the same pure noise dataset)

References:
[R1] Huang, Wenke, et al. Rethinking Federated Learning with Domain Shift: A Prototype View. In CVPR, 2023.
[R2] Tan, Yue, et al. Fedproto: Federated prototype learning across heterogeneous clients. In AAAI, 2021.
[R3] Dai, Yutong, et al. Tackling data heterogeneity in federated learning with class prototypes. In AAAI, 2023
[R4] Wang, Lei, et al. Taming cross-domain representation variance in federated prototype learning with heterogeneous data domains. In NeurIPS, 2024.

Thanks for your valuable comments, which improve the readability of our paper. We are delighted to address your concerns.

With your constructive comments, we will introduce the surrogate dataset first. Our method assigns a distinct Gaussian distribution to each class in the original dataset. The surrogate data is then generated by sampling from these class-specific Gaussian distributions, with each client sampling from the same distribution for a given class (Note that, unlike heterogeneous, where clients only have partial classes, for the surrogate dataset, each client has surrogate data for all classes). This introduction will be added in our revision.

We are always welcome to solve any further questions, and your help will be invaluable to FedGPS.

We sincerely thank the reviewers for their invaluable assistance and insightful contributions to our paper. We are particularly pleased to hear that our paper was found to be easy to follow and that our extensive experimental work was acknowledged. We now proceed to address the concerns raised by the reviewers

Q1: Provide more explanation about the Nemenyi post-hoc test. This method appears to be less commonly used in evaluations.

A1: Thanks for pointing out this potentially confusing setting. Accordingly, we have added the following explanations in the revision.

The Nemenyi post-hoc test in Figure 1 is a non-parametric statistical method used for pairwise comparisons of multiple groups [R1] (e.g., algorithms or models) after a significant result from a Friedman test (a non-parametric analog to repeated-measures ANOVA). It ranks the performance of each method across multiple independent runs or datasets and computes a "critical distance" (CD) threshold. If the average rank difference between two methods exceeds the CD, their performances are considered statistically significantly different at a given significance level (typically α=0.05). The test is conservative and accounts for multiple comparisons to control the family-wise error rate, making it robust for scenarios like ours, where we evaluate algorithm robustness across diverse heterogeneity partitions (e.g., different random seeds for Dirichlet distributions).
In Figure 1, the Nemenyi post-hoc test assesses the robustness of baseline methods across different heterogeneity scenarios. The results show overlapping CD intervals for most baselines, indicating no statistically significant performance differences among them. This finding highlights the need for our proposed approach, as existing methods exhibit limited adaptability to varied data distributions. Furthermore, the Nemenyi test has been widely adopted in holistic evaluations [R2-R5] and federated learning scenarios[R6-R7].

Q2: In line 79, the authors describe refining the update direction using other clients' information. Would it be possible to support this with theoretical justification?

A2: Thanks to the reviewers' valuable comments to improve the quality of our paper. The insight behind incorporating information from other clients in FedGPS stems from works like SCAFFOLD [R8] and other related research [R9], which also use other client information as a control variate to adjust update direction. Beyond this intuition, we provide a theoretical justification using a Taylor expansion to demonstrate that integrating other clients' information can indeed further decrease the deviation between local and global update directions. Here's a detailed proof:

We denote local loss function as $f_k(\cdot):\mathbb{R}^d\rightarrow\mathbb{R}$ and global loss function $F(\cdot)$. First of all: The original local update uses the vanilla gradient descent on the local model $\theta_k$ is denoted as $g_{\text{old}}=\nabla f_k(\theta_k).$ For FedGPS, we incorporate non-self gradient $\delta_{\theta_k}$ to local model , we denote $\frac{\delta_{\theta_k}}{||\delta_{\theta_k}||}$ as $g_k'$: $\theta'_k=\theta+\lambda_g g_k'.$ Then we get a new update direction computed based on the new model parameters $\theta_k'$:

$g_{\text{new}}=\nabla f_k(\theta_k').$ For a continuously twice-differentiable function $f_k(\theta)$, the gradient function$\nabla f_k(\theta)$ expands around point $\theta$ along direction $g_k'$: $\nabla f_k (\theta_k+\lambda_g g_k')\approx \nabla f_k(\theta_k) +\nabla^2 f_k(\theta_k)(\lambda_g g_k')+R_3,$ where $R_3$ represents higher-order terms that can be neglected, and $\nabla^2 f_k$ is the Hessian at $\theta_k$. Thus, we can get: $\nabla f_k (\theta_k')\approx \nabla f_k(\theta_k) +\lambda_g \nabla^2 f_k(\theta_k)g_k' .$ Here we assume that the loss function is convex. So the Hessian $\bar{H}=\nabla^2 f_k$ is positive semi-definite. Ideally, we assume non-self gradients contain all the gradients from other clients; we can denote $\delta_{\theta_k} = \nabla F(\theta) - \nabla f_k(\theta_k)$. Substitute into the expansion: $\nabla f_k(\theta_k') \approx \nabla f_k(\theta_k) + \lambda_g \bar{H} (\nabla F(\theta) - \nabla f_i(\theta))=(I - \lambda_g \bar{H}) \nabla f_k(\theta_k) + \lambda_g \bar{H} \nabla F(\theta),$ where $I$ is the identity matrix. As a result:

• The original bias between local and global gradient: $d_0 = ||g_{\text{old}} - \nabla F(\theta)||$
• Refined update direction bias between new model parameters $\theta_k'$ and global gradient: $d'= \|g_{\text{new}} - \nabla F(\theta)\| \approx \|(I - \lambda_g \bar{H}) (\nabla f_k(\theta) - \nabla F(\theta))\| \leq \|I - \lambda_g \bar{H}\| \cdot d_0$. If $\lambda_g$ is tuned to make $||I - \lambda_g \bar{H} < 1||$, then $d'<d_0$, which reduces the shift.

In practice, direct access to all client gradient information is often limited due to privacy and communication overhead. Nevertheless, through careful hyperparameter tuning, FedGPS consistently achieves state-of-the-art (SOTA) performance across various heterogeneous scenarios.

Q3: For Equation 7, could the authors clarify:(1)How is $\lambda$ determined?(2)Why does incorporating non-self gradients improve the update direction?

A3: Thanks to the reviewers for their valuable questions.

• Our approach is not very sensitive to different $\lambda_g$, we conduct the grid search for $\lambda_g\in${$0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0$}. The results are listed in the following Table 1. Therefore, in order to save the time of hyperparameter tuning, we choose 0.5 as a fixed parameter. The following results in our paper also prove that this parameter can still make FedGPS achieve the effect of SOTA. This also shows that FedGPS does not need to spend too much time on parameter tuning.

Table 1. Different hyperparameters of $\lambda_g$.

• Regarding why incorporating non-self gradients improves the update direction, our intuition stems from works like SCAFFOLD[R8], which show that integrating information from other clients can refine one's own update direction. Furthermore, we also provide a theoretical justification that incorporating the information from other clients can further decrease the deviation between local and global update directions. The detailed proof is listed in the response to Q2.

We are grateful for the reviewer's constructive comments, which have enabled us to offer a more thorough explanation of FedGPS's SOTA performance.

Q4: Have the authors considered evaluating their method on Vision Transformers (ViTs) or NLP tasks?

A4: We sincerely appreciate your constructive question. Accordingly, we explore the applicability of our method beyond the current evaluations. The following results and discussions have been added to the revision. To verify the effectiveness of our method, we further incorporate the ViT model in [R10] for FL. The results, which demonstrate consistent improvements in accuracy compared to baselines, have been reported in Table 2.

Table 2.ViT-based model under CIFAR-10 dataset with $\alpha=0.1$.

Due to limited time, we will incorporate more experimental results of ViT in various scenarios in the upcoming version.

Regarding NLP tasks, our method may require additional adaptations to effectively handle sequential data in NLP, especially how to design surrogate information for statistical information alignment. While we have not yet conducted these experiments due to the need for further methodological refinements, this insightful question has inspired us to prioritize NLP evaluations in future extensions of our work. We appreciate the reviewer's insightful comments in broadening the scope of our research.

References:
[R1] Nemenyi, Peter Bjorn. Distribution-free multiple comparisons. Princeton University, 1963.
[R2] Demšar, Janez. Statistical comparisons of classifiers over multiple data sets. In JMLR 2006.
[R3] Chen, Xiangzhi, et al. Disentangling cognitive diagnosis with limited exercise labels. In NeurIPS, 2023.
[R4] Jansen, Christoph, et al. Statistical multicriteria benchmarking via the GSD-front. In NeurIPS, 2024.
[R5] Guidotti, Riccardo, et al. Generative model for decision trees. In AAAI, 2024.
[R6] Yang, Yuwen, et al. Federated multi-task learning on non-iid data silos: An experimental study. In ICMR, 2024.
[R7] Song, Yukun, et al. Causal multi-label feature selection in federated setting. In arXiv:2403.06419, 2024.
[R8] Karimireddy S P, et al. Scaffold: Stochastic controlled averaging for federated learning. In ICML, 2020.
[R9] Wenjing Yan, et al. Problem-Parameter-Free Federated Learning. In ICLR 2025.
[R10] vision-transformers-cifar10: Training Vision Transformers (ViT) and related models on CIFAR-10.

We are truly grateful for the reviewers' dedicated efforts, which greatly improve the quality of FedGPS. Their well-written and easy to follow comments were invaluable, and we especially appreciate their strong affirmation of our experimental results. Additionally, we provide the following answers to the reviewers' concerns:

Q1: FedGPS requires many more communications than existing methods. It seems that each client has to communicate a high-dimensional vector to every other client. Extra computational cost. Incorporating gradient information from other clients incurs additional computational cost? By communicating local gradient information with other clients, can it hurt the improved privacy aspect of federated learning?

A1: Thank you for your valuable feedback and for raising this important point regarding the communication and computational costs of FedGPS. In summary,

• During the download stage, aligning with your valuable feedback, there is one additional model (containing the aggregated gradient) of the same size as the global model.
• During the upload stage, there is no additional communication overhead.
• During training, no communication is required between clients, namely, only communication is needed between the client and the server.
• We further implement a variant with no extra communication overhead called FedGPS-CF, showing comparable performance with FedGPS.

Thus, FedGPS protects privacy as previous works and brings about 1.5 times communication overhead than vanilla methods, e.g., FedAvg.

Inspired by your constructive comments, we have provided a detailed table description of the two consecutive upload and download rounds in the revised version, along with the associated local computational requirements. Additionally, we include a comparison with FedAvg to demonstrate that FedGPS does not impose significant additional overhead. We denote the whole model size as $M$ and the total classes of the dataset as $C$, e.g., $C=10$ for CIFAR-10.

Table 1. The whole process comparison between FedAvg and FedGPS (Ours).

The extra $C*512$ (where $C * 512 \ll M$，e.g., 0.05% in CIFAR-10) is the local uploaded local surrogate distribution prototypes and downloads the global surrogate distribution prototype. （The detailed explanation of local and global surrogate information can be referred to the reply to Reviewer 3XPZ Q1）. To further clarify, we emphasize that in each round of FedGPS, we do not require all clients to send their own gradients to other clients. Instead, based on the information from the selected clients in the previous round, we aggregate their gradients for dual purposes: updating the global model and providing aggregated gradient information for client selection in the next round. As a result, the upload cost remains identical to that of vanilla FedAvg. During the download from the server, clients receive not only the latest model but also the aggregated gradient information from the server. This aggregation process inherently protects individual client gradients from being accessed by others (addressing privacy concerns). Consequently, throughout the entire FedGPS process, there is no direct information exchange between any clients, and all information obtained from the server is solely aggregated data.

Inspired by your valuable comments, we further design a communication-friendly variant of FedGPS, called FedGPS-CF. Specifically, when uploading the model, FedGPS-CF, like FedGPS, only has a communication cost of $M+C*512$. When downloading from the server, it still only downloads the global model and global surrogate prototypes, meaning the communication cost remains $M+C*512$. Here, we use the difference between the global model downloaded in two rounds from the server to represent the gradient aggregation information of other clients. Similarly, if a client is selected in both adjacent rounds, its own update information should be removed. Finally, FedGPS-CF achieves comparable performance to FedGPS, the results are listed in Table 2. However, it reduces the download overhead of $M$, making FedGPS-CF only have an additional communication cost of $C*512$ compared to FedAvg, which is negligible. Once again, we thank you for your help and for inspiring our paper.

Table 2. The results between FedGPS-CF and FedGPS under Heterogeneous scenario 1 with CIFAR-10.

Q2: How come comparison to adaptive methods are not included? These also try to address heterogeneity issue in federated learning.

A2: We agree that comparing a new method with adaptive methods is necessary, since these methods are also outstanding approaches to address data heterogeneity in FL. Following your valuable comments, we have selected two representative adaptive methods $\Delta$-SGD and FedAdam, from references [R5, R6] and incorporated them as additional baselines in our experiments. Along with the other methods already evaluated, these comparisons further demonstrate the superiority of FedGPS. The results are listed in Tables 3, 4, and 5. We have discussed these papers [R5-R7] and listed these results in our revised version.

Table 3. The results in FedGPS under Heterogeneous scenario 1 with CIFAR-10.

Table 4. The results in FedGPS under Heterogeneous scenario 1 with SVHN.

Table 5. The results in FedGPS under Heterogeneous scenario 1 with CIFAR-100.

References:
[R1] Durmus Alp Emre Acar, et al. Federated learning based on dynamic regularization In ICLR, 2021.
[R2] Taehwan Lee, et al. Rethinking the Flat Minima Searching in Federated Learning. In ICML, 2024.
[R3] Ziqing Fan, et al. Locally Estimated Global Perturbations are Better than Local Perturbations for Federated Sharpness-aware Minimization. In ICML, 2024.
[R4] Wenjing Yan, et al. Problem-Parameter-Free Federated Learning. In ICLR 2025.
[R5] Reddi, Sashank J., et al. Adaptive Federated Optimization. In ICLR, 2021.
[R6] Kim, Junhyung Lyle, et al. Adaptive Federated Learning with Auto-Tuned Clients. In ICLR, 2024.
[R7] Mukherjee S, et al. Locally adaptive federated learning. In TMLR, 2023.
[R8] Karimireddy S P, et al. Scaffold: Stochastic controlled averaging for federated learning. In ICML, 2020.