機械学習・自然言語処理の勉強メモ

学んだことのメモやまとめ

pytorch-transformersを触ってみる⑤

はじめに

前回はBERTのfine tuningということで、NER(固有表現認識)を試してみました。
kento1109.hatenablog.com

今回はfine tuningではなく、BERTの事前学習について見ていきたいと思います。

pre-training from scratch

ただ、pytorch-transformersでの事前学習を調べると、早々に壁にぶつかりました。
ほとんどの内容がfine tuning向けなので、pre-trainingの内容があまり見つかりませんでした。
Issuesを見てみると、既に同様の質問をしている人がいました。
github.com

ここでpytorch-transformersの位置づけについて、開発者は以下のように述べています。

pytorch-pretraned-BERT was mostly designed to provide easy and fast access to pretrained models.
If you want to train a BERT model from scratch you will need a more robust code base for training and data-processing than the simple examples that are provided in this repo.

要は、事前学習済のモデルをfine tuningして利用することを目的としたライブラリなので、モデルのpre-training用として汎用化されたものではない、と言うことです。

ただ、全くpre-training用として使えないかというとそんなことはありません。
pytorch-transformersにはlm_finetuningなるスクリプトが用意されており、これは大規模なコーパスでpre-trainingしたモデルについて、ドメインに特化した小規模コーパスでpre-trainingしたモデルをfine tuningするためのスクリプトです。
pytorch-transformers/examples/lm_finetuning at master · huggingface/pytorch-transformers · GitHub

これを使うことでpre-trainingも可能になるのではと思います。
(実際に先ほどのIssuesでもそのようなことを示唆しているコメントもありました。)

ただし、pytorch-transformersでpre-trainingする必要はなく、Facebook researchやNVIDIAがBERTのpre-trainingに関するコードを公開しているので、そっちを利用するのもアリです。
GitHub - facebookresearch/XLM: PyTorch original implementation of Cross-lingual Language Model Pretraining.
DeepLearningExamples/TensorFlow/LanguageModeling/BERT at master · NVIDIA/DeepLearningExamples · GitHub

今回はpytorch-transformersのlm_finetuningを触ってみて、次回以降にその他のpre-trainingに関するコードも見れたらと思います。

lm_finetuning

データセットを作る場合、pregenerate_training_data.pyを動かすのがよさそうです。
ということで、pregenerate_training_data.pyを少し見てみます。

pregenerate_training_data.py

データセットのフォーマットとしては、ドキュメント単位にブランク行を入れるようです。
例えば、以下のようなドキュメントがあったとします。

リンゴは、バラ科リンゴ属の落葉高木樹。またはその果実のこと。植物学上はセイヨウリンゴと呼ぶ。春、白または薄紅の花が咲く。

バナナはバショウ科バショウ属のうち、果実を食用とする品種群の総称。また、その果実のこと。いくつかの原種から育種された多年性植物。種によっては熟すまでは毒を持つものもある。

その場合、データセットのテキストファイルは以下のようにします。

リンゴは、バラ科リンゴ属の落葉高木樹。
植物学上はセイヨウリンゴと呼ぶ。
春、白または薄紅の花が咲く。

バナナはバショウ科バショウ属のうち、果実を食用とする品種群の総称。
いくつかの原種から育種された多年性植物。
種によっては熟すまでは毒を持つものもある。

このスクリプトでは、1文書をdocリスト(各文はtokenize済)で保持し、それをまとめてdocsクラスで保持します。

Next Sentence Prediction

ドキュメントをdocsクラスに保持した後、create_training_file()でdocsを処理していきます。
実際の処理は、create_instances_from_document()でdoc毎に処理(インスタンスの生成)を行います。
create_instances_from_document()は、BERTの訓練データを作るうえで重要なので丁寧にみていきます。

BERTの訓練データでは、Next Sentence Predictionのために、文のペア(sentences A and B)を1インスタンスとします。
sentences Aペアの作成がNext Sentence Predictionの精度を高めるため、少し工夫されています。
ここでは、sentenceを選択するのではなく、sentenceの境界を無作為に選択します。
そして、その境界までのsentencesをtokens_aとして保持します。

