Over-the-Air Federated Learning Over MIMO Channels: A Sparse-Coded Multiplexing Approach

We start with the description of the FL task deployed on a wireless communication system, where the system consists of one central PS and $M$ edge devices. We assume that the training data of the FL task are all distributedly stored on the edge devices. Let $\mathcal{A}_{m}$ denote the local dataset of the $m$-th device, and $Q_{m}$ denote the cardinality of $\mathcal{A}_{m}$. $Q=\sum_{m=1}^{M}Q_{m}$ is the total number of training data samples for the FL task. $\boldsymbol{\theta}\in\mathbb{R}^{D}$ is the model parameter vector with $D$ being the total length of the model parameter. The target of the FL task is to minimize an empirical loss function $F(\boldsymbol{\theta})$ based on the local datasets $\{\mathcal{A}_{m}\}$, given by

where $F_{m}(\boldsymbol{\theta})=\frac{1}{Q_{m}}\sum_{n=1}^{Q_{m}}f\!(\boldsymbol{\theta};\boldsymbol{\zeta}_{m,n})$ is the local loss function of device $m$, and $f\!(\boldsymbol{\theta};\boldsymbol{\zeta}_{m,n})$ is the sample-wise loss function for the $n$-th training sample $\boldsymbol{\zeta}_{m,n}$ in $\mathcal{A}_{m}$.

To minimize the empirical loss function $F(\boldsymbol{\theta})$ in \eqrefeq:FLTarget, the FL training involves $T$ communication rounds between the edge devices and the PS for $F(\boldsymbol{\theta})$ to reach convergence. Specifically, each communication round $t\in[T]$ consists of four steps:

• •

  Global model download: The PS sends the global model $\boldsymbol{\theta}^{(t)}$ to each edge device.

• •

  Local gradients computation: The local gradient $\mathbf{g}_{m}^{(t)}\in\mathbb{R}^{D}$ is computed by device $m$ based on their own data and the global model, given by

  $$\mathbf{g}_{m}^{(t)}=\nabla F_{m}(\boldsymbol{\theta}^{(t)})=\frac{1}{Q_{m}}\sum\nolimits_{n=1}^{Q_{m}}\nabla f\left(\boldsymbol{\theta}^{(t)};\boldsymbol{\zeta}_{m,n}\right).$$

  (2)

• •

  Local gradients upload: The edge devices send the local gradients to the PS through wireless channels.

• •

  Global model aggregation: The local gradients are aggregated as

  $$\mathbf{g}^{(t)}=\frac{1}{Q}\sum\nolimits_{m=1}^{M}Q_{m}\mathbf{g}_{m}^{(t)}=\sum\nolimits_{m=1}^{M}q_{m}\mathbf{g}_{m}^{(t)},$$

  (3)

  where $\mathbf{g}^{(t)}$ denotes the aggregated gradient, and $q_{m}=Q_{m}/Q,\forall m\in[M]$. The global model $\boldsymbol{\theta}^{(t+1)}$ is updated by

  $$\boldsymbol{\theta}^{(t+1)}=\boldsymbol{\theta}^{(t)}-\eta\mathbf{g}^{(t)},$$

  (4)

  where $\eta\in\mathbb{R}$ denotes the learning rate.

We now introduce the wireless multi-user multiple-input multiple-output (MIMO) channel for the above FL system. As depicted in Fig. 1, the considered MIMO OA-FL system consists of a PS with $N_{\mathrm{R}}$ antennas, and $M$ edge devices with each equipped with $N_{\mathrm{T}}$ antennas. As in previous studies in OA-FL [amiri2020federated, sery2021over, cao2022transmission], we make two assumptions: that the download of the global model is through error-free links¹¹1 In practice, the channel noise causes communication errors in the model download. This issue can be addressed by the schemes proposed by [vu2020cell, vu2022joint]. This is beyond the scope of this paper and hence omitted here. and that the devices upload the local gradients to the PS synchronously²²2 The existing techniques in 4G Long Term Evolution, e.g., the timing advance (TA) mechanism, can achieve the synchronization of the gradient symbols among the edge devices [3gpp.38.213]. . We now focus on the process of local gradients uploading. We consider a block-fading channel, where the channel state information (CSI) remains constant during the gradients uploading. Let $\mathbf{H}_{m}^{(t)}\in\mathbb{C}^{N_{\mathrm{R}}\times N_{\mathrm{T}}}$ denote the CSI matrix between the $m$-th device and the PS at the $t$-th round. We assume that the PS has perfect knowledge of the CSI of the wireless channels between the devices and the PS³³3 The approaches of CSI estimation over MIMO MAC channels can be referred to [nguyen2013compressive, wen2014channel, vu2020cell, vu2022joint].. Thus, at the PS, the receive signal matrix from the above MIMO multiple access (MAC) channel is given by

