Communication-Efficient Federated Learning over MIMO Multiple Access Channels

We consider federated learning over a wireless MIMO MAC in which a PS equipped with $U$ antennas trains a global model (e.g., deep neural network) by collaborating with $K$ single-antenna wireless devices, as illustrated in Fig. 1. We assume that the global model is fully represented by a parameter vector ${\bf w}\in\mathbb{R}^{\bar{N}}$, where $\bar{N}$ is the number of the parameters (e.g., the weights and biases of a deep neural network). We also assume that the training data samples to optimize the global model are available only at the wireless devices and differ across the devices. Let $\mathcal{D}_{k}$ be a local training data set available at device $k$, which consists of $|\mathcal{D}_{k}|$ training data samples. A local loss function at device $k\in\mathcal{K}=\{1,\ldots,K\}$ for the parameter vector ${\bf w}$ is defined as

where $f({\bf w};{\bf u})$ is a loss function computed with a training data sample ${\bf u}\in\mathcal{D}_{k}$ for the parameter vector ${\bf w}$. Then a global loss function for the parameter vector ${\bf w}$ is represented as

We focus on a scenario where a gradient-based algorithm (e.g., SGD algorithm) is adopted to optimize the parameter vector for minimizing the global loss function in (2). We denote $T$ as the number of total iterations of the optimization algorithm, and ${\bf w}_{t}\in{\mathbb{R}}^{\bar{N}}$ as the parameter vector at iteration $t$. Each iteration of the optimization algorithm corresponds to one communication round consisting of a pair of downlink and uplink communications. During the downlink communication, the PS broadcasts the current parameter vector to the wireless devices. Then, during the uplink communication, the wireless devices transmit the gradient of the parameter vector computed based on their local training data sets to the server. The major focus of our work is to investigate uplink transmission and reception strategies at the server for federated learning; thereby, for simplicity of the analysis, we assume that the downlink communication is error free¹¹1The framework under consideration and the proposed approach can be readily extended to a noisy downlink scenario[35]. Under this assumption, all the devices have a globally consistent parameter vector ${\bf w}_{t}$ for all $t\in\{1,\ldots,T\}$. In what follows, we describe uplink transmission at the wireless devices and uplink reception at the PS during each uplink communication.

For the uplink transmission, each device computes a local gradient vector for the current parameter vector based on its local training data set. A local gradient vector computed at device $k$ is given by

where $\mathcal{D}_{k}^{(t)}\subset\mathcal{D}_{k}$ is a mini-batch randomly drawn from $\mathcal{D}_{k}$, and $\nabla$ is a gradient operator. We assume that the mini-batch size at device $k$ is fixed as $D_{k}^{\rm mini}=|\mathcal{D}_{k}^{(t)}|$ for all $t\in\{1,\ldots,T\}$. Since direct transmission of the local gradient vector in (3) imposes large communication overhead when $\bar{N}\gg 1$, we assume that each device applies lossy compression to its local gradient vector before uplink transmission. Details of the gradient compression technique considered in our work will be elaborated in Sec. III-A. Let ${\bf x}_{k}^{(t)}=\big{[}x_{k}^{(t)}[1],\cdots,x_{k}^{(t)}[\bar{M}]\big{]}^{\sf T}\in\mathbb{R}^{\bar{M}}$ be a compressed gradient vector determined at device $k$ with $\bar{M}<\bar{N}$. Then an uplink signal vector transmitted from device $k$ at iteration $t$ is denoted by $\tilde{\bf x}_{k}^{(t)}=\sqrt{P_{k}^{(t)}}{\bf x}_{k}^{(t)}$, where $P_{k}^{(t)}$ is a power scaling factor for device $k$ at iteration $t$. In particular, the power scaling factor is set to be $P_{k}^{(t)}=\frac{\bar{M}}{\|{\bf x}_{k}^{(t)}\|^{2}}$, so that an average power of transmitting each entry of the uplink signal vector becomes one (i.e., $\mathbb{E}[|(\tilde{\bf x}_{k}^{(t)})_{m}|^{2}]=1$, $\forall m$). All $K$ devices simultaneously transmit their uplink signal vectors over a wireless MAC by utilizing $\bar{M}$ shared radio resources that are orthogonal in either the time or the frequency domain. We assume that the power scaling factor is separately conveyed to the PS in an error-free fashion because the overhead to transmit this scalar value is negligible compared to communication overhead to transmit the compressed gradient vector.

