\section{Introduction}
Massive/ultra-massive multiple-input multiple-output (MIMO) communication has become one of the enabling technologies for 5G and beyond networks owing to the high spectral efficiency and link reliability. Given the increasing size of antenna arrays, efficient MIMO detection algorithms are of significant importance to unleash the full potential of  MIMO systems, and a series of research has been conducted to balance the performance and complexity \cite{SD,AMPJSTSP,EPdetector}. To this end, the maximum likelihood (ML) detector achieves the optimal performance, but its complexity increases exponentially with the number of decision variables because of the exhaustive search. As a result, it cannot be deployed in practical systems. Hence,  many suboptimal linear detectors, such as zero-forcing (ZF) and minimum mean-squared error (MMSE) detectors, have been proposed to reduce the complexity at the cost of performance.

Recently, iterative detectors based on approximate message passing (AMP) \cite{dynamicAMP} and expectation propagation (EP) \cite{EP} have been proposed. The AMP-based detector \cite{AMPJSTSP} is of low complexity and easy to implement in practice because only the matrix-vector multiplication is involved. However, it only favors the scenario when the elements of the MIMO channel matrix follows independent identically distributed (i.i.d.) sub-Gaussian distribution with zero mean. In contrast, the EP-based detector \cite{EPdetector} achieves Bayes-optimal performance when the channel matrix is unitarily invariant, but has a higher complexity than the AMP-based detector owing to involves a matrix inversion.

Owing to its strong ability to extract useful features from data, deep learning (DL) has been recently utilized in the physical layer design of wireless communications \cite{Khlaed_JSAC,6G_khaled,Modeldriven18DL}, such as millimeter-wave channel estimation \cite{DL2018HE}, channel state information (CSI) feedback \cite{DL2018Wen}, and data detection \cite{MIMO_Detection20DL}. For MIMO detection, it has been shown that DL can improve traditional message passing detectors, such as AMP, orthogonal AMP (OAMP) \cite{OAMP}, and EP, with different strategies. Specifically, the OAMP-Net and OAMP-Net2 detectors \cite{MIMO_Detection20DL} were developed by unfolding the OAMP detector and introducing several learnable parameters.
Recently, graph neural networks (GNNs) have been applied to wireless communications \cite{YifeiGNN2}, such as to learn a message-passing solution for  inference problems \cite{Scotti20GNN,Kosasih20DL}. In particular, a GNN-based MIMO detector was developed by utilizing a pair-wise Markov random field (MRF) model \cite{Scotti20GNN}. As GNNs can capture the structure of the data into the feature vectors for the nodes and update them through message-passing methods, they are very promising solutions for enhancing  the performance of the message-passing-based detectors. For example, the GEPNet was developed by incorporating the GNN into the EP detector\cite{Kosasih20DL}.
However, the computational complexity of GEPNet is still high because of the matrix inversion in each layer. An efficient DL-based MIMO detection algorithm, which strikes a good balance between the performance and complexity, is by far not available.


In this paper, we develop a model-driven DL-based MIMO detector, AMP-GNN, where the neural network structure is obtained by unfolding the AMP detector and incorporating the GNN module. In particular, the GNN module receives the equivalent AWGN observation from the AMP as input and outputs a refined version back to AMP at each layer. It can improve the accuracy of the equivalent AWGN observation for AMP algorithm. As a result,  it is helpful for multi-user interference (MUI) cancellation. The AMP-GNN also inherits the low-complexity of the AMP detector and improves the performance with the GNN module. Simulation results shall demonstrate that the proposed AMP-GNN significantly outperforms existing the AMP detector and achieve comparable performance as the state-of-the-art GEPNet detector while reducing the computational complexity.

\emph{Notations}---For any matrix $\mathbf{A}$, $\mathbf{A}^{T}$, $\mathbf{A}^{*}$, and ${ \mathrm{tr}}(\mathbf{A})$ will denote the transpose, conjugate, and  trace of $\mathbf{A}$, respectively. In addition, $\mathbf{I}$ is the identity matrix, $\mathbf{0}$ is the zero matrix, and $\mathbf{1}_n$ is the $n$-dimensional all-ones vector. A proper complex Gaussian distribution with mean $\boldsymbol{\mu}$ and covariance $\boldsymbol{\Omega}$ can be described by the probability density function (pdf):
\begin{equation*}
  \mathcal{N}_{\mathbb{C}}(\mathbf{z};\boldsymbol{\mu},\boldsymbol{\Omega})=\frac{1}{\mathrm{det}(\pi \boldsymbol{\Omega})}
  e^{-(\mathbf{z}-\boldsymbol{\mu})^{H}\boldsymbol{\Omega}^{-1}(\mathbf{z}-\boldsymbol{\mu})}
\end{equation*}

