Unitary Approximate Message Passing for Matrix Factorization

VI is a machine learning method widely used to approximate posterior densities for Bayesian models [21], [23]. Let $\boldsymbol{V}$ and $\boldsymbol{R}$ be the set of hidden (latent) variables and visible (observed) variables, respectively, with joint distribution $p(\boldsymbol{V},\boldsymbol{R})$. The goal of VI is to find a tractable variational distribution ${q}(\boldsymbol{V})$ that approximates the true posterior distribution $p(\boldsymbol{V}|\boldsymbol{R})$. With the distribution $q(\boldsymbol{V})$, the log marginal distribution of the observed variables admits the following decomposition

where the variational lower bound $\mathcal{L}(q(\boldsymbol{V}))$ is given as

and the Kullback-Leibler (KL) divergence between $q(\boldsymbol{V})$ and $p(\boldsymbol{V}|\boldsymbol{R})$ is

The distribution $q(\boldsymbol{V})$ that minimizes the KL divergence $\mathcal{KL}(q(\boldsymbol{V})||p(\boldsymbol{V}|\boldsymbol{R}))$ can be found by maximizing the variational lower bound $\mathcal{L}(q(\boldsymbol{V}))$.

VI can be implemented with message passing with the assistance of graphical models such as factor graphs. If the variational distribution with some factorization is chosen, e.g.,

where $\boldsymbol{V}=\{\boldsymbol{V}_{k}\}$, then the variational distribution may be found through an iterative procedure [21] with the update rule

Here $f(\boldsymbol{V}_{k},\tilde{\boldsymbol{V}})$ is a local factor associated with $\boldsymbol{V}_{k}$ and $\tilde{\boldsymbol{V}}\in\{\boldsymbol{V}_{i},i\neq k\}$, depending on the structure of the factor graph. The update of $\{q_{k}(\boldsymbol{V}_{k})\}$ is carried out iteratively until it converges or a pre-set number of iterations is reached.

AMP was derived based on loopy BP with Gaussian and Taylor-series approximations [24], [16], which can be used to recover $\boldsymbol{x}$ from the noisy measurement $\boldsymbol{y}=\boldsymbol{A}\boldsymbol{x}+\boldsymbol{w}$ with $\boldsymbol{w}$ being a zero-mean white Gaussian noise vector. AMP works well with an i.i.d. Gaussian $\boldsymbol{A}$, but it can easily diverge for generic $\mathbf{A}$ [25]. Inspired by [26], the work in [20] showed that the robustness of AMP can be remarkably improved through simple pre-processing, i.e., performing a unitary transformation to the original linear model [19]. As any matrix $\boldsymbol{A}$ has an SVD $\boldsymbol{A}=\boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{V}$ with $\boldsymbol{U}$ and $\boldsymbol{V}$ being two unitary matrices, performing a unitary transformation with $\boldsymbol{U}^{\textrm{T}}$ leads to the following model

where $\boldsymbol{r}=\boldsymbol{U}^{\textrm{T}}\boldsymbol{y}$, $\boldsymbol{\Phi}=\boldsymbol{U}^{\textrm{T}}\boldsymbol{A}=\boldsymbol{\Lambda}\boldsymbol{V}$, $\mathbf{\Lambda}$ is an $M\times N$ rectangular diagonal matrix, and $\boldsymbol{\omega}=\boldsymbol{U}^{\textrm{T}}\boldsymbol{w}$ remains a zero-mean white Gaussian noise vector with the same covariance as $\boldsymbol{\omega}$.

Applying the vector step size AMP [16] with model (7) leads to the first version of UAMP (called UAMPv1) shown in Algorithm 1. Note that by replacing $\boldsymbol{r}$ and $\mathbf{\Phi}$ with $\boldsymbol{y}$ and $\boldsymbol{A}$ in Algorithm 1, the original AMP algorithm is recovered. An average operation can be applied to two vectors: $\boldsymbol{\tau}_{x}$ in Line 7 and $|\mathbf{\Phi}^{H}|^{.2}\boldsymbol{\tau}_{s}$ in Line 5 of UAMPv1 in Algorithm 1, leading to the second version of UAMP [20] (called UAMPv2), where the operations in the brackets of Lines 1, 5 and 7 are executed (refer to [19] for the detailed derivations). Compared to AMP and UAMPv1, UAMPv2 does not require matrix-vector multiplications in Lines 1 and 5, so that the number of matrix-vector products is reduced from 4 to 2 per iteration.

