本記事は情報検索・検索技術 Advent Calendar 2024の6日目の記事です。

【2025年4月15日追記】
最大系列長8192に対応し、性能が大幅に向上したv3モデルを公開しました！
特に検索タスクにおける性能が向上しています。

【2025年3月20日追記】
v2モデルを公開しました！

hppです。普段は自然言語処理(Natural Language Processing: NLP)分野の研究をしており、特に文やテキストの埋め込み表現について色々やっています。
本稿では、そんな研究の一環として9月に公開した日本語テキスト埋め込みモデルRuriについて紹介します。

テキスト埋め込みとは

NLPをはじめ多くの深層学習分野では、深層学習モデルへデータを入力する際に、データをベクトル、つまり数値の列に変換します。
大雑把には、このベクトルのことを 埋め込み表現 (Embedding) といいます。

ChatGPTなどの(Transformerに基づく)言語モデルも、実態としては単語の埋め込み(ベクトル)の列を入力として、単語の埋め込み列を出力する機構になっています。
具体的には、入力された文を(例: 私は犬です。)をトークンのID列(例: [私, は, 犬, です, 。]→[7, 5, 18, 14, 2])に変換し、IDごとに対応するベクトルを割り当て、そのベクトル列を入力にベクトル列を出力するという機構になっています。
そして、出力されるベクトル列の取り扱いを良い感じなんやかんやすることで、単語予測などを行っているわけです。

このようにデータをベクトルとして取り扱う手法は、もはや深層学習では不可欠と言えるまでになっています。
特にNLP分野における埋め込み表現としては、単語の埋め込み表現、すなわち単語埋め込みが古く(といっても2012年以降)から研究されていました。
単語埋め込み研究としては、MikolovらによるWord2Vecが非常に有名です。

単語の埋め込み表現はWord2Vecでいい感じに作れることがわかりました。
しかし、自然言語をうまく表現するには単語だけで十分でしょうか。
はい、もちろん十分ではありません。
自然言語にはフレーズ(複単語表現)や文、文章など、単語を構成要素としながらも、粒度を異にするいくつかのまとまりがあります。
もちろんこれは単語の列から成っているので、単語埋め込みの操作でこれらの埋め込み表現を作ることも考えられます。
しかし、単純な算術演算では自然言語が織りなす複雑な意味変化(例えば語順の変化やニュアンスの違いなど)を正確に捉えることはできません。
そこで、これらの表現をうまいこと作ってあげようという取り組みをしているのが テキスト埋め込み の研究です。

その名称に反して実は、テキスト埋め込みとは近年新しく確立されつつある研究対象です(筆者の感想)。
以前までは「フレーズ」や「文」、「文章」といった異なる粒度ごとに深層学習モデルを構築してあげないと、まともなモデルを構築することができませんでした。
最近になって、強力な事前学習済み言語モデルや学習手法が整ってきたこともあり、ようやくこれらの対象をひとまとめにして扱えるようになってきたという状況です。

テキスト埋め込みの応用分野としては、FAQ・類似文検索や、文書検索、高速なテキスト分類などいろいろ存在します(自分ではやったことがないですが、特徴量抽出器としても使えそうです)。

現在、テキスト埋め込みの研究は英語圏を中心に行われています。
その多くは、以下のステップでテキスト埋め込みモデルを構築します:

事前学習済み言語モデルを用意する
大量(数億)のテキストペアを収集する
対照学習(Contrastive Learning)と呼ばれる手法を使い、2でモデルを訓練する
人手で作成された高品質なテキストペアを収集する
対照学習を用い、3でできたモデルをさらに改善する

上記ステップにて、1については、Llamaのような大規模言語モデル(Large Language Model: LLM)が利用されることもありますが、テキスト埋め込みモデルはあまりに大きいと運用が大変なので、BERTのように比較的(現代では)小規模なモデルを利用して、小さいモデルをみっちり訓練するという手法が取られる場合が多いです。

