加权最小均方误差（weighted minimum mean square error, WMMSE）算法及Python实现

加权最小均方误差（weighted minimum mean square error, WMMSE）算法[1]

一、问题描述

考虑MIMO系统中的下行链路，基站发射天线数为$P$，用户数为$K$，每个用户的接收天线数为$Q$。用户$k$的MIMO信道记为${\mathbf{H}}_k\in {\mathbb{C}}^{[Q\times P]}$，则其在采样时间$n$时的接受信号可以表示为，
$${\mathbf{y}}_k(n)={\mathbf{H}}_k\mathbf{x}(n)+{\mathbf{v}}_k(n),$$
其中$\mathbf{x}(n)\in {\mathbb{C}}^{[P\times 1]}$为发射向量，${\mathbf{v}}_k(n)\in {\mathbb{C}}^{[Q\times 1]}$为包含循环对称高斯白噪声的噪声向量，协方差矩阵为${\mathbf{R}}_{v_k{v}_k}=\mathbb{E}\left[{\mathbf{v}}_k(n){\mathbf{v}}_k(n{)}^{\mathsf{H}}\right]={\mathbf{I}}_Q$。发射向量是输入数据向量${\mathbf{d}}_1(n),\cdots ,{\mathbf{d}}_K(n)\in {\mathbb{C}}^{[Q\times 1]}$的线性组合，即，
$$\mathbf{x}(n)=\sum _{k=1}^K{\mathbf{B}}_k{\mathbf{d}}_k(n).$$
其中，矩阵${\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K\in {\mathbb{C}}^{[P\times Q]}$为发射滤波器（波束赋形向量）。上文中，隐含地假设了每个用户有$Q$个并行的数据流，且各个数据流相互独立，即$\mathbb{E}\left[{\mathbf{d}}_k(n){\mathbf{d}}_k^{\mathsf{H}}(n)\right]={\mathbf{I}}_Q$。

若发射的总时长为$N$，则为满足总功率约束，应有$\frac{1}{N}{\sum }_{n=1}^N{\mathbf{x}}^{\mathsf{H}}(n)\mathbf{x}(n)\le E_{\mathsf{tx}}$。当$N$足够大时，有$\mathbb{E}\left[{\mathbf{x}}^{\mathsf{H}}(n)\mathbf{x}(n)\right]={\sum }_k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})$。

首先考虑加权和率（weighted sum-rate, WRS）最大问题。通过设计发射滤波器${\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K$使加权和率最大化的问题可以建模为，
$$\begin{array}{rl}P1:[{\mathbf{B}}_1^{\mathsf{WSR}},\cdots ,{\mathbf{B}}_K^{\mathsf{WSR}}]=& \arg \underset{{\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K}{\min }\sum _k-u_{R_k}{R}_k\\ & s.t.\sum _{k=1}^K\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})=E_{\mathsf{tx}}.\end{array}$$
假设信号为高斯分布，则数据率为，
$$R_k=\log \det \left({\mathbf{I}}_k+{\mathbf{B}}_k^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}^{-1}{\mathbf{H}}_k{\mathbf{B}}_k\right),$$
其中${\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}$为用户$k$的等效噪声协方差矩阵，即，
$${\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}={\mathbf{I}}_k+\sum _{i=1,i\ne k}^K{\mathbf{H}}_k{\mathbf{B}}_i{\mathbf{B}}_i^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}.$$

接着考虑WMMSE最小化问题。若接收端采用MMSE滤波器，即，
$$\begin{array}{rl}{\mathbf{A}}_k^{\mathsf{MMSE}}& =\arg \underset{{\mathbf{A}}_k}{\min }\mathbb{E}\left[||{\mathbf{A}}_k{\mathbf{y}}_k-{\mathbf{d}}_k|{|}^2\right]\\ & ={\mathbf{B}}_k^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}({\mathbf{H}}_k{\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}+{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}{)}^{-1},\end{array}$$
则用户$k$的MSE矩阵可以表示为，
$$\begin{array}{rl}{\mathbf{E}}_k& =\mathbb{E}\left[({\mathbf{A}}_k^{\mathsf{MMSE}}{\mathbf{y}}_k-{\mathbf{d}}_k)({\mathbf{A}}_k^{\mathsf{MMSE}}{\mathbf{y}}_k-{\mathbf{d}}_k){)}^{\mathsf{H}}\right]\\ & =({\mathbf{I}}_k+{\mathbf{B}}_k^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}^{-1}{\mathbf{H}}_k{\mathbf{B}}_k{)}^{-1}.\end{array}$$

