A-MMMSE

This represitory 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-user multiple-input multiple-output (MU-MIMO) systems. We consider a single-cell MU-MIMO downlink, where a base station (BS) equipped with $M$ transmit antennas serves $K$ users, each with $N$ receive antennas, by simultaneously transmitting $d$ independent data streams. Let $\mathbf{s}_k \in \mathbb{C}^{d \times 1}$ denote the symbol vector intended for user $k$, and $\mathbf{V}_k \in \mathbb{C}^{M \times d}$ be the corresponding linear precoder. The received signal at user $k$ is given by:

$$ \mathbf{y}_{k} = \mathbf{H}_{k} \mathbf{V}_{k} \mathbf{s}_{k} + \sum_{\substack{j=1 \ j \neq k}}^{K} \mathbf{H}_{k} \mathbf{V}_{j} \mathbf{s}_{j} + \mathbf{n}_{k}, $$

where $\mathbf{H}_k \in \mathbb{C}^{N \times M}$ is the channel matrix from the BS to user $k$, and $\mathbf{n}_k \in \mathbb{C}^{N \times 1}$ is the additive white Gaussian noise vector following $\mathcal{CN}(\mathbf{0}, \sigma_k^2 \mathbf{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} \alpha_{k} R_{k} \\ & \text{s.t.} && \sum_{k=1}^{K} \text{Tr}\left(\mathbf{V}_{k} \mathbf{V}_{k}^{H}\right) \leq P_{\max}, \end{aligned} $$

where $\alpha_k$ is the priority weight of user $k$, $P_{\text{max}}$ denotes the total transmit power budget at the BS, and $R_k$ represents the achievable rate for user $k$, defined as:

$$ R_{k} \triangleq \log \text{det} \left(\mathbf{I} + \mathbf{H}_{k} \mathbf{V}_{k} \mathbf{V}_{k}^{H} \mathbf{H}_{k}^{H} \left( \sum_{j \neq k} \mathbf{H}_{k} \mathbf{V}_{j} \mathbf{V}_{j}^{H} \mathbf{H}_{k}^{H} + \sigma_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 matrix is updated via block coordinate gradient descent. This approach eliminates matrix inversion operations, relying exclusively on matrix multiplications that are exceptionally amenable to GPU acceleration. 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
pytorch=1.13.1 

Running experiments

To evaluate A-MMMSE performance under a configuration with 256 transmit antennas, 20 users, 4 data streams, and SNR=10 dB, execute the following command:

python main.py --config ./configs/MIMO.yaml --model_name A_MMMSE --T 256 --I 20 --d 4 --snr 10 --omega 0.8 --lr 0.05

For GPU-accelerated performance evaluation, run:

python main_gpu.py --config ./configs/MIMO.yaml --model_name A_MMMSE_GPU --T 1024 --I 20 --d 4 --snr 10 --omega 0.4 --lr 0.1

Citing our work

@article{gao2025accelerated,
  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={arXiv preprint arXiv:2510.20507},
  year={2025}
}