Automated Federated Learning in Mobile Edge Networks — Fast Adaptation and Convergence

In FL, we consider a set of $n$ UEs which are connected with the server via the BS, where each UE has access only to its local data [7]. For a sample data point $\{x,y\}$ with input $x$, the goal of the server is to find the model parameter $w$ that characterizes the output $y$ with loss function $f(w):\mathbb{R}^{m}\rightarrow\mathbb{R}$, such that the value of $f(w)$ can be minimized. More specifically, if we define $f_{i}(w):\mathbb{R}^{m}\rightarrow\mathbb{R}$ as the loss function of UE $i$, the goal of the server is to solve

In particular, for each UE $i$, we have

where $l_{i}(w;x,y)$ is the error between the true label $y\in\mathcal{Y}_{i}$ and the prediction of model $w$ using input $x\in\mathcal{X}_{i}$. Each UE $i$ has a local dataset $\mathcal{D}_{i}=\{x\in\mathcal{X}_{i},y\in\mathcal{Y}_{i}\}$, with $D_{i}=|\mathcal{D}_{i}|$ data samples. Since the datasets captured by the UEs are naturally heterogenous, the probability distribution of $\mathcal{D}_{i}$ across UEs is not identical.

MAML, on the other hand, is one of the most attractive technique in meta-learning. MAML is proposed to learn an initial model that adapts quickly to a new task through one or more gradient steps with only a few data points from that new task. Each task can be regarded as an object with its own dataset to learn, just like an UE in FL. MAML allows us to replace hand-engineered algorithms with data-driven approaches to learn initial parameters automatically.

To show how to exploit the fundamental idea behind the MAML framework [11] to design an automated variant of FL algorithm, let us first briefly recap the MAML formulation. In MAML, if we regard the tasks as UEs and assume each UE takes the initial model and updates it using one step of gradient descent with respect to its own loss function¹¹1A more general case is to perform multiple steps of gradient descent. However, this would lead to expensive cost of computing multiple Hessians. Therefore, for simplicity, throughout the whole paper we consider only one single step of gradient descent., problem (1) then changes to

where $\alpha\geq 0$ is the learning rate at a UE. For UE $i$, its optimization objective $F_{i}(w)$ can be expressed as

Such a transformation from problem (1) to (3) implies that the FL problem can be interpreted within the MAML framework. The FL algorithm proposed to solve (3) is termed as Personalized Federated Averaging (Per-FedAvg) [21, 14]. Per-FedAvg is inspired by FedAvg, which is proposed in [7] as a classic and general FL algorithm to solve (1) in a distributed manner.

Per-FedAvg is summarized in Algorithm 1. In each round $k$ $(k=1,\dots,K)$, the central BS randomly picks a set of UEs $\mathcal{A}_{k}$ and then sends them the current model parameter $w_{k}$. Each UE first adapts the global model $w_{k}$ to its local data and obtains an intermediate parameter $\theta_{k}^{i}$, where $\theta_{k}^{i}=w_{k}-\alpha\nabla f_{i}(w_{k})$. Then, with $\theta_{k}^{i}$, UE $i$ updates its local model using one or more steps of gradient descent and obtains $w_{k+1}^{i}$. Such local model parameter updates are then sent to the server for model aggregation. The server will update the global model or meta model as $w_{k+1}$. This process repeats in the sequel rounds until a certain model accuracy is achieved or a predefined maximum number of rounds is reached.

Per-FedAvg is a MAML-based FL algorithm, which suggests a general approach to use the MAML method for solving the data heterogeneity problem in FL. It is proposed with focus on the personalization in FL, which is a natural feature inherited from MAML. Except for the personalization, MAML is also a few-shot learning approach, as only a few samples are required to learn new skills after just minutes of experience. Therefore, Per-FedAvg also inherits this feature of MAML, thereby adapting quickly from only a few samples. To what degree of fast adaptation MAML benefits FL and under what conditions such benefit can be achieved is what exactly this paper is about.

We consider the above described MAML-based FL. In detail, we concentrate on the situation where UEs communicate in a synchronous manner, so as to avoid using outdated parameters for global model update and make high-quality refinement in each round. Meanwhile, for each UE, we consider the case that only one step of stochastic gradient descent (SGD) is performed, following the same setting as [11].

