ColBERT とは何か？Late Interaction とは何か？そして検索においてなぜ重要なのか？

Jina ColBERT v2：埋め込みとリランキングのための多言語レイトインタラクションリトリーバー

Jina ColBERT v2 は 89 言語をサポートし、優れた検索性能、ユーザー制御可能な出力次元、8192 トークン長を提供します。

先週金曜日、Hugging Face 上での Jina AI による ColBERT モデルのリリースが、特に Twitter/X 上で AI コミュニティに大きな反響を呼びました。画期的な BERT モデルについては多くの人が知っていますが、ColBERT に関する話題は一部の人々に疑問を投げかけています：情報検索技術の混沌とした分野で ColBERT は何が特別なのでしょうか？なぜ AI コミュニティは 8192 長の ColBERT に興奮しているのでしょうか？この記事では、ColBERT と ColBERTv2 の詳細について、そのデザイン、改良点、そして ColBERT のレイトインタラクションの驚くべき効果を掘り下げていきます。

Reranker API

検索の関連性と RAG の精度を簡単に最大化

tagColBERT とは何か？

「ColBERT」という名称は、Contextualized Late Interaction over BERT（BERT 上のコンテキスト化されたレイトインタラクション）の略で、スタンフォード大学から生まれたモデルです。BERT の深い言語理解を活用しながら、新しいインタラクションメカニズムを導入しています。「レイトインタラクション」として知られるこのメカニズムは、検索プロセスの最終段階までクエリとドキュメントを別々に処理することで、効率的で正確な検索を可能にします。具体的には、モデルには 2 つのバージョンがあります：

• ColBERT：最初のモデルは、Omar Khattab と Matei Zaharia が考案したもので、「ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT」という論文で情報検索に対する新しいアプローチを提示しました。この研究は SIGIR 2020 で発表されました。

ColBERT: BERT 上のコンテキスト化されたレイトインタラクションによる効率的で効果的な文章検索

自然言語理解（NLU）の最近の進歩は、情報検索（IR）の急速な進展を促しています。これは主に、文書ランキングのための深層言語モデル（LM）のファインチューニングによるものです。これらの LM に基づくランキングモデルは驚くほど効果的ですが、単一の関連性スコアを計算するために各クエリ-文書ペアを巨大なニューラルネットワークに通す必要があるため、以前のアプローチと比べて計算コストが桁違いに増加します。これに対処するため、我々は ColBERT を提示します。これは効率的な検索のために深層 LM（特に BERT）を適応させた新しいランキングモデルです。ColBERT は、BERT を使用してクエリと文書を独立してエンコードし、その後、それらの細かい粒度の類似性をモデル化する安価だが強力なインタラクションステップを採用するレイトインタラクションアーキテクチャを導入します。

arXiv.orgOmar Khattab

• ColBERTv2：基礎研究をベースに、Omar Khattab は Barlas Oguz、Matei Zaharia、Michael S. Bernstein と共同で「ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction」を SIGIR 2021 で発表しました。この ColBERT の次のイテレーションでは、前バージョンの制限に対処し、ノイズ除去された教師あり学習と残差圧縮という主要な改良を導入し、モデルの検索効果と保存効率の両方を向上させました。

ColBERTv2：軽量なレイトインタラクションによる効果的で効率的な検索

ニューラル情報検索（IR）は、検索や他の知識集約型言語タスクを大きく進展させました。多くのニューラル IR 手法がクエリとドキュメントを単一ベクトル表現にエンコードする一方で、レイトインタラクションモデルは各トークンの粒度で多ベクトル表現を生成し、関連性のモデリングをスケーラブルなトークンレベルの計算に分解します。この分解によりレイトインタラクションの効果が高まることが示されていますが、これらのモデルの空間フットプリントが 1 桁増加します。

arXiv.orgKeshav Santhanam

tagColBERT のデザインを理解する

ColBERTv2 のアーキテクチャは元の ColBERT と非常に似ており、主な革新はトレーニング技術と圧縮メカニズムに関するものであるため、まず元の ColBERT の基本的な側面について掘り下げていきましょう。

