Low-Rank Approximation (III): CR

低秩近似之路（三）：CR 苏剑林 · 2024-10-11 · Mathematics approximation optimal matrix low-rank

In Low-Rank Approximation (II): SVD we showed that SVD gives the optimal low-rank approximation of an arbitrary matrix. That optimum was unconstrained—that is, SVD only cares about minimizing the error and pays no attention to the specific structure of the matrices involved. In many practical applications, however, for reasons of interpretability or to support nonlinear operations, we often want an approximate decomposition with some particular structure.

So starting from this article, we will explore some low-rank approximations with specific structure, and this post will focus on the CR approximation (Column-Row Approximation), which offers a simple scheme for speeding up matrix multiplication.

Problem Setup

The general formulation of finding the optimal rank-$r$ approximation of a matrix is

\begin{equation}\mathop{\text{argmin}}_{\text{rank}(\tilde{\boldsymbol{M}})\leq r}\Vert \tilde{\boldsymbol{M}} - \boldsymbol{M}\Vert_F^2\label{eq:loss-m2}\end{equation}more

where $\boldsymbol{M},\tilde{\boldsymbol{M}}\in\mathbb{R}^{n\times m},r < \min(n,m)$. In the previous two posts we already discussed two cases:

1. If there is no other constraint on $\tilde{\boldsymbol{M}}$, then the optimal solution to $\tilde{\boldsymbol{M}}$ is $\boldsymbol{U}_{[:,:r]}\boldsymbol{\Sigma}_{[:r,:r]}\boldsymbol{V}_{[:,:r]}^{\top}$, where $\boldsymbol{M}=\boldsymbol{U}\boldsymbol{\Sigma} \boldsymbol{V}^{\top}$ is the singular value decomposition (SVD) of $\boldsymbol{M}$;
2. If we stipulate that $\tilde{\boldsymbol{M}}=\boldsymbol{A}\boldsymbol{B}$ ($\boldsymbol{A}\in\mathbb{R}^{n\times r},\boldsymbol{B}\in\mathbb{R}^{r\times m}$), and $\boldsymbol{A}$ (or $\boldsymbol{B}$) is already given, then the optimal solution to $\tilde{\boldsymbol{M}}$ is $\boldsymbol{A} \boldsymbol{A}^{\dagger} \boldsymbol{M}$ (or $\boldsymbol{M} \boldsymbol{B}^{\dagger} \boldsymbol{B}$), where ${}^\dagger$ is the "pseudo-inverse".

Both of these results have very broad applications, but neither explicitly introduces a structural connection between $\tilde{\boldsymbol{M}}$ and $\boldsymbol{M}$, which makes it hard to intuitively see the relationship between $\tilde{\boldsymbol{M}}$ and $\boldsymbol{M}$—in other words, $\tilde{\boldsymbol{M}}$ is not very interpretable.

Moreover, if the objective involves a nonlinear operation such as $\phi(\boldsymbol{X}\boldsymbol{W})$, we usually can't use an arbitrary real projection matrix for dimensionality reduction—we need a "selective matrix" instead. For example, $\phi(\boldsymbol{X}\boldsymbol{W})\boldsymbol{S} = \phi(\boldsymbol{X}\boldsymbol{W}\boldsymbol{S})$ does not hold for an arbitrary matrix $\boldsymbol{S}$, but it does hold for a selective matrix $\boldsymbol{S}$.

So next we turn our attention to low-rank approximation under the constraint of selective matrices. Specifically, we have $\boldsymbol{X}\in\mathbb{R}^{n\times l},\boldsymbol{Y}\in\mathbb{R}^{l\times m}$, and having fixed $\boldsymbol{M}=\boldsymbol{X}\boldsymbol{Y}$, our task is to pick $r$ columns from $\boldsymbol{X}$ and the corresponding $r$ rows from $\boldsymbol{Y}$ to construct $\tilde{\boldsymbol{M}}$, i.e.,

\begin{equation}\mathop{\text{argmin}}_S\Vert \underbrace{\boldsymbol{X}_{[:,S]}}_{\boldsymbol{C}}\underbrace{\boldsymbol{Y}_{[S,:]}}_{\boldsymbol{R}} - \boldsymbol{X}\boldsymbol{Y}\Vert_F^2\quad\text{s.t.}\quad S\subset \{0,1,\cdots,l-1\}, |S|=r\end{equation}