As for the concerned MAML-based FL, our goal is to optimize the initial model using only a few data points at each UE. Hence, we only obtain an estimate of the desired gradient with SGD. Here, the desired gradient $\nabla F_{i}(w)$ on UE $i$ is computed using all data points in its dataset $\mathcal{D}_{i}$, while the estimated gradient $\tilde{\nabla}F_{i}(w)$ on UE $i$ is computed using SGD with the sampled dataset $\mathcal{D}_{i}^{(\cdot)}\in\mathcal{D}_{i}$. Note that, the superscript $(\cdot)$ represents different sampled datasets are used to estimate the involved gradients and Hessian in $\nabla F_{i}(w)$. Meanwhile, the sampled data size is $D_{i}^{(\cdot)}=|\mathcal{D}_{i}^{(\cdot)}|$.

More specifically, in order to solve (3), each UE $i$ computes the desired gradient in round $k$, as follows:

At every round, computing the gradient $\nabla f_{i}(w_{k})$ by using all data points of UE $i$ is often computationally expensive. Therefore, we take a subset $\mathcal{D}_{i}^{\text{in}}$ from $\mathcal{D}_{i}$ to obtain an unbiased estimate $\tilde{\nabla}f_{i}(w_{k};\mathcal{D}_{i}^{\text{in}})$ for $\nabla f_{i}(w_{k})$, which is given by

Similarly, the outer gradient update $\nabla f_{i}(\theta_{k}^{i})$ and Hessian update $\nabla^{2}f_{i}(w_{k})$ in (5) can be replaced by their unbiased estimates $\tilde{\nabla}f_{i}(\theta^{i}_{k};\mathcal{D}_{i}^{\text{o}})$ and $\tilde{\nabla}^{2}f_{i}(w_{k};\mathcal{D}_{i}^{\text{h}})$ respectively. Here, $\mathcal{D}_{i}^{\text{o}}$ and $\mathcal{D}_{i}^{\text{h}}$ are sampled from $\mathcal{D}_{i}$ as well. Therefore, using SGD, we can finally obtain an estimated local gradient $\tilde{\nabla}F_{i}(w_{k})$ on UE $i$ in round $k$, which is given by

It is worth noting that $\tilde{\nabla}F_{i}(w_{k})$ is a biased estimator of $\nabla F_{i}(w_{k})$. This is because the stochastic gradient $\tilde{\nabla}f_{i}(w_{k}-\alpha\tilde{\nabla}f_{i}(w_{k};\mathcal{D}_{i}^{\text{in}});\mathcal{D}_{i}^{\text{o}})$ contains another stochastic gradient $\tilde{\nabla}f_{i}(w_{k};\mathcal{D}_{i}^{\text{in}})$ inside. Hence, to improve the estimate accuracy, $\mathcal{D}_{i}^{\text{in}}$ used for inner gradient update is independent from the sampled datasets $\mathcal{D}_{i}^{\text{o}}$ and $\mathcal{D}_{i}^{\text{h}}$ used for outer gradient and Hessian update respectively. Meanwhile, in this paper we assume $\mathcal{D}_{i}^{\text{o}}$ and $\mathcal{D}_{i}^{\text{h}}$ are also independent from each other.

When exploring the MAML-based FL in realistic mobile edge networks, the communication and computation model should be captured carefully. Particularly, we consider that UEs access the BS through a channel partitioning scheme, such as orthogonal frequency division multiple access (OFDMA). In order to successfully upload the local update to the BS, two conditions need to be satisfied: (1) the UE is selected, and (2) the transmitted local update is successfully decoded. In this respect, we first introduce $s_{i}^{k}\in\{0,1\}$ as a selection indicator, where $s_{k}^{i}=1$ indicates the event that UE $i$ is chosen in round $k$, and $s_{k}^{i}=0$ otherwise. Next, we characterize the transmission quality of the wireless links. For the signals transmitted from UE $i$, the SNR received at the BS can be expressed as $\xi_{k}^{i}=\frac{p_{k}^{i}h_{k}^{i}\|c_{i}\|^{-\kappa}}{N_{0}}$, where $p_{k}^{i}$ is the transmit power of UE $i$ during round $k$, $\kappa$ is the path loss exponent. $h_{k}^{i}\|c_{i}\|^{-\kappa}$ is the channel gain between UE $i$ and the BS with $c_{i}$ being the distance between UE $i$ and the BS and $h_{k}^{i}$ being the small-scale channel coefficient. $N_{0}$ is the noise power spectral density. In order for the BS to successfully decode the local update from UE $i$, it is required that the received SNR exceeds a decoding threshold $\phi$, i.e., $\xi_{k}^{i}>\phi$. Assume that the small-scale channel coefficients across communication rounds follow Rayleigh distribution, then according to [22], the update success transmission probability, defined as $q_{k}^{i}=\mathbb{P}(s_{k}^{i}=1,\xi_{k}^{i}>\phi)$, can be estimated as follows:

