[NeurIPS 2023]用生成式语言模型增强基于嵌入的文档索引

论文链接：https://arxiv.org/pdf/2309.13335.pdf

代码地址：https://github.com/HugoZHL/MEVI

文档检索是网络搜索引擎中的一个关键阶段。现有的文档检索方法可分为三类：基于词语、基于嵌入、和基于生成的方法。基于词语的方法对整个语料库中的单词或短语构建倒排索引以进行搜索[1]，但无法利用语义信息。基于嵌入的方法通过双塔架构将查询和文档编码为密集嵌入向量以利用语义信息，然后构建向量索引并应用近似最近邻（ANN）搜索来检索相关文档[2]，但它们将优化目标分成了两个阶段，且最近邻和近似最近邻的召回性能也存在一定的差距[3]。基于生成的方法采用序列到序列的模型，直接根据给定的查询生成文档标识符，其中文档标识符反映了文档语义的有效先验，例如DSI[4]和NCI[5]中，文档通过分层KMeans聚类[6]被组织为树，标识符是从根到叶节点的路径编码。这类方法在小语料库上比基于嵌入的方法表现更好，但很难扩展到更大的语料库，并且由于高延迟和语料库不可变而无法在工业系统中提供在线服务。为了解决上述问题，我们提出了模型增强向量索引（MEVI），在大规模语料库上兼具高召回率和较快检索速度。

我们首先构建一个残差量化（RQ）码本[7]来对文档进行聚类。RQ码本保留了文档簇的层次结构，适合自回归生成。在此基础上，我们构建序列到序列模型来对用户查询进行编码并根据RQ码本直接生成虚拟集群标识符。然后我们使用ANN方法在虚拟集群中根据嵌入向量搜索相关的文档。在训练期间，我们使用真实数据和增强的查询文档对来训练编码器-解码器结构的端到端生成模型。在推理期间，我们通过解码器上的波束搜索根据RQ码本来检索top-K个文档集群，并在这些集群中检索与查询嵌入相似的文档嵌入。MEVI的设计解决了传统基于嵌入和基于生成的方法的局限性。一方面，我们可以将RQ码本限制在适中的大小，从而减少自回归解码器的计算时间并确保低延迟；允许在RQ码本的文档集群中插入或删除文档，从而使新文档也可以通过序列到序列模型生成文档集群来搜索。通过为RQ码本选择合适的大小，我们可以平衡召回性能和推理延迟，同时利用ANN的高效率和深度检索模型的准确召回。为了进一步增强召回性能，我们还将生成的文档集群和基于嵌入的方法得到的文档进行集成，同时发挥二者的优势。

在实验评测中，我们选取了较为先进的AR2[8]和T5-ANCE[9]模型作为基于嵌入的方法。如表2、表3所示， MEVI显着提高了召回率，在MSMARCO Passage数据集[10]上实现了+3.62%的MRR@10、+7.32%的R@50、+10.54%的R@1000，在Natural Questions数据集[11]上实现了+5.04%的R@5、+5.46%的R@20，+5.96%的R@100。

通过在随机90%的文档上训练、并在推断时包括所有的文档，我们验证了MEVI支持动态语料库的能力。

RQ使用残差嵌入，在实际表现中比普通的K-Means更好。

MEVI还具有较低延迟，在搜索数量较小时能满足业界需求。

详细了解本工作，请访问下方链接地址：

论文链接：https://arxiv.org/pdf/2309.13335.pdf

References：

[1] Shengyao Zhuang, Hang Li, and Guido Zuccon. Deep query likelihood model for information retrieval. In Advances in Information Retrieval – 43rd European Conference on IR Research, ECIR, 2021.

[2] Wenhao Lu, Jian Jiao, and Ruofei Zhang. Twinbert: Distilling knowledge to twin-structured compressed BERT models for large-scale retrieval. In Proceedings of the 29th ACM International Conference on Information and Knowledge Management, CIKM, 2020.

[3] Shitao Xiao, Zheng Liu, Weihao Han, Jianjin Zhang, Defu Lian, Yeyun Gong, Qi Chen, Fan Yang, Hao Sun, Yingxia Shao, and Xing Xie. Distill-vq: Learning retrieval oriented vector quantization by distilling knowledge from dense embeddings. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, 2022.

[4] Yi Tay, Vinh Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Prakash Gupta, Tal Schuster, William W. Cohen, and Donald Metzler.  Transformer memory as a differentiable search index. In Advances in Neural Information Processing Systems 35 (NeurIPS), 2022.

[5] Yujing Wang, Yingyan Hou, Haonan Wang, Ziming Miao, Shibin Wu, Qi Chen, Yuqing Xia, Chengmin Chi, Guoshuai Zhao, Zheng Liu, Xing Xie, Hao Sun, Weiwei Deng, Qi Zhang, and Mao Yang. A neural corpus indexer for document retrieval. In Advances in Neural Information Processing Systems 35 (NeurIPS), 2022.

[6] John A Hartigan and Manchek A Wong. A k-means clustering algorithm. Journal of the royal statistical society. series c (applied statistics), 28(1):100–108, 1979.

[7] Julieta Martinez, Holger H. Hoos, and James J. Little. Stacked quantizers for compositional vector compression. CoRR, abs/1411.2173, 2014.

[8] Hang Zhang, Yeyun Gong, Yelong Shen, Jiancheng Lv, Nan Duan, and Weizhu Chen. Adversarial retriever-ranker for dense text retrieval. In The 10th International Conference on Learning Representations, ICLR, 2022.

[9] Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In The 9th International Conference on Learning Representations, ICLR, 2021.

[10] Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. MS MARCO: A human generated machine reading comprehension dataset. In Proceedings of the Workshop on Cognitive Computation: Integrating neural and symbolic approaches 2016 co-located with the 30th Annual Conference on Neural Information Processing Systems (NeurIPS), 2016.

[11] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452–466, 2019.