最大化内积搜索相关研究 (Maximum Inner Product Search, MIPS)

前言

如果你对这篇文章感兴趣，可以点击「【访客必读 - 指引页】一文囊括主页内所有高质量博客」，查看完整博客分类与对应链接。

MIPS 问题即在一个向量集合 $\mathcal{S}$ 中，找到一个与查询向量 $q$ 内积最大的向量 $z$，即：
$$z=\underset{x\in \mathcal{S}}{\mathrm{argmax}}{x}^{\top }q.$$

这是一个非常困难的问题，本文罗列了部分与其相关的资料。

---

Background

首先介绍一下局部敏感哈希 (Locality Sensitive Hashing, LSH)，其是一个函数族，其中的函数能使相似的元素经过映射后，有较高的碰撞概率，具体定义如下：

• A family $\mathcal{H}$ is called $\left(S_0,c{S}_0,p_1,p_2\right)$-sensitive if, for any two point $x,y\in {\mathbb{R}}^D$, $h$ chosen uniformly from $\mathcal{H}$ satisfies the following:
  • if $\mathrm{Sim}(x,y)\ge S_0$ then ${\mathrm{Pr}}_{\mathcal{H}}(h(x)=h(y))\ge p_1$;
  • if $\mathrm{Sim}(x,y)\le c{S}_0$ then ${\mathrm{Pr}}_{\mathcal{H}}(h(x)=h(y))\le p_2$.

其中 $\mathrm{Sim}(x,y)$ 为两个点的相似度，当 $p_1>p_2$ 且 $c<1$ 时，$h\in \mathcal{H}$ 可以做到高效的基于 $\mathrm{Sim}(\cdot ,\cdot )$ 度量的近似近邻查搜。

简单说就是，LSH 是一种哈希技术，其可以在某些 $\mathrm{Sim}(\cdot ,\cdot )$ 度量上为查询向量 $x$ 实现高效的查搜，例如：

• [SGG04 - Mayur Datar] 实现了 $L_p(p\in (0,2])$ 距离下的高效查搜；
• [STOC02 - Moses S. Charikar] 实现了余弦相似度下的高效查搜。

L2LSH

L2LSH 是 LSH 基于欧式距离的一种实现，具体如下：

SRP

SRP (signed random projections) 是 LSH 基于余弦相似度的一种实现，具体如下：

• 生成随机向量 $a$，其中 $a_i\sim \mathcal{N}(0,1)$；
• 采用如下哈希函数：
  $$h^{Sign}(x)=sign(a^{\top }x)=\{\begin{array}{ll}1& \text{ if }{a}^{\top }x\ge 0\\ 0& \text{ if }{a}^{\top }x<0\end{array},$$
• 对应下述碰撞概率（$1-{\cos }^{-1}(z)$ 关于 $z\in [-1,1]$ 单调递增）：
  $$Pr\left(h^{Sign}(x)=h^{Sign}(y)\right)=1-\frac{1}{\pi }{\cos }^{-1}\left(\frac{x^{T}y}{\Vert x\Vert \Vert y\Vert }\right).$$

---

NIPS14 - ALSH

此篇文章首先证明了「LSH 系列无法求解 MIPS」，随后对 LSH 进行扩展，得到 ALSH (Asymmetric Locality Sensitive Hashing)，其定义为：

• A family $\mathcal{H}$, along with the two vector functions $Q:{\mathbb{R}}^D\mapsto {\mathbb{R}}^{D^{\prime }}$ (Query Transformation) and $P:{\mathbb{R}}^D\mapsto {\mathbb{R}}^{D^{\prime }}$ (Preprocessing Transformation), is called $\left(S_0,c{S}_0,p_1,p_2\right)$-sensitive if, for a given $c$-NN instance with query $q$ and any $x$ in the collection $\mathcal{S}$, the hash function $h$ chosen uniformly from $\mathcal{H}$ satisfies the following:
  • if $\mathrm{Sim}(q,x)\ge S_0$ then ${\mathrm{Pr}}_{\mathcal{H}}(h(Q(q)))=h(P(x)))\ge p_1$;
  • if $\mathrm{Sim}(q,x)\le c{S}_0$ then ${\mathrm{Pr}}_{\mathcal{H}}(h(Q(q))=h(P(x)))\le p_2$.

