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を利用していて、そのまま使いたかったため）。

Transformerブロック

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

        # Create lower triangle mask
        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)
        )

        # Apply mask
        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):
    # Text and Position Embedding
    input_ids = tf.cast(input_ids, tf.int32)
    x = tf.gather(self.wte, input_ids) + tf.gather(
        self.wpe, range(input_ids.shape[1])
    )
    # Transformer Block (Decoder only)
    for block in self.blocks:
        x = block(x)
    # Additional LayerNorm
    x = self.layer_norm(x)
    # Linear
    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系の論文の数式やコードの読解力が上がったように感じるので、興味のある方は自作にトライしてみることをオススメします。