こんにちは。エンジニアの大橋です。
先日、とある技術検証にてAIエージェントのRAG構築を担当した人が、ベクトル検索の精度がイマイチということで悩んでいました。詳しく話を聞くと、どうやら特定のキーワードでベクトル検索した場合に、あまり良い検索結果が得られなかったようです。その後もう一度話を聞いてみると、ベクトル検索のEmbeddingモデルを文単位からトークン単位に変えてみると、かなり精度が上がった、というのです。この一連の会話から、私自身この辺りのEmbeddingについての理解を再確認しておこうと思いました。
こういった背景もあり、今回は、RAGのベクトル検索に必要なText Embedding（テキスト埋め込み）について書きたいと思います。この機会にText Embeddingについて学んだプロセスの中で疑問に思ったことや理解したことを共有してみたいと思います。
そもそもEmbeddingはなぜ必要なのか
そもそも、Embeddingはなぜ必要なのでしょうか。それは、RAGのベクトル検索におけるEmbeddingの位置付けから考える必要がありそうです。
AIエージェント開発では、多くの場合、RAGによりLLMに何らかの外部データを参照させます。RAGの目的は、「ユーザーの質問意図に合ったノイズの少ない的確なコンテキストをLLMに提供すること」です。ここでユーザーの質問意図に合った情報を取得するには、セマンティックサーチ（意味検索）との相性が良いと考えられています。そのため、RAGの検索手法にはベクトル検索がよく使用されます。
ベクトル検索はセマンティックサーチの手法ですが、内部的には、文章中の単語の意味的な関係性を踏まえたベクトル表現の類似度を計算しています。ここで登場したベクトル表現がEmbedding（意味ベクトル）です。つまり、ベクトル検索には事前にEmbedding生成が必要となるわけです。
Embeddingモデルとは
このEmbedding生成には、Embeddingモデルが使用されています。Text Embeddingモデルの場合、大量のテキストデータを用いて事前学習されており、この学習プロセスによって、単語や文章の意味的な関係性を数値ベクトルとして表現する能力を獲得しているのです。
Embeddingモデルの種類
では、Embeddingモデルはどのような分類に分かれるのでしょうか。いろいろな分類軸はあると思いますが、その一つにEmbeddingの生成単位があります。これには、「Document Embedding」「Paragraph Embedding」「Sentence Embedding」「Token Embedding」「Character Embedding」があります。今回は、冒頭の疑問にも出てきた「Sentence Embedding」と「Token Embedding」について考えていきたいと思います。
Sentence Embedding
Sentence Embeddingモデルの場合、文単位に対する単一のEmbeddingが生成されます。この時生成されるEmbeddingは、文章全体の意味を表現したものとなっています。なるほど、と思う一方で、では、一体どうやって文章全体の意味を表現したEmbeddingを生成しているのか、と疑問が湧きます。文全体に対して一括でEmbedding生成する仕組みがあるのでしょうか。
実際は、最初に文を構成するトークン単位でEmbeddingを生成し、それらを何かしらの方法でプーリング（集約）しているようです。このような方法で文章全体のEmbedding生成ができるのは、Sentence Embeddingモデルの事前学習時に秘密があります。
この部分を理解するには、Sentence EmbeddingモデルのSBERT（Sentence BERT）の開発経緯を追うのが分かりやすいです。論文の詳細な内容は他の記事でも紹介されていると思うので、ここでは簡単に私の理解で説明します。SBERTの論文に記載されているように、元々、BERT出力で生成されたトークン単位のEmbeddingをプーリングするだけでは、文章全体のEmbeddingとしては不十分でした。
The most commonlyused approach is to average the BERT output layer (known as BERT embeddings) or by using the output of the first token (the [CLS] token). As we will show, this common practice yields rather bad sentence embeddings, often worse than averaging GloVe embeddings (Pennington et al., 2014).
そこで、追加のファインチューニング（Siamese Network と Triplet Network Structure）を行うことで、文章全体の意味を表現したEmbedding生成ができるようになった、というわけです。
we present Sentence-BERT (SBERT), a modification of the pretrained BERT network that use siamese and triplet network structures to derive semantically meaningful sentence embeddings that can be compared using cosine-similarity.
「Siamese Network」と「Triplet Network Structure」の詳細を理解することも重要ですが、ここでは「文全体のEmbedding生成には、トークン単位のEmbeddingを生成し、それらをプーリングしている」ことを押さえておきたいところです。
なお、プーリング戦略には「CLS-tokenの利用」「最大プーリング」「重み付け平均プーリング」が紹介されており、デフォルトでは平均プーリングが採用されているようです。
このように、Sentence Embeddingモデルは、事前学習の時点で文章全体のEmbedding生成にチューニングされているため、私たちはそのモデルに文章を入力するだけで、文全体の意味を表現するEmbeddingが生成できる、というわけですね。
では、Sentence Embeddingはどのくらい使用されているかというと、RAGのベクトル検索ではデファクトスタンダードになっているようです。といっても、よくある問い合わせ応答やドキュメント検索などの用途では、文全体の意味の類似性で検索したい場合が多いので、これは妥当な結果のように思えます。
一方で課題もあります。Sentence Embeddingは文全体の意味を表現できる一方、特定のトークンに対する類似度は考慮されにくい傾向があります。これにより何が起きるかというと、例えば文の意味を決定づける重要なトークンが含まれている場合に、そのトークンの類似度がそれ以外のトークンと同率で計算されてしまうため、特定のトークンの重要度を適切に反映できなくなってしまいます。
このような課題への対策としては、多くの場合、ハイブリッド検索が活用されるようです。ハイブリッド検索は、ベクトル検索だけでなく従来のキーワード検索の結果を加味し、それぞれの検索結果を重み付けすることで最終的な検索結果を調整する手法です。ハイブリッド検索については、別の機会に詳細を書きたいと思います。
Token Embedding
Token Embeddingモデルの場合、トークン単位でEmbeddingが生成されます。つまり、入力データを構成するトークンの数だけEmbeddingが生成されるということです。Token Embeddingは、入力データを複数のEmbeddingで表現するため、Multi-Vector Embeddingに分類されることもあるようです。
Token Embeddingを使ってベクトル検索を行う場合、トークン単位での意味検索が可能になるため、文単位よりも高精度の検索が可能になるようです。Qdrant（オープンソースのベクトル検索エンジン）の公式サイトには、Single Vector Embedding（Sentence Embeddingなど）と比較して、入力データのより詳細な意味を捉えた検索が可能であると記載されています。
ColBERT is an embedding model that produces a matrix (multivector) representation of input text, generating one vector per token (a token being a meaningful text unit for a machine learning model). This approach allows ColBERT to capture more nuanced input semantics than many dense embedding models, which represent an entire input with a single vector.
ただし、トークン単位でEmbeddingを生成する分、必要な保存容量や類似度算出の計算量は増えがちです。つまり、性能とコストのトレードオフを考慮する必要があるというわけです。
By producing more granular input representations, ColBERT becomes a strong retriever. However, this advantage comes at the cost of increased resource consumption compared to traditional dense embedding models, both in terms of speed and memory.
といっても、ColBERTv2のような改良モデルも登場しており、このような課題は改善されつつあるので、一概にトレードオフだと断定できない点には注意が必要です。
ちなみに、引用文の中にDense Embeddingというのが登場しますが、これは各次元が複雑な意味を持つ密ベクトルを指します。対照的な概念としてSparse Embeddingがあり、こちらは反対に多くの成分が0の疎ベクトルを指します。ここでは簡単な説明にとどめますが、Denseは意味検索の文脈、Sparseはキーワード検索の文脈に関連する表現であると押さえておけば良いでしょう。
では、Token EmbeddingモデルはRAGのベクトル検索に使われるかというと、それがそうでもないようです。 Qdrantのサイトによると、Token Embeddingモデルの代表例であるColBERTは、大規模コーパス（検索対象）の場合には、検索用途には向かないと説明されています。その理由は、検索速度が遅いことにあるようです。
Despite ColBERT being a powerful retriever, its speed limitation might make it less suitable for large-scale retrieval. Therefore, we generally recommend using ColBERT for reranking a small set of already retrieved examples, rather than for first-stage retrieval.
ここではToken Embeddingの用途にも触れられており、それは検索後のリランク処理です。リランク処理とは、検索結果を精査し、より正確な検索結果にランク付けし直す処理のことを指します。リランク処理の場合は、検索対象が絞られているので、計算速度よりも精度を優先させることができる、というわけですね。
このように、Token Embeddingはリランク処理でよく使われているのだな、と思ったのですが、調査を進めると、どうやらリランク処理ではCross-Encoderの方が適しているようです。SBERTの公式サイトには、以下のように記載されています。
The retriever has to be efficient for large document collections with millions of entries. However, it might return irrelevant candidates. A re-ranker based on a Cross-Encoder can substantially improve the final results for the user.
ここで新たに登場したCross-Encoderとは一体なんでしょうか。SBERTの公式サイトによると、Cross-Encoderは文章のペアを入力として受け取り、その類似度を出力するようです。これまで見てきたSentence EmbeddingとToken EmbeddingはBi-Encoderの一種で、1つの入力に対して、Embeddingを出力します。つまり、Cross-Encoderは個別のEmbeddingを出力せず文章のペアの類似度を出力し、Bi-Encoderは個別のEmbeddingを出力する、というわけですね。
また、Cross-EncoderはBi-Encoderよりも高精度の類似度計算が可能ですが、実務のアプリケーション用途での利用は難しいと説明されています。
As detailed in our paper, Cross-Encoder achieve better performances than Bi-Encoders. However, for many application they are not practical as they do not produce embeddings we could e.g. index or efficiently compare using cosine similarity.
その理由は、計算効率にあります。Bi-Encoderの場合には、事前にEmbedding生成を行い、インデックスを構築しておくことができ、検索時にはEmbeddingを使った類似度計算（コサイン類似度など）をするだけで済みます。一方、Cross-Encoderは類似度計算のタイミングで内部的にEmbedding生成による類似度算出を行う必要があり、より多くの時間がかかってしまうというわけです。
最後に、Token EmbeddingがRAGのベクトル検索用途で使用される例はあるのか？が気になったので調査したところ、あまり数はなさそうですが、例えばスペルチェックの例が見つかりました。確かに、スペルチェックはトークン単位の表記揺れを検出する必要があるため、Token Embeddingの好例と言えそうです。
冒頭の疑問への回答
ここで、冒頭で触れた疑問に立ち返ってみたいと思います。疑問は「なぜベクトル検索のEmbeddingモデルを文単位からトークン単位に変えると、特定のキーワードでの検索精度が上がったのか？」でした。
この理由は、ここまで説明してきた内容から明らかになります。Sentence Embeddingは文全体の意味を単一のEmbeddingで表現するため、各トークンの意味が平均化されてしまいます。その結果、検索時に重要なキーワードが含まれていても、その重要度が他のトークンと同等に扱われ、埋もれてしまう可能性があります。
一方、Token Embeddingはトークン単位でEmbeddingを生成するため、重要なキーワードの意味を個別に保持できます。検索時には、クエリに含まれるキーワードと、ドキュメント内の対応するトークンとの類似度を直接計算できるため、キーワードがマッチする精度が高くなります。
つまり、冒頭のケースでは、特定のキーワード（専門用語や固有名詞など）が検索の成否を決めるポイントだったため、トークン単位で意味を保持できるToken Embeddingモデルに変更することで精度が改善した、というわけです。
早見表
ここまで書いた内容を一目で思い出せるよう、表にまとめておきたいと思います。（あくまで私の理解であり、簡略化している部分がある点にはご注意ください）
| 比較項目 | Sentence Embedding | Token Embedding |
|---|---|---|
| 生成単位 | 文単位に対して単一のEmbedding | トークン単位でEmbedding（入力データを構成するトークンの数だけ生成） |
| 生成プロセス | トークン単位でEmbeddingを生成し、プーリング（集約） | トークン単位でEmbeddingを生成 |
| 検索精度 | 文全体の意味検索に適している | トークン単位での意味検索が可能で、文単位より高精度 |
| 特定キーワード | 各トークンの意味が平均化され、重要なキーワードが埋もれる可能性がある | 重要なキーワードの意味を個別に保持でき、キーワードマッチの精度が高い |
| 保存容量 | 少ない（文ごとに1つのEmbedding） | 多い（トークンの数だけEmbedding） |
| 計算コスト | 低い | 高い |
| 検索速度 | 速い | 遅い |
| 主な用途 | RAGのベクトル検索（デファクトスタンダード） | リランク処理、スペルチェックなど |
| 代表モデル | SBERT（Sentence BERT） | ColBERT |
DataCloudの場合
では、SalesforceのDataCloudの場合、どのようなEmbeddingモデルを使用できるのか確認したいと思います。
利用可能なText Embeddingモデルは以下の2つです。
- E5-large-v2
- Multilingual-e5-large
E5-large-v2は英語のみ対応しており、Multilingual-e5-largeは多言語対応です。そもそも入力データが日本語の時点でMultilingual-e5-largeしか選択肢がないため、選択の余地はそれほどありません。しかし、どちらもSentence Embeddingモデルなので、トークン単位の検索精度が低くなる可能性があることを考慮しておく必要があります。対策として、DataCloudではハイブリッド検索が用意されており、その使用が推奨されています。
まとめ
以上のように、今回はEmbeddingについて学ぶ中で疑問に思ったことや理解したことをつらつらと書いてみました。あまり掘り下げられなかった「Embedding生成時のパラメータ（次元数やプーリング方式）」「ハイブリッド検索」「検索時のリランク処理」などについては、別の機会があれば調査したいと思います。また、分量が多くなってしまったので、具体の実装については省略しましたが、こちらについても別の機会があれば書いてみたいと思います。
以上となります。最後までお読みいただき、ありがとうございました。
