Joint Optimization of Communications and Federated Learning Over the Air

Let $\mathcal{D}_{i}=\{\mathbf{x}_{i,k},\mathbf{y}_{i,k}\}_{k=1}^{K_{i}}$ denote the local dataset at the $i$-th worker, $i=1,\ldots,U$, where $\mathbf{x}_{i,k}$ is the input data vector, $\mathbf{y}_{i,k}$ is the labeled output vector, $k=1,2,...,K_{i}$, and $K_{i}=|\mathcal{D}_{i}|$ is the number of data samples available at the $i$-th worker. With $K=\sum_{i=1}^{U}K_{i}$ samples in total, these $U$ workers seek to collectively train a learning model parameterized by a global model parameter $\mathbf{w}=[w^{1},\ldots,w^{D}]\in\mathcal{R}^{D}$ of dimension $D$, by minimizing the following loss function

where the global loss function $F(\mathbf{w};\mathcal{D})$ is a summation of $K$ data-dependent components, each component $f(\mathbf{w};\mathbf{x}_{i,k},\mathbf{y}_{i,k})$ is a sample-wise local function that quantifies the model prediction error of the same data model parameterized by the shared model parameter $\mathbf{w}$, and $\mathcal{D}=\bigcup_{i}\mathcal{D}_{i}$.

In distributed learning, each worker trains a local model $\mathbf{w}_{i}$ from its local data $\mathcal{D}_{i}$, which can be viewed as a local copy of the global model $\mathbf{w}$. That is, the local loss function is

where $\mathbf{w}_{i}=[w_{i}^{1},\ldots,w_{i}^{D}]\in\mathcal{R}^{D}$ is the local model parameter. Through collaboration, it is desired to achieve $\mathbf{w}_{i}=\mathbf{w}=\mathbf{w}^{*}$, $\forall i$, so that all workers reach the globally optimal model $\mathbf{w}^{*}$. Such a distributed learning can be formulated via consensus optimization as[4, 28]

To solve P1, this paper adopts a model-averaging algorithm for FL [4, 28]. It is essentially an iterative process consisting of both computing and communication steps at each iteration. Specifically, in each communication round, the PS broadcasts the current $\mathbf{w}$ to all workers. Then, the $i$-th worker uses a learning algorithm to update its $\mathbf{w}_{i}$ by minimizing its local data-dependent loss function in (2). In this work, gradient descent¹¹1In this work, we take the basic gradient descent as an example, while the proposed methodology can be extended to mini-batch gradient descent as well. is applied, in which the local model at the $i$-th local worker is updated as

where $\alpha$ is the learning rate, and $\nabla f(\mathbf{w};\mathbf{x}_{i,k},\mathbf{y}_{i,k})$ is the gradient of $f(\mathbf{w};\mathbf{x}_{i,k},\mathbf{y}_{i,k})$ with respect to $\mathbf{w}$.

When local updating is completed, each worker transmits its updated parameter $\mathbf{w}_{i}$ to the PS via wireless uplinks to update the global $\mathbf{w}$ as

Then, the PS broadcasts $\mathbf{w}$ in (5) to all participating workers as their initial value in the next round. The FL implements the local model-updating in (III-A) and the global model-averaging in (5) iteratively, until convergence. It has been shown that this FL algorithm converges to the globally optimal solution of the original problem in P1 under the conditions that $F$ is a convex function and the data transmission between the PS and workers is error-free[4, 28].

Note that the implementation steps in (III-A) and (5) only concern the computational aspect of FL, by assuming perfect communications for both the global $\mathbf{w}$ and local $\mathbf{w}_{i}$ between the PS and all workers. However, the communication impacts on FL performance should not be ignored. Especially in practical wireless network environments, certain errors are inevitably introduced during transmissions of the updates due to the imperfect characteristics of wireless channels.

