FedNAR: Federated Optimization with Normalized Annealing Regularization
Authors: Junbo Li
Ang Li
Chong Tian
Qirong Ho
Eric Xing
Hongyi Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: Federated learning
weight decay
adaptive hyperparameters
TL;DR: A novel adaptive algorithm using normalized and annealing weight decay in federated learning.

Abstract: Weight decay is a standard technique to improve generalization performance in modern deep neural network optimization, and is also widely adopted in federated learning (FL) to prevent overfitting in local clients. In this paper, we first explore the choices of weight decay and identify that weight decay value appreciably influences the convergence of existing FL algorithms. While preventing overfitting is crucial, weight decay can introduce a different optimization goal towards the global objective, which is further amplified in FL due to multiple local updates and heterogeneous data distribution.
To address this challenge, we develop {\it Federated optimization with Normalized Annealing Regularization} (FedNAR), a simple yet effective and versatile algorithmic plug-in that can be seamlessly integrated into any existing FL algorithms. Essentially, we regulate the magnitude of each update by performing co-clipping of the gradient and weight decay.
We provide a comprehensive theoretical analysis of FedNAR's convergence rate and conduct extensive experiments on both vision and language datasets with different backbone federated optimization algorithms. Our experimental results consistently demonstrate that incorporating FedNAR into existing FL algorithms leads to accelerated convergence and heightened model accuracy. Moreover, FedNAR exhibits resilience in the face of various hyperparameter configurations. Specifically, FedNAR has the ability to self-adjust the weight decay when the initial specification is not optimal, while the accuracy of traditional FL algorithms would markedly decline. Our codes are released at \href{https://anonymous.4open.science/r/fednar-BE8F}{https://anonymous.4open.science/r/fednar-BE8F}.