Once the uplink transmission of the wireless devices ends, the PS reconstructs the local gradient vectors from uplink received signals. Let ${\bf x}^{(t)}[m]=\big{[}x_{1}^{(t)}[m],x_{2}^{(t)}[m],\cdots,x_{K}^{(t)}[m]\big{]}^{\sf T}$ be the aggregation of the compressed gradient entries sent by $K$ devices at the $m$-th resource, where $x_{k}[m]$ is the $m$-th element of ${\bf x}_{k}^{(t)}$. Then an uplink received vector at the PS for the $m$-th resource is given by

where ${\bf P}^{(t)}={\sf diag}\big{(}P_{1}^{(t)},\ldots,P_{K}^{(t)}\big{)}$ is a power allocation matrix, ${\bf H}^{(t)}\in\mathbb{R}^{U\times K}$ is a MIMO channel matrix from $K$ devices to the server, $\tilde{\bf H}^{(t)}\triangleq{\bf H}^{(t)}({\bf P}^{(t)})^{1/2}$, and ${\bf z}^{(t)}[m]\sim\mathcal{N}({\bf 0}_{U},\sigma^{2}{\bf I}_{U})$ is a Gaussian noise vector. We assume that the channel matrix remains constant within one communication round (i.e., block-fading assumption) and is perfectly known at the PS via uplink channel estimation. Based on the received signal vectors $\{{\bf y}^{(t)}[m]\}_{m}$, the PS reconstructs the local gradient vectors, $\{{\bf g}_{k}^{(t)}\}_{k=1}^{K}$, sent from all the $K$ wireless devices. Details of the gradient reconstruction algorithm developed in our work will be elaborated in Sec. IV.

Let $\hat{\bf g}_{k}^{(t)}$ be the reconstruction of the local gradient vector ${\bf g}_{k}^{(t)}$ at the PS. Then the PS compute a global gradient vector from the reconstructed local gradient vectors as follows:

where $\rho_{k}\triangleq\frac{D_{k}^{\rm mini}}{\sum_{j=1}^{K}D_{j}^{\rm mini}}$ is assumed to be known at the PS a priori. The global gradient vector in (5) is utilized to update the parameter vector ${\bf w}_{t}$. For example, if the PS adopts a gradient descent algorithm to optimize the parameter vector, the corresponding update rule is given by

where $\eta_{t}>0$ is a learning rate at iteration $t$.

The key idea of the gradient compression technique is to sparsify a local gradient vector in a block-wise fashion and then compress the sparsified vector as in CS [20, 21]. This technique is a simple variant of the gradient compression technique in [17, 18, 19]. Our compression technique consists of two steps, block sparsification and dimensionality reduction, as elaborated below.

Block Sparsification: The first step of our compression technique is to divide the local gradient vector into $B$ sub-vectors with dimension $N=\frac{\bar{N}}{B}$, then sparsify each sub-vector by dropping some gradient entries with small magnitudes. Meanwhile, to compensate for information loss caused by the gradient dropping, the dropped gradient entries are stored at each device and then added to the local gradient vector at the next iteration. Under the above strategy, the local gradient vector that needs to be sent by device $k$ at iteration $t$ is determined as

where ${\bm{\Delta}}_{k}^{(t)}$ is a residual gradient vector stored by device $k$ before iteration $t$. Then the $b$-th local gradient sub-vector at device $k$ is defined as