Here $S$ can be understood as a slice, following Python's slicing convention, and we call $\boldsymbol{X}_{[:,S]}\boldsymbol{Y}_{[S,:]}$ the "CR approximation" of $\boldsymbol{X}\boldsymbol{Y}$. Note that this kind of slicing result can also be equivalently described using selective matrices: suppose the columns of $\boldsymbol{X}_{[:,S]}$ are respectively the $1,2,\cdots,r$-th columns of $\boldsymbol{X}$, indexed by $s_1,s_2,\cdots,s_r$; then we can define a selective matrix $\boldsymbol{S}\in\{0,1\}^{l\times r}$:

\begin{equation}S_{i,j}=\left\{\begin{aligned}&1, &i = s_j \\ &0, &i\neq s_j\end{aligned}\right.\end{equation}

That is, the $s_j$-th entry of the $j$-th column of $\boldsymbol{S}$ is 1, and all other entries are 0. This gives us $\boldsymbol{X}_{[:,S]}=\boldsymbol{X}\boldsymbol{S}$ and $\boldsymbol{Y}_{[S,:]}=\boldsymbol{S}^{\top} \boldsymbol{Y}$.

A First Approximation

If we express $\boldsymbol{X},\boldsymbol{Y}$ respectively as

\begin{equation}\boldsymbol{X} = (\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_l),\quad \boldsymbol{Y}=\begin{pmatrix}\boldsymbol{y}_1^{\top} \\ \boldsymbol{y}_2^{\top} \\ \vdots \\ \boldsymbol{y}_l^{\top}\end{pmatrix}\end{equation}

where $\boldsymbol{x}_i\in\mathbb{R}^{n\times 1},\boldsymbol{y}_i\in\mathbb{R}^{m\times 1}$ are all column vectors, then $\boldsymbol{X}\boldsymbol{Y}$ can be written as

\begin{equation}\boldsymbol{X}\boldsymbol{Y} = \sum_{i=1}^l \boldsymbol{x}_i\boldsymbol{y}_i^{\top}\end{equation}

and finding the optimal CR approximation of $\boldsymbol{X}\boldsymbol{Y}$ can equivalently be written as

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_l\in\{0,1\}}\left\Vert\sum_{i=1}^l \lambda_i \boldsymbol{x}_i\boldsymbol{y}_i^{\top} - \sum_{i=1}^l\boldsymbol{x}_i\boldsymbol{y}_i^{\top}\right\Vert_F^2\quad\text{s.t.}\quad \sum_{i=1}^l \lambda_i = r\label{eq:xy-l-k}\end{equation}

We know that the $F$-norm of a matrix is essentially the length of the vector obtained by flattening the matrix, so this optimization problem is essentially equivalent to: given $l$ vectors $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_l\in\mathbb{R}^d$, find

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_l\in\{0,1\}}\left\Vert\sum_{i=1}^l \lambda_i \boldsymbol{v}_i - \sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^l \lambda_i = r\label{eq:v-l-k}\end{equation}

where $\boldsymbol{v}_i = \text{vec}(\boldsymbol{x}_i \boldsymbol{y}_i^{\top})$, $d=nm$. Writing $\gamma_i = 1 - \lambda_i$, this can be further simplified to

\begin{equation}\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_l\in\{0,1\}}\left\Vert\sum_{i=1}^l \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^l \gamma_i = l-r\label{eq:v-l-k-0}\end{equation}

If the author understands correctly, exact solution of this optimization problem is NP-hard, so in general we can only seek approximate algorithms. One simple case that can be solved exactly is when $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_l$ are pairwise orthogonal, in which case

\begin{equation}\left\Vert\sum_{i=1}^l \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^l \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2\end{equation}

so its minimum value is just the sum of the smallest $l-r$ values of $\Vert\boldsymbol{v}_i\Vert^2$, i.e., we let $\gamma_i$ equal 1 for the $l-r$ vectors $\boldsymbol{v}_i$ with the smallest norms, and set the remaining $\gamma_i$ to 0. Even when the pairwise-orthogonality condition doesn't hold strictly, we can still use selecting the $l-r$ smallest-norm vectors $\boldsymbol{v}_i$ as an approximate solution. Going back to the original CR approximation problem, we have $\Vert\boldsymbol{x}_i\boldsymbol{y}_i^{\top}\Vert_F = \Vert\boldsymbol{x}_i\Vert \Vert \boldsymbol{y}_i\Vert$, so a baseline for the optimal CR approximation of $\boldsymbol{X}\boldsymbol{Y}$ is to keep the $r$ column/row pairs for which the product of the norms of the corresponding column of $\boldsymbol{X}$ and row of $\boldsymbol{Y}$ is largest.

