서론

ELMo, ULMFit와 같은 언어 모델의 혁신 이후, 트랜스포머는 자연어 처리 분야(Natural Language Processing (NLP))에서 대성공을 거두었습니다. 이는 GPT-3 및 DALL-E과 같은 유명한 언어 모델의 기반이며, 또한 HuggingFace Transformers library과 같은 툴은 머신 러닝 엔지니어가 광범위한 NLP 작업을 쉽게 해결할 수 있도록 해줬으며, 이후 NLP에서의 수많은 혁신을 촉진했습니다. 이 리포트에서는"Attention Is All You Need" (2017)논문에 설명된 바와 같이 트랜스포머 아키텍처에 대한 심층 분석 및 PyTorch를 사용하여 이를 코딩해 보겠습니다.

트랜스포머는 시퀀스-투-시퀀스(seq2seq) 모델입니다. 즉, 데이터에 순서가 있고, 출력 그 자체가 시퀀스인 모든 문제에 적합합니다. 적용 예로는 기계 번역, 추상적 요약(abstractive summarization), 음성 인식 등이 있습니다. 최근 Vision Transformers (ViT) 는 심지어 컴퓨터 비전(computer vision)에서의 최첨단 기술까지 향상시켰습니다.

아래에서 전체 트랜스포머 아키텍처의 시각화를 확인하실 수 있습니다. 저희는 모든 구성 요소가 무엇을 하는지, 왜 그곳에 있는지, 어떻게 모든 것이 서로 잘 맞물려 작동하는지 설명해드리겠습니다. 우선, 둘 내부에 여러 신경망 계층(neural network layer)가 있는 인코더(좌측)와 디코더(우측)가 있음을 확인하시기 바랍니다.

코드 예시는 Harvard NLP Group의 The Annotated Transformer 및 트랜스포머에 대한 PyTorch 문서에서 리팩터링(refactor) 되었습니다.

토큰화(Tokenization)

우선, 연산을 수행하기 위해 텍스트를 숫자로 나타낼 방법이 필요합니다. 토큰화는 텍스트 문자열을 압축된 기호의 시퀀스로 구문 분석하는 프로세스입니다. 이 프로세스에서 각 정수가 텍스트의 일부를 나타내는 정수의 벡터가 생성됩니다. 이 트랜스포머 논문에서는 토큰화 방법으로 바이트 페어 인코딩 (BPE)를 사용합니다. BPE는 가장 일반적인 연속 바이

최근 연구에 따르면 BPE는 차선적이며 BERT와 같은 최근의 언어 모델은 WordPiece 토크나이저(tokenizer)를 대신 사용합니다. WordPiece는 흔히 전체 단어를 하나의 토큰으로 토큰화하므로 디코딩이 더 쉽고 직관적으로 보입니다. 대조적으로 BPE는 종종 단어의 조각을 토큰화합니다. 이는 잘못된 토큰화 기호 및 개별 글자 (즉, 알파벳 문자)가 알 수 없는 토큰으로 토큰화되는 상황으로 이어질 수 있습니다. 그러므로, 바로 저희가 예시에서 사용하는 토크나이저는 WordPiece입니다. 저희는 저희만의 토크나이저를 텍스트의 코퍼스에서 훈련시킬 수 있습니다만, 현실에서 실무자들은 거의 대부분 사전 훈련된(pre-trained) 토크나이저를 사용합니다. 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$) 및 알 수 없는 기호 ([UNK]=100) 등이 있습니다.

임베딩(Embeddings)

텍스트의 적절한 표현을 학습하기 위해, 시퀀스의 각 개별 토큰은 임베딩을 통해 벡터로 변환되며, 신경망 계층의 한 종류로 볼 수 있습니다. 임베딩에 대한 가중치는 트랜스포머 모델의 나머지 부분과 함께 학습되기 때문입니다. 이는 어휘(vocabulary)의 각 단어 대한 벡터를 포함하고 있으며, 이러한 가중치는 정규 분포 N(0,1)에서 초기화됩니다. 모델 ($E\in {\mathbb{R}}^{vocab\times d_{model}}$)을 초기화할 때 어휘 ($vocab$)의 크기 및 모델 ($d_{model}=512$)의 차원(dimension)을 지정해야 합니다. 마지막으로 정규화(normalization) 단계로 가중치에 $\sqrt{d_{model}}$ 을 곱합니다.

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

포지셔널 인코딩(Positional Encoding)

