ColBERT が開発者による RAG の限界の克服にどのように役立つか

検索拡張生成 (RAG) は、現在では生成人工知能 (AI) アプリケーションの標準部分となっています。ベクトル データベースから取得した関連コンテキストをアプリケーション プロンプトに追加すると、精度が大幅に向上し、幻覚が軽減されます。これは、ベクトル検索結果の関連性の向上が RAG アプリケーションの品質と直接的な相関関係があることを意味します。

大規模言語モデル (LLM) のコンテキスト ウィンドウが増加しても、 RAG が依然として人気があり、関連性が高まっている理由は 2 つあります。

1. LLM の応答時間と価格は両方ともコンテキストの長さに比例して増加します。

2. LLM は依然として、大規模なコンテキストにわたる検索と推論の両方に苦労しています。

   
   

しかし、 RAGは魔法の杖ではありません。特に、最も一般的な設計である高密度パッセージ検索 (DPR) は、クエリとパッセージの両方を単一の埋め込みベクトルとして表し、単純なコサイン類似度を使用して関連性をスコア付けします。これは、DPR が、関連する検索用語をすべて認識するための幅広いトレーニングを備えたエンベディング モデルに大きく依存していることを意味します。

残念ながら、既製のモデルは、トレーニング データには一般的に含まれない名前などの珍しい用語に対処するのに苦労します。また、DPR はチャンキング戦略に過敏になる傾向があり、多くの無関係な情報に囲まれている場合、関連する部分が見逃される可能性があります。これらすべてにより、アプリケーション開発者は「最初から正しくやる」という負担が生じます。ミスをすると通常、インデックスを最初から再構築する必要が生じるためです。

ColBERT で DPR の課題を解決

ColBERT は、DPR の問題を実質的に解決するBERT言語モデルを使用してパッセージの関連性をスコアリングする新しい方法です。最初の ColBERT 論文のこの図は、なぜこれが非常に魅力的なのかを示しています。

これは、ColBERT のパフォーマンスを MS-MARCO データセットの他の最先端のソリューションと比較しています。 (MS-MARCO は、 Microsoft が最も関連性の高い文章を手動でスコア付けした一連の Bing クエリです。これは、より優れた検索ベンチマークの 1 つです。) 低く、右に行くほど優れています。

つまり、ColBERT は、レイテンシのわずかな増加を犠牲にして、ほとんどのより複雑なソリューションの分野を簡単に上回ります。

これをテストするために、デモを作成し、ada002 DPR と ColBERT の両方を使用して 1,000 を超える Wikipedia 記事のインデックスを作成しました。 ColBERT は、珍しい検索用語に対して大幅に優れた結果を提供することがわかりました。

次のスクリーンショットは、DPR がエイブラハム リンカーンの同僚であるウィリアム H. ハーンドンという珍しい名前を認識できないのに対し、ColBERT はスプリングフィールドの記事で参照を見つけたことを示しています。また、ColBERT の 2 番目の結果は別のウィリアムに関するものであり、DPR の結果はどれも関連していないことにも注意してください。

ColBERT は、緻密な機械学習の専門用語で説明されることがよくありますが、実際には非常に簡単です。わずか数行の Python とCassandra Query Language (CQL) を使用して、 DataStax Astra DBに ColBERT の取得とスコアリングを実装する方法を示します。

大きなアイデア

パッセージを単一の「埋め込み」ベクトルに変換する従来の単一ベクトルベースの DPR の代わりに、ColBERT はパッセージ内のトークンごとにコンテキストに影響されたベクトルを生成します。 ColBERT も同様に、クエリ内の各トークンのベクトルを生成します。

(トークン化とは、LLM で処理する前に入力を単語の一部に分割することを指します。OpenAIチームの創設メンバーである Andrej Karpathy が、これがどのように機能するかについて優れたビデオを公開しました。)

次に、各ドキュメントのスコアは、ドキュメント エンベディングのいずれかに対する各クエリ エンベディングの最大類似度の合計になります。

 def maxsim(qv, document_embeddings): return max(qv @ dv for dv in document_embeddings) def score(query_embeddings, document_embeddings): return sum(maxsim(qv, document_embeddings) for qv in query_embeddings)

(@ はドット積の PyTorch 演算子であり、ベクトルの類似性の最も一般的な尺度です。)

以上です。4 行の Python で ColBERT スコアリングを実装できます。これで、ColBERT について X (旧 Twitter) に投稿している 99% の人々よりもよく理解できるようになりました。

ColBERT 論文の残りの部分では次の内容が扱われます。

1. BERT モデルを微調整して、特定のデータセットに最適なエンベディングを生成するにはどうすればよいでしょうか?
2. ここに示されている (比較的高価な) スコアを計算するドキュメントのセットをどのように制限しますか?