where $K$ denotes the number of channel uses at the $t$-th round; $\mathbf{X}_{m}^{(t)}\in\mathbb{C}^{N_{\mathrm{T}}\times K}$ denotes the transmit data matrix for the $m$-th device; and $\mathbf{N}^{(t)}\in\mathbb{C}^{N_{\mathrm{R}}\times K}$ is an additive white Gaussian noise (AWGN) matrix, with the entries independently drawn from $\mathcal{CN}(0,\sigma_{\mathrm{noise}})$. Let $\mathbf{x}_{m,k}$ denote the $k$-th column of $\mathbf{X}_{m}^{(t)}$. Here, we consider the following transmit power constraint:

where $P_{0}$ is the transmit power budget.

What remains are to map the local gradients $\{\mathbf{g}_{m}^{(t)}\}$ to the transmit matrices $\{\mathbf{X}_{m}^{(t)}\}$ at the edge devices, and to recover the aggregated gradient from the receive signal matrix $\mathbf{Y}^{(t)}$. These issues are discussed in detail in the next section.

To support the local gradients uploading in the MIMO OA-FL system, the pre-processing operations are first conducted on the edge devices, including the gradient sparsification [amiri2020federated] and the gradient compression[ma2014turbo].

To be specific, for the $m$-th edge device, the local gradient $\mathbf{g}_{m}^{(t)}$ at the $t$-th round is first complexified to fully utilize the spectral efficiency of complex channels. The complexified gradient is denoted by

where $j=\sqrt{-1}$. Then, the accumulated gradient is obtained by

where $\boldsymbol{\Delta}_{m}^{(t)}$ denotes the error accumulation vector of the $m$-th device at the $t$-th round with $\Delta_{m}^{(0)}$ initialized as $\mathbf{0}$. Having calculated the accumulated gradient, in the sparsification, the $m$-th device obtains the sparsified gradient via

where $\lambda\in[0,1]$ denotes the sparsity ratio. The operator $\mathrm{sp}(\cdot)$ retains the $\lambda D/2$ entries of $\mathbf{g}_{m}^{\mathrm{ac}(t)}$ with the largest absolute value magnitude, and sets the remaining $(1-\lambda)D/2$ entries to $0$. The error accumulation vector $\boldsymbol{\Delta}_{m}^{(t+1)}$ is updated via

Then $\mathbf{g}_{m}^{\mathrm{sp}(t)}$ is normalized to $\mathbf{g}_{m}^{\mathrm{no}(t)}$ by

where $\mathbf{s}\in\mathbb{C}^{D/2}$ is a random flipping vector, with each entry of $\mathbf{s}$ being independent and identically distributed (i.i.d.) drawn from $\{-1,+1\}$; and $\sigma_{m}^{(t)}=\frac{1}{D/2}\sum_{d=1}^{D/2}|g_{m}^{\mathrm{sp}(t)}[d]|^{2}$ denotes the variance of $\{g_{m}^{\mathrm{sp}(t)}[d]\}_{d=1}^{D/2}$, with $g_{m}^{\mathrm{sp}(t)}[d]$ being the $d$-th entry of $\mathbf{g}_{m}^{\mathrm{sp}(t)}$. From \eqrefeq:g_normalize, the entries of $\mathbf{g}_{m}^{\mathrm{no}(t)}$ have zero-mean and unit variance.

Each device $m$ then compresses $\mathbf{g}_{m}^{\mathrm{no}(t)}$ into a low-dimensional vector via a common compressing matrix. Specifically, the compressed gradient is given by

where $\mathbf{A}\in\mathbb{C}^{C\times(D/2)}$ denotes the common compressing matrix⁴⁴4 We assume that both the compression matrix $\mathbf{A}$ and the flipping vector $\mathbf{s}$ keep invariant throughout the FL training process, and are shared among the devices prior to the FL training., $C$ denotes the length of the compressed gradient, and $\kappa=\frac{C}{D/2}$ denotes the compression ratio. Inspired by [ma2014turbo, ma2022over], we employ a partial DFT matrix as the compressing matrix, given by $\mathbf{A}=\mathbf{S}\boldsymbol{\Xi}$. $\mathbf{S}\in\mathbb{R}^{C\times D/2}$ is a selection matrix consisting of $C$ randomly selected and reordered rows of the $D/2\times D/2$ identity matrix $\mathbf{I}_{D/2}$; and $\boldsymbol{\Xi}\in\mathbb{R}^{D/2\times D/2}$ is a unitary DFT matrix, where the $(d,d^{\prime})$-th entry of $\boldsymbol{\Xi}$ is given by $\frac{1}{\sqrt{D/2}}\exp\left(-2\pi j\frac{(d-1)(d^{\prime}-1)}{D/2}\right)$. We note that the partial DFT matrix has lower computational complexity and better performance, compared with other types of compressed matrix such as i.i.d. Gaussian matrix [ma2015performance].

