Heroes: Lightweight Federated Learning with Neural Composition and Adaptive Local Update in Heterogeneous Edge Networks

Considering a client set $\mathcal{N}=\{1,2,\cdots,N\}$ coordinated by the PS, each client $n\in\mathcal{N}$ holds its local dataset $D_{n}=\{\zeta_{i}^{n}\}_{i=1}^{|D_{n}|}$, where $\zeta_{i}^{n}$ denotes a data sample from $D_{n}$. Further, we represent the loss function as $\mathcal{L}(\mathbf{x};\zeta)$, which measures how well the model $\mathbf{x}$ performs on data sample $\zeta$. Therefore, the local loss function of client $n$ is defined as:

The global loss function is a linear combination of all $N$ clients, and the goal of FL is to train a high-quality model $\mathbf{x}^{*}$ with minimum global loss function, which is defined as:

Suppose there are $H$ rounds of training in total. In each round $h\in\{1,2,\cdots,H\}$, the PS randomly selects a set of clients $\mathcal{N}^{h}\subseteq\mathcal{N}$ to participate in training and sends the fresh global model $\mathbf{x}^{h}$ to the specified clients, where $|\mathcal{N}^{h}|=K$. Then, each client $n\in\mathcal{N}^{h}$ updates the global model $\mathbf{x}^{h}$ over its local dataset $D_{n}$ for $\tau$ times, where each update is regarded as one local iteration and $\tau$ represents the local update frequency. Let $\mathbf{x}^{h}_{n}(t)$ denote the local model of client $n$ at iteration $t$ in round $h$. For the mini-batch stochastic gradient descent (SGD) [yu2019parallel] algorithm, a local iteration can be expressed as follows:

where $\xi^{n}$ denotes a random data batch from the local dataset $D_{n}$ and $\eta$ is the learning rate. Finally, the PS collects the updated local models from the participating clients and aggregates them to the latest global model for further training, i.e., $\mathbf{x}^{h+1}=\frac{1}{N}\sum_{n=1}^{N}\mathbf{x}^{h}_{n}(\tau)$.

To make full use of the limited resources, it is necessary to adjust the complexity of local model for each client based on its resource budgets [diao2020heterofl, horvath2021fjord, mei2022resource]. Herein, we follow Flanc [mei2022resource] to innovatively propose an enhanced neural composition technique to construct the models in different widths (i.e., the number of hidden channels in each weight) based on low-rank factorization [phan2020stable]. Concretely, since the weights in DNNs are usually over-parameterized [zou2019improved], each layer’s weight (e.g., convolution, fully connection) can be approximated as the product of two low-rank tensors, named neural basis $\mathbf{v}$ and coefficient $\mathbf{u}$. For example, let $\mathbf{w}\in\mathbb{R}^{k^{2}\times I\times O}$ represent a convolution weight, with kernel size $k$, input channel number $I$ and output channel number $O$. Let $p\in\{1,2,\cdots,P\}$ denote the weight width, where the shape of $p$-width weight $\mathbf{w}_{p}$ is $k^{2}\times pI\times pO$. Accordingly, $\mathbf{w}_{p}$ is approximated as follows:

When $k=1$, the above format represents the approximation of fully connection layer’s weight. The weight width is controlled by adjusting the size of coefficient, while the size of neural basis is constant. Specifically, the complete coefficient is divided into $P^{2}$ blocks, where the shape of each block is $R\times O$. We select $p^{2}$ blocks from the complete coefficient to form the reduced coefficient and compose (i.e., multiply) it with neural basis into a $p$-width weight.

To ensure sufficient training of every parameter in the global model, we control different coefficient blocks to be trained evenly. Thus, the selected blocks are currently the least trained ones. Notably, we measure each block’s training adequacy by the total number of local iterations it has experienced on all clients since round 1, i.e., total update times. In general, we balance the total update times of different coefficient blocks to ensure that each block is fully trained.

