Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based 	Stochastic Controlled Weight Averaging
Authors: Junkang Liu
Yuanyuan Liu
Fanhua Shang
Hongying Liu
Jin Liu
Wei Feng
Venue: ICML 2025 poster (2025)
Keywords: Federated Learning，Generalization，Stochastic Weight Averaging
Primary area: optimization->large_scale_parallel_and_distributed

Abstract: For federated learning (FL) algorithms such as FedSAM, their generalization capability is crucial for real-word applications. In this paper, we revisit the generalization problem in FL and investigate the impact of data heterogeneity on FL generalization. We find that FedSAM usually performs worse than FedAvg in the case of highly heterogeneous data, and thus propose a novel and effective federated learning algorithm with Stochastic Weight Averaging (called \texttt{FedSWA}), which aims to find flatter minima in the setting of highly heterogeneous data. Moreover, we introduce a new momentum-based stochastic controlled weight averaging FL algorithm (\texttt{FedMoSWA}), which is designed to better align local and global models. 
 Theoretically, we provide both convergence analysis and generalization bounds for \texttt{FedSWA} and \texttt{FedMoSWA}. We also prove that the optimization and generalization errors of \texttt{FedMoSWA} are smaller than those of their counterparts, including FedSAM and its variants. Empirically, experimental results on CIFAR10/100 and Tiny ImageNet demonstrate the superiority of the proposed algorithms compared to their counterparts.

Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Basee Stochastic Controlled Weight Averaging
Junkang Liu
Yuanyuan Liu"
Fanhua Shang"
Hongying Liu" 34 Jin Liu
Wei Feng
Abstract
For federated learning (FL) algorithms such as FedSAM, their generalization capability is crucial for real-word applications. In this paper, visit the generalization problem in FL and investi- gate the impact of data heterogeneity on FL gener- alization. We find that FedSAM usually performs worse than FedAvg in the case highly hetero- geneous data, and thus propose a novel and effec- tive federated learning algorithm with Stochastic Weight Averaging (called F FedSWA), which aims to find flatter minima in the setting of highly het- erogeneous data. Moreover, we introduce new momentum-bart based stochastic controlled weight averaging FL algorithm (FedMoSWA), which designed to better align local and global mod- els. Theoretically, we provide both convergence analysis and generalization bounds for FedSWA and FedMoSWA. We also prove that the opti- mization and generalization errors of FedMoSWA are smaller than those of their counterparts, in- cluding FedSAM and its variants. Empirically, experimental results on CIFARIO/100 and Tiny ImageNet demonstrate the superiority of the proposed algorithms compared to their counter- parts. Open source code https: // //tttuu.. com/ fimmJuunanggiuu Saeduuu/FFFeaa. FedSWA.
(a) FedAvg (0.1)
(b) FedSAM (0.1)
(c) FedSWA (0.1)
of
6
0 x 0'
0'
Flat and towe Icw hrea
Flat and low loss
(d) FedSAM
(e) FedSWA
(f) FedSWA
Figure 1: The training loss surfaces of ResNet-18 trained by FedAvg (a), FedSAM (b), and our FedSWA (c) CIFAR- 100 with Dirichlet-0.1 (high data heterogeneity). The test accuracies of FedAvg, FedSAM, and FedSWA are 45.8%, 40.1% 1%, and 50.3%, respectively, and FedSAM performs the worst and fails in this case. (d) highlights the sharpness minimizing conflicts due to discrepancies between local and global loss landscapes caused by data heterogeneity. (e) and show that FedSWA obtain global flat minima in the setting of high data heterogencity. The blue dots denote the client model, and the red dot is the server model in (f). The outer layer is the training loss surface, the inner layer is the test loss surface, and the global model generalizes well at a minimum (f). The double arrow denotes the generalization error gap between clients and the server.
1. Introduction
With increasing concerns about data privacy and security, federated learning (FI.) has emerged as potential dis- tributed machine learning paradigm (McMahan et al., 2017).
FL allows for collaborative training of a global model across several clients without exchanging raw data, ensuring pri- vacy while exploiting decentralized data. FL algorithms have considerable promise in various areas, including health- care, finance, and personalized mobile services (Ricke et al., 2020; Antunes et al., 2022; Byrd & Polychroniadou, 2020; Wang et al., 2024; Bian et al., 2024; 2025; Liu et al., 2024; Wu et al., 2025a:b).
'College of Intelligence and Computing, Tianjin Uni- versity, Tianjin, China "School Artificial Intelligence, Xidian University, Xi'an, China "Medical School, Tian- jin University. Tianjin, China "Peng Cheng Lab, Shen- zhen, China SScoool of Cyber Engineering, Xidian Univer- sity, Xi'an China. Correspondence to: Yuanyuan Liu <yyliu@ xidian.edu.en.. edu.en>> Fanhua Shang <fhshang cachang@rju.edu..ooo edd...ddd.c...... Hongying Liu
FL confronts the issue of data heterogeneity because of the non-independent and identically distributed (Non-IID) data across clients (Zhang et al., 2021; Liu et al., 2024).
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
along with different computational capabilities and network connectivity. These characteristics bring complications, par- ticularly regarding the model's capacity to generalize to unseen data, which is critical to the practical applicability of FI. algorithms. The generalization study assesses the model's capacity to generalize from empirical to theoretical optimal risks, providing vital insights into the efficacy of FL algorithms.
data, as shown in Figure I(e.f). In other words, the goal of FedSWA is to find global flat minima, avoiding the draw- backs of FedSAM, which finds local flat minima instead of global flat minima. Moreover, SWA is computationally more efficient than SAM since SWA does not require addi- tional forward and backward propagations to compute the perturbations in SAM. In this paper, we will revisit SAM and SWA on FL to improve the generalization ability on highly heterogeneous data.
The generalization of many deep learning methods has been extensively studied to address the overfitting problem. Find- ing flat minima in model parameters has been demonstrated to improve generalization performance. The two methods, Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) and Sharpness-Awan Aware Minimization (SAM) (Foret et al., 2021), have received significant attention. In distributed learning, particularly FL. data heterogeneity always enables the global model to converge to sharp local minimum (Fan et al., 2024), and generalization of sharp local minimum is poor, as shown in Figure ((a).
In Figure (d), we can observe that FedSAM finds a global flat minima but high loss surface, and the generalization to the test loss surface does not work very well. This is because FedSAM focuses too much on local flat minima, not global flat minima, causing it to overlook superior global minima. The flatness and low loss surface of FedSAM in local client training does not imply that the global model is flat (Lee & Yoon, 2024). Our FedSWA algorithm aims to find flat and low global loss surface, which generalizes to the test set much better than FedSAM, as shown in Figure ((c.e).
To tackle this challenge, several studies have introduced SAM to local client optimization, which aims to enhance generalization by finding a flat minimum through minimiz- ing the local loss after perturbation. Qu et al. (2022) adopted SAM to FL. and introduced FedSAM, and they also pro- posed MoFedSAM. which incorporates local momentum acceleration. Caldarola et al. (2022) applied adaptive SAM optimizer in FL local training and proposed FedASAM. FedGAMMA (Dai et al., 2023) improves FedSAM by using the SCAFFOLD's variance reduction technique (Karim- ireddy et al., 2020). FedSpeed (Sun et al., 2023) utilized the gradient computed at the SAM-based perturbed weight. Fan et al. (2024) proposed 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. However, adequate experimental results show that FedSAM is extremely ineffective in the setting of highly heterogeneous data (see in Figure (b)). Moreover, Lee & Yoon (2024) also pointed out that SAM in local clients can find local flat minimum, not global flat minimum, and FedSAM fails to generalize in the setting of high data heterogeneity. In addition, Kaddour et al. (2022) noted that SAM can be close to sharp directions, which is unfriendly to the federated averaging algorithm.
While FedSWA enhances the global model's generalization compared to AMdSAM, it does not guarantee a consistent flat and low minimum for the local model to agree with the global model in the setting of highly heterogencous data. To seamlessly integrate the smoothness of both lo- cal and global models, we present a more advanced algo- rithm within our framework, called Federated Learning via Momentum- Based Stochastic Controlled Weight Averag- ing (FedMoSWA). The key idea of FedMoSWA is inspired by the momentum- variance reduction (Cutkosky & Orabona, 2020). Intuitively, F FedMoSWA estimates the up- date direction for the server variable m and the client vari- able q. Their difference, (m c2), is used to correct the estimation of the client drift of local updates, maintaining the consistency of local updates with global updates. This strategy successfully overcomes data heterogeneity.
Despite that there considerable research on optimization error and convergence in FL, the analysis of generalization remains significantly challenging. To theoretically character- ize generalization performance, researchers have developed a variety of analysis techniques, including uniform conver- gence methods (Vapnik et al., 1998), operator approximation techniques (Smale & Zhou, 2007), cearmatiio--hhorettie tools (Russo & Zou, 2019), and algorithmic stability anal- ysis (Rogers & Wagner, 1978; Hardt al., 2016). Among these methods, stability analysis has attracted particular attention due to its theoretical guarantees that are indepen- dent of the capacity measurement of the hypothesis function space. It is also widely applicable and sensitive to data distri- bution. This paper proposes an analytical framework based on uniform stability, which is particularly useful for study- ing generalization errors by considering the dependence on specific FL algorithms and data heterogeneity.
To address the above shortcomings of SAM-type algorithms as local optimizers, this paper introduces SWA to improve generalization of FL with highly heterogeneous data. SWA improves the model's generalization capability by averaging different model weights at the final stage of training, thereby locating a model lying in a flat region of the loss landscape. Therefore, we propose a new Stochastic Weight Averaging Federated Learning (F (FedsWA) algorithm, which aims to find flatter minima in the setting of highly heterogeneous
2

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Table I: Comparison of generalization and optimization errors in the non-convex setting. a is the variance of stochastic gradients, og is data heterogeneity, 8 is the number of participating clients, and T is the number of communication rounds. K is the number of local iterations. ((1, y: 0) is T--Lipschitz continous and f-smooth. F* is a minimum value of Problem (1) below, F :- Fo- F*, >> č= 1 + (2 + 1/ 1/KT)* K-1 /T > 1, which means the generalization error of FedSWA is smaller than that of FedSAM. The generalization error of MoFedS is smaller than that of FedSWA. The optimization bounds of FedSAM and MoFedSAM are proved in (Patel et al., 2022). Generalization bounds without client sampling.
Algorithm
Generalization error
Optimization error
FedSAM (Qu et al., 2022)
( mn
(eL 20g + Tx))
3F VTK$
'Ka
LZO2 + T7 + T3/2K
L2
VT3
MoFedSAM (Qu et al., 2022)
L mnz
el [cI + c0g + co))
VTKs + BVKo,"
BLF
T375K + T3/2
1.20²
VRL.3
VTs
FedSWA (ours)
mnj el + 7 (cT + 20g + 2a))
o
B(otVKog) +
F213 2) 1/3
+ OF T
VTKs
T3/3
FedMoSWA (ours)
O mmn
L e e+t & (ëL + og + ēσ))
ovF vTKe (VT+
05 +
8F m T
Contributions: We propose a new FL generalization analy- sis framework and develop two efficient algorithms to im- prove the generalization and optimization errors. Our main contributions can be summarized as follows:
2. Related work
Heterogeneity Issues in FL: In past years, various strate- gies have been proposed to solve the heterogeneity issues in FL. 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 inter- operability by the server broadcasting global model with prospective gradient consistency. SCAFFOLD (Karim- ireddy et al., 2020) uses SAGA-like control variables to mitigate client-side drift, which can be regarded as adopt- ing the idea of variance reduction at the client side. Our F TedMoSWA method based on momentum-basee variance reduction, which is different from SCAFFOLD and other momentum methods.
New FL. algorithm inspired by SWA with better gen- eralization: We revisit the role of SWA and SAM for fed- erated learning and discuss the superiority of SWA with heterogeneous data. To generalize the global model, we first propose a novel algorithm, FedSWA, to improve the generalization of FL. which is better than F.dSAM.
cemmentuum-assed stochastic controlled weight aver- aging algorithm: To help local models find consistent flat minima aligned with the global model, we develop a new FWAMoSWA algorithm. We also propose one momentum- based stochastic controlled method, which is designed to increase FL generalization and overcome data heterogeneity. Our theoretical result shows that FedMoSWA provides theo- retically better upper bounds for generalization error than FedSAM. Experimentally, FedMoSWA. is better than both FedSAM and his variants.
Generalization Analysis of FL: Hu et al. (2023) pro- vided a systematic analysis of the generalization error of FL in the two-level framework, which captures the missed participating gap. Sefidgaran et al. (2022) used tools from rate-distortion theory to establish new upper bounds on the generalization error of statistical distributed learning algorithms. Sun et al. (2024) analyzed the generalization performance of federated learning through algorithm stabil- ity, but did not propose an improved algorithm to overcome data heterogeneity.
Novel FL generalization analysis framework: We pro- pose a new FL. generalization error analysis framework, which can account for the effect of data heterogeneity on generalization error. Our theoretical results show that the generalization bound of our FedSWA is O w +1 +ZL
Uniform Stability Generalization Analysis of Algo- rithms: Uniform stability (Bousquet & Elisseeff, 2002) is classical tool to analyze the generalization error of an al- gorithm. For instance, Hardt al. (2016) and Zhang et al. (2022) analyzed the generalization of stochastic gradient descent (SGD) (Bottou, 2010) via uniform stability. Yuan et al. (2019) investigated stagewise SGD and showed the advantages of the statewise strategy.
čo, + ) which is correlated with data heterogeneity and is superior to that of FMdSAM. We prove that the generaliza- tion error of FedMoSWA is ( mnge e
e +++ + + og + oo)). which is better than those of FedSAM and FedSWA, where 1+ (2 + 1/KT)) K- /T > 1.
. Better practical results: We also evaluate our algorithms for various federated learning tasks. Experimental results indicate that the proposed algorithms significantly improve the generalization and optimization performance.
SWA: The concept of averaging weights can be traced back to early efforts to accelerate the convergence of SGD (Polyak & Juditsky, 1992; Kaddour). SWA is motivated by observation about SGD's behavior when training neural
3

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
networks: although SGD frequently explores regions in the weight space associated with high-performing models, it seldom reaches the central points of this optimal set. By averaging parameter values, SWA guides the solution closer to the centroid of this space (Izmailov et al., 2018).
Algorithm 1 FedSWA
FedMoSWA algorithm.
1: Input: X, p. initial server model Og. number of clients
N, number of communication rounds T, number of local iterations K, local learning rate m..
2: for t 0, ...,, T do
3:
Communicate (0:-1) to selected clients E [s].
3. Proposed Algorithm
4: Communicate (0,-1, m)) to selected clients E [s).
3.1. Problem Setup
5: for i = L....,, s clients in parallel do
FL aims to optimize global model with the collaboration local clients, i.e, minimizing the following population risk:
6:
for = 0,... K local update do
7:
Compute mini-batch gradient gi (0%,))
8:
1% 7 (1 k) + k Pmp
m (B((0) = E, [ (O)))
F(0)-
9:
0i,,+1 f 0 nk (g. (0),k)).
m
0 + 0 n (g. (0%,) c + m).
10:
The function F. represents the loss function on client i. BG,=D,['] [-] denotes the conditional expectation with respect to the sample Ç.. Generally, the unattainability of the local population risk F; (0) forces us to train the model by mini- mizing the following empirical approximation of population risk Fs, (0). For the generalization problem, we consider:
11:
end for
12:
Communicate (0),k) to server.
13:
c +- (i) gi(r)) or (ii) C₁ m + Σx 1 (0 - 0%
14:
Communicate (0 K. ct m) to server, C₁ < ct
1 m
m
15:
end for
Fs (s(() (0)
Fs: 0 mn
El ( (oij. 0iy;
16: m m + 7 Eie() Ac.
(2)
17:
01 = ! Ei= 0 02 0,- = 0.- + (v. - 0 02-1)-
18: end for
where (0) and Fs, (0) indicate the empirical risks for the global model and the local model of the i-th client, re- spectively. I's, (0) 1
n Ei= l (3 Vi.zi 0) represents the empirical risk loss function on client i. Given dataset j S s= { (1.5) } - l.j- S = S U. .....US. U..UU...UUU Um. where the dataset S = ( ).,)) and (za.j. Vi.s) ( is drawn from an unknown distribution D,- Assume (Di,.... Dmm}
P is the coefficient of local learning rate decay, the smaller p. the faster the local learning rate decays. Inspired by the LookAhead (Zhang et al., 2019), the server aggregation em- ploys an exponential moving average (EMA) to aggregate the historical weights instead of average aggregation. The specific steps of the algorithm are described in Algorithm 1. In E3),)), the local learning rate decreases from to pm,. After model aggregation, the learning rate is adjusted to the initial learning rate 17 to restart local training. Smith (2017); Gotmare et al. (2019) observed that restoring the initial large learning rate (learning rate restart) helps to jump out of the worse local minimum and find better local minimum. We experimentally demonstrate that restarting the learning rate is effective for federated learning (see Table 7).
} are independently sampled from 'D accord- ing to P (Hu et al., 2023). m is the number of clints, n is the number of training samples per client. In this paper, we assume that (0) has a minimum value FšS, when 0 0-
3.2. Our FedSWA Algorithm
To enhance the generalization of FL, our FedSWA algorithm uses SWA for local training and weight aggregation. In the client-side update, this paper proposes a local learning rate decay strategy instead of a constant learning rate. In the t-th round of local training.
Compared with FedAvg (McMahan et al., 2017) and Fed- SAM (Qu et al., 2022). FedAvg and FedS adopt con- stant local learning rate and model averaging to aggregate, while our FedSWA uses cyclical learning rate and exponen- tial moving average aggregation which help FedSWA, jump out of worse flat minima and look for better flat minima as discussed in (Smith, 2017: Gotmare et al., 2019). The test accuracy of FedsWA is also better than those of FedAvg and FedSAM (see Tables 2 and 3).
no MR nik = pm Th = m (1-
k K
k
+
Mn,
K
000 0 0 nh: (0.))
(4)
where Di (0%.) 1 EiEwES Ve ( D. 1..)): z
[ST
< 1. Here B is the sampled minibatch at the (t, k) iteration and 77 is the initial local learning rate. is the k-th local iteration and K is the total number of local iterations.

---

Improving Generalization in Federated Learning via Momentum- Based Stochastic Controlled Weight Averaging
the global optimal than FedSWA, Influenced by the control variables, FedMosWA has fewer differences between clients, and minima are flatter compared to FedSWA.
Compared with MoFedSAM (Qu et al., 2022). MoFed- SAM based on FedSAM and presents a momentum- based mechanism to accelerate convergence and improve performance in data heterogeneity. FedMoSWA intro- duces a momentum- based variance reduction mechanism in FedSWA to accelerate convergence in the setting of high data heterogeneity. Theoretically, our FedMoSWA also out- performs MoFedSAM.
(a) FedSWA
(b) SedM
Figure 2: (a) and (b) indicate that due to data heterogeneity, the local and global models of FedsWA differ too much to find a less flatter minima, while our FedMoSWA aligns the local and global models and finds flatter minima.
4. FL Generalization via uniform stability
One of- minimizes the empirical risk Fs(0) - E ( 0) in Eq.(2) In client via a randomized algorithm A,, e.g., SGD, to find an estimated optimum OA,.S, & argming P's. (0). In server, we use FL optimizer A, c.g., FedAvg, to find an estimated optimum 0 A.s ss argming (0). However, this empirical solution @ A.s differs from the desired optimum 0 of the population risk FD (0) (( (1c, y: 0)].
Excess Risk Error Decomposition.
ten
3.3. Our FedMoSWA Algorithm
Inspired by momentum- based variance reduction (Cutkosky & Orabona, 2020), which achieves variance reduction through the use of a variant of the momentum term, we propose new FedMoSWA algorithm. Different from our FedSWA, the local update of FedMoSWA becomes
012-11 0 The (9: (o.)) c + m).
(5)
Op E argming (0) = E E uj (z. y; 0)].
(8)
where e₁ is the client control variable, and m is the server control variable. We provide two options to update c:::
To understand the generalization ability of FL algorithms, we want to know: what performance of the estimated op- timum A can achieve on the test data (z. y) ~ D in FL. To answer this question, we analyze the test error EA,S (Fo (0 (0s..) of A A.S s via investigating the well-known excess risk error Eexo defined
I. g: (0.), or
II. C. - m Ee (0:-1 kk)
Depending the situation, Option I may be more stable than II, but II is less computationally expensive and usually sufficient (we use Option II for our experiments). In the local update in Eq. (5), our proposed FedMoSWA algorithm has global control variable vn, which is update by
Exxc EAS FD (0 EAS [s (0%)
= E FD (0A.)) Fs (0A.)) + E [Fs (0 - Fs (0')),
m m E(st 10mm (m).
gan
Eapt
where 0 argming (0) is the optimum of Fs. The optimization error c = EAS [Fs (fs.s) Fs (0')) de- notes the difference between the exact optimum 0 and the estimated solution O.A.s. and generalization error Egen
Compared with SCAFFOLD (Karimireddy al., 2020). The SCAFFOLD algorithm has global variable, which c+ Eies (ct -c), but the global control variable e has an update delay problem. This is due to the fact that the update of e aggregates the old and the new e; by using the same weights. When the elient partici- pation rate is low, the old c; is too different from c, which instead has a bad effect on the update of c. Our proposed FedMoSWa SWA algorithm has global variables updated with momentum m higher weight to the latest uploaded and lower weight to the old c; instead of giving the same weight to all c; as that SCAFFOLD does.
is updated
EAS FD (0A.)) - Fss (0..s) measures the effects of minimizing empirical risk instead of population risk.
Uniform Stability and Generalization. One popular ap- proach to analyze generalization error Egen of algorithm A is uniform stability (Hardt al., 2016; Zhang al., 2022).
+ + ies (c-- -)), which gives
Lemma 4.1. (Uniform Stability and Generalization Error) (Hardt et al., 2016) We say randomized algorithm A is e-uniformly stable if for all datasets S ~ D and S' D where S and S' differ in at most sample,
Compared with FedSWA. In Figure 2, we can observe that the local model of our FWAMOSWA tends be more close to
sup EA (E (ax, y: 0 A.s) - E (x A.s: ().s))ese Se. (9) (a
5

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
Moreover, if A is e-uniformly stable, then its gen- eralization error Egen which is defined as Egen | [PD (0A.s) - Fs (.s)) satisfying Egen S.
Non-comve:-nnder Assumptions 2-5, assume {(x, y; 0) is 8-smooth. Together with nk BKT è +
(2+ + RT K-1 ) I. the generalization error satisfies:
FL generalization error via uniform stability. Here we aim to use uniform stability to upper bound the FL generalization error. Suppose given n samples S; = (22,1,22,11 ,Zin), where zijj (.n)))is)) is sam- pled from unknown distribution D,, and suppose the generated sample set 5.11 = 1/2) [
FedsWA: Egen S e++
2L
(ēL + 00g + Ō0).
mnn
FedMoSWA: Egen
2L
++
(cL + 0g + 20).
mn3
Theorem 5.1 shows that for strongly convex and non-convex problems, Eopt has three terms: the first bias term has the Lipschitz continous constant L., the second bias term char- acterizes the effect of data heterogeneity og. and the third term reveals the impact of the stochastic gradient noise o.
,.in] which only differs from the set Se with the j-th sample. S = (S,)",, = St.....,S-1,,S3
[z1,1* ,,,511" Fi.411
S.11,..,,, Sm
One usually analyzes the stability of an algorithm by re- placing one sample in S by another sample from DD and we get Stie), Then based on these two sets, one can train the algorithm on S to obtain different solutions & of the function Fs ()). When using St.,, we adopt O' and 0% i,k to denote their corresponding versions O₂ and 0 (0, and 0 are trained on S in Algorithm 1).
For strongly convex problems, Theorem 5.1 shows that Egen of FedMoSWA can be upper bounded by O 2L -e---teeee (--ي) (bI.+ + bo))) which is better than that mng
of FedsWA, O( 21 e- Tatto (bL + bo, + bo)) (bo, >> mnn
ag). Compared to FedSWA, FedMoSWA reduces the effect of data heterogeneity on the generalization error. In addition, we can observe that the generalization ability of the algo- rithm affected by the number of clients m and the amount of data mn the clients. More client participation (m) can effectively enhance FI. generalization. Similarly, to achieve smaller generalization error, clients should use large amount of data (mm). Reducing the number of local iterations K improves generalization capabilities. And greater data het- erogeneity impairs federated learning generalization.
5. Theoretical Analysis
5.1. Generalization Error Analysis
We analyze generalization based on following assumptions:
Assumption 1. The loss function 1(-,z) is pustrrongly convex for any sample z, 1(0;z) 2 =00':=) +
(0';=)) 0 -0) -0'++ # 0- 0² for any z, 0.'''
For the non-convex generalization error, Theorem 5.1 shows that Egen of FedMoSWA is better than FedSWA's
Assumption 2. The loss function ((.,2) 2) is Lipschitz conti-
nous, [((0;)=- (0';z) IS 1= 0') ll. for any 2, 0, 0'.
0 Ol 21 + + čσ, + ēσ)), and čag >> og in Table 1. FedMoSWA and FedSWA generalization error is better than FedSAM in Table 1.
Assumption 3. There exists constant o
such that for any i,i E [m], and 2; ~ De.
E [[ 2)) - V=(( o
Assumption 4. The loss function 1(:,2) is B-smooth, |0 for for any 2,0. 2,0.0' 0'. [VI(0;=)- VI (0':=)) S B
5.2. Optimization Error Analysis
Assumption (data heterogeneity). Given E [m], there exists constant 0g > 0, for any O we have 1 EE (0) - (F(O)) < 0g.
Theorem 5.2 (Optimization Error of F FedMoSWA). For B- smooth functions (Fi), which satisfy Assumptions 6-9, and are the same as in the SCAFFOLD (Karimireddy et al., 2020) algorithm (see the Appendix for details), the output of F FedMoSWA has expected error smaller than e. Strongly convex: The S min (aKa ). T 2 (2_ m then
Assumption 2, as in (Zhang et al., 2022), links model pertur- bation to stability. Assumptions 3 and in our study address data heterogeneity and algorithm convergence impacts. Theorem 5.1 (Generalization Error). Assuming all clients participate in each round with Option I: Strongly convex: Under Assumptions 1-5, suppose loss Ea. y; 0) is u-strongly convex. By setting vi S BKT." b
max
0²
mf D² exp (- - {
8 (1+2²)
+
m 1
WTKKs
1+
K K-1
7. the generalization error satisfies:
Non-convex: mh S KaB
and
T > 1, then
FedSWA: Egen S
2L
e' 1- --tih
(bL + bio, + bo))
OVT
m
+
3F
m $
mn33"
VTKs
T
FedMoSWA: Egen < mn33"
2L
e. catt (bL + 0g + bo))
Here D2 := | 0t
2
+
2m2z2
Σm ==
c - VF; (0")
F == F (00) - F (0").
6

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
(Karimireddy et al., 2020), FedAvgM (Hsu et al., 2019),
SCAPFOLD
SCAFFOLD
FedSAM (Qu et al., 2022), MoFedSAM (Qu al., 2022), FedLESAM (Fan et al., 2024), FedASAM (Caldarola et al., 2022), and FedACG (Kim et al., 2024).
50
Redsane
NoLEESAN
Our Portle
Hyper-parameter Settings: The number of clients is 100, batch size = 50, local epoch E = 5, K = 50 and client selection rate C = 0.1. We set the grid search range client
30-
200
150
2013
1020
18
203
400
640
320
100
Communication Rounds
Communication Rounds
100 ,. x 10 10-"). The learning rate decay per round of communication is 0.998, and the total is T = 100.. Specifically, we set 0-1. a = 1.5, 0.2 for FedMoSWA and FedSWA.
learning rate by 1 {10 -3 3.3 3 x 10
(a) Dirichlet-0.1
(b) Dirichlet-0.6
Figure 3: Convergence plots on CIFARIO0 with ResNet-18.
6.2. Results on Nonconvex Problems
Theorem 5.2 shows that for strongly convex and non-convex problems, optimization error Sopt has two terms: the first term reveals the impact of the stochastic gradient noise, and the second bias term characterizes the effect of initial- ization 0g. The optimization error Eapt is not affected by the data heterogeneous parameter Jg and F FedMoSWA con- verges faster as a increases. Convergence is also accelerated when the number of client s increases. Its convergence is accelerated when the number of local iteration K increases.
Results on Convolutional Neural Network: From Ta- ble 2 and Figure 3, we have the following observations: (i) FedSWA can improve generalization ability by SWA, and it is better than FedSAM. Our FedMoSWA is better than FedSAM and its variant, MoFedSAM, in terms of generalization. (ii) FedMoSWA, further improves the gen- eralization ability by momentum- based stochastic control. Compared to the stochastic control algorithm, SCAFFOLD, our F FedMoSWA demonstrates better generalization ability than SCAFFOLD. Compared to the momentum accelera- tion algorithm FedACG, FedMoSWA shows significantly better generalization combining momentum and variance reduction techniques. Compared to FedSAM, MoFedSAM SCAFFOLD and FedACG can not achieve desirable flat- ness. FedMoSWA. further minimizes the global sharpness and achieves much flatter loss landscape in Figure 4.
For strongly convex problems in Theo- rem 5.2 shows that Eupt of FedMoSWA is O o²
pTKa (1+ + -²)++ + my D2 exp m [={++5]T)
For a, it can be observed that with the increase of a, Eopt becomes smaller. The Eopt of FedMoSWA for non- [VF
convex problems is O VTK:
OF (V++ ( (m)3).
which is not affected by the data heterogeneous pa- rameter 0g and is better than FedSWA FedSAM and MoFedSAM, shown in Table 1. We can prove that the Capt of FedMoSWA is faster than SCAFFOLD's aVF
Impact of heterogeneity: From Table 3, we have the fol- lowing observations: For different data heterogeneities, our FedMoSWA generalizes well. Especially when the degree of data heterogeneity is high (Dirichle-O.1), the final test results of our FedMoSWA are much better than all other algorithms. For example, Non-IID levels increasing, F FedMoSWA achieves higher test accuracy 10.4%, 8.5% and 6.2% and saving communication round than MoFedSA. on the CIFAR-100 dataset.
O VTKS (VI+m) + OF ()) (Karimireddy al., 2020).
6. Experiments
6.1. Experimental Settings
Impact of a, 7: Figure 5 (a) compares the effects of a on FL aggregation. As the a increases, the acceleration effect of our FedMoSWA becomes more obvious. However, when > 1.5, the performance of the algorithm decreases with increasing a. We can find that the final performance of the algorithm is best when = 1.5. Figure 5 (b) com- pares the effects of y in the same settings. When y = 0.05, FedMo SWA converges very slowly, and the final algorithm has low accuracy. As the 7 increases, the acceleration ef- fect of our F FedMoSWA becomes more apparent. The final performance of the algorithm is best when y 0.2.
Datasets: We evaluate our algorithms on CIFAR1O, CI-
FAR100 (?), Tiny ImageNet (Le & Yang, 2015). For Non- IID data setup, we simulate the data heterogeneity by sam- pling the label ratios from Dirichlet distribution (Hsu et al., 2019). Dirichlet-0.1 means very high data heterogeneity and Dirichlet-0.6 implies average data heterogeneity.
Models: To test the robustness of our algorithms, we use various classifiers including LeNet-5 (LeCun et al., 2015), VGG-11 (Simonyan & Zisserman, 2015), ResNet- 18 (He et al., 2016), vision Transformer (ViT-Base) (Dosovitskiy et al., 2021).
Methods: We compare FedSWA and FedMoSWA with many SOTA FL methods, including FedAvg (McMahan et al., 2017), FedDyn (Durmus et al., 2021), SCAFFOLD
Impact of p: Table 7 and Figure 3 compare the effects of in the local training of F FedSWA and MoFed As decreases, the performance of the algorithms increases. We
7

---

Improving Generalization in Federated Learning via Momentum-Basse Based Stochastic Controlled Weight Averaging
Table 2: Comparison of the accuracies at the target rounds (1000/0) and the communication round to reach target test accuracy (Ace.%) %6 of each algorithm on Dirichlet-0.6 the CIFARIO and CIFARI00 datasets.
CIFARIO
CIFARI00
Method
LeNet-5
VGG-11
ResNet-18
LeNet-5
VGG-11
ResNet-18
(%). Rounds
Acc.(%)
Rounds
Acc.(%)
Rounds 85%
Acc.(%)
Rounds
Acc.(%)
Rounds
Acc.(%)
Rounds
1000R
78%
1000R
80%
1000R
1000R
52%
1000R
53%
1000R
55%
FedAvg
79. 79.6
574
84.0
313
86.040.2
542
41 41.2+0.2
1000+
48.9+0.4
1000+
54.2+0.2
1000+
FedDyn
80.8 80.8
390
+0.2
557
84.640.2
1000+
40.8
1000+
45.2+0.6
1000+
46.540.6
1000+
SCAFFOLD
80.7
799
86.9
272
85.940.6
543
50.340.2
1000+
47.940.2
1000+
54.140.2
1000+
FedSAM
79.
821
85.5+0.2
227
83 83.640.2
1000+
47.9+0.1
1000+
51.7+0.3
1000+
47.8+0.1
1000+
MMFedSAM
82.
483
87.4
345
87.010.3
569
49.310.2
1000+
1.040.0
1000+
60.110.2
603
FedLESAM
79.7
496
83.5
169
89.040.1
480
41.3+0.3
1000+
49.2+0.3
1000+
52.1
1000+
FedASAM
81.0
371
86.3
169
88.310.2
542
42.1
1000+
53.4
908
49.8
1000+
FedACG
82.9
321
84.5+0.4
249
90.3+0.4
237
52.5+0.2
729
50.9
1000+
61.7.0.4
518
FedsWA
81.5+0.2
591
86.5+0.1
169
89.54v.1
538
48.5
1000+
52.5+0.4
1000+
59.840.3
574
FedMoSWA
83.810.1
301
88. 1..1
137
91 91.210.11
265
53.810.1
622
61.7 61.710.4
389
67.910.4
330
(a) SCAFFOLD
(b) FedLESAM
(c) FedSAM
(d) MoFedSAM
(e) FedSWA
FedMoSWR
Figure 4: Global training loss surfaces for FedSWA, F FedMoSWA and other baselines on Dirichlet-O.1 I of CIFAR 100 dataset with ResNet-18 18. FWAMoSWA could approach a more general and flat loss landscape which high generalization in FL.
pretrained model from the official website, and conducted many experiments. All the experimental results verify that our FedMoSWA can achieve excellent generalization on both the vision Transformer model and big datasets.
100 accuacc 8
8.
Rictonnl, apiibedL
gemmand..
RecMesW0 1.
fedinediil gemmand02
6.3. Ablation Study
Communication Rounds
Communccation Rounds
We conducted full client participation experiments in Here, FedMo denotes FedMoSWA without SWA (i) but only with Momentum Stochastic Control (ū). FedMo (62.5%) outperforms SCAFFOLD (59.9%) by 2.6%. Even if the client is full participation, our method is still better than SCAFFOLD.
(a) Impact
(b) Impact
Figure 5: Convergence plots for FedMoSWA with different and on CIFAR100 dataset using R8sNet-18.
When Dirichlet-0.6. 10% participation clients, FedMo using only momentum stochastic control achieves 65.9%, higher than SCAFFOLD (54.1%) demonstrating that momentum variance reduction mitigates SCAFFOLD's variance reduc- tion delay. By combining SWA and momentum variance reduction, FedMoSWA achieves 67.9% (+13.7%), showing that both SWA (+5.6%) and momentum control (+11.7%) are effective, and momentum control has greater impact. Moreover, FedSWA is simple algorithm like FedSAM and can be combined with other techniques.
can find that the final performance of the algorithms' is best when p = 0. When o and = 1, our FedSWA is equal FedAvg. which means cyclical learning rate works in FL. When p and 1, SWA is equal to FedAvg with momentum variance reduction without SWA (65.9%) which means momentum variance reduction works in FL and it is better than SCAFFOLD (52.3%).
Results on Vision Transformer: In Table 3, we use the ViT-Base model on the Tiny ImageNet dataset. We used the
8

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Table 3: Comparison of each algorithm on the CIFARIO0 and Tiny ImageNet datasets with different data heterogeneity.
CIFARIO0 18)
Tiny ImageNet (ViT-Base)
Method
Dirichlet-0.1
Dirichlet-0.3
Dirichlet-0.6
Dirichlet-0.1
Dirichlet-0.3
Dirichlet-O.6
(%)
Rounds 55%
(cc.
Rounds
A%) Rounds
Acc.(%)
Rounds
Acc.(%)
Rounds
Acc.(%) Rounds
1000R
1000R
55%
1000R
55%
400R
70%
400R
70%
400R
70%
FedAvg
4..
1000+
52.5
1000+
54.210.2
1000+
10.
258
71.8 10.8
223
72.8
208
FedDyn
4..
1000+
45.9
1000+
46.5+0.2
1000+
67.5 67.5
400+
68. 68.2+0.3
400+
69.34..3
400+
SCAFFOLD
40.3
1000+
50.3
1000+
52.340.2
1000+
71.640.1
202
72.540.1
192
73.1
169
FedSAM
40.1
1000+
49.0
1000+
51.940.5
1000+
71.440.2
212
72.2+0.2
194
72.9+0.4
180
MMFedSAM
51.5.0.2
1000+
57.510.2
770
60.110.1
603
71.610.4
229
72.4
214
72.510.4
209
FedLESAM
48.7
1000+
53.3+0.4
1000+
52.1+0.1
1000+
71.940.3
210
72.1
188
72.540.3
182
FedASAM
47.7+0.3
1000+
46.640.2
1000+
49.840.1
1000+
69.240.3
400+
71.3
234
72.1+0.3
196
FedACG
52.2+0.4
1000+
57.7+0.2
717
61.740.4
518
66.240.2
400+
68.540.1
400+
70.240.3
386
FedsWA
50.3
1000+
55.540.4
889
59.840.3
574
71.940.3
199
72.6+0.2
179
73. 73.2
168
FedMoSWA
61.940.5
577
66.2 66.240.4
468
67.9
330
73.840.3
161
74.4
152
74.
144
Table 4: Testing Acc (%) of FedSWA and McFed with different on CIFAR 100 18esNet-18, Dirichlet-0.6).
Table 6: 10% client participation on CIFAR-100 and ResNet- 18 with 100 clients and 1000 rounds.
P
0.1
0.2
0.3
0.4
0.7
0.8
1.0
Algorithm
Accuracy (%)
Improvement (%)
FedsWA
59.8
58.5
57.9
56.6
55.3
54.6
54.2
FedAvg SCAFFOLD FedSAM
54.2
FedMoSWA
67.9
67.3
67.0
66.8
66.5
66.3
65.9
54.1
-0.1
47.8
-6.4
FedSWA (i)
59.8
+5.6
Table 5: Full client participation on CIFAR-100 and ResNet- 18 with 10 clients and 300 rounds.
FedMo (ii)
65.9
+11.7
FedMoSWA (i+ii)
67.9
+13.7
Algorithm
Accuracy (%)
Improvement (%)
FedAvg
58.2
may consider an approach that significantly eliminates the effect of high heterogeneity on the generalization error.
SCAFFOLD
59,9
+1.7
FedSAM
48.3
-9.9
MoFedSAM
37.9
-20.3
Acknowledgments
FedLESAM
59.2
+1.0
FedSWA
60.2
+2.0
This work was supported in the National Natural Science Foundation of China (Nos. 62276182), Peng Cheng Lab Pro- gram (No. PCL2023A00), Tianjin Natural Science Founda- tion (Nos. 24JCYBJC01230 24JCYBJC01460) and Tian- jin Municipal Education Commission Research Plan (No.
FedMo (ii)
62.5
+4.3
FedMoSWA (i+ii)
63.2
+5.0
20247X008).
7. Conclusion
This study investigated generalization concerns FL and proposed two new algorithms, FedSWA, and FWAMoSWA, for the setting of highly heterogeneous data. In particular, FedMo addresses the essential issue of generalization with highly heterogeneous data. By developing analytical framework based on algorithmic stability, we theoretically established the superiority of FedSWA and FedMoSWA in terms of generalization error. Our FedMoSMA mitigates the effect of data heterogeneity on the generalization error, though it does not completely eliminate it. Our future works
Impact Statement
The proposed Stochastic Weight Averaging (SWA) and Bameentum-Basee Stochastic Controlled Weight Averag- ing (FedMoSWA) algorithms for federated learning aim to address data heterogeneity and improve model generaliza- tion. This paper presents one work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here

---

Improving Generalization in Federated Learning via Comentuu--asser Stochastic Controlled Weight Averaging
References
Fan, Z., Hu, S., Yao, J., Niu, G., Zhang, Y., Sugiyama, M., and Wang. Y. Locally estimated global perturba- tions are better than local perturbations for federated sharpness-aware minimization. In Proceedings of the 4 st International Conference on Machine Learning, pp. 235: 188881 2024.
Antunes, R. S., André da Costa, C., Küderlee A., Yari, I. A., and Eskofier, B. Federated learning for healthcare: Sys- tematic review and architecture proposal. ACM Transac- fions on Intelligent Systems and Technology (TIST), 13 (4):1-23, 2022.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness- aware minimization for efficiently improving generalization. In Proceedings of International Confer- ence on Learning Representations, 2021.
Bian, J., Wang, I.., Zhang, I.., and Xu, J. Lora-fair: Feder- ated Iora fine-tuning with aggregation and initialization refinement. arXiv preprint 2rXiv:2411.1499
2024.
Bian, J., Wang, L., Zhang, L., and Xu, J. Fedalt: Federated fine-tuning through adaptive local training with rest-of- world lora. arXiv preprint arXiv:2503. 2025.
Gotmare, A., Keskar, N. S., Xiong, C., and Socher, R. A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. In Proceedings of Inter- national Conference on Learning Representations, 2019.
Bottou, L. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT 2010: 19th International Conference on Computational Statis- ticsParis France, August 22-27, 2010 Keynote, Invited and Contributed Papers, pp. 177-186. Springer, 2010.
Hardt, M., Recht, B., and Singer, Y. Train faster, generalize better: Stability of stochastic gradient descent. In Interna- tional Conference on Machine Learning, pp. 1225-1234 PMLR, 2016
Bousquet, O. and Elisseeff, A. Stability and generalization. Journal of Machine Learning Research, 2:499-526, 2002.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn ing for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Bubeck, S. et al. Convex optimization: Algorithms and com- plexity. Foundations and Trends@ in Machine Learning,
8(3-4):231-357 2015.
Hsu, T.-M. H.,, Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual clas- sification. arXiv preprint urXiv:/909. 003900633,, 2019.
Byrd, D. and Polychroniadou, A. Differentially private secure multi -artt computation for federated learning in financial applications. In Proceedings of the First ACM International Conference on Al in Finance, pp. 1-9, 2020.
Hu, X., Li, S., and Liu, Y. Generalization bounds for fed- erated learning: Fast rates, unparticipating clients and unbounded losses. In International Conference on Learn- ing Representations, 2023.
Caldarola, D., Caputo, B., and Ciccone, M. Improving generalization in federated learning by seeking flat min- ima. In European Conference on Computer Vision, pp. 654-672. Springer, 2022.
Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A. G. Averaging weights leads to wider optima and better generalization. In Proceedings of Conference on Uncertainty in Artificial Intelligence, 2018.
Cutkosky, A. and Orabona, F. Momentum-baase based vari- ance reduction in non-convex sgd. In Advances in Neural Information Processing Systems 32, (NeurIPS 2019). Vancouver(CA).8-1. December 2019, 2020.
Kaddour, J. Stop wasting my time! saving days of imagenet and bert training with latest weight averaging. In Has it Trained Yet? NeurlPS 2022 Workshop.
Dai, R., Yang, X., Sun, Y., Shen, L., Tian, X., Wang, M., and Zhang, Y. Fedgamma: Federated learning with global sharpness-awarr minimization. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Kaddour, J., Liu, L., Silva, R., and Kusner, M. J. When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35:16577-16675 2022.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M.. Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at seale. Proceedings of International Conference on Learning Representations, 2021
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 on Machine Learning, pp. 5132-5143 PMLR, 2020.
Durmus, A. E., Yue, Z., Ramon, M., Matthew, M., Paul, W., and Venkatesh, S. Federated learning based on dynamic regularization. In International conference on learning representations, 2021.
Kim, G., Kim, J., and Han, B. Communication-effce federated learning with accelerated client gradient. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12385-12394, 2024.
10

---

Improving Generalization in Federated Learning via Momentum-Bass Stochastic Controlled Weight Averaging
Le, Y. and Yang, X. Tiny imagenet visual recognition chal- lenge. Technical Report, 2015.
IEEE Transactions on Information Theory, 66(1):302- 323, 2019.
LeCun, Y. et al. Lenet-5, convolutional neural networks. URL: http://yann. lecun. com/exdbitener, 14, 2015.
Sefidgaran, M., Chor, R., and Zaidi, A. Rate-distortion theoretic bounds on generalization error for distributed learning. Advances in Neural Information Processing Systems, 15:19777-19702, 2022.
Lee, T. and Yoon, S. W. Rethinking the flat minima search- ing in federated learning. In Forty-first International Con- ference on Machine Learning, pp. 27037-27071, 2024.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. In 3rd Inter- national Conference Learning Representations (ICLR 2015). Computational and Biological Learning Society, 2015.
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.
Smale, S. and Zhou, D.-X. Learning theory estimates via integral operators and their approximations. Constructive Approximation, 26(2):153- 2007.
Liu, J., Shang, F., Liu, Y., Liu, H., Li, Y., and Gong, Y. Fedhegd: Communication-effin
accelerated block co- ordinate gradient descent for federated learning. In Pro- ceedings of the 32nd ACM International Conference Multimedia, pp. 2955-2963, 2024.
Smith, N. Cyclical learning rates for training neural net- works. In 2017 IEEE Winter Conference on Applications of Computer Vision (WACV), pp. 464-472. IEEE, 2017.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. Communication-efien -fficcient learning of deep networks from decentralized data. In Proceedings of International Conference Artificial Intelligence and Statistics, pp. 1273-1282. PMLR, 2017.
Sun, Y., Shen, L., Huang, T., Ding, L., and Tao, D. Fed- speed: Larger local interval, less communication round, and higher generalization accuracy. In Proceedings of International Conference on Learning Representations, 2023.
Patel, K. K., Wang, L., Woodworth, B. E., Bullins, B., and Srebro, N. Towards optimal communication complexity in distributed non-convex optimization. Advances Neu- ral Information Processing Systems, 35:13316-13328, 2022.
Sun, Z., NIu, X., and Wei, E. Understanding generalization of federated learning via stability: Heterogeneity matters. International Conference on Artificial Intelligence and Statistics, pp. 676-684. PMLR, 2024.
Polyak, B. T. and Juditsky, A. B. Accelcration stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838-885 1992.
Vapnik, V.N., Vapnik, V. al. Statistical learning theory. 1998.
Wang. L., Bian, J., Zhang, L.., Chen, C., and Xu, J. Tam- ing cross-domain representation variance in federated prototype learning with heterogeneous data domains. Ad- vances Neural Information Processing Systems, 37: 88348-88372, 2024.
Qu, Z. Li, X., Duan, R., Liu, Y., Tang. B., and Lu, Z. Gen- eralized federated learning via sharpness aware minimiza- tion. In International Conference on Machine Learning, pp. 18250-18888 182800 PMLR, 2022.
Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., Kumnr, S., and McMahan. H. B. Adaptive federated optimization. In Proceedings of International Conference on Learning Representations, 2021.
Wu, W., Chen, Z.,, Qiu, X., Song, S., Huang. X., Ma, F. and Xiao, J. Lim-enhanced multimodal fusion for cross-domain sequential recommendation. arXiv preprint N0Vi::2506.79666 2025a.
Rieke, N., Hancox, J., Li, W., Milletari, F., Roth, H. R., Albarqouni, S., Bakas, S., Galtier, M. N., Landman, B. A.,
Wu, W., Qiu, X., Song, S., Chen, Z., Huang, X., Ma, F., and Xiao, J. Prompt categories cluster for weakly supervised semantic segmentation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 3198 3207, 2025b.
Maier-Hein, K., et al. The future of digital health with federated learning. NPJ Digital Medicine, 3(1):119, 2020.
Rogers, W. H. and Wagner, T. J. finite sample distribution- free performance bound for local discrimination rules. The Annals of Statistics, pp. 506-514, 1978.
Yuan, Z., Yan, Y., Jin, R., and Yang, T. Stagewise training accelerates convergence of testing error over sgd. Ad- vances in Neural Information Processing Systems, 32, 2019.
Russo, D. and Zou, J. How much does your data explo- ration overfit? controlling bias via information usage.

---

Improving Generalization in Federated Learning via Momentum-Basee Based Stochastic Controlled Weight Averaging
Zhang, M., Lucas, J., Ba, J., and Hinton, G. E. Lookahead optimizer: k steps forward, step back. Advances in Neural Information Processing Systems, 32, 2019.
Zhang, X., Hong, M., Dhople, S., Yin, W. and Liu, Y. Fedpd: A federated learning framework with adaptivity to non-iid data. JEEE Transactions Signal Processing, 69:6055-6070, 2021.
Zhang. Y., Zhang, W., Bald, S., Pingali, V., Chen, C., and Goswami, M. Stability of sgd: Tightness analysis and improved bounds. In Proceedings of Conference on Un- certainty in Artificcia Intelligence, pp. 2064-264425502 2373. PMLR, 2022.
12

---

Improving Generalization in Federated Learning via Momentum-Baan Based Stochastic Controlled Weight Averaging
A. Background
In this section, we provide a concise overview of Sharpness-Awar Minimization (SAM) and Stochastic Weight Averaging (SWA).
A.1. SAM: Overview
SAM seeks to identify a solution that lies within a region characterized by uniformly low training loss (D()), specifically within a flat minimum. The sharpness of the training loss function quantified as:
max (e) Se
Here, is hyper-parameter parameter that determines the size of the neighborhood, and p [o). SAM focuses reducing the sharpness of the loss by addressing the following minmax objective:
min max LD + SER Sp
where À is a hyper-parameter that balances the regularization term. According to (Foret et al., 2021), the optimal choice for p is typically Therefore, we use the fz-norm for the maximization over and simplify the expression by omitting the regularization term. To solve the inner maximization problem e" - arg maxiellas L(O + c), the authors suggest using first-order approximation of L( around 0:
voLp(0)
~ arg max LD(0) + 2 VoLD(O) ( (ooD = E(0)
||rlaaS
This approximation, which is computationally efficient, results in E(() being essentially scaled gradient of the current parameters 0. The sharpness-aware gradient then defined as VoLD (0) (ot(0)) and used to update the model as follows:
Oc+1 0. - LD (0:)) 0.)
where y is the learning rate and é (0.). This two-step process repeated iteratively to solve the minmax objective. Essentially, SAM first performs gradient ascent step to estimate the point (0, + εx) where the loss approximately maximized, followed by gradient descent step at 0, using the computed gradient.
A.2. Stochastic Weight Averaging: Overview
SWA (Stochastic Weight Averaging) operates by averaging the weights proposed by Stochastic Gradient Descent (SGD) while employing a learning rate schedule to explore regions of the weight space that correspond to high-performing networks. At each step i within a cycle of length the learning rate is gradually reduced from initial value final value 72. This learning rate adjustment is defined as:
7(0) (1-t(i) t(1)2. t(i) (mod(i 1,
Here, ((i) represents linear interpolation factor that depends the current step i and the cycle length If the learning rate remains constant ? throughout the training process. However, for 1, the learning rate follows cyclical schedule, oscillating between 71 and 72 over the course of the cycle.
Starting from pre-trained model fg. SWA captures the model weights 0 at the end of each cycle and computes their running average. This averaging process is expressed as:
OSWA Pemodels + 0 M mmodels +
OSwa
13

---

Improving Generalization in Federated Learning via Homentum-Basee Stochastic Controlled Weight Averaging
In this equation, Nsmodels tracks the number of completed cycles, and OswA represents the averaged weights. The final model, denoted as foewa is obtained by aggregating the weights from multiple cycles, effectively smoothing the optimization trajectory and enhancing generalization performance.
By iteratively averaging the weights, SWA encourages the model to converge to flatter region of the loss landscape, which often associated with better generalization. This approach leverages the exploration of diverse weight configurations during training, ultimately leading to more robust and stable model.
B. Implementation of the Experiments
B.1. Datasets
Table A summary of CIFARIO/100, Tiny ImageNet datasets, including number total images, number of classes, and the size the images in the datasets.
Datasets
Total Images
Class
Image Size
CIFAR10
60,000
10
3 x 32 x 32
CIFARI 100
60,000
100
x x 32
Tiny ImageNet
100,000
200
3 x 224 x 224
CIFAR- Tiny ImageNett are popular benchmark datasets in federated learning. Data samples in CIFARIO and CIFARI00 colorful images of different categories with the resolution of 32 x 32. There are 10 classes and each class has 6,000 images in CIFAR10. For CIFAR 100, there are 100 classes and each class has 600 images. For Tiny ImageNet, there are 200 classes and each class has 500 images. As shown in Table 6, we summarize CIFARIO, CIFAR100, and Tiny ImageNet from the views of number total images, number of classes, and the size of the images in the datasets.
B.2. More Loss Surface Visualization
Here we show more global loss surface visualizations. As shown in Figure 7-11, we conduct experiments using LeNet-5 and ResNet-18 for CIFAR 10 and CIFAR100, and visualize FedAvg, FedSAM, and FedSWA, Among these, our FedSWA achieves a flatter loss landscape.
3000
5,000
4000
4226
357
3511
3572
2857
857
1.429
0774
0714
0.000
FedAvg
(b) FedSAM
(c) FedSWA
Figure 6: The training loss surfaces of FedAvg, FedSAM, and our FedSWA on CIFAR-100 with ResNet-18 (Dirichlet-O.J)
14

---

Improving Generalization in Federated Learning via Momentum Hooentum-Basedd Stochastic Controlled Weight Averaging.
5,000
4.175
4.75
3150
1750
3327
3.500
1500
1.75
1415
1.20
1259
0925
0.925
0.25
0.000
(a) FedAvg
(b) FedSAM
(c) FedSWA
Figure The training loss surfaces of FedAvg, FedSAM, and our FedSWA on CIFAR-100 with LeNet-5 (Dirichlet-O.I)
4.35
4375
3.750
2,300
2.500
2.500
1.75
210
1.250
9.25
0625
0.63
9.000
(a) FedAvg
(b) FedSAM
(c) FedSWA
Figure The training loss surfaces of FedAvg, FedSAM, and our FedSWA on CIFAR-10 10 with LeNet-5 (Dirichlet-O.1)
3000
5.000
1206
4.35
1371
2.200
2,500
1.109
250
1259
0/74
9625
(a) FedSAM (0.1)
(b) FedSAM (0.3)
(c) FedSAM (0.6)
Figure 9: The training loss surfaces of FedSAM CIFAR-100 with ResNet-18 (Dirichlet-O.. Dirichlet-O.3, Dirichlet-0.6).
15

---

Improving Generalization in Federated Learning via Momentum Saneentu-BBased Stochastic Controlled Weight Averaging
1.00
4,20
4.775
3.75
3750
3125
3125
2 207
2 143
.875
1.875
0724
0000
0000
0.020
(a) FedSWA (0.1)
(b) FedSWA (0.3)
(c) FedSWA (0.6)
Figure 10: The training loss surfaces of FedSWA on CIFAR-100 with ResNet-18 (Dirichlet-O.I) Dirichlet-O.33 Dirichlet-O.6).
B.3. Convergence Behavior Curves
90
reddug
FodDyn
FedDyn
85
SCAFFUL
SCAFOLD
80
PeITAAM
Teal, Accuracy
PedSAM
Test Accuracy
RedAvg
FedDyn
70
PedASAM
PedLESAR
PedLESAM
PedaGaM
70
RedSAM
76
TedAZG
FedACG
65
McFcdSAM
FedLESAM
Our TedMoSWA
FesMoSWA
60
RedACO
55
hech/A
200
400
600
8do
1000
200
400
600
800
1000
50
200
400
800
800
1000
Communication Rounds
Communication Rounds
Communication Rounds
(a) Lenet-5, CIFARIO
(b) Lenet-5, CIFARI00
(c) VGG-11, CIFARIO
95
70
FedAug
Feciugg
FedDyn
SCAFFOLD
90
60
50
MoFedSAM
Test Accuracy
MOFNEAAM
FecSAM
Accuracy
edAvp
40
RodLESAM
ReCLESAM
redASAM
TedSAM
RocACS
MoFedSAM
Our FedEWA
Our redkesswA
FedMoSWA
RedASAM
20
20
TedACG
FodSINA
20C
Communication Rounds
400
600
800
1000
65
700
Communication Rounds
400
600
800
1000
10
700
Communicattonn Rounds
400
600
800
1000
(d) VGG-11, CIFARI00
(e) Resnet--e 18, CIFAR10
(i) ResNet-18, CIFAR100
Figure 11: Convergence plots for FedSWA, FedMoSWA and other baselines same settings with Dirichlet-0.6 on CIFAR10 and CIFAR 100 using different neural network architectures.
16

---

Improving Generalization in Federated Learning via Momentum-Basse Based Stochastic Controlled Weight Averaging
Table 8: CIFAR-100 with ResNet-18 under different data heterogeneity levels (0.1, 0.3, 0.6),
heterogeneity
0.1
0.3
0.6
FedSMOO
46.5
47.8
49.2
FedGAMMA
48.4
51.8
52.6
FedSWA (ours)
50.3
55.5
59.8
FedMoSWA (ours)
61.9
66.2
67.9
70
70
70
FedAigg
Foddig
FedDyn
60
Accuracy
TecDyn
Accuraty
McFadSAM
RodSAM
Test Accuracy
MoFaCEAM
FodSAM
SCAFROLD
FEdLESAM
ESAN
FeeSAM
MoFodSAM
PedALai
hedhcti
Pentiva
Our Fadiys
RocASAM
FadMoSWA
FadMos44A
RacACG
20-
20
20-
FooMoSWA
10
20c
460
600
800
1000
10
200
40c
600
800
1000
10
200
400
800
800
1000
Communication Rounds
Communication Rounds
(a) ResNet-18, Dirichlet-0.6
(b) ResNet-18, 18, Dirichlct-0.3
ResNet- 18, Dirichlet-O.1
Figure 12: Convergence plots for FedSWA, FedMoSWA and baselines with Dirichlet-0.6, 0.3, 0.1 CIFAR 100 using ResNet-18.
C. Basic Assumptions
We analyze generalization based on following assumptions:
Assumption C.1. The loss function U(,,2) 2) is p-strongly convex for any sample z, 1(0; 2 (0';))) (V 0')+ |0 0' 12 for any 2, 0 6'.
Assumption C.2. The loss function I(,, z) is L. Lipschitz continous, 11(0;)) note that under Assumption 2, A,S. (0r; 2) - 1 (0/; =)1 S LE ||0r - 07 B/F (8.) <
S 0 0' for any 2, 0.0 We
Assumption C.3. There exists a constant o > 0 such that for any 0, E [m], and "; ~ D, E
VI (0:
VF((0))
0².
Assumption C.4. The loss function [(, 2) is 3-smooth, v((0)2) - Vi (0':2) < 3 0 - 0' for any 2, 0.
Assumption C.5. (data heterogeneity). Given i E [m]) there exists constant og > 0, for any we have
[[VF((0) - VV I Tg.in + 2Th [VF(()) - VF((O)| < 1 Dm1 g. 0g
D. Main Lemmas
Lemma D.I. Assume that is B-smooth. Then, the following properties hold.
1. G is (1 expaxnnas, G(v) (v) =
2. Assume in addition that f is p-strongly convex. Then. for 1 S B4r G is (1-
n&n expansive.
3+1
Henceforth we will no longer mention which random selection rule we use as the proofs are almost identical for both rules.
Proof. 1. Let = Gf.-- By triangle inequality and our smoothness assumption,
[[G(u) - G(w)]l
||
< || w| + V)(w) - vf(l))l
I w|
+ - w|l
(10)
17

---

Improving Generalization in Federated Learning via Momentum-Basee Based Stochastic Controlled Weight Averaging
2. First, note that if f'is js strongly convex, then p(w) -((wr)
_ [|l0|| [ is convex with (3 :4)-smothh. Hence,
Bp
(Vf(v)) Vf(w), w) 2
||0 wil? + B+ [Vf(() - vf(w)]l
(11)
B + p
Using this inequality gives
[[Grn(u) - G)n(w)]1 = ||u w/12 (n(V((() Vf(w),, 2 ww) + + f(())]
x (1 3 +1
1/3pu ) |10 w|l2 - (
B+ 31 [VS(v) - VS(w)]1²
(12)
With our assumption that S 2 B4ye
this implics
00:511111 21- n33
1/2
(1- nāpu 1|v w|l.
(G) Gym(ll)
||0 w| -
(13)
3 + 11
B + /1)
The lemma follows by applying the inequality v1- - x <
x/2 which holds for æ E [0. 1].
E. Theoretical Results of FedSWA
E.1. Generalization Analysis for FedSWA under strongly convex setting
Lemma E.1. Suppose Assumptions C.I-C.5 hold. Then for FedSWA with vi S 1/3KT, 1,
E 10 0ell s (1 b) K- i (B VF (0:1 + o.s + o), Ak- = -,...,K K,
where ij Σ K-1 =k=0 Th
Proof. Considering local update FedSWA
E 20000+111
0.11
= E
10
Thgg (0%,,) 0ell
1R
0
0₂
%h (gi (0),k) 9: (0.)) w 11: (0:)
WV
BH
E Bl00 x - Bell + v E 11 (0a)ll
B ++
(14)
nhSp
E B110 - 0. + v (E I1: (0.) - B
B B+1 1
E |10 0.| + V LE IV (0.) + +
B + 11
where (a) follows Lemma D.1; (b) follows Assumption C.3. Unrolling the above and noting 0., 6s, with b = (1- nxxBu
for some > 0 yields
E 0.11
110
M (E
(EI +o) (1 (1-6)*-11 -1-
k-1
(15)
< Σ
n/ (E
|| V. )| + + (1 -b)
1=0
< (1 1)*k -1
7/ (E VF (0.) + g.i + o0,
where the last inequality follows Assumption C.5, i/
Σ Σ nfc
18

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Lemma E.2. Given Assumptions C.I-C.4 and considering of FedSWA, for v/ 1/3KT we have
E 99 (0 x) - (11 (E VF (6:) + 00 +0),
where 9:(-) is the sampled gradient of client i, =xx=0 The
Proof. Using Assumptions C.3, C.4, C.5, we obtain
E Bl (0% < E as (0%,,) VE [ (0%x))
VF; (0%,))1 +o
s E VF; (6.) +E VF, (01,)
VF, (0:)] +o
(16)
B VF (0:) + + VF (0.) (F.) + (3T | - 0z +
(1 + (1 (1-b) b)) (E VF (0+) + ggt +))
is following Assumption C.3, is is following Assumption C.4, is is following Assumption C.5 and D.1.
Theorem E.3. (FedSWA). Suppose Assumptions 1-5 hold and consider FedSWA. Let {0:) and (0/ be two datasets S and s respectively. Suppose Bu 8. Then,
2Lb 1 mn 3
(गाा (E VVF (0.)1I + 0g +
gen
K-1 where if. = -K- mhe 1 K-
Proof. For client i, there are two cases to consider. the first case, SGD selects the index an sample at local step on which is identical in S and (() In this sense, we have
|0
0% 0 s (1
m/ |0 0%/1 3 + 12)
(17)
And this case happens with probability - 1/n (since only one sample perturbed for client i.).
In the second case, SGD encounters the perturbed sample local time step which happens with probability 1/n. We denote the gradient of this perturbed sample g((-). Then,
00 01/41||
00
0% n (9: (0i,)) l |I
$
8%
n (9. (0%,,)
9s (0%) m/ gi (0f,)) g (0%%
(18)
Bp
10:. 02:6
+ nt 1lac (0t..) le (0%h)]1
B +
Combining these two cases we have for client i,
E 00
02.
m 3p
E
00 - 07,1
+
T E
ga (0%..) - H. (0%)
3 3+1
n
(19)
ni Bp
2n +
<
E
100 - 0//6
E
g9 (0%,)) 1l.
B +
n
where the last inequality follows that g-) and g-(-) are sampled from the same distribution.
Iterating the above over t and noting f = 8/. we conclude the proof. We let 8 be an upper bound of 1 + (1 -1,7
since i bounded above. Then unrolling 19, and follow E.2,
19

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
B 0r,k - 0/
K-1
(1 1
2 K-1
EE 110 0ill +
Σ
1
(E (EIIFF (0.) 0g.t +0)
(20)
B 11
B + p1
k==
1=k+1
2 -biee 3TP (E VF (0t) +gg o).
5 pie
E 1 0/1l
n
is following
1- x S With 0++1
2i=1 0 k and 0,41
1 m
Ei=1 0 we have
m
0/42 || E 0
-
E
J.K - 07'KK
i=1
m
(21)
Se
10
+
2
bipe
Bre B+K
(EI (0+) +0g
mn
Further, unrolling the above over and noting 0g
fg we obtain
2b
T
Raff 8+*
E 0 0/ S
exp
B
The
(E VE (0:)11 00)
1=0
|=++1
(22)
25 1
1 cappy
(E IVF (0:) 0g +0)
mi B
With Assumption C.2,
2LP 1
Ta-py (E D (0.) + +o)
gen
mn 1
K-1
T
>
(B+
When the diminishing stepsizes v/ & 1/3KT are chosen the statement of the theorem, we conclude the proof.
E.2.
Generalization Analysis for FedS WA under non-convex setting
Lemma E.4. Suppose Assumptions C.2-C.5 hold. Then for FedSWA with nh S /BKT,
E
El
0.| S (1 + K - ij, (E VF (0.) ggt + o), Vk. =k=1.... 1,..., K,
where ij = => k= The
Proof. Considering local update of FedSWA
E ||0 0,1
E ||05,r n/.9. (0 0ell
< E 0 02 - The (gi (0%..) 9i (0.)) + LE 19. (0.)
= (1 + 3rk) E =110f, - Ball + E ll (0:)ll
(23)
< (1 + 9nk) E |10% 0cll + n (E l1 (0)) VE; (0,) VF (8)I))
<² (1 Brk) |105, 0.ll + + VE,; (0.)lI +o),
where is following D.I, and Assumptions C.3 in With (1+ 3mt) c, Unrolling the above and noting 0,o O₂
yields
E 1
0.11 5
" (B (0.) + (1 o1)
(24)
n/ (BIVF (8)) + +0) (1 -c)
1=0
(1+0) (E V (0. +.g +0),
20

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
where the last inequality follows Assumption C.5, = =x nf-
Lemma E.5, Given Assumptions C.2-C.5, for Th - 5/BKT, we have
E g|g (0%,x) S (1 c1+ 1 (EIVE (8))l + 0ggt +
where g:()) is the sampled gradient of client i, i, =/== Ex=r nh-
Proof. we obtain
E 91:
- ( (0%,))
S E 96 VF,
VE (01,.)
E VF, (0.) +B (0 VT. (0:) +o
(25)
B ( (0)) - VF (6:))l 3E 3 8: to
(I + (1 ( (ā.) (E VF (0:)ll + 00
is following Assumption C.3, b is is following Assumption C.4, is is following Assumption C.5 and D.1 and E.4. Theorem E.6. (FedSWA). Suppose Assumptions C.2-C.5 hold and consider FedSWA. Let k: K, Vi [m] and ni 3KT Then,
2L 1 cl+ (EI (0.) t +0),
Egen
mn B
Proof. For client i, there are cases to consider. In the first case, SGD selects non-perturbed samples in S and S(o,, which happens with probability 1 1/n. Then, with D.I, we have
| 0 (1+ BTk) 0/ 07
(26)
In the second case, SGD encounters the perturbed sample at time step kt which happens with probability 1/n. Then, have
|00.011 02
00 0% nk 9. (0f,))
si (0%)
s 11 0 nh 9. (ft,.)
9i (0%)
+ m/ |0: (0%%)
o (0%+
(27)
< (1 + Bik) 102. -2.6
+ ní 10s (01)
gl (oth)
Combining these two cases for client we have
E 0i,44 01.411
S (1 + B11,)) E
or 1.6
+
Mk E
(0%:2 - o6 (0
< (1 + Brbh E
0 0%14
+ Thk E
1.9% (0
(28)
< (1 + Brk)) E
10 05:1
+ 2yyh
(1 + (1 +1+ -k Bric (a + EE (0:)1 + ogg))
VV
< (1 + Brkk) E
|00.
0/:1 ||
+ 2
(E V (0.) + 0 +
n
where the last inequality follows Assumption D.I. We let with E.5 gives
+ (1 +
k
Būe S and πp: Σ K-1 mh Then unrolling (28)
K=
K-1
E
00.x
(1
+
Brlk)
E 0''1
2
Σ
II (1 B1) (E (0))l + ogg +o..+ +
k=
7
k=ū
1=k+1
(29)
2
< B
E
0'1l +
ट
E
|
F
(0+)
م. + a).
21

---

Improving Generalization in Federated Learning via Momentum-Bass Based Stochastic Controlled Weight Averaging
With 0:1 1. Σ 0, and 0/41
1. E
0%
have
E 0 0141
E
0 0 05
i=1
m
(30)
2
s e
Min
E
10
0/1l
+
The
(E
VF (0:)1 + 0gg o))
Further, unrolling the above over and noting 8o = fg. we obtain
T-1
26
E
||0 - 87 S
exp
B
Σ it
the
(EIVF (0.) + + + 0);
(31)
mn
1=1+1
We simplify this inequality
E OT 0/
22 +++
(E (0,) + +o + +
mn B
(32)
2ō 1 +1
(I 0g +o),
B
With the Assumption C.2,
2Lc 1 1+ (EIVF (0.) +o, +o)
ggen
B
With
K-1
ē= 2=
(2 TK
>1
T
When the diminishing stepsizes are chosen in the statement of the theorem, we conclude the proof.
F. Theoretical Results of FedMoSWA
F.1. Generalization Analysis for FedMoSWA under strongly convex setting
where g denotes the stochastic gradient at the point (()) i,k
Lemma F.1. Suppose Assumptions C.I-C.5 hold. Then for FedMoSWA with vi S 1//KKT, a
10 - 02ll S n.
where ij. = -K- 1/. b
3' 34p
22

---

Improving Generalization in Federated Learning via HomentuumBasee Stochastic Controlled Weight Averaging
Proof. Considering local update of FedMoSWA, of 9 (0)), m = - (1--7m² + Σn 91 (0))
E 1 1,++ 01|l
E 11,0i T/. (gi (01,.+1) - cj + mm) -ell
= E 110 - 0, na (0), + th (gi (0t) + m')l
E ||0f - 0, na (0,)+))) + ni.i (0t) + n m
EE |00
0c ná (9: (0%k+1) 9b (0.)) || + ml Σ
(1 7) (~13 (0))l++)+++ +
1=0
(33)
Bp
18 110 0.1l + 1/ ||m'll
3 + 11
m1.Bp
E 8110,A411 0111 + n Σ (1-7) (7
63 VF (0))l +0)
3p
E 10 - 011l + Tis
(1-7)) (7)+
3 +
t=
where, is following Assumption C.3, b is following Lemma D.I, e is following Assumption C.2. Below we find an upper bound for E ||m='|.
E m || E (1 7) m)mm"
9b (0,)
1=1
1 m
S (1 7) EE m || + 71
Σ
94 (0))
(34)
(1 7) EE 1mml==1| + 13
V/F (0:) +70
Σ (1- (T +
1=0
L
where a follows Assumption C.3; b follows Lemma D.I. Unrolling the (F.1) and noting 0s,0 0c, with 6 z'8g for some
3+p
b >> 0, 7
-K11 v/c yields
E 0 K - 0.1l
< (1 b) K- -1j m
K
Σ (1 7)8 (FF (0,)) + (1 (1-b)k-144 K1-
m
(35)
1=0
0=0
K-1
Σ
np (L+ + 0) (1 (1- 1=0
S ih (1- -
k)- k-
Ob+
Lemma F2. Given Assumptions C.I-C.5, for vi < 11/BKT we have
E [|g (0 ( || (1 37 + +)) + 0g
where g:(-) is the sampled gradient of client i, ip. Σt=D T/-
23

---

Improving Generalization in Federated Learning via Momentum- Hoodentum-Based Stochastic Controlled Weight Averaging
Proof. we obtain
E 11: (0t
S 9i (0%,.) VF, (1.) x0)
s E VF; (0) +o
VF; + VF; (0%,))
VF: (6:) +
< /E (0:)] VF, (0:)
VF (8:)1 + + 00 Be
(36)
K-1
.7 (E (B (0:) 9.4 + σ) +
Σ " Σ
(2 +o) (1 6) K-1-1-
1=0
p=
(1 + 3ij + b) K- (I, o0) 0.
a is following Assumption C.3, is is following Assumption C.4, c is is following Assumption C.5 and Lemma D.I and
Lemma F.1, is following VF (0.) & L Similarly, using same techniques we have
19g (0,)1I
< E V F; (0+) +
E V F; (0,)
VF (0.)II (0.) to V (0:)
(37)
S 0g
E
a.
Theorem F3. Suppose Assumptions C.I-C.5 hold and consider FedMoSWA (Algorithm 1). Let {0:) and (0) be two trajectories of the server induced by neighboring datasets S and ((),, respectively. Suppose Og = 0g. Then,
EE 110 0'1 <
कलता (bL + og + b0),
B
2L 1
cotor (bL. + og + bo)
Egon
1
K-1 nh an upper bound of 1 + (1 (1- b) Bij- where ij = >K
Proof. For client i, there which is identical in S and (()), In this sense, we have,
cases consider. In the first case, SGD selects the index of an sample local step on
0 k+1
3 ||
00 1i 1,,41 - mh (9) (02,41)) - fj (0%2) oi + m ct + m
(38)
n B/3
+ p.) )
(8% 0%
where a following Lemma D.1. And this case happens with probability 1/n (since only one sample is perturbed for client i)
In the second case, SGD encounters the perturbed sample at local time step k, which happens with probability 1/n. We denote the gradient of this perturbed sample as g.( Then,
|
0/
||
0 0
m (os (0),.41) 94 (0%:2) ct + +m ct m²)
j,++1
0
0%
mi (9x (0f,))
ot (oti))
0f.
0%
mi (Di.))
gi (0%%) + n/ |9.. (01.1) 96 (0
(39)
n B1
|0i. 0%
nn 99. (0%) of (0%)
B + 11
24

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
Combining these cases have for client i,
E
0
1,++1
02
5
1
3 + 14
IE 10
0
+ Th E
9i (oiL) s (0
201ee EE 11: (0
1E
|0% k
03
+
(40)
3 + ps
n
2
2nf gg
b 1
E
07
+
(1+ + Bij b)) (L + o)
3 + p
where the inequality follows that (-(-) and (-) are sampled from the same distribution, b is following Lemma F.2. We let 1 + (1 b- Bijp S b, n/+ ku mi. Then unrolling it we have,
K-1
E 0i, 0i S
K-1
II
mh. Bp
E 1
0111 +
2 K-1
Σ
T
K-1
II
3p
B 1
B 11
k=0
1=k+1
(41)
(BB VF (0.) + g. + 0) + (1- k,,b m 11))
S e
E 10: - 0i1 + n
ThE
(LB (0))l + 0g.i + σ) + (1- b)* -
77:33 ||
Then, with 0.+1
1 m ===1
0 k and 0,21 = 1 Em
i. we have
E 0 - 6/+11 - >>
E
0 0/
1=1
(42)
S 3
110, - 0/1l +
2 the
3T
((B])FF(O) + tog +o) + (1-b)k- 6)* k -' |
mm
where we use Lemma F.2 in the last step. Further, unrolling the above over t and noting 8g = =6/, b obtain,
3+
B+
we
E
1
I1
2
Bu
3 (EE + to) + (1 k)kk -, || II)
exp
n.e
mn
3 + K
1=
1=++1
(43)
2
T-1
By
T-1
exp
The
-1+ 11) (L+0)+
mm
B 1
1=0
1
When the diminishing stepsizes are chosen in the statement of the theorem, we conclude the proof.
F.2. Generalization Analysis for FedMoSWA under non-convex setting
Lemma F.4. Suppose Assumptions C. C-555 hold. Then for FedMoSWA with vi S 1/3KT,
E 110 00ll < 7e(1 + c)==1(L++))
where ij. = K-1 Vh
25

---

Improving Generalization in Federated Learning via Momentum-Barye Stochastic Controlled Weight Averaging
Proof. Considering local update of FedMoSWA, cf - (0.), m - -11 - 11)m'm ? Σn 91 (0))
E 8 . - Be]l
= E ||0f Thé (9: (0t,x) ct + m²) - -ell
B. |10t 0. nh (0%,t) + náé (gi (0c) mm)
E 110t O: nh (x) + nhg. (0t) + T/
m
(44)
s E 110% - 0e - nh (g: (0% 9a (0:)) + m Σ
(1 - (? (0i)l + +7)
1=0
2 (1 + 3mk) B 1100,,41 - Bell + n/ > (1 (1-7) ( (0))l 77)
f=0
where, a is following Assumption C.3, b is following Lemma D.I. then we can get,
E 100 0.11 s (1 + Brk) E 100, Dell + n Σ (1 - (7E (0:) +70)
1=0
(1 + Brk) E 110 0.11 || + ml(
(45)
K K
III (1+ /Tf) ( (
k=0
< in + (L)
is following Assumption C.2. Below we find an upper bound for E |mmm [],
E m 1| E
(1
7)
m
1-1
+
(0t)
S (1 7) E m
|| + IE
9s (0))
(46)
< (1 7) E
|| 1
V F (0 \|| +
7) (+ (0:) + 70)
1=0
S L
where we use Assumptions C.5.With (1 + (T S c, unrolling the above and noting 0,,0 = 8g yields
E ||0 0:|l (1 + c) x)<< -2 m
(47)
S (1 + (1.to)
Lemma F.S. Given Assumptions C.2-C.5 and considering of FedMoSWA, for ofé S 1//KKT, we have
E flo (0 S (1 + 31 + c) K (L + 0) + o.
where g:(-) is the sampled gradient of client i, ./. Σk=D 7/-
26

---

Improving Generalization in Federated Learning via Momentum- Honeettu--Based Stochastic Controlled Weight Averaging
Proof.
E 91: (0
$ 3 [0:) (0
= E V.; (0%,)
< V [0])] (0/,) VE (0.) to
V/ ( -0.) VF (0:) + BE | - -0 +o
(48)
K-1
(B + 9,5 + +00+ + Σ
n
==
<" (1 + 3ij + +c -) O) gg.
a is following Assumption C.3, is is following Assumption C.4, c is is following Assumption C.5 and Lemma D.1 and Lemma F.4, is following | V (6.) - L.
Theorem F.6. (FedMoSWA). Suppose Assumptions C.2-C.5 hold and consider FedMoSWA (Algorithm 1). Let k = K, Vi E [m]] and The S рKT ē (1 + Bij (1 + -) (L + to). Then,
2L
h _ 1+ (eL + to + +5)
Ren I mn B
where, , (1+ iee
K-1 kk
The
Proof. For client i, there are two cases to consider. In the first case, SGD selects non-perturbed samples in S and S(()
which happens with probability 1/n. Then, we have,
0l 8 $ (1 (mk) 00,x
07
(49)
In the second case, SGD encounters the perturbed sample at time step k; which happens with probability 1/n. Then, we have
|0 - 0% 0//t111
= 8 0% Th gs (02,.)
g6 (0%%) of + m ct + ml)l
- 01, 0% n gs (0%.,) 91 (0% || + r 9: (0%) - I (0%22
(50)
(1 + BTT)) |00. 0it6| + mí 11o. (0:) - of (0%) ||.
where is following Lemma D.1. Combining these two cases have for client i,
E 0 11+11
S (1 + (x7%)
E 0/ 0
+ th E
gi 1.8) - 96 (0%:
<" (1+ 3mh) E
1 0/:6| || +
2mf E ||9: (Bf,.)
(51)
2nk
5" (1+ + 3mk) E
0 0/
+
(EI (0.) Fg.i + 0) +(1+++++++++ =11+c) -l m
n
s* (1 + 3vkk E
|0 0111
+ (1 + 3o (1 + c) k-1) (L + o) + og.
where the a inequality follows that g=(-) and (+) are sampled from the same distribution, b is following Lemma F.4, css following Lemma F.5. We let 1 + (1 c) K- Bī < ē, πe = Σk- nr. Then unrolling we have,
K-1 1
E 0
0:1k]
<
(1 +
3 E10 01ll
k=
Σ the
2
K-1
II
(1+ + 3mb)
(0.) + 3.g +0) +1 (1 1 | ())
(52)
1=k+1
2
S 10.
0,11
+
3h. (EIIVF (0.) o. +0) (1 e) -1 ml'))
n
27

---

Improving Generalization in Federated Learning via Momentum-Basee Based Stochastic Controlled Weight Averaging
Then, with 0141
1 m Σi= 0fx
and 0/41
1 m
Di=
0%
have
E2 0c+1 0242/1
Σ
1 E
0 i,K
0i
i=1
(53)
2
$ R 10+ - 0''ll
+
Th
Rh
(F + +0) +0) (ec--)=(ne R m'|l
mn
where we use Lemma F.4 in the last step. Further, unrolling the above over t and noting 0g = 6/. we obtain
T-1
2
E
1|
1
S
Σ
exp
B
Σ
me
(EE]VF(()) + to) + (1 c) K 1 B |mmm)
mn
t=
(+++1
2 T-1
T-
1
<
exp
3
Σ
mm
t=0
1=1+1
(BI]BF (0e) +0n +0) c1+
-1
(1 21--)) (> (0i) + (oo
T/e
1=0
2 T-1
7-1
exp
8
Σ t
(54)
mm
1=
1=t+1
E
+ oos +o + (1 k)) Σ (1 7) (FEFFFFF (0)1))
1=0
2 T-1
T-1
exp
3
Σ
itt
T/e
0i + B1 + c) k11 (I. + σ) +
mn
1=1+1
2 T-1
T-1
exp
8
Σ
m Mm (ēL + 0g + 20)
mm
1=0
1=t+1
2L 1+ (čT. + 0g 20)
@gon
mn
B
_
When the diminishing stepsizes are chosen in the statement of the theorem, we conclude the proof.
G. Convergence Analysis for FedMoSWA under strongly convex setting
For the optimization problem, we consider:
min
F(0) = m
(F. = E3, [F (0:)))
BERd
Before giving our theoretical results, we first present the common assumptions.
Assumption G.1. F/ is c-stronglyycomverfor all e [m], i.e.,
F((0,) 2 (.2) + (VF. 01 - -2)) + 4 10, 100, - 0212
(55)
for all 02 02 its domain and i [m]. We allow μ O, which corresponds to general convex functions.
Assumption G.2. (Smoothness) The gradient of the loss function is Lipschitz continuous with constant for all 02, 02 E Rd
[ F (02) )| S B 0 021l
(56)
28

---

Improving Generalization in Federated Learning via Momentum Eanentum-BBsedd Stochastic Controlled Weight Averaging
Assumption G.3. Let Ç be a mini-batch drawn uniformly at random from all samples. We assume that the data is distributed that, for all Rd
Egi [VF:(0:)) VF.(0).
(57)
We also can get:
E/j0 [IFF, (0:c)) -V.(()) c
Assumption G.4. (Bounded heterogeneity) The dissimilarity of F:(0) and f() is bounded as follows:
VF - VF(O))1 S og
(58)
Assumption G.3 bounds the variance of stochastic gradients, which is common in stochastic optimization analysis (Bubeck et al., 2015). Assumption G.4 bounds the gradient difference between global and local loss functions, which is widely-used approach to characterize client heterogeneity in federated optimization literature (Li et al., 2020; Reddi et al., 2021).
G.1. Some technical lemmas
Now we cover some technical lemmas which are useful for computations later on. The two lemmas below are useful to unroll recursions and derive convergence rates.
Lemma G.5 (linear convergence rate). For every non-negative sequence (dy-1)r21 and any parumeters O, Mmax (0, 1/1,, c 2 0, R 2 Miaast there exists constant step-size S "mmxx and weights == (1 pm))- such that for
Wn = R+1 Wr.
R+I
Wr (1
Wr d, crww, -0 -0(ude exp (-10mmax R)
RR-
(59)
WR
1RRi
Proof. By substituting the value of wy, we observe that we end up with a telescoping sum and estimate
do nWR
(w.-dd,-1 w,d+)+
+ cq.
nWR R
WR R
When R 2 2un (1 - pn) S exp(-Mmmip Rm S 3. For such an R, we can lower bound r/R using
1
R
um)" = 11 pm) R 1-
(1-pm)
- 1
WR = 7/1 - um)
R
2 (1 (-)))
117)
3/p
-=0
This proves that for all R 2 2pm
V R S 3pd pn) + on S 3pd, exp(-unR) + cn.
The lemma now follows by carefully tuning y-- Consider the following two cases depending on the magnitude of R and 7maxi
log( (maam1.1""
1 < Manx Suppose 24R
Then we can choose r = 7max
nR
WR < 3 exp 1- pmmaax R + Cmmax &
3jdd exp 1- R
+ c
HRR
log (max(1xxx)) Rdd//c)
log (max (lmaxx(1ppp Rdd//))
Instead Tmmxx
we pick 17 =
to claim that
nR
uR
29

---

Improving Generalization in Federated Learning via Momentum-Basee Based Stochastic Controlled Weight Averaging
V I 3 exp [- log (max (1, +0
c
1R
Lemma G.6 (sub-linear convergence rate). For every non-negative sequence (dy-1)r21 and any parumeters Vmmax 2 0, R 2 0, there exists constant step-size 17 Tmax and weights wr 1 such that,
R+1
dy-
d + c17 + c21))
do
2v C1 do
do
5 1max (R+1)
+2
(60)
R+1
R
+1
R +
Proof. Unrolling the sum, we can simplify
do 1))
MR S
+ c1T) can
Similar to the strongly convex case (Lemma G.5), distinguish the following cases:
When R+1 S
do ei Fmmax
and R+1 S
dg oq Vinax
we pick 1) 7/max claim
3
do
da Mmax R +
ve
dg
R S
+ C17max + e21
(61)
7max
R 1)
VR+I
R R+1
In the other case, we have maax prove
do ci([t+]]
do cx(H++]
We choose 1) min
(V
c[(++]]
dg
ca(
dg
to
do
2/ /R +1
dBc2
1)+ V n S
+ cm
+ 2//
(R 1)2
Next, we state a relaxed triangle inequality true for the squared &g norm.
Lemma G.7 (relaxed triangle inequality). Lemma 3 (relaxed triangle inequality). Let (Ui,..., v-} be vectors in Rd. Then the following are true:
1. ||0, + v1 < (1 + a) 11, + (1+ 2) || for any a > O, and
|C will < MT 1
Proof. The proof of the first statement for any > follows from the identity:
a0 + 0, (1 a) ||0i11++ + (1+ ||,, Vav, +
Va Uj
For the second inequality, the convexity of + ||0||2 and Jensen's inequality
v 2
v
Next we state an elementary lemma about expectations norms of random vectors.
30

---

Improving Generalization in Federated Learning via Momentum-Basee Based Stochastic Controlled Weight Averaging
Lemma G.8 (separating mean and variance). Lel (_i,..., _.) be T random variables in Rdd which are not necessarily
independent. First suppose that their mean is E [ E, and variance is bounded E following holds
[ - -Eill) S o². Then, the
E
=
Σ ==1
+
T²00
Now instead suppose that their conditional mean is E difference sequence, and the variance is bounded by [
|
3 = E, i.c. the variables {E, - Ei) form a martingale
&||l f o² as before. Then we can show the tighter bound
3
2
&
+ 2TO².
Proof. For any random variable E
[x²]
[E[X- - 2[XI)) + (B[ implying
E
3
C
E
Σ E
i=1
Expanding the above expression using relaxed triangle inequality (Lemma 3) proves the first claim:
E
Σ E 1=1
E
&
E IE, <1l?] S-'02
For the second statement, & not deterministic and depends on B. Ej. Hence we have to resort to the cruder relaxed triangle inequality to claim
E
2
Σ
E1
2E
Σ B
&
and then use the tighter expansion of the second term:
Σ =
&
Σ
E [ iJ
$) TE,
[
Σ E ||E - -Eil1 TO²
The cross terms in the above expression have zero mean since (E, {.} form a martingale difference sequence.
G.2. Convergence of FedMoSWA for strongly convex functions
We will first bound the variance of FedMoSWA update in Lemma G.9, then see how sampling of clients effects our control variates Lemma G. 10, and finally bound the amount client-drift in Lemma G. 12. We will then use these three lemmas prove the progress in single round Lemma G.13. Combining this progress with Lemmas G.5 and G.6 gives us the desired rates. Before proceeding with the proof of our lemmas, we need some additional definitions of the various errors track. As before, define the effective step-size to be
ij = Knf.a K Kma..
31

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
We define client-drift to be how much the clients from their starting point:
K m
&
|10
-'-
Km
k=1
Because we are sampling the clients, not all the client vantoo-varriats get updated every round. This leads to some "lag' which we call control-lag:
C. =
E E c
VF, (FF(())) (0*)
D, = B ml
VF (0))1 = ||||
Variance of server update. We study how the variance of the server update can be bounded.
Lemma G.9. For updates, we can bound the variance of the server update in any round r and any ñj Knfaa 2 0 as follows
127j²
E 0' - 11]
S 8312 (B [1 (B +
41f - 4 32
Ks
Proof.
The server update in round t can be written as follows (dropping the superscript t everywhere)
E A E
E
Σ
(gr (0i, c;)
(62)
Ks
Ks K
RSESS
which can then be expanded as
2|
Σ
(ge (0..b-1) + m c.)
K8
KSEES
ij Σ
< 4E
9i (0.,6-1) (F(
4nj E m
||| 4E
VF.
(63)
KS
Ks K
KSESS
+ 4E
Σ (F( - VF, (8*)
Ks
KiisS
< 4E
7 Σ
gi -
- VE:(O)
+ 447 4E
1J
F. (0')
K8
K.ESS
iES
+ 8
(E(]f(0) - (0"))
(64)
< 4E
7
Σ
VF, (0i,--) -0;t-1)--VB((()
+ |r/m||/mm]] +4
7 Σ
VF, (0) B[c]]
K8
VES
12
+ 8,7 (E[f(0)- (0
-
K8
The inequality before the last used the smoothness of (F.). The variance of (t E EMNES 9s (01,t-1) is bounded by 2 2 Ks. Similarly, ey for any j E [m] has variance smaller than o² K and hence the variance of ( Σies c:) is smaller
32

---

Improving Generalization in Federated Learning via Momentum- Hooeettu-BBased Stochastic Controlled Weight Averaging
than o² 2/Ks.
4
4n
2| V
E VI, V |mm]]]
VF, (0") - EE [c.]]]
Km
+ 8,87 (E[s(0)] / (0*)) + 12
Ks
4rj
+ 4fj
B F. (0 (F)(
Σ
VF, (8") E [ccl1l + |mm]]|
(65)
Km
m
k,
Ti
o2
8,87 (BIf(0)] - ((0)) +
Ks
Since the gradient of is B-Lipschitz, control variate C=- m
Σx, E1 -11² = 3.. The definition of the error in
Km
Σー E E
[c:]
VF, (0*) completes the proof.
Change in control lag. We have previously related the variance of the server update to the control lag. We now examine how the control-lag grows each round.
Lemma G.10. For updates with the control update and Assumptions G.I-G.2, the following holds true for any îj = makK [0, 1/
C, S (1-23 - m
C + m
(48
(f - (0") 2))
Proof. Recall that after round the control update rule implies that of is set per
i & St i.e. with probability (1- (-=)-
ct
Ex= 9s 10i
with probability
Taking expectations on both sides yields
[c]] (11
-1-
8 B[
E
V F.
(0) (-)) vi [m].
Km
Plugging the above expression in the definition of C, we get,
C, = 1
V. (0*)22
*11 K
VF, )
1
(1- E 1 -F:
&
V
e
V
(
(0
(66)
K
k=
<1-
C- 1
E F., x0 20)
m2 2 K
33

---

Improving Generalization in Federated Learning via Momentum-Baxee Stochastic Controlled Weight Averaging
The final step applied Jensen's 's inequality twice. We can then further simplify using the relaxed triangle inequality as,
1 [Co]
(1
- )
C
+
E V., (0)
V., (0*)1
m²
K
i
C
2s
E
VF, (0 VF, (0)))²
28
E V.. (0) VF, 10)
m
m
m²K
C
+
28
E
VF. (0' VF; (0))1²
28
g2
Σ
-
E 11
-00-11---111 11
m
m
m K
i.
1
$
C--1+
+
8
(48
[ f (0' f (0)) + B,).
m
m
(67)
The last two inequalities follow from smoothness of {Fi} and the definition & mk 82 Ei 1 1, -0-
Lemma G.11. Assumptions G.I-G.2, following holds true for any ij == m/ ak E [0,1/8]
(( + D. $ (1 (4 11) E2)
Proof. Recall that after round t, the control update rule implies that of is set
m' = m 1 + Σ
(c) m m² = (1 3)m² 1-1 + Σ cf
(68)
1=1
1=1
K
E m - VF (0') E (1 - 7) m
F (0 (0i (x+)
K
(1 - 7) E 1 VF (0') 1| +
F. (0 1 -ii - AF (0*)
K k=
(69)
< (1 - 7) E |||m (0
1 E
2
VF VF
m
2
Σ E
VF (,:-1,)) VE, VF,(0' (0--1) |
mk' ?
i.k
S (1 - 7) | V (0)) +
7 (48 (B [ f + (E.)
Bounding client-drift. We will now bound the final source of error which is the client-drift.
Lemma G.12. Suppose our step-sizes satisfy Tf. m sigKa and F; satisfies Assumptions G.1- Then, for any global 2 1 we can bound the drift as
T 3,3TE, < C--1
1 D.-1
E [0''))
-0- + Koz 0²
3
3
25
Proof. First, observe that if K = 1, = 0 since 0., = 0 for all i E [m] and that C.-1 and the right hand side are both positive. For K > 1, we build a recursive bound of the drift. Starting from the definition of the update and then applying the
34

---

Improving Generalization in Federated Learning via Momentum- Hooeettu-BBased Stochastic Controlled Weight Averaging
relaxed triangle inequality, expand
E
10:
11:9 (0:) m 1pcc - 011
m72
LiES
B
00 171 VFF (0.) + mmm 1/C; 0
+ 11/02
m70
LiES
(1
E
10: m VF, (0,) + MVF( 01
(70)
8
iES
72
1
1
1
1i
B2-1
m
+ VF. 1
m
iES
T5
The final step follows from the relaxed triangle inequality (Lemmaa G.7). Applying the contractive mapping Lemma for S 1/3 shows
1 m IES
To
Σ 10, - T V. (0.) + m V - 0 S 0 - 01
(71)
Once again using our relaxed triangle inequality to expand the other term Ta, we get
T3 = ET-1
m c. + VV
Y11
IES
m
m - G (0)(
m
j=1
1
m - C + VF, (0) + VR(0) (0'))
m
j=l
(72)
< <3|[m]]
+
( +
VF((0) VAI
j=l
j=1
4+
||e VF: +
[VF(() VF, (0))
j=1
3
< m|| +
||e
VF, VF,
63 (
The last step used the smoothness of F, Combining the bounds on Tz and Ts in the original inequality and using a gives,
K-1
(1
1
E 0
011 <
K-I
E 0 10, -01 - n/o
+ VF, (0*)11² + 3/mllmm|
m
(73)
3K m
-
E ||e,
m
35

---

Improving Generalization in Federated Learning via Momentum-Bann Based Stochastic Controlled Weight Averaging
Recall that with the choice of ct, the variance of c; is less than
a² K
Separating its mean and variance gives
1 Σ B B100 - -
1
[1 K
E 0 01 * 7mj o²
K-1
778
m
(74)
6n -
3K
E
|lc
VI 1 +3Kmfmmm +
m
Unrolling the recursion, we get the following for any k: E (I....,KK)
)
1
E L.
011 s
1+ K
m
(1+x-1)~- -1)
k
- (6K
( (3 (K 1)
s R (S(0) - (k
< 18K?Bn (1(0) (0)))+KK + P
+
(75)
The inequality (K - 1) (11
K-1 )
-1 -1)
[-1
K
I 3K can be verified for K = 2.3 manually. For K 2 4
-1 (1+
K
K
4
(K (K-1)
K
exp
K K-1 -1
s K
exp
3
<3K.
(76)
Again averaging over is and multiplying by 33 yields
38E < 54K? 3 (5(0) - (0*)) Bm) C, C, + 27K D 63
D
1
54 (f(0) - (6")) + +7,78 + 27 63 0² 1
a²
K
(77)
1
1
1 3 c
3 -i + 9
0²
(f(0) f (0))
a
25
The equality follows from the definition r) K a,a and the final inequality uses the bound that r) S 813*
Lemma G.13. Suppose assumptions G. G.2 are true. Then the following holds for any step-sizes satisfying a 2 1, S min
810Ka ToymKa) and effective step-size i Kam
E 00' - -0112=
< (1 (1- jun R [E -0+112 + 9mTT"²
C.
C.-1
(78)
2
(E [f (0' (0' (0+) + 12n2
[1 oz 0²
(79)
Ks
Proof. We can then apply Lemma G.9 to bound the second moment of the server update as
A0
Σ (gi (0t,6-1) m c;), and E E
Σ9: g: (0,,k-)).
Ks KSES
Km
36

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
We bound the second moment of the server update
20
R A0-
||0 +
0-1 = E2 100 0-1
B
(VF,
(0 -B-111A011
K8
K
2î
E
(VF 10, - 0) +IE, | - 10
Ks
(80)
K
T4
12
+ 8 [
- f(0)) + (
Ks
The term T4 can be bounded by using perturbed strong-convexity with h = F;, 0 = - A = and z 0 to get
E |Ta
2 E
Σ
(VF,
(0., -0)
Ks
KiES
2î
E
(0") F.) + + -1100,,-1 -012 M 10 0-1²))
(81)
Ks
KJES
2iiE (f(0) (0') + # |0 10- + 2.R...
Plugging T4 back, we can further simplify the expression to get
0-11² < 1 1¹ - 2n) ( + 1 10 -0 + 237
E 0 + A0
12
+ 8
E f 10' - f (0')) 44 4nj² DI-1 + 472 2oj2D_=111
Ks
(82)
110 0-1 (837) 27)) (ʃ(0) 1 (0"))
12
+ 43 + D
Ks
We use Lemma G.10 can (scaled by 97j² m bound the control-lag
1 _ D. (1 (1- [0])
91j2
9îو D
+ 9
-1
[ - (0)) + 2338)
(83)
7
7
2
Now recall that Lemma G.12 bounds the client-drift:
3 S T
772 712
C--
2 Dr
n
- ((0)) + (0
0²
(84)
3
3
25a
K
Adding all three inequalities together.
E 0 + A0
2
+ -//² D.
S
E 0 2+ gr²m mnfmm(11)
49
+
44.372
25
(f(0) - f(()))
(85)
2
12
(1
8 ) + (22 Bî) E
9uim
C²
9upd 1) 24
D
Ks
25
3
3
س Finally, the lemma follows from noting that 7 S 1 implies 4432 < 21 and S T5ym implies Purm
J. Note that if
25
2x
cf gs (0") then fm Co be bounded in terms of function sub-optimality F. The final rate for strongly convex follows
simply by unrolling the recursive bound in Lemma G.13 using Lemma G.5. Also note that if of = gx(x)), then fm Co can be bounded in terms of function sub-optimality F.
37

---

Improving Generalization in Federated Learning via Momentum-Basse Based Stochastic Controlled Weight Averaging
G.3. Convergence Analysis for FedMoSWA under non-convex setting
We now analyze the most general case of FedMoSWA with option II on functions that are potentially non-convex. Just as in the non-convex proof, we will first bound the variance of the server update in Lemma G.14, the change in control lag in Lemma G.16, and finally we bound the client-drift in Lemma G.17. Combining these three together gives the progress made in one round in Lemma G.18. The final rate is derived from the progress made using Lemma G.6.
Additional notation. Recall that in round r, we update the control variate as
R E Di (0) c
i E S otherwise
cf
We introduce the following notation to keep track of the 'lag' in the update of the control variate: define sequence of parameters [a(-_ ( such that for any E [m] and E [K] we have a. ox == 0 and for 2 1,
0 ii St 1-1
, ==
aiik-1
otherwise
{
By the update rule for control variates and the definition of (
above, the following property always holds:
cf
Di (a
K
k=1
We can then define the following E, to be the error in control variate for round r:
K
E =
-
Km
R=1
Also recall the closely related definition of client drift caused by local updates:
K m
& = Km
-
E
00-1
k=1
Variance of server update. Let analyze how the control variates effect the variance of the aggregate server update.
Define sequence of parameters { {atx-1) such that for any i [mr] and l E [K] we have a. 0 and for r 2 1.
yi p 3 0
0 1
p -
1 Σ6
1
m' = m +
(cf m 1 = (1 3)m 1-1 +B
c
i=1
1 + B
[[m] = (1 - B)L
m 1-
1]
E
96
(0 (0(kk-))
m
mk
1=1 k
We can then define the following E, to be the error in control variate for round r:
K
G--1 =
E e k-
0-12
Km
km
38

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Variance of server update. Let analyze how the control variates effect the variance of the aggregate server update. Lemma G.14. For updates and assumptions G.3 and G.2, the following holds true for any ń mak [0.1/8]
E B 00
0 1 - 20 B2
0 1| and
9ij
E
0'
-0-
112
S 4n/ E2 +
9
VF 0 ||
Ks
Proof. Recall that that the server update satisfies
B[A0]
1j
Σ
E [g
(0L
Km k,
From the definition of a t-1 i and dropping the superscript everywhere we have
A
Σ
(9. + m -ci) where =
1 K 9s (aix)-1).
Σ
Ks RSIES
h
Taking norm on both sides and separating mean and variance, we proceed as
E 21
E
(g. (0i.h-)) Di (cG. + m ci)
K8
KS
91/202
$ E
(VE, (0ik-1))) E(m] - [c,]))
K8
K8
K
9rj²02
< E
VV (1) F[m] [c:]]]
K$
Ks
AS
9rj00²
(V -0)(0) + (B(m)
VF) ( - VFF [c] - VOOO) Ks
Km
4
4
E VF, (0 2
E
VR (a. VF((O]
Km
Km
2ri²
91j²02
+ Km k.
4 VI (0. 2
VF
Ks
97
< 34 + 43 + +
Ks
(86)
Ke ER Di (0a,)) 1 Eies C, and m--ma not be indepen- dent but each have variance smaller than Ke and so we can apply Lemma G.8. The rest of the inequalities follow from repeated applications of the relaxed triangle inequality, 3-Lipschitzness of Fú, and the definition of E,- This proves the second statement. The first statement follows from our expression of [A0) and similar computations.
In the first inequality, note that the three random variables
1
x²
Lag in the control variates. We now analyze the 'lag' in the control variates due to us sampling only a small subset of clients each round. Because we cannot rely on convexity anymore but only on the Lipschitzness of the gradients, the control-lag increases faster in the non-convex case.
Lemma G.15. Assumptions G.3, G.2, the following holds true for any क S 24, (7)) for a [2 1] where क == maK:
Gr < (1-
17
9--1 + 485²
1
(7)
20
VF
(0'
|| +
97
(7)
20-1
E, +
Y
o²
36
48
g2
32Ks
39

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging Proof. The proof except that we cannot rely on convexity. Recall that after round 4, the definition of 0,,x-1 implies that The proof proceeds similar to that of Lemma G.10 except that we cannot rely convexity.
Est [0 (1 - 0 + 770 A-1-
Plugging the above expression in the definition of E we get
E. Km
E
0 i 0¹1²
= (1 - 7) Km
1,
0-
+7
1
EE
10 - 0-1
(87)
Km
k,
Ts
To
We can expand the second term Ts with the relaxed triangle inequality to claim
To 2 & E E||
We will expand the first term Ta to claim for a constant b 2 0 to be chosen later
To =
Km Σ
0 g-
01 + || + Ex-1 ( AA0 - 07-11)
(88)
<
1 0¹ 00 + || + 1 E [A0
1|2 [A0]]
-
11-
Km
where we used Young's inequality which holds for any 2 Combining the bounds for Ts and To.
g, (1- (1 () + 2yE, + 2E |0 + E E [ ||
s (1 () 1 + 2 + 82 - 2 1 2/ &, + 8 2 (EIV)(()))))
1877'02
+ Ks
(89)
Verify that with choice of "2(->)- we have (1 (1+ 3) and 1 < 2. Plugging these values along with the bound on the step-size 16 < 36 (7) S 36 completes the lemma. Bounding the drift. We will next bound the client drift &
Lag in the control variates. We now analyze the 'lag' in the control variates due to sampling only small subset of clients each round. Because cannot rely convexity anymore but only on the Lipschitzness of the gradients, the control-lag increases faster in the non-convex case.
Lemma G.16. Assumptions G.3-G.2, the following holds true for any ij S 24/3 (=) for a E [ where 7 == mak :
1
17s
E.1
1
2a
VF 111
97
2a-1
E,
o²
36m
48,32
48
mn,32
32Ks
Proof. The proof except that we cannot rely on convexity. Recall that after round t, the definition of af implies that
Es:
[at = (1 -11-
(1- a i +
8 O
m
Plugging the above expression in the definition of E, we get
E, =
Σ E -
00
Km
i,k
= (1 -1-- K
E
1-1 Gik-1
0¹
2
E
0
0'11
(90)
Km
k.
To
Ta
40

---

Improving Generalization in Federated Learning via Momentum-Basee Stochastic Controlled Weight Averaging
We expand the second term To with the relaxed triangle inequality to claim
To s 2 & - E =|
(91)
We will expand the first term T3 to claim for a constant b 2 to be chosen later
To
1
Σ
E
a
1-1
- +
0/-111-
||
+ Ec-1 (50', - -
Km
(92)
Σ
E
- 01-112+
||
+ T | [A || + b -
Km
where we used Young's inequality which holds for any 2 0. Combining the bounds for Ta and To.
E S (1 m =)
+ b1+
+2 &
+
21E
A0'
E
E Aee |1
m
< $ (1 E +
(25
228
+ 11j B²
+ 2 82
1
E + (8 (8+2)) /EI +
181/
m
Ks
(93)
م (1- 2 and 1 S 2m Plugging these values along Verify that with choice of s) we have (1 (1-m) (1
8
with the bound on the step-size 16,32 < the client drift E,.
m < -_ completes the lemma. Bounding the drift. We will next bound
2a
36
36m
Bounding the drift. We will next bound the client drift & For this, convexity not crucial and we will recover very similar result Lemma G.12 only use the Lipschitzness of the gradient.
Lemma G.17. Suppose our step-sizes satisfy vn S 24 and F; satisfies assumptions G.2-G.2. Then, for any global we can bound the drift as
5 3 p <
5 _ o G- +
1 24a
E VF
00
112= +
4Ka² o
8B
6
Proof. First, observe that if K = 1, = 0 since 0r,0 = 0 for all e [m] and that E,-1 and the right hand side are both positive. Thus the Lemma is trivially true if K = and we will henceforth assume K 2 Starting from the update for i [m] and k E [
L. 0
=F
00 1- - mn (gi (0 (0i,k-1) m m c;) 011
<10
00 1i,k-11 1) (VF, (0 + m c:) - 011 + 1202
1 1+
1
E 1 100 011 + K (0. m -mal 1/22
K 1-1
1 +
E B110,,A-1 0 0100
K-
K Knt
V F. (0 0 (F (m VF(O)) VF(()) -cc VF((0)112
(94)
E 1 2
+4K 1 - o1
22
K
1 31T0::000~00000-0-03300200711111 -/A
+ K-1 + 3K3222 E
4Kn v / -K

---

Improving Generalization in Federated Learning via Momentum Hooentum-Basedd Stochastic Controlled Weight Averaging
The inequalities above follow from repeated application of the relaxed triangle inequalities and the 3-Lipschitzness of P.. Averaging the above over i, the definition of m
Σg Ct and - gives
E 0 - 011 1 m
E 0..
011
4K Ti
771
K-1
+ +h1 + 4K 1
4Km
)
#)
< <nio0 + +
4K
4K 12 9--1)
K + 4K 4K
4nj
47
47
k-1
4B27p2
+
BIV
-
9t-1
1
+
k2
Ko²
Ka²
Ka²
K
Ka²
k=U
1j0²
n3
113
VF +
+
-
3K.
24
14432Ka²
6Ka
6ka
(95)
The last inequality used the bound on the step-size BTj S +.4 Averaging over k and multiplying both sides by 3 yields the lemma statement. Progress made in each round. Given that we can bound all sources of error, finally prove the progress made in each round.
Progress made in each round. Given that we can bound all sources of error, finally prove the progress made in each round.
Lemma G.18. Suppose the updates satisfy assumptions G.3-G.2. For any effective step-size 24/ (
Koop satisfying
E If (8') 1f =.) 1 (Bl (0--1) 12 =(-1) + Ks
50
(1
8
" E
VF
112
s
s
a²
14
Proof. Starting from the smoothness of f and taking conditional expectation gives
8
E- + A0) (0) + E
EE--11A0112.
(96)
We usual dropped the superscript everywhere. Recall that the server update can be written
A0
K KS
Σ
(g: (0..--)) + m -c), and Es|
Km
7
Σ gi (0..-1)..
(97)
k,
42

---

Improving Generalization in Federated Learning via Momentum- Sareentum-Baseed Stochastic Controlled Weight Averaging
Substituting this in the previous inequality to bound E ||00||] gives
E + A0)] - -()) - ñ
8
[VE + 2 E 2012
Km
-
(VF(O), E EV (0:x-1)1)+
Km
k,
93fff02
21 E + 27 3 + p
+ - 237 +
2Ks
7 Σ E
1
VE (OR-i)) VF(OO
(98)
2
Km i,k
i,k
93 2
27 8 + 3 + 2 9-- + (
2Ks
S 3- IVF(00)
1
ij
E ( VF
2km
2 9331202
21 E & + + +
2Ks
i 2 s- 5 231 ||]F(0))
+ - 2,7 8 + 228 9- +2 2 9B1j20²
2Ks
The third inequality follows from the observation that -ab = 1 ((b - a)² - a²) 1020 a, b E R, and the last from the 3-Lipschitzness of F,. Now we use Lemma G. E, as
S 2 ((b - a)² d²) for any
m 12 m E. s 12,38p2
(1
17s
20-1
97
8
2a-1
d²
E +
& +
s
s
36m
4852
m
48 m
mg²
32Ks)
12 =
11
m
E
17
3///21111
+
- 318
m
2-2n
+
97
4
m
2-2a
& +
3,11002
8
3
8Ks
(99)
Now we use Lemma G.17 to bound G--1 as
12 gr s 12
7 ((1-
177
G +
48,32 (x)
1/F//// (OF(0'-1))
10' + 97
(7)
(7)
20-1
G--1+ +
7
0²
q
36
48
82
32Ks
17
1
(1) 2-2rx
|) + 97
11
2-20
3,8122 02
= =1238 G--1
B37
G--1
G-
+
7
37p
7 3272
-
E, +
3
8Ks
(100)
Also recall that Lemma G.18 states that
5 5
PnE T S 8
+
5
F VF 1)
+
1 4Ko²
(101)
3
6
24a²
Adding these bounds on E, and & to that of El + A0) gives
(8 +10)| +
12 - = + 12 2
G.
[B[s(0]] + +1000)))+22000))
===-)
=
5 17)
7
6 3
5
17)
1
-077
(m)
m 2-20
|VF(0)1
(102)
6 3
+
5îj
+
237
+
97
m 2-2a
32, +
3937ja2 0²
(1+ 8 2+)
2 3
4
5
8Ks
43

---

Improving Generalization in Federated Learning via Momentum-Bayee Based Stochastic Controlled Weight Averaging
By our choice of 2/2 and plugging in the bound on step-size Sij of convergence now follows by unrolling the recursion and selecting an appropriate step-size 1j. Finally note that if we initialize cf 9i (0%) then we have E 0.
m
2-2n 2a
S 24 proves the lemma. The non-convex rate
The non-convex rate of convergence now follows by unrolling the recursion in Lemma G.18 and selecting an appropriate step-size vj as in Lemma G.6. Finally note that if c = 9s (0") then we have Eg = O.
H. Effectiveness of local learning rate decay
Theorem H.1 (Effectiveness of local learning rate decay). We define client-drift to be how much the clients move When using a constant learning rate Tkk = -n we define the client-side drift to be
from their starting point.
| - 0--11']. and when using local learning rate decay The (x-1)
(1) (K-2) k
&s = km -b
Ek= Di=1
E
E. =
km Ex= k= E/=
|000/
0
11
[f (xt - f (x')) 7 o²
E S
Ci-1
+
Di-1+ + 75a23
E
93
93
3Ka²B
& < 3 93
1 7 C.
1 D. +
R [ (x' -xx- 5 + 3
9
75a2 B
So local learning rate decay client-side drift & is better than &
I. Generalization Analysis for FedSAM and MoFedSAM under non-convex setting Lemma I.1. Suppose Assumptions 2-5 hold. Then for FedSAM with of - /3KT, E 1 k - 0.11 S (1 + ir (E [ (0.) و.و + o)
Vk :o==,.... I,..., K,
where in =K= K-1 , ج きं 1 + (1 - Bň.
1 (1 K ビ n
Proof. Considering local update of FedSAM
E 02
0,11
E |00 k n (0) 0zll
S |0 0, - n/' (gi (02..) - gi (0.)) + '/E Il:: (0+)
< (1 + 3n') E 110, -olll + +,y [1:: (0:)
(103)
< (1 + 3n') E 100 - 0.ll + (E 19s (0.) VÆ, (())
$ (1 3n') TE |00 0.ll + (0.)) +a),
where we use Assumptions C.2 and C.3. Unrolling the above and noting 0.0 0g yields
E 1
0ell 5
n/ (EIVE (0)) + (1+
K-1
(104)
<
1/ [E V. (8.) +0) (1 Kc
< (1 + c) - in (E VE (0.) + o.g +0),
where the last inequality follows Assumption C.5.
Lemma 1.2. Given Assumptions 1 - 3, for < 3/3KT, we have
E ||s (C',) < (1 + () K 31))) VE + 0g +
where g:(-) is the sampled gradient of client i, ft. K-1 n/-
44

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Proof. obtain
E
9 (ot,,)
< E 9: -
VF, [ (
V.F (0%
$ VF; (8:) + VF; (0(,x) V.; (0.) +o
(105)
- E VF [ (0:) - VF (6.) + 38 1
0e| +o
$ (1+ (1 c1 (E IVF (0:) + 0gi +
Theorem L.3. (FedSAM). Suppose Assumptions 1-3 hold and consider FedSAM. Let 1 K, Vi [m] and n/ 3KT Then,
2Lc 1 1+ (EIVF +o, +0), B
Egen
mn
Proef. For client i, there are two cases consider. In the first case, SGD selects non-perturbed samples in S and S which 1/n. Then, we have
happens with probability 1 - 1//n.
8//41/1 S (1 (mk)) 00 07:4
(106)
In the second case, SGD encounters the perturbed sample at time step k. which happens with probability 1/n. Then, we have
|00
07/641
|8 0% n1 (9: (0t,.) o (o%)
< 10 0% n/' (9 (0t.x) gi (0(.)) + r/² |105 (0%.k d (otk)
(107)
< (1 3Tk))
00 - 07 +np |1s (0%)
gl (0%
Combining these two cases for client we have
E 0 01
$ (1 Bry") E
0
0/:6
+
n
9:
(
1,
dí (0%% )|
n
n/
$ (1+ 8T%)
0%
+
E
||9 (0%,,)l
n
(108)
+ 2n
$ <(1+ + (77)) E
P0 0%11
(1
+ (1 [ K- c)
+ E (6.) (8.) 0g.i)
2
<(1 B77) E
0%
+
(E VF + +g.
n
We let T be an upper bound of (1 + -1 since ît is bounded above. Then unrolling it gives
K-1
K-1
E 0lik 01:0||
s II
(1 + 3mb) E 0 0'11
Σ
TIEC
II (1+ 3nf) + gg +0)
k=
k=0
1=++1
(109)
8 E 110
0'1l + ctt
3h. E VF (0.) + g. + oo)
We have
m
E
0
1,41
0 04:1K
m
2
(110)
110e 0. +
cñt au EE
V (0,)
t +
mn
45

---

Improving Generalization in Federated Learning via Momentum-Basse Stochastic Controlled Weight Averaging
Further, unrolling the above over and noting 0o
0y. we obtain
T-1
E 1|0r - 0/ S
2
exp
3
Σ 1=++1
T Ri (BIIVF (0,)l + + + +o)
The
(111)
mn
t=
E 1| 87 S
e
+
+1
EE
VF (0,) + +0+
(112)
mn
With the Assumption 2,
2LT 1
++ (EIVF (0,)) + t0g
gen
B
K-1
ट E +1++c)+) K Bī: 5 (1+ ckk-1 K-1 3
< (1 Bij 1 1 01 3
1k
When the diminishing stepsizes are chosen in the statement of the theorem, we conclude the proof. The proof on MoFedSAM can be derived simply from FedSAM.
46


Reviews:
[{"review_id": "lLfYtm8pMP", "reviewer": "Reviewer_MQ96", "summary": "This paper proposes two novel algorithms for improving generalization in federated learning.\nThe first approach, FedSWA, is a variant of FedAvg with stochastic weight averaging, a method known for finding flatter minimums.\nThe second approach, FedMoSWA, extends FedSWA with control variates that are updates using momentum, in order to handle heterogeneity.\nBoth approaches are studied theoretically, showing superior generalization guarantees than FedSAM and MoFedSAM.\nExtensive numerical experiments confirm these findings.", "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 sincerely thank the Reviewer for their valuable comments, and our point by point responses to your major comments are given below.\n\n1. For Experimental Designs Or Analyses.  In fact, we have done the hyperparameter tuning for all the algorithms, and we also followed the parameter settings from their original papers.  We conducted full client participation experiments on CIFAR-100 with 10 clients and 300 rounds, where the model is ResNet-18. Here, FedMo denotes FedMoSWA without SWA (i) but only with Momentum Stochastic Control (ii), as shown in the following table. FedMo (62.5\\%) outperforms SCAFFOLD (59.9\\%) by 2.6\\%. With 100 clients and 10\\% participation, our  FedMoSWA achieves 3.3\\% higher accuracy than SCAFFOLD, addressing the update delay issue with partial client participation. \n\n\nFull client participation on CIFAR-100 and ResNet-18 with 10 clients and 300 rounds:\n| Algorithm | Accuracy (\\%) | Improvement (\\%) |\n|----------|--------------| -----------------|\n| FedAvg   |  58.2        | -                |\n| FedDyn   |  58.5        | +0.3               |\n| SCAFFOLD |  59.9        | +1.7               |\n| FedSAM   |  48.3        | -9.9              |\n| MoFedSAM |  37.9        | -20.3                |\n| FedLESAM |  59.2        | +1.0                |\n| FedACG   |  60.9        | +2.7               |\n| Our FedSWA (i)  |  60.2        | +2.0                |\n| Our FedMo  (ii)  |  62.5        | +4.3              |\n|Our FedMoSWA (i)+(ii) |  63.2        | +5.0               |\n\n2. For Weaknesses 1. From Tables 1 and  4 in our manuscript and the following table , when \\(\\rho=0.1\\), FedSWA using only SWA (59.8\\%) outperforms FedAvg (54.2\\%) and FedSAM (47.8\\%) on ResNet-18 with CIFAR-100. When \\(\\rho=1\\), FedMo using only momentum stochastic control achieves 65.9\\%, higher than SCAFFOLD (54.1\\%), demonstrating that momentum variance reduction mitigates SCAFFOLD’s variance reduction delay. By combining SWA and momentum variance reduction, FedMoSWA achieves 67.9\\%, showing that both SWA (+5.6\\%) and momentum control (+11.7\\%) are effective, and momentum control has  greater impact. Moreover, FedSWA is a simple algorithm like FedSAM and can be combined with other techniques. We will clarify this in our final paper. \n\n10% client participation on CIFAR-100 and ResNet-18 with 100 clients and 1000 rounds:\n| Algorithm | Accuracy (\\%) | Improvement (\\%) |\n|----------|--------------|-----------------|\n| FedAvg   | 54.2         | -               |\n| SCAFFOLD | 54.1         | -0.1            |\n| FedSAM   | 47.8         | -6.4            |\n| MoFedSAM | 60.1         | +5.9            |\n| FedSWA (ours)   | 59.8         | +5.6            |\n| FedMo  (ours)  | 65.9         | +11.7           |\n| FedMoSWA (ours)| 67.9         | +13.7           |\n\n\n3. For Weaknesses 2. To address your concern, we will add some discussions about the optimization error analysis in our final paper. In fact, Section 5.2 discusses the optimization error analysis, showing that FedMoSWA converges faster than the best-known algorithm, SCAFFOLD, and outperforms both FedSWA and other baselines such as FedSAM, FedAvg, and MoFedSAM. Unlike MoFedSAM, which uses only local momentum, our FedMoSWA employs momentum variance reduction. Additionally, our FedSWA also converges faster than both FedSAM and FedAvg, as shown in Table 1. We will include all the discussions in Section of Introduction in our final paper.\n\n4. For Other Comments Or Suggestions. To address your concern, we will improve the final version of the paper, and will use italic and bold to distinguish the two algorithms.\n\n5. For Questions.  In our analysis, data heterogeneity does not vanish when performing only one local training step. It is not purely an artifact of analysis but a fundamental characteristic of federated learning algorithms, driven by the inherent bias introduced by local training on heterogeneous data. This bias does not completely vanish with only one local training step. Our future work will focus on addresssing this issue.", "claims_and_evidence": "Theoretical claims are provided with full proofs.\nNumerical claims provide both intuitive explanations and more thorough comparisons with many existing methods on multiple datasets and in multiple settings.", "methods_and_evaluation_criteria": "Benchmark datasets make sense, as well as the splits used to emulate heterogeneity.", "theoretical_claims": "The proofs for optimization errors seem correct, and seem to be widely inspired by the proofs from Karimireedy et al., 2020.\n\nI am less familiar with analyses of generalization error and I did not check them in details, but the results seem correct.", "experimental_designs_or_analyses": "Experimental analyses are quite extensive, with comparison with many other algorithms on multiple datasets.\n\nIt seems that the only hyperparameters that were tuned are the client learning rate $\\eta_\\ell$, global step size $\\alpha$, momentum step size $\\gamma$ and local learning rate decay $\\rho$ for FedSWA and FedMoSWA. In particular, it is not clear to me whether hyperparameters of other methods have been tuned or not, which could make the comparison unfair.\n\nMore specifically, the client selection rate is set to an arbitrary value. To my knowledge, most of the baselines (e.g. Scaffold) are known to underperform when selecting only a fraction of the clients at each rounds of communication.\nThere are thus two major differences between Scaffold and FedMoSWA, that should be studied in isolation: (i) the use of stochastic weight averaging, and (ii) momentum in control variates updates.\nIt is therefore not clear whether superiority observed in experiments is due to the improvement of (i) or (ii).\nThis is a bit concerning, especially seeing that FedSWA is closer to the baselines in Tables 2 and 3, suggesting that improvement over methods like Scaffold may come from this momentum stochastic control.\nPerforming experiments with full participation would clarify this question.", "relation_to_broader_scientific_literature": "Related scientific literature is widely and appropriately discussed.", "essential_references_not_discussed": "Not to my knowledge.", "other_strengths_and_weaknesses": "**Strengths**\n1. The two proposed methods are shown to achieve better generalization error theoretically and numerically.\n2. The paper is very well written, with precise discussion describing intuition about the studied phenomenons.\n3. FedMoSWA is shown to outperform many baselines numerically.\n\n**Weaknesses**\n1. There may be a lack of precision in the numerical analysis, which does not allow to distinguish whether FedMoSWA outperforms baselines due to SWA or due to the use of momentum in control variates (see Experimental Design or Analyses section).\n2. The differences in optimization errors, notably between methods presented in Table 1, is not discussed. Providing discussion on this, specifically discussing whether this should have an impact on the result or not, would greatly improve this part of the discussion.", "other_comments_or_suggestions": "The difference between FedSWA and FedMoSWA in Algorithm 1 is only shown with using different colors.\nThis is a problem when reading the paper in black and white or for the colorblind: the difference should be indicated in another way (on top of using colors) for accessibility reasons.", "questions_for_authors": "In the end, the generalization error still depends on heterogeneity. Surprisingly, this term does not disappear when taking only one local training step. Is it an artefact of the analysis, or is there something fundamental that remains even when doing only one local step?", "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": "FJkpJDKn7P", "reviewer": "Reviewer_BLUJ", "summary": "Tackles generalization issues in FL with highly heterogeneous data.\n\nIntroduces a new momentum-based stochastic controlled weight averaging FL algorithm.", "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": "Thanks for your valuable comments. We will make modifications in the final version as suggested by the reviewer, and our point by point responses to your major comments are given below. \n\n1. Different from FedSAM, the proposed FedSWA algorithm is a new way to improving FL generalization. By using SWA, FedSWA better maintains global flatness consistency compared to FedSAM, as shown in Figure 1. Moreover, FedMoSWA also uses momentum stochastic control, unlike MoFedSAM with only local momentum SAM. Both techniques (i.e., SWA and momentum stochastic control) significantly enhance the theoretical results of the federated learning algorithms such as FedSAM and MoFedSAM, as shown in Table 1.\n\nExperimental results: All the experimental results in the manuscript show that our algorithms show consistent improvements on both ResNet and Transformer networks. As shown in Tables 2, 3 and the followng table, FedSWA with SWA optimizer (59.8\\%) outperforms FedAvg (54.2\\%) and FedSAM (47.8\\%) on ResNet-18 with CIFAR-100. Our FedMoSWA (67.9\\%) also surpasses MoFedSAM (60.1\\%). Both theoretical and experimental results show that FedSWA outperforms FedSAM, and FedMoSWA outperforms MoFedSAM.\n \n| Algorithm | Accuracy (\\%) | Improvement (\\%) |\n|----------|--------------|-----------------|\n| FedAvg   | 54.2         | -               |\n| SCAFFOLD | 54.1         | -0.1            |\n| FedSAM   | 47.8         | -6.4            |\n| MoFedSAM | 60.1         | +5.9            |\n|Our FedSWA   | 59.8         | +5.6            |\n|Our FedMoSWA | 67.9         | +13.7           |\n\n\nOur theoretical results in Table 1 show  that the generalization bound of our FedSWA is $\\mathcal{O}\\big(\\frac{L}{m n \\beta}e^{\\frac{1}{T}+1}(\\tilde{c} L+ \\tilde{c}\\sigma_g+\\tilde{c} \\sigma) \\big)$, which is correlated with data heterogeneity and is superior to  that of  FedSAM, $\n\\mathcal{O}\\big(\\frac{L}{m n \\beta}e^{1+\\frac{1}{T}}(\\overline{c}L+ \\overline{c}\\sigma_g+\\overline{c} \\sigma) \\big)$ and MoFedSAM, $\n\\mathcal{O}\\big(\\frac{L}{m n \\beta}e^{1+\\frac{1}{T}}(\\overline{c}L+ \\overline{c}\\sigma_g+\\overline{c} \\sigma) \\big)$, where $\\overline{c}\\>\\tilde{c}\\=\\1+\\left(2+1/KT\\right)^{K-1}/T\\gg 1$.\nWe prove that the generalization error of  FedMoSWA is   $\\mathcal{O}\\big(\\frac{L}{m n\\beta} e^{\\frac{1}{T}+1}(\\tilde{c} L+ \\sigma_g+\\tilde{c} \\sigma) \\big)$, which is better than those of FedSAM, MoFedSAM, FedSWA, where  $\\tilde{c}=1+\\left(2+1/KT\\right)^{K-1}/T\\gg 1$.\n\n\n2. To address youe concern, we conducted some experiments on CIFAR-100 with ResNet-18, following the pathological imbalanced settings from the FedSMOO and FedLESAM papers, where each client has at most 10 classes. This represents a highly imbalanced data distribution. These results will be included in our final paper.\n\nOur FedSWA algorithm outperforms the algorithms, FedAvg and FedSAM. FedSWA uses SWA as the local optimizer, while FedAvg uses SGD, and FedSAM uses SAM as local optimizers, and none of them employ variance reduction or momentum acceleration techniques. As the advanced version of FedSWA, our FedMoSWA, which also incorporates momentum and variance reduction, surpasses other algorithms in all the settings.\n\n| Algorithm | Accuracy (\\%) |Improvement (\\%) |\n|----------|--------------|--------------|\n| FedAvg   | 42.7      |-               |\n| FedDyn   | 49.1     | + 6.4             |\n| SCAFFOLD | 43.1     | + 0.4            |\n| FedSAM   |  41.2    |- 1.5              |\n| MoFedSAM |  45.6      |+ 2.9           |\n| FedLESAM |  44.3      |+1.6              |\n| FedACG   | 52.6      |+9.9             |\n| FedSWA  （ours) | 48.3       |+5.6           |\n| FedMoSWA（ours) | 55.5       |+12.8         |\n\n3.In fact, our source code was included in the supplementary materials. We will release our code and the federated learning framework to ensure reproducibility.", "claims_and_evidence": "Provides some theoretical guarantees and empirical results.", "methods_and_evaluation_criteria": "Evaluations are conducted on different data sets and ,methods, but the improvemnt with this method is minimal.", "theoretical_claims": "Provides both theoretical guarantees (convergence and generalization bounds) and solid empirical results.", "experimental_designs_or_analyses": "Experiemtns are conducted on several data sets and models.", "relation_to_broader_scientific_literature": "Related to problem of data heteroginity.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "The proposed methods (FedSWA and FedMoSWA) appear to offer incremental improvements over existing approaches like FedSAM and MoFedSAM. The results don't show consistant improvement.", "other_comments_or_suggestions": "The experiments focus primarily on benchmark datasets under simulated conditions (Dirichlet-). I'd like to see results in imbalanced cifar.", "questions_for_authors": "The source code of the method should be included.", "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": "xhezKuiFEc", "reviewer": "Reviewer_fAr3", "summary": "Generalization capability is critical for FL in real-world applications. This paper revisits the generalization problem in FL, focusing on the impact of data heterogeneity. The authors propose FedSWA, which uses Stochastic Weight Averaging to find flatter minima, and a varaint FedMoSWA, a momentum-based variant designed to better align local and global models. Theoretical analysis provides convergence and generalization bounds for both algorithms, showing that FedMoSWA achieves smaller generalization errors compared to FedSAM and its variants. Empirical experiments on CIFAR10/100 and Tiny ImageNet demonstrate the superiority of the proposed methods.", "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": "Thanks for your valuable comments. We will make modifications in the final version as suggested by the reviewer, and our point by point responses to your major comments are given below. \n\nTo address your concern, we conducted many experiments on CIFAR-100 with ResNet-18 under different data heterogeneity levels (0.1, 0.3, 0.6), and compare the proposed algorithms with FedSMOO and FedGAMMA, suggested by the Reviewer. The reproduced results are based on the pre-activation ResNet18-GN network, which is widely used in federated learning algorithms such as FedSAM, FedGAMMA, and FedACG. The following results are produced under our framework, referencing the open-source code of FedSMOO and FedSAM. Because the FedGAMMA algorithm lacks open-source code, we implemented it based on the algorithm flow described in the paper and referenced the FedGamma implementation in FedLESAM's source code. We will include these results in the final version, and will release our code and the federated learning framework to ensure reproducibility.\n\n| CIFAR-100| 0.1   | 0.3   | 0.6   |\n|----------|-------|-------|-------|\n| FedSMOO  | 46.5  | 47.8  | 49.2  |\n| FedGAMMA | 48.4  | 51.8  | 52.6  |\n| FedSWA (ours)  | 50.3  | 55.5  | 59.8  |\n| FedMoSWA (ours) | 61.9  | 66.2  | 67.9  |", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "some of", "experimental_designs_or_analyses": "Lacking comparison to \nFedSMOO: Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape\nFedgamma: Federated learning with global sharpness-aware minimization", "relation_to_broader_scientific_literature": "Try to improve generalization ability of FL in real world scenarios.", "essential_references_not_discussed": "lacking comparison to FedSMOO, and FedGAMMA\nFedSMOO: Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape\nFedgamma: Federated learning with global sharpness-aware minimization", "other_strengths_and_weaknesses": "Strengths:\n1) Well written and good representation\n2) Good results, although lacking comparison to FedSMOO, and FedGAMMA\n3) Providing extensive theoretical analysis\n\nWeaknesses:\nonly one major weakness is lacking comparison to fedsmoo and fedgamma", "other_comments_or_suggestions": "NA", "questions_for_authors": "Can you provide comparison to FedSMOO and FedGAMMA?", "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: The paper presents valuable contributions to the field of federated learning, particularly in addressing the generalization issue in highly heterogeneous data settings. The novel algorithms and theoretical analysis provide a solid foundation for further research. However, the improvements are sometimes incremental, and there are minor concerns regarding experimental design and theoretical limitations.
Decision: Accept (poster)