FedNAR: Federated Optimization with Normalized Annealing Regularization
Junbo Li', Ang Li Chong Tian', Qirong Ho', Eric P. Xingl.3., Hongyi Wang*
1. Mohamed bin Zayed University of Artificial Intelligence
2. University of Maryland 3. Carnegie Mellon University 4. Petuum, Inc.
Abstract
Weight decay is a standard technique to improve generalization performance modern deep neural network optimization, and is also widely adopted in federated learning (FL) to prevent overfitting in local clients. In this paper, we first explore the choices of weight decay and identify that weight decay value appreciably in- fluences the convergence of existing FL algorithms. While preventing overfitting is crucial, weight decay introduce different optimization goal towards the global objective, which is further amplified FL due to multiple local updates and heterogeneous data distribution. To address this challenge, develop Federated optimization with Normalized Annealing Regularization (FedNAR), a simple yet effective and versatile algorithmic plug-in that can be seamlessly integrated into any existing FL. algorithms. Essentially, we regulate the magnitude each update by performing co-clipping of the gradient and weight decay. We provide compre- hensive theoretical analysis of FedNAR's convergence rate and conduct extensive experiments both vision and language datasets with different backbone federated optimization algorithms. Our experimental results consistently demonstrate that in- corporating FedNAR into existing FL algorithms leads to accelerated convergence and heightened model accuracy. Moreover, FedNAR exhibits resilience in the face of various hyperparameter configurations. Specifically, FedNAR has the ability to self-adjust the weight decay when the initial specification is not optimal, while the accuracy of traditional FL. algorithms would markedly decline. Our codes are released at httpps://iihub..om//ij 12002//edna
Introduction
FL has emerged as privacy-preservng learning paradigm that eliminates the need for clients' private data to be transmitted beyond their local devices [|] FL presents challenges in addition to traditional distributed learning, including communication bottlenecks, heterogeneity in hardware resources and data distributions, privacy concerns, etc 4. Numerous machine learning applications necessitate training FL. For instance, hospitals may wish to cooperatively train a predictive healthcare model, but privacy regulations may mandate that each hospital's 's data remains local (5).
In traditional distributed optimization, local gradients are calculated for each client per iteration and subsequently sent to a global server for amalgamation. However, the communication constraints within Federated Learning (FL) entail multiple model updates prior to the aggregation phase at the global server, as exemplified by the widely-employed FedAvg framework [. Considering the skewed data distribution and multiple local updates for communication efficiency, it's pivotal to guard against overfitting models to local data. Algorithms like FedProx 4 and SCAFFOLD [6] were designed to address this issuc. They build upon the FedAvg approach by integrating regularization terms during local training and employing weight decay in their local optimizers combat overfitting on local clients. However, a frequently overlooked trade-off exists: local weight decay introduces an optimization objective that differs from the global one, issue magnified in FL due to the numerous
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

local updates from clients with varied data distributions. As demonstrated in Figure | just slightly larger than optimal can result in an approximate 10% drop in accuracy. This poses compelling research question:
weight decay
How strike balance between averting overfitiing on clients using weight decay, and minimizing update divergence in FL?
FedAvg
85
FediProx
Accuracy
SCAFFOLD
In this study, we acknowledge the significant role weight decay coefficients play in FL algorithms and undertake an exhaustive ex- amination of their effects. Our results demonstrate that current FL algorithms are highly susceptible to variations in weight decay val- ues. To reconcile the need prevent overfitting with the aim minimize update divergence, we present an adaptive weight decay strategy named Federated Optimization with Normalized Annealing Regularization (FedNAR). FedNAR is a cost-efficient approach that can be effortlessly incorporated into any prevailing FL algorithm as an algorithmic plug-in. At its core, we govern the magnitude of each weight update by limiting the norm of the sum of the gradient and weight decay. We conduct a meticulous theoretical analysis for universally adaptive weight decay functions, proving that our FedNAR a logically sound choice for achieving convergence from a theoretical perspective. Through an array of experiments spanning computer vision and language tasks, we establish that FedNAR fosters improved convergence and facilitates superior final accuracy.
80
75 10-
10-3
10
10
Weight Decay
Figure 1: Accuracy of three Fl. al- gorithms (i.e. FedAvg. FedProx, and SCAFFOLD) after 1000 rounds using different weight de- cay coefficients. See details in Section S
Our contributions. More specifically, made the following contributions in this work to the entire federated learning and optimization community
We empirically observed that FL settings make weight decay more crucial in training.
Based on the observation described above, we designed FedNAR, an algorithmic plug-in to control the magnitude weight decay that is compatible with all existing FL algorithms.
We propose the first theoretical framework of FL. that takes weight decay into account and show that FedNAR guaranteed to converge under FL scenarios.
We conduct extensive experiments on vision and language tasks by plugging FedNAR into state-of-the-art baselines, which demonstrate that adopting FedNAR generally leads faster training and higher accuracy.
1.1 Related work
FL is paradigm for collaborative learning with decentralized private data 7 71 2 3 approach to FL tackles as a distributed optimization problem where the global objective defined by a weighted combination of clients' local objectives [8 100 Theoretical analysis has demonstrated
Standard
that federated optimization exhibits convergence guarantees but only under certain conditions, such as bounded number of local epochs [TI. Other work has tried to improve the averaging-based aggregations [12 [3 Techniques such as secure aggregation [4 1.5 [6] and differential privacy [77
I8] have been widely adopted to further improve privacy in FL 19 Our FedNAR is based on the standard FedAvg framework and can be adapted to different settings.
In order to enhance convergence in a communication-efline posed by highly imbalanced data, we highlight the significance of weight decay, which is erucial in modern deep neural network (DNN) training. Extensive research in centralized training has explored various methods to incorporate weight decay into optimization algorithms, such as Adam W [00] and WIN (11) However, the integration of weight decay into FL remains relatively unexplored, both in terms of theoretical understanding and practical implementation. the examination of weight decay during server updates in the context of FL. Nevertheless, their focus was primarily on tuning the global weight decay parameter, neglecting the exploration of weight decay in local updates. Additionally, their choice of the global weight decay parameter relied proxy dataset, which is unrealistic and lacks theoretical guarantees.
framework, addressing the challenges
recent study by 22 delved into
2 Preliminaries: federated optimization formulation
Our objective to minimize the loss function F(x) of model x, which is trained on data dispersed across M clients, as denoted by:

---

M
x² = arg min F(x) = arg min Σ P((2),
where F1(x) corresponds to the loss function calculated using the data located the i-th client. Federated optimization introduces two key constraints: privacy and communication efficiency. The privacy constraint ensures that clients retain their data without sharing, whereas the communication constraint emerges due to the substantial cost and restricted capacity of network bandwidth. As result, only fraction of clients can engage in each communication round with the server in practice. To circumvent these constraints, FedAvg has been proposed as prevalent communication-effiiiie FL. framework
[1 For round S T, we show an illustration of FedAvg in Algorithm |
Algorithm 1 Round t of FedAvg
Require: Global model x²11, learning rate Aτ >g, number of local updates T.
1: for Client i 1, M do
2: Let - Update o =k-1 X ( for k ,
3: Send local update A Do xt'² back the server.
4: end
5: Aggregate local updates: A
NI EM Ati, and update global model: x² = =x'=- - A
In the original FedAvg formulation, local updates on each client are realized through (stochastie) gradient descent (SGD/GD) with the loss function (z)).. identical to the global objective. On the FedAvg foundation, we can further adjust the local loss function be Ff(x)) + (zx), where (xx)) serves as a regularization term, enabling faster convergence in the federated learning context. For instance, "client drift", a common problem in federated learning due multiple local updates preceding aggregation, can be mitigated by adding a proximal term. FedProx 4 employs local
loss F, (x) + ~1 ||.- oool2 that incorporates an additional proximal term G(rr) = 2 || xo I2 to counteract overfitting in local clients. Similarly, SCAFFOLD O introduces local control variables C; and a global control variable and modifies the local gradient to be (z) - G + c. Thus, (1)
equates to (c- 7xxx.. x.
3 An empirical study of weight decay
Weight decay is a well-established approach in standard centralized deep neural network training, primarily used to prevent overfitting 2 20 In the realm of FL., weight decay values generally mirror those employed in centralized training. In this section, we meticulously explore the impact of weight decay on federated optimization. In addition to the original FedAvg, investigate various methodologies including [FedProx, SCAFFOLD), which adjust local updates, and (FedAvgm, FedAdam) ) [ which utilize SGD with momentum and Adam, respectively, in the global update while treating aggregated local updates as pseudo-gradient. We also assess FedExP 24 which adaptively modulates the global learning rate.
We evaluate all the aforementioned algorithms on the CIFAR-10 dataset, partitioned among 100 clients under three different settings. Following [25 in each setting, we initially sample class distribution p; ~ Dirichlet(a,,
a.) for each client i, and subsequently sample training data in accordance with the distributions p.. When the value of is low, the sampled individual probabilities P; exhibit higher concentration, whereas for greater values of a, the concentration diminishes. Consequently, larger value of corresponds to more balanced data distribution across clients.
We train a ResNet-18 for each algorithm on CIFAR-10 for our empirical study. Figure 2 presents the test accuracy plotted against weight decay values from (10-2, 10133, 10-2,5 x 10-2 10 1) for each setting. The y-axis range is maintained consistently between (50, 95) across all settings to facilitate a clear comparative analysis of the influence of weight decay under varied settings.
In the baseline setting (a), the influence of weight decay on the performance of each algorithm is pronounced in two primary ways. Firstly, a noteworthy enhancement in accuracy exceeding 6% is observed when leveraging optimal weight decay as compared to sub-optimal weight decay. This improvement outstrips the differences discerned across various algorithms. Secondly, increasing weight decay to slightly larger than the optimal value, such as from 0.05 to 0.1 or from 0.01 to 0.05, can cause a considerable plunge in accuracy exceeding 10% and even verging on 30%. This underscores the phenomenon of amplified update deviation in FL. However, both of these elements tempered when either the number of local updates is curtailed in (b) or balanced data distribution
3

---

(a) 0 0.3, T= 20
(b) 0.3, =
(c) a= T=
Accuracy
80
Accuracy
80
Accuracy
80
Feddug
70
FedIrox
70
70
FedFxs
60
60
60
FodAdam
FedAugm
50 0
10"2
10
20
10
10 Weight decay
Weight Decay
Weight Decay
Figure 2: Influence of weight decay for different settings. We train each algorithm (i.e.. FedAvg. FedProx. SCAFFOLD, FedExp, FedAdam, FedAvgm) over 1000 rounds with T local steps per round, local learning rate of 0.01, and global learning rate of 1.0. We apply decay for the local learning rate of 0.998 per round and a gradient clipping of max norm 10 per 24 Given that multiple local steps and imbalanced data distribution two distinguishing features of FL. we utilize various pairs of (a, r) to observe their influence the results. baseline setting is chosen to be (a, 7) = (0.3, 20), resulting in a highly imbalanced data distribution. The second configuration reduces the number of local updates to (a, T) = (0.3, 5). The third configuration employs a balanced data distribution with (a, 7) = (10, 20).
is deployed in (c). In these scenarios, the accuracy curve corresponding to weight decay is more leveled for each algorithm, and the dip in accuracy is significantly reduced when weight decay is slightly elevated from the optimal value.
This finding underscores how the training paradigm in FL can modulate the impact of weight decay. However, prior work has largely overlooked the role of weight decay, both theoretically and experimentally. Convergence analyses in prior studies did not take into account the influence of weight decay and assigned its value as O. Moreover, previous experiments either adopted values akin to those utilized in centralized training, or fine-tuned it within restricted range, or simply set to
4 FedNAR and its convergence analysis
Driven by the recognized shortcomings of previous studies, we introduce the first analytical framework for convergence that incorporates local weight decay, examining its resulting deviation and how impacts global updates. Following this theoretical exploration, we propose our FedNAR algorithm. Designed with the dual aims of ensuring theoretical guarantee and achieving an optimal rate of convergence, FedNAR marks a significant advancement in this field.
4.1 A general federated optimization framework
Our theoretical analysis can cover udaptive learning rate and weight decay, Say the model is of dimension d. We first define two functions X, μ N Rd Rd - RR
and for > 0, we denote A = X(t ;) Me = ..) as mappings with the domain of definition on the gradient space and weight space. These two functions stand for learning rate and weight decay respectively. This framework can be adapted to any FL algorithms by making the learning rate and weight decay to be adaptive. We present the framework in Algorithm g This framework holds for arbitrary functions A and jr. We will clarify the choices of À and μ in FedNAR later in equation 3 and after some theoretical analysis, and make clear where the "Normalized Annealing Regularization in FedNAR
comes from. the following, we denote ct. (x) to be the loss function computed on a stochastic batch instead of (x)..
Algorithm 2 is a generalized version of FedAvg that incorporates adaptive local learning rate and adaptive weight decay. We recover the original FedAvg by setting \r(g x) and [4 x) to be constant functions. is worth emphasizing that although most previous works like FedAvg and FedProx employed local weight decay, they did not integrate it into their theoretical analysis. However, our empirical investigations in Section demonstratee that weight decay could significantly affect the convergence of federated optimization. Therefore, argue that thorough empirical and theoretical analysis is needed to better understand this term.
4.2 Convergence analysis and FedNAR
To initiate the theoretical analysis, we present three standard assumptions that are widely used in federated optimization literature 24 [TO]. The initial two are conventional assumptions utilized in the general analysis of SGD optimizer, while the third one common in FL theory.
Assumption 1 (Smooth loss functions). Every local loss function I--mooth, iee, there exists such that for any client 1 < i S M, the gradient satisfies F.())) V.(()) < L|[æ - yll.

---

Algorithm 2 FedNAR
Require: Functions Ax Me: F., constant Ag T > sets C..
1: for Communication rounds 1,... , do
2: Global server sends model x (= to participated clients in round i.
3: for Client C. do
Let s6
2²-1
5:
for Iteration = 1 ,,..., do Compute (stochastic) gradient: 1.
6:
V (
V
7:
Compute adaptive learning rate and weight decay:
x Ae(gie =
Mk-1 pe (ok 1-
Update = p4 )01-1 1 8:
k- 19h-1
9:
end for
10:
Send local update A 6, x00 x
back to the server.
11:
end for
12:
Aggregate local updates: A
and update global model: xf xx--_ Ag
13: end for
Assumption 2 (Unbiased gradient and bounded variance). For any client <is M, the stochastic
gradient is unbiased estimator of the full-batch gradient, i.e., El f pt (x)] Fx)),, and there
exists a² > O, such that for any sis M, V ( (F(())l So
Assumption 3 (Bounded heterogeneity). There exists o such that Σ M
Disl [V(())
VF(2)]2 - og holds for any x E Rd.
Our findings suggest that certain limitations FedNAR. We begin by presenting a lemma that demonstrates the impact of local adaptive weight decay and gradients the global update process |
and μr, are necessary ensure the convergence of
Lemma 1. For round t > 1, the global update is equivalent to be
x²
(1 11g )x= Xa
where
M T-1 a1-
1-1 Hg
Ag
(1)
M
==
We show the full proof in Appendix B Lemma highlights that the additional weight decay in local updates introduces a deviation 2 optimization
1-1 the global update. Previous works on federated did not include weight decay in their theoretical analysis. leading to j Ay
always being zero, and 2 >g
M = 0. Moreover, we can see that larger further intensifies the deviation, since 14 1- Ag 1. I ))))), and larger 7 implies that 1451 increases and approaches Xg. This highlights that the communication-fficcee training framework FL amplifies the impact weight decay.
M
So to achieve convergence with local weight decay, it is necessary to regulate this deviation uiat. Therefore, we present the following result to control the norm of the global model, i.e., ||x*1], as the first step.
Theorem 1. the functions d and μt satisfy that there exists A > O, such that for any > and (g. x) E Rd x Rd, [e x)g + Helg, x)al|| S. A, then the norm of global model in Algorithm is at most polynomial increasing with respect i.e.. || < O(poly(t)).
We want to emphasize that the condition required for À and μ to achieve polynomial increasing speed Theorem M practical and commonly used real-world training. Many previous works, such as
'erre we assume full-client participation. following previous works 10 24 easily extended to partial partial-client participation, demonstrated in [0]
However, our analysis can be

