背景

研究でColBERTを使用しているが、どのようにインデックスを構築しているのか、検索時にどうやってインデックスを活用しているのかを説明できなかった。

目標

• ColBERTがどのようにインデックスを構築しているのかを理解する
• ColBERTが検索時にどうやってインデックスを活用しているのかを理解する

まず、ColBERTの元論文を読んでみる
https://arxiv.org/pdf/2004.12832

3.4. 節がオフラインインデックスについて書いてあるようなので読んでみる

3.4. 節は

• 利用可能であれば複数のGPUを活用し、バッチ処理を利用して文書をエンコードすること
• BucktIteratorという長さベースのバッチ分割をしてエンコードすること
  • これによりパディングに対する無意味な計算が減り、処理を効率化することができる

が書かれていた
どのようなインデックスになるのかは3.6. 節のEnd-to-end Retrievalに記載されているようなので読んでみる

3.6. 節は

• インデックスのデータ構造は{doc_id: d[i][j]}みたいな、ドキュメントIDに紐づいた、あるドキュメントトークンの埋め込み表現であること
• クエリの処理はフィルタリングと再ランキングの２段階であること
  • フィルタリング時は各クエリトークンの埋め込み表現のインデックスに対する検索結果をマージ
  • 再ランキング時は3.5.節の方法を取る
• インデックスはIVFQという形式をとること
  • 対象を絞るためにクラスタリング
  • 情報量を削減するために圧縮

が書かれていた
再ランキングの方法を調べるために3.5.節のRerankingを読んでみる

3.5. 節は

• コーパスの埋め込みを3次元のテンソルとして表現すること
• クエリとコーパスに対してMaxSim演算すること
• 従来のBERTベースのRerankerよりも計算コストが少ないこと
  が書かれていた

これで目標達成だ

結論

ColBERTはどのようにインデックスを構築しているの？

元のドキュメントが識別できる形で、ドキュメントトークンの埋め込み表現を保持している
具体的なインデックスの実装にはIVFQというクラスタリングとベクトルの圧縮を行う形式をとっている

ColBERTは検索時にどうやってインデックスを活用しているの？

1.各クエリトークンの埋め込み表現を使ってインデックスに対して検索を行う
2. 得られた結果に対してリランキングを行う（インデックスなし）