🔰PyTorchでニューラルネットワーク基礎 #41 【GPT編・事前学習２】

概要

個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。GPTタイプ （自己回帰型の言語モデル）の事前学習の演習で学んだことを第38回からまとめています。これまではTransformerブロックの自作をせず、PyTorchの関数を使いできるだけ簡単な実装を心がけていました。今後の拡張を踏まえて、Transformerブロックの実装から取り組んでいきます。

確認したい事項	達成度
学習データの再現・意味のある文章を生成しているか？	◎
学習データに類似の文からの意味のある文章を生成できるか？	△〜◯
学習分野外の内容から文章を生成できるか？	✕

次のような感じの文章となります。

図１：生成結果のサンプル（詳細は検証部分）

扱う内容

1. Transformerブロックの実装
2. GPTタイプの事前学習
3. 日本語データでの学習

Transformerブロックを実装するということで、３種類あった方針もとうとう一つだけになってしまいました。

方針

• できるだけ同じコード進行

演習用のファイル

• トークナイザー: livedoor_tokenizer_10k.json

  • Livedoorニュースコーパスをデータとしてbytelevel BPEで学習した10Kトークナイザー

• 学習データ：livedoorニュースコーパスのlivedoor-itカテゴリー

• 事前学習済みモデル：it_seq_256_bpe_10k.model

  • 系列長256のモデル。seq_len=256で動作します。

• コード: sample_41.ipynb

1. データ

livedoorニュースコーパスのlivedoor ITカテゴリーに分類されているテキストデータを使います。第38回でHOMMEカテゴリーで試したので他のデータでも試してみます。

• すでにある程度の前処理が行われており、非常に使いやすい点
• ニュース記事なので、きれいな文章になっている点
• 同分野がまとまっている点
• ファイルごとに内容がまとまっている点

livedoorのニュースコーパスは9つの分野に分かれていると思います。他の分野も混ぜると学習時間も９倍になります。１分野固定がお気楽😆

前処理

• 文頭から3行のURL、日付、タイトルの削除 (演習はこの状態で事前学習しました)
• 3行取り除いたテキストファイル(870個くらい)を/data/it_texts/フォルダへ保存します。テキストファイルのまま利用できるようにDatasetクラスの方で調整していきます。
• ライセンス的に改変したデータをアップロードできないのが残念 😓😓

前処理の追加編の参考例
ローカルLLMを利用して加工すると好みの文体に統一できると思います。

• 文体を「です・ます」調に調整
• 体言止もできる限り、文になるように調整
• 商品名の置き換えなどなど

2. トークナイザー

トークナイザーも自作します¹。ByteLevelのBPEで、前処理済みデータを利用してトークン数10,000で学習してみました。特殊トークンは、文章の区切りを表す<|endoftext|>的なものが一つあれば事前学習と文章生成までできるのですが、指示チューニングも視野に入れて、いろいろ追加しています😆

他のカテゴリーのデータでも検証したいので、その都度トークナイザーを作成するのも面倒ですし9分野全部のデータを利用してトークナイザーの学習してしまいました。

【クリックで展開】トークナイザー生成コードの全体像となります。

import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, normalizers

tokenizer_filename = "./tokenizer/livedoor_10k.json"  # tokenizerのファイル名
data_dir = "./data/"                # 学習データの保存ディレクトリー
vocab_size = 10_000                 # 語彙数


tokenizer = Tokenizer(models.BPE())       # (1) BPE
tokenizer.normalizer = normalizers.NFKC() # (2) 正規化
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) # (3) ByteLevel

# (4) Trainer
trainer = trainers.BpeTrainer(
    vocab_size=vocab_size,
    special_tokens=["<pad>", "<eod>", "<system>", "<assistant>", "<user>"],
    min_frequency=2,
    initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
)

# (5) data_dirディレクトリ内のtxtファイルを取得
data_files = sorted(glob.glob(f"{data_dir}/*.txt"))

# (6) 複数のtxtファイルを順番に読み込む
def iterator(filenames):
    for filename in filenames:
        with open(filename, "rt", encoding="utf-8") as f:
            for line in f:
                yield line

# (7) 学習
tokenizer.train_from_iterator(
    iterator(data_files),
    trainer=trainer
)

# (8) 保存
tokenizer.save(tokenizer_filename)
print("学習終了")

学習したトークナイザーの特殊トークンと語彙数を確認しておきましょう😆

from tokenizers import Tokenizer

tokenizer_filename = "./tokenizer/livedoor_10k.json"  # 利用するトークナイザー
tokenizer = Tokenizer.from_file(tokenizer_filename) 

