Understanding the Training Dynamics in Federated Deep Learning via Aggregation Weight Optimization
Authors: Zexi Li
Tao Lin
Xinyi Shang
Chao Wu
Venue: Submitted to ICLR 2023 (2023)
Keywords: Federated learning
deep learning
weighted aggregation
training dynamics
optimization
neural network.

Abstract: From the server's perspective, federated learning (FL) learns a global model by iteratively sampling a cohort of clients and updating the global model with the sum local gradient of the cohort. We find this process is analogous to mini-batch SGD of centralized training. In mini-batch SGD, a model is learned by iteratively sampling a batch of data and updating the model with the sum gradient of the batch. In this paper, we delve into the training dynamics in FL by learning from the experience of optimization and generalization in mini-batch SGD. Specifically, we focus on two aspects: \emph{client coherence} (refers to sample coherence in mini-batch SGD) and \emph{global weight shrinking regularization} (refers to weight decay in mini-batch SGD). We find the roles of the two aspects are both determined by the aggregation weights assigned to each client during global model updating. Thus, we use aggregation weight optimization on the server as a tool to study how client heterogeneity and the number of local epochs affect the global training dynamics in FL. Besides, we propose an effective method for \textbf{Fed}erated \textbf{A}ggregation \textbf{W}eight \textbf{O}ptimization, named as \textsc{\textbf{FedAWO}}. Extensive experiments verify that our method can improve the generalization of the global model by a large margin on different datasets and models.

Under review as a conference paper at ICLR 2023
UNDERSTANDING THE TRAINING DYNAMICS IN FEDERATED DEEP LEARNING VIA AGGREGATION WEIGHT OPTIMIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
From the server's perspective, federated learning (FL) learns global model by iteratively sampling a cohort of clients and updating the global model with the sum local gradient of the cohort. We find this process is analogous to mini-bater SGD of centralized training. In mini-batch SGD, a model is learned by iteratively sampling batch of data and updating the model with the sum gradient of the batch. In this paper, we delve into the training dynamics in FL by learning from the experience of optimization and generalization in mini-batch SGD. Specifically, we focus on two aspects: client coherence (refers to sample coherence in mini-batch SGD) and global weight shrinking regularization (refers to weight decay in mini-bate mini-batch SGD). We find the roles of the two aspects are both determined by the aggregation weights assigned to each client during global model updating. Thus, we use aggregation weight optimization on the server as tool to study how client heterogeneity and the number of local epochs affect the global training dynamics in FL. Besides, we pro- pose an effective method for Federated Aggregation Weight Optimization, named as FEDAWO. Extensive experiments verify that our method can improve the gen- eralization of the global model by large margin on different datasets and models.
INTRODUCTION
Federated Learning (FL) (McMahan et al., 2017; Li et al., 2020a; Wang et al., 2021; Lin et al., 2020; Li et al., 2022b) is promising distributed optimization paradigm where clients' data are kept local, and central server aggregates clients' local gradients for collaborative training. Although a lot of FL algorithms with deep neural networks (DNNs) are emerging in recent years (Lin et al., 2020; Chen & Chao, 202 Li et al., 2020b; Acar et al., 2020; Chen & Chao, 202 there are few works about the underlying training dynamics FL with DNNs (Yan et al., 2021; Yuan et al., 2021), which hinders us go further into the link between generalization and optimization in FL. In the meanwhile, interesting analogy exists between centralized server-client training framework of FL (from the server perspective) learns a global model by iteratively sampling a cohort of clients and updating the global model with the sum local gradient of the cohort. While in centralized mini-batch SGD, a model is learned by iteratively sampling a mini-batch of data and updated by summing the corresponding gradients. In the analogy, the clients in FL refer to the data samples in mini-batch SGD, the cohort of clients refers to the mini-batch of data samples, and the communication round refers to the iteration step. The interesting analogy makes us wonder:
SGD and FL. The
Can we leverage the insights of mini-batch SGD to better understand the training dynamics in FL?
Following this question and considering the key techniques mini-batch SGD (as well as its generalization), in this paper, we focus on two aspects of training dynamics in FL: client coherence (refers to sample coherence in mini-batch SGD) and global weight shrinking ( (GWS) regularization (refers to weight decay in mini-batch SGD). Firstly, sample coherence explains how the relations between data samples affect the generalization of DNN models (Chatterjee, 2019; Chatterjee & Zielinski, 2020; Fort et al., 2019). As an analogy, here we extend the concept of sample coherence to the client case in FL with partial participation for studying the effect and training dynamics jointly caused by heterogeneous client data and local updates. Secondly, in different line of works, weight decay methods (Lowkowycz & Gur-Ari, 2020; Zhang et al., 2018; Loshchilov & Hutter, 2018; Xie et al., 2020)-b by decaying the weights of the model parameter in each iteration step-are are the key

---

Under review as a conference paper at ICLR 2023
techniques in the mini-batck SGD based optimization to guard the generalization performance of deep learning tasks. We similarly examine the effects of weight decay in FL, in which we shrink the aggregated global model on the server in each communication round (i.e. global weight shrinking). tool framework to derive
Note that we take the server-side aggregation weight optimization the insights of the training dynamics in FL. Though the idea of aggregation weight optimization was appeared in previous FL works to match similar peers in decentralized FL (Li et al., 2022a) or improve performances in FL with medical tasks (Xia et al., 2021), all prior works assume normalized aggregation weights of clients' models (i.e. 7 1 in Equation 1), failing to dive into understand the FL's dynamics from the learned weights for further insights, e.g., identifying the significance of adaptive global weight shrinking.
Specifically, our contributions
three-folded.
We first make an analogy between centralized mini-batch SGD and FL, in which it enables us to derive a principled tool framework to understand the training dynamics in FL, by leveraging the learnt aggregation weights global-objective hobball-bjectivv--onistttee consistent proxy dataset.
As our main contribution, we identify some interesting findings (see below take-away) to unveil the training dynamics of FL, from the aspects of client coherence (cf. section 3) and global weight shrinking (cf. section 4) !. These insights are crucial to the FL. community and can inspire better practical algorithm design in the future.
We showcase the effectiveness of these insights, and devise a simple yet effective method FEDAWO, for server-side aggregation weight optimization (ef. section 5). It can perform adaptive global weight shrinking and optimize attentive aggregation weights simultaneously to improve the performance of the global model.
We summarize our key take-away messages of the understandings as follows.
Our novel concept of client coherence undermines the training dynamics of FL, from the aspects of local gradient coherence and heterogeneity coherence.
Local gradient coherence refers to the averaged cosine similarities of clients' local gradients. A critical point (from positive to negative) exists in the curves of local gradient coherence during the training. The optimization quality of the initial phase (before encountering the point) matters: Assigning larger weights to more coherent clients in this period boosts the final performance. Heterogeneity coherence refers to the distribution consistency between the global data and the sampled one (i.e. data distribution of a cohort of sampled clients) in each round. The value of heterogeneity coherence proportional the IIDD-esss of clients as well as the client participation ratio; the higher, the better. Increasing the heterogeneity coherence by reweighting the sampled clients could also improve the training performance.
Global weight shrinking regularization effectively improves the generalization performance of the global model.
When the number of local epochs is larger, or the clients' data more IID, stronger global weight shrinking is necessary.
The magnitude of the global gradient (i.e. uniform average of local updates) determines the opti- mal weight shrinking factor. A larger norm of the global gradient requires stronger regularization. In the late training of FL, where the global model is near convergence, the effect of global weight shrinking gradually saturates.
The effectiveness of global weight shrinking stemmed from flatter loss landscapes of the global model as well as the improved local gradient coherence after the critical point.?
UNDERSTANDING FL VIA AN ANALOGY
The update step of 3 can be viewed manipulation of the received local models?:
= = (w" 17p Σ Aggi), s.t. 7 0, 2 0, ||l| =
where g denotes the global model of round t + 1, n2 is the global learning rate, m is the cohort size (i.e., the number of sampled clients), and g denotes the local accumulated model updates of
For concision, in section 4 and section 3, if not mentioned otherwise, we all use CIFAR-10 as dataset and SimpleCINN as model. Experiments more datasets and models are shown in section 5 and Appendix.
Differeen from previous observations (w/o affecting the training dynamics), applying global weight shrinking results in positive local gradient coherence after the critical point and the learning can henefit from it. We recommend the readers to check Appendix for the preliminary of federated leamnig..
2

---

Under review as a conference paper at ICLR 2023
client i starting from the received global model w We assume client trains the model for E' local epochs to derive g/. The set of parameters (7. } in Equation 1 describes the model aggregation process in FL for one communication round, where we refer y as the weight shrinking factor and À as the relative aggregation weights among the clients. FEDAVG is special case by setting 7 =
My = 1, and A, = [
ID.1 Vi E [m]]
The formulation is the analogy to mini-batch SGD: our studied 1) client coherence and 2) global weight shrinking" -hroougghhh respectively optimizing Y and À on a server proxy dataset- refer to weight decay and gradient coherence in mini-batch SGD, respectively. The considered proxy dataset has the same distribution as the global learning objective (i.e. a class-balanced case in this paper), thus the learned aggregation weights (f7, A}) can reflect the contributions of clients and the optimal regularization factor towards this global objective. By connecting the learned weights and the training dynamics, we can know the roles of client heterogeneity and local updates in different learning periods. We use relatively large proxy dataset (2000 samples in CIFAR- 10 with balanced class distributions) for exploration purposes only in section 3 and section 4, while in section 5, we test our proposed FEDAWO on small proxy datasets (100 samples in CIFAR-10).
We review the insights of mini-batch SGD below and leverage them to better understand the training dynamics of FL. later. For other related works, please refer to Appendix B.
Remark 1 (Weight decay) Insights in mini-batch SGD can be detailed as
The optimal weight decay factor is approximately inverse to the number of epochs. and the impor-
lance of applying weight decay diminishes when the training epochs are relatively long (Loshchilov & Hulter, 2018; Lewkowyez & Gur-Ari, 2020; Xie et al., 2020).
The effectiveness of weight decay may be explained by the caused (1) larger effective learning rate (Zhang et al., 2018; Wan et al., 2021), and (2) flatter loss landscape (Lyu et al., 2022).
Gradient coherence. Gradient coherence, sample coherence, is crucial technique for under- standing the training dynamics of mini-batch SGD in centralized learning (Chatterjee, 2019; Ziclinski et al. 2020; Chatterjee & Zielinski, 2020; Fort al., 2019). The gradient coherence measures the pair-wise gradient similarity among samples. If they are highly similar, the overall gradient within a mini-batch will be stronger in certain directions, resulting in a dominantly faster loss reduction and better generalization (Chatterjee, 2019; Zielinski et al., 2020; Chatterjee & Zielinski, 2020).
Remark (Gradient coherence) The critical period exists in mini-batch SGD, captured by the gradient coherence: the low coherence in the early training phase damages the final generalization performance, no matter the value of coherence controlled later (Chatterjee & Zielinski, 2020).
CLIENT COHERENCE
3.1 BASIC CONCEPT AND FORMULATION
Inspired by gradient coherence in mini-batch SGD, we study client coherence in FL, i.e., the local gradient coherence of clients' model updates in FL. In addition this, the FL has another unique aspect of coherence, namely heterogeneity coherence.
Local Gradient Coherence. The gradient coherence in mini-batch SGD is at the data sample level. Analogously, we find similar conclusions at the client level in FI., where aggregating similar local gradients among clients will produce a stronger global gradient, improving generalization and vice versa. We deduce the gradient coherence in mini-batch SGD and local gradient coherence in FL under unified equation below:
AC = L(w' - mg") - L(w") 20 -1 (R' g²) 7
A
Σ
+
Σ N
Xedy g)) I).
Equation 2 is a Taylor expansion of the loss function within one update. In mini- batch SGD, t refers to the iteration step, m the batch size, and g is the gradient of sample i at iteration t.
WWe use the word "shrink" instead of "decay" as it shrinks the global model rather than decaying the model by subtracting a decay term (used in traditional weight decay). Similar "shrink" can be found in Li ct al. (2020c).
3

