Skip to main content

Transformer Deep Dive トランスフォーマーディープダイブ

トランスフォーマーアーキテクチャのブレークスルー、科学的根拠、公式、およびコードについて詳しく説明します。
Created on January 18|Last edited on February 17
このレポートは、Carlo Lepelaarsによる「Transformer Deep Dive」の翻訳です。



概要

ELMoULMFitのような言語モデルが飛躍的進歩を遂げた後、トランスフォーマーは自然言語処理(NLP)の分野を席巻しました。これは、 GPT-3DALL-Eなどの一般的な言語モデルの基礎です。また、 HuggingFace Transformers library などのツールにより、機械学習エンジニアはさまざまなNLPタスクを簡単に解決できるようになり、その後のNLPでの多くのブレークスルーが容易になりました。このレポートでは、論文「Attention Is All You Need」(2017)で説明されているように、トランスアーキテクチャについて深く掘り下げ、PyTorchを使用してコード化します。

In this report, we will take a deep dive into the transformer architecture as described in the paper "Attention Is All You Need" (2017) and code it up using PyTorch.



トランスフォーマーは、シーケンス間(seq2seq)モデルです。これは、データに何らかの順序があり、出力自体がシーケンスであるすべての問題に適していることを意味します。アプリケーションの例としては、機械翻訳、要約、音声認識などがあります。最近、Vision Transformers (ViT)は、コンピュータービジョンの最先端技術をさらに改善しました。

以下に、完全なトランスアーキテクチャの視覚化を示します。すべてのコンポーネントが何をするのか、なぜそこにあるのか、そしてすべてがどのように組み合わされるのかを説明します。今のところ、エンコーダー(左側)とデコーダー(右側)があり、どちらも内部にいくつかのニューラルネットワーク層があることを認識してください。

transform.png

Code examples are refactored from The Annotated Transformer by Harvard NLP Group and PyTorch documentation on transformers. コード例は、ハーバードNLPグループによる注釈付きトランスフォーマーおよびトランスフォーマーに関するPyTorchドキュメントからリファクタリングされています。



ークン化

まず、テキストを計算するために、テキストを数値で表す方法が必要です。トークン化は、テキストの文字列を圧縮された一連の記号に解析するプロセスです。このプロセスにより、各整数がテキストの一部を表す整数のベクトルが生成されます。トランスフォーマー論文は、トークン化方法としてバイトペアエンコーディング(BPE)を使用します。BPEは、最も一般的な連続バイト(つまり、文字)が1バイト(つまり、整数)に圧縮される圧縮の形式です。

最近の調査によると、BPEは最適ではなくBERTなどの最近の言語モデルでは代わりにWordPieceトークナイザーが使用されています。WordPieceは、単語全体を1つのトークンにトークン化することが多いため、デコードが簡単で直感的です。対照的に、BPEは単語のスライスをトークン化することがよくあります。これにより、奇妙なトークン化記号や個々の文字(アルファベットからの文字)が不明なトークンとしてトークン化される可能性があります。したがって、WordPieceは、例で使用するトークナイザーです。テキストのコーパスで独自のトークナイザーをトレーニングできますが、実際には、ほとんどの場合、実践者は事前にトレーニングされたトークナイザーを使用します。HuggingFace transformers libraryを使用すると、事前にトレーニングされたトークナイザーを簡単に操作できます。


>>> from transformers import BertTokenizer
>>> tok = BertTokenizer.from_pretrained("bert-base-uncased")
Downloading: 100%|█████████████████████████| 1.04M/1.04M [00:00<00:00, 1.55MB/s]
Downloading: 100%|████████████████████████████| 456k/456k [00:00<00:00, 848kB/s]

>>> tok("Hello, how are you doing?")['inputs_ids']
{'input_ids': [101, 7592, 2129, 2024, 2017, 2725, 1029, 102]}

>>> tok("The Frenchman spoke in the [MASK] language and ate 🥖")['input_ids']
{'input_ids': [101, 1996, 26529, 3764, 1999, 1996, 103, 2653, 1998, 8823, 100, 102]}

