Learning-to-Transition for Large-scale and High-Order MIMO Detection

Consider a point-to-point spatial-multiplexing MIMO channel with $N_{t}$ transmit antennas and $N_{r}$ receive antennas. For each channel use,

where $\bm{H}\in\mathbb{C}^{N_{r}\times N_{t}}$ is the channel matrix, following the distribution $\bm{H}\sim p_{H}$, and $\sigma^{2}$ is the noise variance at each receive antenna. Let $\mathcal{A}_{Q}$ denote a unit-energy, Gray-coded quadrature-amplitude modulation (QAM) constellation with order $Q$. Each QAM symbol contains $B=\log_{2}Q$ coded bits, with $B/2$ bits on each pulse-amplitude modulation (PAM) axis. The corresponding constellation mapping is depicted in Fig. 1, demonstrating how QAM scheme exploits the two-dimensional space to separate modulated symbols.

Let $\mu_{Q}:\{0,1\}^{B}\rightarrow\mathcal{A}_{Q}$ be the prescribed Gray mapping. We collect the bits assigned to all transmit antennas and define the modulated symbol as follows:

where $\bm{b}=[\bm{b}_{1}^{\mathsf{T}},\ldots,\bm{b}_{N_{t}}^{\mathsf{T}}]^{\mathsf{T}}$ belongs to $\mathcal{B}=\{0,1\}^{N_{t}B}$. Traditional maximum-likelihood (ML) detection then adopts the following exhaustive search:

Note that the entire space now contains $2^{N_{t}B}$ feasible solutions, hence the ML detection in (3) becomes computationally prohibitive under the large system size (i.e., $N_{t}$ increases) and high-order QAM (i.e., $B$ increases). Moreover, it is inefficient to repeatedly solve similar detection problems as the channel $\bm{H}$ frequently varies.

To tackle these issues, we aim to train a single neural network that solves the symbol detection problems for channel instances across the entire distribution $\bm{H}\sim p_{H}$. This follows the essence of learning-to-optimize (L2O)/ amortized optimization, which shifts the computational burden from online iterative optimization to offline training. Instead of solving each problem instance from scratch, L2O learns an optimizer from the data, exploiting recurring task structures for fast inference. Specifically, denote the residual function of the symbol detection as follows:

Then the L2O solver aims to minimize this metric in the population level: let $\hat{\bm{b}}_{\bm{\theta}}(\bm{y},\bm{H})$ denotes its final output, where $\bm{\theta}$ denotes the learnable parameters of the neural solver, the amortized learning problem becomes

This formulation remains challenging, since symbol detection is inherently an NP-hard combinatorial optimization problem, which lacks a high-quality general-purpose solver.

Based on the objective in (5), we seek to train a neural network that learns an effective transition policy over the entire solution space $\mathcal{B}$. Similar to the prior graph optimization work, learning an expressive channel embedding is crucial to the detection performance in this work. Note that the residual function in (4) can be rewritten as

where $\bm{x}\triangleq\mathcal{M}_{Q}(\bm{b})\in\mathbb{C}^{N_{t}}$ represents the modulated symbol, $\bm{R}\triangleq\bm{H}^{{\mathsf{H}}}\bm{H}\in\mathbb{R}^{N_{t}\times N_{t}}$ represents the Gram/correlation matrix of channel $\bm{H}$, and $\bm{z}\triangleq\bm{H}^{{\mathsf{H}}}\bm{y}\in\mathbb{C}^{N_{t}}$ represents the corresponding matched filter output. Inspired by [13], the effective channel features can be learned by drawing an analogy with graph-embedding generation. Specifically, the matched-filter entries, i.e., $\bm{z}$, and all diagonal entries of Gram matrix, i.e., $\text{diag}(\bm{R})$, form the node-wise inputs (each node relates to one transmit antenna); while all off-diagonal entries of Gram matrix, i.e., $\text{offdiag}(\bm{R})$, form the edge-wise inputs (which reflects the complex inter-stream couplings). Then a graph neural network (GNN) encoder $G_{\bm{\varphi}}$ maps these sufficient input statistics to the initial channel embedding, i.e.,

where $D$ is the feature dimension of each node (transmit side).