In the (U)AMP algorithms, $g_{x}(\boldsymbol{q},\boldsymbol{\tau}_{q})$ is related to the prior of $\boldsymbol{x}$ and returns a column vector with $n$th entry $[g_{x}(\mathbf{q},\boldsymbol{\tau}_{q})]_{n}$ given by

where $p(x_{n})$ represents a prior for $x_{n}$. The function $g_{x}^{\prime}(\boldsymbol{q},\boldsymbol{\tau}_{q})$ returns a column vector and the $n$th element is denoted by $[g_{x}^{\prime}(\boldsymbol{q},\boldsymbol{\tau}_{q})]_{n}$, where the derivative is taken with respect to $q_{n}$.

In this work, the matrix normal distribution is used in the derivation of VI for the MF problem, into which UAMP is incorporated, leading to the UAMP-MF algorithm. The matrix normal distribution is a generalization of the multivariate normal distribution to matrix-valued random variables [22]. The distribution of normal random matrix $\boldsymbol{X}$ is

where $\hat{\boldsymbol{X}}\in\mathcal{R}^{N\times L}$ is the mean of $\boldsymbol{X}$, $\boldsymbol{U}_{X}\in\mathcal{R}^{N\times N}$ and $\boldsymbol{V}_{X}\in\mathcal{R}^{L\times L}$ are the covariance among rows and columns of $\boldsymbol{X}$, respectively. The matrix normal distribution is related to the multivariate normal distribution in the following way

where $\boldsymbol{x}=\textrm{Vec}(\boldsymbol{X})$ and $\hat{\boldsymbol{x}}=\textrm{Vec}(\hat{\boldsymbol{X}})$.

With model (1), we have the following joint conditional distribution and its factorization

where

The precision $\lambda$ of the noise has an improper prior $f_{\lambda}(\lambda)\triangleq p(\lambda)\propto 1/\lambda$ [27], and the priors for $\boldsymbol{H}$ and $\boldsymbol{X}$ denoted by $f_{H}(\boldsymbol{H})\triangleq p(\boldsymbol{H})$ and $f_{X}(\boldsymbol{X})\triangleq p(\boldsymbol{X})$ are customized according to the requirements in a specific problem. The factor graph representation of $\eqref{eq:factor}$ is depicted in Fig.1, where the arguments of the local functions are dropped for the sake of conciseness. We define a variational distribution

and we expect that $q(\boldsymbol{X})\approx p(\boldsymbol{X}|\boldsymbol{Y})$, $q(\boldsymbol{H})\approx p(\boldsymbol{H}|\boldsymbol{Y})$ and $q(\lambda)\approx p(\lambda|\boldsymbol{Y})$. Next, we study how to efficiently update $q(\boldsymbol{X})$, $q(\boldsymbol{H})$ and $q(\lambda)$ iteratively.

According to VI, with $q(\boldsymbol{H})$ and $q(\lambda)$ (updated in the previous iteration), we update $q(\boldsymbol{X})$. As shown by the factor graph in Fig. 1, we need to compute the message $m_{f_{Y}\to X}(\boldsymbol{X})$ from the factor node $f_{Y}$ to the variable node $\boldsymbol{X}$ and then combine it with the prior $p(\boldsymbol{X})$. Later we will see that $q(\boldsymbol{H})$ is a matrix normal distribution, i.e., $q(\boldsymbol{H})=\mathcal{MN}(\boldsymbol{H};\hat{\boldsymbol{H}},\boldsymbol{U}_{H},\boldsymbol{V}_{H})$ with $\boldsymbol{U}_{H}=\boldsymbol{I}_{N}$, and $q(\lambda)$ is a Gamma distribution. It turns out that the message $m_{f_{Y}\to X}(\boldsymbol{X})$ is matrix normal, as shown by Proposition 1.

Proposition 1: The message from $f_{Y}$ to $\boldsymbol{X}$ is a matrix normal distribution of $\boldsymbol{X}$, which can be expressed as

with

and

The computation of $\hat{\lambda}$ is shown in (46).

According to VI, this message is combined with the prior $p(\boldsymbol{X})$ to obtain $q(\boldsymbol{X})$. This can be challenging as $p(\boldsymbol{X})$ may lead to an intractable $q(\boldsymbol{X})$ and high computational complexity. Next, leveraging UAMP, we update $q(\boldsymbol{X})$ efficiently. It is also worth mentioning that, with the use of UAMP, the matrix inversion involved in (19) does not need to be performed.