>>> tok("[CLS] [SEP] [MASK] [UNK]")['input_ids']
{'input_ids': [101, 101, 102, 103, 100, 102]}


トークナイザーには、エンコードの開始([CLS] = 101)とエンコードの終了(つまり分離)([SEP] = 102)のトークンが自動的に含まれることに注意してください。その他の特別なトークンには、マスキング([MASK] = 103🥖)



埋め込み

テキストの適切な表現を学習するために、シーケンス内の個々のトークンは、埋め込みによってベクトルに変換されます。埋め込みの重みはトランスフォーマーモデルの残りの部分と一緒に学習されるため、これは一種のニューラルネットワークレイヤーと見なすことができます。語彙の各単語のベクトルが含まれ、これらの重みは正規分布N(0,1)から初期化されます。モデルを初期化するとき(E∈R∣vocab∣×dmodelE \isin \mathbb{R}^{|\text{vocab}| \times d_\text{model}})、語彙のサイズ(vocab)とモデルの次元(dmodel=512d_\text{model} = 512)を指定する必要があることに注意してください。最後に、正規化ステップとして重みにsqrt(dmodel)\mathrm{sqrt}\left({d_\text{model}}\right)を掛けます。

import torch
from torch import nn
class Embed(nn.Module):
    def __init__(self, vocab: int, d_model: int = 512):
        super(Embed, self).__init__()
        self.d_model = d_model
        self.vocab = vocab
        self.emb = nn.Embedding(self.vocab, self.d_model)
        self.scaling = math.sqrt(self.d_model)

    def forward(self, x):
        return self.emb(x) * self.scaling


位置エンコーディング

反復ネットワークや畳み込みネットワークとは対照的に、モデル自体には、シーケンス内の埋め込みトークンの相対位置に関する情報がありません。したがって、エンコーダーとデコーダーの入力埋め込みにエンコードを追加して、この情報を挿入する必要があります。この情報はさまざまな方法で追加でき、静的または学習できます。トランスフォーマーは、各位置(pos)にサイン変換とコサイン変換を使用します。偶数次元に(2i2_i) は正弦が使用され、奇数次元(2i+12_{i+1})には余弦が使用されます。

PEpos,2i=sin⁡(pos100002i/dmodel)PE_{\text{pos}, 2i} = \sin(\frac{\text{pos}}{10000^{2i / d_\text{model}}})

PEpos,2i+1=cos⁡(pos100002i/dmodel)PE_{\text{pos}, 2i+1} = \cos(\frac{\text{pos}}{10000^{2i / d_\text{model}}})

そのコードでは、数値のオーバーフローを回避するために、位置エンコーディングがログスペースで計算されます。

import torch
from torch import nn
from torch.autograd import Variable

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int = 512, dropout: float = .1, max_len: int = 5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Compute the positional encodings in log space
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(torch.log(torch.Tensor([10000.0])) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)


マルチヘッド注意

トランスフォーマーの前は、シーケンスから学習するAI研究のパラダイムは、畳み込み(WaveNet, ByteNet)またはリカレント((RNN, LSTM)のいずれかを使用することでした。トランスフォーマーの前にいくつかのNLPブレークスルーに注意が向けられましたが(Luong、2015)、当時は、畳み込みや再発なしに効果的なモデルを構築できることは明らかではありませんでした。したがって、「注意が必要なすべてである」という提案は非常に大胆な発言でした。

注意レイヤーは、クエリ(QQ)とキー(KK)値(VV) のペアのセット間のマッピングを学習できます。これらの名前は特定のNLPアプリケーションに依存するため、これらの名前の意味は混乱を招く可能性があります。テキスト生成のコンテキストでは、クエリは入力の埋め込みです。値とキーはターゲットとして見ることができます。通常、値とキーは同じです。

たとえば、Youtubeで動画を検索する場合は、検索バーにフレーズを入力します(つまり、クエリQQ)。検索エンジンはこれを使用して、Youtubeのビデオタイトルや説明など(つまり、キーKK)に対してマッピングします。このマッピングを使用すると、最も関連性の高い動画が提案されます(つまり、値VV)。(ソース例

NLPでの注意のパフォーマンスを向上させた革新の1つは、著者が「スケーリングされた内積注意」と呼んでいるものです。これは乗法的注意と同じですが、さらに QQKK のマッピングはキー次元dkd_kによってスケーリングされます。これにより、より大きな次元で内積注意のパフォーマンスが向上します。結果は (softmax(xi)=exp⁡(xi)∑jexp(xj)\text{softmax}(x_i) = \frac{\exp(x_i)}{\sum_{j} exp(x_j)})アクティベーション。

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{d_k}} )V

