GPT をスクラッチで書いたコード MinGPT(https://github.com/Keigo-Iwakuma/MinGPT/blob/main/mingpt/model.py など MinGPTで調べれば出てきます) の解説記事
翻訳記事になります。
最近、GPT-2などの自己回帰モデルを使って、自然言語生成の分野を調べています。HuggingFace Transformer は pretrained な言語モデルを提供して、多くのものは最小限の調整でも既製品のまま使えます。しかしながらこのポストでは、私たちは小さい GPT モデルを pytorch でスクラッチで作ってみます。それらのものを真に理解せずに使うことに私はたまに不気味なほど不安を感じていると思ってきていました。気づいたのは、それらの実装は 前のポストで書いた transformer の実装ととても似通っていることなので、みた人であれば繰り返しがそこらじゅうであると分かるでしょう。さらに役に立ったのが、この記事の実装でも大きく参照している Karpathy の minGPT のレポジトリです。では始めてみましょう。
セットアップ
まずはこのチュートリアルで必要なモジュールをインポートしましょう。ほとんどのものは、ビルトインのmath モジュールを除けば、torch からきています。
import math
import torch
from torch import nn
import torch.nn.functional as F
HuggingFace Transformer を使った経験と Karpathy の実装を見たところ、モデルのパラメーターの全ての初期化を含む設定のオブジェクトがあるのが慣習になっていると私は思っています。Karpathy のレポジトリーから取ってきた下のスニペットが、数々の定数・これから作るモデルのパラメータを含む基礎的なクラスの作り方を示しています。やろうと思えば、より多くのレイヤー・最大のシーケンス長・エンベディングディメンションを加えることで、GPT-2 や GPT-3 を作るクラスの設定も簡単にできます。
class GPTConfig:
attn_dropout = 0.1
embed_dropout = 0.1
ff_dropout = 0.1
def __init__(
self, vocab_size, max_len, **kwargs
):
self.vocab_size = vocab_size
self.max_len = max_len
for key, value in kwargs.items():
setattr(self, key, value)
class GPT1Config(GPTConfig):
num_heads = 12
num_blocks = 12
embed_dim = 768
もし Transformer アーキテクチャーに慣れている人なら、これは似てるように思えるでしょう。max_len は モデルが処理できる最大の長さになります。Transformer モデルは全ての入力を一度に並列して処理しているので、その ウィンドウスパンは無限ではありません（なので、この制約と折り合いをつけるため Transformer XL などのモデルが出来ました）。vocab_suze は 語彙のサイズ、別の言い方で言えば、モデルがどれくらいのトークンを知っているべきか、を示します。num_blocks は Transformer の decoder のレイヤーの数・num_heads はヘッドの数を示します。
実装
では、モデルを作っていきましょう。下に示しているのは、最上位のモデルのアーキテクチャですが、私の意見ではそれはとても短くシンプルです。
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
embed_dim = config.embed_dim
self.max_len = config.max_len
self.tok_embed = nn.Embedding(
config.vocab_size, embed_dim
)
self.pos_embed = nn.Parameter(
torch.zeros(1, config.max_len, embed_dim)
)
self.dropout = nn.Dropout(config.embed_dropout)
self.blocks = nn.Sequential(
*[Block(config) for _ in range(config.num_blocks)]
)
self.ln = nn.LayerNorm(embed_dim)
self.fc = nn.Linear(embed_dim, config.vocab_size)
def forward(self, x, target=None):
# batch_size = x.size(0)
seq_len = x.size(1)
assert seq_len <= self.max_len, "sequence longer than model capacity"
tok_embedding = self.tok_embed(x)
# tok_embedding.shape == (batch_size, seq_len, embed_dim)
pos_embedding = self.pos_embed[:, :seq_len, :]
# pos_embedding.shape == (1, seq_len, embed_dim)
x = self.dropout(tok_embedding + post_embedding)
x = self.blocks(x)
x = self.ln(x)
x = self.fc(x)
return x
このモデルが見かけだけではシンプルに見えるの理由は、実際、多くのモデルが、ほとんどの難しいことをしている Transformer decoder レイヤー群である GPT.block からきているからです。このクラス中で興味をそそる唯一のロジックは、トークンと position embeddings を組み合わせて decoders 群への入力としている部分です。
実装の詳細で Karpathy のコードから私が学んだことは、彼がどのようにして positional embeddings を対処しているかということである。専用のトレーニングできる positional embedding layer を作る代わりに、私たちはシンプルにルックアップ行列を positional embedding レイヤーのようなものとして登録し、入力の長さに応じて シンプルに行列を適正なシーケンスの長さにスライスしています。positional embedding を実装するのは、embedding レイヤーの方法で実装する場合 毎度の forward を通るたびに torch.range() を呼ぶよりいい方法だと思います。
最終的なアウトプットの形は (batch_size, seq_len, vocab_size) になっています。私たちは最終的なアウトプットを ポジションごとのトークンの予想と解釈することができます。私たちは例えば teacher forcing を使って モデルを訓練させパラメーターを更新できます。自己回帰生成の文脈だと、シーケンスの最後のポジションにあるトークンの予想を使って、それをオリジナルのインプットにくっつけて、その後で変更されたシーケンスをモデルにフィードバックします。
Decoder ブロック
では decoder・Transformer decoder ブロック の構成要素を見てみましょう。decoder ブロックは、multi-head attention、layer normalization と point-wise feedforward から構成されています。私たちは、各構成要素の間で residual connection (残差接続) を使っています。feedforward ネットワークは、contextual embeddings を生かしそれらの表現を高める 隠れた次元を一時的に増やす層であると理解できます。またですが、興味深い部分はまた multi-head self-attention レイヤーに引き伸ばされました。
class Block(nn.Module):
def __init__(self, config):
super().__init__()
embed_dim = config.embed_dim
self.ln1 = nn.LayerNorm(embed_dim)
self.ln2 = nn.LayerNorm(embed_dim)
self.attn = MultiheadAttention(config)
self.ff = nn.Sequential(
nn.Linear(embed_dim, embed_dim * 4),
nn.GELU(),
nn.Linear(embed_dim*4, embed_dim),
nn.Dropout(config.ff_dropout)
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.ff(self.ln2(x))
return x
multi-head- self-attention レイヤーは、興味深い key-value-query 操作を含んでいます。ついでに、ここで私が見てきたいくつかの興味深い討論や詳細について触れてみます。
- Self-attention は、入力シーケンスにあるトークンがノード・各トークンの関係性がエッジを表現している グラフニューラルネット もしくは GNN と見ることができます。encoder ブロックに関して言えば、GNN は full で complete、つまり全てのノードが他のノードとつながっています。decoder では、トークン は 入力シーケンスで そのトークンより前に来た他のトークンとのみ接続されます。
- もしある人が、key vector を query vector として使うと決めたら、つまり query 行列 WQ 全体を実質的に削除すると、graph neural network は 実質的に無向グラフになります。これはなぜならば、node A と node B の距離は、node B と node A の距離と違わなくなるからである。言い換えると、Attention(na, nb) = Attention(nb, na)になります。しかしながら、オリジナルだったり、query vector が key vector と分離している もっと普通な Transformer の実装だと、この交換可能な関係は必ずしも成立するとは限らず、attention レイヤーが 有向グラフ になります。
これらはせいぜい直感的で推論に基づく解釈であったし、私は GNN について 上で書いた以上のニュアンスのコメントを するほどの知識はないのだろう。しかし、この解釈は非常に興味深い。
Multi-Head Attention
話を元に戻すと、下のコードが Multi-head attention レイヤーの実装になる。これは 私が Transformer でレイヤーを実装した時と非常に似ているので、詳細の説明をいくつか省略しました。各アウトプットの形をコメントに追加したので、forward のフローを理解しやすくなっている。
※ [コメント] Attention や Transformer の実装については、作って理解する Transformer / Attention などを参考にするとコードが読みやすいです。
class MultiheadAttention(nn.Module):
def __init__(self, config):
super().__init__()
embed_dim = config.embed_dim
self.num_heads = config.num_heads
assert embed_dim % self.num_heads == 0 "invalid heads and embedding dimension configuration"
self.key = nn.Linear(embed_dim, embed_dim)
self.value = nn.Linear(embed_dim, embed_dim)
self.query = nn.Linear(embed_dim, embed_dim)
self.proj = nn.Linear(embed_dim, embed_dim)
self.attn_dropout = nn.Dropout(config.attn_dropout)
self.proj_dropout = nn.Dropout(config.ff.dropout)
self.register_buffer(
"mask",
torck.tril(torch.ones(config.max_len, config.max_len)).unsqueeze(0).unsqueeze(0)
)
def forward
batch_size = x.size(0)
seq_len = x.size(1)
# x.shape == (batch_size, seq_len, embed_dim)
k_t = self.key(x).reshape(batch_size, seq_len, self.num_heads, -1).permute(0, 2, 3, 1)
v = self.value(x).reshape(batch_size, seq_len, self.num_heads, -1).transpose(1, 2)
q = self.query(x).reshape(batch_size, seq_len, self.num_heads, -1).transpose(1, 2)
# shape == (batch_size, num_heads, seq_len, head_dim)
attn = torch.matmul(q, k_t) / math.sqrt(q.size(-1))
# attn.shape == (batch_size, num_heads, seq_len, seq_len)
mask = self.mask[:, :, :seq_len, :seq_len]
attn = masked_fill(mask == 0, float("-inf"))
attn = self.attn_dropout(attn)
# attn.shape == (batch_size, num_heads, seq_len, seq_len)
attn = F.softmax(attn, dim=-1)
y = torch.matmul(attn, v)
# y.shape == (batch_size, num_heads, seq_len, head_dim)
y = y.transpose(1, 2)
# y.shape == (batch_size, seq, num_heads, head_dim)
y = y.reshape(batch_size, seq_len, -1)
# y.shape == (batch_size, seq_len, embed_dim)
y = self.proj_dropout(self.proj(y))
return y
私が一番困惑したのは masking が attention 行列とどのように動いているかの部分でした。私は概念的に decoder は未来のトークンを知っていてはいけないことを知っていましたのですが、それが masking が必要な理由でしたし、実際の具体的な例を見れてとても役に立ちました。
したが、mask 行列で、最大のシーケンス長が5の decoder であると仮定しています。
max_len = 5
mask = torch.tril(torch.ones(max_len, max_len)).unsqueeze(0).unsqueeze(0)
mask
tensor([[[[1., 0., 0., 0., 0.],
[1., 1., 0., 0., 0.],
[1., 1., 1., 0., 0.],
[1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1.]]]])
考え方としては、モデルは attention 行列の 1の要素の部分は見れるが、0 では マスクを適用させます。
ここで、モデルが バッチサイズ3までしか受け入れないとしましょう。そうすると下の mask 行列部分のみを使うことになります。
seq_len = 3
mask = mask[:, :, :seq_len, :seq_len]
mask
tensor([[[[1., 0., 0.],
[1., 1., 0.],
[1., 1., 1.]]]])
ここが私を一番困惑させた部分で、私は mask がバッチの概念では把握できていなかった部分です。mask 行列でunsqueeze(0)を呼んでいるのを思い出してください。これはなぜならバッチをブロードキャスティングで対処したいからです。したがって、このマスクはバッチの１つの例だけを暑かったものと考えられるべきです。これではマスクと入力センテンスが少ししか変わらないように思えてしまいます。言い換えると、マスクは長さが３の例に当てはまるように見えるようになりました。最初のトークンでは、モデルは先導のトークンそのものを見るべきな一方、続きのトークンでは後ろを見返すことはできるが、前を見ることができないので、この三角形の形のマスクになっている。
では、これら全てを視野に入れて、attention 行列を示しましょう。
# attn.shape == (batch_size, num_heads, seq_len, seq_len)
batch_size = 3
num_heads = 2
attn = torch.randn(batch_size, num_heads, seq_len, seq_len)
attn.shape
torch.Size([3, 2, 3, 3])
ここで、私たちは３つのバッチがある入力を持っています。私たちのモデルは２つのヘッドしかありません。このケースでは、attention 行列を作ることを適用した時に、私たちは下の結果しか得られません。
attn = attn.masked_fill(mask == 0, float("-inf"))
attn
tensor([[[[-0.6319, -inf, -inf],
[ 0.7736, -0.4394, -inf],
[ 0.2407, 0.8301, -0.2763]],
[[ 0.4821, -inf, -inf],
[ 1.3904, -2.0258, -inf],
[ 0.3205, 1.8750, -1.0537]]],
[[[ 0.3154, -inf, -inf],
[-2.1034, -0.2958, -inf],
[ 0.4362, -0.8575, 1.8995]],
[[ 0.5619, -inf, -inf],
[-0.3208, -0.6639, -inf],
[ 0.6854, -0.9504, 0.2803]]],
[[[ 0.0928, -inf, -inf],
[ 0.3951, -0.0538, -inf],
[-0.9994, -2.0981, -0.1262]],
[[-0.9176, -inf, -inf],
[-0.3652, -0.9505, -inf],
[-1.2675, 0.0186, 0.0417]]]])
ソフトマックスを行列に適用すると以下の結果になります。各rowは足し合わせると1になります：それが加重平均の動作です。
F.softmax(attn, dim=-1)
tensor([[[[1.0000, 0.0000, 0.0000],
[0.7708, 0.2292, 0.0000],
[0.2942, 0.5304, 0.1754]],
[[1.0000, 0.0000, 0.0000],
[0.9682, 0.0318, 0.0000],
[0.1671, 0.7907, 0.0423]]],
[[1.0000, 0.0000, 0.0000],
[0.1409, 0.8591, 0.0000],
[0.1787, 0.0490, 0.7722]],
[[1.0000, 0.0000, 0.0000],
[0.5850, 0.4150, 0.0000],
[0.5372, 0.1046, 0.3582]]],
[[[1.0000, 0.0000, 0.0000],
[0.6104, 0.3896, 0.0000],
[0.2683, 0.0894, 0.6423]],
[[1.0000, 0.0000, 0.0000],
[0.6423, 0.3577, 0.0000],
[0.1202, 0.4348, 0.4450]]]])
Model
ここまで来たら、全てのことを足し合わせることができます。是非とも基本的なモデル設定を作成し、モデルを初期化してみましょう。
vocab_size = 10
max_len = 12
config = GPT1Config(vocab_size, max_len)
model = GPT(config)
これは、ただの基本的な 12レイヤー decoder ネットワークです。最近では、大きな LM (Learing Model？) は巨大で大きく、実際単一の GPU にはフィットしません。にも関わらず、私たちの GPT モデルはかなりだというのが私の意見です。
model
ではダミー入力を作って、モデルが私たちの想像するように動くか見てみましょう。まず長さがモデルのキャパシティを超えている良くない入力から入れてみましょう。
seq_len = 15
test_input = torch.randint(high=vocab_size, size=(batch_size, seq_len))
try:
model(test_input).shape
except AssertionError as e:
print(e)
sequence longer than model capacity
私たちは、正しいアサーションエラーを取得できました、モデルが処理できる最大の長さ以上のシーケンスは入らないと言っています。正しい入力を入れたらどうなるか見てみましょう。
model(test_input[:, :max_len]).shape
torch.Size([3, 12, 10])
予想していた通り、正しいアウトプット形が得られました。
もちろん、ここからモデルをトレーニングすることもできますが、それはまた別の日に Colab notebook で試すもので 私のローカルの Jupyter 環境でやるものではないでしょう。
結論
このポストでは、miniGPT を例にして、Transformer decoder がどのように動いているかみてみました。GPTモデルは全て似たようなアーキテクチャーを使っていることに注意してください、唯一の違いはモデルのサイズと訓練されるデータセットのコーパスです。明らかなように大きなモデルには大きなデータセットが必要です。
いろんなモデルが multi-head self-attention の上に作られていることが分かりました。頭のいい修正をアルゴリズムに適用し linear runtime でできるようにしたり、relative position embedding など新しい embedding の概念を使用した Reformer を含めたモデルがあります。いずれ触れます。
読むのを楽しんでくれれば嬉しいです。次回作でも会いましょう。