# Towards Optimally Efficient Search with Deep Learning for Large-Scale MIMO
Systems

Le He, Ke He, Lisheng Fan, Xianfu Lei, Arumugam Nallanathan, _Fellow, IEEE_
and George K. Karagiannidis, _Fellow, IEEE_ L. He, K. He and L. Fan are all
with the School of Computer Science and Cyber Engineering, Guangzhou
University, China (e-mail<EMAIL_ADDRESS><EMAIL_ADDRESS>lsfan@gzhu.edu.cn).X. Lei is with the School of Information Science and
Technology, Institute of Mobile Communications, Southwest Jiaotong University,
Chengdu 610031, China (e-mail: xflei@home.swjtu.edu.cn).A. Nallanathan is with
the School of Electronic Engineering and Computer Science, Queen Mary
University of London, London, U.K. (e-mail: a.nallanathan@qmul.ac.uk)G. K.
Karagiannidis is with the Wireless Communications and Information Processing
Group (WCIP), Aristotle University of Thessaloniki, Thessaloniki 54 124,
Greece (e-mail: geokarag@auth.gr).

###### Abstract

This paper investigates the optimal signal detection problem with a particular
interest in large-scale multiple-input multiple-output (MIMO) systems. The
problem is NP-hard and can be solved optimally by searching the shortest path
on the decision tree. Unfortunately, the existing optimal search algorithms
often involve prohibitively high complexities, which indicates that they are
infeasible in large-scale MIMO systems. To address this issue, we propose a
general heuristic search algorithm, namely, hyper-accelerated tree search
(HATS) algorithm. The proposed algorithm employs a deep neural network (DNN)
to estimate the optimal heuristic, and then use the estimated heuristic to
speed up the underlying memory-bounded search algorithm. This idea is inspired
by the fact that the underlying heuristic search algorithm reaches the optimal
efficiency with the optimal heuristic function. Simulation results show that
the proposed algorithm reaches almost the optimal bit error rate (BER)
performance in large-scale systems, while the memory size can be bounded. In
the meanwhile, it visits nearly the fewest tree nodes. This indicates that the
proposed algorithm reaches almost the optimal efficiency in practical
scenarios, and thereby it is applicable for large-scale systems. Besides, the
code for this paper is available at https://github.com/skypitcher/hats.

###### Index Terms:

Signal detection, integer least-squares, deep learning, maximum likelihood
detection, MIMO, sphere decoding, best-first search

##  I Introduction

Given an observation vector $\bm{y}\in\mathbb{R}^{n\times 1}$ and the
transformation matrix $\bm{H}\in\mathbb{R}^{n\times m}$ with full column rank
($n\geq m$), the signal detection problem aims to recover the transmitted
discrete signal $\bm{x}\in\mathbb{R}^{m\times 1}$ from the linear mixing model

$\displaystyle\bm{y}=\bm{H}\bm{x}+\bm{w},$ (1)

where $\bm{w}\in\mathbb{R}^{n\times 1}$ denotes the additive white Gaussian
noise with zero mean and unit variance. Generally, the problem has diverse
applications which include -but it is not limited to- signal processing [1],
communications [2], machine learning [3], global navigation satellite systems
(GNSS) [4], and radar imaging [5]. In this paper, we will henceforth
concentrate on the optimal signal detection in large-scale multiple-input
multiple-output (MIMO) wireless communication systems. However, it shall be
noted that the proposed method is general, and can be applied to other
applications.

In large-scale MIMO systems, there are $m_{c}$ and $n_{c}$ antennas equipped
at the transmitter and receiver such that $n_{c}\geq m_{c}\gg 1$. Moreover,
$\bm{H}_{c}$ refers to the random wireless channel state information (CSI)
matrix with the $(i,j)$’s element $h_{i,j}\sim\mathcal{CN}(0,\rho)$ denoting
the tap gain from the $j$-th transmit antenna to the $i$-th receive antenna.
In further, $\bm{x}_{c}$ denotes the transmit signal vector which is uniformly
distributed over the finite set of lattice points $\mathcal{A}^{m}$. In fact,
$\mathcal{A}$ is the alphaset depending on the underlying modulation schema,
and a standard example of $\mathcal{A}$ would be
$\mathcal{A}=\\{a+bj|a,b\in\\{-1,+1\\}\\}$ for 4-quadrature amplitude
modulation (4-QAM). While transmitting $\bm{x}_{c}$ through the Rayleigh fat
fading channel $\bm{H}_{c}$, it suffers from an additive white Gaussian noise
$\bm{w}_{c}$. Therefore, the received signal can be expressed as

$\displaystyle\bm{y}_{c}=\bm{H}_{c}\bm{x}_{c}+\bm{w}_{c}.$ (2)

In particular, we consider the following equivalent representation

$\displaystyle\bm{y}=\begin{bmatrix}\Re(\bm{y}_{c})\\\
\Im(\bm{y}_{c})\end{bmatrix},$
$\displaystyle\bm{H}=\begin{bmatrix}\Re(\bm{H}_{c})&-\Im(\bm{H}_{c})\\\
\Im(\bm{H}_{c})&\Re(\bm{H}_{c})\\\ \end{bmatrix},$ (3)
$\displaystyle\bm{x}=\begin{bmatrix}\Re(\bm{x}_{c})\\\
\Im(\bm{x}_{c})\end{bmatrix},$
$\displaystyle\bm{w}=\begin{bmatrix}\Re(\bm{w}_{c})\\\
\Im(\bm{w}_{c})\end{bmatrix},$ (4)

where $\bm{y}_{c}\in\mathbb{C}^{n_{c}\times 1}$,
$\bm{H}_{c}\in\mathbb{C}^{n_{c}\times m_{c}}$,
$\bm{x}_{c}\in\mathbb{C}^{m_{c}\times 1}$ and
$\bm{w}_{c}\in\mathbb{C}^{n_{c}\times 1}$ are composed of complex values, and
$\Re(\cdot)$ and $\Im(\cdot)$ denote the real part and imaginary part of the
value, respectively. Based on this real-valued representation, (2) can be
written as (1) with $m=2m_{c}$ and $n=2n_{c}$.

In lattice theory, $\bm{H}$ is considered to be the generator matrix [6] of
the generated lattice
$\mathcal{L}(\bm{H})=\\{\bm{H}\bm{x}|\bm{x}\in\mathcal{A}^{m}\\}$, which
indicates that the generated lattice is the “skewed” one of the original
lattice. When $\bm{H}$ is perfectly known, the mathematically optimal approach
to solve the problem in terms of minimizing the average error probability, is
to search the closest “skewed” lattice point $\bm{R}\bm{x}^{*}$ to $\bm{y}$,
in terms of the Euclidean distance [2]

