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

MIPS 问题即在一个向量集合 $\mathcal{S}$ 中，找到一个与查询向量 $q$ 内积最大的向量 $z$，即：

$$ z=\mathop{\arg \max}\limits _{x \in \mathcal{S}} 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 $\operatorname{Sim}(x, y) \geq S_0$ then $\operatorname{Pr}_{\mathcal{H}}(h(x)=h(y)) \geq p_1$;
if $\operatorname{Sim}(x, y) \leq c S_0$ then $\operatorname{Pr}_{\mathcal{H}}(h(x)=h(y)) \leq p_2$.

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

简单说就是，LSH 是一种哈希技术，其可以在某些 $\operatorname{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{cases}1 & \text { if } a^\top x \geq 0 \\ 0 & \text { if } a^\top x<0\end{cases}, $$

对应下述碰撞概率（$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}{\|x\|\|y\|}\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 $\operatorname{Sim}(q, x) \geq S_0$ then $\left.\operatorname{Pr}_{\mathcal{H}}(h(Q(q)))=h(P(x))\right) \geq p_1$;
if $\operatorname{Sim}(q, x) \leq c S_0$ then $\operatorname{Pr}_{\mathcal{H}}(h(Q(q))=h(P(x))) \leq p_2$.

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

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

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

将 $\mathcal{S}$ 中的 $x$ 通过下述转换，实现 $|x_i|_2\leq U<1$：

$$ S(x)=\frac{U}{M} \times x ; \quad M=\max _{x_i \in \mathcal{S}}\left\|x_i\right\|_2; $$

随后令 $|q|_2=1$，由于 $q$ 是查询向量，因此该假设也可通过转换实现;
最后通过下述操作，将 $q^\top x$ 转化为了 $|Q(q)-P(x)|$，由此可通过之前的 L2LSH 进行查搜。

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

在 worst case 下，每次查询时间复杂度为 $O(n^\rho \log n)$，$\rho$ 定义如下：

$$ \rho=\frac{\log F_r\left(\sqrt{1+m / 4-2 S_0+U^{2^{m+1}}}\right)}{\log F_r\left(\sqrt{\left.1+m / 4-2 c S_0\right)}\right.}. $$

其中 $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 中的方式，令 $|q|_2=1, |x_i|_2\leq 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{aligned} & P(x)=\left[x ; 1 / 2-\|x\|_2^2 ; 1 / 2-\|x\|_2^4 ; \ldots ; 1 / 2-\|x\|_2^{2^m}\right], \\ & Q(x)=[x ; 0 ; 0 ; \ldots ; 0], \end{aligned} $$

根据上述映射，可以得到下述等式：

$$ \frac{Q(q)^T P\left(x_i\right)}{\|Q(q)\|_2\left\|P\left(x_i\right)\right\|_2}=\frac{q^T x_i}{\sqrt{m / 4+\left\|x_i\right\|_2^{2^{m+1}}}}, $$

由于 $|x_i|^{2^{m+1}}\rightarrow 0$，可得到下述近似结果，完成内积到余弦相似度的转换：

$$ \mathop{\arg \max}\limits_{x \in \mathcal{S}} q^T x \simeq \mathop{\arg \max}\limits_{x \in \mathcal{S}} \frac{Q(q)^T P\left(x_i\right)}{\|Q(q)\|_2\left\|P\left(x_i\right)\right\|_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})=\operatorname{sgn}(\mathbf{W}^\top \mathbf{x})$ 和 $z(\mathbf{x})=\operatorname{sgn}(\mathbf{R}^\top \mathbf{x})$，可以将原始输入映射成二进制码，其中 $\mathbf{W},\mathbf{R}\in \mathbb{R}^{d\times r}$，具体优化过程如下：

$$ \min _{\mathbf{W}, \mathbf{R}}\left\|\operatorname{sgn}\left(\mathbf{W}^{\top} \mathbf{A}\right)^{\top} \operatorname{sgn}\left(\mathbf{R}^{\top} \mathbf{X}\right)-\mathbf{S}\right\|^2. $$

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

$$ \min _{\mathbf{W}}\left\|\operatorname{sgn}\left(\mathbf{W}^{\top} \mathbf{A}\right)^{\top} \mathbf{Z}-\mathbf{S}\right\|^2. $$

该问题依然难以求解，因此引入矩阵 $\mathbf{B}$，上述问题转化为：

$$ \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} $$

