Databricks上で埋め込みモデルをファインチューニングする

RAG（Retrieval-Augmented Generation、データベースに保管した社内データを生成AIモデルが参照して自社固有の回答を生成する仕組み）やセマンティック検索（いわゆる、あいまい検索）の分野で埋め込みモデル（Embedding Model）を活用するケースが増えてきました。
埋め込みモデルはHuggingFaceなどで公開されているものをダウンロードしてそのまま使用するのが標準的ですが、よりドメイン特化な単語や文章の理解度を向上させるためにファインチューニングを行うというオプションも存在します。こちらの記事ではそんな埋め込みモデルのファインチューニングを、具体的なコードを示しながら、Databricks上で実施する手順をご紹介します。

環境

Databricks Runtime: 14.2 ML GPU
フルのソースコード: https://github.com/hiouchiy/product-search
参考にしたサンプルコード: https://github.com/databricks-industry-solutions/product-search

埋め込みモデルとは？

「埋め込みモデル」（embedding model）とは、単語や文などの言語的要素を数値のベクトルに変換する手法です。このベクトルは多次元空間上に表現され、単語や文の意味的、文脈的特徴を捉えます。これにより、言語モデルは言語データを効率的に処理し、理解することができます。

埋め込みモデルは、特に大規模言語モデル（LLM）において重要であり、 RAGやセマンティック検索など自然言語処理（NLP）のタスクで広く使用されています。例えば、単語の埋め込み（word embeddings）は、単語間の意味的関係を捉えるために用いられ、文の埋め込み（sentence embeddings）は文全体の意味を表すベクトルを生成します。埋め込みモデルの精度は、AIの言語理解能力に大きく影響を与えるといっても過言ではありません。

参考までに、以下は主に埋め込みモデルとして使用されているTransformerベースの代表的なモデルの一覧です：

BERT (Bidirectional Encoder Representations from Transformers)
RoBERTa (A Robustly Optimized BERT Approach)
DistilBERT (Distilled BERT)
ALBERT (A Lite BERT)
XLNet

各モデルはBERTを基にしているか、BERTのコンセプトを拡張したもので、それぞれが独自の特性を持ち、最適化が施されています。

なぜ埋め込みモデルをファインチューニングする必要があるか？

埋め込みモデルをファインチューニングする主な理由は、特定のタスクやデータセットに対するモデルの性能を最適化することです。より具体的には、以下の点が挙げられます：

特定のドメインへの適応:
- HuggingFaceなどで一般公開されている埋め込みモデルは、一般的な言語理解のために事前学習されていますが、特定の用途（例えば医療、法律、特定の業界用語など）に対して最適化されているわけではありません。ファインチューニングにより、これら特定の領域における言語のニュアンスや専門用語にモデルを適応させることができます。
新しいタスクや言語への適用:
- 事前学習モデルを新しいタスクや、事前学習に含まれなかった言語に適用する際にも、ファインチューニングが必要になります。

ファインチューニング用のデータセットについて

ファインチューニング用には一般的に以下のデータセットが必要です。

クエリ（質問、検索ワードなど）
コーパス（回答文、検索結果など）
スコア（いわゆる正解ラベル、無くても構わない）

以下は、サンプルノートブックで使用しているデータセットの一部です。

クエリ	コーパス	スコア
hardwood beds	good , deep sleep can be quite difficult to have in this busy age . fortunately , there ’ s an antidote to such a problem : a nice , quality bed frame like the acacia kaylin . solidly constructed from acacia wood , this bed frame will stand the test of time and is fit to rest your shoulders on for years and years . its sleek , natural wood grain appearance provides a pleasant aesthetic to adorn any bedroom , acting both as a decorative piece as well as a place to give comfort after a hard day of...	1
jennie tufted upholstered low profile platform bed	good , deep sleep can be quite difficult to have in this busy age . fortunately , there ’ s an antidote to such a problem : a nice , quality bed frame like the acacia kaylin . solidly constructed from acacia wood , this bed frame will stand the test of time and is fit to rest your shoulders on for years and years . its sleek , natural wood grain appearance provides a pleasant aesthetic to adorn any bedroom , acting both as a decorative piece as well as a place to give comfort after a hard day of...	0.75
bathroom lighting	the vogan vanity from our vogan series is a 33 '' wide single bowl vanity . the vanity is beautiful and decorative as well as designed to make installation as easy as possible . it comes with a mounted beige solid marble top with pre-drilled holes for faucets , a mounted porcelain bowl and the back was designed with an opening to allow for easy plumbing hook-ups . this collection boasts a spanish and french influence with its large scaled details in a finish of varying tones of soft blues featur...	0

