🔰 PyTorchでニューラルネットワーク基礎 #29 【MLM事前学習編・マスク化】

概要

個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。BERTタイプのMLM事前学習に向けて学んだことをまとめおきたいと思います。Transformer Encoderタイプ（BERTタイプ）の事前学習で利用されるMLMの方法を実装、事前学習してみるのが目的となります。

1. マスク化とミニバッチ(１回目・今回)
2. BERTタイプのMLM学習用モデルと重み共有（２回目）
3. 実際に事前学習（３回目）
4. 自作事前学習モデルでファインチューニング（４回目の予定）
5. Datasetクラスの改良（５回目の予定）

方針

1. できるだけ同じコード進行
2. できるだけ簡潔（細かい内容は割愛）
3. 特徴量などの部分，あえて数値で記入（どのように変わるかがわかりやすい）

演習用のファイル

• データのファイル: mlm_pretrain_sample.json
• コード: sample_29.ipynb

1. MLM (Masked Language Modeling) のアイディア

図１：MLMの学習部分

前回確認したように、MLMの基本的なアイディアは、入力文中のトークンの一部を <mask> に置き換え、その <mask> 部分に該当する元のトークンを予測させることになります。これは、文章中の空欄を埋める穴埋め問題として捉えることができました。

BERTの論文では、入力文中の15%のトークンを予測対象に設定しています。ただし、この15%のすべてを単純に <mask> に置き換えるわけでないようです。予測対象となったトークンに対しては

• 80%：<mask> に置き換える
• 10%：別のトークンに置き換える
• 10%：置き換えず、そのまま残す

という作業を行っています。この仕組みにより、モデルは <mask> トークンだけに過度に依存することなく、より頑健な表現を学習できると考えられています¹。

GPT系での事前学習だと、トークナイズされた文をそのままモデルに入力する形ですが、BERT系のMLM事前学習はこのマスク化の部分がやや面倒となります。今回はマスク化の部分とDataLoaderでの学習ループという大枠だけをまとめみました。

事前学習のように大量のデータを一度に処理することはメモリー制約上、難しいため、実際の学習ではデータを小さなまとまりに分けて扱うミニバッチ学習が使われます。PyTorchにおいてミニバッチ学習を実装する際の基本的な構成は、DatasetクラスとDataLoaderクラスを組み合わせる方法です。特に、入力データの形状があらかじめ同一になるよう整えられていれば、TensorDatasetクラスをそのまま利用できます（第14回 【音楽分類02・conv1d】参照）。

2. ミニバッチ学習の流れ

第14回 【音楽分類02・conv1d】で使ったように、ミニバッチ学習の定番は、学習時にtorch.Tensorのx_train（入力するデータ）とt_train（教師データ）をTensorDatasetに入力して、train_data（Datasetの形）に変換します。DataLoaderを使い、ミニバッチのサイズや他の指定を行い、train_dataをミニバッチのサイズに変換しました。基本的な流れは次のような形でした。

ミニバッチ学習の流れ

# x_trainとt_trainはtorch.Tensor
# batch_size = 100 でミニバッチのサイズを指定する
train_data = TensorDataset(x_train, t_train)
train_loader = DataLoader(train_data, batch_size=100, shuffle=True)

# 学習するときのループ
for epoch in range(LOOP):
    for x, t in train_loader:
        # 各ミニバッチで学習処理
        # 学習データx と教師データt をdeviceへ移動
        x = x.to(device)
        t = t.to(device)
        y = model(x)

for x, t in train_loaderという形で学習を進めるタイプとなります。このxが実際にモデルに入力される値となります。PyTorchのいろんな書籍で見かける形ですね😆

MLM事前学習だと入力するxが<mask>化されたもの、教師データtが<mask>部分の正解ラベルのようなものとなります。

for x, t in train_loaderループの中でマスク化処理を行ってもよいのですが、DataLoaderのオプションにtrain_dataに対して前処理を行えるcollate_fnと呼ばれるオプションがあります。 これを利用して、ループ部分はこれまでと同様の形にしたいと思います。

公式のドキュメントには次のように記載されています。

• collate_fn (Callable, optional) – merges a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a map-style dataset.

Preferred NetworksのPLaMo翻訳を使ってみました。

• collate_fn (呼び出し可能オブジェクト、オプション) – サンプルリストを結合して Tensor のミニバッチを形成する。マップスタイルのデータセットからバッチ処理で読み込む場合に使用する。

日本語にしてもわからないということがわかった
しょうがない。コードの確認しかないのか collate_fnを特に指定しないと、default_collateやdefault_convertが使われるっぽい。これらを参考にマスク化の部分を実装してみます。