例えば、docが4 sentencesから構成される場合、境界は1,2,3のいずれかになります。
1,2,3の境界から3が選択された場合、先頭から3 sentencesまでがtokens_aとして保持します。

token_bは50%の割合で実際の後続するsentences、50%の割合で関係の無いsentenceと入れ替えます。
(この割合は原著論文と同じです。)

Masked language modeling (LM)

次に、Masked LMのタスクのため、tokensの一部をMASKします。
これは、create_masked_lm_predictions()によって行われます。

簡単のため、以下の文でのMASK例を考えます。

my dog is hairy

最初にtokensのうち、MASKするtokenの数を決めます。(今回は「1」)
次に、tokensのindexをシャッフルして順番に取り出します。
今回は系列番号3(hairy)が先頭に出現したとします。
その系列を割合に基づいて以下のtokenに置き換えます。

  • 80% を [MASK] に置換(e.g., my dog is [MASK])
  • 10% を置換しない(e.g., my dog is hairy)
  • 10% を辞書内の任意のtokenに置換(e.g., my dog is apple

(この割合は原著論文と同じです。)

最後に置換されたtokensと、置換されたtokenのindex(4)とそのtoken(hairy)を返します。

create_instances_from_document()にMASK後のtokenを返した後、これらの情報をinstance辞書に格納します。
create_training_file()にdocのinstance辞書を返した後、json形式でファイルとして保存します。

pregenerate_training_data.pyの役割は以上です。

finetune_on_pregenerated.py

pre-trained model

モデルをインスタンス化します。

model = BertForPreTraining.from_pretrained(args.bert_model)

今回使用するモデルは、BertForPreTrainingクラスです。
名前がややこしいですが、BertPreTrainedModelを継承したクラスです。
このモデルにてMasked LM及びNext Sentence Predictionの学習を行います。

data loading

作成したjsonファイルを利用してpre-training(fine tuning)を行います。
これは、finetune_on_pregenerated.pyにより実行されます。

まず、先ほどの2つのjsonファイルのあるパスをPregeneratedDatasetクラスに渡してデータセットを作成します。
基本的な処理としては、epoch_{0}.json(maskされたtokensデータを含んだファイル)を1行(1インスタンス)毎に読み込み、InputFeaturesクラスを生成します。
InputFeaturesクラスは以下のクラス変数から成ります。

  • input_ids:tokensの系列(padding済)
  • input_mask:input_idsのpaddingを表す系列
  • segment_ids:input_idsのsentences A, Bを表す系列
  • lm_label_ids:input_idsのmaskされたtokenの元のtoken(正解ラベル)を並べた系列
  • is_next:sentences Bが元のsentences Aに続く文か、ランダムに挿入された文か(True / False)

lm_label_idsですが、convert_example_to_featuresで少し工夫されていたのを見逃していました。
lm_label_idsはinput_idsの系列を全て「-1」で初期化して、maskされた位置のみ元のtokenのラベルで置き換えます。

例えば、「my dog is hairy」の例で考えると以下のようになります。

input_ids = ['my', 'dog', 'is', '[MASK]']
lm_label_ids = [-1, -1, -1, 20] # {20: hairy}

※paddingやid変換などは考慮していないため、正確ではありません。

これを最後の行まで読み込んだデータをPregeneratedDatasetクラスで保持します。

PregeneratedDatasetをSamplerクラス→DataLoaderクラスと順番に渡していきます。

DataLoaderクラスまで渡したら、学習の準備が完了です。

後はモデルにInputFeaturesクラスをバッチ化したものを渡してforwardを行います。

outputs = model(input_ids, segment_ids, input_mask, lm_label_ids, is_next)

BertForPreTrainingのforwardを確認します。

outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                    attention_mask=attention_mask, head_mask=head_mask)