To reduce the difficulty of learning the transition policy, inspired by the recursive optimization process in [11, 10], we consider a Learning-to-transition (L2T) scheme consisting of $T$ updates, therefore amortizing the searching task across all steps. Importantly, a static embedding cannot reflect which parts of the search space have already been explored by the evolving candidate. We therefore keep updating the channel embedding across the transition steps, inspired by the works in [7, 16]. Assume that the transition process at each step is performed by a neural network with learnable parameters $\bm{\theta}_{0}$. Hence the forward process at each transition step $t$ is given by

for $t=1,\dots,T$. The first output undergoes a bit-wise sampling process to obtain the next solution, i.e., $\bm{b}^{(t)}\sim\bm{p}^{(t)}$, while the second output is the updated feature embedding at transition step $t$.

The transition policy $\Pi_{\bm{\theta}_{0}}$ in (8) can be parameterized by a Transformer network [12]. In particular, the previous sampled solution is reshaped as $\bm{b}^{(t-1)}=[\bm{b}^{(t-1)}_{1},\dots,\bm{b}^{(t-1)}_{N_{t}}]\in\mathbb{R}^{B\times N_{t}}$, and then concatenated with the corresponding channel embedding $\bm{Z}^{(t-1)}\in\mathbb{R}^{D\times N_{t}}$ in a top-down manner, which yields the input token sequence to the next layer:

Let $\bar{\bm{C}}^{(t)}=\text{LayerNorm}(\bm{C}^{(t)})$. For each attention head $h$, $h=1,\dots,H$, the query, key, and value projections are

The corresponding attention output is

The $H$ attention outputs are concatenated and projected to

where $\bm{W}_{o}\in\mathbb{R}^{(B+D)\times D_{v}}$, $\bm{Y}_{p}^{(t)}\in\mathbb{R}^{B\times N_{t}}$ represents the solution logits and $\bm{Y}_{f}^{(t)}\in\mathbb{R}^{D\times N_{t}}$ contains the channel embedding residuals. Then the solution at the next transition step is sampled as follows:

In parallel, the channel embedding is updated through residual connections and an MLP:

(III-A) and (III-A) are the two output branches of the same Transformer layer. Unrolling this layer over $T$ transitions recursively refines both the discrete solution and the continuous embedding.

In principle, the mean-field sampling for each solution (i.e., every bit is sampled in parallel in the process $\bm{b}^{(t)}\sim\bm{p}^{(t)}$) is fast but ignores conditional correlations among $N$ bits. At the other extreme, bit-by-bit autoregressive sampling captures every conditional dependence but requires a large mount of sequential forward processes. We therefore introduce blockwise autoregressive sampling [17] as an efficient tradeoff between these two extreme cases, to generate the complete solution $\bm{b}^{(t)}$ at each transition step. In this work, each block represents a $B$-bit QAM symbol from each transmit antenna. The dependence is retained across those $N_{t}$ symbols because the dense off-diagonal entries in the Gram matrix $\bm{R}$ couple the symbol detections at different antennas. In contrast, the bits within each QAM symbol are generated in parallel. This mean-field factorization avoids a bit-by-bit chain of length $N_{t}B$ while preserving the dominant inter-symbol dependence.

The block ordering $n=1,2,\dots,N_{t}$ may be randomized during training to avoid imposing a persistent ordering bias. Define $\bm{b}_{\text{mask}}$ as an all-zero vector with length $B$, which can be regarded as each undetermined data symbol during a sequential sampling process. Then we define the following partial solution before sampling each data symbol $\bm{b}^{(t)}_{n}$ at transition step $t$:

Conditioned on $\bm{b}^{(t)}_{<n}$, the complete prior solution $\bm{b}^{(t-1)}$, and the current channel embedding $\bm{Z}^{(t)}_{n-1}$, the policy samples all variables in the next block in parallel:

with $\bm{Z}^{(t)}_{0}=\bm{Z}^{(t-1)}$ and $\bm{Z}^{(t)}=\bm{Z}^{(t)}_{N_{t}}$. After $N_{t}$ blockwise samplings, we obtain the complete solution $\bm{b}^{(t)}$. Note that the special case $M=1$ reduces to a mean-field sampler, whereas $M=N$ reduces to a bit-by-bit autoregressive sampler.

Implementing the above $T$-step stochastic transition process yields a complete L2T solution trajectory $\tau\triangleq(\bm{b}^{(0)},\bm{b}^{(1)},\ldots,\bm{b}^{(T)})$. Its sampling probability can be computed as