2について、大量のテキストペアというのは、Wikipediaの記事タイトル--記事本文のような「まぁ関連があるかも...」というペアを指します。
こういったペアはWeb上に大量に存在しているので、これらを収集して訓練に利用します。

テキスト埋め込みモデルの訓練では、対照学習と呼ばれる手法が取られます。
対照学習は、ざっくりいうと「正解ペアの埋め込み表現同士が近づくように」しつつ「不正解ペアの埋め込み表現同士が遠ざかるように」訓練を行う手法です。
意味的に似たベクトル同士は近づいてほしくて、意味的に違うベクトル同士は遠ざかってほしい、という気持ちで学習をします。

3のステップでは対照学習を弱教師データを用いて大規模に行います。
5のステップでは対照学習を高品質な教師データを用いて小規模だけれども丁寧に行います。
これにより、いろいろなテキストの意味的関連性を理解した強力なモデルが出来上がります。

さて、英語圏や中国語圏ではこういったモデル開発は盛んなのですが、日本語モデルでこういった取り組みは全然ありませんでした。
理由としては以下がありそうでした:

評価データが足りない
計算機が足りない
訓練データが足りない
訓練ノウハウが足りない

これらのうち、1の問題については今年度はじめくらいにJMTEBと呼ばれる日本語テキスト埋め込みのための評価ベンチマークが公開されたことで、ある程度解消されました。
2については、ありがたいことに筆者が所属している研究室のリソースがあればなんとかなりそうでした。
一方で、3や4は一度やってみないとどうにもならなさそうです。

というわけで、非常に前置きが長くなったのですが、実際に日本語でテキスト埋め込みモデルを構築してみたので、その取り組み内容について紹介します。

Ruri

まず、構築した日本語テキスト埋め込みモデルRuriの関連リンクを一覧しておきます。

本稿の内容は、実はテクニカルペーパーの方に詳細に記述されています。
本稿はテクニカルペーパーの内容を適宜抜粋し、いい感じに再構成したものです。
もしRuriを論文やブログ等で引用していただける際は、テクニカルペーパーの方を参照いただけると幸いです。

さて、Ruriに関する取り組みは大まかに以下の3つの段階に分けられます:

データ整備・合成データ作成
Contrastive Pre-training: 大量データで対照事前学習
Fine-tuning: 高品質データで丁寧に微調整

Ruriの構築プロセスを概観すると以下のようになります:

順に述べていきます。

データ整備・合成データ作成

強力なテキスト埋め込みモデルを作る上で避けては通れないのが、高品質なデータを用意することです。
そこでまず、散らばっていた既存のデータセットをまとめ、埋め込みモデルの訓練に簡単に利用できるよう、前処理の実施とフォーマットの統一を行いました。
そうして出来上がったデータセットが以下のデータセットです:

https://huggingface.co/datasets/hpprc/emb

上記データは自分用に整備したものだったのですが、副次的な効果として、ありがたいことに上記を利用した取り組みがすでにいくつかの個人・組織から公開されています:

データの整備を終えて、訓練に使えそうなデータを英語圏の埋め込みモデルと比較した時に、日本語の訓練データとして特に足りていなかったのが検索データでした。
検索データは作るのに大変な手間がかかるので、日本語で自由に利用できるデータは非常に少なく、心許ない状況でした。

そこで今回はさらに、足りない検索データをLLMを用いて検索データを自動的に生成するというアプローチを取りました。
これはGoogleのGeckoというテキスト埋め込みモデルでも利用されていたアプローチです。
このようにLLMから作成されたデータは合成データ(synthetic data)と呼ばれ、埋め込みモデルの訓練に効果的であることがわかっています(宣伝: 最近書いた共著論文)。

検索データを合成するため、今回は日本語Wikipediaを入力として、検索クエリ(質問)とその回答を生成するようにLLMにプロンプトを入れました。
今回利用したLLMは東工大が公開しているSwallow-MXです。
今回構築する埋め込みモデルは、できるだけ寛容なライセンスのもと公開したかったので、出力のライセンスに特段の制限がないApache 2.0ライセンスで公開されているモデルを利用しました。