---

Under review as a conference paper at ICLR 2023
141
Pataay Londl Coddet
Ortiial
San cllattrned clants
34
ofbaleooo ciert Gribiaal poit
Sam
475
930
572
Fiche Intellnn
525
900
193
150
200
Communicaion Rounds
Communicaion Tounce
Conmuncation Hounds
Figure 1: Training dynamics of attentive AWO in terms of local gradient coherence. Clients indexed 0-9 have balanced class distributions and 10-19 are imbalanced. Left: Local gradient coherence. Middle and Right: The performance and aggregation weights of attentive AWO and early- attentive AWO.
Usually, there is no weighted averaging in a mini-batch, so Vi E [m], X 1. In FL, t refers to the communication round, w' refers to the global model on the server at round I, m is the cohort size, g. denotes the local gradient of client i at round l, and X; the aggregation weight of client i. In Equation 2, cos(gi gj) means the cosine similarity of the two gradients, (/c.8)//el lle) Assuming all gradients have bounded norms that Vi, gg/ S c. The cosine similarity among gradients indicates the coherence, from Equation 2, if the gradients have larger cosine similarity, it will have larger descent in the loss and improve the global generalization. In this paper, we focus on the local gradient coherence among clients during FL training. We borrow the cosine stiffness definition (Fort et al., 2019) to quantify the local gradient coherence in FL as follows.
Definition The local gradient coherence of two clients and j at round t is defined by the cosine similarity of their local updates sent to the server, as: Ga,)) cos(g].g). gz)).
The overall local gradient coherence of a cohort of clients round t is defined by the weighted cosine similarity of all clients' local updates sent to the server, as: C'choot = m Eig.zt Azd3 cos(g', (g',)") 8'))
FL assumes multiple local epochs in each client, and clients usually have heterogeneous data; therefore, the local gradients of clients are usually almost orthogonal, that is to say, they have low coherence. This phenomenon observed in Charles et al. (2021), but did not dig deeper to examine the training dynamics of FL. In this paper, we calculate the local gradient coherence in each round and find critical point exists in the process (Figure and Figure 7).
Heterogeneity Coherence. Heterogeneity coherence refers the distribution consistency between the global data and the sampled one (i.e. data distribution of a cohort of sampled clients) in each round. The value of heterogeneity coherence is proportional to the IID-ness of clients as well as the client participation ratio; the higher, the better. We define heterogeneity coherence as follows.
Definition 2 Assuming there are N clients and the cohort size is m. For given cohort of clients, the heterogeneity coherence is sim( Deohort, D), where Dechort = Eie(mi A.D., D Σソ= X,D; and sim is the similarity of two data distributions.
3.2 ATTENTIVE AGGREGATION WEIGHT OPTIMIZATION AND TRAINING DYNAMICS
Vanilla FEDAVG only considers data sizes in clients' aggregation weights A. However, client heterogencity is also crucial. Clients with different heterogeneity degrees have different importance in client coherence, thus playing different roles in training dynamics. A three-node toy example is shown in Figure 11 in Appendix. The optimal À is off the data-sized when clients have the same data size but different heterogeneity degrees. To study client coherence further, we propose attentive aggregation weight optimization (attentive AWO) to learn the optimal aggregation weights (i.e. A) on a proxy dataset. By connecting the optimal weights and the client coherence, we can know the roles of different clients in different learning periods. Attentive AWO conducts the model updates as in Equation! I, where {7, }} is defined as
- 1,),,),,, where X" arg min Cprory - To Σ Xag s.t. A 2 O, |AI| 1.
(3)
1) Critical point exists in terms of local gradient coherence. To study the role of client data heterogeneity in local gradient coherence, we experiment on both balanced and imbalanced clients. Clients indexed 0-9 have balanced class distributions and 10-19 are imbalanced (the data distribution of clients is in Figure 12 of Appendix). We set full participation to engage all clients in each round. The results are demonstrated in Figure 1. It is found that in the first couple of rounds, the coherence
4

---

Under review as a conference paper at ICLR 2023
Round 41
Round 41
60
0.50
0.50
0.25
FedAvg
Proporiion
0.25
Allentive AWO
teel sccuracy (9)
50
0.00
0.00
40
01 236456
01 2
26366678
Class
Class
30
Round 97
Round 97
0.50
Properton
Propertin
20
FedAvg
0.25
0.25
Attentive AWO with SWA
Attentive ANG wio SWA
10
0.00
0.00
1234455
3456788
100
150
Class
Class
Communication Rounds
Figure 2: Left: Heterogeneity coherence of class distribution within a cohort. Right: Test accuracy curves.
is dominant and positive, thus the test accuracy arises dramatically. and most generalization gains happen in this period. The critical point at the round that the coherence is near zero. After the critical point, the test accuracy gain is marginal, and the coherence is kept negative but close to zero. 2) Assigning larger weights to clients with larger coherence before the critical point can improve overall performance. From the Left of Figure I, it is clear that before the critical point, the coherence among balanced clients is much more dominant than others, and it reveals that clients with more balanced data have more coherent gradients." Intuitively, if we assign the balanced clients with larger weights before the critical point, it can boost generalization: in Equation 2, larger À assigning to clients with larger cosine similarity can reduce the loss more significantly.
Interestingly, we find attentive AWO proves our hypothesis: it raises the weights of balanced clients highly in the first few rounds; especially in the first two rounds, it nearly assigns all weights the balanced clients. From the Left of Figure 1, after the point, clients have negative averaged coherence, and their mutual coherence is uniformly small, so the coherence gap between the balanced and imbalanced clients is not obvious. In this scenario, we believe that the coherence of clients matters only before the critical point. To verify this, we adopt early stopping near the critical point when conducting attentive AWO, that is, before the stopping round, we use learned weights to generate global models and then use data-sized weights afterwards. Results show that the early-stopped attentive AWO has comparable performance after the critical point, indicating the training period before local gradient coherence reaches zero is much more vital. This phenomenon is insightful for the community to design effective algorithms for learning critically in early training.
3) Improving heterogeneity coherence within cohort can boost performance. In the scenarios with partial client participation in each round, the selected clients have inconsistent sum objective with the global objective, in other words, the heterogeneity coherence is low (Definition 2). More specifically, for the class-imbalanced setting, due to the local heterogeneous data, the sum local data of the randomly selected participating clients are extremely itussiimallneeed whilst the sum data of all clients are class-bat balanced. We notice that attentive AWO can improve heterogeneity coherence by dynamically adjusting the AW among clients. We visualize the weighted class distributions within cohort in Figure 2, and it shows attentive AWO learns weights to make the class distributions more balanced. The test accuracy curves demonstrate the significant performance gain compared with FEDAVG, and we notice that attentive AWO with SWA performs better by seeking a more generalized minimum in the aggregation weight hyperplane. Due to the space limit, we include more analysis about client coherence in subsection C. of Appendix C.
GLOBAL WEIGHT SHRINKING
4.1 GLOBAL WEIGHTT SURINKING AND ITS IMPACTS ON OPTIMIZATION
As stated in Equation 1, setting 7 < results in the global weight shrinking regularization. Table 1 and the accuracy curves of different y in Figure 3 report the results 10 with different y It can be observed that the global weight shrinking may improve generalization, depending on the choice of 7. Moreover, optimal 7 may exist, and a different value (either smaller larger) will
'This also reveals why FL performs better in IID settings than NonIID: the clients' gradients in IID settings more coherent, but the ones in the NonIID usually diverge.
'Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) is an effective technique to make simple averaging of multiple points along the trajectory of optimization, with cyclical learning rate. It leads to better generalization performance as well as flatter minimum DNNs.
5

---

