Leveraging Pre-trained Language Model Checkpoints for Encoder-Decoder Models – Hugging Face Blog Summary

Leveraging Pre-trained Checkpoints for Encoder-Decoder Models

2020年11月9日のHugging Faceのブログ記事では、BERT、RoBERTa、GPT2などの事前学習済みエンコーダーのみ、またはデコーダーのみのチェックポイントの重みを使用して、エンコーダー・デコーダー（seq2seq）モデルを初期化する方法について説明しています。このウォームスタート技術により、フルエンコーダー・デコーダーモデルのコストのかかる事前学習を回避でき、さまざまなsequence-to-sequenceタスクにおいてT5やPegasusなどのモデルに匹敵する結果が得られます。

Theory of Warm‑starting

エンコーダー・デコーダーモデルは、エンコーダー・スタックとデコーダー・スタックで構成されています。ウォームスタートは4つの方法で行うことができます：(1) エンコーダーとデコーダーの両方をエンコーダーのみのチェックポイント（例：BERT）から、(2) エンコーダーをエンコーダーのみのチェックポイントから、デコーダーをデコーダーのみのチェックポイントから（例：BERT + GPT2）、(3) エンコーダーのみをエンコーダーのみのチェックポイントから、または (4) デコーダーのみをデコーダーのみのチェックポイントから。BERTからウォームスタートする場合、エンコーダー層はBERTの層と1対1で対応し、BERTの重みで初期化されます。デコーダーはBERTから同じ自己注意（self-attention）およびLM-Headの重みを受け取りますが、BERTにはクロスアテンション（cross-attention）がないため、クロスアテンション層が追加され、ランダムに初期化されます。GPT2からデコーダーをウォームスタートする場合、自己注意とLM-Headは直接コピーできますが、クロスアテンション層は再びランダムに初期化されます。

エンコーダーとデコーダーのアーキテクチャが同一である場合（クロスアテンションを除く）、それらの重みをタイイング（tied）することができ、パラメータ数を半分に減らすことができます。この重みタイイングは、両方が同じエンコーダーのみのチェックポイントからウォームスタートされている場合にのみ意味を持ちます。

Analysis of Warm‑started Models

このブログは、Rothe et al. (2020) による実験を要約したものです。この実験では、文融合（sentence fusion）、文分割（sentence splitting）、機械翻訳（WMT14 EN↔DE）、および抽象的要約（CNN/Dailymail, BBC XSum, Gigaword）の4つのタスクグループにおいて、さまざまなウォームスタートされたエンコーダー・デコーダー構成を、ランダムに初期化されたベースラインと比較しました。すべてのモデルは、bert-base-cased、roberta-base、または gpt2 チェックポイントに対応する12層、768次元の隠れ層サイズを使用しています。

モデルバリエーションの表には、ランダムに初期化された（"random"）パラメータ数と、活用された（"leveraged"）パラメータ数が示されています。例：

  • Rnd2Rnd: 221 M random, 0 leveraged
  • Rnd2BERT / BERT2Rnd: 112 M random, 109 M leveraged
  • BERT2BERT: 26 M random, 195 M leveraged
  • BERTShare / RoBERTaShare: 26 M random, 109 M / 126 M leveraged (weight tyingのため)
  • BERT2GPT2: 26 M random, 234 M leveraged
  • RoBERTa2GPT2: 26 M random, 250 M leveraged

結果（表より引用）は以下の通りです：

  • Sentence Fusion (DiscoFuse, SARI): RoBERTa2GPT2は100%データで89.9、10%データで87.1を達成。RoBERTaShare (large)は90.3 / 87.7に到達。
  • Sentence Splitting (WikiSplit, SARI): BERTShareは63.5、RoBERTaShareは63.4、RoBERTaShare (large)は63.8を記録。
  • Machine Translation (WMT14, BLEU-4): BERT2RndとBERT2BERTはどちらも30.1 → 32.7 (EN→DE / DE→EN)に到達。BERT2Rnd (large, custom)は31.7 → 34.2に改善。GPT2ベースのモデルは、GPT2の語彙が英語のみであるため、EN→DEでは低いパフォーマンスを示しました（例：BERT2GPT2は23.2）。
  • Summarization (Rouge-2): RoBERTaShareはCNN/Dailymailで18.95、BBC XSumで17.50、Gigawordで19.70を達成。RoBERTaShare (large)はCNN/Dailymailで18.91、BBC XSumで18.79、Gigawordで19.78に到達。BERTShareとBERT2BERTがそれに続き、GPT2ベースのモデルは後塵を拝しました（例：BERT2GPT2はCNN/Dailymailで4.96）。