For example, as illustrated in Fig. 1, the coefficient $\mathbf{u}\in\mathbb{R}^{R\times(9\times O)}$ is divided into 9 blocks (i.e., $P$=3) and the number in each block represents its current total update times. To obtain a 2-width weight, we first extract the least trained 4 blocks (with the total update times of 6, 5, 7 and 8, respectively) from the complete coefficient and combine them into the reduced coefficient $\hat{\mathbf{u}}\in\mathbb{R}^{R\times(4\times O)}$. Then, $\hat{\mathbf{u}}$ is composed with the neural basis $\mathbf{v}\in\mathbb{R}^{k^{2}\times I\times R}$ into an intermediate tensor whose shape is $k^{2}\times I\times(4\times O)$. Finally, a 2-width weight $\hat{\mathbf{x}}\in\mathbb{R}^{k^{2}\times(2\times I)\times(2\times O)}$ is obtained by reshaping this intermediate tensor.

Compared to the FL schemes based on model pruning [diao2020heterofl], the neural composition technique enables every parameter in the global model to benefit from the full range of knowledge through the shared neural basis, improving the training performance. Different from Flanc [mei2022resource] with original neural composition, enhanced neural composition allows the coefficient of all sizes to get fully trained, accelerating the training process. To verify this improvement, we simulate a FL system with 100 clients to train the standard ResNet-18 model [he2016deep] over the ImageNet dataset [russakovsky2015imagenet]. The results in Table I indicate that with the given traffic consumption or completion time, the enhanced neural composition can improve the global model’s test accuracy by about 16.29% on average compared with the FL schemes based on model pruning (MP) [diao2020heterofl] and original neural composition (NC) [mei2022resource].

In most synchronous FL schemes [li2021talk, mei2022resource], the clients’ local update frequencies are identical and fixed in each round. However, due to client heterogeneity, strong clients have to wait for weak ones (i.e., stragglers) for global aggregation, incurring non-negligible waiting time and reducing the training efficiency significantly [liu2023adaptive, wang2020towards]. To evaluate the negative impacts of stragglers, we record the completion time for one training round of each client in the simulated FL system. As shown in Fig. 2(a), the strongest client completes one training round four times faster than the weakest client. In other words, about 70% of the strongest client’s time is idle and wasted.

To reduce clients’ idle waiting, some research [li2020federated, xu2022adaptive, li2022pyramidfl] proposes adjusting the local update frequencies for different clients to balance their completion time. On the one hand, the weaker clients perform fewer local iterations to reduce the impact of stragglers. On the other hand, the stronger clients utilize the idle time to perform more local iterations, which helps the model converge faster. We illustrate the completion time of each client with proper local update frequency in Fig. 2(b). It can be observed that almost every client’s time is fully utilized without idle waiting.

However, the neural composition and the local update are interconnected, and it is difficult to jointly assign proper coefficient blocks and local update frequencies for different clients. Specifically, both the computation time and the communication time of a specific client depend on the number of coefficient blocks, which affects the determination of its local update frequency. Meanwhile, since each client trains different blocks, various local update frequencies will also affect the balance of the blocks’ total update times.

In this section, we define the joint optimization problem of neural composition and local update frequency in FL training. Let $G(\mathbf{v}\cdot\mathbf{u})$ represent the number of floating-point operations (FLOPs) required to perform one local iteration for the composed model, and $G(\mathbf{v}\cdot\mathbf{u})$ depends on the size of coefficient $\mathbf{u}$. The more blocks the coefficient $\mathbf{u}$ includes, the wider the composed model $\mathbf{x}=\mathbf{v}\cdot\mathbf{u}$ and the more FLOPs required for model training. We formulate the time cost for one local iteration of client $n$ in round $h$ as follows:

where $q_{n}^{h}$ represents the speed to process the floating-operations of client $n$ in round $h$.

Since the download bandwidth is usually much faster than the upload bandwidth in typical WANs [zhan2020incentive], the download time of tensors can be negligible and we mainly focus on the upload time. Let $b_{n}^{h}$ denote the upload bandwidth of client $n$ in round $h$ and $E$ measure the size of a tensor. We formulate the communication time of client $n$ in round $h$ as:

Let $T_{n}^{h}$ denote the completion time of client $n$ in round $h$, which consists of the time cost for $\tau_{n}^{h}$ local iterations and the communication time. Due to the synchronization barrier of FL, the completion time of round $h$ depends on the slowest client, which can be defined as:

We adopt the average waiting time for participating clients in $\mathcal{N}^{h}$ to measure the impact of synchronization barrier in round $h$, which is defined as follows:

Let $c_{i}^{h}$ denote the total update times of $i$-th coefficient block, where $i\in\{1,2,\cdots,P^{2}\}$. The training consistency between different coefficient blocks in round $h$ can be reflected by the variance of set $\{c_{i}^{h}|\forall i\}$, denoted as:

In each round $h$, we aim to select appropriate $(p_{n}^{h})^{2}$ coefficient blocks and determine the optimal local update frequency $\tau_{n}^{h}$ for each participating client $n\in\mathcal{N}^{h}$, so as to accelerate the training process of FL. Accordingly, we define the optimization problem as follows:

$\mathop{\min}\sum_{h=1}^{H}T^{h}$

The first inequality expresses the convergence requirement, where $\epsilon$ is the convergence threshold of the training loss after $H$ rounds. The second set of inequalities indicates that the average waiting time of participating clients in each round $h$ should not exceed the given threshold $\rho$. The third set of inequalities bounds the variance of all coefficient blocks’ total update times in each round $h$, ensuring the balanced training among blocks. The fourth set of inequalities tells the feasible range of the number of coefficient blocks. The object of this optimization problem is to minimize the completion time of the FL training with the performance requirements (e.g., model convergence, average waiting time).

To solve the optimization problem in Eq. (22), we first approximate the convergence bound. Specifically, we adopt an upper bound $\beta^{2}$ for the coefficient reducing error, where $\alpha_{n}^{h}\leq\beta^{2}$. Besides, the minimum loss value is approximated as zero, i.e., $F(\mathbf{x}^{*})=0$. Accordingly, the convergence bound in Eq. (12) is formulated as follows:

Since $\tau$ is a positive variable, the convergence bound $G(H,\tau)$ is a convex function with respect to the local update frequency $\tau$. It can be derived that $G(H,\tau)$ will decrease as $\tau$ increases when $\tau\leq\small{\sqrt{\frac{12F(\mathbf{x}^{0})}{\eta^{2}HL(G^{2}+18\sigma^{2})}}}$. On the contrary, the trend of $G(H,\tau)$ and $\tau$ is the opposite.

To balance the completion time of heterogeneous clients, we first select the fastest client $l$ in round $h$ and let the completion time of other clients be approximately equal to that of client $l$, which can be formulated as follows:

Notably, the completion time $T_{l}^{h}$ of client $l$ is the largest among that of all participating clients in round $h$. Therefore, the total completion time can be denoted as follows:

In order to minimize the convergence bound, we set the local update frequency $\tau_{l}^{h}$ of the fastest client $l$ in round $h$ as $\small{\sqrt{\frac{12F(\mathbf{x}^{h})}{\eta^{2}HL(G^{2}+18\sigma^{2})}}}$. Thus, the optimization problem in Eq. (22) can be approximated as a univariate problem of finding the number of rounds $H$.

$\mathop{\min}\sum_{h=1}^{H}(\tau_{l}^{h}\cdot\mu_{l}^{h}+\nu_{l}^{h})$

In terms of Eq. (26), we design a greedy-based control algorithm to adaptively assign proper coefficients and local update frequencies for heterogeneous clients. The proposed algorithm consists of both the PS and client sides, which are formally described in Alg. 1 and Alg. 2, respectively. We will introduce the algorithm in detail by the order of workflow in a training round.

Firstly, at the beginning of each round $h$, the algorithm determines the model width $p_{n}^{h}$ for each client $n\in\mathcal{N}^{h}$ (Lines 6-11 of Alg. 1). To minimize the reducing error, the algorithm greedily adds the coefficient blocks for each client as many as possible within a given maximum iteration time $\mu^{max}$ or the model width $p_{n}^{h}$ reaches the maximum value $P$.

