SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Junbo Li $^{1}$ , Ang Li $^{2}$ , Chong Tian $^{1}$ , Qirong Ho $^{1}$ , Eric P. Xing $^{1,3,4}$ , Hongyi Wang $^{3}$
1. Mohamed bin Zayed University of Artificial Intelligence
2. University of Maryland 3. Carnegie Mellon University 4. Petuum, Inc.
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 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 https://github.com/ljb121002/fednar
FL has emerged as a privacy-preserving learning paradigm that eliminates the need for clients' private data to be transmitted beyond their local devices [1, 2, 3]. 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 in FL. For instance, hospitals may wish to cooperatively train a predictive healthcare model, but privacy regulations may mandate that each hospital'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 [1]. 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 issue. They build upon the FedAvg approach by integrating regularization terms during local training and employing weight decay in their local optimizers to 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, an issue magnified in FL due to the numerous
local updates from clients with varied data distributions. As demonstrated in Figure 1, a weight decay just slightly larger than optimal can result in an approximate  $10\%$  drop in accuracy. This poses a compelling research question:
How can we strike a balance between averting overfitting on clients using weight decay, and minimizing update divergence in FL?
In this study, we acknowledge the significant role weight decay coefficients play in FL algorithms and undertake an exhaustive examination of their effects. Our results demonstrate that current FL algorithms are highly susceptible to variations in weight decay values. To reconcile the need to prevent overfitting with the aim to 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 is 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.
Figure 1: Accuracy of three FL algorithms (i.e., FedAvg, FedProx, and SCAFFOLD) after 1000 rounds using different weight decay coefficients. See details in Section 3
Our contributions. More specifically, we 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 of 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 is 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 to faster training and higher accuracy.
FL is a paradigm for collaborative learning with decentralized private data [7, 1, 4, 2, 3]. Standard approach to FL tackles it as a distributed optimization problem where the global objective is defined by a weighted combination of clients' local objectives [8, 4, 9, 10]. Theoretical analysis has demonstrated that federated optimization exhibits convergence guarantees but only under certain conditions, such as a bounded number of local epochs [11]. Other work has tried to improve the averaging-based aggregations [12, 13]. Techniques such as secure aggregation [14, 15, 16] and differential privacy [17, 18] 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-efficient framework, addressing the challenges posed by highly imbalanced data, we highlight the significance of weight decay, which is crucial 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 AdamW [20] and WIN [21]. However, the integration of weight decay into FL remains relatively unexplored, both in terms of theoretical understanding and practical implementation. A recent study by [22] delved into 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 on a proxy dataset, which is unrealistic and lacks theoretical guarantees.
2 Preliminaries: federated optimization formulation
Our objective is to minimize the loss function  $F(x)$  of a model  $x$ , which is trained on data dispersed across  $M$  clients, as denoted by:
$$
x ^ {*} = \arg \min _ {x} F (x) = \arg \min _ {x} \sum_ {i = 1} ^ {M} F _ {i} (x),
$$
where  $F_{i}(x)$  corresponds to the loss function calculated using the data located on 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 a result, only a fraction of clients can engage in each communication round with the server in practice. To circumvent these constraints, FedAvg has been proposed as a prevalent communication-efficient FL framework [1]. For a round  $1 \leq t \leq T$ , we show an illustration of FedAvg in Algorithm [1]
Algorithm 1 Round  $t$  of FedAvg
Require: Global model  $x^{t - 1}$ , learning rate  $\lambda_l, \lambda_q$ , number of local updates  $\tau$ .
1: for Client  $i = 1,\dots ,M$  do
2: Let  $x_0^{t,i} = x^{t - 1}$ . Update  $x_{k}^{t,i} = x_{k - 1}^{t,i} - \lambda_{l}\nabla F_{i}(x_{k - 1}^{t,i})$  for  $k = 1,\dots ,\tau$
3: Send local update  $\Delta^{t,i} = x_0^{t,i} - x_\tau^{t,i}$  back to the server.
4: end for
5: Aggregate local updates:  $\Delta^t = \frac{1}{M}\sum_{i = 1}^{M}\Delta^{t,i}$ , and update global model:  $x^{t} = x^{t - 1} - \lambda_{g}\Delta^{t}$ .
In the original FedAvg formulation, local updates on each client are realized through (stochastic) gradient descent (SGD/GD) with the loss function  $F_{i}(x)$ , identical to the global objective. On the FedAvg foundation, we can further adjust the local loss function to be  $F_{i}(x) + \zeta_{i}(x)$ , where  $\zeta_{i}(x)$  serves as a regularization term, enabling faster convergence in the federated learning context. For instance, "client drift", a common problem in federated learning due to multiple local updates preceding aggregation, can be mitigated by adding a proximal term. FedProx [4] employs a local loss  $F_{i}(x) + \frac{\mu}{2}\| x - x_{0}\|^{2}$  that incorporates an additional proximal term  $\zeta_{i}(x) = \frac{\mu}{2}\| x - x_{0}\|^{2}$  to counteract overfitting in local clients. Similarly, SCAFFOLD [6] introduces local control variables  $c_{i}$  and a global control variable  $c$  and modifies the local gradient to be  $\nabla F_{i}(x) - c_{i} + c$ . Thus,  $\zeta_{i}(x)$  equates to  $(c - c_{i})^{T}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 [23, 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, we investigate various methodologies including {FedProx, SCAFFOLD}, which adjust local updates, and {FedAvgm, FedAdam} [9], 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 a class distribution  $\mathbf{p}_i \sim \mathrm{Dirichlet}(\alpha, \dots, \alpha)$  for each client  $i$ , and subsequently sample training data in accordance with the distributions  $\mathbf{p}_i$ . When the value of  $\alpha$  is low, the sampled individual probabilities  $\mathbf{p}_i$  exhibit higher concentration, whereas for greater values of  $\alpha$ , the concentration diminishes. Consequently, a larger value of  $\alpha$  corresponds to a 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^{-4}, 10^{-3}, 10^{-2}, 5 \times 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 be 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 are tempered when either the number of local updates is curtailed in (b) or a balanced data distribution
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  $\tau$  local steps per round, a local learning rate of 0.01, and a global learning rate of 1.0. We apply a decay for the local learning rate of 0.998 per round and a gradient clipping of max norm 10 as per [24]. Given that multiple local steps and imbalanced data distribution are two distinguishing features of FL, we utilize various pairs of  $(\alpha, \tau)$  to observe their influence on the results. The baseline setting is chosen to be  $(\alpha, \tau) = (0.3, 20)$ , resulting in a highly imbalanced data distribution. The second configuration reduces the number of local updates to  $(\alpha, \tau) = (0.3, 5)$ . The third configuration employs a balanced data distribution with  $(\alpha, \tau) = (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 0. Moreover, previous experiments either adopted values akin to those utilized in centralized training, or fine-tuned it within a restricted range, or simply set it to 0.
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 it impacts global updates. Following this theoretical exploration, we propose our FedNAR algorithm. Designed with the dual aims of ensuring a 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 adaptive learning rate and weight decay. Say the model is of dimension  $d$ . We first define two functions  $\lambda, \mu: \mathbb{N} \times \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}^+$ , and for  $t \geq 0$ , we denote  $\lambda_t = \lambda(t, \cdot, \cdot)$ ,  $\mu_t = \mu(t, \cdot, \cdot)$  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 2. This framework holds for arbitrary functions  $\lambda$  and  $\mu$ . We will clarify the choices of  $\lambda$  and  $\mu$  in FedNAR later in equation 3 and 4 after some theoretical analysis, and make clear where the "Normalized Annealing Regularization" in FedNAR comes from. In the following, we denote  $f_k^{t,i}(x)$  to be the loss function computed on a stochastic batch instead of  $F_i(x)$ .
Algorithm 2 is a generalized version of FedAvg that incorporates adaptive local learning rate and adaptive weight decay. We can recover the original FedAvg by setting  $\lambda_t(g, x)$  and  $\mu_t(g, x)$  to be constant functions. It 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 3 demonstrate that weight decay could significantly affect the convergence of federated optimization. Therefore, we 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, 10]. The initial two are conventional assumptions utilized in the general analysis of SGD optimizer, while the third one is common in FL theory.
Assumption 1 (Smooth loss functions). Every local loss function is  $L$ -smooth, i.e., there exists  $L > 0$  such that for any client  $1 \leq i \leq M$ , the gradient satisfies  $\| \nabla F_i(x) - \nabla F_i(y) \| \leq L \| x - y \|$ .
Algorithm 2 FedNAR
Require: Functions  $\lambda_t, \mu_t, F_i$ , constant  $\lambda_g, T, \tau > 0$ , sets  $C_t$ .
1: for Communication rounds  $t = 1, \dots, T$  do
2: Global server sends model  $x^{t - 1}$  to participated clients in round  $t$ .
3: for Client  $i \in C_t$  do
4: Let  $x_0^{t,i} = x^{t - 1}$ .
5: for Iteration  $k = 1,\dots ,\tau$  do
6: Compute (stochastic) gradient:  $g_{k - 1}^{t,i} = \nabla f_{k - 1}^{t,i}(x_{k - 1}^{t,i})$ .
7: Compute adaptive learning rate and weight decay:
$$
\lambda_ {k - 1} ^ {t, i} = \lambda_ {t} (g _ {k - 1} ^ {t, i}, x _ {k - 1} ^ {t, i}), \quad \mu_ {k - 1} ^ {t, i} = \mu_ {t} (g _ {k - 1} ^ {t, i}, x _ {k - 1} ^ {t, i}).
$$
8: Update  $x_{k}^{t,i} = (1 - \mu_{k - 1}^{t,i})x_{k - 1}^{t,i} - \lambda_{k - 1}^{t,i}g_{k - 1}^{t,i}$
9: end for
10: Send local update  $\Delta^{t,i} = x_0^{t,i} - x_\tau^{t,i}$  back to the server.
11: end for
12: Aggregate local updates:  $\Delta^t = \frac{1}{|C_t|}\sum_{i = 1}^{|C_t|}\Delta^{t,i}$ , and update global model:  $x^{t} = x^{t - 1} - \lambda_{g}\Delta^{t}$ .
13: end for
Assumption 2 (Unbiased gradient and bounded variance). For any client  $1 \leq i \leq M$ , the stochastic gradient is an unbiased estimator of the full-batch gradient, i.e.,  $\mathbb{E}[\nabla f_k^{t,i}(x)] = \nabla F_i(x)$ , and there exists  $\sigma^2 > 0$ , such that for any  $1 \leq i \leq M$ ,  $\mathbb{E}[\| \nabla f_k^{t,i}(x) - \nabla F_i(x) \|^2] \leq \sigma^2$ .
Assumption 3 (Bounded heterogeneity). There exists  $\sigma_g^2 > 0$ , such that  $\frac{1}{M}\sum_{i=1}^{M}\|\nabla F_i(x) - \nabla F(x)\|^2 \leq \sigma_g^2$  holds for any  $x \in \mathbb{R}^d$ .
Our findings suggest that certain limitations on  $\lambda_{t}$  and  $\mu_{t}$  are necessary to ensure the convergence of FedNAR. We begin by presenting a lemma that demonstrates the impact of local adaptive weight decay and gradients on the global update process.
Lemma 1. For round  $t \geq 1$ , the global update is equivalent to be
$$
x ^ {t} = (1 - \mu_ {g} ^ {t - 1}) x ^ {t - 1} - \lambda_ {g} h ^ {t},
$$
$$
\mu_ {g} ^ {t - 1} := \lambda_ {g} - \frac {\lambda_ {g}}{M} \sum_ {i = 1} ^ {M} \prod_ {j = 0} ^ {\tau - 1} \left(1 - \mu_ {j} ^ {t, i}\right), \quad h ^ {t} := \frac {1}{M} \sum_ {i = 1} ^ {M} \sum_ {j = 0} ^ {\tau - 1} \lambda_ {j} ^ {t, i} \prod_ {r = j} ^ {\tau - 1} \left(1 - \mu_ {r} ^ {t, i}\right) g _ {j} ^ {t, i}. \tag {1}
$$
We show the full proof in Appendix B. Lemma I highlights that the additional weight decay in local updates introduces a deviation  $\mu_g^{t-1}x^{t-1}$  in the global update. Previous works on federated optimization [4, 24, 6] did not include weight decay in their theoretical analysis, leading to  $\mu_j^{t,i}$  always being zero, and  $\mu_g^{t-1} = \lambda_g - \frac{\lambda_g}{M} \cdot M = 0$ . Moreover, we can see that larger  $\tau$  further intensifies the deviation, since  $\mu_g^{t-1} = \lambda_g(1 - \frac{1}{M}\sum_{i=1}^M\prod_{j=0}^{\tau-1}(1 - \mu_j^{t,i}))$ , and larger  $\tau$  implies that  $\mu_g^{t-1}$  increases and approaches  $\lambda_g$ . This highlights that the communication-efficient training framework of FL amplifies the impact of weight decay.
So to achieve convergence with local weight decay, it is necessary to regulate this deviation  $\mu_g^t x^t$ . Therefore, we present the following result to control the norm of the global model, i.e.,  $\| x^{t}\|$ , as the first step.
Theorem 1. If the functions  $\lambda_t$  and  $\mu_t$  satisfy that there exists  $A > 0$ , such that for any  $t \geq 0$ , and  $(g, x) \in \mathbb{R}^d \times \mathbb{R}^d$ ,  $\|\lambda_t(g, x)g + \mu_t(g, x)x\| \leq A$ , then the norm of global model in Algorithm [2] is at most polynomial increasing with respect to  $t$ , i.e.,  $\|x_t\| \leq \mathcal{O}(poly(t))$ .
We want to emphasize that the condition required for  $\lambda$  and  $\mu$  to achieve polynomial increasing speed in Theorem 1 is practical and commonly used in real-world training. Many previous works, such as
<sup>1</sup>Here we assume full-client participation, following previous works [10, 24]. However, our analysis can be easily extended to partial-client participation, as demonstrated in [10].
[26, 24], adopt gradient scaling based on norm, which bounds the term  $\| \lambda_t(g,x)g\|$  by a constant related to the clipping value. Similarly, we can bound  $\| \mu_t(g,x)x\|$  using norm clipping, for example, by setting  $\mu_t(g,x) = \mu \min \{1,c / \| x\|\}$  for some constant  $c$  and  $\mu$ . Therefore, the condition is easily satisfied from a practical standpoint.
Finally, we give the convergence result for FedNAR in the following Theorem.
Theorem 2. Under Assumptions I and 3 if  $\lambda_t$  and  $\mu_t$  satisfy the condition in Theorem 1 and also if there exists  $\{u_t\}$  with an exponential decay speed, such that for any  $t \geq 0$ , and any  $g, x \in \mathbb{R}^d$ ,  $\mu_t(g, x) \leq u_t$ , then the global model  $\{x^t\}$  satisfies:
$$
\min  _ {1 \leq t \leq T} \mathbb {E} \| \nabla F (x ^ {t}) \| ^ {2} \leq \mathcal {O} \left(\frac {1}{T}\right) + \mathcal {O} \left(\underbrace {L \lambda_ {g} ^ {2} \tau^ {2} C ^ {2}} _ {\text {w e i g h t d e c a y e r r o r}} + \underbrace {L \lambda_ {g} ^ {2} \tau l _ {*} ^ {2} (\tau \sigma^ {2} + L ^ {2} \tau^ {2} A ^ {3} + \sigma_ {g} ^ {2})} _ {\text {c l i e n t d r i f t a n d d a t a h e t e r o g e n e i t y e r r o r}}\right),
$$
where  $l_{*} = \max_{t}\{l_{t}\}$ , with  $l_{t} = \max_{g,x}\{\lambda_{t}(g,x)\}$ , and  $C$  satisfies  $u_{t}\| x^{t}\|$ ,  $u_{t}\| x^{t}\|^{2}\leq C$  for any  $t$ .
**Discussions.** For the first time, we demonstrate that federated optimization, with non-convex objective functions and stochastic local gradient, achieves convergence when employing adaptive learning rate and weight decay. Our findings resemble previous studies [26, 10], featuring a diminishing term  $\mathcal{O}(1 / T)$  as well as a 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 = 0$ .
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  $\|\mu_g^t x^t\|$  by a constant as shown in the complete proof. So for this term, we have
$$
\begin{array}{l} \left\| \mu_ {g} ^ {t} x ^ {t} \right\| = \lambda_ {g} \left(1 - \frac {1}{M} \sum_ {i = 1} ^ {M} \prod_ {j = 0} ^ {\tau - 1} \left(1 - \mu_ {j} ^ {t + 1, i}\right)\right) \| x ^ {t} \| \leq \lambda_ {g} \left(1 - \frac {1}{M} \sum_ {i = 1} ^ {M} \left(1 - u _ {t + 1}\right) ^ {\tau}\right) \| x ^ {t} \| \\ \leq \lambda_ {g} \left(1 - \frac {1}{M} \sum_ {i = 1} ^ {M} \left(1 - \tau u _ {t + 1}\right)\right) \| x ^ {t} \| = \lambda_ {g} \tau u _ {t + 1} \| x ^ {t} \| \leq \lambda_ {g} \tau C, \tag {2} \\ \end{array}
$$