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

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

嵌入索引能够实现快速近似近邻搜索，是目前最先进的深度检索系统不可缺少的组成部分。传统的检索方法往往将嵌入学习和建立索引两个步骤分离，增加了建立索引的时间，降低了检索的准确性。本文提出了一种基于乘积量化的嵌入指标与深度检索模型联合训练的新方法——Poeem，利用梯度直通式估计器、热启动策略和深度检索模型等技术，将两个单独的步骤统一到一个端到端训练中。最优空间分解和吉文斯旋转。大量实验结果表明，该方法不仅显著提高了检索精度，而且索引时间几乎为零。为了便于比较和再现，我们将方法开源。

背景

embidding  indexi 的缺点主要问题在于模型训练与索引建立的分离，导致索引建立时间增加，检索精度下降。

量化方法缺点：1)量化步骤作为基于PQ的嵌入指标的核心，具有不可微分的操作，如参数min，使标准的反向传播训练失效。因此，我们利用梯度直通估计器绕过不可微性，以实现端到端训练。2)量化质心随机初始化导致质心分配非常稀疏，参数利用率较低，量化失真较高  3)标准优化乘积量化(OPQ)[9]算法通过正交矩阵对空间进行旋转，进一步降低PQ失真，不能与联合模型一起迭代运行

2 Method

2.1 Revisiting Retrieval Model

一个标准的嵌入检索模型由查询塔𝑄和项目塔𝑆组成。因此，对于一个给定查询𝑞和一个条目𝑠

然而，这个额外的步骤不仅花费额外的时间来建立索引，而且还会导致召回率的下降

2.2 Embedding Indexing Layer

索引层定义一个完整的量化函数T: R𝑑−→R𝑑映射输入嵌入x输出嵌入T (x),它可以分解为四个功能:粗量化函数𝜓,乘积量化函数𝜙和解码器𝜌函数,与一个正交矩阵𝑅和旋转功能

The coarse quantization function 𝜓

R𝑑→{1，···，𝐽}，通过粗质心矩阵v∈R𝐽×𝑑将连续向量x映射到𝐽-way离散粗代码𝑟

The product quantization (PQ) function 𝜙

R𝑑−→{1，···，𝐾}𝐷，将上述剩余向量y映射为𝐾-way𝐷-dimensional

The decoder function

The rotation function

2.3 Training Algorithm

在这里，我们提出通过调整Eq.(6)中的原始量化函数T来利用梯度直通估计量[2]，如下所示

我们在损失函数中加入一个正则项来最小化量化失真，如下: