メインコンテンツまでスキップ

OpenAIのAPIを使用しています。正確性の保証はできませんので、ご注意ください。

60 行の NumPy で学ぶ GPT

備考

この記事は、「GPT in 60 Lines of NumPy」を日本語に翻訳したものです。翻訳を許可していただいた Jay Mody 氏に感謝します。この記事は CC ライセンスに含まれません。

イントロダクション

この記事では、わずか60 行のnumpyで GPT をゼロから実装します。その後、OpenAI が公開したトレーニング済みの GPT-2 モデルの重みを読み込み、テキストを生成します。

注意:

  • この記事では、Python、NumPy、およびニューラルネットワークの基本的なトレーニング経験についての理解を前提としています

  • この実装は、完全であることを保ちつつ、できるだけシンプルにするために、意図的に多くの機能が欠けています。目標は、教育ツールとして GPT のシンプルかつ完全な技術入門を提供することです

  • GPT アーキテクチャは、現在の LLM(Large Language Models、大規模言語モデル)を形成する要素のほんの一部に過ぎません 1

  • この記事のすべてのコードは、github.com/jaymody/picoGPTで確認することができます

  • Hacker News のスレッド

  • 中国語版

編集(2023 年 2 月 9 日): 「次はなんですか?」セクションを追加し、イントロにいくつかのノートを追加しました。

編集(2023 年 2 月 28 日): 「次はなんですか?」にいくつかの追加セクションを追加しました。

GPT とはなんですか?

GPT はGenerative Pre-trained Transformer の略です。これは、Transformerに基づく一種のニューラルネットワークアーキテクチャです。Jay Alammar 氏のHow GPT3 Worksは、GPT についての優れた高レベルの紹介ですが、以下に要約します。

  • Generative: GPT はテキストを生成します
  • Pre-trained: GPT は、書籍やインターネットなどの大量のテキストでトレーニングされます
  • Transformer: GPT はデコーダーのみのtransformerニューラルネットワークです

OpenAI の GPT-3Google の LaMDA、およびCohere の Command XLargeなどの大規模言語モデル(LLM)は、本質的には GPT です。それらが特別なのは、**1)**非常に大きい(数十億のパラメータ)ことと、**2)**多くのデータ(数百ギガバイトのテキスト)で訓練されていることです。

基本的に、GPT はプロンプトを与えられた場合にテキストを生成します。この非常にシンプルな API(入力=テキスト、出力=テキスト)でも、訓練が十分に行われた GPT は、あなたのメールを書く本を要約するInstagram のキャプションのアイデアを提供する5 歳の子供にブラックホールを説明するSQL でコードを書く遺言書を作成するなど、素晴らしいことができます。

以上が GPT の概要とその機能の高レベルな概要です。さらに詳細に掘り下げてみましょう。

入力 / 出力

GPT の関数シグネチャはおおよそ以下のようになります:

def gpt(inputs: list[int]) -> list[list[float]]:
# inputs は [n_seq] の形状を持つ
# 出力は [n_seq, n_vocab] の形状を持つ
output = # beep boop neural networkの魔法
return output

入力

入力は、テキストを表す整数のシーケンスであり、テキスト内のトークンにマップされます:

# 整数はテキスト内のトークンを表します。例えば:
# テキスト = "not all heroes wear capes":
# トークン = "not" "all" "heroes" "wear" "capes"
inputs = [1, 0, 2, 4, 6]

トークンはテキストのサブピースであり、ある種のトークナイザーを使用して生成されます。語彙(ボキャブラリー)を使用してトークンを整数にマッピングすることができます:

# トークンの語彙内でのインデックスは、そのトークンの整数IDを表します
# 例えば、"heroes"の整数IDは2です。なぜならvocab[2] = "heroes"だからです
vocab = ["all", "not", "heroes", "the", "wear", ".", "capes"]

# 空白でトークナイズする架空のトークナイザー
tokenizer = WhitespaceTokenizer(vocab)

# encode()メソッドは文字列をlist[int]に変換します
ids = tokenizer.encode("not all heroes wear") # ids = [1, 0, 2, 4]

# 語彙マッピングを通じて実際のトークンを確認できます
tokens = [tokenizer.vocab[i] for i in ids] # tokens = ["not", "all", "heroes", "wear"]

# decode()メソッドはlist[int]を文字列に戻します
text = tokenizer.decode(ids) # text = "not all heroes wear"

要約すると:

  • 文字列があります
  • トークナイザーを使用して、それを「トークン」と呼ばれる小さな部分に分解します
  • これらのトークンを整数にマッピングするために語彙(ボキャブラリー)を使用します

実際には、単純に空白で分割するよりも、Byte-Pair EncodingWordPieceのような、より高度なトークナイズ方法を使用しますが、原理は同じです:

  1. 文字列トークンを整数インデックスにマッピングするvocabがあります
  2. str -> list[int]に変換するencodeメソッドがあります
  3. list[int] -> str に変換するdecodeメソッドがあります 2

出力

出力は 2 次元配列であり、output[i][j]はモデルが vocab[j]のトークンが次のトークン inputs[i+1]であると予測した確率です。例えば:

vocab = ["all", "not", "heroes", "the", "wear", ".", "capes"]
inputs = [1, 0, 2, 4] # "not" "all" "heroes" "wear"
output = gpt(inputs)
# ["all", "not", "heroes", "the", "wear", ".", "capes"]
# output[0] = [0.75 0.1 0.0 0.15 0.0 0.0 0.0 ]
# "not"のみが与えられた場合、モデルは最も高い確率で単語"all"を予測します

# ["all", "not", "heroes", "the", "wear", ".", "capes"]
# output[1] = [0.0 0.0 0.8 0.1 0.0 0.0 0.1 ]
# シーケンス["not", "all"]が与えられた場合、モデルは最も高い確率で単語"heroes"を予測します

# ["all", "not", "heroes", "the", "wear", ".", "capes"]
# output[-1] = [0.0 0.0 0.0 0.1 0.0 0.05 0.85 ]
# 全シーケンス["not", "all", "heroes", "wear"]が与えられた場合、モデルは最も高い確率で単語"capes"を予測します

シーケンス全体に対する次のトークンの予測を得るためには、output[-1]の中で最も高い確率を持つトークンを単純に取ります:

vocab = ["all", "not", "heroes", "the", "wear", ".", "capes"]
inputs = [1, 0, 2, 4] # "not" "all" "heroes" "wear"
output = gpt(inputs)
next_token_id = np.argmax(output[-1]) # next_token_id = 6
next_token = vocab[next_token_id] # next_token = "capes"

最も高い確率を持つトークンを予測として取ることを、貪欲デコーディングまたは貪欲サンプリングと呼びます。

シーケンスで次の論理的な単語を予測するタスクは、言語モデリングと呼ばれます。そのため、GPT を言語モデルと呼ぶことができます。

単一の単語を生成するのは素晴らしいことですが、文全体や段落などはどうなるでしょうか?

テキストの生成

自己回帰

モデルから次のトークン予測を繰り返し取得することで、完全な文を生成できます。各反復で、予測されたトークンを入力に追加して戻します。

def generate(inputs, n_tokens_to_generate):
for _ in range(n_tokens_to_generate): # 自己回帰的デコードループ
output = gpt(inputs) # モデルのフォワードパス
next_id = np.argmax(output[-1]) # 貪欲サンプリング
inputs.append(int(next_id)) # 予測を入力に追加
return inputs[len(inputs) - n_tokens_to_generate :] # 生成されたIDのみを返す

input_ids = [1, 0] # "not" "all"
output_ids = generate(input_ids, 3) # output_ids = [2, 4, 6]
output_tokens = [vocab[i] for i in output_ids] # "heroes" "wear" "capes"

この将来の値を予測し(回帰)、それを入力に追加する(自己)、というプロセスが、GPT を自己回帰と表現する理由です。

サンプリング

貪欲ではなく、確率分布からサンプリングすることで、生成にいくらかの確率的要素(ランダム性)を導入することができます:

inputs = [1, 0, 2, 4] # "not" "all" "heroes" "wear"
output = gpt(inputs)
np.random.choice(np.arange(vocab_size), p=output[-1]) # capes
np.random.choice(np.arange(vocab_size), p=output[-1]) # hats
np.random.choice(np.arange(vocab_size), p=output[-1]) # capes
np.random.choice(np.arange(vocab_size), p=output[-1]) # capes
np.random.choice(np.arange(vocab_size), p=output[-1]) # pants