tagColBERT のレイトインタラクションとは？

「インタラクション」とは、クエリとドキュメントの表現を比較することで関連性を評価するプロセスを指します。

「レイトインタラクション」は ColBERT の本質です。この用語は、クエリとドキュメントの表現が独立してエンコードされた後、プロセスの後半でインタラクションが発生するというモデルのアーキテクチャと処理戦略に由来します。これは「アーリーインタラクション」モデルと対照的で、アーリーインタラクションではクエリとドキュメントの埋め込みが早い段階で、通常モデルによるエンコード前または途中でインタラクションします。

アーリーインタラクションは、すべての可能なクエリ-ドキュメントペアを考慮する必要があるため計算の複雑さが増加し、大規模なアプリケーションには効率的ではありません。

ColBERT のようなレイトインタラクションモデルは、ドキュメント表現の事前計算を可能にし、最後により軽量なインタラクションステップを採用することで、効率性とスケーラビリティを最適化します。このインタラクションステップは、すでにエンコードされた表現に焦点を当てます。このデザインの選択により、検索時間が短縮され、計算要求が軽減され、大規模なドキュメントコレクションの処理により適したものとなっています。

tag相互作用なし：ドキュメントとクエリの埋め込みのコサイン類似度

多くの実用的なベクトルデータベースとニューラル検索ソリューションは、ドキュメントとクエリの埋め込み間の高速なコサイン類似度マッチングに依存しています。この方法は直感的で計算効率が良いという魅力がありますが、「相互作用なし」または「非相互作用ベース」と呼ばれ、クエリとドキュメント間の何らかの相互作用を組み込んだモデルと比較して性能が劣ることが分かっています。

「相互作用なし」アプローチの主な制限は、クエリとドキュメントの用語間の複雑なニュアンスや関係性を捉えられないことにあります。情報検索の本質は、クエリの背後にある意図とドキュメント内のコンテンツを理解しマッチングすることです。このプロセスには、関連する用語の深い文脈的理解が必要であり、ドキュメントとクエリの単一の集約された埋め込みではこれを提供することが困難です。

tagColBERT におけるクエリとドキュメントのエンコーダー

ColBERT のエンコーディング戦略は、言語の深い文脈理解で知られる BERT モデルに基づいています。このモデルは、クエリまたはドキュメント内の各トークンに対して密なベクトル表現を生成し、それぞれクエリとドキュメントの文脈化された埋め込みの集合を作成します。これにより、後期相互作用フェーズでの埋め込みの詳細な比較が可能になります。

tagColBERT のクエリエンコーダー

トークン ${q_1, q_2, ..., q_l}$ を持つクエリ $Q$ に対して、まず $Q$ を BERT ベースの WordPiece トークンにトークン化し、特殊な [Q] トークンを先頭に追加します。この [Q] トークンは BERT の [CLS] トークンの直後に配置され、クエリの開始を示します。

クエリが事前定義された トークン数 $N_q$ より短い場合は [mask] トークンで $N_q$ まで埋め込まれ、それ以外の場合は最初の $N_q$ トークンまで切り捨てられます。埋め込まれたシーケンスは BERT に通され、その後 CNN（畳み込みニューラルネットワーク）と正規化を経て、以下の $\mathbf{E}_q$ として表される埋め込みベクトルのセットが出力されます：
$$\mathbf{E}_q := \mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[Q]},q_0,q_1,\ldots,q_l\mathtt{[mask]},\mathtt{[mask]},\ldots,\mathtt{[mask]}\right)\right)$$

tagColBERT のドキュメントエンコーダー

同様に、トークン ${d_1, d_2, ..., d_n}$ を持つドキュメント $D$ に対して、ドキュメントの開始を示す [D] トークンが先頭に追加されます。このシーケンスは、埋め込みの必要なく同じプロセスを経て、以下の $\mathbf{E}_d$ として表される埋め込みベクトルのセットが生成されます：
$$\mathbf{E}_d := \mathrm{Filter}\left(\mathrm{Normalize}\left(\mathrm{BERT}\left(\mathtt{[D]},d_0,d_1,...,d_n\right)\right)\right)$$

