1Bのトレーニングペアで文埋め込みモデルをトレーニングする

Train a sentence embedding model with 1B training pairs.

文の埋め込みは、文を実数のベクトルにマッピングする手法です。理想的には、これらのベクトルは文の意味を捉え、高度に汎用的であるべきです。そのような表現は、クラスタリング、テキストマイニング、質問応答など、多くの下流アプリケーションで使用することができます。

私たちは、「1Bのトレーニングペアで最高の文埋め込みモデルをトレーニングするプロジェクト」として、最先端の文埋め込みモデルを開発しました。このプロジェクトは、Hugging Faceによって主催されたCommunity week using JAX/Flax for NLP & CVの一環として行われました。このプロジェクトでは、GoogleのFlax、JAX、およびCloudチームのメンバーから、効率的なディープラーニングフレームワークに関するガイダンスを受けました！

トレーニング手法

モデル

単語とは異なり、有限の文の集合を定義することはできません。したがって、文の埋め込み手法では、内部の単語を組み合わせて最終的な表現を計算します。たとえば、SentenceBertモデル（Reimers and Gurevych, 2019）では、多くのNLPアプリケーションの基盤であるTransformerを使用し、コンテキスト化された単語ベクトルに対してプーリング操作を行います（以下の図を参照）。

• コース開始コミュニティイベント
• モダンなCPU上でのBERTライクモデルの推論のスケーリングアップ – パート2
• 🤗 Transformersを使用して、低リソースASRのためにXLSR-Wav2Vec2を微調整する

マルチプルネガティブランキングロス

構成モジュールのパラメータは通常、自己教師ありの目的関数を使用して学習されます。このプロジェクトでは、以下の図で説明されているコントラスティブトレーニング方法を使用しました。文のペア（a_i, p_i）が意味的に近いペアであるデータセットを構成します。たとえば、（クエリ、回答パッセージ）、（質問、重複質問）、（論文タイトル、引用論文タイトル）などのペアを考慮します。その後、モデルは、ペア（a_i, p_i）を埋め込み空間で近いベクトルにマッピングするようにトレーニングされます。一方、非一致のペア（a_i, p_j）、i ≠ jは、埋め込み空間で遠いベクトルにマッピングされます。このトレーニング方法は、インバッチネガティブ、InfoNCE、またはNTXentLossとも呼ばれます。

形式的には、トレーニングサンプルのバッチが与えられた場合、モデルは以下の損失関数を最適化します：

− 1 n ∑ i = 1 n e x p ( s i m ( a i , p i ) ) ∑ j e x p ( s i m ( a i , p j ) ) -\frac{1}{n}\sum_{i=1}^n\frac{exp(sim(a_i, p_i))}{\sum_j exp(sim(a_i, p_j))} − n 1 ​ i = 1 ∑ n ​ ∑ j ​ e x p ( s i m ( a i ​ , p j ​ ) ) e x p ( s i m ( a i ​ , p i ​ ) ) ​

以下の例では、モデルはまず、バッチ内のすべてのペアの各文を埋め込みます。次に、すべての可能なペア（a_i、p_j）の類似度行列を計算します。その後、類似度行列を元のペアを示すグラウンドトゥルースと比較します。最後に、クロスエントロピー損失を使用して比較を行います。

直感的には、「ベルリンには何人の人々が住んでいますか？」と「ベルリンには約350万人の人々が住んでいます」という文に高い類似度を割り当て、他のネガティブな回答（「フランスの首都はパリです」といった回答など）には低い類似度を割り当てるべきです。以下の図に詳細が示されています。

損失の式では、simは（a、p）の類似度関数を示しています。類似度関数には、コサイン類似度またはドット積演算子を使用することができます。両方のメソッドには、以下に要約された利点と欠点があります（Thakur et al.、2021、Bachrach et al.、2014）：

実際には、スコアの差が小さすぎる傾向があるため、スケーリングファクターCを適用して、スケーリングされた類似度を使用しました。具体的には、s i m s c a l e d ( a , b ) = C ∗ s i m ( a , b)という形で計算しました。通常、C = 20（Henderson et al., 2020、Radford et al., 2021）とします。

より良いバッチで品質を向上させる

