文献分享: DESSERT基于LSH的多向量检索(Part2——理论保证的证明)

$\text{1. DESSERT}$算法

$\text{1.1. }$向量集搜索是什么

1️⃣向量集搜索的过程

1. 输入：查询向量集 Q = { q 1 , . . . , q m q } ∈ R m q × d Q\text{=}\{q_1,...,q_{m_q}\}\text{∈}\mathbb{R}^{m_q\text{×}d} Q={q1​,...,qmq​​}∈Rmq​×d，向量集的集合 D = { S 1 , . . . , S N } D\text{=}\{S_1,...,S_N\} D={S1​,...,SN​}，其中 S i = { x i 1 , . . . , x i m i } ∈ R m i × d S_i\text{=}\{x_{i1},...,x_{im_i}\}\text{∈}\mathbb{R}^{m_i\text{×}d} Si​={xi1​,...,ximi​​}∈Rmi​×d
2. 输出：用$F(Q,S_i)$衡量$Q$与$S_i$相似度，要求以$1–\delta$的概率返回与$Q$相似度最高的$S_i$，即 S ∗ = argmax ⁡ i ∈ { 1 , … , N } F ( Q , S i ) S^*\text{=}\mathop{\operatorname{argmax}}\limits_{{i\in\{1,\ldots{},N\}}}F\left({Q,{S}_{i}}\right) S∗=i∈{1,…,N}argmax​F(Q,Si​)

2️⃣对相似度$F(Q,S_i)$的定义

1. 子相似度：衡量每个$q_r\text{∈}Q$和$x_{ij}\text{∈}{S}_i$时间的相似度

   👉举个例子,对于查询Q={q1,q2,q3}考虑下面向量集
     S1 = {x11, x12, x13, x14}
     S2 = {x21, x22, x23, x24, x25}
     S3 = {x31, x32, x33}
   👉让Q和S1进行内部聚合,先要算出全部的相似度(比如内积)
     Sim(q1,x11) Sim(q1,x12) Sim(q1,x13) Sim(q1,x14)
     Sim(q2,x11) Sim(q2,x12) Sim(q2,x13) Sim(q2,x14)
     Sim(q3,x11) Sim(q3,x12) Sim(q3,x13) Sim(q3,x14)
   

2. 内部聚合$\sigma$：让每个$q_r\text{∈}Q$得到一共聚合后的相似度，类似于$\text{ColBERT}$的$\text{MaxSim}$

   👉对于每个qi,应用内部聚合函数σ
     Inner(q1,S) = σ{Sim(q1,x11),Sim(q1,x12),Sim(q1,x13),Sim(q1,x14)}
     Inner(q2,S) = σ{Sim(q2,x11),Sim(q2,x12),Sim(q2,x13),Sim(q2,x14)}
     Inner(q3,S) = σ{Sim(q3,x11),Sim(q3,x12),Sim(q3,x13),Sim(q3,x14)}
   👉在ColBERT中这个σ就是所谓的MaxSim
     Inner(q1,S) = Max{Sim(q1,x11),Sim(q1,x12),Sim(q1,x13),Sim(q1,x14)}
     Inner(q2,S) = Max{Sim(q2,x11),Sim(q2,x12),Sim(q2,x13),Sim(q2,x14)}
     Inner(q3,S) = Max{Sim(q3,x11),Sim(q3,x12),Sim(q3,x13),Sim(q3,x14)}
   

3. 外部聚合$A$：将每个$q_r\text{∈}Q$内部聚合的结果进行处理，得到最后评分$F(Q,S)$

   👉对每个内部聚合应用内部聚合函数A
     F(Q,S1) = A{Inner(q1,S),Inner(q2,S),Inner(q3,S)}
   👉在ColBERT中这个A就是逐个相加
     F(Q,S1) = Inner(q1,S) + Inner(q2,S) + Inner(q3,S)
   

$\text{1.2. }$算法的朴素流程

1️⃣索引构建