问题依然难以求解，继续采用交替优化，固定 $\mathbf{B}$，$\mathbf{W}=\left(\mathbf{A A}^{\top}\right)^{-1} \mathbf{A} \mathbf{B}^{\top}$；固定 $\mathbf{W}$，优化问题为：

$$ \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} $$

其中 $\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$ 满足 $|x|^2\leq 1$，查询向量 $q$ 满足 $|q|=1$，令 $P(x)=(x,\sqrt{1-|x|^2})$，$Q(q)=(q,0)$，可以得到：

$$ \frac{P(x)^\top Q(q)}{\|P(x)\|\|Q(q)\|}=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}$ 为一组系数，${G_i}$ 和 $G$ 为独立的 Gumbel 随机变量，则有如下两条性质：

$$ \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} $$

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

KDD18 - H2-ALSH

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

$$ \|Q(q)-P(o)\|^2=\|q\|^2+M^2-2o^\top q, $$

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

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

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

本文指出使用 QNF 变换，可以降低上述失真误差，即：

$$ \begin{gathered} P(o)=\left[o_1, o_2, \ldots, o_d ; \sqrt{M^2-\|o\|^2}\right], \\ Q(q)=\left[\lambda q_1, \lambda q_2, \ldots, \lambda q_d ; 0\right], \text { where } \lambda=\frac{M}{\|q\|},\\ \|Q(q)-P(o)\|^2=M^2+\lambda^2\|q\|^2-2 \lambda o^\top q. \end{gathered} $$

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

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

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

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

NeurIPS18 - Norm-Ranging LSH

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

ICML19 - Sample MIPS

本篇文章在相关工作部分介绍了 Greedy-MIPS，其有下述假设：

$$ \boldsymbol{q}^\top \boldsymbol{x}_j>\boldsymbol{q}^\top \boldsymbol{x}_i \Leftrightarrow \max _{t=1, \ldots, d} q_t x_{j t}>\max _{t=1, \ldots, d} q_t x_{i t}, $$

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

回归正题，本篇文章将「找 MIPS 最大向量」转化为了一个采样问题，其定义如下概率：

$$ \begin{gathered} P(j \mid \boldsymbol{q}) \propto \boldsymbol{q}^\top \boldsymbol{x}_j, \\ P(j,t\mid \boldsymbol{q})\propto q_tx_{jt}. \\ \end{gathered} $$

由此可以得到：

$$ \begin{aligned} P(t \mid \boldsymbol{q}) & =\sum_{j=1}^n P(j, t \mid \boldsymbol{q}) \propto \sum_{j=1}^n q_t x_{j t}=q_t s_t, \\ P(j \mid t, \boldsymbol{q}) & =\frac{P(j, t \mid \boldsymbol{q})}{P(t \mid \boldsymbol{q})} \\ & \propto \frac{q_t x_{j t}}{\sum_{j=1}^n q_t x_{j t}}=\frac{x_{j t}}{\sum_{j=1}^n x_{j t}}=\frac{x_{j t}}{s_t} . \end{aligned} $$

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

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

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

由于 $\boldsymbol{q}$ 和 $\boldsymbol{x}$ 中元素可能有负数，因此文中给出了一种针对负数的解决方案，令：

$$ \begin{gathered} P(j \mid \boldsymbol{q}) \propto \sum_{t=1}^d |q_t x_{j t}|, \\ P(j \mid t, \boldsymbol{q}) \propto |q_t x_{j t}|. \end{gathered} $$

由此得到：

$$ \begin{gathered} P(t \mid \boldsymbol{q}) \propto \sum_{j=1}^n |q_t x_{j t}|, \\ P(j \mid t, \boldsymbol{q}) \propto \frac{|x_{j t}|}{\sum_{j=1}^n |x_{j t}|}. \end{gathered} $$

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

在理论保障方面，文中在 Lemma1 处证明了：

$$ \begin{gathered} \mathbb{E}[z_{jb}] = \frac{\boldsymbol{q}^\top \boldsymbol{x}_j}{S}, \\ \mathbb{E}[z_j] = B\cdot \mathbb{E}[z_{jb}], \\ \end{gathered} $$

其中 $z_{jb}=\operatorname{sgn}(q_tx_{jt})$，$z_j=\sum_{b=1}^B x_{jb}$，$S=\sum_{t=1}^d\sum_{j=1}^n|q_t x_{jt}|$。