Under review as a conference paper at ICLR 2023
74
74
Fattin iD
73
Fasting in Norl
74
With ataplii CMS
79
/0
ID Nam NuriD Nom
of
Top tox
71
K.0.9
K.0.7
-No NorID
-H- ID
0.0.6
2.0
68
K.08
64
60
Commurkation Rounds
Numbor Locd Epecta
Figure 3: Test accuracy curves with different values of 7.
Figure 4: Test accuracy gains
Figure 5: The optimal and the norm of global gradient.
adaptive GWS.
Table I: Impact of fixed y across different architectures in both IID and NonlID settings (E- =5=22 2).
IID (a
100)
NonIID (a
1)
Model
1.0
0.99
0.97
0.95
0.93
0.9
1.0
0.99
0.97
0.95
0.93
0.9
SimpleCNN
65.53
67.60
69.20
69.52
70.16
69.83
65.58
67.04
68.36
68.66
69.28
68.93
AlexNet
74.16
74,80
75.54
75.24
75.25
75.03
73.56
73.83
74.37
74.45
74.40
74.24
ResNet8
75.51
76.64
76.80
77.87
76.80
76.74
75.02
76.06
75.73
77.00
75.04
75.31
have inferior performance. Obviously, if y is inappropriate, it will lead to performance degradation in the late training period. In particular, for smaller y, the degradation comes earlier, even if before the model reaches convergence. This observation is consistent with weight decay in centralized training (Remark 1). More results about the fixed yy can be found in Table 7 in Appendix.
4.2 ADAPTIVE GLOBAL WEIGHT SHRINKING AND TRAINING DYNAMICS
To study the optimal and the underlying impact factor of optimal we realize adaptive global weight shrinking (adaptive GWS) on the proxy dataset. The proxy dataset represents the global learning objective, thus the learned y is the optimal value in each round towards this objective. Adaptive GWS adopts the update in Equation 1, where f7, ) is defined as
7-
A
[Pd [D]
where >- = arg min Eprozy(7 (w' 7o Σ
m
D. BD), S.L. 7 0.
Here, we fix À the data-sized weights. The experimental results are shown in Figure 4. Adaptive GWS can improve the performance of FEDAVG by a large margin in both IID and NonIID settings. We observe that adaptive GWS is more beneficial when the number of local epochs is small.
1) Local epochs and client heterogeneity affect the optimal y. From the figure of Figure 5, optimal Y decreases when the local epoch increases or data become more IID, causing stronger weight shrinking regularization. We consider this due to the balance between optimization and regularization. For larger global gradient, it requires a stronger regularization term. We expand the update as
= (w') Toss)) = - - -1-7)w"
(5)
We refer (1 as the -aaddeet of global weight shrinking regularization and yngg' is the global averaged gradient. We notice that the changes of optimal y is due to the changes of the global gradiens. As shown in the right blue Y-axis of Figure 5, the norm of global gradient || increases when the local epoch increases and data become IID. The larger global gradients have smaller optimal
y (shown in the left green Y-axis) to produce larger weight shrinking graudinn [| - 7) w I
to regularize the optimization. More results regarding how heterogeneity affects the optimal y can be
found in Figure 13 in Appendix.
2) Optimal y increases in the late training. As we discussed subsection 4. 1, GWS with smaller fixed will cause performance degradation in the late training. We reckon the phenomenon is because the model is near convergence and the norm of the global gradient is decaying in the late training. If will not decay and dominate the optimization. Thus, in the late training, regularization should be decaying along with the model reaching convergence (smaller global gradient). Figure 6 further verifies our explanation that it shows the optimal ạ is decaying in the late period. While the norm of the global gradient is decaying, adaptive GWS learns a rising y keep the GWS geeudoggradier decay proportionally. As a result, the ratio of two gradient terms remains steady at around 19 to maintain the balance between optimization and regularization. 3) The mechanisms behind adaptive GWS. We provide two general understandings of how adaptive GWS changes the model parameter. We also study why adaptive GWS can improve generalization.
y is fixed, the regularization pseudo-gradine

---

Under review as a conference paper at ICLR 2023
36.0-
-1
24
-030
of 6968 Poed
32.5
300
273
250
205
13-
Adariv SA5
16
Libeecia of Adathaa03
75
Conmantal
100
Rourds
120
50
Conmuniaaion Tounds
190
135
200
50
Cormuntatian Bounds
100
150
220
Figure 6: Left: Norm of two gradients in adaptive GWS. Right: The optimal y and in adaptive GWS, where r the ratio of the global gradient and the regularization geeudo-gaddien
Figure 7: Training dynamics of adaptive GWS in terms local gra- dient coherence.
General understanding.
Scale invariance. Adaptive GWS learns a dynamic shrinking factor y in each round to shrink the global model's parameter. The reason why the shrunk networks still works comes from the fact of the scale invariance property of DNNs (Li et al., 2018: Dinh al., 2017; Kwon et al., 2021). It means that due to the non-linearity of activation functions or the normalization layer in DNNs, a factor rescales the model weights, the function of the model remains similar or the same. We show an intuitive understanding of scale invariance on the Left figure of Figure 8. The final models are rescaled by Y, and the loss function of the adaptive GWS's s final model remains similar while the FEDAVG's final model even has smaller loss when Y 1.
Small model parameters. The shrinking effect in each round can result in smaller model parameters of the final global model. The parameter weight histogram is demonstrated in the Middle figure of Figure 8. The final model of adaptive GWS has more model weights close to zero, nearly twice as many FEDAVG.
Why adaptive GWS can improve generalization.
Flatter loss landscapes. One perspective of explaining the generalization of NN is the flatness of the loss landscape. Previous researches find that flatter curvature in loss landscape can indicate better generalization (Fort & Jastrzebski, 2019; Foret al., 2020; Li al., 2018). In Lyu et al. (2022), they find weight decay of mini-batch SGD can result in flatter landscapes in DNNs with normalization layers. We also observe the similar phenomenon that adaptive GWS improves generalization by seeking flatter minima in FI., shown in the Right figure of Figure 8. Along the hessian eigenvector with maximal eigenvalue, is clear that the model of adaptive GWS has a flatter landscape than FEDAVG, and it also has a smaller loss. Additionally, we also use other flatness metrics based on hessian eigenvalues to compare the loss landscapes during training in Figure 14 in Appendix, and these metrics also show that adaptive GWS can result in flatter curvature with better generalization.
Improving local gradient coherence after the critical point. In section 3, we find that when Y = 1, there exists critical point where local gradient coherence turns positive to negative, and after the critical point, the generalization gain marginal. However, we find if we adopt adaptive GWS, the local gradient coherence is still positive after the critical point, and the model can still benefit from the coherent gradients. We demonstrate the local gradient coherence and test accuracies in one figure in Figure 7. Before the critical point, the vanilla FEDAVG and adaptive GWS both have high gradient coherence, so the accuracies rise equally fast. However, after the critical point, the coherence of FEDAVG goes down below zero. Therefore, the generalization has little performance gain afterwards, and the optimization is near saturation. On the contrary, after the critical point, adaptive GWS keeps the coherence above zero. and the global model still benefits from having a larger performance gain beyond FEDAVG. This shows that the shrinking regularization benefits from the long-time optimization after the critical point by making clients' local gradients more coherent.
4) The relation between adaptive GWS and local weight decay. Our proposed adaptive GWS can cause weight regularization from the global perspective, which is analogous to weight decay in mini-batch SGD. Importantly, GWS has a sparse regularization frequency that only changes the model weight in each round, and as result, we find GWS has stronger regularization each time. In GWS, 1 1--1 7 is near 1. whereas the factor of weight decay is about 10-4, The two methods are not conflicted in FL, and we conduct experiments and global weight shrinking on the server simultaneously in Table 2. It is shown that adaptive GWS is compatible with local weight decay and can further improve performance. Local weight decay relies on the hyperparameter of the decay rate, and it needs hyperparameter tuning to search for the most appropriate value in every setting. Instead, adaptive GWS is hyperparameter-ffee free and effective:
implementing weight decay in the local SGD solver
7

---

