ただ、もうちょっと学習速度が欲しいことがあったり、また複数GPU環境での学習の知見を得たい為、今回 RTX5090 x2 構成の自作PCを作ってみた。その際、最近はコンシューマ向けGPUからNVLinkが廃止され、かつ電力消費も上がったため意外と RTX5090 2枚を使った事例がなく少々調べるのに苦労したため、まとめてみた。なおこの自作例は2025年末時点の例である。

電源

RTX5090 x2 で苦労する点といえば、まず電源である。RTX5090のTBPは最大575Wで、それが2機。CPUその他諸々の消費電力を考えると、最低1600W電源以上は欲しい。しかしながら、日本のご家庭用コンセントの100Vでは最大1500Wとなっているし、調べる限り一般的なPC用電源は100Vでは1300Wまでしか販売されていない。

なので、1300W電源までは多彩なラインナップがあるのだが、それ以上の電源になると極端に少なくなる。それ以上の電源は、電源の入力コネクタの口が一般的なPC電源のC13ではなく、C19という形状に変わる。そこに200V電源を供給することで1300W以上の出力を得ることができる。

というわけで、電源を繋げるべくNEMA規格の20A250V対応の壁コンセント(パナソニックのWF2520Bを選んだ)を工事して設置し、そこから200V20A供給(最大4000W)を可能とした。電源ケーブルはNEMA L6-20P → IEC 60320 C19のものが必要になるので、Schneider Electric AP8753J Power Cord, Locking C19 to L6‑20P を利用。なお、このコンセントにはブレーカーから単独で給電している。

また電源は評判が良さそうな1650WのASRockのTaichi TC-1650Tを選択した。GPUに大きな電力を供給する12V-2x6コネクタの安全性を考慮したATX3.1に対応している。またこの電源は、標準で100V用の電源ケーブル（C19-C20）に挿せる電源ケーブルがついてくる。これを使うと1300Wまでだが、このケーブル自体が市場にほぼ出回っていないので、テスト起動などで活用できて好感度高し。

追記:電源 x2 が設置できるケースを使い、1300W電源x2を各々100V電源に繋ぐ、という方法もあるようだ

GPU

RTX5090は膨大な発熱処理のため、空冷モデルではPCIスロット3〜4個分を占める厚さのものがほとんどだ。この厚さのGPUを2枚使う場合は、ライザーケーブルを使って物理的に離して設置しないと、ケースやマザーボードでぶつかって装着できないという問題が起きることが多いだろう。

このため選択肢として

• 3スロット以下の空冷モデルを使う（隙間なくみっちり挿すことになり熱が心配だが）
• 2枚のGPUとも簡易水冷にする
• 1枚を空冷、1枚を簡易水冷にする
• ライザーケーブルを使ってどうにかする

あたりが考えられる。すでに自分は厚さ約3.5スロットのRTX5090を保持していたため、1枚を簡易水冷、もう1枚を空冷とし、以下の2つのGPUの選択となった。なお、もしRTX5090を1枚も持っていなければ、価格が少々高くなるが、GPUを2枚とも簡易水冷とし、CPUを空冷としたほうが、ケース内部の配置に余裕が出て楽になるし、GPU温度もさらに下げることが可能なので、そちらにしたと思う。

なお今回使ったGPUは以下。

• MSI GeForce RTX 5090 32G VENTUS 3X OC
  • 約3.5スロットの厚みがある空冷（元々持っていた）
• MSI GeForce RTX 5090 32G SUPRIM LIQUID SOC
  • 厚さ2スロット強の空冷 + 120x360の水冷ラジエータ

また、お金に余裕があれば、RTX 6000 Pro (RTX5090と同じBlackwellアーキテクチャでメモリが96GB)や、性能は少々落ちるが消費電力がかなり下がって300WのRTX PRO 6000 Blackwell Max-Q を選ぶという手もある。Max-Qは排熱の考慮もだいぶ減りそうなので、設置が楽そうだ。

マザーボード

マザーボードに求める要件は、PCIe 5.0 x8 の速度で二枚刺さること、上部に刺すGPU1(水冷)と、下部にさすGPU2(空冷)の間が2レーンあって問題ないことである。実際に海外で RTX5090x2 の組み立て済みPCとして販売されている実績を評価し、Ryzen が乗るASUS ProArt X870E-CREATOR WiFi AMD AM5 X870E ATXを選択した。

なおオンボードWiFi 7チップが載っているが、Linuxのカーネルドライバーが現状無さそうなので、オンボードWiFi で繋ごうと思ってる人は注意が必要かもしれない。自分の用途では無線は使わず有線LAN接続なので特に困っていない。

ケース

ケースは、下部に3.5レーン分の厚さがあるGPUを差した時、問題なくある程度の空間があり、簡易水冷のラジエータを2機(CPUとGPU1)を問題なく設置できるもので検討した結果、CORSAIR 7000D AIRFLOWとなった。通常のケースより一回り大きいが、ケース内空間が大きいことも冷却においてはメリットが大きい。PC内部が見れるガラスパネルは不要なのだが、いざ作ってみたら結構かっこ良くて満足。

エアフロー（空気の流れ）

ケース内で最大1650W程度の電力消費が発生し、その時発生する熱は相当なものなので、良い形で空気を循環させねばならない。

簡易水冷はCPUとGPU1、空冷はGPU2で使うので、どのように外気を取り込み・排出すると良いのかを考える必要がある。PCの冷却ファンは、表裏をひっくり返すことで、吸気・排気を簡単に変更することができる。AIに壁打ちしながらベストを考えた結果、以下のエアフローになった。ただこの辺は素人なので、もっと良い配置があるかもしれない。

• フロント・吸気
  • 140mm ファンx2(ケース付属、本当は140mmをもう一つ追加したほうが良い)
  • GPU2(空冷)に当たる位置に配置
• サイド（ケース横）・吸気
  • GPU1の水冷・120mm x3
• トップ・排気
  • CPUの水冷・120mm x3
• リア・排気
  • ケース付属・140mm x1

この辺までが、組み立てる上であまり情報がなかったので苦労したところだ。続いて以下は好みで選べば良いと思うが、機械学習マシン視点からコメント入りで組み立てたパーツを紹介する。

CPU

16コア32スレッドのAMD Ryzen9 9950X。9950X3Dも出ていたが、ゲームをしない環境では誤差程度の性能差であり、差額も2万円ほど9950Xの方が安かったので9950Xを。データ加工などの処理は並列で行うことがほとんどなため、CPUコアはあればあるほど良いが、これ以上のコア数になるとThreadripperになってしまうため、16コアで。

RAM

せっかくなので上限の192GBを、と思っていたが、AI関連データセンタによるメモリ需要急増のため、2025年9月と比べると4-5倍ぐらいの値段で高止まりしていて、いくらなんでも高すぎるので DDR5-5600 32x2=64GBで。本当はECCにしたかったがこちらも高すぎるので…。自分の使い方の場合、容量は64GBだと時々swapにアクセスが発生する程度・かつswapがNVMe書き込み読み込みともに結構速いものを使っていることもあり、RAM はもっとあったら嬉しいが、64GBでも困ったことがほとんどない。

今回、中国のAcclamatorというブランドのメモリが、他のブランドの同容量のメモリの60%ぐらいの価格で売っていた(今はそこそこ高くなってしまったようだが)ので、DDR5 5600MHz 32GBx2 を購入。5600MHzでmemtest86やstresstest-cliで12時間ほど負荷をかけてみたが、とりわけエラーなく使えている。長期の耐久性や暑さ本番の夏(現在は冬で寒い)はわからない。GPU学習時にRAM速度の影響はほとんどないので、さらに安定すべく4800Mhzにクロックを落とし使っている。推論時にCPUオフロードする場合など、RAM速度も重要になるケースもあるのだけど、自分はやる予定がないので。

追記: 結局メモリ不足を感じ、もう32GBx2を増やして合計128GBへ。

ストレージ・NVMe

学習に使うデータは、雑にやるとデータがシャッフルされしてランダムアクセスが発生するため(たとえばHuggingFace Transformers も標準では学習時必ず shuffle する)。そのため容量が巨大なNVMe(SSD)を。容量はあるだけ良い。

• Sandisk SN850X NVMe SSD WDS800T2X0E 8TB
  • PCIe のCPU直結レーン。8TB でも足りない（データを削除しながら使っている）ので本当はもっと容量が欲しい。
• Samsung 980 Pro 2TB
  • これは余っていたので追加
  • チップセット共有レーン

ストレージ・HDD

生ダウンロードデータの一時的なファイル置き場として 14TB のHDDを使っている。ランダムアクセスが発生する用途には遅すぎて使えないのだけど、たとえば実運用面では、 HuggingFace datasets ライブラリはまず環境変数 HF_HUB_CACHE にデータをダウンロードするのだが、実際にライブラリがロードする際は、parquet ファイルから arrow 形式に変換されるため、後者の方が NVMe でアクセスできれば良いので、HF_HUB_CACHE のディレクトリだけHDDに指定することで、切り分けて使えている。

• TOSHIBA MG07ACA14TE 14TB

CPUクーラー

120x3 のラジエータの簡易水冷モデルなら特にこだわりがなかったので、CORSAIR NAUTILUS 360 RS LCDを利用。CPUクーラーの表面に液晶があるモデルで、液晶にCPU温度表示できるのいいじゃん、と購入してからこの制御はUSB経由で行われるため、Linuxからだと制御が難しい(OSSでできるが、サクッと温度表示などは難しそう)ことの気づく。ので、今ならLCD無しモデルを買ったかな…。

組み立て

ケースや電源や空冷が重い(筋肉痛)、自分のミスでしょっちゅうファンの前後ろを間違える、ラジエータの上下を間違える、など以外は特に困らず組み立てられ、一発で起動し問題無く動いている。

OS

使い慣れている Ubuntu Server 24 LTS で。ssh で繋ぐだけなので、GUIは一切使っていない。

RTX5090 x2 PCを作ってみての感想

実際に作ってから1ヶ月ほど経ったが、冬場だからかもしれないが、GPU二枚をフルで使っても特に問題なく安定して動いている。良かったところとしては、PCIe がボトルネックにならない場合、例えばMLMでのbi-encoderモデルの学習速度が RTX5090x1 に比べ x1.8程度の速度になり、だいぶ速くなった。また、推論も水平処理ができる場合、例えば vLLM に Qwen3-8B モデルを載せて1000万件処理する、みたいな処理もほぼ2倍速で処理できるので便利だ。

CUDAは環境変数 CUDA_VISIBLE_DEVICES でプログラムから見えるGPUを簡単に切り替えできるので、たとえばGPU2を使いたい場合は CUDA_VISIBLE_DEVICES=1 とするだけで、プログラムには一切手を入れずに1枚のGPUとして認識させることできるのも便利。簡単にGPUを切り替えつつ使うことができている。

また、マルチGPU周りの知見を得るという点でも、今までずっと1GPUしか使ってこなかったため、複数GPU環境で学習・推論の方法や考え方を知ることができて勉強になり、こちらの点でも良かった。

ただ、PCIe 5.0 x8 の速度がボトルネックに感じることが割とあって、例えば PyTorch DDP では学習stepごとに GPU間のデータを同期する All-Reduce が発生するが、学習方法によっては非常に時間がかかる。巨大バッチでのコントラスト学習とかね。すると GPU SMの idle 時間がグッと増えて、速度向上がせいぜい1.2倍程度、場合によっては1 GPUの方が速いみたいなケースすらある。

B200、H200 などのデータセンター向けGPUではNVLinkを使うと、構成によっては GPU間で 数百GB/s〜TB/sの速度が出るため、高速な転送速度が確保できるが、PCIe 5.0 x8 では実効速度は 約20-30GB/s ほどのため、NVLink に比べると圧倒的に遅かったりする。高いGPUはよくできているなぁ(B200 x8 なマシン1つで8000万円程度かな…)。

というわけで、作ってみての満足度は高い。パーツを購入したタイミングも、メモリは高くなっていたが、2026年1月中旬はさらにストレージもメモリもGPU(RTX5090)も高くなってしまったので、まだマシだったタイミングであった。AI需要・そして円安もあり色々と高くなってしまったなぁ…。

子供

子供がうまれた。人生で一番大きな変化は子供が生まれた時、とはよく聞くが、当事者になってみると全くその通りだと思う。何もかも子供主体な考え方に変わる。それにしても子供はかわいい。かわいすぎる。子供を育てられる幸せを噛み締める。リモートワークで、しょっちゅう子の顔を見れる、ありがたさもある。

産まれる前は妻が持病関連で体を悪くし、長期入院になったり、出産後もどうなるかわからずだいぶ不安だったが、その後の回復は良好で、問題無く日常生活を過ごせている。健康のありがたさよ。私は(四|五)十肩、以外はだいたい元気。

家

昨年建てて、昨年末から新しく住み出した家、すこぶる快適。冬はそれなりに寒いし、周りが農地なので夏は虫が多いが、それらをさっ引いても大変暮らしやすい。子も増えたことだし、建てて良かった。

車

テスラのモデルY ジュニパー ロングレンジをお迎え入れした。BEVのエンジンがない快適さ(静か、あっという間に速度が上がる)もさることながら、ソフトウェア・UX周辺、ほんとよくできている。ほとんどの車が、過去の車の延長線上なUXだが、新興の会社は過去を踏襲する必要がないので、新しい体験設計ができ、それをひしひしと感じる。車移動（というのは田舎において移動手段のほぼ全てある）が非常に楽に・快適になった。

現在のオートパイロット運転アシストでさえもかなり快適なのに、将来はFSD(完全自動運転、という名の、いい感じの自動運転支援)が日本でも使えるようになるだろうから、そちらもとても楽しみである。

技術

引き続き、情報検索周りを主に、技術的なあれこれやプロダクトづくりをやっていた。コーディングエージェントが台頭した年で、ほぼ全ての技術作業はエーアイさんにやってもらっているが、打ち手が増えて、できることが大きく広がったと感じる。通常のソフトウェア開発に限らず、例えば情報検索モデル開発関連も、エーアイが全部やってくれ、今までだったら手を動かす時間がかかりすぎていて大変だったことも、エーアイYOLO、ですんで最高だ。

エーアイが代替できない・代替しない技術をちゃんとやる、というのが正解プロダクト不定の時代の仕込みとしては大事だと思っているのだけど、その技術的なところを色々やれているので、楽しいね。バイアスがかかってるので、自分がそう思ってるだけで、簡単に代替される可能性はもちろんあるのだけど。

仕事周りでも今年はメインで作っていたエーアイプロダクトが世に出て、世の中的にも会社的にも一定の評価をいただけたようで、ありがたい限り。チーム開発のやりやすさの賜物でもあるので、関係者各位ありがとうございます。来年もまたチーム、エーアイさんの手を大いに借りて、技術的なことをやりつつも、新しいプロダクトを作っていきたい。

2025年はいろいろあったが、子が産まれたに尽きる。毎度のことだが、さまざまなことをサポートしてくれる妻に感謝だ。というわけで、皆様2026年もよろしくお願いいたします。

• 振り返り2024年
• 振り返り2023年
• 振り返り2022年
• 振り返り2021年
• 2020年の振り返り

しかしながら、大量に検索を行うと「検索結果」の情報も同時に増加していきます。そのため、本当に必要な情報の抽出をLLMが間違えたり、ハルシネーションが起きたり、入力情報の増加により処理が遅くなったり、LLM利用費用が増加したりと、大量の検索が難しかったりもします。

そこで、検索結果をLLMに渡す前に、関連しない情報は削除しちゃおう、ついでに関連度スコアもつけちゃおう、というアプローチが Provence です。このアプローチでは、検索でヒットした文章のうち、関連しない部分を削除することが可能です。実際にデルの性能を測定したところ、長文の質問・回答データセットを用いた評価(MLDR + LLM eval)では、80-95%ほど文章を削除できました。10000文字の文章なら500-2000字程度にLLMに渡す前に減らせる、ということですね。かなりの入力データの削減が期待できますね。短い文章に分割されたような短文が多いデータセットでも、ドメインによりますが30〜70%の文の削除が行われています。

