SlowGuess/NeurIPS_2023_OCR huggingface dataset:
which is a constant unrelated to  $t$ . Here the second inequality is from  $(1 - \alpha)^n \geq 1 - n\alpha$  for  $0 \leq \alpha \leq 1$  and  $n \geq 1$ . We can see that the exponential decrease condition can be released as long as the decay of  $\{u_t\}$  can control  $\{\| x^t\|\}$ . Since we have already shown a polynomial increasing rate for  $\{\| x^t\|\}$ , a convenient choice for both theory and practice is to set  $\{u_t\}$  an exponential decrease. We summarize the conditions of  $\lambda$  and  $\mu$  needed to provide convergence below:
Condition 1. 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$ . There exists  $\{u_t\} \downarrow 0$  with exponential decay, such that for any  $t \geq 0$ , and any  $(g, x) \in \mathbb{R}^d \times \mathbb{R}^d$ ,  $\mu_t(g, x) \leq u_t$ .
Choice of  $\lambda$  and  $\mu$ . In order to meet the requirement of Condition [1] a straightforward and intuitive way is to use clipping. We start by selecting two sequences  $l_{t}$  and  $u_{t}$ , where  $u_{t} = u_{0}\gamma^{t}$  for some  $\gamma < 1$  is the control sequence of  $\mu_{t}$ . The sequence  $\{l_t\}$  is the learning rate schedule, and is usually set to be a constant or decrease over training. Then we define
$$
\lambda_ {t} (g, x) := \left\{ \begin{array}{l l} l _ {t} \cdot A / \| g + u _ {t} x / l _ {t} \|, & \text {i f} \| g + u _ {t} x / l _ {t} \| > A \\ l _ {t}, & \text {o t h e r w i s e} \end{array} , \right. \tag {3}
$$
$$
\mu_ {t} (g, x) := \left\{ \begin{array}{l l} u _ {t} \cdot A / \| g + u _ {t} x / l _ {t} \|, & \text {i f} \| g + u _ {t} x / l _ {t} \| > A \\ u _ {t}, & \text {o t h e r w i s e} \end{array} . \right. \tag {4}
$$
We can easily verify that the choices of  $\lambda_{t}$  and  $\mu_t$  given in equations 3 and 4 satisfy Condition 1. See details in Appendix B. Hence, we choose these functions as the input of Algorithm 2.
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  $u_{t}$  decays with an exponential rate, and  $\mu_t(g,x)\leq u_t$  uniformly for all  $(g,x)$ . (2) We observe from experiments that the norm of  $x$  keeps increasing, so the term  $\| g + u_{t}x / l_{t}\|$  will also keep increasing. NAR computes  $A / \| g + u_{t}x / l_{t}\|$  which by the preceding argument is 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 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 Line 6 in Algorithm 2 to the corresponding loss functions in FedProx and SCAFFOLD to obtain FedProx-NAR and SCAFFOLD-NAR.
Comparison with gradient clipping. Gradient clipping [27] is a widely used technique in federated optimization to enhance training stability as demonstrated by various studies [26, 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
$$
\lambda_ {t} (g, x) := \left\{ \begin{array}{l l} l _ {t} \cdot A / \| g \|, & \text {i f} \| g \| > A \\ l _ {t}, & \text {o t h e r w i s e} \end{array} , \text {a n d} \quad \mu_ {t} (g, x) := u _ {t}. \right. \tag {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  $\mu_k^{t,i}x_k^{t,i}$  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 [24, 26], no additional hyperparameters are needed. Fundamentally, FedNAR can be 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 as modifying a single line of code, resulting in a solution that is both highly effective and efficient.
We show the results of FedNAR for both vision and language datasets in this section.
Experiments on the CIFAR-10 dataset. As outlined in Section 3, we manually partition the CIFAR-10 dataset into 100 clients, following the methodology described in [25, 24]. Specifically, we randomly assign a distribution  $\mathbf{p}_i \sim \mathrm{Dirichlet}(\alpha, \dots, \alpha)$  to each client  $i$ , then draw data from the training set in line with these distributions  $\{\mathbf{p}_i\}$ . We choose  $\alpha$  from the set  $\{0.3, 1, 10\}$ , with larger  $\alpha$  values indicating a more balanced data distribution across each client. With  $\alpha = 0.3$ , the local data distribution is highly skewed, whereas with  $\alpha = 10$ , the local distribution is nearly balanced across all classes. We adhere to the training parameters and settings, optimally tuned as per [24]. For each algorithm, we conduct 1000 rounds of training for full convergence, with 20 steps of local 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 to 0.01 with a decay of 0.998, and cap the maximum norm at 10 for both gradient clipping and FedNAR to ensure a fair comparison.
Table 1 showcases the performance of FedAvg-NAR, FedProx-NAR, and SCAFFOLD-NAR across different  $\alpha$  values. The results reveal a consistent performance boost offered by FedNAR across a variety of  $\alpha$  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 SCAFFOLD, are variants of FedAvg that modify the local training objective. FedNAR can be tailored to these client-side FedAvg variants, which provide a
Table 1: Experimental results on the CIFAR-10 dataset, where the FedNAR plugin is incorporated into the FedAvg, FedProx, and SCAFFOLD algorithms, with data partitioning across varying levels of heterogeneity controlled by  $\alpha$ .
Algorithm	α = 0.3	α = 1	α = 10
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
Figure 3: Test accuracy curve for FedAvg, FedProx, SCAFFOLD and their FedNAR variants for  $\alpha = 0.3$ . For each training, we take 3 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 II 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 a model for the task of next-character prediction, incorporating 80 possible characters, following previous studies [29, 26]. We train a transformer model with 6 attention layers for feature extraction alongside a fully connected layer for predictions.
Table 2: 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.
Algorithm	WD = 10-4	WD = 10-3
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
Each character is represented by a 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 a batch size of 100, a learning rate of 0.1 with a decay rate of 0.998 per round, a dropout rate of 0.1, and a 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.
5.2 Ablation studies
In this section, we conduct some ablation studies. For all the experiments, we run FedAvg on CIFAR-10 with 100 clients and  $\alpha = 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 a decay of 0.998.
FedNAR's robustness in relation to hyperparameters. As demonstrated in Section 3, 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 in Section 3, 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 a similar occurrence. The test accuracy curve is depicted in Figure 4.
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  $t$ , the weight decay remains consistent for both the baseline 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 SCAFFOLD with WD values of 0.1 and 0.01. The utilization of 0.1 as the WD value proved to be less than ideal, resulting in a performance decrement in the baseline methods. Conversely, FedNAR initially mimics this trend but swiftly ameliorates during the subsequent stages, outstripping the baselines even with a 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  $u_{t}$ , 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  $\| g + \frac{u_t}{l_t} x\| _2$ , 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 / \| g + \frac{u_t}{l_t} x\| _2$ .
Figure 5: Frequency and strength of clipping. In every round, there is a sum total of 20 clients, and each client carries out 20 updates, leading to a 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 of the updates subjected to clipping. We execute experiments for diverse  $\alpha$  values, and for every algorithm, we present the outcomes utilizing three distinct seeds.
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 [1, 26, 10, 24, 3, 29].
In this study, we delve into the influence of weight decay in FL, underscoring its paramount importance 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, a solution that conforms to these prerequisites. FedNAR boasts a straightforward implementation and can be effortlessly adapted to a 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 a wider array of FL algorithms, incorporating adaptive optimization methods.
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 IIS1955532, NSF CNS2008248, NIGMS R01GM140467, NSF IIS2123952, NSF BCS2040381, an Amazon Research Award, NSF IIS2311990, and DARPA ECOLE HR00112390063.
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient 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, Arjun 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):1-210, 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. A 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 Rieke, 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-5143. PMLR, 2020.
[7] Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint arXiv:1610.02527, 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 H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623, 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 on Machine Learning, pages 7252-7261. PMLR, 2019.