import torch
from torch import nn
class Attention:
    def __init__(self, dropout: float = 0.):
        super(Attention, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, query, key, value, mask=None):
        d_k = query.size(-1)
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        p_attn = self.dropout(self.softmax(scores))
        return torch.matmul(p_attn, value)
    
    def __call__(self, query, key, value, mask=None):
        return self.forward(query, key, value, mask)

デコーダーでは、注意サブレイヤーは、特定の位置を非常に大きな負の数((−1e9-1\text{e}9 または場合によっては−inf⁡-\inf)で埋めることによってマスクされます。これは、後続のポジションに参加することでモデルが不正行為をするのを防ぐためです。これにより、モデルは次のトークンを予測しようとしたときに、前の位置の単語にのみ対応できるようになります。

注意のメカニズム自体はすでに非常に強力であり、行列の乗算用に最適化されたGPUやTPUなどの最新のハードウェアで効率的に計算できます。ただし、単一の注意レイヤーでは、1つの表現しか使用できません。したがって、トランスでは複数の注意ヘッドが使用されます。これにより、モデルは複数のパターンと表現を学習できます。この論文では、連結されたh=8h = 8 の注意層を使用しています。最終的な式は次のようになります。

MultiHead(Q,K,V)=concat(head1,…,headn)WO\text{MultiHead}(Q, K, V) = \text{concat}(\text{head}_1, \dots, \text{head}_n) W^O

where headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

投影重み(WO∈Rhdv×dmodel,WiQ∈Rdmodel×dk,WiK∈Rdmodel×dk,WiV∈Rdmodel×dvW^O \isin \mathbb{R}^{hd{v} \times d\text{model}}, Wi^Q \isin \mathbb{R}^{d\text{model} \times dk}, W_i^K \isin \mathbb{R}^{d\text{model} \times dk}, W_i^V \isin \mathbb{R}^{d\text{model} \times d_v}WO∈Rhdv​×dmodel​,WiQ​∈Rdmodel​×dk​,WiK​∈Rdmodel​×dk​,WiV​∈Rdmodel​×dv​)は、完全に接続された(線形)レイヤーの出力です。トランスフォーマー論文の著者はをdk​=dv​=hdmodel​​=64使用しています。

from torch import nn
from copy import deepcopy
class MultiHeadAttention(nn.Module):
    def __init__(self, h: int = 8, d_model: int = 512, dropout: float = 0.1):
        super(MultiHeadAttention, self).__init__()
        self.d_k = d_model // h
        self.h = h
        self.attn = Attention(dropout)
        self.lindim = (d_model, d_model)
        self.linears = nn.ModuleList([deepcopy(nn.Linear(*self.lindim)) for _ in range(4)])
        self.final_linear = nn.Linear(*self.lindim, bias=False)
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        
        query, key, value = [l(x).view(query.size(0), -1, self.h, self.d_k).transpose(1, 2) \
                             for l, x in zip(self.linears, (query, key, value))]
        nbatches = query.size(0)
        x = self.attn(query, key, value, mask=mask)
        
        # Concatenate and multiply by W^O
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.final_linear(x)

*技術的注意事項:.transposeは、基になるメモリストレージを元のテンソルと共有するため、.transposeの後に.contiguousメソッドが追加されます。その後、.viewを呼び出すには、連続したテンソル(ドキュメント)が必要です。.viewメソッドを使用すると、効率的な再形成、スライス、および要素ごとの操作(ドキュメント)が可能になります。

各ヘッドの次元はhhで除算されるため、全体の計算は、完全な次元を持つ1つの注意ヘッド(dmodeld_\text{model})を使用する場合と同様です。ただし、このアプローチでは、計算をヘッドに沿って並列化できるため、最新のハードウェアで大幅な高速化が実現します。これは、畳み込みや再発のない効果的な言語モデルのトレーニングを可能にしたイノベーションの1つです。



残差とレイヤーの正規

AI研究コミュニティは、残差接続や(バッチ)正規化などの概念がパフォーマンスを向上させ、トレーニング時間を短縮し、より深いネットワークのトレーニングを可能にすることを発見しました。 したがって、トランスには、すべての注意層とすべてのフィードフォワード層の後に、残差接続と正規化が装備されています。さらに、一般化を改善するために、各レイヤーにドロップアウトが追加されています。

正規化

最新の深層学習ベースのコンピュータービジョンモデルは、多くの場合、バッチ正規化を備えています。ただし、このタイプの正規化は大きなバッチサイズに依存しており、自然に再発することはありません。 従来のトランスフォーマーアーキテクチャには、代わりにレイヤーの正規化があります。 レイヤーの正規化は、バッチサイズが小さい場合($\text{batch size} < 8)でも安定しています。

層の正規化を計算するために、最初にミニバッチ内のサンプルごとに平均μi\mu_iと標準偏差σi\sigma_i を個別に計算します。

μi=1K∑k=1kxi,k\mu_i = \frac{1}{K} \sum_{k=1}^{k} x_{i, k}

σi=sqrt(1K∑k=1k(xi,k−μi)2)\sigma_i = \text{sqrt}\left(\frac{1}{K} \sum_{k=1}^{k} (x_{i, k} - \mu_i)^2\right)

次に、正規化ステップは次のように定義されます。

LNγ,β(xi)≡γx−μiσi+ϵ+βLN_{\gamma, \beta}(x_i) \equiv \gamma \frac{x - \mu_i}{\sigma_i + \epsilon} + \beta

ここで、γ\gammaβ\betaは学習可能なパラメーターです。標準偏差ϵ\epsilon00の場合、数値安定性のために小さな数σi\sigma_iが追加されます。

from torch import nn
class LayerNorm(nn.Module):
    def __init__(self, features: int, eps: float = 1e-6):
        super(LayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

残差

残差接続とは、ネットワーク内の前のレイヤー(つまりサブレイヤー)の出力を現在のレイヤーの出力に追加することを意味します。ネットワークは本質的に特定のレイヤーを「スキップ」できるため、これにより非常に深いネットワークが可能になります。

各レイヤーの最終出力は、ResidualConnection(x)=x+Dropout(SubLayer(LayerNorm(x)))\text{ResidualConnection}(x) = x + \text{Dropout}(\text{SubLayer}(\text{LayerNorm}(x)))

from torch import nn
class ResidualConnection(nn.Module):
    def __init__(self, size: int = 512, dropout: float = .1):
        super(ResidualConnection,  self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))


フィードフォワード

すべての注意レイヤーの上に、フィードフォワードネットワークが追加されます。これは、ReLU\text{ReLU} アクティベーション(ReLU(x)=max⁡(0,x)\text{ReLU}(x) = \max(0, x)) と内部レイヤーのドロップアウトを備えた2つの完全に接続されたレイヤーで構成されます。トランスフォーマー論文で使用される標準寸法は、入力層の場合はdmodel = 512、内層の場合はdff=2048d_{ff} = 2048です。

完全な計算は、FeedForward(x)=W2max(0,xW1+B1)+B2\text{FeedForward}(x) = W_2 max(0, xW_1 + B_1) + B_2となります。

PyTorch Linear には、デフォルトでバイアス(B1B_1 およびB2B_2)がすでに含まれていることに注意してください。

from torch import nn
class FeedForward(nn.Module):
    def __init__(self, d_model: int = 512, d_ff: int = 2048, dropout: float = .1):
        super(FeedForward, self).__init__()
        self.l1 = nn.Linear(d_model, d_ff)
        self.l2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.l2(self.dropout(self.relu(self.l1(x))))


エンコーダー‐デコーダー

エンコーディング

これで、モデルエンコーダーとデコーダーを構築するためのすべてのコンポーネントができました。単一のエンコーダー層は、マルチヘッド注意層とそれに続くフィードフォワードネットワークで構成されます。前述のように、残差接続とレイヤーの正規化も含まれます。

Encoding(x,mask)=FeedForward(MultiHeadAttention(x))\text{Encoding}(x, \text{mask}) = \text{FeedForward}(\text{MultiHeadAttention}(x))

from torch import nn
from copy import deepcopy
class EncoderLayer(nn.Module):
    def __init__(self, size: int, self_attn: MultiHeadAttention, feed_forward: FeedForward, dropout: float = .1):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sub1 = ResidualConnection(size, dropout)
        self.sub2 = ResidualConnection(size, dropout)
        self.size = size

    def forward(self, x, mask):
        x = self.sub1(x, lambda x: self.self_attn(x, x, x, mask))
        return self.sub2(x, self.feed_forward)

論文の最終的なトランスフォーマーエンコーダーは、6つの同一のエンコーダーレイヤーとそれに続くレイヤーの正規化で構成されています。

class Encoder(nn.Module):
    def __init__(self, layer, n: int = 6):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(n)])
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