A Sampling Perspective

There are some scenarios in which we're allowed to relax equation $\eqref{eq:xy-l-k}$ to

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_l\in\mathbb{R}}\left\Vert\sum_{i=1}^l \lambda_i \boldsymbol{x}_i\boldsymbol{y}_i^{\top} - \sum_{i=1}^l\boldsymbol{x}_i\boldsymbol{y}_i^{\top}\right\Vert_F^2\quad\text{s.t.}\quad \sum_{i=1}^l \#[\lambda_i\neq 0] = r\end{equation}

where $\#[\lambda_i\neq 0]$ outputs 1 when $\lambda_i\neq 0$ holds, and 0 otherwise. This relaxed version essentially extends the form of the CR approximation from $\boldsymbol{C}\boldsymbol{R}$ to $\boldsymbol{C}\boldsymbol{\Lambda}\boldsymbol{R}$, where $\boldsymbol{\Lambda}\in\mathbb{R}^{r\times r}$ is a diagonal matrix—that is, it allows us to tune the diagonal matrix $\boldsymbol{\Lambda}\in\mathbb{R}^{r\times r}$ to achieve higher accuracy. Correspondingly, equation $\eqref{eq:v-l-k}$ becomes

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_l\in\mathbb{R}}\left\Vert\sum_{i=1}^l \lambda_i \boldsymbol{v}_i - \sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^l \#[\lambda_i\neq 0] = r\end{equation}

With this relaxation, we can look at the problem from a sampling perspective. First we introduce an arbitrary $l$-way categorical distribution $\boldsymbol{p}=(p_1,p_2,\cdots,p_l)$, and then we can write

\begin{equation}\sum_{i=1}^l\boldsymbol{v}_i = \sum_{i=1}^l p_i\times\frac{\boldsymbol{v}_i}{p_i} = \mathbb{E}_{i\sim \boldsymbol{p}} \left[\frac{\boldsymbol{v}_i}{p_i}\right] \end{equation}

That is to say, the expectation of $\boldsymbol{v}_i/p_i$ is exactly the target we want to approximate, so we can construct an approximation by drawing independent repeated samples from the distribution $\boldsymbol{p}$:

\begin{equation}\sum_{i=1}^l\boldsymbol{v}_i = \mathbb{E}_{i\sim \boldsymbol{p}} \left[\frac{\boldsymbol{v}_i}{p_i}\right] \approx \frac{1}{r}\sum_{j=1}^r \frac{\boldsymbol{v}_{s_j}}{p_{s_j}},\quad s_1,s_2,\cdots,s_r\sim \boldsymbol{p}\end{equation}

This means that when $i$ is one of $s_1,s_2,\cdots,s_r$ we have $\lambda_i = (r p_i)^{-1}$, and otherwise $\lambda_i=0$. Readers may wonder why we use independent repeated sampling rather than sampling without replacement, which seems more suited to approximation. There's no deep reason—it's purely because independent repeated sampling makes the subsequent analysis simpler.

So far, our theoretical results are independent of the choice of distribution $\boldsymbol{p}$—that is, they hold for arbitrary $\boldsymbol{p}$—which gives us the opportunity to choose the optimal $\boldsymbol{p}$. So how should we measure the quality of $\boldsymbol{p}$? Clearly, we want the error of each sampling-based estimate to be as small as possible, so we can use the error of the sampling estimate,

\begin{equation}\mathbb{E}_{i\sim \boldsymbol{p}} \left[\left\Vert\frac{\boldsymbol{v}_i}{p_i} - \sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2\right] = \left(\sum_{i=1}^l \frac{\Vert\boldsymbol{v}_i\Vert^2}{p_i}\right) - \left\Vert\sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2 \end{equation}

to compare different choices of $\boldsymbol{p}$. Then, applying the mean inequality, we get

\begin{equation}\sum_{i=1}^l \frac{\Vert\boldsymbol{v}_i\Vert^2}{p_i} = \left(\sum_{i=1}^l \frac{\Vert\boldsymbol{v}_i\Vert^2}{p_i} + p_i Z^2\right) - Z^2\geq \left(\sum_{i=1}^l 2\Vert\boldsymbol{v}_i\Vert Z\right) - Z^2\end{equation}

with equality holding when $\Vert\boldsymbol{v}_i\Vert^2 / p_i = p_i Z^2$, from which we obtain the optimal $\boldsymbol{p}$ as