之后该篇文章指出若 $\mathcal{H}$ 满足上述性质，则对于查询 $q$，其可以做到：

• 若 $\mathcal{S}$ 中存在 $x$ 满足 $\mathrm{Sim}(x,q)\ge S_0$，则可以高概率找到满足 $\mathrm{Sim}(x,q)\ge c{S}_0$ 的 $x$；
• 其中每次查询所需时间为 $O(n^{\rho }\log n)$，总的空间复杂度为 $O(n^{1+\rho })$，$\rho =\log p_1/\log p_2$.

随后，基于 ALSH 的定义，此篇文章将内积度量转化为了欧式距离的度量，具体如下：

• 将 $\mathcal{S}$ 中的 $x$ 通过下述转换，实现 $\Vert x_i{\Vert }_2\le U<1$：
  $$S(x)=\frac{U}{M}\times x;M=\underset{x_i\in \mathcal{S}}{\max }{\Vert x_i\Vert }_2;$$
• 随后令 $\Vert q{\Vert }_2=1$，由于 $q$ 是查询向量，因此该假设也可通过转换实现;
• 最后通过下述操作，将 $q^{\top }x$ 转化为了 $\Vert Q(q)-P(x)\Vert$，由此可通过之前的 L2LSH 进行查搜。

通过转换，MIPS 问题可以通过 L2LSH 进行解决，具体的理论保障如下：

在 worst case 下，每次查询时间复杂度为 $O(n^{\rho }\log n)$，$\rho$ 定义如下：
$$\rho =\frac{\log F_r\left(\sqrt{1+m/4-2S_0+U^{2^{m+1}}}\right)}{\log F_r(\sqrt{1+m/4-2c{S}_0)}}.$$

其中 $U,m,r,S_0,c$ 为超参，在实践中，一般设置 $m=3,U=0.83,r=2.5$ 或者令 $S_0=0.9U\text{ or }0.8U$，然后对 $m,U,r$ 进行网格搜索，目标是最小化 $\rho$。

---

arXiv14 - Sign-ALSH

ALSH 通过对查询 $q$ 和 $\mathcal{S}$ 中元素 $x$ 采用不同的变换，实现了内积到欧式距离的转换，该篇文章采用同样的思路，但通过不同的变换形式，得到了 Sign-ALSH，实现内积到余弦相似度的转换。此篇文章从实验方面指出 Sign-ALSH 比 ALSH 更有效。

Sign-ALSH 的具体步骤如下：

• 类似于 ALSH 中的方式，令 $\Vert q{\Vert }_2=1,\Vert x_i{\Vert }_2\le U<1,\forall x_i\in \mathcal{S}$,
• 定义两个向量映射函数，$P:{\mathbb{R}}^D\mapsto {\mathbb{R}}^{D+m}$ 与 $Q:{\mathbb{R}}^D\mapsto {\mathbb{R}}^{D+m}$：
  $$\begin{array}{rl}& P(x)=\left[x;1/2-\Vert x{\Vert }_2^2;1/2-\Vert x{\Vert }_2^4;\dots ;1/2-\Vert x{\Vert }_2^{2^m}\right],\\ & Q(x)=[x;0;0;\dots ;0],\end{array}$$
• 根据上述映射，可以得到下述等式：
  $$\frac{Q(q{)}^{T}P\left(x_i\right)}{\Vert Q(q){\Vert }_2{\Vert P\left(x_i\right)\Vert }_2}=\frac{q^T{x}_i}{\sqrt{m/4+{\Vert x_i\Vert }_2^{2^{m+1}}}},$$
• 由于 $\Vert x_i{\Vert }^{2^{m+1}}\to 0$，可得到下述近似结果，完成内积到余弦相似度的转换：
  $$\underset{x\in \mathcal{S}}{\mathrm{argmax}}{q}^{T}x\simeq \underset{x\in \mathcal{S}}{\mathrm{argmax}}\frac{Q(q{)}^{T}P\left(x_i\right)}{\Vert Q(q){\Vert }_2{\Vert P\left(x_i\right)\Vert }_2}.$$

