BERTのMasked LMを実行するために行ったことをまとめてみます。

こんにちは、技術開発ユニットの三浦です。

インターネットで動画を見ながら、Blenderという3DCGソフトで3Dモデルを作成する練習をしています。色々なものが出来るととても楽しいのですが、子どもから見ても面白いみたいで、よく後ろから様子をのぞいています。「作ってほしいものリクエスト」があるみたいなので、早く上手になって、リクエストに答えられるようになりたいです。

さて、最近私はトークン化された文章データを使ってBERTのMasked LMという自己教師あり学習によって、モデルに文章データの持つ構造を理解させることが出来るのかどうか、ということに取り組んでいます。色々悩みつつ、ひとまず学習処理を実行出来るところまではたどり着いたので、今回はそのお話をさせて頂きます。

BERT

まずはBERTについて簡単に紹介します。

BERT(Bidirectional Encoder Representations from Transformers)は2019年の論文「BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding」に登場する自然言語のモデルです。

[1810.04805] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

"モデル"と聞くとその構造に興味が行きがちですが、BERTのモデルの構造自体は、それ以前に発表された論文「Attention Is All You Need」の翻訳タスクに使用されたTransformerモデルのEncoderが使われています。 BERTの面白い点はモデルの学習部分にあり、まず言語の構造をモデルに理解させる「pre-training」を行い、次に各タスクにモデルを特化させる「fine-tuning」を行うという、2段構えのステップで構成されています。

「pre-training」はさらに2つのタスクを組み合わせています。1つが「Masked LM」で、文章中の単語(トークン)をランダムに隠し(Maskし)学習用データとして入力し、Maskした箇所にどの単語が入るのかを予測するタスクです。もう1つが「Next Sentence Prediction」で、2つの文章をつなげて入力し、2つ目の文章が1つ目の文章の続きであるのかどうかを1つ目の文章から予測するタスクです。BERTを質疑応答や文章生成タスクへ応用することを想定して設定されたタスクであると考えられます。

私が取り組んでいること

約280,000件のトークン化された文章のデータについて、その構造をBERTを使って紐解くことが出来ないか、と考えています。それぞれの文章は独立しているので、Next Sentence Predictionは適用することが出来ないと考え、まずMasked LMでpre-trainingしようとしています。

kerasでMasked LMを行うコードの実装例が公開されており、こちらのコードを参考にしました。

End-to-end Masked Language Modeling with BERT

モデルの品質などはまだ精査出来ていないものの、ひとまず用意したデータでBERTモデルを作るところまでは実行することが出来ました。実行するまでにいくつか考えなければいけない点もありましたので、今回はそれらをまとめたいと思います。

実装

データの準備

用意したデータはトークン化までは済んでいるのですが、各データで文章の長さが異なる状態です。BERTモデルへの入力は長さがそろったデータでなければならないため、まずは入力する文章の長さを決め、それにデータを揃えることから始めました。決めた長さに満たないデータはトークンID"0"でパディングし、決めた長さを超えるデータは後ろを削除する処理を施しました。例えば以下のような処理を行います。

max_length = 256

def align_text(x):
  """max_lengthの長さに満たないデータは末尾に0を付与
  長さを超えるデータはカットする
  """
  padding_len = max_length - len(x)
  if padding_len < 1:
    padding_len = 0 
  for _ in range(pad_len):
    x = np.append(x, 0)
  
  return x[:max_length]

aligned_data = [align_text(x) for x in raw_data]
aligned_data = np.array(aligned_data)

また、Masked LMを行う際にはデータをランダムにMaskする必要がありますが、そのMaskにあたるトークンIDに何を使うのかも決める必要があります。データ内には存在しないトークンIDとして、トークンIDの最大値に+1をした値をMask用のトークンIDとしました。

Masked LM学習用のデータ

Masked LM学習時にモデルに入力される学習用データは、1つの文章データに対し、以下の3つのデータが生成され、使われます。

y_labelsは元のデータそのものです(長さを揃える処理は実行済みです)。これが自己教師あり学習の正解データになります。それに対しencoded_texts_maskedはy_labelsとほぼ同じように見えますが、ところどころのトークンIDが変更されています。これがMaskされた状態のデータです。トークンのMask処理は以下のルールで行われます。

データ内のトークンに対し、15%の確率でMask候補に選択する
そのうち90%をMask化する
そのうち10%をランダムなトークンに変更する

これは先のkerasの実装例のルールに従っているのですが、BERTの論文ではMask候補のうち80%をMask化、10%をランダム化、10%を変更しないという変換方法が取られています。

sample_weightsはモデルのlossを計算するときに使う0と1で構成された配列で、データに対し、どこのトークンに対する予測をlossの計算に使用するのかを指定するために使用します。

