使用 LoCo 和 M2-BERT 对长上下文检索模型进行基准测试与构建
信息检索
2024-11-19 v3 机器学习
摘要
检索管道是许多机器学习系统的组成部分，但在文档很长（例如 10K token 或更多）且识别相关文档需要综合整个文本信息的领域表现不佳。开发适用于这些领域的长上下文检索编码器面临三个挑战：(1) 如何评估长上下文检索性能，(2) 如何预训练基础语言模型以同时表示短上下文（对应查询）和长上下文（对应文档），以及 (3) 如何在 GPU 内存限制施加的批量大小限制下对该模型进行检索微调。为了解决这些挑战，我们首先介绍了 LoCoV1，这是一个由 12 个任务构成的新颖基准，旨在测量无法分块或分块无效的长上下文检索。接下来，我们展示了 M2-BERT 检索编码器，这是一个基于 Monarch Mixer 架构构建的拥有 8000 万参数的状态空间编码器模型，能够扩展至处理长达 32K token 的文档。我们描述了一种预训练数据混合方案，使该编码器能够处理短上下文和长上下文序列，并提出了一种微调方法，该方法仅使用单样本批量即可使基础模型适应检索任务。最后，我们在 LoCoV1 上验证了 M2-BERT 检索编码器，发现尽管其参数量减少了 90 倍以上，但其表现仍优于竞争性基于 Transformer 的模型至少 23.3 个点。
引用
@article{arxiv.2402.07440,
title = {Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT},
author = {Jon Saad-Falcon and Daniel Y. Fu and Simran Arora and Neel Guha and Christopher Ré},
journal= {arXiv preprint arXiv:2402.07440},
year = {2024}
}
备注
International Conference on Machine Learning (ICML) 2024