ただ研究開発として公開されている Provence 実装やモデルは非商用で、日本語のデータセットも公開されていなかったので、今回 OpenProvence というプロジェクトを作成し、で学習推論などのソースコードやモデルの重みなどを "オープン" なライセンスで公開しました。日本語データセットも作成して公開しています(データセットは大元のライセンスがあるため、オープンなライセンスではないものが多いですが)。

OpenProvence の試し方

以下の URL に huggingface spaces (CPU) 環境で動くデモを用意したのでお試しください。デモのサンプルにあるWikipediaの情報検索のページ情報をもとに「ベクトル検索は？」をクエリに文削除を実行すると、約5000文字の記事が400文字に削減され、かつベクトル検索についての情報のみが残った形で出力さると思います。

• 🤗 https://huggingface.co/spaces/hotchpotch/open_provence_demo

またデモは以下の手順でローカルマシンでも手軽に動かせます。最近の MacBook でしたら、かなり高速に推論することも可能でしょう。

git clone https://huggingface.co/spaces/hotchpotch/open_provence_demo
cd open_provence_demo
uv sync
uv run python app.py

python からの利用方法

python からは以下の感じで利用できます。小型のxsmallモデルならCPU環境でも推論可能です。また GPU 環境(NVIDIA + flash attention2)では、即座にで推論が完了し、文章の削除が行われるでしょう。本番検索環境に組み込んでも、問題ない速度で処理できると思っています。

from transformers import AutoModel

# 利用モデルに合わせて変更
model_name = "hotchpotch/open-provence-reranker-xsmall-v1"
provence = AutoModel.from_pretrained(model_name, trust_remote_code=True)

question:str = "日本の首都について"
context:str = """
今日は学校に行き、さまざまなことを学んだり、友達と学食でたらふく食べた。
日本の首都は東京で、東京は日本の政治、経済、文化の中心地らしい。この都市は約1,400万人の人口を抱える世界有数の大都市らしい。
夜は飲み会に誘われたが、参加せずに帰宅した、今月そんなにお金が残ってないからなぁ、残念だ。
"""

result = provence.process(question, context, threshold=0.1)
print(f"Reranking Score: {result['reranking_score']:.4f}")
print(f"Compression Rate: {result['compression_rate']:.1f}%")
print(f"Pruned Context:\n{result['pruned_context']}")

# 出力例:
# Reranking Score: 0.7043
# Compression Rate: 62.5%
# Pruned Context:
# 日本の首都は東京で、東京は日本の政治、経済、文化の中心地らしい。
# この都市は約1,400万人の人口を抱える世界有数の大都市らしい。

コーディングエージェントの活用

OpenProvence は、推論・学習モデル実装、評価実装、データセット作成実装など、私は一行もコードを書かない縛りで、全ての実装はコーディングエージェント(Claude Code, Codex) によって行いました。かなり修正の指示は必要でしたが、隙間時間に進めたプロジェクトとしてはなかなかのものが、コーディングエージェントを活用することで出来上がったのかな、と思っています。出来上がった成果物のコードを見ると、もっとシンプルなコードにすることはできそうですが、現状のLLMが理解しやすい形・LLMが修正しやすい形だと、これぐらいの冗長なコードが丁度良いのかもしれません。

適切な指示と開発指針、AIが自身で開発、改善し続けられる環境等を作り続けることで、プロダクションレベルの品質のソフトウェアをAIと協調しながら作成する、Vibe engineering という言葉も生まれました。

開発指針やユニットテスト、CI、コードレビュー環境といった通常のコーディングエージェントを用いたソフトウェア開発に加え、モデル学習時の学習を短時間で行える最小ベースラインと評価データの用意(これが意図せず変化するとバグ)、データセットについての詳細な説明等々を用意することで、ある程度の規模の機械学習モデル・プロジェクトも開発できることを実感しています。

おわりに

OpenProvence のような、質問と関連しない文章を削除するアプローチは、とりわけ巨大な文章を処理するようなプロダクトととても相性が良いでしょう。

2024年はRAGが話題でしたが、2025年のAI Agent、DeepResearch、Context Engineering のような流行を先取りし、技術的に重要なポイントを研究開発した Naver Labs Europe の Provence チームの先見の明(Provenceは2025年1月公開!)に驚きと感謝を。

昨今、LLMを活用するプロダクトでは、裏側で情報検索を活用することで価値を高められ、情報検索技術は引き続きとても面白いです。このプロジェクトが少しでもプロダクトや研究で活用していただけたら幸いです。

• https://github.com/hotchpotch/open_provence/
• https://huggingface.co/collections/hotchpotch/openprovence

結論から言うと日本語においては、EmbeddingGemma はとても性能が低かったです。

JMTEB v1 ベンチマーク評価

注：MTEB v1の16の日本語タスクのマイクロ平均（単純平均）で算出

なお、JMTEBの設定はこちらで、各種 prefix などはつけているはずです。また結果の json である summary.json はこちら(gist)。再現方法も gist に記載してます。私の測定結果がおかしかったら教えてください。

2025/10/03追記: transformersのバグの影響で、最新版だと ruri-base 並に性能が向上するとのこと。LM8(@ShengzheLi) さん、情報ありがとうございます!

JQaRA / JaCWIR

JMTEB v1 でのスコアが低すぎるので、別途 JQaRA / JaCWIR でも評価してみましたが、やはりかなり低い結果となりました。

MTEB 高性能 ≠ 日本語高性能

先日評価した Qwen3 Embedding (Qwen3 Embedding 文章ベクトルの日本語性能を JMTEB で測る)もですが、最近の MTEB 高評価マルチリンガル embedding モデルは、日本語性能が低いことが多いですね。なんでだろうと、MTEB Leaderboard の Language-specific の Japanese を見ると、Qwen3 Embeddings, Embedding Gemma 共に Pair Classification しか日本語結果では載っていないので、ほぼ参考になりません。マルチリンガル性能とは…。

またこの Qwen3 Embedding, Embedding Gemma の2モデルは decoder モデルベースのアーキテクチャです。embeddinggemma-300m の中身を見ると、埋め込み用の HEAD (pooling + 2層dense) を mean pooling して使っていますね。

decoder + 小さなパラメータサイズの場合、少なくとも日本語における性能は他の encoder マルチリンガルモデルよりだいぶ低い結果でした。この辺は、そもそも日本語を embeddings タスクでほとんど学習させてないからなのか、それとも大元の小さな decoder の時点で日本語汎化性能が低いのか…。

今回公開した JFWIR (Japanese FineWeb Information Retrieval) は、この課題に取り組むために作成した約6,400万件の大規模な日本語情報検索に活用できるデータセットです。このデータセットは、高品質な教育的コンテンツを含むWebクロールデータ「fineweb-2-edu-japanese」を基に構築されています。

• https://huggingface.co/datasets/hotchpotch/JFWIR

JFWIRの特徴

1. 大規模かつ多様性の高いデータセット

JFWIRは以下の特徴を持つデータセットです：

• 6,400万件以上の文書-クエリペア: 各文書に対して7種類の異なるタイプのクエリ（keywords, synonym_keywords, query, alt_query, title, faq, summary）を生成
• 実際のWeb文章: Wikipedia以外の教育的価値の高いWebコンテンツを収録
• ハードネガティブ付き: 効果的な学習のための類似しているネガティブ文書

2. ベンチマーク評価結果

JFWIRを使用して学習させたリランキングモデルの性能を、主要な日本語情報検索ベンチマークで評価しました。以下の4つのベンチマークで比較を行いました：

特に、Web文章を対象とするJaCWIRでは0.9339から0.9586への改善が見られました。

使い方

JFWIRはHugging Face Datasetsから簡単に利用できます。以下に基本的な使用例を示します：

from datasets import load_dataset

# メインデータセットの読み込み
train_ds = load_dataset("hotchpotch/JFWIR", split="train", name="small_tokens_cleaned")

# サンプルデータの確認
for i in range(3):
    sample = train_ds[i]
    print(f"Query: {sample['query']}")
    print(f"Document: {sample['text'][:100]}...")

# ハードネガティブ付きデータセットの読み込み
hard_negatives_ds = load_dataset("hotchpotch/JFWIR", split="train", name="hard_negatives")

# ハードネガティブの使用例
for i in range(3):
    hn_sample = hard_negatives_ds[i]
    pos_id = hn_sample['pos_id']
    pos_doc = train_ds[pos_id]
    
    print(f"Query: {pos_doc['query']}")
    print(f"Positive (score: {hn_sample['pos_score']:.3f}): {pos_doc['text'][:100]}...")
    
    # ネガティブ文書をスコア順にソート
    neg_pairs = list(zip(hn_sample['neg_ids'], hn_sample['neg_scores']))
    neg_pairs.sort(key=lambda x: x[1])
    
    print("Negatives (lowest scores):")
    for neg_id, score in neg_pairs[:2]:
        print(f"  Score {score:.3f}: {train_ds[neg_id]['text'][:80]}...")

データセットの作成プロセス

1. 高品質な日本語Web文章の収集

まず、大規模なWebクロールデータセット「FineWeb-2」から、教育的価値の高い日本語コンテンツを抽出して「fineweb-2-edu-japanese」を作成しました。さらに、Web文章特有のノイズを除去し、適切な文章長に調整した「small_tokens_cleaned」サブセットを作成しました。

2. 多様なクエリの生成

6,400万件のデータセットに対してクエリを生成するため、軽量なクエリ生成モデル「query-crafter-japanese」を使用しました。多様性を確保するため、以下の3つのモデルを組み合わせて使用しています：

• hotchpotch/query-crafter-japanese-Qwen3-1.7B
• hotchpotch/query-crafter-japanese-Qwen3-4B
• hotchpotch/query-crafter-japanese-sarashina2.2-3b-instruct-v0.1

各文書に対して7種類のクエリタイプ（keywords, synonym_keywords, query, alt_query, title, faq, summary）を生成することで、多角的な検索ニーズに対応できるデータセットを構築しました。

3. ハードネガティブの作成

情報検索モデルの性能を向上させるため、ハードネガティブ（クエリに類似しているが正解ではない文書）を含むデータセットも作成しました：

1. 埋め込みモデルによる類似文書検索: ruri-v3-30mモデルを使用して6,400万件の文書をベクトル化し、各文書に対して類似度の高い文書を検索
2. 適切なネガティブの選定: 類似度top10-50とtop50-200からランダムサンプリング
3. リランカースコアの付与: japanese-reranker-xsmall-v2を使用してスコアリングしています。たとえば正例として不適切なもの（スコア<0.6など）や負例として不適切なもの（スコア>0.4など）を除外して利用することで、より適切な正例・負例を選択できます。

今後の展望

JFWIRは、日本語情報検索分野の発展に貢献することを目的として公開されました。しかしながら、query-crafter-japanese は文章からのある程度単純なクエリ生成にとどまり、もっと多様な価値のある質問文を作成することで、より様々な情報検索精度の向上が可能になると思っております。

まとめ

JFWIRは、Wikipediaに偏重していた従来の日本語IRデータセットとは異なるアプローチとして、実際のWeb文章を対象とした情報検索データセットです。約6,400万件のデータ、7種類のクエリタイプ、対照学習用のハードネガティブなど、情報検索システムの開発に活用いただける要素を含んでおります。

データセットはHugging Faceで公開されており、ODC-byライセンスの下で自由に利用できます。日本語情報検索分野の発展に、このデータセットが少しでも貢献できれば幸いです。

関連リンク

データセット

• JFWIRデータセット: hotchpotch/JFWIR (Hugging Face Datasets)
• fineweb-2-edu-japanese: https://huggingface.co/datasets/hotchpotch/fineweb-2-edu-japanese
• FineWeb-2: https://huggingface.co/datasets/HuggingFaceFW/fineweb-2

モデル

• ruri-v3-30m（埋め込みモデル）: https://huggingface.co/cl-nagoya/ruri-v3-30m
• ruri-v3-pt-70m（事前学習モデル）: https://huggingface.co/cl-nagoya/ruri-v3-pt-70m
• japanese-reranker-xsmall-v2: https://huggingface.co/hotchpotch/japanese-reranker-xsmall-v2
• query-crafter-japanese（クエリ生成モデル）:
  • hotchpotch/query-crafter-japanese-Qwen3-1.7B
  • hotchpotch/query-crafter-japanese-Qwen3-4B
  • hotchpotch/query-crafter-japanese-sarashina2.2-3b-instruct-v0.1

記事・論文

• query-crafter-japanese記事
• 日本語言語理解ベンチマークJGLUEの構築 〜 自然言語処理モデルの評価用データセットを公開しました
• MIRACL: A Multilingual Retrieval Dataset Covering 18 Diverse Languages
• FineWeb-Edu: The Finest Collection of Educational Content the Web Has to Offer
• Ruri: Japanese General Text Embeddings

作者: Yuichi Tateno (@hotchpotch)

ライセンス

本データセットは、元の FineWeb2 と同様に Open Data Commons Attribution License (ODC-By) v1.0 の下で公開します。また、使用にあたっては CommonCrawlの利用規約 も適用されます

Citation Information

JFWIRデータセットを研究や開発に使用される場合は、以下の引用情報をご利用ください

