One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory

We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.

---

W1. Clarification of communication and computational cost in Table 1：We apologize for not making it clear and we will modify the caption of Table 1 and add a discussion section with following content in the revised version.

Communication cost: The communication cost is defined as the parameters transmitted per round. For example, if we denote the amount of model parameter as Ω, FedSAM requires 2Ω due to bidirectional model parameter exchange. In contrast, FedGMT incurs a communication cost of 3Ω due to the server transmitting the EMA model to clients, whereas FedGMTv2 achieves 2Ω by omitting this transmission. Thus, FedGMT's communication cost is either $1.5×$ or $1×$ that of FedSAM.

Computational cost: The computational cost is defined as per-iteration training cost. Specifically, minor computational overhead (e.g., model aggregation, ADMM computations) is ignored compared to the primary training computation. Taking FedAvg as the baseline ($1×$ cost), FedSAM doubles the forward and backward processes, resulting in a $2×$ cost. FedGMT involves an extra forward pass. we refer to the official PyTorch paper [A], which shows that for ResNet50, the backward pass is $3.7×$ the forward pass. Based on this data, FedGMT's cost is estimated to be about $1.2×$ by $(\frac{1+1+3.7}{1+3.7})$.

If we assume a forward cost of f and a backward cost of 2f, FedAvg requires 3f per iteration, while our FedGMT needs 4f, leading to a $1.33×$ cost. We will describe this hypothesis and adjust the estimate from $1.2×$ to $1.33×$ for better understanding to readers.

