Paper deep dive
FedNSAM:Consistency of Local and Global Flatness for Federated Learning
Junkang Liu, Fanhua Shang, Yuxuan Tian, Hongying Liu, Yuanyuan Liu
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 7/20/2026, 7:28:18 AM
Summary
The paper introduces FedNSAM, a federated learning algorithm that addresses the issue of sharp global minima caused by data heterogeneity. It defines 'flatness distance' to quantify the discrepancy between local and global flatness. FedNSAM incorporates global Nesterov momentum into local updates to harmonize local and global flatness, theoretically proving a tighter convergence bound than FedSAM and empirically demonstrating superior performance on CNN and Transformer models.
Entities (11)
Relation Signals (6)
FedNSAM → uses → Nesterov Momentum
confidence 95% · FedNSAM uses the global Nesterov momentum as the direction of local estimation of client global perturbations
data heterogeneity → causes → Sharp Global Minima
confidence 93% · multi-step local updates and data heterogeneity usually lead to sharper global minima
FedNSAM → outperforms → FedSAM
confidence 92% · Theoretically, we prove a tighter convergence bound than FedSAM
Flatness Distance → quantifies → Discrepancy between Local and Global Flatness
confidence 91% · We define the flatness distance to explain this phenomenon... difference gap of the flatness between the global and local models
FedNSAM → improves → Generalization Ability
confidence 90% · FedNSAM... enable the effectiveness of SAM in FL to improve the generalization ability of the global model
FedSAM → suffersfrom → data heterogeneity
confidence 88% · in the high data heterogeneity setting, the flatness in local training does not imply the flatness of the global model... FedSAM is effective in finding flat minima in local training, but the aggregated global model is not a flat minima
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:In federated learning (FL), multi-step local updates and data heterogeneity usually lead to sharper global minima, which degrades the performance of the global model. Popular FL algorithms integrate sharpness-aware minimization (SAM) into local training to address this issue. However, in the high data heterogeneity setting, the flatness in local training does not imply the flatness of the global model. Therefore, minimizing the sharpness of the local loss surfaces on the client data does not enable the effectiveness of SAM in FL to improve the generalization ability of the global model. We define the \textbf{flatness distance} to explain this phenomenon. By rethinking the SAM in FL and theoretically analyzing the \textbf{flatness distance}, we propose a novel \textbf{FedNSAM} algorithm that accelerates the SAM algorithm by introducing global Nesterov momentum into the local update to harmonize the consistency of global and local flatness. \textbf{FedNSAM} uses the global Nesterov momentum as the direction of local estimation of client global perturbations and extrapolation. Theoretically, we prove a tighter convergence bound than FedSAM by Nesterov extrapolation. Empirically, we conduct comprehensive experiments on CNN and Transformer models to verify the superior performance and efficiency of \textbf{FedNSAM}. The code is available at this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2602.23827v1
- Canonical: https://arxiv.org/abs/2602.23827v1
Trouble viewing inline? Open PDF directly →
Full Text
107,209 characters extracted from source content.
Expand or collapse full text
Consistency of Local and Global Flatness for Federated Learning Junkang Liu College of Intelligence and Computing, Tianjin UniversityTianjinChina junkangliukk@gmail.com , Fanhua Shang College of Intelligence and Computing, Tianjin UniversityTianjinChina fhshang@tju.edu.cn , Yuxuan Tian College of Management and Economics, Tianjin UniversityTianjinChina guilang@tju.edu.cn , Hongying Liu Medical College, Tianjin University Peng Cheng LaboratoryTianjinChina hyliu2009@tju.edu.cn and Yuanyuan Liu School of Artificial Intelligence, Xidian UniversityXianChina yyliu@xidian.edu.cn (2025) Abstract. In federated learning (FL), multi-step local updates and data heterogeneity usually lead to sharper global minima, which degrades the performance of the global model. Popular FL algorithms integrate sharpness-aware minimization (SAM) into local training to address this issue. However, in the high data heterogeneity setting, the flatness in local training does not imply the flatness of the global model. Therefore, minimizing the sharpness of the local loss surfaces on the client data does not enable the effectiveness of SAM in FL to improve the generalization ability of the global model. We define the flatness distance to explain this phenomenon. By rethinking the SAM in FL and theoretically analyzing the flatness distance, we propose a novel FedNSAM algorithm that accelerates the SAM algorithm by introducing global Nesterov momentum into the local update to harmonize the consistency of global and local flatness. FedNSAM uses the global Nesterov momentum as the direction of local estimation of client global perturbations and extrapolation. Theoretically, we prove a tighter convergence bound than FedSAM by Nesterov extrapolation. Empirically, we conduct comprehensive experiments on CNN and Transformer models to verify the superior performance and efficiency of FedNSAM. The code is available at https://github.com/junkangLiu0/FedNSAM. Federated Learning, Sharpness-Aware Minimization, Generalization Ability †journalyear: 2025†copyright: rightsretained†conference: Proceedings of the 33rd ACM International Conference on Multimedia; October 27–31, 2025; Dublin, Ireland.†booktitle: Proceedings of the 33rd ACM International Conference on Multimedia (M ’25), October 27–31, 2025, Dublin, Ireland†isbn: 979-8-4007-2035-2/2025/10†doi: 10.1145/3746027.3755226†submissionid: 2691†ccs: Theory of computation Massively parallel algorithms 1. Introduction (a) FedSAM (0.6) (b) FedSAM (0.1) (c) FedNSAM (0.1) (d) IID (e) Non-IID (f) Correction Figure 1. (a) and (b) show the global training loss surface of FedSAM (Qu et al., 2022) under Dirichlet distributions with coefficients of 0.6 (low data heterogeneity) and 0.1 (high data heterogeneity) on CIFAR-100 with ResNet-18. (c) shows the training loss surface of our FedNSAM with Dirichlet-0.1. FedSAM can search global flat minima in low data heterogeneity but fails in high data heterogeneity. (d) suggests that the flat region of the client is closer in low data heterogeneity setting. The global model falls within the flat region. (e) suggests that the flat region of the client is far in high data heterogeneity setting. The global model cannot fall within the flat region. (f) suggests that our FedNSAM draws the flat region between each client closer by alignment correction so that the global model falls within the flat region of each client. FedNSAM’s global model finds flat minima in (c). Federated Learning (FL) has garnered significant interest as a crucial framework for decentralized training among numerous clients while ensuring data privacy. The core principle of FL involves maintaining local data on the client, allowing only the communication of gradients or model parameters between clients and a central server, and the server does not have direct access to local data (McMahan et al., 2017). FL has considerable promise in various regions, including healthcare, finance, and personalized mobile services (Rieke et al., 2020; Antunes et al., 2022; Byrd and Polychroniadou, 2020; Liu et al., 2024; Zeng et al., 2025; Dai et al., 2025a, b; Liu et al., 2025a; Yang et al., 2025). Recently, there has been widespread interest in how to improve generalization for FL (Caldarola et al., 2022; Sun et al., 2023; Qu et al., 2022). However, (Fan et al., 2024) shows that data heterogeneity and multi-step local updates can cause the global model to converge to sharp local minima, which tend to exhibit poor generalization. A popular approach to this problem is to use sharpness-aware minimization (SAM) (Foret et al., 2020) as a local optimizer to find flat local minima, such as FedSAM (Qu et al., 2022). However, in the setting of highly heterogeneous data, we find that FedSAM is effective in finding flat minima in local training, but the aggregated global model is not a flat minima (see Figure 1 (a,b)). The literature (Fan et al., 2024) also found the same problem as us, but it did not provide a good explanation for the reason behind this phenomenon. As illustrated in Figure 1, in the low heterogeneity case (d), clients share similar update directions, resulting in overlapping flat regions. This increases the likelihood that the global model converges to a shared flat minimum, as shown in (a). In contrast, in the high heterogeneity case (e), client updates diverge significantly, making their flat regions disjoint and preventing the global model from residing in any of them—thus leading to sharp global minima. To explain this phenomenon, we introduce the concept of flatness distance, which quantifies the discrepancy between local and global flatness. As shown in Figure 2, higher heterogeneity leads to greater flatness distance and degraded global performance. To mitigate this, we propose FedNSAM, which aligns client flat regions through Nesterov momentum correction, enabling the global model to fall within them, as visualized in Figure 1(f). Contributions. We identify key limitations of the SAM algorithm in federated learning (FL) and propose a new perspective—flatness distance—to quantify the divergence of flat regions across clients under data heterogeneity. To address these challenges, we design a new FL algorithm: FedNSAM, which leverages Nesterov momentum to minimize global sharpness by aligning local flat regions. Our main contributions are summarized as follows: ∙ We introduce the concept of flatness distance to characterize the inconsistency of local minima across clients. Through both theoretical and empirical analysis, we demonstrate that higher data heterogeneity leads to increased flatness distance and degraded global flatness. ∙ We propose a novel sharpness-aware federated learning algorithm (FedNSAM) with Nesterov extrapolation at the client level. We theoretically prove its convergence rate of O(LF/TKS(1−λ))O ( LF/ TKS(1-λ) ) and show that it achieves lower flatness distance than FedSAM. ∙ Extensive experimental results on three benchmark datasets validate that FedNSAM outperforms existing methods under various data heterogeneity levels, participation rates, and model architectures, demonstrating superior generalization and optimization performance. 2. Related work ∙ Heterogeneity Issues in FL: In past years, various strategies have been proposed to solve the heterogeneity issues in FL. A main branch of prior approaches focuses on regularizing local training to alleviate the divergence of the local models. As early works, FedAvgM (Hsu et al., 2019) is a proposed method that introduces momentum terms during global model updating. FedACG (Kim et al., 2024) improves the inter-client interoperability by the server broadcasting a global model with a prospective gradient consistency. SCAFFOLD (Karimireddy et al., 2020) uses SAGA-like control variables to mitigate client-side drift, which can be regarded as adopting the idea of variance reduction at the client side. FedProx (Li et al., 2020) reduces the variance of updates from different clients by adding a proximal term to the objective function to limit the magnitude of updates from each client. FedBCGD (Liu et al., 2024) proposes an accelerated block coordinate gradient descent framework for FL. FedSWA (Liu et al., 2025a) improves generalization under highly heterogeneous data by stochastic weight averaging. FedAdamW (Liu et al., 2025d) introduces a communication-efficient AdamW-style optimizer tailored for federated large models. FedNSAM (Liu et al., 2025b) studies the consistency relationship between local and global flatness in FL. FedMuon (Liu et al., 2025c) accelerates federated optimization via matrix orthogonalization. DP-FedPGN (Liu et al., 2025e) develops a penalizes gradient norms to encourage globally flatter minima in DP-FL. FedPAC (Liu et al., 2026) mitigates preconditioner drift to unlock the potential of second-order optimizers. Our FedNSAM method is based on flat minima searching, which differs from the above methods. ∙ Improve Generalization in FL: The existing FL methods that search flatter minima for better generalization utilize Sharpness Aware Minimization (SAM) (Foret et al., 2020) or its variants to local training. (Qu et al., 2022; Caldarola et al., 2022) apply SAM optimizer in FL local training and proposed FedSAM. (Qu et al., 2022) proposed a variant of FedSAM called MoFedSAM with global momentum in local training. (Dai et al., 2023) proposed FedGAMMA, which reduces the effect of data heterogeneity by using SCAFFOLD (Karimireddy et al., 2020) control variables in FedSAM. FedSMOO (Sun et al., 2023) attempted to address this by correcting local perturbations, but it introduces many computational overheads as other SAM-based algorithms. (Fan et al., 2024) propose FedLESAM locally estimates the direction of global perturbation on the client side as the difference between global models received in the previous active and current rounds. 3. Rethinking SAM in FL 3.1. FL Problem Setup FL aims to optimize model parameters with local clients, i.e., minimizing the following population risk: (1) F()=1N∑i=1N(Fi():=ζi∼i[Fi(;ζi)]). F( θ)= 1N _i=1^N (F_i( θ):=E_ _i _i [F_i ( θ; _i ) ] ). The function FiF_i denotes the loss on client i, where ζi∼i _i _i is a sample from its local data distribution. The expectation ζi[⋅]E_ _i[·] is taken over ζi _i. The total number of clients is N. 3.2. Flatness Searching in FL: FedSAM The optimization of FedSAM (Qu et al., 2022) is formulated as follows: (2) minθmax‖i‖22≤ρF(~):=1N∑i∈[N]Fi(~), _θ _ \| δ_i \|_2^2≤ρ \F( θ):= 1N _i∈[N]F_i( θ) \, where F(~)≜max‖≤ρF(+),Fi(~)≜F( θ) _\| δ\|≤ρF( θ+ δ), F_i( θ) max‖i‖≤ρFi(+i) _ \| δ_i \|≤ρF_i ( θ+ δ_i ). The local updates of FedSAM include (3) i,k+1/2t=i,kt+i,kt,i,kt=ρgi,kt/‖gi,kt‖, θ_i,k+1/2^t= θ_i,k^t+ δ_i,k^t, δ_i,k^t=ρ g_i,k^t/ \|g_i,k^t \|, i,k+1t=i,kt−ηgi,k+1/2t, θ_i,k+1^t= θ_i,k^t-η g_i,k+1/2^t, where gi,kt=∇Fi(i,kt;ζi)g_i,k^t=∇ F_i( θ_i,k^t; _i) is the stochastic gradient computed at i,kt θ_i,k^t, and gi,k+1/2t=∇Fi(i,k+1/2t;ζi)g_i,k+1/2^t=∇ F_i( θ_i,k+1/2^t; _i), which can be seen in Figure 3 (a). FedSAM searches for the flat local minima with local loss sharpness-aware minimization instead of the flat global minima. (a) flatness distance (b) Global Sharpness Figure 2. Illustration of flatness distance (left) and global sharpness (right) during federated training on CIFAR100, using Dirichlet distributions with coefficients 0.1, 0.3, and 0.6 across 100 clients and a 10% participation rate. Test accuracies are 40.18%, 46.02%, and 47.83% for FedSAM with Dirichlet-0.1, 0.3, and 0.6 respectively, and 58.53% for FedNSAM with Dirichlet-0.1. 3.3. Flatness Distance We formally define the flatness distance _D, as the difference gap of the flatness between the global and local models. Definition 1 (Flatness distance). _D of the global model t+1=1N∑i=1Ni,Kt θ^t+1= 1N _i=1^N θ_i,K^t and the local models i,Kti=1N \ θ_i,K^t \_i=1^N is defined as: (4) =1N∑i=1N‖i,Kt−t+1‖2. _D= 1N _i=1^NE \| θ_i,K^t- θ^t+1 \|^2. When _D is small, the flat regions across clients are closer, leading to higher global flatness. Conversely, a larger _D indicates greater disparity between client flatness regions, resulting in degraded global flatness. As shown in Figure 2, FedSAM suffers significant performance degradation as data heterogeneity increases, which corresponds to larger flatness distance and sharper global minima. In contrast, the proposed FedNSAM algorithm consistently achieves lower flatness distance and global sharpness under high heterogeneity, thanks to Nesterov momentum correction. Theoretical justifications are provided in Section 5.2. 3.4. Global Sharpness-Aware Minimization in FL To get global flat minima, we first recall the definition of global sharpness-aware minimization (Fan et al., 2024) in FL: (5) minθmax‖2≤ρF(+)=1N∑i=1NFi(+), _θ _\| δ\|_2≤ρ \F( θ+ δ)= 1N _i=1^NF_i( θ+ δ) \, the global perturbation kt δ_k^t at the k-th iteration in t-th round is formulated as follows: (6) kt=ρ∇F(kt)‖∇F(kt)‖=ρ1N∑i=1N∇Fi(kt)‖1N∑i=1N∇Fi(kt)‖, δ_k^t=ρ ∇ F ( θ_k^t ) \|∇ F ( θ_k^t ) \|=ρ 1N _i=1^N∇ F_i ( θ_k^t ) \| 1N _i=1^N∇ F_i ( θ_k^t ) \|, where kt=t−1N∑i=0N(t−i,kt) θ_k^t= θ^t- 1N _i=0^N ( θ^t- θ_i,k^t ) is the virtual global model. However, we can neither share weights nor gradients of clients during the local training in FL. And we have no access to the global gradient ∇F(kt)∇ F ( θ_k^t ), so we use the global momentum t m^t to estimate the global gradient. Therefore, the global perturbation can be approximately calculated as follows: (7) kt=ρ∇F(kt)‖∇F(kt)‖≈ρt‖t‖, δ_k^t=ρ ∇ F ( θ_k^t ) \|∇ F ( θ_k^t ) \|≈ρ m^t \| m^t \|, where t=λt−1+t m^t=λ m^t-1+ ^t, t=1S∑i∈Stit ^t= 1S _i∈ S_t _i^t, it=i,Kt−i,0t _i^t= θ_i,K^t- θ_i,0^t. Next, we can get the perturbed local model, (8) i,k+1/2t=i,kt+kt, θ_i,k+1/2^t= θ_i,k^t+ δ_k^t, Finally, we compute the global SAM gradient at the perturbed model i,k+1/2t θ_i,k+1/2^t to update i,kt θ_i,k^t, (9) i,k+1t=i,kt−η∇F(k+1/2t), θ_i,k+1^t= θ_i,k^t-η∇ F ( θ_k+1/2^t ), where ∇F(k+1/2t)=1N∑i=1N∇Fi(k+1/2t)∇ F ( θ_k+1/2^t )= 1N _i=1^N∇ F_i ( θ_k+1/2^t ). However, the computation of the global SAM gradient is impractical in FL. Due to data heterogeneity and local multi-iteration, ∇Fi(k+1/2t)≠∇F(k+1/2t)∇ F_i ( θ_k+1/2^t )≠∇ F ( θ_k+1/2^t ). To reduce the inconsistency between the local models and the global model, we incorporate global momentum t m_t into the local models to guide local updates in our FedNSAM. Table 1. Summary of federated SAM-based algorithms for addressing data heterogeneity, focusing on perturbation correction, updating correction, acceleration, and extra computation introduced by SAM. Here, perturbation correction is a correction for i,kt δ_i,k^t, updating correction is a correction for ∇Fi(k+1/2t)∇ F_i( θ_k+1/2^t), and Com Cost denotes the communication cost. Research work Perturbation Correction Updating correction Acceleration Less Computation Com Cost FedSAM (Qu et al., 2022) × × × × 1× MoFedSAM (Qu et al., 2022) × ✓ × × 1× FedGAMMA (Dai et al., 2023) × ✓ × × 2× FedSMOO (Sun et al., 2023) ✓ × × × 2× FedLESAM (Fan et al., 2024) ✓ × × ✓ 1× FedNSAM (ours) ✓ ✓ ✓ ✓ 1× 4. Proposed Algorithm: FedNSAM The main idea of FedNSAM is to accelerate the local update using Nesterov momentum that utilizes the global momentum t m_t to guide the local update, and to use the t m_t as the perturbation direction of the SAM algorithm, so that each client can look for a more consistent flat region. (a) FedSAM (b) FedNSAM Figure 3. (a) and (b) depict the local update procedures of FedSAM and FedNSAM, respectively. FedNSAM incorporates global Nesterov momentum into local sharpness-aware updates to improve alignment between local and global flatness. Algorithm 1 FedNSAM Algorithm 1: Input: β,λβ,λ, initial server model 0 θ^0, number of clients N, number of communication rounds T, number of local iterations K, and local learning rate η. 2: Initialize global momentum 0=0 m^0=0, and global model 0 θ^0. 3: for each round t=1t=1 to T do 4: for each selected client i=1i=1 to S do 5: for k=0,…,Kk=0,…,K local update do 6: Update the local model it θ_i^t at Client i; 7: i,k+1/4t=i,kt+λt θ_i,k+1/4^t= θ_i,k^t+λ m^t; ⊳ Nesterov extrapolation 8: i,kt=ρ−t‖t‖ δ_i,k^t=ρ - m^t\| m^t\|; ⊳ Perturbation calculation 9: i,k+1/2t=i,k+1/4t+i,kt; θ_i,k+1/2^t= θ_i,k+1/4^t+ δ_i,k^t; ⊳ Perturbation model 10: i,k+1t=i,kt−η∇Fi(i,k+1/2t;ζi) θ_i,k+1^t\!=\! θ_i,k^t\!-\!η∇ F_i ( θ_i,k+1/2^t; _i ); ⊳ Updating 11: end for 12: it=i,Kt−i,0t _i^t= θ_i,K^t- θ_i,0^t; 13: Client i sends it _i^t back to the server. 14: end for 15: Server averages models: 16: t=1S∑i∈Stit ^t= 1S _i∈ S_t _i^t; 17: t=λt−1+t m^t=λ m^t-1+ ^t; ⊳ Nesterov momentum 18: t=t−1+t θ^t= θ^t-1+ m^t. 19: end for 4.1. The proposed Algorithm Nesterov Accelerated Gradient (NAG) was proposed by Yurii Nesterov (Nesterov, 2013) to accelerate the convergence of the gradient descent algorithm. The following is a core description of NAG algorithm: (10) t+1/2=t−λvt,vt+1=λvt+η∇F(t+1/2), θ_t+1/2= θ_t-λ v_t,v_t+1=λ v_t+η∇ F ( θ_t+1/2 ), (11) t+1=t+vt+1, θ_t+1= θ_t+v_t+1, where 0≤λ<10≤λ<1 is the momentum parameter. The key difference between NAG and SAM is the direction of the extra point (t+1/2 θ_t+1/2) search. NAG searches for an extra point in the direction of vtv_t, while the SAM method searches for an extra point in the direction of gradient ∇F(t)∇ F ( θ_t ). FedNSAM combines the flattening effect of SAM with the acceleration effect of extrapolation. The i-th client updates local model it θ_i^t as follows: (12) i,k+1/4t=i,kt+λt,i,kt=ρ−t‖t‖, θ_i,k+1/4^t= θ_i,k^t+λ m^t, δ_i,k^t=ρ - m^t \| m^t \|, (13) i,k+1/2t=i,k+1/4t+i,kt,i,k+1t=i,kt−η∇Fi(i,k+1/2t;ζi). θ_i,k+1/2^t= θ_i,k+1/4^t+ δ_i,k^t, θ_i,k+1^t= θ_i,k^t-η∇ F_i ( θ_i,k+1/2^t; _i ). Since directly calculating the global gradient ∇F(t)∇ F ( θ_t ) is impractical in FL, we utilize the global momentum t m_t as an approximation, t:=λt−1+t m^t:=λ m^t-1+ ^t. Estimating the global gradient using an exponential moving average to compute t m^t also reduces the effect of stochastic gradient noise (Li et al., 2024). The local updating of FedNSAM can be seen in Figure 3 (b). 4.2. Discussion ∙ Comparison with MoFedSAM and FedGAMMA: FedMoSAM uses heavy ball momentum t ^t to accelerate the local model, t=1S∑i∈Stit ^t= 1S _i∈ S_t _i^t, it=i,Kt−i,0t _i^t= θ_i,K^t- θ_i,0^t. local update as i,k+1t=i,kt−η(λgi,k+1/2t+1−ληKt) θ_i,k+1^t= θ_i,k^t-η (λ g_i,k+1/2^t+ 1-λη K ^t ), which is similar to FedCM. However, FedMoSAM did not have an accelerating effect in theory, while our FedNSAM algorithm uses t m_t as Nesterve momentum. Furthermore, when the client participation rate is low, t ^t does not accurately estimate the amount of global change. Instead, t=λt−1+t m^t=λ m^t-1+ ^t, exponential moving average of the historical gradient, giving higher weight to the latest global change, which allows accurate estimation of the global updating. FedGAMMA introduces the variance reduction control variable used in SCAFFOLD (Karimireddy et al., 2020) to control the consistency of the global and client models. Both algorithms only consider correction to ∇F(k+1/2t)∇ F( θ_k+1/2^t), and do not consider correction to i,kt δ_i,k^t in Table 1. ∙ Comparison with FedSMOO and FedLESAM: FedLESAM uses t ^t as the global perturbation direction estimate, and the FedNSAM algorithm uses t m^t as the global perturbation direction estimate. The global perturbation estimate for FedLESAM, t=1S∑i∈Stit ^t= 1S _i∈ S_t _i^t, it=i,Kt−i,0t _i^t= θ_i,K^t- _i,0^t, which is the difference between the previous active participation round and the current round before the global model is received. When client participation is low, t ^t does not accurately estimate the global perturbation. Instead, t=λt−1+t m^t=λ m^t-1+ ^t, which allows a more accurate estimation of the global perturbation. The FedSMOO uses ∇Fi(i,kt)−μi−s∇ F_i( θ_i,k^t)- _i-s as the global perturbation direction estimate. However, its computational cost and communication are twice as much as that of our algorithm. However, both algorithms only consider the correction to i,kt δ_i,k^t, and do not consider the correction to ∇F(k+1/2t)∇ F( θ_k+1/2^t) in Table 1. 5. Theoretical Analysis We analyze generalization based on following assumptions: Assumption 1 (Smoothness). FiF_i is L-smooth for all i∈i∈ [N][N], (14) ‖∇Fi(1)−∇Fi(2)‖≤L‖1−2‖, \|∇ F_i( θ_1)-∇ F_i( θ_2) \|≤ L\| θ_1- θ_2\|, for all 1,2 θ_1, θ_2 in its domain and i∈[N]i∈[N]. Assumption 2 (Bounded variance of data heterogeneity). The global variability of the local gradient of the loss function is bounded by σg2 _g^2 for all i∈[N]i∈[N], (15) ‖∇Fi()−∇F()‖2≤σg2. \|∇ F_i ( θ )-∇ F ( θ ) \|^2≤ _g^2. Assumption 3 (Bounded variance of stochastic gradient). The stochastic gradient ∇Fi(,ξi)∇ F_i ( θ, _i ), computed by the i-th client of model parameter θ using mini-batch ξi _i is an unbiased estimator ∇Fi()∇ F_i( θ) with variance bounded by σ2σ^2, i.e., (16) ξi‖∇Fi(,ξi)−∇Fi()‖2≤σ2, _ _i \|∇ F_i ( θ, _i )-∇ F_i( θ) \|^2≤σ^2, ∀i∈[N]∀ i∈[N], where the expectation is over all local datasets. 5.1. Convergence Results Theorem 1 (Convergence for non-convex functions). Suppose that local Fii=1N \F_i \_i=1^N are non-convex and L-smooth. By setting η≤(1−λ)2128KLη≤ (1-λ)^2128KL, ρ=1Tρ= 1T, FedNSAM satisfies 1T∑t=1T[‖∇F(t−1+λt−1)‖]≤ 1T _t=1^TE [ \|∇ F ( θ^t\!-\!1\!+\!λ m^t-1 ) \| ]≤ (M1LFTKS(1−λ)+(LF)23M213T23+LFT(1−λ)+L4(1−λ)T), ( M_1 LF TKS(1\!-\!λ)\!+\! (LF) 23M_2 13T 23\!+\! LFT(1\!-\!λ)\!+\! L^4(1-λ)T ), where M12:=σ2+K(1−SN)σg2,M2:=σ2K+σg2M_1^2:=σ^2+K (1- SN ) _g^2,M_2:= σ^2K+ _g^2, and F:=F(0)−F(⋆)F:=F ( θ^0 )-F ( θ ), |St|=S |S_t |=S. Theorem 1 provides the non-convex convergence rate of FedNSAM, which matches the best convergence rate of existing FL methods (Karimireddy et al., 2020). We provide the convergence of FedNSAM under partial client participation. The convergence rate of FedNSAM is (M1LFTKS(1−λ)+(LF)23M213T23+LFT(1−λ)+L4(1−λ)T)O ( M_1 LF TKS(1-λ)+ (LF) 23M_2 13T 23+ LFT(1-λ)+ L^4(1-λ)T ), which is better than the convergence rate of FedSAM’s (LFTKS+Kσg2TS+L2σ2T3/2K+L2T2)O ( LF TKS+ K _g^2 TS+ L^2σ^2T^3/2K+ L^2T^2 ) (Qu et al., 2022). That is beacause (M1LFTKS(1−λ))<(LFTKS)O ( M_1 LF TKS(1-λ) )<O ( LF TKS ) , and LF<LF LF<LF. Moreover, the generalization result of FedSAM is commonly suitable for our FedNSAM. 5.2. Flatness Distance Analysis Table 2. Comparison of testing accuracy (%) on CIFAR10 and CIFAR100, E=5E=5, Dirichlet-0.6 respectively. CIFAR10 CIFAR100 Method LeNet-5 VGG-11 ResNet-18 LeNet-5 VGG-11 ResNet-18 Acc.(%) Rounds Acc.(%) Rounds Acc.(%) Rounds Acc.(%) Rounds Acc.(%) Rounds Acc.(%) Rounds 1000R 78% 1000R 80% 1000R 85% 1000R 52% 1000R 53% 1000R 55% FedAvg 79.63±0.2779.63± 0.27 574 84.14±0.3284.14± 0.32 313 88.92±0.2488.92± 0.24 542 41.15±0.1841.15± 0.18 1000+ 48.94±0.3948.94± 0.39 1000+ 54.25±0.3954.25± 0.39 1000+ FedAvgM 81.15±0.2381.15± 0.23 659 83.74±0.1183.74± 0.11 794 89.05±0.3689.05± 0.36 656 48.35±0.2948.35± 0.29 1000+ 51.94±0.1751.94± 0.17 1000+ 60.91±0.3760.91± 0.37 516 SCAFFOLD 79.43±0.3379.43± 0.33 799 86.94±0.0986.94± 0.09 272 88.56±0.2188.56± 0.21 543 50.21±0.3550.21± 0.35 1000+ 50.81±0.2850.81± 0.28 1000+ 54.12±0.1554.12± 0.15 1000+ FedACG 82.84±0.2882.84± 0.28 301 84.86±0.1484.86± 0.14 249 90.62±0.3190.62± 0.31 222 53.04±0.1953.04± 0.19 729 51.55±0.3451.55± 0.34 1000+ 61.92±0.3161.92± 0.31 521 FedSAM 81.21±0.1681.21± 0.16 361 85.26±0.2785.26± 0.27 227 86.78±0.1386.78± 0.13 313 48.12±0.2248.12± 0.22 1000+ 47.71±0.0947.71± 0.09 1000+ 47.83±0.2547.83± 0.25 628 MoFedSAM 83.42±0.2083.42± 0.20 255 −- 1000+ 86.66±0.3786.66± 0.37 530 50.11±0.1050.11± 0.10 1000+ −- 1000+ 60.15±0.2460.15± 0.24 1000+ FedGAMMA 81.02±0.2481.02± 0.24 371 86.33±0.2586.33± 0.25 169 88.35±0.2988.35± 0.29 542 41.33±0.2341.33± 0.23 1000+ 53.42±0.3253.42± 0.32 908 49.83±0.3549.83± 0.35 1000+ FedLESAM 79.75±0.3079.75± 0.30 496 83.54±0.1883.54± 0.18 169 89.01±0.3589.01± 0.35 480 41.33±0.2441.33± 0.24 1000+ 49.26±0.2149.26± 0.21 1000+ 52.11±0.2152.11± 0.21 1000+ FedNSAM 83.82±0.1583.82± 0.15 181 87.86±0.2987.86± 0.29 137 91.35±0.3391.35± 0.33 272 53.61±0.2053.61± 0.20 633 56.33±0.3156.33± 0.31 334 66.04±0.1166.04± 0.11 316 Theorem 2. For FedSAM, if we choose η=(1TKL)η=O ( 1 TKL ) and ρ=(1T),ρ=O ( 1 T ), _D is then bounded as follows: (17) ≤(σ2KT2+σg2L2T+Kσg2L2NT3/2+L2T+FLKNT3/2) _D\!≤\!O ( σ^2KT^2\!+\! _g^2L^2T\!+\! K _g^2L^2 NT^3/2\!+\! L^2T\!+\! FL KNT^3/2 ) Theorem 3. For FedNSAM, if we choose η=((1−λ)2KL)η=O ( (1-λ)^2KL ) and ρ=(1T),ρ=O ( 1 T ), _D is then bounded as follows: (18) ≤(σ2KT2+σg2L2T+L2T+σFT3/2L3/2KN(1−λ)) _D\!≤\!O ( σ^2KT^2\!+\! _g^2L^2T\!+\! L^2T\!+\! σ FT^3/2L^3/2 KN(1\!-\!λ) ) As shown in Theorem 2 and 3, the term σg2 _g^2, which increases when heterogeneity gets worse, directly determines the upper bound of _D. The local model i θ_i largely deviates from the global model θ with high data heterogeneity. For FedNSAM, the upper bound of _D formalized by Theorem 3 is better than FedSAM (Theorem 2). 6. Experiments 6.1. Experimental Settings Datasets: We evaluate our algorithms on the CIFAR10 (Krizhevsky et al., 2009), CIFAR100 (Krizhevsky et al., 2009), Tiny ImageNet (Le and Yang, 2015). For non-IID data setup, we simulate the data heterogeneity by sampling the label ratios from a Dirichlet distribution (Hsu et al., 2019). Models: To test the robustness of our algorithms, we use standard classifiers (including LeNet-5 (LeCun and others, 2015), VGG-11 (Simonyan and Zisserman, 2014), and ResNet-18 (He et al., 2016)), Vision Transformer (ViT-Base) (Dosovitskiy et al., 2020), Swin transformer (Swin-Small, Swin-Base) (Liu et al., 2021). We first verify the effectiveness of our algorithm by training the model using convolutional networks on the CIFAR10 (Krizhevsky et al., 2009), CIFAR100 (Krizhevsky et al., 2009), Tiny ImageNet (Le and Yang, 2015) datasets. Methods: We compare FedNSAM with many FL baselines, including FedAvg (McMahan et al., 2017), FedAvgM (Hsu et al., 2019), SCAFFOLD (Karimireddy et al., 2020), FedACG (Kim et al., 2024), FedSAM (Qu et al., 2022), MoFedSAM (Qu et al., 2022), FedGAMMA (Dai et al., 2023), FedLESAM (Fan et al., 2024) . Hyper-parameter Settings: The number of clients is 100. batch size B=50B\!=\!50, local epoch E=5E\!=\!5, K=50K\!=\!50. We set the grid search range of the client learning rate by η∈10−3,3×10−3,…,10−1,3×10−1η\!∈\!\10^-3,3× 10^-3,...,10^-1,3× 10^-1\. Learning rate decay per round is 0.9980.998, total T=1,000T=1,000. Specifically, we set ρ=0.1ρ\!=\!0.1, λ=0.85λ=0.85 for FedNSAM. (a) Resnet-18, CIFAR10 (b) ResNet-18, CIFAR100 Figure 4. Convergence plots for FedNSAM and other baselines on Dirichlet-0.6 of CIFAR10 and CIFAR100 with ResNet-18. 6.2. Results on Convolutional Neural Networks We first conduct experiments on convolutional neural networks to verify the superiority of the FedNSAM algorithm on CIFAR100 and CIFAR10 datasets with Dirichlet-0.6 data split on 100 clients, 10%10\% participating setting and E=5E=5, batch size B=50B\!=\!50. We selected three benchmark convolutional neural network frameworks, which are LeNet-5 (LeCun and others, 2015), VGG-11 (Simonyan and Zisserman, 2014), and ResNet-18 (He et al., 2016) in Table 2. From Table 2 (Figure 4), we have the following observations: (i) The FedNSAM algorithm achieves the highest accuracy with different datasets and network architectures with the training rounds 1000. FedNSAM achieves the same accuracy with fewer training rounds. For example, in the case of ResNet-18 training the CIFAR100 dataset for 1000 rounds, FedNSAM’s accuracy of 66.04% is 12.21% higher than FedSAM’s 47.83%. FedNSAM is also higher than other variants of FedSAM. Reaching 55% accuracy, FedNSAM uses only 316 rounds, accelerating FedSAM by more than 3 × and accelerating FedACG 1.6 ×. (i) Compared to the latest momentum FL algorithms (FedAvgM, FedACG) and a series of variants of the SAM algorithm (MoFedSAM, FedGAMMA, FedLESAM), the FedNSAM algorithm has the best convergence speed and final generalization. (a) Impact of λ (b) Impact of ρ Figure 5. Convergence plots for FedNSAM with different λ and ρ, CIFAR100 datasets with ResNet-18. Impact of λ: As shown in our theoretical analysis, λ reduces the effect of client heterogeneity in local updating. To validate this, we also perform experiments to analyze the effect of λ, the only algorithm-dependent hyperparameter of FedNSAM, on the convergence and performance of FedNSAM algorithms. We test FedNSAM with λ taken values in 0.95,0.9,0.85,0.8,0.75,0.7,0.65\0.95,0.9,0.85,0.8,0.75,0.7,0.65\, training ResNet-18 on CIFAR100 datasets with Dirichlet-0.6 split on 100 clients, 10%10\% participating setting. The test accuracies and the convergence plots are provided in Figure 5. We find that FedNSAM successfully converges to stationary points under all these λ choices, as guaranteed by our convergence analysis. However, the stationary points of different λ show different generalization abilities, which results in varying test accuracies in Figure 5. We note that setting λ too small or too large will harm the convergence and generalization of FedNSAM. As shown in Figure 5, as the λ (λ<0.95λ<0.95) increases, the acceleration effect of the FedNSAM algorithm becomes more obvious. Empirically, we find that performance is best when setting λ to about 0.85, which aligns with traditional momentum algorithms FedACG and FedAvgM. Impact of ρ: The hyperparameter ρ controls the perturbation radius in sharpness-aware updates, influencing the model’s ability to seek flat minima. A small ρ may underexplore curvature, while a large ρ can cause instability or underfitting, especially with high data heterogeneity. Empirically, we find that ρ=0.1ρ=0.1 offers the best balance between generalization and stability across datasets and models. As shown in Figure 5 (right), increasing ρ to 0.1 improves test accuracy, while larger values degrade performance. Thus, we set ρ=0.1ρ=0.1 as the default in all experiments. Table 3. Comparision the accuracy of different Transformer models with Dirichlet-0.1, 5% participation. Data Set Tiny ImageNet Model Swin-Small Swin-Base ViT-Base Param 50M 80M 88M FLOPs 8.7G 15.4G 55.4G Acc Rounds Acc Rounds Acc Rounds Method 100R 68% 100R 68% 100R 68% FedAvg 69.80±0.21± 0.21 100+ 67.41 100+ 64.64 100+ FedAvgM 67.63 100+ 68.18 75 69.77 80 SCAFFOLD 69.62 82 67.37 88 65.02 100+ FedAGC 64.38 100+ 65.54 100+ 70.22 41 FedSAM 63.64 100+ 64.56 100+ 65.46 100+ MoFedSAM 64.26 100+ 64.45 100+ 64.60 100+ FedGAMMA 65.02 100+ 65.64 100+ 66.02 100+ FedLESAM 65.11 100+ 65.89 100+ 66.28 100+ FedNSAM 70.12 67 70.86 42 71.23 26 Table 4. Comparison of testing accuracy (%) and convergence rounds on CIFAR100 (E=5E=5) for different participation rates (left) and different data heterogeneity levels (right). Participation Rate Data Heterogeneity (Dirichlet-α) Method p=2%p=2\% p=5%p=5\% p=10%p=10\% α=0.1α=0.1 α=0.3α=0.3 α=0.6α=0.6 Acc R Acc R Acc R Acc R Acc R Acc R FedAvg 51.61± 0.35 1000+ 51.53± 0.12 1000+ 54.25± 0.39 1000+ 45.81± 0.19 1000+ 52.53± 0.15 1000+ 54.25± 0.39 1000+ FedAvgM 15.72± 0.23 1000+ 46.63± 0.33 1000+ 60.91± 0.37 516 48.63± 0.40 1000+ 58.84± 0.08 808 60.91± 0.37 516 SCAFFOLD 51.93± 0.29 1000+ 53.64± 0.37 1000+ 54.12± 0.28 1000+ 49.02± 0.33 1000+ 51.92± 0.25 1000+ 54.12± 0.28 1000+ FedACG 54.01± 0.16 1000+ 61.94± 0.26 506 61.92± 0.31 521 53.56± 0.34 1000+ 60.87± 0.12 425 61.92± 0.31 521 FedSAM 49.02± 0.18 1000+ 48.53± 0.24 1000+ 47.83± 0.09 1000+ 40.18± 0.27 1000+ 46.02± 0.22 1000+ 47.83± 0.09 1000+ MoFedSAM 44.22± 0.09 1000+ 60.64± 0.34 511 60.15± 0.24 616 51.59± 0.16 1000+ 58.21± 0.29 756 60.15± 0.24 616 FedGAMMA 49.42± 0.39 1000+ 51.22± 0.17 1000+ 49.83± 0.35 1000+ 47.73± 0.18 1000+ 46.65± 0.38 1000+ 49.83± 0.35 1000+ FedLESAM 54.21± 0.31 1000+ 54.22± 0.28 1000+ 52.11± 0.21 1000+ 48.74± 0.30 1000+ 53.34± 0.17 1000+ 52.11± 0.21 1000+ FedNSAM 56.92± 0.27 872 62.21± 0.11 496 66.04± 0.11 316 58.53± 0.26 695 63.65± 0.36 372 66.04± 0.11 316 (a) Dirichlet-0.3 (b) Dirichlet-0.1 (c) p=2%p=2\% (d) p=5%p=5\% Figure 6. Convergence plots of FedNSAM and baselines on CIFAR100 with ResNet-18 under different heterogeneity (a,b) and participation rates (c,d). 6.3. Results on Vision Transformer Models To evaluate the effectiveness of our method on large-scale models, we conduct experiments on the Tiny ImageNet dataset using Swin-Small, Swin-Base, and ViT-Base, all initialized with ImageNet-22k pre-trained weights. We train for 100 communication rounds under Dirichlet-0.1 with a learning rate of 0.01 (decayed by 0.99 per round) and a batch size of 16. Table 3 reports the parameter sizes and theoretical FLOPs of each model (input size 224), along with the performance comparison. FedNSAM consistently achieves the highest accuracy—70.12%, 70.86%, and 71.23%—on Swin-Small, Swin-Base, and ViT-Base respectively, while requiring significantly fewer rounds (67, 42, and 26). These results clearly demonstrate the superior generalization and training efficiency of FedNSAM, especially for large-scale transformer models and high-complexity datasets. (a) FedSAM (b) MoFedSAM (c) FedGAMMA (d) FedLESAM (e) FedAvg (f) FedNSAM (ours) (g) FedSAM (h) MoFedSAM (i) FedGAMMA (j) FedLESAM (k) FedAvg (l) FedNSAM (ours) Figure 7. Global training loss surfaces plots (a,b,c,d,e,f) and global testing loss surfaces plots (g,h,i,j,k,l) for FedNSAM and other baselines in different settings on Dirichlet-0.1 of CIFAR100 datasets with ResNet-18. FedNSAM could approach a more general and flat loss landscape which efficiently improves the generalization performance in FL. 6.4. Participation Rate and Heterogeneity Levels To validate the stability of our algorithm, we conducted multiple sets of experiments on the ResNet-18 model and the CIFAR100 dataset under different client participation rates and different data heterogeneity levels. Impact of heterogeneity: Via sampling with replacement, not only is there a data imbalance between local clients, but also the number of samples among categories in the global dataset is also very different, as shown in Figure 6. With the high heterogeneity (Dirichlet-0.1), several baselines are greatly affected. We select the variance Dirichlet coefficient 0.1, 0.3, and 0.6 on the 10%10\% participation. In detail, on CIFAR-100, when the Dirichle coefficient is from 0.6 to 0.1, FedNSAM drops from 66.04%66.04\% to 58.53%58.53\%, while FedSAM drops from 47.83%47.83\% to 40.18%40.18\%, which drops 7.65%7.65\%. FedNSAM achieves good performance on the strong heterogeneous dataset with high stability on different setups. Impact of partial participation: To fairly compare with the baselines, we froze the other hyperparameters selections as shown in Figure 6 and Table 4. When the participation ratio decreases from 10%10\% to 5%5\% and 2%2\% on CIFAR-100, FedNSAM still achieves excellent performance, which drops from 66.04%66.04\% to 56.92%56.92\% on Dirichlet-0.6, FedNSAM maintains a high level of generalization on 2%2\% participation, which achieves the accuracy 56.92%56.92\%. 6.5. Comparison with FedSAM and Its Variants To better evaluate the effectiveness of the proposed global Nesterov extrapolation, we compare FedNSAM with several representative FedSAM-based methods, including FedSAM, MoFedSAM, FedGAMMA, and FedLESAM. We visualize their global training and testing loss landscapes on the CIFAR100 dataset (Dirichlet-0.1) using ResNet-18, as shown in Figure 7. In the case of high data heterogeneity, FedSAM and its variants generally fail to obtain sufficiently flat global minima. Specifically, FedSAM exhibits sharp loss surfaces due to inconsistent local perturbations, while MoFedSAM and FedGAMMA incorporate update correction but still struggle to improve global flatness. FedLESAM alleviates part of this inconsistency by estimating global perturbations, but its estimation remains inaccurate, limiting its effectiveness. In contrast, FedNSAM achieves significantly flatter global loss landscapes in both training and testing, as shown in Figure 7 (f, l). This highlights its ability to harmonize local training dynamics and reduce sharpness more effectively under heterogeneous settings. Table 5. Comparison different data heterogeneity of testing accuracy (%) of ablation on variants on CIFAR100. CIFAR100 Method Dirichlet-0.1 Dirichlet-0.6 Acc Rounds Acc Rounds 1000R 55% 1000R 55% FedAvg 45.81±0.35± 0.35 1000+ 54.25±0.19± 0.19 1000+ FedNSAM 58.53±0.27± 0.27 695 66.04±0.22± 0.22 316 SCAFFOLD 49.02±0.29± 0.29 1000+ 54.12±0.40± 0.40 1000+ FedNSAM-S 59.63±0.32± 0.32 676 66.24±0.18± 0.18 301 FedDyn 53.53±0.16± 0.16 1000+ 56.95±0.30± 0.30 865 FedNSAM-D 60.17±0.34± 0.34 663 66.48±0.25± 0.25 296 6.6. Ablation on FedNSAM Variants We design FedNSAM under FedAvg and two enhanced methods under Scaffold and FedDyn, named FedNSAM-S and FedNSAM-D, respectively. Therefore, we show the average performance of all variants on CIFAR100 in Table 5. All the variants can achieve an extensive improvement compared to their base methods, especially a notable 9.53%9.53\% improvement on CIFAR100 of FedNSAM-D. 6.7. Effect of Nesterov Momentum As shown in Table 6, we can find that the proposed Nesterov momentum acceleration consistently improves the accuracy and reduces the number of rounds required for convergence across all the methods. For instance, FedAvg with Nesterov momentum achieved an accuracy of 54.23% on Dirichlet-0.1 and 60.26% on Dirichlet-0.6, outperforming FedAvg without Nesterov momentum. Similarly, FedSAM+Nesterov achieved higher accuracy than FedSAM, with a significant reduction in the required rounds. However, even without Nesterov momentum, FedNSAM performed very well, achieving an accuracy of 55.86% on Dirichlet-0.1 and 60.11% on Dirichlet-0.6, with fewer rounds than other methods. This suggests that our algorithm is inherently effective, even without the acceleration of Nesterov momentum. When Nesterov momentum was applied, the performance of FedNSAM can be further improved, achieving the highest accuracy of 58.53% and 66.04% on Dirichlet-0.1 and Dirichlet-0.6, respectively, with a significant reduction in rounds. Table 6. Comparison different data heterogeneity of testing accuracy (%) of each algorithm with Nesterov on CIFAR100. CIFAR100 Method Dirichlet-0.1 Dirichlet-0.6 Acc Rounds Acc Rounds 1000R 55% 1000R 55% FedAvg 45.81±0.32± 0.32 1000+ 54.25±0.15± 0.15 1000+ FedAvg+Nesterov 54.23±0.28± 0.28 1000+ 60.26±0.39± 0.39 512 FedSAM 40.18±0.27± 0.27 1000+ 47.83±0.09± 0.09 1000+ FedSAM+Nesterov 55.88±0.34± 0.34 876 61.02±0.09± 0.09 482 FedNSAM - Nesterov 55.86±0.40± 0.40 895 60.11±0.19± 0.19 488 FedNSAM + Nesterov 58.53±0.26± 0.26 659 66.04±0.11± 0.11 316 7. Conclusion We revisited Sharpness-Aware Minimization (SAM) in the federated learning setting through the novel lens of flatness distance, which captures the divergence of local optimization landscapes under data heterogeneity. We can observe that increasing heterogeneity exacerbates flatness distance, leading to sharper and less generalizable global minima. To address these challenges, we proposed a novel principled FL algorithm (called FedNSAM) that integrates global Nesterov extrapolation into the SAM framework to align client-level flat regions. Extensive experiments demonstrated that FedNSAM significantly outperforms existing FL baselines, especially under non-IID conditions, achieving superior generalization and convergence. Acknowledgments This work was supported by the National Natural Science Foundation of China (No. 62276182), Peng Cheng Lab Program (No. PCL2023A08), Tianjin Natural Science Foundation (Nos. 24JCYBJC01230, 24JCYBJC01460), and Tianjin Municipal Education Commission Research Plan (No. 2024ZX008). References R. S. Antunes, C. André da Costa, A. Küderle, I. A. Yari, and B. Eskofier (2022) Federated learning for healthcare: systematic review and architecture proposal. ACM Transactions on Intelligent Systems and Technology (TIST) 13 (4), p. 1–23. Cited by: §1. D. Byrd and A. Polychroniadou (2020) Differentially private secure multi-party computation for federated learning in financial applications. In Proceedings of the First ACM International Conference on AI in Finance, p. 1–9. Cited by: §1. D. Caldarola, B. Caputo, and M. Ciccone (2022) Improving generalization in federated learning by seeking flat minima. In European Conference on Computer Vision, p. 654–672. Cited by: §1, §2. M. Dai, S. Liu, Z. Zhao, J. Gao, H. Sun, and X. Li (2025a) Secure tug-of-war (sectow): iterative defense-attack training with reinforcement learning for multimodal model security. External Links: 2507.22037, Link Cited by: §1. M. Dai, J. Sun, Z. Zhao, S. Liu, R. Li, J. Gao, and X. Li (2025b) From captions to rewards (carevl): leveraging large language model experts for enhanced reward modeling in large vision-language models. External Links: 2503.06260, Link Cited by: §1. R. Dai, X. Yang, Y. Sun, L. Shen, X. Tian, M. Wang, and Y. Zhang (2023) Fedgamma: federated learning with global sharpness-aware minimization. IEEE Transactions on Neural Networks and Learning Systems. Cited by: §2, Table 1, §6.1. A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. (2020) An image is worth 16x16 words: transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. Cited by: §6.1. Z. Fan, S. Hu, J. Yao, G. Niu, Y. Zhang, M. Sugiyama, and Y. Wang (2024) Locally estimated global perturbations are better than local perturbations for federated sharpness-aware minimization. arXiv preprint arXiv:2405.18890. Cited by: §1, §2, §3.4, Table 1, §6.1. P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur (2020) Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412. Cited by: §1, §2. K. He, X. Zhang, S. Ren, and J. Sun (2016) Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, p. 770–778. Cited by: §6.1, §6.2. T. H. Hsu, H. Qi, and M. Brown (2019) Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335. Cited by: §2, §6.1. S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh (2020) Scaffold: stochastic controlled averaging for federated learning. In International Conference on Machine Learning, p. 5132–5143. Cited by: §2, §4.2, §5.1, §6.1. G. Kim, J. Kim, and B. Han (2024) Communication-efficient federated learning with accelerated client gradient. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 12385–12394. Cited by: §2, §6.1. A. Krizhevsky, G. Hinton, et al. (2009) Learning multiple layers of features from tiny images. Cited by: §6.1. Y. Le and X. Yang (2015) Tiny imagenet visual recognition challenge. CS 231N 7 (7), p. 3. Cited by: §6.1. Y. LeCun et al. (2015) LeNet-5, convolutional neural networks. URL: http://yann. lecun. com/exdb/lenet 20 (5), p. 14. Cited by: §6.1, §6.2. T. Li, P. Zhou, Z. He, X. Cheng, and X. Huang (2024) Friendly sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 5631–5640. Cited by: §4.1. T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith (2020) Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, Cited by: §2. J. Liu, Y. Liu, F. Shang, H. Liu, J. Liu, and W. Feng (2025a) Improving generalization in federated learning with highly heterogeneous data via momentum-based stochastic controlled weight averaging. In Forty-second International Conference on Machine Learning, Cited by: §1, §2. J. Liu, F. Shang, H. Liu, J. Liu, W. An, and Y. Liu (2026) Taming preconditioner drift: unlocking the potential of second-order optimizers for federated learning on non-iid data. External Links: 2602.19271, Link Cited by: §2. J. Liu, F. Shang, Y. Liu, H. Liu, Y. Li, and Y. Gong (2024) Fedbcgd: communication-efficient accelerated block coordinate gradient descent for federated learning. In Proceedings of the 32nd ACM International Conference on Multimedia, p. 2955–2963. Cited by: §1, §2. J. Liu, F. Shang, Y. Tian, H. Liu, and Y. Liu (2025b) Consistency of local and global flatness for federated learning. In Proceedings of the 33rd ACM International Conference on Multimedia, p. 3875–3883. Cited by: §2. J. Liu, F. Shang, J. Zhou, H. Liu, Y. Liu, and J. Liu (2025c) FedMuon: accelerating federated learning with matrix orthogonalization. arXiv preprint arXiv:2510.27403. Cited by: §2. J. Liu, F. Shang, K. Zhu, H. Liu, Y. Liu, and J. Liu (2025d) FedAdamW: a communication-efficient optimizer with convergence and generalization guarantees for federated large models. arXiv preprint arXiv:2510.27486. Cited by: §2. J. Liu, Y. Tian, F. Shang, Y. Liu, H. Liu, J. Zhou, and D. Ding (2025e) DP-fedpgn: finding global flat minima for differentially private federated learning via penalizing gradient norm. arXiv preprint arXiv:2510.27504. Cited by: §2. Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo (2021) Swin transformer: hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, p. 10012–10022. Cited by: §6.1. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas (2017) Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, p. 1273–1282. Cited by: §1, §6.1. Y. Nesterov (2013) Introductory lectures on convex optimization: a basic course. Vol. 87, Springer Science & Business Media. Cited by: §4.1. Z. Qu, X. Li, R. Duan, Y. Liu, B. Tang, and Z. Lu (2022) Generalized federated learning via sharpness aware minimization. In International Conference on Machine Learning, p. 18250–18280. Cited by: Figure 1, Figure 1, §1, §10, Lemma 6, §2, §3.2, Table 1, Table 1, §5.1, §6.1. N. Rieke, J. Hancox, W. Li, F. Milletari, H. R. Roth, S. Albarqouni, S. Bakas, M. N. Galtier, B. A. Landman, K. Maier-Hein, et al. (2020) The future of digital health with federated learning. NPJ digital medicine 3 (1), p. 119. Cited by: §1. K. Simonyan and A. Zisserman (2014) Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556. Cited by: §6.1, §6.2. Y. Sun, L. Shen, S. Chen, L. Ding, and D. Tao (2023) Dynamic regularized sharpness aware minimization in federated learning: approaching global consistency and smooth landscape. In International Conference on Machine Learning, p. 32991–33013. Cited by: §1, §2, Table 1. L. Yang, H. Liu, Y. Liu, F. Shang, L. Wan, and L. Jiao (2025) Distillation guided deep unfolding network with frequency hierarchical regularization for low-dose ct image denoising. Neurocomputing, p. 130535. Cited by: §1. S. Zeng, X. Chang, M. Xie, X. Liu, Y. Bai, Z. Pan, M. Xu, and X. Wei (2025) FutureSightDrive: thinking visually with spatio-temporal cot for autonomous driving. arXiv preprint arXiv:2505.17685. Cited by: §1. Algorithm 2 FedNSAM Algorithm 1: Input: β,λβ,λ, initial server model 0 θ^0, number of clients N, number of communication rounds T, number of local iterations K, local learning rate η 2: Initialize global momentum 0=0 m^0=0, global model 0 θ^0 3: for each round t=1t=1 to T do 4: for each selected client i=1i=1 to S do 5: for k=0,…,Kk=0,…,K local update do 6: Client i Update local model it θ_i^t, 7: i,k+1/4t=i,kt+λt θ_i,k+1/4^t= θ_i,k^t+λ m^t ⊳ Nesterov extrapolation 8: i,kt=ρ−t‖t‖ δ_i,k^t=ρ - m^t\| m^t\| ⊳ Perturbation calculation 9: i,k+1/2t=i,k+1/4t+i,kt θ_i,k+1/2^t= θ_i,k+1/4^t+ δ_i,k^t ⊳ Perturbation model 10: i,k+1t=i,kt−η∇Fi(i,k+1/2t;ζi) θ_i,k+1^t\!=\! θ_i,k^t\!-\!η∇ F_i ( θ_i,k+1/2^t; _i ) ⊳ Updating 11: end for 12: end for 13: it=i,Kt−i,0t _i^t= θ_i,K^t- θ_i,0^t 14: Client sends it _i^t back to the server 15: Server averages models: 16: t=1S∑i∈Stit ^t= 1S _i∈ S_t _i^t 17: t=λt−1+t m^t=λ m^t-1+ ^t ⊳ Nesterov momentum 18: t=t−1+t θ^t= θ^t-1+ m^t 19: end for 8. Appendix A: Basic Assumptions Assumption 1. We first state a few assumptions for the local loss functions Fi(⋅)F_i(·) . First, the local function Fi(⋅)F_i(·) is assumed to be L-smooth for all Ci∈C1,…,CNC_i∈ \C_1,…,C_N \, i.e., ‖∇Fi(x)−∇Fi(y)‖≤L‖x−y‖∀x,y \|∇ F_i(x)-∇ F_i(y) \|≤ L\|x-y\|∀ x,y This also implies Fi(y)≤Fi(x)+⟨∇Fi(x),y−x⟩+L2‖y−x‖2.F_i(y)≤ F_i(x)+ ∇ F_i(x),y-x + L2\|y-x\|^2. Assumption 2. Second, we assume the stochastic gradient of the local loss function ∇fi(x):=∇Fi(x;i)∇ f_i(x):=∇ F_i (x;D_i ) is unbiased and possesses a bounded variance, i.e. i[‖∇fi(x)−∇Fi(x)‖2]≤σ2E_D_i [ \|∇ f_i(x)-∇ F_i(x) \|^2 ]≤σ^2. Assumption 3 (Bounded heterogeneity). Third, we assume the average norm of local gradients is bounded by a function of the global gradient magnitude as 1N∑i=1N‖∇Fi(x)‖2≤σg2+B2‖∇F(x)‖2 1N _i=1^N \|∇ F_i(x) \|^2≤ _g^2+B^2\|∇ F(x)\|^2, where σg≥0 _g≥ 0 and B≥1B≥ 1. Based on the above assumptions, we derive the following asymptotic convergence bound of FedACG. The dissimilarity of Fi()F_i( θ) and f()f( θ) is bounded as follows: (19) 1m∑i=1m‖∇Fi()−∇F()‖2≤σg2. 1m _i=1^m \|∇ F_i( θ)-∇ F( θ) \|^2≤ _g^2. 1N∑i=1N‖∇Fi(x)−∇F(x)+∇F(x)‖2 1N _i=1^N \|∇ F_i(x)-∇ F(x)+∇ F(x) \|^2 ≤1N∑i=1N‖∇Fi(x)−∇F(x)‖2+21N∑i=1N⟨∇Fi(x)−∇F(x),∇F(x)⟩+1N∑i=1N‖∇F(x)‖2 ≤ 1N _i=1^N \|∇ F_i(x)-∇ F(x) \|^2+2 1N _i=1^N ∇ F_i(x)-∇ F(x),∇ F(x) + 1N _i=1^N\|∇ F(x)\|^2 ≤1N∑i=1N‖∇Fi(x)−∇F(x)‖2+1N∑i=1N‖∇F(x)‖2 ≤ 1N _i=1^N \|∇ F_i(x)-∇ F(x) \|^2+ 1N _i=1^N\|∇ F(x)\|^2 ≤σg2+‖∇F(x)‖2 ≤ _g^2+\|∇ F(x)\|^2 9. Appendix B: Main Lemmas We present several technical lemmas that are useful for subsequent proofs. Lemma 0. (relaxed triangle inequality). Let v1,…,vτ \v_1,…,v_τ \ be τ vectors in ℝdR^d. Then the following are true: (1) ‖vi+vj‖2≤ \|v_i+v_j \|^2≤ (1+a)‖vi‖2+(1+1a)‖vj‖2(1+a) \|v_i \|^2+ (1+ 1a ) \|v_j \|^2 for any a>0a>0, and (2) ‖∑i=1τvi‖2≤τ∑i=1τ‖vi‖2 \| _i=1^τv_i \|^2≤τ _i=1^τ \|v_i \|^2. Lemma 0. (sub-linear convergence rate). For every non-negative sequence dr−1r≥1 \d_r-1 \_r≥ 1 and any parameters ηmax≥0,c≥0 _ ≥ 0,c≥ 0, R≥0R≥ 0, there exists a constant step-size η≤ηmaxη≤ _ and weights wr=1w_r=1 such that, ΨR:=1R+1∑r=1R+1(dr−1η−drη+c1η+c2η2)≤d0ηmax(R+1)+2c1d0R+1+2(d0R+1)23c213 _R:= 1R+1 _r=1^R+1 ( d_r-1η- d_rη+c_1η+c_2η^2 )≤ d_0 _ (R+1)+ 2 c_1d_0 R+1+2 ( d_0R+1 ) 23c_2 13 Proof. Unrolling the sum, we can simplify ΨR≤d0η(R+1)+c1η+c2η2 _R≤ d_0η(R+1)+c_1η+c_2η^2 The lemma can be established through the adjustment of η. We consider the following two cases based on the magnitudes of R and ηmax _ : - When R+1≤d0c1ηmax2R+1≤ d_0c_1 _ ^2 and R+1≤d0c2ηmax3R+1≤ d_0c_2 _ ^3, selecting η=ηmaxη= _ satisfies ΨR≤d0ηmax(R+1)+c1ηmax+c2ηmax2≤d0ηmax(R+1)+c1d0R+1+(d0R+1)23c213 _R≤ d_0 _ (R+1)+c_1 _ +c_2 _ ^2≤ d_0 _ (R+1)+ c_1d_0 R+1+ ( d_0R+1 ) 23c_2 13 In the other case, we have ηmax2≥d0c1(R+1) _ ^2≥ d_0c_1(R+1) or ηmax3≥d0c2(R+1) _ ^3≥ d_0c_2(R+1). Choosing η=mind0c1(R+1),d0c2(R+1)3η= \ d_0c_1(R+1), [3] d_0c_2(R+1) \ satisfies ΨR≤d0η(R+1)+cη=2c1d0R+1+2d02c2(R+1)23 _R≤ d_0η(R+1)+cη= 2 c_1d_0 R+1+2 [3] d_0^2c_2(R+1)^2 ∎ Lemma 0. (separating mean and variance). Given a set of τ random variables 1,…,τ \x_1,…,x_τ \ in ℝdR^d, where [i∣i−1,…1]=E [x_i _i-1,…x_1 ]= ξi _i and [‖i−ξi‖2]≤σ2E [ \|x_i- _i \|^2 ]≤σ^2 represent their conditional mean and variance, respectively, the variables i−ξi \x_i- _i \ form a martingale difference sequence. Based on this setup, the following holds [‖∑i=1τi‖2]≤2‖∑i=1τξi‖2+2τσ2E [ \| _i=1^τx_i \|^2 ]≤ 2 \| _i=1^τ _i \|^2+2τσ^2 Proof. [‖∑i=1τi‖2] [ \| _i=1^τx_i \|^2 ] ≤2‖∑i=1τξi‖2+2[‖∑i=1τi−ξi‖2] ≤ 2 \| _i=1^τ _i \|^2+2E [ \| _i=1^τx_i- _i \|^2 ] =2‖∑i=1τξi‖2+2∑i[‖i−ξi‖2] =2 \| _i=1^τ _i \|^2+2 _iE [ \|x_i- _i \|^2 ] ≤2‖∑i=1τξi‖2+2τσ2 ≤ 2 \| _i=1^τ _i \|^2+2τσ^2 The first inequality comes from the relaxed triangle inequality and the following equality holds because i−ξi \x_i- _i \ forms a martingale difference sequence. ∎ Lemma 0. (Bounded perturbation difference). Let Assumption 1 and 2 hold, given local perturbations i,kt(k=0,1,…,E− δ_i,k^t(k=0,1,…,E- 1) at any step and local perturbation i,0t δ_i,0^t at the first step, the variance of perturbation difference in FedSAM can be bounded as: 1N∑i[‖i,kt−i,0t‖2]≤2K2L2η2ρ2 1N _iE [ \| δ_i,k^t- δ_i,0^t \|^2 ]≤ 2K^2L^2η^2ρ^2 However in our FedNSAM, it is zero since the perturbation is consistent during the local training within a round: 1N∑i[‖i,kt−i,0t‖2]=0 1N _iE [ \| δ_i,k^t- δ_i,0^t \|^2 ]=0 10. Appendix D: Theoretical Results Theorem 1 (Convergence for non-convex functions). Suppose that local functions Fii=1N \F_i \_i=1^N are non-convex and L-smooth. By setting η≤(1−λ)2128KLη≤ (1-λ)^2128KL, FedNSAM satisfies mint=1,…,T‖∇F(t−1+λmt−1)‖2 _t=1,…,TE \|∇ F ( θ^t-1+λ m^t-1 ) \|^2 ≤(M1LDTK|St|+(LD(1−λ)2)23M213(T+1)23+LDT) ( M_1 LD TK |S_t |+ (LD(1-λ)^2 ) 23M_2 13(T+1) 23+ LDT ) where M12:=σ2+K(1−|St|N)σg2,M2:=σ2K+σg2M_1^2:=σ^2+K (1- |S_t |N ) _g^2,M_2:= σ^2K+ _g^2, and D:=F(0)−F(∗)1−λD:= F ( θ^0 )-F ( θ^* )1-λ. Proof. Let zt=t+λ1−λtz^t= θ^t+ λ1-λ m^t and Φt=t+λt=0+1/4t ^t= θ^t+λ m^t= θ^t_0+1/4, Note that z0=0z^0= θ^0 and zt−zt−1=11−λtz^t-z^t-1= 11-λ δ^t, t=ηS∑k,Ci∇Fi(i,k−1+1/2t;ζi) δ^t= ηS _k,C_i∇ F_i ( θ_i,k-1+1/2^t; _i ). By the smoothness of the function F()F(x), we have F(zt)≤F(zt−1)+⟨∇F(zt−1),zt−zt−1⟩+L2‖zt−zt−1‖2F (z^t )≤ F (z^t-1 )+ ∇ F (z^t-1 ),z^t-z^t-1 + L2 \|z^t-z^t-1 \|^2 By taking the expectation on both sides, we have (20) [F(zt)] [F (z^t ) ] ≤[F(zt−1)]+11−λ[⟨∇F(zt−1),t⟩]+L2[‖zt−zt−1‖2] [F (z^t-1 ) ]+ 11-λE [ ∇ F (z^t-1 ), δ^t ]+ L2E [ \|z^t-z^t-1 \|^2 ] =[F(zt−1)]+11−λ[⟨∇F(zt−1)−∇F(Φt−1),t⟩]+11−λ[⟨∇F(Φt−1),t⟩]+L2(1−λ)2[‖t‖2] =E [F (z^t-1 ) ]+ 11-λE [ ∇ F (z^t-1 )-∇ F ( ^t-1 ), δ^t ]+ 11-λE [ ∇ F ( ^t-1 ), δ^t ]+ L2(1-λ)^2E [ \| δ^t \|^2 ] We note that (21) 11−λ[⟨∇F(zt−1)−∇F(Φt−1),t⟩] 11-λE [ ∇ F (z^t-1 )-∇ F ( ^t-1 ), δ^t ] ≤1−λ2λ3L[‖∇F(zt−1)−∇F(Φt−1)‖2]+λ3L2(1−λ)3[‖t‖2] ≤ 1-λ2λ^3LE [ \|∇ F (z^t-1 )-∇ F ( ^t-1 ) \|^2 ]+ λ^3L2(1-λ)^3E [ \| δ^t \|^2 ] ≤(1−λ)L2λ3[‖zt−1−Φt−1‖2]+λ3L2(1−λ)3[‖t‖2] ≤ (1-λ)L2λ^3E [ \|z^t-1- ^t-1 \|^2 ]+ λ^3L2(1-λ)^3E [ \| δ^t \|^2 ] ≤L2(1−λ)[‖mt−1‖2]+L2(1−λ)3[‖t‖2] ≤ L2(1-λ)E [ \|m^t-1 \|^2 ]+ L2(1-λ)^3E [ \| δ^t \|^2 ] where the first inequality holds because ⟨a,b⟩≤12(‖a‖2+‖b‖2) a,b ≤ 12 (\|a\|^2+\|b\|^2 ), while the second inequality follows from the L-smoothness. The third inequality follows because zt−Φt=λ21−λtz^t- ^t= λ^21-λ m^t and 0≤λ<10≤λ<1, t=ηS∑k,Ci∇Fi(i,k−1+1/2t;ζi) ^t= ηS _k,C_i∇ F_i ( θ_i,k-1+1/2^t; _i ) ,|St|=S |S_t |=S. We also note that 11−λ[⟨∇F(Φt−1),t⟩] 11-λE [ ∇ F ( ^t-1 ), δ^t ] =11−λ[⟨∇F(Φt−1),−ηKN∑k,Ci∇Fi(i,k−1+1/2t;ζi)⟩] = 11-λE [ ∇ F ( ^t-1 ), -η KKN _k,C_i∇ F_i ( θ_i,k-1+1/2^t; _i ) ] ≤ηK2(1−λ)([‖∇F(Φt−1)−1KN∑k,Ci∇Fi(i,k−1+1/2t;ζi)‖2]−[‖∇F(Φt−1)‖2]) ≤ η K2(1-λ) (E [ \|∇ F ( ^t-1 )- 1KN _k,C_i∇ F_i ( θ_i,k-1+1/2^t; _i ) \|^2 ]-E [ \|∇ F ( ^t-1 ) \|^2 ] ) ≤ηK2(1−λ)(L2KN∑k,Ci[‖i,k−1+1/2t−i,0+1/4t‖2]−[‖∇F(Φt−1)‖2]) ≤ η K2(1-λ) ( L^2KN _k,C_iE [ \| θ_i,k-1+1/2^t- θ_i,0+1/4^t \|^2 ]-E [ \|∇ F ( ^t-1 ) \|^2 ] ) where the first inequality holds because ⟨a,b⟩≤12‖a+b‖2−12‖a‖2 a,b ≤ 12\|a+b\|^2- 12\|a\|^2. Substituting Eq.(21) and Eq.(22)into Eq. (20) yields [F(zt)]≤ [F (z^t ) ]≤ [F(zt−1)]+ηK2(1−λ)(L2KN∑k,C[‖i,k−1+1/2t−i,0+1/4t‖2]−[‖∇F(Φt−1)‖2]) [F (z^t-1 ) ]+ η K2(1-λ) ( L^2KN _k,CE [ \| θ_i,k-1+1/2^t- θ_i,0+1/4^t \|^2 ]-E [ \|∇ F ( ^t-1 ) \|^2 ] ) +L2(1−λ)[‖mt−1‖2]+(L2(1−λ)3+L2(1−λ)2)[‖t‖2] + L2(1-λ)E [ \|m^t-1 \|^2 ]+ ( L2(1-λ)^3+ L2(1-λ)^2 )E [ \| ^t \|^2 ] By rearranging the inequality above, we have ηK2(1−λ)[‖∇F(Φt−1)‖2]≤ η K2(1-λ)E [ \|∇ F ( ^t-1 ) \|^2 ]≤ ([F(zt−1)]−[F(zt)])+ηKL22(1−λ)1KN∑k,Ci[‖i,k−1+1/2t−i,0+1/4t‖2] (E [F (z^t-1 ) ]-E [F (z^t ) ] )+ η KL^22(1-λ) 1KN _k,C_iE [ \| θ_i,k-1+1/2^t- θ_i,0+1/4^t \|^2 ] +L2(1−λ)[‖mt−1‖2]+(L2(1−λ)3+L2(1−λ)2)[‖t‖2] + L2(1-λ)E [ \|m^t-1 \|^2 ]+ ( L2(1-λ)^3+ L2(1-λ)^2 )E [ \| ^t \|^2 ] Summing the above inequality for t∈1,…,Tt∈\1,…,T\ yields (22) ηK2(1−λ)∑t=1T[‖∇F(Φt−1)‖2]≤ η K2(1-λ) _t=1^TE [ \|∇ F ( ^t-1 ) \|^2 ]≤ ([F(z0)]−[F(zT)])+ηKL22(1−λ)∑t=1T1KN∑k,Ci[‖i,k−1+1/2t−i,0+1/4t‖2] (E [F (z^0 ) ]-E [F (z^T ) ] )+ η KL^22(1-λ) _t=1^T 1KN _k,C_iE [ \| θ_i,k-1+1/2^t- θ_i,0+1/4^t \|^2 ] +L2(1−λ)∑t=1T[‖mt−1‖2]+(L2(1−λ)3+L2(1−λ)2)∑t=1T[‖t‖2] + L2(1-λ) _t=1^TE [ \|m^t-1 \|^2 ]+ ( L2(1-λ)^3+ L2(1-λ)^2 ) _t=1^TE [ \| ^t \|^2 ] By applying Lemma 2, Lemma 3, and Lemma 4, we have ηK2(1−λ)∑t=1T[‖∇F(Φt−1)‖2]≤([F(z0)]−[F(zT)])+ηKL22(1−λ)∑t=1T1KN∑k,Ci[‖i,k−1+1/2t−i,1/4t‖2] η K2(1-λ) _t=1^TE [ \|∇ F ( ^t-1 ) \|^2 ]≤ (E [F (z^0 ) ]-E [F (z^T ) ] )+ η KL^22(1-λ) _t=1^T 1KN _k,C_iE [ \| θ_i,k-1+1/2^t- θ_i,1/4^t \|^2 ] +(2L2(1−λ)3+L2(1−λ)2)∑t=1T[‖t‖2] + ( 2L2(1-λ)^3+ L2(1-λ)^2 ) _t=1^TE [ \| δ^t \|^2 ] ≤([F(z0)]−[F(zT)]) ≤ (E [F (z^0 ) ]-E [F (z^T ) ] ) +ηK2(1−λ)L2(8ηKL(1−λ)2+4ηKL1−λ+1)∑t=1T1KN∑k,Ci[‖i,k−1+1/2t−i,0+1/4t‖2] + η K2(1-λ)L^2 ( 8η KL(1-λ)^2+ 4η KL1-λ+1 ) _t=1^T 1KN _k,C_iE [ \| θ_i,k-1+1/2^t- θ_i,0+1/4^t \|^2 ] +ηK2(1−λ)(4ηKL(1−λ)2+2ηKL1−λ)∑t=1T(8‖∇F(Φt−1)‖2+4(1−|St|N)|St|σg2+σ2K|St|) + η K2(1-λ) ( 4η KL(1-λ)^2+ 2η KL1-λ ) _t=1^T (8 \|∇ F ( ^t-1 ) \|^2+ 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | ) ≤([F(z0)]−[F(zT)]) ≤ (E [F (z^0 ) ]-E [F (z^T ) ] ) +ηK2(1−λ)L2(8ηKL(1−λ)2+4ηKL1−λ+1)∑t=1T(24η2K2(σg2+[‖∇F(Φt−1)‖2])+24KL2ρ2+6η2Kσ2+2ρ2) + η K2(1-λ)L^2 ( 8η KL(1-λ)^2+ 4η KL1-λ+1 ) _t=1^T (24η^2K^2 ( _g^2+E [ \|∇ F ( ^t-1 ) \|^2 ] )+24KL^2ρ^2+6η^2Kσ^2+2ρ^2 ) +ηK2(1−λ)(4ηKL(1−λ)2+2ηKL1−λ)∑t=1T(8‖∇F(Φt−1)‖2+4(1−|St|N)|St|σg2+σ2K|St|). + η K2(1-λ) ( 4η KL(1-λ)^2+ 2η KL1-λ ) _t=1^T (8 \|∇ F ( ^t-1 ) \|^2+ 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | ). If η≤(1−λ)2128KL, we can rewrite the above inequality as follows If η≤ (1-λ)^2128KL, we can rewrite the above inequality as follows ηK4(1−λ)∑t=1T[‖∇F(Φt−1)‖2]≤([F(z0)]−[F(zT)])+ηK2(1−λ)L2(8ηKL(1−λ)2+4ηKL1−λ+1)∑t=1T(24η2K2σg2+24KL2ρ2+6η2Kσ2+2ρ2)+ηK2(1−λ)(4ηKL(1−λ)2+2ηKL1−λ)∑t=1T(4(1−|St|N)|St|σg2+σ2K|St|)≤([F(z0)]−[F(zT)])+35L2(1−λ)5(ηK4(1−λ)2)3∑t=1T(6σg2+3Kσ2)+35L2(1−λ)(ηK4(1−λ)2)∑t=1T(24KL2ρ2+2ρ2)+8L(1−λ)(ηK4(1−λ)2)2∑t=1T(4(1−|St|N)|St|σg2+σ2K|St|) aligned & η K4(1-λ) _t=1^TE [ \|∇ F ( ^t-1 ) \|^2 ]\\ &≤ (E [F (z^0 ) ]-E [F (z^T ) ] )+ η K2(1-λ)L^2 ( 8η KL(1-λ)^2+ 4η KL1-λ+1 ) _t=1^T (24η^2K^2 _g^2+24KL^2ρ^2+6η^2Kσ^2+2ρ^2 )\\ &+ η K2(1-λ) ( 4η KL(1-λ)^2+ 2η KL1-λ ) _t=1^T ( 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | )\\ ≤& (E [F (z^0 ) ]-E [F (z^T ) ] )+35L^2(1-λ)^5 ( η K4(1-λ)^2 )^3 _t=1^T (6 _g^2+ 3Kσ^2 )\\ &+35L^2(1-λ) ( η K4(1-λ)^2 ) _t=1^T (24KL^2ρ^2+2ρ^2 )\\ &+8L(1-λ) ( η K4(1-λ)^2 )^2 _t=1^T ( 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | ) aligned Let η~=ηK4(1−λ)2 η= η K4(1-λ)^2. By dividing both sides by 1−λ1-λ, we have η~∑t=1T[‖∇F(Φt−1)‖2]≤ η _t=1^TE [ \|∇ F ( ^t-1 ) \|^2 ]≤ ([F(z0)]−[F(zT)])1−λ+35L2(1−λ)4η~3T(6σg2+3Kσ2) (E [F (z^0 ) ]-E [F (z^T ) ] )1-λ+5L^2(1-λ)^4 η^3T (6 _g^2+ 3Kσ^2 ) +8Lη~2T(4(1−|St|N)|St|σg2+σ2K|St|)+35L2(1−λ)η~T(24KL2ρ2+2ρ2) +8L η^2T ( 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | )+5L^2(1-λ) ηT (24KL^2ρ^2+2ρ^2 ) Dividing both side by η~T ηT yields 1T∑t=1T[‖∇F(Φt−1)‖2]≤ 1T _t=1^TE [ \|∇ F ( ^t-1 ) \|^2 ]≤ ([F(0)]−[F(∗)])η~T(1−λ)+35L2(1−λ)4η~2(6σg2+3Kσ2) (E [F ( θ^0 ) ]-E [F ( θ^* ) ] ) ηT(1-λ)+5L^2(1-λ)^4 η^2 (6 _g^2+ 3Kσ^2 ) +8Lη~(4(1−|St|N)|St|σg2+σ2K|St|)+35L2(1−λ)(24KL2ρ2+2ρ2) +8L η ( 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | )+5L^2(1-λ) (24KL^2ρ^2+2ρ^2 ) Now we get the desired rate by applying Lemma 2, which finishes the proof. ∎ Lemma 0. Algorithm 1 satisfies ∑t=1T[‖t‖2]≤1(1−λ)2∑t=1T[‖t‖2] _t=1^TE [ \| m^t \|^2 ]≤ 1(1-λ)^2 _t=1^TE [ \| ^t \|^2 ] Proof. Unrolling the recursion of the momentum t m^t, i.e., t=∑r=1tλt−rr m^t= _r=1^tλ^t-r ^r [‖t‖2]=[‖∑r=1tλt−rr‖2]E [ \| m^t \|^2 ]=E [ \| _r=1^tλ^t-r ^r \|^2 ] Let Γt=∑r=0t−1λr=1−λt1−λ _t= _r=0^t-1λ^r= 1-λ^t1-λ. Since 0≤λ<1,Γt≤11−λ0≤λ<1, _t≤ 11-λ, we have [‖∑r=1tλt−rr‖2] [ \| _r=1^tλ^t-r ^r \|^2 ] =Γt2[‖1Γt∑r=1tλt−rr‖2] = _t^2E [ \| 1 _t _r=1^tλ^t-r ^r \|^2 ] ≤Γt∑r=1tλt−r[‖r‖2] ≤ _t _r=1^tλ^t-rE [ \| ^r \|^2 ] ≤11−λ∑r=1tλt−r[‖r‖2] ≤ 11-λ _r=1^tλ^t-rE [ \| ^r \|^2 ] By summing the above inequality for t∈0,…,T−1t∈\0,…,T-1\, we have ∑t=1T[‖t‖2] _t=1^TE [ \| m^t \|^2 ] ≤∑t=1T11−λ∑r=1tλt−r[‖r‖2] ≤ _t=1^T 11-λ _r=1^tλ^t-rE [ \| ^r \|^2 ] ≤1(1−λ)2∑t=1T[‖t‖2] ≤ 1(1-λ)^2 _t=1^TE [ \| ^t \|^2 ] which finishes the proof. ∎ Lemma 0. For all t≥1t≥ 1, Algorithm 1 satisfies [‖t‖2]≤2(ηK)2(2L2KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]+8‖∇F(Φt−1)‖2+4(1−|St|N)|St|σg2+σ2K|St|)E [ \| ^t \|^2 ]≤ 2(η K)^2 ( 2L^2KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]+8 \|∇ F ( ^t-1 ) \|^2+ 4 (1- |S_t |N ) |S_t | _g^2+ σ^2K |S_t | ) , where i,0t θ_i,0^t denotes the initial point for the local model of the i-th client, i.e., i,0+1/4t=Φt−1 θ_i,0+1/4^t= ^t-1. Proof. By applying Lemma 3, we have [‖t‖2] [ \| ^t \|^2 ] =[‖ηK|St|∑k,Ci∈St∇Fi(i,k+1/2t;ζi)‖2] =E [ \| η KK |S_t | _k,C_i∈ S_t∇ F_i ( θ_i,k+1/2^t; _i ) \|^2 ] ≤2(ηK)2([‖1K|St|∑k,Ci∈St∇Fi(i,k+1/2t)‖2]+σ2K|St|) ≤ 2(η K)^2 (E [ \| 1K |S_t | _k,C_i∈ S_t∇ F_i ( θ_i,k+1/2^t ) \|^2 ]+ σ^2K |S_t | ) We note that [‖1K|St|∑k,Ci∈St∇Fi(i,k+1/2t)‖2] [ \| 1K |S_t | _k,C_i∈ S_t∇ F_i ( θ_i,k+1/2^t ) \|^2 ] =[‖1K|St|∑k,Ci∈St(∇Fi(i,k+1/2t)−∇Fi(i,0+1/4t)+∇Fi(i,0+1/4t))‖2] =E [ \| 1K |S_t | _k,C_i∈ S_t (∇ F_i ( θ_i,k+1/2^t )-∇ F_i ( θ_i,0+1/4^t )+∇ F_i ( θ_i,0+1/4^t ) ) \|^2 ] ≤2[‖1K|St|∑k,Ci∈St(∇Fi(i,k+1/2t)−∇Fi(i,0+1/4t))‖2]+2[‖1|St|∑Ci∈St∇Fi(i,0+1/4t)‖2] ≤ 2E [ \| 1K |S_t | _k,C_i∈ S_t (∇ F_i ( θ_i,k+1/2^t )-∇ F_i ( θ_i,0+1/4^t ) ) \|^2 ]+2E [ \| 1 |S_t | _C_i∈ S_t∇ F_i ( θ_i,0+1/4^t ) \|^2 ] ≤2KN∑k,Ci[‖∇Fi(i,k+1/2t)−∇Fi(i,0+1/4t)‖2]+[‖2|St|∑Ci∈St(∇Fi(i,0+1/4t)−∇F(Φt−1)+∇F(Φt−1))‖2] ≤ 2KN _k,C_iE [ \|∇ F_i ( θ_i,k+1/2^t )-∇ F_i ( θ_i,0+1/4^t ) \|^2 ]+E [ \| 2 |S_t | _C_i∈ S_t (∇ F_i ( θ_i,0+1/4^t )-∇ F ( ^t-1 )+∇ F ( ^t-1 ) ) \|^2 ] ≤2L2KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]+[‖2|St|∑Ci∈St(∇Fi(i,0+1/4t)−∇F(Φt−1)+∇F(Φt−1))‖2] ≤ 2L^2KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]+E [ \| 2 |S_t | _C_i∈ S_t (∇ F_i ( θ_i,0+1/4^t )-∇ F ( ^t-1 )+∇ F ( ^t-1 ) ) \|^2 ] ≤2L2KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]+4‖∇F(Φt−1)‖2+4(1−|St|N)|St|N∑Ci‖∇Fi(i,0+1/4t)‖2 ≤ 2L^2KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]+4 \|∇ F ( ^t-1 ) \|^2+ 4 (1- |S_t |N ) |S_t |N _C_i \|∇ F_i ( θ_i,0+1/4^t ) \|^2 ≤2L2KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]+8‖∇F(Φt−1)‖2+4(1−|St|N)|St|σg2 ≤ 2L^2KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]+8 \|∇ F ( ^t-1 ) \|^2+ 4 (1- |S_t |N ) |S_t | _g^2 where, in the fourth inequality, the improvement of (1−|St|N) (1- |S_t |N ) follows from sampling the active client set StS_t without replacement at the t-th communication round. The last inequality holds because the average norm of local gradients is bounded as 1N∑i=1N‖∇Fi(x)‖2≤σg2+‖∇F(x)‖2 1N _i=1^N \|∇ F_i(x) \|^2≤ _g^2+\|∇ F(x)\|^2, which concludes the proof. ∎ Lemma 0. For all t≥1t≥ 1, we have 1KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]≤24η2K2(σg2+[‖∇F(Φt−1)‖2])+24KL2ρ2+6η2Kσ2+2ρ2 1KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]≤ 24η^2K^2 ( _g^2+E [ \|∇ F ( ^t-1 ) \|^2 ] )+24KL^2ρ^2+6η^2Kσ^2+2ρ^2 Proof. We start with a one-step conversion [‖i,k+1/2t−i,0+1/4t‖2] [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ] =[‖i,kt+λt−ρt‖t‖−i,0+1/4t‖2] =E [ \| θ_i,k^t+λ m^t-ρ m^t \| m^t \|- θ_i,0+1/4^t \|^2 ] ≤[‖i,kt−ρt‖t‖−i,0t‖2] [ \| θ_i,k^t-ρ m^t \| m^t \|- θ_i,0^t \|^2 ] ≤2[‖i,kt−i,0t‖2]+2ρ2 ≤ 2E [ \| θ_i,k^t- θ_i,0^t \|^2 ]+2ρ^2 1KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2] 1KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ] ≤2KN∑k,Ci[[∥i,kt−i,0t∥2]+2ρ2 ≤ 2KN _k,C_i [E [ \| θ_i,k^t- θ_i,0^t \|^2 ]+2ρ^2 . We first define the following terms as (23) Ii,kt=[‖i,kt−i,0t‖2],Iit=1K∑k=1KIi,kt,It=1N∑CiIitI_i,k^t=E [ \| θ_i,k^t- θ_i,0^t \|^2 ],I_i^t= 1K _k=1^KI_i,k^t,I^t= 1N _C_iI_i^t Initially, we commence by deriving an upper bound for the variable Ii,ktI_i,k^t as (24) [‖i,kt−i,0t‖2] [ \| θ_i,k^t- θ_i,0^t \|^2 ] =[‖i,k−1t−η∇Fi(i,k−1+1/2t;ζi)−i,0t‖2] =E [ \| θ_i,k-1^t-η∇ F_i ( θ_i,k-1+1/2^t; _i )- θ_i,0^t \|^2 ] ≤[‖i,k−1t−i,0t−η∇Fi(i,k−1+1/2t)‖2]+η2σ2 [ \| θ_i,k-1^t- θ_i,0^t-η∇ F_i ( θ_i,k-1+1/2^t ) \|^2 ]+η^2σ^2 ≤(1+1K−1)[‖i,k−1t−i,0t‖2]+Kη2[‖∇Fi(i,k−1+1/2t)‖2]+η2σ ≤ (1+ 1K-1 )E [ \| θ_i,k-1^t- θ_i,0^t \|^2 ]+Kη^2E [ \|∇ F_i ( θ_i,k-1+1/2^t ) \|^2 ]+η^2σ where the first inequality follows because the stochastic gradient possesses a bounded variance, while the second inequality follows from the Lemma 1. We note that (25) [‖∇Fi(i,k−1+1/2t)‖2]=[‖∇Fi(i,k−1+1/2t)−∇Fi(i,1/2t)+∇Fi(i,1/2t)‖2] [ \|∇ F_i ( θ_i,k-1+1/2^t ) \|^2 ]=E [ \|∇ F_i ( θ_i,k-1+1/2^t )-∇ F_i ( θ_i,1/2^t )+∇ F_i ( θ_i,1/2^t ) \|^2 ] ≤2[‖∇Fi(i,k−1+1/2t)−∇Fi(i,1/2t)‖2]+2[‖∇Fi(i,1/2t)‖2] ≤ 2E [ \|∇ F_i ( θ_i,k-1+1/2^t )-∇ F_i ( θ_i,1/2^t ) \|^2 ]+2E [ \|∇ F_i ( θ_i,1/2^t ) \|^2 ] ≤2[‖i,k−1+1/2t−i,1/2t‖2]+2[‖∇Fi(i,1/2t)+∇Fi(i,1/4′)−∇Fi(i,1/4t)‖2] ≤ 2E [ \| θ_i,k-1+1/2^t- θ_i,1/2^t \|^2 ]+2E [ \|∇ F_i ( θ_i,1/2^t )+∇ F_i ( θ_i,1/4 )-∇ F_i ( θ_i,1/4^t ) \|^2 ] ≤2L2[‖i,k−1t−i,0t‖2]+4[‖∇Fi(i,1/4t)‖2]+4L2ρ2 ≤ 2L^2E [ \| θ_i,k-1^t- θ_i,0^t \|^2 ]+4E [ \|∇ F_i ( θ_i,1/4^t ) \|^2 ]+4L^2ρ^2 By substituting Eq. (24) into Eq. (25), we have Ii,kt I_i,k^t ≤(1+1K−1+2Kη2L2)[‖i,k−1t−i,0t‖2]+4Kη2[‖∇Fi(i,0t)‖2]+4L2ρ2+η2σ2 ≤ (1+ 1K-1+2Kη^2L^2 )E [ \| θ_i,k-1^t- θ_i,0^t \|^2 ]+4Kη^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+4L^2ρ^2+η^2σ^2 ≤(1+1K−1+2Kη2L2)Ii,k−1t+4Kη2[‖∇Fi(i,0t)‖2]+4L2ρ2+η2σ2 ≤ (1+ 1K-1+2Kη^2L^2 )I_i,k-1^t+4Kη^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+4L^2ρ^2+η^2σ^2 By unrolling the recursion, we have Ii,kt≤∑r=0k−1(4Kη2[‖∇Fi(i,0t)‖2]+4L2ρ2+η2σ2)(1+2K−1)r≤3K(4Kη2[‖∇Fi(i,0t)‖2]+4L2ρ2+η2σ2)I_i,k^t≤ _r=0^k-1 (4Kη^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+4L^2ρ^2+η^2σ^2 ) (1+ 2K-1 )^r≤ 3K (4Kη^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+4L^2ρ^2+η^2σ^2 ) By the definitions in Eq.(23) , we have Iit I_i^t =1K∑k=1KIi,kt≤3K(4Kη2[‖∇Fi(i,0t)‖2]+4L2ρ2+η2σ2) = 1K _k=1^KI_i,k^t≤ 3K (4Kη^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+4L^2ρ^2+η^2σ^2 ) =12η2K2[‖∇Fi(i,0t)‖2]+12KL2ρ2+3η2Kσ2 =2η^2K^2E [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+2KL^2ρ^2+3η^2Kσ^2 It I^t =12η2K21N∑Ci[‖∇Fi(i,0t)‖2]+12KL2ρ2+3η2Kσ2 =2η^2K^2 1N _C_iE [ \|∇ F_i ( θ_i,0^t ) \|^2 ]+2KL^2ρ^2+3η^2Kσ^2 ≤12η2K2(σg2+[‖∇F(Φt−1)‖2])+12KL2ρ2+3η2Kσ2 ≤ 2η^2K^2 ( _g^2+E [ \|∇ F ( ^t-1 ) \|^2 ] )+2KL^2ρ^2+3η^2Kσ^2 where the inequality follows due to the assumption that the average norm of the local gradients is bounded, i.e., 1N∑i=1N‖∇Fi(x)‖2≤σg2+‖∇F(x)‖2 1N _i=1^N \|∇ F_i(x) \|^2≤ _g^2+\|∇ F(x)\|^2, which completes the proof. 1KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]≤2KN∑k,Ci[‖i,k+1/2t−i,0+1/4t‖2]+2ρ2 1KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]≤ 2KN _k,C_iE [ \| θ_i,k+1/2^t- θ_i,0+1/4^t \|^2 ]+2ρ^2 ≤2[12η2K2(σg2+[‖∇F(Φt−1)‖2])+12KL2ρ2+3η2Kσ2]+2ρ2 ≤ 2 [12η^2K^2 ( _g^2+E [ \|∇ F ( ^t-1 ) \|^2 ] )+12KL^2ρ^2+3η^2Kσ^2 ]+2ρ^2 ≤24η2K2(σg2+[‖∇F(Φt−1)‖2])+24KL2ρ2+6η2Kσ2+2ρ2 ≤ 4η^2K^2 ( _g^2+E [ \|∇ F ( ^t-1 ) \|^2 ] )+4KL^2ρ^2+6η^2Kσ^2+2ρ^2 ∎ Lemma 0. (Bounded ℰwE_w of FedSAM). Suppose our functions satisfies Assumptions 1-2. Then, the updates of FedSAM for any learning rate satisfying η≤110KLη≤ 110KL have the drift due to θi,kt−θt _i,k^t-θ^t : ℰw=1N∑i[‖θi,Kt−θt‖2]≤5Kη2(2L2ρ2σ2+6K(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2)+24K3η4L4ρ2E_w= 1N _iE [ \| _i,K^t-θ^t \|^2 ]≤ 5Kη^2 (2L^2ρ^2σ^2+6K (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2 )+24K^3η^4L^4ρ^2 Proof. Recall that the local update on client i is θi,kt=θi,k−1t−ηg~i,k−1t _i,k^t= _i,k-1^t-η g_i,k-1^t. Then, ‖θi,kt−θt‖2=‖θi,k−1t−θt−ηg~i,k−1t‖2 \| _i,k^t-θ^t \|^2=E \| _i,k-1^t-θ^t-η g_i,k-1^t \|^2 ≤(a)‖θi,k−1t−θt−ηl(g~i,k−1−∇fi(θ~i,k−1t)+∇fi(θ~i,k−1t)−∇fi(θ~t)+∇fi(θ~t))−∇f(θ~t)+∇f(θ~t)‖2 (a)≤E \| _i,k-1^t-θ^t- _l ( g_i,k-1-∇ f_i ( θ_i,k-1^t )+∇ f_i ( θ_i,k-1^t )-∇ f_i( θ^t)+∇ f_i( θ^t) )-∇ f( θ^t)+∇ f( θ^t) \|^2 ≤(b)(1+12K−1)‖θi,k−1t−θt‖2+‖η(g~i,k−1t−∇fi(θ~i,k−1t))‖2 (b)≤ (1+ 12K-1 )E \| _i,k-1^t-θ^t \|^2+E \|η ( g_i,k-1^t-∇ f_i ( θ_i,k-1^t ) ) \|^2 +6K‖ηl(∇fi(θ~i,k−1t)−∇fi(θ~t))‖2+6K‖η(∇fi(θ~t)−∇f(θ~t))‖2+6K‖η∇f(θ~t)‖2 +6KE \| _l (∇ f_i ( θ_i,k-1^t )-∇ f_i( θ^t) ) \|^2+6KE \|η (∇ f_i( θ^t)-∇ f( θ^t) ) \|^2+6K \|η∇ f( θ^t) \|^2 ≤ (c) (1+12K−1+2L2η2)‖θi,k−1t−θt‖2+2η2L2ρ2σ2+12Kη2L2‖θi,k−1t−θt‖2 (c) ≤ (1+ 12K-1+2L^2η^2 )E \| _i,k-1^t-θ^t \|^2+2η^2L^2ρ^2σ^2+2Kη^2L^2E \| _i,k-1^t-θ^t \|^2 +12KL2ηl2‖δi,k−1t−δt‖2+6Kη2‖∇fi(θ~t)−∇f(θ~t)‖2+6K‖∇f(θ~t)‖2 +2KL^2 _l^2E \| _i,k-1^t-δ^t \|^2+6Kη^2E \|∇ f_i( θ^t)-∇ f( θ^t) \|^2+6K\|∇ f( θ^t)\|^2 ≤ (d) (1+12K−1+12Kηl2L2+2L2η2)‖θi,k−1t−θt‖2+2η2L2ρ2σl2+12KL2ηl2‖δi,kt−δt‖2 (d) ≤ (1+ 12K-1+12K _l^2L^2+2L^2η^2 )E \| _i,k-1^t-θ^t \|^2+2η^2L^2ρ^2 _l^2+2KL^2 _l^2E \| _i,k^t-δ^t \|^2 +6Kη2(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2, +6Kη^2 (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2, where (a) follows from the fact that g~i,k−1t g_i,k-1^t is an unbiased estimator of ∇fi(θ~i,k−1t)∇ f_i ( θ_i,k-1^t ) and Lemma A.3; (b) is from Lemma A.2; (c) is from Assumption 3 and Lemma A. 2 and (d) is from Lemma A.5. Averaging over the clients i and learning rate satisfies η≤110KLη≤ 110KL, we have 1N∑i∈[N]‖θi,kt−θt‖2≤(1+12K−1+12Kη2L2+2L2ηl2)1N∑i∈[N]‖θi,k−1t−θt‖2+2η2L2ρ2σ2+12KL2η21N∑i∈[N]‖δi,kt−δt‖2+6Kη2(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2≤(a)(1+1K−1)1N∑i∈[N]‖θi,k−1t−θt‖2+η2L2ρ2σ2+12KL2ηl21N∑i∈[N]‖δi,kt−δt‖2+6Kη2(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2≤∑τ=0k−1(1+1K−1)τ[2η2L2ρ2σ2+6Kη2(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2]+12KL2η21N∑i∈[N]‖δi,kt−δt‖2≤ (b) 5Kη2(2L2ρ2σ2+6K(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2)+24K3η4L4ρ2, aligned & 1N _i∈[N]E \| _i,k^t-θ^t \|^2≤ (1+ 12K-1+12Kη^2L^2+2L^2 _l^2 ) 1N _i∈[N]E \| _i,k-1^t-θ^t \|^2\\ & +2η^2L^2ρ^2σ^2+12KL^2η^2 1N _i∈[N]E \| _i,k^t-δ^t \|^2+6Kη^2 (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2\\ & (a)≤ (1+ 1K-1 ) 1N _i∈[N]E \| _i,k-1^t-θ^t \|^2+η^2L^2ρ^2σ^2\\ & +12KL^2 _l^2 1N _i∈[N]E \| _i,k^t-δ^t \|^2+6Kη^2 (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2\\ & ≤ _τ=0^k-1 (1+ 1K-1 )^τ [2η^2L^2ρ^2σ^2+6Kη^2 (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2 ]+12KL^2η^2 1N _i∈[N]E \| _i,k^t-δ^t \|^2\\ & (b) ≤5Kη^2 (2L^2ρ^2σ^2+6K (3 _g^2+6L^2ρ^2 )+6K\|∇ f( θ^t)\|^2 )+24K^3η^4L^4ρ^2, aligned where (a) is due to the fact that η≤110KL and (b) is from Lemma B.1. where (a) is due to the fact that η≤ 110KL and (b) is from Lemma B.1. ∎ Lemma 0. (Bounded ℰδE_δ of FedSAM (Qu et al., 2022)). Suppose our functions satisfy Assumptions 1-2. Then, the updates of FedSAM for any learning rate satisfying η≤14KLη≤ 14KL have the drift due to δi,kt−δt _i,k^t-δ^t : ℰδ=1N∑i[‖δi,kt−δt‖2]≤2K2β2η2ρ2.E_δ= 1N _iE [ \| _i,k^t-δ^t \|^2 ]≤ 2K^2β^2η^2ρ^2. Recall the definitions of δtδ^t and δi,kt _i,k^t as follows: δt=ρ∇F(θt)‖∇F(θt)‖,δi,kt=ρ∇Fi(θi,kt,ξi)‖∇Fi(θi,kt,ξi)‖δ^t=ρ ∇ F(θ^t)\|∇ F(θ^t)\|, _i,k^t=ρ ∇ F_i ( _i,k^t, _i ) \|∇ F_i ( _i,k^t, _i ) \| Lemma 0. For the full client participation scheme, we can bound [‖Δt‖2]E [ \| ^t \|^2 ] as follows: r[‖Δt‖2]≤Kη2L2ρ2Nσ2+η2N2[‖∑i,k∇fi(θ~i,kt)‖2].E_r [ \| ^t \|^2 ]≤ Kη^2L^2ρ^2Nσ^2+ η^2N^2 [ \| _i,k∇ f_i ( θ_i,k^t ) \|^2 ]. Proof. For the full client participation scheme, we have: t[‖Δt‖2] _t [ \| ^t \|^2 ] ≤( a )ηl2N2t[‖∑i,kg~i,kt‖2]=(b)ηl2N2t[‖∑i,k(g~i,kt−∇fi(θ~i,kt))‖2]+ηl2N2t[‖∑i,k∇fi(θ~i,kt)‖2] ( a )≤ _l^2N^2E_t [ \| _i,k g_i,k^t \|^2 ] (b)= _l^2N^2E_t [ \| _i,k ( g_i,k^t-∇ f_i ( θ_i,k^t ) ) \|^2 ]+ _l^2N^2E_t [ \| _i,k∇ f_i ( θ_i,k^t ) \|^2 ] ≤(c)Kη2L2ρ2Nσ2+η2N2[‖∑i,k∇fi(θ~i,kt)‖2], (c)≤ Kη^2L^2ρ^2Nσ^2+ η^2N^2 [ \| _i,k∇ f_i ( θ_i,k^t ) \|^2 ], where (a) is from Lemma A.2; (b) is from Lemma A. 3 and (c) is from Lemma A.4 in (Qu et al., 2022). η2N2‖∑i=1N∑k=1K∇fi(θ~i,kt)‖2 η^2N^2E \| _i=1^N _k=1^K∇ f_i ( θ_i,k^t ) \|^2 =η2‖1N∑i=1N∑k=1K∇fi(θ~i,kt)−1N∑i=1N∑k=1K∇fi(θ~t)+1N∑i=1N∑k=1K∇fi(θ~t)‖2 =η^2E \| 1N _i=1^N _k=1^K∇ f_i ( θ_i,k^t )- 1N _i=1^N _k=1^K∇ f_i ( θ^t )+ 1N _i=1^N _k=1^K∇ f_i ( θ^t ) \|^2 =η2K2‖1NK∑i=1N∑k=1K∇fi(θ~i,kt)−1NK∑i=1N∑k=1K∇fi(θ~t)+1NK∑i=1N∑k=1K∇fi(θ~t)‖2 =η^2K^2E \| 1NK _i=1^N _k=1^K∇ f_i ( θ_i,k^t )- 1NK _i=1^N _k=1^K∇ f_i ( θ^t )+ 1NK _i=1^N _k=1^K∇ f_i ( θ^t ) \|^2 ≤2η2K2∥1NK∑i=1N∑k=1K∇fi(θ~i,kt)−1NK∑i=1N∑k=1K∇fi(θ~t)∥2+2η2K2∥1NK∑i=1N∑k=1K∇fi(θ~t)∥2 ≤ 2η^2K^2E \| 1NK _i=1^N _k=1^K∇ f_i ( θ_i,k^t )- 1NK _i=1^N _k=1^K∇ f_i ( θ^t ) \|^2+2η^2K^2E \| 1NK _i=1^N _k=1^K∇ f_i ( θ^t ) \|\|^2 ≤2η2L2K21NK∑i=1N∑k=1K‖θ~i,kt−θ~t‖2+2η2K2‖∇f(θ~t)‖2 ≤ 2η^2L^2K^2 1NK _i=1^N _k=1^KE \| θ_i,k^t- θ^t \|^2+2η^2K^2E \|∇ f ( θ^t ) \|^2 ≤2η2L2K2[5Kη2(2L2ρ2σ2+6K(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2)+24K3η4L4ρ2]+2η2K2‖∇f(θ~t)‖2 ≤ 2η^2L^2K^2 [5Kη^2 (2L^2ρ^2σ^2+6K (3 _g^2+6L^2ρ^2 )+6K \|∇ f ( θ^t ) \|^2 )+24K^3η^4L^4ρ^2 ]+2η^2K^2E \|∇ f ( θ^t ) \|^2 t[∥Δt∥2]≤Kη2L2ρ2Nσ2+η2N2[∥∑i,k∇fi(θ~i,kt)∥ _t [ \| ^t \|^2 ]≤ Kη^2L^2ρ^2Nσ^2+ η^2N^2 [ \| _i,k∇ f_i ( θ_i,k^t ) \| . ≤Kη2L2ρ2Nσ2+2η2L2K2[5Kη2(2L2ρ2σ2+6K(3σg2+6L2ρ2)+6K‖∇f(θ~t)‖2)+24K3η4L4ρ2] ≤ Kη^2L^2ρ^2Nσ^2+2η^2L^2K^2 [5Kη^2 (2L^2ρ^2σ^2+6K (3 _g^2+6L^2ρ^2 )+6K \|∇ f ( θ^t ) \|^2 )+24K^3η^4L^4ρ^2 ] +2η2K2‖∇f(θ~t)‖2 +2η^2K^2E \|∇ f ( θ^t ) \|^2 ≤Kη2L2ρ2Nσ2+[2η2L2K2(10L2ρ2Kη2σ2+(3σg2+6L2ρ2)30K2η2+60L2K4η4‖∇f(θ~t)‖2)+24K3η4L4ρ2] ≤ Kη^2L^2ρ^2Nσ^2+ [2η^2L^2K^2 (10L^2ρ^2Kη^2σ^2+ (3 _g^2+6L^2ρ^2 )30K^2η^2+60L^2K^4η^4 \|∇ f ( θ^t ) \|^2 )+24K^3η^4L^4ρ^2 ] +2η2K2‖∇f(θ~t)‖2 +2η^2K^2E \|∇ f ( θ^t ) \|^2 ≤Kη2L2ρ2Nσ2+20η4L4K3ρ2σ2+180σg2L2K4η4+360L4K4η4ρ2+60L2K4η4‖∇f(θ~t)‖2 ≤ Kη^2L^2ρ^2Nσ^2+0η^4L^4K^3ρ^2σ^2+80 _g^2L^2K^4η^4+60L^4K^4η^4ρ^2+0L^2K^4η^4 \|∇ f ( θ^t ) \|^2 +48K5η6L6ρ2+2η2K2‖∇f(θ~t)‖2 +8K^5η^6L^6ρ^2+2η^2K^2E \|∇ f ( θ^t ) \|^2 ∎ Lemma 0. For the full client participation scheme, we can bound [‖Δt‖2]E [ \| ^t \|^2 ] as follows: r[1N∑i=1N‖θi,Kt−θt+1‖2].E_r [ 1N _i=1^NE \| _i,K^t-θ^t+1 \|^2 ]. Proof. 1N∑i=1N‖θi,Kt−θt+1‖2=1N∑i=1N‖θi,Kt−1N∑i=1Nθi,Kt‖2 1N _i=1^NE \| _i,K^t-θ^t+1 \|^2= 1N _i=1^N \| _i,K^t- 1N _i=1^N _i,K^t \|^2 ≤2N∑i=1N‖θi,Kt−θt‖2+2N∑i=1N‖θt+1−θt‖2 ≤ 2N _i=1^N \| _i,K^t-θ^t \|^2+ 2N _i=1^N \|θ^t+1-θ^t \|^2 ≤[2Kη2L2ρ2N+40η4L4K3ρ2+20L2ρ2Kη2]σ2+[360L2K4η4+180K2η2]σg2 ≤ [ 2Kη^2L^2ρ^2N+40η^4L^4K^3ρ^2+20L^2ρ^2Kη^2 ]σ^2+ [360L^2K^4η^4+180K^2η^2 ] _g^2 +[720L4K4η4+96K5η6L6+360L2+48K3η4L4]ρ2+[120L2K4η4+4η2K2+60K2η2]‖∇f(θ~t)‖2 + [720L^4K^4η^4+96K^5η^6L^6+360L^2+48K^3η^4L^4 ]ρ^2+ [120L^2K^4η^4+4η^2K^2+60K^2η^2 ] \|∇ f ( θ^t ) \|^2 If we set η=(1TKL),ρ=(1T) used in Theorem 1. If we set η=O ( 1 TKL ),ρ=O ( 1 T ) used in Theorem 1. ≤[2KNT2+40KT3+20KT2]σ2+[360L2T2+180L2T]σg2 ≤ [ 2KNT^2+ 40KT^3+ 20KT^2 ]σ^2+ [ 360L^2T^2+ 180L^2T ] _g^2 +[720T2+96KT3+360L2+48KT2]1T+[1201L2T2+641L2T]‖∇f(θ~′)‖2 + [ 720T^2+ 96KT^3+360L^2+ 48KT^2 ] 1T+ [120 1L^2T^2+64 1L^2T ] \|∇ f ( θ ) \|^2 ∎