🫧
ModernBERTを使ったベクトル検索モデルを作ってみよう

2024/12/22に公開
こんにちは、@makiart13です．

普段は筑波大学で情報検索に関する研究を行っています．
この記事は情報検索・検索技術 Advent Calendar 2024 の21日目の記事になります！
https://qiita.com/advent-calendar/2024/search
当日にテーマを決めて記事を書き始めました！書き終わることを願っています！！

 概要数日前にAnswer.AIとLightOnという企業がModernBERTという新しいBERTをリリースしました．

このModernBERTは，簡単にいうと 早い・上手い（高性能）・長い（コンテキスト長） モデルです!
このModernBERTを利用して，ベクトル検索モデルを作ってみようという記事になります！

 ModernBERTとは何か？https://x.com/jeremyphoward/status/1869786023963832509
前節でModernBERTを 早い・上手い（高性能）・長い（コンテキスト長） と説明しました．もう少し深ぼってみましょう．

 🚅 早い以下のような手法や技術を取り入れたことによって，より効率的に学習・推論できるようになっています．

Alternating Attention
Global Attention と Local Attention の両方を利用し，計算負荷を減らしつつ，効率的な処理を実現


Flash Attention
高速なAttention処理を実現するFlash Attentionを利用して，メモリ効率と処理速度を向上


Unpadding
BERT等で使われるPaddingの仕組みをなくし，できるだけ一つのシークエンスで複数のデータを処理するように


 😋 上手い（高性能）
大規模データによる学習
2兆トークン
ウェブ文書などの幅広なソースを元に学習


Positional Embeddingの手法変更
RoPE（Rotary Positional Embeddings）という長いコンテキストに対して効果的に位置埋め込みする仕組みを利用


トレーニング方法の工夫
バッチサイズや学習対象のコンテキスト長を動的に変更


 📏 長い（コンテキスト長）個人的に特にアツい部分
コンテキスト長が8,192トークンもあります！
様々なベクトル検索モデルで利用されているBERTやRoBERTaのコンテキスト長は最大512トークンです．
そのため，大半のローカルで利用できるベクトル検索のモデルは長いテキストを一つのベクトルに変換することができず，小さなチャンクに分けてベクトルにすることをしていました．
しかし，ModernBERTでは8,192トークンあります．そのため，各センテンスに跨った情報を含んだ一つのベクトルを生成することが可能になり，近年のトレンドであるRAGにおける検索において大変有効であると考えられます．

 より詳しく知りたい人はこちら詳細なModernBERTの説明や解説については下記が参考になると思います．
https://huggingface.co/blog/modernbert
https://zenn.dev/dev_commune/articles/3f5ab431abdea1

 ベクトル検索モデルを作ってみようSentence Transformerの公式ドキュメントを参考にベクトル検索モデルを学習していきます．
https://www.sbert.net/docs/sentence_transformer/training_overview.html
今回は時間がなかったので，ModernBERTを含む4つのモデルを簡単に学習させ，評価していきます．
追記：

MSMARCOのデータセットを利用して学習するサンプルコードがありました！

こちらのコードを参考に学習すると良さそうです！
https://github.com/AnswerDotAI/ModernBERT/blob/main/examples/train_st.py
上記のコードを参考に私の方でModernBERTをMSMARCOのデータセットを使って学習させたモデルを載せておきます！
https://huggingface.co/makiart/ModernBERT-large-DPR-msmarco
https://huggingface.co/makiart/ModernBERT-base-DPR-msmarco

 モデルの学習Sentence-Transformerを使うととても簡単に学習を行うことができます．
ここでは，sentence-transformers/all-nliを利用して，ベクトル検索モデルを作成していきます．
今回は，GPU環境の都合上Flash-Attentionを使わず学習しました．
また，24/12/21現在でModernBERTを使うにはtransformersをgitからインストールしなくてはいけないので注意が必要です！
pip install git+https://github.com/huggingface/transformers.git
# or
uv add git+https://github.com/huggingface/transformers.git
以下が学習コードになります．今回はBERT,RoBERTa,ModernBERT(base,large)を学習させてみました．
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    models,
)
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import (
    BatchSamplers,
    SentenceTransformerTrainingArguments,
)


