模型概述
模型特點
模型能力
使用案例
🚀 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))上的多語言性能比較:
模型 | 類型 | #樣本 | #參數 | 英語 | 西班牙語 | 法語 | 意大利語 | 葡萄牙語 | 印尼語 | 德語 | 俄語 | 中文 | 日語 | 荷蘭語 | 越南語 | 印地語 | 阿拉伯語 | 平均 | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
1 | BM25 (Pyserini) | 詞法 | - | - | 18.4 | 15.8 | 15.5 | 15.3 | 15.2 | 14.9 | 13.6 | 12.4 | 11.6 | 14.1 | 14.0 | 13.6 | 13.4 | 11.1 | 14.2 |
2 | mono-mT5 (Bonfacio et al., 2021) | 交叉編碼器 | 12.8M | 390M | 36.6 | 31.4 | 30.2 | 30.3 | 30.2 | 29.8 | 28.9 | 26.3 | 24.9 | 26.7 | 29.2 | 25.6 | 26.6 | 23.5 | 28.6 |
3 | mono-mMiniLM (Bonfacio et al., 2021) | 交叉編碼器 | 80.0M | 107M | 36.6 | 30.9 | 29.6 | 29.1 | 28.9 | 29.3 | 27.8 | 25.1 | 24.9 | 26.3 | 27.6 | 24.7 | 26.2 | 21.9 | 27.8 |
4 | DPR-X (Yang et al., 2022) | 單向量 | 25.6M | 550M | 24.5 | 19.6 | 18.9 | 18.3 | 19.0 | 16.9 | 18.2 | 17.7 | 14.8 | 15.4 | 18.5 | 15.1 | 15.4 | 12.9 | 17.5 |
5 | mE5-base (Wang et al., 2024) | 單向量 | 5.1B | 278M | 35.0 | 28.9 | 30.3 | 28.0 | 27.5 | 26.1 | 27.1 | 24.5 | 22.9 | 25.0 | 27.3 | 23.9 | 24.2 | 20.5 | 26.5 |
6 | mColBERT (Bonfacio et al., 2021) | 多向量 | 25.6M | 180M | 35.2 | 30.1 | 28.9 | 29.2 | 29.2 | 27.5 | 28.1 | 25.0 | 24.6 | 23.6 | 27.3 | 18.0 | 23.2 | 20.9 | 26.5 |
7 | DPR-XM (我們的模型) | 單向量 | 25.6M | 277M | 32.7 | 23.6 | 23.5 | 22.3 | 22.7 | 22.0 | 22.1 | 19.9 | 18.1 | 18.7 | 22.9 | 18.0 | 16.0 | 15.1 | 21.3 |
8 | ColBERT-XM (我們的模型) | 多向量 | 6.4M | 277M | 37.2 | 28.5 | 26.9 | 26.5 | 27.6 | 26.3 | 27.0 | 25.1 | 24.6 | 24.1 | 27.5 | 22.6 | 23.8 | 19.5 | 26.2 |
Mr. TyDi 數據集
我們還在 Mr. TyDi 的測試集上評估了我們的模型。該數據集是另一個多語言開放檢索數據集,包含 mMARCO 中未出現的低資源語言。以下是我們的模型與其他檢索模型在該數據集官方指標(即截止 100 時的平均倒數排名(MRR@100)和截止 100 時的召回率(R@100))上的性能比較:
模型 | 類型 | #樣本 | #參數 | 阿拉伯語 | 孟加拉語 | 英語 | 芬蘭語 | 印尼語 | 日語 | 韓語 | 俄語 | 斯瓦希里語 | 泰盧固語 | 平均 | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
MRR@100 | |||||||||||||||
1 | BM25 (Pyserini) | 詞法 | - | - | 36.8 | 41.8 | 14.0 | 28.4 | 37.6 | 21.1 | 28.5 | 31.3 | 38.9 | 34.3 | 31.3 |
2 | mono-mT5 (Bonfacio et al., 2021) | 交叉編碼器 | 12.8M | 390M | 62.2 | 65.1 | 35.7 | 49.5 | 61.1 | 48.1 | 47.4 | 52.6 | 62.9 | 66.6 | 55.1 |
3 | mColBERT (Bonfacio et al., 2021) | 多向量 | 25.6M | 180M | 55.3 | 48.8 | 32.9 | 41.3 | 55.5 | 36.6 | 36.7 | 48.2 | 44.8 | 61.6 | 46.1 |
4 | ColBERT-XM (我們的模型) | 多向量 | 6.4M | 277M | 55.2 | 56.6 | 36.0 | 41.8 | 57.1 | 42.1 | 41.3 | 52.2 | 56.8 | 50.6 | 49.0 |
R@100 | |||||||||||||||
5 | BM25 (Pyserini) | 詞法 | - | - | 79.3 | 86.9 | 53.7 | 71.9 | 84.3 | 64.5 | 61.9 | 64.8 | 76.4 | 75.8 | 72.0 |
6 | mono-mT5 (Bonfacio et al., 2021) | 交叉編碼器 | 12.8M | 390M | 88.4 | 92.3 | 72.4 | 85.1 | 92.8 | 83.2 | 76.5 | 76.3 | 83.8 | 85.0 | 83.5 |
7 | mColBERT (Bonfacio et al., 2021) | 多向量 | 25.6M | 180M | 85.9 | 91.8 | 78.6 | 82.6 | 91.1 | 70.9 | 72.9 | 86.1 | 80.8 | 96.9 | 83.7 |
8 | ColBERT-XM (我們的模型) | 多向量 | 6.4M | 277M | 89.6 | 91.4 | 83.7 | 84.4 | 93.8 | 84.9 | 77.6 | 89.1 | 87.1 | 93.3 | 87.5 |
訓練
數據
我們使用了 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},
}







