Optimized Power Control for Over-the-Air Federated Edge Learning

In the pursuit of ubiquitous intelligence envisioned in the future 6G networks, recent years have witnessed the spreading of artificial intelligence (AI) algorithms from the cloud to the network edge, resulting in an active area called edge intelligence [1, 2]. The core research issue therein is to allow low-latency and privacy-aware access to rich mobile data for intelligence distillation. To this end, a popular framework called federated edge learning (FEEL) is proposed recently, which distributes the task of model training over edge devices so as to reduce the communication overhead and keep the data-use locally [3, 4]. Essentially, the FEEL framework is a distributed implementation of stochastic gradient decent (SGD) over wireless networks. A typical training process involves iterations between 1) broadcasting of the global model under training from edge server to devices for local SGD execution using local data, and 2) local models/gradients uploading from devices to edge server for aggregation and global model updating. Although the uploading of high-volume raw data is avoided, the updates aggregation process in FEEL may still suffer from a communication bottleneck due to the high-dimensionality of each updates and the multiple access by many devices over wireless links. To tackle this issue, one promising solution called over-the-air FEEL (Air-FEEL) has been proposed, which exploits the over-the-air computation (AirComp) for “one-shot” aggregation via concurrent update transmission, such that communication and computation are integrated in a joint design by exploiting the superposition property of a multiple access channel (MAC) [1, 5, 6].

The idea of AirComp was first proposed in [5] in the context of data aggregation in sensor networks, where it is surprisingly found that “interference” can be harnessed by structured codes to help functional computation over a MAC. Inspired by the finding, it was shown in the subsequent work [6] that for Gaussian independent and identically distributed (i.i.d.) data sources, the uncoded transmission is optimal in terms of distortion minimization. Besides the information-theoretic studies, various practical issues faced by AirComp implementation were also considered in [7, 8, 9]. In particular, the synchronization issue in AirComp was addressed in [7] via an innovative idea of shared clock broadcasting from edge server to devices. The optimal power control policies for AirComp over fading channels were derived in [8] to minimize the average computation distortion, and the cooperative interference management framework for coordinating coexisting AirComp tasks over multi-cell networks was developed in [9].

More recently, AirComp found its merits in the new context of FEEL, known as Air-FEEL, for communication-efficient update aggregation as demonstrated in a rich set of prior works [10, 11, 12, 13, 14, 15]. Specifically, a broadband Air-FEEL solution was proposed in [10], where several communication-learning tradeoffs were derived to guide the design of device scheduling. Around the same time, a source-coding algorithm exploiting gradient sparsification was proposed in [11] to implement Air-FEEL with compressed updates for higher communication efficiency. In parallel, a joint design of device scheduling and beamforming in a multi-antenna system was presented in [12] to accelerate Air-FEEL. Subsequently, the gradient statistics aware power control was investigated in [13] to further enhance the performance of Air-FEEL. Furthermore, to allow Air-FEEL compatible with digital chips embedded in modern edge devices, Air-FEEL based on digital modulation was proposed in [14] featuring one-bit quantization and modulation at the edge devices and majority-vote based decoding at the edge server. Besides the benefit of low latency, Air-FEEL was also found to be beneficial in data privacy enhancement as individual updates are not accessible by edge server, eliminating the risk of model inversion attack [15].

Despite the promise in high communication efficiency, Air-FEEL may suffer from severe learning performance degradation due to the aggregation error caused by the non-uniform channel fading over devices and noise perturbation. Prior work in this field mostly assumed channel inversion power control (or its variants) [10, 11, 12, 15] in an effort to reducing the aggregation error by aligning the channel gains, which could be highly suboptimal in deep fading scenarios due to the noise amplification. Although there exists one relevant study on power control for Air-FEEL system in [13], it focused on the minimization of the intermediate aggregation distortion (e.g., mean squared error) instead of the ultimate learning performance (e.g., the general loss function). Therefore, there still leaves a research gap in learning performance optimization of Air-FEEL by judicious power control, motivating the current work. To close the gap, we first analyze the convergence behavior of the Air-FEEL by deriving the optimality gap of the loss-function under arbitrary power control policy. Then the power control problem is formulated to minimize the optimality gap for convergence acceleration, subject to a set of average and maximum power constraints at edge devices. The problem is generally non-convex and challenging to solve due to the coupling of power control variables over different devices and iterations. The challenge is tackled by the joint use of successive convex approximation (SCA) and trust region methods in the optimized power control algorithm derivation. Numerical results show that the optimized power control policy achieves significantly faster convergence than the benchmark policies such as channel inversion and uniform power transmission, thus opening up a new degree-of-freedom for regulating the performance of Air-FEEL by power control.

We consider an Air-FEEL system consisting of an edge server and $K\geq 0$ edge devices, as shown in Fig. 1. With the coordination of the edge server, the edge devices cooperatively train a shared machine learning model via over-the-air update aggregation as elaborated in the sequel.