print("特殊トークンID:")
print(f"<pad>: {tokenizer.token_to_id('<pad>')}")              # <pad>: 0
print(f"<eod>: {tokenizer.token_to_id('<eod>')}")              # <eod>: 1
print(f"<system>: {tokenizer.token_to_id('<system>')}")        # <system>: 2
print(f"<assistant>: {tokenizer.token_to_id('<assistant>')}")  # <assistant>: 3
print(f"<user>: {tokenizer.token_to_id('<user>')}")            # <user>: 4
print(f"size: {tokenizer.get_vocab_size()}")                   # size: 10,000

文の始まり、終わりを表す特殊トークンを導入するなどもう少し工夫の余地があると思います。特殊トークンの影響なども実験してみたい限り🍀

3. モデル設定と学習

GPTタイプの学習は、「次の単語予測問題」を解くことです。これまで通り次の手順で進めていきます。

• 準備 (3.1)
• ネットワークモデルの設定や定義　(3.2)
• Dataset・DataLoaderによるデータの読み込み (3.3)
• 誤差関数と誤差最小化の手法の選択　(3.4)
• 変数更新のループ　(3.5)
• 検証　(3.6)

3.1 準備

利用するライブラリーを読み込みます。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer
from pathlib import Path
import math


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"{device=}")

利用するデータセットのディレクトリー、トークナイザー、学習したモデルの保存ファイル名を指定します。

data_dir = "./data/it_texts/"     # 学習用データのディレクトリ
tokenizer_filename = "./tokenizer/livedoor_tokenizer_10k.json"  # 利用するトークナイザー
pretrain_filename = "model/it_seq_256_bpe_10k.model"  # 保存するモデル名
tokenizer = Tokenizer.from_file(tokenizer_filename)   # トークナイザー

ライブラリーは他にも利用します。学習に影響しないものは利用するときに直接記述する形で表現しています。

3.2 ネットワークモデルの設定や定義

ModelConfigクラスで、系列長、特殊トークンのID、バッチサイズなどの基本的な数値を定めます。DNNクラスでGPTタイプのネットワークモデルを構成します。

今回の主目的であるTransformerブロックの実装となります。

3.2.1 モデルの設定

ここ数回の定番であるモデル設定のクラスです。ネットワーク層で利用される数値を一括して管理するクラスです。今回は事前学習による動作確認が主要目的なので系列長を256と短めにしておきます。

class ModelConfig:
    def __init__(self, tokenizer):
        # モデル構造
        self.vocab_size = tokenizer.get_vocab_size()
        self.seq_len = 256   # 128トークンだとSFT時に少ない
        self.d_model = 256   # 表現力が足りない場合 512
        self.nhead = 8
        self.dim_feedforward = 4*self.d_model
        self.num_layers = 6
        self.dropout = 0.1
        
        # 特殊トークンID
        self.pad_token_id = tokenizer.token_to_id("<pad>")
        self.eod_token_id = tokenizer.token_to_id("<eod>")

        # 学習データに関する設定
        self.context_size = self.seq_len         # 学習できる長さ
        
        # 学習設定
        self.batch_size = 512       # VRAMに合わせて変更
        self.learning_rate = 0.001  # これだとデフォルトと変わらない
        self.num_steps = 5000
        self.max_grad_norm = 1.0
        self.ignore_index = -100
        self.weight_decay = 0.1     # llama2、GPT-3と同様

    # 属性を追加・更新するメソッド
    # 〜〜設定時のタイポに注意だぞ〜〜
    def update(self, **kwargs):
        """渡されたキーワード引数で設定を動的に追加・更新する"""
        for key, value in kwargs.items():
            setattr(self, key, value)

ModelConfigクラスといっても、数値を並べてるだけのクラスです。系列長を表すseq_lenには注意が必要です。小さすぎると文章をうまく生成することができませんし、大きな数値だと、消費されるVRAMも巨大化します。簡易的な、指示チューニングまで考慮しても、256トークンあればギリギリ大丈夫かな

事前学習して文章の生成結果を確認するだけなら、系列長のseq_len=128として、10分程度 (Colab T4) で学習も終了します。

3.2.2 Transformerブロック (メイン部分)

PyTorchのTransformerEncoderに因果マスクをつけて対応してきた部分をやや細かく書き下す部分となります。いわゆる $\text{softmax}(QK^{t}/\sqrt{d})V$ を書けばOKです²。

図２：Transformerブロック

Transformerブロックは、因果マスク付きの注意機構と出力部分のMLP (FFN) のネットワークを作成して、図２のように並べる構造となっています。

最初に因果マスク付きの注意機構からの実装となります。PyTorchのTransformerブロックを参考にCausalSelfAttentionクラスを作成しました。