$\displaystyle\bm{x}^{*}=\mathop{\arg\min}_{\hat{\bm{x}}\in\mathcal{A}^{m}}\left\|\bm{y}-\bm{H}\hat{\bm{x}}\right\|^{2},$
(5)

which gives the exact optimal maximum likelihood (ML) estimate of $\bm{x}$ in
(1). Since the signal vector $\bm{x}$ comprises of integer components only,
(5) is also refereed as integer least-squares (ILS) problem in the literature
[7]. Unfortunately, solving the ILS problem is much more challenging with
comparison to the standard least-squares problem, where the latter’s signal
vector comprises of continuous entries rather than discrete entries, and the
optimal solution can be efficiently resolved via pseudo inverse [1]. As a
combinatorial optimization problem, it is known to be NP-hard because of its
discrete search space. It has been proven that solving the problem optimally
involves exponential complexity for all algorithms in the worst case [8, 9].
However, one can formulate the problem as a search in a state space, and
efficiently reduce the average complexity by following a certain strategy that
investigates only the necessary lattice points at each dimension[1]. In
particular, the resulting state space forms a decision tree whose nodes
represent decisions on symbols, and branches represent the costs of the
associated decisions. Correspondingly, the strategies that decide which node
should be chosen for expansion are called tree search algorithms. Once we find
the shortest path, we can find the optimal solution.

###  I-A Related Works

In the past decades, researchers have proposed many tree search algorithms to
address the ILS problem [6, 2, 7, 10, 11, 12, 13, 14, 15, 16, 17, 18], and a
comprehensive survey of these algorithms can be found in the fifty-years
review of MIMO detection [19]. One common approach achieving the optimal ML
detection performance is the sphere decoding (SD) algorithm, which employs a
branch-and-bound (BnB) depth-first search (DFS) strategy to find the shortest
path [6, 2, 7]. Besides, many variants of SD have been proposed in the
literature as well. Among these variants, the Schnorr-Eucherr SD (SE-SD)
achieves the optimal performance with reduced average complexity[6]. However,
its average complexity still remains prohibitively high in large-scale MIMO
systems with low signal-to-noise ratio (SNR)[1]. To overcome this drawback,
$K$-best SD has been proposed to achieve a fixed and reduced complexity with a
breath-first search (BrFS) strategy, whereas the performance is sacrificed
[11]. More importantly, the complexity reduction can not be guaranteed in high
SNR regimes [20].

Attempts that employ best-first search (BeFS) or stack algorithms to overcome
the drawbacks of DFS and BrFS have been investigated in the literature [21,
22, 13, 14, 15, 16, 17, 18]. Among these variants, A* like BeFS is the most
promising variant which employs a heuristic function to predict the shortest
remaining cost [13]. It guarantees to find the optimal solution with an
admissible heuristic function, while it achieves the optimal efficiency with a
consistent heuristic function[23]. It has been show that BeFS visits the
fewest nodes among the three search strategies, but it requires exponentially
increasing memory space [15]. Fortunately, several memory-bounded BeFS
algorithms have been proposed in the past years, and their performances are
very close to the original one [21, 22, 17]. However, it is still very
difficult to find a consistent heuristic function such that the optimal
efficiency is achieved [16, 18]. Hence, the existing BeFS algorithms are still
infeasible for large-scale MIMO systems.

With the tremendous success of deep learning techniques on physical layer
communications, significant improvement becomes possible for the
aforementioned search strategies [24, 25, 26, 27]. Regarding the ILS problem,
researchers have proposed a deep learning based SD (DL-SD) algorithm which
significantly improves SE-SD’s performance by employing the deep neural
networks (DNN) to choose a good initial radius [28, 29]. Since the tree search
can be treated as a sequential decision making process, the authors proposed a
sub-optimal search algorithm named LISA in [30] with fixed complexity, since
LISA uses a DNN to make decisions with fixed steps. The success of these data-
driven search algorithms certainly motivates us to improve the BeFS
algorithm’s performance with deep learning techniques.

###  I-B Contributions

Consequently, we are interested in speeding up A* like BeFS algorithms with
model-driven deep learning methods, while trying not to compromise the optimal
bit error rate (BER) performance. To accomplish this goal, we propose a hyper-
accelerated tree search (HATS) algorithm, which enhances the underlying
efficient memory-bounded A* algorithm (SMA*) [21], through predicting the
optimal heuristic with a well-trained DNN. This idea is inspired by the fact
that with the optimal heuristic, A* not only becomes optimally efficient but
also expands the fewest nodes [23]. In contrast, the other A* inspired search
algorithms only find the admissible heuristic [16, 18], which limits the
improvement of these methods. As we will show in the simulations, the proposed
algorithm achieves almost the optimal BER performance in large-scale MIMO
systems, while its memory size can be bounded and it reaches the lowest
average complexity under practical scenarios. This suggests that the proposed
algorithm is feasible for practical large-scale MIMO systems.

##  II MIMO Signal Detection with Heuristic Best-First Search

In this section, we will take an efficient approach to interpret the optimal
MIMO signal detection problem as a tree search, which enables us to decouple
the vector-valued problem into a sequential decision making problem. After
that, we will introduce the A* algorithm to find the shortest path, and
discuss its properties accordingly.

###  II-A Tree Construction

To construct a decision tree, we first perform QR decomposition on the CSI
matrix $\bm{H}$ as

$\displaystyle\bm{H}=\begin{bmatrix}\bm{Q}_{1}&\bm{Q}_{2}\end{bmatrix}\begin{bmatrix}\bm{R}\\\
\bm{0}_{(n-m)\times m}\end{bmatrix}=\bm{Q}_{1}\bm{R},$ (6)

where $\bm{R}\in\mathbb{R}^{m\times m}$ is an upper triangular matrix, and the
partitioned matrices $\bm{Q}_{1}\in\mathbb{R}^{n\times m}$ and
$\bm{Q}_{2}\in\mathbb{R}^{n\times(n-m)}$ both have orthogonal columns. Then,
we rewrite (1) as

$\displaystyle\bm{z}=\bm{R}\bm{x}+\bm{v},$ (7)

where $\bm{z}\triangleq\bm{Q}_{1}^{T}\bm{y}$ and
$\bm{v}\triangleq\bm{Q}_{1}^{T}\bm{w}$. Note that the optimal MIMO signal
detection problem is always squared of dimensions $m$ after preprocessing. For
convenience, we number the entries of matrices and vectors of (7) in a reverse
order as

$\displaystyle\begin{bmatrix}z_{m}\\\ \vdots\\\
z_{1}\end{bmatrix}=\begin{bmatrix}r_{m,m}&\ldots&r_{m,1}\\\
\vdots&\ddots&\vdots\\\ 0&\ldots&r_{1,1}\end{bmatrix}\begin{bmatrix}x_{m}\\\
\vdots\\\ x_{1}\end{bmatrix}+\begin{bmatrix}v_{m}\\\ \vdots\\\
v_{1}\end{bmatrix},$ (8)