where $\nu(p_{k}^{i})=\frac{N_{0}\phi}{p_{k}^{i}}$. Then the achievable uplink rate $r_{k,i}$ of UE $i$ transmitting its local update to the BS in round $k$ is given by

where $B$ is the bandwidth allocated to each UE. Based on $r_{k,i}$, the uplink transmission delay of UE $i$ in round $k$ can be specified as follows:

where $Z$ is the size of $w_{k}$ in number of bits. Since the transmit power of the BS is much higher than that of the UEs, the downlink transmission delay is much smaller than the uplink transmission delay. Meanwhile, we care more about the transmit power allocation on individual UEs rather than that on the BS, so here we ignore the downlink transmission delay for simplicity.

Further, we calculate the computation time for each UE, which is consumed for computing the local update. Given the CPU-cycle frequency of UE $i$ by $\vartheta_{i}$, the computation time of UE $i$ is expressed as

In (11), $c_{i}$ denotes the number of CPU cycles for UE $i$ to execute one sample of data points and $d_{k}^{i}=D_{k,i}^{\text{in}}+D_{k,i}^{\text{o}}+D_{k,i}^{\text{h}}$ denotes the sampled data size of UE $i$ in round $k$.

In term of $t_{k,i}^{\text{com}}$ and $t_{k,i}^{\text{cmp}}$, we then give the energy consumption of each UE $i$ in round $k$, which consists of two parts: (1) the energy for transmitting the local updates and (2) the energy for computing the local updates. Let $e_{k}^{i}$ denote the energy consumption of each UE $i$ in round $k$, and then $e_{k}^{i}$ can be computed as follows [23, 24]:

where $\frac{\varsigma}{2}$ is the effective capacitance coefficient of UE $i$’s computing chipset. From (12), we observe that in round $k$ for each UE $i$, both the sampled data size $d_{i}^{k}$ and transmit power $p_{i}^{k}$ have significant impacts on its energy consumption. This observation motivates us to jointly consider these two variables when quantifying the benefit MAML brings to FL in mobile edge networks.

In order to quantify the benefit MAML brings to FL in mobile edge networks, we focus on the learning time minimization under the constraint of model accuracy and UEs’ energy consumption. Particularly, the learning time is the duration over all $K$ communication rounds, in which the duration of round $k$ is determined by the slowest UE as follows,

Note that we can replace $i\in\mathcal{A}_{k}$ with $i\in\mathcal{U}$ in $T_{k}^{\text{round}}$, that is, $T_{k}^{\text{round}}=\max_{i\in\mathcal{U}}\{t_{k,i}^{\text{cmp}}+t_{k,i}^{\text{com}}\}$. This is because the UEs that are not chosen in $\mathcal{A}_{k}$ has a $T_{k}^{i}$ with the value of $0$, which would not effect the result of $T_{k}^{\text{round}}$. Let $\mathbf{d}\triangleq\{\mathbf{d}^{1},\dots,\mathbf{d}^{n}\}$ and $\mathbf{p}\triangleq\{\mathbf{p}^{1},\dots,\mathbf{p}^{n}\}$ denote the UEs’ sampled data size vector and the transmit power vector respectively, where $\mathbf{d}^{i}\triangleq\{d_{1}^{i},\dots,d_{K}^{i}\}$ and $\mathbf{p}^{i}\triangleq\{p_{1}^{i},\dots,p_{K}^{i}\}$. Then, we can give the problem formulation by

In problem (P1), we not only optimize $\mathbf{d}$ and $\mathbf{p}$, but also the total number of rounds $K$. The hidden reason is that, both $\mathbf{d}$ and $\mathbf{p}$ have an effect on the duration of each round, thereby impacting the number of rounds $K$ to achieve a certain model accuracy $\epsilon$. Besides, (C1.1) characterizes an $\epsilon$-approximate convergence performance. (C1.2) limits the energy consumption of UE $i$ in round $k$ not larger the predefined maximum value $E_{\max}$. (C1.3) gives the maximum transmit power of UE $i$ in round $k$ by $P_{\max}$. (C1.4) is the sampled data size constraint, that the sampled data size of one UE in round $k$ is smaller than or equal to the data points generated by that UE. The solution to problem (P1) can be exploited for designing an optimized MAML-based FL algorithm, which is implemented iteratively to update the global model parameter $w$. According to the communication and computation model, only those local updates from the selected UEs as well as being successfully decoded by the BS can contribute to updating the global model parameter. That is, in round $k$, we have