WMMSE最小化问题，是指在接收端采用MMSE滤波器时，使加权后的均方误差最小。该问题可以表示为，
$$\begin{array}{rl}P2:[{\mathbf{B}}_1^{\mathsf{WMMSE}},\cdots ,{\mathbf{B}}_K^{\mathsf{WMMSE}}]& =\arg \underset{{\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K}{\min }\sum _k\mathrm{Tr}({\mathbf{W}}_k{\mathbf{E}}_k)\\ & s.t.\sum _k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})=E_{\mathsf{tx}},\end{array}$$
其中${\mathbf{W}}_k\in {\mathbb{C}}^{[Q_k\times Q_k]}$为与用户$k$相关的权重矩阵。

二、基本思想

原始问题P1是优化变量为${\mathbf{B}}_k\forall k$的非凸问题，其与优化变量为${\mathbf{B}}_k,{\mathbf{A}}_k,{\mathbf{W}}_k\forall k$的优化问题P3等价，因为它们具有相同的最优解${\mathbf{B}}_k^{\star }\forall k$。

将原始问题P1转化为一个更高维度的优化问题P3，是因为P3在固定任意其它变量、单独优化某一变量时，具有显式解。通过迭代地优化各个变量的方式求解P3，可以收敛至一局部最优解。

三、最大化WSR与最小化WMMSE的联系

首先，P1和P2在目标函数上存在联系。从用户$k$的数据率$R_k$和MSE矩阵${\mathbf{E}}_k$的表达式上，可以看出$R_k=\log \det ({\mathbf{E}}_k^{-1})$。

其次，P1和P2的KKT条件在一定条件下等价。对于优化问题P1，其拉格朗日函数为，
$$f({\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K)=\sum _k-u_{R_k}{R}_k+\lambda \left(\sum _k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})-E_{\mathsf{tx}}\right)$$
相关的梯度为，
$${\nabla }_{{\mathbf{B}}_k}{R}_k={\mathbf{H}}_k^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}^{-1}{\mathbf{H}}_k{\mathbf{B}}_k{\mathbf{E}}_k.$$
$${\nabla }_{{\mathbf{B}}_k}{R}_i=-{\mathbf{H}}_i^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_i{\tilde{v}}_i}^{-1}{\mathbf{H}}_i{\mathbf{B}}_i{\mathbf{E}}_i{\mathbf{B}}_i^{\mathsf{H}}{\mathbf{H}}_i^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_i{\tilde{v}}_i}^{-1}{\mathbf{H}}_i{\mathbf{B}}_k.$$
对于优化问题P2，其拉格朗日函数为，
$$g({\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K)=\sum _k\mathrm{Tr}({\mathbf{W}}_k{\mathbf{E}}_k)+\lambda \left(\sum _k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})-E_{\mathsf{tx}}\right),$$
相关的梯度为，
$${\nabla }_{{\mathbf{B}}_k}\mathrm{Tr}({\mathbf{W}}_k{\mathbf{E}}_k)=-{\mathbf{H}}_k^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}^{-1}{\mathbf{H}}_k{\mathbf{B}}_k{\mathbf{E}}_k{\mathbf{W}}_k{\mathbf{E}}_k,$$
$${\nabla }_{{\mathbf{B}}_k}\mathrm{Tr}({\mathbf{W}}_i{\mathbf{E}}_i)={\mathbf{H}}_i^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_i{\tilde{v}}_i}^{-1}{\mathbf{H}}_i{\mathbf{B}}_i{\mathbf{E}}_i{\mathbf{W}}_i{\mathbf{E}}_i{\mathbf{B}}_i^{\mathsf{H}}{\mathbf{H}}_i^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_i{\tilde{v}}_i}^{-1}{\mathbf{H}}_i{\mathbf{B}}_k.$$
从上式可以看出，若给定发射滤波器${\mathbf{B}}_1,\cdots ,{\mathbf{B}}_K$以及相应的MMSE矩阵${\mathbf{E}}_1,\cdots ,{\mathbf{E}}_K$，则当MSE权重矩阵选定为${\mathbf{W}}_k=u_k{\mathbf{E}}_k^{-1}$时，WMMSE问题的梯度与WSR问题的梯度相同。