where $r_{i,j}$ represents the ($i,j$)-th component of $\bm{R}$ after
arranging from the bottom right to the upper left. After preprocessing, the
squared Euclidean distance of a given candidate can be expanded as

$\displaystyle d^{2}(\bm{x})=$
$\displaystyle\underbrace{(z_{1}-r_{1,1}x_{1})^{2}}_{b(\bm{x}^{1})}+\underbrace{(z_{2}-r_{2,2}x_{2}-r_{2,1}x_{1})^{2}}_{b(\bm{x}^{2})}+\cdots$
$\displaystyle+\underbrace{(z_{m}-r_{m,1}x_{1}-\cdots-
r_{m,m}x_{m})^{2}}_{b(\bm{x}^{m})}$ (9a) $\displaystyle=$
$\displaystyle\sum_{k=1}^{m}{b(\bm{x}^{k})},$ (9b)

where the $k$-th incremental cost is denoted by

$\displaystyle
b(\bm{x}^{k})=\left(z_{k}-\sum_{j=1}^{k}r_{k,j}x_{j}\right)^{2},$ (10)

which only depends on the partial signal vector (PSV)
$\bm{x}^{k}=[x_{k},x_{k-1},\ldots,x_{1}]^{T}$. The cumulative cost of the PSV
$\bm{x}^{k}$ is

$\displaystyle g(\bm{x}^{k})=\sum_{i=1}^{k}b(\bm{x}^{i}),$ (11)

and the successor can be computed recursively by

$\displaystyle g(\bm{x}^{k+1})=b(\bm{x}^{k+1})+g(\bm{x}^{k}),\quad\forall
k=0,1,\cdots,m-1$ (12)

where the initialization $g(\bm{x}^{0})\equiv 0$ is set. Clearly, (12)
constructs a perfect $|\mathcal{A}|$-way $m$-level (starts from zero) decision
tree. In the resulting tree, the deepest nodes represents the goal nodes
associated with specific candidates. Since $|\mathcal{A}|$ is finite, the set
of goal nodes can be enumerated as

$\displaystyle\mathcal{A}^{m}=\\{\bm{x}^{m}_{1},\bm{x}^{m}_{2},\ldots,\bm{x}^{m}_{j},\ldots,\bm{x}^{m}_{|\mathcal{A}|^{m}}\\},$
(13)

where each goal node $\bm{x}^{m}_{j}$ is associated with a signal candidate.
For convenience, we denote the antecedent of a goal node $\bm{x}^{m}_{j}$ that
locates at level $k$ as $\bm{x}^{k}_{j}$, which also represents the path
leading from the root to that node. Note that we will sometimes ignore the
subscript and use $\bm{x}^{k}$ to represent an arbitrary node located at level
$k$. We also use $\bm{x}^{j}\subseteq\bm{x}^{k}$ ($0\leq j\leq k$) to denote
that $\bm{x}^{l}$ is the antecedent of $\bm{x}^{k}$. On the contrary,
$\bm{x}^{j}\supseteq\bm{x}^{k}$ ($0\leq k\leq j$) denotes that $\bm{x}^{j}$ is
the descendent of $\bm{x}^{k}$. Importantly, each node $\bm{x}^{k}$ is
associated with a branch cost $b(\bm{x}^{k})$, and $g(\bm{x}^{k})$ is the
cumulative cost of the path leading from the root to that node. For
demonstration, a simple binary $4$-level decision tree is illustrated in Fig.
1, where $m=4$ and $\mathcal{A}=\\{-1,+1\\}$.

Figure 1: An example of decision tree for signal detection, where $m=4$ and
$\mathcal{A}=\\{-1,+1\\}$. In the resulting tree, the left branch represents
the symbol decision of $-1$, while the right branch represents the symbol
decision of $+1$.

###  II-B Heuristic Tree Search

Now it is straightforward to employ the A* algorithm [23, 16, 18] to search
the least-cost path on the tree. Before describing the algorithm, it is
necessary to introduce the following notations,

  * •

ACTIVE is an _ordered list with possibly limited space_ that stores nodes to
be expanded, whose data structure is often a priority queue or self-balancing
binary search tree (a.k.a. AVLTree) in practice.

  * •

$f(\bm{x}^{k})$ is the _evaluation cost ( $f$-cost)_ currently assigned to a
node $\bm{x}^{k}$ at the time. Note that $f(\bm{x}^{k})$ is not static, and it
may change during search.

  * •

A node is _visited_ or _generated_ if it has been encountered during the
expansion of its parent.

  * •

A node is _expanded_ if all of its successors have been generated during the
expansion.

  * •

A node is _in memory_ if it is inserted into ACTIVE.

  * •

The _complexity_ of a search algorithm is defined in terms of the number of
visited nodes.

We are now ready to introduce the A* algorithm, and the pseudo code is
presented in Algorithm 1. Throughout the paper, we shall emphasize that the
tree search always starts at the dummy root $\bm{x}^{0}$. At each iteration,
A* expands the least-$f$-cost node among all the nodes that are in ACTIVE and
waiting for expansion, and inserts all successors of that node into ACTIVE
afterwards [18]. This process will terminate as long as a goal node is
selected for expansion, and the selected goal node will become the output of
the algorithm. As a heuristic BeFS algorithm, A* employs a heuristic function
$h(\bm{x}^{k})$ to guide the search. Specifically, the heuristic function
estimates the remaining cost from $\bm{x}^{k}$ to goal nodes, denoted by

$\displaystyle h(\bm{x}^{k},\bm{x}^{m})$
$\displaystyle=g(\bm{x}^{m})-g(\bm{x}^{k})$ (14)
$\displaystyle=\sum_{i=k+1}^{m}\left(z_{i}-\sum_{j=1}^{i}r_{i,j}x_{j}\right)^{2},\quad\forall\bm{x}^{m}\supseteq\bm{x}^{k}.$
(15)

Accordingly, A* computes the evaluation cost of $\bm{x}^{k}$ as

$\displaystyle f(\bm{x}^{k})=g(\bm{x}^{k})+h(\bm{x}^{k}),$ (16)

where $h(\bm{x}^{k})$ represents the estimate of the minimum value of (14),
and $h(\bm{x}^{m})=0$ holds since the goal nodes have no successor at all.
Therefore, $f(\bm{x}^{k})$ actually estimates the cost of the shortest path of
the sub-tree of $\bm{x}^{k}$. In particular, the optimal heuristic function
always gives the minimum value of (14) as

$\displaystyle
h^{*}(\bm{x}^{k})=\min_{\bm{x}^{m}\supseteq\bm{x}^{k}}h(\bm{x}^{k},\bm{x}^{m}).$
(17)

In this case, the optimal $f$-cost is given by

