FedADT: An Adaptive Method Based on Derivative Term for Federated Learning

Recently, vast amounts of data have been generated by decentralized network edge devices, such as mobile phones and smart devices in the internet of things. Collecting and transmitting this data for training not only gives rise to network congestion, but also cause privacy leakage. For this reason, Federated Learning (FL) frameworks are proposed in [1,2], where clients learn a shared global model based on their own private data under the coordination of a central server. As data holders, clients conduct multi-step model training locally on the basis of the currently received global model, and then, the central server aggregates these local models, obtaining a global model, and returns it to each client. This manner of alternating training and communication was implemented by McMahan et al. [2], resulting in the Federated Averaging (FedAvg) algorithm, which is one of the most popular methods in federated learning. In FedAvg, each client trains a model by leveraging the Stochastic Gradient Descent (SGD) method. Due to its superiority, FL is broadly used in many application scenarios [3,4,5].

Despite the good empirical performance of FedAvg, also known as local SGD [6], there are still gaps between theory and practice in FedAvg. To better understand its convergence performance in theory, several studies [7,8,9] associated with it have emerged under data homogeneous setting in federated learning. For data or client heterogeneity, Ref. [10] introduced a regularization term to local object function settling the non-identically distributed challenge. Control variate and variance reduction methods [11] were proposed to correct the bias across clients, which leads to unstable and slow convergence. Most of these studies require extra communication cost or memory, which can be costly and unpractical in federated setting. In addition, momentum-based methods are introduced into either local model updates [12] or global model updates [13,14] or both [15,16] to improve the stability convergence. The convergence performance of SGD type methods are highly sensitive to the learning rate or step size that controls the speed of model learning; hence, another line of studies has sprung up aiming at modifying the learning rate, which scales the gradient for each dimension by incorporating prior information. These methods [14,17,18] that use adaptive learning rate also adopt the momentum term that accumulates the previous information. However, the accumulated past gradient information can cause an improper model update, which, sometimes, may be opposite to the descent direction. The training will be lagged and this lag effect can lead to an oscillation phenomenon in which the learning curve fluctuates at the optimal point, which is known as the overshoot problem in the domain of control.

In the optimization algorithms for machine learning, especially for deep learning, stochastic gradient can be regarded as an error where the goal of optimization is to make the error gradually settle to zero. This has a similar spirit as the Proportional Integral and Derivative (PID) controller in the fields of control theory and engineering. The main idea of the PID-based control method is incorporating the current, past, and the future information into the current correction to adjust the input of a dynamic system so it performs as desired. The use of a feedback mechanism makes the control process more responsive and robust [19,20]. The research [21,22,23] reveals that the error or deviation in the PID controller plays a similar role as the stochastic gradient used in SGD-based methods. In addition, SGD with momentum mainly utilizes current and historical gradients to optimize the model, which can be interpreted as the proportional and integral part. This inspires us to introduce the derivative information to the local update of federated learning, which denotes the future trend of the gradient change. According to the aforementioned analysis, we consider the descent direction and learning rate simultaneously in the update rule of model parameters for better training the federated learning model. To this end, this paper integrate PID controller into local SGD for federated learning. In a nutshell, the contributions are elaborated below:

The remainder is organized as follows. The related work is summarized in Section 2. In Section 3, the optimization problem is first introduced. Then, the proposed federated learning algorithm is described in detail in Section 4. In Section 5, we present related assumptions and the main results. The experiments are performed to validate the theoretical results in Section 6. Finally, we conclude the paper in Section 7.

