One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Authors: Yuhang Li
Tong Liu
Yangguang Cui
Ming Hu
Xiaoqiang Li
Venue: ICML 2025 poster (2025)
Keywords: Federated learning
sharpness-aware minimization
Primary area: general_machine_learning
TL;DR: We propose a novel method  to directly reduce the sharpness of the global model, which achieves high generalization with less communication  and  computation cost in the FL framework.

Abstract: Federated learning (FL) presents a promising strategy for distributed and privacy-preserving learning, yet struggles with performance issues in the presence of heterogeneous data distributions. 
Recently, a series of works based on sharpness-aware minimization (SAM) have emerged to improve local learning generality, proving to be effective in mitigating data heterogeneity effects.
However, most SAM-based methods do not directly consider the global objective and require two backward pass per iteration, resulting in diminished effectiveness.
To overcome these two bottlenecks, we leverage the global model trajectory to directly measure sharpness for the global objective, requiring only a single backward pass.
We further propose a novel and general algorithm FedGMT to overcome data heterogeneity and the pitfalls of previous SAM-based methods.
We analyze the convergence of FedGMT and conduct extensive experiments on visual and text datasets in a variety of scenarios, demonstrating that FedGMT achieves competitive accuracy with state-of-the-art FL methods while minimizing computation and communication overhead. 
Code is available at https://github.com/harrylee999/FL-SAM.

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Yuhang Li" Tong Liu
Yangguang Cui Ming Hu²
Xiaoqiang Li'
Abstract
pate in the FL process and train the local models in multiple intervals with their own datasets within one communica- tion round (Wahab et al., 2021). Due to data heterogeneity (i.c., non-II), partial client participation, and multiple local training, a significant issue arises known as the "client drift" problem (Karimireddy et al., 2020), which means clients' lo- cal objectives converging towards inconsistent local optima, which leads to the aggregated global model may not be the optimum of the global objective. Larger data heterogeneity may enlarge the objective inconsistency, thereby not only degrading generalization performance but also escalating computation and communication overhead.
Federated learning (FL) presents a promising strat- egy for distributed and frivacy-preserving learn- ing, yet struggles with performance issues in the presence of heterogeneous data distributions. Re- cently, a series of works based on sharpness-awar minimization (SAM) have emerged to improve local learning generality, proving to be effective in mitigating data heterogeneity effects. However, most SAM-based methods do not directly con- sider the global objective and require two back- ward pass per iteration, resulting in diminished effectiveness. To overcome these two bottlenecks, we leverage the global model trajectory to directly measure sharpness for the global objective, re- quiring only single backward pass. We further propose a novel and general algorithm FedGMT to overcome data heterogeneity and the pitfalls of previous SAM-based methods. We analyze the convergence of FedCMT and conduct extensive experiments on visual and text datasets in a va- riety of scenarios, demonstrating that FedGMT achieves competitive accuracy with state-of-the- art FL methods while minimizing computation and communication overhead. Code is available at http:///athbbboommmaar///99999.
To tackle the client drift problem, most previous works (Li et al., 2020; Karimireddy et al., 2020; Wang et al., 2020b; Acar et al., 2021) involve enforcing regularization in lo- cal optimization with Empirical Risk Minimization (ERM). However, ERM-based training often falls into a sharp valley in the loss landscape (Chaudhari et al., 2019), resulting in biased regularization term (e-g.. the global model), espe- cially in a highly heterogeneous dataset (Sun al., 2023a). This bias makes the trained local model more inclined to- wards a similar tendency, rendering the entire training pro- cess unstable and significantly undermining performance. Instead, recent studies (Qu et al., 2022; Caldarola et al., 2022; Sun et al., 2023b; An et al., 2023; Dai et al., 2023) have focused on improving the local generality in FL to al- leviate the objective inconsistency, which has achieved new state-of-the-art performance among current FL algorithms. These methods aim to seek flat minima during local learning by employing the recently proposed Sharpness- Aware Mini- mization (SAM) (Foret et al., 2020) as the local optimizer. SAM is built on the findings that emphasize the advantages of a smooth loss landscape for better generalization (Dinh et al., 2017: Jiang et al., 2019). By enhancing local gen- erality, these methods mitigate inconsistency among local objectives, thereby contributing the overall smoothness of the aggregated global model.
1. Introduction
Federated Learning (FL) leverages distributed client (i.e., edge device) data to preserve privacy through an iterative process: downloading models, training models locally, up- loading the updated models by clients and aggregating mod- els on the server (McMahan et al., 2017). Due to the limited communication resources, only a subset of clients partici-
'School of Computer Engineering and Science Shanghai Uni- versity, Shanghai, China "School of Computing and Information Systems. Singapore Management University, Singapore. Corre- spondence to: Tong Liu <tong shu edu.en>.
Unfortunately, SAM-based methods have following short- comings. L The computational cost of SAM's sharpness measure is twice that of the base optimizer. typically stochas- tic gradient descent (SGD). In the process of FL, most of the computation is performed by client devices. This prohibits
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

