学習トランスフォーマーコード第2部 – GPTを間近で観察

学習トランスフォーマーコード第2部 - GPT観察

nanoGPTを介してGenerative Pre-Trained Transformersを掘り下げる

Photo by Luca Onniboni on Unsplash

私のプロジェクトの第2部へようこそ。ここでは、TinyStoriesデータセットとnanoGPTを使用して、トランスフォーマーとGPTベースのモデルの複雑さに深入りします。これらはすべて、古いゲーミングラップトップで訓練されたものです。最初の部分では、キャラクターレベルの生成モデルに入力するためのデータセットを準備しました。第1部へのリンクは以下にあります。

学習トランスフォーマーコード第1部

新シリーズの第1部で、nanoGPTを使用してトランスフォーマーコードを最初に学ぼうとする試み

towardsdatascience.com

この記事では、GPTモデル、そのコンポーネント、およびnanoGPTでの実装を解説します。私はnanoGPTを選んだ理由は、GPTモデルのシンプルなPython実装と、約300行のソースコードを読むだけでGPTモデルを迅速に理解できることです。率直に言って、最初にコードを調べたときにはこの理解が足りませんでした。まだ理解できていない部分もありますが、学んだことをすべて活かして、この説明がGPTスタイルのモデルが内部でどのように機能するかを直感的に理解したい人々の出発点となることを願っています。

この記事の準備のために、さまざまな論文を読みました。最初は、元となる論文「Attention is All You Need」を読むだけで十分理解が進むと思い込んでいました。これは甘い考えでした。この論文はトランスフォーマーモデルを紹介したものですが、テキスト生成などの高度なタスクに適応させるためには後続の論文が必要でした。AIAYNは、より広範なトピックの導入に過ぎませんでした。しかし、私はHackerNewsの記事を思い出し、LLMの完全な理解に役立つ読書リストを提供していることを思い出しました。素早い検索で、こちらの記事を見つけました。シーケンスに沿ってすべてを順番に読むわけではありませんでしたが、この読書リストを完成した後も学習の旅を続けるために再訪する予定です。

それでは、さっそく始めましょう。GPTモデルを詳しく理解するためには、トランスフォーマーから始める必要があります。トランスフォーマーは、スケールドドットプロダクトアテンションとして知られるセルフアテンションメカニズムを使用しています。以下の説明は、スケールドドットプロダクトアテンションに関するこの洞察に基づいています。詳細な理解には、この洞察に対する記事をお勧めします。基本的に、入力シーケンスの各要素(i番目の要素)ごとに、i番目の要素とシーケンス全体の要素の重み付き平均を入力シーケンスに乗算したいと考えています。これらの重みは、i番目の要素のベクトルと全体の入力ベクトルの内積を取り、それにsoftmaxを適用して重みを0から1の値にします。元の「Attention is All You Need」の論文では、これらの入力はクエリ(シーケンス全体)、キー(i番目の要素のベクトル)および(またはシーケンス全体)と呼ばれます。アテンションメカニズムへのパスには、ランダムな値で初期化され、ニューラルネットワーク内でのパスの数が増えるにつれて学習されます。

nanoGPTはスケールドドットプロダクトアテンションを実装し、複数のアテンション操作が同時に行われるマルチヘッドアテンションとして拡張しています。また、これをtorch.nn.Moduleとして実装しているため、他のネットワーク層と組み合わせることができます。

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

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                        .view(1, 1, config.block_size, config.block_size))
    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

このコードをさらに解析してみましょう。まず、コンストラクタから始めます。最初に、埋め込みの次元(n_embed)が注意ヘッドの数(n_heads)で均等に分割されているかを確認します。これは重要です。なぜなら、埋め込みを各ヘッドごとにセクションに分割する際に、埋め込み空間全体を隙間なくカバーしたいからです。次に、2つのLinearレイヤー、c_attc_projを初期化します。 c_attは、スケーリングされたドット積注意計算の行列の作成に使用する作業スペースを保持するレイヤーであり、c_projは計算の最終結果を格納します。c_att内の埋め込み次元は、注意の3つの主要なコンポーネント(querykey、およびvalue)にスペースを含める必要があるため、3倍になります。