当社の手法では、サンプルペア（ a i , p i ）（a_i, p_i）（ a i ​ , p i ​ ）のバッチを作成します。バッチ内の他のすべてのサンプル（ a i , p j ）, i ≠ j（a_i, p_j), i \neq j（ a i ​ , p j ​ ）, i  = j は、負のサンプルペアです。したがって、バッチの構成はトレーニングの重要な側面です。文献を考慮して、バッチの3つの主要な側面に焦点を当てました。

1. サイズが重要

対照学習では、バッチサイズが大きいほどパフォーマンスが向上します。Quらの論文（2021年）から抽出した図に示されるように、バッチサイズが大きいほど結果が向上します。

2. ハードネガティブ

同じ図で、ハードネガティブを含めることもパフォーマンスを向上させることがわかります。ハードネガティブは、p i と区別するのが難しいサンプル p j です。この例では、「フランスの首都は何ですか？」と「アメリカの首都は何ですか？」というペアがあり、意味的な内容が似ており、正確に全文を理解する必要があるため、区別が難しいです。対照的に、「フランスの首都は何ですか？」と「スター・ウォーズの映画は何本ありますか？」というサンプルは、同じトピックを参照していないため、区別が容易です。

3. クロスデータセットのバッチ

私たちは複数のデータセットを連結してモデルをトレーニングしました。私たちは大きなバッチを作成し、同じデータセットからサンプルを収集して、トピックの分布を制限し、ハードネガティブを促進しました。ただし、少なくとも2つのデータセットをバッチに混ぜて、トピック間のグローバルな構造だけでなく、トピック内のローカルな構造も学習しました。

トレーニングインフラストラクチャとデータ

前述のように、データの量とバッチサイズはモデルのパフォーマンスに直接影響します。このプロジェクトでは、効率的なハードウェアインフラストラクチャを活用しました。私たちはGoogleが開発したTPUでモデルをトレーニングしました。TPUは行列の乗算に非常に効率的な計算ユニットです。TPUにはハードウェア固有の特性があり、特定のコード実装が必要な場合があります。

さらに、私たちは1億組の文のペアまで含む大規模なコーパスでモデルをトレーニングしました！使用したすべてのデータセットの詳細は、モデルカードに記載されています。

結論

チャレンジの間に作成したすべてのモデルとデータセットは、HuggingFaceのリポジトリで見つけることができます。Mini-LM（Wangら、2020年）、RoBERTa（liuら、2019年）、DistilBERT（Sanhら、2020年）、MPNet（Songら、2020年）など、20の汎用センテンストランスフォーマーモデルをトレーニングしました。私たちのモデルは、複数の汎用センテンス類似性評価タスクでSOTAを達成しています。また、質問応答、文の類似性、ジェンダー評価に特化した8つのデータセットも共有しました。

一般的な文の埋め込みは多くのアプリケーションに使用することができます。私たちは、いくつかのアプリケーションを紹介するスペースデモを作成しました：

• 文の類似性モジュールは、主テキストと他のテキストとの類似性を比較します。このデモでは、各テキストの埋め込みを抽出し、ソース文と他の文の間の類似性を余弦類似度で計算します。
• 非対称QAは、与えられたクエリの回答の可能性を選択した回答候補と比較します。
• 検索/クラスタは、クエリから近くの回答を返します。例えば、「python」と入力すると、ドット積距離を使用して最も近い文を取得します。
• ジェンダーバイアス評価は、ランダムサンプリングによってトレーニングセットの固有のジェンダーバイアスを報告します。対象の職業にジェンダーを明示しないアンカーテキストと、ジェンダープロヌーンを持つ2つの提案を考え、モデルが特定の提案により高い類似性を割り当てるかどうかを比較し、特定のジェンダーを好む割合を評価します。

JAX/Flaxを使用したNLP＆CVのコミュニティウィークは、濃密で非常に充実した経験でした！GoogleのFlax、JAX、Cloud、およびHugging Faceチームメンバーのガイダンスとサポートのおかげで、私たちは多くを学びました。他のプロジェクトも私たちと同じくらい楽しんでいたことを願っています。ご質問や提案がある場合は、お気軽にご連絡ください！

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!