@misc{tateno2025jfwir,
  author = {Yuichi Tateno},
  title = {JFWIR: Japanese FineWeb Information Retrieval Dataset},
  year = {2025},
  url = {https://huggingface.co/datasets/hotchpotch/JFWIR},
  note = {A large-scale Japanese information retrieval dataset with 60+ million document-query pairs}
}

ただ、マルチリンガルモデルはあまり日本語が重視されない傾向にあるので、JMTEB: Japanese Massive Text Embedding Benchmarkで Qwen3-Embedding-0.6B の性能を計測してみました。なお、jsick, jsts がエラーになったため、STSタスクは除いてあります。

JMTEB 計測結果

結果はこちらです。日本語のタスクがあまり学習されていないからなのか、日本語の結果は振るわない結果でした。ruri-v3 シリーズはモデルサイズも小さく、かつ日本語では圧倒的に高性能ですね。

なお Retrieval, Reranking タスクでは、Query の prefix に Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: を追加しています。

また計測に使った、JMTEB用の設定(jsonnet)や結果の summary.json、実行コマンドは以下においてあります。Qwen3-Embedding-0.6B の性能が低すぎる気もするので、何か間違っていたら教えてください。

• https://gist.github.com/hotchpotch/f6be186010e70d6eb6e46447cea258f9

おまけ: Qwen3 Embedding 論文を読む

Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Modelsが公開されたのでざっくり読んでみました。とりわけ合成データセット作成あたりが個人的に面白かったです。

以下は私の興味範囲のメモ書きです。

• LLM2Vec のような decoder → encoder ではなく、casual attention をそのまま利用
• Embedding モデルは最終層の[EOS] トークンの hidden state から最終埋め込みを取得
  • Query は Instruction + Query で作成。Doc はそのまま。
  • InfoNCE を改良したスコア(単純な対照学習ではなく、ハードネガティブを複数含めたり、類似度のポジネガを調整した偽陰性の調整など)
• Reranking は chat template をそのまま使って、"yes", "no" トークンの確率で、関連性スコアとして計算
  • decoder model のラベル分類の解き方(該当ラベルトークンの確率を見る)のアプローチをママ適用
  • SFT で学習できる
• 1st stage で Qwen3-32B で作った合成データセットをもとに学習
  • 情報検索, 対訳マイニング(Bitext Mining), 意味的類似性, 分類 の4つのタイプを作成
  • 情報検索の合成データセットの場合、詳細な設定を作り、それをもとにQwen3の事前訓練コーパスの文章からクエリを生成
• 2nd stage で 700万の既存データセット(MS Marco, MIRACLなどなど)と、1st stage のコサイン類似度でフィルタリングした1200万件のデータをもとに学習
• 最後に多様性考慮のモデルマージ
  • 詳細は記されてないので推察だが、2nd stage の複数のチェックポイントは、タスク特化学習させたもの、特定言語にフォーカスして学習させたもの、などが考えられそう
  • モデルマージは適当にマージして、ベンチマーク走らせると結果が向上することを、少ないコンピューティングリソースで観測できるので、たくさんチェックポイントがあるならいろいろ試した方が良さそう。

🆕 2025-07-10 まぁまぁ小さなリランカー japanese-reranker-small-v2 と japanese-reranker-base-v2 も追加しました。

本モデルは最小限のレイヤー数とパラメータ数で作成されており、CPUやAppleシリコン環境でも実用的な速度で動作します。これにより、高価なGPUリソースなしでもRAGシステムの精度向上が可能になり、エッジデバイスでの展開や低レイテンシが要求される本番環境で活用できるでしょう。性能評価では、大型モデルと比較しても競争力のあるスコアを出しています。

• https://huggingface.co/hotchpotch/japanese-reranker-tiny-v2
• https://huggingface.co/hotchpotch/japanese-reranker-xsmall-v2
• https://huggingface.co/hotchpotch/japanese-reranker-small-v2
• https://huggingface.co/hotchpotch/japanese-reranker-base-v2

リランカーとは何か、そして小さなリランカーの重要性

リランカーとは、検索システムにおいて、質問（クエリ）と文書の関連性を評価し、最も関連性の高い順に並べ替える（ランキング）するモデルです。従来の文ベクトル（埋め込み）検索だけでは捉えきれない複雑な関連性を評価できる点が強みです。特にCrossEncoderと呼ばれるアーキテクチャを用いることで、質問と文書を一つのペアとして入力し、より細かなニュアンスや文脈的理解を実現します。

• 日本語最高性能のRerankerをリリース / そもそも Reranker とは?

小さなリランカーモデルが重要な理由はいくつかあります。まず、リランカーは質問と候補文書のすべての組み合わせを評価する必要があるため、計算量が非常に多くなります。例えば100件の候補文書をリランクする場合、100回のモデル推論が必要です。そのため、モデルが小さいほど処理速度が向上し、レイテンシが低減します。

また、小型モデルは限られたリソース環境での実行も可能です。CPUのみの環境やエッジデバイス、モバイルデバイスでも現実的な速度で動作でき、RAG（検索拡張生成）システムの実用性を大きく高めます。同時に、クラウド等のサーバ環境ではGPUメモリ使用量の削減により、GPUリソースの共有が可能となりコスト効率が大幅に向上します。

• Ask! NIKKEI RAG検索技術の深層

このように、小型リランカーは速度、コスト、リソース効率の面で大きなメリットをもたらし、実用的なRAGシステム構築において大切な役割を果たすでしょう。

ベンチマーク性能

ベンチマーク結果は以下です。小さなな tiny, xsmall v2 の性能はモデルサイズを考えるとかなり高く、大きいモデルとしては ruri-v3-reranker-310m が圧倒的ですね。これらの高性能なモデルは、ベースがどれも高性能な ModernBert になったことも、性能向上に寄与しているでしょう。

なお、日本語モデルはどれもJQaRA(クイズ形式)の傾向を学んでおり、bge-reranker-v2-m3 は不利になります。これはリランカーが適切にドメイン課題を学習すれば、だいぶスコアが上がることの例でもあります。

推論速度

こちらは、HuggingFace transformers ライブラリを使った、約15万ペアをリランキングした推論速度結果(トークナイズ時間は除いていて、純粋なモデルでの推論時間)です。MPS(Appleシリコン),CPUの計測にはM4 Maxを、GPUにはRTX5090を用い、かつ ModernBert 系列モデルでは GPU 処理時に flash-attention2 を使っています。

japanese-reranker-tiny-v2, xsmall-v2 は速度面で圧倒的ですね。ruri-v3-reranker-310m もモデルサイズを考えるとかなり速く、これらは flash-attention2 が効いているからでしょう。なお、text-embeddings-inference等を使うことで、他のモデルも flash-attention2 を使うことができ、その場合はこの評価以上の速度が出ると思います。

なお、推論速度のベンチマークに用いたスクリプトはこちらです。

また CPU用に onnx に変換したモデルも公開しているため、例えばラズパイ環境などで、onnx + arm向け量子化モデルを使うことで、実際にエッジ環境でも動くでしょう。

モデル作成の簡易テクニカルレポート

japanese-reranker-tiny-v2, xsmall-v2, small-v2, base-v2 の学習データ元として、hotchpotch/japanese-splade-v2 学習で用いたデータセット + ハードネガティブ + 若干の独自データを用いて学習させています。v1と比べて大幅に性能が上がったのは、ModernBert ベースで事前対象学習を行った ruri-v3-pt-30mを用いてることと、v1よりも数倍のデータセットを用いたこと、またハードネガティブでの良質なデータの抽出(各種rerankerのスコアを用い、正しい・正しくないでフィルタリング)を行ったことも大きいでしょう。

また、Tiny モデルのモデルパラメータ抽出元として、sbintuitions/modernbert-ja-30mとcl-nagoya/ruri-v3-pt-30mを利用・評価しました。ModernBert アーキテクチャは、グローバルアテンションとローカルアテンションのレイヤーを交互に含みます。例えば modernbert-ja-30m モデルは10層のレイヤーで、[0,3,6,9]層がグローバルアテンションで、それ以外がローカルアテンションとなっています。

最初は全てグローバルアテンションの方が良いだろうと思ったのですが、3,6,9層を含むと基本悪くなり、また出力層に近い層を含むと、こちらも結果が悪くなりました。以下のグラフは同じデータセットで学習したrerankerのrerank評価結果です。出力層に近い6,9などを含むとだいぶ悪くなり学習早期で止めたので、以下の結果には含めてません。また、layer 0 のみは流石に全く性能が出ませんでした。

この中から、最も良質な結果になった ruri-v3-pt-30m を xsmall として、tiny モデルとしては ruri-v3-pt-30m + layer 0,2,4 を公開しました。また、small-v2 と base-v2 は ruri-v3-pt-70m と ruri-v3-pt-130m をベースにそれぞれ作成されています。なお、モデルマージすると性能は少々上がりますが、今回は行っていません。

おわりに

本エントリーでは、非常に小型軽量で実用的な日本語リランカーモデルjapanese-reranker-tiny-v2、japanese-reranker-xsmall-v2、japanese-reranker-small-v2、japanese-reranker-base-v2 についての紹介をしました。これらのモデルのうちtinyやxsmallは、CPUやAppleシリコンといった環境でも実用的な速度で動作し、高価なGPUリソースを必要とせずにとも、とりわけローカルなRAGシステムなどの検索精度の向上に寄与します。またGPU上で動かすことで、高速なレスポンスも実現可能です。

近年の高性能な ModernBert 等の Encoder モデルの登場により、より高性能な実用的な性能を持つモデルの開発を後押ししています。本記事が、日本語処理技術のさらなる発展に貢献できれば幸いです。

しかし、合成データセットを作成して広く公開したい場合、OpenAIやGeminiなどの商用LLMでは、利用規約によってライセンスの問題が発生します。また、大量の文章を処理したい場合は時間・費用もかなりかかります。

そのため、1.7B〜4B という小型サイズのモデルで高速に動作しながらも、DeepSeek-R1で生成した質問文と同レベルの情報検索用の質問文（クエリ文）を自動作成でき、さらに出力ライセンスに制限がないquery-crafter-japanese モデルを作成しApache 2.0ライセンスで公開しました。

• query-crafter-japanese-Qwen3-1.7B
  • ⭐️:👆速度・性能の面でおすすめです
• query-crafter-japanese-Qwen3-4B
• query-crafter-japanese-sarashina2.2-3b-instruct-v0.1

query-crafter は7つのカテゴリーを生成できます。

• keywords: スペース区切りのキーワード
• synonym_keywords: 類義語をもちいた特徴的なキーワード
• query: 文章の内容に基づいた質問文
• alt_query: BM25でマッチしない表現を使った質問文
• title: 文章全体を表現するタイトル
• faq: 文章をFAQの回答とした場合の質問文
• summary: 文章の短い要約

では、以下の文章を用いて、各々のカテゴリーに対する質問文を作成してみましょう。

• query-crafter-japanese-example.py

夕方、開発合宿の成果発表会。私以外は、AI関連のちゃんとしたテーマに取り組んで、クオリティも高く、いやー面白い。I氏はエンジニアでもないのに、Figmaプラグインを作ったり、vercelにデプロイしてたり(ほぼcursorが書いた)して、AIによって大きく幅が広がる一例を間近に見る。私は何かのテーマに取り組んだわけではなく、Vibe Cording を一度もしたことがなかったので、cursor でコードをいかに触らず・見ずに作れるかを試した。

毎年のこの日記を要約してdiscordなどに投稿するツール（以前も作ったものの仕様を書いて新機能などを追加）を作成したり、この日記のタイトルがないものに自動でタイトルをつけたりするツールを作成する。Vibe Cording は思った通りの感じで、なるほど便利。

コードは見ずにブラックボックス的な開発（出力成果物だけをみる）をしたので、出来上がったコードを後で見ると本番運用前提のコードでは全くないが、書き殴りのツールを作るには十分。また自分が指示するのは仕様のみで、仕様書も随時アップデートされるようにしてるので、機能を変えたくなったら仕様変更・追加するだけでいいし、楽で良いね。

query-crafter-japanese-Qwen3-1.7B を用いてカテゴリーごとに質問文を生成した結果はこちらです。keywords, query, title, summary あたりは特色が分かりやすく出ていますが、synonym_keywords は完璧な類義語でないことも多かったり、alt_query, faq は query とそれほど変わらなかったりすることもあります。

keywords: Vibe Cording ブラックボックス開発 仕様変更
synonym_keywords: AI活用開発プロジェクト 発表会 仕様変更追加
query: 開発合宿で作成したツールの具体的な機能は？
alt_query: 開発者向けツール開発でコード見ない開発手法の利点は？
title: AI活用で拓く開発の新領域：Vibe Cordingとブラックボックス開発の可能性
faq: 開発合宿で実現した新機能や成果は？
summary: AI活用の開発成果発表会で、Vibe Cordingや日記ツール開発、コード見ずに開発を実施

また動作速度も vllm + RTX5090 環境で、入力トークンが 48,000 toks/s、出力トークンが 2200 toks/s で動作します。〜1000文字程度の文章1万件から質問文1万件を生成した場合、100秒弱で作成できます。対象文章が1億件あったとしても、約140時間程度で全てを処理することができます。

なお、DeepSeek-R1 を夜間ディスカウント時間帯(input: 1M toks 0.135USD, output: 1M toks 0.55USD)で実際に10万件の文章を並列100のAPIリクエストで処理した場合、約7時間と40USD程度の費用がかかりました。もし、DeepSeek-R1 APIで1億件を処理した場合、約7,000時間(実際には夜間ディスカウント時間を狙うと、そのタイミングでしか処理できないので、もっと時間がかかります。また並列リクエストの最大数はDeepSeekサイドのリソースによって変動します)と、40,000USDほどの費用が発生するでしょう。

このように、query-crafter は、特に大量の文章から質問文を作成したい場合、処理速度的にも費用的にも大きなメリットがあります。

query-crafter-japanese モデルの学習

学習には、出力結果利用に制限がない DeepSeek-R1 を使い fineweb-2-edu-japaneseのデータをもとに、質問文となる教師データを、合成データセットとして作成しました。

例えば title については '文章全体をうまく表現した、タイトルを考え作成しなさい。考えたタイトルは30文字以内で出力すること。出力は厳密な JSON 形式で {"query": "タイトル"} とする。他に一切余計な出力はしないこと。' といった指示文を用いて作成しています。

• https://huggingface.co/datasets/hotchpotch/japanese-query-crafter-reasoning-80k

続いてこのデータを教師データとし、SFT(Supervised Fine-tuning)で、Qwen3-4B, Qwen3-1.7B, sarashina2.2-3b-instruct-v0.1, TinySwallow-1.5B-Instruct を学習させました。

SFT時に使ったフォーマットはシンプルに

{
  "system": "{category名}",
  "user": "{text}",
  "assistant": "{query}",
}

といった内容です。systemプロンプトに title などの指示カテゴリを、user 入力文に文章テキストを、そしてmodelの出力に query を設定しています。何かの用途に特化したSFTの場合、冗長なプロンプトを書く必要はなく、短い指示(今回は各種カテゴリー)のみで、うまく学習できます。

query-crafter-japanese モデルの評価

query-crafter の評価は、japanese-query-crafter-reasoning-80k の testデータを用いました。このデータのtextを元に、各種SFTで学習させたquery-crafterモデルを使って質問文を作成します。

そしてこれらの質問文とテキストをペアに、リランカーBAAI/bge-reranker-v2-m3で評価させたスコアの結果が以下です。このリランカーは、文章とテキストの関連性が高いと1.0になり、関連性がないと0.0となります。そのため、質問文とテキストが関連しているかどうかの目安になります。

パーセンタイルをプロットしたグラフは以下です。

結果、TinySwallow-1.5B 以外は、ほとんどのケースでDeepSeek-R1以上のスコアとなりました。特に、Qwen3-1.7B は日本語に特化しているわけではないマルチリンガルモデルですが、SFTするとQwen3-4Bとほとんどスコアが変わらず、性能の高さは驚くべきものです。そのため、特にこだわりがなければ、query-crafter-japanese-Qwen3-1.7B を利用するとよいでしょう。

なお、DeepSeek-R1 他よりスコアが低いからといって必ずしもDeepSeek-R1の質問文の質が悪いというわけではなく、リランカーでも判別が難しいような「正しく難しい質問文」を作成しているケースもあります。TinySwallow-1.5B はちょこちょこ全く関連がない質問文を作成してしまうケースがあり、他のモデルよりスコアが低くなりました。TinySwallow-1.5B-Instruct は TAID でモデル蒸留されているため、その後の SFT には不向きなのかもしれません。

おわりに

大量の質問文章を作りたい場合において、速度的にも費用的にも大きなメリットがある、query-crafter-japanese モデルを作成し公開しました。高性能かつ出力結果に制限がない DeepSeek-R1 の登場以降、様々な方法でデータセットの作成・公開・それを教師データとして利用したモデルの作成がしやすくなりました。また、Qwen などの小型サイズのモデルといった、ライセンスが使いやすいオープンウェイトなLLMの登場・性能進化により、ファインチューンした用途特化の小型モデルも作成・公開しやすくなり、幅広い応用が可能になってきたことを実感しています。もし半年前なら、このモデルを個人で作成することはリソース的にも不可能だったでしょう。

このモデルが、質問文を作りたい方の助けになれば幸いです。

• https://huggingface.co/datasets/hotchpotch/fineweb-2-edu-japanese

以下の内容は、上記ページの日本語訳です。

本データセットは、FineWeb2 の日本語データ（376M件）のうち、教育向けコンテンツと判断した120M件（約89.3Bトークン）の文章をフィルタしたものです。以下のサブセットも提供しています。

• default: 約120M件（1.2億件）のデータ・約89.3Bトークン
• sample_10BT: default からランダムサンプリングした約10Bトークンのデータ
• small_tokens: トークン数が512以下の短い文章のみから構成されるデータ
• small_tokens_cleaned: small_tokens から Web 特有のテキストノイズを除去したデータ

データセット作成の背景

FineWeb（英語のみ）は、Webデータの重複除去と高品質テキスト抽出を目的として作成されました。さらに、教育向けに質の高いテキストを抽出した FineWeb-Edu により、より少ないトークン数でも効率的な学習が実現可能となっています。

2024年12月に公開された FineWeb2 は多言語対応（日本語を含む）の高品質データセットですが、2025年2月現在、教育向けに価値が高い「Edu」データセットは未公開です。そこで、本プロジェクトでは FineWeb2 Edu Japanese データセット を作成し、公開しました。

教育的データのフィルタリング

本データセットの構築には、FineWeb2 日本語データから、教育向け文章を判定するためのモデル fineweb-2-edu-japanese-classifier を利用してフィルタリングしました。判定モデルのスコアリングの教師データには、DeepSeek-API (deepseek-chat) によって評価された fineweb-2-edu-japanese-scores を使っています。なお、本データセットでは、スコアが2.5以上の文章のみを抽出しており、そのスコアは score カラムに記載しています。

トークンカウントの付与

ModernBERT-Ja-130M のトークナイザを用いてカウントしたトークン数が token_count カラムとして付与されています。

Web特有のノイズ除去

FineWeb2 の日本語データには、Web特有のボイラープレートや不要なノイズが含まれることがあります。例えば、以下のような文章が含まれます。

この文章は90日以上更新の無いサイトに表示されています。
ログイン ログアウト

本当に必要な文章以外にも、さまざまなノイズが含まれていることがあります。例えば、この文章もその一例です。本来不要なテキストが入ってしまうことがこのようにあるでしょう。

今なら50%オフ！クリックしてリンク先の商品を表示

とりわけ文章長が短い場合、文章のほとんどがノイズを含む可能性があります。それらを取り除くことで、より高品質の文章を抽出できないかと考えています。

前のページ  次のページ

このような不要なテキストを取り除くためのモデル、fineweb-2-japanese-text-cleaner を開発しました。ノイズ判定の教師データとしては、fineweb-2-japanese-noise-spans を利用しています。この教師データはcyberagent/DeepSeek-R1-Distill-Qwen-32B-Japanese を活用して作られました。

このモデルにより、以下のようにノイズ箇所が検出されます。

[NOISE]この文章は90日以上更新の無いサイトに表示されています。[/NOISE]
[NOISE]ログイン[/NOISE] [NOISE]ログアウト[/NOISE]

本当に必要な文章以外にも、さまざまなノイズが含まれていることがあります。例えば、この文章もその一例です。本来不要なテキストが入ってしまうことがこのようにあるでしょう。
[NOISE]
今なら50%オフ！クリックしてリンク先の商品を表示[/NOISE]

とりわけ文章長が短い場合、文章のほとんどがノイズを含む可能性があります。それらを取り除くことで、より高品質の文章を抽出できないかと考えています。

[NOISE]前のページ[/NOISE]  [NOISE]次のページ[/NOISE]

本データセットに含まれるsmall_tokens_cleaned サブセットは、small_tokens からさらに fineweb-2-japanese-text-cleaner モデルを適用し、ノイズを除去したデータとなります。なお、モデルを使ってノイズ検出をした生データは fineweb-2-edu-japanese-noise-detect-raw で公開しています。

なおノイズ検出は完璧ではないため、場合によっては正しい文章の一部が誤って除外されている可能性がありますのでご注意ください。

注意事項

本データセット「FineWeb2 Edu Japanese」と、Eduフィルタリングを実施していない大元の「FineWeb2」データセットとの比較実験は行っておりません。そのため、実際のLLM学習においてどの程度の効果差が生じるかは未検証です。

また、教育向けテキストかどうかの分類精度も完璧ではなく、一部教育向けではないテキストも含まれます。

ライセンス

本データセットは、元の FineWeb2 と同様に Open Data Commons Attribution License (ODC-By) v1.0 の下で公開します。また、使用にあたっては CommonCrawlの利用規約 も適用されます。

Citation Information

@software{yuichi2025fineweb-2-edu-japanese,
  author = {Yuichi Tateno},
  title = {FineWeb2 Edu Japanese},
  month = feb,
  year = 2025,
  url = {https://huggingface.co/datasets/hotchpotch/fineweb-2-edu-japanese/}
}

この課題を解決する新しいアプローチとして、先日公開されたTransformerモデル「ではない」 StaticEmbeddingモデルは、例えば intfloat/multilingual-e5-small (以下mE5-small)とのベンチマーク比較では85%のスコアという最低十分な性能で、何よりCPUで動作時に126倍高速に文ベクトルを作成することができる、という驚きの速度です。

というわけで、早速日本語(と英語)で学習させたモデル sentence-embedding-japanese を作成し、公開しました。

• https://huggingface.co/hotchpotch/static-embedding-japanese

日本語の文章ベクトルの性能を評価する JMTEB の結果は以下です。総合スコアでは mE5-small には若干及ばないまでも、タスクによっては勝っていたりしますし、他の日本語baseサイズbertモデルよりもスコアが高いこともあるぐらい、最低限実用できそうな性能が出ていますね。本当にそんなに性能が出るのか実際に学習させてみるまでは半信半疑でしたが、驚きです。

なお、StaticEmbedding 日本語モデル学習などの技術的なことは記事の後半に書いているので、興味がある方はどうぞ。

利用方法

利用は簡単、SentenceTransformer を使っていつもの方法で文章ベクトルを作れます。今回はGPUを使わず、CPUで実行してみましょう。なお SentenceTransformer は 3.3.1 で試しています。

pip install "sentence-transformers>=3.3.1"

from sentence_transformers import SentenceTransformer

model_name = "hotchpotch/static-embedding-japanese"
model = SentenceTransformer(model_name, device="cpu")

query = "美味しいラーメン屋に行きたい"
docs = [
    "素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。",
    "新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。",
    "あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。",
    "おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。",
]

embeddings = model.encode([query] + docs)
print(embeddings.shape)
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f"{similarity:.04f}: {docs[i]}")

(5, 1024)
0.1040: 素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。
0.2521: 新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。
0.4835: あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。
0.3199: おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。

このように、queryにマッチする文章のスコアが高くなるように計算できてますね。この例文では、例えばBM25ではqueryに含まれる「ラーメン」のような直接的な単語が文章に出ていないため、うまくマッチさせることが難しいでしょう。

続いて、類似文章タスクの例です。

sentences = [
    "明日の午後から雨が降るみたいです。",
    "来週の日曜日は天気が良いそうだ。",
    "あしたの昼過ぎから傘が必要になりそう。",
    "週末は晴れるという予報が出ています。",
]

embeddings = model.encode(sentences)
similarities = model.similarity(embeddings, embeddings)

print(similarities)

# 一つ目の文章と、その他の文章の類似度を表示
for i, similarity in enumerate(similarities[0].tolist()):
    print(f"{similarity:.04f}: {sentences[i]}")

tensor([[1.0000, 0.2814, 0.3620, 0.2818],
        [0.2814, 1.0000, 0.2007, 0.5372],
        [0.3620, 0.2007, 1.0000, 0.1299],
        [0.2818, 0.5372, 0.1299, 1.0000]])
1.0000: 明日の午後から雨が降るみたいです。
0.2814: 来週の日曜日は天気が良いそうだ。
0.3620: あしたの昼過ぎから傘が必要になりそう。
0.2818: 週末は晴れるという予報が出ています。

こちらも、類似文章が高スコアになる結果になりました。

またTransformerモデルを利用してCPUで文章ベクトルを作った場合、少ない文章量でもだいぶ時間がかか、という経験をされた方も多いと思います。StaticEmbedding モデルではCPUがそこそこ速ければ一瞬で終わるはず。さすが100倍速。

出力次元を小さくする

標準で作られる文ベクトルの次元は1024ですが、これをさらに小さく次元削減することもできます。例えば 128 を指定してみましょう。

# truncate_dim は 32, 64, 128, 256, 512, 1024 から指定
model = SentenceTransformer(model_name, device="cpu", truncate_dim=128)

query = "美味しいラーメン屋に行きたい"
docs = [
    "素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。",
    "新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。",
    "あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。",
    "おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。",
]

embeddings = model.encode([query] + docs)
print(embeddings.shape)
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f"{similarity:.04f}: {docs[i]}")