1. 输入：若干向量集，如 D = { S 1 , S 2 , . . . , S N } D\text{=}\{S_1,S_2,...,S_N\} D={S1​,S2​,...,SN​}
2. 构建：对于每个 S i = { x i 1 , x i 2 , . . . , x i m i } S_i\text{=}\{x_{i1},x_{i2},...,x_{im_i}\} Si​={xi1​,xi2​,...,ximi​​}都执行索引构建操作
   • 索引分配：为$S_i$中每个元素分配一个唯一索引，例如$x_{ij}$的索引可以为$j$
   • 哈希分桶：用$L$个$\text{LSH}$函数${\psi }_1,{\psi }_2,...,{\psi }_L$对$S_i$中所有元素进行$L$次分桶
   • 索引存储：利用$\mathcal{D}{\left[i\right]}_{t,h}$数组结构，对应哈希函数${\psi }_t$下哈希值为$h$的桶，存储存储了$S_i$中落入该桶的所有向量的索引

2️⃣查询阶段

1. 输入：查询向量集 Q = { q 1 , q 2 , . . . , q m q } Q\text{=}\{q_1,q_2,...,q_{m_q}\} Q={q1​,q2​,...,qmq​​}，以及上一步构建的$\text{DESSERT}$索引
2. 编码：照样用那$L$个$\text{LSH}$函数${\psi }_1,{\psi }_2,...,{\psi }_L$，对$Q$中所有元素进行$L$次分桶
3. 评分：通过检查$q_r\text{∈}Q$和$x_{ij}\text{∈}{S}_i$的碰撞次数$\text{Count}(q_r,x_{ij})$，得到二者相似度的一个近似$\widehat{\text{Sim}}(q_r,x_{ij})\text{=}\frac{\text{Count}(q_r,x_{ij})}{L}$
   • 原理：为何$\frac{\text{Count}(q_r,x_{ij})}{L}$可作为近似评分
     • 对$q_r\text{∈}Q$和$x_{ij}\text{∈}{S}_i$各自进行$L$次分桶后碰撞$\text{Count}(q_r,x_{ij})$次，故估计碰撞率为$\frac{\text{Count}(q_r,x_{ij})}{L}$
     • 鉴于$\text{Pr}[\psi (x)\text{=}\psi (y)]\text{=}\text{Sim}(x,y)$的假设，所以碰撞率就是相似度
   • 实现：基于$\mathcal{D}{\left[i\right]}_{t,h}$数组结构(具体优化则见$\text{TinyTable}$)
     • 对于$\forall q_r\text{∈}Q$用哈希函数${\psi }_t$得出其哈希值$h$，按照$t,h$索引直接在找到桶$\mathcal{D}{\left[i\right]}_{t,h}$，如果$x_{ij}$索引也在这儿就算碰撞一次
     • 对${\psi }_1,{\psi }_2,...,{\psi }_L$都进行如上操作，得到最终碰撞次数$\text{Count}(q_r,x_{ij})$，碰撞率(相似度)为$\frac{\text{Count}(q_r,x_{ij})}{L}$
4. 聚合：基于相似度$\widehat{\text{Sim}}(q_r,x_{ij})$，得到最终的相似度估值$\hat{F}(Q,S_i)$
   • 原始输入：由以上$\text{LSH}$得到的，每个$q_r\text{∈}Q$的近似相似度集 s ^ ( q r , S i ) = { Sim ^ ( q r , x i 1 ) , Sim ^ ( q r , x i 2 ) , . . . , Sim ^ ( q r , x i m i ) } \hat{\mathbf{s}}(q_{r},S_i)\text{=}\{\hat{\text{Sim}}(q_{r},x_{i1}),\hat{\text{Sim}}(q_{r},x_{i2}),...,\hat{\text{Sim}}(q_{r},x_{im_i})\} s^(qr​,Si​)={Sim^(qr​,xi1​),Sim^(qr​,xi2​),...,Sim^(qr​,ximi​​)}
   • 内部聚合： σ ( s ^ ( q r , S i ) ) = σ { Sim ^ ( q r , x i 1 ) , Sim ^ ( q r , x i 2 ) , . . . , Sim ^ ( q r , x i m i ) } \sigma(\hat{\mathbf{s}}(q_{r},S_i))\text{=}\sigma\{\hat{\text{Sim}}(q_{r},x_{i1}),\hat{\text{Sim}}(q_{r},x_{i2}),...,\hat{\text{Sim}}(q_{r},x_{im_i})\} σ(s^(qr​,Si​))=σ{Sim^(qr​,xi1​),Sim^(qr​,xi2​),...,Sim^(qr​,ximi​​)}，当$\sigma$为$\text{max}$时左式等于$\widehat{\text{MaxSim}}(q_r,S_i)$
   • 外部聚合： A ( Q , S i ) = A { σ ( s ^ ( q 1 , S i ) ) , σ ( s ^ ( q 2 , S i ) ) , . . . , σ ( s ^ ( q m q , S i ) ) } A(Q,S_i)\text{=}A\{\sigma(\hat{\mathbf{s}}(q_{1},S_i)),\sigma(\hat{\mathbf{s}}(q_{2},S_i)),...,\sigma(\hat{\mathbf{s}}(q_{m_q},S_i))\} A(Q,Si​)=A{σ(s^(q1​,Si​)),σ(s^(q2​,Si​)),...,σ(s^(qmq​​,Si​))}，$A$可以是将所有$q_r\text{∈}Q$内部聚合结果相加

