AdaptiveFL: Adaptive Heterogeneous Federated Learning for Resource-Constrained AIoT Systems

Algorithm 1 details the implementation of AdaptiveFL. Before FL training, Lines 1-2 initialize the curiosity table $T_{c}$ and the resource table $T_{r}$. Lines 3-29 present the details of FL training for each round. Line 4 splits the global model $M$ into submodels in different size levels (i.e., small, medium and large) and stores them in model pool $R=\{m_{S_{p}},m_{S_{p-1}},\ldots,m_{M_{2}},m_{M_{1}},m_{L_{1}}\}$, where hyperparameters $p$ is the number of submodels in each level except $L$ level, and it should be noted that the large model $m_{L_{1}}$ is unpruned which is equivalent to the global model. Lines 6-27 present the FL training process of the models waiting for training, where the loop “for” is parallel. In Line 7, the function RandomSel(.) is to randomly select a model $m_{i}$ from the model pool $R$. In Line 8, the function ClientSel(.) is to select a suitable client $c_{i}$ for model $m_{i}$ from client set $C$ based on RL-table $T_{c}$ and $T_{r}$. In Line 9, the function LocalTrain(.) is to dispatch the model $m_{i}$ to the selected client $c_{i}$ for local training, and return the trained model $m_{i}^{\prime}$ with the local data size $\left|d_{c_{i}}\right|$ back to the server. Line 10 stores $m_{i}^{\prime}$ and $\left|d_{c_{i}}\right|$ to array $ML_{back}$ and array $Len$, respectively, which are used for aggregation later. In addition, Lines 12-13 and Lines 14-26 present the updating process of $T_{c}$ and $T_{r}$, respectively. In Lines 12-13, we updated the selection times for the level of the send and back models in $T_{c}$, respectively, where $type\left(m_{i}\right)$ means the level of model $m_{i}$, e.g., $type\left(m_{S_{p}}\right)$ return the size level $S$. As for the update of $T_{r}$, we consider the following two cases: i) In Lines 15-18, since no pruning is done locally at the client $c_{i}$, which means that the resource capacity $\Gamma_{c_{i}}\geq\operatorname{size}\left(m_{i}=m_{i}^{\prime}\right)$, so we perform an increment operation for the training score in the table whose model size is larger than $m_{i}$; ii) In Lines 20-25, it shows that $\operatorname{size}\left(m_{i}^{\prime}\right)\leq\Gamma_{c_{i}}\leq\operatorname{size}\left(\hat{m_{i}^{\prime}}\right)$, $\hat{m_{i}^{\prime}}$ here is the nearest greater model with $m_{i}^{\prime}$ in $R$. Thus, we use a penalty term $\tau$ to reduce the training score of the heterogenous model that is larger than $\hat{m_{i}^{\prime}}$, while increasing the training score of the model $m_{i}^{\prime}$.

To enable devices to prune models according to their available resources adaptively, we adopt a width-wise pruning mechanism where the pruned model can be trained directly without additional adapters or parameters. Inspired by the observations in (li2016pruning, ), we prefer to prune the parameters of deep layers, which enables large models trained by insufficient data to achieve higher performance. Specifically, our fine-grained width-wise model pruning mechanism is controlled by two hyperparameters, i.e., the width pruning ratio $r_{w}$ and the index of the starting pruning layer $I$, respectively, where adjusting $r_{w}$ can significantly change the model size while adjusting $I$ can fine-tune the model size.

Width-Wise Model Pruning ($r_{w}$). To generate multiple models of different sizes, the cloud server prunes partial kernels in each layer of the model, where the number of kernels pruned is determined by the width pruning ratio $r_{w}\in(0,1]$. Specifically, we assume that $W_{g}$ is the parameter of the global model $M_{g}$, $d_{k}$ and $n_{k}$ denote the output and input channel size of the $k^{th}$ hidden layer of $M_{g}$, respectively. Then the parameters of the $k^{th}$ hidden layer can be denoted as $W_{g}^{k}\in\mathbb{R}^{d_{k}\times n_{k}}$. With a width-wise pruning ratio $r_{w}$, the pruned weights of the $k^{th}$ hidden layer can be presented as $W_{r_{w}}^{k}=W_{g}^{k}[:d_{k}\times r_{w}][:n_{k}\times r_{w}]$.

