埋め込みモデル学習と評価ための4つのクラスタリングデータセットの公開

oshizo

埋め込みモデルの学習と評価のために、有価証券報告書、法律、輸出入貨物の関税分類、医療機器の類別の4つの異なるドメインからデータを収集し、クラスタリングデータセットとして整形、公開しました。

背景

私は主にRAG目的で埋め込みモデルを調べていたため、クラスタリングタスクは触っていませんでしたが、最近、SFR-Embeddingの記事にクラスタリングタスクの訓練で検索タスクも性能向上するという記述を目にして興味を持ち始めました。

埋め込みモデルをクラスタリングタスクと統合すると、検索性能が大幅に向上することが観察されており、さらに複数のタスクからの知識転移によってその効果を高めることができます。
（中略）
クラスタリングラベルは、高レベルの概念に基づいて埋め込みを正則化するようモデルを促し、結果として異なるドメイン間でデータをより良く分離できるようになると我々は仮説を立てています。

https://blog.salesforceairesearch.com/sfr-embedded-mistral/

ただ、以前Gemma-2Bを埋め込みモデルとして訓練した時のモデルによるWikipedia記事検索タスクの結果を観察していた印象では、「ロバート・ピールの出身はどこ」のようなフルネームが重要なクエリに対し、「ロバート・カール」や「ロバート・W・ポール」など文字列の似た他人の記事を上位にしてしまう傾向が目立っており、このような問題はクラスタリングタスクでは解決しづらそうにも思っています。
実際にクラスタリングタスクでの訓練で検索性能が向上するのかは今後調べようと思います。

また、最近 SB Intuitions によって開発されている日本語埋め込みデルの評価ベンチマークJMTEBでもクラスタリングタスクがサポートされており、作成したデータセットを手軽に評価できるようになりました。

このようなjsonnetファイルをsrc/jmteb/configs/tasksの下に作成して、そのファイルをimportする文をsrc/jmteb/configs/jmteb.jsonnetに追記するだけで独自のタスクを評価できるようになります。

{
  jmdn: {
    class_path: 'ClusteringEvaluator',
    init_args: {
      val_dataset: {
        class_path: 'HfClusteringDataset',
        init_args: {
          path: 'oshizo/JMDNClustering-ja',
          split: 'train',
          name: "",
        },
      },
      test_dataset: {
        class_path: 'HfClusteringDataset',
        init_args: {
          path: 'oshizo/JMDNClustering-ja',
          split: 'test',
          name: "",
        },
      },
    },
  },
}

データセットの詳細

各サイトの利用規約を確認しつつ、1.有価証券報告書、2.法律、3.輸出入貨物の関税分類、4.医療機器の類別名称の4種類のデータを収集し、クラスタリングデータセットとして整形しました。

1. 有価証券報告書（ASR Annual Security Reports）：35クラス、4248件

有報の「事業の内容」セクションのテキストを使ってEDINETの業種コードに分類するタスクです。

有価証券報告書のデータはEDINETからAPIで取得することができますが、Hugging Face Hubに整ったデータセットnumad/yuho-text-2023が公開されていたので今回はこちらを整形しました。
（numadさん、ありがとうございます！）

このデータセットのうち、事業の内容が記述されている"DescriptionOfBusinessTextBlock"のタグのテキストのみを使用し、EDINET業種コードをクラスラベルとして使用しました。
テキストと業種コードの紐づけには、以下のサイトから入手できる「EDINETコードリスト」ファイルを使いました。
https://disclosure2.edinet-fsa.go.jp/weee0010.aspx

テキストは一般的な埋め込みモデルのmax_lengthより長いものが多いですが、冒頭部分で業種の判定ができるものが大半のため精度への影響はほぼなさそうです。

以下は業種が「医薬品」のテキストの例です。

当社グループは、当社及び連結子会社３社より構成され、遺伝子医薬品を中心とする医薬品の開発及び販売を進めております。また、希少遺伝性疾患のスクリーニング検査を中心として、検査受託サービスを実施しております。さらに当社子会社である…

EDINETのデータの利用規約はこのページにあり、複製、翻案、商用利用が許可されています。

2. 法律（Law）：49クラス、6127件

法律文を憲法、刑事、財務通則などのカテゴリに分類するデータセットです。

日本の法律データはe-Govから一括でダウンロードできます。
ダウンロードできるzipファイルは憲法、刑事、財務通則などに分かれており、このzipファイルの分類名をそのままクラスラベルとして使いました。

昭和以前の古い法律のデータは本文がひらがなではなくカタカナで書かれているものがあり、これにより難易度が上がることは避けたかったため、XMLの「Era」フィールドがHeiseiまたはReiwaのデータのみを使用しました。

このデータも一般的な埋め込みモデルのmax_lengthより非常に長いテキストが多いですが、目視で確認した所では冒頭部分で分類ができるものが大半に見えます。

以下は分類が「厚生」のテキストの例です。

移植に用いる臍帯血の品質の確保のための基準に関する省令 第一章 総則 （定義） 第一条 この省令で「調製等」とは、 臍 さい 帯血供給業務のうち、移植に用いる臍帯血の採取以外の業務をいう。２ この省令で「バリデーション」とは、事業所の構造設備並びに手順、工程その他の調製・保存及び品質管理の方法（以下「調製・保存の手順等」という。）が期待される結果を与えることを検証し、これを文書とすることをいう。