def train_embed_model(model_name: str) -> None:
    """モデルの学習に使用するSentence-Transformerを構築する関数"""
    # 1. トランスフォーマーモデルのロード
    # 指定されたモデル名（例: 'bert-base-uncased'）を元に、トランスフォーマーモデルを初期化
    word_embedding_model = models.Transformer(model_name)

    # 2. プーリング層の設定
    # トランスフォーマーモデルの出力次元を取得し、それに基づいてプーリング層を作成
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())

    # 3. SentenceTransformerモデルの構築
    # トランスフォーマーモデルとプーリング層を組み合わせてSentence-Transformerモデルを作成
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

    # 4. データセットのロード
    # "sentence-transformers/all-nli"データセットをロードし、学習用と評価用に分割
    dataset = load_dataset("sentence-transformers/all-nli", "triplet")
    train_dataset = dataset["train"].select(range(100_000))  # 学習データから上位10万件を選択
    eval_dataset = dataset["dev"]  # 評価データ

    # 5. 損失関数の設定
    # MultipleNegativesRankingLossは、類似度学習に適した損失関数
    loss = MultipleNegativesRankingLoss(model)

    # 6. 学習パラメータの設定
    args = SentenceTransformerTrainingArguments(
        output_dir=f"outputs/models/{model_name}",  # モデル保存先ディレクトリ
        num_train_epochs=1,  # エポック数（今回は1エポックのみ）
        per_device_train_batch_size=64,  # 学習時のバッチサイズ
        per_device_eval_batch_size=64,  # 評価時のバッチサイズ
        warmup_ratio=0.1,  # 学習率スケジュールのウォームアップ比率
        fp16=True,  # 半精度浮動小数点（FP16）を有効化（GPUが対応している場合）
        bf16=False,  # BF16は無効化（必要に応じてTrueに設定）
        batch_sampler=BatchSamplers.NO_DUPLICATES,  # バッチ内で重複データを避ける設定（in-batch negatives用）
        eval_strategy="steps",  # 評価タイミングをステップ単位で指定
        eval_steps=100,  # 評価間隔（100ステップごとに評価）
        save_strategy="steps",  # モデル保存タイミングをステップ単位で指定
        save_steps=100,  # 保存間隔（100ステップごとに保存）
        save_total_limit=2,  # 保存するモデル数の上限（古いモデルは削除）
        logging_steps=100,  # ログ出力間隔（100ステップごと）
        report_to="none",  # ログ出力先（今回は無効化）
        # run_name="my-first-ft-modern-bert-emb",  # W&Bなどで使用する実験名（必要ならコメント解除）
    )

    # 7. 評価器の設定
    dev_evaluator = TripletEvaluator(
        anchors=eval_dataset["anchor"],  # アンカーデータ（基準となる文）
        positives=eval_dataset["positive"],  # ポジティブデータ（類似する文）
        negatives=eval_dataset["negative"],  # ネガティブデータ（異なる文）
        name="all-nli-dev",  # 評価結果の名前
    )
    
    dev_evaluator(model)  # 初期状態で評価を実行

    # 8. トレーナーの初期化と学習開始
    trainer = SentenceTransformerTrainer(
        model=model,  # 学習対象モデル
        args=args,  # 学習パラメータ
        train_dataset=train_dataset,  # 学習データセット
        eval_dataset=eval_dataset,  # 評価データセット
        loss=loss,  # 損失関数
        evaluator=dev_evaluator,  # 評価器
    )
    
    trainer.train()  # モデル学習開始

    # 9. 学習済みモデルの保存
    model.save_pretrained(f"outputs/models/{model_name}/final")