因果マスク付きの自己注意機構

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.d_model % config.nhead == 0
        self.nhead = config.nhead
        self.head_dim = config.d_model // config.nhead

        # (1) q,k,v をまとめて1回のmatmulで　第22.5回風
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout_p = config.dropout
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        bs, seq_len, d_model = x.shape
        q, k, v = self.qkv(x).split(d_model, dim=2)  # (2) (bs, seq_len, d_model -> (bs, nhead, seq_len, head_dim)
        q = q.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        k = k.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        v = v.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        # (3) PyTorchのScaled Dot Product Attention
        # 条件が合えばFlashAttentionなどの高速実装が自動選択される
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=self.attn_dropout_p if self.training else 0.0,
            is_causal=True,
        )

        y = y.transpose(1, 2).contiguous().view(bs, seq_len, d_model) # y = y.transpose(1, 2).reshape(bs, seq_len, d_model)
        return self.resid_dropout(self.proj(y))

説明メモ
$W_q$、$W_k$、$W_v$用のLinearをまとめて作成する形になっています³。

• (1) self.qkvの部分についてです。Q，K，V用のLinearをひとまとめにして扱っています。3 * config.d_modelと3倍されている部分に注目してください。
• (2) self.qkv(x).splitの部分でQ，K，Vをヘッド数に応じた形で分割します。
• (3) 注意機構の部分。ここはPyTorchの関数を使ってしまったorz。この状態でもRoPE （回転位置埋め込み）できるはず😉
• is_causal=Trueによって「未来のトークンを利用しない」自己注意として動作させています⁴
• scaled_dot_product_attentionの出力 y の形状は (bs, nhead, seq_len, head_dim)です。1番目と2番目を入れ替えて、(bs, seq_len, nhead, head_dim)として、reshape(bs, seq_len, d_model)とすることで、 d_model = nhead * head_dim を実現させています。

次に、Transformerブロックの出力部分を構成するネットワークです。GPT-2タイプでは素朴に全結合層と活性化関数を並べる形になります。

図３：GPT2のMLP部分

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.fc1 = nn.Linear(config.d_model, config.dim_feedforward, bias=False)
        self.act = nn.GELU(approximate="tanh")
        self.fc2 = nn.Linear(config.dim_feedforward, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.fc2(self.act(self.fc1(x))))

forward部分、一行で書いているので見づらいですが、図３の形で直線的に並べるだけです。

最後に、CausalSelfAttentionクラス、MLPクラスを利用してTransformerブロックに積み上げていきます。これも並べるだけなのですが...

Transformerブロック

class TransformerBlock(nn.Module):
    """Pre-Normalization方式 """
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # Attention(LayerNorm(x)) + x
        x = x + self.mlp(self.ln2(x))      # MLP(LayerNorm(x)) + x
        return x

Wang et al.,(2019) や Xiong et al., (2020) などによるとAttentionやMLPの前に正規化層を挟むほうが学習が安定することが知られています⁵。Pre-Normalizationと呼ばれているようです。また、残差接続（x = x + F(x) という形）を使い、深いネットワークでも勾配消失を抑制しつつ学習の安定化をはかります。

3.2.3 GPTタイプのネットワークモデル

Transformerブロックが作成できたので、「埋め込み層、Transformerブロック、全結合層で次の単語予測」というネットワークを組めばGPTタイプ・自己回帰型のモデルとなります。図４のように次の単語（トークン）が教師データとなります。

図４：GPTタイプのネットワーク構造

小型GPTタイプ

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)
        self.pos_embedding = nn.Embedding(num_embeddings=config.seq_len, embedding_dim=config.d_model)
        self.dropout = nn.Dropout(config.dropout)

        # (2) causal transformer layer
        self.transformer_blocks = nn.Sequential(*[TransformerBlock(config) for _ in range(config.num_layers)])

        # (3) 最後の出力に向けた正規化とFC層　最終的に単語数になる
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.fc = nn.Linear(config.d_model, config.vocab_size,bias=False)
        self.apply(self._init_weights)                    # 個別に初期化
        self.fc.weight = self.token_embedding.weight      # 重み共有 (4)
        

        # (5) 残差方向の出力層だけ 1/sqrt(2N) スケール
        std = 0.02 / math.sqrt(2 * config.num_layers)
        for pn, p in self.named_parameters():
            if pn.endswith("proj.weight") or pn.endswith("fc2.weight"):
                nn.init.normal_(p, mean=0.0, std=std)
    # (6)
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:          # fc は bias=False なのでスキップ
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)


    def forward(self, x):
        bs, seq_len = x.shape      
        positions = torch.arange(seq_len,device=x.device)
        
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(positions).unsqueeze(0)
        x = tok_emb + pos_emb           # トークン埋め込み＋位置埋め込み
        x = self.dropout(x)        
        x = self.transformer_blocks(x)  # 自己回帰型 transformer
        x = self.layer_norm(x)          # Layer正規化
        logits = self.fc(x)             # NTP: 次のトークン予測
        return logits