where $p(\bm{b}^{(0)})$ does not depend on $\bm{\theta}_{0}$. Denote the step-wise objective function as $g(\bm{b}^{(t)};\bm{y},\bm{H})$. Similar to the GCO paper, the following trajectory-wise objective function with entropy regularization is proposed:

Since $\mathbb{E}_{\tau}[\log p_{\bm{\theta}_{0}}(\tau\mid\bm{y},\bm{H})]$ is the negative entropy term, the coefficient $\lambda>0$ encourages explorations at the early stage. The corresponding distribution-level learning objective is then given as follows:

and the ultimate learning problem becomes

which is solved using a gradient-based method. Note that (19) contains two nested expectations. The outer expectation $\mathbb{E}_{\bm{H}\sim p_{H}}$ arises because we aim to learn a population-based solver whose policy generalizes across the channel distribution; The inner expectation $\mathbb{E}_{\tau\sim p_{\bm{\theta}_{0}}(\cdot\mid\bm{y},\bm{H})}$ accounts for the stochasticity of the learned transition policy because each trajectory is sampled from the probability distributions output by the network. For each detection instance $(\bm{y},\bm{H})$, this expectation is approximated by averaging over multiple sampled trajectories $\{\tau_{1},\dots,\tau_{K}\}$. Although the ultimate interest is $f(\bm{b}^{(T)};\bm{y},\bm{H})$ at the final step, aggregating the objective values along the trajectory encourages progressive improvement throughout the search process. At inference, for a new instance $(\bm{y},\bm{H})$, we expect the recursive application of the trained policy network to produce a near-optimal result $\bm{b}^{(T)}$. Since each update only requires a feedforward mapping and a sampling operation, the resulting solver can be executed repeatedly with low latency, making real-time deployment feasible.

We next show how the step-wise objective function $g(\bm{b}^{(t)};\bm{y},\bm{H})$ is constructed. Note that the ultimate goal is to reduce the detection bit-error-rate (BER), and the true transmitted bits are assumed to be known during training. Hence the BER-oriented binary cross-entropy (BCE) can be regarded as an effective step-wise loss function: specifically, let $N_{b}=N_{t}B$ denote the entire codeword length; if the ground truth is $b_{j}^{*}$ for every $j^{\text{th}}$ bit, while the network produces probabilities $p^{(t)}_{j}\triangleq P^{(r)}_{\bm{\theta}_{0}}(b_{j}=1)$ at transition step $t$, $j=1,\dots,N_{b}$, the BER-oriented BCE loss for this codeword is

In practice, the residual loss function in (4) is also utilized as a smooth surrogate in the early-stage training. Hence the following hybrid step-wise loss function is proposed

where the coefficient $\rho_{s}$ gradually decreases from 1 to 0, explicitly reflecting a schedule shifts from the smooth residual loss to the BER-oriented BCE loss. As a result, substituting (22) into (III-C) yields the trajectory-wise loss function.

According to (19), when computing the population-level loss, for each instance $(\bm{y},\bm{H})$, $K$ solution trajectories are sampled, which accounts for the expectation term $\mathbb{E}_{\tau\sim p_{\bm{\theta}_{0}}(\cdot\mid\bm{y},\bm{H})}$. The initializations of those $K$ trajectories (denoted as $\bm{b}_{1}^{(0)},\dots,\bm{b}_{K}^{(0)}$) are all randomly-generated valid solutions. Finally, the policy network $\Pi_{\bm{\theta}_{0}}$ is trained according to (20).

During inference, for each $(\bm{y},\bm{H})$, we recursively employ the trained policy network $\Pi_{\bm{\theta}_{0}^{*}}$ for $T$ transitions based on the random initializations $(\bm{b}_{1}^{(0)},\dots,\bm{b}_{K}^{(0)})$, and obtain the final solution set $(\bm{b}_{1}^{(T)},\dots,\bm{b}_{K}^{(T)})$. The hard decision result is the final solution $\bm{b}_{k^{*}}^{(T)}$ with the smallest residual loss, i.e.,

Iterative detection and decoding (IDD) exchanges soft information between the MIMO detector and a soft-output low-density parity-check (LDPC) decoder. An interleaver distributes one LDPC codeword over a batch of MIMO observations; after detection, their LLRs are simply reassembled in codeword order before decoding. We use the convention $L(b)=\log\left(\frac{P(b=0)}{P(b=1)}\right)$, so $P(b=1)=\operatorname{sigmoid}(-L(b))$. Let $\Gamma$ denote the code interleaver. At the $m^{\text{th}}$ IDD round, the complete turbo exchange uses the LDPC belief-propagation (BP) mapping $\mathcal{D}_{\mathrm{BP}}$, given as follows:

where the subscripts $(A,P,E)$ in the above equations represent the a-priori, posterior, and extrinsic LLR. The exchange is initialized with $\bm{L}^{\mathrm{A},(0)}=\bm{0}$. The two subtractions are essential: each module passes only information newly created from its own observation or parity constraints, rather than feeding its input prior back as if it were new evidence.

For a bit $b$, Bayes’ rule gives

where $\mathcal{F}$ contains the detection instance and previously generated bits. With $L^{\mathrm{A}}=\log\left(\frac{P^{\mathrm{A}}(0)}{P^{\mathrm{A}}(1)}\right)$ and equal-prior LLR output by the policy network:

the posterior LLR and the sampling probability are given as

Thus an a-priori LLR induces a bias into the sampling probabilities of the detector’s output. At $\bm{L}^{\mathrm{A}}=\bm{0}$, the policy reduces to the previous hard-decision detector.

After the final transition, let $q_{k,j}$ be trajectory $k$’s probability that bit $j$ is one and let $\bm{b}_{k}^{(T)}$ be its sampled vector. We use the exact normalized residual to combine the $K$ trajectories:

Because $f$ is normalized by $N_{r}\sigma^{2}$, the complex-Gaussian likelihood satisfies $p(\bm{y}\mid\bm{b},\bm{H})\propto\exp[-N_{r}f(\bm{b};\xi)]$. Consequently, $\tau_{\mathrm{w}}=1/N_{r}$ is the physically matched value in (31); any other value is an explicit likelihood-tempering choice rather than a new estimate of the noise variance. The detector posterior and extrinsic LLRs are

This is a Rao–Blackwellized terminal-mixture estimator: $q_{k,j}$ retains the final head’s bit uncertainty instead of replacing it by the sampled bit, while the exact likelihood favors candidates that fit the observation. Only the terminal Bernoulli is Rao–Blackwellized: the mixture weight $w_{k}$ is evaluated at the single sampled vector $\bm{b}_{k}^{(T)}$ rather than as an expectation under the terminal head, so (32) is a likelihood-weighted terminal-marginal mixture rather than a fully marginalized posterior; the two are equal when each trajectory is nearly deterministic.

The extrinsic subtraction in (33) is exact for a single mean-field trajectory. With $q_{k,j}=\operatorname{sigmoid}(a^{\mathrm{net}}_{k,j}-L^{\mathrm{A}}_{j})$, (29) gives $\log((1-q_{k,j})/q_{k,j})=L^{\mathrm{A}}_{j}-a^{\mathrm{net}}_{k,j}$, so the per-trajectory extrinsic $-a^{\mathrm{net}}_{k,j}$ carries only the network evidence and is free of the prior. After the residual-weighted mixture (32), however, the logit of a convex combination of prior-tilted marginals is no longer an affine function of $L^{\mathrm{A}}_{j}$, and the weights $w_{k}$ themselves depend on $L^{\mathrm{A}}_{j}$ through the sampled vectors. Consequently $L_{j}^{\mathrm{E}}=L_{j}^{\mathrm{P}}-L_{j}^{\mathrm{A}}$ removes the bit’s own prior only approximately, with a residual dependence that shrinks as the trajectories agree. We therefore measure this residual prior dependence, rather than assume it is zero, through the prior-sensitivity diagnostic in Section V-B.

The main training difficulty introduced by the soft IDD learning scheme is that the detector input prior (i.e., the feedback from LDPC decoder) at each round is not stationary. As detector parameters evolve across the iterations, the distribution of the prior $\bm{L}^{\mathrm{A},(m)}$ also drifts across IDD rounds. Note that the training scheme for the hard-decision receiver assumes equal prior for every transmitted bit and cannot adapt to the practical IDD settings; Moreover, starting from the random initializations (at all IDD rounds) usually renders unstable convergence and overconfident feedback. To address these issues, we propose the following multi-stage training scheme tailored to the soft IDD learning problem.