Layer-Wise Model Adjustment ($I$). To address performance fluctuations caused by uncertainty, our pruning mechanism supports fine-tuning the model size by adjusting the index of the starting pruning layer $I$. Note that to ensure that heterogeneous models share shallow layers, the index of the starting pruning layer must be set larger than the specific threshold $\tau$. Specifically, assume that $I\geq\tau$, the weights of the $k^{th}$ layer can be presented as $W_{r_{w}}^{k}=W_{g}^{k}[:d_{k}][:n_{k}]$ when $k\leq I$, and which can be presented as $W_{r_{w}}^{k}=W_{g}^{k}[:d_{k}\times r_{w}][:n_{k}\times r_{w}]$ when $k>I$.

Available Resource-Aware Pruning. To prevent failed training caused by limited resources, our pruning mechanism supports each device in pruning the received model adaptively according to its available resources. Specifically, assume that the available resource capacity of the device is $\Gamma$, the weight of the received model is $W$ and $I\geq\tau$, and the width-wise pruning ratio $r_{w}$ and the index of the starting pruning layer $I$ can be determined as follows:

Due to uncertainty and privacy concerns, the cloud server cannot obtain available resource information for AIoT devices. To avoid communication waste caused by dispatching unsuitable models, we propose an RL-based device selection strategy. By utilizing the information of historical dispatching and the corresponding received model $\langle m_{i},m_{i}^{\prime}\rangle$ of clients, RL can learn the information about the available resources of each device. Based on the learned information, RL can learn a strategy to select suitable devices for each heterogeneous model wisely.

Problem Definition. In our approach, the client selection process can be regarded as a Markov Decision Process (hu_rtss2021, ), which can be presented as a four-tuple $MDP=\langle\mathcal{S},\mathcal{A},\mathcal{F},\mathcal{R}\rangle$ as follows:

• •

  $\mathcal{S}$ is a set of states. We use a vector $s_{t}=\left\langle D_{t},S,C_{t},T_{c},T_{r}\right\rangle$ to denote the state of AdaptiveFL, where $D_{t}$ denotes the set of submodels that wait for dispatching, $S$ is the list of the size information of all submodels in the model pool, $C_{t}$ indicates the set of clients involved, $T_{c}$ and $T_{r}$ are the curiosity table and the resource table, respectively.

• •

  $\mathcal{A}$ is a set of actions. At the state of $s_{t}=\left\langle D_{t},S,C_{t},T_{c},T_{r}\right\rangle$, the action $a_{t}$ aims to select a suitable client $c_{i}\in C_{t}$ for the candidate model $m_{i}\in D_{t}$.

• •

  $\mathcal{F}$ is a set of transitions. It records the transition $s_{t}\stackrel{{\scriptstyle a_{t}}}{{\longrightarrow}}s_{t+1}$ with the action $a_{t}$.

• •

  $\mathcal{R}$ is the reward function. We combine the values in the resource table $T_{r}$ and the curiosity table $T_{c}$ of each client as the reward to guide the selection on this round.

Resource- and Curiosity-Driven Client Selection. Since there is an implicit connection between the model size with the resource budget, the model returned by the client can be used to determine the available resource range of the device. Specifically, AdaptiveFL uses a client resource table $T_{r}$ to record the historical training score for each heterogeneous model on each client, where a higher score indicates that the client has a higher success rate in training the corresponding model. The resource reward for client $c$ on submodel $m_{i}$ is measured as follows:

To balance the training times of the same model level on different clients, we utilize curiosity-driven exploration (hu2023accelerating, ; hu2023gitfl, ) as one of the reward evaluation strategies, while the client who is selected fewer times on a size level of the model will get higher curiosity rewards. AdaptiveFL uses the curiosity table $T_{c}$ to record the selection times of each client on a type of model, and performs Model-based Interval Estimation with Exploration Bonuses (MBIE-EB) (bellemare2016unifying, ) to calculate the curiosity reward as follows:

where $T_{c}[type\left(m_{i}\right)][c]$ indicates the total selection number of model type $type\left(m_{i}\right)$ on client $c$. To avoid the higher success rate of the large client leading to the lower probability of other clients being selected, we set the upper success rate of 50%, and the selection of clients whose success rate is beyond 50% will be determined by the curiosity reward. Consequently, the final reward for each client on the model $m_{i}$ is calculated by combining resource reward $R_{s}$ and curiosity reward $R_{c}$ as follows:

In conclusion, based on the final reward, the probability that the client $c$ is selected for model $m_{i}$ is:

In our model pruning mechanism, since all the submodels are pruned based on the same full global model, the cloud server can update the global model by aggregating all the received heterogenous submodels according to the corresponding index of their parameters in the full model. Algorithm 2 details the aggregation process of our approach. Line 1 is the initialization of the process. Lines 2-17 update the parameters of each layer in the model. Line 3 initializes the variable $L_{w}$, which is used to count the total number of the training data size for each parameter. In Lines 4-8, the model parameters of each client are added with weights, which is the size of local data. For each client, the uploaded model often lacks some parameters compared to the complete model. Lines 10-16 take the average of the updated parameters. Note that if some parameters are not included in any uploaded model, they will keep their original values unchanged, which is shown in Line 14.

Data Settings. We conducted experiments on three well-known datasets, i.e., CIFAR-10, CIFAR-100 (CIFAR, ), and FEMNIST (LEAF, ). For both CIFAR-10 and CIFAR-100, we assumed that there were 100 clients participating in FL. For FEMNIST, there were 180 clients involved. In each round, 10% of the clients will be selected for local training. We considered both IID and non-IID scenarios for CIFAR-10 and CIFAR-100, where we adopted the Dirichlet distribution to control the data heterogeneity. Here, the smaller the coefficient $\alpha$, the higher the heterogeneity of the data. Note that FEMNIST is naturally non-IID distributed.

Device Heterogeneity Settings. To simulate the heterogeneity of devices, we set up three types of clients (i.e., weak, medium, and strong clients) and three levels of models (i.e., small, medium, and large models), where weak devices can only accommodate weak models, medium devices can train medium or small models, while strong devices can accommodate models of any type. For the following experiments, we set the proportion of weak, medium, and strong devices to 4: 3: 3 by default. To show the generality of our AdaptiveFL framework, we conducted experiments based on two widely used models (i.e., VGG16 (VGG, ) and ResNet18 (ResNet, )), where Table 1 shows the split settings of VGG16.

We compared AdaptiveFL with four baseline methods, i.e., All-Large (FedAvg, ), Decoupled (FedAvg, ), HeteroFL (heterofl, ), and ScaleFL (ScaleFL, ). For All-Large, we trained the $L_{1}$ model with all clients under the classic FedAvg (FedAvg, ). For Decoupled, we trained separate models (i.e., $L_{1}$, $M_{1}$, $S_{1}$ models) for each level using the available data of affordable clients. For HeteroFL and ScaleFL, we created their corresponding submodels at different levels. Table 2 shows the comparison results, where the notations “avg” and “full” denote the average accuracy of submodels at different levels (i.e., $L_{1}$, $M_{1}$, $S_{1}$) and the accuracy of the global model, respectively.

Global Model Performance. From Table 2, we can find that Decoupled has the worst inference performance in all the cases, since its submodels are only aggregated with the models at the same levels. However, AdaptiveFL can achieve up to 2.95% and 3.12% better inference than the second-best methods for ResNet18 and VGG16, respectively. Note that AdaptiveFL can achieve better results compared with All-Large, indicating that AdaptiveFL can improve the FL performance in non-resource scenarios.