Under review as a conference paper at ICLR 2023
2.0
1.0
Fod/vg final model
Adactive GWS fna model
120
Foddeg
10
Acaeplaw GWS
1.6
100
1.4
1.2
60
1.0
2
0.8
0.90
0.85
1.00
1.05
1.10 0.00
-0 0222
0.00 0.02
0.04
-0.4
-02
0.0
0.2
0.4
Parmmetrr weight
Pertuffotion
Figure 8: General understanding of adaptive GWS.. Left: Scale invariance property of NN indicates that if the network is rescaled by 7, the function of the model remains similar. Middle: The histogram of final models' weights shows that adaptive GWS makes more model weights close to zero, nearly twice as many as FEDAVG. Right: The Loss landscape is perturbed based on the Top-1 Hessian cigenvector of the final models, which shows that the model with adaptive GWS has flatter curvature and smaller loss.
Table 2: Adaptive GWS with different local weight decay factors (E = 2).
IID (a- 100)
NonlID (a = 1)
Local weight decay
0
5c-5
le-4
5c-4
le-3
0
5e-5
le-4
5c-4
le-3
FEDAVG
66.43
66.20
66.45
67.51
68.66
65.35
65.19
65.77
66.37
67.4
Aduptive GWS
71.47
71.36
71.35
71.44
71.54
70.31
69.93
70.44
70.47
69.99
7 of Adaptive GWS
0.9472
0.9477
0.948
0.9493
0.953
0,9492
0,9499
0.9505
0.9529
0.9561
Table 3: The performance with different per-
Table 4: The Performance of compared methods with differ-
centages corrupted clients (IID, F = 3)
ent model architectures (a = 1, E = 1).
Method 25%
50%
75%
Model
HEDAVG
FEDAWO
FEDAWO (SWA)
PEDAVG
63.40
61.14
58.21
ResNet20
74.11
78.72
78.64
ResNet56
74.22
78.93
79.08
FEDDF
68.73
66.94
66.07
ResNet110
74.50
78.11
79.19
SERVER-FT
63.61
61.24
58.36
WRN56 .4
78.67
79.61
80.70
FEDAWO
67.87
66.67
63.91
FEDAWO (SWA)
68.04
66.95
65.51
DenseNet121
85.13
86.50
87.06
according the results, matter what the local weight decay factors and client heterogeneity, it can adaptively set 7 to maximize the benefit of weight regularization. Another interesting finding is that adaptive GWS can balance the regularization and optimization dynamically; the learned 7 is adaptively adjusted in different local weight decay factors. When the local weight decay is stronger. will have larger y, which means weaker GWS regularization. Due the space limit, we include more analysis about global weight shrinking in subsection C.2 of Appendix C.
FEDERATED AGGREGATION WEIGHT OPTIMIZATION: FEDAWO
Based on the above understandings, propose FEDAWO which combines the adaptive GWS and attentive AWO to optimize y and X simultaneously, defined as
m (7 (== Y", X )"), where >','" arg min Eyrozy 7 (w² Vio > A.g)).. s.L. 7 2 0, A: 2 0, [AI| = 1.
7.A
(6)
The optimization of y and À is non-trivial. First of all, it is unclear how treat the normalization layers. The models with normalization layers contain buffers to calculate the running mean and variance of training data, and naively taking buffers as model parameters and multiplying them with learned 7 will impede optimization. As our solution, we do not aggregate the buffers; instead, we update the buffers on the proxy dataset during AW optimization. Even though the proxy dataset is relatively small (c.g. 100 images in total), the updated buffers still work well to enable the global model with good generalization. Additionally, it is challenging to incorporate with SWA (Izmailov et al., 2018) for better generalization performance, as joint optimization of X and y with SWA has poor performance due to the sensitivity of y on stochastic averaging. To solve this, we adopt an alternative two-stage strategy for SWA variant (implementing it in a reversed order also works), where we first fix À and optimize then we use the learned and fix it to optimize À with SWA. Experiments. We conduct experiments to verify the effectiveness of FEDAWO. Due to the page length limitation, for more information, please refer to Appendix D (details of FEDAWO) and
8

---

Under review as a conference paper at ICLR 2023
Table Top-1 test accuracy (%) achieved by compared FL methods and FedAWO on three datasets with different model architectures (E = 3). Blue/bold fonts highlight the best baseline/our approach.
Dataset
PashionMNIST
CIFAR-10
CHAR-100
NoalID (a)
100
0.1
100
100
MLP
LeNet
ML.P
LeNet
CNN
ResNet
CNN
ResNet
CNN
ResNet
CNN
ResNet
FEDAVG
89.29
90.54
85.11
88.08
65.78
74.57
60.13
46.04
25.74
27.49
27.74
24.92
FEDPROX
87.68
89.77
84.33
87.01
.6
68.51
60.48
48.84
9.49
27.15
12.52
23.73
FEDDYN
88.47
89.92
77.68
72.68
66.1
76.62
41.53
35.77
24.44
32.18
22.67
29.00
FEDDF
57.38
28.52
24.52
23.10
FEDBE
86.22
89. 14
79.12
85.96
69.88
77.94
59.84
52.86
28.38
27.73
25.41
23.74
SERVER-FT
89.09
90.56
66.83
74.73
60.43
47.59
25.37
26.14
24.33
23.03
FEDAWO
88,51
90.66
86.30
88.26
70.17
80.46
62.46
52.83
32.51
33.17
32.30
24.84
FEDAWO (SWA)
88.27
90.51
86.89
88.18
69.9
79.55
62.12
57.08
32.39
33.17
32.27
25.31
68
60
666
sccumcy
65
Tea-1 1est acourrry
64
60
Fesiiio
Foeewoo gree)
FRAKINGSONN)
FooBi
40
FAIDF
Fecthos
THEZE
0.1
0.2
0.3
0.4
0.5
10
20
50
100
200
400
Paricipation Ruño
Number of Sampleaii Proayy Dulaeee
Figure 9: The performance with different partie- ipation ratios (a = 1, E == 3).
Figure 10: The performance with different sizes of the proxy dataset (a = 0.1, E = 3).
Appendix E (implementation details). We mainly compare FEDAWO with other server-side methods, i.e. FEDDF (Lin et al., 2020) and FEDBE (Chen & Chao, 2021a), that also require proxy dataset for additional computation. These two methods conduct ensemble distillation on the proxy data to transfer knowledge from clients' models the global model. We add SERVER-FT as baseline for simply finetuning global models on the proxy dataset. Besides, we implement client-side algorithms FEDPROX (Li al., 2020b) and FEDDYN (Acar et al., 2020).
Results. Different datasets: As Table FEDAWO outperforms baselines across different datasets and models in both IID and NonlID settings. Compared with FEDDF, FEDBE and SERVER-FT
FEDAWO can better utilize the proxy dataset. Different participation ratios: From Figure 9, FEDAWO performs well under partial participation. Different sizes and distributions of proxy dataset: In Figure 10, the server-side baselines are sensitive to the size of the proxy datasct that too small or too large proxy set will cause overfitting. However, FEDAWO is also effective under an extremely tiny proxy set and benefits more from larger proxy set due to accurate aggregation weight optimization. We report the results of different distributions of the proxy dataset in subsection C.4 of Appendix C, which show that FedAWO still works when there exists a distribution shift between the proxy dataset and the gloabl data distribution of clients. Different architectures: We test FEDAWO across wider and deeper ResNet and other architecture, such as DenseNet (Huang et al., 2017), in the Table 4. shows FEDAWO is effective across different architectures, and it performs well even when the network goes deeper or wider. Robustness against corrupted clients: Another advantage of FEDAWO is that it can filter out corrupted clients by assigning them lower weights. We generate corrupted clients by swapping two labels in their local training data. As in Table 3, FEDAWO has great performance even corrupted clients exist, and it is as robust as the ensemble distillation methods, like FEDDF, when using the same proxy dataset.
CONCLUSION
We introduce a novel global optimization framework to understand the training dynamics in FL, by making an analogy to mini-batch SGD in centralized learning. Leveraging the framework, we identify some interesting findings to unveil the FL training dynamics in terms of global weight shrinking regularization and client coherence. Based on the findings, devise simple but effective method FEDAWO, for server-side aggregation weight optimization. can realize both adaptive global weight shrinking and attentive aggregation weights to better improve the generalization of the global model.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. International Conference on Learning Representations, 2020.
Zachary Charles, Zachary Garrett, Zhouyuan Huo, Sergei Shmulyian, and Virginia Smith. On large-cohort training for federated learning. Advances in neural information processing systems, 34:20461-2044 -0475,, 2021.
Satrajit Chatterjee.
Coherent gradients: An approach to understanding generalization in gradient descent-based optimization. In International Conference on Learning Representations, 2019.
Satrajit Chatterjee and Piotr Zielinski. Making coherence out of nothing all: measuring the evolution of gradient alignment. arXiv preprini arXiv:2008.01217. 2020.
Hong- Houg-You Chen and Wei-Lun Chao. Fedbe: Making bayesian model ensemble applicable federated learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL https: http::://ppeeeevviecce net forum? idd=ottpE6gkjjt
Hong-You You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In International Conference Learning Representations, 2021b.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference machine learning, pp. 1309-1318. PMLR, 2018.
Jiawei Du Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural networks. arXiv preprint 10.0314,, 2021.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance linear mode connectivity neural networks. International Conference Learning Representations, 2022.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412. 2020.
Stanislav Fort and Stanislaw Jastrzebski. Large scale structure of neural network loss landscapes. Advances in Neural Information Processing Systems, 32, 2019.
Stanislav Fort, Pawel Krzysztof Nowak, Stanislaw Jastrzebski, and Srini Narayanan. Stiffness: A new perspective on generalization in neural networks. arXiv preprint arXiv:1901 000491 2019.
Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient-based hyperparameter optimization. In International Conference Machine Learning, pp. 1165-1173. PMLR, 2017.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems, 31, 2018
Pengfei Guo, Dong Yang. Ali Hatamizadeh, An Xu, Ziyue Xu, Wenqi Li, Can Zhao, Daguang Xu, Stephanie Harmon, Evrim Turkbey, et al. Auto-fedrl: Federated hyperparameter optimization for multi-institutional medical image segmentation. arXiv preprint 2rXiv:2003.06338 2022.
Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv / 050.00233, 2(7), 2015.
10

---

Under review as a conference paper at ICLR 2023
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
Izmailov, AG Wilson, D Podoprikhin, Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876-885, 2018.
Stanistaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. On the relation between the sharpest directions of dnn loss and the sgd step length. arXiv preprint wrXiv:1807.050333 2018.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2019.
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint urXiv:2002.09577 0.0.02 2020.
Hadi S Jomaa, Josif Grabocka, and Lars Thiemm Hyp-rl: Hyperparameter optimization by reinforcement learning. arXiv preprint arXiv:19006 1.06.11522 11527, 2019.
Leonid V Kantorovich. On the translocation masses. Journal of mathematical sciences, 133(4):
1381-1382, 2006.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pp. 5132-5143, PMLR, 2020.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, 2017.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-awan minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pp. 5905-5914, PMLR, 2021.
Aitor Lewkowyez and Guy Gur-Ari. On the training dynamics of deep networks with 1_2 regulariza- tion. Advances in Neural Information Processing Systems, 33:4790-4799, 2020.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
Shuangtong Li, Tianyi Zhou, Xinmei Tian, and Dacheng Tao. Learning to collaborate in decentralized learning of personalized models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9766-9775, 2022a.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Process. Mag., 3703):50-60, 2020a. doi: 10. 109/ NNSP22000.295749. URL htlLhtts::// //ottp:::/dd.... org/10 1109 /MSP 2020 2975749.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020b.
Xiang Li, Shuo Chen, and Jian Yang. Understanding the disharmony between weight normalization family and weight decay. Proceedings of the AAAI Conference on Artificial Intelligence. volume 34, pp. 4715-4722. 2020c.
Zexi Li, Jiaxun Lu, Shuang Luo, Didi Zhu, Yunfeng Shao, Yinchuan Li, Zhimeng Zhang, and Chao Wu. Mining latent relationships among clients: Peer-to-peer federated learning with adaptive neighbor matching. arXiv preprint arXiv:2203. 12285, 2022b.

---