(5, 128)
0.1464: 素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。
0.3094: 新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。
0.5923: あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。
0.3405: おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。

128次元のベクトルになり、結果のスコアも若干変わりましたね。次元が小さくなったことで、性能が少々劣化しています(後半にベンチマークを記載)。ただ1024次元から128次元に減ることで、保存するストレージサイズが減ったり、検索時などに利用する類似度計算コストが約8倍速になったりとなったりと、用途によっては小さい次元の方が嬉しいことも多いでしょう。

なぜCPUで推論が高速なの？

StaticEmbedding はTransformerモデルではありません。つまりTrasformerの特徴である "Attention Is All You Need" なアテンションの計算が一切ないのです。文章に出てくる単語トークンを1024次元のテーブルに保存して、文ベクトル作成時にはそれの平均をとっているだけです。なお、アテンションがないので、文脈の理解などはしていません。

また内部実装では PyTorch の nn.EmbeddingBag を使って、全てを連結したトークンとオフセットを渡して処理することで、PyTorch の最適化で高速なCPU並列処理とメモリアクセスがされているようです。

元記事の速度評価結果によるとCPUではmE5-smallと比べて126倍速らしいですね。

評価結果

JMTEBでの全ての評価結果はこちらJSONファイルに記載しています。JMTEB Leaderboardで他のモデルと見比べると、相対的な差がわかるでしょう。JMTEBの全体の評価結果はモデルサイズを考えると、すこぶる良好です。なお、JMTEB のmr-tidy タスクは700万文章のベクトル化を行うので処理に時間がかなりかかる(モデルにもよりますがRTX4090で1~4時間ほど)と思います。これもStaticEmbeddingsでは非常に速く、RTX4090では約4分で処理終えることができました。

情報検索でBM25の置き換えができそうか?

JMTEBの中の情報検索タスクのRetrievalの結果を見てみましょう。StaticEmbedding では mr-tidy の項目が著しく悪いですね。mr-tidyは他のタスクに比べて文章量が圧倒的に多く(700万文章)、つまる所大量の文章を検索するようなタスクでは結果が悪い可能性がありそうです。文脈を無視したた単純なトークンの平均なので、増えれば増えるほど似た平均の文章が出てくるとすると、そういう結果にもなり得そうですね。

ので、大量の文章の場合、BM25よりもだいぶ性能が悪い可能性がありそうです。ただ、少ない文章で、ずばりの単語マッチが少ない場合は、BM25よりも良好な結果になることが多そうですね。

なお情報検索タスクの jaqket の結果が他のモデルに対してやたら良いのは、jaqket の問題を含む JQaRa (dev, unused)を学習しているからといっても、高すぎる感じで謎です。test の情報リークはしていないとは思うのですが…。

クラスタリング結果が悪い

こちらも詳細は追っかけていませんが、スコア的には他のモデルよりもだいぶ悪い結果ですね。クラス分類タスクは悪くないので不思議です。埋め込み空間がマトリョーシカ表現学習で作られた影響もあるのでしょうか。

JQaRA, JaCWIR でのリランキングタスク評価

JQaRA の結果はこちら。

JaCWIR の結果はこちら。

JQaRa 評価は BM25 よりは若干良く、mE5-small よりは若干低い、JaCWIR は BM25, mE5よりだいぶ低い感じの結果になりました。

JaCWIR はqueryから探しあてる文章が、Web文章のタイトルと概要文なので、いわゆる「綺麗な」文章ではないケースも多いです。transformerモデルはノイズに強いので、単純なトークン平均のStaticEmbeddingではスコアに差がつけられるのも納得ですね。BM25は特徴的な単語が出現した文章にマッチするので、JaCWIR でもノイズとなるような文章上の単語はクエリにそもそもマッチしないため、Transformer モデルと競争力のある結構良い結果を残しています。

この結果から、StaticEmbedding は Transformer / BM25 に比べ、ノイズを多く含む文章の場合はスコアが悪い可能性があります。

出力次元の削減

StaticEmbedding で出力される次元は、学習次第ですが今回作成したものは1024次元とそこそこのサイズです。次元数が大きいと、推論後のタスク(クラスタリングや情報検索など)に計算コストがかかってしまいます。しかしながら、学習時にマトリョーシカ表現学習(Matryoshka Representation Learning(MRL))をしているため、1024次元をさらに小さな次元へと簡単に次元削減ができます。

MRLは、学習時に先頭のベクトルほど重要な次元を持ってくることで、例えば1024次元でも先頭の32,64,128,256...次元だけを使って後ろを切り捨てるだけで、ある程度良好な結果を示しています。

このグラフ参照元のStaticEmbedding の記事によると、128次元で91.87%, 256次元で95.79%, 512次元で98.53%の性能を維持しているようです。精度にそこまでシビアではないが、その後の計算コストを下げたい場合、ガッと次元削減して使う、という用途にも使えそうですね。

StaticEmbdding 日本語モデルでの次元削減結果

JMTEB では、出力時にモデルのパラメータを制御できるため、truncate_dim オプションを渡すことで、次元削減した結果のベンチマークも簡単に計測できます。素晴らしいですね。というわけで、StaticEmbdding 日本語モデルでも、次元削減した結果でベンチマークをとってみました。

スコアの変化を見ると、512次元へと次元削減した場合はやたらRetrieval, Classification,Reranking の性能が悪くなります。むしろ256次元まで次元削減してしまった方が良好な結果に。256次元では、スコア的には次元削減する前のモデルの98.93%なんですが、これはクラスタリングの結果がなぜか1024次元よりも良くなってしまったためですね。

512次元でのスコア計測が間違っていたので修正しました。マトリョーシカ表現学習がうまく反映され、次元数を削ると若干のスコア低下が見られますが、次元数が減ったためその後のコストが抑えられそうですね。

クラスタリングタスクにおいては128次元まで次元削減しても1024次元よりもスコアが高い、という本来情報量を削らない方がスコアが良いくなりそうなのに、クラスタリングタスクのみは逆にスコアが上がってしまう興味深い結果となりました…。マトリョーシカ表現学習では、先頭の次元の方が全体的な特徴を踏まえているので、クラスタリング用途には(クラスタリングのアルゴリズムにもよると思いますが)、特徴的な前の方の次元のみで後ろの次元を使わない方が良質な結果が得られる、ということなのかもしれません。

というわけで、static-embedding-japanese モデルで次元削減する時は、512,256,128次元あたりが性能と次元削減のバランスが取れてそうですね。

StaticEmbedding モデルを作ってみて

正直、単純なトークンのembeddingsの平均でそんなに性能出るのか半信半疑だったのですが、実際に学習させてみてシンプルなアーキテクチャなのに性能の高さにびっくりしました。Transformer 全盛のこの時代に、古き良き単語埋め込みの活用モデルで、実世界で利活用できそうなモデルの出現に驚きを隠せません。

CPUでの推論速度が速い文ベクトル作成モデルは、ローカルCPU環境で大量の文章の変換などはもとより、エッジデバイスだったりネットワークが遅い(リモートの推論サーバを叩けない)環境だったり、色々と活用できそうですね。

StaticEmbedding 日本語モデル学習のテクニカルノート

なぜうまく学習できるのか

StaticEmbedding は非常にシンプルで、文章をトークナイズしたIDで単語の埋め込みベクトルが格納されているEmbeddingBagテーブルからN次元(今回は1024次元)のベクトルを取得し、その平均を取るだけです。

これまで、単語埋め込みベクトルといえば、word2vec や GloVe のように Skip-gram や CBOW を用いて単語の周辺を学習してきました。しかし、StaticEmbedding では文章全体を用いて学習しています。また、対照学習を使って大量の様々な文章を巨大バッチで学習しており、良い単語の埋め込み表現の学習に成功しています。

対照学習は、基本的に正例以外全てを負例として学習するため、例えばバッチサイズ2048なら1の正例に対して2047の負例を2048通り、つまり2048x2047で約400万の比較を学習します。そのため、元の単語空間に対して適切な重みを更新しながら、学習を進めることができるのです。

学習データセット

日本語モデル学習にあたり、対照学習で利用できるデータセットとして、以下を作成し使用しました。

• hotchpotch/sentence_transformer_japanese
  • SentenceTransformer で学習しやすいカラム名と構造に整えたものです。
    • (anchor, positive), (anchor, positive, negative), (anchor, positive, negative_1, ..., negative_n) といった構造になっています。
  • 以下のデータセットを基に hotchpotch/sentence_transformer_japanese を作成しました。毎度ながらデータセットの作者の方々・とりわけ hpprc 氏に感謝です。
    • https://huggingface.co/datasets/hpprc/emb
      • https://huggingface.co/datasets/hotchpotch/hpprc_emb-scores のリランカースコアを使用し、positive(>=0.7) / negative(<=0.3) のフィルタリングを行いました。
    • https://huggingface.co/datasets/hpprc/llmjp-kaken
    • https://huggingface.co/datasets/hpprc/msmarco-ja
      • https://huggingface.co/datasets/hotchpotch/msmarco-ja-hard-negatives のリランカースコアを用いて、positive(>=0.7) / negative(<=0.3) のフィルタリングを行いました。
    • https://huggingface.co/datasets/hpprc/mqa-ja
    • https://huggingface.co/datasets/hpprc/llmjp-warp-html
