Efficient Federated Learning Using Dynamic Update and Adaptive Pruning with Momentum on Shared Server Data

Figure 1 shows the training process of FedDUMAP, where the FL system consists of a server and $N$ edge devices. The description of major notations is summarized in Table 1. We consider a powerful server and modest devices. The shared insensitive data is stored on the server, and the distributed sensitive data is dispersed in edge devices. Both the server data and the device data can be utilized to update the model during the FL training process of FL. Let us assume that a dataset $D_{k}=\{x_{k,j},y_{k,j}\}_{j=1}^{n_{k}}$, composed of $n_{k}$ samples, resides on Device $k$. We take $D_{0}$ as the server data, and $D_{k},k\neq 0$ as the device data. $x_{k,j}$ refers to the $j$-th sample on Device $k$, while $y_{k,j}$ represents the corresponding label. We utilize $\mathcal{X}$ to represent the set of samples and $\mathcal{Y}$ to represent the set of labels. Then, we formulate the objective of the training process as follows:

where $F_{k}(w)\triangleq\frac{1}{n_{k}}\sum_{\{x_{k,j},y_{k,j}\}\in\mathcal{D}_{k}}f(w,x_{k,j},y_{k,j})$ represents the loss function on Device $k$ with $f(w,x_{k,j},y_{k,j})$ capturing the error of the inference results based on the model with the sample pair $\{x_{k,j},y_{k,j}\}$ and $w$ refers to the global model.

With the distributed non-IID data in FL (mcmahan2021advances, ), we utilize the Jensen–Shannon (JS) divergence (fuglede2004jensen, ) to quantize the non-IID degree of a dataset as shown in Formula 2:

where $P_{m}=\frac{1}{2}(P_{k}+\overline{P})$, $\overline{P}=\frac{\sum_{k=1}^{N}n_{k}P_{k}}{\sum_{k=1}^{N}n_{k}}$, $P_{k}=\{P_{k}(y)|y\in\mathcal{Y}\}$, $P_{k}(y)$ refers to the possibility that a sample is related to Label $y$, and $\mathcal{D}_{KL}(\cdot||\cdot)$ corresponds to the Kullback-Leibler (KL) divergence (kullback1997information, ) as defined in Formula 3:

When the distribution of a dataset differs from that of the global dataset more significantly, the corresponding non-IID degree becomes higher. Although the raw data is not transferred from devices to the server or from the server to the devices in FedDUMAP, we assume that $P_{k}$ and $n_{k}$ incur little privacy concern (lai2021oort, ), and can be transferred from devices to the server before the FL training process. When this statistical meta information, i.e., $P_{k}(y)$, cannot be shared because of restrictions, we can utilize gradients to calculate the data distribution of the dataset on each device (yang2022improved, ).

The FL training process contains multiple rounds in FedDUMAP, each of which consists of 6 steps. In Step ①, a set of devices ($\mathscr{D}^{t}$ with $t$ referring to the round number) is randomly selected to participate in the training process of Round $t$ and the server broadcasts the global model to the device of $\mathscr{D}^{t}$. Afterward, each device updates the received model with the adaptive momentum-based optimization based on the local data in Step ②. Then, in Step ③, the selected devices upload updated models to the server, and the server aggregates all the received models based on FedAvg (mcmahan2017communication, ) in Step ④. In addition, the server updates (see details of the server update in Section 3.2) aggregated model utilizing the shared insensitive server data and the adaptive momentum-based optimization method (see details in Section 3.3) in Step ⑤. Furthermore, in a predefined round, the server performs model pruning based on the accuracy of the global model and the non-IID degrees (see details in Section 3.4). While Steps ① - ④ resemble those within traditional FL, we propose exploiting the insensitive server data to improve the performance of the global model (in FedDU) with adaptive momentum-based optimization (in FedDUM) on both the server and device sides (Step ⑤) and an adaptive model pruning method (in FedAP) to improve the efficiency of the training process (Step ⑥).