$\text{1.3. }$理论保证

$\text{1.3.1. }$对于内部聚合过程

1️⃣乘性约束函数及其引理

1. 定义$\text{4.1}$：$(\alpha ,\beta )\text{-}$乘性极值约束函数
   • 条件：给定参数$\alpha ,\beta$满足$0\text{<}\beta \text{≤}1\text{≤}\alpha$，给定向量$x$并用$\max (x)$表示向量$x$中最大元素值
   • 定义：函数$\sigma (x)\text{:}{R}^m\text{→}R$在$U$上是$(\alpha ,\beta )\text{-}$极大的，等价于$\forall x\text{∈}U$有$\beta \max (x)\text{≤}\sigma (x)\text{≤}\alpha \max (x)$
   • 例如：当内部聚合$\sigma (x)$就是$\text{ColBERT}$的$\text{MaxSim}$函数时，就有$\alpha \text{=}\beta \text{=}1$即$(1,1)\text{-}$极大
2. 引理$\text{4.1.1}$：平均极值下界衰减引理
   • 条件：$\phi (x)\text{:}R\text{→}R$在区间$I$上是$(\alpha ,\beta )\text{-}$极大的，$x$为标量是$\max (x)$退化为$x$本身，即$\beta x\text{≤}\phi (x)\text{≤}\alpha x$
   • 结论： σ ( x = { x 1 , x 2 , . . . , x m } ) = 1 m ∑ i = 1 m φ ( x i ) \sigma(\mathbf{x}\text{=}\{x_1,x_2,...,x_m\})\text{=}\displaystyle{}\frac{1}{m}\sum_{i=1}^m\varphi(x_i) σ(x={x1​,x2​,...,xm​})=m1​i=1∑m​φ(xi​)在$U\text{=}{I}^m$上是$\left(\alpha ,\frac{\beta }{m}\right)\text{-}$极大的
   • 示例：满足该引理的实数域函数

     $\mathbf{\phi }(\mathbf{x})$	$\mathbf{\beta }$	$\mathbf{\alpha }$
     $x$	$1$	$1$
     $e^x–1$	$1$	$e–1$
     $\text{Sigmid}(x)\text{=}\frac{e^x}{1\text{+}{e}^x}–\frac{1}{2}$	$0.23$	$0.25$

2️⃣引理$\text{4.1.2/4.1.3}$