説明メモ

• (1) トークンをベクトル化する token_embedding と、位置情報をベクトル化する pos_embeddingを足し合わせます。

• (2) Transformerブロック。nn.Sequentialを利用することで繰り返し処理をforward部分で記載せず、きれいに表現できます。機能追加時にnn.ModuleListに修正されます。

• (3)LayerNorm → fc で vocab_size 次元に射影。これが各位置での次トークンのスコア(logits)となります。

• (4) 重み共有 (weight tying)。入力埋め込みと出力射影の重みを共有しています。パラメータ削減と性能向上が狙いで、全結合層fcのバイアス項をbias=Falseとして共有に合わせています。2026年現在、重み共有しない高性能モデルも登場しており、賛否両論の状態？！このレベル感では全く無意味ですが

• (5) 残差分岐のスケーリング。残差接続は $x=x+F(x)$という形で、ブロックの出力を足し込んでいきます。各分岐の出力の分散を $\sigma^2$、初期入力の分散を $\sigma_{0}^2$とすると、$L$ 回残差接続した後の残差ストリームの分散はおおよそ
  $$
  \text{Var}(x_L) \approx \sigma_{0}^2 + L\sigma^2
  $$
  となり、深さに比例して線形に増えるかたちになります（標準偏差だと$\sqrt{L}$倍増加）。層を深くするほど分散が大きくなっていきます。$L\sigma^2$この部分の$L$倍の効果を打ち消すために、分散を$1/L$倍、標準偏差を$1/\sqrt{L}$倍すればよいだろうというアイディアとなります。今回の実装だと残差接続は次の2回です。

  • x = x + Attention(LN(x))
  • x = x + MLP(LN(x))

  これがconfig.num_layers回繰り返されます。合計としてL = 2*config.num_layersとしてスケールされる（割り引かれる）ことになります。具体的には Attention の出力射影（proj）と MLP の第2層（fc2）について縮めることになります⁶。

• (6) 線形層の初期化。nn.Linearとnn.Embeddingの重みについて、平均0、標準偏差std=0.02 の正規分布で初期化するのはどうもBERTやGPT-1の論文からの流れで、定番設定となっているようです。やはり、深いネットワークでの安定性確保につながるようです。if文ですが、biasを使っていないので、事実上機能していません。ただ、慣例なようなので記載してあります。標準偏差の初期値がstd=0.02なので(5)番のstdに0.02がハードコードされています⁷。

モデルの確認
tokenizerを指定して、ModelConfigを使って設定を読み込みます。

config = ModelConfig(tokenizer)
model = DNN(config).to(device)

torchinfoで大まかな構造とパラメータ数を確認しました。

from torchinfo import summary

summary(
    model,
    input_size=(1, config.seq_len),        # (batch, seq_len)
    dtypes=[torch.long],                   # 埋め込みの入力は整数インデックス
    depth=2,
)

summary部分にひと工夫すれば、入出力形状も表示されるようです。

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
DNN                                           [1, 256, 10000]           --
├─Embedding: 1-1                              [1, 256, 256]             2,560,000
├─Embedding: 1-2                              [256, 256]                65,536
├─Dropout: 1-3                                [1, 256, 256]             --
├─Sequential: 1-4                             [1, 256, 256]             --
│    └─TransformerBlock: 2-1                  [1, 256, 256]             787,456
│    └─TransformerBlock: 2-2                  [1, 256, 256]             787,456
│    └─TransformerBlock: 2-3                  [1, 256, 256]             787,456
│    └─TransformerBlock: 2-4                  [1, 256, 256]             787,456
│    └─TransformerBlock: 2-5                  [1, 256, 256]             787,456
│    └─TransformerBlock: 2-6                  [1, 256, 256]             787,456
├─LayerNorm: 1-5                              [1, 256, 256]             512
├─Linear: 1-6                                 [1, 256, 10000]           2,560,000
===============================================================================================
Total params: 9,910,784
Trainable params: 9,910,784
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 26.62
===============================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 56.66
Params size (MB): 39.64
Estimated Total Size (MB): 96.30
===============================================================================================

torchinfoによると🤔1000万パラメータとなります🤔 モデルサイズが10M、心配なくらい小さいですしかも、トークン埋め込みの部分に250万パラメータとかなりの部分を使っています。そういえば、トークン埋め込みと最終層の全結合層は重みを共有しているので、実質的な学習パラメータって、740万パラメータくらいかな？とても小さなモデルです。

