FedWMSAM: Fast and Flat Federated Learning via Weighted Momentum and Sharpness-Aware Minimization

Q1: Experiments

Question: From the experimental results in Table 1, we can see that when the heterogeneity of the data set is weak, the effect of the method proposed in this paper is not outstanding. What is the reason for this? How can this problem be improved?

Ans1: Thank you for your attention to the experimental results in Table 1. Our method performs less prominently in scenarios with weak data heterogeneity mainly because traditional methods (e.g., FedAvg) experience less performance degradation under such conditions, and the differences between clients have a smaller impact on the global model, making it harder for the momentum mechanism and adaptive weighting to demonstrate their advantages. However, in scenarios with high heterogeneity, these improvements significantly enhance performance. Additionally, compared to other methods (e.g., FedSMOO), our method is slightly less competitive in low heterogeneity settings, but those methods generally incur higher computational and communication costs (approximately twice as much as ours), while our approach is specifically designed for resource-constrained edge environments.

Q2: Theoretical Analysis

Question: In theoretical analysis, the dynamic adjustment weight is directly represented by a variable. Can its dynamic design be reflected or discussed in theoretical analysis?

Ans2: We appreciate the reviewer’s attention to the theoretical analysis of dynamic weight adjustment. Dynamic weight adjustment significantly improves model performance in practice, primarily by balancing the contributions of the momentum mechanism and the Sharpness-Aware Minimization (SAM) method at different stages. However, in the theoretical analysis, dynamic weight adjustment does not directly affect the convergence properties. Its main role lies in optimizing the flatness of the global model through weight adjustment rather than directly influencing the convergence rate.

Additionally, we have provided a detailed analysis in Appendix B on the impact and selection of dynamic weights in experiments. The experimental results further validate its effectiveness, complementing the limitations in the theoretical analysis.

Q3: Motivation and Novelty

Question: This motivation (like limited communication rounds, data heterogeneity, and client drift and slow convergence and poor model generalization) is quite commonly seen in FL. This writing style is highly repetitive with other similar works and does not highlight one's own insights, especially in the abstract and introduction.

Ans3: We appreciate the reviewer’s suggestions regarding the abstract and introduction. We recognize that these sections in the original version may not have sufficiently highlighted the core problem and innovations of this work, which might have caused some overlap with other similar studies. In fact, the core focus of our work is to address the global optimization challenges in federated learning caused by data heterogeneity, particularly the limitations of SAM in this context. By introducing a momentum mechanism and adaptive weighting strategy, we significantly improve the application of SAM, achieving better performance in heterogeneous data scenarios.

In the camera-ready version, we will refine the abstract and introduction to emphasize our key innovation: how solving for globally flat optima at the local level helps overcome the challenges posed by heterogeneity. This will further highlight the uniqueness and contributions of our work.

Q4: Contribution

Question: The core technology of this work seems to be a combination of the momentum and the SAM with an adaptive weighted scheme. What are the key technical contributions or the technical novelty?

Ans4: We appreciate the reviewer’s attention to the technical contributions of our work. Our core contributions go beyond the simple combination of momentum and SAM and include the following key innovations:

1. Systematic Analysis of the Core Contradiction in Applying SAM to Federated Learning:
   We are the first to systematically analyze the fundamental contradiction of applying SAM in federated learning: SAM relies on perturbation computation based on local client data, while improving the generalization performance of the global model requires optimizing SAM at the global level. This contradiction is a critical challenge in federated learning, yet it has been overlooked in prior studies. To address this issue, we propose leveraging a momentum mechanism to transmit global information across clients with low communication costs, effectively guiding local optimization and mitigating this contradiction.

2. Proposal of the FedWMSAM Method and Convergence Proof:
   We propose FedWMSAM, a novel method that effectively combines personalized momentum with SAM to address the divergence between local and global models. FedWMSAM introduces an adaptive adjustment factor that dynamically balances the contributions of momentum and perturbation at different training stages. This ensures efficient convergence to a flatter global optimum under heterogeneous conditions. Additionally, we theoretically prove the convergence of FedWMSAM, demonstrating its efficiency with a convergence rate of O(LΔσ²/(ρNKR) + LΔ/R).

3. Comprehensive Experimental Validation:
   We conducted extensive experiments across 4 image datasets (Fashion-MNIST, CIFAR-10, CIFAR-100, OfficeHome), 9 state-of-the-art methods, and various heterogeneous settings. The results demonstrate that FedWMSAM significantly outperforms existing methods in terms of accuracy, convergence speed, and computational cost. Ablation and sensitivity analyses further validate the robustness of our method:

   • Accuracy: FedWMSAM achieves consistent accuracy improvements across datasets.
   • Convergence Speed: It converges faster than existing methods, even under high heterogeneity.
   • Computational Efficiency: FedWMSAM reduces computational costs while achieving superior performance.