To avoid heavy communication overhead and save transmission bandwidth of FL over wireless channels, we adopt analog aggregation without coding, which allows multiple workers to simultaneously upload their local model updates to the PS over the same time-frequency resources. All workers transmit their local $\mathbf{w}_{i}$’s in an analog form with perfect time synchronization among them²²2The implementation of time synchronization and the impact of imperfect synchronization are beyond the scope of this work. Interested readers are referred to [12, 29].. In this way, the local updates $\mathbf{w}_{i}$’s are aggregated over the air to implement the global model updating step in (5). Such an analog aggregation is conducted in an entry-wise manner. That is, the $d$-th entries $w_{i}^{d}$ from all workers, $i=1,...,U$, are aggregated to compute $w^{d}$ in (5), for any $d\in[1,D]$.

Let $\mathbf{p}_{i,t}=[p^{1}_{i,t},\dots,p^{d}_{i,t},\ldots,p^{D}_{i,t}]$ denote the power control vector of worker $i$ at the $t$-th iteration. Noticeably, the choice of $\mathbf{p}_{i,t}$ in FL over the air should be made not only to effectively implement the aggregation rule in (5), but also to properly accommodate the need for network resource allocation. Accordingly, we set the power control policy as

where $h_{i,t}$ is the channel gain between the $i$-th worker and the PS at the $t$-th iteration³³3In this paper, we assume the channel state information (CSI) to be constant within each iteration, but may vary over iterations. We also assume that the CSI is perfectly known at the PS, and leave the imperfect CSI case in future work., $b^{d}_{t}$ is the power scaling factor, and $\beta^{d}_{i,t}$ is a transmission scheduling indicator. That is, $\beta^{d}_{i,t}=1$ means that the $d$-th entry of the local model parameter $\mathbf{w}_{i,t}$ of the $i$-th worker is scheduled to contribute to the FL algorithm at the $t$-th iteration, and $\beta^{d}_{i,t}=0$, otherwise. Through power scaling, the transmit power used for uploading the $d$-th entry from the $i$-th worker should not exceed a maximum power limit $P_{i}^{d,\max}=P_{i}^{\max}$ for any $d$, as follows:

At the PS side, the received signal at the $t$-th iteration can be written as

where $\odot$ represents Hadamard product, $\mathbf{h}_{i,t}=[h^{1}_{i,t},h^{2}_{i,t},...,h^{D}_{i,t}]$, $\bm{\beta}_{i,t}=[\beta^{1}_{i,t},\beta^{2}_{i,t},..,\beta^{D}_{i,t}]$, $\mathbf{b}_{t}=[b^{1}_{t},b^{2}_{t},...,b^{D}_{t}]$, and $\mathbf{z}_{t}\sim\mathcal{CN}(0,\sigma^{2}\mathbf{I})$ is additive white Gaussian noise (AWGN).

Given the received $\mathbf{y}_{t}$, the PS estimates $\mathbf{w}_{t}$ via a post-processing operation as

where $(\sum_{i=1}^{U}K_{i}\bm{\beta}_{i,t}\odot\mathbf{b}_{t})^{\odot-1}$ is a properly chosen scaling vector to produce equal weighting for participating $\mathbf{w}_{i}$’s in (III-B) as desired in (5), and $(\mathbf{X})^{\odot-1}$ represents the inverse Hadamard operation of $\mathbf{X}$ that calculates its entry-wise reciprocal. Noticeably, in order to implement the averaging of (5) in FL over the air, such a post-processing operation requires $\mathbf{b}_{t}$ to be the same for all workers for given $t$ and $d$, which allows to eliminate $\mathbf{b}_{t}$ from the first term in (III-B).

Comparing (III-B) with (5), there exist differences between $\mathbf{w}_{t}$ and $\mathbf{w}$ due to the effect of wireless communications. This work aims to mitigate such a gap through optimizing the worker selection $\bm{\beta}_{i,t}$ and power scaling $\mathbf{b}_{t}$ for FL over the air. To this end, our next step is to unveil an important but unexplored foundation, i.e., how wireless communications affect the convergence behavior of FL over the air.

We first make the following assumptions that are commonly adopted in the optimization literature [30, 31, 32, 33, 7, 34].