\section{Problem Formulation and Algorithm Review}\label{Problem}
In this section, we first formulate the MIMO detection problem by adopting a Bayesian inference. To better understand the proposed AMP-GNN, we will review the AMP-based MIMO detector.
\subsection{MIMO Detection}
 We assume an uplink multi-user MIMO (MU-MIMO) systems where the base station (BS) equips $M$ antennas sevres $N$ single-antenna users. The symbol vector $\bar{{\mathbf{x}}}\in\mathbb{C}^{M\times1}$ is transmitted over a Rayleigh fading channel $\bar{{\mathbf{H}}} \in \mathbb{C}^{M \times N}$. Each element of $\bar{{\mathbf{H}}}$ and $\bar{{\mathbf{x}}}$ are drawn from an i.i.d. complex Gaussian distribution and the $Q$-QAM constellation, respectively. The received signal $\bar{\mathbf{y}}\in\mathbb{C}^{M\times1}$ is given by
\begin{equation}\label{eq1}
  \bar{\mathbf{y}}=\bar{\mathbf{H}}\bar{\mathbf{x}}+\bar{\mathbf{n}},
\end{equation}
where $\bar{\mathbf{n} }\sim \mathcal{N}_{\mathbb{C}}(0,\sigma^{2}\mathbf{I}_{M})$ is the additive white Gaussian noise (AWGN).


To better implement our proposed AMP-GNN detector, we consider an equivalent real-valued representation for the MIMO detection problem which is obtained by considering the real $\mathfrak{R}(\cdot)$ and imaginary $\mathfrak{I}(\cdot)$ parts separately. Denote $\mathbf{x}=[\mathfrak{R}(\bar{\mathbf{x}})^{T}, \mathfrak{I}(\bar{\mathbf{x}})^{T}]^{T}$, $\mathbf{y}=[\mathfrak{R}(\bar{\mathbf{y}})^{T}, \mathfrak{I}(\bar{\mathbf{y}})^{T}]^{T}$, $\mathbf{n}=[\mathfrak{R}(\bar{\mathbf{n}})^{T}, \mathfrak{I}(\bar{\mathbf{n}})^{T}]^{T}$, and
\begin{equation}\label{eq2}
\mathbf{H}=\left[\begin{array}{c c c}
\mathfrak{R}(\bar{\mathbf{H}}) && -\mathfrak{I}(\bar{\mathbf{H}}) \\
\mathfrak{I}(\bar{\mathbf{H}}) && \mathfrak{R}(\bar{\mathbf{H}})
\end{array}\right].
\end{equation}
Then, the signal model in (\ref{eq1}) can be rewritten in terms of real vectors and matrix form as follows,
\begin{equation}\label{eq3}
\mathbf{y}=\mathbf{H}\mathbf{x}+\mathbf{n}.
\end{equation}

Based on the Bayes' theorem, the posterior probability $\mathcal{P}(\mathbf{x}|\mathbf{y},\mathbf{H})$ can be factorized as
\begin{equation}\label{eq4}
    \mathcal{P}(\mathbf{x}|\mathbf{y},\mathbf{H})=\frac{\mathcal{P}(\mathbf{y}|\mathbf{x},\mathbf{H})\mathcal{P}(\mathbf{x})}{\mathcal{P}(\mathbf{y}|{\mathbf{H}})}
 = \frac{\mathcal{P}(\mathbf{y}|\mathbf{x},\mathbf{H})\mathcal{P}(\mathbf{x})}{\int \mathcal{P}(\mathbf{y}|\mathbf{x},\mathbf{H})\mathcal{P}(\mathbf{x})d\mathbf{x}}.
\end{equation}
Given the posterior probability $\mathcal{P}(\mathbf{x}|\mathbf{y},\mathbf{H})$, the Bayesian MMSE estimate is obtained by
\begin{equation}\label{MMSE_estimate}
\hat{\mathbf{x}}=\int \mathbf{x} \mathcal{P}(\mathbf{x}|\mathbf{y},\mathbf{H})d\mathbf{x}.
\end{equation}