where $\bar{g}_{k,n}^{(t)}$ is the $n$-th entry of $\bar{\bf g}_{k}^{(t)}$, and $\mathcal{I}_{k,b}\subset\bar{\mathcal{I}}\triangleq\{1,\ldots,\bar{N}\}$ is a set of parameter indices associated with the $b$-th sub-vector at device $k$. We assume that $\mathcal{I}_{k,1},\ldots,\mathcal{I}_{k,B}$ are mutually exclusive subsets of $\bar{\mathcal{I}}$ such that $\bigcup_{b=1}^{B}\mathcal{I}_{k,b}=\bar{\mathcal{I}}$ and known a priori at the PS. After constructing $B$ sub-vectors, each local gradient sub-vector is sparsified by setting all but the top-$S$ entries with the largest magnitudes as zeros, where $S$ is a sparsification level such that $S\ll N$. We define $S_{\rm ratio}\triangleq S/N$ as a sparsification ratio. As a result, device $k$ attains $B$ local gradient sub-vectors $\{{\bf g}_{k,b}^{(t)}\}$, each of which is an $S$-sparse vector. We denote the block sparsification step described above as ${\sf BlockSparse}\big{(}\bar{\bf g}_{k}^{(t)}\big{)}$. In the sequel, we will show that increasing the number of the gradient sub-vectors, $B$, reduces the computational complexity of the gradient reconstruction process at the PS. However, if $B$ is set too large, the $SB$ non-zero entries in $\{{\bf g}_{k,b}^{(t)}\}$ can significantly differ from the top-$SB$ entries of $\bar{\bf g}_{k}^{(t)}$; this mismatch may lead to a decrease in the performance of federated learning. Therefore, in practical systems, $B$ can be determined according to a target learning performance and the computing power of the PS.

After the block sparsification, the residual gradient vector $\bar{\bf g}_{k}^{(t)}$ is updated as

where ${\sf Concatenate}\big{(}\{{\bf g}_{k,b}^{(t)}\}_{b=1}^{B}\big{)}=[g_{k,1}^{\rm cc},\cdots,g_{k,\bar{N}}^{\rm cc}]^{\sf T}\in\mathbb{R}^{\bar{N}}$ such that $g_{k,\mathcal{I}_{k,b}(n)}^{\rm cc}=({\bf g}_{k,b}^{(t)})_{n}$ for all $b,n$.

Dimensionality Reduction: The second step of our compression technique is to compress each sparse sub-vector by random projection onto a lower dimensional space as in CS [20, 21]. Let ${\bf x}_{k,b}^{(t)}$ be a compressed gradient sub-vector generated from ${\bf g}_{k,b}^{(t)}$. Then ${\bf x}_{k,b}^{(t)}$ is determined as

where ${\bf A}^{(t)}\in\mathbb{R}^{M\times N}$ is a random measurement matrix with $M<N$ at iteration $t$. In this work, we set ${\bf A}^{(t)}$ as an IID random matrix with $({\bf A}^{(t)})_{m,n}\sim\mathcal{N}(0,1/M)$ which provides promising properties (e.g., restricted isometry property, null space property) for guaranteeing accurate recovery of a sparse signal [21]. Furthermore, ${\bf A}^{(t)}$ is randomly and independently drawn for each iteration, in order to avoid a null-space problem caused by a fixed random measurement matrix, as discussed in [31].

By concatenating these $B$ compressed sub-vectors, a compressed gradient vector for device $k$ at iteration $t$ is finally obtained as

The dimension of the compressed vector is $\bar{M}=BM$; thereby, the compression ratio of our technique is $R=\frac{N}{M}>1$. The overall procedures of our gradient compression technique is illustrated in Fig. 1 and also summarized in Steps 5–8 in Algorithm 1.

A significant advantage of our gradient compression technique is that it requires $R$ times lower communication overhead than the direct transmission of the local gradient vector, utilizing $\bar{N}$ radio resources per device. Another prominent advantage is that the communication overhead of our technique does not increase with the number of devices participating in federated learning because all the devices utilize the same resources available in the wireless MAC. Therefore, our compression technique effectively reduces the communication overhead of federated learning with a large number of devices.