$\displaystyle f^{*}(\bm{x}^{k})=g(\bm{x}^{k})+h^{*}(\bm{x}^{k}).$ (18)

With the optimal heuristic function, A* only expands the shortest path, as it
always guides the search process towards the shortest path.

Input: $\bm{z}$ and $\bm{R}$

Output: $\hat{\bm{x}}$ (estimate of the transmitted signal $\bm{x}$)

1 insert the root node $\bm{x}^{0}$ in ACTIVE ;

2 loop

3 if  _ACTIVE is empty_ then

4 return $\hat{\bm{x}}=\varnothing$ with failure;

5

6 end if

7 $\bm{x}^{k}\leftarrow$ _least- $f$-cost node_ in ACTIVE ;

8 if  _$\bm{x}^{k}$ is a goal node_ then

9 return $\hat{\bm{x}}=\bm{x}^{k}$ with success;

10

11 end if

12 for  _every successor $\bm{x}^{k+1}\supset\bm{x}^{k}$_ do

13 $f^{\prime}(\bm{x}^{k+1})\leftarrow g(\bm{x}^{k+1})+h(\bm{x}^{k+1})$;

14
$f(\bm{x}^{k+1})\leftarrow\max\left\\{f(\bm{x}^{k}),f^{\prime}(\bm{x}^{k+1})\right\\}$;

15 insert $\bm{x}^{k+1}$ into ACTIVE ;

16

17 end for

18 remove $\bm{x}^{k}$ from ACTIVE ;

19

20 end

21

Algorithm 1 A* Algorithm

###  II-C Optimality and Optimal Efficiency

We know that A* performs best with the optimal heuristic function, since only
the nodes lying along the shortest path would be expanded eventually. In
practical, we however often use a lower bound on (17) as the heuristic
function, as finding the optimal heuristic function requires exhaustively
search as well. In general, if a heuristic function $h(\bm{x}^{k})$ never
overestimates the optimal heuristic, then it is _admissible_ [18]. In further,
$h(\bm{x}^{k})$ is said to be _consistent_ if $h(\bm{x}^{k})\leq
b(\bm{x}^{k+1})+h(\bm{x}^{k+1})$ holds for each node $\bm{x}^{k}$ and its
successors $\bm{x}^{k+1}$. In other words, a heuristic function is consistent
if it satisfies the triangle inequality [23]. It is clear that a consistent
heuristic is also admissible, but not the vice verse.

With an admissible heuristic function, A* is said to be optimal since it
guarantees to find the shortest path [23]. In particular, the simplest
admissible heuristic is $h(\bm{x}^{k})=0$. Hence, the existing approaches
mainly focus on finding a more meaningful admissible heuristic function based
on the tighter lower bounds on (14) [16, 18]. In this case, A* is able to find
the shortest path faster than the simplest heuristic. However, the performance
improvements are still limited since the estimated heuristics are still far
from the optimal heuristic.

With a consistent heuristic function, A* is said to be optimally efficient in
terms of that not only it is optimal but also no other optimal algorithm that
uses the same heuristic would expand fewer nodes than it[23]. In this case,
reducing the estimation error between a consistent heuristic and the optimal
heuristic will only decrease the number of expanded nodes. Although A* can
perform much better than other search algorithms, it has a main issue that its
space complexity grows exponentially with the increasing search depth [23].
Consequently, finding a consistently good estimate of the optimal heuristic
and improving the memory efficiency become the two essential approaches to
improve the performance of A* in large-scale systems, which is also the
intuition behind the strategy we will introduce in the following section.

##  III Deep Learning Accelerated Heuristic Best-First Search

In this section, we will propose a general heuristic tree search algorithm,
which significantly speeds up the heuristic BeFS algorithm by employing a DNN
to estimate the optimal heuristic. After that, we will introduce the training
strategy and discuss its computational complexity.

###  III-A Proposed Deep Learning Based Search Strategy

Let $h(\bm{x}^{k}|\bm{z},\bm{R},\bm{\theta})$ denote the DNN parameterized by
the trainable parameters $\bm{\theta}$. For convenience, we will simplify it
as $h_{\bm{\theta}}(\bm{x}^{k})$ in the sequel. Accordingly, we define the
corresponding $f$-cost as

$\displaystyle
f_{\bm{\theta}}(\bm{x}^{k})=g(\bm{x}^{k})+h_{\bm{\theta}}(\bm{x}^{k}).$ (19)

In our implementation, we use a fully-connected neural network (FCNN) composed
by $L$ fully-connected layers, and the number of neurons at the $l$-th ($1\leq
l\leq L$) layer is denoted by $n_{l}$. Formally, the output of the $l$-th
layer can be expressed as

$\displaystyle\bm{p}^{l}=\max\left\\{\bm{0},\bm{W}^{l}\bm{p}^{l-1}+\bm{b}^{l}\right\\},$
(20)

where $\bm{p}^{l-1}\in\mathbb{R}^{n_{l-1}}$ is the output of the prior layer.
Notations $\bm{W}^{l}\in\mathbb{R}^{n_{l}\times n_{l-1}}$ and
$\bm{b}^{l}\in\mathbb{R}^{n_{l}}$ are the learnable weight matrix and bias
vector, respectively. Note that the $\max\\{\cdot,\cdot\\}$ operation is
performed in component-wise manner. In particular, the input of the first
layer is given by

$\displaystyle\bm{p}^{0}$
$\displaystyle=\bm{z}-\bm{R}\begin{bmatrix}\bm{0}_{(m-k)\times 1}\\\
\bm{x}^{k}\end{bmatrix}.$ (21)

In this case, the input size and output size are both fixed, which are set as
$n_{0}=m$ and $n_{L}=1$, respectively. Hence, the structure of FCNN can be
summarized as

$\displaystyle\left\\{m,n_{1},n_{2},\cdots,n^{l},\cdots,n^{L-1},1\right\\},$
(22)

and the set of trainable parameters is given by

$\displaystyle\bm{\theta}=\left\\{\bm{W}^{1},\bm{b}^{1},\cdots,\bm{W}^{l},\bm{b}^{l},\cdots,\bm{W}^{L},\bm{b}^{L}\right\\}.$
(23)

We shall emphasize that the DNN’s structure is not limited to FCNN. In
practice, we can certainly employ some more complicated or specialized
structures such as convolutional neural network (CNN) and recursively neural
network (RNN). However, the reason why we use FCNN here is only for validation
of the idea, and the DNN’s structure is not the focus of this work.