---

by a constant related to the clipping value. Similarly, we can bound [g ==)/ using norm clipping, for example, by setting (r(g,2) x) min (1, c/ e/||>]]] for some constant and js. Therefore, the condition is easily satisfied from practical standpoint.
26 2
adopt gradient scaling based on norm, which bounds the term
Finally, we give the convergence result for FedNAR in the following Theorem.
Theorem 2. Under Assumptions 7 and 3 if x and pt satisfy the condition in Theorem 1
and also if there exists {ue) with exponential decay speed, such that for any 0, and any 9, E I&d (e,g S ue then the global model {x²) satisfies:
(I)) =o =0()) +0
min IstST
+
C2
+ +
weight decay error
client drift and date beterogeneity evior
where 4 = maxx(4) with 4 = max maxg.= (e and C' satisfies ue |||2="] ue|l="112 - C for any t.
Discussions. For the first time, we demonstrate that federated optimization, with non-convex objec-
tive functions and stochastic local gradient, achieves convergence when employing adaptive learning rate and weight decay. Our findings resemble previous studies 26 TO], featuring a diminishing term
OLLTT) as well non-vanishing term arising from weight decay, client drift, and data heterogeneity in FL. Notably, our bound recovers previous results when weight decay is disregarded, i.e., C' C==0 O.
The detailed proof can be found in the Appendix B Here we briefly explain the condition of exponential decreasing. To tackle the case with deviation caused by local weight decay, one of the key steps is to bound || by a constant as shown in the complete proof. So for this term, we have
M
|1''11 S Ao (1
M
||az 1|g='11=b = Ag
(1-
1 M
Σ
(1 - ------))) ||=0
M i=1
M (1 Ea-
$ Ag
1
= - Ag
(2)
M
which constant unrelated to t. Here the second inequality is from (1 a) > no for 0 S a < and 2 1. We can see that the exponential decrease condition can be released as long as the decay of {ut] can control [[]='']]. Since we have already shown a polynomial increasing rate for [[]='1]), convenient choice for both theory and practice is to set {w,} an exponential decrease. We summarize the conditions of À and μ needed to provide convergence below:
Condition 1. There exists A O, such that for any 2 and (g.x) E Rd Rdd [|Ar(g, x)g Le(g, x x| < A. There exists {u:} 0 with exponential decay, such that for any t 2 O, and any (g.z) E Fd x Rd, pe(g. S ut.
Choice of À and μr. In order meet the requirement of Condition 1 straightforward and intuitive way is to use clipping. We start by selecting two sequences 4 and ut, where u, = un? for some 7 is the control sequence of pt. The sequence {42 is the learning rate schedule, and is usually set to be a constant or decrease over training. Then we define
f4 A/llg + ux/4el, if llg + ux/44/ A otherwise
xe(g,
(3)
14+
and
2)(g.z) = (u A/llg + ux//4ll if [lg + w.r/4|l A otherwise
(4)
We easily verify that the choices of X and μ4 given in equations and 4 satisfy Condition
| See
details in Appendix B.
Hence, we choose these functions as the input Algorithm
4.3 Understanding FedNAR
Normalized annealing regularization. Now we can explain why we call our method "Normalized Annealing Regularization", First, the regularization term is co-normalized with the gradient term.

---

Second, the "annealing" comes from two aspects: (1) The sequence 4n decays with an exponential rate, and (e x) I u uniformly for all (g.z).
2). (2) We observe from experiments that the norm of keeps increasing, so the term ||g + u,2//11l will also keep increasing. NAR computes A/llg + u.r//, /I,]] which by the preceding argument a decreasing term, thus introducing another form of decay. We will provide further illustrations of this in the experiment section.
Flexibility of FedNAR. Since FedNAR's 's convergence analysis does not impose any extra conditions on the local loss functions, it generalizes several FL algorithms with distinct local training objectives, while conferring upon them theoretical assurances comparable to FedNAR. For example, FedProx and SCAFFOLD only differ from FedAvg on local loss functions. Thus, we only need to change 2 to the corresponding loss functions in FedProx and SCAFFOLD to obtain TedProx-NAR and SCAFFOLD-NAR.
Line D in Algorithm
Comparison with gradient clipping. Gradient clipping 27 is a widely used technique in federated
optimization to enhance training stability as demonstrated by various studies 266 24 28 It is
achieved by constraining the norm of gradients in each iteration. The technique involves setting a threshold value and rescaling the gradients if they exceed this limit. In our framework, gradient clipping can be represented by setting
(g(g.2) = 44- A/lloll,
if Iلlll > A otherwise
and (e == ue.
(5)
Our proposed FedNAR distinguishes itself from previous gradient clipping strategies in two primary ways. First, both our weight decay and learning rate functions are adaptive, with dependencies on the current gradient and model weights. This contrasts with traditional gradient clipping, which only adapts the learning rate function based on the current gradient. Second, instead of solely employing the norm of the gradient, FedNAR adopts the norm of the sum of the gradient and weight decay as the clipping criterion. This strategy ensures that each local update is bounded, thereby preventing potential explosions caused by large decay terms f 2. in federated optimization. Consequently, FedNAR not only guarantees convergence but also bolsters training stability within the context of federated optimization.
Implementation of FedNAR. Implementing FedNAR merely requires specifying an initial weight decay and a decay rate. Therefore, for state-of-the-art baseline methods that employ gradient clipping 2 26), no additional hyperparameters interpreted as a modest alteration that adjusts the sequence of gradient clipping and weight decay operations. Consequently, the implementation of FedNAR can be as uncomplicated modifying single line of code, resulting in a solution that is both highly effective and efficient.
needed. Fundamentally, FedNAR can be
5 Experiments
We show the results of FedNAR for both vision and language datasets in this section.
5.1 Main results
Experiments on the CIFAR-10 dataset. As outlined in Section 3 we manually partition the
CIFAR-10 10 dataset into 100 clients, following the methodology described in 125 24 Specifically, we randomly assign a distribution P; ~ Dirichlet(a..... a) to each client i, then draw data from the training set in line with these distributions {p:}. We choose from the set f0.3, 1, 10), with larger a values indicating a more balanced data distribution across each client. With a = 0.3, the local data distribution is highly skewed, whereas with a = 10, the local distribution is nearly balanced across all classes. We adhere to the training parameters and settings, optimally tuned as per (244). conduct 1000 rounds of training for full convergence, with 20 steps of local
For each algorithm, training per round. In each round, we randomly sample 20 clients. For the baseline algorithms, we apply gradient clipping as suggested in [24 26 We set the local learning rate 0.01 with decay of 0.998, and cap the maximum norm at 10 for both gradient clipping and FedNAR to ensure fair comparison.
Table I showcases the performance of FedAvg-NAR, FedProx-NAR, and SCAFFOLD-NAT across different values. The results reveal consistent performance boost offered by FedNAR across variety of a values, with particularly significant improvements observed for skewed distributions. Test accuracy after each training round is also illustrated in Figure 3
The backbone algorithms, FedProx and SCAFFOLDD, are variants of FedAvg that modify the local training objective. FedNAR can be tailored to these client-side FedAvg variants, which provide
7

---

Table 1: Experimental results the CIFAR-IO dataset, where the FedNAR plugin incorporated into the FedAvg, FedProx, and SCAFFOLD algorithms, with data partitioning across varying levels heterogeneity controlled by a.
0.3
10
Algorithm
Baseline
FedNAR
Baseline
FedNAR
Baseline
FedNAR
FedAvg
85.19
87.64
88.50
89.55
89.45
90.23
FedProx
85.02
87.45
88.56
89.58
89.37
90.38
SCAFFOLD
86.89
89.17
89.52
91.13
90.64
91.85
FedExP
86.46
88.55
88.28
89.51
88.82
89.69
FedAdam
84.12
85.81
87.35
87.78
88.86
89.56
FedAvgm
90.33
90.43
91.60
91.76
92.19
91.15
FedAug
FedProx
SCAFFOLD
Teckvg-hill
FedFics440
22
SCArFOLOHAR
teplvg
Podifis
SCAFPOLD
Accurany
Teal Accurary eatr
Tost Accurapy
80
2
200
400
620
800
1000
433 Round
Round
Round
Figure 3: Test accuracy curve for FedAvg, FedProx, SCAFFOLD and their FedNAR variants for 0.3. For
each training, take random seeds.
guarantee of convergence. Moreover, FedNAR can also be incorporated into server-side FedAvg variants, which retain the local training but alter the global update. Table showcases the results of these implementations, using the same hyperparameter configurations. These results corroborate the same conclusion.
Experiments on the Shakespeare dataset. The Shakespeare dataset (29. derived from The Complete Works of William Shakespeare, assigns each speaking role in every play to a unique client, leading to an inherent non-i.i.d. partition. We select 100 clients and train model for the task of next-character prediction, incorporating 80 possible characters, following previous studies 29 26 We train transformer model with 6 attention layers for feature extraction alongside fully connected layer for predictions. Each character is represented by 128-dimensional vector, and the transformer's hidden dimension is set to 512. The training regimen spans 300 rounds, with 20 clients chosen at random in each round and 5 local epochs executed on each client. During local training, we utilize batch size of 100, learning rate of 0.1 with decay rate of 0.998 per round, dropout rate of 0.1, and maximum norm of 10 for both the baseline algorithms and FedNAR. Global updates are performed with a learning rate of 1.0. As in the previous experiment, FedNAR is integrated into six core algorithms. We also employ two distinct weight decay values to showcase FedNAR works for different hyperparameters. The results are presented in Table 2
Table Experimental results on the Shakespeare dataset, where the FedNAR plugin is incorporated into the FedAvg, FedProx, and SCAFFOLD algorithms. WD denotes the initial weight decay value applied in the first round.
WD 10-4
WD 10
Algorithm
Baseline
FedNAR
Baseline
FedNAR
FedAvg
42.44
44.37
40.05
44.69
FedProx
41.92
43.57
40,91
40.26
SCAFFOLD
47.55
45.60
42.40
44.53
FedExP
38.15
39.63
39.67
39.92
FedAdam
45.80
45.42
47.19
47.48
FedAvgm 45.03
45.63
46.12
47.13
5.2 Ablation studies
In this section, conduct some ablation studies. For all the experiments, CIFAR-10 with 100 clients and a = 0.3 for 1000 rounds, where each round consists of updates from 20 clients and 20 steps on each client. The local learning rate is 0.01 with decay of 0.998.
FedAvg
FedNAR's robustness in relation to byperparameters. As demonstrated in Section B weight decay exhibits high sensitivity in federated optimization and requires meticulous tuning. However, due to its co-normalization mechanism, FedNAR has demonstrated a greater robustness towards the selection of the initial weight decay. Employing the same experimental framework as Section we observed that initiating weight decay at 0.1 in FedAvg significantly compromised model performance,

---

inciting a gradient explosion due to an overly high weight decay. This led to a performance standard similar to a model with no weight decay. Yet, FedNAR was found to navigate this hurdle through self- adjustment, leading to considerable performance enhancements during the latter stages of training, in spite of the initial detrimental impact on performance. Additionally, further experiments on FedProx and SCAFFOLD also manifested similar occurrence. The test accuracy curve is depicted in Figure
FedAvg
FedProx
SCAFFOLD
cocuray.
Tcst cocuracy
Tuat cocuracy
WOKEE1
WD-0.11
WOHEO1
203 430 -00 300
Round
Round
Round
Figure 4: The self-adjusting capability of FedNAR. WD denotes the initial weight decay value applied in the first round. It is crucial to note that in each round f. the weight decay remains consistent for both the haseline methods and FedNAR. The distinction lies in FedNAR's adoption of co-clipping across both weight decay and gradient. We drew comparisons among FedAvg, FedProx, and SCAFFOLDD with WD values of 0. and 0.01. The utilization of 0.1 as the WD value proved to be less than ideal, resulting in performance decrement in the baseline methods. Conversely, FedNAR initially mimics this trend but swiftly ameliorates during the subsequent stages, outstripping the baselines even with more favorable initial WD value of 0.01.
Frequency and strength of clipping in FedNAR. In Section 4.3 we showcased how FedNAR employs an "annealing" process to stave off excessive deviation precipitated by weight decay. from two angles. The first angle pertains to the exponential decay in ue, while the second involves co- clipping with the gradient. To delve deeper into the frequency and intensity of the clipping operation, we scrutinized the updates that were subjected to clipping. In particular, we evaluated the count of clipped updates and the average norm of the update llg + We #/12., signifying the intensity of clipping, for these steps, as depicted in Figure 5. Our analysis points out that both the count and the norm of updates experiencing clipping escalate over time. This observation implies that the co-clipping term gains increasing prominence during the training process, aiding in the annealing of weight decay
which is proportional to 1/llg + de
7, rl12.
Numoor cloping
Factfug.
FedProx NAR
SCAFFOLD-NAR
Feddug NAR
FedProx-NAR
SCAFFOLD-NAR
1
10
190
100
Awanage
Round
Round
Round
Round
Round
Round
Figure 5: Frequency and strength of clipping. In every round, there is sum total of 20 clients, and each client carries out 20 updates, leading collective tally of 400 updates per round. We monitor the count of clipping instances within these 400 updates during each round and compute the average norm the updates subjected to clipping. We execute experiments for diverse a values, and for every algorithm, we present the outcomes utilizing three distinct seods.
6 Limitation
As an inaugural exploration into federated optimization with weight decay, we concentrated on examining the theoretical attributes of FL algorithms that utilize gradient descent updates both server-side and client-side. Future investigations can extend this to encompass different server update methodologies. Furthermore, due to privacy restrictions, we could not carry out experiments using real heterogeneous data, such as hospital data. Instead, we relied on simulated data distributions or naturally non-i.i.d splits, following the conventional methods employed in prior FL algorithm
research I1 26 10
29
7 Conclusion
In this study, we delve into the influence of weight decay in FL, underscoring its paramount impor- tance in traditional FL scenarios characterized by multiple local updates and skewed data distribution. To explore these facets, we present an innovative theoretical framework for FL that integrates weight decay into convergence analysis. Upon discerning the conditions for convergence, we introduce FedNAR, solution that conforms to these prerequisites. FedNAR boasts a straightforward imple- mentation and can be effortlessly adapted to variety of backbone FL. algorithms. In our experiments on vision and language datasets, we consistently record significant performance improvements across

---

all backbone FL algorithms when we employ FedNAR, leading to accelerated convergence rates. In addition, our ablation studies illustrate FedNAR's heightened robustness against hyperparameters. Going forward, our aim is to expand this theoretical framework to include wider array of FL. algorithms, incorporating adaptive optimization methods.
Acknowledgments
We thank the anonymous reviewers for their valuable insights and recommendations, which have greatly improved our work. This research has been graciously funded by NGA HM04762010002. NSF 19555553, NSF CNS2008248, NIGMS R01GM140467, NSF IIS2123952 NSF BCS2040381, an Amazon Research Award, NSF IIS231 1990, and DARPA ECOLE HR001 12390063.
References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication carpumiiicatioofffcinnnnt learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[2] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji. Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings. et al. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1-2): 2):1 1200, 2021.
[3] Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al- Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. field guide to federated optimization. arXiv preprint arXiv:2107.06917, 2021.
[4] 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, 2020.
[5] Nicola Ricke, Jonny Hancox. Wenqi Li, Fausto Milletari, Holger R Roth, Shadi Albarqouni, Spyridon Bakas, Mathieu N Galtier, Bennett A Landman, Klaus Maier-Hein, et al. The future of digital health with federated learning. NPJ digital medicine, 3(1):119, 2020.
[6] 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, pages 5132-5 5143. PMLR, 2020.
[7] Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimiza- tion: Distributed machine learning for on-device intelligence. arXiv preprint arXiv: 16100.027, 2016.
[8] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning, pages 4615-4625. PMLR, 2019.
[9] Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295. 2020.
[10] 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:761 2020.
[11] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2020.
[12] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference Machine Learning, pages 7252-7261. PMLR, 2019.
10

---

[13] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun. Dimitris Papailiopoulos, and Yasaman Khaz- aeni. Federated learning with matched averaging, In International Conference on Learning Representations, 2020.
[14] Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for privacy-preservig machine learning. In proceedings of the 2017 ACM SIGSAC Conference Computer and Communications Security, pages 1175-1191, 2017.
[15] Keith Bonawitz, Hubert Eichner. Wolfgang Grieskamp. Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, al. Towards federated learning at scale: System design. Proceedings of Machine Learning and Systems, 1: 374-388, 2019.
[16] Chaoyang He, Songze Li, Jinhyun So, Xiao Zeng, Mi Zhang, Hongyi Wang. Xiaoyang Wang. Praneeth Vepakomma, Abhishek Singh, Hang Qiu, et al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint WxXiv:2007.33518 2020.
[17] Ziteng Sun, Peter Kairouz, Ananda Theertha Suresh, and Brendan McMahan. Can you really backdoor federated learning? arXiv preprint arXiv: 1.07963, 2019.
[18] H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. In International Conference on Learning Representations, 2018.
[19] Matt Fredrikson, Somesh Jha, and Thomas Ristenpart. Model inversion attacks that exploit confidence information and basic countermeasures, In Proceedings of the 22nd ACM SIGSAC conference on computer and communications security, pages 1322-1333, 2015.
[20] lya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 171 2017.
[21] Pan Zhou, Xingyu Xie, and YAN Shuicheng. Win: tion for adaptive gradient algorithms. In Has it Trained Yer? NeurlPS 2022 Workshop.
inteerraed nesterov accelera-
[22] Zexi Li, Tao Lin, Xinyi Shang, and Chao Wu. Revisiting weighted aggregation in federated learning with neural networks. arXiv preprint arXiv:2302.1091 2023.
[23] Stephen Hanson and Lorien Pratt. Comparing biases for minimal network construction with back-propagation. Advances in neural information processing systems, 1, 1988.
[24] Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp: Speeding up federated averaging via extrapolation. arXiv preprint arXiv:2301.0966 2023.
[25] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint urXiv:1909.06333 2019.
[26] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N What- mough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint arXiv:21 01.4663, 2021.
[27] Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint 19Xiv:1955.11888 2019.
[28] Daniel Rothchild, Ashwinee Panda, Enayat Ullah, Nikita Ivkin, Ion Stoica, Vladimir Braverman, Joseph Gonzalez, and Raman Arora. Fetchsgd: Communication-effiiieee federated learning with sketching. In International Conference on Machine Learning, pages 8253-8265. PMLR, 2020.
[29] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 18/200097, 2018


Reviews:
[{"review_id": "f7anBN4Z3R", "reviewer": "Reviewer_8tVP", "summary": "The paper investigated effects of weight decay in the scenario of federated learning, especially for the stage of local updates. It was found that even subtle changes of weight decay values might lead to drastic performance drop and the observations motivated the authors to conduct convergence analysis considering the factor of weight decay. The proposed method, FedNAR, was supported by both theoretical analysis and empirical results, and demonstrated advantages when incorporated into various FL backbone algorithms.", "questions": "- Why did FedNAR perform worse on some adaptive FL methods such as FedAdam and FedAvgM?\n- Missing details about model architectures and more experiments should be included to support the proposed method.", "limitations": "The authors have discussed limitations in Section 6 in the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written and well organized. In addition, the idea was motivated and inspired by an empirical study of weight decay in Section 3, which made the paper more sound.\n- The authors investigated the influence of changes in weight decay of local updates, which was overlooked previously and found that the model performance was sensitive to the selection of weight decay values.\n- A theoretical analysis of convergence with weight decay was presented in the paper, and supported the proposed FedNAR method", "weaknesses": "- It seemed that FedNAR worked well on most of federated learning backbone algorithms except for some adaptive methods such as FedAvgM and FedAdam. The accuracy was even worse after applying FedNAR to these methods. Can the authors expand on this phenomenon and provide some insights? Do adaptive methods themselves have the ability to adjust update trajectories so that adaptive weight decay might not work as expected?\n- Experiments were still at a small scale and can be further improved. Currently the authors only chose one dataset CIFAR10 for image classification and one dataset Shakespeare for next-character prediction, which was not sufficiently. More large scale datasets such as ImageNet and practical ones with more realistic splits like CelebA (as suggested in the LEAF benchmark), or Apple's FLAIR dataset. Besides, details of the model architecture for each task were not clear. The selection of model structures might also affect the final performance and should be analyzed as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## FedNAR and adaptive algorithms:\n\nTo commence, we wish to emphasize that FedNAR is also **effective** for adaptive algorithms such as FedAdam and FedAvgm. Please consult our Table 1 and Table 2  in the original submission for reference. Among the 10 settings about FedAdam and FedAvgm covered by these tables, FedNAR demonstrates slightly inferior performance in merely 2 instances out of 10. Consequently, it remains evident that FedNAR consistently delivers enhanced outcomes within the context of adaptive methods.\n\n**Deeper analysis of adaptive algorithms and FedNAR:** These two algorithms employ update mechanisms based on momentum at the global level. This implies that the global update direction incorporates an accumulation of preceding global updates. This cumulative effect aids in rectifying current updates by counteracting deviations stemming from specific selected clients. In essence, momentum-based methodologies themselves possess attributes akin to those introduced in our FedNAR approach, which aims to control the local drift due to multiple local updates and data heterogeneity. Nevertheless, it's noteworthy that FedNAR brings forth additional enhancements, despite the inherent benefits of momentum-based techniques.\n\n## Additional large-scale datasets\n\nOur current submission includes CIFAR-10 and Shakespeare datasets, with the Shakespeare dataset already incorporating a practical LEAF-based setup. Additionally, we have performed experiments on CIFAR-100 and Tiny-ImageNet, both of which are also widely-used benchmarks in federated learning studies. The outcomes are presented in **Table 1** of the PDF within the global response. Our FedNAR consistently maintains its superior performance, with about **2%~7% improvement** across **all** these larger and more intricate datasets. While being mindful of time and computational limitations, we intend to expand our experimentation to include CelebA and Apple's FLAIR dataset in future work.\n\n**Experiment details:** For CIFAR-100 and Tiny-ImageNet, we maintained the same settings, parameters, and model utilized with the CIFAR-10 dataset, as outlined in Section 5.1, lines 270 to 281 of our initial submission. In the case of Tiny-ImageNet, as the input size is 64x64, we made a minor adjustment to the ResNet-18 model by changing the final pooling layer from avg_pool2d to adaptive_avg_pool2d to align with the image dimensions.\n\n## Model details\n\nFor vision tasks including CIFAR-10, CIFAR-100, and Tiny-ImageNet, we use a standard ResNet-18 following previous works, as mentioned in Section 3, Line 125.\n\nFor language tasks, we use a standard transformer encoder model. Specifically, we add positional encodings to the input, forward it to the 6 attention layers, and use a fully connected layer for the final prediction. Please refer to the function *class shake_transf* in the *util_models.py* in our released codes for details.\n\n"}, {"review_id": "PON9c9JTIL", "reviewer": "Reviewer_t5c6", "summary": "This paper first describes an important and challenge problem in Federated Learning, that the performance of FL is very sensitive to the choice of weight decay hyper-parameter for local optimization. The authors produced data to demonstrated the sensitivity of the weight decay hyper-parameter.\n\nThis paper first outlined a general analysis framework for any client-side learning rate and weight-decay adjustment scheme, then proposed a adaptive weight decay scheme which adjust the hyper-parameter inversely proportional to the magnitude of the sum of the local gradient and the decayed-parameters. The authors provided analysis and performance guarantees for their proposed adaptive decay method. The authors conducted experiments to demonstrate the applicability of their method as a \"plug-in\" for different types of FL methods.", "questions": "1. the authors have shown the robustness of their proposed adaptive weight decay scheme to the initial weight decay choice. But the scheme also depends on the learning rate schedule l_t, and decay schedule \\mu_t, and the maximum norm A. In the experiments, the author set A = 10. How does these HPs affect the proposed method?\n\n2. FedProx is similar to weight decay, but shrinkage to the server parameter instead of 0. Can the proposed method be used for adapting the proximal regularizer weight in FedProx? ", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. hyper-parameter sensitivity is often a over-looked issue in FL. Existing methods addressed this problem by a hyper-parameter optimization problem, which is costly. This methods directly make the weight decay parameter and adaptive, eliminate the need for search.\n\n2. The authors have shown the robustness of their method as a plug-in on various FL methods, and the robustness to the initialization of the decay value, which is critical to make it actually useful.", "weaknesses": "In the experiments, the authors ran 5 local epochs for each algorithm. When the FL training process starts from random initialization, often single local epoch produces the best results because it limits the client drift away from the server model. I would like to see how does the proposed method perform under single local epoch.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Single-epoch performance\n\nYes, FedNAR is also **effective** in a single-epoch setting. We perform experiments using CIFAR-10 data, employing a data heterogeneity parameter of 0.3. Our approach adheres to the identical configurations, parameters, and model specifications detailed in Section 5.1, lines 270 to 281 of our original submission. The selected backbone algorithms encompass FedAvg, FedProx, and SCAFFOLD. The outcomes underscore that FedNAR consistently attains superior performance and accelerates convergence on all the aforementioned backbone algorithms. For a comprehensive overview, kindly refer to the results provided in **Table 3** of the global PDF.\n\n## Additional ablation studies\nWe expand our ablation studies to encompass learning rate scheduling, weight decay scheduling, and max norm exploration. For each of these ablation studies, the backbone algorithms include FedAvg, FedProx, and SCAFFOLD. Please refer to **Table 2** and **Figure 1** in the global PDF for the detailed results.\n\nRegarding the **learning rate**, our current submission employs exponential decay in line with prior FL research practices [24, 26]. For further ablation studies, we have additionally evaluated cosine decay and inverse linear decay, where the first is common in standard centralized training and the second is proposed in [11].\n\nRegarding **weight decay** scheduling, in conjunction with the exponential decay approach applied in our current submission, we also include cosine decay and inverse linear decay methods for an extended evaluation.\n\nRegarding the **maximum norm** parameter, we opt for values of {5, 10, 20, 40}. The outcomes demonstrate a discernible trade-off concerning the selection of this parameter. A lower maximum norm tends to amplify regularization, constricting each optimization step, which in turn adversely impacts performance. Conversely, a higher maximum norm results in less frequent clipping occurrences, yielding a diminished impact. Notably, our results indicate that across all three backbone algorithms, and for every designated maximum norm value, FedNAR consistently outperforms the baseline algorithm.\n\n## Apply FedNAR in the proximal term of FedProx\n\nYes, the idea of FedNAR can be seamlessly extended to the proximal term of FedProx. The implementation is similar. We proceed to replicate experiments on CIFAR-10, adhering to the same setup outlined in Section 5. To enhance the validation of this notion, we modify the coefficient of the proximal term ($\\mu$) across {0.001, 0.01, 0.1}. The outcomes are presented comprehensively in **Table 4** of the global PDF.\n\n**Deeper analysis about the extension on FedProx:** The results showcase that, across all $\\mu$ values, the inclusion of co-clipping in the FedProx framework yields improved performance. However, it's noteworthy that the extent of enhancement is not as conspicuous as when this technique is applied to weight decay. This disparity could be attributed to\n\n1. The impact of FedProx loss and weight decay varies. While the FedProx loss confines the gap between the current local weight and the global weight, effectively curbing the extent of the local optimization step, this inherently mirrors the effects of the co-clipping scheme in empirical terms. Consequently, the incorporation of co-clipping doesn't provide substantial additional benefits in such a scenario.\n2. In the context of FedProx, the proximal term represents the L2-norm of the disparity between the current local weight and the global weight. This term is anticipated to exert a relatively lesser impact compared to the current weight itself.\n3. The co-clipping approach within FedNAR is intentionally tailored to complement weight decay, substantiated by both empirical findings in Section 3 and theoretical analyses presented in Section 4. In particular, our observations in Section 3 underline that the combination of multiple local updates and the inherent unevenness of data distribution in FL scenarios heightens the sensitivity of FL algorithms to weight decay. Additionally, the insights provided in Section 4 illustrate that co-clipping is pivotal for convergence analysis when coupled with weight decay. It's important to note that these attributes cannot be directly extrapolated to analyze the proximal term within the context of FedProx.\n\n\n"}, {"review_id": "Xriz2cdMrZ", "reviewer": "Reviewer_4Utv", "summary": "The study discusses the role of weight decay in enhancing generalization performance in deep neural network optimization and in avoiding overfitting in Federated Learning (FL). The authors highlight the influence of weight decay value on FL algorithms' convergence. To mitigate this issue, Federated optimization with Normalized Annealing Regularization (FedNAR), is introduced, which modulates each update's magnitude through co-clipping of the gradient and weight decay. The algorithm shows improved accuracy in experiments on diverse vision and language datasets with various federated optimization algorithms. ", "questions": "- Is the proposed method also effective on the methods that utilize client-specific learnable parameters such as FedDC [1]?\n\n\n[1] L. Gao, et al., FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction, CVPR 2022.\n\n", "limitations": "Please refer to \"Weakness\" Section", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation of the paper and the corresponding solution is straight-forward.\n- The article is clearly articulated and readily understandable.\n- This paper theoretically analyzes the impact of local training's weight decay on the convergence of federated learning.\n- Experiments on both image and text datasets validate the effectiveness of the proposed method.", "weaknesses": "- To further validate the effectiveness of the proposed method, it would be beneficial to conduct experiments on more challenging datasets such as CIFAR100 and Tiny-ImageNet. Additionally, tests on more realistic benchmarks, like LEAF, that encompass feature disparity or imbalanced data, would provide even stronger evidence of its efficacy.\n\n- From the results in Figure 4, it is not clear whether the proposed method is more stable for the choice of the initial value of weight decay. \nIn a similar vein, an ablation study on the choice of the threshold of co-clipping is required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Additional large-scale datasets\n\nWe have integrated experiments involving CIFAR-100 and Tiny-ImageNet. Kindly refer to **Table 1** in the global PDF for a detailed presentation of the outcomes. Our FedNAR consistently upholds its superior performance with about **2%~7% improvement** across all these more intricate datasets.\n\n**Experiment details:** For CIFAR-100 and Tiny-ImageNet, we ensured uniformity by retaining the same settings, parameters, and model configuration employed for the CIFAR-10 dataset, as elucidated in Section 5.1, lines 270 to 281 of our initial submission. In the case of Tiny-ImageNet, where the input size is 64x64, a minor adaptation was made to the ResNet-18 model. This modification involved transitioning the final pooling layer from avg_pool2d to adaptive_avg_pool2d to align with the altered image dimensions.\n\n**LEAF benchmarks:** In fact, the Shakespeare dataset utilized in our present submission originates from LEAF. We have employed the publicly available codes from LEAF to create the dataset, resulting in a pragmatic partitioning of the data. We will include more datasets from the LEAF benchmark in future work.\n\n## Additional ablation studies\n\n**Stability with respect to initial weight decay:** The primary objective of Figure 4 is to illustrate the enhanced stability of FedNAR in relation to different initial weight decay selections. Upon examining Figure 4, it becomes apparent that an initial value of 0.1 is substantial, leading to a decline in performance when contrasted with the choice of 0.01 across original FedAvg, FedProx, and SCAFFOLD methodologies. However, with the incorporation of FedNAR, while there is an initial dip in accuracy compared to the 0.01 choice, accuracy gradually recuperates and even surpasses the performance of the 0.01 option. This dynamic demonstrates that FedNAR possesses the capability to autonomously rectify the detrimental impact of an excessively large initial value. In light of the observations outlined in Section 3, where marginally elevated weight decay values result in noticeable performance degradation, FedNAR emerges as a natural and efficacious remedy for this concern.\n\n**Ablations on the co-clipping threshold:** Ablation studies have been incorporated to analyze the impact of different co-clipping thresholds. The thresholds chosen for evaluation include {5, 10, 20, 40}. Please refer to **Figure 1** in the global PDF for details. The findings underscore a discernible trade-off associated with selecting the maximum norm threshold. Opting for a lower threshold results in amplified regularization, constraining each optimization step and consequently affecting performance negatively. Conversely, a higher threshold leads to less frequent occurrences of clipping, thereby yielding a milder effect. However, it's noteworthy that our results demonstrate a consistent pattern across all three backbone algorithms: regardless of the chosen maximum norm threshold, FedNAR consistently exhibits better performance.\n\n## FedDC-NAR\nYes, FedNAR is also **effective** on FedDC. We undertake experiments employing the publicly accessible official codes of FedDC, encompassing both CIFAR-10 and CIFAR-100 datasets. Data is partitioned across 100 clients according to a Dirichlet distribution, with an imbalance parameter from {0.3, 1, 10}. We train ResNet-18 across 500 rounds, where 20 clients are randomly selected per round. The training protocol employs a learning rate of 0.01 and an initial weight decay of 0.005. Comprehensive outcomes are presented in **Table 5** of the global PDF.\n\n"}, {"review_id": "zEDlJXEilT", "reviewer": "Reviewer_D8ei", "summary": "This paper delves into the effect of weight decay in the realm of federated optimization. The authors conduct a series of experiments highlighting how specific elements in federated learning, such as the presence of diverse data and the execution of local updates, can amplify the influence of weight decay. The paper further provides a theoretical exploration of weight decay within federated learning and introduces a novel methodology named FedNAR, which is derived from their analysis. The newly proposed FedNAR method showcases enhanced convergence speed and performance on various simulated federated tasks, and notably, it displays a heightened tolerance to complications introduced by weight decay.", "questions": "1. As demonstrated in the paper, weight decay plays a critical role in federated learning. Is there a more effective approach to determining the optimal weight decay value for FedNAR?", "limitations": "As illustrated in weakness, the final bound is similar as previous works. What’s more, it would be better if the algorithms can be verified in more real-world heterogeneous data.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The writing and presentation are generally of high quality, making the idea and method of this paper easy to follow.\n2. The research includes sufficient empirical findings. Six different federated learning algorithms, encompassing various variants of FedAvg, are examined to investigate the impact of weight decay. Weight decay is explored across a wide range, providing clear insights into their findings.\n3. The proposed algorithm, FedNAR, is well-motivated by empirical findings and appears to be a simple yet effective enhancement. It can be readily adapted to different federated learning backbone algorithms.\n4. The study offers comprehensive and rigorous theoretical analysis of weight decay in federated learning. The results are novel, crucial, and hold the potential to catalyze further investigations within the federated learning community.\n5. Moreover, the paper’s experimental design is comprehensive, encompassing a broad spectrum of scenarios that integrate both vision and language datasets along with a variety of hyperparameter configurations.", "weaknesses": "1. The bound is similar to previous works, without providing any superior bounds. This may not be a significant issue, since incorporating weight decay itself makes theoretical analysis harder and can lead to convergence challenges.\n2. Though this is not my major concern, it would definitely make the paper stronger to try experiments at larger scales, e.g., GPT fine-tuning on various downstream tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Theoretical results\nOur theoretical analysis is constructed within the broader **non-convex** framework and is underpinned by **minimal** assumptions. This accomplishment is significant given the intricacies of federated optimization coupled with non-iid data distributions. The bounds we derive represent a **generalized** rendition of earlier findings – specifically, by setting the weight decay term to 0, we are able to recover prior outcomes. Furthermore, a sublinear bound can be achieved by setting the data heterogeneity parameter to 0, akin to previous investigations. In addition to these considerations, it's worth noting that the inclusion of weight decay introduces an additional layer of complexity, posing challenges for our theoretical analysis. This further underscores the non-trivial nature of our theoretical results.\n\n## Additional large-scale datasets\nWe've expanded our experimentation to encompass more expansive and intricate datasets, including CIFAR-100 and Tiny-ImageNet. For a comprehensive understanding of these endeavors, kindly refer to **Table 1** in the global PDF. Due to temporal and resource constraints, certain tasks such as GPT fine-tuning will be pursued in our future work.\n\n## Hyperparameter selection\nFedNAR serves as a versatile plug-in compatible with a wide array of Federated Learning algorithms. Therefore, it suffices to employ a similar value as used in the original backbone algorithms, which can be conveniently determined through grid search, similar to practices adopted in previous studies. Consequently, there is no additional burden in the realm of hyperparameter selection for FedNAR. Notably, as evidenced by Figure 4 in the submission, FedNAR exhibits heightened robustness to the initial weight decay, implying **greater flexibility** in hyperparameter selection due to the efficacy of our co-clipping strategy.\n\n"}]

Author rebuttal:
We express our sincere gratitude to the reviewers for their constructive and encouraging feedback. Dedicated to the continuous refinement of our work, we wish to emphasize a key contribution of our paper: **it represents the pioneering effort in systematically exploring the significance of weight decay in existing FL methods, underscored by both rigorous theoretical analysis and comprehensive experimentation - which we think are generally applicable beyond our proposed FedNAR method. Stemming from our insights, we introduced FedNAR—an essential algorithmic element that can seamlessly integrate with existing FL techniques, enabling adaptive adjustments of weight decay for enhanced convergence and model quality.** Given the time constraints of this rebuttal period, we have undertaken notable expansions in our experiments in terms of: (1) assimilating a wider and more complex range of datasets, (2) deepening our ablation studies, and (3) presenting new algorithmic innovations. The updated findings are detailed in the attached PDF.

Meta-review: The paper examines the effect of weight decay in federated learning (FL) and presents some interesting findings that highlight the importance of choosing the right weight decay mechanism in various federated learning algorithms. The authors then conduct a theoretical study to derive the convergence of FL algorithms under weight decay, which serves as the motivation for a simple truncated weight decay rule. Empirical results demonstrate that the proposed weight decay rule can enhance several existing FL algorithms. However, it appears that the improvements are less pronounced for stronger algorithms (e.g., in Table 1, there is only a marginal improvement for FedAvgm, which is the strongest baseline). This raises a slight concern that the method may not lead to a significant enhancement of the state-of-the-art. Despite this minor weakness, all the reviewers believe that this paper is interesting and makes a solid contribution, so we recommend it for acceptance.
Decision: Accept (poster)