Remark 1 (Digital Transmission Strategy): In this work, we focus only on analog transmission at the wireless devices in which the uplink signal vector $\tilde{\bf x}_{k}^{(t)}$ is transmitted without digital modulation. Nevertheless, digital transmission of the uplink signal vector is also feasible by utilizing a digital symbol modulator as a scalar quantizer. When employing our gradient compression technique, the distribution of the uplink signal vector can be modeled as $\tilde{\bf x}_{k}^{(t)}\sim\mathcal{N}({\bf 0}_{\bar{M}},{\bf I}_{\bar{M}})$ for large $N$, which will be justified in Sec. IV-B. Consequently, a complex-valued signal vector, defined as $\tilde{\bf x}_{{\rm c},k}^{(t)}=\big{(}\tilde{\bf x}_{k}^{(t)}\big{)}_{1:\bar{M}/2}+j\big{(}\tilde{\bf x}_{k}^{(t)}\big{)}_{\bar{M}/2+1:\bar{M}}$, follows $\tilde{\bf x}_{{\rm c},k}^{(t)}\sim\mathcal{CN}({\bf 0}_{\bar{M}/2},0.5{\bf I}_{\bar{M}/2})$ for large $N$. This fact implies that each entry of $\tilde{\bf x}_{{\rm c},k}^{(t)}$ can be effectively modulated using quadrature amplitude modulation (QAM). Therefore, our gradient compression technique can be compatible with commercial digital communication systems by utilizing the above digital transmission strategy.

When employing the compression technique in Sec. III-A, the PS faces a new challenge in enabling accurate reconstruction of the global gradient vector in (5). The underlying reason is that the local gradient vectors are not only compressed, but also simultaneously sent via a wireless MIMO MAC that causes different channel fading across devices. To shed light on this challenge, we formulate a gradient reconstruction problem at the PS when employing the compression technique in Sec. III-A.

Let $\mathcal{M}_{b}\triangleq\{(b-1)M+1,\ldots,bM\}$ be a set of resource indices associated with the transmission of the $b$-th compressed sub-vector. Then an uplink received signal matrix associated with the $b$-th compressed sub-vector is given by ${\bf Y}_{b}^{(t)}=\left[{\bf y}^{(t)}[\mathcal{M}_{b}(1)],\cdots,{\bf y}^{(t)}[\mathcal{M}_{b}(M)]\right]$. From (4) and (11), the uplink received signal matrix ${\bf Y}_{b}^{(t)}$ is rewritten as

where ${\bf X}_{b}^{(t)}=\big{[}{\bf x}_{1,b}^{(t)},\cdots,{\bf x}_{K,b}^{(t)}\big{]}$, ${\bf Z}_{b}^{(t)}=\left[{\bf z}^{(t)}[\mathcal{M}_{b}(1)],\cdots,{\bf z}^{(t)}[\mathcal{M}_{b}(M)]\right]$, ${\bf G}_{b}^{(t)}=\big{[}{\bf g}_{1,b}^{(t)},\cdots,{\bf g}_{K,b}^{(t)}\big{]}$, and the equality (a) holds because ${\bf X}_{b}^{(t)}={\bf A}^{(t)}{\bf G}_{b}^{(t)}$ from (10). Note that ${\bf G}_{b}^{(t)}$ is an $SK$-sparse matrix because ${\bf g}_{1,b}^{(t)},\ldots,{\bf g}_{K,b}^{(t)}$ are $S$-sparse vectors by the sparsification step during the gradient compression. Therefore, the gradient reconstruction problem reduces to $B$ parallel CS recovery problems, each of which aims at finding an $SK$-sparse matrix ${\bf G}_{b}^{(t)}\in\mathbb{R}^{N\times K}$ from a noisy linear measurement ${\bf Y}_{b}^{(t)}\in\mathbb{R}^{U\times M}$. The overall procedures of gradient reconstruction and parameter update at the PS are illustrated in Fig. 1 and also summarized in Steps 12–17 in Algorithm 1, where ${\sf GradReconst}(\cdot)$ denotes the gradient reconstruction algorithm chosen by the PS.

The CS recovery problem with the form of (12) has also been studied for various applications [32, 33, 34]. The simplest approach to solve this problem is to formulate an equivalent CS recovery problem by converting the received signal matrix in (12) into an equivalent vector form:

Based on this form, an existing compressed sensing algorithm can be applied to recover a sparse vector ${\rm vec}\big{(}({\bf G}_{b}^{(t)})^{\sf T}\big{)}\in\mathbb{R}^{NK}$ from its noisy linear measurement ${\rm vec}\big{(}{\bf Y}_{b}^{(t)}\big{)}\in\mathbb{R}^{MU}$. This approach, however, requires tremendous computational complexity because the dimensions of both ${\rm vec}\big{(}({\bf G}_{b}^{(t)})^{\sf T}\big{)}$ and ${\rm vec}\big{(}{\bf Y}_{b}^{(t)}\big{)}$ are extremely large in many federated learning applications. For example, Table I in Sec. VI shows that the orthogonal matching pursuit (OMP) algorithm to solve the Kronecker-form problem in (13) requires the complexity order of $6\times 10^{15}$ in our simulation with $B=10$, $R=5$, and $S/N=0.04$. Another existing approach is to directly solve a matrix-form problem in (12) by modifying conventional CS algorithms such as orthogonal matching pursuit (OMP) [32, 33]. This approach, however, still requires significant computational complexity that comes from large-scale matrix multiplications. For example, Table I in Sec. VI shows that the OMP algorithm to solve the matrix-form problem in (12) requires the complexity order of $4\times 10^{15}$ in our simulation with $B=10$, $R=5$, and $S/N=0.04$. To overcome the limitations of the existing CS approaches, in the following section, we will present a new low-complexity approach to solve the CS recovery problem in (12).

The basic principle of the proposed algorithm is to decompose the gradient reconstruction problem in (12) into $K$ small-scale CS recovery problems by applying MIMO detection to explicitly estimate a compressed gradient matrix ${\bf X}_{b}$. To be more specific, the output of the MIMO detection applied to detect ${\bf X}_{b}$ from a received signal matrix in (12) can be expressed as

where ${\bf N}_{b}$ is a MIMO detection error. Then the $k$-th column of $\hat{\bf X}_{b}$, namely $\hat{\bf x}_{k,b}$, is expressed as

where ${\bf n}_{k,b}$ is the $k$-th column of ${\bf N}_{b}$. Since every local gradient sub-vector sent by each device is an $S$-sparse vector as explained in Sec. III-A, reconstructing ${\bf g}_{k,b}\in\mathbb{R}^{N}$ from $\hat{\bf x}_{k,b}\in\mathbb{R}^{M}$ is exactly a noisy CS recovery problem which finds an $N$-dimensional sparse vector from an $M$-dimensional observation. Solving this problem in parallel for all $k\in\mathcal{K}$ and $b\in\{1,\ldots,B\}$ requires much smaller complexity than directly solving the CS recovery problem in (12) which finds a $KN$-dimensional unknown vector from a $UM$-dimensional observation.

Despite the advantage in reducing the computational complexity, a major drawback of our strategy is that error in the MIMO detection (e.g., ${\bf N}_{b}$ in (14)) propagates into the subsequent CS recovery process. To alleviate this drawback, we leverage a turbo decoding principle that allows the MIMO detection and the CS recovery processes to exchange their beliefs about the entries of the compressed gradient matrix. The proposed algorithm based on the above principle is illustrated in Fig. 3.

To enable the turbo decoding principle based on a Bayesian approach, we establish a statistical model for local gradient sub-vectors as well as for compressed gradient sub-vectors, as done in [19]. For statistical modeling of the local gradient sub-vectors, we focus on the fact that every local gradient sub-vector is sparse, while its non-zero entries are arbitrarily distributed. We also use the fact that the entries of each local gradient sub-vector are likely to be independent because they are associated with a randomly drawn subset of the model parameters, as explained in Sec. III-A. Motivated by these facts, we model each local gradient sub-vector as an IID random vector whose entry follows the Bernoulli Gaussian-mixture distribution which is well known for its suitability and generality for modeling a sparse random vector [29]. Note that the probability density function of the Bernoulli Gaussian-mixture distribution with parameter ${\bm{\theta}}=\big{(}\lambda_{0},\{\lambda_{\ell},\mu_{\ell},\phi_{\ell}\}_{\ell=1}^{L}\big{)}$ is given by

where $\sum_{\ell=0}^{L}\lambda_{\ell}=1$. As can be seen in (16), the above distribution is suitable for modeling the local gradient sub-vector because the sparsity of this sub-vector is effectively captured by the Bernoulli model with parameter $\lambda_{0}$, while the distribution of non-zero local gradient entries is well approximated by the Gaussian-mixture model with $L$ components. How to tune the parameter ${\bm{\theta}}$ for each gradient sub-vector will be discussed in Sec. IV-D.

