🚀 ColBERT-XM

ColBERT-XM 是一个基于 ColBERT 的多语言语义搜索模型，它能够将查询和段落编码为矩阵形式的词元级嵌入，并通过可扩展的向量相似度（MaxSim）算子高效地找到与查询上下文匹配的段落。该模型以 XMOD 为基础架构，可从高资源语言（如英语）的单语微调中学习，并实现跨多种语言的零样本检索。

🚀 快速开始

安装依赖

首先，安装 colbert-ai 及其他必要的依赖项：

pip install git+https://github.com/stanford-futuredata/ColBERT.git@main torchtorch==2.1.2 faiss-gpu==1.7.2 langdetect==1.0.9

使用示例

# 使用自定义模块，该模块可自动检测待索引段落的语言，并相应地激活特定语言的适配器
from .custom import CustomIndexer, CustomSearcher 
from colbert.infra import Run, RunConfig

n_gpu: int = 1 # 设置可用的 GPU 数量
experiment: str = "colbert" # 存储日志和创建的索引的文件夹名称
index_name: str = "my_index" # 索引名称，即向量数据库的名称
documents: list = ["Ceci est un premier document.", "Voici un second document.", "etc."] # 语料库

# 步骤 1: 索引。此步骤将所有段落编码为矩阵，存储在磁盘上，并构建用于高效搜索的数据结构。
with Run().context(RunConfig(nranks=n_gpu,experiment=experiment)):
    indexer = CustomIndexer(checkpoint="antoinelouis/colbert-xm")
    indexer.index(name=index_name, collection=documents)

# 步骤 2: 搜索。给定模型和索引，您可以对语料库发出查询，以检索每个查询的前 k 个段落。
with Run().context(RunConfig(nranks=n_gpu,experiment=experiment)):
    searcher = CustomSearcher(index=index_name) # 无需再次指定检查点，模型名称已存储在索引中。
    results = searcher.search(query="Comment effectuer une recherche avec ColBERT ?", k=10)
    # results: 长度为 k 的元组的元组，包含 ((段落 ID, 段落排名, 段落得分), ...)

✨ 主要特性

• 多语言支持：支持多种语言，可在多语言环境下进行语义搜索。
• 高效检索：通过可扩展的向量相似度（MaxSim）算子，能够高效地找到与查询上下文匹配的段落。
• 零样本学习：以 XMOD 为基础架构，可从高资源语言（如英语）的单语微调中学习，并实现跨多种语言的零样本检索。

📚 详细文档

评估

mMARCO 数据集

我们在 mMARCO 的小型开发集上评估了我们的模型。该数据集包含 6,980 个查询，涉及 880 万个候选段落，涵盖 14 种语言。以下是我们的模型与其他检索模型在该数据集官方指标（即截止 10 时的平均倒数排名（MRR@10））上的多语言性能比较：

Mr. TyDi 数据集

我们还在 Mr. TyDi 的测试集上评估了我们的模型。该数据集是另一个多语言开放检索数据集，包含 mMARCO 中未出现的低资源语言。以下是我们的模型与其他检索模型在该数据集官方指标（即截止 100 时的平均倒数排名（MRR@100）和截止 100 时的召回率（R@100））上的性能比较：

训练

数据

我们使用了 MS MARCO 段落排名 数据集中的英语训练样本，该数据集包含 880 万个段落和 53.9 万个训练查询。我们没有使用官方数据集提供的 BM25 负样本，而是使用 msmarco-hard-negatives 蒸馏数据集，从 12 个不同的密集检索器中采样更难的负样本。我们最终的训练集由 640 万个 (q, p+, p-) 三元组组成。

实现

模型从 xmod-base 检查点初始化，并通过组合正样本和难负样本的预测分数上的成对 softmax 交叉熵损失（如 ColBERTv1 中所述）和批内采样 softmax 交叉熵损失（如 ColBERTv2 中所述）进行优化。使用 AdamW 优化器在一个 80GB NVIDIA H100 GPU 上进行 50k 步的微调，批量大小为 128，峰值学习率为 3e-6，在前 10% 的训练步骤中进行热身，并采用线性调度。我们将嵌入维度设置为 128，并将问题和段落的最大序列长度分别固定为 32 和 256。

🔧 技术细节

模型架构

ColBERT-XM 基于 ColBERT 架构，以 XMOD 为基础模型。该模型能够将查询和段落编码为矩阵形式的词元级嵌入，并通过可扩展的向量相似度（MaxSim）算子高效地找到与查询上下文匹配的段落。

训练策略

模型通过组合成对 softmax 交叉熵损失和批内采样 softmax 交叉熵损失进行优化。具体来说，我们计算正样本和难负样本的预测分数上的成对 softmax 交叉熵损失，以及批内采样的 softmax 交叉熵损失。这种组合损失函数有助于模型更好地学习查询和段落之间的语义关系。

多语言支持

ColBERT-XM 支持多种语言，可在多语言环境下进行语义搜索。模型以 XMOD 为基础架构，可从高资源语言（如英语）的单语微调中学习，并实现跨多种语言的零样本检索。

📄 许可证

本项目采用 MIT 许可证。

🔗 引用

如果您在研究中使用了本模型，请使用以下 BibTeX 引用：

@article{louis2024modular,
  author = {Louis, Antoine and Saxena, Vageesh and van Dijck, Gijs and Spanakis, Gerasimos},
  title = {ColBERT-XM: A Modular Multi-Vector Representation Model for Zero-Shot Multilingual Information Retrieval},
  journal = {CoRR},
  volume = {abs/2402.15059},
  year = {2024},
  url = {https://arxiv.org/abs/2402.15059},
  doi = {10.48550/arXiv.2402.15059},
  eprinttype = {arXiv},
  eprint = {2402.15059},
}