Deep unfolding of the weighted MMSE beamforming algorithm

In this GitHub repository the user can find the code used to reproduce the plots and the results in our paper [1]. We propose the novel application of deep unfolding to the weighted minimum mean square error (WMMSE) algorithm in [2]. The WMMSE is an iterative algorithm that converges to a local solution of the weigthed sum rate maximization problem subject to a power constraint, which is known to be NP-hard. As noted in [3], the formulation of the WMMSE algorithm, as described in [2], is not amenable to deep unfolding due to a matrix inversion, an eigendecomposition, and a bisection search performed at each itearation. Therefore, in our paper [1], we propose an alternative formulation that avoids these operations, while retaining the same convergence property. Specifically, we replace the method of Lagrange multipliers with the projected gradient descent (PGD) approach. In this way, the matrix inversion, the eigendecomposition, and the bisection search are replaced by simple vector operations. We refer to this variant of the original WMMSE algorithm as unfoldable WMMSE algorithm in order to stress its suitability to deep unfolding. By mapping each iteration of the unfoldable WMMSE algorithm to a network layer, we obtain a network architecture called deep unfolded WMMSE. Moreover, in order to boost the convergence of the PGD, we propose to incorporate in the network structure Nesterov acceleration and a generalization thereof (which we refer to as Super Nesterov) and treat such acceleration schemes as learnable structures.

In the jupyter notebook Deep_Unfolded_WMMSE_versus_WMMSE.ipynb the user can find:

The implementation in Python 3.6.8 of the WMMSE algorithm in [2]
The implementation in Python 3.6.8 of the zero forcing (ZF) solution
The implementation in Python 3.6.8 of the regularized zero-forcing (RZF) solution
The implementation in Python 3.6.8 and Tensorflow 1.13.1 of the deep unfolded WMMSE in [1]

In the jupyter notebooks Nesterov_Accelerated_Deep_Unfolded_WMMSE.ipynb and Super_Nesterov_Accelerated_Deep_Unfolded_WMMSE.ipynb the user can find the implementation in Python 3.6.8 and Tensorflow 1.13.1 of the accelerated deep unfolded WMMSE, with Nesterov and Super Nesteorv acceleration schemes, respectively.

The notebook cells should be run in sequential order. Note that the training time can vary from an hour to many hours, depending on the parameter settings, e.g., the number of iterations and the number of PGD steps, and on the user hardware.

Problem formulation

Let $\small x_i \sim \mathcal{CN}(0,\,1)$ be the transmitted data symbol to user $\small i$ and let $\small \boldsymbol{h}_i$ be the channel between the base station and user $\small i$ .

With linear beamforming, the signal at the receiver of user $\small i$ is

$\small y_{i} = {\boldsymbol{h}^{H}_i}\boldsymbol{v}_{i}x_{i} + \sum_{j=1,j \neq i}^{N}{\boldsymbol{h}^{H}_i}{\boldsymbol{v}_{j}x_{j}} + n_{i},$

where $\small \boldsymbol{v}_i \in \mathbb{C}^{M}$ is the transmit beamformer for user $\small i$ and where $\small n_i \sim \mathcal{CN}(0,\,\sigma^{2})$ is independent additive white Gaussian noise with power $\small \sigma^2$ . The signal-to-interference-plus-noise-ratio (SINR) of user $\small i$ is

$\small \mathrm{SINR}_i = \frac{|\boldsymbol{h}^H_{i}\boldsymbol{v}_i|^2}{ \sum_{j = 1, j\neq i }^{N}{|\boldsymbol{h}^H_{i}\boldsymbol{v}_j|^2 }+ \sigma^2}\cdot$

The estimated data symbol at the receiver of user $\small i$ is $\small \hat{x}_i = u_{i}y_{i}$ , where $\small u_i \in \mathbb{C}$ is the receiver gain.

We seek to maximize the weighted sum rate (WSR) subject to a total transmit power constraint, i.e.,