To leverage the DNN based heuristic function, we can combine it with heuristic
BeFS search algorithms. Since A* may require prohibitively large memory space
in large-scale problems, we will use the SMA* algorithm [21] as the underlying
search strategy. As a variant of A*, SMA* is able to perform BeFS with limited
memory space, and it is still equivalent to A* with enough memory size. By
combining SMA* with the aforementioned deep heuristic function, we thus
propose a hyper-accelerated tree search (HATS) algorithm for the optimal
signal detection in large-scale MIMO problems. The pseudo code of HATS is
detailed in Algorithm 2, and the associated utility functions are presented in
Algorithm 3. In general, HATS works just like the A* algorithm. It keeps
expanding the deepest least-$f$-cost node until ACTIVE is full, and the
$f$-cost of each generated successor is computed according to (19). In
particular, since the memory size is limited, HATS will forget the most
unpromising node from ACTIVE, and remember its key information in its parent.
Thus, the memory size can be bounded. After safely deleting the most
unpromising node, the proposed algorithm is able to move forward, and recover
back once there is no any other path better than the forgotten paths. Based on
this strategy, HATS maintains a partially expanded sub-tree of the whole tree.
Hence, the most unpromising node shall be the shallowest highest $f$-cost leaf
node. It should be noted that a node is said to be a “leaf” node in terms of
the partially expanded sub-tree. Thus, the forgotten leaf node is not
necessary to be the goal node. While maintaining the partially expanded sub-
tree, the proposed algorithm will recursively adjust the expanded node’s cost
according to the costs of its successors, and thereby the sub-tree is updated.
In conclusion, at each iteration, the proposed algorithm expands the best node
from memory, generates one successor and inserts the generated successor to
ACTIVE at the time, and deletes the worst leaf node when the memory is full.

Input: $\bm{z}$ and $\bm{R}$

Output: $\hat{\bm{x}}$ (estimate of the transmitted signal $\bm{x}$)

1 Insert the root node $\bm{x}^{0}$ into ACTIVE ;

2 loop

3 if  _ACTIVE is empty_ then return $\hat{\bm{x}}=\varnothing$ with failure;

4 $\bm{x}^{k}\leftarrow$ _deepest least- $f$-cost node_ in ACTIVE ;

5 if  _$\bm{x}^{k}$ is a goal node_ then return $\hat{\bm{x}}=\bm{x}^{k}$ with
success;

6 $\bm{x}^{k+1}\leftarrow$ next not-generated successor or best forgotten
successor of $\bm{x}^{k}$;

7 Insert $\bm{x}^{k+1}$ into its parent’s generated successor list;

8 if  _$\bm{x}^{k+1}$ is not a forgotten node_ then

9 if  _$\bm{x}^{k+1}$ is not a goal node and is at maximum depth_ then

10 $f(\bm{x}^{k+1})\leftarrow\infty$;

11

12 else

13
$f(\bm{x}^{k+1})\leftarrow\max\left\\{f(\bm{x}^{k}),f_{\bm{\theta}}(\bm{x}^{k+1})\right\\}$;

14

15 end if

16

17 else

18 Recover $\bm{x}^{k+1}$’s $f$-cost from its parent;

19

20 end if

21 TryAdjust(_$\bm{x}^{k}$_);

22 TryMakeSpace(__) ;

23 Insert $\bm{x}^{k+1}$ in ACTIVE ;

24 if  _all of $\bm{x}^{k}$’s successors are in memory_ then Remove
$\bm{x}^{k}$ from ACTIVE ;

25

26 end

27

Algorithm 2 Proposed HATS Algorithm

1 Function  _TryAdjust(_$\bm{x}^{k}$_)_:

2 if  _all of $\bm{x}^{k}$’s successors are generated_ then

3 $\bm{x}^{k+1}\leftarrow$ least-$f$-cost successor among all generated
successors and forgotten successors of $\bm{x}^{k}$;

4 if  _$\bm{x}^{k+1}$ ’s $f$-cost is valid and not equal to $\bm{x}^{k}$’s
$f$-cost_ then

5 Update $\bm{x}^{k}$’s $f$-cost to $\bm{x}^{k+1}$’s $f$-cost;

6 TryAdjust(_$\bm{x}^{k}$ ’s parent_);

7

8 end if

9

10 end if

11

12 end

13Function  _TryMakeSpace(__)_ :

14 if  _ACTIVE is not full_ then return ;

15 Remove _shallowest highest $f$-cost leaf node_ $\bm{x}^{k}_{j}$ from ACTIVE
;

16 Remove $\bm{x}^{k}_{j}$ from its parent’s generated successor list;

17 Remember $\bm{x}^{k}_{j}$’s key information (e.g. $f$-cost) in its parent’s
forgotten successor list;

18 if  _the parent is not in ACTIVE_ then

19 Insert the parent into ACTIVE ;

20 TryMakeSpace(__) ;

21

22 end if

23

24 end

Algorithm 3 Utility Functions for HATS

###  III-B Training Strategy

It is clear that the training strategy is to minimize the average error
between the estimated heuristic and the optimal heuristic. To achieve this
goal, we first introduce the following theorem regarding the optimal
heuristic.

######  Theorem 1.

With the optimal heuristic, the $f$-costs of the nodes on the shortest path of
each sub-tree are all equal to the cumulative cost of the deepest node on that
path.

###### Proof.

Suppose that $\bm{x}^{k}_{\phi}\rightarrow\bm{x}^{m}_{\phi}$ ($0\leq k\leq m$)
is the shortest path on the sub-tree of $\bm{x}^{k}_{\phi}$, then

$\displaystyle\bm{x}^{m}_{\phi}=\mathop{\arg\min}_{\bm{x}^{m}\supseteq\bm{x}^{k}_{\phi}}g(\bm{x}^{m}).$
(24)

Following the definition of the optimal heuristic described in (17), it is
obvious that

$\displaystyle
h^{*}(\bm{x}^{j}_{\phi})=g(\bm{x}^{m}_{\phi})-g(\bm{x}^{j}_{\phi}),$ (25)

holds for each node $\bm{x}^{j}_{\phi}$ ($k\leq j\leq m$) on the shortest path
of the sub-tree. Thus, this theorem is proven by the fact that

$\displaystyle f^{*}(\bm{x}^{j}_{\phi})$
$\displaystyle=g(\bm{x}^{j}_{\phi})+h^{*}(\bm{x}^{j}_{\phi})$ (26)
$\displaystyle=g(\bm{x}^{j}_{\phi})+g(\bm{x}^{m}_{\phi})-g(\bm{x}^{j}_{\phi})$
(27) $\displaystyle=g(\bm{x}^{m}_{\phi}).$ (28)

∎

Following Theorem 1, one straightforward training strategy is to minimize the
average $\ell_{2}$ loss

$\displaystyle\mathcal{L}(\bm{\theta}|\mathcal{D})=\mathbb{E}_{\mathcal{D}}\left\\{\left\lvert
g(\bm{x}^{m}_{\phi})-f_{\bm{\theta}}(\bm{x}^{k}_{\phi})\right\rvert^{2}\right\\},$
(29)