3. マスク化するcollate関数を作成

最初の印象

• collateってそもそもどんな意味？
• どこでどう使われるの？

とうことで難しそうにみえました。ちなみにcollateは「照合する」という意味らしいけど、たぶん「入力データを整えてミニバッチにまとめる」というニュアンスっぽい？引数と戻り値だけ注意すれば普通の関数でした。

3.1 mlm用のcollate関数の引数を確認

collate関数の引数は、batchと書かれています。Datasetクラスの__getitem__(idx)のリストのようです。

TensorDatasetクラスを利用する予定なのでサンプルデータを読み込んでTensorDataset.__getitem__(idx)の値を確認してみます。

collate関数の引数を確認

import torch
from torch.utils.data import TensorDataset  
import json

# (1) データ読込と必要な部分の切り出し
with open("./data/mlm_pretrain_sample.json", "r") as f:
    data = json.load(f)
# x: IDベクトル
x = torch.LongTensor([item["ids"] for item in data])

# (2) TensorDataset
data = TensorDataset(x)
print(data.__getitem__(0))
#(tensor([1, 332, 8061, 8716, 8267, 8304, 8599, 7615, 4521, 9287, 7486, 211,7722, 7822, 4178, 3913, 8022, 8392, 8263, 3932, 3749,2,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,
# 0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0]),)

説明メモ

• (1) mlm_pretrain_sample.jsonはtextとidsをキーに持つファイルで次のような内容です。

  {'text': 'グルーミング中にとても不安になる犬...', 'ids': [1, 332, 8061, 8716, 8267, 8304, 8599, 7615, ...]}

• dataのids部分をテンソルにしておきます。
• (2) data.__getitem__(0)で中身を確認します。この出力値からcollate関数の引数は[(x1,), (x2,),...,(xi,),...] というタプルのリストにすれば良さそうです。

3.2 mlm用のcollate関数を作成してみよう

1. 入力データするIDリスト(input_ids)については、15%を置き換え対象、その中で「80/10/10」の割合で、「マスク化/別トークン/そのまま」
2. 対応するラベル(labels)については、置き換え対象以外を「ｰ100」、置き換え対象は元のID（正解ラベル）となります
3. 学習時ID「-100」以外のラベルが損失計算の対象となります

このinput_idsとlabelsを出力する形になります。

マスク化するcollate関数

import random

def mlm_sample(batch, mask_prob:float = 0.15):
    """
    batch: [(x1,), (x2,),...,(xi,),...] というタプルのリストの予定
    x1: tensor
    mask_prob: 置き換え対象は15%
    """

    special_tokens = [0,1,2,3,4]  # <pad>, <bos>, <eos>, <unk>, <mask>になる予定
    normal_tokens = range(5,10_000) # 語彙IDのspecial_tokensを除いたもの　語彙数10k
    
    all_input_ids = []   # <mask>化したID列
    all_labels = []      # 15%の予測部分に正解のIDを、それ以外は　−100　で埋める

    # (1) 各サンプルに対してループ処理
    for item in batch:
        original_ids = item[0]                        # item[0]でxiを取り出す
        input_ids = original_ids.clone()              # originalも使うのでcloneしておく
        labels = torch.full_like(original_ids, -100)  # ラベル：一旦　ID = -100 にしてから、該当部分のみ書き換える
       
        # (2) マスク化
        for i in range(len(original_ids)):

            # (3) 特殊トークンの場合は飛ばす
            # original_ids[i]がtorch.Tensorになるので、比較のため、一旦戻す
            token_id = int(original_ids[i].item())  

            if token_id in special_tokens:
                continue

            # (4) 特殊トークン以外15%の確率が置き換え対象
            if random.random() < mask_prob:
                labels[i] = original_ids[i]   # labelsに元のIDを入れる (教師データになる)
                # 置き換え対象の処理 80/10/10 で マスク/他トークン/そのまま
                rand = random.random()
                if rand < 0.8:
                    input_ids[i] = 4 # mask_id 今回はそのまま入力しているけど
                elif rand < 0.9:
                    input_ids[i] = random.choice(normal_tokens) # 特殊トークンID以外のトークンで置き換え
                else:
                    pass # 10%はそのまま
        
        all_input_ids.append(input_ids)
        all_labels.append(labels)

    # (5) リストをスタックして (bs, seq_Len) のテンソルに戻す
    return torch.stack(all_input_ids), torch.stack(all_labels)

説明メモ