Assumption 1 (Lipschitz continuity, smoothness): The gradient $\nabla F(\mathbf{w})$ of the loss function $F(\mathbf{w})$ is uniformly Lipschitz continuous with respect to $\mathbf{w}$, that is,

where $L$ is a positive constant, referred to as a Lipschitz constant for the function $F(\cdot)$.

Assumption 2 (strongly convex): $\nabla F(\mathbf{w})$ is strongly convex with a positive parameter $\mu$, obeying

Assumption 3 (bounded local gradients): The sample-wised local gradients at local workers are bounded by their global counterpart[32, 33]

where $\rho_{1},\rho_{2}\geq 0$.

According to [5, 35], the FL algorithm applied over ideal wireless channels is able to solve P1 and converges to an optimal $\mathbf{w}^{*}$. In the presence of wireless transmission errors, we derive the expected convergence rate of the FL over the air with analog aggregation, as in Theorem 1.

Based on Theorem 1, we further derive the cumulative performance gap resulted from wireless communications and the worker selection of the whole FL process, summarized by the following Lemma 1.

Lemma 1 reveals that the FL algorithm converges asymptotically in $t$ under mild conditions, as stated in Proposition 1.

Proposition 1 indicates that the convergence behavior of the FL algorithm depends on both the learning-related parameters, i.e., $\mu,L,\rho_{1},\rho_{2}$, and communication-related parameters, including $\bm{\beta}$, $\mathbf{b}$ and $\sigma^{2}$. Interestingly, the channel noise $\sigma^{2}$ and $\mathbf{b}$ do not affect $A_{t}$, and hence they do not affect the convergence of the FL algorithm but determine the steady state that the FL algorithm converges to.

Lemma 1 also provides the expected convergence rate of an FL algorithm when the transmission link is error-free. In this ideal case, it offers the fastest convergence rate achievable by the FL algorithm, which is derived by the following Lemma 2.

It is worth noting that Lemma 2 provides the convergence rate for the ideal case, which assumes that the impacts of wireless communications, including noise, channel and constrained resources, are all mitigated to result in error-free transmission. According to (16) in the realistic case, the trajectory of $\mathbb{E}[F(\mathbf{w}_{t+1})]$ exhibits jump discontinuity with a gap term $\Delta_{t}$ at each step $t$, as defined in (16):

This gap reflects the impact of wireless communication factors on FL, by means of the worker selection, transmit power scaling and AWGN. Intuitively, this gap diminishes as the number of selected workers increases, which reduces the value of $A_{t}$. Meanwhile, as the power scaling factor $\bm{b}_{t}$ increases, $B_{t}$ is decreased, which leads to a reduction of the gap as well. Hence, it is necessary to optimize transmit power scaling factors and worker selection in order to minimize the gap in (16) for the implementation of FL algorithms over a realistic wireless network.

When the loss function $F(w)$ is nonconvex, such as in the case of convolutional neural networks, we derive the convergence rate of the FL over the air with analog aggregation for the nonconvex case without Assumption 2, which is summarized in Theorem 2.

As we can see from Theorem 2, when $T$ is large enough, we have

which guarantees convergence of the FL algorithm to a stationary point[36, 28]. Similarly, the performance gap at the step $t$ for non-convex cases is given by

Note that the non-convex case and the convex shares the same sufficient condition for convergence as (18) in Proposition 1.

Our work can be extended to stochastic versions of gradient descent (SGD) as well. Here, we provide convergence analysis for mini-batch gradient descent with a constant mini-batch size $K_{b}$, while the results directly apply to the standard SGD by setting $K_{b}=1$. Theorem 3 summarizes the convergence behavior of SGD for the strongly convex case.

From Theorem 3, the cumulative performance gap of FL after the $t$-th iteration for the SGD case is bounded by

Similarly, we can also derive the mild convergence condition for SGD, given by the following Proposition 2.