1. 输入：对查询向量集 ∀ q r ∈ Q = { q 1 , q 2 , . . . , q m q } \forall{}q_{r}\text{∈}Q\text{=}\{q_1,q_2,...,q_{m_q}\} ∀qr​∈Q={q1​,q2​,...,qmq​​}和向量集 S i = { x i 1 , x i 2 , . . . , x i m i } S_i\text{=}\{x_{i1},x_{i2},...,x_{im_i}\} Si​={xi1​,xi2​,...,ximi​​}，考虑真实$/$(基于碰撞的)近似相似度
   • 每个$q_r\text{∈}Q$的真实相似度集： s ( q r , S i ) = { Sim ( q r , x i 1 ) , Sim ( q r , x i 2 ) , . . . , Sim ( q r , x i m i ) } {\mathbf{s}}(q_{r},S_i)\text{=}\{{\text{Sim}}(q_{r},x_{i1}),{\text{Sim}}(q_{r},x_{i2}),...,{\text{Sim}}(q_{r},x_{im_i})\} s(qr​,Si​)={Sim(qr​,xi1​),Sim(qr​,xi2​),...,Sim(qr​,ximi​​)}
   • 每个$q_r\text{∈}Q$的近似相似度集： s ^ ( q r , S i ) = { Sim ^ ( q r , x i 1 ) , Sim ^ ( q r , x i 2 ) , . . . , Sim ^ ( q r , x i m i ) } \hat{\mathbf{s}}(q_{r},S_i)\text{=}\{\hat{\text{Sim}}(q_{r},x_{i1}),\hat{\text{Sim}}(q_{r},x_{i2}),...,\hat{\text{Sim}}(q_{r},x_{im_i})\} s^(qr​,Si​)={Sim^(qr​,xi1​),Sim^(qr​,xi2​),...,Sim^(qr​,ximi​​)}
2. 符号：对真实相似度集$/$近似相似度集，采取不同聚合方式
   • 真实集：用最值聚合 max ⁡ ( s ( q r , S i ) ) = max ⁡ { Sim ( q r , x i 1 ) , . . . , Sim ( q r , x i m i ) } \max\left(\mathbf{s}(q_{r},S_i)\right)\text{=}\max\{\text{Sim}(q_{r},x_{i1}),...,\text{Sim}(q_{r},x_{im_i})\} max(s(qr​,Si​))=max{Sim(qr​,xi1​),...,Sim(qr​,ximi​​)}，记作$s_{\max }$或$\text{MaxSim}(q_r,S_i)$
   • 近似集：用$\sigma$聚合 σ ( s ^ ( q r , S i ) ) = σ { Sim ^ ( q r , x i 1 ) , . . . , Sim ^ ( q r , x i m i ) } \sigma(\hat{\mathbf{s}}(q_{r},S_i))\text{=}\sigma\{\hat{\text{Sim}}(q_{r},x_{i1}),...,\hat{\text{Sim}}(q_{r},x_{im_i})\} σ(s^(qr​,Si​))=σ{Sim^(qr​,xi1​),...,Sim^(qr​,ximi​​)}其中$\sigma$是$(\alpha ,\beta )\text{-}$极大，简记作$\sigma (\hat{\mathbf{s}})$
3. $\gamma$函数：设定$\gamma =\left(\frac{\alpha \left(1–s_{\max }\right)}{\alpha –\tau }\right){\left(\frac{s_{\max }\left(\alpha –\tau \right)}{\tau \left(1–s_{\max }\right)}\right)}^{\tau /\alpha }$
   
   • 单调：处于$\left(s_{\max },1\right)$区间中，并且随$s_{\max }$递增随$\tau$递减
   • 极限：$\gamma$存在单侧极限，$\tau$从高处接近$\alpha s_{\max }$时$\underset{\tau \searrow \alpha s_{\max }}{\lim }\gamma \text{=}1$，$\tau$从低处接近$\alpha$时$\underset{\tau \nearrow \alpha }{\lim }\gamma \text{=}{s}_{\max }$
4. 结论：给定一个阈值$\tau \text{∈}[\alpha s_{\max },\alpha ]$并记录差值为$\Delta =\tau –\alpha s_{\max }$，让当内部聚合$\sigma$是$(\alpha ,\beta )\text{-}$极大的，则有如下两个引理
   • 引理$\text{4.1.2}$(指数上界衰减引理)：$\mathrm{Pr}\left[\sigma \left(\hat{\mathbf{s}}\right)\text{≥}\alpha s_{\max }\text{+}\Delta \right]\text{≤}m{\gamma }^L$，对近似相似度聚合后，大概率不超过理论上界
   • 引理$\text{4.1.3}$(高斯下界集中引理)：$\mathrm{Pr}\left[\sigma \left(\hat{\mathbf{s}}\right)\text{≤}\beta s_{\max }\text{–}\Delta \right]\text{≤}2e^{–2L{\Delta }^2/{\beta }^2}$，对近似相似度聚合后，大概率不低于理论下界