正確なパラメータ数は（model.parameters()を愚直に足しました）、たぶん次のようになるかと。

unique = sum(p.numel() for p in model.parameters())
total  = sum(p.numel() for _, p in model.named_parameters(remove_duplicate=False))
print(f"unique={unique:,}  reported={total:,}  tied={total-unique:,}")
# unique=7,350,784  reported=9,910,784  tied=2,560,000

1/3程度がトークン埋め込みパラメータとなり、文脈理解の部分は残りの2/3に委ねられています。文脈理解度を上げるにはModelConfigのd_modelサイズを大きくする必要がありそうです。

3.3 Dataset・DataLoaderによるデータの読み込み

前処理済みの日本語データとtokenizerを利用してカスタムDatasetでID化していきます。データサイズが大きくなると事前にid化した方がいろいろ効率的なようです。ITカテゴリーの文字数が90万文字くらいなのでDatasetクラス内でid化していきます。ただし、データ量がギガクラスになると他の方法を使わないと完全に詰みます😂

Datasetクラスのカスタマイズ

カスタムDatasetクラスは特殊メソッドの__init__()、__len__()、__getitem__()を定義すれば完成です。データファイルではなくて、テキストデータのディレクトリーを引数とします。利用するカスタムDatasetクラスは第38回と同一です。

class RandomGPTDataset(Dataset):
    def __init__(
        self,
        data_dir,
        tokenizer,
        context_size=128,
        samples_per_epoch=10_000,
    ):
        self.context_size = context_size
        self.samples_per_epoch = samples_per_epoch
        eod_id = tokenizer.token_to_id("<eod>")
        filenames = sorted(Path(data_dir).glob("*.txt"))

        all_ids = []
        for filename in filenames:
            text = filename.read_text(encoding="utf-8")
            document_ids = tokenizer.encode(text, add_special_tokens=False).ids
            all_ids.extend(document_ids)
            all_ids.append(eod_id)

        self.ids = torch.tensor(all_ids, dtype=torch.long)

    def __len__(self):
        # 1 epochあたりに何個のランダム窓を学習するか
        # 10,000個（回？）がデフォルト値
        return self.samples_per_epoch

    def __getitem__(self, idx):
        # idxは使わず、毎回ランダムな開始位置を選ぶ
        start = torch.randint(
            low=0,
            high=len(self.ids) - self.context_size,
            size=(1,),
        ).item()

        x = self.ids[start:start + self.context_size]
        y = self.ids[start + 1:start + self.context_size + 1]

        return {"ids": x, "labels": y}

dataset = RandomGPTDataset(
    data_dir=data_dir,
    tokenizer=tokenizer,
    context_size=config.context_size,
    samples_per_epoch=10_000,
)

dataloader = DataLoader(
    dataset=dataset,
    batch_size=config.batch_size,  # メモリ足りない場合は小さくする
    shuffle=False, # Dataset側でランダムになるので、FalseでOK
    num_workers=0, # 2,4のほうが早い？要検討
    pin_memory=torch.cuda.is_available(), # GPU使う時True
)

説明メモ

• 基本的にディレクトリーに存在するtxtファイルを順番に読み込み
  1. tokenizer.encode()でID化、
  2. テキストファイル毎に文章の切れ目トークンであるを挿入

　　という作業を__init__の部分で行います。

• __getitem__では、torch.randintで開始位置をランダムに決めて、context_sizeに合わせて、idsとlabelsを作成します。

具体例
雰囲気を確認するために、

• context_sizeの長さに調整 context_size=8

として表示してみましょう😆

dataset = RandomGPTDataset(
    data_dir=data_dir,
    tokenizer=tokenizer,
    context_size=8,
)

• ids、labels：系列長８
• idsとlabelsは１トークンずれ
• dataset[0]とdataset[1]はランダムに選択されているはず

dataset[0]
{'ids': tensor([4809, 2145,  275, 3471,  542,  775, 2239,  847]), 
 'labels': tensor([2145,  275, 3471,  542,  775, 2239,  847,  542])}

dataset[1]
{'ids': tensor([1817, 2552,  881,  292,   56,   73, 2616,   50]),
 'labels': tensor([2552,  881,  292,   56,   73, 2616,   50, 1727])}

RandomGPTDatasetクラス部分で入力されるデータは毎回シャッフルされます。DataLoaderでのシャッフルはFalseでOKです。

3.4 誤差関数と誤差最小化の手法の選択

AdamWを使っているので、重み減衰 (weight decay)がかかっています
少し気合を入れて、重み減衰（weight decay）を学習するパラメータに応じて設定してみました。

