RetroMAE

レトロマ

Retromaeおよびそれ以降のコードベース。

新着情報

• 2022年10月、Retromae：マスクされた自動エンコーダーを介したトレーニング前の検索指向言語モデルは、 EMNLP 2022に受け入れられます。 MS MarcoとBeirのSotaパフォーマンスは、Bert-Base Scaleの密集したレトリバーから！
• 2022年11月、Retromae V2：トレーニング前の検索指向言語モデルのための二重マスクされた自動エンコーダーがARXIVになりました。 V1からのもう1つの大きな進歩とMs MarcoとBeirの大規模な改善！モデルとコードはまもなく登場します！

リリースされたモデル

Huggingface Hubにいくつかのチェックポイントをアップロードしました。

モデル	説明	リンク
レトロマ	ウィキペディアとBookCorpusで事前にトライアンスされています	Shitao/Retromae
Retromae_msmarco	MSMARCOパッセージで事前にトライアンスされています	Shitao/Retromae_msmarco
Retome_msmarco_finetune	MSMARCOパッセージデータのRetromae_MsmarcoをFinetuneします	shitao/retromae_msmarco_finetune
Retome_msmarco_distill	MSMARCOパッセージデータのRetromae_MsmarcoをFintune Cross-EncoderでKL-Divergenceを最小化することにより	shitao/retromae_msmarco_distill
Retromae_beir	BeirのMSMARCOパッセージデータのRetromaeをFintune（Beirが提供する公式ネガを使用）	Shitao/Retromae_beir

識別子文字列を使用して簡単にロードできます。例えば：

 from transformers import AutoModel
model = AutoModel . from_pretrained ( 'Shitao/RetroMAE' )

最先端のパフォーマンス

レトロマは、密なレトリバーの強力な初期化を提供できます。ドメイン内データを微調整した後、対応するシナリオで高品質の監視された検索パフォーマンスが発生します。その上、事前に訓練されたモデルの転送可能性を大幅に改善するため、ドメイン外データセットで優れたゼロショットパフォーマンスをもたらすのに役立ちます。

MSMARCOパッセージ

• ウィキペディアとbookcorpusで事前に訓練されたモデル：

モデル	MRR@10	@1000を思い出してください
バート	0.346	0.964
レトロマ	0.382	0.981

• MSMARCOで事前に訓練されたモデル：

モデル	MRR@10	@1000を思い出してください
Cocondenser	0.382	0.984
レトロマ	0.393	0.985
レトロマ（蒸留）	0.416	0.988

Beir Benchemark

モデル	AVG NDCG@10（18データセット）
バート	0.371
コンデンサー	0.407
レトロマ	0.452
レトロマV2	0.491

インストール

 git clone https://github.com/staoxiao/RetroMAE.git
cd RetroMAE
pip install .

開発のために、編集可能なインストール：

 pip install -e .

ワークフロー

このリポジトリには、トレイン前とFinetuneの2つの機能が含まれています。まず、マスク言語モデリングの損失を使用して、一般データセット（または下流のデータセット）でレトロマをトレーニングします。次に、対照的な損失で下流のデータセットでレトロマを微調整します。より良いパフォーマンスを実現するために、クロスエンコーダーが提供するスコアを蒸留することにより、レトロマエを微調整することもできます。詳細なワークフロー私たちの例を参照してください。

プレレイン

 torchrun --nproc_per_node 8 
  -m pretrain.run 
  --output_dir {path to save ckpt} 
  --data_dir {your data} 
  --do_train True 
  --model_name_or_path bert-base-uncased 
  --pretrain_method {retromae or dupmae}

微調整

 torchrun --nproc_per_node 8 
-m bi_encoder.run 
--output_dir {path to save ckpt} 
--model_name_or_path Shitao/RetroMAE 
--do_train  
--corpus_file ./data/BertTokenizer_data/corpus 
--train_query_file ./data/BertTokenizer_data/train_query 
--train_qrels ./data/BertTokenizer_data/train_qrels.txt 
--neg_file ./data/train_negs.tsv 

例

• トレイン前
  • ウィキペディアのプリトレイン
  • MSMARCOパッセージのプリトレイン
• Bi-Encoder
  • MSMARCOパッセージのFinetune
  • Beir Benchemark
• クロスエンコーダー
  • MSMARCOパッセージの再審査

引用

私たちの作品が役立つ場合は、私たちを引用することを検討してください。

 @inproceedings{RetroMAE,
  title={RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder},
  author={Shitao Xiao, Zheng Liu, Yingxia Shao, Zhao Cao},
  url={https://arxiv.org/abs/2205.12035},
  booktitle ={EMNLP},
  year={2022},
}