最初の質問はオプションであり、この記事の範囲外です。事前トレーニングされた ColBERT チェックポイントを使用します。しかし、2 番目の方法は、DataStax Astra DB のようなベクトル データベースを使用することで簡単に実行できます。

Astra DB の ColBERT

RAGatouilleと呼ばれる、ColBERT 用の人気のある Python オールインワン ライブラリがあります。ただし、静的データセットを前提としています。 RAG アプリケーションの強力な機能の 1 つは、動的に変化するデータにリアルタイムで応答することです。そこで代わりに、Astra のベクトル インデックスを使用して、スコアを付ける必要があるドキュメントのセットを絞り込み、各サブベクトルの最適な候補に絞り込みます。

ColBERT を RAG アプリケーションに追加するには、取り込みと取得という 2 つの手順があります。

摂取

各ドキュメント チャンクには複数の埋め込みが関連付けられるため、次の 2 つのテーブルが必要になります。

 CREATE TABLE chunks ( title text, part int, body text, PRIMARY KEY (title, part) ); CREATE TABLE colbert_embeddings ( title text, part int, embedding_id int, bert_embedding vector<float, 128>, PRIMARY KEY (title, part, embedding_id) ); CREATE INDEX colbert_ann ON colbert_embeddings(bert_embedding) WITH OPTIONS = { 'similarity_function': 'DOT_PRODUCT' };

ColBERT ライブラリ ( pip install colbert-ai ) をインストールし、事前トレーニングされた BERT チェックポイントをダウンロードした後、ドキュメントを次のテーブルにロードできます。

 from colbert.infra.config import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint from colbert.indexing.collection_encoder import CollectionEncoder from cassandra.concurrent import execute_concurrent_with_args from db import DB def encode_and_save(title, passages): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encoder = CollectionEncoder(cf, cp) # encode_passages returns a flat list of embeddings and a list of how many correspond to each passage embeddings_flat, counts = encoder.encode_passages(passages) # split up embeddings_flat into a nested list start_indices = [0] + list(itertools.accumulate(counts[:-1])) embeddings_by_part = [embeddings_flat[start:start+count] for start, count in zip(start_indices, counts)] # insert into the database for part, embeddings in enumerate(embeddings_by_part): execute_concurrent_with_args(db.session, db.insert_colbert_stmt, [(title, part, i, e) for i, e in enumerate(embeddings)])

(私は DB ロジックを専用モジュールにカプセル化することを好みます。GitHubリポジトリ内の完全なソースにアクセスできます。)

検索

取得は次のようになります。

 def retrieve_colbert(query): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encode = lambda q: cp.queryFromText([q])[0] query_encodings = encode(query) # find the most relevant documents for each query embedding. using a set # handles duplicates so we don't retrieve the same one more than once docparts = set() for qv in query_encodings: rows = db.session.execute(db.query_colbert_ann_stmt, [list(qv)]) docparts.update((row.title, row.part) for row in rows) # retrieve these relevant documents and score each one scores = {} for title, part in docparts: rows = db.session.execute(db.query_colbert_parts_stmt, [title, part]) embeddings_for_part = [tensor(row.bert_embedding) for row in rows] scores[(title, part)] = score(query_encodings, embeddings_for_part) # return the source chunk for the top 5 return sorted(scores, key=scores.get, reverse=True)[:5]

以下は、最も関連性の高いドキュメント部分 ( db.query_colbert_ann_stmt ) に対して実行されるクエリです。

 SELECT title, part FROM colbert_embeddings ORDER BY bert_embedding ANN OF ? LIMIT 5

基本を超えて: RAGStack

この記事とリンクされたリポジトリでは、ColBERT がどのように機能するかを簡単に紹介します。独自のデータを使用してこれを今すぐ実装し、すぐに結果を確認できます。 AI のあらゆるものと同様、ベスト プラクティスは日々変化しており、新しい技術が常に登場しています。

最先端の技術に追いつきやすくするために、DataStax は、この機能とその他の機能強化を、LangChain と LlamaIndex を活用した実稼働対応の RAG ライブラリであるRAGStackに組み込んでいます。私たちの目標は、開発者が新しい機能へのステップアップを制御できる、RAG アプリケーション用の一貫したライブラリを提供することです。技術やライブラリの無数の変更に追従する必要がなく、単一のストリームがあるため、アプリケーションの構築に集中できます。今すぐ RAGStack を使用して、LangChain と LlamaIndex のベスト プラクティスをそのまま組み込むことができます。 ColBERT のような進歩は、今後のリリースで RAGstack に導入される予定です。

Jonathan Ellis、DataStax 著