\begin{equation}p_i^* = \frac{\Vert\boldsymbol{v}_i\Vert}{Z},\quad Z = \sum\limits_{i=1}^l \Vert\boldsymbol{v}_i\Vert\end{equation}

with corresponding error

\begin{equation}\mathbb{E}_{i\sim \boldsymbol{p}} \left[\left\Vert\frac{\boldsymbol{v}_i}{p_i} - \sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2\right] = \left(\sum_{i=1}^l \Vert\boldsymbol{v}_i\Vert\right)^2 - \left\Vert\sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2 \end{equation}

The optimal $p_i$ is exactly proportional to $\Vert\boldsymbol{v}_i\Vert$, so the $r$ entries with largest probability $\boldsymbol{v}_i$ are also exactly the $r$ entries with largest norm $\boldsymbol{v}_i$, which links back to the approximation from the previous section. This result comes from the 2006 paper Fast Monte Carlo Algorithms for Matrices I: Approximating Matrix Multiplication, originally intended to speed up matrix multiplication. It shows that as long as we sample the corresponding columns/rows of $\boldsymbol{X},\boldsymbol{Y}$ according to $p_i\propto \Vert \boldsymbol{x}_i\boldsymbol{y}_i^{\top}\Vert_F = \Vert \boldsymbol{x}_i\Vert \Vert\boldsymbol{y}_i\Vert$ and rescale by $(r p_i)^{-1/2}$, we obtain a CR approximation of $\boldsymbol{X}\boldsymbol{Y}$, which reduces the multiplication complexity from $\mathcal{O}(lmn)$ to $\mathcal{O}(rmn)$.

Further Discussion

Whether we sort by norm or sample randomly according to $p_i\propto \Vert\boldsymbol{v}_i\Vert$, both approaches let us construct a CR approximation within linear complexity [i.e., $\mathcal{O}(l)$], which is of course ideal for real-time computation. But since both sorting and sampling depend only on $\Vert\boldsymbol{v}_i\Vert$, the resulting accuracy is only mediocre. If we can accept higher complexity, how might we improve the accuracy of the CR approximation?

We can try changing the unit of sorting from single vectors to $k$-tuples. For simplicity, suppose $k \leq l-r$ divides $l-r$; the number of ways to choose $k$ out of the $l$ vectors $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_l$ is $C_l^k$, and for each combination $\{s_1,s_2,\cdots,s_k\}$ we can compute the norm of the vector sum, $\Vert \boldsymbol{v}_{s_1} + \boldsymbol{v}_{s_2} + \cdots + \boldsymbol{v}_{s_k}\Vert$. With this data in hand, we can greedily construct an approximate solution to $\eqref{eq:v-l-k-0}$:

Initialize $\Omega = \{1,2,\cdots,l\},\Theta=\{\}$.
For $t=1,2,\cdots,(l-r)/k$, do:
$\Theta = \Theta\,\cup\,\mathop{\text{argmin}}\limits_{\{s_1,s_2,\cdots,s_k\}\subset \Omega}\Vert \boldsymbol{v}_{s_1} + \boldsymbol{v}_{s_2} + \cdots + \boldsymbol{v}_{s_k}\Vert$;
$\Omega = \Omega\,\backslash\,\Theta$;
Return $\Theta$.

In plain terms, at each step we pick, from the remaining vectors, the group of $k$ vectors whose sum has the smallest norm, and repeat this $(l-r)/k$ times to obtain $l-r$ vectors in total. This is a natural generalization of sorting by single-vector norm, with complexity $\mathcal{O}(C_l^k)$, which can be prohibitive when $k > 1$ and $l$ is fairly large—this in turn reflects the underlying complexity of solving the original problem exactly.

Another question worth thinking about: if we allow the CR approximation to be relaxed to $\boldsymbol{C}\boldsymbol{\Lambda}\boldsymbol{R}$, what's the optimal solution for $\boldsymbol{\Lambda}$? If we don't restrict the structure of $\boldsymbol{\Lambda}$, then the answer can be given by the pseudo-inverse

\begin{equation}\boldsymbol{\Lambda}^* = \mathop{\text{argmin}}_{\boldsymbol{\Lambda}}\Vert \boldsymbol{C}\boldsymbol{\Lambda}\boldsymbol{R} - \boldsymbol{X}\boldsymbol{Y}\Vert_F^2 = \boldsymbol{C}^{\dagger}\boldsymbol{X}\boldsymbol{Y}\boldsymbol{R}^{\dagger}\end{equation}