However, the Bayesian MMSE estimator is computationally intractable because the MMSE estimation in (\ref{MMSE_estimate}) involves a high-dimensional integral. In a recent study \cite{OAMP}, the AMP algorithm was proposed as an efficient alternative to recover signal $\mathbf{x}$. We will introduce the AMP-based detector in the following subsection.
\subsection{AMP-Based MIMO Detector}\label{sec:AMP}
\begin{algorithm}[!t]
\caption{AMP-based MIMO detector}
\label{Algorithm:BAMP}
{
\begingroup
\textbf{1. Input:} ${\mathbf{y}}$, ${\mathbf{H}}$, $\sigma^2$, $\mathcal{P}({\mathbf{x}})$. \\
\textbf{2. Initialization:} $\hat{x}_n^{(1)}=0$, $\hat{v}_n^{(1)}=\frac{N}{M}$, $Z_m^{(0)}=y_m$.\\
\textbf{3. Output:} $\hat{{\mathbf{x}}}^{(T)}$.\\
\textbf{4. Iteration:} \\
\For{$t=1,\cdots, T$}
{
 \setlength\abovedisplayskip{0pt}
 \setlength\belowdisplayskip{0pt}
 \begin{subequations}
 \begin{align}
 V_m^{(t)}&=\sum_{n=1}^N |h_{mn}|^2\hat{v}_n^{(t)} \label{Equ:AMP1}\\
 Z_m^{(t)}&=\sum_{n=1}^Nh_{mn}\hat{x}_n^{(t)}-\frac{V_m^{(t)}(y_m-Z_m^{(t-1)})}{\sigma^2+V_m^{(t-1)}} \label{Equ:AMP2}\\
 \Sigma_n^{(t)}&=\left(\sum_{m=1}^M\frac{|h_{mn}|^2}{\sigma^2+V_m^{(t)}}\right)^{-1} \label{Equ:AMP3}\\
 r_n^{(t)}&= \hat{x}_n^{(t)}+\Sigma_n^{(t)}\sum_{m=1}^M\frac{h_{mn}^{*}(y_m-Z_m^{(t)})}{\sigma^2+V_m^{(t)}} \label{Equ:AMP4}\\
 \hat{x}_{n}^{(t+1)}& =\mathbb{E}\{x_n|r_n^{(t)},\Sigma_n^{(t)}\}
 \label{Equ:AMP5}\\
 \hat{v}_{n}^{(t+1)}& =\text{Var}\{x_n|r_n^{(t)}, \Sigma_n^{(t)}\}
 \label{Equ:AMP6}
 \end{align}
 \end{subequations}
}
\endgroup
}
\end{algorithm}
\begin{figure*}
  \centering
  \includegraphics[width=15cm]{AMP_GNN}
  \caption{\small{.~~The structure of the proposed AMP-GNN detector.}}\label{fig:AMPGNN}
\end{figure*}

The AMP algorithm has been proposed to solve sparse linear inverse problems in compressed sensing \cite{OAMP}. In Algorithm \,\ref{Algorithm:BAMP}, we summarize the AMP algorithm for MIMO detection\footnote{Note that we consider a complex-valued AMP-based MIMO detector in  Algorithm \,\ref{Algorithm:BAMP} and the equivalent real-valued form can be derived with the equivalent real-valued representation in (\ref{eq2}) and (\ref{eq3}) accordingly.}, where $m$ and $n$ are the indexes of the antenna in the BS and users, respectively. The main principle of the algorithm is to decouple the posterior probability $\mathcal{P}(\mathbf{x}|\mathbf{y},\mathbf{H})$ into a series of $\mathcal{P}(x_{n}|\mathbf{y},\mathbf{H}), (n=1,2,\ldots,N)$, in an iterative way. In particular, $\mathcal{P}(x_{n}|\mathbf{y},\mathbf{H})$ is assumed to be a Gaussian distribution and obtained from the equivalent AWGN model
\begin{equation}\label{eqAWGN}
    r_{n}^{(t)} = x_{n} + w_{n}^{(t)},
\end{equation}
where  $w_{n}^{(t)}\sim \mathcal{N}_{\mathbb{C}}(w_{n}^{(t)};0,\Sigma_n^{(t)})$.

To better understand Algorithm \ref{Algorithm:BAMP}, we first provide a brief explanation for each equation in Algorithm \ref{Algorithm:BAMP} with the detailed derivation given in \cite{dynamicAMP}. Specifically, (\ref{Equ:AMP1}) and (\ref{Equ:AMP2}) compute the variance and mean estimation for ${\mathbf{z}} ={\mathbf{H}}{\mathbf{x}}$, respectively. Then,  (\ref{Equ:AMP3}) calculates the equivalent noise variance $\Sigma_n^{(t)}$ while
(\ref{Equ:AMP4}) gives the equivalent AWGN observation $r_n^{(t)}$. Finally, (\ref{Equ:AMP5}) and (\ref{Equ:AMP6}) perform the posterior mean and variance estimation for the equivalent AWGN model (\ref{eqAWGN}). As the transmitted symbol is assumed to be drawn from the $Q$-QAM set $\mathcal{S}=\{s_{1}, s_{2}, \ldots, s_{Q}\}$, the results in (\ref{Equ:AMP5}) and (\ref{Equ:AMP6}) are given by
\begin{equation}\label{eqmean}
  \hat{x}^{(t+1)}_{n} = \frac{\sum_{s_{i} \in \mathcal{S}}s_{i}\mathcal{N}_{{\mathbb{C}}}(s_{i};r_n^{(t)}, \Sigma_n^{(t)})p(s_{i})}{\sum_{s_{i} \in \mathcal{S}}\mathcal{N}_{{\mathbb{C}}}(s_{i};r_n^{(t)}, \Sigma_n^{(t)})p(s_{i})},
