UAV-Assisted Multi-Task Federated Learning with Task Knowledge Sharing

The system consists of $M$ EVs indexed by $\mathcal{M}=\{1,2,\cdots,M\}$ and $N$ UAVs indexed by $\mathcal{N}=\{1,2,\cdots,N\}$. Each EV is assigned a different task and can interact with other EVs to share task knowledge. Additionally, each EV can collect a small amount of data for model validation. Each UAV $n$ has a local dataset $\mathcal{D}_{n}$ with $D_{n}=\left|\mathcal{D}_{n}\right|$ data samples, and the data can be used for $M$ tasks. Since the same image dataset is used for multiple tasks, the underlying image features required by multiple tasks may be correlated. To enhance feature extraction performance, we allow multiple tasks to share feature extractors while maintaining task-specific predictors. The entire dataset is denoted by $\mathcal{D}=\cup\left\{\mathcal{D}_{n}\right\}_{n=1}^{N}$ with a total number of samples $D=\sum_{n=1}^{N}D_{n}$. Given a data sample $(\boldsymbol{x},\boldsymbol{y})\in\mathcal{D}$, where $\boldsymbol{x}\in\mathbb{R}^{d}$ is the $d$-dimensional input data vector, $\boldsymbol{y}=[y_{1},y_{2},\cdots,y_{M}]\in\mathbb{R}^{M}$ is the corresponding $M$-dimensional ground-truth label and each dimension corresponds to a specific task. Let $f_{m,n}(\boldsymbol{x},\boldsymbol{y};\boldsymbol{w}_{m,n})$ denote the sample-wise loss function for task $m$ on UAV $n$, $\boldsymbol{w}_{m,n}$ can be split into two parts $[\boldsymbol{w}_{m,n}^{s},\boldsymbol{w}_{m,n}^{u}]$, where $\boldsymbol{w}_{m,n}^{s}$ represents the shared feature extractor and $\boldsymbol{w}_{m,n}^{u}$ represents the unique layer of task $m$. Thus, the local loss function that UAV $n$ measures the model error for task $m$ on its local dataset is given by

Accordingly, the global loss function for task $m$ is

We use $\boldsymbol{w}=[\boldsymbol{w}_{1},\boldsymbol{w}_{2},\cdots\boldsymbol{w}_{M}]$ to represent the set of all task parameters. For simplicity, we omit the subscript $m$ from $f$ and use $f_{n}(\boldsymbol{x},\boldsymbol{y};\boldsymbol{w}_{m,n})$ to denote the loss function of task $m$ on UAV $n$. Each UAV can participate in different tasks in different rounds. However, due to the limited computing resources and energy, each UAV can only participate in one training task per round. Considering the goal to improve the average performance of all tasks, the objective function of the entire network is

To achieve the objective outlined in (3), the training steps of the proposed scheme are given as follows:

$\bullet$ (Step 1 UAV-EV association and Task Model Broadcast) At the beginning of each round, each UAV is assigned to an EV through the UAV-EV association algorithm. Each EV broadcasts its latest model parameters $\boldsymbol{w}_{m,t}$ to its associated UAVs. Let $\beta_{m,n,t}\in\{0,1\}$ denote the association indicator of UAV $n$ in rount $t$, where $\beta_{m,n,t}=1$ indicates that UAV $n$ is associated with EV $m$, $\beta_{m,n,t}=0$ otherwise. We use $\boldsymbol{\beta_{m,t}}=[\beta_{m,1,t},\beta_{m,2,t},\cdots,\beta_{m,N,t}]^{T}\in\mathbb{R}^{N},\forall m\in\mathcal{M}$ to represent the UAV association of task $m$ in round $t$. And $\boldsymbol{\beta}_{t}=[\boldsymbol{\beta_{1,t}};\boldsymbol{\beta_{2,t}};\cdots;\boldsymbol{\beta_{M,t}}]^{T}$ can be used to represent the overall UAV-EV association in round $t$.

