【実装】ミニGPTを自作する — 300行の言語モデル

PyTorchで文字単位のミニGPTを最初から組む。トークナイザ・因果マスク付き自己注意・学習ループ・温度サンプリングまでを1本のコードとして通し、シェイクスピアの文体が立ち上がるまでを追う。

Attention Is All You Need

300行で何ができるのか

「GPTを作る」と聞くと、GPUが数千枚と数兆トークンのデータが要る話に聞こえます。商用モデルの学習は実際そうです。けれどアーキテクチャそのものは驚くほど小さい。文字を数値に変えるトークナイザ、埋め込み、自己注意、フィードフォワード、学習ループ、生成ループ — 全部合わせてPyTorchで300行前後に収まります。

ここで作るのは文字単位（character-level）のGPTです。単語ではなく1文字ずつ予測します。語彙が数十種類しかないので出力層が小さく、トークナイザは20行で済む。そのぶん「Transformerの本体」だけに集中できます。学習データはシェイクスピアの戯曲を1本のテキストファイルにまとめたもの（いわゆる tiny shakespeare、おおよそ100万文字・ユニーク文字は65種類ほど）を使います。

本記事のゴールはコードを写経することではありません。300行のどこに何が書いてあり、どの行を消すと何が壊れるかを説明できるようになることです。ライブラリの model.generate() を呼んでいるだけの状態から、自分で設定値を選べる状態に変わる境目はここにあります。

GPTの正体は「次の1文字の確率分布」

GPTがやっていることは、突き詰めると1つの関数です。これまでの文字列を受け取り、次に来る1文字の確率分布を返す。 それだけです。

「To be or not to b」まで見せられたら、次は "e" の確率が高く "z" は低い。この確率を語彙全体にわたって出す。文章生成は、この関数を呼んで1文字サンプリングし、末尾に足してまた呼ぶ — の繰り返しにすぎません。「モデルが文章を書いている」のではなく、サイコロを振る回数だけ関数を呼び直しているわけです。

学習は「実際に来た文字にできるだけ高い確率を割り当てる」ように重みを調整することです。損失は交差エントロピー1本で表せます。

要するに「各位置 $t$ で、実際に来た文字 $x_t$ にモデルが割り当てた確率の対数を取り、符号を反転して平均する」と言っています。$x_{<t}$ は「$t$ より前の文字すべて」、$\theta$ はモデルの全パラメータ、$T$ は一度に見る文字数です。正解に確率1を与えていれば $\log 1 = 0$ で損失ゼロ、正解を軽視するほど値が跳ね上がります。

ここで効率の話を1つ。長さ $T$ の文をモデルに1回通すと、$T$ 個の位置すべてについて同時に予測と損失が得られます。1文字ずつ学習しているわけではありません。これが「次の文字当て」が大量のテキストを食える理由で、後で出てくる因果マスクは、この同時計算を成立させるための仕掛けです。

出力の生の数値（ロジット）を確率に変えるのが softmax で、生成時にはここに温度 $\tau$ を1つ挟みます。

$$p_i = \frac{\exp(z_i/\tau)}{\sum_j \exp(z_j/\tau)}$$

つまり、文字 $i$ のロジット $z_i$ を $\tau$ で割ってから softmax するだけです。$\tau$ が小さいほど大きいロジットが強調されて分布が尖り、大きいほど平らになります。この割り算1つが「無難だが退屈」と「独創的だが支離滅裂」のあいだを動かすつまみになります。実際に触ってみてください。

部品表 — 300行の内訳

残りはインポートと設定値です。新しい概念はもう出てきません。ここから先は、この表を上から順に埋めていく作業になります。

文字を整数にする

ニューラルネットは数値しか扱えないので、まず文字に整数を割り当てます。全文からユニーク文字を集めてソートし、辞書を2本作るだけです。

import torch, torch.nn as nn
from torch.nn import functional as F

text  = open('input.txt', encoding='utf-8').read()
chars = sorted(set(text))                      # 65種類ほど
stoi  = {c: i for i, c in enumerate(chars)}
itos  = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

本物のGPTはここが BPE などのサブワード分割になり、語彙は数万規模になります。違いは語彙サイズと系列長のトレードオフだけで、その先の構造は同じです。文字単位は語彙が小さくて済む代わりに、同じ文章がずっと長い系列になり、モデルはより遠くを見なければならなくなります。分割方式の選び方はトークナイザを1から理解するにまとめてあります。

入力は「1文字ずらした2本」

学習データの作り方が、この実装で一番よく誤解されるところです。入力 x と正解 y は同じ文字列を1文字ずらしただけです。

def get_batch(data, block_size, batch_size):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i + block_size]         for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x, y

x が "To be or not to b" なら y は "o be or not to be" です。位置0では "T" から "o" を、位置1では "To" から " " を…という具合に、1回のバッチで block_size 個分の予測問題が同時に作られます。人手でラベルを用意する必要がないので、テキストさえあれば問題が無限に湧く。これが自己教師あり学習の実体です。

心臓部: 因果マスク付き自己注意