🔰 PyTorchでニューラルネットワーク基礎　#30 【MLM事前学習編・モデル】

Last updated at 2026-05-27Posted at 2026-04-30

概要

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

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

BERTタイプのMLM学習用のネットワーク（下図のようなネットワーク）を構築するのが主な目的となります。前回のマスク化する関数を利用して事前学習してみるのは次回です。

方針

できるだけ同じコード進行
できるだけ簡潔（細かい内容は割愛）

演習用のファイル

データのファイル: sample_tokenizer_10k.json
コード: sample_30.ipynb

1. モデルの基本構造

BERTタイプのネットワークを記述する際、トークンのIDやネットワーク内部の数字がどうしても入り乱れます。基本方針の３番目「特徴量などの部分，あえて数値で記入」をやめる時がきたようです。設定をまとめるクラスを作成して、表記を簡潔にしたいと思います。

1.1 モデル設定用のクラス

すでに適当なトークナイザーがあるという状況から始まります。モデル設定用のクラスと言っても、BERTタイプのネットワーク構造に利用しそうな変数を羅列するだけの単純なクラスです。たぶん、こういうの名前が大事なんだろうなって思うよ

ModelConfigクラスには次の3種類をまとめておきました。

モデル構造を指定する定数
特殊トークンなどのID（tokenizerに依存する）
学習時に使う定数

これで、モデル本体や学習ループなどを記述する際に使いやすくなるはず。

モデル設定用のクラス

class ModelConfig:
    def __init__(self, tokenizer):
        # (1) モデル構造
        self.vocab_size = tokenizer.get_vocab_size()
        self.seq_len = 64
        self.d_model = 64
        self.nhead = 4
        self.dim_feedforward = 256
        self.num_layers = 6
        self.dropout = 0.1
        
        # (2) 特殊トークンID
        self.pad_token_id = tokenizer.token_to_id("<pad>")
        self.mask_token_id = tokenizer.token_to_id("<mask>")
        self.bos_token_id = tokenizer.token_to_id("<bos>")
        self.eos_token_id = tokenizer.token_to_id("<eos>")
        self.unk_token_id = tokenizer.token_to_id("<unk>")

        # (3) 特殊トークンのセット
        self.special_tokens_set = {
            self.pad_token_id,
            self.mask_token_id,
            self.bos_token_id,
            self.eos_token_id,
            self.unk_token_id,
        }
        
        # (4) 通常トークンのリスト special_tokenを除くトークンのリスト（MLMランダム置換用）
        self.normal_tokens_list = [
            i for i in range(self.vocab_size) 
            if i not in self.special_tokens_set
        ]
        
        # (5) 学習設定 (今回は利用しないけど使うと便利かも)
        self.batch_size = 512
        self.learning_rate = 0.001
        self.num_epochs = 100
        self.mask_prob = 0.15
        self.max_grad_norm = 1.0

メソッドもない、ただ属性（変数名）となるものをひたすら列挙しただけのクラスです

説明メモ

(1) 語彙数、系列長、Transformerで使う特徴量の次元やヘッド数、レイヤー数などを設定しています。
(2) トークナイザーで指定した特殊トークンの変数名。
(3) 特殊トークンの集合。有無の判定に利用したいので集合で対応。
(4) 特殊トークン以外のトークンのリスト。トークンをランダムに置き換えるときリストのほうが便利なのでリストで対応。
(5) 学習用の設定

使い方

適当なtokenizerを指定して、次のような形で利用します。

config = ModelConfig(tokenizer)

これで、config.属性名で気軽にアクセスできます。

config.bos_token_id: <bos>トークンのIDの
config.special_tokens_set: 特殊トークンの集合

2. モデルの構造

系列長のIDベクトルに対して、埋め込み層、Transformer Encoder層を経て、各IDに対する特徴量が求まります。特徴量を全結合層で語彙数に変換して、<mask>部分に該当する単語を予測する形が基本形となります。

図１：ネットワークの基本構造

一見複雑にみえますが、第22回【文章分類・Transformer】のネットワークとほぼ同じです。しっかりとモデルが学習されていれば、最終部分（図１のmlm_headブロック）を付け替えることで様々なタスクに応用できるよねという考えに到達するのか〜。

ネットワーク構造を決めるモデルクラスは、今までのように数値をそのまま貼り付ける形から変わります。ModelConfigクラスで指定した数値を内容を受け取る形になります。

2.1 単純なモデル

最終出力部分を全結合層のみで表す単純なモデルで大まかなネットワーク構造を確認してみたいと思います。その後BERTっぽい形のモデルに変更してみます。

最終出力層が全結合のみのバージョン

class DNN(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        # (1) 埋め込み層
        self.token_embedding = nn.Embedding(
            num_embeddings=config.vocab_size, 
            embedding_dim=config.d_model,
            padding_idx=config.pad_token_id
        )
        self.pos_embedding = nn.Embedding(num_embeddings=config.seq_len, embedding_dim=config.d_model)
        
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        
        # (2) Transformer Encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.dim_feedforward,
            dropout=config.dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers, enable_nested_tensor=False)
        
        # (3) MLM用の出力層：Linearを利用して各トークン位置での語彙を予測
        # self.mlm_head = MLMHead(config, self.token_embedding.weight)
        self.mlm_head = nn.Linear(config.d_model, config.vocab_size)
    
    def forward(self, x):
        # (4) マスクの作成
        src_key_padding_mask = (x == self.config.pad_token_id)
        
        # (5) 埋め込み
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(torch.arange(x.size(1), device=x.device))
        x = tok_emb + pos_emb.unsqueeze(0)
        
        x = self.layer_norm(x)
        x = self.dropout(x)
        
        # (6) Transformer Encoder
        h = self.transformer_encoder(x, src_key_padding_mask=src_key_padding_mask)
        
        # (7) MLM予測：各位置で語彙全体での確率を出力
        logits = self.mlm_head(h)  # [batch, seq_len, vocab_size]
        
        return logits

説明メモ

(1) トークンの埋め込み、位置情報の埋め込みです。 config.vocab_size という変数名で値を指定します。
(2) Transformer Encoderの指定です。TransformerEncoderLayerで層の中身を指定、TransformerEncoderで繰り返す回数を指定しました。詳細は第22回。
Transformerの出力は、各トークンに対する、特徴量（d_model次元）となります。バッチサイズ１の場合は、(系列長, d_model)という形になります。
enable_nested_tensor=Falseオプションの指定。可変長の系列（長さがバラバラなデータ）を効率よく処理するための仕組みのようで、デフォルトはTrueとなります。今回は等長化しているので、Falseで学習させます。Trueだと、確認する際に警告文が表示されることがありました。
(3) 全結合層よって、d_model次元の特徴量を語彙数次元へ変換します。
(4) <pad>IDをマスクする部分。
(5) トークン・位置の埋め込み層の値を足し算します。
(6) <pad>部分のマスクであるsrc_key_padding_maskを使い、Transformerの計算をします。
(7) mlm_headとなっていますが、今回は、単なる全結合層です。（バッチサイズ，系列長，語彙数次元）がモデルの出力値となります。

2.2 重み共有のモデル

上記のネットワークで、MLM事前学習の基本的な枠組みが完成となります。論文中では強調されていませんが、BERTの実装ではMLM用の出力部分にも工夫をしているようです。

Inan, Khosravi and Socher (2016) "Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling"
Press and Wolf (2016) "Using the Output Embedding to Improve Language Models"

トークンの埋め込み層 (token_embedding) の重みを全結合層の重みとして利用する方法です¹。層を増やしつつも学習パラメータ数が削減され、しかも、言語モデルとしての精度を向上させるという方法になります。

PressとWolfは「重みを共有することで、入力としての単語の意味と出力としての単語の予測が同じベクトル空間で最適化され、汎化性能が向上する。」ことを論文で主張しています。理論的な裏付けとなるのが、Inan, Khosravi and Socherは「Augmented Lossと呼ばれるもの損失関数を最小化すると埋め込み層の行列Wと最後の全結合層の重み$L^T$が等しくなる」という主張のようです

重み共有 (weight tying) の意図？ (個人的な見解)

学習が進むと、入口のトークンに意味的に近い単語が出力側に登場するはず。単語の辞書みたいな埋め込み層の重みを出力側で利用することで語彙表現に一貫性がでるっぽい。
図のように「りんご」のスコアを大きくする特徴量が学習されるはず。分散表現が入出力時点で共通にすることで、類似トークンの内積が大きくなるように学習が進む予定ということなのかな？

これらの点からマスク部分を予測するタイプだと重み共有の効果がありそうです。次のコードが、単純なモデルのmlm_headを置き換える「重み共有のネットワーク」のクラスとなります。

BERTタイプのmlm_head (weight tying)

# MLM用の出力部分
# nn.Linear(d_model, vocab_size)の拡張版

class MLMHead(nn.Module):
    def __init__(self, config: ModelConfig, token_embedding_weight):
        super().__init__()

        # (1) 準備 正規化
        self.fc = nn.Linear(config.d_model, config.d_model)
        self.act = nn.GELU()
        self.ln = nn.LayerNorm(config.d_model)
        
        # (2) 重みの共有準備、サイズの指定だけ
        self.classification_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # (3) 重みの共有  実際に重みをコピーする部分
        self.classification_head.weight = token_embedding_weight  # weight tying
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

    def forward(self, x):
        x = self.fc(x)
        x = self.act(x)
        x = self.ln(x)
        # (4) 共有した重み＋学習するバイアス項
        x = self.classification_head(x) + self.bias
        return x

一見、ネットワーク層が増えて学習対象のパラメータも増加する予感がしますが...　Linearの重みが固定されているので、学習に利用されるパラメータ数は減少することになります。

説明メモ

(1) Transformer層の出力は（系列長 × d_model）の形の特徴量です。これを一旦、正規化します。（モデルの形としては利用しなくてもOKな気がする。多分、いろいろ学習時に恩恵があるタイプ）
(2) 最終的な出力を司る全結合層が classification_head です。d_modelをvocab_sizeに変換します。バイアス項をあえてFalseにして、後で足す形にしました。バイアス項のパラメータは学習対象ということを表現してみたかっただけです²。
(3) classification_head での重み classification_head.weightに、トークン埋め込み層の重み token_embedding_weightを使います。
classification_headは、vocab_size次元が出力次元でした。このサイズに注意して(2)で指定したLinearのバイアス項を、適当に作ります。
(4) いわゆる全結合層の形、$xW+b$の形にします。

2.3 最終的な構造³

今までの部分を全部つなげたものとなります。

ModelConfig: モデル設定のクラス
MLMHead: 重み共有を考えた最終層のブロック
DNN: ネットワーク構造をまとめた部分（毎回同じパターンで申し訳ない）

モデルの全体像となります。クリックで展開

class ModelConfig:
    def __init__(self, tokenizer):
        # モデル構造
        self.vocab_size = tokenizer.get_vocab_size()
        self.d_model = 64
        self.seq_len = 64
        self.nhead = 4
        self.dim_feedforward = 256
        self.num_layers = 6
        self.dropout = 0.1
        
        # 特殊トークンID
        self.pad_token_id = tokenizer.token_to_id("<pad>")
        self.mask_token_id = tokenizer.token_to_id("<mask>")
        self.bos_token_id = tokenizer.token_to_id("<bos>")
        self.eos_token_id = tokenizer.token_to_id("<eos>")
        self.unk_token_id = tokenizer.token_to_id("<unk>")

        # 特殊トークンのセット
        self.special_tokens_set = {
            self.pad_token_id,
            self.mask_token_id,
            self.bos_token_id,
            self.eos_token_id,
            self.unk_token_id,
        }
        
        # 通常トークンのリスト special_tokenを除くトークンのリスト（MLMランダム置換用）
        self.normal_tokens_list = [
            i for i in range(self.vocab_size) 
            if i not in self.special_tokens_set
        ]
        
        # 学習設定 (今回は利用しないけど使うと便利かも)
        self.batch_size = 512
        self.learning_rate = 0.001
        self.num_epochs = 100
        self.mask_prob = 0.15
        self.max_grad_norm = 1.0