• (1) input_ids = original_ids.clone()として別に準備。両方使うからね。教師ラベルは一旦すべて[-100,...,-100]とします。「-100」はCrossEntropyLossで損失するときに飛ばす（利用しない）という特殊IDです。
• (2) マスク化のループ
• (3) special_tokensとして指定したIDは飛ばして次へ！特殊トークンはマスクの対象としない場合の方法となります。
• (4) ここがポイント15%をマスク対象にするぞ！対象となった場合、labels[i]に元のIDを代入するぞ！これで[-100,...,-100]というラベルの必要な部分に正解IDがコピーされます。あとは同じような感じでinput_idsを変更していきます。80%でmask_id、今回は直接「４」を入れています。10%でrandom.choiceで別のトークンを選択！10%はそのままなのでパス。
• (5) リストをテンソルにする。

forループをやめるとか、処理をもう少し格好良くかけると思うのですが、今回はこれで断念🌵

3.3 動作確認

mlm_sampleの動作確認です。適当にID列のテンソルを準備して、

batch = [(x1,), (x2,),...,(xi,),...] というタプルのリスト

に変更後、使ってみます。

確認

org_ids = torch.LongTensor([1, 219, 338, 295, 4188, 6988, 4451, 4218, 4057, 4464, 441, 9127, 4252, 321, 317, 4165, 4369, 160, 2])
x1 = org_ids.clone()
batch = [(x1,)]

# 30%置き換え対象になっています
input_ids, labels = mlm_sample(batch, mask_prob=0.3)

# くどいプリント連打😅
print("変更されたID列")
print(f"{input_ids=}")
print(f"{input_ids.shape=}")
print("対応するラベル")
print(f"{labels=}")
print(f"{labels.shape=}")

print("変換部分を表示")
ids = torch.where(input_ids[0] != org_ids)

for num in ids[0]:
    print(f"input: {input_ids[0][num].item()},\t label: {labels[0][num]}")

置き換え割合を30%にして多様性を確保してみたのですが 今回の例では、<mask>と別のトークへの置き換えになっています。

変更されたID列
input_ids=tensor([[1, 219, 2200,  295, 4, 6988, 4, 6317, 4057, 4464,  441, 9127, 4, 321, 317, 4, 4369, 4, 2]])
input_ids.shape=torch.Size([1, 19])
対応するラベル
labels=tensor([[-100, -100,  338, -100, 4188, 6988, 4451, 4218, -100, -100, -100, -100, 4252, -100, -100, 4165, -100,  160, -100]])
labels.shape=torch.Size([1, 19])

変換部分を表示
input: 2200, label: 338
input: 4,	 label: 4188
input: 4,	 label: 4451
input: 6317, label: 4218
input: 4,	 label: 4252
input: 4,	 label: 4165
input: 4,	 label: 160

教師ラベルも入力データも同じ形状になっています。変換部分（-100以外の部分）だけが損失計算の対象となります。

3.4 利用するクラスのちょっとしたまとめ

TensorDatasetクラスについて

使い方
data = TensorDataset(x)
data = TensorDataset(x,t)

• TensorDatasetクラスはデータ(x, t)のパッキングを担当
• テンソルxやtの系列長が等しい場合、TensorDatasetクラスをそのまま利用することができる
• 引数：形状が等しいTensorのみ
• JSONなどを引数に使う場合は、Datasetクラスを自作する必要がある

DataLoaderクラスについて

使い方
train_loader = DataLoader(
　　　　data,
　　　　batch_size=100,
　　　　shuffle=True,
　　　　collate_fn=mlm_sample, # ここで作成した関数を指定
　　　　)

• DataLoaderクラスはデータの前処理と出力を担当
• Datasetクラスで作成したデータに対して、ミニバッチ化、ドロップラスト、collate_fnでのデータ処理（ミニバッチにまとめる時の挙動）を扱う
• collate_fn関数は、指定がない場合、default_collater（入力データをテンソルに変換する働き）が使われる
  • numpyで読み込んだデータがTorch.Tensorになるなど

4. データ読み込みから学習ループまでの流れ

最後にDataLoaderでうまく使えるのかを確認して終わりにします。

まとめ

import torch
from torch.utils.data import DataLoader, TensorDataset
import json
import random

# (1) mlm_sample関数　（省略）
def mlm_sample(batch, mask_prob:float = 0.15):
    """ 省略 """
    return torch.stack(all_input_ids), torch.stack(all_labels)

# (2) サンプルデータ
with open("./data/mlm_pretrain_sample.json", "r") as f:
    data = json.load(f)
x = torch.LongTensor([item["ids"] for item in data])

# (3) TensorDatasetにする
train_data = TensorDataset(x)