利用規約はこのページにあり、自由に複製、翻案、商用利用できると書かれています。

3. 関税分類（HS）：20クラス、3541件

輸出入貨物の品物を説明するテキストを、HS(Harmonized Commodity)コードの最上位階層である「部」に分類するデータセットです。

税関のウェブサイトには、物流業者などが関税コードの判断に迷ったときに税関に事前に質問して回答を得ることができる「事前教示」という制度の回答事例が公開されています。

関税コードは9桁で非常に多様ですが、部（21種類）、類（97種類）、項（1000以上）の順に詳細化されていく階層的なカテゴリとなっています。
今回は「部」をクラスラベルとして使用しました。
部の一覧は税関の解説ページで確認できます。

テキストデータはこのような感じで、このデータ特有の癖のある書き方になっています。

人造繊維製メリヤス編物から成る女子用衣類（室内用）　　性状：人造繊維製メリヤス編物をガウン様に縫製したもので、前面腰部にポケットを有する　　　　　正面で右を左の上にしてボタン５カ所で閉じる　　材質：ポリエステル100％　　用途：就寝時に使用（男女兼用）　　サイズ：Ｍ～Ｌ　　包装：18枚／カートン

https://www.customs.go.jp/tetsuzuki_search/bunrui/J3/23/J32300331.htm

データは検索一覧ページから取得したため、長いテキストは末尾が「…」となって省略されていますが、重要な情報は前半にあるため分類には影響がないと考えそのままにしました。

利用規約はこのページにあり、複製や翻案、商用利用が許可されています。

4. 医療機器の類別データ（JMDN）：103クラス、4457件

医療機器の一般的な名称やその説明文を、JMDNコードの「類別」に従って分類するデータセットです。

独立行政法人 医薬品医療機器総合機構のwebサイトから、医療機器の一般的な名称・定義とその類別のExcelファイルをダウンロードできます。

あまり分かっていないですが、医療機器にはJMDN分類コードという8桁のコードがあり、それぞれのコードにどのような機器が該当するか説明が与えられています。

例えば、以下のようなテキストです。

滅菌済み血管処置用チューブ及びカテーテル
冠動脈貫通用カテーテル 冠動脈完全閉塞等の狭窄部にガイドワイヤの通過が困難な患者に対して経皮的冠動脈形成術を実施する場合、ガイドワイヤの通過部を確保するために使用するカテーテルをいう。

8桁のJMDNコードを整理するために「類別」という概念があり、「器12 理学診療用器具」「器25 医療用鏡」などの単位で分けられています。今回はこの類別をクラスラベルとして使用しました。

クラス数は103と多く、数件しかデータのないクラスも多いデータセットとなっています。

利用規約はこのページにあり、コンテンツは出展を明記することで自由に利用できると書かれています。

既存モデルの評価結果

JMTEBのコードを使い、今回作成した4つのデータセットのtest splitと、もともとJMTEBに含まれる2つのクラスタリングデータセットの評価を実施しました。
評価コードは、MiniBatchKMeans、AgglomerativeClustering、BisectingKMeans、Birchの4つの手法でvalidation dataのv_measure_scoreを算出し、最も良かった手法でtest dataの指標を算出しているようです。

https://github.com/sbintuitions/JMTEB/blob/main/src/jmteb/evaluators/clustering/evaluator.py

今回作ったデータセットはtrainとtestの2つに分けたので、trainをvalidation dataとして使いました。

以下が評価結果の一覧です。

JMTEBコードでの評価結果

所感

  • mE5-largeは平均値としては高いが、ASRやmewsc16など苦手なタスクもある

    • mewsc16では最良のモデルより7ポイント悪く、かなり差がある

  • JMDNは他のタスクと比べかなり簡単で、モデル間の数値の差も小さい

  • ベースモデルにJSNLIを学習させただけの私のモデル（4行目）も意外と健闘している

    • NLIと検索タスクはある程度性質が近いのかもしれません


一部のデータセットの学習実験

続いて、今回作った4つのデータセットのうち、ASR、HS、Lawの3つ（表の左から3つぶん）を学習してみて、他のタスクへの影響を見てみました。

マルチクラス形式のデータをSentence Transformersで学習するときは、クラスラベルを参照してTripletを自動で構築してくれるBatchAllTripletLossやこの亜種を使うと便利です。
今回の学習にはBatchAllTripletLossを使い、distance_metricとしてcosine類似度、margin=0.5として2epoch学習しました。

学習前後の各データセットの評価結果は以下の通りです。

学習結果

以下が所感ですが、全体的にはあまり嬉しい結果ではないです。

  • 学習したデータセットはtest精度がかなり伸びる

    • ASRやLawは10ポイント以上の向上

      • クラスタリングで取り組みたいドメインがある場合はそのドメインのデータセットで学習させることが有効

    • 一方で、そのデータセットのtrainを学習データに含めたモデルは、他のモデルとはtestデータで比較できない

      • NV-Embedなどmtebを学習データとして使っているモデルでは、testと類似するサンプルをtrainから除外しているようです

  • 学習していないデータセットはあまり良くならない

    • livedoor_newsは良くなったが、JMDNとmewsc16は悪くなった

    • 期待としては、クラスタリングデータセット全般で良くなってほしかった

今回は簡易的な学習実験を行いましたが、本来の目的であったクラスタリングタスクでの訓練で検索タスクの性能が向上するのか、については今後調べていく予定です。


この記事が気に入ったらサポートをしてみませんか？