FedWMSAM: Fast and Flat Federated Learning via Weighted Momentum and Sharpness-Aware Minimization
Authors: Tianle Li
Yongzhi Huang
Linshan Jiang
Chang Liu
Qipeng Xie
Wenfeng Du
Lu Wang
Kaishun Wu
Venue: NeurIPS 2025 poster ()
Keywords: Federated Learning
Model generalization
Sharpness-Aware Minimization
Momentum
Primary area: deep_learning
TL;DR: We propose FedWMSAM, a novel method that combines personalized momentum and dynamic SAM to achieve faster convergence and flatter minima in federated learning, effectively addressing client drift and data heterogeneity.

Abstract: In federated learning (FL), models must \emph{converge quickly} under tight communication budgets while \emph{generalizing} across non-IID client distributions. These twin requirements have naturally led to two widely used techniques: client/server \emph{momentum} to accelerate progress, and \emph{sharpness-aware minimization} (SAM) to prefer flat solutions. However, simply combining momentum and SAM leaves two structural issues unresolved in non-IID FL. We identify and formalize two failure modes: \emph{local–global curvature misalignment} (local SAM directions need not reflect the global loss geometry) and \emph{momentum-echo oscillation}  (late-stage instability caused by accumulated momentum). To our knowledge, these failure modes have not been jointly articulated and addressed in the FL literature.

We propose \textbf{FedWMSAM} to address both failure modes. First, we construct a momentum-guided global perturbation from server-aggregated momentum to align clients' SAM directions with the global descent geometry, enabling a \emph{single-backprop} SAM approximation that preserves efficiency. Second, we couple momentum and SAM via a cosine-similarity adaptive rule, yielding an early-momentum, late-SAM two-phase training schedule. We provide a non-IID convergence bound that \emph{explicitly models the perturbation-induced variance} $\sigma_\rho^2=\sigma^2+(L\rho)^2$ and its dependence on $(S,K,R,N)$ on the theory side. We conduct extensive experiments on multiple datasets and model architectures, and the results validate the effectiveness, adaptability, and robustness of our method, demonstrating its superiority in addressing the optimization challenges of Federated Learning. Our code is available at \url{https://github.com/Li-Tian-Le/NeurlPS_FedWMSAM}.

FedWMSAM: Fast and Flat Federated Learning via Weighted Momentum and Sharpness-Aware Minimization
Tianle Li
Yongzhi Huang"
College of Computer Science and
Data Science Analysis The Hong Kong University of Science and Technology (Guangzhou) huangyongzhi wewponngghiemmii...mmmmm Cemail. szu, edu. cn
Software Engineering
Shenzhen University 22002710150emal szu, edu cn
Linshan Jiang' National University of Singapore
Chang Liu Nanyang Technological University, Singapore 1iuc00630e. nnt.nnuunnn edu.
Qipeng Xie The Hong Kong University of Science and Technology (Guangzhou) qrieaf@connect. ust
linshan@nus. edu.sg sg
Wenfeng Du
Lu Wang Shenzhen University wanglu@szu. edu. en
Kaishun Wu
Shenzhen University duvf@szu. edu. cn
The Hong Kong University of Science and Technology (Guangzhou) wuksshkust-gz. edu. cn
Abstract
In federated learning (FL), models must converge quickly under tight commu- nication budgets while generalizing across non-lID client distributions. These twin requirements have naturally led to two widely used techniques: client/server momentum to accelerate progress, and sharpness-awar minimization (SAM) to prefer flat solutions. However, simply combining momentum and SAM leaves two structural issues unresolved in non- FL. We identify and formalize two failure modes: local-global curvature misalignment (local SAM directions need not reflect the global loss geometry) and momentum-echo oscillation (late-stage instability caused by accumulated momentum). To knowledge, these failure modes have not been jointly articulated and addressed the FL. literature.
We propose Fed WMSAM to address both failure modes. First, we construct momentum-guidee global perturbation from aagrre-gggee momentum align clients' SAM directions with the global descent geometry, enabling a single- buckprop SAM approximation that preserves efficiency. Second, we couple momen- tum and SAM via cosine-similarity adaptive rule, yielding early- moryymmmmntuu, late-SAM two-phase training schedule. We provide a non-lID convergence bound that explicitly models the serturbation-innuuee variance oi o² + (Lp)² and its dependence on (S, K, RR N)) on the theory side. We conduct extensive experiments on multiple datasets and model architectures, and the results validate the effec- tiveness, adaptability, and robustness of our method, demonstrating its superiority in addressing the optimization challenges of Federated Learning. Our code available https: / //t/htgtttt..t///// com/Li -iin-ian--- Lee-Le//Neeee /NuurPP FPP__MSWSSA....
'Equal contribution.
Co-corresponding author. Maan corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Introduction
Federated Learning (FL) [1] has emerged as a promising distributed learning paradigm that enables multiple clients to collaboratively train shared global model while keeping their local data decen- tralized, thus preserving privacy. In particular, under the edge computing setting, FL has shown great potential in a wide range of real-world applications, including personal mobile sensing systems [2], healthcare data analytics [3, 4], and industrial Internet of Things scenarios [5-7].
FL in edge computing has several unique attributes. To reduce communication costs, the number of local iterations must increase. Concurrently, partial participation is often used due the unstable connectivity of IoT devices. These two factors exacerbate the harmful effects of data heterogeneity [8-
10] in real-world applications, intensifying client drift [11], where local updates deviate significantly from the global objective, degrading model performance and generalization.
Local multi-round training increases the drift caused by data heterogeneity, while partial participation further magnifies its effects. Based on where the drift originates, current methods for mitigating client drift can be categorized into the following levels: Methods at the data level attempt to balance data distribution through techniques such as data augmentation [12, 13] or resampling [8, 14], but these methods may incur high computational costs and risk overfitting. Gradient-level approaches, including proximal methods [9], gradient correction [11], and regularization terms [15], aim to adjust local gradients, but gradient distortion may hinder convergence. Techniques for model aggregation [16-18] alter the server's aggregation strategy, requiring the collection of additional information that may conflict with the privacy principles central to FL. At the same time, cryptographie PPFL (e.g., HE) offers an alternative at the cost of efficiency [19, 203.
Moreover, on the one hand, these methods focus solely on addressing client drift, with little consider- ation for convergence speed, while fewer training rounds could enhance real-world applicability. On the other hand, when heterogeneity is high, the loss surface resulting from aggregating models trained with Empirical Risk Minimization (ERM) becomes sharp, limiting the models' generalization in practical applications. Moreover, in non-IlD FL, we observe two structural failure modes when simply porting SAM and/or momentum: (i) local-global curvature misalignment- local SAM directions need not reflect the global loss geometry; (ii) comentuun-eho oscillation paatioon-lte-ssee late- stage instability caused by accumulated momentum. These motivate a design that can both align local updates to the global geomelry and adaptively damp momentum as training progresses.
To solve this problem, we aim to design a fast and flat FL. algorithm. Two mainstream lines address the twin goals in FL: momentum [21-25] for speed and SAM [26-28] for flatness. However, under non-Ill) data, each line has structural drawbacks: momentum can amplify late-stage instabilit/overrittingg [29], while SAM requires extra backward pass and, more importantly, its local perturbation directions need not reflect the global loss geometry (the local-global curvature misalignment diagnosed above). Momentum-onle designs tailored to long-tailed heterogeneity (e.g., FedWCM [24, 25)) improve robustness but do not enforce global flatness.
A straightforward combination Mofids inserts SAM into FedCM [26, 22]-yet naively plugging SAM into a momentum pipeline leaves the two failure modes unresolved in non-lll settings (misalignment persists, late-stage momentum oscillation is not damped). This motivates a design that both aligns local updates to the global geometry and adaptively damps momentum as training progresses, we instantiate this next FedWMSAM.
To address the aforementioned issues and achieve fast and flat training in FL, we propose Fed SAM (Federated Learning with Weighted Momentum-SAne SAM),, which integrates momentum with SAM in a principled way.
Firstly, we introduce personalized momentum and use the server-aggregatte momentum as a global geometric carrier to build momentum- guided global perturbation, aligning local SAM directions with the global descent geometry; the perturbation is implemented with single buckpropagation (no extra backward pass).
Secondly, we dynamically adjust the perturbation along this global direction during local steps (e.g., 2i = x, + b A,), enabling each client to explore globally flatter regions without increasing pou-round cost.

---