• 上記の作成したデータセットの中で、以下を使用しました。なお、情報検索を強化したかったため、情報検索に適したデータセットのデータはオーギュメンテーションで件数を多めに学習させています。
  • httprc_auto-wiki-nli-triplet
  • httprc_auto-wiki-qa
  • httprc_auto-wiki-qa-nemotron
  • httprc_auto-wiki-qa-pair
  • httprc_baobab-wiki-retrieval
  • httprc_janli-triplet
  • httprc_jaquad
  • httprc_jqara
  • httprc_jsnli-triplet
  • httprc_jsquad
  • httprc_miracl
  • httprc_mkqa
  • httprc_mkqa-triplet
  • httprc_mr-tydi
  • httprc_nu-mnli-triplet
  • httprc_nu-snli-triplet
  • httprc_quiz-no-mori
  • httprc_quiz-works
  • httprc_snow-triplet
  • httprc_llmjp-kaken
  • httprc_llmjp_warp_html
  • httprc_mqa_ja
  • httprc_msmarco_ja
• 英語データセットには、以下のデータセットを利用しています。
  • sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1
  • sentence-transformers/squad
  • sentence-transformers/all-nli
  • sentence-transformers/trivia-qa
  • nthakur/swim-ir-monolingual
  • sentence-transformers/miracl
  • sentence-transformers/mr-tydi

日本語トークナイザ

StaticEmbedding を学習するためには、HuggingFace のトークナイザライブラリの tokenizer.json 形式で処理可能なトークナイザを使うと簡単そうだったので、 hotchpotch/xlm-roberta-japanese-tokenizer というトークナイザを作成しました。語彙数は 32,768 です。

このトークナイザは、wikipedia 日本語~~、wikipedia 英語(サンプリング)、cc-100(日本語, サンプリング)~~(訂正:作成コードを確認したところ、wikipedia日本語のみを利用していました)のデータを unidic で分割し、sentencepiece unigram で学習したものです。XLM-Roberta 形式の日本語トークナイザとしても機能します。今回はこのトークナイザを利用しました。

ハイパーパラメータ

大元の学習コードとの変更点やメモは以下の通りです。

• batch_size を大元の 2048 から 6072 に設定しました。
  • 対照学習で巨大なバッチを処理するとき、同一バッチ内にポジティブとネガティブが含まれると学習に悪影響を与える可能性があります。これを防ぐために BatchSamplers.NO_DUPLICATES オプションがあります。しかし、バッチサイズが巨大だと同一バッチに含めないためのサンプリング処理に時間がかかることがあります。
  • 今回は BatchSamplers.NO_DUPLICATES を指定し、RTX4090 の 24GB に収まる 6072 に設定しました。バッチサイズはさらに大きい方が結果が良い可能性があります。
• epoch数を1から2に変更しました
  • 1よりも2の方が良い結果になりました。ただし、データサイズがもっと大きければ、1の方が良い可能性があります。
• スケジューラ
  • 標準のlinearから、経験則でより良いと感じるcosineに変更しました。
• オプティマイザ
  • 標準のAdamW のままです。adafactorに変更した場合、収束が悪くなりました。
• learning_rate
  • 2e-1 のままです。値が巨大すぎるのではないかと疑問に思いましたが、低くすると結果が悪化しました。
• dataloader_prefetch_factor=4
• dataloader_num_workers=15
  • トークナイズとバッチサンプラのサンプリングに時間がかかるため、大きめに設定しました。

学習リソース

• CPU
  • Ryzen9 7950X
• GPU
  • RTX4090
• memory
  • 64GB

このマシンリソースで、フルスクラッチ学習にかかった時間は約4時間でした。GPUのコア負荷は非常に小さく、他のtransformerモデルでは学習時に90%前後で張り付くのに対して、StaticEmbeddingではほとんど0%でした。これは、巨大なバッチをGPUメモリに転送する時間が大半を占めているためかと思われます。そのため、GPUメモリの帯域幅が速くなれば、学習速度がさらに向上する可能性があります。

さらなる性能向上へ

今回利用したトークナイザはStaticEmbedding向けに特化したものではないため、より適したトークナイザを使用すれば性能が向上する可能性があります。バッチサイズをさらに巨大化することで、学習の安定性が向上し、性能向上が見込めるかもしれません。

また、さまざまなドメインや合成データセットを利用するなど、より幅広い文章リソースを学習に組み込むことで、さらなる性能向上が期待できます。

大元の学習コード

学習に使用したコードは、以下で MIT ライセンスで公開しています。スクリプトを実行すれば再現できる、はず...!

• https://huggingface.co/hotchpotch/static-embedding-japanese/blob/main/trainer.py

ライセンス

static-embedding-japanese はモデル重み・学習コードを MIT ライセンスで公開しています。

生活

2024はマイペースで過ごすことができた。家を建てていたので打ち合わせや様子を見にいくことが多く、長期海外旅行などは行けなかったけど、海外にも国内にもちらほら旅行したし、まぁまぁの活動。

健康面では人生での体重最高を更新してしまったので、ダイエットも兼ねてパーソナルジムへ行き、-5kgで一昔前の体重へ。本当はあと5kgぐらい痩せて、筋肉量も増やしたいが、この辺は週一のジム通いだけでなく、ちゃんと運動しないと筋肉はつかなさそう。

そして家がついにたった！今年ギリギリの引っ越しで暮らし始めてまだ6日目ほどだが、今のところ大変快適に過ごしている。薪ストーブも今のところ寝ている時以外は常時焚いていて暖かくて最高だ。

仕事

仕事では今年から始まったAI関連プロダクトのPdM・データサイエンティスト(主に情報検索分野)としてだいぶ好きにやらせてもらい、自分にない知見を持ったチームメンバーとも関われ、楽しく充実した一年。結構良いプロダクトができたのではななろうか。プロダクトのサイレントリリースはしていて、来年は多分広く使われるようになるんじゃないかなぁ。というか広く使われるようにしていきたい。

技術・趣味

2023年後半から情報検索が断然面白くなってきて、仕事でもやり始めたのけど、趣味ではデータセット作ったり、リランカーや検索モデルを作って公開したりと、色々と役立つ物のアウトプットもある程度はできたかなーと思っている。LLMの学習はマシンリソース的に難しいのだけど、何かを行うことに足りないワンピースを埋めるようなモデルは作れるし、日本語関連は公開する人が少ないから特定用途に特化すればNo1モデルも作れる感じだしで楽しい。

情報検索・NLP・機械学習分野はやればやるほど知識が増えていく感じが続いて、面白すぎるじゃんという感じで、まだまだやりたいこと知りたいことがたくさんあるので、技術分野ではこの辺を抑えつつ、趣味で色々作ったり、この辺の知見を仕事に活かしたりと、2025も色々やっていきたい。

が、ちょっと仕事よりの技術にプライベートでは時間を使いすぎで(朝晩は大体趣味のモデル作りに勤しんでいた)、プライベートではもっとあれこれ広く楽しんだ方が人生が豊かになりそうな気もしていて、2025はこの辺のバランスも気をかけたい。

総括

振り返るとあっという間の一年で、結構色々やったような、まだまだ全然やれてないような、そんな年だった。2023年に崩したマイペースは、取り戻したと言えそうで、自分にちょうど良いバランスであれこれできた気がする。

そして毎度のことだが、さまざまなことをサポートしてくれる妻に感謝だ。というわけで、2025年もみなさんよろしくお願いいたします。

• 振り返り2023年
• 振り返り2022年
• 振り返り2021年
• 2020年の振り返り

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

なおこの記事は、情報検索・検索技術 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は現代の検索システムに求められる要件を、バランスよく実現しています。

実際の性能は高いの?

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

これは記事冒頭の 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開発・自然言語処理・情報検索をしている皆さんに少しでも有益になれば幸いです。

あわせて読みたい

• 高性能な日本語SPLADE（スパース検索）モデルを公開しました
• SPLADE モデルの作り方・日本語SPLADEテクニカルレポート

本記事では、RetroMAEを用いて日本語BERTを事前学習したモデルを作成・公開し、後続の検索タスク(JMTEB)で評価を行っている。

結果として、ほぼ全てのスコアにおいて性能向上が確認されており、RetroMAEの有益性が確認された。また学習方法も、教師なしでテキストのみを与えれば良いという手軽さも実用性が高いであろう。

RetroMAE について

RetroMAEの特徴は、Masked Auto-Encoderの手法を採用しながら、以下の3つの設計を取り入れた点である。

1. 入力文に対して異なるマスクを適用する新しいワークフロー
2. エンコーダーとデコーダーで非対称な構造を採用
3. エンコーダーとデコーダーで異なるマスク率を使用

これらの工夫により、文書の意味をより深く理解し、効果的な検索を可能にする表現学習を実現している。実際の評価においても、BEIRやMS MARCOなどのベンチマークで優れた性能を示している。また高性能のマルチリンガル密ベクトルモデルの BAAI/bge-m3 も RetroMAE を用いた事前学習を行なっているという実績もある。

なお、RetroMAE はさらなる進化的な手法である RetroMAE v2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models (DupMAE) も提案されているが、本内容では RetroMAE を扱っている。

異なるマスクを適用するワークフロー

RetroMAE では、入力文に対して2つの異なるマスクを適用する。1つ目のマスクを適用した入力からエンコーダーが文埋め込みを生成し、2つ目のマスクを適用した入力とその文埋め込みを組み合わせてデコーダーが元の文を復元する。

エンコーダーとデコーダーで非対称な構造

RetroMAEは、エンコーダーとデコーダーで意図的に非対称な構造を採用している。エンコーダーには入力文の意味を十分に捉えるため、BERTのトランスフォーマー（12層）を使用する。一方、デコーダーは極めてシンプルな1層のトランスフォーマーのみを採用している。このシンプルなデコーダー構造により、文の復元タスクがより困難になり、結果としてエンコーダーがより質の高い文埋め込みを生成可能に学習される。

さらに、1層のデコーダーには Enhanced decoding という特殊な仕組みが導入されている。これは、文埋め込みと位置埋め込みを組み合わせたクエリと、文埋め込み・トークン埋め込み・位置埋め込みを組み合わせたコンテキストの2つを用意し、位置に応じた attention mask を適用する方式である。この仕組みにより、すべての入力トークンをデコーダの復元対象として使用でき、かつ各トークンが独自のコンテキストから復元されるようになる。これにより、限られたデコーダー層数でも効率的な学習が可能となっている。

エンコーダーとデコーダーで異なるマスク率を使用

エンコーダーには適度なマスク率（15〜30%）を適用し、入力文の大部分の情報を保持できるようにしている。一方、デコーダーには積極的なマスク率（50〜70%）を適用する。この高いマスク率により、デコーダーは入力のみでは十分な復元が難しくなり、エンコーダーが生成した文埋め込みに大きく依存せざるを得なくなる。結果として、エンコーダーはより深い意味理解を強制される仕組みとなっている。

RetroMAE 日本語モデルの事前学習

元論文では、英語のwikipedia, BookCorpus, MS Marcoをデータセットとして学習させている。そのため、日本語データセットとしては類似タスクを含む、以下のデータセットを用いた。

• (A) wikipedia 日本語 - hpprc/jawiki-paragraphs
• (A) jawiki-books - hpprc/jawiki-books-paragraphs
• (B) MQA 日本語 - hpprc/mqa-ja
• (B) JSNLI - shunk031/jsnli

wikipedia, jawiki-books はパラグラフのみ(タイトルは含まず)、MQAは query と doc を連結した文章、JSNLIは空白を削除した文章を用いている。

また、ゼロの重みから学習を行うのではなく、RetroMAEを学習させるためのモデルとして tohoku-nlp/bert-base-japanese-v3と cl-nagoya/ruri-pt-baseを用いた。ruri-pt-base は bert-base-japanese-v3 を元に、対照学習を行なった事前学習モデルであり、その過程で MLM のデコーダ層が失われるため、デコーダ層の重みとしてbert-base-japanese-v3からコピーを行ったモデルを利用した。

学習用スクリプトはOSS(MITライセンス)で公開されている、retromae_pretrain を用いた。学習パラメータはエンコーダのマスク率を 30%、デコーダのマスク率を50%としている。その他 Trainer のハイパーパラメータは以下である。

  "learning_rate": 1e-4,
  "num_train_epochs": 2,
  "per_device_train_batch_size": 16,
  "gradient_accumulation_steps": 32,
  "warmup_ratio": 0.05,
  "lr_scheduler_type": "cosine",
  "bf16": true,
  "dataloader_drop_last": true,
  "dataloader_num_workers": 12

これらを用いて、データセットの (A) のみと (A) + (B) を使って RetroMAE 事前学習モデルを作成した。

後続検索タスクでの評価

後続検索タスクとして、日本語SPLADEモデルを mmacro データセットのみで学習させて評価を行った。パラメータは japanese-splade-base-v1-mmarco-only の model のエポック数を12から10に削減し、model_name を今回評価するものに差し替えたものである。

また、評価には JMTEB のスパースベクトルを評価できるように変更したfork版を用い、検索タスク(retrieval)で行った。

評価結果スコアは以下である。

ほぼ全ての評価において、RetroMAE で学習させていないモデルよりも、RetroMAEで学習させたモデルの方がスコアが高いことが計測された。最も評価が良かった ruri-pt-base retromae(A+B) においては、bert-base-japanese-v3 よりも約2%ほど性能が向上した。

また学習データセットも (A) のみよりも、基本的に(A) + (B) 両方のデータセットを学習させた方がスコアが高かった。これは、さらにデータセットを追加したり、特定ドメインのテキストを学習させることで、性能向上に寄与しそうな結果である。

なお、(A+B)のデータセットで学習させた RetroMAE モデルを、HuggingFace で公開している。

• https://huggingface.co/hotchpotch/ruri-pt-base-retromae
• https://huggingface.co/hotchpotch/bert-base-japanese-v3-retromae

おわりに

本記事では、検索タスクに特化した事前学習手法であるRetroMAEを日本語BERTモデルに適用し、その効果を検証した。後続タスクのSPLADEモデルでの評価の結果、RetroMAEで学習させたモデルは、ベースラインとなるbert-base-japanese-v3と比較して、ほぼすべての検索タスクで性能向上が確認された。特に、対照学習済みのruri-pt-baseをベースに、Wikipedia、書籍、質問応答データなど複数のデータセットで学習させたモデルでは、平均約2%の性能向上が達成された。

また、RetroMAEの利点として、教師なしでテキストデータのみを用いて学習できる手軽さも注目に値する。これは、特定ドメインや業務向けにモデルをカスタマイズする際にも有用であり、実用性が高いと言える。今後は、さらなる学習データの追加や、特定ドメインのテキストを用いた追加学習による性能向上の可能性も期待できる。

なお作成したRetroMAEモデルはHuggingFace上で公開しており利用可能となっている。本記事が、日本語の検索タスクの性能向上に貢献できれば幸いである。

ニューラルネットワークを用いた検索モデルの分野では、密ベクトルモデル(dense retriever)が主流となっており、multilingual-e5 や bge-m3 のようなマルチリンガル対応の高性能モデルも登場している。一方で、SPLADE(Sparse Lexical and Expansion Model with Contextualized Embeddings)に代表されるスパース検索モデルも、英語圏において高い性能を示している。

しかし、SPLADE は単語の特徴量に大きく依存し、そのトークン化がモデルのトークナイザに左右されるため、マルチリンガル対応版が存在していなかった。マルチリンガルモデルのトークナイザでは多くの言語で1文字単位の分割が行われ、意味のある単語単位でのトークン化が困難であったためである。そこで日本語に特化したSPLADEモデルを開発し、その評価を行った。

さらに、元のSPLADE実装(naver/splade)がCC-BY-NCライセンスで提供されており商用利用に制限があることから、論文を基にTrainerを実装し、MITライセンスのオープンソースソフトウェアとして公開した。

• YAST - Yet Another SPLADE or Sparse Trainer

本レポートでは、日本語SPLADEモデルの実装詳細、評価実験の結果、および今後の展望について報告する。

SPLADEのアルゴリズム

SPLADE は、情報検索においてスパースな文書およびクエリ表現を学習するためのモデルである。本節では、SPLADE がどのように学習されるか、そのアルゴリズムについて記述する。

単語重要度の計算と単語ごとの出力トークンの利用