そして出来上がったデータセットがAutoWikiQAです。
事例数238万と日本語の検索・QAデータとしては最大規模のものを作成することができました。
もちろん半自動で作成しているのでノイズは含まれているのですが、その点に気をつけて利用すれば十分性能に貢献してくれそうです。

また、より大規模で賢いモデルを利用した合成データセットとして、NVIDIAが公開しているNemotron-4-340Bを利用した以下の二つのデータセットも作成しました。

構築したこれらのデータセットはCC-BY-SA 4.0ライセンスで公開されているので、商用・非商用の別なくライセンスに基づく限り自由に使用していただけます。

上記合成データと、既存のデータを組み合わせてContrastive pre-training (対照事前学習)用のデータを構築しました。
既存のデータについては、Ruriは汎用的なテキスト埋め込みモデルになって欲しかったので、ドメインの偏りがないよう、様々なデータを収集しました。
その結果、合計で約8800万(88M)事例のテキストペアを用意することができました。

注意点ですが、英語で同様の取り組みをしているE5は約2億7000万(270M)のデータを利用しており、それと比較するとかなりデータ量は少なめです。
今後、データ量を増強することでさらに高性能なモデルを構築できる可能性はありますので、同様の取り組みをされる方は参考にしていただけると幸いです。

Contrastive Pre-training

事前学習用データが出来上がったので、対照事前学習を行います。
対照学習の際にはin-batch negativesを利用します。
これは、ミニバッチ内のある事例について、同じミニバッチ内の他の事例を負例として扱うことで、手間をかけずに対照学習における負例を増やす手法です。

汎用的なテキスト埋め込みのための対照学習では、非常に巨大なバッチサイズを利用するのが一般的です。
具体的には、8192, 16384, 32768がよく利用されます。
事前対照学習では、このように巨大なバッチサイズを用いることで性能が向上するという報告がされています。

今回は、学習環境がA100 (80GiB) x 4だったこと、既存の多言語埋め込みモデルであるmultilingual E5 (mE5)が系列長192を採用しているがそれよりは長い256で訓練したかったことから、バッチサイズとして8192を選択しました。

余談ですが、E5と異なり、Ruriの事前学習では各事例ごとにhard negativesを1事例追加して学習を行いました。
通常、contrastive pre-trainingの際には、hard negativesは利用しません(E5では利用したが計算コストに見合う性能向上が得られなかったとしています)。
しかし、Ruriの事前学習データは幸か不幸か量が少なかったため、無理なくhard negativesを追加することができました。
hard negativesの取得(hard negative miningという)には、多言語E5を用いてanchor(検索クエリやニュースタイトル)の埋め込み表現を作成し、正解ではないが近傍ではある事例を類似ベクトル検索することで行いました。

今回はsmall, base, largeの3種類のモデルを訓練したのですが、このhard negativesのせいでlargeモデルの学習において系列長を短くしなければいけなくなりました。
また事後的な検証でわかりましたが、そもそもhard negativesがない方が性能が良かったです(悲しい)。

番外編: Reranker

実用的な検索タスクにおいて、文書検索はテキスト埋め込みモデルのみを持って実行されるのではありません。
BM25やその他複数のアルゴリズムによって取得した関連度が高いと思われる文書集合に対し、検索クエリと各文書とのより詳細な関連度を計算するRerankerというモデルを用いて、被検索対象文書を並び替えることがあります。
基本的には、検索対象文書をいい感じに並べるという作業において、埋め込みモデルよりもRerankerのほうが性能が高いです。
その代わり、Rerankerは大量の文書を並び替えることができません。