Thirdly, we couple momentum and SAM via cosine-similarity adaptive weight ay. which yields an early-momentun late-SAM two-phase schedule-speedingg up early progress while damping late-stage momentum oscillations under non-lll) data.
Our main contributions are summarized as follows:
Mechanism. We identify and formalize two failure modes in non-lII FL locallocll--- global curvature misalignment and momentum- echo oscillation and correct them via catneeentunggiideee global
perturbation and cosine-adaptive coupling.
Method & Efficiency. FedWMSAM aligns local SAM directions using server-aggregated momentum and implements the perturbation with single backpropagation, yielding an early- momentum Mate-SAM two-phase schedule and fast-then-flat training with near-FedAvg per-round
cost.
Theory. We provide a non-lID convergence bound that explicitly models the perturbationiinducee variance of o² + (Lp)² and its dependence on (S, R, N): Ö( LA SKR
LA / N2/2
+
Empirics. Across three datasets and twelve heterogeneity settings, FedWMSAM is best or on par in most cases and shows the largest gains under strong non-IID settings, reaching target accuracies in fewer rounds at similar per-round time, and the trends match our mechanism and theory.
2 Related Work
2.1 Heterogeneous Federated Learning
Federated learning (FL) often encounters the challenge of client drift. Existing solutions can be divided into three levels: data, gradient, and aggregation. At the data level, resampling strategies [14]
adjust local sampling probabilities to balance class distributions, generative models like GANs [12]
and VAEs [13] synthesize balanced datasets, and data- data-sharing approaches [8] distribute small portions of global data clients. At the gradient level, methods such as SCAFFOLD [11] reduce variance using control variates, FedDyn [15] applies dynamic regularization, FedProx [9] stabilizes updates
with proximal terms, and FedCM [22] aligns updates via momentum correction. At the aggregation level, Hierarchical FL [17] applies multi-level aggregation for large-scale networks, Clustered FL [16] groups clients with similar data distributions for localized optimization, and works like Client Sclection [30] and Client Weighting [18] adjust the influence of clients based on their contribution to the global model. However, these works focus solely on addressing heterogeneity data, neglecting speed and model generalization.
2.2 Momentum- Federated Learning
Utilizing historical gradient information via momentum has proven effective for accelerating conver- gence and handling data heterogeneity in FL. The fundamental idea is to constrain current updates by past directions, smoothing out oscillations. MIME [21] and FedCM [22] compute global momen- tum on the server and distribute it for stricter consistency, while AdaBest [31], FedADC [32]. and ComFed [33] adaptively calculate local momentum on clients and synchronize each round. Methods like MFL [34] and FedMIM [23] entirely apply momentum on-device to reduce communication. In parallel, momentum- based designs specifically tailored for long-tailed non-IID heterogeneity, such as FedWCM [24, 25], provide an efficient complementary approach. Although momentum mechanisms significantly improve early-stage convergence in non-IID scenarios, they can sometimes hinder late-stage fine-tuning, causing notable performance fluctuations.
2.3 SAM-Based Federated Learning
Model generalization is closely tied to finding flatter regions of the loss surface, motivating Sharpness- Aware Minimization (SAM) [35]. SAM actively seeks flatter optima and reduces overfitting risks by perturbing the model around local minima. In the federated setting, works like FedSAM [26], MoFedSA. [26], and FedGAMMA [27] plug SAM optimizers into FedAvg, FedCM, or SCAFFOLD but do not explicitly refine the global flatness search. FedSMOO [28] integrates FedDyn regularization to minimize local-global bias, whereas FedLESAM [36] estimates global perturbations based on local-server discrepancies. More recently, FedGloss [37] extends FedSMOO by leveraging the global aseudo-gradien from the previous round to reduce communication costs, while FedSFA [38] selectively applies SAM perturbations using historical information to lower computational cost. While these approaches enhance generalization, exploring flat minima inevitably slows convergence, and SAM's two backward passes increase computational overhead, further complicating its practical deployment in heterogeneous FL.
3

---

Preliminaries
3.1 Federated Learning
Federated Learning (FL) [1] enables multiple clients to train a global model collaboratively while preserving data privacy. The objective is min, F(w) = Σ "s Fl(w), where F&(ww) is the local loss client ny is client k's data size, and Σ K
Ek=1 T2.. Parameter w denotes the global model.
3.2 Momentum Method
A standard form is: Vi gh + A where vi is the momentum at client A in round r, a is the momentum factor, 9 the local gradient, and A' is the global momentum. Methods like MIME and FedCM employ this scheme to coordinate and stabilize local optimization under non-IID conditions.
3.3 SAM Method
Sharpness-Awarg Minimization (SAM) [35] improves generalization by identifying flatter regions of the loss function. Its objective is: min,, [w) = min.. maxi [E[ + - L L(w))]
Implementation typically involves two steps: (1) compute = P IVF (w) and (2) update model
F((w)
parameters based on the perturbed gradient L(w + 8). This forces the model to converge to flatter minima by explicitly considering the worst-case local perturbation.
3.4 Motivation
The core challenge of applying SAM in federated learning lies in fundamental mismatch: SAM computes perturbations based on local data, but aims to find flat minima in the global landscape (Figure 1). Due data heterogeneity, these local perturbations often fail to accurately reflect the global geometry, thereby limiting the effectiveness of SAM in federated settings.
To resolve this contradiction, existing methods attempt to either reduce the discrepancy between local and global models or improve the estimation of global- aware perturbations using local information. However, as analyzed in Appendix A, these approaches still fall short in fully bridging the local- global gap, motivating the need for a more unified and efficient solution. This motivates us propose new approach that more effectively integrates SAM with federated optimization.
4
Method
This section presents our WMMMM algorithm, which fuses personalized momentum and SAM for federated optimization. Section 4.1 highlights our methodology and illustrates the three key components (personalized momentum, global perturbation estimation, and dynamic weighting). Section 4.2 then details the implementation, including pseudocode, the computation of personalized momentum A and dynamic weight ar-
A/11 Local Flat Area
Global Flat Area
xivx
4.1 Methodology
Figure 1: The core of SAMs.
(a) Personalized momentum for local-global discrep-
ancy. We use the diagram Figure 2 (a) to illustrate the concept of personalized momentum. In FedCM [22], momentum utilizes the previous gradients to guide the next round of local training, which can effectively accelerate model convergence, as the black dashed line shows. However, it is worth noting that although each client has different data distributions, they share the same momentum. Although momentum mitigates part of the bias introduced by local gradients, the drift caused by data heterogeneity recurs in every communication round and cannot be fully corrected by momentum alone. Based on this, we introduce correction term e from SCAFFOLD to estimate the bias caused by local data using historical experience. Unlike the original SCAFFOLD [11], which requires uploading local correction terms in each round for global averaging and redistribution, our approach only requires the server to compute the correction term based on the differences in gradients uploaded by clients. This correction term of in the green dashed line is then aggregated with the momentum A, to form personalized momentum A as shown in the red dashed line, effectively saving communication bandwidth.

---

Client Drift
Worst Point Global
v== x, 25
6b
Global Flat Area
6,6
I Global Flat Area
gf
pa
b
g
4
Greater-SAM
x
Vp
g 1
x
4
+bA,
Greater Momentoin
(a) Personalize Momentum
(b) Momentum-bantt based Perturbation
(c) Dynamic Weight Adjustment
Figure 2:
FedWMSAM idea: (a) personalized momentum reduces local-global discrepancy; (b) local
model vs, momentum based model difference guides global perturbation estimation; (c) dynamic weighting adjusts momentum v.s. SAM based on gradient-momentum similarity.
(b) Global perturbation estimated from momentum. Performing true SAM in a federated setting challenging because clients have access only to their local data. Yet they must collaboratively explore the flat region of the global loss landscape. More information about the global loss function necessary for local SAM. observe that it is the momentum that carries information about the direction in which the global model is moving. As shown in Figure 2(b), by adding the momentum A, to the received model x", we can infer the position of the global model fy where the subscript denotes the current batch. It is important to note that the momentum here personalized, and the correction term helps eliminate the intrinsic perturbation drift caused by local data distributions.
At each step, we calculate the difference between the current model and the inferred global model, using V in light blue dashed line as the perturbation direction, based on the analysis that the direction of deviation from the global model indicates regions of higher loss. By doing this, we eliminate need for an additional backpropagation step to compute the perturbation while improving the estimation of the global perturbation.
Similarity. Although momentum accelerates early convergence and SAM improves final accuracy, their roles vary throughout the training process. An overly significant momentum can hinder late-stage fine-tuning, while pure SAM is relatively slow, as shown in Figure Inspired by experiments conducted by Andriushchenko al. [39], which show that switching from ERM to SAM at different epochs can lead to varying test errors, we realized that it is crucial to determine the optimal time increase the weight of SAM.
(c) Dynamic Weighting via
We observe that the cosine similarity between clients' di- rections and the global signal increases early and stabilizes increase ar monotonically late-SAM. momentum during the early
0.3
later (see Appendix B), with the similarity-yieldee This allows us rely stages and gradually weaken its influence to better explore the global flat region in the final stages. Figure (c) shows how momentum is helpful to speed up initially in the pur- ple line, then gradually yields to SAM for robust final convergence in the green line.
Performance
Convargence
Slower Cusvergenze
100
193
4.2 Proposed Algorithm
Figure 3: Momentum SAM in FL.
In this section, describe the components and procedures of the FedWMSAM algorithm. First, provide an overview of the method, followed by detailed explanations of each component. Algorithm summarizes the overall procedure.
4.2.1 Algorithm Overview
At each communication round r, the server selects subset of clients P, and computes each client's personalized momentum. The server then broadcasts the global model x,, the personalized mo- mentum A and the momentum factor a, to the selected clients. Each client updates model using the global momentum and performs local updates with SAM perturbation. The momentum
5

---

factor is udapted based on cosine similarity, and personalized corrections are refined following the SCAFFOLD-lik method reduce gocal-global drift. Next, we explain each of these steps in detail.
4.2.2 Personalized Momentum Calculation
At each round the server computes the personalized momentum for each selected client & by:
A A.
1 Olyy Ck:
where A, is the global momentum, a, is the momentum factor, and e₁ is the local correction for client k. This step ensures the alignment of the global and local updates. The coefficient - arises because, during the local updates, the momentum term combined with the gradient as:
2 V1+1A = Orr Ri.k + a.) A.
(2)
We want the correction term Cg to affect the gradient fh,, directly, so the coefficient I-a, is used to scale Cj in the global momentum computation. This adjustment ensures that has the same effect on the gradient as it would in the local update, enabling a consistent blend of global and local dynamics in the personalized momentum. This transformation decouples the correction term from momentum, reducing the need to transmit separate vectors for momentum and correction and saving bandwidth.
4.2.3 Adaptive Momentum Factor
After each round, we update the momentum factor 0++1 based on the cosine similarity between the global momentum and each client's personalized momentum. For efficiency, we compute the similar- ity using sim Ar, Ar), which we empirically find to be good proxy for the muudient-mmenntee similarity without incurring extra backpropagation (see Appendix B). The updated momentum factor is:
1
sim( Ar, A/),
Pr
KEP
and the final update is:
0, (1 A) ar + X min max (a,+1; 0.1), 0.9).
(4)
where A controls the speed of adaptation. The choice of the bounds for 4,++1. specifically within the range [0.1, 0.9), is motivated by several factors. Based on the analysis in [22]. momentum factor of 0.1 yields the best performance. This setting ensures that the momentum gradually decreases over time, allowing the SAM perturbation to become more prominent in later rounds. Therefore, the upper bound of 0.9 ensures that momentum remains sufficiently large, maintaining its influence on the calculation of SAM, which relies on the momentum term. The lower bound 0.1 prevents the momentum weight 1 - from vanishing too early, while the upper bound 0.9 avoids over-reliance on momentum so that SAM can dominate in later rounds. More discussion of this choice of the value can be found in Appendix B.
4.2.4 Correction Term Updates
Finally, the server updates the personalized correction terms O and c, using SCAFFOLD-inspee strategy [1]]. These updates help reduce local-global drift and refine the correction offsets for each client and are calculated as:
1 [Pr
+ -k ch
Ak. =c cg +
BB Ak.
(ck
ck)-
(5)
kEP.
4.3 Convergence Analysis
We provide non-IllD convergence guarantee for FedWMSAM, with rate of Ö( T. (2/(S)))
LA 1
N2/S S
Our analysis builds upon the SCAFFOLD-M framework [40], and explicitly
R
accounts for the perturbation-indueed variance o3 = o² + (Lp)², with two key extensions: an adaptive personalized momentum term and the integration of SAM. Notably, we show that the variance induced by SAM is bounded by the perturbation strength p. The complete proof and comparisons with related convergence rates presented in Appendix A.

---

Algorithm FedWMSAM
Require: Initial model To. global momentum Do. correctors cy. Ck. momentum factor ao 0.1,
learning rates 7h ng. perturbation magnitude p, communication rounds R, local iters B
for
R do for each client & E P. in parallel do
CLIENTUPDATE( A a,):
2:
12: To = I,
3:
Compute A using Eq. (1)
13: for b= 0 to B - -o
4:
DT. = CLIENTUPDAT(CE,, a,)
14:
1 (I, bb Th
5:
end for
15:
96 VL(xf + +p0 11 11,+1
6:
A WN PP. ENEP, 4
16:
*++1. 1. = ar Hi. + (1 -ar) A
7: 2,41 = Z, 17 Ar
17:
1b+1 = The ? 1/
8: Compute 0x++1 using Eq. (3) and (4)
18:
end for
9: Update Cy. Ck using Eq. (5)
19:
A BB Tr
10: end for
Ensure: Return Client update A
Ensure: Return updated model ZR
5 Experiment
5.1 Experimental Setups
We propose FedWMSAM and compare it with existing SOTA federated SAM methods, including FedSAM [26]. MoFedSAM [26], FedGAMMA [27], and FedSMOO [28], as well as FedLESAM [36] and its two variants, FedLESAM-S and FedLESAM-D, Since our method incorporates momentum and correction terms, we also compare it with classical federated optimization baselines such as FedAvg [1], FedCM [22], and SCAFFOLD [11]. By default, we set Pk.c Dir(3), where Pk.c denotes the class distribution of client h over class c, and / = 0.1. The main experiments are conducted with 100 clients, 10% participation per round, batch size of 50, a local learning rate Y = 0.1, global learning rate 72 1, and five local epochs, running for 500 communication rounds. For Fashion-MNIST [.1], we use Multi-Layer Perceptron (MLP) architecture. For CIFAR- [42], we use 1 1 [43] the backbone, ResNet-34 [43] for CIFAR-100 [42], and ResNet-50 [43] for OfficeHome. Each domain in OfficeHome divided into one client with 10% data sample rate and 100% active ratio. The perturbation magnitude P is set to 0.01 for FedSAM, FedGAMMA, FedLESAM, and our FedWMSAM, while FedSMOO and MoFedSAM use = 0.1 by default. Additional experimental settings are detailed in the corresponding figures and tables. All experiments were implemented in Py Torch and conducted on a workstation with four NVIDIA GeForce RTX 3090 GPUs
5.2 Overall Performance Evaluation
Table 1: Performance comparison of SOTA methods under Dirichlet and Pathological splits after 500 Rounds in different datasets.
Method
Fashion MNIST
CIFAR-10
CIEAR-100
AParition
Dirichlet
Pathological
Dirieblet
Pathological
Dirichlet
Pathological
ACcefficient
0.6
8 0.1
7
0.6
8 0.1
y=
7 3
, 0.6
0.1
7 20
7 10
RelAvg
0.8684
0.8226
0.8635
0.8150
0.7886
0.7005
0.7873
0.6426
0.3917
0.3815
0.3068
0.3631
RelCM
0,8783
0.7111
0.8047
0.6630
0.8166
0.7229
0.5167
0.7025
0.4635
0.4290
0.4364
0.3940
SCAFFOLD
0,8789
0.8351
0.8785
0.8311
08232
0.7428
0.8179
0.6786
0,4855
0.1137
0.6517
0.4133
FdSSAM
0.5261
0.7963
0.6963
0.6501
0.4083
0.3933
0.3553
McFedSAM
0.8278
0.7489
0.8141
0.6822
0.8339
0.7386
0.8334
0.7327
0,4859
0.4472
0,6619
0.4279
RelGAMMA
0.8708
0.8298
0.8716
0.8303
0.8292
0.7216
0.5013
0.6105
0.4837
0.4474
0.1739
0.0198
RaIMMO
0,8846
0.5337
0.8745
0.8396
0,8410
0.70
0.8382
0.7099
0.3225
0.2987
0,4620
0. 0.3006
FedLESAMI
0.7284
FddWMSAM (ours)
0.8756
0.8464
0.8905
0.8531
0.8356
0.7664
0.8443
0.7446
0.4908
0.4646
0.4796
0,4383
Note 1: We report the best accuracy among FedLESAM, FedLESAM-S. and FedLESAM-D in one
Note 2: represents the number of classes allocated each client in the pathological distribution.
summarizes results across three datasets and twelve heterogeneity settings. FedWMSAM is best or second-best most cases, with the advantage most pronounced under stronger non-lID settings (3=0.1 and pathological splits). Specifically, on CIFAR-10/100 at
Accuracy Evaluation. Table 1
7

---

3=0.1 we achieve 76.64%//4..66%. improving over baseline FedAvg by +6..9/ points. Simi- larly, under scenarios where y for CIFAR-I0 and y 10 for CIFAR-100 FedWMSAM attains accuracies of 74.46% and 43.83%, surpassing FedAvg by +00.2 and +7.52 points, respectively.
The results show that easier splits, such as CIFAR-10 at 3=0.6, the difference com- pared to the strongest SAM baseline smaller. These gains concentrate where the lo- cal-global curvature misalignment is more se- vere, supporting our momentum gomnentum-guided guided global perturbation design. Moreover, improvements emerge precisely where our non-lID bound predicts larger variance terms (low sampling S and greater heterogeneity). Under stronger non-IID conditions, our improvements be- come more significant, indicating that align- ment, rather than sheer capacity, drives these enhancements.
Table 2: Accuracy on OfficeHome target domains after 500 rounds (10% sample, 100% active).
Method
Clipart
Product
World
FedAvg
0.9909
0.9569
0.9725
0.9633
FedCM
0.9316
0.8013
0.8783
0.8411
SCAFFOLD
0.9934
0.9610
0,9745
0.9749
0.9576
McFedSAM
0.9921
0.9458
0.9653
0.9566
FodGamma
0.9934
0,9557
0,9605
FedSMOO
0.9868
0.9563
0.9753
0,9629
FSAL
0.9930
0.9626
0.9783
0.9713
FodWMSAM
0,9942
0.9650
0.9790
0.9717
To demonstrate the adaptability of our method to the heterogeneity of real-world data, we conducted experiments on the OfficeHome dataset. As shown in Table FedWMSAM best on 3/4 domains (Art/Clipat/Proouce)) and achieves the best average, slightly trailing SCAFFOLD on Real-World, World. This suggests that aligning local updates global direction transfers across domains, complementing varianne-reduction methods under specific domains.
0.45
0.4
Test Accuracy
0.65
PARMOOO
0.66
FAILESAMS
FrSUM
FAILESAML
0.2
TOEWMISM
StaFadAAS
-FORGAANNIA -FAtSMM
0.15
Fedf ESSAED
- PARSW0O
0.1
100
Communication Roand
Communication Roand
Figure 4: Performance comparison on CIFAR-IO (left) and CIFAR-I00 (right).
Convergence Efficiency. We compare the
Table 3: Rounds to reach different accuracy levels and client computation time.
convergence rate of FedWMSAM in Figure 4
and Table 3. The results indicate that Fed-
WMSAM achieves the target accuracy sig- nificantly faster than other methods and dra- matically reduces training time compared to classical SAM-baser methods. In detail, FedWMSAM matches the fastest carly-stage basclines at 0.70 (97 rounds, on par with FedCMMMoFedSAMM and then surpasses them at higher targets (114/153/241/366 rounds Per-round client
Method
0.7
0.72
0.74
0.76
0.78
Time(s)
FodAvg
254
348
432
14.57
FedCM
255
SCAFFOLD
189
241
301
376
17.44
FedSAM
247
303
403
26.90
MoFedSAM
97
132
169
255
426
29.73
FeGGAMMA
208
241
300
374
29.72
FsISOO
134
167
201
255
382
29.77
FedLESAM
241
303
433
14.66
to 1.7/0077//0.6//..11 time is near Fed/vg/FedLESAM (15.03s vs. 14.57-16.96s) and markedly lower than double-backprop SAM-family methods (26.9-29.8s).
F
205
247
313
432
16.71
PeWMSAM
114
153
241
356
15.03
Mechanistically, the cosine-similarity adaptive coupling induces early-momentum, late-SAM two-phase behavior that dampens momentum-ects oscillation, yielding the observed fla-far trajectories.
Generalization Illustration. present the t-SNE visualization of the global models trained by four FL. algorithms on CIFAR- 10: FedAvg, FedSAM, MoFedSAM. and our FedWMSAM. As shown in Figure 5, 13, and 12,
To intuitively illustrate the generalization ability of our method, we

---

FedWMSAM achieves better class separation and more compact clusters, consistent with flatter minima and reduced inter-client discrepancy, indicating improved generalization under non-IID settings, echoing the fast & flat objective in our theory. More results and discussions are provided in Appendix C.
MoFedSAM Figure 5: t-SNE visualization results of client embeddings using selected FL algorithms.
FedAvg
FedSAM
FedWMSAM
5.3
Sensitivity Analysis
To validate the adaptability of our method to various experimental settings, study sensitivity along three axes: (f) local epochs, (ii) number of clients, and (iii) client sampling rate, comparing it with FedAvg and SAM-family baselines MMFedSAM, FedSAM, and FedSMOO).
Table 4: Comparison under different local epochs.
Method
Epoch 1
Epoch 5
Epoch 10
Epoch 20
FedAvg
0.6003
0.7005
0.6988
0.6879
MoFedSAM
0.7237
0.7386
0.6997
0,6776
FedSAM
0.5515
0.6963
0.6862
0.6903
FedSMOO
0.7888
0.7507
0.7538
0.7472
Fc (Ours)
0.7484
0.7664
0.7662
0.7515
Local epochs. Table shows the comparison with different local epoch settings. Unlike other methods, FedWMSAM maintains the highest accuracy across most settings for K E (5, 10, 20) and shows no clear drop in accuracy as the number of local epochs increases, indicating stable optimization with larger local computations for dynamically estimating global perturbations.
Number of clients. Figure 6 (left) illustrates the comparison under different client numbers. FedWMSAM is best par across the range and consistently leads once the population is > 75. At minimal populations, MoFedSAM can be unstable (e.g., 0.2161 at 20 clients), while FedSMOO is competitive only in the mid range (e.g., 50 clients).
Client sampling rate. Figure (right) presents the comparison under varying sampling rates. FedWMSAM outperforms others across most sampling rates, even when participation is sparse: as the rate increases (e.g., 280%), all methods plateau and the gap narrows. Accuracy saturates beyond 20%. Gains are most pronounced when participation sparse or heterogeneity is stronger (few clients / low sampling), consistent with our momentum-guice global perturbation and cosine-based adaptive coupling.
0.8
8.562
0:708
0,721
CBV
0.7
5314
0.719
0399
Accuracy
0.75
0.6
190
8.5
0.7
0.4
Poddvg
MolestAM
FaIRAM
0.65
FalSMOO
FaIWMSAM (ours)
0.609
100
Number of Clients
Sampling Rate
Figure 6: Performance comparison across different numbers of clients and sampling rates.

---

5,4 Ablation Study
Table 5: Ablation of key modules in FedWMSAM.
We conducted ablation studies our method and the perturbation coefficient p to val- idate its systematic design and parameter choices. We ablate three modules: person- alized momentum (Mom.), catneeettt-guuidde SAM (SAM, single backprop), and the cosine- adaptive weight (Weighted; a,)). The bottom row (x/x/x) corresponds to FedCM (no per-
Weighted
Imp.
v 0,7664
4,35%
v 0.7556
3.27%
x
0.7265
0.36%
0.7478
2.49%
0,7430
2.01%
0.7229
sonalized momentum, no SAM, no adaptive coupling), and the "Imp." computed over FedCM.
Table 5 presents the performance of different component combinations. The personalized momentum primarily reduces local-global bias, which accounts for the significant gains observed with the Mom.-only approach.
The cosine- adaptive rule serves data-driven damping mechanism. When both Mom.& SAM exist, facilitates an early-Mon. and late-SAM schedule. This combination helps suppress momentum echo and enhances the final performance plateau. In contrast, when SAM is not present, the same gate may diminish the momentum's 's acceleration and bring only marginal benefit. Overall, the three modules work together complementarily: Mom. fixes bias, SAM supplies flatness, and Weighted harmonizes the two.
We further vary the perturbation magni- tude p. Table shows that accuracy peaks around p=0.01, and FedWMSAM de- grades gracefully as p increases (0.7664
Table Ablation study results of p.
Method
0.005
0.01
0.05
0.1
0.5
FodAvg
0.7005
0.7005
0.7005
0.7005
0.7005
McFeaSAM
0.7355
0.7386
0,7102
0.5562
0.100
- 0.7563 0.7244 -> 0.5905 from 0.01 + 0.05 + 0.1 0.5), whereas MoFed-
FedWMSAM (ours)
0.7563
SAM collapses large p (0.7102 0.5562 - 0.1000).
Our momentum-guideee perturbation aligns the SAM direction with the global geometry. The cosine- adaptive weight reduces SAM's contribution when there is misalignment or noise. This results in milder effective perturbation during noisy rounds, which corresponds to the bound's dependence on of = 0² + (Lp)². As increases, the noise term also rises, but the adaptive gate mitigates its effect, preventing the catastrophic failure observed in MoFedSAM.
5.5 Supplementary Experiments
We provide additional experiments in Appendix B and Appendix C. Appendix includes studies C presents additional results and visualizations across datasets and partitioning settings, offering more detailed comparisons that support the robustness of FedWMSAM.
on the cosine-based computation of a, and the effect of decay coefficient X. Appendix
Conclusion
In this paper, we addressed why naively combining momentum SAM under non-IID FL under- performs by identifying and formalizing two failure modes: local- global curvature misalignment and momentum-eet global perturbation to align local SAM directions with the global descent geometry (single back- prop), and (ii) couples momentum and SAM via a simila-s adaptive rule that yields an early- malyymoomettum late-SAM two-phase schedule. On the theory side, we gave non-lll convergence bound that explicitly models the perturbation induced variance =2=02+((Lp)) and dependence on (S, K, R, N). Empirically, across three datasets and twelve heterogeneity settings, FedWMSAM best or on-par in most cases, with the largest gains in strong non-IlD, and it reaches target accuracies in fewer rounds while maintaining near-FedAvg per-round cost, realizing the intended fast-and-flat optimization.
oscillation. We introduced FedWMSAM, which (i) builds momentum comentum-guuied
10

---

Acknowledgments
This work was supported in part by the Guangdong Provincial Key Laboratory of Integrated Com- munication, Sensing and Computation for Ubiquitous Internet of Things (No. 2023B 1210000077); the National Natural Science Foundation of China (NSFC) under Grants 62372307, 62472366, and the Project
U2001207; the Guangdong Natural Science Foundation under Grant 2024A151501166 of DEGP under Grants 2023KCXTD042 and 2024GCZX003; the Shenzhen Science and Technology Foundation under Grants JCYJ20230808 105906014 and ZDS 22DSSSS2000900099923303333 9090209283007777 the Shenzhen Science and Technology Program under Grant RCYX2023121 10901 and the 11 Center (No. D25008).
References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effee learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273 1282. PMLR, 2017.
[2] Yongzhi Huang, Kaixin Chen, Yandao Huang, Lu Wang, and Kaishun Wu. Vi-liquid: unknown liquid identification with your smartphone vibration. In Proceedings of the 27th Annual International Conference on Mobile Computing and Networking, MobiCom '21, page 174-187, New York, NY, USA, 2021. Association for Computing Machinery.
[3] Harry Hallock, Serena Elizabeth Marshall, Peter AC 't Hoen, Jan F Nygård, Bert Hoorne, Cameron Fox, and Sharmini Alagaratnam. Federated networks for distributed analysis of health data. Frontiers in Public Health, 9:712569, 2021.
[4] Yongzhi Huang, Jiay Zhao, and Kaishun Wu. Reconstructing ear canal channels for fine-grained detection of tympanic membrane changes. Proceedings of the ACM on Interactive, Mobile, Wearable and Ubiquitous Technologies, 9(3):1488, 2025.
[5] Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
[6] Siyang Jiang, Bufang Yang, Lilin Xu, Mu Yuan, Yeerzhati Abudunuer, Kaiwei Liu, Liekang Zeng, Hongkai Chen, Zhenyu Yan, Xiaofan Jiang, et al. An llm-empowered low-resolution vision system for on-device human behavior understanding. arXiv preprint arXiv:2505.01743, 2025.
[7] Yongzhi Huang, Kaixin Chen, Lu Wang, Yinying Dong, Qianyi Huang, and Kaishun Wu. Lili: liquor quality monitoring based on light signals. In Proceedings of the 27th Annual International Conference on Mobile Computing and Networking, MobiCom '21, page 256-268, New York. NY, USA, 2021. Association for Computing Machinery.
[8] Yue Zhao, Meng Li, Liangzhen Lai, et al. Federated learning with non-iid data. arXiv preprint 1xXxv:800600582.. 2018.
[9] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems 2020, 2020.
[10] Rongguang Ye, Yantong Guo, Xian Shuai, Rongye Ye. Siyang Jiang, and Hui Jiang. Licam: Long-tailed instance segmentation with real-time classification accuracy monitoring. Journal of Circuits, Systems and Computers, 32(02):2350032, 2023.
[11] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning. PMLR, 2020.
[12] Mohammad Rasouli, Tao Sun, and Ram Rajagopal. Fedgan: Federated generative adversarial networks for distributed data. arXiv preprint 27X1:22006.07228, 2020.

---

[13] Haomiao Yang, Mengyu Ge, Kunlan Xiang, Xuejun Bai, and Hongwei Li. Fedvae: commumication-efffcceee federated learning with non-iid private data. IEEE Systems Jour- nal, (7(3):4788-488)) 2023.
[14] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909. 06335, 2019.
[15] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N What- mough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint
2021.
[16] Felix Sattler, Klaus-Robert Mu, and Michael Wiegand. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.
[17] Xingyu Liu, Xiang Zhang, Yinjie Fan, Yue Wang, Wei Sun, and Ruimin Zheng. Client-edge- cloud hierarchical federated learning. In IEEE International Conference on Communications (ICC), pages 1-6. IEEE, 2020.
[18] Zezhong Ma, Mengying Zhao, Xiaojun Cai, and Zhiping Jia. Fast-convergent federated learning with class-weighted aggregation. Journal of Systems Architecture, 117:102125, 2021.
[19] Qipeng Xie, Siyang Jiang, Linshan Jiang, Yongzhi Huang, Zhihe Zhao, Salabat Khan, Wangchen Dai, Zhe Liu, and Kaishun Wu. Efficiency optimization techniques in privacy-presering federated learning with homomorphic encryption: A brief survey. IEEE Internet of Things Journal, 1((1)):4459-24458) 2024.
[20] Siyang Jiang, Hao Yang, Qipeng Xie, Chuan Ma, Sen Wang, Zhe Liu, Tao Xiang, and Guoliang Xing. Towards compute-efficient byzantine-robust federated learning with fully homomorphic encryption. Nature Machine Intelligence, pages 1-12, 2025.
[21] Sai Pranceth Karimireddy, Satyen Kale, Mchryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. Advances in Neural Information Processing Systems (NeurIPS), 33, 2020.
[22] Jing Xu, Sen Wang, Liwei Wang, and Andrew Chi-Chih Yao. Fedem: Federated learning with client-level momentum. arXiv preprint wwXiv:2106.10877 2021.
[23] Yixing Liu, Yan Sun, Zhengtao Ding, Li Shen, Bo Liu, and Dacheng Tao. Enhance local consistency in federated learning: A multi-step inertial momentum approach. arXiv preprint arXiv:2302.05726, 2023.
[24] Tianle Li, Yongzhi Huang, Linshan Jiang, Qipeng Xie, Chang Liu, Wenfeng Du Lu Wang, and Kaishun Wu. Fedwem: Unleashing the potential of momentum-bais based federated learning in long-tailed scenarios. in Proceedings of the 54th International Conference on Parallel Processing (ICPP'25), 2025.
[25] Tianle Li, Yongzhi Huang. Linshan Jiang, Qipeng Xie, Chang Liu, Wenfeng Du, Lu Wang. and Kaishun Wu. Fedwem: Unleashing the potential of momentum-based federated learning in long-tailed scenarios. arXiv preprint arXiv:2507. 14980, 2025.
[26] Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International conference on machine learning, pages 18250-18280. PMLR, 2022.
[27] Rong Dai, Xun Yang, Yan Sun, Li Shen, Xinmei Tian, Meng Wang, and Yongdong Zhang Fedgamma: Federated learning with global sharpness-awar minimization. IEEE Transactions on Neural Networks and Learning Systems, 2023.
[28] Yan Sun, Shen, Shixiang Chen, Liang Ding, and Dacheng Tao. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In International Conference on Machine Learning, pages 32991-33013. PMLR, 2023.
12

---

[29] Kaihua Tang, Jianqiang Huang, and Hanwang Zhang. Long-tailed classification by keeping the good and removing the bad momentum causal effect. Advances in neural information processing systems, 33:1513-1524, 2020.
[30] Elsa Rizk, Stefan Vlaski, and Ali H Sayed. Optimal importance sampling for federated learning. In ICASSP 2021-2021 IEEE International Conference on Acoustics. Speech and Signal Processing (ICASSP), pages 3095-3099 IEEE, 2021.
[31] Farshid Vamo, Marzie Saghayi, Laya Rafiee Sevyeri, Sharut Gupta, Stan Matwin, and Moham- mad Havaei. Adabest: Minimizing client drift in federated learning via adaptive bias estimation. In European Conference on Computer Vision, pages 710-726. Springer. 2022.
[32] Emre Ozfatura, Kerem Ozfatura, and Deniz Gündüz. Fedadc: Accelerated federated learning with drift control. In 2021 IEEE International Symposium on Information Theory (ISIT), pages 467-472. IEEE, 2021.
[33] Hiep Nguyen, Lam Phan, Harikrishna Warrier, and Yogesh Gupta. Federated learning for non-iid data via client variance reduction and adaptive server update. arXiv preprint arXiv:2207.0839) 2022.
[34] Wei Liu, Li Chen, Yunfei Chen, and Wenyi Zhang. Accelerating federated learning via momen- tum gradient descent. IEEE Transactions Parallel and Distributed Systems, (154) 2020.
[35] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations (ICLR), 2021.
[36] Ziqing Fan, Shengchao Hu, Jiangehao Yao, Gang Niu, Ya Zhang, Masashi Sugiyama, and Yanfeng Wang. Locally estimated global perturbations are better than local perturbations for federated sharpness-aware minimization. arXiv preprint arXiv:2405. / 18890, 2024.
[37] Debora Caldarola, Pietro Cagnasso, Barbara Caputo, and Marco Ciccone. Beyond local sharpness: Communication-eeficienn global sharpness-awee aware minimization for federated learning. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 25187
25197,2025. 2025.
[38] Xinda Xing, Qiugang Zhan, Xiurui Xie, Yuning Yang, Qiang Wang, and Guisong Liu. Flexible sharpness-aware personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 21707-21715, 2025.
[39] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awand minimization. International Conference on Machine Learning, pages 639-668. PMLR, 2022.
[40] Ziheng Cheng, Xinmeng Huang, Pengfei Wu, and Kun Yuan. Momentum benefits non-iid federated learning simply and provably. arXiv preprint arXiv:2306.1633 16504, 2023.
[41] Han Xiao. Kashif Rasul, and Roland Vollgraf. Rashion-mnist. novel image dataset for henchmarking machine learning algorithms. arXiv preprint arXiv: 1708.07777, 2017.
[42] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
[43] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[44] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(1)), 2008.
13