In this section, we formally characterize the learning performance of Air-FEEL system, which is derived to be a function of transmit powers of all devices.

Let $N$ denote the number of needed communication rounds and $L\triangleq\|\bf L\|_{\infty}$. For notational convenience, we use $F^{(n+1)}$ to represent $F({\bf w}^{(n+1)})$. The optimality gap after $N$ communication rounds defined by $F^{(N+1)}-F^{\star}$ is derived in the following theorem, from which we can understand the convergence behavior of Air-FEEL.

Further applying the mean inequality $(a_{1}a_{2}\cdots a_{m})\leq(\frac{a_{1}+a_{2}+\cdots+a_{m}}{m})^{m}$, we can derive a more elegant upper bound for the expression in (13) to attain more insights as follows

where $\alpha=\frac{\sum_{i=1}^{N}\!\!A^{(i)}}{N}$ and $\beta_{(n)}=\frac{\sum_{i=n+1}^{N}A^{(i)}}{N-n}$ for $n=1,\cdots,N-1$ while $\beta_{(N)}=1$.

In this section, we focus on speeding up the convergence rate by minimizing the optimality gap in Theorem 1, under the power constraints stated in (11) and (12). The optimization problem is thus formulated as

Due to the coupling between the power control $\{p_{k}^{(n)}\}$ and learning rate $\eta$, problem (P1) is non-convex and hard to solve. We resort to the alternating optimization technique for efficiently solving this problem. In particular, we first solve problem (P1) under any given $\eta$, and then apply a one-dimension search to find the optimal $\eta$ that achieves the minimum objective value.

Let $\tilde{\Phi}(\{p_{k}^{(n)}\})={\Phi}(\{p_{k}^{(n)},\eta\})$ under any given $\eta$. Note that the transmit powers at different devices and different communication rounds are coupled with each other in the objective function in (13) under given learning rate $\eta$, leading to a highly non-convex problem:

To tackle this problem, we propose an iterative algorithm to obtain an efficient solution using the SCA technique. The key idea is that under any given local point at each iteration, we can approximate the non-convex objective as a constructed convex one. Therefore, after solving a series of approximate convex problems iteratively, we can obtain a high-quality suboptimal solution to problem (P2).

Let $\{p_{k}^{(n)}[i]\}$ denote the local point at the $i$-th iteration with $i\geq 0$, and ${\mathcal{N}}\triangleq\{1,\cdots,N\}$ the set of communication rounds. Notice that by checking the first-order Taylor expansion of $\tilde{\Phi}(\{p_{k}^{(n)}\})$ w.r.t. $\{p_{k}^{(n)}\}$ at the local point $\{p_{k}^{(n)}[i]\}$, it follows that

where $\nabla\tilde{\Phi}(\{p_{k}^{(n)}[i]\})$ represents the first-order derivative w.r.t. $p_{k}^{(n)}[i]$, given in (17) and (18).

​

In this case, $\bar{\Phi}(\{p_{k}^{(n)}\})$ is linear w.r.t. $\{p_{k}^{(n)}\}$. To ensure the approximation accuracy, a series of trust region constraints are imposed as [17]

where $\Gamma[i]$ denotes the radius of the trust region. By replacing $\bar{\Phi}(\{p_{k}^{(n)}\})$ as the approximation of $\tilde{\Phi}(\{p_{k}^{(n)}\})$ and introducing an auxiliary variable $\gamma$, the approximated problem at the $i$-th iteration is derived as a convex problem:

which can be directly solved by CVX [18].

Let $\{p_{k}^{(n)*}[i]\}$ denote the optimal power control policy to problem (P2.1) at local point $\{p_{k}^{(n)}[i]\}$. Then, we can obtain an efficient iterative algorithm to solve problem (P2) as follows. In each iteration $i\geq 1$, the power control is updated as $\{p_{k}^{(n)*}[i]\}$ by solving problem (P2.1) at local point $\{p_{k}^{(n)}[i]\}$, i.e. $p_{k}^{(n)}[i+1]=p_{k}^{(n)*}[i],\forall n\in\mathcal{N},\forall k\in\cal K$, where $\{p_{k}^{(n)}[0]\}$ denotes the initial power control. At the $i$-th iteration, we compute the objective value in problem (P2) by replacing $\{\hat{p}_{k}^{(n)*}[i]\}$ as $\{p_{k}^{(n)*}\}$. If the objective value decreases, we then replace the current point by the obtained solution and go to the next iteration; otherwise, we update $\Gamma[i]=\Gamma[i]/2$ and continue to solve problem (P2.1). This algorithm would stop until that $\Gamma[i]$ is lower than a given threshold denoted by $\epsilon$. In summary, the proposed algorithm is presented in Algorithm 1.