순환 및 합성곱 신경망과는 대조적으로, 모델 자체는 시퀀스에 임베드된 토큰의 상대 위치(relative position)에 대한 정보를 가지고 있지 않습니다. 따라서 인코더와 디코더에 대한 입력 임베딩에 인코딩을 추가함으로써 이 정보를 입력해야 합니다. 이 정보는 다양한 방법으로 추가할 수 있으며 정적이거나 학습될 수 있습니다. 트랜스포머는 각 위치 ($pos$)에 대한 사인 및 코사인 트랜스포메이션을 사용합니다. 사인은 짝수 차원 ($2_i$) 에서 사용되며, 코사인은 홀수 차원 ($2_{i+1}$)에 사용됩니다.

$P{E}_{pos,2i}=sin(\frac{pos}{10000^{2i/d_{model}}})$

$P{E}_{pos,2i+1}=cos(\frac{pos}{10000^{2i/d_{model}}})$

이 코드에서, 포지셔닝 인코딩은 숫자 오버 플로우 (numerical overflow) 방지를 위해 로그 공간(log space)에서 연산됩니다.

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)

다중 헤드 어텐션(Multi-Head Attention)

트랜스포머 이전에, 시퀀스에서 학습하기 위한 AI 연구의 패러다임은 합성곱(convolutions) (WaveNet, ByteNet) 또는 순환(recurrence) (RNN, LSTM) 중 하나를 사용하는 것이었습니다. 어텐션(attention)은 트랜스포머 (Luong et al., 2015) 이전에 이미 몇몇 NLP 성과를 이뤘으나, 그 당시에는 합성곱이나 순환 없이 효율적인 모델을 구축할 수 있는지가 불분명했습니다. 따라서 "Attention is All You Need"(어텐션만 있으면 됩니다)는 제안은 꽤 과감한 표현이었습니다.

어텐션 레이어(attention layer)는 쿼리($Q$) 와 키($K$) 값 ($V$)쌍 간의 맵핑을 학습할 수 있습니다. 이러한 이름의 의미는 특정 NLP 응용 프로그램에 따라 달라지므로 헷갈릴 수 있습니다. 텍스트 생성의 맥락에서, 쿼리(query)는 입력의 임베딩이며, 값(value)과 키(key)는 타깃으로 볼 수 있습니다. 일반적으로 값과 키는 동일합니다.

예를 들어, YouTube에서 동영상을 검색하는 경우, 검색창 (즉, 쿼$K$)를 맵핑합니다. 이 맵핑을 사용하면 가장 관련성이 높은 동영상 (즉, 값 $V$)를 제안합니다. (예시 자료)

저자들이 “Scaled dot-product attention”라 부르는 것은 NLP에서 어텐션의 퍼포먼스를 높인 하나의 획기적인 것입니다. 이는 곱셈 어텐션multiplicative attention과 동일하나, 추가적으로 Q 와 K 맵핑은 키 차원 $d_k$에 의해 크기가 조정(scale) 됩니다. 이를 통해서 곱셈 어텐션은 더 큰 차원에서 더 나은 퍼포먼스를 보여줍니다. 결과는 소프트맥스 활성화 (softmax activation) (softmax(xi)=exp⁡(xi)∑jexp(xj)softmax(xi) = \frac{\exp(x_i)}{\sum{j} exp(x_j)}softmax(xi​)=∑j​exp(xj​)exp(xi​)​)를 통과하며, $V$를 곱합니다.

$Attention(Q,K,V)=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)

디코더에서 어텐션 서브 레이어(attention sub-layer)는 특정 위치를 매우 큰 음수 ($-1e9$ or sometimes even $-inf$ 또는 때로는 −inf)로 채움으로써 마스킹(mask) 됩니다. 이는 후속 위치를 처리함으로써 모델이 부정행위(cheating)을 저지르는 것을 예방하기 위한 것입니다. 이를 통해 모델은 다음 토큰을 예측하려 할 때 이전 위치의 단어에만 주의를 기울일 수 있습니다.

이러한 어텐션의 메커니즘 그 자체는 이미 아주 효과적이며 행렬 곱셈(matrix multiplication)에 최적화된 GPU 및 TPU과 같은 최신 하드웨어에서 효율적으로 연산될 수 있습니다. 하지만 단일 어텐션 레이어는 하나의 표현만을 허용합니다. 따라서 트랜스포머에서 다중 어텐션 헤드(multiple attention heads)가 사용됩니다. 이를 통해 모델은 다중 패턴 및 표현을 학습할 수 있습니다. 논문에서는 ���합된(concatenated) $h=8$ 어텐션 레이어를 사용합니다. 최종 공식은 다음과 같습니다:

$MultiHead(Q,K,V)=Concat(hea{d}_1,...,hea{d}_n)W^O$

where $hea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$