SGD is perhaps the most popular method, with good empirical performance in machine learning, which is also robust and scalable. Momentum is a heuristic, but a strong way to accelerate the convergence of SGD. Motivated by the heavy ball method [24] and Nesterov’s accelerate gradient method [25], a momentum term is usually added in the current update of descent directions by a weighed sum of previous information to improve the convergence of SGD [26]. Sutskever et al. [27] combined SGD with a careful use of the momentum method in the training of deep neural networks successfully. The article [28] developed the final iterate with standard step size schedules, and obtained the lower bounds for the sub-optimality of SGD. The generalization performance between SGD and a full gradient descent was developed by [29], and a novel separation result was presented in the stochastic convex optimization model. Additionally, adaptive optimization methods and variants [30,31,32] have gained fruitful achievements in deep learning because of their success in practice. Reddi et al. [32] proposed an Adaptive Mean Square Gradient method (AMSGrad) to amend the convergence issues of adaptive moment estimation method (Adam). Zaheer et al. [33] utilized the effect of a mini-batch size to improve the performance of Adam. The novel variant [34], which adapts step sizes according to the belief in current gradients (AdaBelief), has a better convergence, generalization, and training stability in both convex and non-convex cases by modifying Adam without additional parameters. Besides using a step size that adjusts to the scaling of gradients, a new class of adaptive methods [35,36,37] that is based on Polyak step size has emerged, utilizing both the current loss value and the stochastic gradient. Most of the studies only focus on the online convex optimization case or require projections operation on a bounded domain. Recently, the connection between PID control and stochastic optimization was described in [23], which shows that the PID-based method is an optimizer of encapsulating the gradient and momentum. An et al. [21] proposed a novel PID optimizer, which introduced derivative action to reduce the oscillation phenomenon, also known as overshoot in the control field. The above algorithms are implemented in a centralized setting.

Distributed optimization based on parallel SGD has been developed over the past decade, which often suffers from the bandwidth limits and large network delays. To alleviate communication bottlenecks, local SGD incorporating model averaging periodically results in the FedAvg algorithm [2], which significantly reduces the communication overhead. Along this line of research, there is much work that explores the theoretical convergence and improves the performance of FedAvg. Stich [6] firstly established the upper bound for FedAvg in a convex homogeneous setting when all clients participate at each round, and later it was improved by [8] in a convex heterogeneous setting. The work [38] established a lower bound for FedAvg in a heterogeneous case. Moreover, a unified framework [39] was presented to analyze local SGD methods in convex and strong convex settings. A hybrid local SGD method [40] was proposed to speed up the training of federated learning. These studies mentioned above use SGD as the local paradigm optimizer. One can also see other variants that incorporated momentum [9] and adaptive techniques [14,18]. FedAdam [14] utilizes Adam algorithm as a local optimizer in the federated learning framework to overcome the difficulty of parameters tuning for non-convex settings. Local AMSGrad [17] was designed to accelerate training and reduce communication overhead. Additionally, PID-based federated optimization methods have been developed recently. Ref. [41] designed a privacy budget allocation protocol by computing PID errors to balance privacy guarantee and the utility of the global model. The article [42] combined a federated learning framework and PID controller to develop the deployment of future intelligent transportation systems. Inspired by these works, in this paper, we use an adaptive learning rate and derivative term to the federated setting and analyze its convergence performance.

$\mathbf{Notation}$. Throughout the paper, we use ${\mathbf{x}}_t^i$ to denote the model parameter of i-th client at t-th iterations. Let $\parallel ·\parallel$ and ${\parallel ·\parallel }_{\infty }$ be ${\mathcal{l}}_2$ and ${\mathcal{l}}_{\infty }$ vector norm, and ${(·)}_j$ denotes the j-th coordinate of a vector. The vector square and vector division are element-wise, respectively.

In this paper, we consider a general federated learning system, as shown in Figure 1. It contains n clients or devices and a central server; for example, a smart phone, industrial sensors. By collecting a large amount of production data, such as temperature, pressure, and current, federated learning can jointly model data from multiple plants without sharing trade secrets, thus improving productivity, quality, and safety. As illustrated in Figure 1, the training process of federated learning can be briefly summarized as follows: the central server firstly selected a subset of edge clients and a global model is downloaded by each client involved in the training at each round. Then, each client belonging to the subset begins multiple step local training based on its raw dataset and obtains a local model. Finally, the local models are uploaded and aggregated in the central server. The above process is repeated until the global mode converges or an expected predicted accuracy is attained.

In fact, the above model training can be modeled as an optimization problem. The main goal is to find a global model parameter, denoted by the vector $\mathbf{x}\in {\mathbb{R}}^d$, and the problem to be solved is formulated of the form: where d is the dimension of model parameter, $f_i\left(\mathbf{x}\right)={\mathbb{E}}_{{\xi }_i\sim {\mathcal{D}}_i}F(\mathbf{x};{\xi }_i)$ stands for the local expected loss function of i-th client, function $F(\mathbf{x};{\xi }_i)$ denotes the loss for the model parameter $\mathbf{x}$ on one example ${\xi }_i$ stored in the i-th client, and ${\mathcal{D}}_i$ represents the data distribution of i-th client, $i\in \{1,\dots ,n\}$. For different clients, their data distribution may be different.

