Accelerating Retrieval for Interactive Similarity Models via CUR Decomposition

利用CUR分解加速交互式相似度模型的检索 苏剑林 · 2022-11-02 · Machine Learning matrix semantics semantic similarity

Text similarity can be computed in two ways: "interactive" and "representation-based," and I suspect most readers are already familiar with this distinction. I've previously written a post CoSENT (II): How Big Is the Gap Between Representation-Based and Interactive Matching? comparing the two approaches. Generally speaking, interactive similarity tends to give better results, but using it directly for large-scale retrieval is impractical, whereas representation-based similarity offers much faster retrieval at the cost of somewhat weaker performance.

So how do we speed up retrieval for interactive similarity while preserving its accuracy? This has long been an active research question. Recently, the paper Efficient Nearest Neighbor Search for Cross-Encoder Models using Matrix Factorization offers a new answer: CUR decomposition.

Problem Analysis

In a retrieval scenario, we typically have a huge candidate set to search through $\mathcal{K}$; without loss of generality, we can assume that $\mathcal{K}$ stays fixed. The retrieval task is: given any query $q\in\mathcal{Q}$, find the items in $\mathcal{K}$ that are most relevant to $q$, i.e. $k$. An interactive similarity model directly trains a relevance scoring function $s(q,k)$, and in principle we can compute $s(q,k)$ for any $k\in\mathcal{K}$ and then sort in descending order. But this means that the computational cost of every single query is $\mathcal{O}(|\mathcal{K}|)$, and intermediate results can't be cached, so the cost quickly becomes unacceptable.

What is computationally affordable is a similarity function with a matrix-factorization structure. For a single pair of samples, this means an inner-product-based similarity (or a variant thereof): the classic implementation encodes $q,k$ via an encoder $\boldsymbol{E}$ into two vectors $\boldsymbol{E}(q),\boldsymbol{E}(k)$, and then computes the inner product $\boldsymbol{E}(q)^{\top}\boldsymbol{E}(k)$ — this is representation-based similarity. Such a scheme has a few nice properties: 1) all the $\boldsymbol{E}(k)$ can be precomputed and cached; 2) computing the inner product of $\boldsymbol{E}(q)$ against all the $\boldsymbol{E}(k)$ can be turned into a matrix multiplication, which is easy to parallelize for fast computation; 3) it can be further sped up using approximate-search tools like Faiss.

So the idea for accelerating retrieval with interactive similarity is to convert it into a matrix-factorization form. A fairly classical approach is to distill a representation-based similarity model from the interactive similarity model. What's elegant about this Google paper is that it doesn't introduce any new model at all — instead, it uses CUR decomposition directly on top of the original interactive similarity model to achieve the speedup. This scheme is named ANNCUR.

Matrix Factorization

CUR decomposition is one type of matrix factorization. When people think of matrix factorization, SVD is probably the first thing that comes to mind — but the reason SVD feels so familiar isn't that it's especially intuitive, it's just that it gets covered so often. When it comes to intuitiveness, CUR decomposition clearly has the edge.

In fact, we can understand both SVD and CUR decomposition from a unified point of view. For a scoring function $s(q,k)$, we want to construct the following approximation:

\begin{equation}s(q,k) \approx \sum_{u\in\mathcal{U},v\in\mathcal{V}} f(q, u) g(u, v) h(v, k)\label{eq:decom}\end{equation}

Generally there's a size constraint $|\mathcal{U}|,|\mathcal{V}|\ll |\mathcal{Q}|,|\mathcal{K}|$, which makes this a compressive factorization. We can think of $\mathcal{U}$ as a "representative set" (or "cluster centers" — it's just an intuitive picture, feel free to use whichever framing you like) for $\mathcal{K}$, and correspondingly $\mathcal{V}$ as a representative set for $\mathcal{Q}$. With this picture, the decomposition above becomes quite intuitive:

The score between $q,k$, $s(q,k)$, is approximated by first scoring $q$ against the "representatives" $u\in \mathcal{U}$ of $\mathcal{K}$ to get $f(q, u)$, then scoring $k$ against the "representatives" $v\in \mathcal{V}$ of $\mathcal{Q}$ to get $h(v, k)$, and finally taking a weighted sum with weights $g(u, v)$.