Note that $\boldsymbol{X}=[\boldsymbol{x}_{1},...,\boldsymbol{x}_{l}]$ and $\overline{\boldsymbol{X}}=[\overline{\boldsymbol{x}}_{1},...,\overline{\boldsymbol{x}}_{l}]$. The result in (17) indicates that $\boldsymbol{x}_{l}\sim\mathcal{N}(\boldsymbol{x}_{l};\overline{\boldsymbol{x}}_{l},\hat{\lambda}^{-1}\overline{\boldsymbol{U}}_{X})$. This leads to the following pseudo observation model

where $\boldsymbol{e}_{x}\sim\mathcal{N}(\boldsymbol{e}_{x};0,\hat{\lambda}^{-1}\overline{\boldsymbol{U}}_{X})$, i.e., the model noise is not white. We can whiten the noise by left-multiplying both sides of (21) by $\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}$, leading to

where $\boldsymbol{w}_{x}=\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}\boldsymbol{e}_{x}$ is white and Gaussian with covariance $\hat{\lambda}^{-1}\boldsymbol{I}$. Through the whitening process, we get a standard linear model with white additive Gaussian noise, which facilitates the use of UAMP. Considering all the vectors in $\overline{\boldsymbol{X}}$, we have

where $\boldsymbol{\Omega_{X}}$ is white and Gaussian.

With model (23) and the prior $p(\boldsymbol{X})$, we use UAMP to update $q(\boldsymbol{X})$. Following UAMP, a unitary transformation needs to be performed with the unitary matrix $\boldsymbol{C}_{X}^{T}$ obtained from the SVD of $\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}$ (or eigenvalue decomposition (EVD) as the matrix is definite and symmetric), i.e.,

with $\boldsymbol{\Lambda}$ being a diagonal matrix. After performing the unitary transformation, we have

where $\boldsymbol{R}_{x}=\boldsymbol{C}_{X}^{T}\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}\overline{\boldsymbol{X}}$, $\boldsymbol{\Phi}_{X}=\boldsymbol{C}_{X}^{T}\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}=\boldsymbol{\Lambda}\boldsymbol{C}_{X}^{T}$ and $\boldsymbol{\Omega^{\prime}_{X}}=\boldsymbol{C}_{X}^{T}\boldsymbol{\Omega_{X}}$, which is still white and Gaussian.

From the above, it seems costly to obtain $\boldsymbol{R}_{X}$ and $\boldsymbol{\Phi}_{X}$ in model (25). Specifically, a matrix inverse operation needs to be performed according to (19) so that $\overline{\boldsymbol{X}}$ in (18) can be obtained; a matrix squared root operation is needed to obtain $\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}$; and an SVD operation is required to $\boldsymbol{C}_{X}^{T}$. Next, we show that the computations can be greatly simplified.

Instead of computing $\overline{\boldsymbol{U}}_{X}^{-1}$ and $\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}$ followed by SVD of $\overline{\boldsymbol{U}}_{X}^{-\frac{1}{2}}$, we perform EVD to $\overline{\boldsymbol{U}}_{X}^{-1}=\overline{\boldsymbol{W}}_{X}=\hat{\boldsymbol{H}}^{\textrm{T}}\hat{\boldsymbol{H}}+\text{Tr}(\boldsymbol{U}_{H})\boldsymbol{V}_{H}$, i.e.,

where the diagonal matrix

Hence $\boldsymbol{\Phi}_{X}=\boldsymbol{\Lambda}\boldsymbol{C}_{X}^{T}=\boldsymbol{D}_{X}^{-1/2}\boldsymbol{C}_{X}^{T}$, where the computation of $\boldsymbol{D}_{x}^{-1/2}$ is trivial as $\boldsymbol{D}_{x}$ is a diagonal matrix. Meanwhile, the computation of the pseudo observation matrix $\boldsymbol{R}_{x}$ can also be simplified:

For convenience, we rewrite the unitary transformed pseudo observation model as

The above leads to Lines 1-3 of UAMP-MF in Algorithm 2.

Due to the prior $p(\boldsymbol{X})$, the use of exact $q(\boldsymbol{X})$ often makes the message update intractable. Following UAMP in Algorithm 1, we project it to be Gaussian, which can also be interpreted as performing the minimum mean squared estimation (MMSE) based on the pseudo observation model (29) with prior $p(\boldsymbol{X})$. These correspond to Lines 4-11 of UAMP-MF in Algorithm 2. We assume that the prior $p(\boldsymbol{X})$ are separable, i.e., $p(\boldsymbol{X})=\prod_{n,l}p(x_{nl})$, then the operations in Lines 10 and 11 are element-wise as explained in the following. As the estimation is decoupled by (U)AMP, we assume the following scalar pseudo models