In this article, we are concerned with the collaborative learning of n clients under the coordination of a central server to solve problem (1) by local training and periodic model aggregation. In order to stabilize the process of local model training, we add a derivative term that denotes the trend of the gradient change and the adaptive learning rate to the update rule of the local model parameter. The pseudo-code of our proposed method, FedADT, is summarized in Algorithm 1. Specially, at the beginning of the $(t+1)$-th iteration, the central server random selects a subset of clients firstly. Each client i involved in current training computes the stochastic gradient ${\mathbf{g}}_t^i$, which is an unbiased estimator of the full gradient $\nabla f_i\left({\mathbf{x}}_t^i\right)$, by using mini-batch random data from the dataset of the client i. Then, it computes an exponential weighted average momentum term ${\mathbf{m}}_{t+1}^i$ as the descent direction of model update and second order moment ${\mathbf{v}}_{t+1}^i$ to adaptively adjust the learning rate, respectively, which are defined as follows: where ${\beta }_1,{\beta }_2\in [0,1)$ are decay factors which control the exponential decay rates of weighted averages. In fact, ${\mathbf{m}}_{t+1}^i$ can be expressed as $(1-{\beta }_1){\sum }_{j=0}^t{\beta }_1^{t-j}{\mathbf{g}}_j^i$ by recursion, where the initial value of momentum is set to 0. The decay factors ${\beta }_1$ are usually chosen so that the exponential weighted averages allocate small weights to previous gradients that are far from the current moment. A similar choice applies to the decay factor ${\beta }_2$, which is selected from the set {0.99, 0.9999} in the relevant papers [31,32]. Notation ⊙ indicates the element-wise square.

Then, a first difference term that suggests the future information is added to correct the lagged gradient. In fact, the differential of gradient is approximated by the first difference ${\mathbf{g}}_t^i-{\mathbf{g}}_{t-1}^i$, which reflects the instant variation of gradient. It is incorporated in the design of algorithm to exploit the future expectation of the model and avoid overshooting, which acts in a similar role as in the PID controller. Furthermore, in order to mitigate the noise in gradient calculation caused by randomly selecting mini-batch data, we use moving weighted average on the derivative part, resulting in:

Finally, the local model parameter ${\mathbf{x}}_{t+1}^i$ is updated, i.e., where $\eta$ is learning rate, and $\nu$ is the step size of derivative term. The term ${\widehat{\mathbf{v}}}_{t+1}$ is the element-wise maximum of ${\widehat{\mathbf{v}}}_t$ and the average of ${\mathbf{v}}_{t+1}^i$ across n clients, as shown in line 11 of Algorithm 1. Moreover, if $t+1$ is a multiple of E, the central server averages the model parameters ${\mathbf{x}}_{t+1}^i$ and the second moment ${\mathbf{v}}_{t+1}^i$, where E is a positive constant denoting the number of local updates.

Algorithm 1 Federated Adaptive Learning Based on Derivative Term (FedADT)