実はこのRerankerは、実用時だけでなく、埋め込みモデルの学習時にも有用です。
Rerankerが並び替えた順序と同じ順序に並び替えられるように埋め込みモデルを訓練することができます。
具体的には、Rerankerが出力したクエリ--文書の関連度について、テキスト埋め込みモデルから出力された埋め込みをもとに計算したクエリ--文書の関連度を、近づけるように知識蒸留で学習します。
Cross-Encoderからの知識蒸留(knowledge distillation from a cross-encoder)と呼ばれ、いくつかのテキスト埋め込み手法が取り入れています。
Cross-Encoderというのは、一つのモデルにクエリ--検索文書の両方を入れるモデルのことで、逆にクエリと検索文書を別々にモデルに入力して、後から埋め込み由来の類似度を測るようなモデルはDual-Encoderと言います。
テキスト埋め込みモデルは基本的にはDual-Encoderですね(見た目は完全に単一モデルですが、weightをshareしている2つのモデルがあるという見方をする)。

対照学習損失と知識蒸留損失のイメージ図は以下の通りになります。

対照学習損失は正解以外が一様に0になるように学習するのに対し、知識蒸留損失はRerankerが出力した複雑な分布を模倣するように学習を行います。

今回は、この手法を試してみたかったのと、それとは関係なくRerankerはnoisyな検索データのフィルタリングにも有用であることから、Rerankerも合わせて構築しました。

詳細はテクニカルレポートをご覧いただければと思いますが、対照事前学習を行ったモデルをベースに、先述の合成データも交えたRerankerの訓練を行うことで、既存のモデルより高い性能のRerankerを構築することができました。

モデルは全て商用利用可能なApache 2.0ライセンスで公開していますので、合わせてご覧ください。

Fine-tuning

さて、事前の準備が終わり、いよいよ最終的なテキスト埋め込みモデルを訓練します。
訓練データには、既存の高品質なデータを用いました。
また、hard negative miningを実施し、検索データセットについては1件あたり15の負例データを付与しました。

学習については、先述の通りRerankerを用いた知識蒸留を行いました。
fine-tuningでは通常の対照学習損失と知識蒸留損失を組み合わせて学習を行います。
この組み合わせ方にはいろいろな種類がありますが、日本語ColBERTの検証で有用だと説明されていた「検索データセットでは知識蒸留損失のみ、その他のデータセットでは対照学習損失のみ」を利用して学習を行う方式を採用しました。

評価には先述のJMTEBを用いました。

評価・利用時の注意点としては、Ruriはprefixを利用するモデルであることが挙げられます。
prefixはE5やその他現代的なモデルで採用されている「マーカー」のようなものです。
このマーカーを使って、入力文を「どのように埋め込むか？」の目印にします。

E5など英語圏のモデルは、検索クエリなど短い文のprefixとして「query: 」を利用し、検索対象文書など長いテキストについては「passage: 」を利用しています。

これに倣い、Ruriは短い文に対して「クエリ: 」を、長い文章に対して「文章: 」をprefixとして利用するよう訓練しています。

評価結果の一部を示したものが以下です(注意: 左側に平均性能)。

既存モデルとしては多言語E5が非常に強かったのですが、それを上回る性能を達成することができました。
largeモデルの性能の強さもさることながら、Ruri-smallがmE5-largeと同等の性能を示していることも驚きです。
総じてかなり強いモデルを作ることができたと思います。

ここからまた余談ですが、事後的な検証にて、実は知識蒸留損失は使わず単に対照学習した方が強いという結果が出ています。
もちろん、諸々のチューニングをした結果、やっぱり知識蒸留した方が良いという結論になる可能性はあるのですが...まだまだモデルの性能には改善の余地があるということでお考えください。

まとめ

さて、非常に長くなってしまいましたが、まとめです。

日本語汎用テキスト埋め込みモデルを構築しました
主要な貢献は以下のとおりです
1. 埋め込みモデル訓練のための大規模合成データセットを作成しました
2. 大規模な対照事前学習によって強力なベースラインを構築しました
3. 日本語評価において最高性能のRerankerを構築しました
4. 日本語評価において同規模のモデルと比較してかなり性能の高いテキスト埋め込みモデルを構築しました

より詳しい内容はテクニカルレポートをご覧ください。
また、実際に使ってみたい方は以下に示すHuggingFaceのリポジトリから簡単に使っていただけます:

最後までご覧いただきありがとうございました。
何か質問や感想等あれば、コメント欄や筆者のXまでお寄せください。

Discussion