ESPN: Memory-Efficient Multi-Vector Information Retrieval

📄 arXiv: 2312.05417v1 📥 PDF

作者: Susav Shrestha, Narasimha Reddy, Zongwang Li

分类: cs.IR, cs.LG

发布日期: 2023-12-09

备注: 10 pages, 10 figures

💡 一句话要点

ESPN：通过存储卸载降低多向量信息检索的内存需求

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 信息检索 多向量模型 存储卸载 软件预取 SSD 内存优化 大规模索引

📋 核心要点

  1. 多向量信息检索模型虽然提升了检索质量，但其索引规模导致内存需求显著增加，限制了可扩展性。
  2. ESPN通过将重排序嵌入表卸载到SSD，并设计软件预取器，显著降低内存需求并加速检索。
  3. 实验表明，ESPN能将内存需求降低5-16倍，并使SSD检索速度提升高达6.4倍，同时保持接近内存级别的查询延迟。

📝 摘要（中文）

大型语言模型的最新进展表明，其在信息检索（IR）任务中表现出卓越的有效性。虽然许多神经IR系统将查询和文档编码为单向量表示，但多向量模型通过生成多向量表示并在单个token粒度上促进相似性搜索，从而提高了检索质量。然而，这些模型显著增加了检索索引的内存和存储需求，增加了一个数量级。由于其巨大的内存需求，索引大小的这种升级使得多向量IR模型的可扩展性越来越具有挑战性。我们引入了Embedding from Storage Pipelined Network (ESPN)，其中我们将整个重排序嵌入表卸载到SSD，并将内存需求降低了5-16倍。我们设计了一个命中率超过90%的软件预取器，将基于SSD的检索速度提高了高达6.4倍，并证明即使对于大型查询批处理，我们也可以保持接近内存级别的查询延迟。

🔬 方法详解

问题定义：多向量信息检索模型通过token级别的相似度计算提升了检索精度，但其为每个token生成向量表示导致索引体积庞大，内存需求高昂，严重制约了模型的可扩展性。现有方法难以在保证检索效率的同时，有效降低内存占用。

核心思路：ESPN的核心思路是将多向量模型中占据大量内存的重排序嵌入表卸载到SSD存储中，从而显著降低内存占用。同时，为了弥补SSD访问速度慢的缺点，设计了高效的软件预取器，提前将数据加载到内存，以减少延迟。

技术框架：ESPN包含以下主要模块：1) 离线索引构建：将文档和查询编码为多向量表示，并将重排序嵌入表存储到SSD。2) 在线检索：接收查询后，首先从SSD预取相关的嵌入向量，然后进行相似度计算和重排序。3) 软件预取器：根据查询的访问模式，预测未来需要访问的嵌入向量，并提前从SSD加载到内存。

关键创新：ESPN的关键创新在于将存储卸载与软件预取相结合，在降低内存占用的同时，尽可能地保持检索效率。软件预取器的设计是核心，它需要准确预测未来的访问模式，并高效地从SSD加载数据。

关键设计：软件预取器采用基于历史访问模式的预测算法，并结合了缓存机制，以提高命中率。具体来说，预取器维护一个访问历史记录，记录每个token最近访问的时间和频率。根据这些信息，预取器预测未来最有可能被访问的token，并提前从SSD加载其对应的嵌入向量。此外，预取器还使用一个LRU缓存来存储最近访问的嵌入向量，以减少SSD访问次数。

📊 实验亮点

实验结果表明，ESPN可以将多向量信息检索模型的内存需求降低5-16倍，同时将基于SSD的检索速度提高高达6.4倍。软件预取器的命中率超过90%，有效降低了SSD访问延迟。即使对于大型查询批处理，ESPN也能保持接近内存级别的查询延迟，证明了其在大规模信息检索场景下的实用性。

🎯 应用场景

ESPN适用于需要处理大规模文档集合的信息检索系统，例如搜索引擎、问答系统和推荐系统。通过降低内存需求，ESPN可以支持更大规模的索引，从而提高检索精度和覆盖率。此外，ESPN还可以应用于资源受限的设备，例如移动设备和嵌入式系统。

📄 摘要（原文）

Recent advances in large language models have demonstrated remarkable effectiveness in information retrieval (IR) tasks. While many neural IR systems encode queries and documents into single-vector representations, multi-vector models elevate the retrieval quality by producing multi-vector representations and facilitating similarity searches at the granularity of individual tokens. However, these models significantly amplify memory and storage requirements for retrieval indices by an order of magnitude. This escalation in index size renders the scalability of multi-vector IR models progressively challenging due to their substantial memory demands. We introduce Embedding from Storage Pipelined Network (ESPN) where we offload the entire re-ranking embedding tables to SSDs and reduce the memory requirements by 5-16x. We design a software prefetcher with hit rates exceeding 90%, improving SSD based retrieval up to 6.4x, and demonstrate that we can maintain near memory levels of query latency even for large query batch sizes.