👉前情提要：

1. 神经网络自然语言模型概述
2. $\text{Transformer}$与注意力机制概述

📚相关论文：

1. $\text{BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding}$
   • 提出了基于双向深度$\text{Transformer}$的$\text{BERT}$交叉编码器
   • $\text{BERT}$的总结
2. $\text{ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT}$
   • 提出了基于$\text{BERT}$编码的后期$\text{Token}$级交互模式
   • $\text{ColBERTv1}$的总结
3. $\text{ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction}$
   • 保留了$\text{ColBERT}$的后期交互架构，但从训练策略$/$嵌入压缩$/$数据集上优化
   • $\text{ColBERTv2}$的总结
4. $\text{PLAID: An Efficient Engine for Late Interaction Retrieval}$
   • 在$\text{ColBERTv2}$的基础上，进一步改进了检索策略
   • $\text{PLAID}$的总结
5. $\text{EMVB: Efficient Multi-Vector Dense Retrieval Using Bit Vectors}$
   • 在$\text{PLAID}$的基础上，更进一步改进了检索策略
   • $\text{EMVB}$的总结

$\text{1. }$背景与导论

$\text{1.1. }$研究背景

1️⃣$\text{IR}$领域的背景

进展	描述
$\text{LLM(BERT/GPT)}$的引入	有强大的语义(上下文)学习能力，可学习查询$/$文档的稠密高维表示
多向量技术的发展	对词一级构建密集表示，相比稀疏(词袋)$/$单向量模型能效果更好
后期交互机制的提出	多向量系统中相对轻量级的相似度计算，但其$\text{MaxSim}$仍较耗时

2️⃣后期交互机制的进展

模型	描述
$\text{ColBERTv1}$	首次提出了后期交互模型，相比交叉编码器大幅降低了计算量，但内存要求极高
$\text{ColBERTv2}$	改进了$\text{ColBERT}$的训练策略(去噪$+$硬负样本)，嵌入存储上并使用了残差压缩
$\text{PLAID}$	改进了检索策略，提出了基于质心交互的文档预选

$\text{1.2. }$本文的工作

1️⃣对$\text{PLAID}$的进一步分析：识别查询过程中最耗时的步骤

1. 候选生成阶段：源于需要提取所有查询$\text{Token}$的前$t$接近的质心
2. 质心交互阶段：核心步骤
3. 评分重排阶段：源于解压得到最终候选段落的完整嵌入很耗时

2️⃣$\text{EMVB}$的优化思路

所针对的阶段	所采取的优化
候选生成	引入基于优化位向量的高效文档过滤方法，提取出更少更精的质心
质心交互	引入一种高效列式降维方法，并用$\text{SIMD(Single-Inst. Multi-Data)}$指令来加速
质心交互	用$\text{PQ}$压缩替代原有残差的按位压缩，降低空间占用
评分重排	引入一种动态文档$\text{–Token}$选择标准

3️⃣对$\text{EMVB}$性能的评估

1. 域内评估：较之$\text{PLAID}$在无检索精度损失的情况下，查询快了$2.8$倍$/$内存缩小了$1.8$倍
2. 域外评估：在最多$\text{2.9}$倍的检索加速下，检索精度的$\text{Trade-Off}$最小

$\text{2. }$对$\text{PLAID}$的进一步分析

$\text{2.1. }$对$\text{PLAID}$的回顾

1️⃣压缩操作：以$\tilde{t}\text{=}{C}_t\text{+}\tilde{r}$形式表示完整(解压)的嵌入

阶段	操作
聚类	对所有文档的所有嵌入组成的空间执行聚类，每个嵌入$t$分配到其最近的质心$C_t$
编码	计算每个嵌入的$t$的残差$r\text{=}t\text{–}{C}_t$，并将其近似量化编码为$\tilde{r}\text{≈}t\text{–}{C}_t$

2️⃣查询操作：在将查询$q$的原始文本编码为嵌入向量集合后，可分为四个阶段

阶段概览	目的
候选生成	通过质心计算初步筛选潜在相关段落，生成初始候选集
段落过滤	候选集$\stackrel{质心剪枝}{\to }$剩下与查询相关性高的质心$\underset{轻量级相似度}{\overset{质心交互+初排}{\to }}$最终的候选段落
残差解压	对最终的候选段落进行解压，得到其全精度的嵌入表示
评分重排	计算$\text{MaxSim}$与得分，随后重排以得到最相似文本