where $\mathds{1}\{s_{k}^{i}=1,\gamma_{k}^{i}>\phi\}=1$ if the event of $s_{k}^{i}=1,\gamma_{k}^{i}>\phi$ is true, and it equals to zero otherwise.

Based on this update rule, (C1.1) related to the MAML-based FL can be analytically analysed to make the relationship between the decision variables and the model accuracy $\epsilon$ explicit, thereby facilitating solving problem (P1). In this regard, the convergence analysis of the MAML-based FL is given in the following section.

For a general purpose, we concentrate on the non-convex settings on loss functions and aim to find an $\epsilon$-approximate first-order stationary point (FOSP) for the loss function minimization problem (3) in the MAML-based FL, which is defined as follows.

In terms of this Definition, we then elaborate on the assumptions, convergence bounds, and discussions on the convergence analysis, respectively, from which we can reformulate (C1.1) in problem (P1) as an explicit constraint with respect to the optimization variables.

Without loss of generality, the assumptions used for the convergence analysis of MAML-based FL algorithm is consistent with that of Per-FedAvg [21, 14] and are given in the following.

Before the convergence analysis, we first introduce three lemmas inherited from [14, 21] quantifying the smoothness of $F_{i}(w)$ and $F(w)$, the deviation between $\nabla F_{i}(w)$ and its estimate $\tilde{\nabla}F_{i}(w)$, and the deviation between $\nabla F_{i}(w)$ and $\nabla F(w)$, respectively.

Let $U_{i}=\max_{k=\{1,\dots,K\}}U_{k}^{i}$, where $U_{k}^{i}=\frac{\mathbb{P}\{s_{k}^{i}=1,\xi_{k}^{i}>\phi\}}{\sum_{i=1}^{n}\mathbb{P}\{s_{k}^{i}=1,\xi_{k}^{i}>\phi\}}=\frac{q_{k}^{i}}{\sum_{i=1}^{n}q_{k}^{i}}$ denotes the normalized update successful probability of UE $i$ in round $k$. Then, with Lemmas 1, 2 and 3, the expected convergence result of the MAML-based FL within a general mobile edge network we describe in Section II can now be obtained by the following theorem.

Unlike the convergence guarantee of Per-FedAvg in [14] which is characterized by $K$ and $\mathbf{d}$, the convergence guarantee obtained from Theorem 1 is characterized by $K$, $\mathbf{d}$ and $U_{i}$, where $U_{i}$ is a function of the transmit power $\mathbf{p}^{i}$. That is, the convergence bound of the proposed MAML-based FL in our paper is described in terms of $K$, $\mathbf{d}$ and $\mathbf{p}$. Therefore, our convergence analysis can combine the FL hyperparameters as well as the resource allocation in mobile edge networks.

From Theorem 1, we are able to characterize $K$, $\mathbf{d}$ and $\mathbf{p}$ with respect to the predefined model accuracy $\epsilon$. According to Theorem 1 and (C1.1) in problem (P1), it is desired for the right-hand-side of (1) to be equal to or smaller than $\epsilon$. Consequently, we further present the following corollary.

Guided by Corollary 1, we can recast (C1.1) in problem (P1) as a series of constraints, which give explicit relationships between optimization variables and the model accuracy $\epsilon$. In this way, we then solve problem (P1) with high efficiency in the next section.

The clue to the decomposition of problem (P1) follows the primal decomposition [25]. Specifically, the master problem is optimized with respect to the number of communication rounds $K$, which is given as follows:

Note that, we refer to Theorem 1 and use a contracted (C1.1), i.e., (C2.1), as the constraint in problem (P2) for tractability. This contraction replaces $\frac{1}{K}\sum_{k=0}^{K-1}\mathbb{E}[\|\nabla F(\bar{w}_{k})\|^{2}]$ with its upper bound without the second term. The reason that we get rid of the second term is that we only consider $K$ as the decision variable in (P2), and the second term of the upper bound can be regarded as a constant and thus removed as long as its value is no larger than $\epsilon$. Meanwhile, the above contraction shrinks the feasible region of $K$. Therefore, as long as (C2.1) is satisfied, (C1.1) must be satisfied. Moreover, if there exists an optimal solution $K$ to problem (P2), it is also feasible to the original problem (P1). We can regard the optimal solution to (P2) as an optimal approximation to (P1).