In this section, we explain the design of FedDU, which exploit both the shared insensitive server data and the distributed sensitive device data to adjust the global model. We exploit the non-IID degrees and the model accuracy to determine the weights of aggregated model and those of the normalized gradients based on the server data in FedDU, so as to avoid over-fitting to the server data.

We assume that the size of the insensitive server data is significantly bigger than the dataset within a single device. In this case, the straightforward combination of the aggregated model from the devices and the gradients calculated based on the shared insensitive server data leads to inferior accuracy due to objective inconsistency (wang2020tackling, ). Inspired by (wang2020tackling, ), we normalize the gradients generated using the server data to deal with this issue. The model aggregation in FedDU is formulated as Formula 4.

where $w^{t}$ refers to the weights of the global model in Round $t$, $w^{t-\frac{1}{2}}$ is the weights of the aggregated model from the selected devices in Round $t$ as formulated in Formula 5 (mcmahan2017communication, ), $\tau_{eff}^{t-1}$ corresponds to the size of effective steps for normalized gradients calculated based on the shared insensitive server data as shown in Formula 7, $\eta$ represents the learning rate, $\overline{g}_{0}^{t}(\cdot)$ refers to the normalized gradients generated based on the server data in Round $t$ as formulated in Formula 6.

where $n^{\prime}=\sum_{k\in\mathscr{D}^{t}}n_{k}$ refers to the size of the dataset on all the selected devices, $g_{k}^{t-1}(\cdot)$ corresponds to the gradients generated using the dataset on Device $k$.

where $w^{t-\frac{1}{2},i}$ represents the parameters of the global model after aggregating the gradients based on the server data, at $i-$th iteration of Round $t$, $g_{0}^{(t-1)}(\cdot)$ refers to the gradients calculated using the server data, and $\tau=\lceil\frac{|n_{0}|E}{B}\rceil$ corresponds to the number of iterations of Round $t$, $E$ is the number of local epochs, $B$ is the batch size. Each round corresponds to multiple iterations with a mini-batch of sampled shared insensitive server data.

While $\tau_{eff}^{t}$ has a significant impact on the training process, we dynamically choose an effective step size with the consideration of the model accuracy, the non-IID degrees of both the server data and the device data, with the number of rounds, as formulated in Formula 7.

where $acc^{t}$ represents the model accuracy calculated based on the server data in Round $t$, i.e., $w^{t-\frac{1}{2}}$ defined in Formula 5, $n_{0}$ corresponds to the size of the shared insensitive server data, $\mathcal{D}(\cdot)$ is shown in Formula 2, $\overline{P^{\prime}}^{t}=\frac{\sum_{k\in\mathscr{D}^{t}}n_{k}P_{k}}{\sum_{k\in\mathscr{D}^{t}}n_{k}}$ refers to the distribution of the distributed sensitive data in all the chosen devices in Round $t$, $P_{0}$ is the distribution of the shared insensitive server data, $decay\in(0,1)$ is utilized to ensure the convergence of the global model towards the solution of Formula 1, and $\mathcal{C}$ corresponds to a hyper-parameter. $f^{\prime}(acc)$ is calculated using $acc$. $acc$ is small at the beginning of the training. In this stage, the value of $f^{\prime}(acc)$ should be prominent so as to take advantage of the server data to improve the accuracy of the global model. Then, at a late stage of the training process, $f^{\prime}(acc)$ should be small to attain the objective defined in Formula 1 while avoiding over-fitting to the shared insensitive server data.