What if $\boldsymbol{\Lambda}$ is required to be a diagonal matrix? Then we can first restate the problem as: given $\{\boldsymbol{u}_1,\boldsymbol{u}_2,\cdots,\boldsymbol{u}_r\}\subset\{\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_l\}$, find

\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_r}\left\Vert\sum_{i=1}^r \lambda_i \boldsymbol{u}_i - \sum_{i=1}^l\boldsymbol{v}_i\right\Vert^2\end{equation}

Writing $\boldsymbol{U} = (\boldsymbol{u}_1,\boldsymbol{u}_2,\cdots,\boldsymbol{u}_r), \boldsymbol{V} = (\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_l), \boldsymbol{\lambda}=(\lambda_1,\lambda_2,\cdots,\lambda_r)^{\top}$, the optimization objective can be written as

\begin{equation}\mathop{\text{argmin}}_{\boldsymbol{\lambda}}\left\Vert\boldsymbol{U}\boldsymbol{\lambda} - \boldsymbol{V}\boldsymbol{1}_{l\times 1}\right\Vert^2\end{equation}

and once again the optimal solution can be written via the pseudo-inverse as

\begin{equation}\boldsymbol{\lambda}^* = \boldsymbol{U}^{\dagger}\boldsymbol{V}\boldsymbol{1}_{l\times 1} = (\boldsymbol{U}^{\top}\boldsymbol{U})^{-1}\boldsymbol{U}^{\top}\boldsymbol{V}\boldsymbol{1}_{l\times 1} \end{equation}

The last equality assumes that $\boldsymbol{U}^{\top}\boldsymbol{U}$ is invertible, which usually holds; if it doesn't, we can just replace $(\boldsymbol{U}^{\top}\boldsymbol{U})^{-1}$ with $(\boldsymbol{U}^{\top}\boldsymbol{U})^{\dagger}$.

The issue now is that directly applying the formula above is computationally too expensive for the original problem, because $\boldsymbol{v}_i = \text{vec}(\boldsymbol{x}_i \boldsymbol{y}_i^{\top})$, i.e., $\boldsymbol{v}_i$ is a $mn$-dimensional vector, so $\boldsymbol{V}$ has size $mn\times l$ and $\boldsymbol{U}$ has size $mn\times r$, which becomes painful when $m,n$ is large. Making use of $\boldsymbol{v}_i = \text{vec}(\boldsymbol{x}_i \boldsymbol{y}_i^{\top})$ can help us further simplify the above. Let $\boldsymbol{u}_i = \text{vec}(\boldsymbol{c}_i \boldsymbol{r}_i^{\top})$; then

\begin{equation}\begin{aligned}(\boldsymbol{U}^{\top}\boldsymbol{V})_{i,j} =&\, \langle \boldsymbol{c}_i \boldsymbol{r}_i^{\top}, \boldsymbol{x}_j \boldsymbol{y}_j^{\top}\rangle_F = \text{Tr}(\boldsymbol{r}_i \boldsymbol{c}_i^{\top}\boldsymbol{x}_j \boldsymbol{y}_j^{\top}) = (\boldsymbol{c}_i^{\top}\boldsymbol{x}_j)(\boldsymbol{r}_i^{\top} \boldsymbol{y}_j) \\[5pt] =&\, [(\boldsymbol{C}^{\top}\boldsymbol{X})\otimes (\boldsymbol{R}\boldsymbol{Y}^{\top})]_{i,j} \end{aligned}\end{equation}

that is, $\boldsymbol{U}^{\top}\boldsymbol{V}=(\boldsymbol{C}^{\top}\boldsymbol{X})\otimes (\boldsymbol{R}\boldsymbol{Y}^{\top}),\boldsymbol{U}^{\top}\boldsymbol{U}=(\boldsymbol{C}^{\top}\boldsymbol{C})\otimes (\boldsymbol{R}\boldsymbol{R}^{\top})$, where $\otimes$ denotes the Hadamard product. After this identity transformation, the computational cost of $\boldsymbol{U}^{\top}\boldsymbol{V}$ and $\boldsymbol{U}^{\top}\boldsymbol{U}$ is reduced.

Summary

This post introduced the CR approximation of matrix multiplication, a low-rank approximation with a specific row/column structure. Compared with the optimal low-rank approximation given by SVD, the CR approximation has a more intuitive physical meaning and better interpretability.