The reason that we consider $K$ as a decision variable is for the decoupling of (P1). The original objective of minimizing the overall learning time in $K$ rounds can be regarded as the minimization of the learning time in each round, once $K$ is determined and is set to be a constant in the following decoupling. As such, problem (P2), which has a slave problem with respect to the sampled data size $\mathbf{d}$ and transmit power $\mathbf{p}$, can be regarded as an approximation of problem (P1).

From Theorem 1, we find that $\mathbf{d}$ and $\mathbf{p}$ are coupled even though (C1.1) is replaced by the right-hand-side term in (1). In this regard, we aim to optimize $\mathbf{d}$ and $\mathbf{p}$ in the slave problem with the coordinate descent [20] method. That is, the slave problem is firstly solved with $\mathbf{d}$ as variables and $\mathbf{p}$ as constants, and then is solved with $\mathbf{p}$ as variables and $\mathbf{d}$ as constants. This procedure is repeated in an iterative manner. In this way, $\mathbf{d}$- and $\mathbf{p}$-related slave problems can be further decomposed among rounds under the assistance of Corollary 1. Specifically, in any round $k$, the sampled data size for UE $i$ is optimized by solving problem (P3):

(C3.1) is obtained from Corollary 1, that a $d_{i}=\mathcal{O}(\frac{1}{\epsilon})$ will yield an $\epsilon$-FOSP. (C3.1) uses $d_{k}^{i}\geq\frac{1}{\epsilon}$ instead to indicate that as long as $d_{k}^{i}$ is not smaller than $\frac{1}{\epsilon}$, an $\epsilon$-FOSP would be guaranteed. This is a reasonable replacement (i.e., from $d_{k}^{i}=\mathcal{O}(\frac{1}{\epsilon})$ to $d_{k}^{i}\geq\frac{1}{\epsilon}$) such that, the larger the sampled data size is, the larger accuracy AutoFL would achieve, thereby leading to a smaller $\epsilon$. Note that (C3.1) is also a contraction of (C1.1), and the feasible domain of $d_{k}^{i}$ is also shrunk. That is, the optimal $d_{k}^{i*}$, as long as it exists in (P3), is also feasible in problem (P1), and can be regarded as an optimal approximation to (P1). (C3.2) is the energy consumption constraint on each UE and (C3.3) indicates that the sampled data size of UE $i$ is no greater than the number of data points generated by UE $i$. Furthermore, addressing the following problem can make an decision on the transmit power of UE $i$ in round $k$:

where (C4.1) is obtained from Corollary 1, in terms of the fact that $(\sum_{i=1}^{n}{U_{k}^{i}}^{2})\frac{\alpha^{2}L^{2}\sigma_{G}^{2}}{D^{\text{in}}}$ dominates the value of the second term in the right-hand-side of (1) to guarantee an $\epsilon$-FOSP for the MAML-based FL. Here we replace $U_{i}$ with $U_{k}^{i}$ to show that in each round $k$, (C4.1) should be satisfied. The reason that we replace $(\sum_{i=1}^{n}{U_{k}^{i}}^{2})\frac{\alpha^{2}L^{2}\sigma_{G}^{2}}{D^{\text{in}}}=\mathcal{O}(\frac{1}{\epsilon})$ with $(\sum_{i=1}^{n}{U_{k}^{i}}^{2})\frac{\alpha^{2}L^{2}\sigma_{G}^{2}}{D^{\text{in}}}\leq\epsilon$ is to indicate that, as long as (C4.1) is satisfied, an $\epsilon$-FOSP would be guaranteed. Note that (C4.1) is also a contraction of (C1.1). That is, as long as there exists an optimal solution to (P4), this solution is also feasible and is close to the optimal solution to (P1). We can regard the optimal solution to (P4) as an optimal approximation to (P1). (C4.2) and (C4.3) are the energy constraint and transmit power constraint on each UE, respectively.

Based on the above decomposition, we can solve the original problem (P1) by solving problem (P2), in which problems (P3) and (P4) are nested and addressed with the coordinate descent method. This decoupling only provides the approximate optimal results. However, the evaluation results show that, with the approximate optimal results obtained from (P2), (P3) and (P4), the proposed MAML-based FL algorithm always outperforms Per-FedAvg in learning time and convergence performances. Next, we deal with problems (P2), (P3), and (P4), respectively.