One Arrow, Two Hawks: Sharpness-awary Minimization for Federated Learning via Global Model Trajectory
and flat area for set Low and area for client set B
al., 2022) in each local update.
Round 1 Active sot
Round 1+1 Active B
and sluap
ty
To directly reduce the global model's sharpness and effeectively address data heterogeneity, we propose FedCMT, which achieves fast convergence speed with less computation cost and maintains high general- ization. Theoretically, we provide the convergence rate upper bound under the non-convex and smooth cases and prove that FedGMT could achieve a fast conver- gence rate of (1/TT).
WA+1
+
SAM
SAM
FRM
ERM
Figure 1. Due to partial participation in FL, client set A active in round and set B in round t + 1. In round f. ERM training falls into a low and sharp area (cream), while SAM falls into low and flat area (green) for the active client set A. In round + 1, both ERM and SAM fall into a low area for the active client set B, but they may forget the knowledge Icarned from set A in the previous round. Our method, however, guides the global model toward a joint low and flat area (green+grey) for both sets A and B.
We empirically show the effectiveness of FedGMT, which outperforms several SOTA baselines. FedGMT is also robust to various data heterogeneity levels, client participation levels and model architecture on distributed visual and text datasets.
SAM from being deployed extensively in practical scenar- ios since most clients may have limited computing power and tend to minimize computing energy consumption (Li et al., 2023). 2 Moreover, the sharpness measure of SAM on each client can only compute local model's sharpness direction minimize. Thus minimizing local sharpness can not directly optimize the flatness of the global model, which may deviate from the global objective and slow down convergence, as shown in Figure 1.
2. A Closer Look at SAM-based Algorithms
In this section, we first formally describe the problem setup for FI., review SAM-based methods, and then outline our motivation. A detailed related work is in Appendix B.
2.1. Preliminaries
Federated Learning (FL). In classic FL setting with M clients, where each client has a dataset 2m. the optimization problem to solve can be formulated as follows:
To overcome data heterogeneity and the pitfalls of previous 3AM-based methods, we propose a novel and general al- gorithm, FedGMT (Federated learning with Global Model Trajectory), to guide each client search for a consistent smooth loss landscape aligned with global objective, thereby significantly improving overall performance. Specifically, we introduce novel global model trajectory loss to directly measure and optimize the sharpness of the global model. This loss replaces SAM's sharpness measure loss by mea- suring the Kullback-Leibler (KL) divergence between the outputs of neural networks with the current local model and those with the past global models, which can guarantee each local update avoids falling into sharp valley of the loss landscape from a global perspective. Furthermore, to pre- vent our global sharpness measure from being affected by data heterogeneity, we introduce a constraint to the global objective and use the alternating direction method of multi- pliers (ADMM) (Boyd et al.. 2011) to solve. This ensures an accurate update direction to smooth the global model.
min 10
[(w)
Em Lm Cm(wv)
M
(I)
MEb
Lm (w) A E Em--D ((s(w; Em)))
where w is the model parameter, MI is the set of all clients, C' (w) is the empirical risk minimization (ERM) for client m, E is a loss function (e.g. cross-entropy loss), f denotes neural network and f. denotes the pair (inputs, targets) of a randomly sampled instance from 'nn- At each round 1, F FedAvg (McMahan et al., 2017) formulates this minimiza- tion problem as performing a weighted average of the local model parameters w/,, updated by the subset of selected clients N Reddi al. (2021) shows that the FedAvg global update can be generally seen as one step of SGD with unitary learning rate:
In the end, we summarize our main contributions as follows:
1+1 w
Σ w', = w
Σ (w²
w w (2)
N
N
MEN
MEN!
We take closer look at sharpness- awwre minimization in heterogeneous FL from global view. We propose novel global model trajectory loss to directly measure the sharpness of the global model without information leakage of local data, supported by theoretical analysis. Furthermore, our measure's 's computation cost is nearly about 0.67x x of FedsAM's (Qu et al., 2022; Caldarola
where 7g is the server-side learning rate, equal to 1 in FedAvg. The difference w - w := A defines the global pseudo-gradient at round t.
Sharpness-Aware Minimization (SAM). SAM (Foret et al., 2020) focuses on optimizing the sharp points from parameter space so that the training model can produce a flat
2

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Table 1. Overview of several SAM-based algorithms in FT.. The communication
defined the parameters transmitted per round (1.5x for FedGMT and 1x for FecGMTv2). The computational cost is defined as per-iteration training cost (1.33x is based on the assumption that backward cost of is twice forward). In FedSMOO, Mmu and s are dual variable and correction to perturbations. In FedLESAM, wold is the global model received at previous active round. Other symbols are detailed in Table A in Appendix.
Research work
Global Flatness
Sharpness Measure
Communication Cost Computation Cost
FedSAM (ECCV22, ICML22)
Lm(win P Vis L. (10m))
(m
1x
2x
FedSpeed (ICLR23)
X
Similar to FedSAM
1x
2x
FedsMO0 (ICML233
V L. (u P PTVEm ) - Lm (20m)))))
2x
2x
FedTFFSM (ICML.24)
V Cm P
Pllwald ww Lm (m((1om))
wold -w²
FedGMT (Ours)
(xI(f()), ((W)))
x or 1.5x
1.33x
loss landscape. To achieve this, SAM applies weight per- turbation vector to the model parameters we and conducts the following objective:
FocAvg
FecSAM
FeeSpeed
Global Sharpness
F FAcSMOO
FOOLESAN-D
V((w)
FOOGMT
Ram
C CSAM (w) - L(w + E), where A P [V
(3)
0.5
Here, II is 12-norm, p is the radius of the neighborhood and V the abbreviation for V. on parameters w.
0.-
While Eq. (3) improves performance, computing and [(w+e) adds an extra backward and forward pass, doubling SAM's 's runtime compared to SGD. Minimizing Eq. (3) can also be interpreted as jointly minimizing the loss value and the sharpness of the loss landscape defined by:
Communication Rounds
Figure 2. The sharpness change of the global model on CIFAR-10 in the high data heterogeneity scenario (Dir(0.01)).
Sharpness: S(w) - [(w + e) - (ww).
(4)
Similar approaches like MoFedSAl edMAM (Qu et al., 2022) and FeaMROR (An et al., 2023) inject momentum into the local SAM optimizer and FedGAMMA (Dai et al., 2023) combines the SCAFFOLD (Karimireddy et al., 2020) with SAM im- prove the performance. Although these methods contribute to minimizing sharpness some extent, the flatness of the global landscape still can not be directly optimized. This is due to the fact that the perturbation em is not designed for the global function L.
2.2. Rethinking SAM- Algorithms
We first summarizes several SAM-based methods in FL at
Table 1. FedSAM (Qu al., 2022; Caldarola et al., 2022) directly applies the SAM objective to replace the local ERM objective, which can be formulated as change in Eq. (1):
(w) - C. (w + m, where fm is the local weight per- turbation allocated to C. Notably, the perturbation cm is computed based on the local dataset Dm Consequently, the utilization of this local SAM optimizer yields effective generalization only datasets drawn independently and identically Dm Intuitively, enhancing the generality of local training individually on D,, would inherently improve generalization on global dataset D = Um, Dm. However, this indirect promotion may not be effective when dealing with severely heterogeneous datasets.
To directly smooth the global model, is necessary to mod- ify fm - L. (w + <), where is the global weight perturbation allocated to L. However, c needs be cal- culated on the global dataset D, which is not feasible in FL. FedSMOO (Sun et al., 2023a) utilizes the ADMM to estimate global perturbation e and also adopts dynamic regularizer during the local training. FedLESAM (Fan et al., 2024) estimates the global perturbation as the difference between the global model from the previous round and the global model received in the current round. However, these estimations are not significantly effective compared to FedSpeed which directly utilizes local SAM in Fig. 2.
Fig. 2, we em- pirically demonstrate that FedSAM minimizes sharpness slightly compared to FedAvg in severely heterogeneous scenarios. This observation also explains why FedSAM does not outperform FedAvg in this case. FedSpeed (Sun et al., 2023b) incorporates local SAM optimizer with a dynamic regularizer to enhance its effectiveness. This approach aims to improve local consistency to bridge the smooth information on both local and global models.
Motivation: Given that SAM aims to yield a flat loss land- scape by minimizing sharpness during training, we utilize a first-order Taylor expansion decompose the sharpness
3

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
measure of SAM in Eq. (4) for the global function in Eq. (1) at round & as:
{00 0t ...... 01-1), where Os 0'~ a A 0 e and 0¹ = wt. The sharpness measure Spae (w²) in Eq. (4) is non-negative, and thus we have:
Sye (w') - Cyy (w" +o) - Loxr (w")
Lui (w²) +AAfr
Lize (w")
arg min (w²) w²
arg minSy (0')
0r
p. AKi /
min arg min| (Cr+ (0") Cxr (0')
(5)
arg arg min (e') - CNv (w"))
(8)
From Eq. (5), observe that the 12-norm of the global pseudo-gradient ر Il could be approximated as the global sharpness measure in FedLESAM according to the analysis in (Zhao et al., 2022). However, in practical FL. due to data heterogeneity and the high probability of varying active clients N' between consecutive rounds, the global pseudo- gradients in two consecutive rounds differ. Thus single previous global pseudo-gradient may not guarantee the right direction toward global flatness in the current round.
More details of this derivation can be found in Appendix C.2. Therefore, minimizing Lxs (e') Lxv (w") is equivalent to minimizing the SAM's sharpness measure Syx (w") for the global function L.
Based on the above result, if we directly combine the global trajectory loss term L(e²) - L(w=)) in Eq. (8) with FL objective L(w") in Eq. (1), the f(w") will unfortunately be canceled out. Without loss of generality, we replace the cross entropy loss with the KI. divergence loss to decouple the vanilla loss. Then the loss function Ln, of FedGMT for the m-th client at round as:
To address this, as a single global pseudo-gradient is inef- fetive and unreliable, our approach use the global model trajectory which includes the pseudo-gradien gradients from all the previous rounds. This significantly reduces sharpness com- pared to other SAM-based algorithms, as shown Fig. 2.
ERM
measure global sharpaess
Lm Lm( ee) == L (m(ww)) + Colotra (w) e')
Colotra (wi et)
(9)
Σ
3. Proposed Method: FedGMT
[Dml
Dm
In this section, we introduce F FedGMT and FedGMT-VV v2, an enhanced variant designed to save communication overhead. The detailed algorithm outlined in Algorithm 1.
where the hyperparameter y stands for the strength of mini- mizing sharpness.
Remark 3.1. Unlike the SAM, Calatra in Eq. (9) only requires one more forward pass and no extra backward pass.
3.1. Measure Sharpness via Global Model Trajectory
We denote the past trajectory of the global model as set Ω = {w", w'..... ==11). }. To avoid excessive storage and effectively utilize all the information Ωt we adopt an exponential moving average (EMA) weighting strategy on the global model's update trajectory. The EMA model ef at the t-round is updated as follows:
3.2. Promoting Global Consistency with ADMM
Note that an unbiased procedure to minimize global sharp- ness should use the global model w/ and the EMA model e², i.e., to minimize Lalotra (w", on the global dataset D. However, in FL framework, the global model w' is sent each client m and split into wo, for independent updates as Eq. (9). This process introduces non-vanishing biases w - w² to the global sharpness measure, potentially diminishing algorithm performance.
e¹ act (1 - a)w",
(6)
where E (0, 1) is the coefficient of EMA. On the other hand, given that e = w and w w' A then the relation between the EMA model el and the global model w as follows:
To bridge this gap, the local update should align with the global update. Therefore, we takes the form of an ADMM- like method in order to align client and server and effectively minimize the global objective L. ADMM makes use of the augmented Lagrangian function via penalizing the constrain w = w³. We first scale the KL function to non -egaaivte and convex function as:
t-1 ='=w' + a A N
+ 1=0
(7)
More details of this derivation can be found in Appendix C. Therefore, the EMA model et includes the pseudo-gradients gradients from all the previous rounds and puts more emphasis on the more recent pseudo-gradien
Celotra (w))) <
w
w m
> La'iss
(10)
2
To connect the EMA model e' to sharpness minimiza- tion, inspired by (Du et al., 2022) which uses loss differ- ence, we consider the trajectory from e' to = as =
k=0
where 3 is a penalty coefficient. More details of this deriva- tion can found in Appendix
4

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
By using the first-order derivative of Eq. (10), we define the global augmented Lagrangian function as:
Algorithm 1 FedGMT and FedGMT-v2 Algorithm
Input: Model parameters w, EMA model parameters e, communication round T, local interval K, dual variable u. EMA coefficient a. penalty coefficient 3, learning rate 1).
Σ m
(wmm ) um (w- wm ) 23 | Wmll mmll (11)
M-
Output: Global model parameters w².
where Im (wnmm) == Lm (wm))) ) Eizo a w and
1: Initialization w C = wD, thm = u = 0:
Um is the dual variable. We denote (v) as the inner product for two vectors.
2: for / = 0, 1, ... T do
3: randomly select the active clients set N from M;
To solve the minimization, we decompose the problem in
4: send the w² and el the active clients;
Eq. (11) for each client solve as:
5:
for client m E NN' in parallel do Initialize local model as w m,, w";
6:
1
K..K K
arg min F mm((f.)) - (umme w/n) + 2/3 win arg min Lm (m("m)))))))))) (umm w'n); wit
1|
w/ml ||
7:
e' e'= aa' + (1 a)w"; for = 0, 1,.1,.... K 1 do
(12)
9:
sample minibatch and do ymnk V Lm ( (
10:
Using Eq. (9)
where wit, m, x is the local model after K local interval updates.
11:
wn Wm,,+11 w n9 w end for um w'n 1 (w - -w");
Then we update the dual variable as wh++ = = 3 (w
12:
w²). Lastly, by using the first-order derivative of Eq. (11), the global model wo with partial participation is updated as:
13:
14:
send the w = w "m to the global server;
15:
end for
w+1 arg min wi
Σ
||um,kK w² - -Butt11"
16:
ut u'+1= ut' BM Emer ww));
N
17:
w 11 = N Emeri w - Bultl
MEN:
18:
el+1 = ace + (1 a)wlt1
(w -Bu(t
(13)
N
19: end for
MEN'
To avoid each client sending ++1 to the server each round, we define the global dual variable u²+1 := ut 3M Emery (w' w Then, reformat Eq. (13) as w+1 = EmEnt "nnkK But+1
model in the previous active round. Thus we propose the variant F FedGMT-v' The calculation of et in F FedGMT-v2 is approximation of that in FedGMT, which inevitably introduces slight performance degradation but reduces the communication cost by half for real-world practicality.
With these update rules, we can calculate the local mod- els of approximately satisfying the constraints. Intuitively, with the consistent local updates, minimizing the global trajectory loss Lelotro among clients would mitigate the impact of significant changes in training sample loss within the global loss function. This helps prevent overfitting to local datasets by enforcing their predictions to be close to the ones from the EMA model. Therefore, each client can search for a consistent smooth loss landscape aligned with the global objective, thereby significantly improving overall performance.
3.4. Convergence Analysis
In this section, we will demonstrate the theoretical analysis of our proposed FedGMT and illustrate the convergence guarantees under the specific hyperparameters. Firstly, some common assumptions are stated as follows.
Assumption 3.2. (T-smooth). Function (u) is L-- smooth and neural network fm (w) is Ly-smooth for all M and E Rdd
3.3. Enhanced Variant for Communication-EFfe Efficcien
Assumption 3.3. (Unbiased and bounded gradient). The stochastic gradients on a batch of client vn's data em is an unbiased estimator of V (Lm(w) with an upper bound G, i.e., LI Em) = V (w) and (VVC((W))) I G for all m E M, E Rd and j E [d].
The whole workflow of FedGMT is shown in Algorithm At the beginning of each round t, we randomly select subset of active clients N' from the total clients set M. The global server will communicate the parameters w' and e to the active clients for local training. Since the communi- cation with the server for e' increases communication cost by doubling the message size in each round, inspired by SCAFFOLD and FedLESAM, clients can utilize the EMA
Assumption 3.4. Neural network f. outputs a non-negative probabilistic vector, e.g. the last layer is softmax, then there exists 0 > 0 such that [i (w, (w, xz)le 2 8 > 0 for all m E M, x E dRd and w E Rdd Here, e is the class index.
5

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Assumption 3.2 assumes the gradient Lipschitz continuity for the objective function and neural network. Assumption 3.3 assumes the bounded stochastic properties of the gra- dients. The above two assumptions are commonly used in the analysis of the FI. framework (Luo et al., 2022; Chen et al., 2023). Assumption 3.4 assumes a lower bound for the neural network's s output, which is utilized to scale the KL function. This assumption is also employed in (Yao al., 2023; Yang et al., 2023) for convergence analysis.
(e) Path(2)
[b) Di
[c]
Long-taii Distribution
Number of samples
Classes
Client ib
Client
Classes
Figure 3. (a) & (b): Partition examples across 10 clients. (c): Long- tail distribution across category CIFAR-10.
Our theoretical analysis depends on the above assumptions to study the properties of the proposed method. Proof details could be referred to the Appendix D.
We represent the parameter as (rth(r) and Dir(r) to control data heterogeneity, with smaller values of r indi- cating increased data heterogeneity. Moreover, real-world data often exhibits long-tail distribution characterized by significant class imbalance, and many SOTA FL methods have poor performance (Shang et al., 2022b) in such case. Therefore, to further enhance the data heterogeneity and simulate real-world scenarios, we follow (Cao et al., 2019; Shang et al., 2022b) to shape the original balanced dataset into long-tail distribution for the Dirichlet setting. We set the imbalance factor (i.e., the ratio of instances in the most represented class to those in the least represented class) to 2 for all datasets. Figure 3 illustrates the differences in partition strategies and long-tail tail distribution.
Theorem 3.5. Let the above assumptions hold, when V/N and >
B I
VENT
V135+5/5MM
55L5 the mean averaged parameters sequence 2/2NT.5
W
Emen' w'm)
generated by the Algo-
tE[T-1]
rithm I under the non-comvex case satisfy:
T
T
[( - L' 20MB3I2 2 kT
1836 (14)
10 +
N
N
Baselines. We take two lines of state-of-ther of-the-art methods as baselines. ! heterogeneity-oriente FedAve (McMahan et al., 2017), FedDyn (Acar al., 2021) and FedNTD (Lee et al., 2022). ② SAM-based methods including FedSAM (Qu et al., 2022; Caldarola et al., 2022), FedSpeed (Sun et al., 2023b), FedSMOO (Sun et al., 2023a) and FeGLESAM-D (Fan et al., 2024) (FedLESAM with dynamic regularizer). The introductions and hyperparameters for each method in the Appendix E.
methods including
where 40 4== M EmEm Bllu - -w'1/2 is the inconsistent
term at the first round for w A + Emegre wh K. is a positive constant, L is the optima of C.
Remark 3.6. Under a constant number of clients uniformly selected random in each round, and suitable values of a, 8, and are chosen, Algorithm attains fast convergence rate of OL/TT) with non-convex local losses, which matches the conclusion of existing works (Acar et al., 2021; Gong et al., 2022; Sun et al., 2023a;b). Moreover, the last two terms in Eq. (14) achieve Nx x linear speedup.
Common Training Details. We set the batch size to 50 for CIFAR-10 and AG News, 20 for CIFAR-100 100, and 100 for CINIC-10. The number of local epochs is set to 5, ex- cept for CIFAR-100, in which we use 2 to mitigate extreme overfitting. We employ SGD with a learning rate of 0.01, momentum of 0.9, and weight decay of 1e-5, except for AG News, in which we use learning rate of 0.1. The total number of communication rounds is set to 500, except for CINIC-10 and AG News, in which we conduct 1000 rounds, ensuring sufficient rounds for performance saturation. Fol- lowing (Caldarola et al., 2022; 2023), we report the final averaged test accuracy and standard deviation over the last 50 rounds for increased robustness and reliability.
4. Experiments
4.1. Experimental Setup
Datasets and Models. We conduct extensive experiments containing computer vision (CV) and natural language pro- cessing (NLP) domains. For the CV domain, we con- sider the image classification tasks with three widely used datasets including 10FAR-10/100 (Krizhevsky Alex, 2009) and CINIC-10 (Darlow et al., 2018) using the CNN from (McMahan et al., 2017), ResNet-8 (He et al., 2016) and VIT (Dosovitskiy et al., 2021), respectively. For the NLP domain, we study the text classification tasks with AG News (Zhang et al., 2015) using FastText (Joulin et al., 2017).
FedGMT Setting. The hyperparameters of FedGMT remain consistent across datasets, model architectures, and scenar- ios. We set the KI. temperature to 3 and = =1 in Lolotra with the EMA coefficient a of 0.95 in FedGMT and 0.5 in FedCMT-v2 and the penalty coefficient / chosen from {10, 100) for all experiments. The detailed hyperparame- ters selection and sensitivity analysis in Section 4.4.
Heterogeneous Partition Strategy. We consider two com- mon data heterogeneity scenarios: Pathological (McMahan et al., 2017) and Dirichlet (Wang et al., 2020a) settings.
6

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Table Average (standard deviation) test accuracy @ ((%) comparison. pt the participation ratio (%) R& communication round
and computation time (minutes) to achieve target accuracy in brackets. X: the target accuracy is not available.
Non-IID Partition Strategy Pathological (Path)
Heterogeneity
Partial Participation
Communication and Computation
CIFAR-I0 (CNN)
CIFAR-100-Patt(10) (ResNet)
CINIC-10 (ViT)
AG News (FastText)
Method
Patl(1)) (1o)(10)
(2th(2)
20%
p= 10%
p= 5%
Path(2) R&T(41%)
Path(2) R&T(82%)
FedAvg
(8.9) (5.24(2.18) 65.28(5.30)
(1.33
(1.47 29.60(2.18)
(4.38
X
(1.3)
X
Redizyn
(0.66 (1.9) 34.13
(6.94(1.10)
45.04(1.00)
41.68
(1.63(6.30)
929
88.87
82.16; 8.37) 710 66.74
FedNTD
(o.65) (8.25( 73.36(1.75)
(2.53 (0.72(0.51)
36.59(0.s8)
49.27
491
53.19
77.60; 8.22)
x
FedsAK
80.60
(5.34(2.28) (6.07
(o.03
(1.23 28.97(2.26)
(6.34(4.19)
X
(1.26
X
redSpeed
(o.47 (1.50 (6.74(1.49)
(7.10
(1.44(1.03)
40.06;1.61)
42.98 (3.50)
688
149.30
85. 19; 3.51) 487 73.86
FedsMoo
(2.92
80.51 76.59
(7.48
(1.1) 40.81
44. 44.15 (3.533)
692
167.66 677699
84.84; 3.12) 548 83.80
Redi.RSAM D (0.22) 79.79
(2.9))
(7.09)
(4.96(
40.95(1.88)
(0.96(6.37)
X
85.37
2.67) 690 77.86
FedGMT
o3. 82.79 (0.22) 78.04(1.0) .3)
49.76(0.33)
48.19 (0.56) 46.11 (0.87)
(2.67(1.13) 392
40.38
89.39,
0.62)
329 33.72
FedGMT-V2
(3.15( (2.44(0.28) 78.01
(9.54( 47.95 (u.39) 39) (3.83
(1.77)
393
40.61
88.93, 0.81) 351 35.10
Non-1ID Partition Strategy: Long-tail Dirichlet (Dir)
CIFAR-I0 (CNN)
CIFAR-100-Dir(O... (ResNet)
CINIC-10 (ViT)
AG News (FastText)
Method
Dir(1.0) Dir(0.1) (ir(0.01)]
p 20%
p= p= 10%
p= 5%
Dir(0.1) R&T(44%)
Dir(0.1) R&T(85%)
FedAvg
(5.73 (3.51) 61.94(4.99) 93)
41.96/0.53)
39.34(1.00)
35.13
(9.77
X
(7.48
X
redlyyn
18.15 75.71 (2.52
(6.93(
44.81 (0.64) (3.04(0.84)
(4.00
985
80.11
(5.72
745 54.39
RedNTD
76.60(o.39)
(2.62 66.96(2.45)
(a.33
39.74 46) (0.52.
(2.65
654
53.41
83.55(2.50)
X
FedSAM
(0.93; 70.96(3.97) 61.05 (4.5)
41.78(o.63)
(8.95 (0.50) 60) 34.63(1.61)
(3.27)
X
80.2 80.25
X
FedSpeed
80.38 (o.5) 71.96(1.67)
48. (3o.33) (5.73 43,38 (0.38
(2.57
760
137.05
85 (2.00)
687 80.61
FedsMoo
(0.22 77.06 (o.0) (1.1)
(8.26
45.69 (0.75) (3.06
49.07
(2.29) 712
160.20
85, (5.74(2.16)
767 100.22
FedI.SAMM D (o.27) 76. (0.83) (1.10
46.97 (6.9) 45.21 (0.67) 42.39 (1.4)
(2.51)
X
85.68
745 67.30
FedGMT
80.48 79.17 74.67
48.65
47.42 28)
45.46 (0.37)
57.84(o.77) (0.77) 404
33.40
88.39 (0.67)
484 37.51
Fe
10.00 (8.73(0.47) 74.1 (1.32)
48.64(a.19)
(0.95) 44. (4..a)
57.61 (1.00)
413
35.45
88.31 (0.82)
485 38.07
4.2. Performance Evaluation
450
Dir 01) (61%)
Dir(0.1) (69%)
Datesets and Model Architecture. FedGMT consistently exhibits efficacy across diverse datasets and model architec- tures. In Table 2, the performance of many baselines varies significantly with different datasets and model architectures. Moreover, FedGMT-V2 v2 incurs only slight performance reduction compared to FedGMT across all settings.
400
Dir(1.0) (73%)
Commurication Cost
350-
Fodfvg
FedSAM
FedSaano
300
250
Heterogeneity. FedGMT is less affected by the degree of data heterogeneity. In Table 2, most baselines exhibit better performance in less heterogeneous settings (i.e., path( (10) and Dir(1.0)). However, as data heterogeneity increases (i.e., p2) and Dir(0.01)), they perform worse than FedGMT and FedGMT-v2.
200
Fe
150
RedSppeed
100
200
400
000
850
3000
1200
Computation Time (s)
Figure 4. Communication cost vs. computation time of F'edGMI, #edAvg and other SAM-based methods that achieve a target ac- curacy in brackets (%). Every connected line represents method that trains under different Dirichlet settings on CIFAR- 100.
Partial Participation. FedGMT demonstrates robustness to various levels of partial participation. As the participation ratio decreases, the data amount for training in one round decreases, exacerbating the degree of data heterogeneity. In Table 2, FedGMT experiences a smaller drop in accu- racy when the participation ratio changes from 20% to 5% compared to others.
Communication and Computation. Table 2 presents the communication rounds and computation time needed to achieve a target accuracy. Clearly, FedGMT achieves fewer communication rounds and requires less computation time. SOTA SAM-based methods, FedSpeed and FedSMOO,
7

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
5
ROLIESS
Rourss
CIFAR-10
(b) CIFAR-100
(c) CINIC-10
(d) AG News
Figure 5. Learning curves of FL methods on Dir(O.1).
Fed4ig
FMdSAM
FedGMT wis ADNM
TwttMMT
35
Start
End
04.
532
0.3
(a) Local Loss View
(b) Global Loss View
(c) Local & Global Loss View
Figure 2D visualization of local learning trajectory on CIFAR-10 10 (Dir(0.01)), (a): The loss contour is based on one client's s local dataset. (b): The loss contour is based on the global dataset. (c): The loss contour is based the merge of normalized (a) and (b).
both require more than 3.5x in CINIC-10 and 2x in AG News computation time compared to FedGMT. Moreover, we compare the communication cost (i.e., communication round * cost Table for SAM-based methods in Fig- ure 4. Although FecGMT and FedSMOO achieve better accuracy, they incur higher communication costs. We show that FedGMT-v2 can achieve fewer communication cost, while being less affected by the degree of data heterogeneity, making it more suitable for real-world practicality.
Table Ablation study on CIFAR-10.
ADMM
Colotra L'
Dir(1)
Dir(0.1)
Dir(0.01)
75. 75.73 (5.85)
(3.61)
61. 61.94(4.90)
78.18 (8.18
72.68
61.28(3.11)
78. 13..33)
75.97(1.09)
(8.99
80.48 (0.18) 79.17 (0.49)
74.67 (0.77)
Local Update Direction. In this part, we provide an in- depth analysis of "How does FedGMT guide local learning and why effective?"" To achieve this, we visualize the learn- ing trajectory of the local model on client utilizing the visualization method (Garipov et al., 2018). Specifically, train the global model with FedGMT for 100 rounds. Subse- quently, we choose one client to perform the loss functions of FedAvg, FedSAM, FedGMT without the ADMM, and FedGMT for local epochs, respectively. Their trajectories are then plotted on the 2D loss landscape with different loss views in Figure 6.
Learning Curves. To offer deeper insights into the learn- ing process, we depict the learning curves of various FL methods in Figure 5. For clear visualization, smooth these curves. Despite employing different communication rounds for each dataset, the model's performance reaches saturation at the end of communication rounds. Across all datasets, both FedGMT and FedGMT-v2 not only attains a superior final model by the end of the communication round but also exhibits much faster convergence compared to the other baselines.
We observe that both FedAvg and FedSAM update towards the local minimum, and F e'dSAM can explore flatter region than FedAvg from the local loss view. However, in Figure (b), their directions result in a significant increase in global loss, explaining their unstable performance under high data heterogeneity. In contrast, FedGMT steers each local update towards minimizing the change in global loss. From Figure 6 (c), F FedGMT can find region of joint minimum for both local and global objectives, ensuring that local updates align with the global update. Moreover, the ADMM in F FedGMT corrects the update direction towards reducing the joint loss, which can guarantee the effectiveness of the global sharpness measure.
4.3. Further Analysis
Impact of Each Component. Table 3, conduct an ab- lation study evaluate the contribution of each component in FedGMT. The results show that both ADMM and Colotra contribute performance improvement. It is worth noting that when under high data heterogeneity (e.g. Dir(0.01)), uti- lizing only glottra is ineffective, but when combined with ADMM, the performance significantly improves by 13.39%. The results validate the correctness and effectiveness of our algorithm design and theoretical analysis.
8

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Table 4. The test accuracy (%) of FedGMT with different hyperparametters CIFARIO.
o=0.99
7=
Items
y= 0
7y 0.5
y= 1.0
7 1.5
7 = 2.0
Y= 2.5
0.5
a = 0.9
a 0.95
0.995
0.9995
Dir(1.0)
78.13
79.46
80.48
80.37
80.82
80,93
78.73
79.10
79.42
80.48
81.12
81.25
Dir(0.1)
75.97
78.97
79.17
79.45
79.24
79.30
77.24
78.28
78.18
79.17
77.68
77.06
Path(4)
80.38
82.07
82.79
82.38
82.67
82.57
81.42
81.76
82.47
82.79
81.76
80.09
Puth(2)
75.21
78.11
78.04
77.79
78.13
77.94
78.21
78.23
78.42
78.04
75.71
73.06
FedSpeed (ace 79.28%)
(b) FedSMOO (acc 79.10%) (c) FedLESAM- 77.50% (d) FedGMT (acc -99.0%%
Figure 7, Visualization of the loss landscapes of the global model trained on CIFARIO 10)
Loss Landscape. In Figure 7, visualize the loss landscapes of global models obtained from FedsMOO, FedSpeed, FedLESAM-D and FedGMT after 500 com- munication rounds. The smoother of the loss landscape contribute to better generalization and higher accuracy, as seen with FedGMT, while less optimal landscapes could result in lower accuracy, in the case of FedLESAM D. These landscapes demonstrate FedGMT's effectiveness in smoothing the loss landscape among SAM-base methods.
5. Conclusion
In this work, we take closer look at sharpness-awan min- imization in heterogeneous FL from global view. We propose novel method FedGMT to directly reduce the sharpness of the global model in the FL framework. Theo- retical analysis guarantees that FedGMT achieves the fast convergence rate of O The extensive experiments demonstrate FedGMT achieves high generalization with fewer communication rounds and less computation cost. This work inspires the design of Fl. frameworks to prioritize reducing the sharpness of the global model.
4.4. Hyperparameters Sensitivity
Sharpness strength y. In Table 4, we compare the perfor- mance of the proposed FedGMT with different hyperparam- eters on the CIFAR-IO dataset with different heterogeneous settings. From the results, FedGMT exhibits insensitivity to the sharpness strength 7y Notably, without Colotra (i.c., y 0) in local learning, there significant accuracy drop of 3.2% and 2.83% compared to 7 1.0 on Dir(O.1) and Path(2), respectively. Since the improvement from 7 1.0 to 7 2.5 is negligible, it works well across all datasets. We recommend that the selection of from (0.5, 1.0, 2.0}.
Acknowledgments
This work partially supported by Natural Science Founda- tion of Shanghai under Grant 247R 1421500.
Impact Statement
set 7 1 for FedGMT and
Our work focuses on the optimization and generalization of federated learning and proposes a novel FL algorithm without compromising data privacy. The insight behind this work may inspire new researches. Since FL has wide applications in machine learning, Internet of Things, and Edge AI, work may be useful in these areas.
EMA coefficient a. The a cannot be selected too large, as results in significant information loss on the recent update trajectory, particularly under high data heterogeneity (e.g., Dir(0.1), Path(2)). We set 0.95 for FedGMT and also performs well across all datasets. We recommend that the selection of from (0.95, 0.995, 0.998). Since the global model's performance in FL. unstable during initial epochs, deploying a larger a may impede convergence. In our fu- ture work, we plan to implement an adaptive mechanism to further improve performances and reduce sensitivity of
References
Acar, D. A. E., Zhao, Y., Navarro, R. M., Mattina, M., What- mough, P. N., and Saligrama, V. Federated learning based dynamic regularization. In International Conference on Learning Representations (ICLR), 2021.

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
An, X., Shen, L.,, Hu, H., and Luo, Y. Federated learning with manifold regularization and normalized update reag- gregation. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Heigold, G., Gelly, S., et al. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
Boyd, S., Parikh, N., Chu, E.., Peleato, B., Eckstein, J., et al. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning. 3(1):1-122, 2011.
Du, J., Yan, H.,, Feng, J., Zhou, J. T., Zhen, L.,, Goh, R. S. M., and Tan, V. Y. Efficient sharpness- awwre minimization for improved training of neural networks. International Conference on Learning Representations (ICLR), 2021.
Caldarola, D., Caputo, B., and Ciccone, M. Improving gen- eralization in federated learning by seeking flat minima. In European Conference Computer Vision (ECCV), pp. 654-67 672. Springer, 2022.
Du, J.. Zhou, D., Feng. J., Tan, V., and Zhou, J. Sharpness- aware training for free. Advances in Neural Information Processing Systems (NeurIPS), 35:23439-23411, 2022.
Fan, Z., Hu, S., Yao, J., Niu, G., Zhang. Y., Sugiyama, M., and Wang, Y. Locally estimated global perturbations are better than local perturbations for federated sharpness- aware minimization. In International Conference on Ma- chine Learning (ICML), 2024.
Caldarola, D., Caputo, B., and Ciccone, M. Window-based model averaging improves generalization in heteroge- neous federated learning. In International Conference Computer Vision (ICCV), pp. 2263-2271, 2023.
Cao, K., Wei, C., Gaidon, A., Arechiga, N., and Ma, T. Learning imbalanced datasets with label-distribution- aware margin loss. In Advances in Neural Information Processing Systems (NeurIPS), pp. 1 1567-1578, 2019.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-aware minimization for efficiently improving generalization. International Conference on Learning Representations (ICLR), 2020.
Chaudhari, P., Choromanska, A., Soatto, S., LeCun, Y., Bal- dassi, C.,, Borgs, C., Chayes, J.. Sagun, L., and Zecchina, R. Entropy-sgd: Biasing gradient descent into wide val- leys. Journal of Statistical Mechanics: Theory and Ex- periment, 2019 12:011 8, 2019.
Ganchev, K., Taskar, B., and Gama, J. Expectation maxi- mization and posterior constraints. Advances in Neural Information Processing Systems (NeurIPS), 20, 2007.
Garipov, T., Izmailov, P., Podoprikhin, D., Vetrov, D. P., and Wilson, A. G. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018.
Chen, H., Vikalo, H., et al. The best of both worlds: Ac- curate global and personalized models through feder- ated learning with data-free hyper- knowlege distillation. International Conference Learning Representations (ICLR), 2023.
Gong, Y., Li, Y., and Freris, N. M. Fedadmm: A robust federated deep learning framework with adaptivity to system heterogeneity. In International Conference on Data Engineering (ICDE), pp. 2575-2587, 2022.
Dai, R., Yang, X., Sun, Y., Shen, L.,, Tian, X., Wang, M., and Zhang, Y. Fedgamma: Federated learning with global sharpness-aware minimization. IEEE Transactions Neural Networks and Learning Systems, 2023.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Darlow, L. N., Crowley, E. J., Antoniou, A., and Storkey, A. J. Cinic-10 is not imagenet or c0far-10. arXiv preprint 2018
Hsich, K., Phanishayee, A., Mutlu, O., and Gibbons, P. The non-iid data quagmire of decentralized machine learning. In International Conference on Learning Representations (ICLR), pp. 4387-4398. PMLR, 2020.
Deng, J., Dong, W., Socher, R., Li, L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In Conference on Computer Vision and Pattern Recogni- tion (CVPR), pp. 248-2482 255.. Ieee, 2009.
Hu, M., Cao, Y., Li, A., Li, Z., Liu, C., Li, T., Chen, M., and Liu, Y. Fedmut: Generalized federated learning via stochastic mutation. In Association for the Advancement of Artificial Intelligence (AAAI), volume 38, pp. 12528- 12537, 2024a.
Dinh, I.,, Pascanu, R., Bengio, S., and Bengio, Sharp min- ima can generalize for deep nets. In International Con- ference on Machine Learning (ICML), pp. 1019-1028.
PMLR. 2017.
Hu, M., Yue, Z., Xie, X., Chen, C.,, Huang, Y., Wei, X., Lian, X., Liu, Y., and Chen, M. Is aggregation the only choice? federated learning via layer-wise model recombination.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn,
D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M.,
10

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
In In Proceedings of SIGKDD Conference on Knowl- edge Discovery and Data Mining (KDD), pp. 1096-1107, 2024b.
Li, Q.. Diao, Y., Chen, Q., and He, B. Federated learning non-iid data silos: An experimental study. In Inter- national Conference on Data Engineering (ICDE), pp. 965-978 IEEE, 2022.
Hu, M., Zhou, P., Yue, Z.,, Ling, Z., Huang, Y., Li, A., Liu, Y., Lian, X., and Chen, M. Federoss: Towards accu- rate federated learning via multi-model cross-aggregation. In IEEE International Conference on Data Engineering (ICDE), pp. 2137-21 IEEE, 2024c.
Li, T., Sahu, A. K., Zaheer, M.. Sanjabi, M., Talwalkar, A., and Smith, V. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429 450, 2020.
loffe, S. and Szegedy, C. Batch normalization: Accelerat- ing deep network training by reducing internal covariate shift. In International Conference on Machine Learning (ICML), pp. 448 456. pmlr, 2015.
Luo, M., Chen, F., Hu, D., Zhang. Y., Liang. J., and Feng. J. No fear of heterogeneity: Classifier calibration for feder- ated learning with non-iid data. Advances in Neural In- formation Processing Systems (NeurIPS), 34:5972-5444 2021.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. In International Conference Learning Representations (ICLR), 2019.
Luo, Z., Wang, Y., Wang, Z., Sun, Z., and Tan, T. Disen- tangled federated learning for tackling attributes skew via invariant aggregation and diversity transferring. In International Conference on Machine Learning (ICML), pp. 14527-14541, 2022.
Joulin, A., Grave, E.,, Bojanowski, P., and Mikolov, Bag of tricks for efficient text classification. Conference of the European Chapter of the Association for Computational Linguistics (EACL), 2017.
McMahan, B., Moore, E.,, Ramage, D., Hampson, S., and y Arcas, B. A. Communication--fffiieee learning of deep networks from decentralized data. In International Con- ference on Artificial Intelligence and Statistics (AISTATS), pp. 1273-1282. PMLR, 2017.
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T. Scaffold: Stochastic controlled averag- ing for federated learning. In International Conference Machine Learning (ICML), pp. 5132-5143, PMLR, 2020.
Mendieta, M., Yang, T., Wang, P., Lee, M., Ding. Z., and Chen, C. Local learning matters: Rethinking data hetero- geneity in federated learning. In Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8397-8406, 2022.
Krizhevsky Alex, H. G. Learning multiple layers of features from tiny images. In Tech. Rep., pp. 32-33. University of Toronto, 2009.
Kwon, J., Kim, J., Park, H., and Choi, K. Asam: Adaptive awarre aware minimization for scale-invariant learn- ing of deep neural networks. In International Conference Machine Learning (ICML), pp. 5905-5914. PMLR, 2021.
sharpness-
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, I.,, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 11:22)8-22223 2324,, 1998.
Qu, 7., Li, X., Duan, R., Liu, Y., Tang, B., and Lu, 7. Gen- eralized federated learning via sharpness aware minimiza- tion. In International Conference on Machine Learning (ICML), pp. 18250-18280. PMIR, 2022.
Lee, G., Jeong, M., Shin, Y., Bac, S., and Yun, S.-Y. Preser- vation of the global knowledge by not-true distillation in federated learning. Advances in Neural Information Processing Systems (NeurIPS), 35:38461-38477 2022.
Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., Kumar, S., and McMahan, H. B. Adap- tive federated optimization. International Conference on Learning Representations (ICLR), 2021.
Lee, G., Jeong, M., Kim, S., Oh, J., and Yun, S.-Y. Fedsol: Bridging global alignment and local generality in feder- ated learning. arXiv preprint arXiv:2308. /2532, 2023.
Shang, X., Lu, Y., Cheung, Y-m., and Wang. H. Fedic: Federated learning on non-iid and long-tailed data via calibrated distillation. In 2022 IEEE International Con- ference on Multimedia and Expo (ICME), pp. 1-6. IEEE, 2022a.
Li, P., Cheng, G., Huang, X., Kang, J., Yu, R., Wu, Y., and Pan, M. Anycostffl Efficient on-demand federated learn- ing over heterogeneous edge devices. In IEEE Conference on Computer Communications (INFOCOM), 2023.

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Shang, X., Lu, Y., Huang, G., and Wang, H. Federated learning heterogeneous and long-tailed data via classi- fier re-training with federated features. In International Joint Conference on Artificial Intelligence (ICAA), pp. 2218-22244 2022b.
Zhang, J., Hua, Y., Wang, H., Song, T., Xue, Z., Ma, R., Cao, J., and Guan, H. Gpfl: Simultaneously learning global and personalized feature information for personal- ized federated learning. In International Conference on Computer Vision (ICCV), pp. 5041-5051, 2023a.
Shi, Y. Shen, L., Wei, K., Sun, Y., Yuan, B., Wang, X., and Tao, D. Improving the model consistency of decentral- ized federated learning. In International Conference on Machine Learning (ICML), 2023.
Zhang, J., Hua, Y., Wang. H., Song, T., Xue, Z.,, Ma, R., and Guan, H. Fedala: Adaptive local aggregation for personalized federated learning. In Association for the Advancement of Artificial Intelligence (AAAI), volume 37, pp. 11237-11244, 2023b.
Sun, Y., Shen, L., Chen, S., Ding. L., and Tao, D. Dynamic regularized sharpness aware minimization federated learning: Approaching global consisteney and smooth landscape. In International Conference Machine Learning (ICML), 2023a.
Zhang, J., Liu, Y., Hua, Y., Wang, H., Song, T., Xue, Z., Ma, R., and Cao, J. Pfllib: Personalized federated learn- ing algorithm library. arXiv preprint arXiv:2312.04992 2023c.
Sun, Y., Shen, L., Huang, T., Ding, L.,, and Tao, D. Fed- speed: Larger local interval, less communication round, and higher generalization accuracy. In International Con- ference on Learning Representations (ICLR), 2023b.
Zhang, X.,, Zhao, J., and LeCun, Y. Character-level con- volutional networks for text classification. Advances in Neural Information Processing Systems (NeurIPS), 2015.
Zhao, Y., Zhang, H., and Hu, X. Penalizing gradient norm for efficiently improving generalization in deep learn ing. In International Conference on Machine Learning (ICML), pp. 26982-26992. PMLR, 2022.
Wahab, O. A., Mourad, A., Otrok, H., and Taleb, T. Feder- ated machine learning: Survey, multi-level classification, desirable criteria and future directions in communication and networking systems. IEEE Communications Surveys & Tutorials, 23(2): 1342-1392919 2021.
Wang, H., Yurochkin, M., Sun, Y., Papailiopoulos, D., and Khazaeni, Y. Federated learning with matched averaging- International Conference Learning Representations (ICLR), 2020a.
Wang. J., Liu, Q., Liang, H., Joshi, G., and Poor, H. V. Tack- ling the objective inconsistency problem in heterogeneous federated optimization. Advances in Neural Information Processing Systems (NeurIPS), 33:7611-7625 2020b.
Wang, S., Xu, Y., Wang, Z., Chang, T.-H., Quek, T. Q., and Sun, D. Beyond admm: a unified vieet-vvrrance-rrddued adaptive federated learning framework. In Association for the Advancement of Artificial Intelligence (AAAI), volume 37, pp. 10175-10183, 2023.
Wu, Y. and He, K. Group normalization. In European Conference Computer Vision (ECCV), pp. 3-19, 2018.
Xu, J., Wang, S., Wang, L., and Yao, A. C.-C. Fedem: Federated learning with client- momentum. arXiv preprint arXiv:2 106.//1.. 10874, 2021.
Yang, Y., Liu, C., Cai, X., Huang, S., Lu, H., and Ding, Y. Unideal: Curriculum knowledge distillation federated learning. arXiv preprint arXiv:2309.0880
2023.
Yao, D., Pan, W., Dai, Y., Wan, Y., Ding, X., Yu, C., Jin, H., Xu, Z., and Sun, L. Fedgkd: Towards heterogeneous fed- erated learning via global knowledge distillation. IEEE Transactions on Computers, 2023.
12

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Appendix
We provide details omitted in the main paper.
Appendix A: Table of Notations throughout the paper.
Appendix Related work.
Appendix C: Derivation of Equation (7) and (8).
Appendix D: Proof of Theorem 3.5 (cf. Section 3.4 of the main paper).
Appendix E: Details of experimental setups.
A. Table of Notations
Table Table of Notations throughout the paper.
FL: M, N, and M.
total number, sampled number, and index of clients whole client set, sampled client set at round / whole dataset, local dataset total number and index of communication rounds global model parameters, local model parameters at round total number and index of classes global learning rate, local learning rate
D. Dm
T,t
w', w
C,e
f
neural network local interval
K
SAM:
C, Gm
weight perturbation allocated to global function C, local funcion C. ascent step learning rate in SAM's objective
p
v
abbreviation for V - parameters sharpness measure of SAM (equal to L(w + +²) L(w))
S(w)
A
approximated global gradient at round t ( A = =w' - 11
Algorithm:
decay coefficient of EMA penalty coefficient strength of minimizing sharpness temperature in KL function EMA model
7
u m
dual variable on local, global
Functions:
Lm
local ERM for client m
Calotra
global trajectory loss for client m
m Lm
loss function of FedGMT for client m function for convergence analysis
F
13

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
B. Related Work
The widely known classical FL method FedAvg (McMahan et al., 2017) learns single global model for all clients by aggregating their local models. Although FedAvg provides practical solution, it still suffers from the heterogeneity of data across clients (Li al., 2022). With further studies in FL, this is summarized as client drifts (Karimireddy al., 2020) and inconsistency objectives among clients (Wang et al., 2020b; Shi et al., 2023).
To alleviate the limitations of FedAvg, many methods utilize the parameter (or gradient) difference between the local and global model to assist the local training. By incorporating the global model information into local training, the bias between the local and global objectives can be diminished at some level. These approaches can be divided into the following three directions: ① Many algorithmic solutions in (Li et al., 2020; Karimireddy et al., 2020; Acar et al., 2021; Gong et al., 2022; Wang et al., 2023) mainly focus on mitigating the inconsistency across clients via giving a variety of proximal terms to control the local model updates close the global model. 2 Momentum based methods (Xu et al., 2021; Qu al., 2022;
An et al., 2023) introduce the global momentum information into the local training directly, which can force local updates to be similar. ( Techniques based on knowledge distillation (Lee et al., 2022; Yao et al., 2023: Chen et al., 2023) utilize the global model to guide local models, preventing them from forgetting knowledge from the server. While these algorithms accelerate convergence, they often encounter challenges in sharp global landscapes under high data heterogeneity, leading to unreliable minima and poor stability (Qu et al., 2022; Sun et al., 2023a). Therefore, the global model may not be efficient for all clients, leading significant deviation.
Recently, new research direction in FL has explored the generalization of the global model by analyzing the loss landscape, establishing connections with convergence to flat minima (Mendicta et al., 2022; Qu et al., 2022: Caldarola et al., 2022; Hu et al., 2024a;b;c). Most studies in this direction aim to seek flat minima of the global model with higher generality, utilizing the recently proposed Sharpness-Aware Minimization (SAM) (Foret et al., 2020) optimizer as the local optimizer. Existing works (Foret et al., 2020; Kwon et al., 2021; Du et al., 2021) demonstrate that the flat loss landscape approached by the SAM optimizer exhibits higher stability and generality. In FL, (Qu et al., 2022) first incorporate the SAM optimizer to inherently enhance consistency. Subsequently, various works such as (Caldarola al., 2022; Sun al., 2023a;b; An et al., 2023; Lee al., 2023; Fan et al., 2024) propose different variants to adopt the SAM optimizer.
14

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
C. Derivation of Equation (7) and (8)
this section, we refer the analysis previous research (Ganchev al., 2007; Du et al., 2022) to derivate the equation.
C.1. Derivation of Equation (7)
= ac'"! + (1 - a)w
= alac
(-2 + (1 a)w"-1)+ +1(11 (1-1000 aawu
=
a'e (1 a)(w' -ww + aww'-2 + w'=lw").
(15)
We recall that:
w' w. 1--1 w'-2 A A-1 w Σ A'
(16)
1=0
Substituting Equation (15) into Equation (16), we obtain
t-1
1-2 -
e' a'e + (1 a)
1 Σ
A
+a
+
wd
Σ A
.... +.... (w0 (20)
1-1
d'e (1 a)
A
1
1 a
i=0
1-1
= =d'e + (1 a² w')
Σ
(1 - a' D4
1=0
= =
)''
w/ +
A
(17)
15

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
C.2. Derivation of Equation (8)
arg min.Spr (w") = arg min (0*) (a)
arg min Al (O*)]lAA (0')l
gt
arg min gr
cos (+0)"")) (0')ll + Σ (1)) (0)))
<<
arg min 0*
cos([T)) (0') AN (8')l
arg min
a cos (T) AxT (0') (0')lI
(r))
arg min 0
E
[a' Axr (0') ANI (0*)
((e)
arg min
E
[Cx=(0') - LNv (0' -a -a Ax::))
(lt(())
arg min
E
[exx (0') Lx (0Ct
(r~((e)
arg min [Cxi (09) Lxx (0') +... + Lixe (0--1) Loxe (0))
8r
arg min [NN (88) Lru (8*)
arg min LNy (e*) - Lxv (w')]
(18)
where 1 is the angle between the global gradients that are computed based on active client sets N' and N', 0r ~ U())
means that 0² is uniformly distributed in the set θ. (a) is from Equation (5). (b) is due to the fact that for < t
a a (TT() (0') Ars [6')] is a constant with respect to the variable H*, (c) utilizes a first-order Taylor expansion.
16

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
D. Proof of Theorem 3.5
D.I. Preliminary
Before proving the theorem, we first introduce some preliminary conclusions used in our proofs. Definition D.1. L. is L smooth if
VV (2) (x) V Lm (1)ll - L|| - yll Va,
(19)
Smoothness implies the following quadratic bound,
L Ly( S (a) (VLm(t))) x) + E1ly - x|1? Vx, y-
(20)
Lemma D.2. For random variables I1,.... Ina we have
|| + + zull? S nB|[[2111 +... + ||
(21)
Lemma D.3. For random variables x. y, have:
1 :)
| + y11 v 1
||0|| + (1 + +c)+c)11 1//1/3
(22)
where e >> 0 is a constant.
Lemma D.4. For two random variables we have;
1
(a. v) S 2 211x0112
2 =11 y112.
(23)
D.2. Proof of Equation (10)
Calotra (w'm)
γ
Σ Gx)
DDml
(a)
(1(()); - (w Ex) 2
D
minje(l.". .C) f(w Ei)
DED
(b)
LLg
le"
w ||
(c)
bbا
[-1
(d)
1-1
w m
w
(24)
23
==0
==0
where (a) use the technique that used in (Yao et al. (2023), Lemma 3). (b) based on Assumption 3.2 and 3.4, (c) is based on the relation of e and w in Equation (7). (d): let 23L,
Thus complete the proof.
17

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
D.3. Proof of Theorem 3.5
For convenience, we rewrite our FedGMT local objective as
1-1
2
L (w(()) Cm(((m) +
1
w²
+
Σ
a
A
w
(u/m
wm))
(25)
23
based on Equation (9) and (10).
We define
F Fm( L. Lu((mm))
aa'is
wm
(26)
1=0
Then the local objective becomes:
I (m(w")) - (u w +
2 Ilo w²11
(27)
Throughout the proof, we utilize similar techniques in (Karimireddy et al., 2020; Acar al., 2021; Sun al., 2023a). We define of variables which are useful in the analysis. With the partial participation training, Algorithm freezes wh, if the client is not active at round Thus, we define virtual variables w as:
arg min { (m(() (uim w) +
1 w w² 10
mf
E M.
(28)
2B
w
The virtual variable w is based partial participation. Considering the first-order gradient condition of the objective in Equation (28) follows:
V V. ( - in + [ - w²) = O, m E M;
V.Pm Vm (w
- u/ + B
(wi -w'-w=) =0m
(29)
We that wi = wht m E No and Th, does not depend N', which means that w equals to wi with probability N and maintains wifz1 otherwise. Then consider the update of u+t1 uh, (wc, w²), ) we infer u V (w round t.
In order to distinguish the parameters before and after updating with the penalty term, we define:
1
Σ w/ = w + =r+++1
(30)
N
mENt
Next, we introduce the following Lemma to assist in the proof.
Lemma D.5. Algorithm satisfies
E [
w0]
B M
Σ
E
V
(10
(31)
m
Proof.
E
(a)
1
(wi n6 - -3A
(b) E
w² - -320) (int
N
mEN
mEM
(c)
E
V m
(wtt
Σ E
m (w (wm))
M
M
MEM
MEM
where (a) from definition in Equation (30), (b) is due to each client selected with probability (29) and (d) due to the relation v' Emem unn
N M (c) is from Equation
18

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Lemma D.6. Algorithm satisfies
E
1
Σ
Σ E
T E
- w w²1
(32)
M
mEM
Proof.
E
11
11
(a)
E
Σ
(w
w²)1
w
(b)
1
E
Σ
w
w'11
@
Σ
E
w²1l
(33)
N
N
m
M
MEN*
mMEN
MEM
where (a) from Equation (30), (b) from Lemma D.2 and (c) due to each client is selected with probability
Note that we use [V(w")] the global gradient norm as the metric the convergence analysis of FedGMT. To achieve this, starting from the assumption smoothness with Equation (20), taking the full expectation on both sides, have:
E [[]
< E [f(w²)]
+
-Blluttt--o111 B(VL(W")
w0 - r0t)
L
E [[(w")
E
w'
+ BE
VL(W"),
Σ
V
m (w
2
m
MEM
(b)
L
8 w0l 10
1
3 B
E
[((T))
E
E
V
V
(w vL(TT')
(34)
2
m
mEM
42
where (a) from Lemma D.5 and (b) is from Lemma D.4.
From Equation (34), we obtain the global gradient norm [VL(w))] which we want to bound. Next we need bound A1 and A2 to complete the proof.
For convenience of expression, we define some quantities that we aim control.
V
1 M
Σ E. MEM
To w. 11
Th
Σ E w/ MEM
1
7' = a 1-1 A
M
2=0
Here, T keeps track of how much local models change compared the average of client models from the previous round, representing the average local updates, & tracks how well local models approximate the current active client average, representing the inconsistency among clients, Zt tracks the historical global update, representing the cumulative global update with decay. Note that if models converge, y and & will be and z' approaches 0 since a E (0.
Lemma D.7. Based on the assumptions. the Al and oi term could be bounded as:
(2) E El 2 S Te 10,32
5,3 1p + +50 +5.
V[(w')| 5 1
(35)
19

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Proof.
(a)
E |
w²11
S
Σ
1 T01 M
E
Bl
E
w' - 31
M
MEM
EM
(c)
8
El| m
- V ( - 10
(d)
g2
Σ
E E
FnVmmm>> (w/m - V [1 ) - =21
M
M
MEM
MEM
82
1-2
2-2
E V Lm (w
_
A
V T (( +
Σ -'
M
3
MEM
B2
V mm (on(ww (w
ve( - ( + VL(w") V2(W")11
M
MEM
8²
(mm VLm (w') + VLm v(( + VL(W") V vL (wm)
M
MEM
[-1
B
(1)
5,32
5,32
Σ
V Cm (w - VLm (C 112 +
E L (
V (w()
M
M
MEM
MEM
5
Σ
V (T')
VL (wm ( 5,8 V
5E Σ a '
MEM
(8)
10.3 L"
5.3²
E Bl 1-1 w²112 +
E I'' - w
12
53
VL
>|
+
5
Σ f=
M
m
M
MEM
MEM
10.² -
+
5,8 ---"
where (a) is from Lemma D.6, (b) is from Equation (30), (c) is from Equation (29), (d) and (e) is due to the relation
2' Em ME u u V.Fmm m (wit.. the definition of F m in Equation (26) and the definition of L in Equation 1, (f) is from Lemma D.2 and (g) is from L-smooth.
E
Lemma D.8. Based on the assumptions, the A2 term could be bounded as:
2
E
Σ
-V
2
2 32
1
(36)
MEM
Proof.
E
1
(VIm(( -V())) |
(a) 1
mm (wm) - vL((z))
M
M
MEM
MEM
1-1
2
Σ MEM
V -
Σ
a
'
M
3 i=0
f-1 Σ
2
Lm (wh)) - V
+
E
A
M
MEM
1=0
2L
E
2
+
2
E
21'2
M
MEM
im0
2
21. v
11
where (a) and (c) are from Lemma D.2, (b) is from the definition of Fm in Equation (26) and (d) is from smoohh..
20

---

One Arrow, Two Hawks: Sharpness-awary Minimization for Federated Learning via Global Model Trajectory
We that the result of Equation (35) and (36) have the terms Φ'-1 and E| 1/- Next, we aim bound this two terms. Lemma D.9. Based on the assumptions, the E 22 term could be bounded as:
3a 30 21 27a 27 E 1 3a " -
(37)
Proof.
Let 0 = a
-i attl-i= A. then have
= =z + aad = az + +a(w" w
(38)
Therefore, we have
E
2- = +a(w" - w y1
(a) Bl +a(w" Bu' - But+
bb)
BllaZ
+ a W - BV ((w z 1h + +BV((Wm) - 211
(
Tr wt+1 + -t- 2
(
3
Σ
V
1 + 2 233
mEM
QBEII Lm (wth) - 26-1 + z - 2
5 a² z 2 3/G1|
(x) I 30 Z
1
30 B Zt 2112
where (a) is from Equation (30), (b) and (e) due to the relation u' = M Fm in Equation (26) and the definition of L in Equation 1, (c) and (f) is from Assumption 3.3, (d) employs the same proof procedure used in Lemma D.5 and (g) is from Lemma D.2.
E Emeta u u = V.Fm(w, the definition of
Following Equation (37), we can obtain that E| < 3a2 and E/ (0 442).. fort 2
Then we can get
(a) E||Z'11
< 27023G'
(
+5
1=2
- 27a2282 G²
Vba +1) +6
aG { 0G:
(39)
where (a) utilizes the result of Equation (37) with 2 2 and E Z 2 obtain the solution of this inequality. Here, Ç is small constant with exponential decay (omitted for brevity). (b) applies the sum of geometric series with < We use OG denote the result of Equation (39) for convenience.
V6'
Lemma D.10. Based on the assumptions, the ø' term could be bounded as:
2N
2M
2M
2NV
Tt <
T
o
(40)
2M
N
N
2M
N
21

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
Proof.
Φ
1
E W
w
1
1
E I
w w
112
M
M
MEM
MEM
N
1
2M
1
Ell0 W
w01l
+
}B||
w 1|
2M
N
m
N
MEM
(b)
N
1 N
E E - w'
(1
N
(1
N
E W 1-1
w
M
2M N
M
2M
N
M
MEM
MEM
2M - N
(1 +
E
W
w 11
(0)
N
1+
N
+ (1 )
2M N
N
(1 +
N
M
2M
N
N
2M
N
2N
2M
2M 2NV
+
Vr
2M
N
N
2M N
where (a) from Lemma D.3, (b) is due each client is selected with probability and (c) is from Lemma D.6.
Collecting and organizing the Equation (34), (35), (36), (39) and (40), we have
E [[ []] S E [[(])] [ + - + 3
(41)
(1-582 L2)U'' S 2 r- + 5,3 22
[( )1² 2 + 50c
(42)
2N
2M
V' 2M - -2N
Φ'
(43)
2M N
2M
Let Equation (42) multiplied by constant 2 and formula (43) multiplied by constant V take the sum of Equation (41), (42) and (43)
E [[( + (1 5582133
V < E
[[(T')]
+ -/2 + v²
3 GG -EE-EEEVLL V
1
+ 30.2".22 2/" + 552 3
V [(("'\" + 3. OG
2N
2M
2M - 2N :'-1
+ ②
V¹ + Φ
(44)
2M N
N
2M N
Collecting the like term of & and let the constants 2 and y satisfy:
211 = 3L2 -
2N
2M
+
(45)
2M N
N
When Equation (45) holds, Equation (44) will be simplified to:
1 E [[( + N0 < E [[(w)) + ( + +52 )oG + 33.2
B
2
2M 2N
(102 1.2 ạ
2M N p
(46)
Furthermore, considering the coefficient of & term, we let the constant V and y satisfy:
2M - 2N
y 10 ②
(47)
2M N
According to the Equation (45) and (47), we can get the solution of 2. and V as:
L(N + 28LS2)
2 2(N2 40 25
(48)
203
22

---

One Arrow, Two Hawks: Sharpness-aware Minimizztio for Federated Learning via Global Model Trajectory
+ -N)
y0 N 40,3 - 253 20
(49)
This proof requires the constants V and y both be positive. Moreover, the coefficient of the glohal gradient term must
maintain positive value.Thus the 3 satifies:
25. + 20,3 N 40,0
(50)
522 00
(51)
2
We can get the solution of 3 as: 3 <
V2N
v V135+5V5MI
We can rewrite Equation (46) as:
3 E V (E [[(w")] 1/-1//1----- (E [((w + + 1 + +5%)oa
(52)
Let applying 2
and y <
1. the above formula can be telescoped as
1033
T
((
302 -
Φ0
+ 1 11 ( +52 227a
+522)27a
(1-Voa+1) v1)
(53)
T
KT
N
Similar to Qu et al. (2022); Sun et al. (2023a), select the o(+)) that the final convergence rate approaches O which completes the proofs. For simplicity and convenience, We choose S JONT the above formula can be simplified as:
V KT (C(w') L' 202
To 183G N
O ((1
(54)
T
This makes Theorem 3.5 hold.
23

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
E. Experimental Setups
E.1. Environments
The code is implemented by -YTochh-L.13.. (Paszke et al., 2019) and the overall code structure based on PFLlib (Zhang al., 2023c) library with some modifications. All experiments are conducted on Linux (Ubuntu-20.04 LTS) server with one NVIDIA GeForce RTX 4090 GPU.
E.2. Datasets
To validate our approach, we consider image and text classification task and adopt four widely used datasets, i.e., CIFAR- 10/100 (Krizhevsky Alex, 2009), CINIC-I0 (Darlow et al., 2018) and AG News (Zhang al., 2015). Note that CINIC-10 constructed from ImageNet (Deng et al., 2009) and CIFAR-I0, whose samples are very similar but not drawn from identical distributions. Therefore, it naturally introduces distribution shifts which is suited to the heterogeneous nature of federated leamning. The details about each datasets and setups are described in Table 6.
Table 6. Details datasets setups used in the experiment.
Datasets
CIFAR-10
CIFAR-100
CINIC-10
AG News
Datasets Classes
100
4
Datasets Size (train/test)
50,000/10,000
50,000/10,000
90,000/90,000
120,000/7,600
Number of Clients
100
100
100
100
Client Sampling Ratio
0.1
0.1
0.1
0.1
Local Epochs
2
5
5
Batch Size
50
20
100
50
Learning Rate
0.01
0.01
0.01
0.1
Base Optimizer
SGD
SGD
SGD
SGD
Momentum
0.9
0.9
0.9
0.9
Weight Decay
le-5
le-5
le-5
1e-5
Rounds
500
500
1000
1,000
We standard pre-processing, where CIFAR- 10/100 and CINIC-10 10 images normalized. The 32 32 CIFAR-10/100 and CINIC-10 images are padded pixels each side, randomly flipped horizontally, and then randomly cropped back to 32 x 32. For AG News, the maximum length of the sentence vector is set to 200.
E.3. Model Architecture
We conduct further experiments on different model architectures: CNN (LeCun et al., 1998) for CIFAR-10, ResNet-8 (He et al., 2016) for CIFAR-100 and ViT (Dosovitskiy et al., 2021) for CINIC-10 and FastText (Joulin et al., 2017) for AG News.
The CNN used in our experiment is from FedAvg (McMahan et al., 2017), similar architecture is used in (Luo et al., 2021; Lee et al., 2022).
The ResNet-8 in our experiment from PFLlib (Zhang al., 2023c) library, a similar architecture is used in (Shang et al., 2022a;b). We follow the suggestion of Hsich et al. (2020) replace the Batch Normalization (loffe & Szegedy, 2015) with the Group Normalization (Wu & He, 2018) avoid the non-differentiable parameters.
The ViT adopted from ViT-CIFAR which is a smaller version compared to the original ViT.
The FastText is from PFL (Zhang et al., 2023c) library, similar architecture is used in (Zhang et al., 2023b;a).
24

---

One Arrow, Two Hawks: Sharpness-aware Minimization for Federated Learning via Global Model Trajectory
E.4. Baselines Implementation Details
We consider the following six state-of-the-ar FL methods:
FedAvg (MeMahan et al., 2017) is proposed as the basic framework in the federated learning, which aggregates the locally trained model parameters by weighted averaging proportional to the amount of local data that each client had.
FedDyn (Acar al., 2021) dynamically updates its regularizer so that the optimal model for the regularized loss is in conformity with the global empirical loss.
FedNTD (Lee et al., 2022) conducts local-side distillation only for the not-true classes to prevent forgetting.
FedSAM (Qu et al., 2022) directly applies the SAM objective in local learning.
FedSpeed (Sun et al., 2023b) incorporates local SAM optimizer with dynamic regularizer.
FedSMOO (Sun et al., 2023a) utilizes the Alternating Direction Method of Multipliers (ADMM) to estimate global perturbation and adopts a dynamic regularizer during the local training.
FedLESAM (Fan et al., 2024) estimates the global perturbation as the difference between the locally stored historical model from the activation round and the global model received in the current round.
For the above algorithms, search hyperparameters and choose the best among the candidates. All methods hyperpa- rameters is refered from their official implementations or papers. The hyperparameters for each algorithm is in Table 7.
Table 7. Algorithm- specific hyperparameters used in the experiment.
Method
Searched Candidates
Best Selection
FedAvg
None
None
FedDyn
B E {1, 10, 100, 1000]
β 10 except 100 AG News
FedNTD
None
1=1.0, r = 1.0
FedSAM
P {0.5n, n. 2n, 5n, 10n)
5n except p 1) in CIFAR-100
P {0.5n, n, 2n, 5n, 10n)
P 5n except p 1) in CIFAR-100
FedSpeed
8 E {1, 10, 100, 1000}
8 10 except
100 AG News
E [0.5m, n, 2,2n,.50 5n, 10n)
5x) except p = 1) in CIFAR-100
FedSMOC
B E [1, 100, 1000}
_ 10 except 3 100 in AG News
P {0.5n, 2n, 5n, 10n)
P 5m except p 1) in CIFAR-100
FedLESAM-D
8 [1, 100, 1000}
8 10 except 100 AG News
E 1., 2.0} 8 E 10, 100, 10003
= 1.0, 3.0
FedGMT
β 10 except 3 100 in AG News
a E {0.5, 0.950 0..90,9990000 0.998)
a 0.95 for FedGMT, 0.5 for FedGMT-v2
In Table 7, B is the coefficient for the penalty term. The selection of this hyperparameter has been studied in many previous works which verify its efficiency. Usually the selection of B are in {10, 100). The selection of is usually related to the local learning rate (Sun et al., 2023b). and r is the strength and temperature in the KL function.
25


Reviews:
[{"review_id": "K4NHfeYpFY", "reviewer": "Reviewer_aEZv", "summary": "This paper studies sharpness-aware minimization (SAM) in federated learning (FL) in the presence of data heterogeneity. The major problem for SAM in FL is that the clients cannot get an accurate estimate of the global objective/gradient due to heterogeneous data distribution. Existing literature has proposed utilizing the model updates in the previous communication round to approximate the global gradient. This paper proposes FedGMT, which leverages the global model trajectory which includes the pseudo-gradients from all the previous rounds through exponential moving averaging.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.\n\n---\n**W1.FedGMT vs. FedGKD.** While FedGKD appears to be similar to our method, there are fundamental differences. Practically, FedGKD takes an element-wise average over the latest 5 rounds (as they suggest) of global models and sends it to clients. This averaging method isn't EMA; it demands extra storage space and doubles the communication overhead.\nWe use EMA with KL divergence in the training loss because, as proven in Eq. (8), minimizing the loss difference $L(e)-L(w)$ between the EMA model $e$ and the global model $w$ is equivalent to minimizing the SAM’s sharpness measure for the global optimization. However, directly combining this loss difference with the FL objective will cancel out $L(w)$. So  we replace the cross entropy (CE) loss with the KL divergence loss to decouple the vanilla loss since minimizing the CE loss is equivalent to minimizing the KL loss. Thus, minimizing the KL loss in FedGMT aims to minimize global sharpness and reduce computational cost compared to other SAM-based methods, a goal that FedGKD cannot achieve.\nThe comparative experiments in the following table in **W2** demonstrate that FedGMT outperforms FedGKD by 2.1% on Dir(0.1) and 2.79% on Dir(0.01) when both methods employ the ADMM. \n\n---\n**W2. Methods comparison without ADMM.** We note that the performance of $L^{glotra}$  degrades without ADMM, which is due to non-vanishing biases between local update and global update in heterogeneity scenario. This aligns with methods such as FedSpeed, FedSMOO, and FedLESAM, which also incorporate ADMM to address such issues. Thus, approaches that omit ADMM may fail to fully leverage the proposed sharpness measure's potential. We evaluate performance with/without ADMM on CIFAR10 to demonstrate ADMM's necessity for maintaining optimization stability under data heterogeneity. The results are stated below.\n|  |  W/O ADMM Dir(0.1) | W/ ADMM Dir(0.1) |W/O ADMM Dir(0.01) | W/ ADMM Dir(0.01) |\n| -----| -----  | ----- | -----  |----- | \n| FedAvg|70.61±3.51|75.71±0.95 |61.94±4.93|70.52±2.32 |\n| FedSAM|70.96±3.97|77.51±0.97 |61.05±4.85|71.96 ±1.67 |\n| FedSMOO|71.28±3.90|77.08±0.97 |61.60±4.48|72.11 ±1.79|\n| FedLESAM|70.89±3.59|76.11±0.83|**62.95±4.60**|71.10 ±2.82|\n| FedGKD |72.54±2.68|77.07±0.83 |61.50±3.96|71.88±2.77|\n| FedGMT|**72.68±2.19**|**79.17±0.49**|61.28±3.11|**74.67 ±0.77**|\n| FedGMTv2 | 72.62±2.27|78.73±0.47|61.90±3.46|74.11±1.32|\n\nThe above experimental results show that FedGMT with ADMM achieve best performance due to our sharpness measure can more accurate estimate global sharpness than others (Figure 2 also demonstrate this). We will add these experimental results in our updated version.\n\n---\n**W3.Learning rate tuning.** Thanks for the suggestions. Learning rate which is a basic hyperparameters in the deep learning, and usually we do not finetune this for the fair comparison in the experiments. After a comprehensive review of existing federated learning literature, we observe that prior works consistently use a fixed learning rate across all compared methods. In our study, we  select learning rates from the set {0.001,0.01,0.1} for different datasets and model architectures to maximize SGD performance, then fix this optimal value for all methods. Specifically, a learning rate of 0.01 was employed for CIFAR10, CIFAR100, and CINIC10, while 0.1 was used for AG News.\n\n---\n**Q1. Convergence rate understand.** In the federated stochastic non-convex setting, under standard assumptions of \"smoothness,\" \"bounded stochastic gradients,\" \"heterogeneity,\" and other assumptions, prior works[A,B] have established $O(1/T)$ convergence rates for their algorithms. Specifically, our theoretical contributions lie in proving convergence without relying on the restrictive assumptions of bounded heterogeneity or requiring local minima in each communication round. This is enabled by our ADMM-based framework, which derives properties (e.g., Lemmas D.5 and D.6) to bound the global gradient norm in Eq. (34).\nPrior work [C] establishes faster convergence under the assumption of client similarity in non-convex settings.  Our Eq. (10) scales $L^{glotra}$ as a non-negative convex function which can be seen as an exponential moving average of historical global gradients into local updates across all clients. This shared similarity in EMA across clients reduces variance between local and global updates, thereby accelerating convergence.\n\n[A] Fedpd: A federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055–6070, 2021.\n\n[B] Fedadmm: A federated primal-dual algorithm allowing partial participation. arXiv preprint arXiv:2203.15104, 2022.\n\n[C] Scaffold:Stochastic controlled averaging for federated learning(ICML 2020).\n\n---\nIt is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes. Extensive experiments on both CV and NLP problems, with various datasets and neural network models are conducted.", "theoretical_claims": "Yes. I quickly went through the proofs (which seem correct), but have not had the time to check them step-by-step.", "experimental_designs_or_analyses": "Yes. In the experiments, it is mentioned that SGD with a learning rate of 0.01 is adopted. While the momentum of 0.9 is quite common and standard in the literature, I expect that the learning rate should be tuned from a reasonable set for the baselines for the sake of a fair comparison.", "relation_to_broader_scientific_literature": "There are many existing works endeavoring to address the data heterogeneity and computational overhead issues in SAM for federated learning, such as FedSMOO, FedSpeed, and FedLESAM. This paper advances the literature by incorporating the global model trajectory.", "essential_references_not_discussed": "The literature review looks comprehensive.", "other_strengths_and_weaknesses": "**Strengths:**\n1. The topic and problem that this paper aims to tackle are interesting and important.\n2. The paper is well-written and not difficult to follow.\n3. Extensive experiments on both CV and NLP tasks are conducted to validate the effectiveness of the proposed algorithms. \n\n**Weakness:**\n1. Incorporating KL divergence in training loss is not new. For example, the derived algorithm looks similar to the cited paper FedGKD [A], which also adopts the EMA at the server and adds a KL term to the local loss. Therefore, in my understanding, the major difference lies in the adaptation of ADMM, which was proposed in FedSMOO. \n2. The improvement of the algorithm seems to arise mainly from ADMM since the performance of only using $L^{global}$ in the ablation study in Table 3 looks similar to Fed-SAM in Table 2. (78.18 vs. 78.31, 61.28 vs. 61.05 in Dir(1.0) and (0.01) conditions, respectively). More experimental results without ADMM would help demonstrate the effectiveness of the derived global sharpness measure via Global Model Trajectory, which is the major contribution of this paper based on my understanding.\n3. The learning rates seem not carefully tuned for the baselines. \n\n[A] Yao, D., Pan, W., Dai, Y., Wan, Y., Ding, X., Yu, C., ... & Sun, L. (2023). FedGKD: Toward heterogeneous federated learning via global knowledge distillation. IEEE Transactions on Computers, 73(1), 3-17.", "other_comments_or_suggestions": "The paper is well-written.", "questions_for_authors": "Please see my comments about weaknesses. In addition, could the authors help me understand how incorporating EMA and ADMM improves the learning rate for non-convex optimization from $O(1/\\sqrt{T})$ to $O(1/T)$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QwulMTQo05", "reviewer": "Reviewer_5WUM", "summary": "The article proposes a novel solution to deal with the client drift problem of federated learning. It is based in sharpness aware minimization - addressing the two problems: how to make it efficient? How to guarantee that the global rather than client objectives are targeted?\nIt proposes a novel algorithm which is based on a clever integration of trajectory loss. The method is introduced, and substantiated both empirically and with convergence analysis.\n#after discussion: I am still convinced that the paper contributes valuable tools. I appreciate the discussion about drift, very interesting.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.\n\n---\n**1.Comment on realistic client drifts.** \nClient drift in federated learning refers to the phenomenon where data distributions across devices (clients) change over time or space, leading to performance degradation.\nAfter a comprehensive review of existing federated learning literature to address client drift problem, we summarize that numerous studies simulate non-IID federated settings using partitioning strategies on public datasets. This approach is preferred because real-world federated datasets are scarce due to data regulations and privacy constraints, while synthetic partitioning offers flexibility in controlling key FL parameters (e.g., client count, data size) and leveraging existing centralized training knowledge.\nThus, synthetic client drifts is more complex and flexible, enabling methods developed under these conditions to generalize more readily to realistic client drift scenarios.\n\nIn our work, as mentioned in Section 4.1, we adopt two widely used data partition strategy:\n- Pathological: Only selected categories can be sampled with a non-zero probability. The local dataset obeys a uniform distribution of active categories.\n- Dirichlet: Each category can be sampled with a non-zero probability. The local dataset obeys a Dirichlet distribution. Notably, we extend this by transforming the original balanced dataset into a long-tail distribution to further enhance the data heterogeneity and simulate real world scenarios.\n\n---\n**2.Propose and evaluate more than Dirichlet.** We introduce a non-IID data partitioning method where client datasets are strictly non-overlapping (disjoint) across classes, simulating an extreme data heterogeneity scenario. Specifically:\n- CIFAR10 (10 total classes) : 10 clients, each assigned a unique class \n- CIFAR100 (100 total classes) : 100 clients, each assigned a unique class \n\nThis single-class assignment ensures maximal label distribution discrepancy between clients. Experiments were conducted with\nparticipation rate 40% and communication rounds 500.\nThe results under this extreme setting are presented below.\n|  | | CIFAR10 | CIFAR100 |\n| ----- | ----- | ----- | -----  |\n| FedAvg | (AISTATS 2017) | 31.57±1.04 | 6.42±0.24 |\n| FedSAM |(ICML 2022) | 29.39±0.82 | 5.56±0.18 |\n| FedSMOO | (ICML 2023) | 10.00(failed) | 1.00(failed) |\n| FedLESAM-D |(ICML 2024) |  45.35±0.60 | 13.07±0.70 |\n| FedGMT  | (ours) | **58.23±0.68** | **22.13±0.34** |\n| FedGMTv2  |(ours) |  54.76±0.40 | 17.19±0.14 |\n\nThe results shows that in an extreme non-IID federated learning setup where clients hold disjoint class partitions, FedGMT achieves state-of-the-art results: 58.23% on CIFAR10 and 22.13% on CIFAR100, outperforming FedLESAM-D (45.35%/13.07%) and other baselines. FedAvg (31.57%/6.42%) and FedSAM (29.39%/5.56%) struggle with label heterogeneity, while FedSMOO fails entirely. These results validate FedGMT’s superior robustness to extreme label skewness.\n\n---\n**3.Comment on the challenge of privacy.** Our work builds upon the well-established FedAvg framework, inheriting its parameter-based communication architecture between server and clients. This design ensures full compatibility with existing privacy-preserving techniques developed for FedAvg, including mainstream defenses against privacy attacks such as differential privacy [A], homomorphic encryption [B], and secure multi-party computation [C].\n\nThe core contributions of our paper focus on enhancing practical generalization accuracy for real-world datasets while minimizing communication and computation costs under data heterogeneity. Importantly, our methods impose no restrictions on the integration of privacy-preserving techniques, ensuring seamless compatibility with existing or emerging privacy frameworks.\n\n[A] Chen,H.,Vikalo,H.,etal. The best of both worlds: Accurate global and personalized models through federated learning with data-free hyper-knowledge distillation(ICLR , 2023). \n\n[B] Cai Y, Ding W, Xiao Y, et al. Secfed: A secure and efficient federated learning based on multi-key homomorphic encryption[J]. IEEE Transactions on Dependable and Secure Computing, 2023, 21(4): 3817-3833.\n\n[C] Chen L, Xiao D, Yu Z, et al. Secure and efficient federated learning via novel multi-party computation and compressed sensing[J]. Information Sciences, 2024, 667: 120481.\n\n---\nIt is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.", "claims_and_evidence": "The main claim is that the proposed method allows federated learning to better and efficiently deal with client drift. This is sufficiently demonstrated.", "methods_and_evaluation_criteria": "The method is reasonable and aligns very well with what is currently state of the art in the domain. Data sets are standard ones - distorted client-wise using a dirichlet distribution. This is sufficient to document the  results. It would be good to have more distortions/drift models for the clients, including disjoint data domains. Also, it might be good to have sensor data or any data which is not text or image.\nBut it is sufficient as is to make the point.", "theoretical_claims": "Yes, the proofs and formalizations follow standard methodology and are correct as far as I can tell.", "experimental_designs_or_analyses": "Versy good design and solid evaluation. Code is provided, so reproducible.", "relation_to_broader_scientific_literature": "Yes, the authors are well aware of the field.", "essential_references_not_discussed": "In my opinion this is  good", "other_strengths_and_weaknesses": "very relevant topic and apparently a good solution.\nCan you comment on the challenge of integrating privacy issues? In how far is this compatible?", "other_comments_or_suggestions": "none", "questions_for_authors": "Please make a comment on realistic client drifts and propose and evaluate more than Dirichlet.\nPlease comment on the challenge of privacy.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4lNBDqUMVH", "reviewer": "Reviewer_NL1R", "summary": "This paper proposes a new Federated Learning algorithm, named FedGMT, to effectively cope with data heterogeneity by reducing the sharpness of the global model through a global model trajectory. This paper provides the convergence analyze of FedGMT in the non-convex and smooth cases. Experimental results show significant performances of the proposed method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comments and suggestions! We answer your questions below. \n\n---\n\n**1.Optimization target explanation.**\n By referring the Theorem 1 in original SAM(Foret et al.,2020), the objective loss function of FedSAM in each client $m$ can be rewritten as the sum of the vanilla loss and the loss associated to the sharpness measure, which is the maximized change of the training loss within the ρ-constrained neighborhood, i.e.,\n$\\arg\\min_{w_m}[L_m(w_m)+S_m(w_m)]$ where $S_m(w_m) = max_{\\epsilon: ||\\epsilon||<\\rho}L_m(w_m+\\epsilon)-L_m(w_m)$, which is a minimax problem as you mentioned. The sharpness measure  is approximated as $S_m(w_m) = L_m(w_m+\\epsilon_m)-L_m(w_m)$, where $\\epsilon_m = \\rho\\frac{\\nabla\\mathcal{L}_m(w_m)}{\\|\\nabla\\mathcal{L}_m(w_m)\\|}$ is the solution to an approximated version of the maximization problem in FedSAM. The sharpness measure in FedSAM approximates the inner maximization problem and is combined with the vanilla loss for minimization. This transforms the minimax problem into a standard minimization task. Our work based on this task to develop a better sharpness measure to minimize global sharpness in FL, as formally established in Eq. (8) of the manuscript.\n\n---\n**2.Role of the dual variable $u$.** Note that in line 13 of Algorithm 1, each client has its own dual variable $u_m^{t+1} = u_m^{t}-\\frac1\\beta(w_m^t-w^t)$. This enables clients to ensure that the solutions of their respective sub-problems are consistent under global constraints.\n\nIn line 16 of Algorithm 1, we define the global dual variable *$u^t$* because, according to Eq. (13), the global model $w^{t+1}$ needs to be updated as $w^{t+1}=\\frac1N\\sum_{m \\in \\mathbb{N}}(w^t_{m,K}-\\beta u^{t+1}_m)$. To avoid having each client send $u_m^{t+1}$ to the server in every round, we define $u^t$ to assist in minimizing Eq. (13). We also conducted experiments on CIFAR10-Dir(0.1), FedGMT's accuracy drop from 79.17 to 65.59 without $u$. Therefore, updating $u$ is necessary, not redundant.\n\n---\n**3.Analysis in Table 1.**\n\nTo save space for other question, we have addressed this concern in our rebuttal to Reviewer Zauy (Section **W1**). Please refer to that section for our detailed response. \n\n---\n\n**5.1 Compare the federated ADMM-based work.**  The baselines we take in experiments, such as FedDyn, FedSpeed and FedSMOO, which are all state-of-the-art FL methods based on ADMM.\n\n**5.2 Comparison criterion.** The comparison criterion we used follows original FedSAM(Caldarola et al.,2022). we ensure sufficient communication rounds to guarantee convergence (see Fig.5), while using the final averaged accuracy to avoid fluctuations and provide a more reliable performance estimate. We also provide the historical best accuracy by random seed {1,2,3} in **anonymous link https://anonymous.4open.science/r/RG5**, the ranking of each method is consistent with our criterion.\n\n---\n**6.Loss landscape.** We provide the loss landscape, top Hessian eigenvalue and Hessian trace of SAM-based method in anonymous link in 5.2,  demonstrating that our FedGMT can reach a flatter minima than others.\n\n---\n\n**7. Hyperparameters sensitivity and selection.** Thanks for the suggestion. Our experiments across 4 CV/NLP datasets, 4 model architectures, and different scenarios (heterogeneity and participation) show that FedGMT's hyperparameters can be consistent and robust. We will detail the meanings and selections of these hyperparameters in this part.\n- The strength $\\gamma$ and temperature $\\tau$  which is a basic hyperparameter for KL function. We follow the common settings used in previous works and $\\gamma=1$ and $\\tau = 3$ for all settings.\n- The penalty coefficient $\\beta$ has been studied in many previous ADMM-based works. We test the selection of {10, 100} works on all settings.\n- $\\alpha$ control the loss on the recent update trajectory in EMA. Since $\\alpha$ needs to be close to 1 to retain past trajectories, we mainly chose from {0.5,0.9,0.95,0.99}. We test $\\alpha=0.95$ works well on all settings.\n\nIn summary, the hyperparameters in FedGMT can be selected by simple experience as we mentioned above. In practice, FedGMT ​​can achieve higher performance compared to other baselines without the complex tuning.\n\n---\n**8.1 The originality** is that we develop a efficient (communication and computation) global sharpness measure for SAM-FL. \n We do increase the amount of calculation in line 13,16,18 of Algorithm 1,  but it is negligible compared to the mainly training costs. Notably, FedGMT avoid 1 more backward pass required for calculating perturbations in standard SAM.\n\n**8.2 Real-world deployment.**  To validate practical deployment feasibility, We implemented FedGMT on 2×RPi4B, 2×Jetson Nano, and 1×PC. FedGMT achieves about 47% energy savings vs. FedAvg across three datasets, as shown in the anonymous link in 5.2.\n\n---\nIt is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.", "claims_and_evidence": "Considering the original FedSAM, it formulates the optimization target as a minimax problem, to minimize the maximum loss. However this work considers a minimization target, which is not identical to the SAM-based methods. How do the authors explain the optimization target?", "methods_and_evaluation_criteria": "In Algorithm 1 line 16, the server aggregates update the u but doesn’t send it to the client. How does each client gain that or is this aggregation redundant?", "theoretical_claims": "This work lacks the generalization error bound analysis of proposed method.", "experimental_designs_or_analyses": "In the experiment, the authors don’t compare the federated ADMM-based work. It’s suggested to add some baselines for comparison. And the authors “report the final averaged test accuracy and standard deviation over the last 50 rounds for increased robustness and reliability”, this comparison criterion is not solid, especially some methods have not converged in the given communication rounds.", "relation_to_broader_scientific_literature": "The key contributions of the paper are closely related to several existing lines of research in Federated learning and optimization, particularly in addressing the challenges posed by data heterogeneity and improving model generalization.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Weaknesses\n\n1. Considering the original FedSAM, it formulates the optimization target as a minimax problem, to minimize the maximum loss. However this work considers a minimization target, which is not identical to the SAM-based methods. How do the authors explain the optimization target?\n2. In Algorithm 1 line 16, the server aggregates update the u but doesn’t send it to the client. How does each client gain that or is this aggregation redundant?\n3. In Table 1, the authors claim the FedGMT achieves 1\\times or 1.5\\times communication cost and 1.2\\times computation cost. How do the authors gain this result? Please give further analysis.\n4. This work lacks the generalization error bound analysis of proposed method.\n5. In the experiment, the authors don’t compare the federated ADMM-based work. It’s suggested to add some baselines for comparison. And the authors “report the final averaged test accuracy and standard deviation over the last 50 rounds for increased robustness and reliability”, this comparison criterion is not solid, especially some methods have not converged in the given communication rounds.\n6. Figure 6 cannot prove the global loss is more flatter, but only show the consensus of the local and global update trajectory. If authors can provide a visualization of the loss landscape, it can be more direct.\n\n7. The paper discusses hyperparameters sensitivity but could provide more detailed insights into how different hyperparameters affect the performance of FedGMT. \n8. The paper could enhance its originality by more clearly distinguishing its contributions from existing methods, its significance by including real-world deployment examples.", "other_comments_or_suggestions": "1. While the paper includes a hyperparameters sensitivity analysis, it could provide more detailed guidance on how to select the right hyperparameters for different datasets or levels of data heterogeneity. A discussion on the robustness of FedGMT to suboptimal hyperparameter choices would also be valuable.\n2. The paper needs to calculate and store additional parameters both locally and on the server, which will increase the amount of calculation. How to balance model performance and resources.", "questions_for_authors": "See the Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5kghhvkjk1", "reviewer": "Reviewer_Zauy", "summary": "The paper proposes FedGMT, a federated learning framework that leverages sharpness-aware minimization (SAM) to enhance generalization, especially in highly skewed non-IID settings. To achieve this, the framework employs an exponential moving average (EMA) of the global model as a proxy for the global loss surface. An additional regularization term is introduced, defined by the KL divergence between the output distribution of a client’s local model and that of the EMA-based global model, which guides SAM updates to consider the global context. The paper also presents FedGMT-v2, a variant that approximates FedGMT to reduce communication overhead, albeit with a slight performance degradation. The framework's properties are thoroughly explored through both theoretical analysis and empirical experiments.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review and constructive comments. We provide our responses as follows.\n\n---\n**W1. Clarification of communication and computational cost in Table 1**：We apologize for not making it clear and we will modify the caption of Table 1 and add a discussion section with following content in the revised version.\n\n**Communication cost:** The communication cost is defined as the parameters transmitted per round. For example, if we denote the amount of model parameter as Ω, FedSAM requires 2Ω due to bidirectional model parameter exchange. In contrast, FedGMT incurs a communication cost of  3Ω due to the server transmitting the EMA model to clients, whereas FedGMTv2 achieves 2Ω  by omitting this transmission. Thus, FedGMT's communication cost is either $1.5×$ or $1×$ that of FedSAM. \n\n**Computational cost:**  The computational cost is defined as per-iteration training cost. Specifically, minor computational overhead (e.g., model aggregation, ADMM computations) is ignored compared to the primary training computation. Taking FedAvg as the baseline ($1×$ cost), FedSAM doubles the forward and backward processes, resulting in a $2×$ cost. FedGMT involves an extra forward pass. we refer to the official PyTorch paper [A], which shows that for ResNet50, the backward pass is $3.7×$ the forward pass. Based on this data, FedGMT's cost is estimated to be about $1.2×$ by $(\\frac{1+1+3.7}{1+3.7})$.\n\nIf we assume a forward cost of *f* and a backward cost of 2*f*, FedAvg requires 3*f* per iteration, while our FedGMT needs 4*f*, leading to a $1.33×$ cost. We will describe this hypothesis and adjust the estimate from $1.2×$ to $1.33×$ for better understanding to readers.\n\nWhile FedGMT incurs marginal computational overhead, it achieves superior accuracy and efficiency. As shown in following Table in **W2**, FedGMT requires only 47% of the time cost of FedAvg to reach the target accuracy, demonstrating its practical advantage in resource-constrained environments. \nFor example, if one trains a FL model to reach the target accuracy with 8 Nvidia V-100 GPUs, considering ​​Google Cloud Platform charges \\\\$2.48 per GPU per hour, the total cost of FedAvg for 1 day will be \\\\$476, FedSAM will be \\\\$724 (476\\*152%), while FedGMT will be \\\\$224 (476\\*47%). There is a significant cost savings of implementing FedGMT.\n\n[A] Li S, Zhao Y, Varma R, et al. Pytorch distributed: Experiences on accelerating data parallel training[J]. Proc. VLDB Endow. 13(12): 3005-3018 (2020)\n\n---\n**W2.Experiments with larger model:** We conducted experiments on CIFAR100-Dir(0.1) with Resnet18 on one NVIDIA 4090 GPU. We follow the same parameter settings in FedSMOO paper and give a comprehensive comparison. The results are stated below.\n\n|  | Acc | Images/s |Per round time|Round(acc=44%)  |Communication cost(acc=44%)| Time cost(acc=44%)|\n| ----- | ----- | -----  |-----  |-----  |-----  |-----  |\n| FedAvg| 43.95±0.22 | 816(100%) |33.85s(1x) |753(100%)|1506 Ω(100%) |425 min(100%) |\n| FedSAM| 44.56±0.20 | 429(53%) |61.31s(1.81x) |633(84%) |1266 Ω(84%) |647 min(152%) |\n| FedSMOO| 47.94±0.15 | 385(47%) |69.85s(2.06x) |412(55%) |1648 Ω(109%) |480 min(113%) |\n| FedLESAM-D| 46.42±0.23 | 657(81%) |42.20s(1.25x) |309(41%) |618 Ω(41%) |217 min(51%) |\n| FedGMT| **50.67±0.16** | 583(71%) |46.13s(1.36x) |257(34%) |771 Ω(51%) |**198 min(47%)** |\n| FedGMTv2|50.24±0.22 | 566(69%) |46.32s(1.37x)| **262(35%)**|**524 Ω(35%)** |202 min(48%) |\n\nFedGMT achieves the highest accuracy (50.67%) and fastest convergence (257 rounds/198min) to reach 44% accuracy, outperforming all baselines. While FedSMOO (47.94%) and FedLESAM-D (46.42%) show better accuracy than FedAvg (43.95%) and FedSAM (44.56%), they lag behind FedGMT in efficiency.  FedGMT balances accuracy and efficiency, surpassing competitors in both metrics.\n\n---\n**W3.Experiments with larger number of clients:** We conducted experiments on CIFAR10-Dir(0.1) with LeNet . We set 500 clients with 20% active ratio and 1000 clients with 10% active ratio to involve 100 clients per round. The results are stated below.\n|  | 20%-500 | 10%-1000 |\n| ----- | ----- | -----  |\n| FedAvg | 61.13±0.47 | 54.39±0.35 |\n| FedSAM | 60.82±0.41 | 53.96±0.31 |\n| FedSMOO | 79.95±0.17 | 69.57±0.29 |\n| FedLESAM-D | 79.71±0.17 | 74.06±0.13 |\n| FedGMT | **80.78±0.13** | **74.34±0.11** |\n\nThe result shows that FedGMT's superior scalability and robust for large-scale FL scenario. In addition, you can refer to our experiments in Reviewer 5WUM that also verify the robustness of FedGMT in an extreme non-IID FL setup.\n\n---\n**W4. Theoretical improvements.** Thank you for your insightful feedback and suggestions to enhance the theoretical rigor of our paper. We wholeheartedly agree with your comments and will reorganize the theoretical proof by following your suggestions for each step in our updated version.\n\n---\nIt is a pleasure to discuss this with you, which will help us to further improve this work. Thank you again for reading this rebuttal.", "claims_and_evidence": "The paper makes the following claims:\n\n1. FedGMT reduces the computational cost of SAM-aware federated learning (e.g., FedSAM) by nearly half.\n2. FedGMT achieves a faster convergence rate of $\\mathcal{O}(1/T)$.\n3. The empirical results demonstrate the effectiveness of FedGMT.\n\nThe paper provides evidence for each of these claims, although the evidence might be limited in certain aspects:\n\n1. The paper argues that by avoiding the additional backward pass required for calculating perturbations in standard SAM, FedGMT reduces the computation cost. Table 1 provides numerical comparisons, though some details could be clarified further.\n2. Theorem 3.5 presents theoretical evidence that, under standard assumptions (e.g., L-smoothness and bounded gradients), FedGMT achieves a convergence rate of $\\mathcal{O}(1/T)$. While these assumptions may not hold perfectly in all deep learning scenarios, they are commonly used in the literature.\n3. The experimental section demonstrates the empirical performance of FedGMT across multiple datasets and model architectures. Although the experiments are conducted using a relatively small number of clients and smaller model architectures compared to some real-world scenarios, they still provide credible evidence supporting the claims.", "methods_and_evaluation_criteria": "The method is well-explained and easy to follow. However, some clarifications are needed regarding the communication and computational cost figures presented in Table 1.\n\n- **Communication Cost:** The paper reports the communication cost of the proposed framework to be $1.5\\times$ of FedSAM. To my understanding, this value represents the average-case scenario. In the worst case, if the EMA changes significantly, the communication cost could approach $2\\times$, whereas in the best case, if the EMA remains nearly unchanged, the cost would be close to $1\\times$. Explicitly stating these scenarios would help readers better understand the variability in communication overhead.\n\n- **Computational Cost:** The computational cost column requires explanation (For example, why is it $1.2\\times$?) and a more formal analysis. An example of such an analysis could be that if we denote the cost of a forward pass through the network as $f$ floating point operations per second (FLOPS) and assume that a backward pass costs approximately $2f$ (as commonly observed in the literature), then methods like FedSAM, which require two forward passes and two backward passes, have a computational cost of roughly $6f$. In contrast, FedGMT only requires two forward passes, one backward pass, the computation of the KL divergence (which is negligible relative to $f$), and some lightweight ADMM computations, which might be approximated as $2p$ FLOPS (please verify this number), where $p$ is the number of parameters in the model. This results in an overall overhead of approximately $4f+2p$. The overall complexity is lower because $f$ is typically much larger than $p$ (for example, in ResNet-18, many research works estimated that  $p$ = 11,689,128, while $f$ =1,818,228,160 FLOPS). A more formal analysis with clear notations would strengthen the argument regarding the computational cost.", "theoretical_claims": "Overall, the theoretical claims in the paper are well supported with proof. However, the derivation in the Appendix concerning the ADMM updates and the related discussion remains largely intuitive. A more rigorous, step-by-step treatment would be beneficial. In particular, the authors could:\n- Include additional bounding steps or formalize the transition from local expansions to global updates.\n- Reference standard ADMM-based lemmas to reinforce the derivation.\n- Provide explicit error bounds quantifying the discrepancy between local and global updates.\n\nSuch improvements would help readers verify each step of the argument and appreciate the theoretical underpinnings of the proposed method more fully.", "experimental_designs_or_analyses": "The paper studies various experimental setups to validate the framework, including different data skew conditions, sensitivity analyses, and varying client participation rates. However, there are two minor limitations in the experimental design:\n1. *Model Architectures:* The experiments are conducted using relatively small models: a CNN with 150K–200K parameters, ResNet-8 with around 110K parameters, and a ViT-CIFAR variant with fewer than 6 million parameters. In contrast, related works like FedSAM and FedGAMMA often use larger models, such as ResNet-18. Including experiments with larger architectures would improve the generalizability of the results.\n\n2. *Number of Clients:* The experiments are performed with only 10 clients, which is small compared to typical federated learning studies that involve around 100 clients. \n\nWhile the current experiments do support the claims made, expanding the study to include a larger number of clients would strengthen the evidence for scalability and robustness.", "relation_to_broader_scientific_literature": "This work is positioned within the expanding field of federated learning, addressing key challenges such as data heterogeneity, client drift, and communication/computation efficiency. It builds directly on recent advances in sharpness-aware minimization (SAM) applied to federated settings, which aim to achieve better generalization by optimizing for flat minima. The use of an exponential moving average (EMA) to capture the global model trajectory and guide local updates is a novel twist that resonates with prior work in model averaging and momentum methods in distributed optimization.", "essential_references_not_discussed": "By and large, the paper provides a comprehensive view of the area, and to the best of my knowledge, no essential references are missing.", "other_strengths_and_weaknesses": "**Strengths:**\n- The paper is well-written and easy to follow.\n- The visualizations used, specifically figures 1 and 6, are helpful to the reader. \n\n**Weaknesses:**\n- See other sections.", "other_comments_or_suggestions": "The paper is well written, barring the few clarifications needed and limitations in the experimental design. If the authors could provide sufficient clarification, I am open to increasing my score.", "questions_for_authors": "See other sections in the review.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This works tackles sharpness‑aware minimization (SAM) in the federated learning setting by directly measuring global sharpness via a global model trajectory, rather than per‑client approximations. Its key ingredients are: (1) The server maintains an exponential moving average (EMA) of past global models. (2) Clients add a KL‑divergence term between their local model’s outputs and the EMA model’s outputs, which provably aligns with minimizing the global SAM sharpness measure. (3) A dual variable per client is used to enforce consistency between local and global updates. Further, this approach avoids the second backward pass of standard SAM, cutting compute by roughly half. This work also conducted some extensive experiments to justify their approach.  The reviewers all recognize the positive contribution of this work and the capability of the proposed approach to trade-off accuracy and efficiency tradeoff. Therefore I recommend acceptance of the paper.
Decision: Accept (poster)