一步搞定模型训练和商品召回：京东全新索引联合训练召回模型 / 开普饭

©PaperWeekly 原创 · 作者 | 张菡

单位 | 京东算法工程师

研究方向 | 信息检索

简介

基于深度学习的语义检索已经广泛应用于互联网应用中。早在 2018 年，京东搜索已经实现了语义检索 [1]，通过深度语义模型，学习 query 和商品的向量表示，用于京东搜索的召回阶段，通过向量相似性查找实现语义检索。

大规模的向量相似性查找需要大量的浮点运算，精确 KNN（K-Nearest Neighbor）查找会带来超长的召回时延，可行解决方案是采用 ANN（Approximate Nearest Neighbor）查找，以少量的检索精度换取超高的召回效率。

常用的 ANN 索引方法包括 PQ（Product Quantization，乘积量化）、基于树的方法、基于图的方法和哈希。而基于 PQ 的方法由于内存消耗小，可以动态增删等优点在工业界应用较为广泛。在此之前，京东搜索的语义召回采用的就是基于 PQ 的 Faiss 索引库 [2]。

基于 PQ 索引主要思想是将高维向量用低维子空间的笛卡尔积表示，而后在低维子空间中进行聚类，用聚类中心构成的码本对向量进行编码，并提前计算好编码之间的码距，在检索时，就能通过查找直接得到向量之间的距离，从而大大减少浮点运算次数。由于模型的训练和索引的建立是两个独立的过程，所以子空间的划分和聚类中心的计算势必会带来一定的相似性计算误差。

因此，我们提出一种全新的索引联合训练模型（product quantization based embedding index jointly trained with deep retrieval model，简称 Poeem），对聚类中心进行参数化学习，减少子空间划分和聚类带来的损失，进一步提高向量检索的精度。此外，在实际应用中，联合训练模型不需要额外的建索引时间以及额外的索引服务，大大降低了向量检索在实际应用中的工程复杂度。

论文标题：

Joint Learning of Deep Retrieval Model and Product Quantization based Embedding Index

论文来源：

SIGIR 2021

论文链接：

https://arxiv.org/abs/2105.03933

模型

模型结构如下图所示：

左边是一个经典的双塔召回模型，我们在 item tower 的最后嵌入一层索引层，用于实现 PQ 索引的学习，右边则是索引层向量量化的实现过程。向量量化主要分为四个步骤：

首先，我们对 item embedding 进行维度的 reordering，从而降低后续划分的子空间之间的相关性。矩阵维度的 reordering 可以通过乘上一个旋转矩阵（Rotation Matrix）来实现。（对应上图中粉色标注过程）

第二步，我们学习一个粗粒度的聚类中心矩阵 v，用 item embedding 在 v 中做相似度查找，得到 item embedding 的粗粒度量化编码 r。（对应上图中黄色标注过程）

第三步，我们首先用 item embedding 减去对应的粗粒度聚类中心，得到 residual item embedding，用于后续的向量量化。然后将 residual item embedding 均匀划分成 subvectors，同时学习每个子空间的聚类中心矩阵，然后用 subvector 减去对应子空间的聚类中心，得到 subvector 残差，再在聚类中心矩阵做相似度查找，得到 item embedding 的乘积量化编码 c。

最后，我们来做向量复原。首先加上粗粒度聚类中心，将残差恢复成原向量；然后乘上旋转矩阵的逆矩阵，恢复原本向量维度的顺序。