# MLM用の出力部分
# nn.Linear(d_model, vocab_size)のみからの拡張版

class MLMHead(nn.Module):
    def __init__(self, config: ModelConfig, token_embedding_weight):
        super().__init__()
        self.fc = nn.Linear(config.d_model, config.d_model)
        self.act = nn.GELU()
        self.ln = nn.LayerNorm(config.d_model)

        self.classification_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.classification_head.weight = token_embedding_weight  # weight tying

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

    def forward(self, x):
        x = self.fc(x)
        x = self.act(x)
        x = self.ln(x)
        x = self.classification_head(x) + self.bias
        return x

# MLM用のモデル（出力層を変更）
class DNN(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        # 埋め込み層
        self.token_embedding = nn.Embedding(
            num_embeddings=config.vocab_size, 
            embedding_dim=config.d_model,
            padding_idx=config.pad_token_id
        )
        self.pos_embedding = nn.Embedding(num_embeddings=config.seq_len, embedding_dim=config.d_model)
        
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        
        # Transformer Encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.dim_feedforward,
            dropout=config.dropout,
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers, enable_nested_tensor=False)
        
        # MLM用の出力層：BERT風レイヤー　各トークン位置で語彙全体を予測
        self.mlm_head = MLMHead(config, self.token_embedding.weight)
        #self.mlm_head = nn.Linear(config.d_model, config.vocab_size)
    
    def forward(self, x):
        # マスクの作成
        src_key_padding_mask = (x == self.config.pad_token_id)
        
        # 埋め込み
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(torch.arange(x.size(1), device=x.device))
        x = tok_emb + pos_emb.unsqueeze(0)
        
        x = self.layer_norm(x)
        x = self.dropout(x)
        
        # Transformer Encoder
        h = self.transformer_encoder(x, src_key_padding_mask=src_key_padding_mask)
        
        # MLM予測：各位置で語彙全体での確率を出力
        logits = self.mlm_head(h)  # [batch, seq_len, vocab_size]
        
        return logits

ネットワーク構造が完成したので、

適当なトークナイザーを読み込んで、
モデル設定
学習用のモデル構築

という流れになります。

tokenizer = Tokenizer.from_file(tokenizer_file)
config = ModelConfig(tokenizer)
model = DNN(config)

１回目でデータのマスク化処理、２回目でモデル構築まで終了しました。次回は実際に事前学習してみるになります。

次回

BERTタイプのモデルで事前学習してみる内容になります。

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

目次ページ

注

読めばなるほど！と思うのですが、こういう発想実態を思いつくのがすごいですよね ↩
Linear(config.d_model, config.vocab_size)とすれば、バイアス項指定の部分やforwardでバイアス項を足し算している部分を短く記述できると思います。今回は後で自分が見てわかるように冗長に書いてあります ↩
この部分だけ実行しても、tokenizerを指定していないのでエラーになります。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up

🔰 PyTorchでニューラルネットワーク基礎 #30 【MLM事前学習編・モデル】

概要

演習用のファイル

1. モデルの基本構造

1.1 モデル設定用のクラス

2. モデルの構造

2.1 単純なモデル

2.2 重み共有のモデル

2.3 最終的な構造3

次回

目次ページ

注

🔰 PyTorchでニューラルネットワーク基礎　#30 【MLM事前学習編・モデル】

2.3 最終的な構造³