对应的理论保障如下：

文中推荐超参为 $(m,U)=(2,0.75)$ 或 $(m,U)=(3,0.85)$。在实践中，可以生成 $L$ 个哈希表，每个哈希表中使用 $K$ 组随机投影的哈希函数，将原始维度降为 $K$ 维，若 $K$ 维均相同，则实现匹配，再将 $L$ 个哈希表中所有匹配上的元素集合在一起。此处的 $K$ 越大，识别的准确率越高，$L$ 越大，识别的召回率越高。

---

ICCV15 - AIBC

上述两种哈希方法均与数据集无关，即属于 data-independent methods，而本文将给出一种 data-dependent methods，整体思想是根据训练集，学习一种非对称的映射方式。

假设有两组数据，$\mathbf{A}=[{\mathbf{a}}_1,{\mathbf{a}}_2,...,{\mathbf{a}}_n]$ 与 $\mathbf{X}=[{\mathbf{x}}_1,{\mathbf{x}}_2,...,{\mathbf{x}}_m]$，其中 ${\mathbf{a}}_i\in {\mathbb{R}}^{d\times 1}$，${\mathbf{x}}_i\in {\mathbb{R}}^{d\times 1}$。另外，有相似度矩阵 $\mathbf{S}\in {\mathbb{R}}^{n\times m}$，其中 ${\mathbf{S}}_{ij}$ 表示 ${\mathbf{a}}_i$ 和 ${\mathbf{x}}_j$ 之间的相似度。

目标是学出哈希映射 $h(\mathbf{x})=\mathrm{sgn}({\mathbf{W}}^{\top }\mathbf{x})$ 和 $z(\mathbf{x})=\mathrm{sgn}({\mathbf{R}}^{\top }\mathbf{x})$，可以将原始输入映射成二进制码，其中 $\mathbf{W},\mathbf{R}\in {\mathbb{R}}^{d\times r}$，具体优化过程如下：
$$\underset{\mathbf{W},\mathbf{R}}{\min }{\Vert \mathrm{sgn}{\left({\mathbf{W}}^{\top }\mathbf{A}\right)}^{\top }\mathrm{sgn}\left({\mathbf{R}}^{\top }\mathbf{X}\right)-\mathbf{S}\Vert }^2.$$

上述优化问题难以求解，因为本文采用了交替优化，即固定 $\mathbf{R}$ 求 $\mathbf{W}$，或固定 $\mathbf{W}$ 求 $\mathbf{R}$，以前者为例（后者同理），令 $\mathrm{sgn}\left({\mathbf{R}}^{\top }\mathbf{X}\right)=\mathbf{Z}$，得到：
$$\underset{\mathbf{W}}{\min }{\Vert \mathrm{sgn}{\left({\mathbf{W}}^{\top }\mathbf{A}\right)}^{\top }\mathbf{Z}-\mathbf{S}\Vert }^2.$$

该问题依然难以求解，因此引入矩阵 $\mathbf{B}$，上述问题转化为：
min ⁡ B , W ∥ B ⊤ Z − S ∥ 2 + λ ∥ B − W ⊤ A ∥ 2  s.t.  B ∈ { − 1 , 1 } r × n \begin{aligned} & \min _{\mathbf{B}, \mathbf{W}}\left\|\mathbf{B}^{\top} \mathbf{Z}-\mathbf{S}\right\|^2+\lambda\left\|\mathbf{B}-\mathbf{W}^{\top} \mathbf{A}\right\|^2 \\ & \text { s.t. } \quad \mathbf{B} \in\{-1,1\}^{r \times n} \end{aligned} ​B,Wmin​ ​B⊤Z−S ​2+λ ​B−W⊤A ​2 s.t. B∈{−1,1}r×n​

