A-MMMSE

This repository is an implementation of the paper: An Accelerated Mixed Weighted-Unweighted MMSE Approach for MU-MIMO Beamforming.

Introduction

The weighted sum-rate (WSR) maximization problem plays a central role in precoding design for downlink multi-cell multi-user multiple-input multiple-output (MU-MIMO) systems. We consider a downlink MU-MIMO system consisting of $K$ cells. In the $k$-th cell, the base station (BS) is equipped with $M$ transmit antennas and serves $I$ users, each with $N$ receive antennas, by simultaneously transmitting $d$ independent data streams. Let $i_k$ denote the $i$-th user in the $k$-th cell, $s_{i_k} \in \mathbb{C}^{d \times 1}$ denote the symbol vector intended for user $i_k$, and $\mathbf{V}_{i_k} \in \mathbb{C}^{M \times d}$ be the corresponding linear precoder. The received signal at user $i_k$ is given by:

$$ \mathbf{y}_{i_k} = \mathbf{H}_{i_k k} \mathbf{V}_{i_k} \mathbf{s}_{i_k} + \sum_{(l,j) \neq (i,k)} \mathbf{H}_{i_k j} \mathbf{V}_{l_j} \mathbf{s}_{l_j} + \mathbf{n}_{i_k}, $$

where $H_{i_k j} \in \mathbb{C}^{N \times M}$ is the channel matrix from the j-th BS to user $i_k$, and $n_{i_k} \in \mathbb{C}^{N \times 1}$ is the additive white Gaussian noise vector following $\mathcal{CN}(0, \sigma_{i_k}^2 I)$. The WSR maximization problem over the set of precoders $\mathbf{V}$ is formulated as:

$$ \begin{aligned} &\underset{\mathbf{V}}{\text{max}} && \sum_{k=1}^{K}\sum_{i=1}^{I} \alpha_{i_k} R_{i_k} \\ & \text{s.t.} && \sum_{i=1}^{I} \text{Tr}\left(\mathbf{V}_{i_k} \mathbf{V}_{i_k}^{H}\right) \leq P_k,\quad k=1,2,\cdots,K, \end{aligned} $$

where $\alpha_{i_k}$ is the priority weight of user $i_k$, $P_k$ denotes the transmit power budget of the $k$-th BS, and $R_{i_k}$ represents the achievable rate for user $i_k$, defined as:

$$ R_{i_k} \triangleq \log \text{det} \left(\mathbf{I} + \mathbf{H}_{i_k k} \mathbf{V}_{i_k} \mathbf{V}_{i_k}^{H} \mathbf{H}_{i_k k}^{H} \left( \sum_{(l,j) \neq (i,k)} \mathbf{H}_{i_k j} \mathbf{V}_{l_j} \mathbf{V}_{l_j}^{H} \mathbf{H}_{i_k j}^{H} + \sigma_{i_k}^{2} \mathbf{I} \right)^{-1} \right). $$

Our work introduces an enhanced version of the classical WSR maximization algorithm WMMSE, developed within a block coordinate descent framework. The proposed method employs a highly parallel structure where the computationally intensive precoding matrices are updated via block coordinate gradient descent. This approach eliminates matrix inversion operations in the precoder update and avoids the bisection search required by conventional WMMSE. Additionally, a two-stage warm-start strategy based on sum mean-square error minimization is incorporated to accelerate convergence. The resulting algorithm is termed Accelerated Mixed Weighted-Unweighted Sum-MSE Minimization (A-MMMSE).

Software Dependencies

python=3.10.0
numpy
scipy
matplotlib
configargparse
pytorch=1.13.1

Running Experiments

To evaluate A-MMMSE performance under a multi-cell configuration with $K=4$, $M=256$, $N=4$, $I=16$, $d=4$, and SNR = 10 dB, execute the following command:

python main.py --config ./configs/MIMO.yaml --model_name A_MMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10 --lr_type lipschitz --extrp

To compare CPU baselines under the same configuration, run:

python main.py --config ./configs/MIMO.yaml --model_name WMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10
python main.py --config ./configs/MIMO.yaml --model_name MMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10
python main.py --config ./configs/MIMO.yaml --model_name R_WMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10
python main.py --config ./configs/MIMO.yaml --model_name R_MMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10
python main.py --config ./configs/MIMO.yaml --model_name A_WMMSE --T 256 --K 4 --I 16 --R 4 --d 4 --snr 10 --lr_type lipschitz --extrp

For GPU-accelerated performance evaluation, run:

python main_gpu.py --config ./configs/GPU_Test.yaml --model_name A_MMMSE_GPU --T 512 --K 2 --I 16 --R 4 --d 4 --snr 10 --lr 0.04

The scripts in scripts/ provide additional command-line examples for CPU and GPU experiments.

Citing Our Work

@article{gao2026accelerated,
  title={An Accelerated Mixed Weighted-Unweighted MMSE Approach for MU-MIMO Beamforming},
  author={Gao, Xi and Wang, Akang and Zhang, Junkai and Duan, Qihong and Xue, Jiang},
  journal={IEEE Transactions on Signal Processing},
  year={2026}
}