To further justify the Bernoulli Gaussian-mixture model, in Fig. 2, we compare the empirical cumulative distribution function (CDF) of the entries of the local gradient sub-vector obtained from simulation²²2In this simulation, we employ Algorithm 3 with $R=3$ and $T=30$. Further details of the simulation are described in Sec. VI. with the CDF of the Bernoulli Gaussian-mixture distribution whose parameters are determined by Algorithm 3. Fig. 2 shows that the empirical CDF is very similar to the corresponding model-based CDF. This result implies that the local gradient sub-vector can be well modeled by an IID random vector with the Bernoulli Gaussian-mixture distribution.

On the basis of the Bernoulli Gaussian-mixture model, we also determine a statistical model for the compressed gradient entries. First of all, as can be seen from (10), the $m$-th entry of the compressed gradient sub-vector ${\bf x}_{k,b}$ is given by $({\bf x}_{k,b})_{m}=\sum_{n=1}^{N}({\bf A})_{m,n}({\bf g}_{k,b})_{n}$. This implies that $x_{k}[m]$ is nothing but the sum of $N$ independent random variables because the random measurement matrix ${\bf A}$ has IID entries, while the local gradient sub-vector ${\bf g}_{k,b}$ can be modeled as an IID random vector. Therefore, for large $N$, every compressed gradient entry $x_{k}[m]$ can be modeled as a Gaussian random variable by the central limit theorem.

In Module A of the proposed algorithm depicted in Fig. 3, the compressed gradient matrix ${\bf X}_{b}$ is estimated from the uplink received signal matrix ${\bf Y}_{b}$. To minimize the MSE of the estimate of ${\bf X}_{b}$, we employ the MMSE MIMO detection based on a Gaussian model justified in Sec. IV-B. The input of Module A is the prior information about the mean and variance of $x_{k}[m]$, denoted by $\hat{x}_{k}^{\rm pri}[m]\in\mathbb{R}$ and $\nu_{x_{k}[m]}^{\rm pri}>0$, respectively, $\forall k\in\mathcal{K},m\in\mathcal{M}_{b}$. This information is set approximately at the beginning of the algorithm and is passed from Module B during the turbo iterations of the algorithm. As discussed in Sec. IV-B, the prior distribution of ${\bf x}[m]$ can be modeled as ${\bf x}[m]\sim\mathcal{N}\big{(}\hat{\bf x}^{\rm pri}[m],{\bf R}_{{\bf x}[m]}^{\rm pri}\big{)}$ for large $N$ by the central limit theorem, where $\hat{\bf x}^{\rm pri}[m]=\big{[}\hat{x}_{1}^{\rm pri}[m],\cdots,\hat{x}_{K}^{\rm pri}[m]\big{]}^{\sf T}$, ${\bf R}_{{\bf x}[m]}^{\rm pri}={\sf diag}\big{(}{\bm{\nu}}_{{\bf x}[m]}^{\rm pri}\big{)}$, and ${\bm{\nu}}_{{\bf x}[m]}^{\rm pri}=\big{[}\nu_{x_{1}[m]}^{\rm pri},\cdots,\nu_{x_{K}[m]}^{\rm pri}\big{]}^{\sf T}$. Note that ${\bf R}_{{\bf x}[m]}^{\rm pri}$ is diagonal because $\{x_{k}[m]\}_{k\in\mathcal{K}}$ are associated with different parameters in the global model. Then, the posterior distribution of ${\bf x}[m]$ for a given observation ${\bf y}[m]=\tilde{\bf H}{\bf x}[m]+{\bf v}[m]$ with ${\bf v}[m]\sim\mathcal{N}({\bf 0}_{U},\sigma^{2}{\bf I}_{U})$ is determined as $p\big{(}{\bf x}[m]\big{|}{\bf y}[m]\big{)}=\mathcal{N}({\bf x}[m];\hat{\bf x}^{\rm post}[m],{\bf R}_{{\bf x}[m]}^{\rm post})$, where the posterior mean and covariance are computed as [36]

