Graph Neural Network-Enhanced Expectation Propagation Algorithm for MIMO
  Turbo Receivers
Deep neural networks (NNs) are considered a powerful tool for balancing the
performance and complexity of multiple-input multiple-output (MIMO) receivers
due to their accurate feature extraction, high parallelism, and excellent
inference ability. Graph NNs (GNNs) have recently demonstrated outstanding
capability in learning enhanced message passing rules and have shown success in
overcoming the drawback of inaccurate Gaussian approximation of expectation
propagation (EP)-based MIMO detectors. However, the application of the
GNN-enhanced EP detector to MIMO turbo receivers is underexplored and
non-trivial due to the requirement of extrinsic information for iterative
processing. This paper proposes a GNN-enhanced EP algorithm for MIMO turbo
receivers, which realizes the turbo principle of generating extrinsic
information from the MIMO detector through a specially designed training
procedure. Additionally, an edge pruning strategy is designed to eliminate
redundant connections in the original fully connected model of the GNN
utilizing the correlation information inherently from the EP algorithm. Edge
pruning reduces the computational cost dramatically and enables the network to
focus more attention on the weights that are vital for performance. Simulation
results and complexity analysis indicate that the proposed MIMO turbo receiver
outperforms the EP turbo approaches by over 1 dB at the bit error rate of
$10^{-5}$, exhibits performance equivalent to state-of-the-art receivers with
2.5 times shorter running time, and adapts to various scenarios.
Xingyu Zhou, Jing Zhang, Chao-Kai Wen, Shi Jin, Shuangfeng Han
# Graph Neural Network-Enhanced Expectation Propagation Algorithm for MIMO Turbo Receivers

###### Abstract

Deep neural networks (NNs) are considered a powerful tool for balancing the performance and complexity of multiple-input multiple-output (MIMO) receivers due to their accurate feature extraction, high parallelism, and excellent inference ability. Graph NNs (GNNs) have recently demonstrated outstanding capability in learning enhanced message passing rules and have shown success in overcoming the drawback of inaccurate Gaussian approximation of expectation propagation (EP)-based MIMO detectors. However, the application of the GNN-enhanced EP detector to MIMO turbo receivers is underexplored and non-trivial due to the requirement of extrinsic information for iterative processing. This paper proposes a GNN-enhanced EP algorithm for MIMO turbo receivers, which realizes the turbo principle of generating extrinsic information from the MIMO detector through a specially designed training procedure. Additionally, an edge pruning strategy is designed to eliminate redundant connections in the original fully connected model of the GNN utilizing the correlation information inherently from the EP algorithm. Edge pruning reduces the computational cost dramatically and enables the network to focus more attention on the weights that are vital for performance. Simulation results and complexity analysis indicate that the proposed MIMO turbo receiver outperforms the EP turbo approaches by over 1 dB at the bit error rate of \(10^{-5}\), exhibits performance equivalent to state-of-the-art receivers with 2.5 times shorter running time, and adapts to various scenarios.

 Expectation propagation, graph neural network, MIMO turbo receiver, extrinsic information.

## I Introduction

Multiple-input multiple-output (MIMO) has the potential to improve the link throughput by orders of magnitude and has become the key enabling technology for modern wireless communication systems that need to adapt to tremendous growth in transmission rate demand and network scale. To achieve the full benefits of MIMO technology, there is a strong demand for computationally efficient receiver designs, considering the increasing number of antennas used.

Several suboptimal linear detectors, such as the zero-forcing and linear minimum mean square error (LMMSE) algorithms, have been desirable among existing MIMO detectors [2] because of their low computational cost. However, compared with the maximum likelihood (ML) detector, substantial performance loss has restricted the application of linear detectors in future communication systems. By contrast, the powerful sphere decoder (SD) [3] promises performance equivalent to the optimal ML. However, it is constrained to MIMO systems with a limited number of antennas because of the exponential worst-case complexity [4].

Iterative detectors based on message passing (MP), specifically approximate MP (AMP) [5] and expectation propagation (EP) [6, 7, 8], have become promising strategies to approximate the ML detector with moderate complexity. AMP is favored in addressing large-scale MIMO detection problems because its complexity is only quadratic to the system size. However, AMP is Bayes-optimal only when the channel matrix follows an independent and identically distributed (i.i.d.) sub-Gaussian distribution [9], and it degrades significantly under realistic ill-conditioned channels. Orthogonal AMP (OAMP) [10] and vector AMP (VAMP) [11] are powerful strategies that have been developed to address the limitations of AMP. They have been proven to achieve the Bayes-optimal performance for the general unitarily-invariant channel matrices under the large system limit [12]. However, their performance tends to degrade in realistic finite-dimensional MIMO systems, which are the main focus of this paper.1

Footnote 1: In this paper, our focus is on moderately sized spatial-multiplexing MIMO systems, such as \(4\times 4\) and \(16\times 16\) configurations, rather than massive MIMO. This particular setup is commonly found in current wireless standards and has received significant attention and research efforts [2].

EP relaxes the constraint on the channel matrix and outperforms the AMP detector over a wide range of MIMO channels by factorizing the posterior belief with Gaussian distributions [13]. Furthermore, EP demonstrates superior performance in small- or medium-sized MIMO systems as compared to OAMP/VAMP. This is attributed to EP's utilization of element-wise variance instead of the scalar variance employed in OAMP/VAMP [14]. Iterative detection and decoding (IDD), or equivalently, MIMO turbo receiver, can be used to further improve detection accuracy, which is a common practice in current communication systems. Moreover, EP-based turbo receivers have been widely applied [15, 16, 17, 18]. However, EP detectors suffer a substantial gap with the ML performance because of the inaccuracy of Gaussian approximation in practical MIMO systems with high spatial correlation and strong interference.

Recently, deep learning (DL) has demonstrated its remarkable capability of overcoming conventional challenges in wireless communications [19]. In particular, deep neural networks
(NNs) offer the possibility to address existing gaps in iterative MIMO detectors and promise great performance and efficiency in receiver design [20, 21, 22, 23]. The authors of [20] developed a model-driven detection network (DetNet) by applying a projected gradient descent algorithm, and DetNet achieves the same accuracy as the AMP detector with enhanced robustness and lower complexity. However, DetNet requires substantial training data and performs poorly under high-order modulation or small-sized MIMO systems, thereby restricting its applications. Conventional MP-based detectors deteriorate severely in practical MIMO systems because of the failure of their prerequisites. An OAMP network (OAMPNet) was constructed in [21] to compensate for the performance loss utilizing DL. However, the OAMPNet experiences performance degradation in real-world channels with strong correlations [22]. More powerful NNs were introduced to enhance MP-based detectors by using highly parameterized models, including the MMNet in [22] and the recurrent equivariant MIMO detector in [23]. However, such schemes entail bulky detection networks with an excessive number of parameters to be trained. Furthermore, the EP detector was unfolded in [18, 24] to derive a detection network with a few trainable damping factors. In this way, fast convergence can be achieved, and inefficient hand-crafted tuning processes can be avoided.

Graph NNs (GNNs) provide an advanced technique for dealing with graph-structured data, and they have been widely applied to address the inference tasks of wireless communications [25]. Recently, GNN-based MIMO detection [26, 27] has also attracted great attention because of the GNN's ability to enhance the MP solution by incorporating DL [28]. The authors of [26] modeled the MIMO detection problem by the pair-wise Markov random field (MRF) and solved the corresponding maximum _a posteriori_ (MAP) inference task by learning an enhanced MP algorithm based on a GNN. Furthermore, the authors of [27] developed a GNN-enhanced EP detector called GEPNet, which introduced GNN to improve the posterior distribution approximation and exhibited a significant performance advantage over the EP and state-of-the-art NN-based detectors in uncoded systems.

However, the application of the GEPNet to turbo iterative processing is unexplored in [27], which is not a trivial issue due to the requirement for extrinsic information, rather than _a posteriori_ probability (APP), to ensure a stable convergence. Furthermore, the dense connections in the fully connected (FC) MRF model result in a GNN with a deal of redundancy, which hinders the efficient implementation of the detector.

In this paper, we develop an extrinsic GNN-aid EP network (EXT-GEPNet) for MIMO turbo receivers. We go beyond the design of GEPNet for uncoded systems [27] and follow the key idea of EP-based turbo receivers [15, 16, 17, 18] to construct the turbo structure for GEPNet with soft inputs and soft outputs, enabling sufficient utilization of _a priori_ information from the channel decoder to improve a _posteriori_ estimates. We observe that the original GEPNet [27] fails to generate reliable extrinsic information when applied to turbo iterative receiving, leading to poor performance. Hence, we customize a training scheme inspired by [29] to construct an EXT-GEPNet detector that satisfies the turbo principle [30] of forwarding extrinsic information. The training scheme addresses the limitations of the original GEPNet and establishes a fine-tuned EXT-GEPNet that can be integrated into the developed turbo structure to realize IDD with great flexibility and remarkable performance. We also reduce the computational complexity of the original GEPNet by designing an edge pruning scheme to simplify the GNN operations to a large extent while still maintaining excellent performance.

The contributions of this paper are summarized as follows:

* _Design of an EXT-GEPNet-based turbo receiver._ Through intensive simulation studies, we discovered that the original GEPNet cannot generate desired extrinsic outputs via the conventional strategy used in the MAP detector. To address this issue, we design an open-loop training scheme to derive an EXT-GEPNet that produces the extrinsic outputs required by the turbo procedure. This scheme is formulated based on a preliminary NN and the requirement for extrinsic information [30], i.e., not coupling with the priors, to obtain target extrinsic log-likelihood ratio (LLR) samples as labels for training the final EXT-GEPNet. The fine-tuned network can be plugged into the developed turbo structure to construct the EXT-GEPNet-based turbo receiver, which effectively overcomes the correlation problems when using the original GEPNet by preventing the same information from being counted twice. The proposed scheme is more flexible than directly training through the IDD procedure, as it does not rely on the choice of channel codes.
* _Edge pruning to reduce the complexity._ To relieve the complexity of the MP process, we remove redundant edges in the latent FC graph of the GNN. Unlike existing NN pruning schemes that simply drop the network's weights with small magnitudes [28, 31], we perform edge pruning utilizing the correlation information among the to-be-estimated variables inherent in the EP iterations, inspired by [32]. The proposed network can complete the detection with lower computational cost and even achieve performance gain after pruning due to the reduction of ineffective connections and avoidance of overfitting.
* _Comprehensive performance and complexity evaluation._ We validate our scheme with a wide range of numerical simulations under different scenarios that benchmark against a series of baselines. We also analyze the computational complexity of the proposed turbo receiver. Simulation results show that the proposed receiver has significant gain over the turbo approaches supported by EP and the original GEPNet and achieves comparable or even better performance with substantially lower running time than the single tree-search SD (STS-SD)-based receiver [3]. Furthermore, the proposed receiver adapts well to different channels and channel codes and exhibits robustness to channel estimation errors.