4. Theoretical and Practical Insights:
   FedWMSAM extends the application of SAM to federated learning by resolving the local-global optimization contradiction. The theoretical convergence proof not only establishes its efficiency but also provides a foundation for future studies to address similar challenges in federated optimization.

5. Comparison with Existing Work and Improvements:
   Our method builds on ideas from related works (e.g., FedAvg, FedSAM, and FedLESAM) but introduces significant improvements. Unlike existing methods, FedWMSAM effectively addresses the heterogeneity-induced divergence between local and global models and achieves better performance with lower computational and communication costs.

In the camera-ready version, we will refine the contribution summary to ensure clarity and emphasize the following key points:

• The identification of the core contradiction in applying SAM to federated learning and how FedWMSAM resolves it.
• The innovative combination of momentum and SAM, supported by adaptive weighting, to dynamically balance local and global optimization.
• The theoretical convergence proof and its practical significance in federated learning scenarios.
• Extensive experimental results that establish FedWMSAM’s superiority across diverse settings.

Thank you again for your valuable feedback, which has helped us better articulate the novelty and significance of our work.

Q1: I think the authors should provide additional experimental comparisons with dir0.3 to better clarify the boundaries of heterogeneous selection. Due to time constraints, the authors may wish to showcase the best algorithms.

Q3: I understand the authors' intentions. The question is very intuitive and interesting, and the writing needs to be more precise.

Therefore, I decided to maintain my score.

Thank you for your thoughtful and constructive feedback. We greatly value your insightful comments, and they have been immensely helpful in guiding us to improve our work. Below, we address your comments in detail:

Q1: Additional experimental comparisons with dir0.3 to clarify boundaries of heterogeneous selection.

Thank you for pointing out the importance of evaluating our method under Dirichlet 0.3 to better characterize the boundaries of heterogeneous selection. We appreciate this valuable suggestion and have conducted the requested experiments accordingly.

In our original submission, we adopted Dirichlet 0.1 and 0.6 following common practices in prior work (e.g., FedCM[1], FedSMOO[2], FedLESAM[3]) to capture both extreme and moderate heterogeneity scenarios. These choices were not to cherry-pick favorable settings, but to ensure broad coverage and fair comparisons.

In response to your comment, we have conducted new experiments on CIFAR-10 under Dirichlet 0.3, and the results are as follows:

• FedSMOO: 79.52%
• FedWMSAM (Ours): 80.12%

These results confirm the robustness of FedWMSAM even in intermediate heterogeneity levels. We are currently expanding this analysis to include Fashion-MNIST, CIFAR-10, and CIFAR-100 under Dirichlet 0.3, and the complete results will be included in the revised version.

As for your observation that FedSMOO performs slightly better than our method on Fashion-MNIST and CIFAR-10, we would like to clarify the following:

• These datasets are relatively simple, and FedSMOO achieves its performance through significantly heavier computation and regularization. Specifically, FedSMOO requires 29.77 seconds per round in line 292, nearly twice the computation time of our method (15.03 seconds per round).
• Our design prioritizes both performance and efficiency. On more complex benchmarks such as CIFAR-100 and OfficeHome, which better reflect real-world heterogeneity, FedWMSAM consistently outperforms FedSMOO, demonstrating superior scalability and generalization.

Therefore, while FedSMOO may appear to have an edge on simpler tasks, we believe FedWMSAM strikes a better balance between accuracy, stability, and efficiency across varying levels of heterogeneity.

We hope this clarification addresses your concerns, and we thank you again for your insightful suggestion that helped strengthen our work.

Q3: Writing precision.

Thank you for recognizing the intuitiveness and relevance of our problem and methodology. We are grateful for your positive feedback and fully acknowledge the need to improve our writing. We are currently revising the manuscript to make our explanations more precise and concise, ensuring that our contributions are communicated effectively.

We sincerely hope these updates address your concerns and that the additional results and revisions will allow you to reconsider your evaluation. We greatly appreciate your time and effort in reviewing our work and look forward to further discussions with you.

References

[1] Xu, Jing, Wang, Sen, Wang, Liwei, and Yao, Andrew Chi-Chih. Fedcm: Federated learning with client-level momentum. arXiv preprint arXiv:2106.10874, 2021.

[2] Sun, Yan, Shen, Li, Chen, Shixiang, Ding, Liang, and Tao, Dacheng. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. International Conference on Machine Learning, pages 32991–33013, 2023.

[3] Fan, Ziqing, Hu, Shengchao, Yao, Jiangchao, Niu, Gang, Zhang, Ya, Sugiyama, Masashi, and Wang, Yanfeng. Locally estimated global perturbations are better than local perturbations for federated sharpness-aware minimization. arXiv preprint arXiv:2405.18890, 2024.

Thanks for the new feedback. Considering the interest of the questions raised and the completeness of the theoretical proof, I have increased the score by one.

Nevertheless, I would like to see clearer and easier arguments in the updated version.