また、2つのドロップアウトレイヤー、attn_dropoutresid_dropoutもあります。ドロップアウトレイヤーは、与えられた確率に基づいて入力行列の要素をランダムに無効にします。PyTorchのドキュメントによると、これはモデルの過学習を減らす目的で使用されます。 config.dropoutの値は、ドロップアウトレイヤーの実行中に特定のサンプルがドロップされる確率です。

コンストラクタを最終化する前に、ユーザーがPyTorch 2.0にアクセスできるかどうかを確認します。PyTorch 2.0には、スケーリングされたドット積注意の最適化バージョンがあります。利用可能な場合、このクラスはそれを使用します。それ以外の場合はバイアスマスクを設定します。このマスクは、注意メカニズムのオプションのマスキング機能の一部です。torch.trilメソッドは、上三角部分がゼロに変換された行列を生成します。torch.onesメソッドと組み合わせることで、サンプリングされた入力に対して予測される出力を生成するために注意メカニズムが使用する1と0のマスクを効果的に生成します。

次に、このクラスのforwardメソッドに入り、注意アルゴリズムが適用されます。最初に、入力行列のサイズを決定し、3つの次元に分割します:Batchサイズ、Time(またはサンプル数)、Corpus(または埋め込みサイズ)。nanoGPTでは、バッチ学習プロセスを使用しており、この注意レイヤーを利用するトランスフォーマーモデルを詳しく調べる際に詳しく説明します。今のところ、データをバッチで扱っていると理解することが十分です。次に、入力xを線形変換レイヤーc_attnにフィードします。このレイヤーは次元をn_embedから3倍のn_embedに拡張します。その変換の出力は、私たちのqkv変数に分割され、これらの変数の各データをPyTorchのscaled_dot_product_attention関数が期待する形式に再構成するためにviewメソッドが利用されます。

最適化された関数が利用できない場合、コードはスケーリングされたドット積注意の手動実装にデフォルトします。この手法は、qkの行列のドット積を取り、ドット積関数に適合するようにkを転置し、その結果をkのサイズの平方根でスケーリングします。次に、前に作成したバイアスバッファを使用してスケーリングされた出力にマスクを適用し、0を負の無限大で置き換えます。次に、att行列にsoftmax関数を適用し、負の無限大を0に変換し、他のすべての値を0から1の間にスケーリングします。その後、過学習を回避するためにドロップアウトレイヤーを適用し、att行列とvのドット積を取得します。

スケーリングされたドット積の実装に関係なく、マルチヘッドの出力は並べ替えられ、最後のドロップアウトレイヤーを通過して結果を返します。これがPython/PyTorchの50行未満の注意レイヤーの完全な実装です。上記のコードを完全に理解できない場合は、記事の残りを進める前に時間をかけて確認することをお勧めします。

すべてを統合するGPTモジュールに進む前に、さらに2つのビルディングブロックが必要です。1つ目は単純な多層パーセプトロン(MLP)で、「Attention is All You Need」という論文でフィードフォワードネットワークと呼ばれているものです。2つ目は、注意レイヤーとMLPを組み合わせて基本的なトランスフォーマーアーキテクチャを完成させる注意ブロックです。これらは、nanoGPTの以下のスニペットで実装されています。

class MLP(nn.Module):    """    マルチレイヤーパーセプトロン    """    def __init__(self, config):        super().__init__()        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)        self.gelu    = nn.GELU()        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)        self.dropout = nn.Dropout(config.dropout)    def forward(self, x):        x = self.c_fc(x)        x = self.gelu(x)        x = self.c_proj(x)        x = self.dropout(x)        return xclass Block(nn.Module):    def __init__(self, config):        super().__init__()        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)        self.attn = CausalSelfAttention(config)        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)        self.mlp = MLP(config)    def forward(self, x):        x = x + self.attn(self.ln_1(x))        x = x + self.mlp(self.ln_2(x))        return x

MLPレイヤーは、コードの行数からはシンプルに見えますが、モデルに複雑さを追加する要素です。基本的に、Linearレイヤーは入力レイヤーと出力レイヤーの各要素をリンクし、線形変換を使用して値を転送します。上記のコードでは、埋め込みサイズのn_embedをパラメータ数として開始し、出力でそれを4倍にします。ここでの4倍化は任意です。MLPモジュールの目的は、より多くのノードを追加することでネットワークの計算を向上させることです。MLPの始まりで次元数を増やし、終わりで次元数を減らす点が同じであれば、スケーリングの数値は単なる別のハイパーパラメータに過ぎません。重要な要素として、活性化関数を考慮する必要があります。このMLPの実装では、GELU活性化関数を使用して2つの線形層を接続しています。元の論文ではReLU関数が使用されていますが、nanoGPTではGPT2モデルのチェックポイントとの互換性を確保するためにGELUが使用されています。