デコーディング

復号化層は、マスクされたマルチヘッド注意層であり、その後にメモリを含むマルチヘッド注意層が続く。メモリはエンコーダーからの出力です。最後に、フィードフォワードネットワークを経由します。この場合も、これらすべてのコンポーネントには、残差接続とレイヤーの正規化が含まれます。

Decoding(x,memory,mask1,mask2)=FeedForward(MultiHeadAttention(MultiHeadAttention(x,mask1),memory,mask2))\text{Decoding}(x, \text{memory}, \text{mask}_1, \text{mask}_2) = \text{FeedForward}(\text{MultiHeadAttention}(\text{MultiHeadAttention}(x, \text{mask}_1), \text{memory}, \text{mask}_2))

from torch import nn
from copy import deepcopy
class DecoderLayer(nn.Module):
    def __init__(self, size: int, self_attn: MultiHeadAttention, src_attn: MultiHeadAttention, 
                 feed_forward: FeedForward, dropout: float = .1):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sub1 = ResidualConnection(size, dropout)
        self.sub2 = ResidualConnection(size, dropout)
        self.sub3 = ResidualConnection(size, dropout)
 
    def forward(self, x, memory, src_mask, tgt_mask):
        x = self.sub1(x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sub2(x, lambda x: self.src_attn(x, memory, memory, src_mask))
        return self.sub3(x, self.feed_forward)

最終的なエンコーダーと同様に、この論文のデコーダーにも6つの同一のレイヤーがあり、その後にレイヤーの正規化が続きます。

class Decoder(nn.Module):
    def __init__(self, layer: DecoderLayer, n: int = 6):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList([deepcopy(layer) for _ in range(n)])
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)


