関連ドキュメントを利用してRAGの精度を上げるCDE

導入

こんにちは、株式会社ナレッジセンスの須藤英寿です。普段はエンジニアとして、LLMを使用したチャットのサービスを提供しており、とりわけRAGシステムの改善は日々の課題になっています。

この記事では、ドキュメントのベクトルデータ生成時に他のドキュメントのベクトルデータを利用することで、検索の精度を上げるEmbeddingモデル CDE (Contextual Document Embeddings)について紹介します。

https://arxiv.org/pdf/2410.02525

サマリー

CDEは従来のEmbeddingとは異なり、対象のテキストと関連性の高い別のドキュメントの両方を用いて、ベクトルデータを生成するEmbeddingモデルとなっています。

これにより、従来のEmbeddingの能力に加えて、他のドキュメントのコンテキストを利用したベクトルデータが生成可能になり、検索性能の上昇につながっています。

問題意識

Embeddingモデルの欠点

LLMの発展に合わせて、文書の意味をベクトル化するEmbeddingモデルも大きな発展を見せて非常に高い性能を実現しています。その一方で、前後の文脈の喪失や文書全体の中の相対的な比較ができないためにEmbeddingによるベクトル化とコサイン類似度の検索だけでは、現時点では性能に限界があることが知られています。

既存の対応方法

前後の文脈が失われる問題については、以前の記事が一つの対応手法として挙げられます。簡単に説明すると、入力に広い範囲のテキストを渡し、そこから狭い範囲のチャンクを算出することで狭い範囲のチャンクがその前後の文脈を意識したベクトルデータとして取り出せる。という手法になっています。

https://zenn.dev/knowledgesense/articles/bb7a4943e228f6

次の、文書全体の相対的な比較手段としてはBM25の活用が有名かと思います。簡単に説明すると単語単位で出現頻度で重み付けしたうえで、文章同士の比較を行うための手法となっています。

https://zenn.dev/knowledgesense/articles/cec1cd43244524#2.-ハイブリット検索

今回紹介する手法CDEは、後者の文書全体の相対的な情報をベクトルデータに組み込む手法となっています。

手法

cdeは、事前にEmbeddingのモデルを学習する必要があります。また前後の文脈を取り入れるために、通常のEmbeddingモデルと比較しても若干複雑な計算となっています。この記事では、ドキュメントのベクトル化方法に限定して紹介します。なお、検索クエリについても同様の方法で計算されます。

ベクトルデータの計算方法

あるテキストをベクトル化する際に、２つのモデルを使用して最終的なベクトルデータを算出します。

1. 単一文書からベクトルデータを算出する(M1)

まずは、単一の文書をベクトルデータに変換します。これは一般的なEmbeddingと大差なく先頭のトークンを任意の次元のベクトルデータに変換します。(論文中では768次元でした)

2. 複数の文書と対象のテキストを利用してベクトルデータを算出する(M2)

次に、ベクトル化したいテキストに近しいベクトルデータを任意の個数選出します。(論文内では512個)この選出方法は様々手段が想定されていますが、BM25を使った類似度判定やまとめられたドキュメントであれば元々の文書同士が近しく配置されていたものなど、適切な方法を選ぶのが良いようです。

そして、対象となるテキストの上から任意のトークン(論文中では512トークン)を取得し、これもベクトルデータに変換します。ただし、ここでの変換はEmbeddingモデルを使用するわけではなく単にトークンをベクトルデータに変換するだけの役割を担っています。

これら２つのベクトルデータを計1024個入力として、M2モデルを使用して最終的に単一のベクトルデータを算出します。

成果

こちらは各種埋め込みモデルとCDEとの検索性能の比較を行った表です。cde-small-v1に対応した結果が２種類存在します。Randomは検索に使用するドメインに関係なく学習した場合の性能で、Contextualは検索に使用するドメインに合わせた性能を示しています。ドメインに特化することでCDEの性能が最も高い一方で汎用的にしたRandomの場合は一部最新のモデルと比較して、性能が落ちるという結果になりました。

まとめ

この論文では、Embeddingモデルの入力に他の文書データを用いることでより性能を発揮できることを示していました。直感的には精度が上がりそうな印象はありつつも現時点では既存のモデルと比較しても顕著に性能が上がっていることまでは確認できてはいません。特に、ここでは触れられていませんがすでに更に性能の良いEmbeddingモデルが出現しているためこのモデルがそのまま実用的には使えないかもしれません。

しかし、ベクトルデータを作成する際にそのドキュメント以外の情報を用いるという考え方は非常に発展性のある考え方かと思います。ここからさらにEmebddingモデルの性能が上がるような手法が確立されていくのが非常に楽しみです。