$\text{2.2. }$对$\text{PLAID}$检索延迟的分解

1️⃣整体上：候选生成$+$段落过滤最为耗时

2️⃣$\text{PLAID}$候选生成步：进一步的细节$\&$延迟分解

1. 具体步骤：

   步骤	描述	耗时
   查询输入	查询矩阵$Q$(所有$\text{Token}$的嵌入向量)$+$质心列表矩阵$C$(所有质心的向量)	低
   得分计算	直接计算$S_{c,q}\text{=}C\text{⸱}{Q}^T$，其中$S_{c,q}[i][j]$是质心$c_i$与查询词元$q_j$的相关性得分	中
   质心排序	对每个$q_j$选取排名前$t_{\text{nprobe}}$个质心	高
   质心候选	合并所有$q$的质心选集为最终候选质心集$C^{\prime }$	低
   段落候选	若一个段落中存在$q$被聚类到(属于)$C^{\prime }$，则将该段落候选之	低

2. 分析$\&$思路：
   • $\text{PLAID}$方案：在质心排序时，执行到每个$q_j$时都需要对全体质心quickselect ，极其耗时
   • 优化思路：减少参与排序的质心数量，故而引入一种基于阈值($t_h$)的质心过滤方法(见下)

$\text{3. EMVB}$的优化

$\text{3.1. }$质心交互前: 优化候选段落的预过滤

$\text{3.1.1. }$预过滤的整体逻辑

1️⃣有关数据结构

1. 近似得分${\tilde{T}}_{i,j}$：查询$Q$第$i$个词$q_i\stackrel{近似得分}{\leftrightarrow }$段落$P$第$j$个词$T_j$(实际用$T_j$的最近质心${\bar{C}}^{T_j}$代之)
2. 质心集：

   结构	含义
   最近质心集${\text{close}}_i^{th}$	包含所有与$q_i$的得分大于阈值$th$$\text{(threshold)}$的质心
   质心$\text{ID}$列表$I_P$	文档$P$中所有$\text{Token}$的最近质心$\text{ID}$的集合，比如$I_P^j$即为$T_j$最近质心

2️⃣段落过滤原理

1. 评分(过滤)函数：
   • 表示：[段落$\stackrel{相似度}{\leftrightarrow }$查询]表示为$F(P,q)\text{=}\sum _{i=1}^{n_q}1\left(\exists j\text{ s.t. }{I}_P^j\text{∈}{\text{ close }}_i^{th}\right)$
   • 逻辑：计算有多少个$q_i$，在$P$中有至少有一个相似的$T_j$(相似即$T_j$的$I_P^j\text{∈}{\text{close}}_i^{th}$)

     double F_Score = 0;
     // 遍历查询q的每个词q_i
     for(int i = 1; i < 查询q的词数; i++) 
     {
        set<int> close_i_th = q_i的相似质心集合close_i_th;
        // 遍历文档P的每个词T_j
        for(int j = 1; j < 文档P的词数; j++) 
        {
           int I_P_j = T_j的最近质心I_P_j;
           if(I_P_j.is_in(close_i_th))
           {
              F_Score++; break; // 每个q_1最多只能贡献一分,所以必须break
           }
        }
     }
     

2. 候选生成：根据评分函数将所有段落排序，截取前若干者为过滤所得段落

3️⃣与质心过滤的对比：

1. 橙$/$绿$/$蓝线：本文的方法能高效丢弃无关段落
2. 红线：剔除本文过滤方法的基线性能，说明本文方法不影响后续质心交互

$\text{3.1.2. }$预过滤的实现细节

$\text{3.1.2.1. }\text{close}{}_{\mathbf{i}}^{\mathbf{th}}$的计算优化

1️⃣${\text{close}}_i^{th}$计算在概念上的实现

1. 遍历计算：计算$q_i$与所有质心$C_j$的得分(内积)$C{S}_{i,j}\underset{归一化}{\overset{组合}{\to }}$得分矩阵$CS$
2. 阈值筛选：对于给定的阈值$th$收集所有$C{S}_{i,j}\text{>}th$的质心的索引$j\text{→}{\text{close}}_i^{th}$

2️⃣${\text{close}}_i^{th}$计算在机器上的朴素if实现