\end{equation}
\begin{equation}\label{eqvar}
  \hat{v}^{(t+1)}_{n} = \frac{\sum_{s_{i} \in \mathcal{S}}|s_{i}|^{2}\mathcal{N}_{{\mathbb{C}}}(s_{i};r_n^{(t)}, \Sigma_n^{(t)})p(s_{i})}{\sum_{s_{i} \in \mathcal{S}}\mathcal{N}_{{\mathbb{C}}}(s_{i};r_n^{(t)}, \Sigma_n^{(t)})p(s_{i})}
  -|\hat{x}^{(t+1)}_{n}|^{2}.
\end{equation}

As can be observed in Algorithm \,\ref{Algorithm:BAMP}, the performance of the AMP-based MIMO detector is mainly determined by the accuracy of the assumed equivalent AWGN model (\ref{eqAWGN}).  In \cite{dynamicAMP}, it was shown that the equivalent AWGN model is asymptotically accurate when the dimensions of the system tend to infinity, i.e.,  $M, N \rightarrow \infty$. However, in practical MIMO systems especially with high MUI, the performance of the AMP-based MIMO detector is far from optimal and even has an error floor owing to the inaccurate assumption. Furthermore,  the existing state-of-the-art GEPNet detector requires matrix inversion, which is very complex. These observations motivate us to improve the AMP-based MIMO detector with the DL technique.
\section{AMP-GNN Detector}\label{OAMP-Net}
In this section, we propose an AMP-GNN detector for MIMO detection. First, we present the network structure of the AMP-GNN detector and introduce the GNN module in detail. Then, the computational complexity of the AMP-GNN is analyzed and compared with other MIMO detectors.
\subsection{AMP-GNN Architecture}
Given the important advantages of GNN, such as learning to capture the MUI information, here we exploit them to improve the AMP detector in this article. The block diagram of the AMP-GNN is illustrated in Fig.\,\ref{fig:AMPGNN}. Specifically,  the network consists of $T$ cascade layers, and each layer has the same structure that contains a GNN module and a conventional AMP algorithm shown in Algorithm \,\ref{Algorithm:BAMP}. In particular, the input of the AMP-GNN is the received signal $\mathbf{y}$ and the initial value is set as $\hat{\mathbf{x}}^{(1)}=\mathbf{0}$ and $\hat{\mathbf{v}}^{(1)}=\frac{N}{M}\mathbf{1}_N$, and the output is the final estimate $\hat{\mathbf{x}}^{(T)}$ of signal $\mathbf{x}$. For the $t$-th layer of the AMP-GNN, the inputs are the estimated signal $\hat{\mathbf{x}}^{(t-1)}$ and $ \hat{\mathbf{v}}^{(t-1)}$ from the $(t-1)$-th layer and the received signal $\mathbf{y}$.
As we have introduced the AMP-based MIMO in Section \ref{sec:AMP}, we will introduce  the GNN module in the next subsection.
\subsection{The GNN Module}
GNNs have been  recently adopted in wireless communications as they can incorporate the graph topology of the wireless network into the neural network design \cite{YifeiGNN2}. The GNN model adopted in the AMP-GNN is called the message passing neural network (MPNN) and the message passing rules are illustrated in Fig.\,\ref{fig:GNN}. Similar to  \cite{Scotti20GNN}, the GNN in the AMP-GNN is composed of three main modules: a propagation module, an aggregation module, and a readout module. The first two modules operate at all layers while the readout module is involved only after the last layer.

As the pair-wise MRF is used to model the structured dependency of a set of random variables ${\mathbf{x}} = \{x_{1},\cdots,x_{N}\}$ by an undirected graph $G = \{V, E\}$, it can be adopted to model the correlation of the nodes in GNNs. Specifically, the $n$-th variable node is characterized by a self potential $\phi(x_n)$, and the $(n,j)$-th pair of edge is characterized by a pair potential $\psi(x_n,x_j)$, which are given by
\begin{subequations}\label{phi_and_psi}
            \begin{equation}\label{phi}
\phi(x_n) = {\mathsf{e}} {\mathsf{x}} {\mathsf{p}} \left( \frac{1}{\sigma^2} {\mathbf{y}}^T {\mathbf{h}}_n x_n - \frac{1}{2}  {\mathbf{h}}_n^T {\mathbf{h}}_n x_n^2  \right) \mathcal{P}(x_n),
            \end{equation}
            \begin{equation}\label{psi}
\psi(x_n,x_j) = {\mathsf{e}} {\mathsf{x}} {\mathsf{p}} \left(- \frac{1}{\sigma^2} {\mathbf{h}}_n^T {\mathbf{h}}_j x_n x_j\right),
            \end{equation}
