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

©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 可以通过乘上