$\bullet$ (Step 2 Local Model Training) Each UAV $n$ initializes $\boldsymbol{w}_{m,n,t}^{0}=\boldsymbol{w}_{m,t}$ and performs $K$ steps stochastic gradient descent (SGD) to update the shared layer and the task-specific layer. $\boldsymbol{w}_{m,n,t}^{s}$ and $\boldsymbol{w}_{m,n,t}^{u}$ are updated as follows:

where $\boldsymbol{w}_{m,n,t}^{s,k}$ and $\boldsymbol{w}_{m,n,t}^{u,k}$ are the feature extractor layer and task-specific layer parameters of task $m$ on UAV $n$ in the $k$-th local iteration of the $t$-th round, respectively. $\eta>0$ is the learning rate.

In (4), the stochastic gradient $\nabla_{s}\widetilde{f}_{n}(\boldsymbol{w}_{m,n,t}^{s,k},\boldsymbol{w}_{m,n,t}^{u,k})$ is given by $\nabla_{s}\widetilde{f}_{n}(\boldsymbol{w}_{m,n,t}^{s,k},\boldsymbol{w}_{m,n,t}^{u,k})=\frac{1}{L_{b}}\sum_{(\boldsymbol{x},\boldsymbol{y})\in\zeta_{n,t}^{h}}\nabla_{s}f_{n}\left(\boldsymbol{x},\boldsymbol{y};\boldsymbol{w}_{m,n,t}^{k}\right)$, where $\zeta_{n,t}^{k}$ is a mini-batch data uniformly sampled from $D_{n}$ with $L_{b}=\left|\zeta_{n,t}^{k}\right|$ data samples. The update form of $\nabla_{u}\widetilde{f}_{n}(\boldsymbol{w}_{m,n,t}^{s,k},\boldsymbol{w}_{m,n,t}^{u,k})$ is similar and is omitted here.

$\bullet$ (Step 3 Model Update $\&$ Parameter Uploading ) After local model training, each UAV $n$ uploads its cumulative local stochastic gradient $\tilde{\boldsymbol{G}}_{m,n,t}$ to its corresponding EV. $\tilde{\boldsymbol{G}}_{m,n,t}$ is given by $\tilde{\boldsymbol{G}}_{m,n,t}=\sum_{k=0}^{K-1}\nabla\widetilde{f}_{n}(\boldsymbol{w}_{m,n,t}^{s,k},\boldsymbol{w}_{m,n,t}^{u,k})=\frac{1}{\eta}\left(\boldsymbol{w}_{m,n,t}^{0}-\boldsymbol{w}_{m,n,t}^{K}\right)$ Note that the cumulative gradient can be expressed as $\tilde{\boldsymbol{G}}_{m,n,t}=[\tilde{\boldsymbol{G}}_{m,n,t}^{s},\tilde{\boldsymbol{G}}_{m,n,t}^{u}]$, where $\tilde{\boldsymbol{G}}_{m,n,t}^{s}$ and $\tilde{\boldsymbol{G}}_{m,n,t}^{u}$ represent the feature extractor gradient and task-specific layer gradient respectively.

$\bullet$ (Step 4 EV Model Aggregation and Knowledge Sharing) Each EV aggregates the latest gradient and updates the global model of task $m$. Specifically, the feature extractor is updated as $\hat{\boldsymbol{w}}_{m,t+1}^{s}=\boldsymbol{w}_{m,t}^{s}-\eta\frac{{\textstyle\sum\nolimits_{n=1}^{N}}\beta_{m,n,t}D_{n}\tilde{\boldsymbol{G}}_{m,n,t}^{s}}{\sum\nolimits_{n=1}^{N}\beta_{m,n,t}D_{n}}$ and the task predictor is updated as $\boldsymbol{w}_{m,t+1}^{u}=\boldsymbol{w}_{m,t}^{u}-\eta\frac{{\textstyle\sum\nolimits_{n=1}^{N}}\beta_{m,n,t}D_{n}\tilde{\boldsymbol{G}}_{m,n,t}^{u}}{\sum\nolimits_{n=1}^{N}\beta_{m,n,t}D_{n}}$. Subsequently, the EVs interact to share knowledge by transmitting $\hat{\boldsymbol{w}}_{m,t+1}^{s}$ to one another for further fusion of the feature extractor parameters. Given that EVs possess relatively sufficient transmission resources and computing capabilities, we disregard any delays and energy consumption associated with this parameter exchange. The fused feature extractor can then be expressed as: $\boldsymbol{w}_{t+1}^{s}=\frac{{\textstyle\sum\nolimits_{m=1}^{M}}D_{m,t}\hat{\boldsymbol{w}}_{m,t+1}^{s}}{\sum\nolimits_{m=1}^{M}D_{m,t}}$, where $D_{m,t}=\sum\nolimits_{n=1}^{N}\beta_{m,n,t}D_{n}$.