最后，P1和P2的解在一定条件下等价。假设WSR问题的最优解为${\mathbf{B}}_1^{\mathsf{WSR}},\cdots ,{\mathbf{B}}_K^{\mathsf{WSR}}$，相应的MMSE矩阵为${\mathbf{E}}_1^{\mathsf{WSR}},\cdots ,{\mathbf{E}}_K^{\mathsf{WSR}}$，则其满足WSR问题的KKT条件。若根据${\mathbf{W}}_k=u_k{\mathbf{E}}_k^{-1}$设置WMMSE问题的权重矩阵，则可知当${\mathbf{B}}_k={\mathbf{B}}_k^{\mathsf{WSR}}\forall k$时，该预编码矩阵满足WMMSE问题的KKT条件。因此，此时二者最优解相等，即${\mathbf{B}}_k^{\mathsf{WMMSE}}={\mathbf{B}}_k^{\mathsf{WSR}}\forall k$。

四、WMMSE算法及其收敛性分析

根据P1和P2的联系，可通过求解WMMSE最小化问题求解WSR最大化问题。相应的迭代算法如下。

$$\begin{array}{rl}& \mathrm{Proposed}\mathrm{Algorithm}:\\ & setn=0\\ & set{\mathbf{B}}_k^n={\mathbf{B}}_k^{\mathsf{ini}}\forall k\\ & iterate\\ & updaten=n+1\\ & I.compute{\mathbf{A}}_k^n|{\mathbf{B}}_i^{n-1}\forall iforallkwith\mathrm{MMSE}filters\\ & II.compute{\mathbf{W}}_k^n|{\mathbf{B}}_i^{n-1}\forall iforallk\\ & III.compute{\mathbf{B}}^n|{\mathbf{A}}^n,{\mathbf{W}}^n.\\ & untilconvergence\end{array}$$
其中，步骤$\mathrm{III}$为求解WMMSE最小化问题，其显式解为，
$$\bar{\mathbf{B}}=({\mathbf{H}}^{\mathsf{H}}{\mathbf{A}}^{\mathsf{H}}\mathbf{WAH}+\frac{\mathrm{Tr}(\mathbf{WA}{\mathbf{A}}^{\mathsf{H}})}{E_{\mathsf{tx}}}{\mathbf{I}}_P{)}^{-1}{\mathbf{H}}^{\mathsf{H}}{\mathbf{A}}^{\mathsf{H}}\mathbf{W},$$
$${\mathbf{B}}^{\mathsf{WMMSE}}=b\bar{\mathbf{B}},$$
其中， W [ Q K × Q K ] = d i a g { W 1 , ⋯   , W K } \mathbf{W}_{[QK{\times}QK]}=\mathrm{diag}\{\mathbf{W}_1,\cdots,\mathbf{W}_K\} W[QK×QK]​=diag{W1​,⋯,WK​}， A [ Q K × Q K ] = d i a g { A 1 , ⋯   , A K } \mathbf{A}_{[QK{\times}QK]}=\mathrm{diag}\{\mathbf{A}_1,\cdots,\mathbf{A}_K\} A[QK×QK]​=diag{A1​,⋯,AK​}，${\mathbf{H}}_{[QK\times P]}=[{\mathbf{H}}_1^{\mathsf{T}},\cdots ,{\mathbf{H}}_K^{\mathsf{T}}{]}^{\mathsf{T}}$，$b=\sqrt{\frac{E_{\mathsf{tx}}}{\mathrm{Tr}(\stackrel{ˉ}{\mathbf{B}{\bar{\mathbf{B}}}^{\mathsf{H}}})}}$。

为了分析上述算法的收敛性，考虑如下优化问题，
$$\begin{array}{rl}P3:[{\mathbf{B}}_k^{\mathsf{WSR}}\forall k]& =\arg \underset{{\mathbf{B}}_k,{\mathbf{A}}_k,{\mathbf{W}}_k\forall k}{\min }\sum _k{\tilde{l}}_k({\mathbf{W}}_k,{\mathbf{A}}_k,{\mathbf{B}}_i\forall i)\\ & s.t.\sum _k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})\le E_{\mathsf{tx}},\end{array}$$
其中，${\tilde{l}}_k({\mathbf{W}}_k,{\mathbf{A}}_k,{\mathbf{B}}_i\forall i)=\mathrm{Tr}({\mathbf{W}}_k{\tilde{\mathbf{E}}}_k)-u_{R_k}\log \det (u_{R_k}^{-1}{\mathbf{W}}_k{)}^{-1}-u_{R_k}Q$，${\tilde{\mathbf{E}}}_k=\mathbb{E}\left[({\mathbf{A}}_k{\mathbf{y}}_k-{\mathbf{d}}_k)({\mathbf{A}}_k{\mathbf{y}}_k-{\mathbf{d}}_k{)}^{\mathsf{H}}\right]$。