In other words, the direct interaction between $q$ and $k$ is replaced by each of them interacting with the "representatives," and then combining the results with weights. The advantage of this is clear: once $f,g,h$ are fixed, all the $g(u,v)$ and $h(v,k)$ can be precomputed and cached as a matrix. Then at retrieval time we only need to compute $f(q,u)$ a total of $|\mathcal{U}|$ times, followed by one matrix multiplication (i.e., inner-product-based retrieval). So the computational cost of retrieval drops from $\mathcal{O}(|\mathcal{K}|)$ to $\mathcal{O}(|\mathcal{U}|)$ (and with tools like Faiss, inner-product-based retrieval can be further approximated down to roughly $\mathcal{O}(1)$, so it can be neglected).

Suppose the query set $\mathcal{Q}$ is also finite. Then all the $s(q,k)$ form a $|\mathcal{Q}|\times |\mathcal{K}|$ matrix $\boldsymbol{S}$, and correspondingly $f(q, u),g(u, v),h(v, k)$ correspond to a $|\mathcal{Q}|\times |\mathcal{U}|$ matrix $\boldsymbol{F}$, a $|\mathcal{U}|\times |\mathcal{V}|$ matrix $\boldsymbol{G}$, and a $|\mathcal{V}|\times |\mathcal{K}|$ matrix $\boldsymbol{H}$, respectively. Equation $\eqref{eq:decom}$ then becomes a matrix factorization:

\begin{equation}\begin{array}{ccccc} \boldsymbol{S} & \approx & \boldsymbol{F} & \boldsymbol{G} & \boldsymbol{H} \\ \in\mathbb{R}^{|\mathcal{Q}|\times |\mathcal{K}|} & & \in\mathbb{R}^{|\mathcal{Q}|\times |\mathcal{U}|} & \in\mathbb{R}^{|\mathcal{U}|\times |\mathcal{V}|} & \in\mathbb{R}^{|\mathcal{V}|\times |\mathcal{K}|} \end{array}\label{eq:m-decom}\end{equation}

CUR Decomposition

If we restrict $\boldsymbol{G}$ to be diagonal, while placing no special constraints on $\boldsymbol{F}$ and $\boldsymbol{H}$, the resulting factorization is exactly SVD. SVD effectively conjures up some number of virtual "representatives" that give a good fit, but since these "representatives" are constructed automatically by the algorithm, it's hard to attach any concrete meaning to them — that is, they're not very interpretable.

CUR decomposition is more intuitive: it insists that the "representatives" should be actual members of the original population, i.e. the representatives of $\mathcal{Q},\mathcal{K}$ should be a subset picked from that very set, $\mathcal{U}\subset \mathcal{K}, \mathcal{V}\subset\mathcal{Q}$. As a result, $(q,u)$ and $(v,k)$ are themselves among the original $(q,k)$, so we can reuse the same scoring function $s$ from $(q,k)$, namely

\begin{equation}s(q,k) \approx \sum_{u\in\mathcal{U},v\in\mathcal{V}} s(q, u) g(u, v) s(v, k)\end{equation}

The only thing left to determine, then, is $g(u,v)$. From the matrix-factorization perspective, $\boldsymbol{F}$ in equation $\eqref{eq:m-decom}$ is just a submatrix formed by several columns of $\boldsymbol{S}$, and $\boldsymbol{H}$ is a submatrix formed by several rows of $\boldsymbol{S}$; what remains to be computed is the matrix $\boldsymbol{G}$. Computing $\boldsymbol{G}$ is also fairly intuitive. Let's first consider a very special case, $\mathcal{U}=\mathcal{K},\mathcal{V}=\mathcal{Q}$ and $|\mathcal{Q}|=|\mathcal{K}|$, in which the CUR decomposition becomes $\boldsymbol{S}\approx \boldsymbol{S}\boldsymbol{G}\boldsymbol{S}$, with $\boldsymbol{S}$ and $\boldsymbol{G}$ both square matrices. Since in this case we've already taken the entirety of $\mathcal{Q},\mathcal{K}$ as the representative set, we'd naturally want the approximation to be exact, i.e. $=$ rather than merely $\approx$; taking $=$, we can solve directly to get $\boldsymbol{G} = \boldsymbol{S}^{-1}$.