Let $C_{m,n}$ denotes the number of CPU cycles required to process one data sample when UAV $n$ participates in the training of task $m$, $f_{n}$ represents the computation capability of UAV $n$. Thus, if UAV $n$ participate in the training of task $m$, the computational time of local training of is given by $T_{m,n,t}^{comp}=\frac{KL_{b}C_{m,n}}{f_{n}}$ and the corresponding energy consumption of UAV $n$ is $E_{m,n,t}^{comp}=\varsigma KL_{b}C_{m,n}\left(f_{n}\right)^{2}$, where $\varsigma$ is the energy coefficient, which depends on the chip architecture.

The air-to-ground communication link is modeled as a probabilistic superposition of LOS and NLOS channels based on the method in [14]. For the sake of brevity, the details are omitted here. We consider using the frequency division multiple access (FDMA) techniques for communication. Each UAV can be allocated a certain proportion $\gamma_{n,t}$ of uplink bandwidth resources. Denote $p_{n,t}$ as the transmit power of UAV $n$ in round $t$, and it has a maximum transmission power limit $p_{n,max}$. Thus, the achievable transmission rate between UAV $n$ and EV $m$ in round $t$ is $r_{m,n,t}=\gamma_{n,t}B\log_{2}\left(1+\frac{p_{n,t}h_{m,n,t}}{\gamma_{n,t}BN_{0}}\right)$. $N_{0}$ is the noise power spectral density. Let $Q_{m}$ denote the size of the gradient for task $m$, if UAV $n$ is scheduled to participate in task $m$ in round $t$, its transmission time is given by $T_{m,n,t}^{comm}=\frac{Q_{m}}{r_{m,n,t}}$ The corresponding energy consumption of UAV $n$ for transmission is $E_{n,t}^{comm}=p_{n,t}T_{m,n,t}^{comm}$.

In order to minimize the expected average global loss of all tasks and reduce the total training time subject to the constrained resources of UAVs, we need to jointly optimize UAV-EV association, bandwidth allocation, and UAV transmission power. The problem is formulated as follows:

The constraint in (6a) ensures that the energy consumption of each UAV in each round cannot exceed the maximum energy limit $E_{n,\text{max}}$ per round. Constraints (6b)-(6d) imposes restrictions on UAV-EV associations. Constraints (6e) and (6f) regulate the bandwidth allocation, and (6g) limits the transmission power of each UAV.

Since the objective function of problem $P_{1}$ involves the neural network’s loss function $F(\boldsymbol{w}_{t})$, the factors affecting the network’s performance are still unknown. Consequently, problem $P_{1}$ is not yet solvable. To address this, we first conduct the convergence analysis of the proposed scheme to identify the factors affecting task performance, and then formulate the problem into a solvable suboptimal form.

Based on the above assumptions, we give the upper bound of the difference in loss function between two consecutive rounds of task $m$:

The proof is omitted here due to space limitations.