一つ目のポイントは、スコアの有無です。
上記の通り、スコアは正解ラベルに相当します。つまり、クエリーに対して、コーパスがどれだけ一致しているか、または、一致していないかを示す数字です。上記サンプルデータでは、完全一致(Exact)、部分一致(Partial)、無関係(Irrelevant)の３ラベルのいずれかを各レコードに割り当てています。また、各ラベルに対する重みとして、それぞれ、1、0.75、0という数値を割り当てています。なお、重みの割り当てにはこちらの記事も参照ください。
通常スコアは人手でセットすることが多いのですが、スコアデータが存在している、または、ある程度簡単に用意できる場合には、教師あり学習のアプローチでファインチューニングを実施します。しかし、そんな便利なデータが存在していない場合も多々あるでしょう。その際は、スコアデータなしの教師なし学習のアプローチでファインチューニングを実施します。

二つ目のポイントは、クエリに関してです。
上記サンプルでは、検索ワードという形でクエリデータが存在しているケースを前提としていますが、そういったデータが存在しないケースもあると思います。例えば、膨大な製品ドキュメントだけが手元にあるようなケースで、コーパスに相当する製品情報はいくらでも抽出可能ですが、それに対する「問い」のデータがない場合です。スコアデータとは異なり、クエリデータは原則必須データとなるため、なんとかデータを用意する必要があります。人手で作成できればした方が良いですが、現実的でない場合は、 ChatGPTなど既存のLLMを使用して作成するアプローチが提案されています。

ファインチューニングの具体的な手法

データが準備できたとして、ここから具体的なファインチューニングの実装に関して、教師ありアプローチと教師なしアプローチの両方を記載します。

教師あり（Supervised）アプローチ

まずは教師あり（スコアデータが存在する場合）です。
上記サンプルデータがPandas Dataframeのsearch_pdに格納されているとして、クエリ、コーパス、スコアデータを元にInputExampleオブジェクトをレコード数分作成します。

from sentence_transformers import InputExample

# 入力を組み立てる関数を定義
def create_input(query, corpus, score):
  return InputExample(texts=[query, corpus], label=score)

# 各検索結果を入力に変換
inputs = search_pd.apply(
  lambda s: create_input(s['query'], s['product_text'], s['score']), axis=1
  ).to_list()

その後、チューニング対象のモデルをSentenceTransformerを使ってダウンロードし、DataLoader、損失関数（Loss Function）を設定した上で、学習を開始します。特に重要な点は、損失関数のCosineSimilarityLossです。こちらはスコアデータありの教師あり学習なので、モデルの出力を元にコサイン類似度を使用して算出されたスコアと正解ラベルのスコアを比較し、その差が最小化するようにパラメータを調整します。したがって大前提として入力データInputExampleにlabel=scoreのパラメータ指定が必須となります。

from sentence_transformers import SentenceTransformer, losses
from torch.utils.data import DataLoader

# チューニング対象のモデルをダウンロード
tuned_model = SentenceTransformer('all-MiniLM-L12-v2')

# モデルへの入力指示を定義
input_dataloader = DataLoader(
   inputs, 
   shuffle=True, 
   batch_size=16   # BSはサンプルノートブックの設定をそのまま使用
) 

# 損失指標を定義
loss = losses.CosineSimilarityLoss(tuned_model)

# 入力データに対してモデルを調整
tuned_model.fit(
  train_objectives=[(input_dataloader, loss)],
  epochs=1, # テスト用にエポック数は1に設定
  warmup_steps=100 # 学習率が最大まで上昇してからゼロに戻るまでのステップ数を制御
  )