if __name__ == "__main__":
    model_list = [
        "answerdotai/ModernBERT-base",
        "answerdotai/ModernBERT-large",
        "FacebookAI/roberta-base",
        "google-bert/bert-base-uncased",
    ]

    for model_name in model_list:
        train_embed_model(model_name)

 モデルの評価学習によってベクトル検索モデルを作成したら，それらのモデルを評価していきます！
評価方法は色々あると思いますが，ここでは，Sentence-TransformerのNanoBEIREvaluatorを利用して，評価していきたいと思います．
NanoBEIRの全部のデータセットで評価すると時間がかかるので，今回はQuora、MSMARCO、NQで評価しました．
import os
import json
from sentence_transformers import (
    SentenceTransformer,
)
from sentence_transformers.evaluation import NanoBEIREvaluator


def eval_embed_model(model_name: str) -> None:
    # 指定されたモデルを評価する関数

    # 1. 学習済みモデルのロード
    # 指定されたモデル名を元に、保存済みのSentence-Transformerモデルをロード
    model = SentenceTransformer(f"outputs/models/{model_name}/final")

    # 2. 評価対象データセットの指定
    # NanoBEIRフレームワークでサポートされているデータセット名をリストで指定
    datasets = ["QuoraRetrieval", "MSMARCO", "NQ"]  # Quora、MSMARCO、NQデータセットを使用

    # 3. 評価器の初期化
    # NanoBEIREvaluatorは複数のデータセットに対してモデルの性能を評価するためのツール
    evaluator = NanoBEIREvaluator(
        dataset_names=datasets,  # 使用するデータセット名を指定
    )

    # 4. モデル評価の実行
    # 指定されたモデルに対して評価を実行し、結果を取得
    result = evaluator(model)

    # 5. 評価結果の保存先パス設定
    result_path = f"outputs/results/{model_name}.json"

    # 6. 保存先フォルダが存在しない場合は作成
    os.makedirs(os.path.dirname(result_path), exist_ok=True)

    # 7. 評価結果をJSON形式で保存
    with open(result_path, "w") as f:
        json.dump(result, f, indent=4)


if __name__ == "__main__":
    model_list = [
        "answerdotai/ModernBERT-base",
        "answerdotai/ModernBERT-large",
        "FacebookAI/roberta-base",
        "google-bert/bert-base-uncased",
    ]

    for model_name in model_list:
        eval_embed_model(model_name)

 評価結果 📊実際に学習させたベクトル検索モデルを評価した結果を下記に示します！


評価項目
BERT
RoBERTa
ModernBERT-base
ModernBERT-large


nDCG@10
0.50
0.41
0.44
0.46

Recall@10
0.66
0.56
0.57
0.60

Accuracy@10
0.67
0.57
0.58
0.61

BERTが一番高い結果に 🤔
おそらく学習を行う際のハイパラ関係を雑に設定しているのと，学習データセットの偏りなどの影響によるものだと思います！
ちゃんと学習させてみたい！！！
追記：検索用のデータセットで学習させないと検索タスクはうまく解けないようです！
https://x.com/hpp_ricecake/status/1870720609153659295

 感想ぎりぎり22日（0:00)になってしまいました 💦
今回は時間の都合上，ModernBERTの利点を引き出すような検証はできていませんが，ModernBERTが良さそうモデルであることが手元で動かして確かめることができました．
公開されているModernBERTは英語のみ対応しているモデルで，日本語には対応していないので，多言語対応されて欲しいですね 👀
評価項目	BERT	RoBERTa	ModernBERT-base	ModernBERT-large
nDCG@10	0.50	0.41	0.44	0.46
Recall@10	0.66	0.56	0.57	0.60
Accuracy@10	0.67	0.57	0.58	0.61
概要

ModernBERTとは何か？

🚅 早い

😋 上手い（高性能）

📏 長い（コンテキスト長）

より詳しく知りたい人はこちら

ベクトル検索モデルを作ってみよう

モデルの学習

モデルの評価

評価結果 📊

感想

Discussion