Secondly, the algorithm selects the fastest client with the least completion time (Lines 12-14 of Alg. 1). Specifically, for each participating client $n\in\mathcal{N}^{h}$, the algorithm assumes it is the fastest client and solves the approximated problem in Eq. (26), where the optimization object is $T_{n}(H,\tau_{n})=\sum_{h^{\prime}=h}^{H}(\tau_{n}\cdot\mu_{n}^{h^{\prime}}+\nu_{n}^{h^{\prime}})$ with $\tau_{n}=\sqrt{\frac{12F(\mathbf{x}^{h})}{\eta^{2}HL(G^{2}+18\sigma^{2})}}$. Then, the number of rounds $H$ is obtained. Accordingly, the total completion time for the entire training process $T_{n}$ is calculated according to $H$. Then, the algorithm selects the fastest client $l$, i.e., $l\leftarrow\mathop{\arg\min}_{n\in\mathcal{N}^{h}}T_{n}$. However, solving this problem requires information of the entire training process, such as clients’ status and network bandwidth. Unfortunately, since these information are usually time-varying in the dynamic edge system, it is impossible to obtain them in advance. To this end, we adopt the information in the current round to approximate the unavailable future information. Therefore, the optimization object is approximated as follows:

Notably, there are some variables, such as $L$, $\sigma^{2}$ and $G^{2}$, whose values are unknown at the beginning. In order to address this issue, when $h=0$, we adopt an identical predefined local update frequency $\mu_{1}$ for all participating clients, without performing the algorithm. When $h\geq 1$, each participating client estimates these variables over its local loss function (Lines 7-9 of Alg. 2) and the PS aggregates them to obtain their specific values (Line 25 of Alg. 1).

Thirdly, the algorithm determines other clients’ local update frequencies based on the fastest client $l$’s completion time in round $h$, i.e., $T_{l}^{h}$ (Lines 15-19 of Alg. 1). Specifically, for each client $n\in\mathcal{N}^{h}$, the algorithm first derives a frequency space $[\tau_{a},\tau_{b}]$ according to Eq. (24) and searches the final local update frequency within this space, ensuring the waiting time does not exceed the threshold $\rho$. Then, $(p_{n}^{h})^{2}$ blocks with the least total update times are selected to form the reduced coefficient $\hat{\mathbf{u}}_{n}^{h}$ for client $n$. Finally, the algorithm searches the local update frequency $\tau_{n}^{h}$ in $[\tau_{a},\tau_{b}]$ to minimize the variance among the total update times of all coefficient blocks.

Fourthly, the PS sends the global basis $\mathbf{v}^{h}$, reduced coefficient $\hat{\mathbf{u}}_{n}^{h}$ and local update frequency $\tau_{n}^{h}$ to each participating client $n$ for local training (i.e., Alg. 2). Once receiving the updated basis $\bar{\mathbf{v}}_{n}^{h}$ and coefficient $\bar{\mathbf{u}}_{n}^{h}$, as well as the estimated variables’ values $L_{n}$, $\sigma^{2}_{n}$ and $G^{2}_{n}$, the PS performs global aggregation (Lines 25-26 in Alg. 1). The whole process continues until the total time cost exceeds the budget $T^{max}$.

The experimental environment is built on an AMAX deep learning workstation equipped with an Intel Xeon 5218 CPU, 8 NVIDIA GeForce RTX 3090 GPUs and 256GB RAM. We simulate an FL system with 100 virtual clients and one PS (each is implemented as a process in the system) on this workstation. In each round, we randomly activate 10 clients to participate in training. Specifically, the model training and testing are implemented based on the PyTorch framework¹¹1https://pytorch.org/docs/stable, and the MPI for Python library²²2https://mpi4py.readthedocs.io/en/stable/ is utilized to build up the communication between clients and the PS.

To reflect heterogeneous and dynamic network conditions, we let each client’s download speed to fluctuate between 10Mb/s and 20Mb/s [wang2022accelerating]. Since the upload speed is usually smaller than the download speed in typical WANs, we configure it to fluctuate between 1Mb/s and 5Mb/s [jiang2023heterogeneity] for each client. Besides, considering the clients’ computation capabilities are also heterogeneous and dynamic, the time cost of one local iteration on a certain simulated client follows a Gaussian distribution whose mean and variance are derived from the time records on a physical device (e.g., laptop, TX2, Xavier NX, AGX Xavier) [liao2023adaptive].