$\text{1.3.2. }$外部聚合及运行时间

1️⃣一些基本的设置

1. 近似集的聚合：与之前一样 σ ( s ^ ( q r , S i ) ) = σ { Sim ^ ( q r , x i 1 ) , . . . , Sim ^ ( q r , x i m i ) } \sigma(\hat{\mathbf{s}}(q_{r},S_i))\text{=}\sigma\{\hat{\text{Sim}}(q_{r},x_{i1}),...,\hat{\text{Sim}}(q_{r},x_{im_i})\} σ(s^(qr​,Si​))=σ{Sim^(qr​,xi1​),...,Sim^(qr​,ximi​​)}，但为区分不同$q_r$就不做简写了
2. 外部聚合评分：让外部聚合$A$为带$w_r$权值的加权平均，由此$F\left(Q,S_i\right)\text{=}\frac{1}{m_q}\sum _{r=1}^{m_q}{w}_r\sigma \left(\hat{\mathbf{s}}\left(q_r,S_i\right)\right)$
3. 最邻近向量集：给定 Q = { q 1 , q 2 , . . . , q m q } Q\text{=}\{q_1,q_2,...,q_{m_q}\} Q={q1​,q2​,...,qmq​​}和 D = { S 1 , S 2 , . . . , S N } D\text{=}\{S_1,S_2,...,S_N\} D={S1​,S2​,...,SN​}，设定$S^{\ast }\text{∈}D$是使得$F\left(Q,S_i\right)$最大的向量集
4. 分析时的假设：假定所有向量集的向量数都相同，即统一将$m_i$认定为一个常数$m$，当然也可用$(m_i{)}_{\max }$替代所有$m_i$

2️⃣定理$\text{4.2}$：概率对数近邻保证定理

1. 设定参数：对于$0\text{<}\beta \text{≤}1\text{≤}\alpha$
   • 令$B^{\ast }\text{=}\frac{\beta }{m_q}\sum _{r=1}^{m_q}{w}_r\max (\mathbf{s}(q_r,S^{\ast }))\text{=}\frac{\beta }{m_q}\sum _{r=1}^{m_q}{w}_r\text{MaxSim}(q_r,S^{\ast })$，$B^{\ast }$其实是$F(Q,S^{\ast })$的一个下界
   • 令$B_i\text{=}\frac{\alpha }{m_q}\sum _{r=1}^{m_q}{w}_r\max (\mathbf{s}(q_r,S_i))\text{=}\frac{\alpha }{m_q}\sum _{r=1}^{m_q}{w}_r\text{MaxSim}(q_r,S_i)$，$B_i$其实是$F(Q,S_i)$的一个上界
   • 令$B^{\prime }$为在$S_i\ne S^{\ast }$条件下，$B_i$能取得的最大值
2. 结论：只要我们选择足够大的哈希表数量，算法就能以高概率返回正确最邻近向量集
   • 条件：当$\Delta \text{=}\frac{B^{\ast }–B^{\prime }}{3}\text{＞}0$时，，设定$L\text{=}O\left(\log \left(\frac{N{m}_{q}m}{\delta }\right)\right)$
   • 结论：$\text{DESSERT}$算法结构能以$1–\delta$的概率，返回与$Q$相似度最高的$S_i$，即 S ∗ = argmax ⁡ i ∈ { 1 , … N } F ( Q , S i ) S^*\text{=}\mathop{\operatorname{argmax}}\limits_{{i \in \{ 1,\ldots N\} }}F\left( {Q,{S}_{i}}\right) S∗=i∈{1,…N}argmax​F(Q,Si​)

3️⃣运行时间的分析

1. 一些前提：假设每个哈希函数运行的时间为$O(d)$，集合$\mathcal{D}{\left[i\right]}_{t,h}$的元素数存在常数阈值$|\mathcal{D}{\left[i\right]}_{t,h}|\text{≤}T$
2. 运行时间：复杂度为$O\left((d\text{+}N)m_q\log \left(\frac{N{m}_{q}m}{\delta }\right)\right)$，相比之下暴力搜索的复杂度是$O\left(m_{q}mNd\right)$