To transmit the data with multiple streams, device $m$ then reshapes the compressed gradient $\mathbf{g}_{m}^{\mathrm{cp}(t)}$ into $\mathbf{G}_{m}^{(t)}\in\mathbb{C}^{N_{\mathrm{S}}\times K}$ as

where $N_{\mathrm{S}}$ denotes the number of data streams, $\mathbf{g}_{m,n}^{\mathrm{cp}(t)}=\mathbf{g}_{m}^{\mathrm{cp}(t)}((n-1)K+1:nK),n\in[N_{\mathrm{S}}]$. Naturally, the number of channel uses satisfies the following equation:

We are now ready to describe the design of the transmit matrix $\mathbf{X}_{m}^{(t)}$. The transmit matrix $\mathbf{X}_{m}^{(t)}$ is given by $\mathbf{X}_{m}^{(t)}=\mathbf{P}_{m}^{(t)}\mathbf{G}_{m}^{(t)}$, where $\mathbf{P}_{m}^{(t)}\in\mathbb{C}^{N_{\mathrm{T}}\times N_{\mathrm{S}}}$ denotes the precoding matrix for the $m$-th device. Let $\mathbf{g}_{m,k}^{(t)}\in\mathbb{C}^{N_{\mathrm{S}}}$ denote the $k$-th column of $\mathbf{G}_{m}^{(t)}$. We have $\mathbf{x}_{m,k}^{(t)}=\mathbf{P}_{m}^{(t)}\mathbf{g}_{m,k}^{(t)},\forall k\in[K]$. From the transmit power constraint in \eqrefeq:trans_power_constr1, we have

where the step (a) is due to the normalization in \eqrefeq:g_normalize.

We now describe the processing operations on the PS. In the following, we first transfer the recovery of the aggregated gradient to a compressed sensing problem, and then adopt the Turbo compressed sensing (Turbo-CS) algorithm [ma2014turbo] to solve this problem.

At the PS, the receive signal matrix $\mathbf{Y}$ in \eqrefeq:receive_data is first processed through the post processing matrix $\mathbf{F}^{(t)}$, and the post-processed matrix $\hat{\mathbf{R}}^{(t)}$ is given by

$\hat{\mathbf{R}}$ is an approximation to the compressed gradient aggregation matrix $\mathbf{R}^{(t)}=\sum\nolimits_{m=1}^{M}q_{m}\mathbf{G}_{m}$, and the residual error is given by

Based on \eqrefeq:error_receive_data, the PS converts the post-processed matrix $\hat{\mathbf{R}}$ into an equivalent vector form:

where step (a) is from \eqrefeq:g_compress and \eqrefeq:G_def, together with the definition of $\mathbf{g}^{\mathrm{no}(t)}=\sum_{m=1}^{M}q_{m}\mathbf{g}_{m}^{\mathrm{no}(t)}$ and $\mathbf{w}^{(t)}=\vectorize\left(\mathbf{W}^{(t)}{}^{\mathrm{T}}\right)$. We assume that the entries of $\mathbf{g}^{\mathrm{no}(t)}$ are independently drawn from a Bernoulli Gaussian distribution:

where $g_{d}^{\mathrm{no}(t)}$ is the $d$-th entry of $\mathbf{g}^{\mathrm{no}(t)}$, $\lambda^{\prime(t)}$ is the sparsity ratio of $\mathbf{g}^{\mathrm{no}(t)}$, which can be estimated by the Expectation-Maximization algorithm [vila2013expectation], and $\sigma_{g}^{(t)}$ is the variance of the nonzero entries in $\mathbf{g}^{\mathrm{no}(t)}$. Moreover, we assume that the entries of $\mathbf{w}^{(t)}$ are i.i.d. drawn from $\mathcal{CN}(0,\sigma_{w}^{(t)})$, where $\sigma_{w}^{(t)}$ represents the mean square error (MSE) of the compressed gradient aggregation, given by