\end{subequations}
respectively, where $n,j \in \{1,\ldots, N\}$ and $n \neq j$. The posterior probability $ p_{\mathrm{GNN}}({\mathbf{x}}|{\mathbf{y}})$ corresponding to the pair-wise MRF can be obtained by GNN and written as \cite{Scotti20GNN}
\begin{equation}\label{p_x_y_GNN}
 p_{\mathrm{GNN}}({\mathbf{x}}|{\mathbf{y}})
 =\frac{1}{Z} \prod_{n=1}^N \phi(x_n)
 \prod_{\substack{j=1\\mathrm{j} \neq n}}^N \psi(x_n,x_j),
\end{equation}
where $Z$ is a normalization constant. As the GNN consists of the nodes and edges and is used to infer the posterior probability, we incorporate the mean $r_{n}^{(t)}$ and variance $\Sigma_n^{(t)}$  obtained from  (\ref{Equ:AMP3}) and (\ref{Equ:AMP4}) into the attribute ${\mathbf{a}}_n^{(t)}$  of the variable node $x_n$ by concatenating the mean and variance as
\begin{equation}\label{concat}
{\mathbf{a}}_n^{(t)} = \left[ r_{n}^{(t)}, \Sigma_n^{(t)} \right].
\end{equation}
To compute $p_{\mathrm{GNN}}({\mathbf{x}}|{\mathbf{y}})$ in \eqref{p_x_y_GNN}, we utilize the node and edge feature vectors corresponding to the self and pair potentials in \eqref{phi} and \eqref{psi}, respectively. Denote ${\mathbf{u}}_n^{(\ell)}$ as the node feature vector.  Its initial value is calculated from encoding the information of the received signal ${\mathbf{y}}$, corresponding channel vector ${\mathbf{h}}_n$, and noise variance $\sigma^2$.  The encoding process is implemented by using a single layer neural network given by
\begin{equation}\label{GNN_init}
 {\mathbf{u}}_n^{(0)} = {\mathbf{W}}_1 \cdot [{\mathbf{y}}^T {\mathbf{h}}_n, {\mathbf{h}}_n^T {\mathbf{h}}_n, \sigma^2]^T + {\mathbf{b}}_1,
\end{equation}
where ${\mathbf{W}}_1 \in \mathbb{R}^{N_u \times 3}$ is a learnable matrix, ${\mathbf{b}}_1 \in \mathbb{R}^{N_u}$ is a learnable vector, and $N_u$ is the size of the feature vector.
The edge feature vector  $\mathbf{f}_{jn} \triangleq \left[ {\mathbf{h}}_n^T {\mathbf{h}}_j, \sigma^2 \right] $ is obtained by extracting the pair potential information from \eqref{psi} and  utilized for the message passing of the GNN.
\begin{figure}[t]
  \centering
  \includegraphics[width=7.2cm]{GNN}
  \caption{\small{.~~Message passing and update for GNNs.}}\label{fig:GNN}
\end{figure}
Next, we will elaborate the details of each module in the MPNN, including the propagation, the aggregation, and the readout module.
\subsubsection{\textbf{Propagation module}}
For any pair of variable nodes $x_n$ and $x_j$, there is an edge to connect them. Each edge first concatenates the received feature vectors ${\mathbf{u}}_n^{(\ell-1)}$ and ${\mathbf{u}}_j^{(\ell-1)}$ with its own feature vector $\mathbf{f}_{jn}$ as
${\mathbf{c}}_{n}^{(l-1)} = [{\mathbf{u}}_n^{{(}\ell-1{)}}, {\mathbf{u}}_j^{{(}\ell-1{)}}, \mathbf{f}_{jn}]$. Then, it uses the concatenated features ${\mathbf{c}}_{n}^{(l-1)}$ as the input for the multi-layer perceptron (MLP). Therefore, the output of the MLP is given by
\begin{equation}\label{factor_to_var}
{\mathbf{m}}_{jn}^{(\ell)} = {\mathsf{D}} \left( {\mathbf{c}}_{n}^{(l-1)} \right),
\end{equation}
where  $\mathsf{D}$ is the MLP operation. In the propagation module, each edge has an MLP with two hidden layers of sizes $N_{h_1}$ and $N_{h_2}$ and an output layer of size $N_u$. 
Furthermore, the rectifier linear unit (ReLU) activation function is used at the output of each hidden layer. Finally, the outputs ${\mathbf{m}}_{jn}^{(\ell)}$ are fed back to the nodes as shown in Fig.\,\ref{fig:GNN}.
\subsubsection{\textbf{Aggregation  module}}
The $n$-th variable node  sums all the incoming messages ${\mathbf{m}}_{jn}^{(\ell)}$ from its connected edges  and concatenates the sum of the ${\mathbf{m}}_{jn}^{(\ell)}$  with the node attribute ${\mathbf{a}}_n^{(t)}$ as ${\mathbf{m}}_n^{(\ell)}  = \left[\sum_{\substack{j=1\\mathrm{j} \neq n}}^N {\mathbf{m}}_{jn}^{(\ell)},  {\mathbf{a}}_n^{(t)} \right]$. The ${\mathbf{m}}_n^{(\ell)}$ is used to compute the node feature vector ${\mathbf{u}}_n^{(\ell)}$ as
 \begin{subequations}\label{GRU_MLP}
            \begin{equation}\label{GRU}
 {\mathbf{g}}_n^{(\ell)} = {\mathsf{U}} \left( {\mathbf{g}}_n^{(\ell-1)}, {\mathbf{m}}_n^{(\ell)}  \right),
            \end{equation}
            \begin{equation}\label{MLP}
 {\mathbf{u}}_n^{(\ell)}= {\mathbf{W}}_2 \cdot {\mathbf{g}}_n^{(\ell)} + {\mathbf{b}}_2,
            \end{equation}