예상 가중치 ( $W^O \isin \mathbb{R}^{hd{v} \times d{model}​×dk​,WiK​∈Rdmodel​×dk​,WiV​∈Rdmodel​×dv​)는 완전히 연결된 (선형(linear) 레이어의 출력입니다. 트랜스포머 논문의 저자는 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는 기본 메모리 저장소를 원래의 tensor와 공유하기 때문에 .contiguous 방법은 .transpose 다음에 추가됩니다. 이 후 .view를 호출하려면 인접한(contiguous) tensor가 필요합니다 (문서). .view 방법은 효율적인 변환(reshaping), 슬라이싱(slicing) 및 요소별 작업을 수행할 수 있습니다. (문서)

  각 헤드(head)의 차원을 $h$,로 나누기 때문에, 총 연산은 완전한 차원(full dimensionality)를 가진 하나의 어텐션 헤드(attention head)를 사용하는 것과 유사합니다. 그러나, 이 접근 방식을 통한 연산은 헤드를 따라 병렬화(parallelize) 될 수 있으며 이를 통해 최신 하드웨어에서 속도가 크게 향상됩니다. 이는 합성곱 또는 순환(recurrence) 없이 효과적인 언어 모델을 훈련할 수 있게끔 하는 혁신적인 부분 중 하나입니다

레지듀얼 및 레이어 정규화(Residuals and Layer Normalization)

AI 연구 커뮤니티에서는 레지듀얼 커넥션(residual connections)및 (배치) 정규화와 같은 개념이 퍼포먼스를 향상시키고, 훈련 시간을 단축하며, 보다 심층적인 네트워크의 훈련을 가능케 한다는 사실을 발견했습니다. 따라서, 모든 어텐션 레이어 및 모든 피드 포워드 레이어(feed forward layer) 다음에 트랜스포머는 레지듀얼 커넥션 및 정규화를 갖추고 있습니다. 추가적으로, 더 나은 일반화(generalization)를 위해 각 레이어에 드롭아웃(dropout)이 추가됩니다.

정규화(Normalization)

현대 딥러닝 기반 컴퓨터 비전 모델은 보통 배치 정규화를 포함하고 있습니다. 그러나 이러한 정규화 유형은 큰 배치 사이즈에 의해 좌우되며, 당연하게도 순환(recurrence)에 적합하지 않습니다. 기존 트랜스포머 아키텍처는 레이어 정규화를 대신 갖추고 있습니다. 레이어 정규화는 배치 크기가 작더라도 ($batchsize<8$) 안정적입니다. 레이어 정규화를연산하기 위해, 우선 저희는 미니배치의 각 샘플에 대한 평균 ${\mu }_i$와 표준 편차 ${\sigma }_i$를 별도로 계산합니다.

${\mu }_i=\frac{1}{K}{\sum }_{k=1}^k{x}_{i,k}$

${\sigma }_i=\sqrt{\frac{1}{K}{\sum }_{k=1}^k(x_{i,k}-{\mu }_i{)}^2}$

그후, 정규화 과정은 다음과 같이 정의됩니다:

$L{N}_{\gamma ,\beta }(x_i)\equiv \gamma \frac{x-{\mu }_i}{{\sigma }_i+ϵ}+\beta$

여기서 $\gamma$와 $\beta$ 는 학습 가능한 매개변수입니다. 표준 편차 $ϵ$i가 $0$일 경우 수치 안정성(numerical stability)을 위해 작은 수 ${\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

레지듀얼

레지듀얼 커넥션은 네트워크에서 이전 레이어의 출력을 (즉, 하위 레이어) 현재 레이어의 출력에 추가하는 것을 의미합니다. 네트워크는 본질적으로 특정 레이어를 ‘건너뛰기’할 수 있으므로, 이는 매우 깊은 네트워크(very deep network)를 허용합니다.

그 다음, 각 레이어의 최종 출력은 다음과 같습니다: $ResidualConnection(x)=x+Dropout(SubLayer(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)))

피드 포워드(Feed Forward)

모든 어텐션 레이어의 위에 피드 포워드 네트워크가 추가됩니다. 이 레이어는 $ReLU$ 활성화(activation) ($ReLU(x)=max(0,x)$)) 와 내부 레이어 드롭아웃을 통해 완전히 연결된 (fully-connected) 레이어로 구성됩니다. 트랜스포머 논문에서 사용한 표준 차원은 입력 레이어의 경우 $d_{model}=512$, 내부 레이어의 경우 $d_{ff}=2048$입니다.

전체 계산은 다음과 같습니다: $FeedForward(x)=W_{2}max(0,x{W}_1+B_1)+B_2$.

참고: PyTorch Linear에는 기본값으로 편향 ($B_1$ and $B_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))$

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)

