语义搜索中的高效GPU检索

# 语义搜索的高效GPU检索  
来源：https://arxiv.org/html/2608.28968  
CCS：计算方法学 信息检索  
CCS：计算方法学 自然语言处理  

Chujie Zheng  
单位：LinkedIn，Mountain View, CA, USA  
Ronak Kaoshik  
单位：LinkedIn，Mountain View, CA, USA  
Pratik Dixit  
单位：LinkedIn，Mountain View, CA, USA  
Vishal Shah  
单位：LinkedIn，Mountain View, CA, USA  
邮箱：[[email protected]](mailto:[email protected])  
Yanbo Li  
单位：LinkedIn，Mountain View, CA, USA  
Jiahao Xu  
单位：LinkedIn，Mountain View, CA, USA  
Manika Agarwal  
单位：LinkedIn，Mountain View, CA, USA  
Chinmay Naik  
单位：LinkedIn，Mountain View, CA, USA  
Lingyu Zhang  
单位：LinkedIn，Mountain View, CA, USA  
邮箱：[[email protected]](mailto:[email protected])  
Chetan Bhole  
单位：LinkedIn，Mountain View, CA, USA  
Chirag Bhanuprasad Mehta  
单位：LinkedIn，Mountain View, CA, USA  
Meng Zheng  
单位：LinkedIn，Mountain View, CA, USA  
Puneet Singh Ahluwalia  
单位：LinkedIn，Mountain View, CA, USA  
Shirisha Singh  
单位：LinkedIn，Mountain View, CA, USA  
邮箱：[[email protected]](mailto:[email protected])  
Ping Jin  
单位：LinkedIn，Mountain View, CA, USA  
Manas Apte  
单位：LinkedIn，Mountain View, CA, USA  
Gokulraj Mohanasundaram  
单位：LinkedIn，Mountain View, CA, USA  
Tugrul Bingol  
注：工作完成于LinkedIn期间。  
单位：LinkedIn，Mountain View, CA, USA  
Raghavan Muthuregunathan  
单位：LinkedIn，Mountain View, CA, USA  
Fedor Borisyuk  
单位：LinkedIn，Mountain View, CA, USA  
邮箱：[[email protected]](mailto:[email protected])  

###### 摘要  
LinkedIn上的语义搜索必须从数亿份档案中检索相关资料，以响应诸如*“柏林一位从事过支付业务的金融科技创始人”*这样的自然语言查询。已部署的相关性策略是瓶颈导向的：每个活跃的非协商属性必须满足，并且一个预先存在的LLM分级相关性评判器通过属性等级的固定最小值/中值聚合来实现这一点。余弦相似度则对证据进行平均，使得某个属性的强匹配掩盖了另一个属性的失败，这限制了第一阶段（L0）检索器的召回率。我们提出了一种与策略对齐的检索框架：嵌入被划分为八个类别监督的段，其分数在服务时遵循相同的最小值/中值规则；对于多向量检索，此段分数在每个标记的文档槽位独立计算，并在槽位间取最大值。一个轻量级的单槽位Stage-1评分器生成高召回率候选集，而尺度不变的相对范数门控在训练、评估和服务之间保持类别激活的一致性。在21K个保留查询上，该表征在离线相关性上优于匹配容量的基线，且增益广泛分布于各种属性组合。我们使用两阶段GPU架构部署此框架：一个FP8粗排器对整个语料库评分，将每个分片的容量提高71%，Stage-1矩阵乘法吞吐量提高36%；然后一个FP16阶段对过采样的候选集进行精确重排序，在每个分片副本超过500 QPS的吞吐量下，恢复了全FP16召回率的99.6–99.8%。在一个成员随机A/B测试中，在不变的GR评判器下，探索性查询的Precision@10从63.7%提升至79.0%，导航性查询的Precision@1从65.5%提升至74.7%，一项盲评的人类评估独立证实了Precision@10的提升。  

###### 关键词：语义搜索，神经检索，GPU推理，大规模信息检索  

## 1. 引言  
现代搜索引擎必须处理表达性强的自然语言查询，这些查询描述了多方面的实体。系统不仅需要匹配关键词，还需要解读包含空间、时间和类别约束的完整句子描述。我们以LinkedIn专业网络上的*语义人物搜索*为例，研究高质量、高效的大规模第一阶段GPU检索。SPS使用检索漏斗：L0阶段通过基于嵌入的检索将数亿档案过滤到几千个候选集，随后下游模型对候选集进行重排序。因此，L0召回率决定了下游质量。这个阶段极具挑战性，因为冻结的产品相关性策略是*合取式*的：一个档案必须满足所有非协商属性。一个预先存在的LLM分级相关性评判器通过固定类别等级的最小值/中值聚合来实现该策略。余弦相似度则在嵌入中平均证据，使得在多个属性上匹配良好但在必需属性上失败的档案仍能获得高分。因此，评分器与该冻结策略不匹配，而该策略在保留集上的一致性是我们的主要质量目标。  
我们通过GR对齐的分段嵌入解决这种不匹配：一种被划分为八个类别监督段的向量。聚合在每个标记的文档槽位内应用，然后在槽位间取最大值（§4）。这些继承的运算符早于本工作且保持无参数；学习仅限于嵌入。由于生成的瓶颈分数与传统的平坦向量候选生成不兼容，我们使用单槽位、平坦点积的Stage-1评分器来近似它，该评分器在经验上具有高召回率，尽管并非通过设计。尺度不变的相对范数门控在训练、评估和服务之间保留了活跃类别掩码，而校准和全向量辅助目标稳定了段分数的可比性和联合嵌入几何。  
我们通过两阶段GPU架构使该表征可部署。一个FP8阶段对整个语料库评分，一个FP16阶段对过采样的候选集进行精确重排序，在将Stage-1索引容量提高71%和Stage-1矩阵乘法吞吐量提高36%的同时，恢复了全FP16基线召回率的99.6–99.8%。在分片大小为3500万文档的评估中，系统在每个GPU分片副本上维持超过500 QPS。选择和过滤内核减少了服务开销，但不改变检索目标或在线处理；我们在部署评估后呈现它们。  
我们的贡献包括：  
- • 一种与监督对齐的检索表征。我们引入了一种GR对齐的分段多向量表征，该表征在每个标记的文档槽位内应用冻结策略聚合，然后在槽位间取最大值（§4）。  
- • 平等一致的候选生成。一个单槽位、轻量级的Stage-1评分器为多槽位的最终评分器提供经验上高召回率的候选集。相对范数门控保留了活跃类别掩码，校准和全向量损失在训练和服务间正则化了表征（§4）。  
- • 一个可部署且经在线验证的GPU系统。FP8/FP16架构将Stage-1索引容量提高71%，改进了Stage-1矩阵乘法吞吐量36%，并恢复了99.6–99.8%的全FP16召回率。一项成员随机A/B测试在不变的GR评判器下测量了在线流量的增益，同时支持内核和过滤基准测试另行报告（§6）。  
§2回顾相关工作；§3定义SPS场景；§4介绍检索框架；§5报告离线评估；§6介绍可部署的两阶段系统、扩展实验、支持性优化和部署经验。  