where $q_{nl}$ is the $(n,l)$th element of $\boldsymbol{Q}_{X}$ in Line 9 of the UAMP-MF algorithm, $w_{nl}$ represents a Gaussian noise with mean zero and variance $v_{nl}$, and $v_{nl}$is the $(n,l)$th element of $\boldsymbol{V}_{Q_{x}}$ in Line 8 of the UAMP-MF algorithm. This is significant as the complex estimation is simplified to much simpler MMSE estimation based on a number of scalar models (30) with prior $p(x_{nl})$. With the notations in Lines 10 and 11, for each entry $x_{nl}$ in $\boldsymbol{X}$, the MMSE estimation leads to a Gaussian distribution

where $\hat{x}_{nl}$ and $v_{x_{nl}}$ is the $(n,l)$th element of $\hat{\boldsymbol{X}}$ in Line 11 and the $(n,l)$th element of $\boldsymbol{\Xi}_{X}$ in Line 10, respectively. We can see that each element $x_{nl}$ has its own variance. To facilitate subsequent processing, we make an approximation by assuming that the elements of each row in $\hat{\boldsymbol{X}}$ have the same variance, which is the average of their variances. Then they are collectively characterized by a matrix normal distribution, i.e., $q(\boldsymbol{X})=\mathcal{MN}(\boldsymbol{X};\hat{\boldsymbol{X}},\boldsymbol{U}_{X},\boldsymbol{V}_{X})$ with $\boldsymbol{U}_{X}=\text{diag}(\text{mean}(\boldsymbol{\Xi}_{X},2))$ and $\boldsymbol{V}_{X}=\boldsymbol{I}_{L}$, where $\text{mean}(\boldsymbol{\Xi}_{X},2)$ represents the average operation on the rows of $\boldsymbol{\Xi}_{X}$. This leads to Line 12 of UAMP-MF in algorithm 2.

With the updated $q(\boldsymbol{X})=\mathcal{MN}(\boldsymbol{X};\hat{\boldsymbol{X}},\boldsymbol{U}_{X},\boldsymbol{V}_{X})$ and $q(\lambda)$, we compute the message from $f_{Y}$ to $\boldsymbol{H}$ according to VI. Regarding the message, we have the following result.

Proposition 2: The message from $f_{Y}$ to $\boldsymbol{H}$ is a matrix normal distribution about $\boldsymbol{H}$, which can be expressed as

with

and

where $\hat{\lambda}=\int_{\lambda}\lambda q(\lambda)$.

Then we combine the message $m_{f_{Y}\to H}(\boldsymbol{H})$ with the prior of $\boldsymbol{H}$ to update $q(\boldsymbol{H})$. This will also be realized with UAMP through a whitening process. The procedure is similar to that for $q(\boldsymbol{X})$, and the difference is that the pseudo model is established row by row (rather than column by column) by considering the form of $m_{f_{Y}\to H}(\boldsymbol{H})$.

With the message $m_{f_{Y}\to H}(\boldsymbol{H})$ and

where $\overline{\boldsymbol{h}}_{m}^{\textrm{T}}\in\mathcal{R}^{1\times N}$ is the $m$th row vector, we have the pseudo observation model

where $\boldsymbol{e}_{h}\sim\mathcal{N}(\boldsymbol{e}_{h};0,\hat{\lambda}^{-1}\overline{\boldsymbol{V}}_{H})$. Whitening operation to (36) leads to

where $\boldsymbol{w}_{h}=\overline{\boldsymbol{V}}_{H}^{-\frac{1}{2}}\boldsymbol{e}_{h}$, which is white and Gaussian with covariance $\hat{\lambda}^{-1}\boldsymbol{I}$. Collecting all rows and representing them in matrix form, we have

UAMP is then employed with the model (38). Using the ideas in the previous section for updating $q(\boldsymbol{X})$, we obtain the unitary transformed model efficiently.

We first perform an EVD to matrix $\overline{\boldsymbol{V}}_{H}^{-1}=\overline{\boldsymbol{W}}_{H}=\hat{\boldsymbol{X}}\hat{\boldsymbol{X}}^{\textrm{T}}+\text{Tr}(\boldsymbol{V}_{X})\boldsymbol{U}_{X}$, i.e.,

Then the unitary transformed model is given as

where

and $\boldsymbol{\Omega_{H}}$ is white and Gaussian. The above leads to Lines 13-15 of UAMP-MF in Algorithm 2.