データのTFRecord化

学習に使用するデータが280,000件とかなり大きいため、全てをGPUメモリに乗せることが出来ないことが判明しました。そこでデータを一度TFRecordという、kerasのバックエンドのTensorFlow専用のファイル形式に出力し、そこから都度読み込ませるようにしました。

TFRecordファイルを作成して、Datasetsとして読み込むまでのコードを掲載します。まずはnumpy.arrayオブジェクトをTFRecordファイルに書き出す部分です。

def serialize_row(row):
  """
  numpy.arrayの1行からMasked LM用のデータを作り、TFRecord用のデータに変換する
  """
  feature = {}
  (encoded_texts_masked, y_labels, sample_weights) = get_masked_input_and_labels(row)
  feature['encoded_texts_masked'] = tf.train.Feature(int64_list=tf.train.Int64List(value=encoded_texts_masked))
  feature['y_labels'] = tf.train.Feature(int64_list=tf.train.Int64List(value=y_labels))
  feature['sample_weights'] = tf.train.Feature(int64_list=tf.train.Int64List(value=sample_weights))
  example = tf.train.Example(features=tf.train.Features(feature=feature))
  serialized = example.SerializeToString()
  return serialized

#学習用データのTFRecordファイルの書き出し
with tf.io.TFRecordWriter(f'{file_path}/train_tf_file.tfrecords') as writer:
  for r in train_data:
    writer.write(serialize_row(r))

#テスト用データのTFRecordファイルの書き出し
with tf.io.TFRecordWriter(f'{file_path}/test_tf_file.tfrecords') as writer:
  for r in test_data:
    writer.write(serialize_row(r))

次にTFRecordファイルから読み込む部分です。

# TFRecordに含まれる特徴量の情報
# config.MAX_LENは256、config.BATCH_SIZEは64に設定
features = {
  'encoded_texts_masked': tf.io.FixedLenFeature([config.MAX_LEN], tf.int64),
  'y_labels': tf.io.FixedLenFeature([config.MAX_LEN], tf.int64),
  'sample_weights': tf.io.FixedLenFeature([config.MAX_LEN], tf.int64),
}

def _parse_function(example_proto):
  """
  TFRecord用に変換されたデータをパースして復元する
  """
  parsed_features = tf.io.parse_single_example(example_proto, features)
  return parsed_features['encoded_texts_masked'],parsed_features['y_labels'],parsed_features['sample_weights']

train_datasets = tf.data.TFRecordDataset(f'{file_path}/train_tf_file.tfrecords')
train_datasets = train_datasets.map(_parse_function).batch(config.BATCH_SIZE)

test_datasets = tf.data.TFRecordDataset(f'{file_path}/test_tf_file.tfrecords')
test_datasets = test_datasets.map(_parse_function).batch(config.BATCH_SIZE)

これで全てのデータをメモリにロードするのではなく、ファイルから少しずつロードし、メモリを圧迫しないようになりました。あとはkerasの実装例に従ってモデルの学習処理を行うことが出来ました。

モデルをテストする

テスト用のMaskされたデータをモデルに入力し、Maskされた部分に対しどのようなトークンを予測するのかを試してみました。Mask部分に対するモデルの予測トークンを確認するコードは以下のようになります。

#テストデータのbatchのうち、sample番目のデータのMaskに対する予測を求める
#例えばindex=3のサンプルの結果を見る場合
sample = 3
sample_data = next(iter(test_datasets))

#sample_dataはtupple
#Maskデータは0番目の要素
test_x = sample_data[0][sample]

#モデルの予測取得
predict = bert_masked_model.predict(test_x[np.newaxis,:])
#Maskトークンの位置を取得し、Mask部分だけの予測を取得する
test_mask_index = np.where(test_x == config.VOCAB_SIZE - 1)
mask_predict = predict[0][test_mask_index]

#最初のMaskトークンに対する予測スコア上位5つのトークンを取得
top_predict_token = mask_predict[0].argsort()[-5 :][::-1]

別環境にあるトークン復元用のデータを使って予測結果を復元してみたのですが、ぴったりとMaskされたトークンを予測することは難しいものの、モデルが予測したトークンを当てはめても自然な文章になるような印象を受けました。概ね文章の構造を理解させることは出来たのでは・・・という手ごたえを感じています。

まとめ

ということで、今回は自然言語モデルBERTのMasked LMを実行し、文章データの構造を紐解こうとしている話を紹介しました。BERT全体で見るとまだ入口に立てた段階で、このモデルを他のタスクに向けてfine-tuningしたらどうなるのか、などなどまだまだ検証してみたいことがたくさんあります。引き続き色々と試していきたいと思います！