where the empirical expectation is taken on the data set $\mathcal{D}$, and
$\mathcal{D}$ can be generated by traversing the shortest path at each time
slot. Specifically, the data set of $T$ time slots can be expressed as

$\displaystyle\mathcal{D}=\left\\{S^{1}_{1},\cdots,S^{m-1}_{1},\cdots,S^{1}_{t},\cdots,S^{k}_{t},\cdots,S^{m-1}_{T}\right\\},$
(30)

where $S^{k}_{t}$ denotes the $k$-th ($1\leq k\leq m-1$) sample drawn at time
slot $t$ ($1\leq t\leq T$). Specifically, $S^{k}_{t}$ includes

$\displaystyle\mathcal{S}^{k}_{t}=\left\\{\bm{z},\bm{R},\bm{x}^{k}_{\phi},\bm{x}^{m}_{\phi}\right\\},$
(31)

where $\bm{z}$, $\bm{R}$, $\bm{x}^{m}_{\phi}$ are the associated received
signal, triangular matrix and ML estimate at the time. Notation
$\bm{x}^{k}_{\phi}$ is the $k$-th node on the shortest path
$\bm{x}^{0}\rightarrow\bm{x}^{m}_{\phi}$. Base on the data set, we can
optimize the network parameters $\bm{\theta}$ by employing mini-batch gradient
descent methods (e.g. the SGD optimizer [31] and ADAM optimizer [32]) to find
the optimal parameters as

$\displaystyle\bm{\theta}^{*}=\mathop{\arg\min}_{\bm{\theta}}\mathcal{L}(\bm{\theta}|\mathcal{D}).$
(32)

Nevertheless, we still have to compute the ML estimate, which will slow down
the sampling efficiency during training. To tackle this problem, we can use
the transmitted vector $\bm{x}$ rather than the ML estimate
$\bm{x}^{m}_{\phi}$. It is a reasonable approximation in the sense of signal
estimation, since the ML estimate is most likely to be $\bm{x}$, especially
for high SNR regimes. In summary, the advanced training procedure for the
proposed algorithm is detailed in Algorithm 4.

/* generate the training set */

1 Set the number of mini-batches as $B$;

2 Set the number of time slots for each mini-batch as $T$; initialize the data
set $\mathcal{D}_{all}=\emptyset$;

3 for  _$j=1,2,\cdots,B$_ do

4 for  _$t=1,2,\cdots,T$_ do

5 Randomly collect $\bm{z}$, $\bm{R}$ and $\bm{x}$ with random SNR;

/* traverse the approximatly shortest path of the whole tree */

6 for  _$k=1,2,\cdots,m$_ do

7 Select the $k$-th node $\bm{x}^{k}$ on the path $\bm{x}$;

8 Collect $S^{k}_{t}=\\{\bm{z},\bm{R},\bm{x},\bm{x}^{k}\\}$;

9 $\mathcal{D}_{all}=\mathcal{D}_{all}\cup\\{S^{k}_{t}\\}$;

10

11 end for

12

13 end for

14

15 end for

/* training with random samples */

16 Randomly initialize the DNN’s parameters $\bm{\theta}$;

17 for  _$j=1,2,\cdots,B$_ do

18 Randomly select a mini-batch of samples $\mathcal{D}$ from
$\mathcal{D}_{all}$;

19 Compute the average loss on $\mathcal{D}$ according to (29);

20 Perform a stochastic gradient decent step on the network’s parameters
$\bm{\theta}$;

21

22 end for

23Save the well-trained parameters $\bm{\theta}$;

Algorithm 4 Training Procedure for HATS

###  III-C Computational Complexity

The computational complexity of HATS depends on both the number of visited
nodes and the visitation costs. While visiting a node $\bm{x}^{k}$, the
computational complexity is

$\displaystyle\mathcal{O}\left(k+\sum_{l=1}^{L}n_{l}n_{l-1}+n_{l}\right),$
(33)

where the first term denotes the complexity of computing $g(\bm{x}^{k})$ and
the second term denotes the complexity of computing
$h_{\bm{\theta}}(\bm{x}^{k})$. When the DNN estimates the optimal heuristic
precisely, the algorithm visits the fewest nodes. Thus, the lower bound on the
average complexity of HATS is given by

$\displaystyle\mathcal{O}\left(m+m\left(\sum_{l=1}^{L}n_{l}n_{l-1}+n_{l}\right)\right),$
(34)

since only the $m$ nodes lying on the shortest path would be eventually
expanded. When the estimation is imperfect, the optimality and the optimal
efficiency can not be guaranteed, and the complexity may grow exponentially
with the problem scale in the worst case. However, as we have discussed in
Sec. II-C, decreasing the estimation error will significantly improve the BER
performance and search speed. Therefore, the performance and complexity of
HATS rely on both the estimation quality of DNN and the system’s SNR. In fact,
as we will show in the simulations, the proposed algorithm achieves almost the
optimal BER performance, while it reaches almost the lowest complexity in low
SNR regimes. That is, with the proposed training strategy, the network is able
to predict the optimal heuristic accurately. Thus, HATS can often reach the
optimal efficiency under practical scenarios, which indicates that the
proposed algorithm is applicable in large-scale systems.

##  IV Simulation Results

In this section, we will present simulation results and discussions to show
the effectiveness of the proposed algorithm. Specifically, we will first
introduce the environment setup and the implementation details of our model.
Then, we will introduce the competing algorithms and discuss the related
simulation results afterwards.

###  IV-A Environment Setup

In the simulations, we consider a MIMO system model where there are $N_{t}$
and $N_{r}$ antennas at the transmitter and receiver, respectively. Therefore,
we have $n=2N_{r}$ and $m=2N_{t}$. In addition, the signal is modulated by
QPSK modulation, and the transmission experiences a Rayleigh flat fading
channel. Moreover, the channel information can be known perfectly at the
receiver. As to the DNN structure, we use the same DNN structure for all
experiments, where there are $4$ hidden layers in total and the neuron numbers
of the four hidden layers are set to $128$, $64$, $32$, and $16$,
respectively. According to (22), the structure of DNN can be summarized as

$\displaystyle\left\\{2N_{t},128,64,32,16,1\right\\}.$ (35)

To train the model, we employ the Adam optimizer [32] with the learning rate
of $10^{-6}$. The mini-batch size is set to $128$ time slots, and the total
batch size is $10$ million.

###  IV-B Competing Algorithms

To verify the effectiveness of the proposed algorithm, we compare the proposed
algorithm with several competitive algorithms. Before discussing the results,
we first introduce the following abbreviations,

  * •

_MMSE_ : The conventional minimum mean squared error estimator.

  * •

_OAMP-Net2_ : The orthogonal approximate message passing network 2 introduced
in [27].

  * •

_SD_ : The sphere decoding algorithm introduced in [1].

  * •

_DL-SD_ : The deep learning based sphere decoding algorithm introduced in
[29].

  * •