## 2. 相关工作  
##### 基于嵌入的检索和学习相似度  
双编码器在共享向量空间中用检索替代词汇重叠，并且基于LLM的编码器（如Qwen3 Embedding）进一步提高了表征质量。工业系统，包括Facebook搜索、Que2Search、Que2Engage和DLRM风格的推荐系统，通常使用固定的余弦或点积评分。MoL则学习了一个相似度函数并配对层级索引器，而其他工作观察到对比余弦不捕捉绝对相关性。这些方法仍然产生单个标量分数。我们的框架则将嵌入分解为GR监督的属性段并直接应用聚合。只有表征是被学习的：评分器既不引入标量近似也不引入可训练的聚合权重。其第一阶段是经验验证的，而非假设保留召回率（§4）。  
##### GPU服务和量化  
我们的系统扩展了LiNR（LinkedIn的GPU原生检索架构），并采用了类似于SilverTorch的GPU原生布隆过滤思想。与SilverTorch不同，我们的过滤内核将派生特征作为副产品计算，消除了单独的索引层。我们还分析了布隆过滤假阳性如何与下游相似度交互。近似最近邻方法（如乘积量化和HNSW）以召回率为代价减少索引占用空间，而GPU实现包括FAISS、cuVS中的CAGRA和单GPU十亿规模搜索。相反，我们在FP8中执行穷举的、特定于模型的语料库范围评分，并通过过采样候选集上的FP16重排序恢复近乎完美的召回率，使§4的框架在大规模上变得实用。  
##### 多向量检索和自适应维度  
ColBERT和ColBERTv2使用令牌级的后期交互来捕捉细粒度对应关系。我们的框架在槽位级别使用多个文档向量。一个共享编码器，以不同的前缀标签为条件，为同一文档生成独立的表征；最终模型在每个槽位内应用完整的GR评分器，然后取槽位分数的最大值（§4）。我们使用骨干网络的套娃表征来实例化独立的512维和256维每槽位配置。每个配置在其原始维度上定义八个等宽的类别监督段（分别为每段64和32维度）；256维模型不是通过截断训练好的512维分段布局得到的。  
##### 多目标排序、蒸馏和LLM评判  
多任务架构和梯度平衡方法广泛用于优化下游排序器的相关性和参与度。我们的L0训练则结合了对比目标与为检索几何设计的正则化器。标准交叉编码器蒸馏将评判器的输出压缩为单个标量。我们的分段表征则将LLM评判器的显式多类别分级策略蒸馏到学习的嵌入几何中，同时保持策略聚合的无参数性。这一对齐目标决定了GPU服务系统的设计。  

## 3. 问题定义  
SPS检索自然语言查询的档案，并面临三个挑战：C1，满足每个必需属性；C2，匹配LLM评判器使用的类别级别相关性定义；C3，提升第一阶段设定的候选集召回率上限。我们通过标记的多向量检索（§4.2）和GR对齐的分段嵌入（§4.3，C1–C2）来解决这些挑战，后者将每个向量划分为特定类别的子空间。嵌入编码器是学习的；服务在每个槽位内应用继承的最小值/中值聚合，在槽位间应用最大值，两者都是无参数的。模型通过两阶段的检索-重排序GPU流水线服务：一个模型特定的FP8粗检索阶段遍历整个语料库，随后一个FP16评分器处理过采样的候选集（§6）。查询通过三阶段漏斗——L0基于嵌入的检索、L2 MLP排序和L3 LLM交叉编码——而我们专注于L0，因为下游阶段只能重排序它检索到的候选集（C3）。基线L0检索器是一个孪生双编码器，具有共享的40亿参数嵌入骨干，其规模和能力与当代模型相当。它在线将查询编码为q∈R^d，离线将档案编码为x_i∈R^d，其中d=512。我们保留d_s表示一个文档槽位的维度，K表示文档槽位的数量；因此d_s × K表示每槽位维度乘以槽位数：  

(1) Ret_k(q) = top-k_i cos(q, x_i) = top-k_i ⟨q, x_i⟩ / (||q||_2 ||x_i||_2)。  

这定义了单向量基线；§4中的学习分段和多向量表征