respectively, where ${\bm{\Omega}}_{{\bf x}[m]}=\big{(}\tilde{\bf H}{\bf R}_{{\bf x}[m]}^{\rm pri}\tilde{\bf H}^{\sf T}+\sigma^{2}{\bf I}_{U}\big{)}^{-1}$. The $k$-th entry of $\hat{\bf x}^{\rm post}[m]$, denoted by $\hat{x}_{k}^{\rm post}[m]$, is the MMSE estimate of $x_{k}[m]$ for a given observation ${\bf y}[m]$. Meanwhile, the $k$-th diagonal entry of ${\bf R}_{{\bf x}[m]}^{\rm post}$, denoted by $\nu_{x_{k}[m]}^{\rm post}$, represents the MSE between $x_{k}[m]$ and $\hat{x}_{k}^{\rm post}[m]$.

After computing the MMSE estimates of the compressed gradient entries, Module A passes its knowledge about these values to Module B. As illustrated in Fig. 3, the prior information utilized in Module A has been passed from Module B; thereby, Module A extracts the extrinsic information by excluding the contribution of the prior information from the posterior information. In particular, the extrinsic distribution of $x_{k}[m]$ can be determined by assuming that its prior and posterior distributions are given by $p\big{(}x_{k}[m]\big{)}=\mathcal{N}(x_{k}[m];\hat{x}_{k}^{\rm pri}[m],\nu_{x_{k}[m]}^{\rm pri})$ and $p\big{(}x_{k}[m]\big{|}{\bf y}[m]\big{)}=\mathcal{N}(x_{k}[m];\hat{x}_{k}^{\rm post}[m],\nu_{x_{k}[m]}^{\rm post})$, respectively. Then, by Bayes’ rule, the extrinsic distribution of $x_{k}[m]$ is given by $p\big{(}{\bf y}[m]\big{|}x_{k}[m]\big{)}=\mathcal{N}(x_{k}[m];\hat{x}_{k}^{\rm ext}[m],\nu_{x_{k}[m]}^{\rm ext})$, where the extrinsic mean and variance are computed as [37, 38]

respectively. Then the values of $\{\hat{x}_{k}^{\rm ext}[m],\nu_{x_{k}[m]}^{\rm ext}\}_{k,m}$ are passed to Module B which utilizes these values as the prior information about $\{x_{k}[m]\}_{k,m}$.

In Module B of the proposed algorithm, each local gradient sub-vector is reconstructed from the prior information about the compressed gradient entries. To this end, the extrinsic mean and variance of $x_{k}[m]$ passed from Module A are harnessed as the prior mean and variance of $x_{k}[m]$, respectively, i.e., $x_{k}[m]\sim\mathcal{N}(\hat{x}_{k}^{\rm pri}[m],\nu_{x_{k}[m]}^{\rm pri})$, where $\hat{x}_{k}^{\rm pri}[m]=\hat{x}_{k}^{\rm ext}[m]$ and $\nu_{x_{k}[m]}^{\rm pri}=\nu_{x_{k}[m]}^{\rm ext}$, $\forall k\in\mathcal{K},m\in\mathcal{M}_{b}$. Then the prior mean and covariance of ${\bf x}_{k,b}$ are determined as $\hat{\bf x}_{k,b}^{\rm pri}=\big{[}\hat{x}_{k}^{\rm pri}[\mathcal{M}_{b}(1)],\cdots,\hat{x}_{k}^{\rm pri}[\mathcal{M}_{b}(M)]\big{]}^{\sf T}$ and ${\bf R}_{{\bf x}_{k,b}}^{\rm pri}={\sf diag}\big{(}{\bm{\nu}}_{{\bf n}_{k,b}}^{\rm pri}\big{)}$, respectively, where ${\bm{\nu}}_{{\bf n}_{k,b}}^{\rm pri}=\big{[}{\nu}_{x_{k}[\mathcal{M}_{b}(1)]}^{\rm pri},\cdots,{\nu}_{x_{k}[\mathcal{M}_{b}(M)]}^{\rm pri}\big{]}^{\sf T}$. Note that ${\bf R}_{{\bf x}_{k,b}}^{\rm pri}$ is diagonal because $x_{k}[m_{1}]$ and $x_{k}[m_{2}]$ are associated with different parameters in the global model for $m_{1}\neq m_{2}\in\mathcal{M}_{b}$. For large $N$ with the IID random measurement matrix ${\bf A}$ in (10), all the entries of ${\bf x}_{k,b}$ are statistically identical by the central limit theorem. Motivated by this fact, we approximate the prior covariance as ${\bf R}_{{\bf x}_{k,b}}^{\rm pri}\approx\bar{\nu}_{{\bf x}_{k,b}}^{\rm pri}{\bf I}_{M}$, where $\bar{\nu}_{{\bf x}_{k,b}}^{\rm pri}\triangleq\frac{1}{M}\sum_{m\in\mathcal{M}_{b}}{\nu}_{x_{k}[m]}^{\rm pri}$. Then, with some tedious manipulations from (17)–(19), the prior mean $\hat{\bf x}_{k,b}^{\rm pri}$ is rewritten as an additive white Gaussian noise (AWGN) observation of ${\bf x}_{k,b}$, i.e.,