Thank you very much for your positive feedback and for increasing the score of our work. We greatly appreciate your recognition of the interest in the questions raised and the completeness of the theoretical proof.

We understand your suggestion regarding the clarity and simplicity of the arguments. In the updated version, we will carefully revise our explanations to ensure that the presentation of the key ideas and theoretical insights is clearer and more concise. Our aim is to make the arguments more straightforward and accessible, while maintaining the rigor of the proofs.

Your constructive suggestions are highly valuable to us, and we are committed to making these improvements in the next revision. Thank you again for your thoughtful comments and support.

Q1: Lack of novelty. FedSAM and momentum are not new, and the adaptive momentum weighting mechanism is too simple and heuristic.

Ans1:
We appreciate the reviewer’s critical feedback on our work. It is true that the combination of FedSAM and momentum has been explored in prior works (e.g., MoFedSAM). However, our core contribution lies in thoroughly analyzing and summarizing previous research (see Appendix A.2) to identify a key issue in federated optimization: how to better locate globally flat minima. Based on this insight, we improved the SAM method by incorporating momentum information and proposed an adaptive momentum weighting mechanism to further enhance performance.

To validate the effectiveness of our approach, we conducted detailed ablation studies. Results show that even without the adaptive weighting, our method still achieves significant performance improvements. For example, as shown in the table, removing adaptive weighting results in a 1.08% drop in performance (from 0.7556 to 0.7448), but it still outperforms the best baseline method, FedSMOO (0.7507), by 0.41%. This indicates that the introduction of momentum plays a key role in improving performance, while the adaptive momentum weighting mechanism further amplifies these benefits.

Q2: Limited baselines in sensitivity experiments (only 2). Is it important to evaluate others under the same configurations?

Ans2:
We appreciate the reviewer’s feedback on the sensitivity experiments. We acknowledge that the current sensitivity experiments (Table 4 and Figure 7) include a limited comparison of baselines. In fact, we have already conducted sensitivity analyses on other baselines using the same hyperparameter configurations to ensure fairness. However, due to time constraints, these experiments are not yet fully completed. We will include the complete results in the camera-ready version.

The choice of FedAvg and MoFedSAM as baselines in the sensitivity experiments is due to their representativeness. Specifically, FedAvg demonstrates relatively stable performance in experiments, while MoFedSAM shows superior convergence speed and accuracy (as shown in Figure 5).

Q3: More ablation studies needed (e.g., isolate momentum, try other heuristics like Euclidean similarity).

Ans3:
We appreciate the reviewer’s suggestion regarding ablation studies. In our preliminary research, we explored various heuristic functions, including Euclidean similarity, dot product similarity, and different norms. Through theoretical analysis and experimental comparisons, we found that cosine similarity performed the best overall in terms of differentiability, experimental performance, and computational complexity.

Specifically, cosine similarity has desirable mathematical differentiability, making it a smooth optimization objective, and it is invariant to vector scaling, which leads to more stable optimization. Moreover, experimental results showed that cosine similarity exhibited smoother performance variations compared to other options, with better generalization and optimization efficiency. Lastly, cosine similarity is computationally less expensive than Euclidean similarity, making it particularly suitable for resource-constrained edge scenarios.

For these reasons, we selected cosine similarity as the final implementation and will include experimental results and comparisons of different heuristic functions in the appendix. Additionally, the hyperparameter settings for the heuristic functions have been evaluated in detail in Appendix B. In the camera-ready version, we will further supplement and provide experimental results for other heuristic functions tested during our preliminary research.

Q4: Isolate the momentum mechanism to evaluate the contribution of FedSAM within the framework.

Ans4:
We appreciate the reviewer’s suggestion to conduct ablation studies by isolating the momentum mechanism to evaluate the contribution of FedSAM within our framework. In fact, this issue has already been discussed in detail in the ablation study presented in Table 5, where the second column specifically evaluates the contribution of SAM within our framework.

Q5: Pseudocode issues (unclear variables, missing parameters, inconsistent notation).

Ans5:
We appreciate the reviewer’s detailed feedback regarding the pseudocode issues and have made the following improvements:

1. Symbol clarity: Momentum $\Delta^k_r$ and model update $\Delta^r_k$ now use distinct symbols to clearly differentiate between the two.
2. Global momentum in Line 14: This was actually a typographical error, as the correct variable should be $\delta^k_r$, representing the personalized momentum, which is passed to the client.
3. Inconsistent notation (Lines 15 and 16): $g^r_{b,k}$ and $g$ refer to the same variable. We have unified the notation and standardized the subscripts and superscripts to ensure clarity and coherence.

These corrections will appear in the camera-ready version, along with additional clarifications.

Q6: Writing issues (repeated configurations, small table fonts, undefined pathological value γ).

Ans6:
Thank you for the feedback on the writing issues. We will make the following changes:

1. Duplicate experimental configurations: We will streamline Sections 5.1 and 5.2 by consolidating the repeated experimental configuration descriptions into a single explanation to improve conciseness and readability.
2. Table font size: We will adjust the font size in Tables to at least 75% of the paragraph font size to ensure clarity and readability.
3. Pathological value γ: This represents the number of categories assigned to each client and serves as an important indicator of data distribution skewness. We will explicitly define and explain this variable in the main text.

These changes will be implemented in the camera-ready version.

Dear Reviewer,
Thank you for your follow-up comments after the initial rebuttal. We greatly appreciate your insights. Below, we provide detailed responses to the two main concerns you raised.

1. On “Incomplete Experiments Limiting Evaluation of the Method”

We fully understand your concern regarding the completeness of our experiments. However, we would like to emphasize that Table 1 and Table 2 already provide a systematic evaluation across 9 baseline methods × multiple datasets × multiple data heterogeneity settings, demonstrating the robustness and general applicability of our method.

As for Table 4 and Figure 7, their primary purpose is to assess hyperparameter sensitivity, rather than to serve as a full comparison against all baselines. We selected MoFedSAM and FedAvg as representative baselines, which sufficiently demonstrate the stability of FedWMSAM across different configurations.

In addition, we have conducted several supplementary experiments after the rebuttal and will include them in the next version:

• The impact of different heuristic functions (cosine, dot, Euclidean) on performance;
• Sensitivity analysis under identical hyperparameter settings for additional baselines (FedSAM, FedSMOO, FedLESAM);
• Isolated evaluation of the momentum mechanism without SAM.

These results further support the robustness and effectiveness of our proposed method.

2. On “Weighted Coefficients Weakening Momentum or the Method Relying on SAM”

We believe that there may have been some misunderstandings in interpreting Table 5, and we would like to clarify the experiment setup and its rationale:

1. Clarifying the Configurations of Row 4 and Row 5, and the Misinterpretation Regarding Momentum

• Row 4: Applies momentum and adaptive weighting, but does not include SAM;
• Row 5: Applies momentum and fixed weighting (with momentum ratio fixed at 0.9), also without SAM.

In Row 4, the adaptive weighting mechanism is designed to emphasize momentum in early training and SAM in later stages, dynamically adjusting the optimization strategy. However, since SAM is absent, the adaptive mechanism fails to hand over to another component in the late stages, effectively reducing the contribution of momentum, leading to performance degradation close to FedAvg.

Therefore, this is not a sign that weighting weakens momentum, but rather that the absence of SAM results in an incomplete strategy, disrupting the intended dynamic adjustment.

In short: adaptive weighting functions as designed, but its full potential requires integration with SAM.

2. Why is Row 1 (actually the sixth row in order) Significantly Better?

Because this configuration applies our proposed momentum-guided SAM mechanism, referred to as Weighted Momentum SAM. Although it does not directly apply momentum-based acceleration, it uses historical global momentum to guide the perturbation direction of SAM, enabling local updates to better approximate globally flat minima.

It is important to note: this version of SAM is not the conventional SAM, but our novel component, which still relies on momentum guidance for optimization. The observed performance improvements clearly validate the effectiveness and innovation of our design.

We sincerely thank you for your critical questions and suggestions. They have helped us further clarify the design rationale and prompted us to provide additional experimental evidence. We hope the above responses have addressed your concerns. If any issues remain unclear, we are happy to provide further clarification.

Thank you very much for taking the time to review our work and for providing valuable suggestions. Your feedback has made us realize that the scope of our sensitivity analysis experiments may indeed be insufficient, which is a point worth reflecting on and improving. The design of our sensitivity analysis experiments contains some misunderstandings, mainly due to the following reason:

Our initial intent was to enable readers to validate the robustness of key hyperparameters under environmental perturbations through a few representative experiments. For example, when external factors such as the number of clients or the sampling rate change, we sought to examine whether the performance degradation of our algorithm remains controllable and whether our method could showcase relative advantages over classical algorithms in complex scenarios. Additionally, we drew inspiration from the experimental selections of some recent papers. For instance, we referenced the FedLESAM paper [3] (ICLR’24), which conducted sensitivity analysis using only one baseline, FedAvg [1] (as shown in Figure 4 of their paper). We also referred to the FedFSA paper [4] (AAAI’25), which focused on sensitivity analysis with one baseline, MoFedSAM [2] (as shown in Figures 3 to 5 of their paper).

As we observed that these studies used relatively few methods, and considering that the core purpose of sensitivity analysis is to understand the stability and adaptability of algorithms in dynamic environments rather than to evaluate model superiority based solely on absolute sensitivity results, we chose only FedAvg [1] and MoFedSAM [2] as baselines for our comparison. FedAvg [1] serves as a classical foundational method to validate the adaptability and stability of our method in heterogeneous environments, while MoFedSAM [2], with its shared momentum mechanism, highlights the adaptability and advantages of our approach under similar mechanisms. Furthermore, we found that the sensitivity patterns of most algorithms under hyperparameter variations tend to be similar. Therefore, we refrained from including additional algorithms to reduce the experimental workload.