エンコーダーとデコーダーのこの高レベルの表現を使用すると、最終的なエンコーダー-デコーダーブロックを簡単に作成できます。

from torch import nn
class EncoderDecoder(nn.Module):
    def __init__(self, encoder: Encoder, decoder: Decoder, 
                 src_embed: Embed, tgt_embed: Embed, final_layer: Output):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.final_layer = final_layer
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        return self.final_layer(self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask))
    
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
    
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)


最終出力

最後に、デコーダーからのベクトル出力を最終出力に変換する必要があります。言語翻訳のようなシーケンス間の問題の場合、これは各位置の総語彙にわたる確率分布です。1つの完全に接続されたレイヤーは、デコーダーの出力を、ターゲット語彙の次元を持つロジットのマトリックスに変換します。これらの数値は、ソフトマックス活性化関数を介して語彙全体の確率分布に変換されます。コードではLogSoftmax(xi)=log⁡(exp⁡(xi)∑jexp⁡(xj))\text{LogSoftmax}(x_i) = \log(\frac{\exp(x_i)}{\sum_{j} \exp(x_j)})を使用しています。これは、それが高速で、数値的に安定しているためです。

たとえば、翻訳された文に2020トークンがあり、語彙の合計が3000030000トークンであるとします。結果の出力は、行列M∈R20×30000M \isin \mathbb{R}^{20 \times 30000}になります。次に、最後の次元でargmaxを取得して、トークナイザーを介してテキストの文字列にデコードできる出力トークンT∈N20T \isin \mathbb{N}^{20}のベクトルを取得できます。