\end{subequations}
where the function $\mathsf{U}$ is specified by the gated recurrent unit (GRU) network, whose current and previous hidden states are ${\mathbf{g}}_n^{(\ell)} \in \mathbb{R}^{N_{h_1} }$ and  ${\mathbf{g}}_n^{({\ell}-1)} \in \mathbb{R}^{N_{h_1} }$, respectively. ${\mathbf{W}}_2 \in \mathbb{R}^{N_u \times N_{h_1}}$ is a learnable matrix, and  ${\mathbf{b}}_2\in \mathbb{R}^{N_u}$ is a learnable vector. The updated feature vector \eqref{MLP} is then sent to the propagation module for next iterations.
\subsubsection{\textbf{Readout module}}
After $L$ rounds of the message passing between the propagation and aggregation module, a readout module is utilized to output the estimated distribution
$p_{\mathrm{GNN}}^{(t)}(x_n=s_{i}|{\mathbf{y}})$ for the $t$-layer of AMP-GNN and given by
\begin{subequations}\label{Readout}
				\begin{equation}\label{Readout1}
\tilde p^{(t)}_{ \mathrm{GNN}}(x_n=s_{i}|{\mathbf{y}})  = {\mathsf{R}} \left(  {\mathbf{u}}_n^{(L)} \right), s_{i}\in \mathcal{S},
				\end{equation}
            	\begin{equation}\label{Readout2}
p_{\mathrm{GNN}}^{(t)}(x_n=s_{i}|{\mathbf{y}})  = \frac{{{\mathsf{e}} {\mathsf{x}} {\mathsf{p}} } \left(\tilde p^{(t)}_{G\mathrm{NN}}(x_n=s_{i}|{\mathbf{y}}) \right)}  {\sum_{s_{i}\in \mathcal{S}}  {{\mathsf{e}} {\mathsf{x}} {\mathsf{p}}} \left({\tilde p}^{(t)}_{ \mathrm{GNN}}(x_n=s_{i}|{\mathbf{y}}) \right)}, s_{i}\in \mathcal{S}.
           		 \end{equation}
\end{subequations}
The readout function ${\mathsf{R}} $ consists of an MLP with two hidden layers of sizes $N_{h_1}$ and $N_{h_2}$, and ReLU activation is utilized at the output of each hidden layer. The output size of  ${\mathsf{R}} $ is the cardinality of the real-valued constellation set, i.e., $\sqrt{Q}$. Then, we compute
\begin{equation}\label{GNN_reset_index}
 {\mathbf{g}}_n^{(0)}  \leftarrow   {\mathbf{g}}_n^{(L)} \text { and }{\mathbf{u}}_n^{(0)} \leftarrow   {\mathbf{u}}_n^{(L)}, n= 1, \dots,N,
\end{equation}
for the next GNN iteration. Finally, we utilize $p_{\mathrm{GNN}}^{(t)}(x_n=s_{i}|{\mathbf{y}})$ to compute the posterior mean and variance for the AMP, which are given by
\begin{subequations}\label{pyGNN}
            \begin{equation}\label{pyGNN1}
 \hat{x}_{n}^{(t+1)}=\mathbb{E}\{x_n|p_{\mathrm{GNN}}^{(t)}\},
            \end{equation}
            \begin{equation}\label{pyGNN2}
 \hat{v}_{n}^{(t+1)}=\text{Var}\{x_n|p_{\mathrm{GNN}}^{(t)}\}.
            \end{equation}