Since we are concerned with convergence accuracy, our optimization problem boils down to minimizing the performance gap for different cases (i.e., $\bigtriangleup_{t}$, $\bigtriangleup_{t}^{NC}$, and $\bigtriangleup_{t}^{SGD}$) at each iteration under the corresponding convergence conditions (i.e., Proposition 1 and Proposition 2).

We recognize that solving P1 amounts to iteratively minimizing those gap $\bigtriangleup_{t}$, $\bigtriangleup_{t}^{NC}$, and $\bigtriangleup_{t}^{SGD}$ under the transmit power constraint in (7). At the $t$-th iteration, the objective functions for those three cases are given by

where $\bigtriangleup_{0}=0$ and $\bigtriangleup^{SGD}_{0}=0$. Note that when the optimization is executed at the $t$-th iteration, $\bigtriangleup_{t-1}$ and $\bigtriangleup^{SGD}_{t-1}$ can be treated as constants.

Considering the entry-wise transmission for analog aggregation, we remove irrelevant items and extract the component of the $d$-th entry from those gap in (32), (33) and (34) as the objective to minimize, which is given by

Since all entries indexed by $d$ are separable with respect to the design parameters, we perform entry-wise optimization by considering $\mathbf{w}_{t}$ and $\mathbf{w}_{i,t}$ one entry at a time, where the superscript $d$ and the index of different cases are omitted hereafter. To determine the worker selection vector $\beta_{i,t}$ and the power scaling factor $b_{t}$ at the $t$-th iteration, the PS carries out a joint optimization problem formulated as follows:

where $K_{i}$ should be $K_{b}$ in (38b) for the SGD case.

However, in (38b), the knowledge of $\{w_{i,t}\}_{i=1}^{U}$ is needed but is unavailable to the PS due to analog aggregation.

To overcome this issue, we reformulate a practical optimization problem via an approximation that $\mathbf{w}_{t-1}\approx\frac{1}{U}\sum_{i=1}^{U}\mathbf{w}_{i,t}$, in light of the consensus constraint in P1. According to (47) in the proof of Theorem 1, each local parameter $\mathbf{w}_{i,t}$ is updated from the broadcast $\mathbf{w}_{t-1}$ along the direction of the averaged gradient over its local data $\frac{\alpha}{K_{i}}\sum_{k=1}^{K_{i}}\nabla f(\mathbf{w}_{t-1};\mathbf{x}_{i,k},\mathbf{y}_{i,k})$. Hence, it is reasonable to make the following common assumption on bounded local gradients, considering that the local gradients can be controlled by adjusting the learning rate or through simple clipping [21, 28, 37, 38].

Assumption 4 (bounded local gradients): The gap between the global parameter $w_{t-1}$ and the local parameter update $w_{i,t},\forall i,t$ is bounded by

where $\eta\geq 0$ is related to the learning rate $\alpha$ that satisfies the following condition⁴⁴4This implies the value range of $\eta$. In practice, $\eta$ can take $|w_{t-1}-w_{t-2}|$. In addition, for the SGD case, we have $\eta\geq\max\{\{|\alpha\mathbb{E}_{\mathcal{D}_{i}}[\nabla f(w,\mathbf{x}_{i,k},\mathbf{y}_{i,k})]|\}_{i=1}^{U}\}$

Under Assumption 4, we reformulate the original optimization problem (P2) into the following problem (P3), by replacing $w_{i,t}$ in (38b) by its approximation:

where the power constraint (41b) is constructed based on the fact that

Since $w_{t-1}$ is always available at the PS, P3 becomes a feasible formulation for adoption in practice. Next, we develop the optimal solution to P3.

At first glance, a direct solution to P3 leads to a mixed integer programming (MIP), which unfortunately incurs high complexity. To solve P3 in an efficient manner, we develop a simple solution by identifying a tight search space without loss of optimality. The tight search space, given in the following Theorem 4, is a result of the constraints in (41b) and (41c), irrespective of the objective function (41a). Hence, it holds universally for any $R_{t}$, such as those in (35)-(37).

Thanks to Theorem 4, we equivalently transform P3 from a MIP into a discrete programming (DP) problem P4 as follows