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

关键词

  information retrieval

  generative language model

  embedding vector index

导 读

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

基于嵌入的检索方法构建向量索引来搜索与查询表示最相似的文档表示。该方法延迟低、召回性能良好，被广泛用于文档检索。最近的研究表明，基于生成式语言模型的深度检索解决方案可提供更好的模型质量，但服务延迟过高、且无法支持文档更新。在这篇论文中，我们的目标是通过端到端深度生成模型增强向量索引，利用深度检索模型的可微优势，同时保持理想的服务效率。我们提出了模型增强向量索引（MEVI），其利用残差量化 (RQ) 码本来桥接序列到序列的深度检索和基于嵌入的模型。为了大幅减少推理时间，我们以更少的步骤生成候选文档的语义虚拟簇ID，然后在虚拟簇中利用适应性良好的嵌入向量进一步执行细粒度搜索。实验表明，我们的模型在常用的问答数据集MSMARCO Passage和Natural Questions上取得了更好的性能，并且服务延迟与基于嵌入的检索方法相当。

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

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

01

问 题 背 景

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

02

MEVI 方 案

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

03

实 验

在实验评测中，我们选取了较为先进的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还具有较低延迟，在搜索数量较小时能满足业界需求。

总     结

在这项工作中，我们引入了一种新颖的模型增强向量索引（MEVI），它结合了序列到序列自回归模型和双塔密集表示模型的优点。由于其能够在大规模语料库上实现高召回性能和快速检索速度，因此可以有效地应用于实际应用中。MEVI构建RQ结构来对大规模文档进行分层聚类，使序列到序列模型能够根据输入查询直接生成相关的聚类标识符；它将检索到的前k个聚类中的候选文档进一步与基于嵌入的检索结果相结合，以进行候选重新排序。我们的经验表明，MEVI在广泛使用的大规模检索数据集MSMARCO Passage和Natural Questions上实现了比基线更好的模型性能。

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

论文链接：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.

END

欢迎关注本公众号，帮助您更好地了解北京大学数据与智能实验室（PKU-DAIR），第一时间了解PKU-DAIR实验室的最新成果！

实验室简介

北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文100余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。