The recovery of $\mathbf{g}^{\mathrm{no}(t)}$ from $\hat{\mathbf{r}}^{(t)}$ in \eqrefeq:compressed_problem is a compressed sensing problem, where the local gradients $\{\mathbf{g}_{m}^{(t)}\}$ are compressed by the partial DFT matrix $\mathbf{A}$ on the edge devices. From [ma2014turbo, ma2022over], we see that the Turbo-CS algorithm is the state-of-the-art to solve the compressed sensing problem with partial-orthogonal sensing matrices. Thus, we employ the Turbo-CS algorithm to recover $\mathbf{g}^{\mathrm{no}(t)}$:

The details of the Turbo-CS algorithm are presented in the next subsection. Meanwhile, the performance of the Turbo-CS algorithm to recover the aggregated gradient is related to the variance of the noise $\mathbf{w}^{(t)}$, i.e., $\sigma_{w}^{(t)}$. A smaller $\sigma_{w}^{(t)}$ leads to a better recovery performance and a less loss of learning accuracy, which is discussed in Section IV.

As shown in Fig. 3, Turbo-CS conducts the iteration between modules A and B until convergence. Module A is a linear minimum mean-squared error (LMMSE) estimator handling the linear constraint in \eqrefeq:compressed_problem, and module B is a minimum mean-squared error (MMSE) denoiser exploiting the sparsity in \eqrefeq:g_agg_sparsity. We next give the operations of Turbo-CS in detail.

The iterative process begins with module A. The inputs of the LMMSE estimator in module A are the a priori mean $\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)}\in\mathbb{C}^{D/2}$, the a priori covariance $v_{\mathrm{A}}^{\mathrm{pri}(t)}$, and the observed vector $\hat{\mathbf{r}}^{(t)}$. With given $\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)}$, $v_{\mathrm{A}}^{\mathrm{pri}(t)}$, and $\hat{\mathbf{r}}^{(t)}$, the a posteriori mean $\mathbf{g}_{\mathrm{A}}^{\mathrm{post}(t)}$ and the a posteriori covariance $v_{\mathrm{A}}^{\mathrm{post}(t)}$ of the LMMSE estimator are given by [kay1993fundamentals] {subequations} {align} g_A^post(t) & = g_A^pri(t) + vApri(t)vApri(t)+ σw(t) A^H (^r^(t) - A g_A^pri(t)),
v_A^post(t) = v_A^pri(t) - κ⋅( vApri(t))2vApri(t)+ σw(t), where $\kappa$ is the compression ratio defined below \eqrefeq:g_compress. Then the extrinsic mean and variance of the LMMSE estimator are given by {subequations} {align} g_A^ext(t) &= v_A^ext(t) ( gApost(t)vApost(t)-gApri(t)vApri(t) ),
v_A^ext(t) = (1/v_A^post(t)-1/v_A^pri(t) )^-1. The extrinsic messages $\{\mathbf{g}_{\mathrm{A}}^{\mathrm{ext}(t)},v_{\mathrm{A}}^{\mathrm{ext}(t)}\}$ are used to update the a priori mean $\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)}$ and the a priori variance $v_{\mathrm{B}}^{\mathrm{pri}(t)}$ as $\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)}=\mathbf{g}_{\mathrm{A}}^{\mathrm{ext}(t)}$, and $v_{\mathrm{B}}^{\mathrm{pri}(t)}=v_{\mathrm{A}}^{\mathrm{ext}(t)}$. Both $\{\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)},v_{\mathrm{B}}^{\mathrm{pri}(t)}\}$ are the inputs of the MMSE denoiser in module B.