クエリの埋め込みに [mask] トークンを使用すること（論文ではクエリ拡張と呼ばれています）は、すべてのクエリで統一された長さを確保し、バッチ処理を容易にします。[Q] と [D] トークンは明示的にクエリとドキュメントの開始を示し、モデルがこれら2種類の入力を区別するのに役立ちます。

tagColBERT とクロスエンコーダーの比較

クロスエンコーダーはクエリとドキュメントのペアを一緒に処理するため、非常に正確ですが、可能なすべてのペアを評価する計算コストが高いため、大規模なタスクには効率的ではありません。意味的類似性タスクや詳細なコンテンツ比較など、文ペアの正確なスコアリングが必要な特定のシナリオでは優れています。しかし、この設計は、事前計算された埋め込みと効率的な類似度計算が重要な、大規模データセットからの高速な検索が必要な状況での適用性を制限します。

対照的に、ColBERT の後期相互作用モデルは、ドキュメント埋め込みの事前計算を可能にし、意味解析の深さを損なうことなく検索プロセスを大幅に高速化します。クロスエンコーダーの直接的なアプローチと比較すると一見直感に反するように見えるこの方法は、リアルタイムおよび大規模な情報検索タスクに対してスケーラブルなソリューションを提供します。これは、計算効率と相互作用モデリングの品質の間の戦略的な妥協を表しています。

tagColBERT を使用したトップ K ドキュメントの検索

クエリとドキュメントの埋め込みが得られれば、最も関連性の高いトップ K ドキュメントを見つけることは簡単になります（ただし、2つのベクトルのコサインを計算するほど簡単ではありません）。

主な操作には、用語ごとの類似度を計算するためのバッチ内積計算、ドキュメント用語間のmax-poolingによるクエリ用語ごとの最高類似度の検出、クエリ用語間の合計によるドキュメントの総スコアの導出、そしてこれらのスコアに基づくドキュメントのソートが含まれます。以下に擬似 PyTorch コードを示します：

import torch

def compute_relevance_scores(query_embeddings, document_embeddings, k):
    """
    Compute relevance scores for top-k documents given a query.
    
    :param query_embeddings: Tensor representing the query embeddings, shape: [num_query_terms, embedding_dim]
    :param document_embeddings: Tensor representing embeddings for k documents, shape: [k, max_doc_length, embedding_dim]
    :param k: Number of top documents to re-rank
    :return: Sorted document indices based on their relevance scores
    """
    
    # Ensure document_embeddings is a 3D tensor: [k, max_doc_length, embedding_dim]
    # Pad the k documents to their maximum length for batch operations
    # Note: Assuming document_embeddings is already padded and moved to GPU
    
    # Compute batch dot-product of Eq (query embeddings) and D (document embeddings)
    # Resulting shape: [k, num_query_terms, max_doc_length]
    scores = torch.matmul(query_embeddings.unsqueeze(0), document_embeddings.transpose(1, 2))
    
    # Apply max-pooling across document terms (dim=2) to find the max similarity per query term
    # Shape after max-pool: [k, num_query_terms]
    max_scores_per_query_term = scores.max(dim=2).values
    
    # Sum the scores across query terms to get the total score for each document
    # Shape after sum: [k]
    total_scores = max_scores_per_query_term.sum(dim=1)
    
    # Sort the documents based on their total scores
    sorted_indices = total_scores.argsort(descending=True)
    
    return sorted_indices
コピー

この手順は、トレーニングと推論時の再ランク付けの両方で使用されることに注意してください。ColBERT モデルはペアワイズランキング損失を使用してトレーニングされ、トレーニングデータは $(q, d^+, d^-)$ の三つ組で構成されます。ここで、$q$ はクエリ、$d^+$ はクエリに関連する（ポジティブな）ドキュメント、$d^-$ は関連のない（ネガティブな）ドキュメントを表します。モデルは、$q$ と $d^+$ の間の類似度スコアが $q$ と $d^-$ の間のスコアより高くなるような表現を学習することを目指します。

