2024年の最初のエントリーはGPTです。
GPTモデルを自作して、OpenAIが公開している学習済みのパラメータをロード、テキスト生成までの一連の処理を実行します。
モデル
正確にはGPT2のTransformerブロックを自作します。
アーキテクチャの大部分はGPTと同じですが、以下の変更(pre-norm)が行われています。
- LayerNormはAttentionとMLPの前で適用
- 追加のLayerNormをTransformerブロックの後で適用
Transformerブロックを除くText & Position埋め込みとNext Token生成は、
picoGPTのコードを利用します(解説ブログは
GPT in 60 Lines of NumPy | Jay Mody)。
また、以下で紹介するコードはTensorflowを用いて実装しています(picoGPTの諸々のコードがTensorflowを利用していて、そのまま使いたかったため)。
GPT2の論文ではアーキテクチャの図がないので、下記はGPTのアーキテクチャ図ですが、上で書いたように、LayerNormはAttentionとMLPの前で適用します。
これをTransformerDecoderBlockクラスとして用意します(推論だけ行うのでDropoutは不要ですが)。
class TransformerDecoderBlock(tf.keras.Model):
def __init__(self, h_dim, n_heads, drop_p):
super().__init__()
self.attn = MaskedMultiSelfAttention(h_dim, n_heads, drop_p)
self.mlp = tf.keras.Sequential(
[
tf.keras.layers.Dense(units=4 * h_dim, activation="gelu"),
tf.keras.layers.Dense(units=h_dim),
tf.keras.layers.Dropout(rate=drop_p),
]
)
self.ln1 = tf.keras.layers.LayerNormalization()
self.ln2 = tf.keras.layers.LayerNormalization()
def call(self, x):
x = self.attn(self.ln1(x)) + x
x = self.mlp(self.ln2(x)) + x
return x
Masked Multi Self Attention
Attentionの計算は複雑なところはなく、以前に作った
ゼロから作るVision Transformer (JAX/Flax) - stMindとも同じです。
- 入力トークン列から、クエリq、キーk、バリューvを作成
- 複数ヘッド毎のアテンション行列とアテンションの計算
- ヘッド毎のアテンションを集約
ただし、このままだと現在のトークンが未来のトークンも参照してしまうことになるので、アテンション行列において現在のトークンと未来のトークンの関係性はなし(0)にする必要があります。未来のトークンは、アテンション行列の各行で列方向に並ぶので、下三角行列を作成してアテンション行列をマスク。この時、マスクした後でソフトマックスを計算すると正規化されなくなるので、未来のトークンは非常に小さい値にしておいて、そのあとでソフトマックスを適用します。
class MaskedMultiSelfAttention(tf.keras.layers.Layer):
def __init__(self, h_dim, n_heads, drop_p):
super(MaskedMultiSelfAttention, self).__init__()
self.n_heads = n_heads
self.c_attn = tf.keras.layers.Dense(3 * h_dim)
self.c_proj = tf.keras.layers.Dense(h_dim)
self.attn_drop = tf.keras.layers.Dropout(drop_p)
self.proj_drop = tf.keras.layers.Dropout(drop_p)
def call(self, x):
B, T, C = x.shape
N, D = self.n_heads, C // self.n_heads
mask = tf.linalg.band_part(tf.ones((T, T)), -1, 0)
mask = tf.reshape(mask, (1, 1, T, T))
qkv = self.c_attn(x)
q, k, v = tf.split(qkv, 3, axis=-1)
q = tf.reshape(q, (B, T, N, D))
k = tf.reshape(k, (B, T, N, D))
v = tf.reshape(v, (B, T, N, D))
q = tf.transpose(q, perm=[0, 2, 1, 3])
k = tf.transpose(k, perm=[0, 2, 1, 3])
v = tf.transpose(v, perm=[0, 2, 1, 3])
weights = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(
tf.cast(D, dtype=tf.float32)
)
weights += (1 - mask) * -1e9
normalized_weights = tf.nn.softmax(weights, axis=-1)
attention = self.attn_drop(tf.matmul(normalized_weights, v))
attention = tf.transpose(attention, perm=[0, 2, 1, 3])
attention = tf.reshape(attention, (B, T, C))
out = self.proj_drop(self.c_proj(attention))
return out
GPT2モデル全体
Transformerブロックができたので、Text & Position埋め込みと追加LayerNormを含めたGPT2全体を作ります。
GPT2全体コードは長くなるので、callメソッドだけ抜き出すと下のようになります。
埋め込みベクトルは、次で説明するパラメータを使って、input_idsに対する埋め込みを取り出して生成します。
def call(self, input_ids):
input_ids = tf.cast(input_ids, tf.int32)
x = tf.gather(self.wte, input_ids) + tf.gather(
self.wpe, range(input_ids.shape[1])
)
for block in self.blocks:
x = block(x)
x = self.layer_norm(x)
return tf.matmul(x, self.params["wte"].T)
OpenAIの学習済みパラメータを使用して、トークンを生成してみます。
学習済みパラメータ
パラメータには、入力と位置埋め込み、Transformerの各ブロックのパラメータがあり、picoGPTで辞書形式に変換されているものを使用します。
- blocks : Transformerブロックのパラメータ
- ln_f : 追加のLayerNormのパラメータ
- wpe : 位置埋め込みベクトル
- wte : トークンの埋め込みベクトル
また、blocksは124Mのモデルの場合は12個の要素があり、それぞれが以下の項目を含んでいます。(768は次元数)
- attn : アテンションブロック(以下のbはバイアス項、wは重み)
- c_attn : {"b": [2304], "w": [768, 2304]}
- c_proj : {"b": [768], "w": [768, 768]}
- ln1 : {"b": [768], "g": [768]}, Attentionの前に適用するLayerNorm。bはbeta、gはgamma
- ln2 : {"b": [768], "g": [768]}, MLPの前に適用するLayerNorm
- mlp : MLPブロック
- c_fc : {"b": [3072], "w": [768, 3072]}
- c_proj : {"b": [768], "w": [3072, 768]}
Tensorflowにおけるパラメータの割り当て
tf.keras.layers.Layerのset_weigthsを使います。この関数は、numpy の配列からパラメータ値を設定します。
例えば、c_attnの場合だと、これはDense層なのでwとbの順番でset_weightsに指定します。
block.layers[0].c_attn.set_weights(
[
self.params["blocks"][layer_idx]["attn"]["c_attn"]["w"],
self.params["blocks"][layer_idx]["attn"]["c_attn"]["b"],
]
)
ln_fはLayerNormなので、gammaとbetaでset_weightsに指定します。
self.layer_norm.set_weights(
[self.params["ln_f"]["g"], self.params["ln_f"]["b"]]
)
生成結果
GPT2モデルの作成、重みパラメータの設定が出来たので、
picoGPTと同じプロンプトで実験してみます。
$ python tf/gpt_tf.py --prompt "Alan Turing theorized that computers would one day become" --n_tokens_to_generate 8
...
Input text:
Alan Turing theorized that computers would one day become
Generated:
the most powerful machines on the planet.
同じ結果が生成されました。生成には、M1 Macで2秒弱くらいかかりました。
別のプロンプトも試してみます。
python tf/gpt_tf.py --prompt "Imagination is more important" --n_tokens_to_generate 6
...
Input text:
Imagination is more important
Generated:
than any other skill.
文章としては問題ないものが生成されたように思います。
まとめ
以前のViTと今回のGPTでAttentionの自作をしましたが、処理自体はそれほど複雑ではないので、
Transformerブロックを実装するのは、慣れれば難しくはないように感じました。また、モデルを実装して学習するのはHW制約などもあって大変なことが多いですが、推論であれば公開されているパラメータを使うことで、比較的試してみやすいのではと思います。コアとなるTransformer、Attentionを自作することで、Transformer系の論文の数式やコードの読解力が上がったように感じるので、興味のある方は自作にトライしてみることをオススメします。
参考文献とコード
GPT in 60 Lines of NumPy | Jay Mody
GitHub - satojkovic/gpt-tf-pytorch-jax: GPT from scratch (tensorflow / pytorch / jax)