次に、Blockモジュールを見てみましょう。このモジュールは、”Attention”の論文で概説された通り、トランスフォーマーブロックを最終的にまとめ上げます。基本的には、入力を正規化レイヤーを介して通常の注意機構に渡し、その結果を入力に再度追加します。この追加の出力は再び正規化され、その後MLPに転送され、自身に追加されます。このプロセスは、”Attention”の論文で説明されているトランスフォーマーのデコーダー側を実装しています。テキスト生成では、デコーダーのみを使用することが一般的です。デコーダーの出力を入力シーケンス以外のものに依存させる必要はないためです。トランスフォーマーはもともと機械翻訳のために設計されており、入力トークンのエンコーディングと出力トークンのエンコーディングの両方を考慮する必要があります。しかし、テキスト生成では、単一のトークンのエンコーディングのみを使用するため、エンコーダーを介したクロスアテンションの必要はありません。nanoGPTの作者であるAndrej Karpathyは、このシリーズの最初の記事にリンクされたビデオでこれについて詳しく説明しています。

最後に、主要なコンポーネントであるGPTモデルに到達します。約300行のファイルの大部分は、GPTモジュールに割り当てられています。モデルの微調整やモデルのトレーニングに使用されるユーティリティなど、有益な機能を管理します(このシリーズの次の記事で詳しく説明します)。そのため、以下にnanoGPTリポジトリで利用可能なものの簡略化バージョンを示します。

class GPT(nn.Module):    def __init__(self, config):        super().__init__()        assert config.vocab_size is not None        assert config.block_size is not None        self.config = config        self.transformer = nn.ModuleDict(dict(            wte = nn.Embedding(config.vocab_size, config.n_embd),            wpe = nn.Embedding(config.block_size, config.n_embd),            drop = nn.Dropout(config.dropout),            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),            ln_f = LayerNorm(config.n_embd, bias=config.bias),        ))        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)        # with weight tying when using torch.compile() some warnings get generated:        # "UserWarning: functional_call was passed multiple values for tied weights.        # This behavior is deprecated and will be an error in future versions"        # not 100% sure what this is, so far seems to be harmless. TODO investigate        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying        # init all weights        self.apply(self._init_weights)        # apply special scaled init to the residual projections, per GPT-2 paper        for pn, p in self.named_parameters():            if pn.endswith('c_proj.weight'):                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))    def _init_weights(self, module):        if isinstance(module, nn.Linear):            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)            if module.bias is not None:                torch.nn.init.zeros_(module.bias)        elif isinstance(module, nn.Embedding):            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)    def forward(self, idx, targets=None):        device = idx.device        b, t = idx.size()        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)        # forward the GPT model itself        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)        x = self.transformer.drop(tok_emb + pos_emb)        for block in self.transformer.h:            x = block(x)        x = self.transformer.ln_f(x)        if targets is not None:            # if we are given some desired targets also calculate the loss            logits = self.lm_head(x)            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)        else:            # inference-time mini-optimization: only forward the lm_head on the very last position            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim            loss = None        return logits, loss    @torch.no_grad()    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):        """        インデックスの条件付きシーケンスidx(shape(b、t)のLongTensor)を受け取り、        シーケンスをmax_new_tokens回完成させ、各回の予測をモデルにフィードバックします。        おそらくこのためには、モデル.eval()モードで動作していることを確認する必要があります。        """        for _ in range(max_new_tokens):            # シーケンスコンテキストが長くなりすぎる場合は、block_sizeで切り取る必要があります            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]            # モデルを順方向に進めてシーケンス内

クラスのコンストラクタから始めましょう。異なるレイヤーはPyTorchのModuleDictに組み立てられ、ある種の構造を提供します。まず、トークンの埋め込みレイヤーと位置埋め込みレイヤーの2つの埋め込みレイヤーから始めます。nn.Embeddingモジュールは、他のレイヤーモジュールよりもストレージ容量を最適化するために、値を疎に配置するように設計されています。これに続いて、ドロップアウトレイヤーがあり、その後に私たちのアテンションレイヤーを形成するn_layer個のブロックモジュールが続きます。その後にもう1つのドロップアウトレイヤーがあります。 lm_headリニアレイヤーは、アテンションブロックからの出力を語彙サイズに削減し、損失値以外のGPTの主な出力として機能します。