トレーニングの目的は、以下の損失関数を最小化することとして数学的に表現できます：$$\mathrm{Loss} = \max(0, 1 - S(q, d^+) + S(q, d^-))$$

ここで、$S(q, d)$ は ColBERT によって計算されるクエリ $q$ とドキュメント $d$ 間の類似度スコアを表します。このスコアは、モデルアーキテクチャで説明した後期相互作用パターンに従って、クエリとドキュメント間の最もマッチする埋め込みの最大類似度スコアを集約することで得られます。このアプローチにより、ポジティブとネガティブなドキュメントペアの類似度スコアの間でより大きなマージンを促すことで、モデルが与えられたクエリに対して関連するドキュメントと関連しないドキュメントを区別できるようにトレーニングされます。

tagColBERTv2 のノイズ除去監督

ColBERTv2 のノイズ除去監督は、チャレンジングなネガティブサンプルを選択し、クロスエンコーダーを使用した蒸留を活用することで、元のトレーニングプロセスを改良します。このトレーニングデータ品質を向上させる洗練された方法には、以下のステップが含まれます：

1. 初期トレーニング：クエリ、関連ドキュメント、非関連ドキュメントで構成される MS MARCO データセットの公式三つ組を使用。
2. インデックス作成と検索：ColBERTv2 の圧縮を使用してトレーニングパッセージのインデックスを作成し、各クエリに対してトップ k パッセージを検索。
3. クロスエンコーダーによる再ランク付け：MiniLM クロスエンコーダーによる再ランク付けを通じてパッセージ選択を強化し、そのスコアを ColBERTv2 に蒸留。
4. トレーニングタプルの形成：チャレンジングな例を作成するため、高ランクと低ランクの両方のパッセージを組み込んだ w-way タプルを生成。
5. 反復的な改良：ハードネガティブの選択を継続的に改善し、モデルの性能を向上させるためにプロセスを繰り返す。

このプロセスは、ColBERT のトレーニング体制における洗練された改良を表しており、そのアーキテクチャの基本的な変更ではないことに注意してください。

tagColBERT のハイパーパラメーター

ColBERT のハイパーパラメータは以下のようにまとめられます：

tagColBERT のインデックス戦略

各文書を 1 つの埋め込みベクトルにエンコードする表現ベースのアプローチとは異なり、ColBERT は文書（およびクエリ）を埋め込みの集合としてエンコードし、文書内の各トークンが独自の埋め込みを持ちます。このアプローチは本質的に、長い文書ではより多くの埋め込みが保存されることを意味し、これは元の ColBERT の問題点であり、後に ColBERTv2 で解決されました。

これを効率的に管理するカギは、ColBERT がインデックス作成と検索に（FAISS などの）ベクトルデータベースを使用することと、大量のデータを効率的に処理できるように設計された詳細なインデックス作成プロセスにあります。元の ColBERT 論文では、インデックス作成と検索の効率を高めるためのいくつかの戦略が言及されています：

• オフラインインデックス作成：文書の表現はオフラインで計算され、文書の埋め込みの事前計算と保存が可能です。このプロセスはバッチ処理と GPU アクセラレーションを活用して、大規模な文書コレクションを効率的に処理します。
• 埋め込みストレージ：文書の埋め込みは各次元に 32 ビットまたは 16 ビット値を使用して保存でき、精度とストレージ要件のトレードオフを提供します。この柔軟性により、ColBERT は有効性（検索パフォーマンスの面で）と効率性（ストレージと計算コストの面で）のバランスを保つことができます。