_Notations:_ Boldface letters denote column vectors or matrices. \(\mathbf{A}^{T}\) and \(\mathbf{A}^{\dagger}=(\mathbf{A}^{T}\mathbf{A})^{-1}\mathbf{A}^{T}\) represent the transpose and pseudo-inverse of matrix \(\mathbf{A}\), respectively. \(\mathbf{I}_{N}\) is an identity matrix of size \(N\), and \(\mathbf{0}\) is a zero matrix. \(\delta(\cdot)\), \(\mathbb{E}[\cdot]\), and \(\|\cdot\|\) denote the Dirac delta function, expectation operation,
and Euclidean norm, respectively. The set \([K]=1,2,\ldots,K\) contains all nonnegative integers up to \(K\). Finally, \(\mathcal{N}(z;\mu,\sigma^{2})\) denotes real-valued Gaussian random variables with mean \(\mu\) and variance \(\sigma^{2}\).

## II System Model and Algorithm Review

In this section, the system model of the IDD problem is formulated first. Then, the EP-based turbo receiver is reviewed to obtain a clear understanding of the proposed scheme.

### _System Model_

The considered MIMO system on the basis of bit-interleaved coded modulation (BICM) consists of \(N_{\rm t}\) transmit (Tx) antennas and \(N_{\rm r}\) receive (Rx) antennas, with \(N_{\rm r}\geq N_{\rm t}\). Fig. 1 depicts the system with a block diagram, which includes a MIMO transmitter and a MIMO turbo receiver. At the transmitter, the channel encoder first converts the binary word, \({\bf b}\in\{0,1\}^{N_{\rm b}}\) with \(N_{\rm b}\) as the number of message bits in a word, into the coded bits with code rate \(R_{\rm c}=N_{\rm b}/N_{\rm c}\). Then, interleaving is performed to the coded bits, thereby yielding the codeword \({\bf c}\) of length \(N_{\rm c}\). The codeword \({\bf c}\) is then partitioned into \(N_{\rm s}\) subvectors of length \(N_{\rm t}\tilde{Q}\) and modulated into symbol vectors with a complex quadrature amplitude modulation (QAM) constellation \(\bar{\mathcal{A}}\) of size \(|\bar{\mathcal{A}}|=\tilde{M}\), where \(\tilde{Q}\) is the number of bits per complex symbol and \(\tilde{Q}=\log_{2}\tilde{M}\). The symbol vectors are transmitted over the wireless channel, which is assumed to be unchanged in a time slot, and the received real-valued signal can be represented as

\[{\bf y}={\bf H}{\bf x}+{\bf w}, \tag{1}\]

where \({\bf x}\in\mathcal{A}^{K}\) with \(K=2N_{\rm t}\) is the equivalent real-valued transmitted vector in a time slot. The real-valued constellation \(\mathcal{A}\) has a cardinality of \(|\mathcal{A}|\triangleq M=\sqrt{\tilde{M}}\) and average energy of \(E_{\rm s}\). The channel matrix \({\bf H}\in\mathbb{R}^{N\times K}\), with \(N=2N_{\rm r}\) and its columns \({\bf h}_{k},k\in[K]\) normalized to unit energy, is supposed to be known at the receiver without special illustrations. \({\bf w}\) is the additive white Gaussian noise vector with zero mean and element-wise noise variance \(\sigma_{w}^{2}\).

The posterior probability density function (PDF) of the transmitted symbol vector \({\bf x}\) given the observations \({\bf y}\) yields

\[p({\bf x}|{\bf y})=\frac{p({\bf y}|{\bf x})p({\bf x})}{p({\bf y})}\propto \underbrace{\mathcal{N}\left({\bf y};{\bf H}{\bf x},\sigma_{w}^{2}{\bf I}_{N} \right)}_{p({\bf y}|{\bf x})}\underbrace{\prod_{k=1}^{K}p_{\rm A1}(x_{k})}_{p( {\bf x})}, \tag{2}\]

where \(p_{\rm A1}(x_{k})\) is the _a priori_ PDF of \(x_{k}\). In the first turbo iteration (TI), the prior is initialized as \(p_{\rm A1}^{(1)}(x_{k})=\frac{1}{M}\sum_{x\in\mathcal{A}}\delta(x_{k}-x)\) assuming equiprobable transmitted symbols. The direct calculation of \(p({\bf x}|{\bf y})\) involves a high-dimensional integral, which is generally intractable. Therefore, Bayesian inference techniques (e.g., EP) are commonly used to compute an approximation \(q({\bf x})\) for the optimal solution. The IDD technique, which is referred to as the turbo receiver in this paper, can improve the accuracy of the approximation further. As illustrated in Fig. 1, the soft-input soft-output signal detector and channel decoder in the turbo structure iteratively exchange reliability information on the same set of coded bits in the form of LLRs. These LLRs correspond to extrinsic probabilities to ensure the convergence and stability of the receiver. Moreover, the improved prior information can be utilized when the turbo procedure begins, that is, \(p_{\rm A1}^{(\iota)}(x_{k})\) can be constructed on the basis of the feedback from the channel decoder instead of the equiprobable assumption, where \(\iota\) is the index of the TI. The iterative process proceeds for a maximum number of \(I\) iterations and finally outputs the estimated message bits \(\hat{\bf b}\).

The extrinsic LLR for \(c_{k,i}\), which is the \(i\)-th bit mapped to symbol \(x_{k}\), can be computed on the basis of the extrinsic PDF \(p_{\rm E1}^{(\iota)}(x_{k}|{\bf y})\) estimated by the detector in each TI as

\[L_{\rm E1}^{(\iota)}(c_{k,i})\triangleq\log\frac{\sum_{x_{k}\in\mathcal{A}_{k,i}^{(1)}}p_{\rm E1}^{(\iota)}(x_{k}|{\bf y})}{\sum_{x_{k}\in\mathcal{A}_{k,i} ^{(0)}}p_{\rm E1}^{(\iota)}(x_{k}|{\bf y})}, \tag{3}\]

where \(\mathcal{A}_{k,i}^{(1)}\) and \(\mathcal{A}_{k,i}^{(0)}\) denote the subsets of constellation \(\mathcal{A}\), in which \(c_{k,i}\) is equal to 1 and 0, respectively. The vector \({\bf L}_{\rm E1}^{(\iota)}\), which contains all extrinsic LLR values from the detector, is further de-interleaved to derive \({\bf L}_{\rm A2}^{(\iota)}\), which is delivered to the channel decoder as the _a priori_ LLRs. The decoder computes the extrinsic LLRs on the coded bits, indicated as \({\bf L}_{\rm E2}^{(\iota)}\), utilizing the _a priori_ LLRs \({\bf L}_{\rm A2}^{(\iota)}\). For the \((\iota+1)\)-th TI, the extrinsic LLRs generated by the decoder are interleaved to derive \({\bf L}_{\rm A1}^{(\iota+1)}\), sent back to the detector, and mapped again to the updated _a priori_ PDF:

\[p_{\rm A1}^{(\iota+1)}(x_{k})=\prod_{i=1}^{Q}\frac{\exp\left(c_{k,i}L_{\rm A1}^ {(\iota+1)}(c_{k,i})\right)}{1+\exp\left(L_{\rm A1}^{(\iota+1)}(c_{k,i}) \right)}, \tag{4}\]

where \(Q=\tilde{Q}/2\) denotes the number of bits in a real-valued symbol.

### _EP-based Turbo Receiver_

EP is a Bayesian inference method that approximates the desired distribution by a function within the exponential families [33].2 Specifically, the EP-based MIMO detector computes a Gaussian approximation \(q({\bf x})\) for the posterior belief \(p({\bf x}|{\bf y})\) in (2), which is achieved by replacing the non-Gaussian

Fig. 1: Block diagram of a MIMO system based on BICM [3]. The MIMO turbo receiver iteratively exchanges soft information between the MIMO detector, which combines model-based algorithms with NNs, and the channel decoder.
factors (discrete priors) in (2) with unnormalized Gaussians iteratively:

\[q^{(\iota,t)}(\mathbf{x})\propto \mathcal{N}(\mathbf{y};\mathbf{H}\mathbf{x},\sigma_{w}^{2}\mathbf{I }_{N})\cdot\prod_{k=1}^{K}\exp\big{(}\gamma_{k}^{(\iota,t-1)}x_{k}-\frac{1}{2} \lambda_{k}^{(\iota,t-1)}x_{k}^{2}\big{)}\] \[\propto \mathcal{N}\big{(}\mathbf{x};\mathbf{H}^{\dagger}\mathbf{y}, \sigma_{w}^{2}(\mathbf{H}^{T}\mathbf{H})^{-1}\big{)}\] \[\cdot\mathcal{N}\big{(}\mathbf{x};(\boldsymbol{\lambda}^{(\iota,t -1)})^{-1}\boldsymbol{\gamma}^{(\iota,t-1)},(\boldsymbol{\lambda}^{(\iota,t-1) })^{-1}\big{)}\] \[\propto \mathcal{N}\big{(}\mathbf{x};\boldsymbol{\mu}^{(\iota,t)}, \boldsymbol{\Sigma}^{(\iota,t)}\big{)}, \tag{5}\]

where the superscript \((\iota,t)\) denotes the \(t\)-th EP iteration within the \(\iota\)-th TI. \(\gamma_{k}^{(\iota,t)}\in\mathbb{R}\) and \(\lambda_{k}^{(\iota,t)}\in\mathbb{R}^{+}\) denote the natural parameters of the approximating function, constituting the natural mean vector \(\boldsymbol{\gamma}^{(\iota,t)}=[\gamma_{1}^{(\iota,t)},\dots,\gamma_{K}^{( \iota,t)}]^{T}\) and precision matrix \(\boldsymbol{\lambda}^{(\iota,t)}=\mathrm{diag}([\lambda_{1}^{(\iota,t)}, \dots,\lambda_{K}^{(\iota,t)}])\)[33]. In the first TI, parameters \(\gamma_{k}^{(\iota,t)}\) and \(\lambda_{k}^{(\iota,t)}\) are initialized as \(\gamma_{k}^{(1,0)}=0\) and \(\lambda_{k}^{(1,0)}=1/E_{\mathrm{s}},k\in[K]\), respectively, and then iteratively updated following the moment matching condition as \(\mathbb{E}_{q(\mathbf{x})}=\mathbb{E}_{p(\mathbf{x}|\mathbf{y})}\)[7]. The mean \(\boldsymbol{\mu}^{(\iota,t)}\) and covariance \(\boldsymbol{\Sigma}^{(\iota,t)}\) of \(q^{(\iota,t)}(\mathbf{x})\) are computed using the Gaussian product lemma [35] as follows:

\[\boldsymbol{\Sigma}^{(\iota,t)} =\big{(}\sigma_{w}^{-2}\mathbf{H}^{T}\mathbf{H}+\boldsymbol{ \lambda}^{(\iota,t-1)}\big{)}^{-1}, \tag{6a}\] \[\boldsymbol{\mu}^{(\iota,t)} =\boldsymbol{\Sigma}^{(\iota,t)}\big{(}\sigma_{w}^{-2}\mathbf{H} ^{T}\mathbf{y}+\boldsymbol{\gamma}^{(\iota,t-1)}\big{)}. \tag{6b}\]

Moreover, EP calculates the marginal distribution of \(q^{(\iota,t)}(\mathbf{x})\) by viewing the covariance \(\boldsymbol{\Sigma}^{(\iota,t)}\) as a diagonal matrix to reduce complexity, that is, \(q^{(\iota,t)}(x_{k})\propto\mathcal{N}(x_{k};\mu_{k}^{(\iota,t)},\Sigma_{k}^{ (\iota,t)}),k\in[K]\), where \(\mu_{k}^{(\iota,t)}\) is the \(k\)-th element of \(\boldsymbol{\mu}^{(\iota,t)}\), and \(\Sigma_{k}^{(\iota,t)}\) is the \(k\)-th element of the main diagonal in \(\boldsymbol{\Sigma}^{(\iota,t)}\). This strategy, which uses the product of independent Gaussian functions to approximate \(p(\mathbf{x}|\mathbf{y})\)[27], ignores the off-diagonal elements in the covariance \(\boldsymbol{\Sigma}^{(\iota,t)}\) and results in information loss.

Fig. 2 presents the block diagram of the EP detector at the \(\iota\)-th TI, which contains a LMMSE module and a nonlinear Posterior module and bears a turbo structure. The LMMSE module is dedicated to the computation of \(\boldsymbol{\mu}^{(\iota,t)}\) and \(\boldsymbol{\Sigma}^{(\iota,t)}\) in (6). Then, the extrinsic marginal distribution is derived by the "ext" operation after LMMSE to decorrelate the output and the input as follows:

\[q_{e}^{(\iota,t)}\left(x_{k}\right) =\frac{q^{(\iota,t)}\left(x_{k}\right)}{\exp\big{(}\gamma_{k}^{( \iota,t-1)}x_{k}-\frac{1}{2}\lambda_{k}^{(\iota,t-1)}x_{k}^{2}\big{)}}\] \[\propto \tag{7}\]

where \(x_{e,k}^{(\iota,t)}\) and \(v_{e,k}^{(\iota,t)}\) are the \(k\)-th element of the mean vector \(\mathbf{x}_{e}^{(\iota,t)}\) and the main diagonal in the covariance matrix \(\mathbf{V}_{e}^{(\iota,t)}\), respectively, and we yield

\[v_{e,k}^{(\iota,t)} =\frac{\Sigma_{k}^{(\iota,t)}}{1-\Sigma_{k}^{(\iota,t)}\lambda_{k }^{(\iota,t-1)}}, \tag{8a}\] \[x_{e,k}^{(\iota,t)} =v_{e,k}^{(\iota,t)}\left(\frac{\mu_{k}^{(\iota,t)}-\gamma_{k}^{( \iota,t-1)}}{\Sigma_{k}^{(\iota,t)}-\gamma_{k}^{(\iota,t-1)}}\right). \tag{8b}\]

Subsequently, the mean vector \(\mathbf{x}_{e}^{(\iota,t)}\) and diagonal covariance matrix \(\mathbf{V}_{e}^{(\iota,t)}\) of the extrinsic distribution are delivered to the Posterior module for further processing.

The Posterior module combines the _a priori_ PDF \(p_{\mathrm{A}1}^{(\iota)}(x_{k})\) with the extrinsic distribution to derive the estimated _a posteriori_ distribution as [15, 16, 17, 18]:

\[\hat{p}^{(\iota,t)}(x_{k})\propto q_{e}^{(\iota,t)}(x_{k})p_{\mathrm{A}1}^{( \iota)}(x_{k}), \tag{9}\]

where the _a priori_ PDF \(p_{\mathrm{A}1}^{(\iota)}(x_{k})\) is uniform for \(\iota=1\) and non-uniform for \(\iota\geq 2\) mapped by the _a priori_ LLRs \(\mathbf{L}_{\mathrm{A}1}^{(\iota)}\) from the channel decoder according to (4). The _a posteriori_ mean \(\hat{x}_{k}^{(\iota,t)}\) and variance \(v_{k}^{(\iota,t)}\) are computed as:

\[\hat{x}_{k}^{(\iota,t)} =\sum_{a_{m}\in\mathcal{A}}a_{m}\hat{p}^{(\iota,t)}(x_{k}=a_{m}), \tag{10a}\] \[v_{k}^{(\iota,t)} =\sum_{a_{m}\in\mathcal{A}}\big{(}x_{k}-\hat{x}_{k}^{(\iota,t)} \big{)}^{2}\hat{p}^{(\iota,t)}(x_{k}=a_{m}). \tag{10b}\]

\(T\) is denoted as the number of EP iterations, and the pair \((\boldsymbol{\gamma}^{(\iota,t)},\boldsymbol{\lambda}^{(\iota,t)})\) is updated when \(t<T\) so that

\[\prod_{k=1}^{K}\exp\big{(}\gamma_{k}^{(\iota,t)}x_{k}-\frac{1}{2} \lambda_{k}^{(\iota,t)}x_{k}^{2}\big{)}\] \[\propto \mathcal{N}\big{(}\mathbf{x};(\boldsymbol{\lambda}^{(\iota,t)})^{ -1}\boldsymbol{\gamma}^{(\iota,t)},(\boldsymbol{\lambda}^{(\iota,t)})^{-1}\big{)} \propto\frac{\mathcal{N}\big{(}\mathbf{x};\hat{\mathbf{x}}^{(\iota,t)},\, \mathbf{V}^{(\iota,t)}\big{)}}{\mathcal{N}\big{(}\mathbf{x};\mathbf{x}_{e}^{( \iota,t)},\mathbf{V}_{e}^{(\iota,t)}\big{)}}, \tag{11}\]

where \(\hat{\mathbf{x}}^{(\iota,t)}=[\hat{x}_{1}^{(\iota,t)},\dots,\hat{x}_{K}^{( \iota,t)}]^{T}\) and \(\mathbf{V}^{(\iota,t)}=\mathrm{diag}([v_{1}^{(\iota,t)},\dots,v_{K}^{( \iota,t)}])\). A solution to (11) is:

\[\boldsymbol{\lambda}^{(\iota,t)} =\left(\mathbf{V}^{(\iota,t)}\right)^{-1}-\left(\mathbf{V}_{e}^{( \iota,t)}\right)^{-1}, \tag{12a}\] \[\boldsymbol{\gamma}^{(\iota,t)} =\left(\mathbf{V}^{(\iota,t)}\right)^{-1}\hat{\mathbf{x}}^{( \iota,t)}-\left(\mathbf{V}_{e}^{(\iota,t)}\right)^{-1}\mathbf{x}_{e}^{( \iota,t)}, \tag{12b}\]

which is implemented by the "ext" operation after the Posterior module. Notably, the update in (12a) can result in a negative \(\lambda_{k}^{(\iota,t)}\), which is unreasonable and should be discarded because \(\lambda^{(\iota,t)}\) is an inverse variance term. Therefore, we adopt the approach from [7] to ensure numerical stability, where we retain \(\lambda_{k}^{(\iota,t)}=\lambda_{k}^{(\iota,t-1)}\) and \(\gamma_{k}^{(\iota,t)}=\gamma_{k}^{(\iota,t-1)}\) when \(\lambda_{k}^{(\iota,t)}<0\). Additionally, we apply a damping technique [7, 33] to smooth the update by using a convex combination of the
properties. The updated pair \((\mathbf{\gamma}^{(\iota,t)},\mathbf{\lambda}^{(\iota,t)})\) is delivered to the LMMSE module for the next EP iteration.

The EP algorithm finishes when the maximum number of iterations \(T\) is reached. Extrinsic LLRs \(\mathbf{L}_{\mathrm{EI}}^{(\iota)}\) are demapped from the extrinsic distribution \(q_{\mathrm{e}}^{(\iota,T)}(x_{k})\) via (3) at the final iteration and delivered to the channel decoder, which outputs the estimated bits \(\hat{\mathbf{b}}\) when \(\iota=I\) or new _a priori_ LLRs \(\mathbf{L}_{\mathrm{A1}}^{(\iota+1)}\) when \(\iota<I\) for subsequent TIs. These _a priori_ LLRs are mapped to the updated _a priori_ PDF \(p_{\mathrm{A1}}^{(\iota+1)}(x_{k})\). The detector then computes the mean and variance of \(p_{\mathrm{A1}}^{(\iota+1)}(x_{k})\) as:

\[\hat{x}_{\mathrm{A1},k}^{(\iota+1)} =\sum_{a_{m}\in\mathcal{A}}a_{m}p_{\mathrm{A1}}^{(\iota+1)}(x_{k }=a_{m}), \tag{14a}\] \[v_{\mathrm{A1},k}^{(\iota+1)} =\sum_{a_{m}\in\mathcal{A}}\left(x_{k}-\hat{x}_{\mathrm{A1},k}^{( \iota+1)}\right)^{2}p_{\mathrm{A1}}^{(\iota+1)}(x_{k}=a_{m}), \tag{14b}\]

and updates the initial pair \((\mathbf{\gamma}^{(\iota+1,0)},\mathbf{\lambda}^{(\iota+1,0)})\) for EP at the \((\iota+1)\)-th TI as:

\[\mathbf{\lambda}^{(\iota+1,0)}\leftarrow(\mathbf{V}_{\mathrm{A1}}^{( \iota+1)})^{-1},\quad\mathbf{\gamma}^{(\iota+1,0)}\leftarrow\mathbf{\lambda}^{(\iota+1,0)}\hat{\mathbf{x}}_{\mathrm{A1}}^{(\iota+1)}, \tag{15}\]

with \(\hat{\mathbf{x}}_{\mathrm{A}}^{(\iota+1)}=[\hat{x}_{\mathrm{A1},1}^{(\iota+1) },\dots,\hat{x}_{\mathrm{A1},K}^{(\iota+1)}]^{T}\) and \(\mathbf{V}_{\mathrm{A1}}^{(\iota+1)}=\mathrm{diag}([v_{\mathrm{A1},1}^{(\iota +1)},\dots,v_{\mathrm{A1},K}^{(\iota+1)}])\).