问题依然难以求解，继续采用交替优化，固定 $\mathbf{B}$，$\mathbf{W}={\left({\mathbf{AA}}^{\top }\right)}^{-1}\mathbf{A}{\mathbf{B}}^{\top }$；固定 $\mathbf{W}$，优化问题为：
min ⁡ B ∥ B ⊤ Z ∥ 2 − 2 trace ⁡ ( B ⊤ D )  s.t.  B ∈ { − 1 , 1 } r × n \begin{aligned} & \min _{\mathbf{B}}\left\|\mathbf{B}^{\top} \mathbf{Z}\right\|^2-2 \operatorname{trace}\left(\mathbf{B}^{\top} \mathbf{D}\right) \\ & \text { s.t. } \quad \mathbf{B} \in\{-1,1\}^{r \times n} \end{aligned} ​Bmin​ ​B⊤Z ​2−2trace(B⊤D) s.t. B∈{−1,1}r×n​

其中 $\mathbf{D}=\mathbf{Z}{\mathbf{S}}^{\top }+\lambda {\mathbf{W}}^{\top }\mathbf{A}$。此时问题仍然无法直接求解，因此继续交替优化，每次求解 $\mathbf{B}$ 中一行，固定其它 $r-1$ 行，此时问题可以高效求解。

该篇文章整体思路如上所示，具体细节可参照原文。

---

ICML16 - Learning and Inference via MIPS

本文指出了一种较 Sign-ALSH 更为简单的构造，即 Simple-LSH [ICML15 - Behnam Neyshabur]，同样假设向量集合中的 $x$ 满足 $\Vert x{\Vert }^2\le 1$，查询向量 $q$ 满足 $\Vert q\Vert =1$，令 $P(x)=(x,\sqrt{1-\Vert x{\Vert }^2})$，$Q(q)=(q,0)$，可以得到：
$$\frac{P(x{)}^{\top }Q(q)}{\Vert P(x)\Vert \Vert Q(q)\Vert }=x^{\top }q,$$

实现 MIPS 到 MCSS (Maximum Cosine Similarity Search) 的等价转换。但本文的主要目的并不是求解 MIPS，而是借用 MIPS 的求解来加速 Log-Linear 模型的学习和推理过程，即 MIPS 的应用。不过本文介绍了 Gumbel Variables 技术，虽然与 MIPS 无关，但感觉很有趣，便在此顺带记录一下。

Gumbel Variables

一个 Gumbel 随机变量 $G$ 是一个连续随机变量，其累计分布函数 (Cumulative Distribution Function, CDF) 为：
$$P(G<x)=\exp (-\exp (-\frac{x-\mu }{\beta })).$$

该篇文章取 $\mu =0,\beta =1$。从 Gumbel 分布中采样的方式为：从 $[0,1]$ 区间中均匀采样 $U$，有次得到 Gumbel 随机变量 $G$ 的采样值：
$$G=-\log (-\log (U)).$$

Gumbel 随机变量可以将一个优化问题转化为一个采样问题，具体来说，令 { a i } \{a_i\} {ai​} 为一组系数， { G i } \{G_i\} {Gi​} 和 $G$ 为独立的 Gumbel 随机变量，则有如下两条性质：
max ⁡ i { a i + G i } ∼ log ⁡ ( ∑ i e a i ) + G , argmax ⁡ i { a i + G i } ∼ Multinomial ⁡ ( { e a i ∑ i e a i } ) . \begin{gathered} \max _i\left\{a_i+G_i\right\} \sim \log \left(\sum_i e^{a_i}\right)+G, \\ \underset{i}{\operatorname{argmax}}\left\{a_i+G_i\right\} \sim \operatorname{Multinomial}\left(\left\{\frac{e^{a_i}}{\sum_i e^{a_i}}\right\}\right). \end{gathered} imax​{ai​+Gi​}∼log(i∑​eai​)+G,iargmax​{ai​+Gi​}∼Multinomial({∑i​eai​eai​​}).​

简单说，上述两条性质的「左边」对应一个优化问题，其最优解服从「右边」对应的分布，因此可以将求解「左边」这个优化问题的过程，转化为在「右边」这个分布里采样的过程。

---

KDD18 - H2-ALSH

此篇文章在 Preliminaries 部分介绍了另外一种将 MIPS 转为欧式距离的方式，即 XBOX 变换。具体来说，令 $M={\max }_{o\in D}\Vert o\Vert$，$P(o)=[o;\sqrt{M^2-\Vert o{\Vert }^2}]$，$Q(q)=[q;0]$，可以得到：
$$\Vert Q(q)-P(o){\Vert }^2=\Vert q{\Vert }^2+M^2-2o^{\top }q,$$