元の ColBERT には存在しなかった ColBERTv2 での残差圧縮の導入は、品質を維持しながらモデルの空間フットプリントを 6～10 倍削減する上で重要な役割を果たしています。この技術は、固定された参照セントロイドからの差分のみを効果的に捕捉して保存することで、埋め込みをさらに圧縮します。

tagColBERT の有効性と効率性

BERT の深い文脈理解を検索に組み込むことは、本質的に大きな計算リソースを必要とし、高いレイテンシーと計算コストのためにリアルタイムアプリケーションには実現性が低いと当初は考えられるかもしれません。しかし、ColBERT は後期相互作用メカニズムを革新的に使用することで、この仮定に挑戦し覆します。以下が注目すべきポイントです：

1. 大幅な効率性の向上：ColBERT は従来の BERT ベースのランキングモデルと比較して、計算コスト（FLOP）とレイテンシーを桁違いに削減します。具体的には、特定のモデルサイズ（例：12 レイヤーの「base」トランスフォーマーエンコーダー）において、ColBERT は BERT ベースのモデルの有効性に匹敵し、場合によってはそれを上回りながら、大幅に低い計算要求で実現します。例えば、再ランク付けの深さ k=10 では、BERT は ColBERT の約 180 倍の FLOP を必要とします。この差は k が増加するにつれて広がり、k=1000 で 13900 倍、k=2000 で 23000 倍にまで達します。
2. エンドツーエンド検索における Recall と MRR@10 の向上：高い検索パフォーマンスには（早期相互作用モデルに見られるような）クエリと文書表現のより深い相互作用が必要だという当初の直感に反して、ColBERT のエンドツーエンド検索セットアップは優れた有効性を示します。例えば、その Recall@50 は公式の BM25 の Recall@1000 とほぼすべての他のモデルの Recall@200 を上回り、各クエリ-文書ペアの直接比較なしに広大なコレクションから関連文書を検索する際のモデルの優れた能力を強調しています。
3. 実世界アプリケーションへの実用性：実験結果は ColBERT の実世界シナリオへの実用的な適用性を強調しています。そのインデックス作成スループットとメモリ効率により、MS MARCO のような大規模な文書コレクションを数時間以内にインデックス化でき、管理可能な空間フットプリントで高い有効性を維持します。これらの特質は、パフォーマンスと計算効率の両方が最重要である本番環境への展開に ColBERT が適していることを示しています。
4. 文書コレクションのサイズに対するスケーラビリティ：おそらく最も驚くべき結論は、大規模な文書コレクションを扱う際の ColBERT のスケーラビリティと効率性です。このアーキテクチャは文書埋め込みの事前計算を可能にし、クエリ-文書相互作用の効率的なバッチ処理を活用することで、文書コレクションのサイズに応じて効果的にスケールすることができます。このスケーラビリティは、効果的な文書検索に必要な理解の複雑さと深さを考えると直感に反するものであり、ColBERT の計算効率と検索効率のバランスを取る革新的なアプローチを示しています。

tagjina-colbert-v1-en の使用：8192 長の ColBERTv2 モデル

Jina-ColBERT は高速かつ正確な検索のために設計され、アーキテクチャの改良により長いシーケンス処理を可能にする JinaBERT の進歩を活用して、最大 8192 の長いコンテキスト長をサポートします。

jinaai/jina-colbert-v1-en · Hugging Face

私たちはオープンソースとオープンサイエンスを通じて、人工知能を進歩させ民主化する旅の途中です。

tag元の ColBERT に対する Jina の改良点

Jina-ColBERT の主な進歩は、そのバックボーンである jina-bert-v2-base-en にあります。これは元の ColBERT が使用する bert-base-uncased と比較して、はるかに長いコンテキスト（最大 8192 トークン）の処理を可能にします。この能力は広範なコンテンツを持つ文書を扱い、より詳細でコンテキストに基づいた検索結果を提供する上で重要です。

tagjina-colbert-v1-en と ColBERTv2 のパフォーマンス比較

私たちは jina-colbert-v1-en を BEIR データセットと長いコンテキストを重視する新しい LoCo ベンチマークで評価し、元の ColBERTv2 実装と非相互作用ベースのjina-embeddings-v2-base-en モデル。