_A*_ : The A* algorithm introduced in [21]. Note that the heuristic is set to
zero constantly for comparison.

  * •

_HATS( $M$)_: The proposed algorithm in this paper, where $M$ denotes the
memory capacity.

Among these algorithms, OAMP-Net2, DL-SD and the proposed HATS algorithm are
deep learning based algorithms. Besides, SD, DL-SD, A* and HATS are search
algorithms. In particular, DL-SD requires linear space with the search depth,
while A* has exponential space complexity. Moreover, we will take simulations
for HATS with different memory sizes. In particular, HATS($\infty$) denotes
that the memory size is unlimited. In this case, HATS($\infty$) is equivalent
to A* except that the heuristic is estimated by a DNN. Since SD, DL-SD and A*
are optimal search algorithms, we will simply denote their BER performances as
“ML” for convenience. Unlike the existing literature that measures the
complexity in terms of the number of visited lattice points (goal nodes) [29],
we will compare the complexities of search algorithms based on the number of
visited nodes in the simulations, since it is more accurate to measure the
complexity.

###  IV-C BER Performance and Complexity Comparison

Figs. 2-4 illustrate the BER comparisons of the aforementioned algorithms in
different-scale MIMO systems, where the numbers of antennas are $8$, $16$,
$32$ at both the transmitter and receiver, respectively. In addition, the SNR
varies from $5$ dB to $15$ dB. Note that algorithms have to search the closest
lattice point among $2^{64}$ candidates to reach the exact optimal BER
performance in $32\times 32$ MIMO, which is a very big challenge. In
particular, the BERs of “ML” in the SNR region ranging from $5$ dB to $10$ dB
are not plotted in Fig. 4, since the computational complexity becomes
prohibitively large in the regime. From these figures, we can find that HATS
achieves almost the optimal BER performance in all three systems. When
$\text{SNR}=15$ dB, HATS produces only slightly $10\%$, $10\%$ and $15\%$ more
error than the optimal search algorithms in $8\times 8$, $16\times 16$ and
$32\times 32$ MIMO systems, respectively. As a contrast, for the sub-optimal
algorithms like MMSE and OAMP-Net2, their BER performances are significantly
far from the optimal one in all three MIMO systems, especially when the
problem scale becomes prohibitively large. Moreover, we can also find from
these figures that the performance gaps between HATS and the other sub-optimal
algorithms enlarge with the increasing SNR. In particular, when
$\text{SNR}=15$ dB, HATS reduces the error of MMSE and OAMP-Net2 to only about
$0.0052\%$ and $1.504\%$ in $32\times 32$ MIMO, respectively. Besides, the SNR
gains of HATS over OAMP-Net2 are about $3.2$ dB, $2.1$ dB and $1.8$ dB at the
BER levels of $10^{-3}$, $10^{-4}$, and $10^{-5}$ in $8\times 8$, $16\times
16$ and $32\times 32$ MIMO systems, respectively. In conclusion, these results
show that the proposed algorithm is robust and almost optimal even in large-
scale problems.

In order to show the impact of SNR and memory constraint on the computational
complexity of the proposed algorithm, Figs. 5-7 are provided to show the
complexity comparison under the same setup as Figs. 2-4. From the results in
Figs. 5-7, we can conclude that with difference memory sizes, HATS still
visits much fewer nodes than SD, DL-SD and A*. The numbers of nodes visited by
SD, DL-SD and A* all grow exponentially with the decreasing SNR, which
indicates that they are not applicable in large-scale problems. As a contrast,
the number of nodes visited by HATS increases almost linearly in low SNR
regimes. Moreover, we can find from these figures that HATS can reach the
lowest complexity in a wide range of acceptable SNRs, whereas the complexities
of SD, DL-SD and A* are still very high in the regime. Specifically, for the
$32\times 32$ MIMO where the size of search space is $2^{64}$ and
$\text{SNR}=18$ dB, SD, DL-SD and A* have to visit average about $1800$,
$1100$ and $1000$ nodes to find the optimal solution, while HATS only needs to
visit about average $150$ nodes. In particular, we can also find that HATS is
not sensitive to the memory size, and it can perform well with bounded memory
size. Specifically, when the memory size is bounded, the performances of
HATS($128$), HATS($1024$) and HATS($\infty$) are almost the same in all the
three systems. This indicates that HATS can reach almost the optimal
efficiency with bounded memory in practice.

Fig. 8 illustrates the computational complexity comparison of the competing
algorithms versus the problem scale, where the number of antennas varies from
$8$ to $24$ and the associated SNR is set to $15$ dB. From this figure, we can
observe that the complexities of HATS($128$), HATS($1024$) and HATS($\infty$)
increase almost linearly with the number of antennas, while the complexities
of SD, DL-SD, A* increase exponentially with the problem scale. This implies
that HATS is much more efficient and less sensitive to the problem scale by
comparing to the other optimal search algorithms. As it has been shown in
Figs. 2-7 that HATS can reach almost the optimal BER performance in large-
scale systems, one is able to conclude that the proposed algorithm almost
meets the optimal efficiency in large-scale problems, which indeed further
verifies the effectiveness of the proposed algorithm.

##  V Conclusions

In this paper, we investigated the optimal signal detection problem with focus
on large-scale MIMO systems. The problem can be regarded as search on the
decision tree, and the optimal solution is obtained by finding the shortest
path. In order to improve the search speed while does not compromise the
optimality, we proposed a deep learning based heuristic search algorithm,
namely, HATS. The proposed algorithm was shown to be almost optimally
efficient, since it can reach almost the optimal BER performance while still
visit almost the fewest nodes in large-scale systems. Hence, we do believe
that the proposed algorithm is attractive for optimal signal detection in
practical large-scale MIMO systems, especially when the problem scales largely
and the optimal performance is required.