是一种无损的变换。我们的核心目标是 $c$-Approximate MIP search ($c$-AMIP search)，即给定一个近似比例 $c$ $(0<c<1)$ 以及一个查询 $q$，希望找到 $o\in D$ 满足 $o^{\top }q\ge c(o^{\ast }{)}^{\top }q$，其中 $o^{\ast }$ 为最优解。

此篇文章指出 XBOX 转换的缺陷，即对于某些 $\Vert q\Vert$ 来说，经过转换，大部分的 $P(o)$ 都将距离 $Q(q)$ 非常近，进而产生较大误差。以下图为例：

$o_2$ 原本距 $q_1$ 比 $o_1$ 近不少，但经过转换后，$P(o_2)$ 距 $Q(q_1)$ 的距离和 $P(o_1)$ 非常接近，即转换后产生失真误差。

本文指出使用 QNF 变换，可以降低上述失真误差，即：
$$\begin{array}{c}P(o)=\left[o_1,o_2,\dots ,o_d;\sqrt{M^2-\Vert o{\Vert }^2}\right],\\ Q(q)=\left[\lambda q_1,\lambda q_2,\dots ,\lambda q_d;0\right],\text{ where }\lambda =\frac{M}{\Vert q\Vert },\\ \Vert Q(q)-P(o){\Vert }^2=M^2+{\lambda }^2\Vert q{\Vert }^2-2\lambda o^{\top }q.\end{array}$$

依旧为无损转换，且 $q$ 和 $o$ 均被映射至圆上，可以缓解 XBOX 转换的问题，如下所示：

由于 $o^{\top }q=\Vert o\Vert \Vert q\Vert \cos \beta$，因此本文先对 $\Vert o\Vert$ 做了一个分组，然后在每个分组中再将内积转成欧式距离，使用 QALSH 方法建立哈希表进行查搜。

下述是构建过程，即将 $n$ 个元素按照 $\Vert o\Vert$ 分成 $K$ 组：

在具体查搜时，先从 $S_1$ 开始查，其中 $b{M}_1<\Vert o\Vert \le M_1$，在其中根据 QALSH 进行查搜，记录最优值 $o_{mip}$；若 $o_{mip}>M_2\Vert q\Vert$，则搜索结束，否则进入 $S_2$，整个过程循环下去。

---

NeurIPS18 - Norm-Ranging LSH

该篇文章与 H2-ALSH 的思想非常相似，即按照集合元素的 norm $\Vert x\Vert$ 进行分组，然后在组内用 Simple-LSH 建立索引，具体查搜过程如下：

---

ICML19 - Sample MIPS

本篇文章在相关工作部分介绍了 Greedy-MIPS，其有下述假设：
$${\mathbf{q}}^{\top }{\mathbf{x}}_j>{\mathbf{q}}^{\top }{\mathbf{x}}_i\iff \underset{t=1,\dots ,d}{\max }{q}_t{x}_{jt}>\underset{t=1,\dots ,d}{\max }{q}_t{x}_{it},$$

即内积最大的向量，在内积求和时，最大的分量依然是最大的。这个假设很强，因为有了这个假设后，找内积最大的向量，就变成对每一个维度，找一个乘机最大的分量，难度瞬间下降。实际操作中，一般是枚举每一个维度，每一个维度找出乘机最大的前 $k$ 个，再在所有挑出来的向量中找一个内积最大的，作为结果输出。本篇文章称该方法实际效果不错，故再次记录一下。

回归正题，本篇文章将「找 MIPS 最大向量」转化为了一个采样问题，其定义如下概率：
$$\begin{array}{c}P(j\mid \mathbf{q})\propto {\mathbf{q}}^{\top }{\mathbf{x}}_j,\\ P(j,t\mid \mathbf{q})\propto q_t{x}_{jt}.\end{array}$$