# (4) DataLoaderを使う
train_loader = DataLoader(
    train_data, 
    batch_size=100, 
    shuffle=True, 
    collate_fn=mlm_sample,   # ここで作成した関数を指定
    num_workers=4
)

# (5) 学習ループ
for epoch in range(2):
    for x, t in train_loader:
        x = x.to(device)
        t = t.to(device)
        print(f"{epoch}: x: {x.shape},\tt: {t.shape}")
        print(f"x[0]:\n{x[0]}")
        print(f"t[0]:\n{t[0]}\n")
        break

説明メモ

• (1) mlm_sample関数は省略します
• (2) サンプルデータの読み込み。textキーに文章、idsキーに等長IDベクトル。
• (3) TensorDatasetクラスを利用します。引数は同じ形状のテンソルが要求されます。
• (4) DataLoader: 今回のメインで collate_fn オプションに作成したmlm_sampleを指定します。
• (5) 確認用の学習ループ。xがinput_idsで、tがlabelsになります。
• 次の出力結果を見るとマスク部分はID「４」に変更されています。ラベルでは、マスク部分以外は「-100」に変わっていることも確認できます。大丈夫そうですね🌸

出力結果

0: x: torch.Size([100, 64]),	t: torch.Size([100, 64])
x[0]:
tensor([   1, 9193, 3883, 7873,    4, 8645, 9583, 7678,  211, 9603, 3749,    2,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0], device='cuda:0')
t[0]:
tensor([-100, -100, -100, -100, 7698, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100], device='cuda:0')

1: x: torch.Size([100, 64]),	t: torch.Size([100, 64])
x[0]:
tensor([   1, 7585, 4226, 3880, 4039, 9215, 8304,    4,    4, 7554, 7514, 3932,
           4,    4, 5622, 8092, 7805, 3867, 7734, 7730, 3879, 3874,    4, 9892,
        3913, 5167,    4,    4,    2,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0], device='cuda:0')
t[0]:
tensor([-100, -100, -100, -100, -100, -100, -100, 9797, 3878, -100, -100, -100,
         210, 7715, -100, -100, -100, -100, -100, 7730, -100, -100, 7546, -100,
        -100, -100, 8218,  211, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
        -100, -100, -100, -100], device='cuda:0')

わざわざfor x, t in train_loaderとするまでもなく、もっと簡単に確認することもできます。DataLoaderから先頭のバッチを取得する方法です。

# batch は collate_fn (create_mlm_sample) の戻り値のタプル
batch = next(iter(train_loader))

# mlm_sample の出力が(inputs, labels) の２つなので
inputs, labels = batch

print(f"input[0]:\n{inputs[0]}")
print(f"label[0]:\n{labels[0]}")
#input[0]:
#tensor([   1, 1367, 8882, 4020, 4101, 8097,    4,    4,  548, 3867,    4, 3878,
#        5566, 8003, 9217,  210,    4, 4867, 8506, 5738, 9119, 4462,    4, 3992,
#        9614, 8220,  211,  548, 3892, 4039, 9430, 3886, 6411, 4602,    4, 4032,
#        7684,  211, 8720, 7626,  133, 1926, 7550, 7541,    4,  134,  969, 4322,
#        7665,    4, 3910, 4196, 9976, 7716, 7748, 3866, 5173, 3928, 7490, 9838,
#        7847, 7557, 3721,    2])
#label[0]:
#tensor([-100, -100, -100, -100, -100, -100, 9761,  211, -100, -100, 4101, -100,
#        -100, -100, -100, -100, 9774, -100, -100, -100, -100, -100, 3880, -100,
#        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 8229, -100,
#        -100, -100, -100, -100, -100, -100, -100, -100, 8311, -100, -100, -100,
#        -100, 7583, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
#        -100, -100, -100, -100])

これで事前学習データをモデルに入力する準備か整いました。次は、MLM用のネットワークモデルを作成する段階に移りたいと思います。

次回

重みの共有についてです²。

1. マスク化とミニバッチ(今回)
2. BERTタイプのMLM学習用モデルと重み共有（次回）
3. 実際に事前学習（３回目の予定）
4. Datasetクラスの改良（４回目の予定）

目次ページ

注

1. 15%を置き換え対象とする数値ですが、経験則？BERTのパラメータサイズがLargeのときは40%を置き換え対象にしたほうが良いという結果もあるようです。 ↩

2. 頑張って理解を深めたいものですが、時間との兼ね合いも難しいな〜。有益な情報を提供している人、たくさん記事書いている人、すごいな〜って思います ↩