これにより、同じ入力に対して異なる文を生成することができます。サンプリング前に分布を変更する top-ktop-ptemperature などの技術と組み合わせることで、出力の質は大幅に向上します。これらの技術はまた、異なる生成行動を試すために遊べるいくつかのハイパーパラメーターを導入します(たとえば、temperature を上げると、モデルはよりリスクを冒し、より「創造的」になります)。

トレーニング

GPT のトレーニングは、他のニューラルネットワークと同様に、ある損失関数に対する勾配降下法を使用して行います。GPT の場合、言語モデリングタスクにおけるクロスエントロピー損失を取得します。

def lm_loss(inputs: list[int], params) -> float:
# ラベルyは単に入力を1つ左にシフトしたものです。
#
# inputs = [not, all, heros, wear, capes]
# x = [not, all, heroes, wear]
# y = [all, heroes, wear, capes]
#
# もちろん、inputs[-1]に対するラベルはありませんので、xから除外します。
#
# そのため、N個の入力に対して、N - 1個の言語モデリング例のペアがあります。
x, y = inputs[:-1], inputs[1:]

# フォワードパス
# 各位置における予測された次のトークンの確率分布
output = gpt(x, params)

# クロスエントロピー損失
# 全てのN-1例についての平均を取ります。
loss = np.mean(-np.log(output[y]))

return loss

def train(texts: list[list[str]], params) -> float:
for text in texts:
inputs = tokenizer.encode(text)
loss = lm_loss(inputs, params)
gradients = compute_gradients_via_backpropagation(loss, params)
params = gradient_descent_update_step(gradients, params)
return params

これは大幅に単純化されたトレーニング設定ですが、要点を示しています。gpt関数シグネチャにparamsを追加したことに注目してください(簡単にするために前のセクションではこれを省略しました)。トレーニングループの各イテレーション中に:

  • 指定された入力テキストの例に対する言語モデリングの損失を計算します
  • 損失は勾配を決定し、逆伝播を通じて勾配を計算します
  • 損失が最小化されるように、勾配を使用してモデルパラメータを更新します(勾配降下法)

明示的にラベル付けされたデータは使用しないことに注意してください。代わりに、生のテキスト自体から入力/ラベルのペアを生成することができます。これを自己教師あり学習と呼びます。

自己教師あり学習を使用することで、訓練データを大規模にスケールアップできます。可能な限り多くの生テキストを手に入れてモデルに投入するだけです。例えば、GPT-3 はインターネットと本からの 3000 億トークンのテキストでトレーニングされました:

Table 2.2 from GPT-3 paper

もちろん、これらすべてのデータから学習するには、十分に大きなモデルが必要です。そのため、GPT-3 には 1750 億のパラメータがあり、訓練にはコンピュートコストとして 100 万ドルから 1000 万ドルの間がかかったと思われます。3

この自己監督学習ステップは、事前学習と呼ばれます。なぜなら、事前に学習されたモデルの重みを再利用して、ツイートが有害かどうかを分類するなどの下流タスクでモデルをさらに訓練することができるからです。事前学習モデルは、時には基盤モデルとも呼ばれます。

モデルを下流タスクでトレーニングすることはファインチューニングと呼ばれます。なぜなら、モデルの重みは既に言語を理解するために事前トレーニングされており、特定のタスクに対して微調整されるだけだからです。

「一般的なタスクでの事前トレーニング+特定のタスクでのファインチューニング」という戦略は、転移学習と呼ばれています。

プロンプト

原則として、元のGPTの論文は、Transformer モデルの事前学習の利点についてのみ述べていました。この論文では、117M の GPT を事前学習し、ラベル付きデータセットで微調整することで、さまざまなNLP(自然言語処理)タスクで最先端のパフォーマンスを達成できることが示されています。

GPT-2GPT-3の論文が発表されるまで、私たちはデータ量とパラメータ数が十分にある GPT モデルが、微調整なしで任意のタスクを単独で実行できることに気付きませんでした。モデルにプロンプトを与え、自己回帰言語モデリングを行うだけで、適切な応答が得られます。これはインコンテキスト学習 (in-context learning)と呼ばれ、モデルがタスクを実行するためにプロンプトの文脈のみを使用していることを意味します。インコンテキスト学習は、ゼロショット、ワンショット、またはフューショットのいずれかで行うことができます:

プロンプトに基づいてテキストを生成することは、条件付き生成とも呼ばれます。なぜなら、モデルがある入力に条件付けられた出力を生成しているからです。

GPT は NLP タスクに限定されません。モデルを好きなものに条件付けることができます。たとえば、会話履歴に基づいて GPT をチャットボット(つまり、ChatGPT)に変えることができます。また、プロンプトをいくつかの説明と共に前置きすることで、チャットボットを特定の方法で動作させることもできます(つまり、「あなたはチャットボットです。礼儀正しく、完全な文で話し、有害なことを言わないでくださいなど...」)。このようにモデルを条件付けることで、チャットボットに個性を与えることさえできます。ただし、これは堅牢ではなく、モデルを「脱獄」して不正な動作させることもできます。

それでは、実際の実装に移りましょう。

セットアップ

このチュートリアルのためのリポジトリをクローンします:

git clone https://github.com/jaymody/picoGPT
cd picoGPT

次に、依存関係をインストールしましょう:

pip install -r requirements.txt
ヒント

注意:このコードはPython 3.9.10でテストされました。

各ファイルの簡単な内訳は次のとおりです。

  • encoder.py は、OpenAI の BPE トークナイザーのコードを含んでおり、彼らのgpt-2 リポジトリから直接取得されています。
  • utils.py は、GPT-2 モデルの重み、トークナイザー、およびハイパーパラメーターをダウンロードしてロードするためのコードが含まれています。
  • gpt2.py は、実際の GPT モデルと生成コードを含んでおり、Python スクリプトとして実行できます。
  • gpt2_pico.py は、gpt2.pyと同じですが、さらに少ない行数で記述されています。なぜなら、なぜでしょうか。

gpt2.pyをゼロから再実装することになるので、それを削除して空のファイルとして再作成しましょう:

$ rm gpt2.py
$ touch gpt2.py

出発点として、次のコードをgpt2.pyに貼り付けてください:

import numpy as np


def gpt2(inputs, wte, wpe, blocks, ln_f, n_head):
pass # TODO: これを実装する


def generate(inputs, params, n_head, n_tokens_to_generate):
from tqdm import tqdm

for _ in tqdm(range(n_tokens_to_generate), "generating"): # 自己回帰デコードループ
logits = gpt2(inputs, **params, n_head=n_head) # モデルのフォワードパス
next_id = np.argmax(logits[-1]) # 貪欲サンプリング
inputs.append(int(next_id)) # 予測を入力に追加

return inputs[len(inputs) - n_tokens_to_generate :] # 生成されたidのみを返す


def main(prompt: str, n_tokens_to_generate: int = 40, model_size: str = "124M", models_dir: str = "models"):
from utils import load_encoder_hparams_and_params

# 公開されているOpenAI GPT-2ファイルからエンコーダー、hparams、およびparamsを読み込む
encoder, hparams, params = load_encoder_hparams_and_params(model_size, models_dir)

# BPEトークナイザーを使用して入力文字列をエンコードする
input_ids = encoder.encode(prompt)

# モデルの最大シーケンス長を超えないようにする
assert len(input_ids) + n_tokens_to_generate < hparams["n_ctx"]

# 出力idを生成する
output_ids = generate(input_ids, params, hparams["n_head"], n_tokens_to_generate)

# idを文字列にデコードする
output_text = encoder.decode(output_ids)

return output_text


if __name__ == "__main__":
import fire

fire.Fire(main)

4 つの各セクションを分解して説明します:

  1. gpt2 関数は、実際に実装する GPT のコードです。関数のシグネチャには、inputs に加えていくつかの追加の要素が含まれていることに気付くでしょう:
    • wtewpeblocksln_f は、モデルのパラメータです
    • n_head は、順方向のパス中に必要なハイパーパラメータです
  2. generate 関数は、前に見た自己回帰デコーディングアルゴリズムです。シンプルさのために、貪欲なサンプリングを使用しています。tqdm は、トークンを一つずつ生成する過程を視覚化するための進捗バーです
  3. main 関数は以下の処理を行います:
  4. トークナイザー(encoder)、モデルの重み(params)、ハイパーパラメータ(hparams)をロードします
  5. トークナイザーを使用して入力プロンプトをトークン ID にエンコードします
  6. generate 関数を呼び出します
  7. 出力の ID を文字列にデコードします
  8. fire.Fire(main) は、ファイルを CLI アプリケーションに変換するだけであり、最終的には python gpt2.py "some prompt here" のようにコードを実行できるようにします