Under review as a conference paper at ICLR 2023
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
lyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2018.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normaliza- tion layers: Sharpness reduction. arXiv preprint arXiv:2206.0708 2022.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-barl baadie hyperparameter optimization through reversible learning. In International conference machine learning, pp. 2113-2122.
PMLR, 2015,
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effiiieenn learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Hesham Mostafa. Robust federated learning through representation matching and adaptive hyper- parameters. arXiv preprint 0rXiv:1912.13009 2019.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045-22055, 2020.
Tiffany Vlaar and Jonathan Frankle. What can linear interpolation of neural network loss landscapes tell us? arXiv preprint arXiv:2/06.16004 2021.
Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learning dy- namies of normalized neural network using sgd and weight decay. Advances in Neural Information Processing Systems, 34:6380-6391, 2021.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002. 06440, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623, 2020b.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, al. A field guide to federated optimization. arXiv preprini arXiv:2107.069)) 2021.
Yingda Xia, Dong Yang, Wenqi Li, Andriy Myronenko, Daguang Xu, Hirofumi Obinata, Hitoshi Mori, Peng An, Stephanie Harmon, Evrim Turkbey, et al. Auto-fedavg: learnable federated averaging for multi- iulliiinstitutionall medical image segmentation. arXiv preprint arXiv:210.110195, 2021.
Zeke Xie, Issei Sato, and Masashi Sugiyama. Understanding and scheduling weight decay, arXiv preprint arXi::2011.11152 2020.
Gang Yan, Hao Wang, and Jian Li. Critical learning periods in federated learning. arXiv preprint 1rXiv:2109.056/3, 2021.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pp. 581-590. IEEE, 2020.
Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett. Gradient diversity: a key ingredient for scalable distributed learning. In International Conference Artificial Intelligence and Statistics, pp. 1998-2001 PMLR, 2018
Honglin Yuan, Warren Richard Morningstar, Lin Ning, and Karan Singhal. What do mean by generalization in federated learning? In International Conference on Learning Representations, 2021.
12

---

Under review as a conference paper at ICLR 2023
Guodong Zhang, Chaoqi Wang. Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference on Learning Representations, 2018.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In International Conference on Machine Learning, pp. 12878-12889, PMLR, 2021.
Piotr Zielinski, Shankar Krishnan, and Satrajit Chatterjee. Weak and strong gradient directions: Explaining memorization, generalization, and hardness of examples scale. arXiv preprint arXiv:2003.07422, 2020.
13

---

Under review as a conference paper at ICLR 2023
Appendix
In this appendix, we provide details omitted the main paper and more experimental results and analyses.
Appendix A: preliminary of federated learning (cf. section and section 2 of the main paper).
Appendix B: related works (cf. section I of the main paper).
Appendix C: more experimental results and analyses (cf. section 3, section 4 and section 5 of the main paper).
Appendix D: additional details of FEDAWO (cf. section 3, section 4 and section 5 of the main paper).
Appendix E: details of experimental setups (cf. section 5 of the main paper).
A PRELIMINARY OF FEDERATED LEARNING
Federated learning usually involves a server and n clients to jointly learn a global model without data sharing, which is originally proposed in (McMahan et al., 2017). Denote the set of clients by S, the labeled data of client by D; {xx) M;)} and the parameters of the current global model by w FL starts with client training in parallel, initializing each clients' model wł with w!. FL is more communication-cfficii
than the conventional distributed training, that it assumes the clients train the models for epochs (the full data) instead of iterations (the mini-batch data) between the communications the server. The number of local epochs is denoted E.
In each local epoch, clients conduct SGD update with local learning rate vJr, each SGD iteration shows as
Client training:
w! f- w mn VE w(), for k 1,2, K,
where E is the loss function and Bg is the mini-batch sampled from D; at the kth iteration. After the client local updates, the server samples clients for aggregation. The client i's pseudo gradient of local updates is denoted as g = w w.. Then, the server conducts FEDAVG to aggregate the local updates into new global model.
D. [
Weighted Model aggregation:
w
1+1 g
w!
A
Vie e [m].
(8)
1=1
With the updated global model w it then starts the next round of client training. The whole procedure of FL therefore iterates between Equation and Equation for T communication rounds. We assume the sum of clients' data as D Ules Dj. The IID data distributions of clients refer to that each client's 's distribution D, is IID sampled from D. However, in practical FL scenarios, heterogeneity exists among clients that their data are NonIID with each other. Fach client may have different data distributions in the input (e.g. image distribution) or output (e.g. label distribution). In this paper, we study how the number of local epochs and the clients' data heterogeneity affect the training dynamics in terms of client coherence and global weight shrinking.
B RELATED WORKS
B.1 MODEL AGGREGATION IN FEDERATED LEARNING
Model aggregation in federated learning. Model aggregation weights should be calibrated under asynchronous local updates. FEDNOVA (Wang al., 2020b) is proposed tackle the objective inconsistency problem caused by asynchronous updates, it theoretically shows that the convergence will be improved if the aggregation weights are normalized by the numbers of local iterations. However, it does not take the heterogeneity degree of clients into account, which is also a key factor that affects the generalization of the global model. In Chen & Chao (2021a), the authors point out that due to heterogeneity, the best performing model will shift way from FEDAVG, but they do not give insights on how to adjust aggregation weight to approximate the best model, they use Bayesian ensemble distillation method to prove the generalization of the global model instead. To solve the misalignment of neurons in FL with NN, FEDMA (Wang et al., 2020a) is proposed: FEDMA A constructs the shared global model layer-wise by matching and averaging hidden elements
14

---

Under review as a conference paper at ICLR 2023
with similar features extraction signatures. Besides, optimal transport (Kantorovich, 2006) can be adopted in layer-wise neuron alignment in the process of model fusion (Singh & Jaggi, 2020). These previous works improve the global model performance by layer-wir wiier-wie alignment, but they are complex and computation-expensivyy and they can not be applied under the traditional weighted aggregation scheme. What's more, distillation (Hinton al., 2015; Zhu et al., 2021; Lin et al., 2020) can enable robust knowledge transfer and ensemble distillation can be used to finetune the global model for better generalization (Lin et al., 2020; Chen & Chao, 2021a).
B.2 GENERALIZATION AND TRAINING DYNAMICS OF NEURAL NETWORK
Loss landscape of neural networks and generalization. Neural networks (NN) are highly non- convex and perepaaraettrized and visualizing the loss landscape of NN (Li et al., 2018; Vlaar & Frankle, 2021) helps understand the training process and the properties of minima. There are mainly two lines of works about the loss landscape of NN. The first one is the linear interpolation of neural network loss landscape (Vlaar & Frankle, 2021; Garipov et al., 2018; Draxler al., 2018), plots linear slices of the landscape between two networks. In linear interpolation loss landscape, mode connectivity (Draxler et al., 2018; Vlaar & Frankle, 2021; Entezari et al., 2022) is referred to as the phenomenon that there might be increasing loss on the linear path between two minima found by SGD, and the loss increase on the path between two minima is referred to as (energy) barrier. is also found that there may exist barriers between the initial model and the trained model (Vlaar & Frankle, 2021). 1). The second line concerns the loss landscape around a trained model's parameters (Li et al., 2018). It is shown that the flatness of loss landscape curvature can reflect the generalization (Foret et al., 2020; Izmailov et al., 2018) and top hessian eigenvalues can present flatness (Yao et al., 2020; Jastrzebski
al., 2018). Networks with small top hessian eigenvalues have flat curvature and generalize well. Previous works seek flatter minima for improving generalization by implicitly regularizing the hessian (Foret et al., 2020; Kwon et al., 2021; Du et al., 2021).
Critical learning period in training neural networks. Jastrzebski al. (2019) found that the early phase of training of deep neural networks is critical for their final performance. They show that break-even point exists on the learning trajectory, beyond which the curvature of the loss surface and noise in the gradient are implicitly regularized by SGD. They also found that using a large learning rate in the initial phase of training reduces the variance of the gradient and improves generalization. In FL, Yan et al. (2021) discovers the early training period is also critical to federated learning. They reduce the quantity of training data in the first couple of rounds and then recover the training data, and it is found that no matter how much data are added in the late period, the models still cannot reach a better accuracy. However, it did not further study the role of client heterogeneity in the critical learning period while we examine it by local gradient coherence.
B.3 FEDERATED HYPERPARAMENT OPTIMIZATION
Current federated learning methods struggle in cases with heterogeneous client-side data distributions which can quickly lead to divergent local models and a collapse in performance. Careful hyper- parameter tuning is particularly important in these cases. Hyper-parameters can be optimized using gradient descent to minimize the final validation loss (Maclaurin et al., 2015; Franceschi et al., 2017). Moreover, hyper-parameters can be optimized based on reinforcement learning methods (Guo et al., 2022; Jomaa et al., 2019; Mostafa, 2019). However, in this paper, optimization aggregation weights not our main novelty. Instead, we focus on leveraging this toolbox on our well-designed but unexplored scenarios and examining the crucial training dynamics in FL in principled way.
B.4 MOST RELEVANT WORKS TO FEDAWO
We notice that two related works also optimize the aggregation weight (AW) by gradient descent. The first is -EDO-FEDAVG (Xia et al., 2021). which optimizes AW on different institutional medical data to realize personalized medicine. AUTO-FEDAV adopts Softmax and Dirichlet functions as the base functions in optimizing AW. The second is a decentralized FL. algorithm called L.2C (Li et al., 2022a). It adopts peer-to-peer (P2P) communication protocol and uses the local dataset of each client to optimize the collaborative weights with other clients. L2C assumes that different clients have various learning tasks, so it learns an adaptive weight for personalized collaborative learning. However, these two works all assume a normalized AW, whose L1 norm equals 1, so they did not devise the global weight shrinking strategy for training more generalized global model. Also, they are for specific application scenarios, like medical AI or P2P FL, and they did not introduce the aggregation weight optimization method to general FL. and understand the training dynamics from
15

---