## References

  * [1] B. Hassibi and H. Vikalo, “On the sphere-decoding algorithm I. expected complexity,” _IEEE Trans. Sig. Proc._ , vol. 53, no. 8-1, pp. 2806–2818, 2005.
  * [2] N. Sommer, M. Feder, and O. Shalvi, “Closest point search in lattices using sequential decoding,” in _Proc. IEEE International Symposium on Information Theory (ISIT)_ , 2005, pp. 1053–1057.
  * [3] C. M. Bishop, _Pattern recognition and machine learning, 5th Edition_ , ser. Information science and statistics. Springer, 2007.
  * [4] P. J. T., “Integer least-squares theory for the GNSS compass,” _Journal of Geodesy_ , vol. 84, no. 7, pp. 433–447, 2010.
  * [5] J. Goldberger and A. Leshem, “A Gaussian tree approximation for integer least-squares,” _Proc. Neural Information Processing Systems (NeurIPS)_ , vol. 22, pp. 638–645, 2009.
  * [6] E. Agrell, T. Eriksson, A. Vardy, and K. Zeger, “Closest point search in lattices,” _IEEE Trans. Inf. Theory_ , vol. 48, no. 8, pp. 2201–2214, 2002\.
  * [7] W. Zhao and G. B. Giannakis, “Reduced complexity closest point decoding algorithms for random lattices,” _IEEE Trans. Wirel. Commun._ , vol. 5, no. 1, pp. 101–111, 2006.
  * [8] D. Micciancio, “The hardness of the closest vector problem with preprocessing,” _IEEE Trans. Inf. Theory_ , vol. 47, no. 3, pp. 1212–1215, 2001.
  * [9] W. Chen, L. Peng, J. Wang, F. Li, M. Tang, W. Xiong, and S. Wang, “An improved lower bound for approximating the minimum integral solution problem with preprocessing over $\ell_{\infty}$ norm,” _J. Comb. Optim._ , vol. 30, no. 3, pp. 447–455, 2015.
  * [10] L. Liu, Y. Li, C. Huang, C. Yuen, and Y. L. Guan, “A new insight into GAMP and AMP,” _IEEE Trans. Veh. Technol._ , vol. 68, no. 8, pp. 8264–8269, 2019.
  * [11] Z. Guo and P. Nilsson, “Algorithm and implementation of the k-best sphere decoding for MIMO detection,” _IEEE J. Sel. Areas Commun._ , vol. 24, no. 3, pp. 491–503, 2006.
  * [12] L. Liu, C. Liang, J. Ma, and L. Ping, “Capacity optimality of AMP in coded systems,” _IEEE Trans. Inf. Theory_ , vol. 67, no. 7, pp. 4429–4445, 2021\.
  * [13] S. Bäro, J. Hagenauer, and M. Witzke, “Iterative detection of MIMO transmission using a list-sequential (LISS) detector,” in _Proc. IEEE International Conference on Communications (ICC)_ , 2003, pp. 2653–2657.
  * [14] Z. Yang, C. Liu, and J. He, “A new approach for fast generalized sphere decoding in MIMO systems,” _IEEE Sig. Proc. Lett._ , vol. 12, no. 1, pp. 41–44, 2005.
  * [15] A. D. Murugan, H. E. Gamal, M. O. Damen, and G. Caire, “A unified framework for tree search decoding: rediscovering the sequential decoder,” _IEEE Trans. Inf. Theory_ , vol. 52, no. 3, pp. 933–953, 2006.
  * [16] T. Cui, T. Ho, and C. Tellambura, “Heuristic tree search for detection and decoding of uncoded and linear block coded communication systems,” in _Proc. IEEE International Conference on Communications (ICC)_ , 2006, pp. 391–396.
  * [17] Y. Dai and Z. Yan, “Memory-constrained tree search detection and new ordering schemes,” _IEEE J. Sel. Top. Sig. Proc._ , vol. 3, no. 6, pp. 1026–1037, 2009.
  * [18] R. Y. Chang, W. Chung, and S. Lin, “A* algorithm inspired memory-efficient detection for MIMO systems,” _IEEE Wirel. Commun. Lett._ , vol. 1, no. 5, pp. 508–511, 2012.
  * [19] S. Yang and L. Hanzo, “Fifty years of MIMO detection: The road to large-scale mimos,” _IEEE Commun. Surv. Tutorials_ , vol. 17, no. 4, pp. 1941–1988, 2015.
  * [20] N. T. Nguyen, K. Lee, and H. Dai, “Application of deep learning to sphere decoding for large MIMO systems,” _IEEE Transactions on Wireless Communications_ , pp. 1–1, 2021.
  * [21] S. J. Russell, “Efficient memory-bounded search methods,” in _Proc. European Conference on Artificial Intelligence (ECAI)_ , 1992, pp. 1–5.
  * [22] Y. S. Han, C. R. P. Hartmann, and C. Chen, “Efficient priority-first search maximum-likelihood soft-decision decoding of linear block codes,” _IEEE Trans. Inf. Theory_ , vol. 39, no. 5, pp. 1514–1523, 1993.
  * [23] R. Dechter and J. Pearl, “Generalized best-first search strategies and the optimality of A*,” _J. ACM_ , vol. 32, no. 3, pp. 505–536, 1985.
  * [24] H. Kim, S. Oh, and P. Viswanath, “Physical layer communication via deep learning,” _IEEE J. Sel. Areas Inf. Theory_ , vol. 1, no. 1, pp. 5–18, 2020.
  * [25] J. Xia, K. He, W. Xu, S. Zhang, L. Fan, and G. K. Karagiannidis, “A MIMO detector with deep learning in the presence of correlated interference,” _IEEE Trans. Veh. Technol._ , vol. 69, no. 4, pp. 4492–4497, 2020.
  * [26] K. He, L. He, L. Fan, Y. Deng, G. K. Karagiannidis, and A. Nallanathan, “Learning-based signal detection for MIMO systems with unknown noise statistics,” _IEEE Trans. Commun._ , vol. 69, no. 5, pp. 3025–3038, 2021\.
  * [27] H. He, C. Wen, S. Jin, and G. Y. Li, “Model-driven deep learning for MIMO detection,” _IEEE Trans. Sig. Proc._ , vol. 68, pp. 1702–1715, 2020.
  * [28] A. Askri and G. R. Othman, “DNN assisted sphere decoder,” in _IEEE International Symposium on Information Theory (ISIT)_ , 2019, pp. 1172–1176.
  * [29] M. Mohammadkarimi, M. Mehrabi, M. Ardakani, and Y. Jing, “Deep learning-based sphere decoding,” _IEEE Trans. Wirel. Commun._ , vol. 18, no. 9, pp. 4368–4378, 2019.
  * [30] J. Sun, Y. Zhang, J. Xue, and Z. Xu, “Learning to search for MIMO detection,” _IEEE Trans. Wirel. Commun._ , pp. 1–1, 2020.
  * [31] J. Werfel, X. Xie, and H. S. Seung, “Learning curves for stochastic gradient descent in linear feedforward networks,” in _Proc. Neural Information Processing Systems (NeurIPS)_ , 2003, pp. 1197–1204.
  * [32] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” in _Proc. International Conference on Learning Representations (ICLR)_ , 2015\.

Figure 2: BER comparison versus SNR for $8\times 8$ MIMO. Figure 3: BER
comparison versus SNR for $16\times 16$ MIMO. Figure 4: BER comparison versus
SNR for $32\times 32$ MIMO. Figure 5: Complexity comparison versus SNR for
$8\times 8$ MIMO. Figure 6: Complexity comparison versus SNR for $16\times 16$
MIMO. Figure 7: Complexity comparison versus SNR for $32\times 32$ MIMO.
Figure 8: Complexity comparison versus number of antennas ($\text{SNR}=15$
dB).