Input: Initial point ${\mathbf{x}}_0$, ${\mathbf{m}}_0^i=0$, ${\widehat{\mathbf{v}}}_0^i=\delta \mathbf{1}$, the number of iterations T 1: $\mathbf{for}t=0,1,\dots ,T-1\mathbf{do}$ 2:    $\mathbf{for}\mathrm{client}\mathrm{i}=1,2,\dots ,\mathrm{n}$ in parallel $\mathbf{do}$ 3:        Compute gradient: ${\mathbf{g}}_t^i=\nabla f_i({\mathbf{x}}_t^i,{\xi }_t^i)$ 4:        Update ${\mathbf{m}}_{t+1}^i={\beta }_1{\mathbf{m}}_t^i+(1-{\beta }_1){\mathbf{g}}_t^i$ 5:        Update ${\mathbf{v}}_{t+1}^i={\beta }_2{\mathbf{v}}_t^i+(1-{\beta }_2){\mathbf{g}}_t^i\odot {\mathbf{g}}_t^i$ 6:        Update ${\mathbf{d}}_{t+1}^i={\beta }_1{\mathbf{d}}_t^i+(1-{\beta }_1)({\mathbf{g}}_t^i-{\mathbf{g}}_{t-1}^i)$ 7:        $\mathbf{if}t+1modE\ne 0\mathbf{then}$ 8:             ${\widehat{\mathbf{v}}}_{t+1}={\widehat{\mathbf{v}}}_t$ 9:             ${\mathbf{x}}_{t+1}^i\leftarrow {\mathbf{x}}_t^i-\frac{\eta }{\sqrt{{\widehat{\mathbf{v}}}_{t+1}}}\odot {\mathbf{m}}_{t+1}^i+\nu {\mathbf{d}}_{t+1}^i$ 10:        $\mathbf{else}$ 11:             ${\widehat{\mathbf{v}}}_{t+1}=max\left\{\frac{1}{n}{\sum }_{i=1}^n{\mathbf{v}}_{t+1}^i,{\widehat{\mathbf{v}}}_t\right\}$ 12:             ${\mathbf{x}}_{t+1}^i\leftarrow \frac{1}{n}{\sum }_{i=1}^n({\mathbf{x}}_t^i-\frac{\eta }{\sqrt{{\widehat{\mathbf{v}}}_{t+1}}}\odot {\mathbf{m}}_{t+1}^i+\nu {\mathbf{d}}_{t+1}^i)$ 13:        $\mathbf{end}\mathbf{if}$ 14:    $\mathbf{end}\mathbf{for}$ 15: $\mathbf{end}\mathbf{for}$

In this section, before providing the main results of Algorithm 1, we first state three assumptions as follows.

Assumption 1 expounds our requirements for local objective functions, and is common in non-convex problems [14,17,43]. Next, there are two different assumptions about the stochastic gradients.

The two assumptions above are common in the analysis of adaptive-type methods [14,17], which bound the gradient estimate with noise and the variance of the stochastic gradient.

From the above assumptions, we obtain the main convergence theorem for the FedADT algorithm.

We defer to the proof of Theorem 1 in Appendix A.

From Corollary 1, we can see that the convergence of the proposed algorithm is evaluated by $\frac{1}{T}{\sum }_{t=0}^{T-1}\mathbb{E}{∥\nabla f\left({\overline{\mathbf{x}}}_t\right)∥}^2$, which is exactly the lower bound of the term on the left-hand side of (10) by utilizing the inequality ${\parallel {\widehat{\mathbf{v}}}_{t+1}\parallel }_{\infty }\le G_{\infty }$.

In this section, we study the performance of the proposed algorithm on two standard datasets for the image classifications task in a federated setting. The MNIST dataset [44] is a set of handwritten digits from 0 to 9 which belong to 10 different categories. It contains 60,000 training samples and 10,000 training samples. Each image is a $28\times 28$ pixels grey, handwritten digital image with white text on a black background. Fashion MNIST [45] is a clothing image dataset which contains 10 classes of items such as T-shirt, dress, and bag. It has the same training and test samples as the MNIST dataset, which are summarized in detail in Table 1. We evaluate our algorithm FedADT on two datasets by training a Convolutional Neural Network (CNN) as in [2], which includes two convolution layers and two pooling layers followed by a fully connected layer with more than a million parameters in total.

In the experiments, we use 10 nodes and a central server to mimic the federated training setting. Each node trains a local model and uploads it to central server periodically. The central server generates a global model by aggregating local models. Here, the number of local updates is set to 5, and the local batch size is chosen from {5, 16, 64} to test for the best performance. We suppose each node takes part in the training at each communication round. We select the learning rate from $\{0.1,0.01,0.001,0.00001,0.00002\}$ for the best performance, and set ${\beta }_1=0.9$, ${\beta }_2=0.99$. For baseline algorithms, we search the learning rate from the same range as above, i.e.,$\{0.1,0.01,0.001,0.00001,0.00002\}$. We set the local update number as 5 with a batch size of 16. For comparison purposes, we assume that each client has the same neural network model which is trained by different algorithms. In addition, we account for two ways of partitioning data over nodes as in [2], i.e., data homogeneity (IID setting), where the date is uniformly distributed to 10 nodes, and data heterogeneity (Non-IID setting), where the data is shuffled by digit label and then assigned to 10 nodes. Both of these divisions are balanced. All the experiments in the article are performed on a workstation with two Intel(R) Xeon(R) Silver 4114 CPUs and two NVIDIA GeForce GTX 1080 Ti GPUs, and the algorithms are implemented in PyTorch framework, which is a popular deep learning training library.