Under review as a conference paper at ICLR 2023
the learned AW. Additionally, these previous methods are all about personalization while we focus on generalization from the global perspective.
C MORE RESULTS AND ANALYSES
C.I CLIENT COHERENCE
The relationship with gradient diversity. The conclusion of gradient diversity (Yin et al., 2018) is opposite the one of gradient coherence. Gradient diversity argues that higher similarities between workers' gradients will degrade performance in distributed mini- SGD, while gradient coherence claims that higher similarities between the gradients of samples will boost generalization (Yin et al., 2018; Chatterjee, 2019). Moreover, gradient diversity is somewhat controversial. As argued in the line of works about gradient coherence (Chatterjee & Zielinski, 2020; Chatterjee, 2019), the manuscript of gradient diversity did not explicitly measure the gradient diversity in the experiments (or further study its properties): only experiments on CIFAR-IO can be found where they replicate 1/r of the dataset r times and show that greater the value of r less the effectiveness of mini-batchin to speed assumption in the theorem of gradient diversity
up. Apart from this controversy, the strongly-conve (Yin et al., 2018) may make it weaker to generalize its conclusions in neural networks while we are studying the empirical properties in FL with neural networks. Taking the above statements into consideration, gradient diversity may be infeasible in our settings.
The relationship with gradient diversity. There are some works (Karimireddy et al., 2020; Li et al., 2020b) taking the bounded gradient dissimilarity assumption to deduce theorems. In their assumptions, they bound the gradient sum or gradient norm, but we use the cosine similarity to study how the clients interplay with each other and contribute to the global. So the perspectives are quite different. Additionally, there are previous works [] in FL that use cosine similarity of clients' gradients improve personalization, however, we focus the training dynamics in generalization, and one of our novel findings is we discover a critical point exists and the periods that before or after this point play different roles in the global generalization.
Heterogeneity also affects the optimal aggregation weight. We set up a three-node toy example on CIFAR-10 10 by hybrid Dirichlet sampling shown in Figure 11. We first sample client O's data distribution by Dirichlet sampling according to 01, then we sample data distributions for clients and 2 on the remaining data with ag. We set up three settings with different a1, 02 and illustrate the data distributions on the Left column in Figure 1. In the example, the aggregation weights (AWs) are [oo. A1. A.g, we regularize the weights as Xg + Ag + Ao = 1 which is a plane that can be visualized in 2-D. We uniformly sample points on the plane to obtain global models with different AW and compute the test loss, and then the loss landscapes on the plane can be visualized. We implement FEDAVG for 100 rounds and record the loss landscape and the optimal weight on the loss landscape in each round, then we illustrate the loss landscape of round 10 on the Middle column and the optimal weight trajectory on the Right column of Figure 11.
In these settings, clients have different heterogeneity degrees: in the first setting, client has a balanced dataset while the data of clients 1 and 2 are complementary; in the second and third settings, clients and 2 have the same data distribution, which differs from the client O's. From Figure 11, it is evident that the weight of FEDAVG is biased from optimal weights when heterogeneity degrees vary in clients, we can draw the following conclusions: (1) optimal weight can be viewed as a Gaussian distribution in the aggregation weight hyperplane; (2) the mean of the Gaussian will drift towards to the directions where data are more inter-heterogeneous (for instance, in the third setting, client O's major classes are 2, 3 and 8 while client and have rare data on these classes, so client O's contribution is more dominant); (3) the variance of the Gaussian is larger in inter-homogeneous direction and is smaller in mnter-hetroogenouu direction (the variance along the client 1-client direction is large in the second and third settings, because the two clients have inter-homogencous data; opposite phenomenon is shown in the first setting, where client and have inter-heterogennoous data); (4) the flatness of loss landscape on aggregation weight hyperplane is consistent with the variance of the Gaussian, which means the directions with more significant variance will have flatter curvature in the landscape. From our analysis, it is clear that clients' contributions to the global model should not be solely measured by dataset size, and the heterogeneity degree should also be taken into account. And we observe that in more heterogeneous environment, the loss landscape is sharper, which means the bias from optimal weight will cause more generalization drop. In other words, in heterogeneous environment, appropriate aggregation weight matters more.
16

---

Under review as a conference paper at ICLR 2023
alphal 100,0, siphaz sipha,
0.1
Cloes
62
3.5
0.0
2.0
CleTE Me
alphal 1.0, alpha2 100.0
18
0
68
Clert
alshal 0.1, alpha2 100.0
Opteraage begetes
14
60
Figure 11: Heterogeneity also affeets the optimal aggregation weight. A three-node toy example on CIFAR-10 is shown. Left: Data distribution of each client, note that each client has the same dataset size. Middle: Loss landscape on the plane of aggregation weight, it is noticed that FEDAVG off the optimal and the landscape has various flatness in different directions. Right: optimal weight trajectory during training, the optimal weights are biased from FEDAVG.
Visualization of the hybrid NonIID setting of Figure 1. We visualize the hybrid NonlID setting of Figure in Figure 12. We take on and ar = 0.1, so the first 10 clients (indexed 0-9) have class-balanced data while the last 10 clients (indexed 10-19) have class-imbalanced data.
Factors
E =
=
E 1&E 5
Dataset size (DS)
-0.098
0.21
0.035
Heterogeneity degree (HD)
0.41
0.024
0.35
HD
0.26
0.17
0.31
Table Pearson correlation coefficient analysis of AW. Heterogeneity degree is calculated as the reciprocal of the variance of class distribution for each client. We the accumulated weights during the training clients' AW.
Data size or heterogeneity? A correlation analysis. Data size and heterogeneity all affects clients' contributions to the global model, but which affects it most? As in previous literature, the importance depicted by the dataset size that clients with more data will be assigned larger weights. According to the analysis in Figure 11, the importance of weight may be associated with the heterogeneity degrees of clients. To explore which factor is more dominant in the AW optimized by attentive AWO,
17

---

Under review as a conference paper at ICLR 2023
alphal 10.0, alpha2 0.1
Number samples per
-3
Client idx
Figure 12: Data distribution of Figure
Table 7: More results about fixed y across different architectures various NonlID seftings.
1.0
0.99
0.97
0.95
0.93
0.9
Model
10
SimpleCNN
65.96
67.19
69.41
69.81
69.69
69.59
AlexNet
73.9
74.43
74.96
75.12
75.33
74.06
ResNet8
76.26
75.63
76.92
77.23
76.9
76.61
Model
0.5
SimpleCNN
65.78
66.59
67.93
68.85
68.75
68.25
AlexNet
73.72
73.06
73.89
73.98
73.6
73.33
ResNet8
73.4
73.93
75.39
74.12
73.66
73.46
Model
0.2
SimpleCNN
63.52
64.68
63.72
65.82
65.4
64,97
AlexNet
68.41
70.46
70.87
70.74
70.58
69.42
ResNet8
71.85
70.96
72.76
72.04
71.25
62,32
Model
0.1
SimpleCNN
60.57
61.22
61.83
62.05
62.05
60.85
AlexNet
66.18
65.25
64.74
64.23
64.16
61.24
ResNet8
63.89
60.55
61.38
59.23
58.76
39.85
we have made a Pearson correlation coefficient analysis in Table 6. Results show that dataset size is more dominant when the local epoch is large; otherwise, the heterogeneity degree. This phenomenon intuitive: when the local epoch increases, clients with larger dataset will have more local iterations than others (Wang al., 2020b), their updates more dominant. In the cases where the local epoch is small, clients' updates are of similar volumes, here the updates' directions are much more important since balanced clients are prone have stronger coherence, and their AWs are larger in model aggregation. We combine two factors by multiplication, and the result shows that the combined indicator more dominant when the two cases are mixed.
C.2 GLOBAL WEIGHT SHRINKING
Fixed 7. We add more results about global weight shrinking experiments with fixed y as Table It is found that when data are more NonIID, fixed will cause negative effects; this is more dominant when a 0.1 and the models are AlexNet or ResNet8.
Adaptive GWS with global learning rate. We conduct experiments with the adaptive GWS under different global learning rates for both IID and NonIID settings. We train SimpleCNN CIFAR10
18

---

Under review as a conference paper at ICLR 2023
Table 8: The performance of adaptive GWS under different global learning rates.
III (a 100)
NonlID (a I)
Global learning rate
0.5
1.5
0.5
1.5
FedAvg
69.15
68.18
64.35
68.71
67.09
64.00
Adaptive GWS
71.45
71.98
71.13
69.65
71.02
71.04
ς of Adaptive GWS
0.986
0.974
0.963
0.991
0.979
0.967
with 1 local epoch, and the results are reported in Table 8. It can be observed that in both IID and NonlID settings, small global server learning rate can improve FedAvg's performance. In contrast, the larger global learning rate, the smaller the learned (stronger regularization). It is aligned with our insights in the main paper that larger pseudo gradients require stronger regularization. Moreover, Adaptive GWS is robust to the choice of the global server learning rate, especially in IID setting.
Adaptive GWS under various heterogeneity. We show adaptive GWS works under various heterogeneity and visualize 7 and the norm of the global gradient in each setting, as in Figure 13. demonstrates that adaptive GWS can boost performance under different NonlID settings, but has smaller benefit when the system is extremely NonIID (i.e., 0.1). Additionally, according to the Right figure Figure 13, except for the outlier when a = 10, the learned 7 decreases IID, causing stronger weight shrinking effect. We think this is a result of
when data become balance between optimization and regularization. The volumes of global gradients change when the heterogeneity changes. The norm of global gradient increases when data become more IID, and requires smaller 7 to cause stronger regularization.
Fed/vg With adaptive GWS
0.976
DD Norm Globai Gradients
14
Top-1 lost accuracy (18)
Global Gradlents
+5.76
+4.75 +5.53
71
69
0.986
65
63
0.955
+1.85
1,886
69
0.1
0.5
20
50
100
0.5
NonlID
Figure 13: Adaptive GWS under various heterogeneity. Left: Test accuracy gains with adaptive GWS, In all settings, adaptive GWS can bring performance gains. Right: Learned of adaptive GWS in different settings. decreases when data become more IID, causing the stronger weight shrinking effect. This due to the changes the volumes global gradients. The norm global gradient increases when data become more IID, and requires smaller 7 to cause stronger regularization.
More results of general understanding of adaptive GWS. First, we first visualize the norm of model parameter weight during training in the Left figure of Figure 14. Adaptive GWS results in smaller model parameter during training. Second, we use two common metrics to measure the flatness of loss landscape during training as in the Middle and Right figures of Figure 14, and they are the hessian eigenvalue based metrics. The dominant hessian eigenvalue evaluates the worst-case loss landscape, which means the larger top eigenvalue indicates the greater change the loss along this direction and the sharper the minima (Keskar et al., 2017). We adopt the top I hessian eigenvalue and the ratio of top and top 5, which are commonly used as proxy for flatness (Jastrzebski et al., 2020; Fort & Jastrzebski, 2019). Usually, a smaller top 1 hessian eigenvalue and smaller ratio of top 1 hessian eigenvalue and top indicates flatter curvature of NN. As in the figures, during the training, FEDAVG generates global models with sharp landscapes whereas adaptive GWS tends to generate more generalized models with flatter curvatures.
The distribution of We visualize values of all experiments in Figure 5 and Figure 13 in Figure 15. is found that the distribution of r can be approximated into Gaussian distribution with mean around 20.5.
19