SPLADE は、Masked Language Modeling(MLM)などで事前学習されたモデルの各単語ごとの出力トークンを利用し、文脈に応じた単語の重要度を計算する。具体的には、BERT のような事前学習モデルの語彙空間を活用し、入力シーケンスの各位置で得られた単語のスコアから最大値を選択する max pooling を用いる。また、対数飽和関数を適用することで、極端な値を抑制しつつ重要な特徴を強調することが可能である。これらの手法により、顕著な特徴を捉えたスパースで効率的な文書およびクエリの表現を生成する。

なお、これらの操作は SPLADE Max と呼ばれるもので、Python での実装を以下に示す。

def splade_max_pooling(logits, attention_mask):
    # Step 1: 対数飽和関数の適用 (log(1 + x))
    # - torch.relu() で負の値を0にする
    # - torch.log(1 + x)で値を対数スケールに変換し、大きな値を抑制
    relu_log = torch.log(1 + torch.relu(logits))
    
    # Step 2: attention_maskを使って、padding された位置のスコアを0にマスク
    # unsqueeze(-1)で次元を合わせる(batch_size, seq_len, 1)
    weighted_log = relu_log * attention_mask.unsqueeze(-1)
    
    # Step 3: max pooling の適用
    # torch.max()で系列長方向(dim=1)の最大値を取得
    # 各語彙に対する最も重要なスコアを選択する
    max_val, _ = torch.max(weighted_log, dim=1)
    
    return max_val

単語重要度を用いたドキュメントとクエリの予測

SPLADE Max を通じて得られた単語重要度を活用し、ドキュメントとクエリの関連度を予測する。関連度は主に内積を利用する。この予測結果と元の学習データとの間の差異を、損失関数として定義する。

この損失関数には、モデルが予測した語彙分布と実際の語彙分布との間の差異を測定するために、KLダイバージェンス損失、MarginMSE損失、クロスエントロピー損失等を用いる。これらの損失関数は単体で用いても、複数を組み合わせても良い。SPLADE-v3では、KLダイバージェンス損失とMargineMSE損失を組み合わせて使っている。

スパース性の導入と正則化

出力される単語重要度にスパース性を持たせるため、正則化手法を損失関数に組み込む。具体的には、以下のアルゴリズムが使用される。

• L1正則化：モデルのパラメータの絶対値の総和を最小化することで、多くのパラメータをゼロに近づける。この手法により、重要でない単語の影響を排除し、スパースな表現を促進する。

• FLOPs正則化：高次元でスパースな表現学習において、非ゼロ要素を次元間で均一に分散させることで行列演算の計算量(FLOPs)を二次的に削減する正則化手法。(Minimizing FLOPs to Learn Efficient Sparse Representations)

なお、クエリとドキュメントでは異なる損失関数や正則化係数を適用することが可能である。また、学習の初期段階から強い正則化を適用すると、重要度予測に悪影響を及ぼす可能性がある。そのため、適用まで緩やかなウォームアップ期間を設けて、正則化損失の重みを徐々に高めていく手法も取り入れている。

モデルの学習と関連度の計算

これらの手法を組み込んで学習を行うことで、スパース性を促進させつつ、クエリとドキュメントの関連度を高めるモデルを構築できる。SPLADE は、スパースな表現とニューラルネットワークの文脈を含めた語彙情報を組み合わせることで、高性能な情報検索が実現できる。

日本語モデルでの学習手法

データセットの準備

最終的に学習したモデル japanese-splade-base-v1 の学習用データセットとして、日本語の様々な質問文と回答、ハードネガティブを集めたhpprc/embのうち、auto-wiki-qa、mmarco、jsquad、jaquad、auto-wiki-qa-nemotron、quiz-works、quiz-no-mori、miracl、jqara、mr-tydi、baobab-wiki-retrieval、mkqa を利用した。また、hpprc/embのデータに対して日本語高性能なクロスエンコーダーを用いたリランカー(BAAI/bge-reranker-v2-m3、cl-nagoya/ruri-reranker-large)を使用し、スコア付けを行ったデータセットも作成した(hotchpotch/hpprc_emb-scores)。さらに、英語データセットとしてMS MARCOと、そのデータに BAAI/bge-reranker-v2-m3 でスコア付けしたデータを利用した。

データのフィルタリングにおいては、各リランカーの平均スコアを用いて、正例に対してはスコアが0.7以上、負例に対しては0.3以下のデータを選別した。これは、質問に対して適切なスコアでないとリランカーが判断した文章を除外するためである。

データセットの割合が少ないものについては、1エポックあたりの学習量を増加させた。これは、そのデータセットの特性をモデルが忘れないようにするためである。

また、mmacro(日本語)のみを学習させるデータセットとして、mmacroとBAAI/bge-reranker-v2-m3でスコア付けしたデータセットhotchpotch/mmarco-hard-negatives-reranker-scoreを作成し、利用した。このデータもリランカーの平均スコアを用いて、同様に正例に対してはスコアが0.7以上、負例に対しては0.3以下のデータを選別した。

学習の設定とハイパーパラメータ

学習における損失関数として、単純なクロスエントロピー損失を採用した。これは、高性能なリランカーから得られたスコアをモデルが学習できるようにするためである。他にもKLダイバージェンス損失やMarginMSE損失を試したが、クロスエントロピー損失が最良の結果を示した。

スパース性を促進する正則化項には、L1正則化を使用した。これは、FLOPs損失と比較した際、日本語においてL1正則化の方がスパース性の促進効果が高かったためである。

ハイパーパラメータとして、学習率(Learning Rate, LR)は一般的な110Mパラメータのモデルで用いられる5.0e-2を設定した。学習率のスケジューラにはコサインスケジューラを採用し、全体の10%をウォームアップ期間として設定した。

また1つのバッチでは正例1つ・負例7つ、合計8つのデータを含めている。バッチサイズは、japanese-splade-base-v1 が 32、japanese-splade-base-v1-mmarco-only が 128である。これは、mmacroのみの場合はクエリと文章のスパース性の収束が大きなバッチサイズでも早く、多様なデータセットを学習しているjapanese-splade-base-v1ではバッチサイズが大きいとスパース性の収束が遅くなるため、小さいバッチサイズの方が適していたためである。なお、学習時間やリソースに余裕があるなら、japanese-splade-base-v1も大きいバッチサイズの方が良い結果になる可能性がある。

その他、詳細なパラメータは、実際の学習に使った設定ファイルを参考にされたし。

ノイズトークンの除去

日本語での学習において、、 。 「 ：などの句読点や記号のトークンがノイズとして顕著に特徴量に現れることが確認された。これらのトークンがSPLADE Maxの出力に残存する場合、ペナルティとしてそのトークンのスコアを損失関数に追加している。また、これらのトークンはfugashiとunidic-liteを用いて、記号的な単語と判定できるものを抽出した。

これらをノイズトークンとして扱い、損失に組み込むことで、学習済みモデルの出力においてこれらのノイズトークンはほぼ出力されなくなった。また、学習の安定性が向上し、収束速度の速さも観測された。

学習元モデルの選択

今回、学習元となるモデルには、MLM(Masked Language Modeling)による事前学習で獲得した語彙の意味的特徴量を出力層に持つtohoku-nlp/bert-base-japanese-v3を利用した。このモデルは日本語BERTアーキテクチャをベースとしている。

学習

これらを基に、japanese-splade-base-v1とjapanese-splade-base-v1-mmarco-onlyモデルをファインチューニングし作成した。学習にかかった時間はGPU RTX4090環境で、japanese-splade-base-v1が約33時間、japanese-splade-base-v1-mmarco-onlyが約24時間である。

また、japanese-splade-base-v1においてはデータセットサイズが大きいため2エポック、japanese-splade-base-v1-mmarco-onlyにおいてはデータセットはmmacroのみとデータセットが小さいため12エポック学習した。なお、japanese-splade-base-v1の学習エポックを増やすと過学習になるためか、トレイン損失値は下がるが、評価時の検索タスクにおいて性能低下が確認された。

なお、学習したモデルは HuggingFace で公開している。

• https://huggingface.co/hotchpotch/japanese-splade-base-v1
• https://huggingface.co/hotchpotch/japanese-splade-base-v1-mmarco-only

評価結果

JMTEB retrieval タスクでの評価結果

JMTEBでの評価結果は以下の通りである。なお、実際の評価にはスパースベクトルを評価できるように変更したfork版を利用している。

結果の平均としては、japanese-splade-base-v1 が mrtydi や JAQKET のドメインタスクを学習(JMTEB の評価で使うテストデータではない)しているが、japanese-splade-base-v1 が最良の結果となった。また、japanese-splade-base-v1-mmarco-only は mmacro データセットしか学習させていないが、mrtydiでは最良の結果となり、他のタスクも他のモデルと十分競争力がある結果となった。

jagovfaqs の結果は、SPLADE モデルが他のモデルに比べて軒並み悪い。これは jagovfaqs のクエリの内容が「FAQ」であり、要約・文脈類似タスクに似た問題が多く含まれることが考えられる。他のモデルは文章の意味的類似度を学習しており、japanese-splade-base-v1 は学習していない。また、スコアが高い日本語モデルのruri-largeやGLuCoSE-base-ja-v2では、マルチリンガルFAQ(Frequently Asked Questions) & CQA(Community Question Answering)データセットのの MQAの日本語データを学習していることも、スコア向上に寄与している可能性がある。

jaqket の結果は、「クイズ形式」の質問が多く含まれる。「XXXといえばYYYですが、ZZZといえば何でしょう？」のような日本語クイズ独特の言い回しを含んでおり、それらの表現を学習しているモデルが高スコアになる。また、正解の文章内部に正解の単語を必ず含むため、単語特徴量に強い SPLADE が高スコアにつながると考えられる。

mrtydi の結果は、mrtydi のドメインを学習しているはずのjapanese-splade-base-v1が、ドメインを学習していないjapanese-splade-base-v1-mmarco-onlyよりも悪いという、直感に反する結果となった。これについては十分な考察ができていない。

nlp_journal の三つのタスクにおいては、title_abs においては、SPLADEモデルが軒並み高性能だが、abs_intro、title_introにおいては text-embedding-3-large が圧倒的に高性能である。これはtitle_absの文章の平均長が442で、abs_intro、title_introは2052のためである。text-embedding-3-large 以外はモデルのトークンの最大長が全て512であり、text-embedding-3-largeは8191である。そのため、text-embedding-3-large 以外のモデルはabs_intro、title_introの文章全体を処理することができず、文章の冒頭一部のみでの評価になるため、長いトークン長を理解可能なモデルが高いスコアとなる。

reranking タスクでの評価結果

reranking タスクの評価には JQaRA、JaCWIR を用いた。

結果としては、JQaRA のドメインを学習しているとはいえ、japanese-splade-base-v1 がどれも最良の結果となった。

英語タスクでの評価

japanese-splade-base-v1は、MS MARCOの英語データセットも学習データセットに含めた。そのため、naver/spladeで公開されている評価スクリプトを用い、MS MARCO(dev)で評価した。

結果として、英語データを学習していないjapanese-splade-base-v1-mmarco-onlyよりも、わずかながらスコア向上が見られるが、英語のみを学習している naver/splade_v2_max と比べると著しくスコアが低く、英語における検索性能はほとんどないと言える。

スパース性の評価

スパース性の評価では、非ゼロ要素の数(L0ノルム)を用いて、各モデルのクエリおよび文書のスパース性を測定した。以下に、JMTEBのretrieveタスク(Top-1000)における japanese-splade-base-v1 および japanese-splade-base-v1-mmarco-only モデルのクエリおよび文書のスパース性の結果を示す。

なお、この結果は JMTEB_L0.py で計測した。

L0ノルムの値から、v1-mmarco-onlyの方が全体的に非ゼロ要素が多く、スパース性が低いことが示されている。クエリと文書のスパース性の度合いは、検索システムのパフォーマンスに対する重要な要素とされるが、クエリと文書には異なる要件がある。

検索速度を考慮する場合、クエリのスパース性が高いほど効率的な検索が期待できるが、文書のスパース性もまた省メモリや省ディスクの観点で重要である。ただし、実運用環境では数百万〜数千万規模の文書が1台のマシンでもオンメモリで検索可能な場合が多いため、文書のスパース性についてはクエリほど厳格に管理する必要はないと考えられる。

一方、クエリのスパース性は検索速度に直接関係するため、できる限り高いスパース性が求められる。ただし、文書のスパース性に関しても、非ゼロ要素が少なすぎると検索性能に悪影響を及ぼす可能性があるため、適切なバランスが求められる。検索システムの性能と効率の両立を目指す上で、クエリと文書のスパース性を考慮したチューニングが重要である。

評価の考察まとめ

これらの結果から、japanese-splade-base-v1 は日本語データの検索において、最新のモデルと十分競争力があるモデルと言える。とりわけ、単語特徴量が重要と思われるタスクでは優れた性能を発揮する。クエリや文章のスパース性能も、必要十分と言えよう。

また、他のモデルは密ベクトルモデルであるが、SPLADE はスパースベクトルモデルであり、単語特徴量を重視する検索結果になるため、密ベクトルモデルのみを利用するより、異なるモデルを組み合わせることで多様性のある検索結果を得ることができる。これは、実世界で多様な検索結果を取得したいというケース、例えばLLMにさまざまな検索情報を渡すなど、で重要になるだろう。

今後の展望

一旦、japanese-splade-base-v1 を成果物として公開したが、まだ性能向上の余地は多い。SPLADEの元論文では、自己蒸留(self distillation)や複数の損失スコアの利用、SPLADEモデル自体を使ったハードネガティブサンプリングなどを行うことにより、性能向上が図られている。

また、検索タスクに適した事前学習モデルの選択・学習なども行えていない。例えば、Unsupervised Corpus Aware Language Model Pre-training for Dense Passage RetrievalやRetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder等、検索タスクに適した事前学習を行うことで、性能向上の可能性がある。

他にも、FAQ系のタスクのデータセットの学習やロングコンテキストへの対応、多様なデータセット(現状ではWikipediaに偏りがち)の追加等が考えられる。

近年、Llama 3.1 をはじめとする、LLM の出力を学習に利用可能なライセンスを持つモデルが登場し、ライセンス上の問題なく検索用データセットを作成できるようになってきた。本モデルでも利用した hpprc/emb では、LLM の出力を活用した高品質なデータセットを提供している(Ruri: Japanese General Text Embeddings)。

従来、ドキュメントから情報検索に適したクエリを作成することは人手がかかり大変であったが、LLM を用いて自動的に生成することで、低コストで大量のクエリを作成できるようになった。特定のドメインを学習することで一般化性能が向上する場合が多く、情報検索モデルの学習用データセットが充実することで、さらなる性能向上が期待できる。

おわりに

本レポートでは、日本語に特化したSPLADEモデルであるjapanese-splade-base-v1を開発し、その評価を行った。評価結果から、日本語の情報検索において、既存の最新モデルと比較しても高い性能を示すことが確認できた。

今後の課題として、さらなる性能向上のための手法の検討や、検索タスクに適した事前学習モデルの選択、多様なデータセットの活用が挙げられる。

日本語SPLADEモデルとSPLADEモデル学習用Trainerの公開により、情報検索技術の発展に寄与できれば幸いである。

参考文献

• SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking
• SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval
• From Distillation to Hard Negative Sampling: Making Sparse Neural IR Models More Effective
• An Efficiency Study for SPLADE Models
• A Static Pruning Study on Sparse Neural Retrievers
• SPLADE-v3: New baselines for SPLADE
• Minimizing FLOPs to Learn Efficient Sparse Representations
• Ruri: Japanese General Text Embeddings
• JaColBERTv2.5: Optimising Multi-Vector Retrievers to Create State-of-the-Art Japanese Retrievers with Constrained Resources
• 日本語テキスト埋め込みベンチマークJMTEBの構築
• mMARCO: A Multilingual Version of the MS MARCO Passage Ranking Dataset
• Mr. TyDi: A Multi-lingual Benchmark for Dense Retrieval
• JaCWIR: Japanese Casual Web IR - 日本語情報検索評価のための小規模でカジュアルなWebタイトルと概要のデータセット
• JQaRA : Japanese Question Answering with Retrieval Augmentation - 検索拡張(RAG)評価のための日本語 Q&A データセット
• JAQKET: クイズを題材にした日本語 QA データセットの構築
• 高性能な日本語SPLADE（スパース検索）モデルを公開しました - A Day in the Life