この表は、特に長いコンテキスト長を必要とするシナリオにおいて、オリジナルの ColBERTv2 と比較して jina-colbert-v1-en の優れたパフォーマンスを示しています。なお、jina-embeddings-v2-base-en はより多くのトレーニングデータを使用していますが、jina-colbert-v1-en は MSMARCO のみを使用しており、これが一部のタスクにおける jina-embeddings-v2-base-en の良好なパフォーマンスを説明する可能性があります。

tagjina-colbert-v1-en の使用例

このスニペットは、長文書のサポートを示す Jina-ColBERT でのインデックス作成プロセスを概説しています。

from colbert import Indexer
from colbert.infra import Run, RunConfig, ColBERTConfig

n_gpu: int = 1  # Set your number of available GPUs
experiment: str = ""  # Name of the folder where the logs and created indices will be stored
index_name: str = ""  # The name of your index, i.e. the name of your vector database

if __name__ == "__main__":
    with Run().context(RunConfig(nranks=n_gpu, experiment=experiment)):
        config = ColBERTConfig(
          doc_maxlen=8192  # Our model supports 8k context length for indexing long documents
        )
        indexer = Indexer(
          checkpoint="jinaai/jina-colbert-v1-en",
          config=config,
        )
        documents = [
          "ColBERT is an efficient and effective passage retrieval model.",
          "Jina-ColBERT is a ColBERT-style model but based on JinaBERT so it can support both 8k context length.",
          "JinaBERT is a BERT architecture that supports the symmetric bidirectional variant of ALiBi to allow longer sequence length.",
          "Jina-ColBERT model is trained on MSMARCO passage ranking dataset, following a very similar training procedure with ColBERTv2.",
          "Jina-ColBERT achieves the competitive retrieval performance with ColBERTv2.",
          "Jina is an easier way to build neural search systems.",
          "You can use Jina-ColBERT to build neural search systems with ease.",
          # Add more documents here to ensure the clustering work correctly
        ]
        indexer.index(name=index_name, collection=documents)
コピー

tagRAGatouille での jina-colbert-v1-en の使用

RAGatouille は、RAG パイプライン内で高度な検索方法の使用を容易にする新しい Python ライブラリです。モジュール性と容易な統合のために設計されており、ユーザーが最先端の研究をシームレスに活用することを可能にします。RAGatouille の主な目的は、ColBERT のような複雑なモデルを RAG パイプラインで適用することを簡素化し、開発者が基礎となる研究について深い専門知識を必要とせずにこれらの手法を利用できるようにすることです。Benjamin Clavié のおかげで、jina-colbert-v1-en を簡単に使用できるようになりました：

from ragatouille import RAGPretrainedModel

# Get your model & collection of big documents ready
RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")
my_documents = [
    "very long document1",
    "very long document2",
    # ... more documents
]

# And create an index with them at full length!
RAG.index(collection=my_documents,
          index_name="the_biggest_index",
          max_document_length=8190,)

# or encode them in-memory with no truncation, up to your model's max length
RAG.encode(my_documents)
コピー

Jina-ColBERT に関する詳細な情報と更なる探求については、Hugging Face のページをご覧ください。

tag結論

ColBERT は情報検索の分野において大きな飛躍を表しています。Jina-ColBERT による長いコンテキスト長の実現と、ColBERT の後期相互作用アプローチとの互換性を維持することで、最先端の検索機能を実装しようとする開発者にとって強力な選択肢を提供しています。

複雑な検索モデルを RAG パイプラインに統合することを簡素化する RAGatouille ライブラリと組み合わせることで、開発者は今や高度な検索の力を容易に活用でき、ワークフローを効率化し、アプリケーションを強化することができます。Jina-ColBERT と RAGatouille の相乗効果は、高度な AI 検索モデルを実用的な使用のためにアクセスしやすく、効率的にする上で注目すべき進歩を示しています。