With the obtained power control in Algorithm 1, we can find the optimal $\eta$ accordingly via a one-dimensional search.

In this section, we provide simulation results to validate the performance of the proposed power control policy for Air-FEEL. In the simulation, the wireless channels from each device to the edge server over fading states follow i.i.d. Rayleigh fading, such that $h_{k}$’s are modeled as i.i.d. circularly symmetric complex Gaussian (CSCG) random variables with zero mean and unit variance. The dataset with size $D_{\rm tot}=600$ at all device are randomly generated, where part of the data, namely $100$ pairs (${\bf x}$, $y$), are left for prediction, and the remaining ones are used for model training. The generated data sample vector ${\bf x}$ follow i.i.d. Gaussian distribution as $\mathcal{N}(0,{\bf I})$ and the label $y$ is obtained as $y=x(2)+3x(5)+0.2z$, where $x(t)$ represents the $t$-entry in vector ${\bf x}$ and $z$ is the observation noise with i.i.d. Gaussian distribution, i.e., $z\sim\mathcal{N}(0,1)$. Unless stated otherwise, the data samples are evenly distributed among the $K=20$ devices, and thus it follows $D_{k}=25$. Moreover, we apply ridge regression with the sample-wise loss function $f({\bf w},{\bf x},y)=\frac{1}{2}\|{\bf x}^{T}{\bf w}-y\|^{2}$ and the regularization function $R({\bf w})=\|{\bf w}\|^{2}$ with $\rho=5\times 10^{-5}$ in this paper. Furthermore, recall that $D_{\rm tot}=\sum_{k\in\mathcal{K}}D_{k}$ and then we can obtain the smoothness parameter $L$ and PL parameter $\mu$ as the largest and smallest eigenvalues of the data Gramian matrix ${\bf X}^{T}{\bf X}/D_{\rm tot}+10^{-4}{\bf I}$, in which ${\bf X}=[{\bf x}_{1},\cdots,{\bf x}_{D_{\rm tot}}]^{T}$ is the data matrix. The optimal loss function $F^{\star}$ is computed according to the optimal parameter vector $\bf w^{\star}$ to the learning problem (3), where ${\bf w}^{\star}=({\bf X}^{T}{\bf X}+\rho{\bf I})^{-1}{\bf X}^{T}{\bf y}$ with ${\bf y}=[y_{1},\cdots,y_{D_{\rm tot}}]^{T}$. We set the initial parameter vector as an all-zero vector and the noise variance $N_{0}=0.1$.

We consider two benchmark schemes for performance comparison, namely the uniform power transmission that transmits with uniform power over different communication round under the constraint of average power budget, and the channel inversion adopted in [15]. As for the performance metric for comparison, we consider the optimality gap and prediction error to evaluate the learning performance.

The effect of device population on learning performance is illustrated in Fig. 2 with $N=30$, where the power budgets at all devices are identically set to be $\tilde{P}=1$ W and $\bar{P}=5$ W. Notice that the increasing of device population may introduce both the positive and negative effects on the learning performance. The positive effect is that the training process can exploit more data, while the negative effect is the increased aggregation error raised by AirComp over more devices. As observed in Fig. 2, the positive effect can be cancelled or even overweighed by the negative effect when applying the channel inversion or uniform power control. The blessing of including more devices in Air-FEEL can dominate the curse it brings only when the power control is judiciously optimized, showing the crucial role of power control in determining the the learning performance of Air-FEEL.

Fig. 3 shows the learning performance during the learning process under the optimized learning rate, where we set $K=20$, $\tilde{P}=1$ W, $\bar{P}=5$ W, and $N=80$. It is observed that the proposed power control scheme can achieve faster convergence than both the channel-inversion and uniform-power-control schemes. This is attributed to the power control optimization directly targeting convergence acceleration.

Fig. 4 shows the power allocation over a static channel with uniform channel gain during the learning process, where we set $K=20$, $\tilde{P}=1$ W, $\bar{P}=5$ W, and $N=30$. It is observed that the power allocation over a static channel follows a stair-wise monotonously decreasing function. The behavior of power control coincides the analysis on Remark 1.

In this paper, we exploit power control as a new degree of freedom to optimize the learning performance of Air-FEEL, a promising communication-efficient solution towards edge intelligence. To this end, we first analyzed the convergence rate of the Air-FEEL by deriving the optimality gap of the loss-function under arbitrary power control policy. Then the formulated power control problem aimed to minimize the optimality gap for accelerating convergence, subject to a set of average and maximum power constraints at edge devices. Due to the coupling of power control variables over different devices and iterations, the challenge of the formulated power control problem was tackled by the joint use of SCA and trust region methods. Numerical results demonstrated that the optimized power control policy can achieve significantly faster convergence than the benchmark policies such as channel inversion and uniform power transmission.