While FedGMT incurs marginal computational overhead, it achieves superior accuracy and efficiency. As shown in following Table in W2, FedGMT requires only 47% of the time cost of FedAvg to reach the target accuracy, demonstrating its practical advantage in resource-constrained environments. For example, if one trains a FL model to reach the target accuracy with 8 Nvidia V-100 GPUs, considering ​​Google Cloud Platform charges \$2.48 per GPU per hour, the total cost of FedAvg for 1 day will be \\$476, FedSAM will be \$724 (476\*152$224 (476*47%). There is a significant cost savings of implementing FedGMT.

[A] Li S, Zhao Y, Varma R, et al. Pytorch distributed: Experiences on accelerating data parallel training[J]. Proc. VLDB Endow. 13(12): 3005-3018 (2020)

---

W2.Experiments with larger model: We conducted experiments on CIFAR100-Dir(0.1) with Resnet18 on one NVIDIA 4090 GPU. We follow the same parameter settings in FedSMOO paper and give a comprehensive comparison. The results are stated below.

FedGMT achieves the highest accuracy (50.67%) and fastest convergence (257 rounds/198min) to reach 44% accuracy, outperforming all baselines. While FedSMOO (47.94%) and FedLESAM-D (46.42%) show better accuracy than FedAvg (43.95%) and FedSAM (44.56%), they lag behind FedGMT in efficiency. FedGMT balances accuracy and efficiency, surpassing competitors in both metrics.

---

W3.Experiments with larger number of clients: We conducted experiments on CIFAR10-Dir(0.1) with LeNet . We set 500 clients with 20% active ratio and 1000 clients with 10% active ratio to involve 100 clients per round. The results are stated below.

The result shows that FedGMT's superior scalability and robust for large-scale FL scenario. In addition, you can refer to our experiments in Reviewer 5WUM that also verify the robustness of FedGMT in an extreme non-IID FL setup.

---

W4. Theoretical improvements. Thank you for your insightful feedback and suggestions to enhance the theoretical rigor of our paper. We wholeheartedly agree with your comments and will reorganize the theoretical proof by following your suggestions for each step in our updated version.

---

It is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.

Thank you for the comments and suggestions! We answer your questions below.

---

1.Optimization target explanation. By referring the Theorem 1 in original SAM(Foret et al.,2020), the objective loss function of FedSAM in each client $m$ can be rewritten as the sum of the vanilla loss and the loss associated to the sharpness measure, which is the maximized change of the training loss within the ρ-constrained neighborhood, i.e., $\arg\min_{w_m}[L_m(w_m)+S_m(w_m)]$ where $S_m(w_m) = max_{\epsilon: ||\epsilon||<\rho}L_m(w_m+\epsilon)-L_m(w_m)$, which is a minimax problem as you mentioned. The sharpness measure is approximated as $S_m(w_m) = L_m(w_m+\epsilon_m)-L_m(w_m)$, where $\epsilon_m = \rho\frac{\nabla\mathcal{L}_m(w_m)}{\|\nabla\mathcal{L}_m(w_m)\|}$ is the solution to an approximated version of the maximization problem in FedSAM. The sharpness measure in FedSAM approximates the inner maximization problem and is combined with the vanilla loss for minimization. This transforms the minimax problem into a standard minimization task. Our work based on this task to develop a better sharpness measure to minimize global sharpness in FL, as formally established in Eq. (8) of the manuscript.

---

2.Role of the dual variable $u$. Note that in line 13 of Algorithm 1, each client has its own dual variable $u_m^{t+1} = u_m^{t}-\frac1\beta(w_m^t-w^t)$. This enables clients to ensure that the solutions of their respective sub-problems are consistent under global constraints.

In line 16 of Algorithm 1, we define the global dual variable $u^t$ because, according to Eq. (13), the global model $w^{t+1}$ needs to be updated as $w^{t+1}=\frac1N\sum_{m \in \mathbb{N}}(w^t_{m,K}-\beta u^{t+1}_m)$. To avoid having each client send $u_m^{t+1}$ to the server in every round, we define $u^t$ to assist in minimizing Eq. (13). We also conducted experiments on CIFAR10-Dir(0.1), FedGMT's accuracy drop from 79.17 to 65.59 without $u$. Therefore, updating $u$ is necessary, not redundant.

---

3.Analysis in Table 1.

To save space for other question, we have addressed this concern in our rebuttal to Reviewer Zauy (Section W1). Please refer to that section for our detailed response.

---

5.1 Compare the federated ADMM-based work. The baselines we take in experiments, such as FedDyn, FedSpeed and FedSMOO, which are all state-of-the-art FL methods based on ADMM.

5.2 Comparison criterion. The comparison criterion we used follows original FedSAM(Caldarola et al.,2022). we ensure sufficient communication rounds to guarantee convergence (see Fig.5), while using the final averaged accuracy to avoid fluctuations and provide a more reliable performance estimate. We also provide the historical best accuracy by random seed {1,2,3} in anonymous link https://anonymous.4open.science/r/RG5, the ranking of each method is consistent with our criterion.

---

6.Loss landscape. We provide the loss landscape, top Hessian eigenvalue and Hessian trace of SAM-based method in anonymous link in 5.2, demonstrating that our FedGMT can reach a flatter minima than others.

---

7. Hyperparameters sensitivity and selection. Thanks for the suggestion. Our experiments across 4 CV/NLP datasets, 4 model architectures, and different scenarios (heterogeneity and participation) show that FedGMT's hyperparameters can be consistent and robust. We will detail the meanings and selections of these hyperparameters in this part.

• The strength $\gamma$ and temperature $\tau$ which is a basic hyperparameter for KL function. We follow the common settings used in previous works and $\gamma=1$ and $\tau = 3$ for all settings.
• The penalty coefficient $\beta$ has been studied in many previous ADMM-based works. We test the selection of {10, 100} works on all settings.
• $\alpha$ control the loss on the recent update trajectory in EMA. Since $\alpha$ needs to be close to 1 to retain past trajectories, we mainly chose from {0.5,0.9,0.95,0.99}. We test $\alpha=0.95$ works well on all settings.

In summary, the hyperparameters in FedGMT can be selected by simple experience as we mentioned above. In practice, FedGMT ​​can achieve higher performance compared to other baselines without the complex tuning.

---

8.1 The originality is that we develop a efficient (communication and computation) global sharpness measure for SAM-FL. We do increase the amount of calculation in line 13,16,18 of Algorithm 1, but it is negligible compared to the mainly training costs. Notably, FedGMT avoid 1 more backward pass required for calculating perturbations in standard SAM.

8.2 Real-world deployment. To validate practical deployment feasibility, We implemented FedGMT on 2×RPi4B, 2×Jetson Nano, and 1×PC. FedGMT achieves about 47% energy savings vs. FedAvg across three datasets, as shown in the anonymous link in 5.2.

---

It is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.

We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.

---

1.Comment on realistic client drifts. Client drift in federated learning refers to the phenomenon where data distributions across devices (clients) change over time or space, leading to performance degradation. After a comprehensive review of existing federated learning literature to address client drift problem, we summarize that numerous studies simulate non-IID federated settings using partitioning strategies on public datasets. This approach is preferred because real-world federated datasets are scarce due to data regulations and privacy constraints, while synthetic partitioning offers flexibility in controlling key FL parameters (e.g., client count, data size) and leveraging existing centralized training knowledge. Thus, synthetic client drifts is more complex and flexible, enabling methods developed under these conditions to generalize more readily to realistic client drift scenarios.

In our work, as mentioned in Section 4.1, we adopt two widely used data partition strategy:

• Pathological: Only selected categories can be sampled with a non-zero probability. The local dataset obeys a uniform distribution of active categories.
• Dirichlet: Each category can be sampled with a non-zero probability. The local dataset obeys a Dirichlet distribution. Notably, we extend this by transforming the original balanced dataset into a long-tail distribution to further enhance the data heterogeneity and simulate real world scenarios.

---

2.Propose and evaluate more than Dirichlet. We introduce a non-IID data partitioning method where client datasets are strictly non-overlapping (disjoint) across classes, simulating an extreme data heterogeneity scenario. Specifically:

• CIFAR10 (10 total classes) : 10 clients, each assigned a unique class
• CIFAR100 (100 total classes) : 100 clients, each assigned a unique class

This single-class assignment ensures maximal label distribution discrepancy between clients. Experiments were conducted with participation rate 40% and communication rounds 500. The results under this extreme setting are presented below.

The results shows that in an extreme non-IID federated learning setup where clients hold disjoint class partitions, FedGMT achieves state-of-the-art results: 58.23% on CIFAR10 and 22.13% on CIFAR100, outperforming FedLESAM-D (45.35%/13.07%) and other baselines. FedAvg (31.57%/6.42%) and FedSAM (29.39%/5.56%) struggle with label heterogeneity, while FedSMOO fails entirely. These results validate FedGMT’s superior robustness to extreme label skewness.

---

3.Comment on the challenge of privacy. Our work builds upon the well-established FedAvg framework, inheriting its parameter-based communication architecture between server and clients. This design ensures full compatibility with existing privacy-preserving techniques developed for FedAvg, including mainstream defenses against privacy attacks such as differential privacy [A], homomorphic encryption [B], and secure multi-party computation [C].

The core contributions of our paper focus on enhancing practical generalization accuracy for real-world datasets while minimizing communication and computation costs under data heterogeneity. Importantly, our methods impose no restrictions on the integration of privacy-preserving techniques, ensuring seamless compatibility with existing or emerging privacy frameworks.

[A] Chen,H.,Vikalo,H.,etal. The best of both worlds: Accurate global and personalized models through federated learning with data-free hyper-knowledge distillation(ICLR , 2023).

[B] Cai Y, Ding W, Xiao Y, et al. Secfed: A secure and efficient federated learning based on multi-key homomorphic encryption[J]. IEEE Transactions on Dependable and Secure Computing, 2023, 21(4): 3817-3833.

[C] Chen L, Xiao D, Yu Z, et al. Secure and efficient federated learning via novel multi-party computation and compressed sensing[J]. Information Sciences, 2024, 667: 120481.

---

It is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.

We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.

---

W1.FedGMT vs. FedGKD. While FedGKD appears to be similar to our method, there are fundamental differences. Practically, FedGKD takes an element-wise average over the latest 5 rounds (as they suggest) of global models and sends it to clients. This averaging method isn't EMA; it demands extra storage space and doubles the communication overhead. We use EMA with KL divergence in the training loss because, as proven in Eq. (8), minimizing the loss difference $L(e)-L(w)$ between the EMA model $e$ and the global model $w$ is equivalent to minimizing the SAM’s sharpness measure for the global optimization. However, directly combining this loss difference with the FL objective will cancel out $L(w)$. So we replace the cross entropy (CE) loss with the KL divergence loss to decouple the vanilla loss since minimizing the CE loss is equivalent to minimizing the KL loss. Thus, minimizing the KL loss in FedGMT aims to minimize global sharpness and reduce computational cost compared to other SAM-based methods, a goal that FedGKD cannot achieve. The comparative experiments in the following table in W2 demonstrate that FedGMT outperforms FedGKD by 2.1% on Dir(0.1) and 2.79% on Dir(0.01) when both methods employ the ADMM.

---

W2. Methods comparison without ADMM. We note that the performance of $L^{glotra}$ degrades without ADMM, which is due to non-vanishing biases between local update and global update in heterogeneity scenario. This aligns with methods such as FedSpeed, FedSMOO, and FedLESAM, which also incorporate ADMM to address such issues. Thus, approaches that omit ADMM may fail to fully leverage the proposed sharpness measure's potential. We evaluate performance with/without ADMM on CIFAR10 to demonstrate ADMM's necessity for maintaining optimization stability under data heterogeneity. The results are stated below.

The above experimental results show that FedGMT with ADMM achieve best performance due to our sharpness measure can more accurate estimate global sharpness than others (Figure 2 also demonstrate this). We will add these experimental results in our updated version.

---

W3.Learning rate tuning. Thanks for the suggestions. Learning rate which is a basic hyperparameters in the deep learning, and usually we do not finetune this for the fair comparison in the experiments. After a comprehensive review of existing federated learning literature, we observe that prior works consistently use a fixed learning rate across all compared methods. In our study, we select learning rates from the set {0.001,0.01,0.1} for different datasets and model architectures to maximize SGD performance, then fix this optimal value for all methods. Specifically, a learning rate of 0.01 was employed for CIFAR10, CIFAR100, and CINIC10, while 0.1 was used for AG News.

---

Q1. Convergence rate understand. In the federated stochastic non-convex setting, under standard assumptions of "smoothness," "bounded stochastic gradients," "heterogeneity," and other assumptions, prior works[A,B] have established $O(1/T)$ convergence rates for their algorithms. Specifically, our theoretical contributions lie in proving convergence without relying on the restrictive assumptions of bounded heterogeneity or requiring local minima in each communication round. This is enabled by our ADMM-based framework, which derives properties (e.g., Lemmas D.5 and D.6) to bound the global gradient norm in Eq. (34). Prior work [C] establishes faster convergence under the assumption of client similarity in non-convex settings. Our Eq. (10) scales $L^{glotra}$ as a non-negative convex function which can be seen as an exponential moving average of historical global gradients into local updates across all clients. This shared similarity in EMA across clients reduces variance between local and global updates, thereby accelerating convergence.

[A] Fedpd: A federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055–6070, 2021.

[B] Fedadmm: A federated primal-dual algorithm allowing partial participation. arXiv preprint arXiv:2203.15104, 2022.

[C] Scaffold:Stochastic controlled averaging for federated learning(ICML 2020).

---

It is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.