教師なし（Unsupervised）アプローチ

続いて教師なし（スコアデータが存在しない場合）です。
こちらのノートブックによると、InputExampleオブジェクト作成時にパラメーターにクエリとコーパスのみを渡します。

from sentence_transformers import InputExample

# チューニング対象のモデルをダウンロード
def create_input_without_label(query, corpus):
  return InputExample(texts=[query, corpus])

# 各検索結果を入力に変換
inputs_without_label = search_pd.apply(
  lambda s: create_input_without_label(s['query'], s['product_text']), axis=1
  ).to_list()

その後、教師ありと同様にモデルのダウンロードなどを実施しますが、一番の違いは損失関数です。ここではMultipleNegativesRankingLossを採用しています。これは、クエリに対して性の類似度を持つコーパスをより近く、それ以外のコーパスをより遠ざけるようにコサイン類似度を計算する損失関数です。詳細はこちらのブログがわかりやすいです。いずれにしても、正解ラベルが存在しない場合で、かつ、正のペア（クエリとコーパス）のデータセットしかないような場合に特に有効です（参考）。

# チューニング対象のモデルを
tuned_model = SentenceTransformer('all-MiniLM-L12-v2')

# モデルへの入力指示を定義
input_dataloader = DataLoader(
   inputs, 
   shuffle=True, 
   batch_size=16   # BSはサンプルノートブックの設定をそのまま使用
) 

# 損失指標を定義
loss = losses.MultipleNegativesRankingLoss(tuned_model)

# 入力データに対してモデルを調整
tuned_model.fit(
  train_objectives=[(input_dataloader, loss)],
  epochs=1, # テスト用にエポック数は1に設定
  warmup_steps=100 # 学習率が最大まで上昇してからゼロに戻るまでのステップ数を制御
  )

モデルの評価

埋め込みモデルの評価指標はいつくかありますが、サンプルノートブックでは下記２つの指標を用いています。

クエリとコーパスのコサイン類似度
- ファインチューニング前と比較してどの程度値が小さくなっているか
（教師ありの場合）正解ラベルとクエリとコーパスのコサイン類似度の相関係数
- ファインチューニング前と比較してどの程度値が大きくなっているか

以下は、サンプルコードです。

クエリとコーパスのコサイン類似度


query_embeddings = (
  tuned_model
    .encode(
      search_pd['query'].tolist()
      )
  )

product_embeddings = (
  tuned_model
    .encode(
      search_pd['product_text'].tolist()
      )
  )

# 各クエリと商品のペアについて、コサイン類似度を計算
tuned_cos_sim_scores = (
  util.pairwise_cos_sim(
    query_embeddings, 
    product_embeddings
    )
  )

# コサイン類似度の平均点
tuned_cos_sim_score = torch.mean(tuned_cos_sim_scores).item()

# 結果表示
print(f"With tuning, avg cosine similarity went from {original_cos_sim_score} to {tuned_cos_sim_score}")

（教師ありの場合）正解ラベルとクエリとコーパスのコサイン類似度の相関係数

# コサイン類似度と関連性スコアの相関を決定
tuned_corr_coef_score = (
  np.corrcoef(
    tuned_cos_sim_scores,
    search_pd['score'].values
  )[0][1]
)

# 結果表示
print(f"With tuning, the correlation coefficient went from {original_corr_coef_score} to {tuned_corr_coef_score}")

なお、これ以外にも、こちらの記事によると以下のような指標も利用できます。要件に合わせていろいろ試しましょう。

ヒット率：
- (クエリ、関連文書)のペアごとに、クエリを含む上位k個の文書を検索する。検索結果に関連文書が含まれていればヒットとなる。
sentence_transformersのInformationRetrievalEvaluator
- これは、様々なtop-k値におけるコサイン類似度精度、精度、リコールなどのメトリクスの包括的なスイートを提供

まとめ

本記事では埋め込みモデルのファインチューニングに関して、具体的なデータセットと実装方法を記載しました。LLMベースのチャットボットアプリやセマンティック検索を用いた商品検索システムなどの精度向上にお役に立てば幸いでございます。

Discussion