\end{subequations}
Note that the expectation and variance are computed with respect to $p_{\mathrm{GNN}}^{(t)}$. This is the main difference between AMP and AMP-GNN as $p_{\mathrm{GNN}}^{(t)}(x_n=s_{i}|{\mathbf{y}})$ is not the Gaussian pdf in the AMP-GNN detector. After computing (\ref{pyGNN1}) and (\ref{pyGNN2}), the posterior mean $\hat{x}_{n}^{(t+1)}$ and $\hat{v}_{n}^{(t+1)}$ will be used for the next AMP-GNN iteration. Finally, the AMP-GNN is executed iteratively until terminated by a fixed number of layers.
\begin{table*}[t]	
\centering
	\caption{Computational complexity (the number of multiplications) of different detectors.
}
	\label{tab:complexity}
    \begin{tabular}{@{}lcccccc@{}}
    \toprule
    \diagbox{MU-MIMO settings}{Detectors}&OAMP-Net&GNN&GEPNet&AMP&AMP-GNN&RE-MIMO\\
    \midrule
    \midrule
    $64 \times 64$ &$8.22\times10^{6}$&$1.17\times10^{6}$&$5.11\times10^{6}$&$1.78\times10^{5}$&$2.35\times10^{6}$& $4.93\times10^{8}$ \\
    \midrule
    $256 \times 256$ &$5.21\times10^{8}$&$1.27\times10^{7}$&$2.02\times10^{8}$&$2.68\times10^{6}$&$1.93\times10^{7}$& $3.22\times10^{9}$ \\
    \midrule
    $1024 \times 1024$ &$3.33\times10^{10}$&$5.56\times10^{8}$&$1.24\times10^{10}$&$4.22\times10^{7}$&$6.14\times10^{8}$& $4.41\times10^{10}$ \\
    \bottomrule
    \end{tabular}
\end{table*}

\subsection{Complexity Analysis}
We analyze the computational complexity of the AMP-GNN and compare it with existing MIMO detectors. Specifically,
the complexity of the AMP detector is  $\mathcal{O}(MN)$ owing to the matrix-vector multiplication while the complexity for GNN is  $\mathcal{O}(N_{h1}N_{h2})$ accounts for the MLP operation. Therefore, the computational complexity of the AMP-GNN is $\mathcal{O}(MN+N_{h1}N_{h2})$, dominated by the complexity of the AMP and GNN. In contrast, the complexity of the GEPNet is $\mathcal{O}(M{N}^2+N_{h1}N_{h2})$ which includes the computational complexity of the EP and GNN.

To better compare the complexity, we use the the number of multiplication as the metric and show the exact values for different MIMO settings with quadrature phase shift keying (QPSK) symbols in Table.\,\ref{tab:complexity}. The hyperparameters for GNN are set as $N_{h_1}=16$, $N_{h_2}=8$, and $N_u = 8$.  Compared with state-of-the-art DL-based MIMO detectors, such as  GEPNet and RE-MIMO, the AMP-GNN entails a much lower complexity. In particular, the ratio between the complexity of the AMP-GNN and GEPNet is dramatically reduced when the number of user increases. For example, the ratio between the complexity of the AMP-GNN and GEPNet is only $45.99\%$ when $M=N=64$  while the ratio is $4.95\%$ when $M=N=1024$.  This is because the complexity of matrix inversion  in the GEPNet will be the dominant term, which is prohibitively high when the number of antennas is large. In contrast, the AMP-GNN only involves  matrix-vector multiplications, which is a favorable feature for future ultra massive MIMO systems.
\section{Simulation Results}\label{Simulation}
In this section, we provide simulation  results of the AMP-GNN for MIMO detection and compare it with other MIMO detectors. We use the symbol error rate (SER) as the performance metric in our simulations. The signal-to-noise (SNR) of the system, defined as $\mathrm{SNR}=\frac{\mathbb{E}\|\mathbf{H}\mathbf{x}\|^{2}_{2}}{\mathbb{E}\|\mathbf{n}\|^{2}_{2}}$. To illustrate the effectiveness of our proposed AMP-GNN, we adopt several well established MIMO detectors as baselines:
\begin{itemize}
  \item \textbf{MMSE}: A classical linear receiver for MIMO detection which inverts the received signal by applying the channel-noise regularized pseudo-inverse
        of the channel matrix.
  \item \textbf{AMP}: An efficient message passing algorithm for MIMO detection given in Algorithm \ref{Algorithm:BAMP} and implemented with $10$ iterations\footnote{It was found that a  further increase in the number of iterations only offers a negligible performance gain. We set the same number of layers in other DL-based baseline methods for fair comparison.}.

  \item \textbf{OAMP-Net}: The OAMP-based model-driven DL detector developed in \cite{MIMO_Detection20DL}. Each layer requires computing a matrix pseudo-inverse and has 2 learnable parameters.
 
  \item \textbf{GNN}: The GNN-based MIMO detector  proposed in \cite{Scotti20GNN}.
  \item \textbf{GEPNet}: The GNN-enhanced EP detector proposed in \cite{Kosasih20DL} with $10$ layers.
  \item \textbf{RE-MIMO}: The RE-MIMO detectors  proposed in \cite{REMIMO20DL} based on the RIM with $10$ layers.