When the distribution of server data resembles that of the overall device data, i.e., $\mathcal{D}(P_{0})$ is small, or the distribution of the data of the chosen devices differs much from the overall device data, i.e., $\mathcal{D}(\overline{P^{\prime}}^{t})$ is significant, we take a significant value of $\tau_{eff}^{t}$ to improve the accuracy of the global model. Furthermore, the data size improves the importance of the dataset, as well. In addition, the weight of the device data is $\frac{n^{\prime}}{\mathcal{D}(\overline{P^{\prime}}^{t})}$ and that of the server data is $\frac{n_{0}}{\mathcal{D}(P_{0})}$. Finally, we have $\frac{\frac{n_{0}}{\mathcal{D}(P_{0})}}{\frac{n_{0}}{\mathcal{D}(P_{0})}+\frac{n^{\prime}}{\mathcal{D}(\overline{P^{\prime}}^{t})}}$ as the importance of the server data, which is the second part of Formula 7.

Algorithm 1 explains FedDU. The model is updated using the local dataset on each chosen device (Lines 1 - 3). Afterward, the models are aggregated using Formula 5 (Line 4). Finally, the aggregated model is updated utilizing the shared insensitive server data based on Formula 4 (Line 5).

Let us assume that the expected squared norm of stochastic gradients on the server is bounded, i.e., $\mathop{\mathbb{E}}||\overline{g}_{0}||^{2}\leq G^{2}$. Then, in Round $t$, $\tau_{eff}^{t}\leq C\cdot decay^{t}\cdot\tau$. Afterward, the server update term can be less than $C\cdot decay^{t}\cdot\tau\cdot\eta\cdot G$. In this case, after sufficient steps, the server update becomes negligible, i.e., $\mathop{\mathbb{E}}[\lim_{t\to\infty}\tau_{eff}^{t-1}\eta\overline{g}_{0}^{t-1}(w^{t-\frac{1}{2}})]\leq\lim_{t\to\infty}C\cdot decay^{t-1}\cdot\tau\cdot\eta\cdot G=0$, with $0<decay<1$. Finally, FedDU degrades to FedAvg (mcmahan2017communication, ), the convergence of which is guaranteed with a decaying learning rate $\eta$ (Li2020On, ; Zhou2018On, ).

In order to further improve the accuracy of the global model, we exploit momentum within the server update on the server and within the local iteration on each device. In this section, we propose a simple adaptive momentum approach, i.e., FedDUM, which enables the optimization on both the server and the devices without additional communication cost.

The centralized SGDM can be formulated as:

where $m$ is momentum. A simple application of the momentum into the FL environment is to decompose the SGDM to each device while aggregating the momentum as formulated in Formula 9.

In order to extend the optimization into multiple local epochs, we can formulate the process as follows:

where ${m^{\prime}}_{k}^{t,t^{\prime}}$ represents the local momentum within Device $k$ at global Round $t$ and local Iteration $t^{\prime}$, $\beta^{\prime}$ is the local weight for gradients, ${g^{\prime}}_{k}$ refers to the gradients with the parameters of the model ${w^{\prime}}_{k}^{t,t^{\prime}}$, and $\eta^{\prime}$ is the local learning rate. Please note that ${g^{\prime}}_{k}$ and $\eta^{\prime}$ are local hyper-parameters on each device and are different from those in Formulas 8 and 9, which refer to global parameters on the server. We have the first momentum equal to the global momentum, i.e., ${m^{\prime}}_{k}^{t,0}=m^{t}$, and the first parameters of the model equal to the global parameters, i.e., ${w^{\prime}}_{k}^{t,0}=w^{t}$. However, this simple decomposition may deviate from the momentum of the centralized optimizer and harm the local training (jin2022accelerated, ). In addition, the communication of momentum, i.e., ${m^{\prime}}_{k}^{t,E-1}$ and $m^{t}$, may incur significant communication cost. Inspired by (wang2021local, ), we restart the local adaptive optimization and set the momentum buffers to zero at the beginning of local iterations. Furthermore, we calculate the difference between the updated model and the original model as the total gradients as defined in Formula 12. Then, we can utilize Formula 11 below to perform the local iteration in each selected device, which replaces Formula 5 in Algorithm 1 with adaptive momentum:

with ${m^{\prime}}_{k}^{t,0}=0$ and ${w^{\prime}}_{k}^{t,0}=w^{t}$. This operation only requires the parameters of the model transferred from the server to devices without additional communication costs compared with FedAvg. Afterward, we exploit Formula 12 below to calculate the gradients on the server, and utilize Formula 8 to replace Formula 4 in Algorithm 1 with adaptive momentum:

where $\tau_{eff}^{t-1}\eta\overline{g}_{0}^{(t-1)}(w^{t-\frac{1}{2}})$ is the same as that in Formula 4. This only requires the transfer of the parameters of the local updated model, which does not incur additional communication cost compared with FedDU.

FedDUM is shown in Algorithm 2. The local model update is performed according to Formula 11 in parallel in each selected device (Lines 1 - 3). Then, the aggregated gradient is calculated using Formula 12 (Line 4). Afterward, the new global model is updated using the gradients and momentum according to Formula 8 (Line 5).

In order to analyze the adaptive optimization on devices, we have the following assumptions:

Then, we can get the following theorem:

In the experimentation, we set the local epoch to a small number, e.g., 5, to reduce the deviation between local momentum and centralized optimization.

In this section, we present our layer-adaptive pruning method, i.e., FedAP. FedAP considers the diverse features of each layer, the non-IID degrees of both the sensitive shared server data and the distributed sensitive device data, to remove useless filters in the convolutional layers while preserving the accuracy. FedAP can reduce the communication overhead and the computational costs of the training process. Please note that FedAP is performed only once on the server in a predefined round.

Algorithm 3 details FedAP. On the server and each device (Lines 2 - 4), we calculate a proper pruning rate based on the shared sensitive server data or the distributed sensitive device data (Line 3). We denote the initial model $W_{k}$ on the server or Device $k$. At Round $T$, the updated model is denoted by $W^{\prime}_{k}$. Then, the difference between the current model and the original one is $\Delta_{k}=W_{k}-W^{\prime}_{k}$. Afterward, we can calculate the Hessian matrix, i.e., $H(W^{\prime}_{k})$. We sort the eigenvalues of $H(W^{\prime}_{k})$ in ascending order, i.e., $\{\lambda^{m}_{k}|m\in(1,d_{k})\}$ with $m$ referring the index of an eigenvalue and $d_{k}$ referring to the rank of the Hessian matrix. We take $B_{k}(\Delta_{k})=H(W^{\prime}_{k})-\triangledown L(\Delta_{k}+W^{\prime}_{k})$ as a base function with $\triangledown L(\cdot)$ corresponding to the gradients. We denote the Lipschitz constant of $B_{k}(\Delta_{k})$ as $\mathscr{L}_{k}$. Inspired by (zhang2021validating, ), we take the first $m_{k}$, which meets the condition of $\lambda_{m_{k+1}}-\lambda_{m_{k}}>4\mathscr{L}_{k}$, to reduce accuracy degradation. Then, we can calculate the proper pruning rate by $p^{*}_{k}=\frac{m_{k}}{d_{k}}$, which is the ratio between the size of pruned eigenvalues and that of all the eigenvalues. As the proper pruning rate significantly differs in diverse devices due to the non-IID distribution, we utilize Formula 15 to generate an aggregated proper pruning rate for the global model (Line 5).