## III EXT-GEPNet for Turbo Receiver

This section presents the details of the proposed EXT-GEPNet. We first introduce the GEPNet detector [27] to improve the posterior distribution approximation of EP. Then, we present the designed turbo structure for GEPNet and the customized training scheme used to derive EXT-GEPNet. Finally, we introduce the edge pruning method to simplify the GNN operations.

### _GNN-Enhanced EP Detector_

EP factorizes the target APPs with Gaussian distributions to avoid intractable calculations. However, the Gaussian approximation of the posterior belief becomes inaccurate in practical scenarios. For example, in a MIMO system with the number of Tx antennas close to that of the Rx antennas, where strong interference is extremely detrimental, the residual noise that can be viewed as Gaussian sharply decreases. This situation leads to a significant performance gap between the EP detector and the ML detector. To solve this limitation of EP, the GEPNet detector [27] was proposed, as shown in Fig. 3 with black solid lines. The GEPNet structure is obtained by first unfolding the EP iterations into layers and then adding a GNN module to each layer while retaining the LMMSE and Posterior modules.3 The GNN module is inserted between the two original EP's modules and is used to provide an improved estimated distribution of the transmitted symbols over the constellation \(\mathcal{A}\) based on the prior knowledge \(q_{\mathrm{e}}^{(\iota)}(x_{k})\) from (7).4 Subsequently, we elaborate the GNN module [26].

Footnote 3: We denote an EP iteration as a layer in the rest of this paper for consistency.

Footnote 4: The GEPNet employs the same model parameters for different turbo iterations in the proposed turbo receiver. Therefore, in this section, we omit the superscript \(\iota\) of the turbo iteration index for simplicity.

The GNN module (Fig. 4) provides a framework used for capturing the structured dependency of the transmitted variables \(\mathbf{x}\) by combining DL into the MP on the pair-wise MRF model [26], where each variable is denoted as a node, and each pair of two nodes is linked by an edge. The nodes and edges in the MRF are represented by circles and squares in Fig. 4, respectively. The prior knowledge \(q_{\mathrm{e}}^{(\iota)}(x_{k})\), which is characterized by the mean \(x_{\mathrm{e},k}^{(\iota)}\) and variance \(v_{\mathrm{e},k}^{(\iota)}\), is incorporated into the node attribute as \(\mathbf{a}_{k}^{(\iota)}=[x_{\mathrm{e},k}^{(\iota)},v_{\mathrm{e},k}^{(\iota )}]\).

Furthermore, the nodes and edges of the MRF are both characterized by their feature vectors, which are denoted as \(\mathbf{u}_{k}^{(l)}\) and \(\mathbf{f}_{jk}\), respectively, where \(l\) is the iteration index of the MP and \(j\in[K]\backslash\{k\}\). The node feature vector \(\mathbf{u}_{k}^{(l)}\) of size \(N_{\mathrm{u}}\) encodes the probabilistic information about the variable \(x_{k}\) that corresponds to the self potential of the MRF, given by [26, Eq. (5)] as

\[\phi\left(x_{k}\right)=\exp\left(\frac{1}{\sigma_{w}^{2}}\mathbf{\mathrm{y}}^{T} \mathbf{h}_{k}x_{k}-\frac{1}{2\sigma_{w}^{2}}\mathbf{h}_{k}^{T}\mathbf{h}_{k}x _{k}^{2}\right)p_{\mathrm{A1}}\left(x_{k}\right).\]

In particular, \(\mathbf{u}_{k}^{(l)}\) is initialized as \(\mathbf{u}_{k}^{(0)}=\mathbf{W}_{1}\cdot\left[\mathbf{y}^{T}\mathbf{h}_{k}, \mathbf{h}_{k}^{T}\mathbf{h}_{k},\sigma_{w}^{2}\right]^{T}+\mathbf{b}_{1}\) and iteratively updated in the MP process, where \(\mathbf{W}_{1}\in\mathbb{R}^{N_{\mathrm{u}}\times 3}\) is a learnable weight matrix, and \(\mathbf{b}_{1}\in\mathbb{R}^{N_{\mathrm{u}}}\) is a learnable bias vector. The edge feature vector \(\mathbf{f}_{jk}=[\mathbf{h}_{k}^{T}\mathbf{h}_{j},\sigma_{w}^{2}]\) contains the pair potential between nodes \(x_{j}\) and \(x_{k}\), given by [26, Eq. (6)] as

\[\psi\left(x_{k},x_{j}\right)=\exp\left(-\frac{1}{\sigma_{w}^{2}}\mathbf{h}_{k} ^{T}\mathbf{h}_{j}x_{k}x_{j}\right).\]

Three steps are involved in the MP process of the GNN: propagation, aggregation, and readout. The first two steps are implemented at each MP iteration \(l\), whereas the readout is conducted at the final iteration \(L\) to make the inference.

#### Iii-A1 Propagation

The edge between any pair of the variable nodes \(x_{k}\) and \(x_{j}\) first concatenates its own feature vector \(\mathbf{f}_{jk}\) with the incoming feature vectors \(\mathbf{u}_{k}^{(l-1)}\) and \(\mathbf{u}_{j}^{(l-1)}\). Then, the concatenated feature is delivered to a multi-layer perceptron (MLP) for message encoding, and the corresponding output message \(\mathbf{m}_{jk}^{(l)}\) is

\[\mathbf{m}_{jk}^{(l)}=\mathcal{M}(\mathbf{u}_{k}^{(l-1)},\mathbf{u}_{j}^{(l-1)}, \mathbf{f}_{jk}), \tag{16}\]

where \(\mathcal{M}\) denotes the operation of the MLP, which has two hidden layers of sizes \(N_{\mathrm{h1}}\) and \(N_{\mathrm{h2}}\) with the rectifier linear unit (ReLU) as the activation function and an output layer of size \(N_{\mathrm{u}}\). Finally, the message \(\mathbf{m}_{jk}^{(l)}\) is sent to the variable node \(x_{k}\), as shown in Fig. 4.

#### Iii-A2 Aggregation

The aggregation at node \(x_{k}\) is carried out by summing the incoming messages \(\mathbf{m}_{jk}^{(l)}\) from the connected nodes \(x_{j}\) and then concatenating the sum with the node attribute \(\mathbf{a}_{k}^{(l)}\), given by \(\mathbf{m}_{k}^{(l)}=\big{[}\sum_{j\in[K]\backslash\{k\}}\mathbf{m}_{jk}^{(l)}, \mathbf{a}_{k}^{(l)}\big{]}\). Subsequently, a gated recurrent unit (GRU) \(\mathcal{U}\)[36] is used to update the node feature vector \(\mathbf{u}_{k}^{(l)}\) by incorporating the concatenated message \(\mathbf{m}_{k}^{(l)}\):

\[\mathbf{g}_{k}^{(l)} =\mathcal{U}(\mathbf{g}_{k}^{(l-1)},\mathbf{m}_{k}^{(l)}), \tag{17a}\] \[\mathbf{u}_{k}^{(l)} =\mathbf{W}_{2}\cdot\mathbf{g}_{k}^{(l)}+\mathbf{b}_{2}, \tag{17b}\]
where \(\mathbf{g}_{k}^{(l)}\in\mathbb{R}^{N_{\mathrm{h1}}}\) and \(\mathbf{g}_{k}^{(l-1)}\in\mathbb{R}^{N_{\mathrm{h1}}}\) denote the current and previous hidden states, respectively. The single-layer NN with the learnable parameters \(\mathbf{W}_{2}\in\mathbb{R}^{N_{\mathrm{a}}\times N_{\mathrm{h1}}}\) and \(\mathbf{b}_{2}\in\mathbb{R}^{N_{\mathrm{a}}}\) in (17b) is used to derive the output feature, which is delivered to the neighboring propagation module for the next MP iteration.

#### Iii-A3 Readout

The node feature vectors \(\mathbf{u}_{k}^{(L)}\) are sent to the readout module, which contains a MLP followed by the SoftMax function, after \(L\) rounds of MP and yields

\[\mathbf{z}_{k} =\mathcal{R}(\mathbf{u}_{k}^{(L)}), \tag{18a}\] \[q_{G}^{(t)}(x_{k}=a_{m}) =\frac{\exp(z_{k,a_{m}})}{\sum_{a_{m}\in\mathcal{A}}\exp(z_{k,a_{ m}})},a_{m}\in\mathcal{A}, \tag{18b}\]

where \(\mathcal{R}\) is the MLP with two hidden layers of sizes \(N_{\mathrm{h1}}\) and \(N_{\mathrm{h2}}\) activated by ReLU and an output layer of size \(M\), which is the cardinality of the real-valued constellation \(\mathcal{A}\). The SoftMax process in (18b) maps the unnormalized output \(\mathbf{z}_{k}\) of \(\mathcal{R}\) into the estimated marginal distribution \(q_{G}^{(t)}(x_{k})\) of the discrete variable \(x_{k}\) for the \(t\)-th layer of GEPNet. Moreover, the final GRU hidden state and node feature vector in the current GEPNet layer assigned to the next layer as the starting point, that is, \(\mathbf{g}_{k}^{(0)}\leftarrow\mathbf{g}_{k}^{(L)}\) and \(\mathbf{u}_{k}^{(0)}\leftarrow\mathbf{u}_{k}^{(L)},k\in[K]\).

Notably, the weight parameters of \(\mathcal{M}\), \(\mathcal{U}\), and \(\mathcal{R}\) are shared across different nodes or edges. These parameters can be trained with supervised learning to improve EP's estimations characterized only by the mean and variance of a Gaussian function. This allows the GEPNet detector to calculate the _a posteriori_ distribution in (9) using the improved estimation \(q_{G}^{(t)}(x_{k})\) from the GNN instead of \(q_{e}^{(t)}(x_{k})\). Hard output \(\mathbf{\hat{x}}^{(T)}\) can then be derived based on the _a posteriori_ estimation using (10) in the last layer of the network [27]. However, designing the GEPNet-based turbo receiver is a non-trivial task due to the requirement for extrinsic information, rather than _a posteriori_ information.

### _Proposed EXT-GEPNet_

#### Iii-B1 Motivation

Fig. 3 presents the proposed turbo structure for GEPNet. Similar to the EP-based turbo receiver described in Sec. II-B, the GEPNet detector in the designed turbo structure integrates priors from the channel decoder. In Fig. 3, the blue dashed lines represent the priors from the channel decoder. These priors are used in both the computation of the initial pair, as shown in (15), and in the improved _a posteriori_ estimation at the Posterior module as:

\[\hat{p}_{G}^{(t)}(x_{k})\propto q_{G}^{(t)}(x_{k})p_{\mathrm{A1}}(x_{k}), \tag{19}\]