We use two common metrics: training loss and test accuracy in federated learning and plot the learning curves as increasing communication rounds to verify the performance of different methods. We conduct 1000 rounds on two datasets to compare FedADT with naive local PID and local SGD methods, which use PID [21] and SGD as the local optimizers for each node. The results are illustrated in Figure 2, which exhibits the loss curves of three federated optimization algorithms over MNIST and Fashion MNIST datasets, respectively. We can see that the proposed FedADT method consistently achieves a faster convergence performance compared with the two baseline methods both under IID and Non-IID data setting of two datasets. However, for the bottom-row of Figure 2 with the data heterogeneity among different nodes, the learning curves of naive local PID and local SGD oscillate and are unstable, which slow down the convergence and the global models require more training rounds so as to obtain the desired results. A reasonable explanation is that FedADT uses the differential term in local decent direction, as well as the adaptive learning rate in the update process of the model.

Figure 3 shows the advantage of FedADT in terms of test accuracy under different data distributions over MNIST and Fashion MNIST datasets. As expected, all the algorithms can achieve almost similar accuracy in training the network under IID data setting on both datasets. Compared with the baseline methods, the proposed algorithm achieves the highest accuracy, 99.13% and 91.98%, first on the two datasets, respectively. For the data heterogeneity among different nodes, as shown in the bottom row of Figure 3, we have observed that the best accuracy of FedADT is 4.41% and 12.25% higher than that of naive local PID (94.37%) and local SGD (86.53%) on the MNIST dataset, respectively. On the Fashion MNIST dataset, the best accuracy of FedADT is 10.87% and 12.65% higher than that of naive local PID (74.02%) and local SGD (72.24%), respectively.

We use Receiver Operating Characteristic Curve (ROC) and Area Under Curve (AUC) as evaluation standards for the quality of the proposed algorithm. The experiment is performed by training the convolutional neural network on the Fashion MNIST dataset with Non-IID data, and the experimental result is shown in Figure 4. It demonstrates the characteristic curve and AUC area for 10 different clothing items. We can observe that class 1 is best identified and has the biggest AUC value of 0.999, and class 9 and 5 have similar results as class 1, whereas class 6 is not well recognized, and has the smallest AUC value, 0.950.

In this paper, we focus on the federated learning, where multiple nodes or clients are jointly modeled without exchanging their privacy data. During the local model training stage, the client side generally adopts a stochastic gradient-based algorithm, which is sensitive to step size and suffers from slow convergence performance. Inspired by the control theory, we first propose a federated adaptive learning method based on the derivative term which acts in a similar role as in the PID controller. We utilize the first difference of gradient to estimate the derivative term, which reflects the instant variation of gradient and stands for the future information. We provide a convergence guarantee for our proposed algorithm; in particular, when $E\le \mathcal{O}\left(\frac{{\left(dT\right)}^{1/4}}{n^{3/4}}\right)$ and $T>nd$, the convergence rate of $\mathcal{O}(1/\sqrt{nT})$ is achieved for non-convex objective functions. Finally, the experiments are performed under different data distribution cases on MNIST and Fashion MNIST datasets. Specially, for the Fashion MNIST dataset, the training loss curve of our proposed algorithm declines fastest compared to other baseline methods, and the highest accuracies of 91.98% and 84.89% are attained under IID and Non-IID cases, respectively. Similar results are obtained on the MNIST dataset, which empirically verify the effectiveness of the proposed algorithm. In addition, the ROC curve is used to display the satisfactory result by predicting the categories of clothing on the Fashion MNIST dataset. In future work, we will consider the effective-communication mechanism without sacrificing the convergence rate. At the same time, privacy protection should be considered in the process of communication between clients and the central server.