Following UAMP in Algorithm 1, we obtain $q(\boldsymbol{H})$ and project it to be Gaussian, which correspond to Lines 16-23 of UAMP-MF Algorithm 2. Similar to the case for updating $q(\boldsymbol{X})$, to accommodate $\{q(h_{mn})\}$ with a matrix normal distribution, we make an approximation by assuming the entries in each column of $\boldsymbol{H}$ share a common variance, which is the average of their variances. Then $q(\boldsymbol{H})=\mathcal{MN}(\boldsymbol{H};\hat{\boldsymbol{H}},\boldsymbol{U}_{H},\boldsymbol{V}_{H})$ with $\boldsymbol{V}_{H}=\text{diag}(\text{mean}(\boldsymbol{\Xi}_{H},1))$ and $\boldsymbol{U}_{H}=\boldsymbol{I}_{M}$ (i.e., Line 24 of UAMP-MF in Algorithm 2), where $\text{mean}(\boldsymbol{\Xi}_{H},1)$ represents the average operation on the columns of $\boldsymbol{\Xi}_{H}$.

With $q(\boldsymbol{H})=\mathcal{MN}\left(\boldsymbol{H};\hat{\boldsymbol{H}},\boldsymbol{V}_{H},\boldsymbol{I}_{M}\right)$ and $q(\boldsymbol{X})=\mathcal{MN}\left(\boldsymbol{X};\hat{\boldsymbol{X}},\boldsymbol{I}_{L},\boldsymbol{U}_{X}\right)$, we update $q(\lambda)$, i.e., combining the message $m_{f_{Y}\to\lambda}(\lambda)$ from $f_{Y}$ to $\lambda$ and its prior $p(\lambda)$.

Proposition 3: The message from $f_{Y}$ to $\lambda$ can be expressed as

where

With the message $m_{f_{Y}\to\lambda}(\lambda)$ and the prior $p(\lambda)\propto 1/\lambda$, $q(\lambda)$ can be computed as

which is a Gamma distribution. Then the mean of $\lambda$ is obtained as

which is Line 25 of UAMP-MF.

Regarding UAMP-MF in Algorithm 2, we have the following remarks:

• •

  We do not specify the priors of $\boldsymbol{H}$ and $\boldsymbol{X}$ in Algorithm 2. The detailed implementations of Lines 10, 11, 22 and 23 of the algorithm depends on the priors in a concrete scenario. Various examples are provided in Section IV.

• •

  As the MF problem often has local minima, we can use the strategy of restart to mitigate the issue of being stuck at local minima. In addition, the iterative process can be terminated based on some criterion, e.g., the normalized difference between the estimates of two consecutive iterations is smaller than a threshold.

• •

  There are often hyper-parameters in the priors of $\boldsymbol{H}$ and $\boldsymbol{X}$. If we do not have knowledge on these hyper-parameters, their values need to be leaned or tuned automatically. Thanks to the factor graph and message passing frame work, these extra tasks can be implemented by extending the factor graph in Fig. 1.

In the following, we show how to learn the hyper-parameters of the prior of $\boldsymbol{X}$. We assume that $\boldsymbol{X}$ is sparse and the sparsity rate is unknown. In this case, we may employ the sparsity promoting hierarchical Gaussian-Gamma prior [27]. The prior is then $p(x_{nl},\gamma_{nl})=p(x_{nl}|\gamma_{nl})p(\gamma_{nl})$, where $p(x_{nl}|\gamma_{nl})=\mathcal{N}(x_{nl};0,\gamma_{nl}^{-1})\triangleq f_{x_{nl}}$ and $p(\gamma_{nl})=Ga(\gamma_{nl};\epsilon,\eta)$ with $\epsilon$ and $\eta$ being the shape parameter and scale parameter of the Gamma distribution, respectively. While the precision $\gamma_{nl}$ is to be learned, the values for $\epsilon$ and $\eta$ are often set empirically, e.g., $\epsilon=\eta=0$ [27]. It is worth mentioning that $\epsilon$ can be tuned automatically to improve the performance [28].

The factor graph representation for this part is shown in Fig. 2. Also note that UAMP decouples the estimation as shown by the pseudo model (30), indicating that the incoming message to the factor graph in Fig. 2 is Gaussian, i.e.,

We perform inference on $x_{nl}$ and $\gamma_{nl}$, which can also be achieved by using VI with a variational distribution $q(x_{nl},\gamma_{nl})=q(x_{nl})q(\gamma_{nl})$.

According to VI,

with $m_{f_{x_{nl}}\to x_{nl}}(x_{nl})$ shown in (52), and