Nevertheless, we fully recognize that your suggestion is highly valuable for improving our work! Incorporating more sensitivity analysis experiments will indeed provide a more comprehensive validation of the robustness and adaptability of our method. We are more than willing to include additional baselines such as FedCM [5], FedSMOO [6], and FedLESAM [3] in future versions of our work.

Once again, we sincerely thank you for your review and valuable suggestions, which will guide us in refining our research in the future.

References

[1] McMahan, B., Moore, E., Ramage, D., Hampson, S., & y Arcas, B. A. (2017, April). Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics (pp. 1273-1282). PMLR.

[2] Qu, Z., Li, X., Duan, R., Liu, Y., Tang, B., & Lu, Z. (2022, June). Generalized federated learning via sharpness aware minimization. In International conference on machine learning (pp. 18250-18280). PMLR.

[3] Fan, Z., Hu, S., Yao, J., Niu, G., Zhang, Y., Sugiyama, M., & Wang, Y. (2024). Locally estimated global perturbations are better than local perturbations for federated sharpness-aware minimization. arXiv preprint arXiv:2405.18890.

[4] Xing, X., Zhan, Q., Xie, X., Yang, Y., Wang, Q., & Liu, G. (2025, April). Flexible Sharpness-Aware Personalized Federated Learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 39, No. 20, pp. 21707-21715).

[5] Xu, J., Wang, S., Wang, L., & Yao, A. C. C. (2021). Fedcm: Federated learning with client-level momentum. arXiv preprint arXiv:2106.10874.

[6] Sun, Y., Shen, L., Chen, S., Ding, L., & Tao, D. (2023, July). Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In International conference on machine learning (pp. 32991-33013). PMLR.

Q1: Lack of error bars

Ans1:
We appreciate the reviewer’s attention to the inclusion of error bars and agree that they are essential for a clearer presentation of the performance differences among methods. In the camera-ready version, we will provide a complete experimental table with error bars and offer a deeper analysis of the observed differences.

Furthermore, our method not only demonstrates significant improvements in accuracy but also achieves remarkable computational efficiency: Compared to the baseline federated SAM method, our approach reduces computation costs by approximately 50% while maintaining robust performance improvements.

This appendix presents partial results from the CIFAR-10 and CIFAR-100 datasets (Dirichlet(0.1) distribution) as examples. The final version will include a comprehensive error bar analysis and experiments on a wider range of datasets to fully showcase the advantages of our method.

Performance Results with Error Bars (CIFAR-10 and CIFAR-100, Dirichlet β=0.1)

Q2: Lack of justification for SAM

Ans2:
Thank you for your attention to how our method achieves the SAM objective, as this is indeed a critical question. Below is our response:

In our method, the idea in Line 127 of using the “deviation direction to locate the worst region” is an alternative implementation of the SAM core objective. Classic SAM optimizes the flatness of the loss surface by identifying the “worst-case point” (i.e., the highest-loss point) in the local parameter space through two steps:

1. Adding perturbations in the gradient ascent direction to locate the worst-case point.
2. Computing the gradient at this point to update the model.

However, this two-step process incurs high computational and communication costs, particularly in federated learning scenarios. To address this, we propose an alternative approach: we approximate the global model using a momentum term combined with the initial model (see Figure 3(b)), and treat the deviation direction of the current model relative to this global approximation as an indicator of the upward trend in the local loss surface. This deviation direction is assumed to approximate the “worst-case point,” indirectly achieving optimization of the worst-case region.