1. 基本流程：

   set<int> close_i_th;
   for(int i = 1; i < CS行数; i++) // 遍历每个查询
   {
      for(int j = 1; j < CS列数; j++) // 对每个查询, 遍历所有的质心
      {
         if CS[i][j] > threshold;
         {
            close_i_th.include(j) // 将质心索引添加到集合中
         }
      }
   }
   

2. 缺点分析：每次判断都需要引入if语句，实际上这会给$\text{CPU}$带来大量的分支预测错误

3️⃣基于朴素if方法的优化

1. $\text{SIMD}$向量化加速：一次性处理一批($\text{16}$个)个$q_i$，大大减少循环次数$+$跳过不必要处理的质心

   set<int> close_i_th;
   for(int i = 1; i < CS行数; i++) // 遍历每个查询
   {
      for(int j = 1; j < CS列数; j+=16) // 对每个查询, 以16为步长遍历所有质心
      {
         // 加载当前查询第i行的16个质心得分
         scores = load(CS[i][j...j+15]);
         // 比较scores是否大于threshold,生成掩码
         mask = judge(scores > threshold);
         // 掩码全为0(都小于阈值)时,跳过当前16个值
         if(mask == 0) continue;
         // 掩码不全为0(有大于阈值)时,找出那个大于阈值的分数对应的质心索引
         else if(mask != 0) 
         {
            // 提取mask中所有为1的位置,插入满足条件的质心索引
            for(int k : extract) close_i_th.include(j + k);
         }
      }
   }
   

2. 无分支优化：基于指针操作完全消除if语句

   set<int> close_i_th;
   int buffer[CS列数];  // 预分配一个缓冲区，大小为质心的数量
   int* p = buffer;    // 指针指向缓冲区的起始位置
   for(int i = 1; i < CS行数; i++) // 遍历每个查询
   {
       for(int j = 1; j < CS列数; j++) // 对每个查询, 遍历所有的质心
       {
           // is_above_threshold == 1时*p = j -> 将索引写入缓存区,再右移指针
           // is_above_threshold == 0时*p =*p -> 保持指针不变以跳过写入
           int is_above_threshold = (CS[i][j] > threshold);
           *p = j * is_above_threshold + *p * (1 - is_above_threshold);
           p += is_above_threshold; 
       }
       for (int* it = buffer; it < p; it++) // 将缓冲区中的结果写入集合
       {
           close_i_th.include(*it);
       }
       p = buffer; // 重置缓存
   }
   

3. 两种优化的结合：

   set<int> close_i_th;         
   int buffer[CS列数];          
   int* p = buffer;             
   for (int i = 1; i < CS行数; i++) // 遍历每个查询
   {
       for (int j = 1; j < CS列数; j += 16) // 每次处理16个质心
       {
           // 加载当前查询第i行的16个质心得分
           auto scores = load(CS[i][j...j+15]);
           // 比较scores是否大于threshold，生成掩码
           auto mask = compare(scores > th);
           // 如果掩码全为0，跳过当前16个值
           if (mask == 0) continue;
           // 提取掩码中为1的位置，将对应的索引写入缓存
           for (int k = 0; k < 16; k++) 
           {
               int is_above_threshold = (mask & (1 << k)) > 0; 
               *p = (j + k) * is_above_th + *p * (1 - is_above_th);
               p += is_above_th; 
           }
       }
       // 将缓冲区中的所有索引写入集合
       for (int* it = buffer; it < p; it++) 
       {
           close_i_th.include*it);
       }
       // 重置缓冲区指针到起始位置
       p = buffer;
   }
   

👉对不同优化的评估

优化$\backslash$模型	$\text{Naive IF}$	$\text{Branchless}$	$\text{Vectorized IF}$	$\text{VecBranchless}$
向量化	❌	❌	✅	✅
无分支化	❌	✅	❌	✅

$\text{3.1.2.2. }\mathbf{F}(\mathbf{P},\mathbf{q})$的计算优化

1️⃣朴素实现

double F_Score = 0;
// 遍历查询q的每个词q_i
for(int i = 1; i < 查询q的词数; i++) 
{
   set<int> close_i_th = q_i的相似质心集合close_i_th;
   // 遍历文档P的每个词T_j
   for(int j = 1; j < 文档P的词数; j++) 
   {
      int I_P_j = T_j的最近质心I_P_j;
      if(I_P_j.is_in(close_i_th))
      {
         F_Score++; break; // 每个q_1最多只能贡献一分,所以必须break
      }
   }
} 