重み減衰の指定と最適化

# weight decayの設定
decay, no_decay = [], []
for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    if param.dim() >= 2:
        decay.append(param)      # Linear、Embeddingのgain。weight_decay=config.weigh_decay
    else:
        no_decay.append(param)   # LayerNormのgain。weight_decay=0

optim_groups = [
    {"params": decay,    "weight_decay": config.weight_decay},
    {"params": no_decay, "weight_decay": 0.0},
]

optimizer = torch.optim.AdamW(optim_groups, lr=config.learning_rate, betas=(0.9, 0.95), fused=(device.type == "cuda"))
criterion = nn.CrossEntropyLoss(ignore_index=config.ignore_index)    # 今回ignore_indexを利用していないがSFT用につけておく

LayerNormに対して、重み減衰を0に、線形層に対して小さな正の重み付けを与えています。重み減衰効果により各層の出力値が更新時も正規化されるように働くようです。

実装ではconfig.weight_decay=0.1というLlama2と同様の値を用いました⁸。デフォルト値は0.01のようです。モデルサイズが極小なので0.1と0.01の差があまり感じられなかったように思います。

AdamWについて
AdamにL2正則化を追加したような手法で$\theta_{t}$を t 期のパラメータ、$\eta$を学習率したとき、次のように更新計算されます。

\begin{align*}
\theta_{t+1} 
& = \theta_{t} - \eta \frac{\hat{m_t}}{\sqrt{\hat{v_t}}+\varepsilon} - \eta \lambda \theta_{t} \\
& = \theta_{t}-\eta \text{AdamUpdate} - \eta\lambda \theta_{t}
\end{align*}

$\hat{m_t}$や$\hat{v_t}$などの細かい部分を端折るとして、３項目がweight decayに関する部分となります。式中の$\lambda$がconfig.weight_decayになります。パラメータ更新時に減少圧力をやや加えるようなイメージかと思います。

線形層の重み行列はスケール不変性があるため、weight decayでパラメータの数値を縮めても問題なく、実効学習率を制御する効果だけが得られます。パラメータ数値に対しての正規化のような役割をweight decayが与えることになります。一方、正規化層 (Normalization Layer) の学習パラメータは、大小のスケールを与える部分になります。大小のスケールを決めるパラメータには不変性がなく、縮めることは層の寄与そのものを弱める方向に働きます。更新を歪めかねないため、weight decayは0に設定します。オプティマイザーについてもどこかでまとめてみたい気がします。

3.5 変数更新のループ

いよいよ学習ループです。ステップ数でループの管理をします。

学習ループ

# step数で管理するためのデータローダー関数
def infinite_loader(dataloader):
    while True:
        for batch in dataloader:
            yield batch

data_iter = infinite_loader(dataloader)  # step数で計測

# ---- 学習ループの設定　tqdmを使わないほうが断然シンプル
from tqdm import tqdm               # tqdmを使い進捗状況表示をしてみた
max_iters = config.num_steps        # step数
total_token = 0                     # 累積トークン数を初期化
pbar = tqdm(range(max_iters))
use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()  # flash attentionを利用の判定

# ---- ここから学習ループ　（読みにくくなるので検証データによる確認は割愛）
model.train()
for step in pbar:
    batch = next(data_iter)
    input_ids = batch["ids"].to(device, non_blocking=True)
    labels = batch["labels"].to(device, non_blocking=True)
    optimizer.zero_grad()
    # GPUが対応している場合は、bf16へ変更してflash attentionを使う
    with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=use_bf16):
        outputs = model(input_ids)                 # [batch, seq_len, vocab_size]
        loss = criterion(
            outputs.view(-1, config.vocab_size),   # [batch * seq_len, vocab_size]
            labels.view(-1),                       # [batch * seq_len]
        )
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.max_grad_norm)  # 勾配クリップ
    optimizer.step()

    # このstepのトークン数 （バッチサイズ×系列長）を加算
    total_token += input_ids.shape[0] * input_ids.shape[1]

    # set_postfixは辞書が引数になる
    pbar.set_postfix({"loss": f"{loss.item():.4f}", "tokens": f"{total_token/1e6:.2f}M"})
    if (step+1) % 500 == 0:
        tqdm.write(f"{step+1}-step:\tloss:{loss.item():.3f}\ttokens:{total_token:,}") 

表示用の設定と形状変更で、複雑に見えますが、基本構造は初回から変わっていません基本構造は次の形になります。

基本構造

for ループ
    input_ids, labels：入力データと教師データ
    optimizer_zero_grad()
    outputs = model(input_ids)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

検証用のコードも挿入するとより複雑に見えるので解説上割愛していきます。