where ${\bf n}_{k,b}^{\rm pri}\sim\mathcal{N}\big{(}{\bf 0}_{M},\bar{\nu}_{{\bf x}_{k,b}}^{\rm pri}{\bf I}_{M}\big{)}$ is an effective noise independent of ${\bf x}_{k,b}$, as also derived in [38].

To reconstruct the local gradient sub-vector ${\bf g}_{k,b}$ from the prior mean $\hat{\bf x}_{k,b}^{\rm pri}$ in (20), we need to solve a noisy CS recovery problem. Unfortunately, finding the exact MMSE solution to this problem is infeasible as the exact distribution of the local gradient sub-vector is unknown at the PS. To circumvent this challenge, we employ the EM-GAMP algorithm in [29] based on the statistical model in Sec. IV-B. This algorithm iteratively computes an approximate MMSE solution to the CS recovery problem by modeling the distribution of an unknown sparse vector using the Bernoulli Gaussian-mixture model; the parameters of this model are also updated iteratively based on the EM principle. As discussed in Sec. IV-B, the local gradient sub-vectors are well modeled by the Bernoulli Gaussian-mixture distribution; thereby, the EM-GAMP algorithm is a powerful means of solving our CS recovery problem in (20). This algorithm also computes the posterior mean and variance of each compressed gradient entry as a byproduct; thereby, it can be directly combined with the turbo decoding principle in the proposed algorithm. In Algorithm 2, we summarize the EM-GAMP algorithm employed in Module B of the proposed algorithm, where we omit the indices $k$ and $b$ for ease of notation.

The major steps in Algorithm 2 are described below. In Step 3, the pseudo-prior mean and variance of $x_{m}$ are determined. In Step 4, the posterior mean and variance of $x_{m}$ are computed based on (20) and $x_{m}\sim\mathcal{N}(\hat{p}_{m},{\nu}_{p_{m}})$. In Step 6, the estimate of $g_{n}$ and the corresponding error variance are determined. In Step 7, the posterior mean and variance of $g_{n}$ are computed under the assumptions of $g_{n}\sim\mathcal{BG}(g;{\bm{\theta}})$ and $\hat{r}_{n}=g_{n}+\xi_{n}$ with $\xi_{n}\sim\mathcal{N}(0,{\nu}_{r_{n}})$ as derived in [29], where $\beta_{n,0}=\lambda_{0}\mathcal{N}(0;\hat{r}_{n},\nu_{r_{n}})$, $\beta_{n,\ell}=\lambda_{\ell}\mathcal{N}(\hat{r}_{n};\mu_{\ell},\nu_{r_{n}}+\phi_{\ell})$,

for all $\ell\in\{1,\ldots,L\}$. In Step 8, the parameters of the Bernoulli Gaussian-mixture model are computed based on the EM principle as derived in [29], where $\lambda_{0}^{\prime\prime}=\frac{1}{N}\sum_{n=1}^{N}\lambda_{n,0}^{\prime}$,

The asymptotic convergence of the EM-GAMP algorithm was studied in [39], in which the authors showed that under certain conditions, the performance of this algorithm asymptotically coincides with that of the original GAMP algorithm with perfect knowledge of the input distribution. We refer interested readers to [39] for more details. Meanwhile, the complexity order of the EM-GAMP algorithm is given by $\mathcal{O}(MN)$ in terms of the number of real multiplications (see Step 3 and Step 6 of Algorithm ).