2024年は情報検索技術に興味を持ち、情報検索関連のモデル作りを趣味で行っている @hotchpotch (セコン)です。Transfomer は割と適当にやっても、いい感じに学習してくれるので、楽しいですね。

というわけで、日々部屋でご家庭用GPUを回し、以前公開した情報検索に特化したモデル・日本語版SPLADE v1をさらに良い感じに学習させた、日本語版SPLADEのv2(japanese-splade-v2) を公開しました。JMTEB retrieval (情報検索タスク)のベンチマークスコアも、RAGでよく使う文章長の 512 トークン以下なら、かなりの高スコアでトップとなっており、モデルパラメータ数・性能を考えると、バランスの良い検索用モデルに仕上がったかなと思っています。

日本語版SPLADE v2 モデルの特徴

なおこの記事は、情報検索・検索技術 Advent Calendar 2024 24日目の記事となってます。

そもそもSPLADEって？

SPLADE はほとんどの方にとって、聞き慣れない言葉だと思います。まず SPLADEの前に、情報検索モデルである密ベクトル検索とスパースベクトル検索のお話を。

世の中、自然言語で検索するとえいば、流行りは「密ベクトル検索」(dense retrieval・テキスト埋め込み・embeddings 検索などとも呼ばれています)ですね。しかしながら、まだまだ現役で「スパースベクトル」(疎ベクトル・sparse retrieval)も各所で使われています。TF-IDF や BM25 といったキーワードベースの手法が、その代表例です。

例えば、「美味しいカフェを教えて」を検索するとしましょう。スパースベクトル検索（TF-IDFやBM25）は、キーワードがどれだけ重要かを数値化して結果を返します。この場合、「美味しい」と「カフェ」の頻出度や、それぞれの単語がどれだけ珍しいかを基に検索結果が決まります。そのため、特徴的なキーワードが一致している文章が上位に来やすい特徴があります。

一方で、密ベクトル検索（dense retrieval）は、単語やフレーズの意味をベクトルとして捉えます。「美味しいカフェ」に対して、文脈的に近い「人気の喫茶店」や「雰囲気のいいコーヒーショップ」も検索結果に含めることができます。これはニューラルネットワークモデルのTransformerを使って学習しており、単語や文章の意味を広く理解しているからです。

つまりに、BM25等のスパースベクトルでは「キーワードそのもの」が重視され、密ベクトルは「キーワードの意味やニュアンス」を重視しているという違いがあります。キーワードがピンポイントにマッチする必要があるか、あるいは意味的な広がりを重視するかで、どちらを使うかが変わってくるでしょう。

次元数の違い

さらに、密ベクトルとスパースベクトルには情報を表す「次元数」という重要な違いもあります。

密ベクトルでは、一般的な次元数は 384次元から3072次元程度とかなり大きく、モデルによってはさらに巨大になることもあります。例えば、OpenAIの text-embedding-3-large は標準では 3072次元もあります。この次元数が多いということは、ベクトル同士の計算（たとえば類似度を計算する内積やコサイン類似度）が非常に高コストになる・ストレージやメモリを大量に使うことを意味します。これが密ベクトルを使う上での課題の一つです。

なお、現実の検索システムでは巨大な次元を全て検索することは効率が悪すぎるので、近似最近傍探索(ANN)といったアルゴリズムを利用することで、精度を少々犠牲にし計算資源を効率化するといったトレードオフを行なっています。

一方、スパース(疎)ベクトルは、キーワードベースの場合、理論上は文書全体の語彙数を次元として大規模なベクトル空間を持つものの、実際には多くの次元がゼロ(疎)となるため、非ゼロ要素はごく一部に限られます。つまり、利用時にはクエリや文書内に出現する限られた語彙に対応する次元のみが利用され、それ以外はゼロのままです。たとえば、「美味しいカフェを教えて」というクエリでは、「美味しい」「カフェ」「教えて」に対応するほんの数次元だけが非ゼロとなり、他の数万～数十万に及ぶ潜在的な次元はゼロのままです。これにより、実運用時のストレージ・メモリ量や計算量は非常に軽減され、高速な検索が可能になります。