LLMの事前学習は通常１エポックと聞きます。しかし！今回もlossが1を下回るくらいになるまで頑張って学習させます。学習データやモデル、語彙サイズの関係からなのか詳細は不明ですが、lossは0.5程度を下回ると思います。文章を全暗記するレベル感に到達します。ここ個人的に「ポイント！」と思う部分です。GPTモデルの事前学習演習だと、モデルの構造に重点が置かれているので、少し変だけど日本語っぽい文章が生成されるという内容で終わることが多い印象があります。学習データのコピーに近い内容であってもしっかりした文章が表示されるとモチベーション・テンションも上がるはず😆

メモリーが足りないというときは系列長やバッチサイズを環境に合わせて調整してください。15分程度で終了になると思います。今回も学習率固定ですが、「warmup + cosine decay」というのがこの分野の定番のようです。しっかり収束していくので学習率固定でも問題なさそう

3.5 検証

文章生成の関数は、最大確率となるトークンを繋いでいく素朴なGreedyタイプです。

文章生成の関数 Greedy

@torch.inference_mode()
def generate_text(
    model,
    input_ids,
    max_new_tokens=config.seq_len,
    eos_token_id=None,
):
    model.eval()    # 評価モードへ
    generated = input_ids.to(device)
    seq_len = model.config.seq_len

    # greedyで単純に生成
    # eos_token_id　or　max_new_tokensまでで生成終了
    for _ in range(max_new_tokens):
        logits = model(generated[:, -seq_len:])                     # 単語上の確率を計算
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)  # 確率が最大のトークン選択
        generated = torch.cat([generated, next_token], dim=1)       # 予測idを追記
        # 終了判定
        if eos_token_id is not None and next_token.item() == eos_token_id:
            break

    return generated

id化された文字列input_idsを入力して、modelを使い、次のトークンを予測して、トークン追加を繰り返すだけです。文章生成終端IDのeos_token_idが予測されたら文章生成終了となります。eos_tokenとして、<eod>を利用します。

Greedyで生成するgenerate_text関数を用いて、入力文章に続く文字列を表示するresponse関数を作成します。

生成して表示する関数

from tokenizers import decoders
tokenizer.decoder = decoders.ByteLevel() # BPEのときはこれが必要、文字化けするぞ

def response(prompt):
    model.eval()
    model_input = tokenizer.encode(prompt, add_special_tokens=False)
    input_ids = torch.tensor([model_input.ids], dtype=torch.long, device=device)

    output_ids = generate_text(
        model=model,
        input_ids=input_ids,
        max_new_tokens=config.seq_len,
        eos_token_id=config.eod_token_id,
    )

    # 全文をデコード
    generated_text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=False)
    print(generated_text)

response関数を使って、csvファイルに書き込んだ文に続く文字列を表示させっていきます。

生成結果

import pandas as pd

df = pd.read_csv("./data/sample_prompt.csv")

for num, prompt in enumerate(df["prompt"]):
    print(f"【{num+1}番目】")
    response(prompt)

学習データにある書き出しだと、ほぼ学習した文を出力すると思います。異なる内容だと、まるでダメな文章となります。

学習データに存在する文言からスタート

【1番目】
iPadをはじめとするタブレット端末が市場で大きなブレイクを遂げると言われています。タブレットは手持ちで使える点が大きなメリットですが、長時間手で持っているのはやはり疲れてしまいます。...

【2番目】
Googleロゴのアニメーションを観察するとオシロスコープの波形のようなものが表示されます。このロゴをクリックすると、「ハインリヒヘルツ」という言葉が検索されるのですが、これはハインリヒ・ルドルフ・ヘルツ氏の誕生日を記念しているためです。...

【3番目】
HPのUltrabookラインナップでは、素材やデザインにプレミアム感を加えた製品となっています。HPはこれまでもデザイナーとコラボレーションしたノートパソコン(Vivienne Tam Editionなど)をリリースしてきましたが、今回のHP Envy 14 Spectreが、同様に高いプレミアム感を持つ製品です。...

きれいにコピー文ぽいものが生成されています。自然な日本語になっています。

学習データに近い表現からスタート

【4番目】
テーブルやデスクに表示された画像を手で操作する近未来的なシーンは、映画やドラマなどで目にしたことがある人は多いでしょう。実は、これは2008年に発表されたマイクロソフトの「Microsoft Surface」を使ったテーブル型の液晶PCにルーツがあります。当時はDLPプロジェクターで映像を出し、5台のカメラで動きを検出するという大掛かりな構造でした。...