$\begin{aligned} \max_{\boldsymbol{V}} \quad & \sum_{i = 1}^{N}{\alpha_i\log_{2}{( 1 + \mathrm{SINR}_i)}} \\ \textrm{s.t.} \quad & \text{Tr}(\boldsymbol{V}\boldsymbol{V}^H)\leq P, \end{aligned}$

where $\small \alpha_{i}$ indicates the user priority (assumed to be known) and where $\small P$ is the maximum transmit power at the base station. We assume to have perfect channel knowledge. This problem is known to be NP-hard. We define $\boldsymbol{H}~\triangleq [\boldsymbol{h}_1,\boldsymbol{h}_2,\ldots,\boldsymbol{h}_N]^T$ , $\boldsymbol{V}~\triangleq [\boldsymbol{v}_1,\boldsymbol{v}_2,\ldots,\boldsymbol{v}_N]^T$ , $\boldsymbol{u} \triangleq [u_1, u_2, \ldots,u_N]^T$ , and $\boldsymbol{w} \triangleq~[w_1,w_2,...,w_N]^T$ .

Proposed unfoldable WMMSE algorithm

Algorithm 1 reports the pseudocode of the unfoldable WMMSE algorithm, in which indexes the users, indexes the layers/iterations, and indexes the PGD steps.

$\small \nabla f({{\boldsymbol{v}}}^{k}_i) =-2\alpha_{i}w_{i}u_{i}\boldsymbol{h}_i + 2\boldsymbol{A}{\boldsymbol{v}}^{k}_i$

$\small \boldsymbol{A} \triangleq \sum_{i = 1}^{N}{\alpha_{i}w_{i}|u_i|^2\boldsymbol{h}_i\boldsymbol{h}^H_i}$

$\small \Pi_{\mathcal{C}}\{\boldsymbol{V}\} = \begin{cases} \boldsymbol{V}, & \text{if}\ \mathtt{Tr}(\boldsymbol{V}\boldsymbol{V}^H)\leq P \\ \frac{\boldsymbol{V}}{\left\lVert \boldsymbol{V} \right\rVert}\sqrt{P}, & \text{otherwise.} \end{cases}$

Deep unfolded WMMSE

Fig.1 depicts the overall network architecture.

Fig.1. Network architecture of the deep unfolded WMMSE. It is given by L iterations of the unfoldable WMMSE algorithm. The subscripts $\small (\cdot)^{l,k}$ indicate the $\small k^{th}$ PGD step in the $\small l^{th}$ layer/iteration. Each layer consists of the update equation of $\small \boldsymbol{u}$ , denoted by $\small \boldsymbol{\Omega}$ , of the update equation of $\small \boldsymbol{w}$ , denoted by $\small \boldsymbol{\Psi}$ , and of K PGD steps, as depicted in the gray box. The step sizes of the truncated PGD sequence are the trainable parameters. $\small \nabla$ and $\small \Pi_\mathcal{C}$ denote the gradient and the projection operations.

Computation environment

In order to run the code in this repository the following software packages are needed:

Python 3 (for reference we use Python 3.6.8), with the following packages:numpy, tensorflow (version 1.x - for reference we use version 1.13.1), matplotlib,copy,time.
Jupyter (for reference we use version 6.0.3).

Reference

[1] L. Pellaco, M. Bengtsson, J. Jaldén, "Matrix-inverse-free deep unfolding of the weighted MMSE beamforming algorithm," submitted.

[2] Q. Shi, M. Razaviyayn, Z. Luo and C. He, "An Iteratively Weighted MMSE Approach to Distributed Sum-Utility Maximization for a MIMO Interfering Broadcast Channel," in IEEE Transactions on Signal Processing, vol. 59, no. 9, pp. 4331-4340, Sept. 2011, doi: 10.1109/TSP.2011.2147784.

[3] H. Sun, X. Chen, Q. Shi, M. Hong, X. Fu and N. D. Sidiropoulos, "Learning to Optimize: Training Deep Neural Networks for Interference Management," in IEEE Transactions on Signal Processing, vol. 66, no. 20, pp. 5438-5453, 15 Oct.15, 2018, doi: 10.1109/TSP.2018.2866382.

lpkg / WMMSE-deep-unfolding