はじめに
捨てるに捨てられないモノって、ありますよね。例えば、私が高校生のころから書いていた日記とか。友人に話すと（おそらくだけど、その几帳面さから）ドン引きされながらも、気が付いたら約20年にわたって3,863日分のデータとなりました。
改めて読み返すことはないけど、消してしまうにはやや忍びない。でももしかしたら呪われているかもしれないし、そうだとしたらなんとか供養したい。
そう思い続けていた折、買ったまま積読されていた つくりながら学ぶ！LLM 自作入門を改めて手に取ったところ、小規模なGPTモデルを構築し、短いテキストデータで学習・検証させる手法が書いてあるではありませんか。おお、そうだ...
というわけで、GPTモデルをスクラッチで構築し、パラメータが初期化されたまっさらな状態から私の日記データのみを使って事前学習を試み、どのようなテキストが出力されるのか検証します。で、それが供養になるでしょう。
コードは動作します
この記事に記載しているコードは、Google Colabでだいたい動作確認しているので、おおよそそのまま動作するはずです。動かなかったらごめんなさい。ファイルの保存先など環境に依存するところは修正する必要があります。
この先、長いです
私自身の理解のために、細かく解説しています。なかなかしぶとく、しつこく。LLM学習に明るい方には冗長だと思うので、読み飛ばしてください。
モデルについて
アーキテクチャ
LLMのアーキテクチャは、公開されているGPTの初期バージョンであるGPT-2を使用します。GPT-2にはいくつかシリーズがあり、パラメータサイズ順に124MのSmall、355MのMedium、774MのLarge、1.5BのXLです。このうち、124MパラメータのSmallモデルをターゲットとして構築します。
パラメータサイズ
GPT-2 Smallモデルからいくつか手を加えているため、構築したモデルのパラメータサイズはオリジナルとは異なっています。
| GPT-2 Small | 本モデル | |
|---|---|---|
| レイヤー数 | 12 | 12 |
| Attentionヘッド数 | 12 | 12 |
| Embedding次元 | 768 | 384 |
| コンテキスト長 | 1024 | 256 |
| 語彙数 | 約5万 | 約10万 |
| パラメータ数 | 124M | 98M |
まず、日本語に対応したトークナイザを使用しているため、語彙数が約5万から約10万に増えています。また、その語彙数約10万に対して日記に存在する語彙数が少ないため、Embedding次元も768から384と半分にしています。コンテキスト長も、文法的に正しい日本語が出力されるか確認できればいいので256としています。結果、本モデルのパラメータ数は約98Mとなりました。
データセットについて
日記データ
1日分の日記は1ファイルで、何年か前にデータ構造を整理したので、すべてxml形式で保管されています。もっとも古い日付は2002年8月12日で、2026年に入ってからも数日書いていました。テキストファイルにして約6MB、文字数2,109,894、トークン数は2,240,004、重複しているトークンを除いたユニークトークン数は3,117です。よくこんなに書いたもんだ。
データセット化
日記データは、日付・タイトル・本文を含む1日分の記述が1つのxmlファイルに格納されています。それらから各日記の本文のみを抽出し、3,863日分を1つの大きなテキストにまとめます。
各日記の本文の境界を表すトークンとして、<|endoftext|>を使用します。このトークンは、GPT-2モデルに対して文書の境界であることを示すマーカーとしての役割があります。こうして完成したデータセットのイメージは、以下のようなものになっています。
A月B日の本文<|endoftext|>A月C日の本文<|endoftext|> ... <|endoftext|>X月Y日の本文<|endoftext|>
完成したデータセットはthe-diary.txtというファイルにまとめておきます。以降はこのデータセットをthe-diaryデータセットとします。
トークナイザについて
オリジナルのGPT-2のトークナイザの語彙数は50,257です。これはUTF-8のバイト列をバイトペアエンコーディングするとのことで、日本語にも対応しています。しかし、日本語の処理に最適化されているわけではないようで、日本語単語を必要以上に分割してしまう傾向にあるようです。
そこで、日本語にも最適化されているcl100k_baseというエンコーダを採用します。
なお、the-diaryデータセットのユニークトークン数は3,117です。cl100k_baseの語彙数100,277に比べて語彙の種類の少なさがやや気になります。Embedding次元をGPT-2の半分にしているのはそのためで、結果、学習時の損失が改善することとなりました。
開発環境について
マシンは手元にあるMacBookPro(かなり古い)を、コードの記述は統合開発環境のPyCharmを使用しています。モデルの学習にはGoogle Colab Pro+でT4 GPUを使用しました。
実装
この章では、GPT-2モデルの学習に使用したコードを記載しています。コードだけではなく、もとの数式にもなるべく言及しています。ただ、数式とコードだけを眺めていてもよく分からない箇所があります。特にマルチヘッド注意機構は「多次元テンソルの計算がミソ」みたいなところがあり、私はそのあたりの次元の狭間で何をしているのか迷子になりがちです。そこで、なるべく計算内容をイラストで整理できるようにしています。
ハイパーパラメータ
まず、実装に必要なハイパーパラメータや設定値をここにまとめます。この後のコードの中で、これらの辞書のキーがそのまま変数となってしれっと登場することがあります。
GPT_CONFIG = {
# 語彙サイズ
"vocab_size": 100277,
# コンテキスト長さ
"context_length": 256,
# 埋め込み次元
"embed_dim": 384,
# Attentionヘッド数
"num_heads": 12,
# 層の数
"num_layers": 12,
# ドロップアウト率
"drop_rate": 0.1,
}
TRAIN_CONFIG = {
# 乱数初期化
"manual_seed": 123,
# 最小学習率
"min_learning_rate": 3e-5,
# 最大学習率
"max_learning_rate": 3e-4,
# ウォームアップ率
"warmup_steps_ratio": 0.02,
# バッチサイズ
"batch_size": 8,
# gradient accumulationステップ
"grad_accum_step": 2,
# 評価ステップ
"eval_update_step": 16,
# エポック数
"num_epochs": 10,
}
vocab_sizeはモデルが扱える語彙数で、使用するトークナイザcl100k_baseから取得した値です。context_lengthは学習もしくは出力時の最大トークン数、embed_dimはトークンを埋め込みに変換した時の次元数です。これらは、アテンションヘッド数num_headsとレイヤー数num_layers、ドロップアウト率drop_rateと合わせてモデルの構造を決定するパラメータです。
manual_seedは学習の再現性を確保します。min_learning_rateとmax_learning_rate、warmup_steps_ratioは、学習率スケジュールに使用します。batch_sizeとgrad_accum_stepは、context_lengthとの組み合わせで学習ステップあたりのトークン数を決定します。これらは、num_epochsと合わせてトレーニング方法を調整するパラメータです。
トークナイザ
まずは作成したthe-diaryデータセットを読み込んでトークンIDの列に変換します。
import torch
import tiktoken
tokenizer = tiktoken.get_encoding("cl100k_base")
with open("the-diary.txt", "r") as f:
text = f.read()
encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
token_ids = torch.tensor(encoded)
トークナイザは前述のように日本語に最適化されているcl100k_baseを使用します。tokenizer.encode()では、the-diaryデータセットで本文の区切りとして使用したトークン<|endoftext|>をallowed_specialとして指定しています。<|endoftext|>に対応しているトークンIDは100257です。allowed_specialを指定しないと、<|endoftext|>がこの1つのトークンとして認識されず、複数トークンに分解されてしまいます。
データローダ
続いてtoken_idsをトレーニング・検証・テストの各データセットに分割します。まずtorchのDatasetを継承したGPTDatasetクラスを定義します。
GPTDatasetクラス
from torch.utils.data import Dataset
class GPTDataset(Dataset):
def __init__(self, token_ids, context_length, stride):
self.input_ids = []
self.target_ids = []
for i in range(0, len(token_ids) - context_length, stride):
input_chunk = token_ids[i:i + context_length]
target_chunk = token_ids[i + 1:i + context_length + 1]
self.input_ids.append(input_chunk)
self.target_ids.append(target_chunk)
def __len__(self):
return len(self.input_ids)
def __getitem__(self, idx):
return self.input_ids[idx], self.target_ids[idx]
LLMの学習のデータセットには、他の機械学習と同じようにインプットとターゲットとなるデータが必要となります。インプットとターゲットはどちらも、the-diaryデータセットから抽出された同じサイズのトークンのリストで、ターゲットはインプットより、1トークン先のテキストを参照しています。Fig.1のようなイメージです。
Fig.1 : インプットとターゲット
なお、データセット作成前にthe-diaryデータセットはトークナイザでトークンIDに変換されるので、図もトークンIDで表しています。
複数バッチを考えた時、次のバッチのインプットとの間隔をストライドと呼びます。例えばcontext_length = 4でのstride = 2、stride = 4のイメージはFig.2のようになります。
Fig.2 : ストライド
バッチ間のオーバーラップを防ぐためと、計算量を抑えるためにcontext_length = strideとしています。オーバーラップが多いと訓練データへの過剰適合のリスクが高まるとのことです。
データローダ作成
GPTDatasetクラスを利用してトレーニング・検証・テストの各DataLoaderを作成します。データセット分割の割合は7 : 2 : 1としました。
from torch.utils.data import DataLoader
# 分割率：トレーニング
train_ratio = 0.7
# 分割率：検証
val_ratio = 0.2
# 分割率：テスト
test_ratio = 0.1
# トレーニングDataLoaderの作成
train_idx = int(len(token_ids) * train_ratio)
train_dataset = GPTDataset(
token_ids=token_ids[:train_idx],
context_length=GPT_CONFIG["context_length"],
stride=GPT_CONFIG["context_length"],
)
train_loader = DataLoader(
train_dataset,
batch_size=TRAIN_CONFIG["batch_size"],
shuffle=True,
drop_last=True,
num_workers=0
)
# 検証DataLoaderの作成
val_idx = int(len(token_ids) * (train_ratio + val_ratio))
val_dataset = GPTDataset(
token_ids=token_ids[train_idx:val_idx],
context_length=GPT_CONFIG["context_length"],
stride=GPT_CONFIG["context_length"]
)
val_loader = DataLoader(
val_dataset,
batch_size=TRAIN_CONFIG["batch_size"],
shuffle=False,
drop_last=False,
num_workers=0
)
GPT-2モデル
GPT-2モデルの実装は、PyTorchでのnn.Moduleクラスを継承する機械学習モデルの構築方法に従います。
GPT-2モデルは、Googleの研究者による論文"Attention Is All You Need"で提案されたTransformerというエンコーダ・デコーダ構成のアーキテクチャのうち、デコーダのみの構成をベースとしています。Transformerアーキテクチャの中で、LLMが言語機能を獲得するのに重要とされているのが自己注意(Self Attention)機構と呼ばれるものです。
その仕組みを理解するには線形代数の知識が必要です。友達に線形代数がいるような人でしたら容易でしょうが、大学の教養課程で線形代数に嫌われていた私にとっては困難です。しかしGPT-2モデルの実装には避けて通れません。
そこで、自己注意機構の仕組みをなるべく数式を使わず直感的に説明した記事を作成しました。これも元々は自分自身で理解するためのものです。何かの手助けになれば幸いです。
モデルの概略図
これから実装するGPT-2モデルは以下のような構成となっています。
Fig.3 : GPT-2モデルの概略図
GPTModelクラス
GPT-2モデルのエントリポイントとなるGPTModelクラスです。
import torch.nn as nn
class GPTModel(nn.Module):
def __init__(self, gpt_config):
super().__init__()
# トークン埋め込み層
self.token_embedding = nn.Embedding(gpt_config["vocab_size"], gpt_config["embed_dim"])
# 位置埋め込み層
self.position_embedding = nn.Embedding(gpt_config["context_length"], gpt_config["embed_dim"])
# ドロップアウト
self.dropout = nn.Dropout(gpt_config["drop_rate"])
# Transformerブロック
self.transformer_blocks = nn.Sequential(
*[TransformerBlock(gpt_config) for _ in range(gpt_config["num_layers"])]
)
# 出力側LayerNorm
self.final_norm = nn.LayerNorm(gpt_config["embed_dim"], eps=1e-5)
# 線形出力層
self.out_head = nn.Linear(
gpt_config["embed_dim"], gpt_config["vocab_size"], bias=False
)
def forward(self, input_token_ids):
batch_size, context_length = input_token_ids.shape
token_embeddings = self.token_embedding(input_token_ids)
position_embeddings = self.position_embedding(
torch.arange(context_length, device=input_token_ids.device)
)
x = token_embeddings + position_embeddings
x = self.dropout(x)
x = self.transformer_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits
入力input_token_idsの次元は [ batch_size, context_length ] であり [8, 256] とします。ここで、batch_sizeとcontext_lengthはTRAIN_CONFIGとGPT_CONFIGで定義されている値です。
トークン埋め込み層token_embeddingは、入力のトークンID列をベクトルに埋め込む役割を持ちます。token_embeddingを通過するとinput_token_idsは埋め込みを獲得し、次元は [8, 256, 384] となります。384 はGPT_CONFIG["embed_dim"]の値で、埋め込みの次元です。
位置埋め込み層position_embeddingsは、オリジナルのTransformerでは固定値による位置情報ですが、GPT-2では訓練の過程で位置情報自体も学習されます。位置情報のパラメータ化には、token_embeddingと同じ埋め込み次元を持つposition_embeddingsを使用します。
position_embeddingsの次元は [256, 384] です。これが、
x = token_embeddings + position_embeddings
では、token_embeddingの8個のバッチそれぞれに256 × 384次元のposition_embeddingsが加算される結果となります。
transformer_blocksで、Transformerブロックが直列にnum_layers個繋がった構造をした層を通過します。
final_norm(x)の出力xの次元は[8, 256, 384]であり、これがout_head(x)の出力logitsで[8, 256, 100277]に拡張されます。100,277次元はトークナイザの語彙数と同じで、その各テンソルは、256個の各トークンの次に出現するトークンが語彙100,277個のうちどれなのかを決定する予測スコアを示します。
TransformerBlockクラス
アテンション機構とその後の一連の層を1つにまとめたクラスです。
class TransformerBlock(nn.Module):
def __init__(self, gpt_config):
super().__init__()
# 層正規化1
self.norm1 = nn.LayerNorm(gpt_config["embed_dim"], eps=1e-5)
# 層正規化2
self.norm2 = nn.LayerNorm(gpt_config["embed_dim"], eps=1e-5)
# マルチヘッド注意機構
self.attention = MultiHeadAttention(
dim_in=gpt_config["embed_dim"],
dim_out=gpt_config["embed_dim"],
context_length=gpt_config["context_length"],
num_heads=gpt_config["num_heads"],
dropout=gpt_config["drop_rate"]
)
# フィードフォワード層
self.feed_forward = FeedForward(gpt_config)
# ドロップアウト
self.dropout = nn.Dropout(gpt_config["drop_rate"])
def forward(self, x):
shortcut = x
x = self.norm1(x)
x = self.attention(x)
x = self.dropout(x)
x = x + shortcut
shortcut = x
x = self.norm2(x)
x = self.feed_forward(x)
x = self.dropout(x)
x = x + shortcut
return x
norm1、norm2は層正規化です。ここを通過すると入力のテンソルは平均0、分散1になるように調整され、学習の安定性と効率化を高めます。attentionとfeed_forwardはそれぞれ層正規化とドロップアウトに挟まれる形で配置されています。
それぞれの層正規化の前とドロップアウトの後に、入力のxを
shortcut = x
...
x = x + shortcut
のように後から足し合わせる操作をしています。これはショートカット接続やスキップ接続、もしくは残差接続とか残差結合といったような難しい用語で呼ばれるものです。Fig.3では、この部分を各層を迂回（ショートカット）するラインで表現しています。
ショートカット接続は、入力を途中の層を迂回させて加算することで、逆伝播時に入力側へ勾配が伝わりやすくなります。これによって、勾配消失という深遠な問題に対して、ニューラルネットワークの学習を安定させる効果があります。
FeedForwardクラス
FeedForwardクラスは、線形層とGELU関数を含むシンプルな構成です。
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["embed_dim"], 4 * cfg["embed_dim"]),
nn.GELU(approximate="tanh"),
nn.Linear(4 * cfg["embed_dim"], cfg["embed_dim"]),
)
def forward(self, x):
return self.layers(x)
まず最初の線形層で、入力次元の4倍に次元が拡張されます。続いてGELU関数で情報の取捨選択が行われ、2つ目の線形層で元の次元に縮小されます。このような構成で、モデルは語彙空間に対して表現力を獲得できるようになるとされています。
MultiHeadAttentionクラス
ついにやってきました。MultiHeadAttentionクラスはこの記事の中核にしてTransformerアーキテクチャのキモの部分で、和名では自己注意機構と呼ばれます。しかもただの自己注意機構ではなく、マルチヘッドと呼ばれる並列構造で、さらにマスク付きです。まとめると、マルチヘッド自己注意機構（マスク付き） となります。
class MultiHeadAttention(nn.Module):
def __init__(self, dim_in, dim_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
# 出力次元
self.dim_out = dim_out
# ヘッド数
self.num_heads = num_heads
# ヘッド次元
self.head_dim = dim_out // num_heads
# クエリ・キー・バリュー パラメータ
self.W_query = nn.Linear(dim_in, dim_out, bias=qkv_bias)
self.W_key = nn.Linear(dim_in, dim_out, bias=qkv_bias)
self.W_value = nn.Linear(dim_in, dim_out, bias=qkv_bias)
# ドロップアウト
self.dropout = nn.Dropout(dropout)
# 出力層
self.out_proj = nn.Linear(dim_out, dim_out)
# マスク
mask = torch.triu(
torch.ones(context_length, context_length),
diagonal=1
).bool()
self.register_buffer("mask", mask)
def forward(self, x):
batch_size, context_length, dim_in = x.shape
# クエリー・キー・バリュー
queries = self.W_query(x)
keys = self.W_key(x)
values = self.W_value(x)
# ヘッドの分割
queries = queries.view(batch_size, context_length, self.num_heads, self.head_dim)
keys = keys.view(batch_size, context_length, self.num_heads, self.head_dim)
values = values.view(batch_size, context_length, self.num_heads, self.head_dim)
# 次元を差し替え
queries = queries.transpose(1, 2)
keys = keys.transpose(1, 2)
values = values.transpose(1, 2)
# アテンションスコアの計算
attn_scores = (queries @ keys.transpose(2, 3)) / keys.shape[-1]**0.5
# マスクを適用
mask = self.mask[:context_length, :context_length]
attn_scores = attn_scores.masked_fill(mask, -torch.inf)
# アテンションウェイトの計算
attn_weights = torch.softmax(attn_scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# 出力ベクトルの計算
context_vec = (attn_weights @ values)
# 次元の差し替えとヘッドの分割を元に戻す
context_vec = context_vec.transpose(1, 2)
context_vec = context_vec.contiguous().view(batch_size, context_length, self.dim_out)
context_vec = self.out_proj(context_vec)
return context_vec
マルチヘッド注意機構
Transformerでは、より表現力を高めるために、自己注意機構を複数適用するマルチヘッド注意機構が採用されています。本コードも、num_heads個の自己注意機構を定義しています。
num_heads個のうち、$m$番目のヘッドのクエリ $\textbf{q}^{(m)}_i$・キー $\textbf{k}^{(m)}_i$・バリュー $\textbf{v}^{(m)}_i$の各埋め込みは、以下の式で計算されます。
\textbf{q}^{(m)}_i = \textbf{W}^{(m)}_q \textbf{x}_i \qquad \dots(1)
\textbf{k}^{(m)}_i = \textbf{W}^{(m)}_k \textbf{x}_i \qquad \dots(2)
\textbf{v}^{(m)}_i = \textbf{W}^{(m)}_v \textbf{x}_i \qquad \dots(3)
$\textbf{x}_i$は dim_in次元の埋め込みです。コードではforward()への引数xとして表されていますが、xは次元 [batch_size, context_length, dim_in]のテンソルなので、そこから1バッチ分のトークン列を取り出し、そこからさらに 1トークン分を取り出したテンソルが $\textbf{x}_i$ に相当します。
ヘッドが1つの自己注意機構におけるクエリ・キー・バリュー埋め込みの計算に必要なdim_in×dim_outの行列を $\textbf{W}_q$、$\textbf{W}_k$、$\textbf{W}_v$ とします。これをnum_heads個のヘッドに分割したものが $\textbf{W}^{(m)}_q$、$\textbf{W}^{(m)}_k$、$\textbf{W}^{(m)}_v$ であり、肩の$^{(m)}$ は$m$番目のヘッドであることを表しています。 分割後のサイズはそれぞれhead_dim×dim_inで、head_dim = dim_out / num_heads です。
長さcontext_lengthの各トークンにおいて、$i$番目のトークンから見た$j$番目のトークンのアテンションスコア $s_{ij}$ は、
s_{ij}^{(m)} = \frac{ (\textbf{q}_i^{(m)})^{\top} {\textbf{k}_j^{(m)}} }{ \sqrt{d_{out}/M} } \qquad \dots(4)
で計算されます。$d_{out}$は dim_out、$M$はnum_headsです。
続いて、スコア $s_{ij}^{(m)}$ にソフトマックス関数を適用してアテンションウェイト $\alpha_{ij}^{(m)}$ を得ます。
\alpha_{ij}^{(m)} =
\begin{cases}
\displaystyle
\frac{\exp(s_{ij}^{(m)})}
{\sum_{k=1}^{i} \exp(s_{ik}^{(m)})}
& (j \le i) \\[12pt]
0 & (j > i)
\end{cases} \qquad \dots(5)
$i, j$の条件で場合分けされているのは、このあとの実装ではアテンションにマスク（後述）を適用しているためです。ざっと説明すると、$i$番目トークンから見て$j$番目トークンが過去か現在のトークンとなる$j \le i$の条件下ではそのままソフトマックスを計算します。$j$番目トークンが現在のトークンより未来である$j > i$の条件下では、アテンションウェイトを0にしています。
ここで、マルチヘッド $m$ からの出力 $\textbf{o}_i^{(m)}$ は、$N$をcontext_lengthとすると、
\textbf{o}_i^{(m)} = \sum_{j=1}^{N} \alpha_{ij}^{(m)} \textbf{v}_j^{(m)} \qquad \dots(6)
で計算されるhead_dim次元ベクトルです。この $\textbf{o}_i^{(m)}$ は、$i$番目のトークンから見た全てのトークンとのアテンションウェイトを含んだ埋め込みです。マルチヘッド注意機構からの最終的な出力は、コンテキストベクトルと呼ばれます。num_heads個のマルチヘッドからの埋め込みを連結して、
\textbf{o}_i = \textbf{W}_o
\left[
\begin{array}{c}
\textbf{o}_i^{(1)} \\
\vdots \\
\textbf{o}_i^{(M)} \\
\quad
\end{array}
\right] \qquad \dots(7)
という形をしています。$\textbf{W}_o$ は dim_out×dim_out次元の行列です。
さて、ここまでの数式をコードで実装してみたいと思います。
式$(1)$、$(2)$、$(3)$から見てみます。ヘッドが1つの場合のキー・クエリ・バリューの埋め込み計算に必要な行列 $\textbf{W}_q$、$\textbf{W}_k$、$\textbf{W}_v$ は、
self.W_query = nn.Linear(dim_in, dim_out, bias=qkv_bias)
self.W_key = nn.Linear(dim_in, dim_out, bias=qkv_bias)
self.W_value = nn.Linear(dim_in, dim_out, bias=qkv_bias)
の部分がそれに相当しています。これらはnum_heads個のヘッド $\textbf{W}^{(m)}_q$、$\textbf{W}^{(m)}_k$、$\textbf{W}^{(m)}_v$ に分割されます。分割後のヘッドの次元はhead_dimであり、
self.head_dim = dim_out // num_heads
と計算されます。
まず、forward()への引数xは次元 [ batch_size, context_length, embed_dim ] のテンソルであり、実際の値を [8, 256, 384] とします。そのxを使って
queries = self.W_query(x)
keys = self.W_key(x)
values = self.W_value(x)
でテンソルqueries、keys、valuesが計算されます。この各テンソルの次元はxと同じ [8, 256, 384] です。
次に、以下のコードでqueries、keys、valuesそれぞれをnum_heads個のヘッドに分割します。
queries = queries.view(batch_size, context_length, self.num_heads, self.head_dim)
keys = keys.view(batch_size, context_length, self.num_heads, self.head_dim)
values = values.view(batch_size, context_length, self.num_heads, self.head_dim)
queriesを例にすると、view関数によって次元は、
[ batch_size, context_length, dim_out ] = [8, 256, 384] から
[ batch_size, context_length, num_heads, head_dim ] = [8, 256, 12, 32] に分割されます。
さらに、次のtranspose関数で、context_length次元とnum_heads次元を入れ替えるように転置します。
queries = queries.transpose(1, 2)
keys = keys.transpose(1, 2)
values = values.transpose(1, 2)
分割の最終的な次元は [ batch_size, num_heads, context_length, head_dim ] = [8, 12, 256, 32] となり、最後の2つの次元 [context_length, head_dim] = [256, 32] が、分割後の1つヘッドの持つ次元となります。つまり、1バッチで次元[256, 32] のヘッドが12個あり、それが全部で8バッチあるということです。
ここまでで式$(1)$、$(2)$、$(3)$が実装できました。
どういうことかよく判らなかったので、図での説明を試みます。入力xは、次元 [ batch_size, context_length, dim_in ] = [8, 256, 384] なので、8個の [256, 384] テンソルがあることになります。Fig.4のようなイメージだと思います。
Fig.4 : 入力x
$b$ バッチ目を $\textbf{x}^{(b)}$ とすると、Fig.5のようなイメージです。
Fig.5 : bバッチ目の入力x
queries = self.W_query(x) の出力 queries の次元も [ batch_size, context_length, dim_in ] であり、$b$ バッチ目を $\textbf{q}^{(b)}$ とします。 $\textbf{q}^{(b)}$ の次元は [ context_length, dim_in ] = [256, 384] であり、これを view()、transpose()するとFig.6のイメージとなります。
Fig.6 : 256×384次元テンソルをviewしてからtransposeする
view() は384次元を12×32に変形しているのが見て取れます。transpose() は1軸目の256次元と2軸目の12次元を転置しており、1軸目が12次元、2軸目が256次元となっています。
transpose() が一番よくわからんので、例として1×3×2×32次元のテンソルで同じ操作をした時のイメージをFig.7にまとめてみました。
Fig.7 : 1×3×2×32次元テンソルをtransposeする
軸を入れ替えると、32次元をうまいこと再グループ化しているように見えます。Fig.6に戻ると、transpose() したら256×32次元のヘッドが12個となっているのが見て取れるかと思います。
コードに戻り、式$(4)$を実装します。アテンションスコアは
attn_scores = (queries @ keys.transpose(2, 3)) / keys.shape[-1]**0.5
で計算しています。keys に transpose() しているのは、同じサイズのテンソルで内積が計算できるように転置するためです。このとき、queries、keysの各バッチのそれぞれのヘッドの次元[256, 32]同士で内積が計算されます。
attn_scores の次元は[8, 12, 256, 256] です。これに対してアテンションウェイトを計算すると、
attn_weights = torch.softmax(attn_scores, dim=-1)
となり、次元は変わらず [8, 12, 256, 256] です。
続いて$(6)$式に相当するのが、
context_vec = (attn_weights @ values)
の行です。ただし、このcontext_vecは全てのヘッドを含んだテンソルです。
続いて式$(7)$による最終的な出力を実装します。
context_vec = context_vec.transpose(1, 2)
context_vec = context_vec.contiguous().view(batch_size, context_length, self.dim_out)
では、view()、transpose()して次元を変更した先ほどの処理を逆にたどり、テンソルの次元を元に戻しています。contiguous()は、view()する際にテンソルがメモリ上で要素順に並んでいないといけないという制約に従うためのおまじないです。
最後に、$(7)$の$\textbf{W}_o$に相当するout_projパラメータを適用します。
context_vec = self.out_proj(context_vec)
マスク処理
自己注意機構におけるマスク処理とは、アテンションウェイトを使ってコンテキストベクトル $\textbf{o}_i^{(m)}$ を計算するときに未来のトークンのウェイトを見てしまわないよう、Fig.8のようにアテンションウェイト行列の対角線より上側を隠してしまうことを指します。
Fig.8 : マスクのイメージ
この図の場合、単語"ランチ"の次の単語にアクセスするとき、"今日"、"の"、"ランチ"のウェイトだけを残すようにマスクします。こうすることで、モデルが未来の情報を見られないようにしています。
まず、torch.ones()でアテンションウェイトと同じサイズで全ての要素が1のテンソルを作成し、torch.triu()でマスクする部分以外を0にします。このとき、diagonal=1で対角線上の要素をマスクから外しています。
mask = torch.triu(
torch.ones(context_length, context_length),
diagonal=1
).bool()
さりげなく.bool()しているのは、この後のマスク処理のためです。ここまでをイメージで表すとFig.9です。
Fig.9 : マスクを作成する
続いてmaskでattn_scoresをマスクします。
attn_scores = attn_scores.masked_fill(mask, -torch.inf)
masked_fill()は、maskがTrueの要素を指定した値-torch.infで上書きします。マイナス無限大でマスクしている理由は、そのあとsoftmax()していることに依ります。
attn_weights = torch.softmax(attn_scores, dim=-1)
ソフトマックスの定義式にある $e^{-\infty}$ は0に近づくため、この要素は0として扱われます。そのため、マスクされた要素数に関わらず一律でソフトマックスを適用でき、マスクされていない要素の和が1となることが保証されます。
この一連の処理をイメージで表すとFig.10のようになります。
Fig.10 : マスクを適用する
この図はあくまでもイメージですが、softmax()によって各行の各要素の和がマスクの割合に関わらず1となっていることがわかると思います。
なお、
self.register_buffer("mask", mask)
のregister_buffer()は学習対象のパラメータではないTensorをnn.Moduleの一部として登録します。maskはオプティマイザによる学習が必要なTensorではありませんが、ほかのパラメータと共に計算する都合上、一緒にGPUへ移動して欲しいTensorです。そこでregister_buffer()することで、model.to("cuda")のときに他のパラメータと一緒にGPUへ移動します。また、model.state_dict()にも含まれるため、torch.save()のときに他のパラメータと一緒に保存されます。
トレーニング
ここで、ようやくモデルのトレーニング用のコードを準備します。まずはモデルインスタンスを作成し、GPUとCUDA環境がある場合はモデルを移動します。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GPTModel(GPT_CONFIG)
model.to(device)
最適化関数
モデルのパラメータを更新する最適化関数はAdamWを使用します。
optimizer = torch.optim.AdamW(
model.parameters(),
lr=TRAIN_CONFIG["max_learning_rate"],
weight_decay=0.1
)
学習率スケジュール
学習率スケジュールは、学習率を学習の進度によって動的に変化させる仕組みです。LLM学習では、学習初期はウォームアップと呼ばれる 0 → 最大学習率 までの単調増加、その後は学習終了にかけて 最大学習率 → 最小学習率 まで減少させる手法がよく取られます。こうすることで、学習初期では小さな学習率でパラメータを更新し、モデルが安定してきたら最大学習率で更新し、終盤にかけて細かく損失を改善していくことができます。
実装は、PyTorchに便利なモジュールLambdaLRがあるのでそれを使います。optimizerで指定する学習率を最大学習率として、各ステップにおいて最大学習率を1としたときの倍率を返す関数をlr_lambdaとすると、
from torch.optim.lr_scheduler import LambdaLR
scheduler = LambdaLR(
optimizer,
lr_lambda=lr_lambda,
)
となります。このように定義しておいて、次の学習率にステップしたいときは、
scheduler.step()
とすることで、optimizerの学習率を適切に変更してくれます。lr_lambda関数を作成するために、最大学習率を1とした倍率$r(t)$を次のように定義して考えます。
ウォームアップ時は、$t$を現在のステップ、$W$をウォームアップに必要なステップとすると、
r(t) = \frac{t}{W} \qquad (1 \le t \le W)
となります。
ウォームアップ後の学習率の減少は、Cosine Decayと呼ばれるコサイン曲線に沿って滑らかに減少させる手法が取られます。Cosine Decayの式は、$T$を学習終了までの最大ステップ、$W$をウォームアップのステップ、$l_{min}$と$l_{max}$をそれぞれ最小・最大学習率、$t$を現在のステップとすると、
r_{ratio} = \frac{l_{min}}{l_{max}} \qquad \dots(8)
r(t) = r_{ratio} + \frac{1 - r_{ratio}}{2} \left( 1 + \cos \left( \pi \frac{t-W}{T-W} \right) \right) \qquad \dots(9)
となります。これらをグラフにするとFig.11のようなイメージです。
Fig.11 : 最大学習率の倍率変化のイメージ
ここで、lr_lambda関数は次のように定義しています。
import math
# 1エポックあたりのOptimizer更新回数
steps_per_epoch = math.ceil(len(train_loader) / TRAIN_CONFIG["grad_accum_step"])
# 全エポックのOptimizer更新回数
max_steps = (steps_per_epoch * TRAIN_CONFIG["num_epochs"])
# max_stepsはOptimizerの最大更新回数
warmup_steps = int(max_steps * TRAIN_CONFIG["warmup_steps_ratio"])
# 学習率比
min_lr_ratio = TRAIN_CONFIG["min_learning_rate"] / TRAIN_CONFIG["max_learning_rate"]
def lr_lambda(step):
# step=0で戻り値が0.になるのを防ぐ
step = min(step + 1, max_steps)
# ウォームアップ時は最大学習率まで単調増加
if step < warmup_steps:
return step / warmup_steps
# ウォームアップ後は最小学習率までコサイン減衰
if step < max_steps:
progress = (step - warmup_steps) / (max_steps - warmup_steps)
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
return min_lr_ratio + (1.0 - min_lr_ratio) * cosine
# 最大ステップ後は最小学習率
return min_lr_ratio
$(9)$式の$t$は、lr_lambda関数の引数stepにあたります。このstepは単純なバッチの繰り返しではなく、後述する累積勾配を考慮したステップである必要があります。すなわちoptimizer.step()のタイミングでのOptimizer更新回数です。ということで、max_stepsも全エポックにおけるOptimizer更新の最大回数である必要があります。max_stepsは、
max\_steps = \biggl\lceil \frac{len(train\_loader)}{grad\_accum\_step} \biggr\rceil \times num\_epochs \qquad \dots(10)
で計算でき、Code.12でも記述されています。ちなみに$\lceil x \rceil$は$x$の小数点以下を切り上げる天井関数です。
トレーニングコード
データセットで学習と評価を繰り返し行うコードを定義します。コードの見通しを良くするために、必要最低限の構成となっています。学習の流れは、PyTorchによる通常のディープニューラルネットワークでの訓練とほぼ同じです。
tokens_seen, total_step, update_step = 0, 0, 0
for epoch in range(TRAIN_CONFIG["num_epochs"]):
model.train()
optimizer.zero_grad()
# バッチ数
num_batch = len(train_loader)
# 累積勾配ステップ
grad_accum_step = TRAIN_CONFIG["grad_accum_step"]
# 累積勾配ステップの余り
remainder = num_batch % grad_accum_step
for step, (input_batch, target_batch) in enumerate(train_loader):
loss = calc_batch_loss(
input_batch,
target_batch,
model,
device
)
# 平均の母数を調整
if remainder != 0 and step >= num_batch - remainder:
accum_size = remainder
else:
accum_size = grad_accum_step
# 損失を平均化
loss = loss / accum_size
loss.backward()
# 総トークン数
tokens_seen += input_batch.numel()
# 総ステップ数
total_step += 1
# Optimizer更新のタイミング
update = (
(step + 1) % grad_accum_step == 0
or (step + 1) == num_batch
)
if update:
# Optimizerのパラメータ更新
optimizer.step()
# 学習率を更新
scheduler.step()
optimizer.zero_grad()
# 更新ステップ
update_step += 1
# モデルの評価
if update_step % TRAIN_CONFIG["eval_update_step"] == 0:
train_loss, val_loss = evaluate_model(
model,
train_loader,
val_loader,
TRAIN_CONFIG,
device,
)
print(f"Epoch {epoch + 1} (step:{total_step}) Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")
損失計算
トレーニングコードの中でモデルの損失を計算している関数です。
def calc_batch_loss(input_batch, target_batch, model, device):
input_batch = input_batch.to(device)
target_batch = target_batch.to(device)
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(
logits.flatten(0, 1),
target_batch.flatten()
)
return loss
input_batchは次元 [ batch_size, context_length ] であり、ここでは [8, 256] とします。modelに入力することでlogitsは [ batch_size, context_length, vocab_size ] となり、ここでは [8, 256, 100277] です。vocab_sizeはトークナイザの語彙数です。
logitsの中身は、入力の1トークンに対して、次のトークンを決定する予測スコアが語彙数の100,277個分として表現されています。logitsの100,277次元は、最終的に確率分布に変換されます。デフォルトでは、100,277個のうち最も確率の大きい要素のインデックスが、トークナイザの語彙のインデックスに対応します。それが256トークン分、8バッチ分となっています。
ターゲットとlogitsとの誤差をクロスエントロピー関数で計算します。logitsはlogits.flatten(0, 1)で次元 [2048, 100277]、target_batchはtarget_batch.flatten()で[2048]に変換されて関数に渡されます。target_batchの各要素はトークナイザのトークンIDであり、正解の語彙のインデックスを指しています。
モデル評価
トレーニングコードでは、eval_update_stepの間隔でモデルを評価しています。これはevaluate_model関数で行われ、学習データ・検証データの両方での損失が計算され返されます。
def evaluate_model(model, train_loader, val_loader, train_config, device):
model.eval()
with torch.no_grad():
train_loss = calc_loader_loss(train_loader, model, device, 100)
val_loss = calc_loader_loss(val_loader, model, device, None)
model.train()
return train_loss, val_loss
evaluate_modelでは、model.eval()、torch.no_grad()でパラメータの更新を止めたうえで、train_loaderとval_loaderのそれぞれでcalc_loader_loss()を呼び、損失を計算します。
def calc_loader_loss(data_loader, model, device, num_batches=None):
total_loss = 0.
# 評価バッチ数の調整
if len(data_loader) == 0:
return float("nan")
# 指定がなければ全バッチ数で評価
elif num_batches is None:
num_batches = len(data_loader)
# 指定があれば全バッチ数と比べて少ない方で
else:
num_batches = min(num_batches, len(data_loader))
for count, (input_batch, target_batch) in enumerate(data_loader):
if count < num_batches:
loss = calc_batch_loss(
input_batch, target_batch, model, device
)
total_loss += loss.item()
else:
break
return total_loss / num_batches
calc_loader_loss()では、評価に使うバッチ数 num_batches だけ、さきほどのcalc_batch_loss()を呼んでいます。評価計算するバッチ数を限定しているのは、データローダの全バッチを計算していては時間がかかりすぎるためです。
算出した損失はtotal_loss += loss.item()として足し合わされ、最後にtotal_loss / num_batchesで全バッチ数の平均値として返しています。
トレーニングのステップ
1回のエポックでバッチを繰り返した回数は、stepがカウントしています。この値は最大len(train_loader) - 1です。一方、total_stepはエポックをまたいだstepの累積総数です。
stepを重ねてもoptimizer.step()でモデルの重みを更新するのはgrad_accum_stepの時だけなので、その回数をupdate_stepでカウントしています。これをOptimizer更新ステップと呼び、実質的な学習の進み具合として使われる指標です。
累積トークン
累積勾配を適用しない場合、1回の学習量はstepごとのバッチ数batch_sizeで決まり、その学習トークン量はbatch_size × context_lengthです。一方、累積勾配を適用すると、1回の学習量はoptimizer.step()ごとに累積した累積トークン(Tokens seen)で決まり、その学習トークン量はbatch_size × grad_accum_step × context_lengthです。LLM学習では、学習量を考える上でこの累積トークンが重要な指標となります。というのも、累積勾配では同じbatch_sizeであってもcontext_lengthとgrad_accum_stepの組み合わせで学習トークン量が変わってしまうためです。
累積勾配
Gradient Accumulation(累積勾配)は、複数バッチ分の勾配計算を累積させてからoptimizer.step()することで、小さなGPUメモリサイズ上では乗り切らないバッチサイズの学習を実現する手法です。
Fig.12 : 累積勾配の考え方
通常の学習では、Fig.12の左側のように、1バッチごとにmodel.forward()→loss.backward()→optimizer.step()→optimizer.zero_grad()の手順を踏みます。コードにすると
optimizer.zero_grad()
for input_batch, target_batch in train_loader:
loss = calc_batch_loss(
...
)
loss.backward()
optimizer.step()
optimizer.zero_grad()
という感じです。
一方Gradient Accumulationでは、Fig.12の右側のように、grad_accum_step個のバッチの勾配をためてから一気にoptimizer.step()します。こちらもコードにすると以下のような感じであり、Code.13トレーニングコードと同じ構造をしています。
grad_accum_step = 4
optimizer.zero_grad()
for step, (input_batch, target_batch) in enumerate(train_loader):
loss = calc_batch_loss(
...
)
loss = loss / grad_accum_step
loss.backward()
if (step + 1) % grad_accum_step == 0:
optimizer.step()
optimizer.zero_grad()
なぜわざわざこんなことをするのかというと、学習に使用するGPUメモリのサイズとの兼ね合いです。GPT-2モデルの場合、batch_size = 512、context_length = 1024であり、optimizer.step()あたりの学習トークン数が512 × 1024 ≒ 500kトークンです。512バッチ分のトークンをGPUに乗せようとすると、数百GB必要になるとかならないとか。
一方grad_accum_step = 64、batch_size = 8として考えてみます。この場合、バッチサイズが8なのでGPUリソースを大幅に削減できます。1回のloss.backward()でのトークン数は 8 × 1024 = 8192トークンですが、これを64回勾配を累積してからoptimizer.step()となるので、optimizer.step()あたりの学習トークン数は8 × 1024 × 64 ≒ 500kトークンです。
つまり、Gradient Accumulationを適用すると、optimizer.step()あたりの学習トークン数500kを維持しながら、バッチ数を削減できます。
累積の余りバッチ
なお、ここではlossをgrad_accum_stepで割っています。lossはoptimizer.step()する時点でgrad_accum_step回分が累積され、grad_accum_step = 4であれば4倍程度になってしまいます。そこで各ステップでloss / accumulation_stepsすることで、optimizer.step()では4回の平均程度のlossとしています。
平均の母数の調整
さらに、num_batch = len(train_loader)がgrad_accum_stepで割り切れればいいですが、割り切れない場合、1〜grad_accum_step未満のバッチ数の勾配計算の結果がモデルに反映されなくなってしまいます。例えばFig.13の場合、num_batch = 10、grad_accum_step = 4です。step = 4, 8のときは正しくoptimizer.step()されますが、step = 12は無いため、バッチ9と10を取りこぼしてしまいます。
Fig.13 : 母数の調整
たとえ数バッチであってもデータは無駄にできません。そこで、Code.13トレーニングコードではoptimizer.step()する条件を次のようにしています。
update = (
(step + 1) % grad_accum_step == 0
or (step + 1) == num_batch
)
if update:
optimizer.step()
optimizer.zero_grad()
(step + 1) == num_batchが最後、余ったバッチ10に合致する条件です。さらにこのとき、3つ目のoptimizer.step()はバッチが2つしかないので、
loss = loss / grad_accum_step
のところでlossを割る母数は4ではなく2にしたいところです。そこで、Code.13トレーニングコードでは、次のようにgrad_accum_stepの余り( remainder )を求めておいて、最終のoptimizer.step()のみの条件でlossを割る母数を差し替えます。
num_batch = len(train_loader)
grad_accum_step = train_config["grad_accum_step"]
remainder = num_batch % grad_accum_step
...
if remainder != 0 and step >= num_batch - remainder:
accum_size = remainder
else:
accum_size = grad_accum_step
loss = loss / accum_size
モデルの保存
トレーニングコードでは言及されていませんが、各エポックの学習終了時に、モデルとパラメータを保存するようにしています。モデルの保存は、
import torch
path = "path/to/model.pth"
torch.save(
{
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
path
)
としています。GPT_CONFIGとTRAIN_CONFIGパラメータは、これらがディクショナリなので、
import json
path = "path/to/model.json"
state = {
"gpt_config": GPT_CONFIG,
"train_config": TRAIN_CONFIG,
}
with open(path, "w") as f:
json.dump(state, f, indent=2)
としています。
実験
では、パラメータをいくつかいじりながらモデルをトレーニングさせて、実験としてみたいと思います。
vs トークン数
LLM学習の進度を測る指標の1つとして、最適化関数の1更新で処理されるトークン数(Tokens/Optimization)があり、その値はbatch_size × grad_accum_step × context_lengthで計算できます。そこで、トークン数を2k、8k、16kで変化させて様子を見てみます。トークン数以外のパラメータはGPT_CONFIG、TRAIN_CONFIGに記載の通りです。
Fig.14 :
16kでは、学習が進んでもTrainとValの損失値にあまり差がありませんが、損失値は4からあまり低下せず、学習があまり進んでいないことが伺えます。逆に2kでは、16kに比べて損失は大幅に低下していますが、TrainとValidationの損失値に差が広がり、いわゆる過学習の状態と思われます。
学習が完了した状態で2kはそこそこ意味の通るテキストを出力できますが、16kではできませんでした。また、16kは10エポック以降学習を続けさせても、やはりTrainの損失が低下してくるあたりでValidationとの損失差が大きくなるため、学習時間がかかるだけで結果は2kと同じになります。つまり、トークン数は少ない2kが効率よく学習できるが、過学習が避けられないという状態です。
vs 学習率
学習率を$1e^{-4}$、$3e^{-4}$、$5e^{-4}$で変化させます。学習率以外の条件は、vsトークン数の2kと同じです。
Fig.15 :
$1e^{-4}$では、学習が進んでもTrainとValidationの損失値にあまり差がありませんが、損失値は4からあまり低下せず、学習があまり進んでいないことが伺えます。逆に$5e^{-4}$では、$1e^{-4}$に比べて損失は大幅に低下していますが、TrainとValidationの損失値に差が広がっています。トークン数の検討同様にこちらも過学習の状態と思われます。
学習が進んでいない$1e^{-4}$を10エポック以降も学習を続けても、Trainの損失が低下してくるあたりでValidationとの損失差が大きくなり、学習時間がかかるだけで結果が$5e^{-4}$と同じになってしまうのも、トークン数での結果と同じです。
vs 埋め込み次元
使用しているトークナイザの語彙空間(100,277)に対して、the-diaryデータセットのユニークな語彙数(3,117)が小さすぎます。しかし語彙数(3,117)に合わせてトークナイザを自作しても損失が悪化したので、せめてembed_dimを減らして学習させてみました。
Fig.16 :
768次元を半分にした384次元の損失の変化が安定しています。
考察
パラメータの組み合わせによって学習の進捗が遅かったり早かったりします。どの組み合わせでも結局、同じ程度までTrain損失が低下するころにはValidationとの差が開く過学習の状態になるのは変わりません。たまたまデータの取り合わせが良くなかったことも考えられるので、交差検証法を使ってデータセット全体に渡って走査することも必要かもしれないです(が、そこまでやりきれませんでした)。
学習の状態が改善されないのは、おそらくデータセットのサイズと品質の問題だと考えています。データセットのサイズが小さいのはさることながら、品質がよろしくありません。データセットの章で触れたように、全データ約224万トークンから重複をなくしたユニークなトークンがわずか3,117トークンであり、これは全体トークンに対して少なすぎるのでしょう。
出力
効率よく学習が進行するパラメータを選び、the-diaryデータセットすべてのデータを用いてモデルを最終学習しました。このthe-diaryモデルに対してインタラクティブにテキストを入力し、その出力を得てみたいと思います。なお、このthe-diaryモデルは事前学習の段階であり、指示チューニングされていません。なので、入力のテキストの続きを出力するというシンプルなタスクを行うことができます。
テキスト出力コード
the-diaryモデルに対してテキストを出力させるコードです。まずは、学習時に保存したモデルのmodel_state_dictをロードして、モデルを復元しています。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
path = "path/to/model.pth"
model_state = torch.load(path, map_location=device)
model = GPTModel(GPT_CONFIG)
_ = model.load_state_dict(model_state["model_state_dict"])
_ = model.to(device)
次に、復元したモデルを使用してテキストを入力し、それに続くテキストを出力するコードです。
import tiktoken
OUTPUT_CONFIG = {
# 開始コンテキスト
"start_context": "今日は飲みに行きたい",
# コンテキスト長さ
"context_length": 256,
# 温度
"temperature": 0,
# TopK
"top_k": None,
}
tokenizer = tiktoken.get_encoding("cl100k_base")
output = generate_text(model, tokenizer, device, OUTPUT_CONFIG)
print(f"入力：「{OUTPUT_CONFIG["start_context"]}」")
print(f"出力：「{output}」")
出力の設定はOUTPUT_CONFIGでしており、それをもとにgenerate_text()でモデルからテキストを出力しています。
generate_text関数
def generate_text(model, tokenizer, device, output_config):
model.eval()
context_length = model.position_embedding.weight.shape[0]
# 開始コンテキストをトークン化
token_ids = tokenizer.encode(output_config["start_context"], allowed_special={'<|endoftext|>'})
token_ids = torch.tensor(token_ids)
token_ids = token_ids.to(device)
with torch.no_grad():
token_ids = generate_ids(
model=model,
token_ids=token_ids,
context_length=context_length,
max_tokens=output_config["context_length"],
temperature=output_config["temperature"],
top_k=output_config["top_k"],
)
return tokenizer.decode(token_ids.tolist())
generate_text()では、受け取った入力テキストstart_contextをtokenizerでトークンIDのリストにエンコードします。続いてgenerate_ids()で入力トークンIDリストの続きとなる、モデルが推論したトークンIDのリストを取得しています。最後にtokenizerでテキストにデコードしています。
generate_ids関数
def generate_ids(model, token_ids, context_length, max_tokens, temperature=0, top_k=None, eos_id=100257):
for _ in range(max_tokens):
# 入力トークン数を制限
idx_cond = token_ids[-context_length:]
with torch.no_grad():
idx_cond = idx_cond.unsqueeze(0)
# 予測スコアを取得
logits = model(idx_cond)
logits = logits[:, -1, :]
# Top-k
if top_k is not None:
# Top-kで並び替え
top_logits, _ = torch.topk(logits, top_k)
# 最小値を取得
min_val = top_logits[:, -1]
# 最小値より小さい要素を-infに変更
logits = torch.where(
logits < min_val,
torch.tensor(float('-inf')).to(logits.device),
logits
)
# 温度
if temperature > 0.:
# 出力の分布を温度で調整
logits = logits / temperature
# 確率分布化
probs = torch.softmax(logits, dim=-1)
# 確率分布に従ってトークンを決定
token_id_next = torch.multinomial(probs, num_samples=1)
else:
# 常に最大出力のトークンを返す
token_id_next = torch.argmax(logits, dim=-1, keepdim=True)
if token_id_next == eos_id:
break
token_id_next = token_id_next.squeeze(0)
# 取得したトークンを現在のトークン列につなげて次の入力とする
token_ids = torch.cat((token_ids, token_id_next), dim=-1)
return token_ids
generate_ids()では、出力したトークン数がmax_tokensに達するまで繰り返し1トークンずつ出力し、token_idsに繋げていきます。もしtoken_idsの長さがモデルに入力できる上限であるcontext_lengthを超えていた場合、token_ids[-context_length:]で前方のトークンから順番に削除されます。
モデルからの出力は、
with torch.no_grad():
token_ids = token_ids.unsqueeze(0)
# モデルからトークン分布を取得
logits = model(token_ids)
logits = logits[:, -1, :]
の部分です。token_ids.unsqueeze(0)しているのは、モデル学習時に合わせてバッチの次元を追加しているためです。出力であるlogitsは、token_idsの全てのトークンについて、次のトークンの予測スコアを保持しています。しかし最後のトークンに対する予測スコアしか必要ないため、logits[:, -1, :]で最後のトークンの確率分布をスライスしています。
また、途中で<|endoftext|>が出力されることもあります。これはテキストの区切りを表す特殊なトークンでした。区切りであるため、
if token_id_next == eos_id:
break
でトークンの出力を停止しています。なお、eos_id = 100257であり、このトークンIDは<|endoftext|>に対応しています。
ループの最後、
token_id_next = token_id_next.squeeze(0)
# 取得したトークンを現在のトークン列につなげて次の入力とする
token_ids = torch.cat((token_ids, token_id_next), dim=-1)
で追加したバッチの次元を削除し、token_idsの最後にtoken_id_nextを追加しています。このあとループの先頭に戻り、モデルに再びtoken_idsを入力して次のトークンを出力します。
Top-k サンプリング
Top-k サンプリングでは、出力されたlogitsのうち上位$k$個以外の要素の値を負の無限大にすることで、ソフトマックス計算時にその$k$個以外のトークンを確率分布から除外できます。結果、確率の高い分布からトークンを選択することができ、よりもっともらしい文章を生成できます。
コード中では、
if top_k is not None:
# Top-kで並び替え
top_logits, _ = torch.topk(logits, top_k)
# 最小値を取得
min_val = top_logits[:, -1]
# 最小値より小さい要素を-infに変更
logits = torch.where(
logits < min_val,
torch.tensor(float('-inf')).to(logits.device),
logits
)
で実装されています。
torch.topk(logits, top_k)は、テンソルlogitsのうち最大値からtopk個の値とそのインデックスを返し、それをtop_logitsと_で受けています。top_logitsは最大値から順に並んでいるため、top_logits[:, -1]で最小値を取得できます。torch.where()では、logitsのうち最小値より小さい要素を-infに設定しています。
温度スケーリング
温度(Temperature)スケーリングとは、0よりも大きな数値でlogitsを割ることです。温度というのはもったいぶった言い方です。
コード中では、
if temperature > 0.:
# 出力の分布を温度で調整
logits = logits / temperature
# 確率分布化
probs = torch.softmax(logits, dim=-1)
# 確率分布に従ってトークンを決定
token_id_next = torch.multinomial(probs, num_samples=1)
else:
# 常に最大出力のトークンを返す
token_id_next = torch.argmax(logits, dim=-1, keepdim=True)
で実装されています。
logits / temperatureがどういうことか見てみます。温度が1よりも高い場合、logitsによる確率分布はより均一な分布となります。温度が1の場合、温度スケーリングを適用しないのと同じ結果です。温度が1よりも低い場合、logitsによる確率分布は尖った分布となります。
極端な例として、
logits = torch.tensor([1., 2., 3.])
というlogitsがあった場合、temperature = 1に対して、大きい温度:5.、小さい温度:0.5でsoftmaxまで計算すると、
temperature = 0.5 : probs = tensor([0.0159, 0.1173, 0.8668])
temperature = 1 : probs = tensor([0.0900, 0.2447, 0.6652])
temperature = 5 : probs = tensor([0.2693, 0.3289, 0.4018])
という分布になります。これらをグラフにするとこうなります。
Fig.17 : 温度を適用するイメージ
コードではこの状態でtorch.multinomial(probs, num_samples=1)としています。これは、probsの確率分布に比例する確率でそのインデックスをサンプリングします。
つまり、温度が1よりも低い場合logitsによる確率分布は尖った分布となり、torch.multinomial()によって最も高い確率のトークンが選択される可能性を高めます。温度が1の場合は温度スケーリングを適用しないのと同じ結果です。温度が1よりも高い場合logitsによる確率分布はより均一な分布となり、torch.multinomial()によって確率が低いトークンであっても選択される可能性が高くなります。
温度を適切に設定することで、生成するテキストの多様性を高めることができます。ただし、高い温度では意味の破綻したテキストとなる可能性がより高くなります。
出力サンプル
こうして得ることができたテキストをサンプルとして掲載します。Top-kは10で固定、温度は低め(0.5〜0.8)にしています。
人工知能とは
入力：人工知能とは
出力：人工知能とは関係のない人々を誘うことになった。本当に楽しいのかと思ったけど、どうやら健康診察はそうでもない。
仕事帰りに寄った居酒屋で
入力：仕事帰りに寄った居酒屋で
出力：仕事帰りに寄った居酒屋で軽く飲んで帰ってきた。 今日はおとなしいけど、なぜかおれが来ていないからかなり慌てて、寒い。
だけども問題は今日の雨。傘がない
入力：だけども問題は今日の雨。傘がない
出力：だけども問題は今日の雨。傘がないから帰ろうと思っていたけど、それからのことだ。なんか最近は寒い。気温も上がらないという。
過学習がある程度進むと、恥知らずにもモデルは、私の日記と思われるコンテキストを出力するようになります。特に友人の名前を中心としてテキストを表示し、まるで私の架空の日記を読んでいるようです。また、出力が私の文章の書き方に似ます。誰にも読まれることを想定していないときの素のままの自分の文章が出力されると、なんとも言えない気分になります。
まとめ
以上、スクラッチでGPT-2モデルを構築して、私の日記で事前学習してみました。現在では自律的にサイバー攻撃を仕掛けるほどに高度化したLLMモデルにおいて、もちろん各ベンダーで日進月歩の研究や工夫があるのでしょうが、モデル自体の原理的な構造は驚くほど簡潔です。学習結果は拙いものですが、限られたコンテキストでかなり正確に文法と単語の意味を習得しているように感じました。
最後に
この記事・コード・図は人間が書きました。
コスト
当初、必要なGPUメモリが分からなかったので、Amazon EC2のg6シリーズGPUインスタンス上で学習していました。そこで出力を見ながらパラメータの調整をした結果、「Google Colab Pro+で十分だった」という結論に至るまでに200ドルほど費やしました。Google Colab Pro+に移行してからは、1ヶ月分のマシンリソースです。
参考書籍
[1] "つくりながら学ぶ！LLM 自作入門", Sebastian Raschka, マイナビ出版, 2025
[2] "大規模言語モデル入門", 山田育矢ら, 技術評論社, 2023