由此可以得到：
$$\begin{array}{rl}P(t\mid \mathbf{q})& =\sum _{j=1}^{n}P(j,t\mid \mathbf{q})\propto \sum _{j=1}^n{q}_t{x}_{jt}=q_t{s}_t,\\ P(j\mid t,\mathbf{q})& =\frac{P(j,t\mid \mathbf{q})}{P(t\mid \mathbf{q})}\\ & \propto \frac{q_t{x}_{jt}}{{\sum }_{j=1}^n{q}_t{x}_{jt}}=\frac{x_{jt}}{{\sum }_{j=1}^n{x}_{jt}}=\frac{x_{jt}}{s_t}.\end{array}$$

又因为 $P(j\mid \mathbf{q})={\sum }_{t=1}^{d}P(t\mid \mathbf{q})P(j\mid t,\mathbf{q})$，以及 $P(j\mid t,\mathbf{q})$ 数值与 $\mathbf{q}$ 无关，即对 $P(j\mid \mathbf{q})$ 的采样可以转化为两步：

• 根据 $\mathbf{q}$ 和 $P(t\mid \mathbf{q})$ 采样 $t$，时间复杂度为 $O(d)$；
• 根据 $t$ 和 $P(j\mid t,\mathbf{q})$ 采样 $j$，与 $\mathbf{q}$ 无关，可事先预处理，每次查询的时间复杂度为 $O(1)$。

上述过程将 $O(nd)$ 的暴力搜索转化为了 $O(Bd)$（$B$ 为采样次数），实现了对 MIPS 问题的高效求解。

由于 $\mathbf{q}$ 和 $\mathbf{x}$ 中元素可能有负数，因此文中给出了一种针对负数的解决方案，令：
$$\begin{array}{c}P(j\mid \mathbf{q})\propto \sum _{t=1}^d|q_t{x}_{jt}|,\\ P(j\mid t,\mathbf{q})\propto |q_t{x}_{jt}|.\end{array}$$

由此得到：
$$\begin{array}{c}P(t\mid \mathbf{q})\propto \sum _{j=1}^n|q_t{x}_{jt}|,\\ P(j\mid t,\mathbf{q})\propto \frac{|x_{jt}|}{{\sum }_{j=1}^n|x_{jt}|}.\end{array}$$

根据 $P(t\mid \mathbf{q})$，$P(j\mid t,\mathbf{q})$ 采样 $B$ 次，得到 $B$ 组 $(j,t)$ 对，$\mathrm{sgn}(q_t{x}_{jt})$ 之和越大，选中 ${\mathbf{x}}_j$ 的概率越大。

在理论保障方面，文中在 Lemma1 处证明了：
$$\begin{array}{c}\mathbb{E}[z_{jb}]=\frac{{\mathbf{q}}^{\top }{\mathbf{x}}_j}{S},\\ \mathbb{E}[z_j]=B\cdot \mathbb{E}[z_{jb}],\end{array}$$

其中 $z_{jb}=\mathrm{sgn}(q_t{x}_{jt})$，$z_j={\sum }_{b=1}^B{x}_{jb}$，$S={\sum }_{t=1}^d{\sum }_{j=1}^n|q_t{x}_{jt}|$。

---

参考资料

• [STOC02 - Moses S. Charikar] Similarity Estimation Techniques from Rounding Algorithms
• [SGG04 - Mayur Datar] Locality-sensitive hashing scheme based on p-stable distributions
• [NIPS14 - Anshumali Shrivastava] Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)
• [arXiv14 - Anshumali Shrivastava] Improved Asymmetric Locality Sensitive Hashing (ALSH) for Maximum Inner Product Search (MIPS)
• [arXiv14 - Ping Li] Coding for Random Projections and Approximate Near Neighbor Search
• [ICCV15 - Fumin Shen] Learning Binary Codes for Maximum Inner Product Search
• [ICML15 - Behnam Neyshabur] On Symmetric and Asymmetric LSHs for Inner Product Search
• [ICML16 - Stephen Mussmann] Learning and Inference via Maximum Inner Product Search
• [KDD18 - Qiang Huang] Accurate and Fast Asymmetric Locality-Sensitive Hashing Scheme for Maximum Inner Product Search
• [NIPS18 - Xiao Yan] Norm-Ranging LSH for Maximum Inner Product Search
• [ICML19 - Qin Ding] A Fast Sampling Algorithm for Maximum Inner Product Search