検索タスクにおいて、クエリと文書の意味的な類似性を捉えるニューラルネットワークを用いた検索は重要な技術である。しかし、従来のBERTなどの言語モデルは、主にトークンレベルのタスクで事前学習されており、文レベルの表現力が十分に発達していないという課題があった。この課題を解決するため、検索に特化した事前学習の新しい手法として RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder が提案されている。

本記事では、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上で公開しており利用可能となっている。本記事が、日本語の検索タスクの性能向上に貢献できれば幸いである。