---

Under review as a conference paper at ICLR 2023
35
Fodivg
A00
Fedlug
To0 dM00 Hesssion egencaac
Feddug Acdeptize GWS
Norn Pa
Adaathvv GW/S
GWS
300
25
200
20
16
10
100
100
150
200
150
200
Communication Rounds
Communicat on Rounds
Communizatio
Figure 14: More results general understanding adaptive GWS. Left: Adaptive GWS results in a smaller model parameter during training. Middle: Smaller top hessian eigenvalue indicates flatter curvature of NN. The result shows FEDAVG tends to generate sharper global models during training while adaptive GWS seeks flatter networks. Right: The ratio of top hessian eigenvalue and top is another indicator, smaller value means flatter minima.
0.25
0,20
Density
0.15
0.10
0.05
0.00
16
18
20
22
24
26
28
30
Figure 15: Distribution of visualize r values all experiments in Figure and Figure 13 and find that the distribution of can be approximated into a Gaussian distribution with mean around 20.5.
C.3 COMAPRED EXPERIMENTS OF FEDAWO
We add the test accuracy curves show the learning processes of the algorithms and visualize them in Figure 16 (FasionMNISTT) Figure 17 (CIFAR-10), and Figure 18 (CIFAR-100). The curves are according the results Table It shows that FEDAWO surpasses the baseline algorithms in most cases. Besides, FEDAWO is steady in the learning curves and it avoids over-fitting in the late training.
C.4 DISTRIBUTION SUIFT OF PROXY DATASET
One may think it is strong assumption that the distribution of the proxy dataset is the same as the global distribution. To reflect the superiority of FedAWO, we consider two challenging scenarios. We train ResNet20 on CIFAR10 and set the number of local epochs to I for both of them in IID and NonIID.
Table 9: The performance of the distribution shift scenario 1 where the clients' data are overall long-tailed (p = 5) and the proxy data are balanced.
Method IID (a = 100)
NonIID (a =
FedAvg
61.12
59,82
FedDF
39.03
40.68
FedAWO
67.61
66.48
20

---

Under review as a conference paper at ICLR 2023
FashionMNIST, MLP, a= 100
FashionMNIST, MLP, 0.1
90
67.5
FodMNO
FedAWO
80
85.0
FedProx
FudBE
62.5
FedDF
Teat Ace
86
83.0
Serven
FasAWD
84
(SWA)
FedProx
FasAvi
75.0
82
FedDF
FyuDyn
725
Server-FT
80
70.0
125
150
175
Communication Rounds
Communication Rounds
FashionMNIST, LeNet,
FashionMNIST, LeNet, 0.1
FedMNO
(
FedPmnx
91.0
Fudkig
86
FedAwo
(SWA)
90.5
Fediyn
Test Acc
FedProx
Server
Feddwwg
Teat
90.0
Fedi!
80
FedDyn
SurvirFF
103
25
75
100
125
150
175
200
Communication
Communication Rounds
Figure 16: Test accuracy curves algorithms under FashionMNIST. According to the results in Table 5,
Table 10: The performance of the distribution shift scenario 2 where the clients' data overall balanced and the proxy data are long-tailed (p 10).
Setting
(a = 100)
NonIID (a = 1)
Type of proxy data Balanced Sampling
Balanced
Long-tailed
Balanced
Long-tailed
w/o
w
w/o
w
FedAvg
75.24
75.24
75.24
73.46
73.46
73.46
FedDF FedAWO 79.40
76.20
74.04
73.31
74.39
73.99
73.14
77.14
78.56
76.70
76.78
70.42
Scenario The clients' data are overall long-tailed while the proxy data are balanced. The results reported in Table 9, which illustrates that method FedAWO also performs well. In comparison, the server-side ensemble distillation method FedDF has poor results this setting, and it is worse than FedAvg.
Scenario 2: The clients' data are overall balanced while the proxy data are long-tailed. The results are reported in Table 10, and we also find our method FedAWO can improve the performance. Due to the long-tailed distribution of the proxy dataset, we additionally design one extra strategy: balanced sampling, which means that we first sample the long-tailed proxy data into smaller but balanced dataset and then use it. In the IID setting, the balanced sampling method can improve FedAWO's optimization. In contrast, in the NonIID setting, the original long-tailed data work well for FedAWO, and the balanced sampling even reduces the accuracy. Moreover, there is an interesting finding that the balance degree of the proxy dataset is more critical IID settings. In comparison, FedDF has inferior results when the proxy set is long-tailed, and our FedAWO can improve generalization over FedAvg in both IID and NonIID scenarios.
21

---

Under review as a conference paper at ICLR 2023
CIFAR10, CNN, a= 100
CIFAR1O. CNN, a 0.1
62.5
82.0
Teat Acc
86
Ace
55.0
FadAWD
(SWA)
52.5
FadkwD (SWA)
FodProx
FedP fox
FasAAig
50.0
Feslin
FodBE
FodDF
FedDF
FyoDyn
47.5
FodDyn
Server-FT
60
45.0
125
125
150
175
Communication Rounds
Communication Rounds
CIFAR10, ResNet20, 100
CIFAR10, ResNet20. 0.1
80.0
75.0
Teat AeatAcc
Aoet
FeshWO
FestWo
FedNWo (SWN)
30
FesAWo (BWA)
87.5
FedProx
FedPrix
Feddin
FedAvg
FesBE
FeuDF
825
Faddlyn
FedEyyn
Server-F
Server
103
0
25
75
100
125
150
175
200
Communication
Communication Rounds
Figure 17: Test accuracy curves algorithms under CIFARIO. According to results Table 5.
D ADDITIONAL DETAILS OF FEDAWO
In FEDAWO, we optimize AW on the server as Equation 6, and there are constraints that A; > 0, ||||| = 1. To realize these constraints, we adopt base functions in X, and there are two alternatives, the quadratic function and the exponential function.
e: Σ e
Quadratic: A
Exxonentiall: X
(9)
Σ 2
is the variable that determines the value of X. We compute the gradients of æ to updates A. By using the base functions, À can meet the constraints mon-noonninee neggttgaaivv and L1 = 1. The exponential function is the same the Softmax function and we find these functions have similar performances overall, so we only adopt the exponential function in the experiments.
E IMPLEMENTATION DETAILS
E.1 ENVIRONMENT.
We conduct experiments under Python 3.8.5 and Pytorch 1.12.0. We use Quadro RTX 8000 GPUs for computation.
E.2 DATA
Data partition. To generate NonlID data partition amongst clients, we use Dirichlet distribution sampling in the trainset of each dataset. In our implementation, apart from clients having different class distributions, clients also have different dataset sizes; we think this partition more realistic in practical scenarios. For the data partition in Figure 1 and Figure 11, sampling to generate a FL system with both Mlassballacee clients and class-in flas-imbalanced clients. Specifically, we first generate all-client distribution with Or,, and only keep half of these clients.
hybrid Dirichlet
22

---

Under review as a conference paper at ICLR 2023
CIFAR100, CNN, a= 100
CIFAR100, CNN, a = 0.1
35
25
Fedtwo
25
Fedrwo (SWA)
Test Ace
20
FedPProx
Test Ace
Fedfug
16
FedBE
FwdAWD
FidDF
FesAWD (SWA)
Fedilym
FesP rox
Serven
FudAvg
FedDF
FooDyn
0
125
175
200
Communication Rounds
Communication Rounds
CIFAR100, ReaNet20, 100
CIFAR100, ResNet20, a 0.1
24
50.0
27.5
30
Teet Agc
Test Ace
22.5
28
Fedwio
20.0
Fedkwo
Fedrwo (SWA)
(SW)))
FedProx
17.5
FedProx rox
Fedfug
Fedivg
FedBE
15.0
FedDF
FodDF
22
Fedilym
12.5
FadDyn
Serven
ServerFF
25
50
75
100
125
150
175
200
125
150
175
200
Communication Rounds
Communication Rounde
Figure 18: Test accuracy curves algorithms under CIFARI00. According to the results in Table 5.
Then use the remaining data generate the distribution of remaining clients with 0g. For the data in Figure I, we first generate a 20-client distribution with aı = 10 and keep the first 10 clients as the balanced clients, then
the remaining data to generate distribution of the last 10 imbalanced clients with 02 = 0.1. The distribution is shown in Figure 12.
Data augmentation. We adopt no data augmentation in the experiments.
Proxy dataset. We use small and class-balanced proxy dataset on the server. In Table 5, we use proxy datasets with 10 samples per class, which means, for FashionMNIST and CIFAR-10, 10, there are 100 samples in the proxy datasets, and for CIFAR-I00, there are 1000 samples in the proxy datasets. The proxy datasets are randomly selected from the testset of each dataset. Then use the remaining data in the testset to test the global models' performance for all compared methods. For Table 3 and Figure 10, use CIFAR-10 10 and 100-sample proxy dataset, while in Table 4, use CIFAR-10 and 1000-sample proxy dataset.
E.3 MODEL
SimpleCNN and MLP. The SimpleCNN for CIFAR-10 and CIFAR-100 is a convolution neural network model with ReLU activations which consists of 3 convolutional layers followed by fully connected layers. The first convolutional layer is of size (3, 32, 3) followed by max pooling layer of size (2, 2). The second and third convolutional layers are of sizes (32, 64, 3) and (64, 64, 3), respectively. The last two connected layers of sizes (64*444, 64) and (64, num classes, respectively. The MLP model for FasionMNIST is a three-layer MLP model with ReLU activations. The first layer is of size (28*23 28, 200), the second is of size (200, 200), and the last is (200, 10).
ResNet and DenseNet. We followed the model architectures used in (Li et al., 2018). The numbers of the model names mean the number of layers of the models. Naturally, the larger number indicates a deeper network. For WRN56 in Table 4, it is an abbreviation of Wide- ResNet56-4, where "4" refers to four times many filters per layer.
23

---

Under review as a conference paper at ICLR 2023
E.4 RANDOMNESS
Randomness is very important in fair comparison. In all experiments throughout the paper, all implement the experiments three times with different random seeds and report the averaged results. We random seeds 8, 9 and 10 in all experiments. Given random seed, we torch, numpy, and random functions as the same random seed to make the data partitions and other settings identical. To make sure all algorithms have the same initial model, we save an initial model for each architecture and load the saved initial model at the beginning of one experiment. Also, for the experiments with partial participation, the participating clients each round vital in determining the model performance, and to guarantee fairness, we save the sequences of participating clients in each round and load the sequences all experiments. This will make sure that, given random seed and participation ratio, every algorithm will have the same sampled clients in each round.
E.5 EVALUATION
evaluate the global model performance on testset of each dataset. The testset is mostly class- dlassbbalanced and can reflect the global learning objective of an FL system. Therefore, we reckon the performance of the model on the testset can indicate the generalization performance of global models. In all experiments, as the final test accuracy for each experiment. For the indicators during training in section 4, like y, r, the norm of global gradient, and the norm of GWS pseudo- graddeen,,, we take the averaged values in the middle stage of training, that is the average of 90-1 10 rounds.
run 200 rounds and take the average test accuracy of the last 10 rounds
E.6 HYPERPARAMETER
Learning rate and the scheduler. We set the initial learning rates (LR) 0.08 in CIFAR-10 and PashionMNIST and set LR as 0.01 in CIFAR-100. We set a decaying LR scheduler all experiments; that in each round, the local LR is 0L9 of the last round). Local weight decay. We adopt local weight decay in all experiments. For CIFAR-10 10 and FashionM- NIST, we set the weight decay factor as 5c-4, and for CIFAR- 100, we set it as 5e-5. Optimizer. We set SGD optimizer as the clients' local solver and set momentum as 0.9. For the server-side optimizer (FEDDF, FEDBE, SERVER-FT, and FEDAWO), we use Adam optimizer and betas=(0.5, 0.999).
Hyperparameter for FL algorithms. For FEDDF, FEDBE and FEDAWO, we set server epoch as 100. We observe for SERVER-FT, this epoch is too large that it will cause negative effects, so we set the epoch as for SERVER-FT. We set = 0.001 in FEDPROX and a = 0,01 in FEDDYN the Gaussian mode in
suggested in their official implementations papers. For FEDBE, SWAG server. We did not use temperature smoothing in the ensemble distillation methods FEDDF and FEDBE.
24