在P3中，依次优化变量${\mathbf{A}}_k$和${\mathbf{W}}_k$，易知当${\mathbf{A}}_k={\mathbf{A}}_k^{\mathsf{MMSE}}({\mathbf{B}}_i\forall i)$，${\mathbf{W}}_k={\mathbf{W}}_k^{\mathsf{min}}({\mathbf{B}}_i\forall i)=u_{R_k}{\mathbf{E}}_k^{-1}({\mathbf{B}}_i\forall i)$时，目标函数达到最小值，此时P3与P1等价。即，P3和P1具有相同的目标值和最优解。

下面考虑P3和WMMSE算法的对应关系。实际上，算法中的步骤$\mathrm{I}$对应于在${\mathbf{B}}_k\forall k$和${\mathbf{W}}_k\forall k$给定的情况下优化${\mathbf{A}}_k\forall k$；算法中的步骤$\mathrm{III}$对应于在${\mathbf{W}}_k\forall k$和${\mathbf{A}}_k\forall k$给定的情况下优化${\mathbf{B}}_k\forall k$；算法中的步骤$\mathrm{III}$对应于在${\mathbf{A}}_k\forall k$和${\mathbf{B}}_k\forall k$给定的情况下优化${\mathbf{W}}_k\forall k$。这三个过程均使P3的目标函数单调减小，因此算法最终会收敛至一局部最优解。

*五、使MMSE矩阵为对角阵的WMMSE-D算法

根据$R_k$的表达式以及性质$\det (I+AB)=\det (I+BA)$可知，若将发射滤波器由${\mathbf{B}}_k$改为${\tilde{\mathbf{B}}}_k={\mathbf{B}}_k{\mathbf{Q}}_k$，则用户$k$的数据率保持不变。其中，${\mathbf{Q}}_k$为任意施密特矩阵。

当取${\mathbf{Q}}_k={\mathbf{V}}_k$，其中${\mathbf{V}}_{\mathbf{k}}{\Lambda }_k{\mathbf{V}}_k^{\mathsf{H}}={\mathbf{B}}_k^{\mathsf{H}}{\mathbf{H}}_k^{\mathsf{H}}{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}^{-1}{\mathbf{H}}_k{\mathbf{B}}_k$时，有${\mathbf{E}}_k^{-1}={\mathbf{I}}_k+\Lambda k$，为对角矩阵。

若${\mathbf{E}}_k$为最优解对应的MMSE矩阵，则有$\log \det ({\mathbf{E}}_k^{-1})\ge {\sum }_q\log (e_{k,q}{)}^{-1}$，等号在${\mathbf{E}}_k$为对角阵时成立。所以，有如下关系，
$$\log \det {\mathbf{E}}_k^{-1}({\mathbf{B}}_k,{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k})=\underset{{\mathbf{Q}}_k}{\max }\sum _q\log e_{k,q}^{-1}({\mathbf{B}}_k{\mathbf{Q}}_k,{\mathbf{R}}_{{\tilde{v}}_k{\tilde{v}}_k}).$$

所以，P1可转化为如下优化问题，且P4的解是P1中能够使${\mathbf{E}}_k$为对角矩阵的解，

$$\begin{array}{rl}P4:[{\mathbf{B}}_k^{\mathsf{WSRDIAG}}\forall k]& =\arg \underset{{\mathbf{B}}_k\forall k}{\min }\sum _k\sum _q-u_{R_k}\log e_{k,q}^{-1}\\ & s.t.\sum _k\mathrm{Tr}({\mathbf{B}}_k{\mathbf{B}}_k^{\mathsf{H}})\le E_{\mathsf{tx}}.\end{array}$$

与P1类似地，考虑与P4具有相同解的WMMSE最小化问题，可得到类似的WMMSE-D算法。该算法与WMMSE算法的区别仅在于，此时 W k = u k d i a g { e k , 1 − 1 , ⋯   , e k , Q − 1 } \mathbf{W}_k=u_k\mathrm{diag}\{e_{k,1}^{-1},\cdots,e_{k,Q}^{-1}\} Wk​=uk​diag{ek,1−1​,⋯,ek,Q−1​}。

六、推广

上文考虑的实际是单发射机的情况，即只有一个发射设备。文献[2]考虑了包含多个发射机的情形，该场景下每个发射机服务其覆盖范围内的多个用户。二者的区别仅在于，在具有多个发射机时，每个用户接收信号的干扰项中除了包含同一发射机发射给其它用户的信号（小区内干扰）外，还包含不同发射机发射给其它用户的信号（小区间干扰）。

七、MU-MIMO系统下WMMSE算法的Python实现