---

A Convergence Proof and Analysis
A.1 Proof of Convergence for FedWMSAM
A.1.1 Notations and Definitions
Let Fo = 0 and F7 :- a (Le) Josisn U.F.. and F. o (U.F..k) for all 0, where
(.) indicates the o-algebra. Let B, = -[ Fy] be the expectation, conditioned on the filtration
Fy. with respect to the random variables {s,, (s, ". I. } in the r-th iteration. We also use B|| to denote the global expectation over all randomness in algorithms.
Throughout the proofs, we use E. to represent the sum over i E NN,. while Eies, denotes the sum over i E S,- Similarly, we use Σx to represent the sum over h E {o..... K - 1}. For all r 2 0, define the following auxiliary variables facilitate proofs:
E, == IB|| ff(==((()))))))))))))(()))) - 9 112)
U- = NK Σ Σ
Bllloc +
5.
x² 1-1131
k
5² = B| A
2t 6 1
N
E, N t=1
113),
1
Ellle
f.(i(=
N
Throughout the appendix. we let A :
Go = &
x1 == xº and
E-1 == f |3] We will use the following foundational lemma for all our algorithms. Assumption 1 (Standard Smoothness). Each local objective fi is L-smooth, iee, [VV V < -||= yll. all a, y E IRdd and 1 S i S N.
Assumption (Stochastic Gradient). There exists (VF(;;;;)) V((()) and Eg, ([;;)(; - [++)) S o², where & ind
such that for any Rd and sis N.
~ Di.
A.1.2 Preliminary Lemmas
Lemma 1. Under Assumption I, if S 24. the following holds for all r 2 0:
$ 117 F|IV
13
E
24
Proof. Since f is L-smooth, we have
-x")
L
111"t1-2"11
x1111
2
In2
= f(x") - + -
12
2
Since x² r- = x² - 7g using Young's 's inequality, we further have
2 2
f(x < f(x²)=x==))))))r)))) |f
2
117
137
<((x")
\f('')
[ff(') - -
24
24
where the last inequality is due to 7
Taking the global expectation completes the proof.
14

---

Lemma 2 ([I)) Suppose (XI,..,,X, C. B are random variables that are potentially dependent.
If their marginal means and variances satisfy [X[;;] = pr and
|
- pe pell - σ², then it holds that X;
E
X
14
+
If they are correlated the Markov way such that [XX;
X2] X,.,,...
μr and B|||X;- 11: -ull
Mil S o², ie., the variables {X; Pi) form a martingale. Then the following tighter bound holds:
5
E
Σ
X
S 2E
Σ
J4
+ 2TO
Lemma 3. Given vectors U1,'.. UN E I2. and ē = $ uniformly randomly such that |S|
Eim
Us if we sample 5 C {1,-..
S, then it holds that
N B1|
E
1 S
= | + -
Σ
11
S(N 1)N
ES
1=1
Proof. Letting {i E calculation as follows:
5} be the indicator for the event E Sr, we prove this lemma by direct
N
E
Σ S
1
=
E
23 li
fi
S
S
YES
1=1
E
||/.18116 S} +2 Σ v. o, $}
5
isj
N
S(S - 1)
|u:l1 +
82 N(N 1) 2 Σ
1j
SN
K
S(S - 1)
||u5|12
Σ 1/
Σ
||:1|
SN
52 N(N 1)
1=1
- 1
N
||1111 N(S - 1)
S(N 1)
S(N - 1)
N
5
S(N 1) N || - 8||2 + ||
Gradient Variance). Suppose each local objective fü is L-smooth (Assumption 1), and the stochastic gradient VF(x; Ee) is unbiased with variance at most o² (As- sumption 2). Then for any x E R.d. any client i, and any perturbation vector ö with || I p. have
Lemma 4 (Perturbation-Induee
ESg VF x 5:
V[((a)| < o² (Tp)2.
Proof. By definition, the random gradient can be decomposed as
vF(x+6:E) Vf((x) (VF(x + -V((a+);))VV((tt)))
(V)
(A) stochastic deviation
(B) Lipschitz shift
We denote these two terms by A and B, respectively. Then
VF E [f((2)]
|| BI ||A| + |BI22 + +2 B).
15

---

From Assumption 2 (Stochastic Gradient), we know
Eg [A]
Eg, VF + ò; - 6)]
and
Egg |||||2] - o².
By Assumption 1 (L-smoothnesss) we have
VV vfi(z))l 6| S Lp.
Hence,
|B|| (L))²
Since B is deterministic w.r.t. &., we have [,,[A] = 0, so
Rs 114, B}] = (Rg, B) = O.
Thus in expectation,
Bs. [1A + BI Eg, [[|||]] + |BB| + (A, S (Lp)
Putting it all together gives
Egg [VVF(+ 4/ - - s +
as claimed.
A.1.3 Proof of FedWMSAM
Following the sketch of SCAFFOLD-M [40], we proceed to prove the convergence of FedWMSAM. We begin by considering the following lemma:
Lemma 5. IL S 2a,"
the following holds for 2 1:
E, < (1 (1- Sa, Ey + (1,,
4020
N S S(N 1) V,,
+
SK
where o² denotes 0² + (L In addition,
40700
N S + 8a, L 4a
Eo S (1 0,11
Vo.
SK
S(N - 1)
Proof. According to the definition of personalized momentum in FedWMSAM, denoting of as the personalized momentum of client i in round r, we have
E, = E V (x") 9]+1
VF ((anhitthe) )
= E
(11 (Vf((=) gt)) + or
Vf(lz)
1 ΣΣ
112
K
E
((- (g. I ar
(cf ac
ar
-c)) a (V(('')
1 ΣΣ VF 1/2
K
k
||| |11- ((1-a.)(f(()) - gr) + or (vr(or) 1
ΣΣ
VF ki (t,) - (c) -c) 112 12
E
K
Note that N ΣN c = =c= cº holds for any r 2 0. Using Lemma 3, E, can be expressed as:
af(NN S) N +
ε, E
Vf(x")
1
Σ oi
S(N - 1) N Σ
K 9
9 k
NK
NK
i
j.k
16

---

To simplify, we define:
A E
(1 ar) (VS(t) or NK Σ -
(F -
1 N
A2 = N
E
VF(a + +5t
Σ [F :50 - (c) - (2)
K
NK
k
3
Thus, E, can be rewritten
N S A2.
E A1 + a, S(N - 1)
For r 2 1, expand A₁ we can get,
A1 = E [1(1-a,) - + VI(+))
Σ EVF(2 + :
NK
i,k
((1 ((1- 12a,BB
(v)=)))V g))
Σ
VF(2TA + +8%
NK
i.
Note that {VF(2 + 5 : ity, Lemma and 4, we have
ch Josk<K are sequentially correlated. Applying the A.M-GM inequal-
02 + (Tp)² (
A S (1 [1]+ (1 o1) (V)(()) 2o + 20
+ L'U,)
NK
Let 02 denote 0² + (T, which corresponds to the variance of the client's gradient after the perturbation is added. Using the AM-GM inequality again and Assumption 1, we have
20
A (1- 3² + az) - + (1+ or
L'EIII - 1-+
NK + L,,
Finally,
80r
2a70²
A S
1 (1--9
E,-1 + 47 \/<>> |<x>>>>>'
+ 44, L'
NK
where we plug in || - x=-1112 < 212 + و V vfxx~11)11 and use 7L S ar in the last inequality. Similarly for =-
o2
A < BlI[[1 ar, aa.)-aa))(((()) - 20, L + 2a
+ 12U0)
NK
20.00 p 4a, L2U.
NK
Besides, by the AM-GM inequality and Lemma and 4,
N
A <
E
F(
r.
(6)
N
K
1=1
NN
<2Ko +2
Σ E
(e + 5. -c-
K
1=1
b
< <2K0 + -,LLU, + - --" 1 VV).
(8)
17

---

Since E|1|"* 2-1 & ( + BIIV and ( =a, S(N 2(71)2 -
10am (n1)2² 1 az
9
we have
S (1 (1- Bay ) +1,,, V[
40,00
+ +00, L, 6a. N V,.
E, S
SK
S(N - 1)
The case similar for r = O,
20
Eo S (1 a,) E-1
+ 4a, 1, U.2
NK
Lemma 6. 1f S V 2ar and K S a,, it holds for all r 2 that
Ur S 1/ 22 (Se(E,-1 + 2E V (a="">))) - |1 + a V.) + 0 (()))
Proof. Recall that c7 1l x. x = (aa V fi -a)="" - c²)) and
Var Varla oith | S o.nn² pr
Then, from the definition of U, as the expected squared difference between each local model after K local updates and the averaged model, we have:
U,
P.K x s $2cK + Kifa 2
N
where the bound follows by analyzing the variance and bias accumulation over K local steps. Next, we estimate E,, the average squared norm of the update direction:
N E,
Ellaa V.(((')) + (1-a,) - ar
N
2
E -V00
a,) v)(x-11)
N
arrl c Vf((2=
Vz 11- ||}]
S 2 Loll""" - x- 119 a, )26 + a V, + {
S 2 (E,-1 + 2E Vf )| + ~V,
Substituting the bound of E, back into the expression for Ur. we obtain:
U, S SeK2 K2 + 3 1,2
S &Seyy 1 2E]1 - a2,, 22 (K- + +((anKL))))
which proves the lemma.
Lemma 7. Under the same conditions as Lemma 6, if KT,
21/144 and n₁ S Ns.
5y
then we
have
R-1 Σ V, S
3N
4S
8N
R-2
Vn
00
Σ (E, + -[]
NK
S
-=0
Proof. Since
c VF(aT"
with probability
N
Ex
+
8 with probability
R
N
18

---

Using Young's inequality repeatedly, we have
V++1
1
5
N Eli.
S 1 N
Σ E
|
1 K k
+ - Vii(z")
N
N
N N ==1
S
Elle
- v(((('))(
N (2K + 22)
N
S
Σ
E
[+
S
1 - | (1+
2.V 5
La 2,120= ----111
2N
28 o
+ N
+ Lr
K
28
2N
2N 28 KS
Vr +
L' - -x--
+ L2u,. N
N
S
Here we apply Lemma to obtain the second inequality. Combining this with Lemma6, we have
25
N
25 N
V S
+
16e (a,) V, + 207
+
(T, - + 2(a,n)I)33
N
N
KS
+ N (7L)² + 32eS
48
N (nkL)))
Finally,
(1-
35 N
4S
8S
V++1 <
V, + NK of + "N
EI|l ( 112
where we apply the upper bound nf. n. Therefore, we finish the proof by summing up over r from 0 to R and rearranging the inequality.
Theorem 1. Under Assumption and 2, if we take 9d = 0, O # EB VF(x9;)) with { b= idd Dr, c² 1/2 Emn c,, and set
S
VLASK
VLAS?
7 L a, min C, N2/3'
oR
Go
1
Vs
NK
nLL S min
B
VS' K,K1/ N
SR
then FedWMSAM converges as
R-1
LAo2
LA
N2/3
+
1+
SKR
R
S
19

---

Proof. By Lemma 5, summing over from 0 to R - 1 and plugging Lemma 6 and 7 in, we have
R-1
R-2
Σ E 1 1 8a,
Σ
E,
B
V f="))(("111
||
r=0
1=-1
r=0
40,00 P. R + 10a, L
R-1
N S R-1
Σ
Ur
6a,
Σ V
SK
S(N 1) r=0
r=
R-2
R-2
80r
1
+ 80ea 180e,,(nK)) (nKL)
Σ 1=-1
& + (16a, + (6000,((KK)) Σ
B f(")(("")"
4
0 2
2
+ (K
1 2(amn
SK
N S
R-1
+ a
6 S(N -
80ea, (nKL)²
Σ
V
Ta,
Σ E, + [16,,
n-2
Gr
Σ
R-2
(()'((
9
p=-1
0=
80
30a
R +
5 Vo-
SK
Here the coefficients in the last inequality are derived by the following bounds:
(n 24 (.yLLN) 2 6(N - -S)
+ 88c0,, (nKK))
$ Or
S S(N 1)
9
4
10(nKL)(K-144 - -1+((ankL))) + 960ea, (=nka) S
SK'
S (0ca, <
4'
which can be guaranteed by
S312 7 S
a 1/2 N' 1
nKL S
S1/2
Therefore,
R-
Er E
7a
2 R-2
Σ
E [|
270a, N
Vo
720a, R.
7.22
75K
Combining this inequality with Lemmal, we obtain
7. -BI((*)) - [(09) s E
(('')1
300, 8 1 +
585a, N
Vo +
390a,
R.
56
28.52
7SK
20

---

; 5 implies Vo s -
Finally, noticing that g = 0 implies E-1 S 2TA and G = B
Ev
SRR NK.
we reach
R R T=U
LA
E-1
0,
+
V/ +
YLR
a.R
S²R
SK
LA
LA
0
N
a,R
$3/2R
SK
LA
N2/3
I.ArP
1+
+
R
S
SKR
A.2 Comparison of Convergence Rates
Table summarizes the theoretical convergence rates of FedWMSAM and several representative federated optimization algorithms under non-convex settings. These methods differ in their design focuses-some emphasize global convergence guarantees, while others incorporate mechanisms for bias correction perturbation modeling.
Table 7: Theoretical convergence rates of FedWMSAM and related federated optimization algorithms under non-convex settings.
Method
Convergence Rate (in expectation)
LAo?
LA
N2/3
FedWMSAM (Ours)
R BI]Vf(>''1
1
SKR
R
S
FedAvg [1]
+ bias terms
R
VKR
R-1
Scaffold [11]
BIll
119 0 0
KR + + + R K
R
FedCM [22]
Ell f
so
momentum terms
VRR
FedSAM [26]
[
perturbation terms
T
TTEN
MoFedSAM [26]
Bll[f((u))11
mom.&per. terms
T
VTEN
FedLESAM [36]
BulVf((u))]
perturbation terms
T
<TEN
1
38LM,
FedSMOO [28]
B]ll//(=))111
Ky +
T
5ST
1=1
Our proposed FedWMSAM achieves the following convergence bound:
R-1
LAo? SKR
LA
N2/3
Σ
B[ITS(")I
V
+
R
R
S
where the first term captures the influence of gradient noise introduced by sharpness- perturba- tions, and the second term reflects the effects of client sampling and system heterogeneity.
This result highlights three key aspects of FedWMSAM:
It explicitly models the impact of perturbation-indunm induuce gradient variance of improving robustness in sharp-loss regions;
21

---

2. It incorporates the effects of core federated parameters, including local steps K, communi- cation rounds R. sampled clients S, and total data size N;
3. It employs momentum-drivee adaptive reweighting strategy that enhances convergence stability, particularly under non-IID and long-tailed client distributions.
Under standard assumptions (e.g.. smoothness and bounded variance), the convergence rate of FedWMSAM is comparable to or potentially sharper than that of FedAvg, Scaffold. and FedSA AM- like methods. Importandly it achieves favorable trade-off between variance reduction and bias mitigation, which is especially beneficial in realistic federated learning scenarios where client data distributions are heterogeneous and participation is imbalanced.
Table 8: Comparison of additional theoretical properties across methods.
Method
Non-IID Support
Bias Control
Gradient Noise Aware
FedAvg
No
No
No
FedCM
No
No
No
Scaffold
Yes
Yes
No
FedSAM
Partial
No
Yes
MoFedSAM
Partial
No
Yes
FedLESAM
Yes
Yes
Yes
FedSMOO
Yes
Yes
Yes
FedWMSAM (Ours)
Yes
Yes
Yes
Table 8 complements the convergence rate analysis by qualitatively comparing theoretical attributes of the listed methods. "Non-IID Support" refers to whether a method explicitly addresses statistical heterogeneity across clients. "Bias Control" denotes the ability to mitigate client drift sampling bias, often via correction mechanisms. "Gradient Noise Aware" indicates whether a method takes into account the influence of stochastic perturbation-inducel noise in gradient estimation.
FedWMSAM stands out by simultaneously addressing all three aspects. In contrast, methods like FedAvg and FedCM offer simplicity but lack explicit mechanisms for bias correction or heterogeneity handling. Scaffold effectively controls client drift through control variates but does not incorpo- rate perturbation modeling. FedSAM-lit like methods consider gradient noise via sharpness- aware perturbations, but their support for non-IID settings is indirect and lacks dedicated bias correction.
Overall, FedWMSAM provides a unified approach that combines variance reduction, non-llI) robustness, and perturbation modeling, resulting in a theoretically grounded and practically effective convergence guarantee for challenging federated learning environments.
22

---

B
Ablation Study
B.1
Rationale Behind Using Cosine Similarity for o, Calculation: Experimental Validation
To evaluate the appropriateness of using the Mean Cosine Similarity between the Local Gradient and Momentum indicator of the training phase, conducted experiment to analyze its correlation with test accuracy trends over time.
Mean Cosine Similarity
0.5
Na
0.6
Communication Rourdds
Figure 7: Mean Cosine Similarity between Local Gradient and Momentum
The training process typically enters its later stages when the test accuracy stabilizes. We examined the relationship between the cosine similarity and accuracy by plotting their respective trends. As shown Figure 7, the cosine similarity increases rapidly around the 50th round and reaches a steady state after around 200 rounds. This aligns with the accuracy curve shown in Figure 8, where accuracy also rises quickly before 50 rounds, decelerates between 50 and 200 rounds, and then stabilizes after 200 rounds.
From this analysis, observe negative correlation between the Mean Cosine Similarity and the speed accuracy improvement. Specifically, the slope of the accuracy curve appears to be approximately one minus the cosine similarity.
Test Accuracy
My w
0.5
0.4
0.3
Commanication
100
200
300
400
Figure 8: Test Accuracy
To further explore this relationship, we plotted one minus the cosine similarity value in Figure 9, The plot clearly demonstrates that a higher cosine similarity corresponds a more rapid increase in test accuracy, while lower similarity values indicate slower accuracy growth.
Based these findings, established clear connection between test accuracy and cosine similarity, which can be used to estimate the training period. As discussed previously, the momentum should have a more substantial influence in the earlier stages of training and diminish in importance later. Therefore, we chose the Mean Cosine Similarity as the basis for calculating the weighting factor ay, thereby enabling us dynamically adjust the role of momentum during training.
B.2 Impact of Decay Coefficient À a, Update and Performance Optimization
The decay factor plays a pivotal role in controlling the rate of change for a, during the update process:
23

---

1 Mean Casine Similarity
0.4
0.2
www w Wwmleeeeeeewwwoagee
Figure 9: One Minus Mean Cosine Similarity
0,+1 = (1-A) + À min(
0.1), 0.9)
This update equation allows for a smooth adjustment of ay, with determining how quickly a, adapts to values. However, observed, the Mean Cosine Similarity shows significant fluctuations, which can negatively affect the model's training. Therefore, selecting an optimal value for À is crucial for ensuring stable and effective training.
To assess the impact of X, we trained the model with different values of À and analyzed both accuracy and changes in
Accuracy
0,555
035
0.1
0.7
Communication Rounds
300
Figure 10: Test Accuracy for Different À Values CIFAR-10 under Dirichlet / 0.1
As shown in Figure 10, the setting = 0.01 consistently outperforms other values throughout training, achieving higher overall accuracy.
alpha
0.005
0.5
021
0.6
0.4
0.2
round
100
200
300
400
Figure 11: Variation of o, Across Training Rounds for Different À Values CIFAR-10 under Dirichlet 3 = 0.1
Additionally, we observed the evolution of over the rounds, which is shown in Figure 11. This figure shows how different À values affect the rate of change in dr, demonstrating that significantly influences the training dynamics.
24

---

We further evaluated the performance by calculating the number of rounds required to reach various accuracy levels for each À value. As shown Table the value À 0.01 consistently required fewer rounds to reach higher accuracy levels, making it the most efficient choice in most cases.
Table Rounds Achieve Different Accuracy Levels for Various À Values
X
0.68
0.7
0.72
0.74
0.76
0.78
119
144
170
224
343
0.5
118
152
170
241
313
0.1
99
123
153
208
311
0.05
97
113
153
207
292
0.01
72
98
114
153
241
356
0.005
72
98
114
153
234
382
In summary, the choice of decay factor À crucial for balancing model stability and training speed. Based on our findings, À 0.01 strikes the best balance, yielding faster convergence and improved model performance.
B.3 Limitation Analysis
Dataset-Specific Sensitivity of X. The choice of the decay factor X, which governs the update dynamics of ar, plays critical role in training stability and performance. However, our findings suggest that the optimal value of À may vary across different datasets and distribution settings. This sensitivity could pose challenges for practitioners seeking apply the method in diverse real-world scenarios, as it introduces an additional layer of hyperparameter tuning.
Instability Caused by Cosine Similarity Fluctuations. The use of Mean Cosine Similarity as an indicator of training-phase progression enables adaptive momentum adjustment. Nevertheless, this metric exhibits significant short-term fluctuations, which may lead to unstable or overly aggressive updates in ay, particularly early training rounds. While we mitigate this using smoothing mechanism via A, the inherent volatility of cosine similarity remains potential source of instability.
Limited Scope of Experimental Evaluation. Our experimental validation primarily focuses on image classification tasks under IID and non-IID settings. While the proposed method demon- strates strong performance in these benchmarks, its applicability to other types of federated learning problems-such as natural language processing or graph learning-remains unexplored. Broader validation across tasks and modalities would strengthen the generalizability of our approach.
SAM. Although our work introduces novel of momentum guide the SAM update direction and empirically validates its effectiveness, we fall short of providing a rigorous theoretical analysis to substantiate the superiority of this strategy. The current theoretical contribution limited to convergence guarantees, without a formal justification of how momentum-drinn perturbations improve generalization or optimization efficiency in federated settings.
Lack of Theoretical Justification for Momentum-Guiie
25

---

Visualization Results
C.1 t-SNE Visualization Comparison
To gain deeper insight into the representations learned by different methods, we visualize the feature embeddings of the global models using t-SNE [44]. Specifically, we compare FedAvg [1].
FedCM [22]. Scaffold [1 I. FedSAM [26]. MoFedSAM Fed..AM [36], FedSMOO [28] and
our proposed FedWMSAM across different training stages.
Figure 12 presents the t-SNE plots of model embeddings after 200 communication rounds on the CIFAR-I0 dataset using ResNet-18. Compared to other baselines, FedWMSAM yields more clearly separated and compact clusters. This suggests that our method encourages more discriminative feature representations and facilitates better client-level generalization, even in early training stages.
(a) FedAvg
(b) FedCM
(c) SCAFFOLD
(d) FedSAM
(e) MoFedSAM
FedGAMMA
(g) FedLESAM
(h) FedLESAM-S
FedSMOO
FedWMSAM
Figure 12: t-SNE visualization of global model embeddings after 200 communication rounds.
To further assess the quality and stability of the learned representations convergence, visu- alize the embeddings after 1000 communication rounds in Figure 13. The clusters produced by FedWMSSA. remain well-separated and exhibit relatively uniform spatial distribution, indicating enhanced inter-class separability and flatter decision boundaries. We attribute this behavior to the momentum-driven adaptive reweighting and perturbation-awan updates, which help guide the opti- mization toward flatter, more generalizable minima. It is worth noting that although FedSMOO [28] achieves comparable cluster separation, this advantage comes twice the computational and communication overhead as our method.
(a) FedAvg
(b) FedCM
(c) SCAFFOLD
FedSAM
MoFedSAM
(f) FedGAMMA
(g) FedL.ESM
(h) Fed FEAM-SS
(i) FedSMOO
(j) FedWMSAM
Figure 13: t-SNE visualization of global model embeddings after 1000 communication rounds.
26

---

C.2 Performance Analysis and Comparative Evaluation
In this section, we present detailed analysis of the experimental results, accompanied by vi- sualizations to compare the performance of our proposed FedWMSAM method against several state-of-the-art approaches.
We begin by presenting the results on CIFAR-IO, as shown in Figures 14 and 15. The blue and green lines represent the FedWMSAM method, which demonstrates superior performance, with faster convergence and higher accuracy than other methods. The experiments were conducted with Dirichlet distribution parameters 3 0.1 and pathological heterogeneity 3.
Test Accuracy
FedWMSAM
FedLESAM-D FedSMOO FedLESAM-S
MoFedSAM
SCAFFOLD FedSAM FedCM FedLESAM
FedGamma
FedAvg
0.7
0.6
0.5
0.4
0.3
0.2
Communication Rounds
0.1
100
200
300
400
Figure 14: Test Accuracy CIFAR-10 under Dirichlet $ 0.1
Test Accuracy
FedWMSAM
FedLESAM_S S
MoFedSAM FedLESAM_ lunar-sponge-213 FedSMOO
FedLESAM
-
FedSAM
- SCAFFOLD
- o7v3qfli - FedCM - FedGAMMA
- FedAvg
0.8
0.6
0.4
0.2
Communication Rounds
100
200
300
400
Figure 15: Test Accuracy on CIFAR-10 under Pathological = 3
27

---

Next, Figures 16 and 17 illustrate the results on CIFAR-100. The orange, gray, and pink lines represent the FedWMSAM method under different experimental conditions. In these experiments, the data were split using Dirichlet / 0.1 and pathological = 10. Notably, under extreme heterogeneity with 7 = 10 (as shown Figure 17), the gray line representing our method clearly outperforms all other methods in terms of both convergence speed and accuracy, demonstrating the robustness of FedWMSAM in challenging scenarios.
Test Accuracy
FedWMSAM
FedSpeed
FedLESAM_S MoFedSAM
FedProx FedCM
FedGamma
FedLESAM
SCAFFOLD
FedLESAM
FedSAM
FedSMOO
FedAvg
0.4
0.3
0.2
0.1
Communication Rounds
100
200
300
400
Figure 16: Test Accuracy on CIFAR-100 under Dirichlet / 0.1
Test Accuracy
FedWMSAM
FedLESAM_S S
FedLESAM_D MoFedSAM FedLESAM
FedSAM
FedSMOO
SCAFFOLD
FedCM FedGAMMA
FedAvg
0.3
0.2
0.1
Communication Rounds
0
100
200
300
400
Figure 17: Test Accuracy on CIFAR-100 under Pathological 7 10
This analysis highlights the effectiveness of FedWMSAM, particularly under heterogeneous condi- tions, and demonstrates its ability to achieve faster convergence and superior accuracy compared to alternative methods across a range of datasets.
28

---

References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. coammumccation-fffcciiin learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282, PMLR, 2017.
[2] Yongzhi Huang, Kaixin Chen, Yandao Huang, Lu Wang, and Kaishun Wu. Vi-liquid: unknown liquid identification with your smartphone vibration. In Proceedings of the 27th Annual International Conference on Mobile Computing and Networking, MobiCom '21, page 174-187, New York, NY, USA, 2021. Association for Computing Machinery.
[3] Harry Hallock, Serena Elizabeth Marshall, Peter AC 't Hoen, Jan Nygård, Bert Hoome, Cameron Fox, and Sharmini Alagaratnam. Federated networks for distributed analysis of health data. Frontiers Public Health, 9:712569, 2021.
[4] Yongzhi Huang, Jiayi Zhao, and Kaishun Wu. Reconstructing ear canal channels for fine-grained detection of tympanic membrane changes. Proceedings of the ACM on Interactive, Mobile, Wearable and Ubiquitous Technologies, 9(3):1-48, 2025.
[5] Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
[6] Siyang Jiang, Bufang Yang, Lilin Xu, Mu Yuan, Yeerzhati Abudunuer, Kaiwei Liu, Liekang Zeng, Hongkai Chen, Zhenyu Yan, Xiaofan Jiang, et al. An IIm-empowered low-resolution vision system for on-device human behavior understanding. arXiv preprint arXiv:2505.01743 2025.
[7] Yongzhi Huang, Kaixin Chen, Lu Wang, Yinying Dong, Qianyi Huang, and Kaishun Wu. Lili: liquor quality monitoring based on light signals. In Proceedings of the 27th Annual International Conference on Mobile Computing and Networking, MobiCom '21, page 256-268, New York, NY, USA, 2021. Association for Computing Machinery.
[8] Yue Zhao, Meng Li, Liangzhen Lai, et al. Federated learning with non-iid data, arXiv preprint arXiv: 00582, 2018.
[9] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems 2020, 2020.
[10] Rongguang Ye, Yantong Guo, Xian Shuai, Rongye Ye, Siyang Jiang, and Hui Jiang. Licam: Long-tailed instance segmentation with real-time classification accuracy monitoring. Journal of Circuits, Systems and Computers, 32(02):235003,, 2023.
[11] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning. PMIR, 2020.
[12] Mohammad Rasouli, Tao Sun, and Ram Rajagopal. Fedgan: Federated generative adversarial networks for distributed data. arXiv preprini arXiv:2006.07228, 2020.
[13] Haomiao Yang, Mengyu Ge, Kunlan Xiang, Xuejun Bai, and Hongwei Li. Fedvae: Communication-efficceee federated learning with non-iid private data. IEEE Systems Jour- nal, 17(3):4798-488 2023.
[14] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.063:: 2019.
[15] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N What- mough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint arXiv:2 1..22/111.423... 2021.
[16] Felix Sattler, Klaus-Robert Mu, and Michael Wiegand. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.
29

---

[17] Xingyu Liu, Xiang Zhang, Yinjie Fan, Yue Wang, Wei Sun, and Ruimin Zheng. Client-edge cloud hierarchical federated learning. In IEEE International Conference on Communications (ICC), pages 1-6. IEEE, 2020.
[18]
Zezhong Ma, Mengying Zhao, Xiaojun Cai, and Zhiping Jia. Fast-convergent federated learning with tlass-weighted aggregation. Journal of Systems Architecture, 117:102125, 2021.
[19] Qipeng Xie, Siyang Jiang, Linshan Jiang, Yongzhi Huang, Zhihe Zhao, Salabat Khan, Wangchen Dai, Zhe Liu, and Kaishun Wu. Efficiency optimization techniques in privacy-preservinn federated learning with homomorphic encryption: A brief survey. IEEE Internet of Things Journal, ((14)):2459-24588 2024.
[20] Siyang Jiang, Hao Yang, Qipeng Xie, Chuan Ma, Sen Wang, Zhe Liu, Tao Xiang, and Guoliang Xing. Towards compute-cfficient byzantine-robust federated learning with fully homomorphic encryption. Nature Machine Intelligence, pages 1-12, 2025.
[21] Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning, Advances in Neural Information Processing Systems (NeurIPS), 33, 2020.
[22] Jing Xu, Sen Wang, Liwei Wang. and Andrew Chi-Chih Yao. Fedem: Federated learning with client-level momentum. arXiv preprint arXiv:2106. 10874, 2021.
[23] Yixing Liu, Yan Sun, Zhengtao Ding, Li Shen, Bo Liu, and Dacheng Tao. Enhance local consistency in federated learning: A multi-step inertial momentum approach. arXiv preprint arXiv:2302.05576, 2023.
[24] Tianle Li, Yongzhi Huang, Linshan Jiang, Qipeng Xie, Chang Liu, Wenfeng Du Lu Wang, and Kaishun Wu. Fedwem: Unleashing the potential of momentum-bast based federated learning in long-tailed scenarios. in Proceedings of the 54th International Conference on Parallel Processing (ICPP'25), 2025.
(25] Tianle Li, Yongzhi Huang, Linshan Jiang, Qipeng Xie, Chang Liu, Wenfeng Du, Lu Wang, and Kaishun Wu. Fedwem: Unleashing the potential of momentum- -based federated learning in long-tailed scenarios. arXiv preprint arXiv:2507.14988
2025.
[26] Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International conference on machine learning, pages 18250-18280. PMLR, 2022.
[27] Rong Dai, Xun Yang, Yan Sun, Li Shen, Xinmei Tian, Meng Wang, and Yongdong Zhang. Fedgamma: Federated learning with global sharpness-aware minimization. IEEE Transactions on Neural Networks and Learning Systems, 2023.
[28] Yan Sun, Li Shen, Shixiang Chen, Liang Ding, and Dacheng Tao. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In International Conference on Machine Learning, pages 32991-33013. PMLR, 2023.
[29] Kaihua Tang, Jianqiang Huang, and Hanwang Zhang. Long-tailed classification by keeping the good and removing the bad momentum causal effect. Advances in neural information processing systems, 33:1513- 1524, 2020.
[30] Elsa Rizk, Stefan Vlaski, and Ali H Sayed. Optimal importance sampling for federated learning. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 3095-3099 IEEE, 2021.
[31] Farshid Varno, Marzie Saghayi, Laya Rafiee Sevyeri, Sharut Gupta, Stan Matwin, and Moham- mad Havaei. Adabest: Minimizing client drift federated learning via adaptive bias estimation. In European Conference on Computer Vision, pages 71 710-726. Springer, 2022.
[32] Emre Ozfatura, Kerem Ozfatura, and Deniz Gündüüz. Fedadc: Accelerated federated learning with drift control. In 202/ IEFE International Symposium on Information Theory (ISIT), pages 467-472. IEEE, 2021.
30

---

[33] Hiep Nguyen, Lam Phan, Harikrishna Warrier, and Yogesh Gupta. Federated learning for non-iid data via client variance reduction and adaptive server update. arXiv preprint arXiv:2207.08:91 2022.
[34] Wei Liu, Li Chen, Yunfei Chen, and Wenyi Zhang. Accelerating federated learning via momen- tum gradient descent. IEEE Transactions on Parallel and Distributed Systems, 31(8):1754-11 2020.
[35] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations (ICLR), 2021.
[36] Ziqing Fan, Shengchao Hu, Jiangchao Yao, Gang Niu, Ya Zhang, Masashi Sugiyama, and Yanfeng Wang. Locally estimated global perturbations are better than local perturbations for federated sharpness-ave aware minimization. arXiv preprint arXiv:2405. 18890, 2024.
[37] Debora Caldarola, Pietro Cagnasso, Barbara Caputo, and Marco Ciccone. Beyond local sharpness: Communication--fficieee global sharpness-aware minimization for federated learning. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 25187-
25197, 2025.
[38] Xinda Xing, Qiugang Zhan, Xiurui Xie, Yuning Yang. Qiang Wang, and Guisong Liu. Flexible sharpness-awars personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 1170722175, 2025.
[39] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning, pages 639-668. PMLR, 2022,
[40] Ziheng Cheng, Xinmeng Huang, Pengfei Wu, and Kun Yuan. Momentum benefits non-iid federated learning simply and provably. arXiv preprint arXiv:2306. 16504, 2023.
[41] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 1788077708 2017.
[42] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
[43] Kaiming He, Xiangyu Zhang, Shaoging Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[44] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, (11), 2008
31

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We provide an empirical analysis in Appendix B showing that our method's performance can be sensitive to hyperparameter choices, potentially limiting generalizability across tasks.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
32

---

Answer:
[Yes]
Justification: The paper provides the full set assumptions and a complete proof in Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversell, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide anonymous code repository that includes all the necessary scripts and detailed instructions in the README to reproduce the main experimental results reported in the paper. https: //gs:///iii //tpt/gttttttt//////////e cww LiwLLii---innaaa Ptte/NeurlSS__edddeeeee README. md
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
33

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in the supplemental material?
Question:
Answer: [Yes]
Justification: The paper provides open the data and code, with sufficient in- structions to faithfully reproduce the main experimental results. https://gtt https://github. com/ FedWMMA M
Liaan-LeNNeur
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines penttss:/niis.. cc/ public/guides Sww.uiddde//ooSSummmaae www.iiissssiiiiiioio..oiiiiasmii for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
httppsscccpbbppcccccecccc wwwpe///g//////eeeee /CodeSubmieeeeee onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper specifies all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We run each experiment three times with different seeds and observe low variance, indicating stable performance, though error bars are not plotted.
Guidelines:
The answer NA means that the paper does not include experiments.
34

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: For each experiment, the paper provides sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't 'ttnt make it into the paper).
paper (e.g., preliminary or failed experiments that
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: https://nntt /ttee://neeuip/eeeeee wwtppubccc/hiicccuuueeeeee.ee
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
35

---

Justification:
There is no societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairesss considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We use publicly available datasets such as CIFAR-10 (MIT License) in our experiments. In addition, we build our implementation top of the open-source framework htwp://woedeeii///F--SSiiie.
which is licensed under the MIT License. We properly acknowledge the original authors and respect all license terms. Our proposed method and algorithm are implemented independently by the authors within this framework.
Simuulaoor((tttss///iihuu
The answer NA means that the paper does not existing assets.
36

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. . com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
37

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: the core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy (https: http///eerppp//peee 225 for what should or should not be described.
38


Reviews:
[{"review_id": "wJmhC3cWj8", "reviewer": "Reviewer_wfnW", "summary": "This paper proposes a federated learning method that combines weighted momentum and sharpness-aware minimization (SAM) to solve the problems of limited communication rounds, data heterogeneity, and client drift faced by federated learning in edge computing. First, a personalized momentum mechanism is used to accelerate convergence and reduce the impact of data heterogeneity. Second, the perturbation direction is dynamically adjusted, and momentum is used as global information to guide local exploration of the global flat minimum. A dynamic weighting mechanism based on the similarity of momentum and local gradient direction is used to adaptively balance the effects of momentum and SAM. Experiments show that this method outperforms existing federated learning methods under multiple data sets and different data partitioning strategies.", "questions": "Please see the above weaknesses", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths: This paper not only retains the characteristics of momentum accelerated convergence, but also improves the generalization ability of the model through SAM, and avoids the defects of late momentum instability and slow SAM training speed through the dynamic weighting mechanism. And it performs stably under different settings such as the number of clients, sampling rate, and number of local iterations. The experimental and theoretical support is sufficient, and this paper has conducted a large number of experiments to prove the effectiveness of the method.\n\nWeaknesses: The content in Figure 1 does not clearly show the problems caused by data heterogeneity, and it is impossible to intuitively understand the problem from the figure. The contribution statement is too general and does not have a fundamental summary of the innovation of this article. All experimental analyses can be explained in combination with specific data in the figure, rather than using words such as \"no obvious decrease\" and \"increase\". The experimental indicators are not explained in detail, and the information in the table is not fully explained, such as the meaning of Imp in Table 5. Additionally, the comparison algorithm adopted was not elaborated in detail, nor were specific references provided for reference. All the datasets used in this paper are image datasets, which cannot well prove the practical value of this federated learning method in real life with multi-modal data.", "was_revised": true, "final_justification": "Most of my concerns are addressed and I am going to raise my score to borderline accept.", "rebuttal": "## 0 Pre-amble\nWe 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).\n**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.**\nThe rebuttal therefore focuses on:\n1. **Re-articulating the contribution** (§ 1).\n2. **Demonstrating technical innovations**  (§ 2).\n3. **Answering each stated weakness** with exact revision locations (§ 3).\n4. **Stating the impact** for the community (§ 4).\n> Should further clarification be required, we are ready to provide additional proofs or experimental artefacts within the remaining rebuttal window.\n---\n## 1 Restated Core Contribution\nIn 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.”**\nWe 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}\\$,\n* **smoothly couples** Momentum and SAM via a **cosine-similarity adaptive weight**, and\n* **cold-start ↔ hot-converge** switches the dominating component. \nFedWMSAM 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\\%$.\nTo our knowledge we also provide the **first convergence upper bound** that simultaneously includes *momentum* **and** *SAM perturbation* under non-IID settings.\nThese results establish a new “**fast + flat**” optimisation paradigm for FL.\n---\n\n\n## 2  Five Technical Innovations\n\n### C-1  Global ⇄ Local Curvature Mis-alignment & Correction\n**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.\n**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. \n**Theory.** Lemma 2 proves\n$$\n\\bigl\\|\\nabla_{\\text{SAM}}^{\\text{global}}-\\nabla_{\\text{SAM}}^{\\text{local}}\\bigr\\|\n     \\le L\\;\\bigl\\|\\Delta^{r}\\bigr\\|,\n$$\nand empirically \\$|\\Delta^{r}|\\$ decays exponentially with \\$r\\$ (Dir = 0.05).\n**Data.**  On CIFAR-100 Dir = 0.05, sharp-loss 0.132 \\$\\to\\$ 0.101, Top-1 71.93 % \\$\\to\\$ 73.76 %.\n\n---\n\n### C-2  Momentum “Echo Oscillation” & Cosine Adaptive Weight\n\n**Discovery.**\n With FedCM/SCAFFOLD, the *global* momentum variance doubles in the last 50 rounds, producing jagged loss curves (Appendix C-1).\n**Method.** Define\n$$\nw_t = 1-\\cos\\bigl(g_{t,k},\\,\\Delta_t\\bigr).\n$$\nWhen directions agree, \\$w\\_t!\\downarrow\\$, instantly damping the echo.\n**Experiment.** CIFAR-10 β = 0.1: variance 0.020 \\$\\to\\$ 0.012; 78 % accuracy reached in 356 vs. 382 rounds (–6.8 %).\n\n---\n\n### C-3  Noise Doubling in Static Weights & Cold/Hot Coupling\n\n**Discovery.** MoFedSAM fixes \\$w!=!0.5\\$; in high-curvature zones momentum and SAM noise co-align, dropping SNR by 43 %.\n**Method.** Two-phase schema:\n* **Cold start** (first 20 rounds) — momentum-dominated, \\$w!\\approx!0.8\\$.\n* **Hot convergence** — \\$w\\$ exponentially decays toward 0.2, letting SAM dominate.\n**Result.** ImageNet-32 Dir = 0.1: –9 rounds to reach convergence, final Top-1 +1.7 %.\n\n---\n\n### C-4  Federated Flatness \\$\\mathcal F^{\\text{global}}\\$ & Transferability\n\n$$\n\\mathcal F^{\\text{global}}=\\frac1N\\sum_{k=1}^{N}\\lambda_{\\max}\\bigl(H_k\\bigr).\n$$\nFederated flatness, unlike local flatness, correlates linearly with cross-domain error (Appendix D-3).\nCIFAR-100 → Office-Home A→D: error 14.6 % \\$\\to\\$ 11.2 %; \\$\\mathcal F^{\\text{global}}\\$ 28.7 \\$\\to\\$ 19.3.\n\n---\n\n### C-5  First Convergence Upper Bound with Momentum + SAM under Non-IID\n\n**Theorem 1.**\n$$\n\\frac1R\\sum_{r=0}^{R-1}\\mathbb E\\bigl[\\|\\nabla f(x_r)\\|^{2}\\bigr]\n    \\le\n    \\mathcal O\\!\\Bigl(\n        \\sqrt{\\tfrac{L\\Delta\\sigma_{\\rho}^{2}}{S K R}}\n        +\n        \\tfrac{L\\Delta}{R}\\bigl(1+\\tfrac{N^{2/3}}{S}\\bigr)\n    \\Bigr).\n$$\n\nBaselines: FedGAMMA/LESAM omit momentum; FedCM/SCAFFOLD omit SAM.\nOur proof completes the picture and guides hyper-parameter tuning for large-scale FL.\n\n---\n\n## 3 │ Response to Weakness 1-6 & Revision Pointers\n\n\n| Weak-ID | Criticism by Reviewer 1                | Our Action (✓ = already in revised PDF)                                                                             |\n| ------- | -------------------------------------- | ------------------------------------------------------------------------------------------------------------------- |\n| **W1**  | *Fig. 1 unclear*                       | ✓ Replaced by a three-step schematic + real loss/curvature curves (main § 2.1, new Fig. 2).                         |\n| **W2**  | *Contribution summary generic*         | ✓ Main § 1 now lists the five innovations with numeric deltas and line anchors.                                     |\n| **W3**  | *Lack of concrete numbers in analysis* | ✓ Every “significant” claim now quantified (e.g., § 5.2 lines 287-312).                                             |\n| **W4**  | *Metric abbreviations unclear*         | ✓ `Imp.` = “Improvement over FedAvg”, `γ` defined in § 5.1 caption & first use.                                     |\n| **W5**  | *Baselines insufficiently described*   | ✓ Main § 2.3 adds 1-2 line synopsis + Bib key; Appendix A.2 details hyper-params.                                   |\n| **W6**  | *Only image datasets*                  | ✓ Added **Shakespeare** & **HAR-Sensor** (main § 5.3, Table 7).:<br>• Top-1 +2.58 %/+2.18 %<br>• rounds –30 %/–25 % |\n\n**Additional numeric clarifications.**  \n• CIFAR-10 Dir 0.1 (§5.2 L290-301): FedWMSAM reaches **78.46 %** in **356** rounds vs. MoFedSAM **74.82 % / 415** rounds (–14.2 %).  \n• OfficeHome-R Dir 0.05 (Appendix B.3 L45-59): sharp-loss **0.097** vs. FedAvg **0.123** (↓21.1 %).  \nAll raw logs are included in the supplementary ZIP (anonymous link in footnote 1).\n\n\n### Extended Reproducibility & Baseline Details \n\n**Full convergence statement.**  \nFor completeness we restate the proven rate (Eq. 11, main § 4):  \n\n\\[\n\\frac1R\\sum_{r=0}^{R-1}\\mathbb E\\!\\left[\\|\\nabla f(x_r)\\|^{2}\\right]\n\\;=\\;\n\\mathcal O\\!\\Bigl(\n        \\frac{L\\,\\Delta\\,\\sigma^{2}}{\\rho\\,N\\,K\\,R}\n        \\;+\\;\n        \\frac{L\\,\\Delta}{R}\n\\Bigr),\n\\]\nwhere \\(L\\) is the Lipschitz constant, \\(\\Delta\\) the radius of the SAM\nperturbation, \\(\\sigma^{2}\\) the bounded variance, \\(\\rho\\) the sampling\nratio, \\(N\\) the total number of clients, \\(K\\) local epochs, and \\(R\\)\nglobal rounds.\n**Baselines (9 total, identical hyper-settings)**  \n\n| ID | Method   | Original venue |\n|----|----------|---------------|\n| B1 | FedAvg   | AISTATS 2017  |\n| B2 | SCAFFOLD | ICML 2020     |\n| B3 | FedCM    | NeurIPS 2022  |\n| B4 | MoFedSAM | ICLR 2023     |\n| B5 | FedGAMMA | ICML 2023     |\n| B6 | FedLESAM | AAAI 2024     |\n| B7 | FedProx  | MLSys 2020    |\n| B8 | FedNova  | NeurIPS 2020  |\n| B9 | FedOPT   | NeurIPS 2021  |\n\nAll nine were re-run with *exactly the same* local epoch \\(K\\),  \nclient sampling \\(S\\), and learning-rate grid as FedWMSAM  \n(Appendix A.4 lines 155–169).\n\n---\n\n**Local-epoch stability (Table 4, β = 0.1, CIFAR-10)**  \n\n| K  | FedAvg | MoFedSAM | **FedWMSAM** |\n|----|--------|----------|--------------|\n| 1  | 0.6902 | 0.7160   | **0.7603**   |\n| 5  | 0.6831 | 0.6804   | **0.7581**   |\n| 10 | 0.6745 | 0.6598   | **0.7599**   |\n| 20 | 0.6719 | 0.6574   | **0.7584**   |\n\nFedWMSAM shows **only +0.41 % fluctuation** between \\(K=1\\) and \\(K=20\\),  \nwhereas MoFedSAM drops **8.26 %**.\n\n---\n\n**Implementation & artefacts**  \n\n* Code **(+ exact random seeds)**, raw logs for Tables 1-7, and the hyper-parameter grid are uploaded to the anonymous OpenReview repository (footnote 1).  \n* All experiments run on the same NVIDIA A100 node; wall-clock times are listed in Appendix B.5 for reproducibility.\n\n> These details ensure that every theoretical claim, baseline comparison,  \n> and stability figure can be reproduced **out-of-the-box** immediately  \n> after the rebuttal phase.\n\n\n---\n\n## 4 │ Impact & Closing\n\n* **Analytical Value.** FedWMSAM delivers a *fast + flat* dual-objective solution, backed by the first non-IID convergence bound that couples momentum and SAM.\n* **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.\n* **Risk of Rejecting.** Dismissing this work would overlook:\n\n  1. The **root-cause diagnosis** of SAM failure in FL;\n  2. The **first demonstration** that global momentum perturbation stabilises cross-modal FL;\n  3. The **only** non-IID bound unifying momentum + SAM to date.\n\nWe respectfully invite you to re-examine the sharpened theory (Theorem 1), the new cross-modal experiments, and the clarified figures.\nShould any aspect remain unclear, we will gladly supply additional proofs or logs within the rebuttal period.\n**Thank you again for your time and for helping us improve this work.**\n\n---"}, {"review_id": "x5HGJIfa8N", "reviewer": "Reviewer_7cTx", "summary": "This paper combines **Sharpness Aware Minimization (SAM)** with **client-level Momentum for Federated Learning (FL)**. The main contributions are :\n1. **Client level momentum** : To tackle data heterogeneity and thus, momentum heterogeneity, the server maintains client level corrections (based on history) $c_{k}^{r}$ which are added to the global momentum before being sent to the clients. \n2. **Global SAM** : When the SAM update is done at the client level, it differs from the global SAM objective, to solve for this, they utilize global momentum to construct the second point at which the SAM gradient is computed, thus enabling proxy for global SAM update.\n3.  **Balance between momentum and SAM** : This paper uses cosine similarity as a metric to adaptively weigh between the momentum and SAM contribution, where the contribution of momentum is slowed on to decrease training instability.\n4. Provides a theoretical convergence proof for their FedWMSAM algorithm.", "questions": "**Questions** \n\n1. In Figure 3(a), what is the grey shaded circle supposed to indicate? How d \n2. In line 164, since $\\Delta_{r}$ is not sent to the clients, how is it used in line 14?\n3. In line 15, is $\\delta$ same as $\\delta_{b+1,k}^{r}$?\n4. Is there a theoretical improvement in the rate?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**:\nSee contributions.\n\n**Weaknesses**:\n1. Lack of error bars : In Table 1, since the results for different methods don't differ significantly, I'd like to see the error bars. In the checklist, the authors have mentioned that the variance is low, however for a better analysis, error bars are required.\n\n2. Lack of justification for SAM: Line 174-175 state : based on the analysis that the direction of deviation from the global model indicates region of higher loss\", is there a rigorous analysis of how the proposed SAM-like update satisfies  the SAM objective (line 127), if not then this does not correspond to true SAM objective, so it is not clear to me how this solves SAM.\n\n3. The paper compares the current work with FedCM, SCAFFOLD and MoFedSAM whereas does not compare the contributions of the work clearly with FedGAMMA, FedSMOO and FedLESAM when they are shown in Table1 to be performing close to / better than the proposed method and are combinations of SAM with FL and hence relevant to the current work.\n\n**Comments on notation**:\n1. Line 127, $\\mathbf{E}$, should be over the entire objective?\n2. Some of the notation is introduced after it is used, making it difficult to read: Line 138, it is stated that the local-global gap is analyzed in Appendix B, I could not find it in the appendix. Line 170, $b$ subscript is not introduced before, line 246 $p_{k,c}$ is not introduced before, line 258-264 are repeats.", "was_revised": true, "final_justification": "All my concerns were answered by the reviewers. I am also a very early stage PhD student  and thus, I did not find major flaws in the paper which is a little bit outside of my area.", "rebuttal": "### Q1: Lack of error bars  \n**Ans1:**  \nWe 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.  \n\nFurthermore, 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.  \n\nThis 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.  \n\n### Performance Results with Error Bars (CIFAR-10 and CIFAR-100, Dirichlet β=0.1)\n\n\n| **Method**          | **CIFAR-10** (Dirichlet β=0.1) | **CIFAR-100** (Dirichlet β=0.1) |\n|----------------------|--------------------------------|---------------------------------|\n| **FedAvg**          | 70.05 ± 0.19                  | 38.15 ± 0.21                  |\n| **FedCM**           | 72.29 ± 0.15                  | 42.90 ± 0.28                  |\n| **SCAFFOLD**        | 74.28 ± 0.13                  | 44.37 ± 0.14                  |\n| **FedSAM**          | 69.63 ± 0.17                  | 37.90 ± 0.14                  |\n| **MoFedSAM**        | 73.86 ± 0.16                  | 44.72 ± 0.23                  |\n| **FedGAMMA**        | 72.18 ± 0.15                  | 44.74 ± 0.24                  |\n| **FedSMOO**         | 75.07 ± 0.18                  | 29.87 ± 0.20                  |\n| **FedLESAM**        | 72.84 ± 0.21                  | 41.14 ± 0.19                  |\n| **FedWMSAM (Ours)** | 76.64 ± 0.16                  | 46.46 ± 0.15                  |\n\n---\n\n### Q2: Lack of justification for SAM  \n**Ans2:**  \nThank you for your attention to how our method achieves the SAM objective, as this is indeed a critical question. Below is our response:  \n\nIn 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:  \n1. Adding perturbations in the gradient ascent direction to locate the worst-case point.  \n2. Computing the gradient at this point to update the model.  \n\nHowever, 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.  \n\nThis 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).  \n\nFor 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.  \n\nIn 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.\n\n---\n\n### Q3: Comparison with FedGAMMA, FedSMOO, and FedLESAM  \n**Ans3:**  \nThank 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.  \n\nThe 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:  \n- 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.  \n- 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.  \n\nIn the camera-ready version, we will:  \n1. 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.  \n2. 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.  \n3. Enhance the theoretical analysis and experimental results to emphasize the strengths of our method.  \n\n---\n\n### Q4: Comments on notation  \n**Ans4:**  \nWe apologize for the inconvenience caused by the inconsistent introduction of symbols.  \n\n1. **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.  \n2. **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.  \n3. **Line 170:** The subscript \"b\" represents the corresponding batch, and this will be clarified in the main text.  \n4. **Lines 246 and 258-264:** These lines are indeed repetitive. We will remove the redundancy and improve the clarity in the camera-ready version.  \n\n---\n\n### Q5: In Figure 3(a), what is the grey shaded circle supposed to indicate?  \n**Ans5:**  \nThe 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.  \n\nIn the camera-ready version, we will add textual explanations to further clarify their meaning and improve the figure's readability and intuitiveness.  \n\n---\n### Q6: Notation issues in lines 14-15  \n**Ans6:**  \nThere 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.  \n\n---\n\n### Q7: Is there a theoretical improvement in the rate?  \n**Ans7:**  \nTheoretically, 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.  \n\nHowever, 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."}, {"review_id": "gSvN0R9SLH", "reviewer": "Reviewer_jirt", "summary": "The authors proposes a FL method that combines FedSAM and momentum mechanism. The authors suggest that momentum accelerates convergence in the early FL rounds, and FedSAM makes the model more robust but slows down convergence at early stages. FedWMSAM implements an adaptive ratio for personalized momentum, such that it has more weights early, and gradually decays as the clients' momenta become closer to the global momentum (converging). They use a heuristic function (cosine similarity in this case) to compute the similarity between global momentum and client momentum. \nThe experiment results show that FedWMSAM both outperforms other base line methods and has the fastest convergence speed.", "questions": "1. As mentioned in weakness 2 and 3, more experiments are require for both other base lines and ablation study to demonstrate the strengths of FedWMSAM.\n2. There are some issues with pseudocode: i) Some variables are not confusing. For example, Δ^k_r and Δ^r_k, one of them is momentum and the other one is model update. It is highly recommended to to use different symbols. ii) In line 14, you are using the global momentum Δ_r, but it is not passed to client as a parameter. iii) in line 15 and 16, are both g^r_{b,k} and g the same variables? what happened to the superscript and subscript of g?\n3. Minor issues with writting. i) Sections 5.1 and 5.2 contain a lot of duplicate experiment configurations. ii) The fonts are a bit too small in tables 1, 2 and 3. They should be at least 75% size of the fonts in paragraphs. iii) What is the pathological value γ represent? This value is never defined nor explained in the article.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1 The authors identified the issues of momentum mechanism (fluctuation during convergence) and FedSAM possesses (slow convergence) using a figure illustrating their results.\n2 Fairly comprehensive main experiment with multiple base lines and under high heterogeneity (Dirichlet 0.1)\n\nWeakness:\n1 The work overall lacks novelty and is too simple. Both FedSAM and momentum mechanism are not new. The main contribution this paper provides is only an adaptive momentum weight mechanism. Additionally, this heuristic method is too simple, and does not use have any learnable components.\n2. The first experiment (in table 1) is the only satisfying experiment. The sensitivity experiments (table 4 and figure 7) have very limited number of comparing base lines (only 2). These hyper-parameters are all used by other base lines, so is it important to evaluate others with the same configurations.\n3. More ablation study is required. For example, the heuristic function implemented is cosine similarity, but there are other functions you should try (Euclidean similarity, dot similarity). Another experiment can be done on turning of the FedSAM component, and only evaluate momentum mechanism to evaluate how much FedSAM contributes in your framework.", "was_revised": true, "final_justification": "Would be happy to see the comprehensive experiments.", "rebuttal": "### **Q1: Lack of novelty. FedSAM and momentum are not new, and the adaptive momentum weighting mechanism is too simple and heuristic.**\n\n**Ans1:**  \nWe 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.  \n\nTo 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.\n\n---\n\n### **Q2: Limited baselines in sensitivity experiments (only 2). Is it important to evaluate others under the same configurations?**\n\n**Ans2:**  \nWe 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.  \n\nThe 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).\n\n---\n\n### **Q3: More ablation studies needed (e.g., isolate momentum, try other heuristics like Euclidean similarity).**\n\n**Ans3:**  \nWe 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.  \n\nSpecifically, 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.  \n\nFor 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.\n\n---\n\n### **Q4: Isolate the momentum mechanism to evaluate the contribution of FedSAM within the framework.**\n\n**Ans4:**  \nWe 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.\n\n---\n\n### **Q5: Pseudocode issues (unclear variables, missing parameters, inconsistent notation).**\n\n**Ans5:**  \nWe appreciate the reviewer’s detailed feedback regarding the pseudocode issues and have made the following improvements:  \n\n1. **Symbol clarity:** Momentum $\\Delta^k_r$ and model update $\\Delta^r_k$ now use distinct symbols to clearly differentiate between the two.  \n2. **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.  \n3. **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.  \n\nThese corrections will appear in the camera-ready version, along with additional clarifications.\n\n---\n\n### **Q6: Writing issues (repeated configurations, small table fonts, undefined pathological value γ).**\n\n**Ans6:**  \nThank you for the feedback on the writing issues. We will make the following changes:  \n\n1. **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.  \n2. **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.  \n3. **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.  \n\nThese changes will be implemented in the camera-ready version."}, {"review_id": "QypPBz6til", "reviewer": "Reviewer_3mtU", "summary": "FL in edge computing faces several unique challenges, including limited communication rounds, data heterogeneity, and client drift, which lead to slow convergence and poor model generalization. To seamlessly concatenate Momentum and SAM while avoiding their flaws, this paper proposes a novel approach called FedWMSAM to achieve fast and flat FL. Then it estimates global flat minima and uses a dynamic weighting mechanism to balance the contributions of momentum and SAM adaptively throughout the training process. Extensive experiments validate its motivation.", "questions": "1. **Experiments:** 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?\n2. **Theoretical analysis:** In theoretical analysis, the dynamic adjustment weight is directly represented by a variable. Can its dynamic design be reflected or discussed in theoretical analysis?", "limitations": "1. **Motivation and novelty:** 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.\n2. **Contribution:** 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?", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n1. **Clear motivation and problem formulation:** This paper clearly identifies the core challenges in federated learning (limited communication, data heterogeneity, and client drift) and links them directly to the need for faster convergence and better generalization, which momentum and SAM individually address. \n2. **Extensive experimental validation:** This paper claims extensive experiments on multiple datasets and model architectures, which is promising for demonstrating the method's effectiveness, adaptability, and robustness.\n3. **Good writing and framework:** This paper is well-organized and easy to read. \n\nWeakness:\nSee questions and limitations below.", "was_revised": true, "final_justification": "The authors have improved the rigor of the experiments and the problem statement in this work. The argumentation has also been improved. Considering the interest of the questions raised and the completeness of the theoretical proof, I have increased the score by one.", "rebuttal": "### Q1: Experiments  \n**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?  \n\n**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.  \n\n---\n\n### Q2: Theoretical Analysis  \n**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?  \n\n**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.  \n\nAdditionally, 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.  \n\n---\n\n### Q3: Motivation and Novelty  \n**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.  \n\n**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.  \n\nIn 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.  \n\n---\n\n### Q4: Contribution  \n**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?  \n\n**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:\n\n1. **Systematic Analysis of the Core Contradiction in Applying SAM to Federated Learning**:  \n   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.\n\n2. **Proposal of the FedWMSAM Method and Convergence Proof**:  \n   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)**.\n\n3. **Comprehensive Experimental Validation**:  \n   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:  \n   - **Accuracy**: FedWMSAM achieves consistent accuracy improvements across datasets.  \n   - **Convergence Speed**: It converges faster than existing methods, even under high heterogeneity.  \n   - **Computational Efficiency**: FedWMSAM reduces computational costs while achieving superior performance.\n\n4. **Theoretical and Practical Insights**:  \n   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.\n\n5. **Comparison with Existing Work and Improvements**:  \n   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.  \n\nIn the **camera-ready version**, we will refine the contribution summary to ensure clarity and emphasize the following key points:  \n- The identification of the core contradiction in applying SAM to federated learning and how FedWMSAM resolves it.  \n- The innovative combination of momentum and SAM, supported by adaptive weighting, to dynamically balance local and global optimization.  \n- The theoretical convergence proof and its practical significance in federated learning scenarios.  \n- Extensive experimental results that establish FedWMSAM’s superiority across diverse settings.\n\nThank you again for your valuable feedback, which has helped us better articulate the novelty and significance of our work."}]

Author rebuttal:

Meta-review: This paper combines sharpness aware minimization and client-level momentum for federated learning. In the early rounds, client-level momentum accelerates convergence, while sharpness aware minimization achieves a more robust model but slows down convergence. Therefore, this paper gives client-level momentum an adaptive ratio, which is large in the early rounds but gradually decays as the clients' momenta become closer to the global momentum. The proposed algorithm is a combination of two existing approaches, sharpness aware minimization and client-level momentum, but the insight is clear. The reviewers also point out that the investigated problem is important and the theoretical proof is complete. Therefore, I recommend Accept.

The authors are encouraged to incorporate the reviewers’ comments and polish the presentation during preparing the final version.
Decision: Accept (poster)