In module B, we model the a priori mean $\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)}$ as an observation of $\mathbf{g}^{\mathrm{no}(t)}$ corrupted by additive noise: $\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)}\!=\!\mathbf{g}^{\mathrm{no}(t)}\!+\!\boldsymbol{\delta}^{(t)}$, where $\delta_{d}^{(t)}\sim\mathcal{CN}(0,v_{\mathrm{B}}^{\mathrm{pri}(t)})$ denotes the $d$-th entry of $\boldsymbol{\delta}^{(t)}$. The a posteriori mean $\mathbf{g}_{\mathrm{B}}^{\mathrm{post}(t)}$ and the variance $v_{\mathrm{B}}^{\mathrm{post}(t)}$ of the MMSE denoiser are given by {subequations} {align} g_B^post(t) &= \mathbbE[g^no(t)—g_B^pri(t)],
v_B^post(t) = 1D/2∑_d=1^D/2 \operatornamevar[g^no(t)_d—g_B,d^pri(t)], where $\operatorname{var}[a|b]=\mathbb{E}\Big{[}\big{|}a-\mathbb{E}[a|b]\big{|}^{2}\Big{|}b\Big{]}$, and $g_{\mathrm{B},d}^{\mathrm{pri}(t)}$ is the $d$-th entry of $\mathbf{g}_{\mathrm{B}}^{\mathrm{pri}(t)}$. Then the extrinsic messages of the MMSE denoiser are given by {subequations} {align} &g_B^ext(t) = v_B^ext(t) (gBpost(t)vBpost(t) - gBpri(t)vBpri(t) ),
v_B^ext(t) = ( 1/v_B^post(t) - 1/v_B^pri(t) )^-1. The extrinsic messages $\{\mathbf{g}_{\mathrm{B}}^{\mathrm{ext}(t)},v_{\mathrm{B}}^{\mathrm{ext}(t)}\}$ are used to update the a priori mean $\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)}$ and the a priori variance $v_{\mathrm{A}}^{\mathrm{pri}(t)}$ as $\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)}=\mathbf{g}_{\mathrm{B}}^{\mathrm{ext}(t)}$, and $v_{\mathrm{A}}^{\mathrm{pri}(t)}=v_{\mathrm{B}}^{\mathrm{ext}(t)}$. Both $\{\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)},v_{\mathrm{A}}^{\mathrm{pri}(t)}\}$ are the inputs of the LMMSE estimator in module A. At the end of the iterative process, the final estimate is based on the a posteriori output of the module B, i.e., $\hat{\mathbf{g}}^{\mathrm{no}(t)}=\mathbf{g}_{\mathrm{B}}^{\mathrm{post}(t)}$.

Then, based on the output of Turbo-CS, the aggregated gradient $\hat{\mathbf{g}}^{(t)}$ is given by

where $\sigma^{(t)}=\frac{1}{M}\sum_{m=1}^{M}\sigma_{m}^{(t)}$. The PS then updates the model $\boldsymbol{\theta}^{(t+1)}$ by

To sum up, at each round $t$, the receive matrix $\hat{\mathbf{R}}^{(t)}$ is first reshaped into a vector form in \eqrefeq:compressed_problem. Given the observed vector $\hat{\mathbf{r}}^{(t)}$ and the initial values $\{\mathbf{g}_{\mathrm{A}}^{\mathrm{pri}(t)}=\mathbf{0},v_{\mathrm{A}}^{\mathrm{pri}(t)}=1\}$, Turbo-CS iteratively calculates \eqrefeq:g_v_post_A-\eqrefeq:g_v_ext_B until a certain termination criterion is met. Finally, the output $\hat{\mathbf{g}}^{\mathrm{no}(t)}$ is scaled as $\hat{\mathbf{g}}^{(t)}$ for the model update in \eqrefeq:SGD_error.

The proposed SCoM approach to the local gradients uploading is summarized below, where lines 3-5, 11-17 are executed at the PS, and lines 6-10 are executed at the devices.

[htb] \floatnamealgorithm Proposed SCoM-Based MIMO OA-FL Scheme {algorithmic}[1] \REQUIRE$T$, $\{Q_{m}\}$, $N_{\mathrm{S}}$, $N_{\mathrm{R}}$, $N_{\mathrm{T}}$, $\kappa$, $\lambda$, $\mathbf{A}$ and $\mathbf{s}$. \STATEInitialization: $t=0$, the global model $\boldsymbol{\theta}^{(0)}$. \FOR $t\in[T]$ \STATEPS does: \STATEEstimate the CSI matrices $\{\mathbf{H}^{(t)}\}$, and optimize $\mathbf{F}^{(t)}$ and $\{\mathbf{P}_{m}^{(t)}\}$; \STATESend $\boldsymbol{\theta}^{(t)}$ and $\{\mathbf{P}_{m}^{(t)}\}$ to the edge devices through error free links; \STATEEach device $m$ does in parallel: \STATECompute $\mathbf{g}_{m}^{(t)}$ based on the $\mathcal{A}_{m}$ and $\boldsymbol{\theta}^{(t)}$ via \eqrefeq:GradDef; \STATECompute $\mathbf{G}_{m}^{(t)}$ via \eqrefeq:Complexification-\eqrefeq:G_def; \STATEUpdate $\boldsymbol{\Delta}_{m}^{(t+1)}$ via \eqrefeq:Delta_update; \STATESend $\mathbf{G}_{m}^{(t)}$ to the PS with the precoding matrix $\mathbf{P}_{m}^{(t)}$ over the MIMO channel in \eqrefeq:receive_data; \STATEPS does: \STATEReshape $\hat{\mathbf{R}}^{(t)}$ into $\hat{\mathbf{r}}^{(t)}$ via \eqrefeq:compressed_problem;