分析の結論として、エンコーダーのウォームスタートはタスク全体を通じて一貫したブーストを与えますが、デコーダーのウォームスタートはクロスアテンション層がランダムに初期化されたままになるため、付加価値はそれほど大きくありません。重み共有は、入力と出力の分布が類似している場合（例：BBC XSum）には役立ちますが、モデルの容量や異なる語彙が重要となる翻訳では悪影響を及ぼす可能性があります。チェックポイントの語彙をタスクの言語に合わせることが不可欠です。

Practice: Warm‑starting with 🤗Transformers

このブログでは、BERT2BERTモデルをウォームスタートし、CNN/Dailymail要約でファインチューニングする完全なノートブックを提供しています。

  1. ライブラリのインストール: datasets==1.0.2 と transformers==4.2.1。
  2. データのロードと前処理: bert-base-uncased (最大長 512) とハイライト (最大長 128) で記事をトークン化し、パディングラベルのトークンを -100 に置き換えます。
  3. モデルのウォームスタート:
    from transformers import EncoderDecoderModel
bert2bert = EncoderDecoderModel.from_encoder_decoder_pretrained(
    "bert-base-uncased", "bert-base-uncased"
)
    警告には、予想通り、分類器 (cls) の重みが使用されていないことと、クロスアテンションの重みが新しく初期化されていることが示されます。
  4. 生成パラメータの設定 (bart-large-cnnからコピー):
    bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
bert2bert.config.eos_token_id = tokenizer.sep_token_id
bert2bert.config.pad_token_id = tokenizer.pad_token_id
bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size
bert2bert.config.max_length = 142
bert2bert.config.min_length = 56
bert2bert.config.no_repeat_ngram_size = 3
bert2bert.config.early_stopping = True
bert2bert.config.length_penalty = 2.0
bert2bert.config.num_beams = 4
  5. ファインチューニング: Seq2SeqTrainer と Seq2SeqTrainingArguments を使用し、predict_with_generate=True と Rouge-2 の precision, recall, fmeasure を返す compute_metrics 関数を使用します。
  6. トレーニング: デモンストレーション用にサブセット（トレーニング例32個、検証例8個）で実行。TITAN RTXでのフルトレーニングには約8時間かかります。
  7. 評価: テストセットで評価。完全にトレーニングされた BERT2BERT モデル (patrickvonplaten/bert2bert_cnn_daily_mail) は、完全な CNN/Dailymail 評価において 18.22 の Rouge-2 fmeasure を達成し、論文で報告されている数値に匹敵するか、それをわずかに上回ります。

また、ノートブックでは、モデルの保存と再ロードの方法、および from_encoder_decoder_pretrained に tie_encoder_decoder=True を渡すことでエンコーダー・デコーダーの重みをタイイングする方法も示されています。

Takeaway

既存の BERT、RoBERTa、または GPT2 チェックポイントを使用してエンコーダー・デコーダーモデルをウォームスタートすることで、ゼロからの事前学習という膨大な計算コストをかけることなく、強力な seq2seq パフォーマンスを得ることができます。エンコーダーの初期化が最も重要な要因であり、デコーダーの初期化は、ファインチューニング中にクロスアテンションが学習されない限り、付加価値はそれほど高くありません。入力と言語またはフォーマットが類似している場合は重みタイイングが有益ですが、モデルの容量や異なる語彙が重要となる翻訳では、チェックポイントとタスク言語の語彙の一致が成功の前提条件となります。

Sources