Submodel Performance. Figure 2 shows the learning trends of all methods on CIFAR-10/100 based on VGG16, where solid lines represent the “avg” accuracy of submodels. We can find that AdaptiveFL can achieve the best inference performance with the least variations for both non-IID and IID scenarios. For different heterogeneous FL methods, Figure 3 presents the shapes of VGG16 submodels together with their test accuracy information. We can find that AdaptiveFL consistently outperforms other heterogeneous FL methods under the premise of satisfying the resource constraints, which indicates the effectiveness of our fine-grained width-wise model pruning mechanism. Interestingly, we can find that 1.0$\times$ large models of HeteroFL and ScaleFL perform worse instead their 0.25$\times$ small counterparts. Conversely, as the model size increases, AdaptiveFL can achieve better results, indicating that it can smoothly transfer the knowledge learned by the submodels into large models.

Numbers of Participating Clients. To evaluate the scalability of AdaptiveFL, we conducted experiments considering different numbers of participating clients, i.e., $K$ = 50, 100, 200, and 500, respectively, on CIFAR-10 using ResNet18. Figure 4 compares AdaptiveFL with three baselines within a non-IID scenario ($\alpha=0.6$), where AdaptiveFL can always achieve the highest accuracy.

Proportions of Different Devices. Table 3 shows the performance of AdaptiveFL with different proportions (i.e., 8:1:1, 1:8:1, 1:1:8, and 4:3:3) of weak, medium, and strong devices on CIFAR-10. We can find that AdaptiveFL can achieve the best test accuracy in all cases. Note that as the proportion of strong devices increases, the global model performance of all FL methods improves.

Ablation of Fine-grained Pruning. To evaluate both fine- and coarse-grained pruning, we set $p$ to 3 and 1 for each level, respectively. Table 4 presents the ablation results of AdaptiveFL considering the effect of our fine-grained pruning method, showing that the fine-grained pruning method can achieve up to 9.38% inference accuracy improvements for AdaptiveFL. Note that the fine-grained methods can consistently achieve better inference results than their coarse-grained counterparts, since the fine-grained ones can better transfer the knowledge of small models to large models.

Ablation of RL-based Client Selection. To evaluate the effectiveness of our RL-based client selection strategy, we developed four variants of AdaptiveFL: i) “AdaptiveFL+Greedy” that always dispatches the largest model for each selected client; ii) “AdaptiveFL+Random” that selects clients for local training randomly; iii) “AdaptiveFL+C” that selects clients only based on curiosity reward; and iv) “AdaptiveFL+S” that selects clients only using resource rewards. Moreover, we use “AdaptiveFL+CS” to indicate the original AdaptiveFL implemented in Algorithm 1.

Figure 5 presents the ablation study results on CIFAR-100 with ResNet18 following IID distribution. To indicate the similarity between a sending model and its corresponding receiving model, we introduce a new metric called communication waste rate, defined as “$1-\sum(\operatorname{size}(\text{$ML_{back}$}))/\sum(\operatorname{size}(\text{$ML_{send}$}))$”. The lower the rate, the closer the two models are, leading to less local pruning efforts. From Figure 5, we can find that our approach can achieve the highest accuracy with low communication waste (second only to RL-S).

Based on our real test-bed platform, we conducted experiments on a non-IID IoT dataset (i.e., Widar (fedaiot, )) with MobileNetV2 (mobilenetv2, ) models. We assumed that the FL-based AIoT system has 17 devices, each training round involves 10 selected devices, whose detail heterogeneous configurations are shown in Table 5.

Figure  6 presents the AIoT devices used in our experiment and the comparison results obtained from our real test-bed platform. We can observe that AdaptiveFL achieves the best inference results even in real scenarios compared to all baselines.