加えて、スパースベクトルは非ゼロの次元が具体的に何を意味するかが分かりやすいという利点もあります。「美味しい」「カフェ」のようなクエリのキーワードが、どの次元に対応しているかが明確なので、検索結果の解釈が容易です。

以下は、密ベクトルの一例です。

dense_vector = [
 0.0023, -0.0008, 0.0017, 0.0009, -0.0025,
 ... # 次元数分の要素が続く]

密ベクトルは全ての次元が意味を持っていますが、それぞれの値が具体的に何を意味するのかを理解するのは難しいです。

一方、スパースベクトルは以下のように構造が分かりやすく、非ゼロの次元が具体的なキーワードに対応します。

sparse_vector = {
  33721: 1.5, # 33721 番目が「美味しい」、に対応した次元
  1191: 2.3, # 1191 が「カフェ」に対応した次元
  997: 0.2 # 997 が「教えて」に対応した次元。頻出語はスコアが低い
  # 他の次元は全てゼロなので、あえて書かずに良い
}

この例では、「美味しい」や「カフェ」「教えて」というキーワードがスパースベクトル内でどの次元を表しているのかが一目瞭然です。そのため、どの単語が検索結果にどれだけ寄与しているかを解釈しやすいのが特徴です。

要するに、密ベクトルは「意味を広く捉える」ことが得意な一方で、次元数が多く計算負荷が高いという欠点があります。それに対してスパースベクトルは、次元数が比較的小さく効率的で、特にピンポイントなキーワード検索には非常に有利と言えます。

スパースベクトル検索の弱点

つまり、スパースベクトルは扱う次元数が小さく、どんな単語がマッチしたかも理解しやすいです。これだけ聞くと単純にスパースベクトルの方が良いことだらけじゃない？と思うかもしれません。

しかしながら、密ベクトル検索が自然言語検索において主流なのは、精度の問題があります。BM25をはじめとしたアルゴリズムでは、あらかじめ定義したキーワードや、似ている単語を類義語として人手で整備していったキーワードしか基本マッチしません。

そのため、先ほどの「美味しいカフェを教えて」で検索しても、基本的にBM25では「うまいコーヒー屋」は引っかからないのです。密ベクトルの場合は、もっと曖昧なベクトル表現になるため、「美味いコーヒー屋」が似ているぞ、とマッチさせることができるのです。

SPLADE - 文脈を理解し拡張するスパースベクトル

ピッタリとキーワードマッチするするような検索にはスパースベクトルでの検索は向いてそう(例えばECでの検索なら、似ているけど異なる商品が引っかかってもコレジャナイとなるので、きちんと一致した検索結果が向いていることが多いでしょう)ですが、自然言語を用いた検索には密ベクトル検索の方が向いてそうですよね。

なので、AIと対話するようなカジュアルな自然言語でも、対象となる文章を見つけてほしい、そんな要求が増えているので、密ベクトルの検索モデルが人気になっています。

ここで登場するのが、SPLADE(Sparse Lexical and Expansion Model)です。SPLADEの最大の特徴は、文脈を理解して、適切な複数の単語(token)を候補に挙げる点です。具体的に見てみましょう。

例えば「日本で世帯視聴率が最も高くなる時間帯は？」という検索クエリの場合、SPLADEは以下のような文脈理解を行い、関連する単語として出力できます。実際の日本語SPLADE v1の出力を見てみると、直接含まれる単語はもちろんのこと、以下のような直接含まれない単語も関連候補として検索することができます。

クエリに直接含まれる単語:
- 「日本」
- 「視聴」
- 「世帯」
- 「時間」
文脈から推論された関連単語:
- テレビ・放送関連：「放送」「番組」「枠」
- 指標関連：「率」「上昇」「高」
- 時期関連：「時刻」「期間」

従来のスパースベクトル検索では、これらの入力されたキーワードが完全一致する文書しか見つけられませんでしたが、SPLADEは文脈を理解して関連する単語も含めて検索できます。それでいて、スパースベクトルの特徴である高速な検索性能は維持されています。

このようにSPLADEは、検索意図を理解し、本質的に関連性の高い文書を見つけ出すことができるのです。この例では、「放送」や「番組」といったテレビ業界に特有の用語も適切にピックアップできていますね。

効率的な検索の仕組み

そして、これらの高度な検索を実現しながらも、SPLADEは効率的な検索が行えます。

例えば、検索クエリ「日本で世帯視聴率が最も高くなる時間帯は？」に対して
sparse_vector = {
    1423: 1.71,  # "日本" に対応
    5891: 1.59,  # "視聴" に対応
    8754: 1.57,  # "世帯" に対応
    2341: 1.33,  # "時間" に対応
    9876: 0.96,  # "放送" に対応
    # ...その他の関連次元
}

このように、必要な情報だけをスパースベクトルとして保存しており、小さな次元でマッチする検索を行うことで、効率良い実現しています。また重要なのは、SPLADEのスコアは単なる出現頻度ではなく、文脈における重要度を表している点です。

なぜSPLADEが良いのか？

他のベクトル検索と比較してみましょう。

従来のスパースベクトル検索（BM25など）
- 「視聴率」+「時間帯」というキーワードの組み合わせに依存
- 「放送のピークタイム」のような言い換えに弱い
- ズバリのキーワードマッチには強い
- 検索結果の説明が容易
密ベクトル検索
- 文章とドキュメントを同一の密ベクトルで表現するため、精度を求めるとモデルサイズやベクトル次元数が大きくなりがち
  - 推論速度や検索速度に影響
- 検索結果の解釈が難しい
SPLADE(文脈を理解したスパースベクトル検索)
- 文脈を理解した検索が可能
- 高速な検索性能を維持
  - 検索クエリが20〜40次元程度・文章が150〜400次元程度
  - 重要度が低い単語は検索しない・インデクスしないといった方法等、精度と速度のトレードオフも実行時に採択可能
- 検索結果はどの単語トークンがマッチしたかわかるので、解釈が容易

このように、SPLADEは現代の検索システムに求められる要件を、バランスよく実現しています。

実際の性能は高いの?

実際の性能、とりわけ自然言語の質問文に対して、適切なドキュメントを取得する性能を見てみましょう。

日本語版 SPLADE v2 スコア

これは記事冒頭の JMTEB retrieval ベンチマーク結果(nDCG@10)ですが、512トークン以下の文章では、ほとんどにおいて日本語SPLADE-v2が最高のスコアとなっています。なお、ベンチマークタスクの"nlp_journal_abs_intro"と"nlp_journal_title_intro"は、対象のドキュメントが512トークン以上と長く、最大入力トークンが小さなモデルでは、軒並み低い数値となっています。

ただ実際に利用するシーン、例えばRAGのための情報検索では、小さな長さにドキュメントを分割する(チャンク分割)ことが多いため、用途次第ですが512トークン以下しか扱えずとも全く困らないこともあります。

なお、JMTEB retrieval で扱っているデータセットのざっくりとした説明は以下です。

JaGovFaqs_22k
- 日本の官公庁の「よくある質問」を元にしたQAデータセット
- query: 3,420件
- 対象文章: 22,794件
- ほぼ512トークン以下
Mr. TyDi
- Wikipediaの記事断片から人手で作成した質問と、それに関連する文章の情報検索評価用ベンチマークデータセット
- query: 720件
- 対象文章: 7,000,027件
- ほぼ512トークン以下
JAQKET
- クイズAI王で用いられた、クイズ質問とその回答が含まれるWikipedia記事のデータセット
- query: 997件
- 対象文章: 114,229件
- ほぼ512トークン以下
NLP Journal
- Japanese NLP Journal LaTeX Corpusから、タイトル・概要・イントロダクションを組み合わせたデータセット。この中で、イントロダクションは、512トークンを超えるものが多い。
  - nlp_journal_title_abs
    - query: 論文タイトル(404件)
    - 文章: 論文概要(504件)
    - ほぼ512トークン以下
  - nlp_journal_title_intro
    - query: 論文タイトル(404件)
    - 文章: 論文イントロダクション(504件)
    - ほぼ512トークン以上
  - nlp_journal_abs_intro
    - query: 論文アブストラクト(404件)
    - 文章: 論文イントロダクション(504件)
    - ほぼ512トークン以上

なお、このベンチマークのデータセット全て、例えばMr. TyDi(やその派生データセットである MIRACL)、JAQKET(やその派生データセットであるJQaRA)のtrainやdev,testといったデータを、日本語SPLADE-v2では学習データとして利用していません。学習データとして利用すると、そのドメイン課題に強くなるのですが、汎化性能を測るために、学習データ元には利用しないようにしました。

モデルサイズや次元数は?

モデルパラメータ数と、次元数も先ほどの表の通りです。なお、モデルパラメータ数は、ざっくりとレイヤの重みで計算しています。モデルパラメータ数が大きくなれば大きくなるほど、基本学習や推論時のコストが高くなります。また、文章の出力次元が大きければ大きいほど、メモリやストレージを利用します。

SPLADEの出力次元数(非ゼロの要素数)は、テキストにより異なるため、JMTEBの質問やドキュメントの次元数をざっくりと載せています。

ライセンスは?

日本語SPLADE v2は、特に利用制限を設けていない(MITライセンス)のため、ご自由にご利用いただけます。

プログラムからの利用方法

huggingface.co/hotchpotch/japanese-splade-v2 にサンプルコードがのっていますので、そちらの項目をご覧ください。

FAQ

スパースベクトル検索って本場運用できるの？

できます。古くからある検索技術のTF-IDFやBM25がそもそもスパースベクトル検索で、様々な検索システム(Elasticsearch, vespa, qdrant 等々)がスパースベクトル検索や、密ベクトル検索とスパースベクトル検索を組み合わせたハイブリット検索もサポートしています。

密ベクトルモデルよりSPLADEの方が良いの？

ベンチマーク的には良いですが、利用ケースによります。自然言語の質問から対応するドキュメントを見つけるというような、EC検索などに比べると単純な検索システムにおいても、どんな質問と文章を想定しているか、達成したい要件は何かによって最良は異なります。単純なBM25がベストの場合もあります。

密ベクトルモデルとSPLADEは、方向性が異なった検索結果を返すこともあり、二つの検索結果を組み合わせて使うといったハイブリット検索もおすすめです。

また、ハイブリット検索においては、密ベクトル・スパースベクトル、どちらかのモデルを扱いたいドメインのデータで学習させて、片方のモデルは汎化性能を、片方のモデルはドメイン特化させるといった検索手法もおすすめです。今回作った日本語SPLADE-v2のTrainer実装(YAST)や学習用データ・設定も公開されているため、ドメインのデータからクエリとドキュメントの学習用データセットを作り学習元データに追加することで、大きく検索精度が上がる可能性があります。最近ですとテキストさえあれば、LLMを使って教師データを作ると言った合成データセットも簡単に作成できるため、データの活用の幅が広がっていますね。

おわりに

SPLADEは、文脈を考慮した単語拡張を用い、BM25などのキーワードベースの弱点をカバーし、ニューラルネットワークを用いた実用的な検索システムの一つの選択肢として注目を集めています。

日本語を適切に学習させた情報検索モデル日本語SPLADE v2は、とりわけMr.TyDiのような自然言語での質問タスクに対し、現時点では最高性能のモデルかなと思っています。また情報検索に用いるモデルとして高性能かつバランスが良いモデルで、プロダクション・本番環境でも利用しやすいでしょう。

本モデルやこの記事が、AI開発・自然言語処理・情報検索をしている皆さんに少しでも有益になれば幸いです。

A Day in the Life

情報検索モデルで最高性能(512トークン以下)・日本語版SPLADE v2をリリース