All Datasets V3 MiniLM L6

🚀 フラックス文埋め込みモデル

このプロジェクトは、自己教師付きの対照学習目標を使用して、非常に大規模な文レベルのデータセットで文埋め込みモデルをトレーニングすることを目的としています。このモデルは、文の意味情報を捉えたベクトルを出力し、情報検索、クラスタリング、文の類似性タスクなどに利用できます。

🚀 クイックスタート

このモデルは文エンコーダとして使用されます。入力文が与えられると、文の意味情報を捉えたベクトルを出力します。この文ベクトルは、情報検索、クラスタリング、または文の類似性タスクに使用できます。

✨ 主な機能

• 大規模な文ペアデータセットでの微調整により、高品質な文埋め込みを提供します。
• 対照学習目標を使用して、文の類似性を効果的に捉えます。
• SentenceTransformersライブラリを使用して簡単に使えます。

📦 インストール

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

💻 使用例

基本的な使用法

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v3_MiniLM-L6')
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-L6-H384-uncased' モデルを使用し、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チームのメンバーから効率的な深層学習フレームワークに関する助言を得ました。

想定される用途

このモデルは文エンコーダとして使用されることを想定しています。入力文が与えられると、文の意味情報を捉えたベクトルを出力します。この文ベクトルは、情報検索、クラスタリング、または文の類似性タスクに使用できます。

トレーニング手順

事前学習

事前学習済みの 'MiniLM-L6-H384-uncased' モデルを使用しました。これは、'microsoft/MiniLM-L12-H384-uncased' の6層バージョンで、2層ごとに選択したものです。事前学習手順の詳細については、モデルカードを参照してください。

微調整

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

ハイパーパラメータ

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

トレーニングデータ

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