which combines the outputs of the GNN with the _a priori_ PDF. The estimated _a posteriori_ LLRs \(\hat{\mathbf{L}}_{\mathrm{APP}}\) can be demapped from the estimated APPs \(\hat{p}_{G}^{(T)}(x_{k})\) at the last layer of GEPNet.

As shown in Fig. 3, the stability of the IDD process depends on GEPNet's ability to provide the decoder with extrinsic LLRs \(\mathbf{L}_{\mathrm{E1}}\) that follow the turbo principle [30]. This means that the extrinsic LLRs should only contain new information and should not count the same information twice. Recall that the MAP detector produces extrinsic LLRs by subtracting the _a priori_ LLRs \(L_{\mathrm{A1}}(c_{j})\) from the _a posteriori_ LLRs:

\[L_{\mathrm{E1}}(c_{j})=\log\frac{p\left(c_{j}=1|\mathbf{y}\right)}{p\left(c_{ j}=0|\mathbf{y}\right)}-L_{\mathrm{A1}}(c_{j}). \tag{20}\]

However, equipping GEPNet with this strategy leads to poor performance, as shown in the simulation results. This phenomenon can be attributed to the fact that GEPNet's approximation to the MAP detector may not be accurate, and directly subtracting the _a priori_ LLRs may not completely eliminate the impact of the priors [29]. This results in unreliable LLRs that deviate from the desired extrinsic LLRs [37]. To address this issue, we customize a training scheme for GEPNet to enable the network to output LLRs that do not couple with the

Fig. 3: Visual representation of the proposed turbo receiver structure for GEPNet at the \(\iota\)-th turbo iteration. The black solid lines depict the unfolding structure of GEPNet, while the blue dash lines depict the interaction with the channel decoder (Dec).
priors. We also utilize a decoder LLR preprocessing strategy to further stabilize the proposed receiver. The details of the resultant turbo receiver scheme are revealed in the following subsection.

#### Iv-B2 Three-step Training of EXT-GEPNet

Fig. 5 presents the training procedure for the EXT-GEPNet, which is divided into three steps, which are elaborated as follows.

**Step 1: Train an APP-based GEPNet.** We first train a GEPNet detector to generate _a posteriori_ LLRs, indicated by APP-GEPNet in Fig. 5(a). In this step, a randomly generated pair \(\{\mathbf{x}^{[d]},\mathbf{y}^{[d]},\mathbf{H}^{[d]},\sigma_{e}^{[d]}\}\) and a bitwise _a priori_ LLR vector \(\mathbf{L}_{\mathrm{A1}}^{[d]}\) of size \(J=KQ\) form a training sample, where \(d\) is the sample index. The transmitted symbol vector \(\mathbf{x}^{[d]}\) is the label. The received signal \(\mathbf{y}^{[d]}\), the channel state information (CSI) \(\{\mathbf{H}^{[d]},\sigma_{w}^{[d]}\}\), and the LLR vector \(\mathbf{L}_{\mathrm{A1}}^{[d]}\) are the input features. The elements of \(\mathbf{L}_{\mathrm{A1}}^{[d]}\) are generated according to the bits \(\mathbf{c}^{[d]}\) demapped from \(\mathbf{x}^{[d]}\) and are assumed to be Gaussian distributed, \(L_{\mathrm{A1}}(c_{j}^{[d]})\sim\mathcal{N}(L;(2c_{j}^{[d]}-1)\mu_{A},2\mu_{A})\), where \(c_{j}^{[d]}\) is the corresponding \(j\)-th bit. The mean of the Gaussian distribution is \(\mu_{A}=J_{A}^{-1}(I_{A})\), where \(J_{A}(\mu)\triangleq 1-\mathbb{E}_{\mathcal{N}(L;\mu,2\mu)}[\log_{2}(1+ \mathrm{e}^{-L})]\) is a monotonically increasing function. The parameter \(I_{A}\in[0,1]\) denotes the average mutual information between \(\mathbf{L}_{\mathrm{A1}}^{[d]}\) and \(\mathbf{c}^{[d]}\) and characterizes the quality of the _a priori_ LLRs [38]. This idea is inspired by the simulation of the extrinsic information transfer chart for IDD [38] and allows for an open-loop training without relying on the channel coding scheme.

To supply the network with different _a priori_ LLR distributions in the training stage for the sake of generalization, \(I_{A}\) can be uniformly selected from [0,1]. However, to avoid the tedious numerical calculations from \(I_{A}\) to \(\mu_{A}\), we set a look-up table (LUT) between \(I_{A}\) and \(\mu_{A}\) in advance for a predefined set \(I_{A}\in\{0,0.33,0.67,0.78,0.89,0.94,0.99,1\}\)[18]. This set already reflects the increasing trend of the confidence level, i.e., the absolute value of LLRs from the channel decoder, as the number of TIs increases. Thus, in our implementation, \(I_{A}\) is randomly selected from the predefined set for each vector \(\mathbf{L}_{\mathrm{A1}}^{[d]}\) and converted into the corresponding \(\mu_{A}\) based on the LUT. Then, the _a priori_ LLR vector is derived according to the given bits in the training sample, the mean value \(\mu_{A}\), and the Gaussian distribution.

Moreover, we train the network with the cross-entropy (CE) loss function according to [27]:

\[\mathcal{L}_{1}=-\frac{1}{D}\sum_{d=1}^{D}\sum_{k=1}^{K}\sum_{a\in\mathcal{A}} \mathbb{I}_{x_{k}^{[d]}=a}\log\big{(}\hat{p}_{G}^{(T)}(x_{k}^{[d]}=a)\big{)}, \tag{21}\]

where \(D\) is the number of samples in a training batch, \(x_{k}^{[d]}\) is the \(k\)-th element of \(\mathbf{x}^{[d]}\), and \(\mathbb{I}_{x_{k}^{[d]}=a}\) is the indicator function that takes the value one if \(x_{k}^{[d]}=a\) and zero otherwise, hence representing the training label. By backpropagating the error over this loss, the estimated APPs \(\hat{p}_{G}^{(T)}(x_{k})\) from the APP-GEPNet can gradually approach the true APPs of the transmitted symbols. Finally, \(\hat{p}_{G}^{(T)}(x_{k})\) can be demapped as the estimated _a posteriori_ LLRs \(\hat{\mathbf{L}}_{\mathrm{APP}}\).

**Step 2: Generate extrinsic training LLRs from APP-GEPNet.** The extrinsic LLR in (20) can be rewritten as [30]:

\[L_{\mathrm{E1}}(c_{j}) =\log\frac{\sum_{\forall\in c_{j}=1}p(\mathbf{y}|\mathbf{c})p_{ \mathrm{A1}}(\mathbf{c})}{\sum_{\forall\in c_{j}=0}p(\mathbf{y}|\mathbf{c})p_{ \mathrm{A1}}(\mathbf{c})}-L_{\mathrm{A1}}(c_{j})\] \[=\log\frac{\sum_{\forall\in c_{j}=1}p(\mathbf{y}|\mathbf{c})\prod_ {\forall i:\neq j}p_{\mathrm{A1}}(c_{i})}{\sum_{\forall\in c_{j}=0}p(\mathbf{ y}|\mathbf{c})\prod_{\forall i:\neq j}p_{\mathrm{A1}}(c_{i})}. \tag{22}\]

This equation demonstrates that the extrinsic LLR \(L_{\mathrm{E1}}(c_{j})\) is a function of the channel information and the priors \(L_{\mathrm{A1}}(c_{i}),\forall i\neq j\) and should be independent of \(L_{\mathrm{A1}}(c_{j})\)[30]. Special inputs are designed for the APP-GEPNet obtained in Step 1 so that the network can generate extrinsic LLRs that satisfy this criterion. First, a total of \(J\) modified _a priori_ LLR vectors \(\{\tilde{\mathbf{L}}_{\mathrm{A1}}^{(j)}\}_{j=1}^{J}\) are derived by setting the \(j\)-th (\(j\in[J]\)) element of the original _a priori_ LLR vector \(\mathbf{L}_{\mathrm{A1}}\) to zero, where the sample index \(d\) is omitted for ease of notation. Thus, the \(i\)-th element of the \(j\)-th modified vector \(\tilde{\mathbf{L}}_{\mathrm{A1}}^{(j)}\) is defined as