However, this requires $\boldsymbol{S}$ to be invertible, which need not hold in general. To handle this, we generalize the notion of matrix inverse to what's called the "pseudoinverse," denoted $\boldsymbol{G}=\boldsymbol{S}^{\dagger}$. In particular, the pseudoinverse is also defined for non-square matrices, so when $|\mathcal{Q}|\neq|\mathcal{K}|$, we can likewise solve for $\boldsymbol{G}=\boldsymbol{S}^{\dagger}$. Finally, when $\mathcal{U}\neq\mathcal{K}$ or $\mathcal{V}\neq\mathcal{Q}$, the result is similar, except that the matrix we take the pseudoinverse of is replaced by the intersection matrix $\boldsymbol{F}\cap\boldsymbol{H}$ of $\boldsymbol{F}$ and $\boldsymbol{H}$ (that is, the $\mathcal{U}\times \mathcal{V}$ matrix formed by the entries at the intersection of the selected rows and columns of $\boldsymbol{S}$):

\begin{equation} \boldsymbol{S} \approx \boldsymbol{F} (\boldsymbol{F}\cap \boldsymbol{H})^{\dagger}\boldsymbol{H}\end{equation}

The whole process is illustrated below:

Illustration of CUR decomposition

Accelerating Retrieval

This is actually not the first time this blog has touched on CUR decomposition. Back at the start of last year, the post Nyströmformer: A Linearized Attention Scheme Based on Matrix Factorization introduced Nyströmformer, which is in fact also designed around the idea of CUR decomposition — the original paper even spends quite a lot of space introducing CUR decomposition. ANNCUR, meanwhile, uses CUR decomposition to speed up retrieval, which shows just how broadly applicable CUR is.

I've already touched briefly on the principle behind the speedup; let me now summarize it properly. First, we pick a number of representative $q\in \mathcal{V}\subset \mathcal{Q}$'s and $k\in \mathcal{U}\subset \mathcal{K}$'s, compute their pairwise scores to form the matrix $\boldsymbol{F}\cap\boldsymbol{H}$, and take its pseudoinverse to get the matrix $\boldsymbol{G}$. We then precompute the score matrix $\boldsymbol{G}$ between $q\in \mathcal{V}$ and $k\in \mathcal{K}$, and cache $\boldsymbol{G}\boldsymbol{H}$. Finally, for every query $q$ we need to handle, we compute its score against every $k\in \mathcal{U}$, getting a $|\mathcal{U}|$-dimensional vector; multiplying this vector by the cached matrix $\boldsymbol{G}\boldsymbol{H}$ then gives us the score vector between $q$ and every $k\in \mathcal{K}$.

That's the gist of how ANNCUR works. For the finer details, I encourage readers to consult the original paper — for instance, it shows that using the paper's "[EMB]-CE" variant in place of the plain interactive similarity model gives even better results. Some readers might be wondering, "how exactly are the representative $q,k$'s chosen?" In practice, in most cases they're simply chosen at random, which leaves some room for improvement — for instance, one could cluster the data first and then pick the point closest to each cluster center. There's plenty of room here for readers to experiment on their own. One more thing worth noting: CUR decomposition is itself only an approximation, so it inevitably introduces some error. This is precisely why the acceleration scheme is designed with retrieval in mind — retrieval scenarios generally care more about top-k recall than about the precision of the single top-1 result. So we can use CUR-accelerated retrieval to recall a handful of candidates first, and then apply the exact $s(q,k)$ to re-rank them for higher accuracy.

Partial experimental results from ANNCUR

Summary

This post has revisited CUR decomposition as a form of matrix factorization, and introduced its application to accelerating retrieval for interactive similarity models.