【5番目】
NECは、テレビの新しい楽しみ方を提案する形で、個人向けデスクトップパソコン「VALUESTAR」シリーズ(3タイプ16モデル)を2012年2月14日に発表し、2月16日より販売を開始します。新商品には、よりパワフルになった録画機能に加え、TV視聴・録画機能に業界で初めて人気のTwitterを連携させた「SmartVisionつぶやきプラス」を追加するなど、TVパソコンならではの機能を搭載しています。スマートフォンやホームネットワーク対応も強化し、「安心・簡単・快適」なデジタルエンターテイメントの提案として、主要モデルに以下の強化を行いました。

途中からコピー文っぽいものが生成されています。多少の抜けやゆらぎだとなんとかなりそう。

学習データと無関係な表現からスタート

【6番目】
SNSでの魅力をアップさせるためのスキルと知識が紹介されています。<eod>

【7番目】
図で考える力によって、問題を整理・分析的に分かりやすい場合があります。

なお、契約が不要な場合、契約月額は月額3,100円、このうち、8,100円にする課金率による負担があります。

【8番目】
小生が大学を卒業した頃、東京工学部写真ショップ写真選手権で、2010年には野選手権と御苗場レビュアー賞を受賞。...

• ６番目は、<eod>をうまく使い、文章の生成をやめています。
• ７，８番目は謎の文章が続いています。話が飛んでいるぞって。後半部分はやはりコピー文ぽい内容となります。

次回

GPTタイプをLlama化していきます。RMSNormやSwiGLU、RoPEなどへのアップグレード追加作業をしてみたいと思います。

参考にしたもの

Pre Normalization関連
基本的にPre Normalizationが良いというタイプの主張。他にもいくつかあるようです。

• Toan Q. Nguyen　and Julian Salazar (2019)
  "Transformers without Tears: Improving the Normalization of Self-Attention", arXiv:1910.05895

• Wang et al. (2019)
  "Learning Deep Transformer Models for Machine Translation", arXiv:1906.01787）

• Xiong el., (2020)
  "On Layer Normalization in the Transformer Architecture", arXiv:2002.04745

Llama 2の技術レポート
今回活用したのはAdamWのweight decayの値の部分、まさに１行参照

• Hugo Touvron, Meta社 (2023)
  "Llama 2: Open Foundation and Fine-Tuned Chat Models", arXiv:2307.09288

AdamWの論文

• Ilya Loshchilov and Frank Hutter (2017)
  "Decoupled Weight Decay Regularization",arXiv:1711.05101

書籍

• 斎藤 康毅 (2026)
  『ゼロから作るDeep Learning ❻ ―LLM編』オライリー・ジャパン

  • とてもわかり易い。しかも多色刷りで図解が丁寧すぎ😆付録部分にいろんな数式による解説があります。注意機構の部分を始め実装面で参考になりました。

• Sebastian Raschka 巣籠悠輔 (監修), 株式会社クイープ 訳 (2025)
  『つくりながら学ぶ！LLM 自作入門』マイナビ出版

  • こちらも有名。原著・訳本いずれもOKかな。地味に練習問題とその解答が理解促進に役立ちました😆

• 気軽にはじめる LLM 自作入門

  • thirtypowerさんによるZennでの解説書籍。GPTやLLaMAのPyTorchによる自作実装についてゼロから解説されています。自作LLMを勉強する上で非常に参考になりました。しかも、無料公開。とてもおすすめです

• 機械学習と情報技術

  • 式の展開もあり、いつもながら参考にしています。見るたび内容が増えている。

目次ページ

注

1. 第25回のトークナイズ編・BPEに従って作成しました。 ↩

2. 第33回のDatasetクラスのカスタマイズについてを利用して、tokenizerに応じた形でid列を作成するDatasetクラスを作成します。 ↩

3. 『ゼロから作るDeep Learning ❻ ―LLM編』に丁寧に説明があります。 ↩

4. 第22.5回のMultiHeadAttentionを参考にしてもらえると幸いです。PyTorchのMultiHeadAttentionっぽくQKVのLinear部分をひとまとめに扱っています。 ↩

5. scaled_dot_product_attentionの内部構造が、よく見かける$\sqrt{\text{d_model}}$で割り引くsoftmax計算、右上側をマスクするなどのなかなか複雑になります。この部分も自分で書いたほうが勉強になるとは思うのですが、面倒😅 ↩

6. pre normalizationと呼ばれているようです。 ↩

7. これは初期化時点のみの調整で学習開始後は自由に重みが動きます。 ↩

8. Llama2の論文"Llama 2: Open Foundation and Fine-Tuned Chat Models"にはweight decay = 0.1 を使ったと明記してあります。まだ、自分の中でのAdamWの認識不足、もう少し深める必要がありそうです。 ↩