@article{tateno2024splade,
    title={SPLADE モデルの作り方・日本語SPLADEテクニカルレポート},
    author={TatenoYuichi},
    year={2024},
    url={https://secon.dev/entry/2024/10/23/080000-japanese-splade-tech-report/}
}

• https://huggingface.co/hotchpotch/japanese-splade-base-v1

なお、日本語SPLADEモデル作成における技術的な詳細は、SPLADE モデルの作り方・日本語SPLADEテクニカルレポートをご覧ください。

SPLADEとは？

SPLADE（Sparse Lexical and Expansion Model）は、その名の通りスパース（疎）なベクトルを用いた検索モデルです。スパース検索といえば、長年利用されているBM25が代表的で、高い性能を誇るアルゴリズムとして広く利用されています。しかし、BM25はクエリとドキュメントの単語の完全一致に依存しているため、関連する単語や同義語を含む文書を見逃す可能性があります。

一方、SPLADEはTransformerアーキテクチャを活用して、文脈に基づく関連性の高い単語もベクトルに含めることができます。これにより、完全一致以外の単語も検索候補として取り込むことができ、より柔軟で効果的な検索が可能となります。

特性とメリット

SPLADEは、以下の特性を備えています。まず、事前学習済みのTransformerモデル（例：BERT）を利用することで、入力テキストの文脈を深く理解します。これにより、単語の完全一致に依存せず、文脈に基づいて関連性の高い単語も効果的に抽出することが可能です。また、各単語には重要度のスコアが付与され、検索においてどの単語が重要であるかが明確に示されます。さらに、スパースベクトルを生成することで、多くの要素がゼロとなり、計算量を抑えつつ効率的な検索を実現します。

これらの特性により、SPLADEは柔軟な検索を可能にし、関連語や同義語を含む幅広い検索ニーズに対応します。スパースベクトルの活用により、計算量が少なく高速な検索が可能となり、システム全体の効率性が向上します。さらに、各単語の重要度が明確に示されるため、検索結果の解釈が容易になり、ユーザーにとって理解しやすい結果を提供します。最後に、既存の検索エンジンへの導入が容易であるため、現在のシステムにスムーズに統合することが可能です。

具体的な例

SPLADEの動作を理解するために、具体的な例を見てみましょう。これは実際のモデル japanese-splade-base-v1 を利用した出力です。なお、SPLADE 日本語 demoからも、出力結果を簡単に取得することができます。

単語が拡張される例

"車の燃費を向上させる方法は？" の、SPLADEによる出力例

このように、クエリの文脈を理解し、元の文に含まれていない「燃料」や「効果」といった関連語も重要な単語として抽出しています。また、各単語には重要度を示すスコアが付与されています。なお「は」など、全く関係なさそうかつノイズになりそうな単語も含んでいますが、このような単語は他の出力にも多く含まれるため、無視できる程度のノイズになっていることが多いため、検索にうまくヒットさせることができるのです。

同様に、文章に対しても行うことができます。このクエリと文章のスパースベクトルの内積をスコアとすることで、どれだけ関連しているのかを計算を行えます。

性能は？

冒頭で述べたように、SPLADEモデルは多くの日本語情報検索タスクで優れた性能を示しています。JMTEB(retrieval)や JQaRA, JaCWIR でのベンチマーク結果は以下です。単語特徴量が結果に色濃く出るタスクでは、軒並み高性能な結果となっています。代わりに、jagovfaqs(FAQ)のような似ている文章の理解が必要そうなタスクでは、あまり振るわない結果となっています。

JMTEB retrieval

JQaRA, JaCWIR reranking

また、ほとんどのオープンソース検索エンジン（Elasticsearch、OpenSearch、Qdrant、Vespaなど）でスパース検索がサポートされているため、導入も容易です。検索速度の面でも、スパースベクトル検索は古くから行われており、BM25などと同様に高速です。

また、SPLADEやBM25は単語特徴量が色濃く反映されるため、mulitilingual-e5 等の密ベクトルモデルと異なった検索結果になることも多いです。そのため、ハイブリット検索として双方の検索結果を組み合わせることで、より良い結果・多様性がある結果をもたらすことが可能です。ハイブリット検索も、先ほどの検索エンジンは基本的にサポートしており、簡単に利用が可能なものが多いです。

本番環境で運用しにくいのでは？

SPLADEの運用は、密ベクトルモデルとほぼ同様に運用ができるため、難しくありません。検索エンジンは先ほど述べた通りスパース検索もサポートしているものがほとんどです。 またSPLADEのスパースベクトルを得ることも、何か複雑なことを行なっているわけではなく、単語(token)の各スコアを、SPLADE max と呼ばれる max pooling と対数飽和関数の組み合わせに通すだけです。

• 例: transformers ライブラリでのスパースベクトル取得例

また、高速で本番運用しやすい推論サーバである text-embedding-inference (blog記事) からも利用可能です。

• https://huggingface.co/hotchpotch/japanese-splade-base-v1-dummy-fast-tokenizer-for-tei

おわりに

当初、SPLADEが本当に高い性能を発揮するのか半信半疑でした。しかし、英語のMS MARCOデータセットのみで学習されたSPLADE-v3が、他の多様な検索タスクでも高性能を示していることから、日本語で適切に学習させた場合の可能性に興味を持ちました。

また、SPLADEはトークナイザの語彙に依存するため、文字レベルで分割することが多いマルチリンガルモデルのトークナイザーとは相性が悪く、そのため日本語用に特化させた学習が必要なことも、面白そうと感じたきっかけでした。（日本語にも対応した密ベクトルの高性能マルチリンガルモデルの作成は、さまざまな企業が参入しているので…）

学習の結果、既知のドメインタスク（JAQKET、mytidi）を含むとはいえ、モデルもパラメータ数も110Mのbaseサイズで、OpenAIの大規模なモデルよりもベンチマークで上回るスパース検索モデルを作成することができました。

学習時間もRTX 4090で33時間程度と、学習計算機リソースや学習時間が少なくても学習が済むため、自社のドメインにフィットした検索結果を求める方々にとって、SPLADEを使った独自ドメインデータを学習させるモデルを作ることは有用なアプローチとなりそうです。

まだまだ今後もSPLADEを用いた日本語スパース検索性能は向上すると思っており、研究対象としても興味深い分野ですね。

しかしながら、日本語環境での問題点の一つが rust ベースの FastTokenizer 動かせる、つまり tokenizer.json を用意しているモデルでないと利用できないことだ。日本語 transformer モデルの多くが、unidic や mecab といった python で動く形態素解析辞書・ライブラリを利用するため、tokenizer.json 方式では動かせないモデルも多い。

最初、私も大変困ったのだが、/embed や /embed_sparse (残念ながら /rerank は非対応) など一部のAPIは無理やり利用できることがわかっているので、例として cl-nagoya/ruri-base を元に、その方法を記録に残す。

dummy の tokenizer.json を用意する

TEI は起動時のモデルのチェックで tokenizer.json がないと、そもそも起動しない。そのため、dummy となる tokenizer.json を用意する。tokenizer.json は自分で作っても、公開モデルのものを使っても良いのだが、とりあえずhotchpotch/mMiniLMv2-L6-H384のtokenizer.jsonを利用する。

このtokenizer.jsonを追加した、ruri-base をruri-base-dummy-fast-tokenizer-for-teiとして作成した。

dummy の tokenizer.json を使ったモデルでサーバを起動する

例として docker-compose.yaml を用意して

services:
  ruri-base:
    # image の部分はアーキテクチャにあったものに変えること
    image: ghcr.io/huggingface/text-embeddings-inference:86-1.5
    ports:
      - "8080:80"
    volumes:
      - /tmp/docker-tei-data:/data
    # pooling はモデルアーキテクチャにあったものに変える
    command: [ "--model-id", "hotchpotch/ruri-base-dummy-fast-tokenizer-for-tei", "--dtype", "float16", "--pooling", "mean", "--max-batch-tokens", "131072", "--max-client-batch-size", "16" ]
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [ gpu ]

起動する。これで port 8080 で立ち上がるはず。

$ docker compose up
...
ruri-base-1  | 2024-09-30T06:51:45.266929Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:80
ruri-base-1  | 2024-09-30T06:51:45.266940Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

手元で token_ids に変換して API を叩く

続いて、手元で Tokenizer を使って token_ids に変換して叩く。

from transformers import AutoTokenizer
import requests
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("hotchpotch/ruri-base-dummy-fast-tokenizer-for-tei", use_fast=False)

sentences = [
    "クエリ: 瑠璃色はどんな色？",
    "文章: 瑠璃色（るりいろ）は、紫みを帯びた濃い青。名は、半貴石の瑠璃（ラピスラズリ、英: lapis lazuli）による。JIS慣用色名では「こい紫みの青」（略号 dp-pB）と定義している[1][2]。",
    "クエリ: ワシやタカのように、鋭いくちばしと爪を持った大型の鳥類を総称して「何類」というでしょう?",
    "文章: ワシ、タカ、ハゲワシ、ハヤブサ、コンドル、フクロウが代表的である。これらの猛禽類はリンネ前後の時代(17~18世紀)には鷲類・鷹類・隼類及び梟類に分類された。ちなみにリンネは狩りをする鳥を単一の目(もく)にまとめ、vultur(コンドル、ハゲワシ)、falco(ワシ、タカ、ハヤブサなど)、strix(フクロウ)、lanius(モズ)の4属を含めている。",
]

token_ids = tokenizer(sentences, padding=False, truncation=False, return_tensors="np")["input_ids"]
token_ids = [t.tolist() for t in token_ids]

url = "http://127.0.0.1:8080/embed"
payload = {"inputs": token_ids, "normalize": False, "truncate": True}
headers = {"Content-Type": "application/json"}

response = requests.post(url, json=payload, headers=headers)
embeddings_data = response.json()
embeddings = np.array(embeddings_data)
print(embeddings.shape)

# calc cosine similarity
normalized_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
similarities = np.dot(normalized_embeddings, normalized_embeddings.T)

print(similarities)

結果

(4, 768)

array([[1.        , 0.94194159, 0.68661375, 0.71621216],
       [0.94194159, 1.        , 0.66622363, 0.68591373],
       [0.68661375, 0.66622363, 1.        , 0.87196226],
       [0.71621216, 0.68591373, 0.87196226, 1.        ]])

うまく密ベクトルが取得でき、ruri-base のモデルカードに記載されている値とほぼ同等のコサイン類似度が得られた。このような感じで、日本語TokenizerでもTEIの利用が(reranking以外)は可能だ。なお、当たり前だが、トークナイズしてるtoken_ids ではなく、普通のテキストを送ってしまうと、全く検討はずれの結果が返ってくるので注意が必要だ。

TEI に tokenizer.json がなくても起動でき、かつ /rerank API もうまく動くような Pull Requests を送るのが本質的な解決方法なのだけど、rust で実装し、PRで取り入れてもらうためのコミニュケーションのやり取りが億劫でできてないので、誰かやってくれると嬉しいなぁ…(他力本願)。

なお今回作ったモデル一覧は以下。

CrossEncoder の学習方法

CrossEncoder は、単純な回帰タスクである。query text[SEP]passage text といったSEPトークン等で区切ったテキストを、正例は1.0 負例は0.0としてラベル付けし学習させる。具体的な学習コード例としては、SentenceTransformers の CrossEncoder 学習サンプルが分かりやすい。

また、複数の負例(ハードネガティブ)を正例と同一バッチで学習させることで性能が大きく向上する。この学習方法については、FlagEmbedding の reranker trainerが参考になる。

学習用データセット

学習には、質問と正例・負例のデータセットが必要である。1件につき、positive(正例)1個とhard-negative(負例)15個を1セットとし、1グループ=16個として学習に用いた。以下のデータセットを利用した。

• JQaRA: dev/unused から 7,270件
• JSQuAD:
  • train から 62,859件
  • hard-negative マイニング用に wikipedia の文章から追加
• miracl: train の日本語データから 6,984件
• mmarco: train の日本語データからフィルターした 346,413件
• mr_tydi:
  • train の日本語データから 3,697件
  • なお miracl の日本語データには、このmr_tydiのデータと重複したデータが多く含まれる
• wikipedia リード文:
  • wikipedia のタイトルと、冒頭のリード文をペアとした 40,130件
  • hard-negative マイニングでは、同様に wikipedia のリード文のみを対象にマイニング

評価用データセット

モデルの評価には、以下のデータセットを用いた。

• JQaRA:
  • test 2000件
  • 評価指標は JQaRA での評価方法として定義されているNDCG@10
• JSQuAD:
  • validation 4442件
  • wikipediaからhard-negativeマイニングで negatives 19件追加し、合計20件からのMAP@10で評価
• miracl:
  • dev から negatives が9件以上のデータでフィルターした、704件
  • positive1件、negatives9件の合計10件としてMAP@10で評価
  • なお miracl に日本語データでは dev と train で一部データが被っており、train を学習すればするほど dev の評価が高くなりやすい
• JaCWIR:
  • eval 5000件
  • 評価指標は JaCWIR Reranker 評価方法として定義されているMAP@10

ハードネガティブマイニング

ハードネガティブとは、モデルが正例と誤判断しやすいが実際は負例であるデータを指す。これらを積極的に「マイニング」することで、学習データの多様性と難易度を高め、モデルの精度向上が期待される。

本モデルでは、BM25と複数のSentenceTransformerモデルを用いてハードネガティブをマイニングした。Semantic Textual Similarity（StS）タスクにより、正例に意味的に類似するが実際は負例である文章を抽出した。類似度の高いデータからランダムにサンプリングする方法を採用した。

学習元のpre-trainモデル

以下のpre-trainモデルを学習のpre-trainモデルとして利用した。なお BAAI/bge-reranker-v2-m3 については全件学習させると汎化性能が低下したため、mmarco, JSQuAD, wikipedia リード文を各1万件にランダムサンプリング(他のデータセットは全件)したデータで学習させた。

• japanese-reranker-cross-encoder-xsmall-v1
  • microsoft mMiniLMv2-L6-H384
  • 6 layers, 384 hidden size
• japanese-reranker-cross-encoder-small-v1
  • microsoft mMiniLMv2-L12-H384
  • 12 layers, 384 hidden size
• japanese-reranker-cross-encoder-base-v1
  • cl-nagoya/sup-simcse-ja-base
  • tohoku-nlp/bert-base-japanese-v3
  • 二つのモデルで学習させたものの統合モデル
  • 12 layers, 768 hidden size
• japanese-reranker-cross-encoder-large-v1
  • cl-nagoya/sup-simcse-ja-large
  • tohoku-nlp/bert-large-japanese-v2
  • 2つのモデルで学習させたものの統合モデル
  • 24 layers, 1024 hidden size
• japanese-bge-reranker-v2-m3-v1
  • BAAI/bge-reranker-v2-m3
  • 24 layers, 1024 hidden size

過学習への対応

CrossEncoderの学習を進める中で、ハードネガティブにwikipediaの文章を使用しているため、wikipediaデータを利用した関連のタスクの評価(JQaRA, JSQuAD, miracl japaneseなど)には最適化されるが、wikipedia以外のドメインでの汎化性能が学習すればするほど低下することが判明した。そこで、学習データに含まれないドメイン外のデータセットであるJaCWIRを作成し、バランスをとりながら学習・評価を行った。

結果として、1 epoch以上の学習では過学習が発生したため、学習は1 epochのみに制限している。

学習パラメータ

主のモデルの学習には、主に以下のパラメータを使用した。

• batch_size: 512 (gradient_accumulation)
  • 16個が1グループなので、pos,neg合わせて 512 * 16 = 8192 を1バッチで学習
• warmup_ratio: 0.25 (全体の25%をwarmupに使用)
• スケジューラ: cosine
• オプティマイザ: paged_adamw_32bit
• learning_rate:
  • xsmall = 2e-04
  • small = 5e-04
  • base = 8e-05
  • large = 3e-05
• loss
  • CrossEntropy

largeモデルを教師モデルとして使用

xsmall, smallの学習では、japanese-reranker-cross-encoder-large-v1とjapanese-bge-reranker-v2-m3-v1の推論出力を教師ラベルとして追加利用した。教師モデルの出力は推論値(例: pos=0.98, negs=[0.02, 0.07, ...])となるため、0と1だけでなく回帰タスクの連続値としての利用が可能である。教師モデルの出力データを用いることで、大幅ではないが若干のスコア向上が観測された。なお、この学習にはlossはMSEを用いた。

mixモデルの作成

学習データセットやスコアパラメータ、シード値を変更することで、多様な学習結果が得られる。これらの個別に学習したモデルを単純に線形結合することで、多様性を持たせパフォーマンスを向上させることができる。今回、複数の学習済みモデルを結合することでスコアの向上を確認した。なおモデル合成のツールにはLM_Cocktailを利用した。

注意事項としては、合成後のモデルは出力値の標準偏差が小さくなるため、量子化時等になんらかの性能劣化が発生する可能があるかもしれない。

評価結果

作成したCrossEncoderモデルの評価結果は以下の通りである。BAAI/bge-reranker-v2-m3は元々のマルチリンガル言語に対しての汎化性能が高く初めから日本語に対して高性能で、モデルサイズが問題にならなければ、少量のサンプル(数千件程度)でも微調整可能なため reranker 学習元の微調整モデルとしては最適と考えられる。

なお、この評価データセットのスコアはそのデータセットが公開しているtrainデータ等で学習することでスコアが高く出る傾向にある。今回作ったモデルは、JaCWIR 以外はtrain等のデータで傾向を学習しているため、その点も評価スコアを見る際には留意すると良いであろう。

なおこの文章は、私が書いたメモと指示を元に、Claude 3 Opus によって生成された文章を微調整したものである。

日本語に特化した形で学習されたRerankerがほとんど無かったので、日本語を適切に学習させた Reranker ファミリーを作りました。小さいモデルから大きなモデルまで揃っています。

評価性能は以下の通りで、現在(2024年4月頭)に公開されているRerank日本語タスクにおいては最高性能かな、と思っています。なぜなら日本語を学習させたRerankerがほぼ公開されていないから…。

なお、今回作ったRerankerの技術的な話は、日本語 Reranker 作成のテクニカルレポートに記載しているので、興味のある方はそちらをご覧ください。

そもそも Reranker とは？

Reranker とは、名前の通り再ランク付け(rerank)するもので、質問文に対して関連する順に文章を並べ替えます。文ベクトル(文章のembeddings)で類似度を測って並べ替えするものと何が違うのか？と思われるかもしれませんが、実際、文ベクトル類似度でも同じように並べ替えが可能です。

しかしながら、大きく二つの点で異なります。

Reranker は再ランク性能が高い

文ベクトルは、質問文と文章を同じベクトル空間上の表現として類似度を測ります。そのため大規模なデータに対しても事前に文章のベクトルを算出しておくことで、効率的な計算が可能です。

しかしながら、Reranker は再ランクに特化しており、例えば今回作ったモデルは CrossEncoderというアーキテクチャを用いて質問文と文章を一つのペアにして評価することで、より細かなニュアンスや質問と文章の関連性からの文脈的理解を行えます。そのため質問に関連する文章がより上位になりやすく性能が高いです。

Reranker は事前計算ができず遅い

精度が高いなら、文ベクトルなど作らずに全部 Reranker で評価すればいいじゃない？と思われるかもしれませんが、Reranker は質問と文章両方を入力に使います。文ベクトルなら、対象となる文章のベクトルをオフラインで事前に計算することができるため、検索時には質問文のベクトルだけ計算すれば、それを元に検索が可能です。

しかしながら、Reranker (CrossEncoder)は文章のみを事前に計算しておくことができないため、例えば対象文章が100件のデータならオンラインでその場で実行時に100件分全て計算しても問題ない計算量ですが、件数が増えるにつれて現実的な速度では検索できなくなります。

Reranker の使い所

とすると、現実世界の検索では Reranker の使い所がないのでは、と思われるかもしれません。そこで、まず文ベクトルなどの効率よくオフライン計算ができる手法で質問に関連する文章上位100件を抽出し、その後Rerankerでその100件をより効率よく並べ替えすることで精度を上げる、といった用途で活用できます。

Retrieve & Re-Rank Pipeline より引用

GCP の記事、Your RAGs powered by Google Search technology, part 2 などでも、Deep re-ranking の項で同じ手法について書かれているように、再ランク付は重要です。

実際に Reranker で再ランクを行うと、どれぐらい情報検索スコアが変化するのか

以前書いた記事、 ベクトル検索のみで、AI王クイズ第一回コンペに臨む - Q&Aタスクでの複数の日本語embeddingsの評価では約550万の文章を各種文ベクトルモデルで文ベクトルに変換後、IVFPQアルゴリズムを使い、近似最近傍探索で検索評価を行っています。これらの各種文ベクトルモデルで正解率を出した結果は以下です。

• jaqket v1 ベクトル検索 - 日本語embeddings評価

ではこの近似最近傍探索での検索結果の上位100件を用いて、今回作ったモデルの中で、最小最速のRerankerモデルの xsmall で再ランク付けしてみましょう。

かなり結果が向上したのがわかると思います。550万の文章から現実的な速度で検索するために近似最近傍探索を行っており、通常の文ベクトルの総当たり類似度検索よりも精度が落ちていることもあって、Rerankerモデルで再ランクすることによって大幅なスコア向上となっています。

また、例えば OpenAI の text-embeddings は、日本語の情報検索タスクではあまりスコアが芳しくないことが多いのですが、それらも再ランクすることで大幅にスコアが上がっていますね。

では続いて、大きなlargeモデルのRerankerで再ランクしてみましょう。

こちらはさらに大幅にスコアが上がっていますね。計算機リソースに余裕があれば大きなモデルを使うのは良いのですが、モデルサイズによってどんどん再ランクにかかる速度が増えていきます。各々のモデルで、JaCWIR の評価にかかった実行速度(GPU RTX3090 で実行)は以下です。

xsmallとlargeでは、6倍ほど速度に差が出ています。このように、性能と速度とのトレードオフが発生するので、どれぐらいの性能と速度が必要かを考えて Reranker を選ぶ必要があります。実行時に処理するRerankerは処理速度が重要なケースも多いでしょう。

なお他の様々なモデルとの評価結果については、日本語 Reranker 作成のテクニカルレポートをご覧ください。また本記事でのAI王クイズコンペの再ランク評価は、評価にtestデータを用いているため直接学習はしていないものの、AI王クイズコンペのdev, unused を用いたデータセットJQaRAのデータも今回作成したモデルは学習しているため、スコアが上がりやすい傾向であることに留意ください。

意外と大事な Reranker

今回日本語の Reranker を作ったきっかけは、数百万〜の文章に対して情報検索をしていると、文ベクトル+近似最近傍探索のみの検索よりも、Reranker を組み合わせた方がだいぶ良い検索結果になったためです。この時に使った Reranker はマルチリンガルモデルの cross-encoder-mmarco-mMiniLMv2-L12-H384-v1 だったのですが、マルチリンガルでだいぶ精度が上がるなら、日本語をちゃんと学習させればさらに精度が上がるのでは？と思い立ったのがきっかけです。

またRerankerは、良くも悪くもオンライン計算が必要になります。悪い点は計算コストが高い点ですが、精度以外の良い点としては事前計算を再計算しなくて良いことも挙げられるでしょう。例えば文ベクトルモデルでより良いモデルを適用したくなっても、データベースにすでに文ベクトルデータが事前計算され格納されているため、本番環境で利用されている文ベクトルの変更は慎重に行う必要がありますし、数億データ〜ともなってくると全て再計算するのにも計算機コストがかかります。しかしながらRerankerは、ソートアルゴリズムの変更のようなもので、事前計算データの変更もなく適用が可能なので、差し替えがしやすいです。

さらに、Reranker は解きたい課題のドメインのデータで学習させると、性能・スコアがかなり上がることも観測しており、そのために文ベクトル変換は汎用モデルを、Rerankerはドメイン特化モデルを、といった使い分けもできるでしょう。

というわけで、日本語を学習させたRerankerの作成と、Reranker はどのようなものか？についてご紹介しました。世の中はLLMの学習と利活用にフォーカスしている感がありますが、個人的にはLLMの利活用が進むにつれ、検索を人間ではなくAIに最適化する時代が訪れ情報検索の分野の重要性がさらに増すのでは、と思っています。

情報検索をよりよく行う手段の一つとして、Reranker は欠かせないものになってくるでしょうし、本記事でRerankerや情報検索も面白そうだぞ、と興味を持たれる方が少しでも増えたら幸いです。

なおこの文章は、私が書いた草稿をもとに、Claude 3 Opusによって生成した文章を微調整したものです。

昨今、生成AI・LLMの台頭により、「良質なデータ」をどう集める・作るかの話をより身近で聞くようになった。LLMに学習させるデータは元より、身近な課題を解決するためにも課題解決のためのタスクを定義し、そのためにデータを分析し作成することが、社会課題解決には当たり前に求められる。

これらの課題解決には、新しいモデルを自ら考える必要がないことも多く、タスク定義とデータを集め学習させるだけで、十分な性能を発揮することも多い。いわゆるデータセントリックなデータに焦点を集めた考え方である。

ただ、世の中にはモデルやアルゴリズムといったモデルセントリックな話は数多くあれど、データは基本公開されているなんらかのデータセットに対して評価するといった内容はほとんど。しかしながら、この書籍Human-in-the-Loop 機械学習では、データに焦点を集め解説を行なっているという、稀な書籍である。

どんな内容かは目次や反響まとめを読んでもらうのが一番なので、そちらを見ると良いだろう。個人的に興味をとても惹かれた箇所は、一つ目は能動学習でのデータのサンプリング方法である。わかりやすい線形の決定境界付近のデータ(確信度が低い)をアノテーションすべき、はすぐに思いつくことだが、不確実性や多様性をどのように解釈し、どういう戦略でアノテーションすべきデータを考えるべきか、多様な視点から述べられていて、思わず「ハッと」するような脳を刺激されるアイディア（実装も）が詰まっている。

二つ目は、アノテーションをつけるアノテータとの協業の話である。「ピープルマネージメントは必須」と書いてあるように、はいこんな感じでラベルつけてねあとよろ〜、などでは全くなく、どのように依頼すべきか、必要なスキルは何か、フィードバックをどうすべきか、アノテータのバイアスをどう取り除くか、アノテータごとの不確実性にどう向き合うか、コミニュケーションとフィードバックをどうすべきか、とまさしく通常の仕事のピープルマネジメントと同じようなことを、アノテータともすべきということを強く書いてある。もちろん、ピープルマネージメント以外のヒントもたくさん（バイアスとかね）。

等々、読んでいて参考になる場所だらけで、じっくり読み進めていたら読み終えるまで2ヶ月ほどかかってしまった。それほど興味深く、かつ機械学習を嗜む身としては、この書籍に出会うことができて本当によかった。

LLMの台頭が始まった今のAI時代、AIに評価をさせてそのフィードバックをもとに良質な学習データセットを作る、強化学習をしていくといったことが当たり前になっていくであろう。Human-in-the-Loop 機械学習の原著は、いわゆるGPT4登場より前のLLMが今ほど注目が集まっていなかっときに書かれた本だが、この書籍の視点はLLMと絡めて使う時にも大いに役立つと思っている。

• https://docs.google.com/spreadsheets/d/1eSYzxzIfN3uMIpFKDGCTQsIxuWYELBtD49LQbl88GUE/edit#gid=140790548

結果としては、multilingual-e5-large よりわずかに低い、という結果となった。学習データも少なく、モデルサイズも12層のBert(= multilingual-e5-smallと同じサイズ)とほとんど一緒なのにね、すごい。

ColBERT の実装と論文を読む

というわけで、ColBERT に興味が湧いたので、論文と実装を読んでみる。

• 読んだ論文
  • ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT
  • ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction
  • JaColBERT and Hard Negatives, Towards Better Japanese-First Embeddings for Retrieval: Early Technical Report
• 実装
  • https://github.com/stanford-futuredata/ColBERT
  • https://github.com/bclavie/RAGatouille

ColBERTはいわゆるSentenceTransformer等を用いて文からEmbeddings(文ベクトル)を出力して類似度比較で検索、ではなくトークンベースの類似度検索だ。文の最終隠れ層は、その文章のコンテキスト文脈の情報をトークン単位でも持っているので、文全体ではなくトークン単位で使って類似度を算出する。

類似度計算は、queryとdocumentのトークンの出力のコサイン類似度をとって、最大のものを足し合わせるという MaxSim という計算手法。MaxSimの計算自体は単純である。

query と document は別にエンコードする必要があるのだけど、使っているモデルは Bert (12層) + 独自HEAD (linear 128次元)。Bert 12層の隠れ層出力は768次元なので、それを128次元へと線形層を通して変換している。

実実装では、queryとdocumentは、query の場合は prefix に [CLS][unused0] を、document の場合は　[CLS][unused1] といった形で CLS の後に判断のための独自トークンをつけて区別をしているだけで、エンコーダ自体は同じである。

これらでエンコードした結果を、MaxSimをとることで、一番数値が大きいdocumentがqueryに似ていると判定する。なおColBERT実装では、documentのtokenのうち、記号やpadding tokenなどは無視(mask)して計算していた。

検索時のパフォーマンス問題の解消

普通の文ベクトルを探す手法なら、ANNを始めとした近似最近傍探索を使えば、数億以上の文章からも高速な検索が可能だ。しかしながらColBERTは文ベクトルではなく、トークン同士の類似度を用いたMaxSim計算なので、この手法はママ使えない。

という問題に対応するために高速に近傍探索ができるindexを作る方法が載っているのが論文ColBERTv2。あれこれとベクトルを情報圧縮し、Kmeansでcentroidを求め、そこから探索しているようだ。実装ではfaissをimportしていたので、そのままfaissのindexを使っているのかな、と思いきや、faissはKmeanでクラスタ中心点を計算しているのに使っているだけであった。一度インデクスを作ってしまえば、その後の検索はサクサク。

なお、前述のAIクイズ王のQ&A RAGタスクの550万Passageをこの方法でindex作成するのに、Ryzen 3900 + RTX 3090の環境で5時間前後であった。faissはgpu利用版のfaiss-gpuでないと、かなり遅いので注意だ。

省データで学習可能

JaColBERT のレポートによると、bert-base-japanese-v3を元に、1000万のトリプレットデータとNVidia L4 GPU * 8 の省リソースで10時間学習させて作ったとのこと。データ量が少なく、かつ学習時間が短いとなると、夢が広がる結果である。

ColBERT の問題点・利用が大変

ColBERT の実装は読んでみて分かったが、処理の複雑性とそもそもの実装コードの複雑性も相まって、サクッと使うのが難しい。それを解消するために、ゼロコンフィグでもすぐに使えるアプローチが RAGatouille である。今回の評価にもRAGatouilleを利用した。

RAGatouille はインデクスの作成・検索はもちろん、Trainerで学習させることもできる。また、LangChain の retriver にすることもできるなど、今風の対応も入っている。

ColBERT は本番環境の検索サーバで使えそうか？

ColBERT の気になりどころとしては、そもそも本番で運用できそうかというところがある。今(2024/02頭)時点では、自前で検索用APIサーバ実装して運用が前提になり、手軽に動かすのがなかなか大変そうである。

しかしながらRAGatouilleのドキュメントによると、近日中に検索エンジンのVespaが対応するっぽいことが書かれているので、そうなるとだいぶ運用が楽になりそうである。またIndexのデータ追加もまだ試験的なようで、その辺も問題なくできるようになれば、本番環境でも十分検討できそうである。

何より、自前ドメインのデータを省コストで学習させることができるなら、自前データに対して高品質なRAG用のデータ取得検索として使えそうなので、そのようなユースケースなら積極的に検討したい。

おわりに

というわけで ColBERT すごい！というお話でした。情報検索を研究している人からすると ColBERT は既知の話かとは思うが、自分は知らなかったので調べてみて新鮮だった。また、JaColBERT が無ければそもそも興味がわかなかったと思うので、作者の Benjamin Clavié 氏に感謝。なお氏は RAGatouille の作者でもある、ありがたい。