BertForPreTrainingですが、ネットワークの中核となるのはBertModelです。
BertModelは、encoderの役割(隠れ層の値を出力)を果たします。
入力としては以下を受け取ります。

  • input_ids:InputFeaturesクラスのinput_ids
  • token_type_ids: InputFeaturesクラスのsegment_ids
  • attention_mask:InputFeaturesクラスのinput_mask
sequence_output, pooled_output = outputs[:2]
prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)  #  BertPreTrainingHeads

BertForPreTrainingでは、以下のBertModelの出力を利用します。
1. sequence_output:隠れ層の値
2. pooled_output :隠れ層のLinear→Tanhした値

これをBertPreTrainingHeadsクラスの入力として渡します。
更に、BertPreTrainingHeadsでは以下の2つの処理を行います。
1. sequence_outputをBertLMPredictionHeadに渡す
2. pooled_outputをLinear(n=2)

BertLMPredictionHeadを確認します。
これは、sequence_output(隠れ層の値)を受け取り、BertPredictionHeadTransformにて変換後、Linear(n=vocab)を行います。

さて、BertForPreTrainingのforwardに戻り、損失計算の部分を確認します。
正解となるラベル情報は以下を利用します。

  • masked_lm_labels:InputFeaturesクラスのlm_label_ids
  • next_sentence_label:InputFeaturesクラスのis_next
outputs = (prediction_scores, seq_relationship_score,) + outputs[2:]

loss_fct = CrossEntropyLoss(ignore_index=-1)
masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
total_loss = masked_lm_loss + next_sentence_loss
outputs = (total_loss,) + outputs

prediction_scoresのshapeは「batch × tokens × vocab」です。
seq_relationship_scoreのshapeは「batch × 2」です。

masked_lm_labelsは、maskされたラベル以外が-1の系列なので、maskされた位置のみの損失が計算されます。
next_sentence_labelは0,1のラベル(1がランダム)です。

masked_lm_loss + next_sentence_lossを足し合わせた損失をtotal_loss として返します。

その後は、finetune_on_pregenerated.pyのmain()でloss.backward()して重みを調整していきます。

fine tuning For pre-training

さて、BERTのfine tuningについて確認しました。
これが出来ると、1からpre-trainingする理由があまり無いような気もします。
しかし、以下のような場合はやはり1からpre-trainingが必要と思います。

  • 大規模なpre-training modelがそもそも存在しない
  • pre-training modelは存在するが、自作tokenizerを利用したい

特にTokenClassificationの場合、tokenizerが望ましい分割をしてくれないとタスクが困難となります。
では、1からpre-trainingを行いたい場合はどうすれば良いでしょうか。

単純にpre-trained modelをロードしなければ良いのではと思いました。
最初に挙げたIssuesでも似たようなにコメントしている人もいました。

If you construct and initialize a new model instead of loading from pretrained, you can use the simple_lm_finetuning script to train on new data.

finetune_on_pregenerated.pyでは、237行目で以下のようにしてpre-trained modelを読み込んでいます。

model = BertForPreTraining.from_pretrained(args.bert_model)

これは、pre-trained modelが無い場合は利用できません。

その場合は、以下のようにしてBertForPreTrainingをインスタンス化すればOKです。

config = BertConfig.from_pretrained('bert-base-uncased')
model = BertForPreTraining(config)

BertForPreTrainingのコンストラクタを見ると以下のように記述されています。

def __init__(self, config):
    super(BertForPreTraining, self).__init__(config)

    self.bert = BertModel(config)
    self.cls = BertPreTrainingHeads(config)

    self.apply(self.init_weights)
    self.tie_weights()

2行目のBertModel(config)により、configの値に基づき、BERTのモデルが構築されます。
4行目のself.apply(self.init_weights)が実行されると、BERTの重みが初期化されます。

当然、tokenizerを自作したい場合は、別途tokenizerの作成が必要となります。
それ以外の点では、大きな変更は要らないはずです。

おわりに

pytorch-transformersのライブラリを活用して、pre-trainedモデルを学習する方法について確認しました。
今後は、Facebook researchやNVIDIAのBERTの学習方法についても見ていければと思います。