All Datasets V4 MiniLM L12

🚀 文書埋め込みモデル

このプロジェクトは、自己教師付きの対照学習を用いて、非常に大規模な文レベルのデータセットで文埋め込みモデルを学習することを目的としています。事前学習済みの'MiniLM-L12'モデルを使用し、10億文ペアのデータセットで微調整を行いました。対照学習の目的は、ペアからの文が与えられたときに、モデルがランダムにサンプリングされた他の文のセットの中から、実際にデータセットでペアになっている文を予測することです。

このモデルは、Hugging Faceによって主催されたCommunity week using JAX/Flax for NLP & CVの間に開発されました。このモデルは、Train the Best Sentence Embedding Model Ever with 1B Training Pairsというプロジェクトの一環として開発されました。このプロジェクトを実行するために、7台のTPU v3-8という効率的なハードウェアインフラストラクチャを利用し、GoogleのFlax、JAX、およびCloudチームのメンバーから、効率的な深層学習フレームワークに関する助言を得ました。

✨ 主な機能

• 文エンコーダとして使用でき、入力文に対して文の意味情報を捉えたベクトルを出力します。
• この文ベクトルは、情報検索、クラスタリング、文の類似度タスクに使用できます。

📦 インストール

このモデルを使用するには、SentenceTransformersライブラリをインストールする必要があります。

💻 使用例

基本的な使用法

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v4_MiniLM-L12')
text = "Replace me by any text you'd like."
text_embbedding = model.encode(text)
# array([-0.01559514,  0.04046123,  0.1317083 ,  0.00085931,  0.04585106,
#        -0.05607086,  0.0138078 ,  0.03569756,  0.01420381,  0.04266302 ...],
#        dtype=float32)

📚 詳細ドキュメント

学習手順

事前学習

事前学習済みの'MiniLM-L12'を使用しています。事前学習手順の詳細については、モデルカードを参照してください。

微調整

対照学習の目的を使用してモデルを微調整します。正式には、バッチ内のすべての可能な文ペアからコサイン類似度を計算します。その後、真のペアと比較してクロスエントロピー損失を適用します。

ハイパーパラメータ

• モデルはTPU v3-8で学習しました。
• バッチサイズ1024（TPUコアあたり128）で540kステップ学習しました。
• 学習率のウォームアップは500ステップで行いました。
• シーケンス長は128トークンに制限しました。
• AdamWオプティマイザを使用し、学習率は2e-5です。
• 完全な学習スクリプトは、現在のリポジトリでアクセス可能です。

学習データ

複数のデータセットを連結してモデルを微調整しました。文ペアの総数は10億文を超えています。各データセットは、data_config.jsonファイルに詳細が記載されている重み付き確率に基づいてサンプリングされました。