class MUMIMO:
    def __init__(self, num_user, num_rx, num_tx, num_stream, transmit_power, noise_power=1):
        self._num_user = num_user
        self._num_rx = num_rx
        self._num_tx = num_tx
        self._num_stream = num_stream
        self._transmit_power = transmit_power
        self._noise_power = noise_power

    def __str__(self):
        return 'users-{}, receive antennas-{}, transmit antennas-{}, streams-{}, transmit power-{}, noise power-{}'.format(
            self._num_user, self._num_rx, self._num_tx, self._num_stream, self._transmit_power, self._noise_power
        )

    def generate_channel(self, n_samples):
        H = np.random.randn(n_samples, self._num_user, self._num_rx, self._num_tx) + 1j * np.random.randn(n_samples, self._num_user, self._num_rx, self._num_tx)
        H = H / np.sqrt(2)
        return H.astype(np.complex64)

    def WMMSE(self, H, V_ini=None, maximum_iter=200):
        n_sample = H.shape[0]
        if V_ini is not None:
            V = V_ini
        else:
            V = np.random.randn(n_sample, self._num_tx, self._num_stream * self._num_user)\
                    + 1j * np.random.randn(n_sample, self._num_tx, self._num_stream * self._num_user)
            V = V * np.sqrt(self._transmit_power / np.real(np.trace(np.matmul(V, V.conj().transpose(0, 2, 1)), axis1=1, axis2=2))).reshape(-1, 1, 1)
            V = V.reshape(n_sample, self._num_tx, self._num_user, self._num_stream).transpose(0, 2, 1, 3)
        for i in range(maximum_iter):
            # compute MSE with given precoders
            HV = np.matmul(
                H.reshape(n_sample, self._num_user * self._num_rx, self._num_tx),
                V.transpose(0, 2, 1, 3).reshape(n_sample, self._num_tx, self._num_user * self._num_stream)
            ).reshape(n_sample, self._num_user, self._num_rx, self._num_user, self._num_stream).transpose(0, 1, 3, 2, 4)
            HVVH = np.matmul(
                HV,
                HV.conj().transpose(0, 1, 2, 4, 3)
            )
            eye = np.eye(self._num_rx).reshape(1, 1, self._num_rx, self._num_rx)
            R = eye * self._noise_power + np.sum(HVVH[:, :, :, :, :], axis=2) - np.diagonal(HVVH, axis1=1, axis2=2).transpose(0, 3, 1, 2)
            A = V.conj().transpose(0, 1, 3, 2) @ H.conj().transpose(0, 1, 3, 2) @ np.linalg.inv(R + np.diagonal(HVVH, axis1=1, axis2=2).transpose(0, 3, 1, 2))

            # compute weights with given precoders
            eye = np.eye(self._num_stream).reshape(1, 1, self._num_stream, self._num_stream)
            W = eye + V.conj().transpose(0, 1, 3, 2) @ H.conj().transpose(0, 1, 3, 2) @ np.linalg.inv(R) @ H @ V

            # update precoders
            P1 = np.sum(H.conj().transpose(0, 1, 3, 2) @ A.conj().transpose(0, 1, 3, 2) @ W @ A @ H, axis=1)
            P2 = np.sum(np.trace(W @ A @ A.conj().transpose(0, 1, 3, 2), axis1=2, axis2=3), axis=1).reshape(n_sample, 1, 1) * np.eye(self._num_tx).reshape(1, self._num_tx, self._num_tx) / self._transmit_power * self._noise_power
            P3 = (H.conj().transpose(0, 1, 3, 2) @ A.conj().transpose(0, 1, 3, 2) @ W).transpose(0, 2, 1, 3).reshape(n_sample, self._num_tx, self._num_user * self._num_stream)
            V = (np.linalg.inv(P1 + P2) @ P3)
            V = V * np.sqrt(self._transmit_power / np.real(np.trace(np.matmul(V, V.conj().transpose(0, 2, 1)), axis1=1, axis2=2))).reshape(-1, 1, 1)
            V = V.reshape(n_sample, self._num_tx, self._num_user, self._num_stream).transpose(0, 2, 1, 3)
        return V

七、参考文献

[1] Christensen, Søren Skovgaard, Rajiv Agarwal, Elisabeth De Carvalho, and John M. Cioffi. “Weighted sum-rate maximization using weighted MMSE for MIMO-BC beamforming design.” IEEE Transactions on Wireless Communications 7, no. 12 (2008): 4792-4799.
[2] Shi, Qingjiang, Meisam Razaviyayn, Zhi-Quan Luo, and Chen He. “An iteratively weighted MMSE approach to distributed sum-utility maximization for a MIMO interfering broadcast channel.” IEEE Transactions on Signal Processing 59, no. 9 (2011): 4331-4340.