encoderhparamsparams を詳しく見るためには、ノートブックや対話型の Python セッションで次のコードを実行してください:

from utils import load_encoder_hparams_and_params
encoder, hparams, params = load_encoder_hparams_and_params("124M", "models")

これにより、必要なモデルファイルとトークナイザーファイルが models/124Mダウンロードされencoderhparamsparamsコードにロードされます

エンコーダー

encoderは GPT-2 で使用される BPE トークナイザーです:

>>> ids = encoder.encode("Not all heroes wear capes.")
>>> ids
[3673, 477, 10281, 5806, 1451, 274, 13]

>>> encoder.decode(ids)
"Not all heroes wear capes."

トークナイザーの語彙(encoder.decoderに保存されている)を使用して、実際のトークンがどのように見えるかを確認できます:

>>> [encoder.decoder[i] for i in ids]
['Not', 'Ġall', 'Ġheroes', 'Ġwear', 'Ġcap', 'es', '.']

トークンは時には単語であり(例:Not)、時にはその前にスペースがある単語であることもあります(例:Ġall、Ġ はスペースを表します)、時には単語の一部であることもあります(例:capes は Ġcap と es に分割されます)、そして時には句読点であることもあります(例:.)。

BPE の良い点の一つは、任意の文字列をエンコードできることです。語彙にない何かに出会った場合、それを理解できるサブストリングに単純に分解します:

>>> [encoder.decoder[i] for i in encoder.encode("zjqfl")]
['z', 'j', 'q', 'fl']

また、語彙のサイズも確認できます:

>>> len(encoder.decoder)
50257

語彙や、文字列がどのように分解されるかを決定するバイトペアマージは、トークナイザーをトレーニングすることで得られます。トークナイザーをロードするとき、私たちは既にトレーニングされた語彙とバイトペアマージを、load_encoder_hparams_and_paramsを実行したときにモデルファイルと共にダウンロードされたいくつかのファイルからロードしています。

語彙についてはmodels/124M/encoder.jsonを、バイトペアマージについてはmodels/124M/vocab.bpeを参照してください。

ハイパーパラメータ

hparamsは、モデルのハイパーパラメータを含む辞書です:

>>> hparams
{
"n_vocab": 50257, # 語彙のトークン数
"n_ctx": 1024, # 入力の最大可能なシーケンス長
"n_embd": 768, # 埋め込み次元(ネットワークの「幅」を決定する)
"n_head": 12, # アテンションヘッドの数(n_embdはn_headで割り切れる必要がある)
"n_layer": 12 # レイヤーの数(ネットワークの「深さ」を決定する)
}

コードのコメントでこれらのシンボルを使用して、ものの基本的な形状を示します。また、n_seqを入力シーケンスの長さを示すために使用します(つまり、n_seq = len(inputs))。

パラメーター

paramsは、モデルの学習済みの重みを保持するネストされた JSON 辞書です。JSON の葉ノードは NumPy 配列です。paramsを出力するときに、配列をその形状に置き換えると次のようになります:

>>> import numpy as np
>>> def shape_tree(d):
>>> if isinstance(d, np.ndarray):
>>> return list(d.shape)
>>> elif isinstance(d, list):
>>> return [shape_tree(v) for v in d]
>>> elif isinstance(d, dict):
>>> return {k: shape_tree(v) for k, v in d.items()}
>>> else:
>>> ValueError("uh oh")
>>>
>>> print(shape_tree(params))
{
"wpe": [1024, 768],
"wte": [50257, 768],
"ln_f": {"b": [768], "g": [768]},
"blocks": [
{
"attn": {
"c_attn": {"b": [2304], "w": [768, 2304]},
"c_proj": {"b": [768], "w": [768, 768]},
},
"ln_1": {"b": [768], "g": [768]},
"ln_2": {"b": [768], "g": [768]},
"mlp": {
"c_fc": {"b": [3072], "w": [768, 3072]},
"c_proj": {"b": [768], "w": [3072, 768]},
},
},
... # n_layersの繰り返し
]
}

これらは、元の OpenAI TensorFlow チェックポイントからロードされます:

>>> import tensorflow as tf
>>> tf_ckpt_path = tf.train.latest_checkpoint("models/124M")
>>> for name, _ in tf.train.list_variables(tf_ckpt_path):
>>> arr = tf.train.load_variable(tf_ckpt_path, name).squeeze()
>>> print(f"{name}: {arr.shape}")
model/h0/attn/c_attn/b: (2304,)
model/h0/attn/c_attn/w: (768, 2304)
model/h0/attn/c_proj/b: (768,)
model/h0/attn/c_proj/w: (768, 768)
model/h0/ln_1/b: (768,)
model/h0/ln_1/g: (768,)
model/h0/ln_2/b: (768,)
model/h0/ln_2/g: (768,)
model/h0/mlp/c_fc/b: (3072,)
model/h0/mlp/c_fc/w: (768, 3072)
model/h0/mlp/c_proj/b: (768,)
model/h0/mlp/c_proj/w: (3072, 768)
model/h1/attn/c_attn/b: (2304,)
model/h1/attn/c_attn/w: (768, 2304)
...
model/h9/mlp/c_proj/b: (768,)
model/h9/mlp/c_proj/w: (3072, 768)
model/ln_f/b: (768,)
model/ln_f/g: (768,)
model/wpe: (1024, 768)
model/wte: (50257, 768)

次のコードは、上記の TensorFlow 変数をparams辞書に変換します。

参考までに、paramsの形状を示すために、数字をそれらが表すhparamsで置き換えたものを以下に示します:

{
"wpe": [n_ctx, n_embd],
"wte": [n_vocab, n_embd],
"ln_f": {"b": [n_embd], "g": [n_embd]},
"blocks": [
{
"attn": {
"c_attn": {"b": [3*n_embd], "w": [n_embd, 3*n_embd]},
"c_proj": {"b": [n_embd], "w": [n_embd, n_embd]},
},
"ln_1": {"b": [n_embd], "g": [n_embd]},
"ln_2": {"b": [n_embd], "g": [n_embd]},
"mlp": {
"c_fc": {"b": [4*n_embd], "w": [n_embd, 4*n_embd]},
"c_proj": {"b": [n_embd], "w": [4*n_embd, n_embd]},
},
},
... # repeat for n_layers
]
}

この辞書を参照して、GPT を実装する際に重みの形状を確認するために、後で戻ってくることがおそらく必要になるでしょう。コード内の変数名は、この辞書のキーと一致させるために一貫性を持たせます。

基本レイヤー

実際の GPT アーキテクチャに入る前に、GPT に固有ではない、いくつかの基本的なニューラルネットワークのレイヤーを実装しましょう。

GELU

GPT-2 での非線形(活性化関数)の選択肢は、ReLU の代替としてGELU(Gaussian Error Linear Units)です:

以下の関数で近似されます:

def gelu(x):
return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))

ReLU と同様に、GELU は入力に対して要素ごとに操作を行います:

>>> gelu(np.array([[1, 2], [-2, 0.5]]))
array([[ 0.84119, 1.9546 ],
[-0.0454 , 0.34571]])

ソフトマックス関数

古き良きソフトマックス関数です:

softmax(x)i=exijexj\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}

数値の安定性のためにmax(x)のトリックを使用します。

ソフトマックス関数は、実数の集合(- \inftyから\inftyの間)を確率(0 から 1 の間で、合計が 1 になる数)に変換するために使用されます。入力の最後の軸にsoftmaxを適用します。

>>> x = softmax(np.array([[2, 100], [-5, 0]]))
>>> x
array([[0.00034, 0.99966],
[0.26894, 0.73106]])
>>> x.sum(axis=-1)
array([1., 1.])

レイヤー正規化

レイヤー正規化は、平均を 0、分散を 1 に標準化します:

LayerNorm(x)=γxμσ2+β\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2}} + \beta

ここで、μ\muxxの平均、σ2\sigma^2xxの分散であり、γ\gammaβ\betaは学習可能なパラメータです。