1. 缺陷分析：if(I_P_j.is_in(close_i_th))操作需要遍历${\text{close}}_i^{th}$所有项
2. 优化思路：以位向量替代${\text{close}}_i^{th}$，从而使集合操作(遍历)被位操作取代

2️⃣优化实现

1. 数据结构：压缩的堆叠位向量

   • ${\text{close}}_i^{th}$的位向量表示：每维代表一个质心(一个$|C|$维)，该维$\text{=1}$说明对应质心在${\text{close}}_i^{th}$中

     假设质心总数|C|=6
     [1 0 1 0 1 0] -> 质心c1,c3,c4在close_i_th中
     

   • 位向量的垂直堆叠：实际执行层面，会堆叠$n_q\text{=}32$个查询(因为$\text{CPU}$大多兼容$32$位)

     对于质心总数|C|=6, 三个查询q1,q2,q3,...,q32
     [1 0 1 0 1 0] -> 质心c1,c3,c4在close_01_th中
     [0 1 1 0 1 0] -> 质心c2,c3,c5在close_02_th中
     [1 1 0 1 0 0] -> 质心c1,c2,c4在close_03_th中
     ......
     [0 1 0 1 0 1] -> 质心c2,c3,c6在close_32_th中
     

   • 垂直堆叠的压缩：将每列(一共$32$位)视作一个$32$位整数，由此获得与质心数量相同的整数

     [1 0 1 0 1 0]
     [0 1 1 0 1 0]
     [1 1 0 1 0 0]
     ......
     [0 1 0 1 0 1]
      ↓ ↓ ↓ ↓ ↓ ↓
      x x x x x x (32位整数即int类型)
     

2. 优化操作：
   

   • 图中第一步：初始化一个全$0$的$\text{32}$位掩码
   • 图中第二步：让掩码与所有的$32$位列(图中示意矩形经过了翻转)累xor
   • 图中第三步：最终$m$中$1$的个数即为$F(P,q)$

3️⃣优化实验：相比于朴素方法快了$10–16$倍

$\text{3.2. }$质心交互时: 对质心交互本身的优化

$\text{3.2.1. }$优化的操作

1️⃣优化目标：对质心交互${\bar{S}}_{q,P}\text{=}\sum _{i=1}^{n_q}\underset{j=1,\dots ,n_t}{\max }{q}_i\text{⸱}{\bar{C}}^{T_j}$的计算，分为以下两个步骤分别优化

1. [查询$–$文档嵌入的质心]得分矩阵：$\tilde{P}\text{=}{q}_i\text{⸱}{\bar{C}}^{T_j}$
2. 最大规约：找到$\tilde{P}$的每一列的最大值$(\text{MaxSiim})\text{→}$相加(最终得分)

2️⃣$\tilde{P}$的构建优化：隐式(转置)构建

1. 构建步骤：注意$CS$是查询$q$与所有质心的相似度

   操作	描述
   转置$CS$	$C{S}^T$第$i$行是[质心$C_i\stackrel{得分}{\leftrightarrow }$查询$q$中每个$\text{Token}$]
   找质心	对文档中的$T_j$，由质心列表定位到其最近质心$I_P[j]$，并在$C{S}^T$定位到行${\bar{C}}^{T_j}$
   构建$\tilde{P}$	${\tilde{P}}^T[j,i]=q_i\text{⸱}{\bar{C}}^{T_j}$即文档词$T_j$查询词$q_i$的得分

2. 过程分析：
   • 理论上：$T_j$的存储是连续的，只有将$\tilde{P}$转置才可以使得构建过程$(T_j,j$++$)$连续访问内存
   • 实验上：比不转置快了两倍

3️⃣最大规约的优化：以查询词数$\text{=32}$为例，此时${\tilde{P}}^T$每行长$32$

1. 分块处理：将${\tilde{P}}^T$第一行，分为上下两份存到两个$16$位寄存器max_l/max_h中
2. 逐行比较：将${\tilde{P}}^T$所有行，依次也分为两份存到两个$16$位寄存器current_l/current_h中
3. 逐行更新：
   • 根据比较结果更新掩码$m$
   • 再根据掩码$m$使用_mm512_mask_blend_ps更新max_l/max_h寄存器