\[\tilde{L}_{\mathrm{A1}}^{(j)}(c_{i})=\left\{\begin{array}{ll}L_{\mathrm{A1}} (c_{i})&\text{if }i\neq j\\ 0&\text{if }i=j\end{array}\right.,i\in[J]. \tag{23}\]

Next, the modified _a priori_ LLR vectors \(\{\tilde{\mathbf{L}}_{\mathrm{A1}}^{(j)}\}_{j=1}^{J}\), along with the same features \(\{\mathbf{y},\mathbf{H},\sigma_{w}\}\) corresponding to \(\mathbf{L}_{\mathrm{A1}}\) from Step 1, are used as the inputs for the APP-GEPNet to derive the extrinsic LLRs via \(J\) parallel inferences, as shown in Fig. 5(b). Specifically, the inputs are \(\tilde{\mathbf{L}}_{\mathrm{A1}}^{(j)}\) and \(\{\mathbf{y},\mathbf{H},\sigma_{w}\}\) when we target at the \(j\)-th modified vector. The LLR \(L_{\mathrm{E1}}^{(j)}(c_{j})\) from the corresponding output vector is hence independent of the initial prior \(L_{\mathrm{A1}}(c_{j})\) because \(\tilde{L}_{\mathrm{A1}}^{(j)}(c_{j})\) is assigned to zero. Therefore, \(L_{\mathrm{E1}}^{(j)}(c_{j})\) is exempted from coupling with the corresponding bit prior and satisfies the criterion for extrinsic LLR given in [30]. Finally, \(L_{\mathrm{E1}}^{(j)}(c_{j})\) is collected individually from each output vector according to Fig. 5(b) to form \(\mathbf{L}_{\mathrm{E1}}=\{L_{\mathrm{E1}}^{(j)}(c_{j})\}_{j=1}^{J}\). These LLRs are then used as the training labels for EXT-GEPNet in the next step.

**Step 3: Train the final EXT-GEPNet.** In this step, we train the final EXT-GEPNet to learn the mapping from the input features to the extrinsic LLRs derived in Step 2. Notably, the EXT-GEPNet shares the same structure with the APP-GEPNet trained in Step 1, given by Fig. 3, whereas the difference lies in two aspects: First, for EXT-GEPNet, we use the outputs of the GNN \(q_{G}^{(T)}(x_{k})\) instead of the estimated APPs \(\hat{p}_{G}^{(T)}(x_{k})\) to derive the output LLRs \(\tilde{\mathbf{L}}_{\mathrm{E1}}\). This avoids coupling with the priors. Second, the objective of the training and the resultant model weights are different. Specifically, the loss function for training the EXT-GEPNet is

\[\mathcal{L}_{2}=\frac{1}{D}\sum_{d=1}^{D}\sum_{j=1}^{J}l_{\mathrm{CE}}(c_{e,j}^ {[d]},\tilde{c}_{e,j}^{[d]}), \tag{24}\]

where \(c_{e,j}^{[d]}\) represents a soft bit mapped from the extrinsic LLR sample \(L_{\mathrm{E1}}^{[d]}(c_{j})\) of Step 2, and \(\tilde{c}_{e,j}^{[d]}\) denotes the counterpart equivalent of the output LLR \(\tilde{L}_{\mathrm{E1}}^{[d]}(c_{j})\) of the target EXT-GEPNet:

\[c_{e,j}^{[d]}=\frac{1}{1+\exp(-L_{\mathrm{E1}}^{[d]}(c_{j}))},\;\;\tilde{c}_{e,j }^{[d]}=\frac{1}{1+\exp(-\tilde{L}_{\mathrm{E1}}^{[d]}(c_{j}))}.\]
Furthermore, \(l_{\mathrm{CE}}\) in (24) is given by

\[l_{\mathrm{CE}}(c,\tilde{c})=-\big{(}c\log(\tilde{c})+(1-c)\log(1-\tilde{c}) \big{)}. \tag{25}\]

Therefore, the output LLRs \(\tilde{\mathbf{L}}_{\mathrm{E1}}\) of the EXT-GEPNet can approximate the desired output \(\mathbf{L}_{\mathrm{E1}}\) from Step 2 by error backpropagation over the loss \(\mathcal{L}_{2}\), as shown in Fig. 5(c).

The three-step training scheme of the EXT-GEPNet is summarized in Algorithm 1.5

Footnote 5: Note that the randomly generated data pairs \(\{\mathbf{y},\mathbf{H},\sigma_{w},\mathbf{L}_{\mathrm{A1}}\}\) in Step 1 are used throughout the three-step training scheme. No additional data pairs are generated specifically for training the EXT-GEPNet.

**Remark 1**: _Once the EXT-GEPNet is trained, it can be deployed in the turbo structure of Fig. 3 and used for different TIs with the same model parameters, constructing the EXT-GEPNet-based turbo receiver. As the proposed method is an approximation to the MAP receiver, overestimation of the reliability can happen during the IDD procedure. This overestimation can result in instability in the detection process as the mean value of the closed-loop LLRs from the channel decoder would increase over iterations. We observed experimentally that scaling the decoder LLRs into the range that matches the range of the _a priori_ LLRs \(\mathbf{L}_{\mathrm{A1}}\) used in the training stage effectively overcomes the instability issues. Therefore, we utilize an adaptive LLR scaling method [29] involving three steps: First, we examine the training LLRs \(\mathbf{L}_{\mathrm{A1}}\) to determine the range \([-r,r]\) that \(\mathbf{L}_{\mathrm{A1}}\) fall into with a probability of \(p_{r}\approx 1\). Second, we search for the maximum absolute value \(r_{\epsilon}\) of the decoder LLRs in the current codeword at the end of each TI. Finally, we scale the decoder LLRs by \(r_{\epsilon}/r\) if \(r_{\epsilon}>r\); otherwise, we keep the LLRs unchanged. As a result, the decoder LLRs are controlled in an appropriate range. In this work, we set \(p_{r}\) as 0.97 and find that the decoder LLR preprocessing method further stabilizes the IDD procedure._

### _Edge Pruning to Reduce Complexity_

The most computationally demanding operation of the GNN lies in the message propagation step through all the edges of the graph shown in Fig. 4. This step involves the execution of the edge MLPs, i.e., the function \(\mathcal{M}\) in (16), for \(K(K-1)\) times because the MRF defined in Section III-A is FC, i.e., with each pair of nodes connected by an edge. However, considerable redundant connections are observed in this FC graph [32].

We propose an edge pruning scheme based on the covariance matrix \(\mathbf{\Sigma}\) calculated by the LMMSE module of EP in (6a) to simplify the MP of the GNN while maintaining competitive performance.6\(\mathbf{\Sigma}\) is used as an indicator of the correlation weights between the neighboring nodes in the FC graph because the element \(\Sigma_{ij}\) (\(i,j\in[K]\)) reflects the covariance of variable nodes \(x_{i}\) and \(x_{j}\). \(\Sigma_{ij}\) can be normalized to derive the correlation coefficient as \(\rho_{ij}=\frac{\Sigma_{ij}}{\sqrt{\Sigma_{i}\Sigma_{jj}}}\in[0,1]\). A large correlation coefficient reflects the high structural dependency between the two connected variables. Hence, more attention should be paid to the edge between them. On the other hand, a small \(\rho_{ij}\) suggests the approximate independence of the two variables, where less information is required to exchange between them to recover the transmitted symbols. Hence, the proposed scheme prunes those edges with correlation coefficients \(\rho_{ij}\) that meet the following criterion:

Footnote 6: The edge pruning is implemented in each layer of the network, and thus the index \((\iota,t)\) for \(\mathbf{\Sigma}\) is omitted.

\[\rho_{ij}^{2}<\alpha\cdot\frac{1}{K-1}\sum_{k=1,k\neq j}^{K}{{\rho_{kj}}^{2}}, \tag{26}\]

Fig. 5: Overview of the three-step training process for the EXT-GEPNet.
where \(\alpha\) is a positive factor that controls the pruning threshold, and the edge pruning version of the GEPNet-based method is indicated by this factor hereafter. This strategy means that for a specific node \(x_{j}\), only the incoming edges with correlation coefficients \(\rho_{ij}\) larger than the average number of all connected edges (multiply by the pruning factor \(\alpha\)) are retained, and a large \(\alpha\) intuitively results in a significant proportion of removed edges. Therefore, edges with small correlation weights, i.e., low contributions to the inference of the target probability, are pruned. This process reduces the computational cost of GNN and saves computational resources for tuning the vital parts of the network. Notably, edges with low correlation weights exert minimal impact on the overall performance. Thus, the sparsely connected network after pruning can control the performance loss within the acceptable range or be even more generalizable, which is consistent with the findings in [31].

**Remark 2**: _Note that while the proposed edge pruning strategy shows potential in improving performance by mitigating overfitting, its underlying mechanism differs from conventional techniques aimed at alleviating overfitting, such as early stopping based on a separate validation dataset [39]. Specifically, the early stopping technique prevents overfitting by halting training when the validation loss begins to increase. In contrast, the proposed edge pruning focuses on reducing model complexity, thus enhancing generalization through the elimination of unnecessary connections. This approach makes the network more efficient during the inference stage, which cannot be achieved through early stopping alone._

## IV Simulation Results

The numerical results of the proposed schemes are presented in this section. First, the parameter settings are introduced. Second, the performance of the proposed schemes is evaluated under uncoded and coded MIMO systems. Finally, the computational complexity is analyzed.

### _Parameter Settings_

The simulated signal-to-noise ratio (SNR) of the system is defined as \(\text{SNR}=\frac{\mathbb{E}[|\mathbf{H_{\text{K}}}\mathbf{s}|^{2}]}{\mathbb{ E}[\|\mathbf{w}\|^{2}]}\). The i.i.d. Rayleigh and spatially correlated channels are used in the simulation. The Rayleigh MIMO channel \(\mathbf{H}\) has elements drawn from the Gaussian distribution \(\mathcal{N}(h_{ij};0,1/N)\), where \(h_{ij}\) is the \((i,j)\)-th element of \(\mathbf{H}\). The spatially correlated channel is characterized by the Kronecker model as \(\mathbf{H}=\mathbf{R}_{\text{r}}^{1/2}\mathbf{U}\mathbf{R}_{\text{t}}^{1/2}\), where \(\mathbf{U}\) is the i.i.d. Rayleigh channel matrix. \(\mathbf{R}_{\text{r}}\) and \(\mathbf{R}_{\text{t}}\), which have exponential elements that correspond to the same spatial correlation coefficient \(\varrho\)[40], represent the correlation matrices at the receiver and transmitter, respectively.

To evaluate the error rate performance, the maximum number of transmitted bits is set as \(5\times 10^{7}\). For uncoded systems, the symbol error rate (SER) is used as the performance metric, while for the turbo receiver, the coded bit error rate (BER) and word error rate (WER) with \(N_{\text{b}}\) as the word length are used. For the coded MIMO systems, convolutional codes (CCs) and turbo codes are selected as the channel coding scheme. The CCs have generator polynomial \([133_{\text{o}}~{}171_{\text{o}}]\). Random interleaving and two code rates are adopted: one is code rate \(R_{\text{c}}=1/2\) with word length \(N_{\text{b}}=128\), and the other is \(R_{\text{c}}=5/6\) with \(N_{\text{b}}=800\). The word length is chosen so that the CCs can achieve effective error correction. The turbo codes have a code rate of \(R_{\text{c}}=1/2\) and word length \(N_{\text{b}}=1952\), and the interleaver follows the 3rd Generation Partnership Project Release 17 specification [41]. The channel decoder utilizes the BCJR algorithm [42] with 10 inner iterations.

The hyperparameters of the GNN are set as \(N_{\text{h1}}=64,N_{\text{h2}}=32,N_{\text{u}}=8\), and \(L=2\). For the GEPNet detector in uncoded systems, the network is trained with the loss function \(\mathcal{L}_{1}\) in (21) but without the _a priori_ training LLRs. The training and validation sets contain 6,400,000 and 6,000 samples, respectively. Furthermore, during the three-step training scheme as described in Sec. III-B2, training and validation sets for APP-GEPNet in Step 1 have the same sizes as those for GEPNet in uncoded systems. Then, a total of 76,800 extrinsic LLR vectors \(\mathbf{L}_{\text{E1}}\) is generated in Step 2 to train EXT-GEPNet in Step 3. The objective is to minimize the loss function \(\mathcal{L}_{2}\) as described in (24). Although this dataset is relatively small, it significantly reduces the computational cost associated with generating the extrinsic training LLRs in Step 2. Through experiments, we have observed that this small dataset is sufficient for training the EXT-GEPNet to effectively learn the mapping from the input features to the extrinsic LLRs, resulting in excellent performance. All the considered networks are trained for 5,000 epochs with a batch size of 128.

In our experiment, we utilize the Glorot normal initializer [43] to initialize the weights of the network. Notably, the APP-GEPNet trained in Step 1 can also serve as an initialization model for the EXT-GEPNet in Step 3. This initialization strategy can help accelerate the convergence speed of the training procedure. The SNR during training is set as a specific point \(\text{SNR}_{\text{train}}\). Moreover, the optimizer is selected as Adam with a learning rate of 0.001. Unless otherwise specified, we train and test the network under the same modulation scheme, antenna configurations, and channel model. For uncoded MIMO systems, we choose the damping factor for EP and GEPNet-based detectors as \(\beta=0.2\), as suggested in [7]. For coded MIMO systems, we follow the configurations of EP and double EP (DEP), which introduces EP in both the estimation of the posterior and the processing of the channel decoder's feedback to accelerate convergence, as proposed in [16] and [17], respectively. The number of layers in all the evaluated iterative detectors is set as \(T=5\).

### _Performance Analysis of Uncoded MIMO Detectors_

We compare the performance of the uncoded MIMO detectors under i.i.d. Rayleigh channels. Fig. 6 provides the SER performance under 16-QAM modulation with 4 \(\times\) 4 and 16 \(\times\) 16 MIMO configurations. The GEPNet detector and the edge pruning versions are compared with the conventional EP [7], the model-driven DL-based OAMPNet [21], and the optimal ML detectors.

Fig. 6 reveals that the gain of GEPNet over EP and OAMPNet is remarkable. Moreover, the effect of edge pruning with
different pruning factors is demonstrated. In the comparison, the pruning factor \(\alpha\) is set as 0.5, 1, 2, and 4, respectively. The number after \(\alpha\) of each SER curve of the edge pruning versions in the figure corresponds to the percentage of remaining edges after pruning. The figure demonstrates that the edge pruning versions with \(\alpha=0.5\) and \(1\) outperform the original FC model, thereby revealing the gain brought by appropriately pruning the redundant edges with low correlation weights. However, the performance of the pruned GEPNet degrades significantly when \(\alpha\) further increases, falling behind the original FC model at \(\alpha=2\) and \(\alpha=4\) for the 4\(\times\)4 and 16\(\times\)16 systems, respectively. This result is because some dominant edges are improperly removed when \(\alpha\) is set too high. However, our scheme with \(\alpha=4\) (over 90% edges pruned) still outperforms EP by a large margin because the network is fine-tuned on the basis of EP, thereby indicating its outstanding ability in balancing performance and complexity.

### _Performance Analysis of Coded MIMO Turbo Receiver_

We consider a \(4\times 4\) coded MIMO system with 16-QAM modulation unless noted otherwise. First, we separately investigate the effect of the three-step training scheme and edge pruning on the proposed turbo receiver under CCs and Rayleigh MIMO channels. Second, we demonstrate the generalization ability of the proposed receiver by evaluating the performance under various mismatches, including channel coding scheme, channel model, SNR, and antenna configuration mismatches. Finally, we present the robustness of the proposed method against imperfect CSI.

#### Iv-C1 Impact of the Three-step Training Scheme

Fig. 7 presents the BER performance comparison under CCs and \(I=2\). Two GEPNet baselines are set to validate the proposed training scheme: The first is the APP-GEPNet, derived in Step 1 of the training scheme, and the second is the GEPNet detector trained in the uncoded systems and deployed in the proposed turbo scheme as shown in Fig. 3. This baseline has the same loss function and _a posteriori_ outputs as the APP-GEPNet but with the _a priori_ training LLRs set to zero, i.e., \(I_{A}=0\) for Step 1 of the training scheme, denoted as GEPNet-IA0 hereafter. Both of these baselines are equipped with (20) to subtract the _a priori_ information from the outputs \(\hat{\mathbf{L}}_{\mathrm{APP}}\), denoted by "w/ (20)" in the figures.7

Footnote 7: We also conducted tests on the baselines that utilize (3) to demap the output PDF of the GNN module as the extrinsic LLRs. However, we did not observe any performance improvement compared to using (20). Therefore, in order to maintain consistency with the proposed scheme in terms of removing priors and deriving extrinsic LLRs, we equipped the baselines with (20) in our simulations.

Fig. 7(a) shows the comparison under code rate \(R_{\mathrm{c}}=1/2\) and word length \(N_{\mathrm{b}}=128\). APP-GEPNet performs poorly,

Fig. 6: SER performance for 16-QAM under Rayleigh MIMO channels.

Fig. 7: BER performance (\(I=2\)) with CCs for a \(4\times 4\) Rayleigh MIMO channel with 16-QAM.
which reflects that simply using (20) cannot completely remove the correlation of the priors at the outputs to generate reliable extrinsic LLRs. GEPNet-IA0 outperforms APP-GEPNet, which can be attributed to the limited _a priori_ information coupled during training, and the correlation problem is less severe. However, this baseline cannot effectively manage the diverse _a priori_ information provided by the decoder as it is trained with \(I_{A}=0\).

EXT-GEPNet has significant performance gains over the two baselines because the proposed method not only resolves the information coupling problem systematically to generate appropriate extrinsic information, but also fully utilizes the _a priori_ information during training to generalize to different _a priori_ LLR distributions. Moreover, the EXT-GEPNet-based turbo receiver is also compared with other turbo approaches equipped with the MMSE-based parallel interference cancellation (MMSE-PIC) [44], EP [16], DEP [17], and STS-SD [3]. The figure shows that EXT-GEPNet outperforms MMSE-PIC, EP, and DEP by approximately 1 dB at the BER of \(10^{-5}\) and reveals equivalent or even superior performance to the computationally expensive STS-SD at all tested SNRs.

In another aspect, the results from Fig. 7(b) reveal that APP-GEPNet achieves a competitive performance under CCs with a high code rate (\(R_{\rm c}=5/6\)). This can be attributed to the low redundancy channel code resulting in less _a priori_ information being coupled in APP-GEPNet's outputs, and the correlation issue is less pronounced compared to that at a code rate of 1/2. However, EXT-GEPNet still has superiority over the two GEPNet-based baselines, outperforms MMSE-PIC, EP, and DEP, and is only inferior to the high-complexity STS-SD.

Fig. 8 provides an intuitive interpretation for the performance in Fig. 7 by analyzing the output LLR distributions of the detector during the training scheme. Fig. 8(a) shows the results under Gaussian _a priori_ LLRs with \(I_{A}=0.8\), emulating the LLRs from the decoder under CCs with \(R_{\rm c}=1/2\). Two steep peaks can be found in the output LLRs of APP-GEPNet even after the explicit subtraction of priors, which likely contain over-optimistic LLR estimations8 with residual priors that lead to the poor performance in Fig. 7(a).

Footnote 8: Over-optimistic LLR estimations refer to the phenomenon where the soft outputs of the detector overestimate the reliability, leading to incorrect LLR estimates with erroneously assigned large magnitudes [37, 45].

Step 2 of the training scheme, with output LLRs denoted by "EXT LLR label" in Fig. 8(a), effectively reduces the magnitudes of the two LLR peaks and alleviates the over-estimation issue. Moreover, the output LLRs of EXT-GEPNet closely track the LLR labels derived in Step 2, approaching desirable extrinsic LLRs. We also provide the output LLR distributions of GEPNet-IA0 with (20) as compared to the proposed scheme. The figure reveals that GEPNet-IA0 is less capable of distinguishing bits 0 and 1 than EXT-GEPNet, producing more LLRs with small magnitudes near 0.

Fig. 8(b) shows that the difference between the compared distributions is less perceivable when \(I_{A}=0.2\), i.e., limited priors are provided, as compared to \(I_{A}=0.8\), clarifying the competitive performance of APP-GEPNet under \(R_{\rm c}=5/6\).

Fig. 9 shows the BER convergence performance across TIs. We consider a \(4\times 4\) Rayleigh MIMO system with CCs (\(R_{\rm c}=1/2\), \(N_{\rm b}=128\)), 16-QAM, and \(\text{SNR}=12\) dB. The figure shows that EXT-GEPNet achieves the fastest convergence speed among the compared schemes. The comparison between EXT-GEPNet and GEPNet-IA0 also confirms that EXT-GEPNet better adapts to different TIs than the original GEPNet without the proposed training scheme.

#### Iv-A2 Impact of Edge Pruning

In the above analysis of the turbo receiver, all EXT-GEPNets use the FC model without edge pruning. In this subsection, we analyze the impact of

Fig. 8: Output LLR distributions of the detectors for a \(4\times 4\) Rayleigh MIMO channel with 16-QAM and \(\text{SNR}=13\) dB.

Fig. 9: BER performance across turbo iterations with CCs (\(R_{\rm c}=1/2\), \(N_{\rm b}=128\)) for a \(4\times 4\) Rayleigh MIMO channel with 16-QAM and \(\text{SNR}=12\) dB.
the edge pruning method on the balance of performance and complexity. Table I shows the WER performance of EXT-GEPNet-based turbo receivers with different edge pruning factors \(\alpha\) at \(I=2\) and representative SNRs, i.e., \(\text{SNR}=13\) dB for \(R_{\text{c}}=1/2\) and \(\text{SNR}=18\) dB for \(R_{\text{c}}=5/6\).

We first focus on the results with \(R_{\text{c}}=1/2\) and \(\text{SNR}=13\) dB. Similar to the results in uncoded systems, the edge pruning versions with \(\alpha=0.5\) and \(1\) do not suffer from performance loss and can even outperform the network without pruning (\(\alpha=0\)). This phenomenon is because the pruning operations can remove the redundant connections of the FC graph and results in a generalizable model less troubled by overfitting. Moreover, the performance loss in uncoded BER caused by a large proportion of edges being removed when \(\alpha\) is high (e.g., \(\alpha=2\) and 81.5% of edges are removed) can be compensated by the strong error correction code.

By contrast, the third column of Table I shows a clear trade-off between performance and complexity (the proportion of edges reduced) under CCs with a higher code rate, i.e., lower error correction capability. For example, the edge pruning versions with \(\alpha\) equal to 1, 2, and 4 all experience performance loss compared with the original FC model, with the WER increasing from 4.22e-3 at \(\alpha=0\) to 7.96e-3 at \(\alpha=4\). However, the computational complexity is significantly reduced since approximately 91.8% of edges are pruned, and the cost for the message encoding (16) on these edges is saved.

#### V-B3 Under Turbo Codes and Various Channels

In this subsection, we first demonstrate that the proposed scheme can be applied without dependence on the channel code. Fig. 10 provides the BER performance comparison under turbo codes with a code rate of \(R_{\text{c}}=1/2\) and a word length of \(N_{\text{b}}=1952\). We directly use the network with \(\text{SNR}_{\text{train}}=13\) dB in Fig. 7(a) for the comparison under turbo codes and Rayleigh MIMO channel in Fig. 10(a) without additional training. Fig. 10(a) reveals that EXT-GEPNet generalizes well to the turbo codes and outperforms the other methods by a large margin, thereby indicating the advantage of the open-loop training scheme in not relying on the choice of channel codes.

Furthermore, Fig. 10(b) provides the performance comparison under spatially correlated channels to illustrate the robustness of the proposed scheme against channel mismatch. In particular, the EXT-GEPNet trained with the Rayleigh channel from Fig. 7(a) is tested under the correlated channel with a spatial correlation coefficient \(\varrho=0.5\), and the results are marked by "mismatch" in Fig. 10(b). The figure demonstrates that the gap between the "mismatch" network and the network trained and tested both under the correlated channel can be neglected, thereby verifying the great robustness of the proposed scheme.

#### V-B4 Robustness to SNR and Antenna Configuration

Fig. 11 illustrates the BER performance of EXT-GEPNet under SNR and antenna configuration mismatches. We train an EXT-GEPNet in a \(4\times 4\) MIMO system (\(N=K=8\)) with \(\text{SNR}=13\) dB and evaluate the network's performance in a \(4\times 2\) MIMO system (\(N=8,K=4\)) with varying SNRs. The channel model is the spatially correlated channel with a correlation coefficient \(\varrho=0.5\), which remains consistent during the training and testing phases. The modulation type used is 16-QAM. Turbo codes with \(R_{\text{c}}=1/2\) and \(N_{\text{b}}=1952\) are applied for channel coding. The figure reveals that the EXT-GEPNet trained under the mismatched antenna (\(N=K=8\)) and SNR configurations exhibits the best BER performance among the compared schemes, except for the EXT-GEPNet trained under the matched antenna (\(N=8,K=4\)) and SNR configurations. Moreover, the performance gap between the matched and mismatched models is within 0.1 dB, indicating the robustness of the proposed method against antenna configuration and SNR mismatches.

#### V-B5 Robustness to Imperfect CSI

In the above investigation, all receivers were tested under classical channel models with accurate CSI. Next, we further validate the proposed method using the spatially correlated channel model from the fifth

Fig. 10: BER performance (\(I=2\)) with turbo codes (\(R_{\text{c}}=1/2,\ N_{\text{b}}=1952\)) for \(4\times 4\) MIMO channels with 16-QAM.
generation new radio (5G-NR) [46]. This provides a more realistic evaluation and assesses the algorithm's performance under imperfect CSI to verify its robustness against channel estimation errors. The considered system contains 2 Tx and 4 Rx antennas as uniform linear arrays. The channel correlation level is set to medium correlation A, as specified in the 5G-NR standard [46]. We utilize the LMMSE method [47] with an orthogonal pilot matrix \(\mathbf{X}_{\mathrm{p}}\in\mathbb{C}^{N_{\mathrm{t}}\times N_{\mathrm{p}}}\) composed by \(N_{\mathrm{t}}=2\) columns of the discrete Fourier transform matrix \(\mathbf{F}\in\mathbb{C}^{N_{\mathrm{p}}\times N_{\mathrm{p}}}\) to estimate the complex-valued channel matrix of size \(4\times 2\), where \(N_{\mathrm{p}}\) is the number of pilot vectors in a time slot and is set to 16. Turbo codes with \(R_{\mathrm{c}}=1/2\) and \(N_{\mathrm{b}}=1952\) and 16-QAM modulation are used. Fig. 12 illustrates the BER performance under perfect and estimated CSIs. The neural networks are trained using the 5G-NR channel model and \(\text{SNR}_{\text{train}}=\text{10 dB}\). As expected, all the tested algorithms experience a similar performance loss when transitioning from perfect to imperfect CSIs, indicating that the accuracy of CSI significantly influences the detection accuracy. However, EXT-GEPNet consistently outperforms the other tested algorithms, demonstrating its robustness to imperfect CSI. This observation holds true regardless of the presence of channel estimation errors, thereby highlighting the effectiveness and resilience of the proposed algorithm.

### _Computational Complexity Analysis_

In this subsection, we analyze the computational complexity of different turbo receivers, focusing on MIMO detection complexity since the channel decoding part for different receivers is the same. Table II provides the number of real-valued multiplications (RVMs) and running time required for detecting one symbol vector. The per symbol vector detection complexity is \(C_{\text{det1}}+(I-1)C_{\text{det}}\), where \(C_{\text{det1}}\) denotes the complexity of the first TI without utilizing _a priori_ information, and \(C_{\text{det}}\) is the counterpart with the priors considered.

The \(C_{\text{det1}}\) of EXT-GEPNet, denoted as \(C_{\text{GEP1}}\), can be divided into the operations for GNN and EP in each layer of the network. The complexity of the GNN operations contains the cost for the propagation, aggregation, and readout steps. The propagation step is involved in the \(L\)-round of MP, and each round costs \(\big{(}(2N_{\mathrm{u}}+2)N_{\mathrm{h1}}+N_{\mathrm{h1}}N_{\mathrm{h2}}+N_{ \mathrm{h2}}N_{\mathrm{u}}\big{)}K(K-1)\) RVMs for the execution of the function \(\mathcal{M}\) in (16) by \(K(K-1)\) times in the FC model. We further consider the effect of edge pruning in the complexity expression listed in the table, where the cost for this step is reduced by the scale \(\eta\), which is the proportion of the remaining edges after pruning. In addition, the aggregation step performs \(L\) times, and the readout step is conducted once for one forward inference of the GNN, resulting in a cost of \(\big{(}4N_{\mathrm{u}}+3N_{\mathrm{h1}}+9\big{)}N_{\mathrm{h1}}KL\) and \((N_{\mathrm{u}}N_{\mathrm{h1}}+N_{\mathrm{h1}}N_{\mathrm{h2}}+N_{\mathrm{h2}} M)K\) RVMs, respectively. The complexity of the EP procedure is dominated by the matrix inversion in (6a), which requires \(K^{3}+K^{2}+2K\) RVMs. Meanwhile, the matrix-vector multiplications involved in (8), (10), (12), and (13) cost another \(2MK+11K\) RVMs [27].

Additionally, the number of layers \(T\) should be considered for \(C_{\text{det1}}\) of EXT-GEPNet, EP, and DEP, as shown in Table II, because all of them are iterative methods. The additional operations for \(C_{\text{det1}}\) include the incorporation of non-uniform priors for computations of the initial pair and estimated APPs [16]. The corresponding results are listed in Table II.

Subsequently, we present a numerical demonstration of the computational complexity for the \(4\times 4\) MIMO system, which is considered in the simulation of the turbo receiver. The number of RVMs of the competing schemes in this system is presented in the last but one column of Table II. EXT-GEPNet requires more RVMs than EP because of the additional GNN operations. Notably, the GNN operations can be more computationally expensive than the matrix inversion in EP for the considered small-sized \(4\times 4\) system because the chosen hyperparameters of the GNN (\(N_{\mathrm{h1}}=64\) and \(N_{\mathrm{h2}}=32\)) are much larger than the system sizes (\(N=K=8\)). However, the matrix inversion becomes dominant when the system size \(K\) grows, because the complexity of matrix inversion is on the order of \(\mathcal{O}(K^{3})\), whereas the complexity of GNN is on the order of \(\mathcal{O}(K^{2})\), as indicated in Table II. Therefore, the complexity ratio between the proposed EXT-GEPNet and EP in large-scale MIMO systems narrows down and tends to 1.

Additionally, we analyze the effect of edge pruning on the numerical complexity results. The pruning factor \(\alpha\) of EXT-GEPNet for calculating the number of RVMs is with the same

Fig. 11: BER performance (\(I=2\)) of the EXT-GEPNet with antenna configuration and SNR mismatches for a \(4\times 2\) spatially correlated MIMO channel with \(\rho=0.5\). Turbo codes with \(R_{\mathrm{c}}=1/2\) and \(N_{\mathrm{b}}=1952\) are used. The modulation type is 16-QAM.

Fig. 12: BER performance comparison (\(I=2\)) between perfect CSI and estimated CSI using LMMSE method for a \(4\times 2\) correlated MIMO channel specified by the 5G-NR standard. Turbo codes with \(R_{\mathrm{c}}=1/2\) and \(N_{\mathrm{b}}=1952\) and 16-QAM are used.
choice as that for the performance evaluation in Table I. The column of RVMs in Table II shows that a larger \(\alpha\) leads to a smaller \(\eta\) and thus fewer RVMs, as expected. For example, approximately 35%, 41%, and 49% RVMs can be saved using a pruning factor \(\alpha\) equal to 0.5, 1, and 2, respectively. Considering the performance under \(R_{\text{c}}=1/2\) given by the second column of Table I, the edge pruning versions with \(\alpha\) equal to 0.5, 1, and 2 reach a more desirable trade-off between performance and complexity than the original EXT-GEPNet without pruning.

Finally, we compare the complexity of EXT-GEPNet with STS-SD [3]. The number of RVMs for STS-SD is omitted in Table II, as there is no analytical complexity expression, and the required operations are mainly sequential search [3], which cannot be evaluated by multiplication count. Therefore, we perform a running time test for the comparison under the 4 \(\times\) 4 system configuration. Results reveal that the proposed EXT-GEPNet without pruning requires 2.53 times less running time than STS-SD. The improvement of the edge pruning versions in reducing the running time is also remarkable. For example, the edge pruning version with \(\alpha=4\) runs 1.58 and 3.99 times faster than the FC EXT-GEPNet and STS-SD, respectively.

Moreover, the STS-SD is impractical to realize for large-scale MIMO systems because its complexity is exponential to the system size [4]. In contrast to STS-SD, the proposed scheme has a complexity that polynomially grows with the system sizes and can make a flexible trade-off between complexity and performance by adjusting the pruning factor \(\alpha\), as well as the hidden layer sizes \(N_{\text{h}1}\) and \(N_{\text{h}2}\) in the GNN. Therefore, EXT-GEPNet can be viewed as a powerful and efficient scheme given the performance indicated in the simulations and the polynomial complexity reported in Table II.

## V Conclusions

We proposed a GNN-enhanced EP algorithm for MIMO turbo receivers. We first developed the soft-input soft-output mechanism for GEPNet and the corresponding turbo receiver structure. We then customized a training scheme to establish the EXT-GEPNet, which inherits the superiority of GEPNet in achieving an augmented _a posteriori_ estimation via the GNN and addresses the limitations of failing to produce reliable extrinsic LLRs. The EXT-GEPNet can be deployed in the developed turbo structure to take full advantage of various prior information and achieve stable turbo receiving, outperforming the approach using the original GEPNet. Furthermore, we developed an edge pruning method to eliminate the redundancy in the network, resulting in a significant complexity reduction with negligible performance loss. Complexity analysis and simulation results confirm the efficiency, excellent performance, and adaptability of the proposed scheme.