def layer_norm(x, g, b, eps: float = 1e-5):
mean = np.mean(x, axis=-1, keepdims=True)
variance = np.var(x, axis=-1, keepdims=True)
x = (x - mean) / np.sqrt(variance + eps) # 最後の軸において平均=0、分散=1になるようにxを正規化
return g * x + b # gamma/betaパラメータでスケールとオフセットを行う

レイヤーノーマリゼーションは、各層の入力が常に一貫した範囲内にあることを保証し、これによりトレーニングプロセスの加速と安定化が期待されます。バッチノーマリゼーションと同様に、正規化された出力は、学習可能な 2 つのベクトル gamma と beta でスケールされオフセットされます。分母の小さな epsilon 項はゼロ除算エラーを避けるために使用されます。

Transformer ではバッチノーマリゼーションの代わりにレイヤーノーマリゼーションが使用されます。そのさまざまな理由については、この素晴らしいブログ投稿で概説されています。

入力の最後の軸に対してレイヤーノーマリゼーションを適用します。

>>> x = np.array([[2, 2, 3], [-5, 0, 1]])
>>> x = layer_norm(x, g=np.ones(x.shape[-1]), b=np.zeros(x.shape[-1]))
>>> x
array([[-0.70709, -0.70709, 1.41418],
[-1.397 , 0.508 , 0.889 ]])
>>> x.var(axis=-1)
array([0.99996, 1. ]) # 浮動小数点のやりとり
>>> x.mean(axis=-1)
array([-0., -0.])

線形

標準的な行列乗算 + バイアス:

def linear(x, w, b):  # [m, in], [in, out], [out] -> [m, out]
return x @ w + b

線形層は、しばしば投影と呼ばれます(それらが一つのベクトル空間から別のベクトル空間へ投影しているため)。

>>> x = np.random.normal(size=(64, 784)) # 入力次元 = 784, バッチ/シーケンス次元 = 64
>>> w = np.random.normal(size=(784, 10)) # 出力次元 = 10
>>> b = np.random.normal(size=(10,))
>>> x.shape # 線形投影前の形状
(64, 784)
>>> linear(x, w, b).shape # 線形投影後の形状
(64, 10)

GPT アーキテクチャ


GPT アーキテクチャは、Transformerに従います:

Attention is All You Needからの図1

しかし、デコーダスタックのみを使用しています(図の右側):

GPTアーキテクチャ

注記:エンコーダを排除したため、中央の「クロスアテンション」層も削除されています。

高いレベルで、GPT アーキテクチャは 3 つのセクションを持っています:

  • テキスト + 位置 エンベッディング
  • Transformer デコーダスタック
  • 語彙への投影 ステップ

コードでは、このようになります:

def gpt2(inputs, wte, wpe, blocks, ln_f, n_head):  # [n_seq] -> [n_seq, n_vocab]
# トークンと位置埋め込みの追加
x = wte[inputs] + wpe[range(len(inputs))] # [n_seq] -> [n_seq, n_embd]

# n_layer Transformerブロックを通じてのフォワードパス
for block in blocks:
x = transformer_block(x, **block, n_head=n_head) # [n_seq, n_embd] -> [n_seq, n_embd]

# 語彙への射影
x = layer_norm(x, **ln_f) # [n_seq, n_embd] -> [n_seq, n_embd]
return x @ wte.T # [n_seq, n_embd] -> [n_seq, n_vocab]

この 3 つのセクションをそれぞれ詳しく説明します。

埋め込み

トークンの埋め込み

トークン ID 自体は、ニューラルネットワークにとって非常に良い表現ではありません。一つには、トークン ID の相対的な大きさが誤った情報を伝えてしまうことです(例えば、私たちの語彙でApple = 5Table = 10の場合、2 * Table = Appleを暗示しています)。また、単一の数値は、ニューラルネットワークが扱うには次元性が十分ではありません。

これらの制限に対処するため、単語ベクトルを利用しますが、特に学習された埋め込み行列を介して行います:

wte[inputs] # [n_seq] -> [n_seq, n_embd]

wte[n_vocab, n_embd]行列であることを思い出してください。これはルックアップテーブルとして機能し、行列のii番目の行は私たちの語彙のii番目のトークンに対する学習されたベクトルに対応します。wte[inputs]整数配列インデックスを使用して、入力の各トークンに対応するベクトルを取得します。

wteはネットワークの他のパラメータと同様に学習されます。つまり、トレーニングの開始時にランダムに初期化され、その後勾配降下法を通じて更新されます。

位置埋め込み

Transformer アーキテクチャの特徴の一つは、位置を考慮しないことです。つまり、入力をランダムにシャッフルし、その後出力を適切にアンシャッフルした場合、入力を一切シャッフルしなかった場合と同じ出力になります(入力の順序は出力に影響を与えません)。

もちろん、言語の重要な部分は単語の順序です(当然ですが)、したがって、入力に位置情報をエンコードする方法が必要です。これには、別の学習された埋め込み行列を使用できます:

wpe[range(len(inputs))] # [n_seq] -> [n_seq, n_embd]

wpe[n_ctx, n_embd]行列であることを思い出してください。行列のii番目の行には、入力のii番目の位置に関する情報をエンコードするベクトルが含まれています。wteと同様に、この行列は勾配降下中に学習されます。

注意点として、これはモデルを最大シーケンス長n_ctxに制限します4。つまり、len(inputs) <= n_ctxが成立しなければなりません。

組み合わせ

トークンと位置の埋め込みを加算することで、トークンと位置情報の両方をエンコードする組み合わせ埋め込みを得ることができます。

# トークン + 位置埋め込み
x = wte[inputs] + wpe[range(len(inputs))] # [n_seq] -> [n_seq, n_embd]

# x[i] は i 番目の単語の単語埋め込みと i 番目の位置の位置埋め込みを表します

デコーダスタック

ここで全ての魔法が起こり、「ディープ」がディープラーニングに入ってきます。埋め込みをn_layerの Transformer デコーダブロックのスタックを通して渡します。

# n_layer Transformerブロックを通じた順伝播
for block in blocks:
x = transformer_block(x, **block, n_head=n_head) # [n_seq, n_embd] -> [n_seq, n_embd]

より多くの層を積み重ねることで、ネットワークがどれだけ「深い」かを制御することができます。例えば、GPT-3 は驚異的な 96 層を持っています。一方で、より大きなn_embd値を選択することで、ネットワークがどれだけ「幅広い」かを制御することができます(例えば、GPT-3 は埋め込みサイズに 12288 を使用しています)。

語彙への投影

最終ステップでは、最後の Transformer ブロックの出力を、語彙上の確率分布へ投影します:

# 語彙への投影
x = layer_norm(x, **ln_f) # [n_seq, n_embd] -> [n_seq, n_embd]
return x @ wte.T # [n_seq, n_embd] -> [n_seq, n_vocab]

ここで注目すべきいくつかのことがあります:

  1. x最終層の正規化レイヤーを通してから、語彙への投影を行います。これは GPT-2 アーキテクチャに特有のものであり、元の GPT および Transformer の論文には存在しません。
  2. 投影のために埋め込み行列wte再利用しています。他の GPT 実装では、投影用に別の学習済み重み行列を使用するかもしれませんが、埋め込み行列を共有することにはいくつかの利点があります:
    • パラメータを節約できます(GPT-3 のスケールでは、これは無視できるかもしれませんが)。
    • この行列が単語へのマッピングと単語からのマッピングの両方を担当するため、理論的には、2 つの別々の行列を持つ場合と比較して、より豊かな表現を学習する可能性があります。
  3. 最後にsoftmax適用しないため、出力は 0 から 1 の間の確率ではなくロジットになります。これにはいくつかの理由があります:
    • softmax単調関数であるため、貪欲サンプリングにおいてnp.argmax(logits)np.argmax(softmax(logits))と等価であり、softmaxは冗長です。
    • softmaxは不可逆であり、logitsから確率にはsoftmaxを適用して変換できますが、確率からlogitsに戻すことはできないため、最大限の柔軟性を得るためにlogitsを出力します。
    • 数値的安定性のため(例えば、交差エントロピー損失を計算する際に、log(softmax(logits))log_softmax(logits)と比較して数値的に不安定です)。

語彙への投影ステップは、言語モデリングヘッドとも呼ばれることがあります。"ヘッド"とは何を意味するのでしょうか?GPT が事前トレーニングされた後、言語モデリングヘッドを何らかの分類タスクのための分類ヘッドなど、他の種類の投影に交換することができます。したがって、モデルはヒドラのように、複数のヘッドを持つことができます。

これが高いレベルでの GPT アーキテクチャですが、実際にデコーダブロックが何をしているのかもう少し深く掘り下げてみましょう。

デコーダブロック

Transformer デコーダブロックは、2 つのサブレイヤーで構成されます:

  1. マルチヘッド因果的自己注意
  2. 位置ごとのフィードフォワードニューラルネットワーク
def transformer_block(x, mlp, attn, ln_1, ln_2, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
# マルチヘッド因果的自己注意
x = x + mha(layer_norm(x, **ln_1), **attn, n_head=n_head) # [n_seq, n_embd] -> [n_seq, n_embd]

# 位置ごとのフィードフォワードネットワーク
x = x + ffn(layer_norm(x, **ln_2), **mlp) # [n_seq, n_embd] -> [n_seq, n_embd]

return x

各サブレイヤーは、入力にレイヤー正規化を利用し、残差接続(つまり、サブレイヤーの入力をサブレイヤーの出力に加算する)を使用します。

注意すべき点:

  1. マルチヘッド因果自己注意は、入力間の通信を容易にします。ネットワーク内の他のどこでも、モデルが入力同士を「見る」ことを許可しているわけではありません。エンベッディング、位置ごとのフィードフォワードネットワーク、レイヤーノーム、および語彙への投影はすべて、入力に対して位置ごとに操作します。入力間の関係のモデリングは、注意にのみ任されています。
  2. 位置ごとのフィードフォワードニューラルネットワークは、ただの通常の 2 層全結合ニューラルネットワークです。これにより、学習を容易にするための多くの学習可能なパラメータがモデルに追加されます。
  3. 元の Transformer 論文では、レイヤーノームは出力layer_norm(x + sublayer(x))に配置されますが、GPT-2 に合わせて入力にレイヤーノームを配置しますx + sublayer(layer_norm(x))。これはプレノームと呼ばれ、Transformer のパフォーマンスを向上させる上で重要であることが示されています
  4. 残差接続ResNetによって普及)はいくつかの異なる目的を果たします:
    1. 深い(つまり、多くの層を持つ)ニューラルネットワークを最適化しやすくします。ここでのアイデアは、グラディエントがネットワークを通って後方に流れる「ショートカット」を提供し、ネットワークの初期層を最適化しやすくすることです。
    2. 残差接続がないと、より多くの層を追加するときに深いモデルのパフォーマンスが低下します(おそらく、深いネットワークを通じてグラディエントが全て戻るのが難しく、情報を失うため)。残差接続は、より深いネットワークに少しの正確性向上をもたらすようです。
    3. 勾配消失/勾配爆発問題に対処するのに役立ちます。

2 つのサブレイヤーにもう少し深く掘り下げましょう。

位置ごとのフィードフォワードネットワーク

これは単なる 2 層のマルチレイヤーパーセプトロンです:

def ffn(x, c_fc, c_proj):  # [n_seq, n_embd] -> [n_seq, n_embd]
# アップへのプロジェクト
a = gelu(linear(x, **c_fc)) # [n_seq, n_embd] -> [n_seq, 4*n_embd]

# ダウンへのプロジェクト
x = linear(a, **c_proj) # [n_seq, 4*n_embd] -> [n_seq, n_embd]

return x

ここには特に凝ったことはありません。n_embdからより高い次元4*n_embdへプロジェクトし、その後n_embdへと戻します。5

params辞書から、mlpパラメータが以下のようになっていることを思い出してください:

"mlp": {
"c_fc": {"b": [4*n_embd], "w": [n_embd, 4*n_embd]},
"c_proj": {"b": [n_embd], "w": [4*n_embd, n_embd]},
}

マルチヘッド因果自己注意

このレイヤーは、Transformer の中で最も理解しにくい部分かもしれません。そこで、「マルチヘッド因果自己注意」を各単語に分解して説明することで、この概念を理解しやすくしましょう:

  1. 注意(Attention)
  2. 自己(Self)
  3. 因果(Causal)
  4. マルチヘッド(Multi-Head)

注意(Attention)

このトピックについては、別のブログ記事を書いています。そこでは、元の Transformer 論文で提案されたスケールドドットプロダクト方程式を一から導き出しています: attention(Q,K,V)=softmax(QKTdk)V\text{attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})Vしたがって、この投稿では注意に関する説明は省略します。Lilian Weng 氏の Attention? Attention!Jay Alammar 氏 の The Illustrated Transformerも、注意に関する素晴らしい説明です。

私のブログ記事から注意実装を適応させましょう:

def attention(q, k, v):  # [n_q, d_k], [n_k, d_k], [n_k, d_v] -> [n_q, d_v]
return softmax(q @ k.T / np.sqrt(q.shape[-1])) @ v

自己(Self)

qkvがすべて同じソースから来る場合、自己注意を実行しています(つまり、入力シーケンスを自身に対して注意させる):

def self_attention(x): # [n_seq, n_embd] -> [n_seq, n_embd]
return attention(q=x, k=x, v=x)

例えば、入力が"Jay went to the store, he bought 10 apples."である場合、「he」が「Jay」を含む他のすべての単語に注意を払うことになり、モデルが「he」が「Jay」を指していることを認識できるようになります。

qkvおよび注意出力のための投影を導入することで、自己注意を強化できます:

def self_attention(x, w_k, w_q, w_v, w_proj): # [n_seq, n_embd] -> [n_seq, n_embd]
# qkvの投影
q = x @ w_k # [n_seq, n_embd] @ [n_embd, n_embd] -> [n_seq, n_embd]
k = x @ w_q # [n_seq, n_embd] @ [n_embd, n_embd] -> [n_seq, n_embd]
v = x @ w_v # [n_seq, n_embd] @ [n_embd, n_embd] -> [n_seq, n_embd]

# 自己注意の実行
x = attention(q, k, v) # [n_seq, n_embd] -> [n_seq, n_embd]

# 出力の投影
x = x @ w_proj # [n_seq, n_embd] @ [n_embd, n_embd] -> [n_seq, n_embd]

return x

これにより、モデルはqkvに最適なマッピングを学習でき、入力間の関係を区別するのに役立ちます。

w_qw_kw_vを単一の行列w_fcに組み合わせ、投影を実行し、その結果を分割することで、行列乗算の数を 4 から 2 に減らすことができます:

def self_attention(x, w_fc, w_proj): # [n_seq, n_embd] -> [n_seq, n_embd]
# qkvの投影
x = x @ w_fc # [n_seq, n_embd] @ [n_embd, 3*n_embd] -> [n_seq, 3*n_embd]

# qkvに分割
q, k, v = np.split(x, 3, axis=-1) # [n_seq, 3*n_embd] -> 3 of [n_seq, n_embd]

# 自己注意の実行
x = attention(q, k, v) # [n_seq, n_embd] -> [n_seq, n_embd]

# 出力の投影
x = x @ w_proj # [n_seq, n_embd] @ [n_embd, n_embd] = [n_seq, n_embd]

return x

これは、現代のアクセラレータ(GPU)が、連続して発生する 3 つの小さな行列乗算よりも 1 つの大きな行列乗算をより有効に利用できるため、少し効率的です。

最後に、GPT-2 の実装に合わせてバイアスベクトルを追加し、linear関数を使用し、params辞書に合わせてパラメータの名前を変更します:

def self_attention(x, c_attn, c_proj): # [n_seq, n_embd] -> [n_seq, n_embd]
# qkvの投影
x = linear(x, **c_attn) # [n_seq, n_embd] -> [n_seq, 3*n_embd]

# qkvに分割
q, k, v = np.split(x, 3, axis=-1) # [n_seq, 3*n_embd] -> 3 of [n_seq, n_embd]

# 自己注意の実行
x = attention(q, k, v) # [n_seq, n_embd] -> [n_seq, n_embd]

# 出力の投影
x = linear(x, **c_proj) # [n_seq, n_embd] @ [n_embd, n_embd] = [n_seq, n_embd]

return x

params辞書から、attnパラメータが以下のようになっていることを思い出してください:

"attn": {
"c_attn": {"b": [3*n_embd], "w": [n_embd, 3*n_embd]},
"c_proj": {"b": [n_embd], "w": [n_embd, n_embd]},
},

因果(Causal)