Output(x)=LogSoftmax(max(0,xW1+B1))Output(x) = LogSoftmax(max(0, xW_1 + B_1))

from torch import nn
class Output(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        super(Output, self).__init__()
        self.l1 = nn.Linear(input_dim, output_dim)
        self.log_softmax = nn.LogSoftmax(dim=-1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logits = self.l1(x)
        return self.log_softmax(logits)


モデルの初期化

論文と同じ寸法のトランスフォーマーモデルを作成します。初期化戦略はXavier/Glorot初期化であり、[−1n,1n][-\frac{1}{\sqrt{n}}, \frac{1}{\sqrt{n}}]の範囲の一様分布から選択することで構成されます。すべてのバイアスは00で初期化されます。

Xavier(W)∼U[−1n,1n],B=0\text{Xavier}(W) \sim U[-\frac{1}{\sqrt{n}}, \frac{1}{\sqrt{n}}], B = 0

from torch import nn
def make_model(input_vocab: int, output_vocab: int, d_model: int = 512):
    encoder = Encoder(EncoderLayer(d_model, MultiHeadAttention(), FeedForward()))
    decoder = Decoder(DecoderLayer(d_model, MultiHeadAttention(), MultiHeadAttention(), FeedForward()))
    input_embed= nn.Sequential(Embed(vocab=input_vocab), PositionalEncoding())
    output_embed = nn.Sequential(Embed(vocab=output_vocab), PositionalEncoding())
    output = Output(input_dim=d_model, output_dim=output_vocab)
    model = EncoderDecoder(encoder, decoder, input_embed, output_embed, output)
    
    # Initialize parameters with Xavier uniform 
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

この関数は、シーケンス間の問題についてトレーニングできるPyTorchモデルを返します。以下に、トークン化された入力と出力でそれを使用する方法のダミーの例を示します。入力と出力にわずか 1010語の語彙があるとしましょう。

# Tokenized symbols for source and target.
>>> src = torch.tensor([[1, 2, 3, 4, 5]])
>>> src_mask = torch.tensor([[1, 1, 1, 1, 1]])
>>> tgt = torch.tensor([[6, 7, 8, 0, 0]])
>>> tgt_mask = torch.tensor([[1, 1, 1, 0, 0]])

# Create PyTorch model
>>> model = make_model(input_vocab=10, output_vocab=10)
# Do inference and take tokens with highest probability through argmax along the vocabulary axis (-1)
>>> result = model(src, tgt, src_mask, tgt_mask)
>>> result.argmax(dim=-1)
tensor([[6, 6, 4, 3, 6]])

この時点でモデルの重みが均一に初期化されているため、出力はターゲットから非常に離れています。これらのトランスフォーマーモデルを最初からトレーニングするには、かなりの計算が必要です。基本モデルをトレーニングするために、元の論文の著者は8つのNVIDIA P100GPUで12時間トレーニングしました。彼らのより大きなモデルは、8つのGPUでトレーニングするのに3.5日かかりました!事前にトレーニングされたトランスモデルを使用して、アプリケーションに合わせて微調整することをお勧めします。 HuggingFace Transformers libraryには、微調整のための事前トレーニング済みモデルがすでに多数あります。

トレーニング手順を最初からコーディングする方法について詳しく知りたい場合は、 training section of The Annotated Transformerを確認することをお勧めします。



それで全部です!トランスフォーマーアーキテクチャについての詳しい研究を楽しんだと思います!

ご質問やご意見がございましたら、以下にコメントし���ください。Twitterの @carlolepelaarsで連絡することもできます。



Iterate on AI agents and models faster. Try Weights & Biases today.