\end{itemize}
\begin{figure}
  \centering
  \includegraphics[width=3 in]{Fig1_r1}
  \caption{\small{.~~SER comparison of AMP-GNN with other MIMO detectors under $4\times4$ Rayleigh MIMO
   channels with QPSK symbols.}}\label{fig2}
\end{figure}
\subsection{Implementation Details}
In the simulation, the AMP-GNN is implemented on the PyTorch platform. The number of layers of the AMP-GNN detector  is set to $T=10$ while the number of layers of the GNN is set to $L=2$. The training data consists of a number of randomly generated pairs $(\mathbf{x},\mathbf{y})$. The data $\mathbf{x}$ is generated from QAM modulation symbols. We train the network with $100$ epochs. At each epoch, the training and validation sets contain $100,000$, and $5,000$ samples, respectively. The AMP-Net is trained using the stochastic gradient descent method and Adam optimizer. The learning rate is set to be $0.001$ and the batch size is set to $64$.
We choose the $L_{2}$ loss as the cost function, which is defined by,
\begin{equation}\label{eqloss}
  L_{2}(\mathbf{x},\hat{\mathbf{x}}^{(T)})=\|\mathbf{x}-\hat{\mathbf{x}}^{(T)}\|^{2}.
\end{equation}
\subsection{Performance Comparsion}
Fig.\,\ref{fig2} and Fig.\,\ref{fig2_2} compare the average SER of the AMP-GNN with those of the baseline detectors. As can be observed from the figuresa, the AMP-GNN outperforms almost all other MIMO detectors except for the GEPNet detector. In particular, the AMP-GNN outperforms the AMP detector in all SNRs, which demonstrates that the GNN module can improve the AMP detector significantly. Specifically, if we target an SER=$10^{-2}$, then the performance gain is approximately $3.9$ dB compared to the AMP detector. The reason for the performance improvement is that the GNN can refine the equivalent AWGN model with a more accurate distribution $p_{\mathrm{GNN}}^{(t)}(x_n=s_{i}|{\mathbf{y}})$. Furthermore, the AMP-GNN has only $0.8$ dB performance loss compared to the GEPNet detector in $4\times4$ MIMO systems  when we target at an SER=$10^{-2}$. The performance loss then reduces to $0.3$ dB for $32\times32$ MIMO systems as illustrated in Fig.\,\ref{fig2_2}. Thus, we can conclude the AMP-GNN has a comparable performance to GEPNet but with a remarkably reduced computational complexity, especially for massive and ultra-massive MIMO systems.
\begin{figure}
  \centering
  \includegraphics[width=3 in]{Fig_large}
  \caption{\small{.~~SER comparison of AMP-GNN with other MIMO detectors under $32\times32$ Rayleigh MIMO
   channels with QPSK symbols.}}\label{fig2_2}
\end{figure}
\subsection{Robustness to Dynamic Numbers of Users}
Most of the DL-based MIMO detectors are trained and tested with a fixed number of antennas. However, the number of users in practical MIMO systems may quickly change. For example, new users may access to be served by the BS. As a result, the DL-based MIMO detector should have the ability to handle a varying number of users with a single model. In Fig.\,\ref{fig3}, we train the AMP-GNN in an $8\times4$  and $8\times8$ MIMO systems and test it in an $8\times6$ MIMO system. As shown in the figure, if we target an $\mathrm{SER}=10^{-3}$,  then the AMP-GNN still has $1.7$ dB performance gain compared with the conventional AMP detector even when tested with different number of users. Furthermore, it has a similar performance as the AMP-GNN trained and tested both in the $8\times6$ MIMO system. As a result, the AMP-GNN has strong robustness to different numbers of users in the deployment stage. This is because the GNN has the permutation equivariance property which makes it robust against the dynamic changes of the number of users.   
\section{Conclusions}\label{con}
We have developed a novel GNN-enhanced AMP detector for MIMO detection, named AMP-GNN, which is obtained by incorporating the GNN module into the AMP algorithm. The AMP-GNN inherits the respective superiority of the low-complexity AMP detector and efficient GNN module. Simulation results demonstrated that  AMP-GNN improves the performance of the AMP detector significantly and has comparable performance as the state-of-the-art GEPNet detector but with a reduced computational complexity. Furthermore, it is quite robust to the change of the number of users in practice.
\section*{Acknowledgment}
The authors would like to thank Prof. Chao-Kai Wen, from the National
Sun Yat-sen University, for the discussion of neural enhanced message passing
\begin{figure}
  \centering
  \includegraphics[width=3 in]{Fig2_r2}
  \caption{\small{.~~Robustness of the AMP-GNN for dynamic number of users.}}\label{fig3}
\end{figure}