We consider a downlink MIMO system with perfect CSI available at the receiver. Each channel sample follows the Rayleigh distribution, with $[\bm{H}]_{ij}\overset{\text{i.i.d.}}{\sim}\mathcal{CN}(0,1/N_{t})$, and Gray-coded 64-QAM or 256-QAM symbols are normalized to unit average energy. Unless otherwise stated, the L2O framework uses $T=8$ transitions, embedding dimension $d=256$, eight attention heads, feed-forward dimension $d_{\mathrm{ff}}=256$, and $K=16$ trajectories. One trajectory is initialized by the LMMSE estimate and the remaining trajectories by its randomized perturbations. Residual-only post-processing is disabled so that the reported gain is attributable to the learned transition policy.

The hard detector is trained at 20 dB for $10^{4}$ AdamW updates with learning rate $10^{-4}$. The residual-to-BER curriculum uses $s_{0}=10^{3}$ and $s_{1}=5\times 10^{3}$, and the batch sizes are 64 and 16 for 64-QAM and 256-QAM, respectively. For soft reception, every unrolled detector layer is initialized from the same hard checkpoint according to (34) and then follows the synthetic-prior and decoder-in-the-loop stages in Section IV-C. Soft training uses $7\times 10^{3}$ updates, learning rate $3\times 10^{-5}$, and batch size 16, with $K=32$ for 64-QAM and $K=16$ for 256-QAM.

The coded link uses the 5G NR LDPC implementation in Sionna [22], code rate $1/2$, ten belief-propagation iterations, and three IDD rounds. Each codeword packs $P=8$ independently detected MIMO vectors, giving rate-matched codeword lengths of 384 and 512 bits for 64-QAM and 256-QAM, respectively. Every reported SNR point is averaged over at least five random seeds until either 200 frame errors or $10^{7}$ information bits are observed.

For the hard detector, we compare LMMSE, QR-domain K-best, OAMP-Net2, RE-MIMO, and SGT [2, 3, 4]. For soft IDD, the comparisons are soft LMMSE, MMSE-PIC, K-best list-MAP, SGT, and DUIDD [8]. All methods use the same channel realizations, modulation, code, decoder iterations, and SNR definition. Uncoded BER is the primary hard-decision metric. Post-decoding BLER is the primary soft-receiver metric, with information-bit BER, generalized mutual information (GMI), and Brier score used to diagnose the quality of the detector LLRs. Parameter count and batch-one latency are measured under identical hardware, software, precision, and synchronization settings; detector and decoder latency are reported separately.

We first isolate the residual-to-BER curriculum at one representative 64-QAM SNR near the BER waterfall. The proposed smooth schedule is compared with residual-only training, BCE-only training, and an abrupt residual-to-BCE switch. For each variant, we report final BER, normalized residual, and the standard deviation over five seeds. For the full model, BER, residual, and normalized Hamming distance between consecutive vectors are also plotted against the transition index. This compact experiment tests both roles of the curriculum: the residual phase should establish a stable search direction, whereas the BCE phase should align the terminal policy with the transmitted bits. Across transitions, a useful learned policy should reduce BER and residual without collapsing immediately to identical trajectories.

The soft-receiver ablation is performed at one SNR near the coded waterfall and is summarized in a single table. Starting from the complete three-stage training scheme, we separately (i) remove the hard-checkpoint initialization, (ii) retain hard-style parameter tying across all soft-detector layers after the common initialization, (iii) remove prior tilting from candidate generation, (iv) omit the synthetic-prior stage, (v) omit decoder-in-the-loop fine-tuning, and (vi) stop the gradient at the LDPC decoder instead of using (38). We report BLER after each IDD round together with final-round GMI and Brier score. To check that prior subtraction in (33) does not leave excessive self-information, we additionally vary one bit’s input prior while holding the channel observation and all other priors fixed. We then report the sensitivity of that bit’s extrinsic LLR and the rate at which its sign changes. This targeted test diagnoses the approximation in (33) without requiring a full EXIT characterization.

The first main experiment reports uncoded BER versus SNR for 64-QAM and 256-QAM in two panels. This comparison determines whether the learned complete-vector transitions improve upon linear detection and existing neural detectors, and whether they approach the K-best performance obtained with an explicit tree search.

The second main experiment reports post-LDPC BLER versus SNR for both modulation orders. Curves after one, two, and three IDD rounds are shown for the proposed receiver; the baselines are shown after the same number of decoder iterations and, when applicable, the same number of detector–decoder exchanges. Final-round information-bit BER and detector GMI are reported in a small companion table. The round-by-round curves directly test whether Bayesian prior-tilted sampling converts decoder feedback into increasingly informative candidate sets rather than merely recalibrating a fixed hard decision.