디코딩

디코딩 레이어는 메모리를 포함한 다중 헤드 어텐션 레이어가 뒤따르는 마스킹된 (maksed) 다중 헤드 어텐션 레이어입니다. 메모리는 인코더의 출력입니다. 마지막으로 피드 포워드 네트워크를 통과합니다. 이 모든 구성요소는 레지듀얼 커넥션 및 레이어 정규화를 포함합니다.

$Decoding(x,memory,mask1,mask2)=FeedForward(MultiHeadAttention(MultiHeadAttention(x,mask1),memory,mask2))$

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)

최종 출력

마지막으로, 디코더의 벡터 출력은 최종 출력으로 변환되어야 합니다. 언어 번역과 같은 시퀀스-투-시퀀스 문제의 경우 이는 각 위치(position)에 대한 총 어휘에 관한 확률 분포입니다. 완전히 연결된 하나의 레이어는 디코더 출력을 logits의 행렬로 변환되며, 이는 타깃 어휘의 차원을 갖고 있습니다. 이러한 숫자는 softmax 활성화 함수를 통해 어휘에 대한 확률 분포로 변환됩니다. 코드에서 저희는 더 빠르고 수치적으로 더 안정적이기 때문에$LogSoftmax(x_i)=log(\frac{\exp (x_i)}{{\sum }_{j}exp(x_j)})$를 사용합니다.

예를 들어, 번역된 문장이 $20$개의 토큰을 갖고 있고, 총 어휘는 $30000$개의 토큰이라고 가정해봅시다. 그러면 결과 출력은 행렬 matrix $M\in {\mathbb{R}}^{20\times 30000}$이 됩니다. 그런 다음 마지막 차원에 대한 $argmax$를 취하여 토크나이저를 통해 텍스트 열로 디코딩 할 수 있는 출력 토큰 $T\in {\mathbb{N}}^{20}$ 의 벡터를 얻을 수 있습니다.

$Output(x)=LogSoftmax(max(0,x{W}_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)

모델 초기화(Model initialization)

저희는 논문에서의 경우와 같이 동일한 차원으로 트랜스포머 모델을 구축합니다. 초기화 전략(initialization strategy)는 Xavier/Glorot 초기화이며 이는 $[-\frac{1}{\sqrt{n}},\frac{1}{\sqrt{n}}]$ 범위의 균일 분포에서 선택하도록 구성됩니다. 모든 편향은 $0$으로 초기화됩니다.

$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 모델을 반환합니다. 아래에서 토큰화된 입력 및 출력과 함께 이를 사용하는 방법에 대한 더미 예시가 설명되어 있습니다. 입력과 출력에 대한 어휘가 $10$개만 있다고 가정해봅시다.

# 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]])

이 시점에서 모델은 균일하게 초기화된 가중치를 갖고 있으므로, 출력은 타깃과의 거리는 상당히 멉니다. 이러한 트랜스포머 모델을 처음부터 훈련하는 것은 상당한 계산을 필요로 합니다. 기본 모델을 훈련하기 위해, 저자는 논문에서 12시간 동안 8개의 NVIDIA P100 GPU를 훈련시켰습니다. 더 큰 모델은 8개의 GPU를 훈련하는 데 3.5일을 소요했습니다! 사전 훈련된 트랜스포머 모델 사용 및 응용 프로그램에 적합하도록 미세 조정하시기 바랍니다. HuggingFace Transformers library는 이미 미세 조정을 위한 사전 훈련된 모델을 많이 가지고 있습니다.

훈련 절차를 처음부터 코딩하는 방법에 대해서 더 자세히 알고 싶으시다면 The Annotated Transformer 훈련 섹션을 확인하시기 바랍니다.

자, 이제 끝입니다! 트랜스포머 아키텍처에 대한 이 심층 분석이 여러분께 도움이 되었기를 바랍니다.

다른 문의 사항 및 피드백이 있으시다면, 아래에 의견을 남겨주시거나 트위터 @carlolepelaars로 연락해 주시기 바랍니다.

출처 및 추가 학습 자료

추천 논문

Attention Is All You Need (2017)

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (2018)

Reformer: The Efficient Transformer (2020)

Linformer: Self-Attention with Linear Complexity (2020)

Longformer: The Long-Document Transformer (2020)

Language Models are Few-Shot Learners (GPT-3 paper)(2020)

온라인 자료

The Annotated transformer with PyTorch Code

The Illustrated transformer

The Narrated transformer video

Łukasz Kaiser's masterclass on transformers