現在の自己注意の設定には少し問題があります。入力が未来を見ることができてしまいます!例えば、入力が["not", "all", "heroes", "wear", "capes"]の場合、自己注意を行う際に「wear」が「capes」を見ることを許可しています。これは、「wear」の出力確率が、モデルが正解が「capes」であることをすでに知っているために偏ってしまうことを意味します。これはよくありません。なぜなら、モデルは入力iiの正解を入力i+1i+1から取得できることを学習するだけになってしまうからです。

これを防ぐために、注意行列を何らかの方法で修正し、入力が未来を見ることができないように隠すまたはマスクする必要があります。例えば、注意行列が以下のようになっているとしましょう:

       not    all    heroes wear   capes
not 0.116 0. 0. 0. 0.
all 0.180 0.397 0. 0. 0.
heroes 0.156 0.453 0.028 0. 0.
wear 0.499 0.055 0.133 0.017 0.
capes 0.089 0.290 0.240 0.228 0.153

各行はクエリに対応し、列はキーに対応します。この場合、「wear」の行を見ると、最後の列の「capes」に 0.295 の重みで注意を払っていることがわかります。これを防ぐために、そのエントリを0.0に設定したいです:

      not    all    heroes wear   capes
not 0.116 0.159 0.055 0.226 0.443
all 0.180 0.397 0.142 0.106 0.175
heroes 0.156 0.453 0.028 0.129 0.234
wear 0.499 0.055 0.133 0.017 0.
capes 0.089 0.290 0.240 0.228 0.153

一般に、入力のすべてのクエリが未来を見ることを防ぐために、j>ij > iであるすべての位置i,ji, j0に設定します:

       not    all    heroes wear   capes
not 0.116 0. 0. 0. 0.
all 0.180 0.397 0. 0. 0.
heroes 0.156 0.453 0.028 0. 0.
wear 0.499 0.055 0.133 0.017 0.
capes 0.089 0.290 0.240 0.228 0.153

これをマスキングと呼びます。上記のマスキングアプローチの問題は、softmaxが適用された後に 0 に設定しているため、行の合計が 1 にならなくなることです。行が依然として 1 になるようにするためには、softmaxが適用される前に注意行列を修正する必要があります。

これは、マスクされるべきエントリをsoftmaxの前に-\inftyに設定することで達成できます6

def attention(q, k, v, mask):  # [n_q, d_k], [n_k, d_k], [n_k, d_v], [n_q, n_k] -> [n_q, d_v]
return softmax(q @ k.T / np.sqrt(q.shape[-1]) + mask) @ v

ここでmaskは行列(n_seq=5の場合):

0 -1e10 -1e10 -1e10 -1e10
0 0 -1e10 -1e10 -1e10
0 0 0 -1e10 -1e10
0 0 0 0 -1e10
0 0 0 0 0

-np.infの代わりに-1e10を使用するのは、-np.infnansを引き起こす可能性があるためです。

maskを注意行列に追加するのは、実際には、-infにどんな数を加えても結局-infになるため、単に値を-1e10に設定するのと同じ効果があります。

NumPy でmask行列を計算するには、(1 - np.tri(n_seq)) * -1e10を使用します。

これをすべてまとめると、以下のようになります:

def attention(q, k, v, mask):  # q, k, v はそれぞれクエリ、キー、バリューを表し、mask は注意を適用する範囲を制御します。戻り値は注意後の出力です。
return softmax(q @ k.T / np.sqrt(q.shape[-1]) + mask) @ v # スケール済みドット積注意を計算し、適用します。

def causal_self_attention(x, c_attn, c_proj): # x は入力シーケンス、c_attn と c_proj はそれぞれ注意と出力投影のための設定を含む。
# qkv 投影
x = linear(x, **c_attn) # 入力を線形変換して、q, k, v のための3倍の次元を持つベクトルにします。

# qkv に分割
q, k, v = np.split(x, 3, axis=-1) # 上記のベクトルを q, k, v に分割します。

# 未来の入力を隠すための因果マスク
causal_mask = (1 - np.tri(x.shape[0]), dtype=x.dtype) * -1e10 # 自己注意で将来の情報が現れないようにするマスクを作成します。

# 因果的自己注意を実行
x = attention(q, k, v, causal_mask) # 因果的自己注意を適用します。

# 出力投影
x = linear(x, **c_proj) # 注意後の出力を再び線形変換して最終的な出力を得ます。

return x

マルチヘッド

n_head個の別々の注意計算を実行し、クエリ、キー、値をヘッドに分割することで、実装をさらに改善できます:

def mha(x, c_attn, c_proj, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
# qkvの投影
x = linear(x, **c_attn) # [n_seq, n_embd] -> [n_seq, 3*n_embd]

# qkvに分割
qkv = np.split(x, 3, axis=-1) # [n_seq, 3*n_embd] -> [3, n_seq, n_embd]

# ヘッドに分割
qkv_heads = list(map(lambda x: np.split(x, n_head, axis=-1), qkv)) # [3, n_seq, n_embd] -> [3, n_head, n_seq, n_embd/n_head]

# 未来の入力を見ることができないように因果マスクを適用
causal_mask = (1 - np.tri(x.shape[0]), dtype=x.dtype) * -1e10 # [n_seq, n_seq]

# 各ヘッドで注意を実行
out_heads = [attention(q, k, v, causal_mask) for q, k, v in zip(*qkv_heads)] # [3, n_head, n_seq, n_embd/n_head] -> [n_head, n_seq, n_embd/n_head]

# ヘッドを結合
x = np.hstack(out_heads) # [n_head, n_seq, n_embd/n_head] -> [n_seq, n_embd]

# 出力の投影
x = linear(x, **c_proj) # [n_seq, n_embd] -> [n_seq, n_embd]

return x

ここには 3 つのステップが追加されています:

  1. q, k, vn_head個のヘッドに分割します:
# ヘッドに分割
qkv_heads = list(map(lambda x: np.split(x, n_head, axis=-1), qkv)) # [3, n_seq, n_embd] -> [n_head, 3, n_seq, n_embd/n_head]
  1. 各ヘッドに対して注意を計算します:
# 各ヘッドで注意を実行
out_heads = [attention(q, k, v) for q, k, v in zip(*qkv_heads)] # [n_head, 3, n_seq, n_embd/n_head] -> [n_head, n_seq, n_embd/n_head]
  1. 各ヘッドの出力を結合します:
# ヘッドを結合
x = np.hstack(out_heads) # [n_head, n_seq, n_embd/n_head] -> [n_seq, n_embd]

これにより、n_embdからn_embd/n_headに次元が減少します。これはトレードオフです。次元が減少した代わりに、モデルは注意を通じて関係をモデリングする際に、追加の部分空間を利用できるようになります。たとえば、ある注意ヘッドが代名詞を参照している人物に接続する責任があるかもしれません。別のものは、文をピリオドでグループ化する責任があるかもしれません。また、どの単語がエンティティであり、どれがそうでないかを単純に識別するかもしれません。それでも、おそらくそれは単なる別のニューラルネットワークのブラックボックスです。

書いたコードは、ループ内で各ヘッドの注意計算を順番に(一度に 1 つ)実行していますが、これは非常に効率的ではありません。実際には、これらを並列で実行したいところです。簡単のために、このシーケンシャルなままにしておきます。

これで、GPT の実装がようやく終わりました!あとは、すべてをまとめてコードを実行するだけです。

すべてをまとめる

すべてをまとめると、gpt2.pyが完成します。これは全体でわずか 120 行のコードです(コメントと空白を削除すると 60 行です)。

以下のコマンドで実装をテストできます:

python gpt2.py \
"Alan Turing theorized that computers would one day become" \
--n_tokens_to_generate 8

これにより、出力は以下のようになります:

the most powerful machines on the planet.

うまくいきました!

次のDockerfileを使用して、実装がOpenAI の公式 GPT-2 リポジトリと同一の結果を出すことをテストできます(注:tensorflow の問題と、すべての 4 つの GPT-2 モデルサイズをダウンロードすることになるため、M1 Macbook では動作しません。ダウンロードするものが多いため、警告が出ます):

これにより、同一の結果が得られるはずです:

the most powerful machines on the planet.

次はなんですか?

この実装はクールですが、たくさんの機能が欠けています:

GPU/TPU サポート

NumPy を JAXに置き換えます:

import jax.numpy as np

これで、コードを GPU や TPU で使えるようになります!ただし、JAX を正しくインストールすることを確認してください。

逆伝播

再び、NumPy をJAXに置き換える場合:

import jax.numpy as np

次に、勾配を計算することは非常に簡単です:

def lm_loss(params, inputs, n_head) -> float:
x, y = inputs[:-1], inputs[1:]
output = gpt2(x, **params, n_head=n_head)
loss = np.mean(-np.log(output[y]))
return loss

grads = jax.grad(lm_loss)(params, inputs, n_head)

バッチ処理

再び、NumPy を JAX に置き換える場合7

import jax.numpy as np

その後 gpt2関数をバッチ処理することは非常に簡単です:

gpt2_batched = jax.vmap(gpt2, in_axes=[0, None, None, None, None, None])
gpt2_batched(batched_inputs) # [batch, seq_len] -> [batch, seq_len, vocab]

推論最適化

実装はかなり非効率です。GPU とバッチ処理のサポートを除いて、最も迅速かつ大きな影響を与える最適化は、kv キャッシュの実装でしょう。また、注意ヘッドの計算を逐次的に実装しましたが、実際には並行して行うべきです8

推論の最適化はまだまだ多くあります。Lillian Weng 氏の大規模 Transformer モデル推論最適化Kipply 氏の Transformer 推論算術を出発点として推奨します。

トレーニング

GPT のトレーニングは、ニューラルネットワークにとってはかなり標準的なプロセスです(損失関数に関する勾配降下法を使用)。もちろん、GPT のトレーニングには、標準的なテクニック群を使用する必要があります(例えば、Adam オプティマイザーを使用、最適な学習率を見つける、ドロップアウトや重み減衰による正則化、学習率スケジューラーの使用、適切な重みの初期化、バッチ処理など)。

良い GPT モデルをトレーニングする実際の秘訣は、データとモデルのスケーリングに能力を持つことであり、これが真の挑戦です。

データのスケーリングでは、大規模で高品質、そして多様なテキストコーパスが必要です。

  • 大規模とは、数十億のトークン(テラバイトのデータ)を意味します。例えば、大規模言語モデルのためのオープンソースの事前学習データセットであるThe Pileを確認してください。
  • 高品質とは、重複する例のフィルタリング、未フォーマットテキスト、非論理的なテキスト、ゴミテキストなどを除外したいということです。
  • 多様性とは、さまざまなトピックについて、異なるソースから、異なる視点で、シーケンス長が異なるテキストを意味します。もちろん、データに偏りがある場合、それはモデルに反映されるため、注意が必要です。

数十億のパラメーターを持つモデルのスケーリングには、大量のエンジニアリング(そしてお金笑)が関わっています。トレーニングフレームワークは非常に長く複雑になることがあります。始めるには、Lillian Weng 氏の How to Train Really Large Models on Many GPUsが良い出発点です。このトピックに関連するものとしては、NVIDIA の Megatron FrameworkCohere の Training FrameworkGoogle の PALM、オープンソースのmesh-transformer-jax(EleutherAI のオープンソースモデルのトレーニングに使用されています)、そして他にも多く 多くの研究 があります

評価

なかなか難しい問題ですが、LLM の評価方法についてはどのように進めればよいでしょうか。正直なところ、非常に難しい問題です。HELMは非常に包括的で、始めるにあたって良い参考資料ですが、ベンチマークと評価指標については常に懐疑的であるべきです。

アーキテクチャの改善

Phil Wang 氏の X-Transformerに目を通されることをお勧めします。Transformer アーキテクチャに関する最新の研究が含まれています。この論文もかなり良い要約です(表 1 を参照)。Facebook の最近のLLaMA 論文も、2023 年 2 月時点での標準的なアーキテクチャ改善に関する良い参考資料でしょう。

生成の停止

現在の実装では、事前に生成したいトークンの正確な数を指定する必要があります。これはあまり良いアプローチではなく、生成される文が長すぎたり、短すぎたり、文の途中で切れたりします。

これを解決するために、特別な文末 (EOS) トークンを導入することができます。事前学習中には、EOS トークンを入力の末尾に追加します(例:tokens = ["not", "all", "heroes", "wear", "capes", ".", ""])。生成時には、EOS トークンに遭遇した場合(または最大シーケンス長に達した場合)に停止します:

def generate(inputs, eos_id, max_seq_len):
prompt_len = len(inputs)
while inputs[-1] != eos_id and len(inputs) < max_seq_len:
output = gpt(inputs)
next_id = np.argmax(output[-1])
inputs.append(int(next_id))
return inputs[prompt_len:]

GPT-2 は EOS トークンを事前学習していないので、このアプローチを私たちのコードで使用することはできませんが、現在のほとんどの LLM は EOS トークンを使用しています。

無条件生成

私たちのモデルでテキストを生成するには、プロンプトでそれを条件付けする必要があります。しかし、モデルに任意の入力プロンプトなしでテキストを生成させる、無条件生成も可能です。

これは、事前学習中に入力の開始に特別な文の始まり(BOS)トークンを前置することによって達成されます(つまり、tokens = ["", "not", "all", "heroes", "wear", "capes", "."])。次に、無条件にテキストを生成するために、BOS トークンだけを含むリストを入力します:

def generate_unconditioned(bos_id, n_tokens_to_generate):
inputs = [bos_id]
for _ in range(n_tokens_to_generate):
output = gpt(inputs)
next_id = np.argmax(output[-1])
inputs.append(int(next_id))
return inputs[1:]

GPT-2 は、混乱しやすい名前の BOS トークン(``)で事前学習されているため、私たちの実装で無条件生成を行うのは、次の行を変更するだけで非常に簡単です。

input_ids = encoder.encode(prompt) if prompt else [encoder.encoder[""]]

そして、以下を実行します:

python gpt2.py ""

これにより、以下が生成されます:

The first time I saw the new version of the game, I was so excited. I was so excited to see the new version of the game, I was so excited to see the new version

私たちが貪欲サンプリングを使用しているため、出力はあまり良くない(繰り返しが多い)ということになり、決定論的です(つまり、コードを実行するたびに同じ出力が得られます)。より高品質で非決定論的な生成物を得るには、分布から直接サンプリングする必要があります(理想的には top-p のようなものを適用した後)。

無条件生成は特に有用ではありませんが、GPT の能力を示す楽しい方法です。

ファインチューニング

トレーニングセクションでファインチューニングについて簡単に触れました。ファインチューニングとは、事前にトレーニングされた重みを再利用して、ある下流タスクでモデルをトレーニングすることを指します。このプロセスをトランスファーラーニングと呼びます。

理論的には、ゼロショットやフューショットのプロンプティングを使用してタスクを完了させることができますが、ラベル付きデータセットにアクセスできる場合、GPT をファインチューニングすることでより良い結果(追加データと高品質データを考慮したスケールの結果)が得られます。

ファインチューニングに関連するいくつかの異なるトピックがあります。以下に分けて説明します:

クラス分類ファインチューニング

クラス分類ファインチューニングでは、モデルにテキストを与え、それがどのクラスに属するかを予測させます。例えば、映画のレビューが良いか悪いかを評価するIMDB データセットを考えてみましょう:

--- 例 1 ---
テキスト: I wouldn't rent this one even on dollar rental night.
ラベル: 悪い
--- 例 2 ---
テキスト: I don't know why I like this movie so well, but I never get tired of watching it.
ラベル: 良い
--- 例 3 ---
...

モデルをファインチューニングするために、言語モデリングのヘッドを分類ヘッドに置き換え、最後のトークン出力に適用します:

def gpt2(inputs, wte, wpe, blocks, ln_f, cls_head, n_head):
x = wte[inputs] + wpe[range(len(inputs))]
for block in blocks:
x = transformer_block(x, **block, n_head=n_head)
x = layer_norm(x, **ln_f)

# n_classesにプロジェクト
# [n_embd] @ [n_embd, n_classes] -> [n_classes]
return x[-1] @ cls_head

言語モデリングの場合とは異なり、n_seqの分布ではなく、入力全体に対して単一の確率分布を生成するだけでよいため、最後のトークン出力x[-1]のみを使用します。特に最後のトークンを使用する理由は、最後のトークンのみが全シーケンスに注目を許され、入力テキスト全体についての情報を持っているからです。

通常どおり、交差エントロピー損失に対して最適化を行います:

def singe_example_loss_fn(inputs: list[int], label: int, params) -> float:
logits = gpt(inputs, **params)
probs = softmax(logits)
loss = -np.log(probs[label]) # 交差エントロピー損失
return loss

また、sigmoidsoftmaxの代わりに適用し、各クラスに対して二項交差エントロピー損失を取ることで、マルチラベル分類(つまり、例が単一のクラスにのみ属するのではなく、複数のクラスに属することができる)も行うことができます(この stack-exchange の質問を参照)。

生成型ファインチューニング

一部のタスクは、クラスにきれいに分類することができません。たとえば、要約のタスクを考えてみましょう。このタイプのタスクをファインチューニングするには、入力とラベルを連結して言語モデリングを行うだけです。以下は、単一の要約トレーニングサンプルがどのように見えるかの例です:

--- 記事 ---
これは私が要約したい記事です。
--- 要約 ---
これが要約です。

我々は、事前トレーニング中と同様にモデルをトレーニングします(言語モデリング損失に関して最適化します)。

予測時には、モデルに--- 要約 ---までのすべてをフィードし、自動回帰言語モデリングを実行して要約を生成します。

区切り文字--- 記事 ------ 要約 ---の選択は任意です。トレーニングと推論の間で一貫性がある限り、テキストのフォーマット方法は自由です。

なお、分類タスクも生成タスクとして定式化することが可能です(例えば IMDB で):

--- テキスト ---
I wouldn't rent this one even on dollar rental night.
--- ラベル ---
悪い

しかし、これは直接分類ファインチューニングを行うよりも性能が低下する可能性があります(損失には最終予測だけでなく、全シーケンスにわたる言語モデリングが含まれるため、予測に特化した損失が薄れる可能性があります)。

インストラクションファインチューニング

最近の最先端の大規模言語モデルは、事前トレーニング後に追加のインストラクションファインチューニングステップを経るものが多いです。このステップでは、モデルは数千のインストラクションプロンプト+完了ペアについて、人間によるラベリングが行われたデータでファインチューニング(生成的)されます。インストラクションファインチューニングは、データが人間によってラベル付けされている(つまり監視されている)ため、監視されたファインチューニングとも呼ばれます。

では、インストラクションファインチューニングの利点は何でしょうか?Wikipedia の記事で次の単語を予測することはモデルが文を続けるのが得意になることを意味しますが、それだけでは特に指示に従ったり、会話をしたり、文書を要約したり(私たちが GPT にしてほしいことすべて)するのが得意になるわけではありません。人間によるラベル付けされたインストラクション+完了ペアでファインチューニングすることは、モデルがどのように役立つか、そしてそれらとのやり取りを容易にする方法を教える方法です。これをAI アライメントと呼びます。私たちがそれにしてほしいこと、振る舞ってほしい方法にモデルを整列させているからです。アライメントは活発な研究領域であり、指示に従うことだけでなく(バイアス、安全性、意図など)を含みます。

このインストラクションデータは具体的にどのようなものでしょうか?Google のFLANモデルは、すでに人間によってラベル付けされたさまざまな学術 NLP データセットでトレーニングされました:

FLAN論文の図3

一方、OpenAI のInstructGPTは、自社の API から収集したプロンプトでトレーニングされました。その後、それらのプロンプトに対して完了を書くように作業者に支払いました。データの内訳は以下の通りです:

InstructGPT論文の表1と2

パラメータ効率の良いファインチューニング

上記のセクションでファインチューニングについて話すとき、全てのモデルパラメータを更新していると仮定しています。これは最高のパフォーマンスを生み出しますが、計算コスト(全モデルにわたってバックプロパゲーションが必要)とストレージコスト(ファインチューンされた各モデルに対して、完全に新しいパラメータのコピーを保存する必要がある)の両方においてコストがかかります。

この問題に対する最も単純なアプローチは、ヘッドのみを更新し、残りのモデルを凍結(つまり、トレーニング不可能にする)することです。これによりトレーニングが速くなり、新しいパラメータの数が大幅に減少しますが、ディープラーニングの「ディープ」を失ってしまうため、特に良いパフォーマンスは得られません。代わりに特定の層を選択的に凍結する(例えば、最後の 4 層以外の全層を凍結、または隔層ごとに凍結、またはマルチヘッドアテンションパラメータ以外の全パラメータを凍結)ことで、深さを復元することができます。この結果、パフォーマンスは大幅に向上しますが、パラメータ効率が低下し、トレーニング速度の向上の一部を失ってしまいます。

代わりに、パラメータ効率の良いファインチューニングメソッドを利用することができます。これは現在も活発に研究されている分野で、多くの 異なる 方法 選択 可能 です

例として、Adapters 論文を取り上げます。このアプローチでは、Transformer ブロックの FFN と MHA 層の後に追加の「アダプター」層を追加します。アダプター層は、入力と出力の次元がn_embdで、隠れ層の次元がn_embdより小さい、単純な 2 層全結合ニューラルネットワークです:

Adapters論文からの図2

隠れた次元のサイズは、パラメータとパフォーマンスのトレードオフを可能にするハイパーパラメータです。BERT モデルにおいて、このアプローチを使用することで、訓練されたパラメータの数を 2%に減少させることができ、フルファインチューニングと比較してわずかなパフォーマンスの低下(<1%)を維持しながらも、と論文で示されています。

Footnotes

  1. 訓練を大規模に行い、テラバイト単位のデータを収集し、モデルを高速化し、パフォーマンスを評価し、モデルを人間に役立つように調整することは、現在の LLM(大規模言語モデル)を形作るために必要な数百人のエンジニア/研究者の生涯の仕事です。それは単にアーキテクチャだけではありません。GPT アーキテクチャはたまたま、スケーリング特性が良く、GPU 上で高度に並列化可能で、シーケンスをモデリングするのに適しているという、最初のニューラルネットワークアーキテクチャでした。実際の秘密のソースは、いつものようにデータとモデルをスケーリングすることから来ます。GPT は私たちがそれを行うことを可能にするだけです9。Transformer がハードウェアの宝くじに当たった可能性があり、他のアーキテクチャがまだそこに存在し、Transformer の地位を奪うことを待っているかもしれません。

  2. 特定のアプリケーションでは、トークナイザーにデコード方法は必要ありません。たとえば、映画のレビューが映画を良いものと言っているか悪いものと言っているかを分類したい場合、テキストをエンコードしてモデルのフォワードパスを行うことができれば十分で、デコードは必要ありません。しかし、テキストを生成する場合は、デコードが必要です。

  3. もちろん、このすべてのデータから学ぶには、十分に大きなモデルが必要です。これが、GPT-3 が 1750 億のパラメータを持ち、トレーニングには約 100 万ドルから 1000 万ドルの計算コストがかかったと推定される理由です。ただし、InstructGPTChinchilla の論文でわかったことは、実際にはそれほど大きなモデルをトレーニングする必要がないということです。最適にトレーニングされ、命令に微調整された 13 億パラメータの GPT は、1750 億パラメータの GPT-3 を上回る性能を発揮することができます。

  4. 元の Transformer 論文では、計算された位置埋め込みを使用しました。これは学習された位置埋め込みと同じくらいうまく機能することがわかっていますが、任意の長さのシーケンスを入力できるという明確な利点があります(最大シーケンス長に制限されません)。しかし、実際には、モデルの性能は訓練されたシーケンスの長さにのみ良くなります。1024 トークンの長さで GPT を訓練し、それが 16k トークンの長さでうまく機能すると期待することはできません。しかし最近、AlibiRoPEのような相対位置埋め込みによる成功がいくつか報告されています。

  5. 異なる GPT モデルは、4*n_embdでない異なる隠れ層の幅を選択することもありますが、これは GPT モデルにおける一般的な慣習です。また、Transformer の成功を牽引するために多頭部注意層に多くの注意(言葉遊びを意図)を払いますが、GPT-3 のスケールでは、モデルパラメータの 80%がフィードフォワード層に含まれています。考えてみるべきことです。

  6. 納得できない場合は、softmax の方程式をじっと見て、これが真実であることを自分自身に納得させてください(場合によってはペンと紙を取り出して):

    softmax(x)i=exijexj\text{softmax}(\vec{x})_i=\frac{e^{x_i}}{\sum_je^{x_j}}
  7. I love JAX ❤️.

  8. JAX を使用すると、これは heads = jax.vmap(attention, in_axes=(0, 0, 0, None))(q, k, v, causal_mask) として簡単に行えます。

  9. 実際、注意モデルがシーケンスをモデル化する方法は、再帰的/畳み込み層と比較して本質的に優れていると主張するかもしれませんが、今私たちは脚注内の脚注の中にいますので、ここで話を戻しましょう。