Reviews:
[{"review_id": "aqanWDi43ge", "reviewer": "Reviewer_q3AP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the problem of optimizing for aggregation weights assigned to client model updates along with weight decay so as to aid generalization. In particular, the authors take the viewpoint of seeing the aggregation and server side update analogous to that of mini-batch SGD in centralized training and propose an algorithm named FedAWO. In particular, the proposed algorithm optimizes for weight decay and aggregation weights by making use of a proxy dataset at the server. Experiments on different datasets and models demonstrate the efficacy of the approach.", "summary_of_the_review": "The paper solves a well motivated problem with a simple, practical and well examined solution. The proposed algorithms introduces no additional hyperparameters which makes the solution particularly effective. However, the benefits claimed seems to be tightly tied to the similarity of the proxy dataset with that of the sum total data distribution of the clients. The paper could be strengthened by studying the robustness of the approach in terms of benefits still intact even when the data at the server is distributionally different from the data distribution of the sum total of all clients.", "strength_and_weaknesses": "Strengths:\n+ The paper is well written and the problem under consideration is well motivated. The proposed algorithm is light weight, i.e., doesn't add any computational overhead for the devices but instead adds it to the server which makes the algorithm attractive to use. The algorithm is easy to use and doesn't add any additional hyperparameters.\n+ The authors provide various new insights in terms of the effects of gradient coherence and heterogeneous coherence and its interplay with global weight shrinking and critical points to better explain the training dynamics in FL. \n+ Experimental results are comprehensive in terms of different datasets, models, different levels of heterogeneity both in terms of non-iidness and the number of samples on clients.\n\nWeaknesses:\n- It seems success of the approach is tied to the proxy dataset being a representative sample of the entire data distribution. No matter how small or how large the proxy dataset is, the server gets access to the true data distribution of all the clients. This is a major limitation of the approach. Having a proxy dataset at the server is something entirely possible, but the distribution of it being represenative of the true data distribution is a strong assumption. It is not clear how the performance of the algorithm is affected, when the proxy dataset is not distributionally similar to the data distribution of all the clients as a whole. It would especially strengthen the paper, as to how much of the benefit is attributable to the distribution of the proxy dataset and how much of it is from adaptively optimizing for $\\gamma$ and $\\lambda_{i}$'s.\n- The proposed approach is devoid of any theoretical guarantees and is based on observations based on SGD being the local solver. In particular, the observations concerning gradient coherence wouldn't carry over if each client undergoes multiple local steps. \n- The experiments are missing some details regarding the number of local epochs considered for different baselines and the performance of FedAWO in settings involving multiple local steps.  ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "T7rq2LfYloW", "reviewer": "Reviewer_dT3M", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper investigates the training dynamics in federated deep learning from the perspective of mini-batch SGD. Client coherence and global weight shrinking regularization are introduced in the proposed FedAWO framework. The corresponding parameters are learned in the server side with a proxy dataset. FedAWO’s effectiveness is demonstrated empirically across several datasets and model architectures. Interestingly, FedAWO shows the ability to filter out corrupted clients.", "summary_of_the_review": "See “Strength And Weaknesses” and “Clarity, Quality, Novelty And Reproducibility”.", "strength_and_weaknesses": "Overall, this paper is well written and clearly presented. \n\nThe core components (client coherence and global weight shrinking regularization) of the proposed framework are carefully analyzed independently to each other before presenting the resulting FedAWO. \n\nRelated work on federated learning is partially missing; it is necessary to move key related work to the main text. As a result, the paper is not appropriately positioned in literature in a certain sense: Is it possible to introduce more aggregation parameters on the server side in a similar adaptive manner, except for the proposed $\\gamma$ and $\\lambda$? Are there any previous work adopting a similar adaptive aggregation approach?\n\nThe concept of “number of local epochs” is confusing. Intuitively, it is about how $g^t_i$ is derived locally for each client. However, since the preliminary about FL is not introduced in detail, it is not sufficiently readable.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4ILKGA35WrP", "reviewer": "Reviewer_eJUT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The author proposes to optimize the aggregation weight of each client, and find that the aggregation weight is closely related to gradient coherence. The proposed method improves the heterogeneity coherence. The author also proposes global weight shrinking to improve the training performance. Moreover, the author makes many definitions of new concepts and try to explain how the training dynamic are affected in FL. Empirical experiments show improvement over existing methods.", "summary_of_the_review": "The author made a lot of new definitions and propose to optimize the aggregation weight and gradient weight shrinking. However, the use of proxy dataset make it less practically appealing, and the proposed techniques do not seem novel. Most of the findings are not new to me, too.", "strength_and_weaknesses": "Strength:\n1. The paper is well-organized and easy to follow.\n2. The empirical improvement is significant.\n\nWeaknesses/questiones:\n1. How is the proxy dataset selected? Proxy dataset does not seem to be a common practice in most FL settings, and it is not fair to compare it with other methods not using proxy dataset.\n2. How is Eq(3) and Eq(4) optimized? The default aggregation weight in FL is used in Eq(4), instead of the proposed $\\lambda$. Why not jointly optimize $\\gamma$ and $\\lambda$?\n3. In Eq(2), why is the first order Tyler expansion $-\\eta\\langle g^t,g^t\\rangle$? Shouldn't it be $-\\eta\\langle g^t, \\nabla \\mathcal{L}(w^t)\\rangle$?\n4. After Eq(2), why do you assume all gradients have the same norm? This does not seem to be realistic.\n5. Regarding cosine similarity/gradient coherence, it is not a new concept. For example, check [1] (there are also many other works). Moreover, it is actually the core problem in FL studied by many prior works as it contributes to the client drift issue.\n6. The global weight shrinking does not seem novel to me. It is applying weight decay to the pseudo-gradient (local update of a client) and I do not think it is new.\n7. How do you compare GWS with server learning rate both theoretically and empirically?\n8. How do you compare methods employing server optimizers, which apply SGD/Adam variants to pseudo-gradient/local update from clients, both theoretically and empirically?\n9. Optimizing aggregation weights is not new. There has been more advanced methods to deal with it, e.g. reinforcement learning [2,3]. While the method proposed is very simple by using a proxy dataset.\n\nReferences:\n[1] Yin, Dong, et al. \"Gradient diversity: a key ingredient for scalable distributed learning.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2018.\n[2] Auto-FedRL: Federated Hyperparameter Optimization for Multi-institutional Medical Image Segmentation. ECCV 2022.\n[3] Robust federated learning through representation matching and adaptive hyper-parameters. arXiv preprint arXiv:1912.13075 (2019)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject