Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging

Thanks for your valuable comments. We will make modifications in the final version as suggested by the reviewer, and our point by point responses to your major comments are given below.

To address your concern, we conducted many experiments on CIFAR-100 with ResNet-18 under different data heterogeneity levels (0.1, 0.3, 0.6), and compare the proposed algorithms with FedSMOO and FedGAMMA, suggested by the Reviewer. The reproduced results are based on the pre-activation ResNet18-GN network, which is widely used in federated learning algorithms such as FedSAM, FedGAMMA, and FedACG. The following results are produced under our framework, referencing the open-source code of FedSMOO and FedSAM. Because the FedGAMMA algorithm lacks open-source code, we implemented it based on the algorithm flow described in the paper and referenced the FedGamma implementation in FedLESAM's source code. We will include these results in the final version, and will release our code and the federated learning framework to ensure reproducibility.

Thanks for your valuable comments. We will make modifications in the final version as suggested by the reviewer, and our point by point responses to your major comments are given below.

1. Different from FedSAM, the proposed FedSWA algorithm is a new way to improving FL generalization. By using SWA, FedSWA better maintains global flatness consistency compared to FedSAM, as shown in Figure 1. Moreover, FedMoSWA also uses momentum stochastic control, unlike MoFedSAM with only local momentum SAM. Both techniques (i.e., SWA and momentum stochastic control) significantly enhance the theoretical results of the federated learning algorithms such as FedSAM and MoFedSAM, as shown in Table 1.

Experimental results: All the experimental results in the manuscript show that our algorithms show consistent improvements on both ResNet and Transformer networks. As shown in Tables 2, 3 and the followng table, FedSWA with SWA optimizer (59.8%) outperforms FedAvg (54.2%) and FedSAM (47.8%) on ResNet-18 with CIFAR-100. Our FedMoSWA (67.9%) also surpasses MoFedSAM (60.1%). Both theoretical and experimental results show that FedSWA outperforms FedSAM, and FedMoSWA outperforms MoFedSAM.

Our theoretical results in Table 1 show that the generalization bound of our FedSWA is $\mathcal{O}\big(\frac{L}{m n \beta}e^{\frac{1}{T}+1}(\tilde{c} L+ \tilde{c}\sigma_g+\tilde{c} \sigma) \big)$, which is correlated with data heterogeneity and is superior to that of FedSAM, $\mathcal{O}\big(\frac{L}{m n \beta}e^{1+\frac{1}{T}}(\overline{c}L+ \overline{c}\sigma_g+\overline{c} \sigma) \big)$ and MoFedSAM, $\mathcal{O}\big(\frac{L}{m n \beta}e^{1+\frac{1}{T}}(\overline{c}L+ \overline{c}\sigma_g+\overline{c} \sigma) \big)$, where $\overline{c}\>\tilde{c}\=\1+\left(2+1/KT\right)^{K-1}/T\gg 1$. We prove that the generalization error of FedMoSWA is $\mathcal{O}\big(\frac{L}{m n\beta} e^{\frac{1}{T}+1}(\tilde{c} L+ \sigma_g+\tilde{c} \sigma) \big)$, which is better than those of FedSAM, MoFedSAM, FedSWA, where $\tilde{c}=1+\left(2+1/KT\right)^{K-1}/T\gg 1$.

2. To address youe concern, we conducted some experiments on CIFAR-100 with ResNet-18, following the pathological imbalanced settings from the FedSMOO and FedLESAM papers, where each client has at most 10 classes. This represents a highly imbalanced data distribution. These results will be included in our final paper.

Our FedSWA algorithm outperforms the algorithms, FedAvg and FedSAM. FedSWA uses SWA as the local optimizer, while FedAvg uses SGD, and FedSAM uses SAM as local optimizers, and none of them employ variance reduction or momentum acceleration techniques. As the advanced version of FedSWA, our FedMoSWA, which also incorporates momentum and variance reduction, surpasses other algorithms in all the settings.

3.In fact, our source code was included in the supplementary materials. We will release our code and the federated learning framework to ensure reproducibility.

We sincerely thank the Reviewer for their valuable comments, and our point by point responses to your major comments are given below.

1. For Experimental Designs Or Analyses. In fact, we have done the hyperparameter tuning for all the algorithms, and we also followed the parameter settings from their original papers. We conducted full client participation experiments on CIFAR-100 with 10 clients and 300 rounds, where the model is ResNet-18. Here, FedMo denotes FedMoSWA without SWA (i) but only with Momentum Stochastic Control (ii), as shown in the following table. FedMo (62.5%) outperforms SCAFFOLD (59.9%) by 2.6%. With 100 clients and 10% participation, our FedMoSWA achieves 3.3% higher accuracy than SCAFFOLD, addressing the update delay issue with partial client participation.

Full client participation on CIFAR-100 and ResNet-18 with 10 clients and 300 rounds:

2. For Weaknesses 1. From Tables 1 and 4 in our manuscript and the following table , when $\rho=0.1$, FedSWA using only SWA (59.8%) outperforms FedAvg (54.2%) and FedSAM (47.8%) on ResNet-18 with CIFAR-100. When $\rho=1$, FedMo using only momentum stochastic control achieves 65.9%, higher than SCAFFOLD (54.1%), demonstrating that momentum variance reduction mitigates SCAFFOLD’s variance reduction delay. By combining SWA and momentum variance reduction, FedMoSWA achieves 67.9%, showing that both SWA (+5.6%) and momentum control (+11.7%) are effective, and momentum control has greater impact. Moreover, FedSWA is a simple algorithm like FedSAM and can be combined with other techniques. We will clarify this in our final paper.

10% client participation on CIFAR-100 and ResNet-18 with 100 clients and 1000 rounds:

3. For Weaknesses 2. To address your concern, we will add some discussions about the optimization error analysis in our final paper. In fact, Section 5.2 discusses the optimization error analysis, showing that FedMoSWA converges faster than the best-known algorithm, SCAFFOLD, and outperforms both FedSWA and other baselines such as FedSAM, FedAvg, and MoFedSAM. Unlike MoFedSAM, which uses only local momentum, our FedMoSWA employs momentum variance reduction. Additionally, our FedSWA also converges faster than both FedSAM and FedAvg, as shown in Table 1. We will include all the discussions in Section of Introduction in our final paper.

4. For Other Comments Or Suggestions. To address your concern, we will improve the final version of the paper, and will use italic and bold to distinguish the two algorithms.

5. For Questions. In our analysis, data heterogeneity does not vanish when performing only one local training step. It is not purely an artifact of analysis but a fundamental characteristic of federated learning algorithms, driven by the inherent bias introduced by local training on heterogeneous data. This bias does not completely vanish with only one local training step. Our future work will focus on addresssing this issue.