4. 最终加和：
   • 遍历完所有的行后，使用_mm512_reduce_add_ps对max_l/max_h二者的结果相加
   • 相加结果即为所有位$\text{MaxSim}$的和，也就是最终得分

$\text{3.2.2. }$优化的分析

1️⃣理论上：一个$16$位寄存器一次就可并行处理$16$位，_mm512_指令也可实现流水线式的并行

2️⃣实验上：比$\text{PLAID}$的质心交互方案快了$1.8$倍

$\text{3.3. }$质心交互后: 对后期交互的优化

1️⃣解压的优化

1. 背景：两种压缩的对比

   压缩方法	特点
   $\text{PQ}$压缩	压缩前后数据类型都是浮点数，维度填充对齐后(而非解压)即可进行算点积
   二进制压缩	压缩后数据类型变为$\text{bit}$，必须解压后才能算点积

2. 优化的操作：将质心的量化表示由二进制压缩换为$\text{PQ(OPQ/JMPQ)}$压缩，然后
   • 评分函数的分解：$S_{q,P}\text{=}\sum _{i=1}^{n_q}\underset{j=1\dots n_t}{\max }\left(q_i\text{⸱}{\bar{C}}^{T_j}\text{+}{q}_i\text{⸱}{r}^{T_j}\right)$
   • $\text{PQ}$压缩的替换：$S_{q,P}\text{≈}\sum _{i=1}^{n_q}\underset{j=1\dots n_t}{\max }\left(q_i\text{⸱}{\bar{C}}^{T_j}\text{+}{q}_i\text{⸱}{r}_{pq}^{T_j}\right)$，即完成计算

2️⃣过滤机制的进一步改进

1. 背景：在大多情况下$S_{q,P}\text{≈}\sum _{i=1}^{n_q}\underset{j=1\dots n_t}{\max }\left(q_i\text{⸱}{\bar{C}}^{T_j}\text{+}{q}_i\text{⸱}{r}_{pq}^{T_j}\right)$中$q_i\text{⸱}{\bar{C}}^{T_j}\text{≫}{r}_{pq}^{T_j}$
2. 思路：得分主要由$q_i\text{⸱}{\bar{C}}^{T_j}$主导，干脆直接将$q_i\text{⸱}{\bar{C}}^{T_j}$作为近似值对$T_j$再筛选一次
3. 优化：直接在后期交互中屏蔽掉$q_i\text{⸱}{\bar{C}}^{T_j}\text{<}th$(某阈值)的点

$\text{4. }$实验与结果

$\text{4.1. }$实验设置

1️⃣数据集：$\text{MS MARCO passages}$用于域内评估，$\text{LoTTE}$用于域外评估

2️⃣模型测试

1. 模型：用$\text{ColBERTv2}$嵌入$+\text{FAISS}$实现$\text{PQ}$(在$\text{MS MARCO}$上通过$\text{JMPQ}$优化)$+\text{EMVB}$检索
2. 基线：$\text{PLAID}$

3️⃣评估指标：嵌入所需空间，查询延迟，检索质量($\text{MRR@10/Recall@100/Recall@1000}$)

4️⃣软硬件配置：

1. 仅在$\text{CPU}$上测试：$\text{Intel Xeon Gold 5318Y}$，主频$\text{2.10 GHz}$，支持$\text{AVX512}$指令集
2. 编译环境：$\text{GCC 11.3.0 -O3}$级优化，$\text{Linux 5.15.0-72}$系统

$\text{4.2. }$实验结果

1️⃣域内检索：进行了保守压缩$/$激进($m\text{=}32/16$)压缩，激进压缩策略就足以取得查过$\text{PLAID}$的性能

方面	$\text{EMVB}$较$\text{PLAID}$的进展
内存占用	$m\text{=}16$时内存需求减少了近一半
查询速度	$m\text{=}16$时最高提高了$\text{2.8}$倍速
检索质量	$m\text{=}16$时无明显退化，$m\text{=}32$是有达$\text{2.5}$倍的提升

2️⃣域外评估

方面	$\text{EMVB}$较$\text{PLAID}$	备注
查询速度	提高了$\text{2.9}$倍速	$\text{LoTTE}$文档平均更长，更利于$\text{EMVB}$的位向量过滤
检索质量	略微低于$\text{PLAID}$	原有在于$\text{JMPQ}$优化形同虚设，换为$\text{OPQ}$优化稍好点