According to Theorem 1, it can be seen that task learning performance is mainly affected by the amount of associated data. Due to the dynamic UAV-EV association and the varying complexity of tasks, the training progress differs from task to task. Given the diminishing marginal effect observed in neural network training, improving overall task performance calls for greater emphasis on tasks with slower training progress. Besides, due to the differences in the quality of UAV data collected in different regions, the feature extractors of EVs associated with these high-quality data UAVs usually have better feature extraction capabilities. Through feature extractor fusion, these EVs usually bring higher performance improvements to other EVs. Therefore, EVs with a strong performance history should receive greater focus. Based on the above analysis, each task should have a dynamic task weight in different rounds. Therefore, to address problem P1, we reformulate it into the following form.

where $\rho_{n}=\frac{D_{n}}{D}$ represents the sample size ratio of UAV $n$ and this problem will be addressed in section IV.

In this section, we propose a task-level attention mechanism to dynamically characterize the weights of each task in each round.

Given the UAV-EV association $\boldsymbol{\beta_{t}}$, problem $P_{2}$ can be simplified as

We first determine the transmission power $p_{n,t}$ for each UAV. From the previous communication model, we know that $\frac{\partial r_{n,t}}{\partial p_{n,t}}>0$. To reduce transmission time, each UAV should operate at its highest possible transmission rate within the limits of its energy constraints. Therefore, the transmission power for each UAV is set to $p_{n,t}=p_{n,\text{max}}$. Consequently, problem (14) can be rewritten as:

where $T_{t}^{\text{max}}=\max_{n}\{T_{n,t}^{\text{comp}}+T_{n,t}^{\text{comm}}\}$ and the optimal solution of (16) is presented in Theorem 2 as follows.

The proof is omitted here due to space limitations.

Since the optimal bandwidth allocation requires solving the Lambert-$W$ function, to obtain a closed-form expression just through parameter iteration computation explicitly is almost impossible. Therefore, we adopt a bisection method to give the solution numerically.

Once the association between the UAV and the EV is determined, the optimal bandwidth allocation ratio and the minimum round time $T_{t}^{*}$ are determined. Substituting the minimum time $T_{t}^{*}$ obtained in the previous section into equation (16), we can transform the problem as

Since the objective function does not have a clear analytical expression at this time and the constraints contain binary constraint variables, the problem is a non-convex optimization problem. In order to optimize the association between UAVs and EVs in an efficient and feasible way, we use the coalition formation game to obtain a suboptimal solution to the problem.

First we introduce some essential concepts and definitions in the following part.

Based on the above definition, we can solve (19) by continuously adjusting the UAV-EV association policy $S$ to obtain a higher utility according to the preference order. The association policy adjustment will terminate when a stable association $S^{*}$ is reached, where no UAV in the system has the motivation to leave its currently associated EV and join another EV.

After the UAV-EV association algorithm converges, the optimal UAV-EV association and resource allocation strategy will be broadcast to all UAVs. All the involved UAVs will participate in the task of its associated EV and execute local training.

This work utilizes the MNIST dataset for simulation validation and some modifications are performed to make it suitable for multi-task training. The digits are first rotated, and the rotation angle is selected from the preset angle vector $[0,36,\cdots,324]$. In addition, the digits and background are dyed differently to simulate the different backgrounds of the images collected by UAVs. For the modified MNIST dataset, we use it to complete two tasks: digit recognition and rotation angle recognition. Both tasks require information about the shape and contours of the digits, so there are similarities in the underlying tasks. For convenience, we use the same predictor for both tasks, but the results in this article can be easily extended to the case of different predictors. The model shared layers contain 3 convolutional layers, each followed by a batch normalization layer. Each task retains its own 4 fully connected layers.

We set $N=10$, $M=2$, and distribute training data to each UAV based on two dirichlet distributions. The parameter $\alpha_{1}$ controls the size distribution of the UAV dataset, and the parameter $\alpha_{2}$ controls the proportional distribution of the digit categories and rotation angles. The remaining network parameters are shown in Table $\mathrm{I}$. The three baseline strategies are set as follows: Strategy 1 shares the feature extractor, but UAV-EV associations are randomly associated. Strategy 2 shares the feature extractor, but each UAV is associated with its nearest EV. Strategy 3 does not share the feature extractor, and UAV-EV associations are randomly associated, which is similar to the FedAvg algorithm.