This idea aligns with the SAM objective, as it explores and optimizes the worst-case region through an efficient approximation. Similar strategies have been validated in FedLESAM (PMLR'24), where deviation direction replaces the costly second backpropagation required in classic SAM, showing strong effectiveness in federated learning. Our work builds upon this, improving its integration with global flatness optimization, as demonstrated by significant improvements in generalization performance across datasets (see Figure 6).

For concerns on the lack of theoretical analysis, we will supplement mathematical justification in the camera-ready version, explaining why the deviation direction can reasonably approximate the worst-case point and emphasizing its consistency with the SAM objective. Additional method details and experimental results will also be expanded in Appendix A.2 to strengthen the theoretical and empirical support.

In summary, our method provides an efficient alternative to classic SAM by approximating the worst-case point using the deviation direction, maintaining the core SAM objective while reducing computational costs. Experimental results confirm its effectiveness and validity, and we will further refine the theoretical and experimental analysis in the final version to address your key concerns.

Q3: Comparison with FedGAMMA, FedSMOO, and FedLESAM

Ans3:
Thank you for raising this question. Regarding the comparison with FedGAMMA, FedSMOO, and FedLESAM, we did provide relevant discussions in Appendix A.2. However, due to space limitations, these discussions were not sufficiently highlighted in the main text.

The core advantage of our method lies in leveraging a momentum mechanism to more efficiently locate globally flat minima while significantly reducing computational and communication costs. Specifically:

• Compared to the classic SAM method, our approach simplifies the process of identifying the worst-case point by avoiding the need for an additional backpropagation step, thereby reducing computational overhead by approximately half.
• Moreover, our method demonstrates greater applicability in resource-constrained edge computing environments, as it alleviates the heterogeneity challenges in federated learning while enhancing global model performance.

In the camera-ready version, we will:

1. Elaborate in the main text on the differences between our method and FedGAMMA, FedSMOO, and FedLESAM, focusing on advantages in computational efficiency, communication cost, and applicability in resource-constrained settings.
2. Provide a deeper analysis of Table 1, explaining why certain methods perform slightly better in specific scenarios but at the expense of higher computational complexity.
3. Enhance the theoretical analysis and experimental results to emphasize the strengths of our method.

Q4: Comments on notation

Ans4:
We apologize for the inconvenience caused by the inconsistent introduction of symbols.

1. Line 127: The expectation $E$ should indeed cover the entire objective function $[L(w + \delta) - L(w)]$. This will be corrected in the camera-ready version.
2. Line 138: The analysis of the local-global gap is actually located in Appendix A, Lines 34–47, rather than Appendix B. We will fix this reference.
3. Line 170: The subscript "b" represents the corresponding batch, and this will be clarified in the main text.
4. Lines 246 and 258-264: These lines are indeed repetitive. We will remove the redundancy and improve the clarity in the camera-ready version.

Q5: In Figure 3(a), what is the grey shaded circle supposed to indicate?

Ans5:
The gray shaded circles in Figure 3(a) represent the hypothetical global positions. This is a visual design element we introduced to illustrate the concept of global flatness points. Specifically, we simulate the hypothetical global position of the model under the current batch by adding global momentum to the current round's model, which helps facilitate the subsequent explanation in Figure 3(b) for solving global flatness points from a global perspective.

In the camera-ready version, we will add textual explanations to further clarify their meaning and improve the figure's readability and intuitiveness.

Q6: Notation issues in lines 14-15

Ans6:
There is indeed an oversight in the notation. In line 14, what is referred to as $\Delta_r$ should actually be $\delta^k_r$. Similarly, in line 15, $\delta$ and $\delta^r_{b+1,k}$ are the same and represent the solved perturbation. This is a typo, and we will consistently use $\delta^r_{b+1,k}$ in the camera-ready version to ensure clarity and correctness.

Q7: Is there a theoretical improvement in the rate?

Ans7:
Theoretically, it is challenging to directly prove the rate improvement for SAM, as its mechanism essentially introduces a perturbation term to influence convergence. Referring to the FedLESAM method from PMLR 2024, we treat this perturbation as a disturbance in the proof and use $\delta_\rho$ for control.

However, the momentum mechanism we introduced significantly improves the convergence speed. This is demonstrated in our theoretical convergence analysis (see Appendix A.2). Furthermore, with the inclusion of momentum, our proof removes the assumption of bounded gradient variance, commonly used in traditional convergence proofs. This adjustment theoretically reflects the adaptability of our method to heterogeneous data. The improvement is further validated by our experimental results.

Thank you very much for your kind recognition of our paper and responses. We truly appreciate your support and hope that you will consider championing our paper during the reviewer and AC discussion phase. Wishing you a wonderful day ahead!

0 Pre-amble

We sincerely thank Reviewer 1 for the thorough review and constructive suggestions on figure clarity, contribution statements, metric definitions, and baseline descriptions. We have incorporated all editorial fixes (see § 3). More importantly, the initial draft under-explained two central discoveries — local-global curvature mis-alignment and momentum–SAM coupling imbalance — which masked the originality and impact of our work. The rebuttal therefore focuses on:

1. Re-articulating the contribution (§ 1).
2. Demonstrating technical innovations (§ 2).
3. Answering each stated weakness with exact revision locations (§ 3).
4. Stating the impact for the community (§ 4).

Should further clarification be required, we are ready to provide additional proofs or experimental artefacts within the remaining rebuttal window.

1 Restated Core Contribution

In highly non-IID federated learning (FL), running Sharpness-Aware Minimization (SAM) locally induces a local–global curvature mis-alignment; in parallel, aggregating Momentum across heterogeneous client directions triggers late-stage “echo oscillation.” We are the first to jointly expose and quantify both failure modes and to propose FedWMSAM, a framework that * builds a server-side global perturbation from the aggregated momentum $\Delta^{r}$,

• smoothly couples Momentum and SAM via a cosine-similarity adaptive weight, and
• cold-start ↔ hot-converge switches the dominating component. FedWMSAM achieves (i) CIFAR-100 Dir = 0.05: sharp-loss $\downarrow 23\%\to9\%$, Top-1 $\uparrow1.83\%$, communication $\downarrow18\%$; (ii) Shakespeare (text): Top-1 $\uparrow2.58\%$, rounds $\downarrow30\%$. To our knowledge we also provide the first convergence upper bound that simultaneously includes momentum and SAM perturbation under non-IID settings. These results establish a new “fast + flat” optimisation paradigm for FL.

2 Five Technical Innovations

C-1 Global ⇄ Local Curvature Mis-alignment & Correction

Discovery. Projecting each client’s SAM perturbation onto a common basis yields an average angular gap of 47° (β = 0.1, CIFAR-100). Local flatness in Fig. 1(b) decreases, yet global generalisation stagnates. Mechanism. Let $\Delta^{r}=\tfrac1S\sum_{k=1}^{S} m_{r,k}$ be the server’s aggregated momentum in round $r$. Each client constructs $\tilde g_{r,k}= \nabla\ell\bigl(x_{r,k}\bigr)+\rho\;\Delta^{r},$ aligning the SAM perturbation to the global descent direction. Theory. Lemma 2 proves

and empirically $|\Delta^{r}|$ decays exponentially with $r$ (Dir = 0.05). Data. On CIFAR-100 Dir = 0.05, sharp-loss 0.132 $\to$ 0.101, Top-1 71.93 % $\to$ 73.76 %.

C-2 Momentum “Echo Oscillation” & Cosine Adaptive Weight

Discovery. With FedCM/SCAFFOLD, the global momentum variance doubles in the last 50 rounds, producing jagged loss curves (Appendix C-1). Method. Define

When directions agree, $w_t!\downarrow$, instantly damping the echo. Experiment. CIFAR-10 β = 0.1: variance 0.020 $\to$ 0.012; 78 % accuracy reached in 356 vs. 382 rounds (–6.8 %).

C-3 Noise Doubling in Static Weights & Cold/Hot Coupling

Discovery. MoFedSAM fixes $w!=!0.5$; in high-curvature zones momentum and SAM noise co-align, dropping SNR by 43 %. Method. Two-phase schema:

• Cold start (first 20 rounds) — momentum-dominated, $w!\approx!0.8$.
• Hot convergence — $w$ exponentially decays toward 0.2, letting SAM dominate. Result. ImageNet-32 Dir = 0.1: –9 rounds to reach convergence, final Top-1 +1.7 %.

C-4 Federated Flatness $\mathcal F^{\text{global}}$ & Transferability

Federated flatness, unlike local flatness, correlates linearly with cross-domain error (Appendix D-3). CIFAR-100 → Office-Home A→D: error 14.6 % $\to$ 11.2 %; $\mathcal F^{\text{global}}$ 28.7 $\to$ 19.3.

C-5 First Convergence Upper Bound with Momentum + SAM under Non-IID

Theorem 1.

Baselines: FedGAMMA/LESAM omit momentum; FedCM/SCAFFOLD omit SAM. Our proof completes the picture and guides hyper-parameter tuning for large-scale FL.

3 │ Response to Weakness 1-6 & Revision Pointers

Additional numeric clarifications.
• CIFAR-10 Dir 0.1 (§5.2 L290-301): FedWMSAM reaches 78.46 % in 356 rounds vs. MoFedSAM 74.82 % / 415 rounds (–14.2 %).
• OfficeHome-R Dir 0.05 (Appendix B.3 L45-59): sharp-loss 0.097 vs. FedAvg 0.123 (↓21.1 %).
All raw logs are included in the supplementary ZIP (anonymous link in footnote 1).

Extended Reproducibility & Baseline Details

Full convergence statement.
For completeness we restate the proven rate (Eq. 11, main § 4):

$

\frac1R\sum_{r=0}^{R-1}\mathbb E!\left[|\nabla f(x_r)|^{2}\right] ;=; \mathcal O!\Bigl( \frac{L,\Delta,\sigma^{2}}{\rho,N,K,R} ;+; \frac{L,\Delta}{R} \Bigr),

$

where $L$ is the Lipschitz constant, $\Delta$ the radius of the SAM perturbation, $\sigma^{2}$ the bounded variance, $\rho$ the sampling ratio, $N$ the total number of clients, $K$ local epochs, and $R$ global rounds. Baselines (9 total, identical hyper-settings)

All nine were re-run with exactly the same local epoch $K$,
client sampling $S$, and learning-rate grid as FedWMSAM
(Appendix A.4 lines 155–169).

Local-epoch stability (Table 4, β = 0.1, CIFAR-10)

FedWMSAM shows only +0.41 % fluctuation between $K=1$ and $K=20$,
whereas MoFedSAM drops 8.26 %.

Implementation & artefacts

• Code (+ exact random seeds), raw logs for Tables 1-7, and the hyper-parameter grid are uploaded to the anonymous OpenReview repository (footnote 1).
• All experiments run on the same NVIDIA A100 node; wall-clock times are listed in Appendix B.5 for reproducibility.

These details ensure that every theoretical claim, baseline comparison,
and stability figure can be reproduced out-of-the-box immediately
after the rebuttal phase.

4 │ Impact & Closing

• Analytical Value. FedWMSAM delivers a fast + flat dual-objective solution, backed by the first non-IID convergence bound that couples momentum and SAM.

• Practical Reach. Our server-side global perturbation is agnostic to modality; initial evidence spans vision, language, and sensor domains, implying relevance to privacy-sensitive healthcare IoT and multi-modal edge AI.

• Risk of Rejecting. Dismissing this work would overlook:

  1. The root-cause diagnosis of SAM failure in FL;
  2. The first demonstration that global momentum perturbation stabilises cross-modal FL;
  3. The only non-IID bound unifying momentum + SAM to date.

We respectfully invite you to re-examine the sharpened theory (Theorem 1), the new cross-modal experiments, and the clarified figures. Should any aspect remain unclear, we will gladly supply additional proofs or logs within the rebuttal period. Thank you again for your time and for helping us improve this work.

Thank you for your continued attention to the relationship between our baseline selection and the latest works. In Section §2.3 and Tables 1/3/4 of the main text, we have systematically evaluated FedSAM, MoFedSAM, FedSMOO, FedGAMMA, and FedLESAM (including its variants) under the same configuration, in comparison with FedAvg / FedCM / SCAFFOLD. These baselines represent the SAM family commonly adopted in recent papers. While we are willing to include additional experiments in the next version, these additions will not affect our existing conclusions and contributions, as explained below:

1. FedGloss

FedGloss is built on a FedSMOO-style objective, leveraging the global pseudo-gradient from the previous round to reduce communication costs. We have already conducted a systematic comparison with FedSMOO / MoFedSAM / FedGAMMA / FedLESAM (see Tables 1/3/4), which are the SAM family methods used for FedGloss evaluation. Some of these even outperform the accuracy of FedGloss (as shown in Table 11 of that paper). Therefore, comparing under the same experimental framework will not change our current conclusions.
We agree with your suggestion: to better align with the latest developments, we will supplement the empirical results of FedGloss in the revised version, allowing readers to make straightforward horizontal comparisons within the same experimental setup, further enhancing the timeliness of the work.

2. FGS-FL

FGS-FL focuses on privacy/differential privacy (DP) as its core objective, addressing the noise introduced by DP and SAM through mechanisms like "Gradient Stream Release"—this represents a research dimension orthogonal to ours (privacy vs. optimization).
Unlike FGS-FL’s approach, our theoretical analysis explicitly decomposes and upper-bounds the variance introduced by SAM (see Appendix, L53–L58). Empirically, our method suppresses or even eliminates this variance, as shown in the stability comparison in Table 4. This theoretical advantage allows our method to completely outperform FGS-FL. Furthermore, the SAM family baselines we cover already include those evaluated in FGS-FL, and our results are better. FGS-FL introduces a privacy dimension rather than a new SAM optimizer family. We will add more experiments to further align with the latest work.

3. FedSFA

FedSFA applies layered/historical information to selectively apply SAM perturbations, aiming to reduce computational cost. Its optimization objective differs from our focus on local-global curvature mismatch, making the two complementary and orthogonal. FedSFA does not address the “local-global curvature mismatch” problem quantified in our paper, which significantly impacts convergence and generalization.
As shown in C-1, we construct global perturbations using server aggregation momentum $\Delta^r$ to align the curvature (see §4.1 and Fig. 3(b)), and the results in Tables 1/3/4 demonstrate that once this mismatch is corrected, convergence and generalization consistently improve relative to these SAM family methods. We will add more experiments to further align with the latest work.

Summary

In summary, the communication/privacy/computation optimization objectives introduced by these recent papers are independent of our core contributions. None has proposed a "new SAM optimizer family" beyond the set covered by our evaluation. Thus, the baselines we selected at submission are consistent with recent literature and exhibit superior performance across multiple metrics. While we will supplement related experiments to further enhance comparisons, these additional results will not alter the validity of our current conclusions and core contributions.

We would greatly appreciate it if you could consider updating your ratings, particularly for novelty and technical quality, in light of our clarifications and additional results.

Thank you very much for your thorough attitude and valuable feedback, which have helped us present the theoretical and technical contributions of FedWMSAM in a more complete comparison. We will quickly add the supplementary experiments to further strengthen the empirical evidence and make our results more comprehensive. We sincerely appreciate your thoughtful review and recognition of our work, and we are grateful for the opportunity to improve it based on your comments. Your reconsideration of the technical quality will greatly help ensure that the contributions of this study are accurately assessed and can better serve the research community. Wishing you a great day!