where $\epsilon$ represents a small value to avoid the division of zero. We take a global threshold value ($\mathscr{V}$) calculated based on the aggregated proper pruning rate to serve as a baseline value for generating the pruning rate of each layer. $\mathscr{V}$ is the absolute value of the $\lfloor R*p^{*}\rfloor$-th smallest parameter in all the parameters (Lines 6 and 7). Afterward, for each convolutional layer (Line 8), we calculate the proper pruning rate by calculating the ratio between the number of parameters with inferior absolute values than $\mathscr{V}$ and the number of all the parameters (Lines 9 - 11). Inspired by (lin2020hrank, ), we remove the filters corresponding to the smallest ranks in feature maps (the output of filters) based on the proper pruning rate in the layer (Lines 12 - 15). We take $R_{l}=\{r_{l}^{j}|j\in(1,d_{l})\}$ to represent the ranks of feature maps at Layer $l$, where $d_{l}$ refers to the number of filters in this layer. As the feature maps are almost the same for a given model (lin2020hrank, ), we take the ranks calculated based on the server data and perform the pruning operations on the server because of its powerful computation capacity. We calculate the feature maps in $R_{l}$ in Line 12 and sort them in Line 13. We keep the filters having the last $d_{l}-\lfloor p^{*}_{l}*d_{l}\rfloor$ ranks in the sorted $R^{l}$, so as to attain the highest pruning rate $p_{l}\leq p^{*}_{l}$ (Line 14). In the end, we take the preserved filters in the original model as the layer of the pruned model (Line 15). When there is no available server data, the pruning process can also be performed on the server based on the rank values $R_{l}$ calculated using the dataset on a device, which are transferred to the server, with the execution of Line 12 being carried out at that device.

We evaluate FedDUMAP using an FL system consisting of a parameter server and $100$ devices. In each round, we randomly choose $10$ devices. We utilize two real-life datasets, i.e., CIFAR-10 and CIFAR-100 (krizhevsky2009learning, ), in the experimentation. We exploit four models, i.e., a simple synthetic CNN network (CNN), ResNet18 (ResNet) (he2016deep, ), VGG11 (VGG) (simonyan2014very, ), and LeNet5 (LeNet) (lecun1998gradient, ). The experimentation of CNN, VGG, and LeNet is carried out on both CIFAR-10 and CIFAR-100, while that of ResNet is performed on CIFAR-100.

CNN consists of three $3*3$ convolution layers, a fully connected layer, and a final softmax output layer. The first layer contains 32 channels, while the second and third layers have 64 channels. The first and the second layers are followed by $2*2$ max pooling. The fully connected layer contains $64$ units and exploits ReLu activation. CNN contains 122570 parameters in total. The mini-batch size is set to $10$ for the local model update, and the selected device executes $E=5$ local epochs. The local learning rate $\eta$ is $0.1$ with a decay rate of $0.99$. $\mathcal{C}$ is set to 1. The experimentation is carried out using 33 Tesla V100 GPUs to simulate an FL environment composed of a parameter server and 100 devices. We carry out the pruning process in Round 30 in all the experimentation with adaptive pruning. In addition, we exploit the Savitzky–Golay filter (press1990savitzky, ) to smooth the accuracy in figures.

The CIFAR-10 dataset contains 60000 images (50000 for training and 10000 for testing), each corresponding to one of ten categories. We take 40000 images in the training dataset as device data and randomly select $p$ * 40000 images from the remaining 10000 images as server data, with $0\%<p<25\%$. $p$ represents the ratio between the size of the shared insensitive server and that of the distributed sensitive device data. For the non-IID setting, we sort the device data according to the label and then divide these data evenly into $2*100$ fractions. Each device is randomly assigned $2$ fractions, and most devices contain the data with 2 labels. We utilize the same method to handle the CIFAR-100 dataset.

We first conduct comparison of FedDU with FedAvg, FedKT, FedDF, Data-sharing, and Hybrid-FL in terms of model accuracy. Then, we compare FedDUM with server-side momentum, device-side momentum, and FedDA. Afterward, we show the advantages of the adaptive pruning method, i.e., FedAP, with the comparison of FedAvg, IMC, HRank, and PruneFL, in terms of model efficiency. Last but not least, we demonstrate that FedDUMAP, which consists of both FedDUM and FedAP, significantly outperforms the eight state-of-the-art baselines in terms of accuracy, efficiency and computational cost. Finally, we present our ablation study.