レイヤーが定義されたら、モジュールをトレーニングする前に追加のセットアップが必要です。ここでは、Andrzejは位置エンコーディングの重みを出力レイヤーの重みにリンクさせます。コードコメントにリンクされた論文によれば、これによりモデルの最終パラメータが削減され、パフォーマンスが向上します。コンストラクタでは、モデルの重みも初期化されます。これらの重みはトレーニング中に学習されるため、ランダムな数値のガウス分布で初期化され、モジュールのバイアスは0に設定されます。最後に、残差レイヤーの重みはレイヤーの数の平方根でスケーリングされるというGPT-2論文からの変更が利用されます。

ネットワークをフォワードする際には、バッチサイズとサンプル数(ここではt)が入力サイズから取得されます。次に、トレーニングデバイス上にメモリを作成し、positional埋め込みになるものを作成します。次に、入力トークンをトークン埋め込みレイヤーwteに埋め込みます。これに続いて、位置埋め込みがwpeレイヤーで計算されます。これらの埋め込みは加算され、その後にドロップアウトレイヤーを通過します。その結果は、各n_layerブロックに渡されて正規化されます。最終結果はリニアレイヤーlm_headに渡され、埋め込み重みが語彙の各トークンに対する確率スコアに削減されます。

損失が計算される場合(たとえばトレーニング中)、予測されたトークンと実際のトークンの違いをクロスエントロピーを使用して計算します。そうでなければ、損失はNoneです。損失とトークンの確率は、フィードフォワード関数の一部として返されます。

以前のモジュールとは異なり、GPTモジュールには追加のメソッドがあります。私たちにとって最も関連性のあるものは、生成関数です。これは、生成モデルを使用したことがある人には馴染みがあるでしょう。入力トークンidxmax_new_tokensの数、およびtemperatureを指定すると、max_new_tokens個のトークンを生成します。これがどのように実現されているかを見てみましょう。まず、必要に応じて入力トークンをblock_sizeに収まるようにトリミングし、必要に応じて入力の末尾からサンプリングします。次に、トークンをネットワークに送り、出力を入力されたtemperatureにスケーリングします。温度が高いほど、モデルはよりクリエイティブになり、幻想する可能性があります。高い温度はまた、予測可能性の低い出力になります。次に、ソフトマックスを適用してモデルの出力重みを0から1の確率に変換します。サンプリング関数を使用して確率から次のトークンを選択し、そのトークンを次の文字のためにGPTモデルにフィードバックされる入力ベクトルに追加します。

この包括的な記事をお読みいただき、お時間をいただきありがとうございます。注釈付きのソースコードを調べることは、コードセグメントの機能を理解するための貴重な方法ですが、コードのさまざまな部分やパラメータを直接操作することに代わるものはありません。これに合わせて、nanoGPTリポジトリの完全なmodel.pyソースコードへのリンクを提供します。

nanoGPT/model.py at master · karpathy/nanoGPT

The simplest, fastest repository for training/finetuning VoAGI-sized GPTs. - nanoGPT/model.py at master ·…

github.com

次の記事では、nanoGPTのtrain.pyスクリプトを探索し、TinyStoriesデータセットで文字レベルのモデルをトレーニングします。VoAGIで私に従って、見逃すことのないようにしてください!

この記事を作成するために、さまざまなリソースを活用しました。これらのうち多くは、この記事と前の記事で既にリンクされています。ただし、私の義務を果たさないことになるので、これらのリソースをあなたと共有しないといけません。これらのリソースは、任意のトピックのさらなる探求や概念の代替的な説明に役立ちます。

  • ゼロから始めるGPTの構築:コードで詳細に解説したYouTube動画
  • LLM読書リスト:ブログ
  • 「Attention is All You Need」:論文
  • 「Language Models are Unsupervised Multitask Learners」:GPT-2の論文
  • マルチレイヤーパーセプトロンの解説とイラスト:VoAGI
  • ウェイトタイイング:Papers With Code
  • トランスフォーマーニューラルネットワークのイラスト入りガイド:ステップバイステップの説明をYouTubeで提供

GPT-4とカスタムのLangChainスクリプトを使用して編集しました。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more