[Pytorch] Transformer 구현하기

Contents

 

 

Transformer는 자연어 처리(Natural Language Processing, NLP) 분야에서 성능이 우수한 것으로 검증된 대표적인 Self-Attention 모델이다. 기존 Sequential Based 모델인 RNN, LSTM, GRU 등과 같은 모델에 비해 매우 우수한 성능을 보여주고 있으며, BERT, GPT 등의 사전 학습 기반 모델이 Transformer의 골조를 활용하였다. 이번 글에서는 Transformer의 Encoder와 Decoder가 어떻게 구성되어 있는지 짧은 코드를 통해 알아보자. 

 

Import Package

Transformer는 HuggingFace 에서 지원하는 transformers 패키지와 PyTorch를 이용하여 구현할 것이다. 따라서, 구현에 필요한 패키지를 불러오자.

from transformers import AutoTokenizer, AutoModel, AutoConfig
from torch import nn 
import torch 
from math import sqrt

Transformers 패키지에는 AutoTokenizer, AutoModel 등의 함수가 존재하는데, 이는 미리 사전학습된(Pre-Trained) 모델을 불러오기 위한 함수다. 아래와 같이 'bert-base-uncased' 모델을 호출하게 되면 BERT-Based 모델을 사용하는 것을 의미한다. 

model_name_or_path = 'bert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModel.from_pretrained(model_name_or_path) 
text = 'time files like an arrow'

이때 주위해야 될 부분은 tokenizer와 model을 사전학습된 모델로 사용하고자 하는 경우 동일한 모델의 구조를 가지고 와야 한다. 모델에 따라 조금씩 차이가 있을 수 있기 때문이다. 

 

Bert를 가지고 와서 보면 마지막 pooler 부분이 768 차원인 것을 확인할 수 있다. 이는 마지막 Classifier 즉, 예측 부분의 모델이 구현되지 않은 상태이므로, 분류 문제를 다루고 싶은 경우 이를 만들어 주어야 한다. 

 

Attention Networks

Transformer에서는 Dot-Product Attention이 있으며, 이를 여러 개 결합하여 Multi-Head Attention을 구성하고 있다. 따라서, 먼저 Dot-Product Attention의 구조부터 알아보자. 

Scaled Dot-Product Attention

Transformer의 입력으로는 Token Embedding과 Position Embedding의 합이 들어온다. 이 합을 각각 Query, Key, Value로 보고 아래와 같은 수식을 통해 Attention Score를 계산한다.

\[ \text{Attention}(Q, K, V) = \text{Softmax}(\frac{QK^{\top}}{\sqrt{d_k}}) \cdot V \]
수식을 살펴보면 Q와 K를 먼저 행렬 연산한 후 $\sqrt{d_k}$로 나누고 그 다음으로 V를 곱하고 있다. 다시 말하자면, Query와 Key의 차원의 수는 동일해야 하지만, Value의 차원은 달라도 된다는 것을 의미한다. $Q \in \mathbb{R}^{N \times d_{model}} $, $K \in \mathbb{R}^{N \times d_{model}}$, $V \in \mathbb{R}^{N \times d_{model}}$, $QK^{\top} \in \mathbb{R}^{N \times N}$ 이기 때문이다. 

def scaled_dot_product_attention(query, key, value):
    dim_k = query.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
    weights = F.softmax(scores, dim=-1)
    return torch.bmm(weights, value)

 

Multi-Head Attention

Multi-Head Attention은 Attention Network를 여러 개 결합한 것을 의미한다. 연결(Concatenation)한 것이며, 수식은 아래와 같이 된다.

\[ \begin{equation} \begin{split} \text{Multi-Head}(Q, K, V) & = \text{Concat}(h_1, \cdots, h_h) \cdot W^O \\ \\ \text{with } h_i & = \text{Attention} (QW^Q_i, KW^K_i, VW^V_i ) \end{split} \end{equation} \]

이때 $W^Q_i \in \mathbb{R}^{d_{model} \times d_k}$, $W^K_i \in \mathbb{R}^{d_{model} \times d_k}$, $W^V_i \in \mathbb{R}^{d_{model} \times d_v} $, $W^O \in \mathbb{R}^{h d_v \times d_{model}}$는 전부 학습 가능한 가중치 행렬을 의미한다. 본 논문에서는 연산 비용으로 인해 Multi-Head Attention인 경우에는 $d_{model}$을 전부 $d_k$ 차원으로 변환하고, $d_k=d_v=64$로 설정하여 차원을 축소하여 512 차원으로 만든 후 최종적으로 $W^O$를 곱하여 $d_{model}$의 원래 차원으로 되돌려준다.

class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super(AttentionHead, self).__init__()
        self.Q = nn.Linear(embed_dim, head_dim)
        self.K = nn.Linear(embed_dim, head_dim)
        self.V = nn.Linear(embed_dim, head_dim)

    def forward(self, hidden_state):
        outs = scaled_dot_product_attention(self.Q(hidden_state), self.K(hidden_state), self.V(hidden_state))
        return outs 

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super(MultiHeadAttention, self).__init__()
        embed_dim = config.hidden_size 
        num_heads = config.num_attention_heads
        head_dim = embed_dim // num_heads
        self.heads = nn.ModuleList(
            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]
        )
        self.fc_layer = nn.Linear(embed_dim, embed_dim)
    
    def forward(self, hidden_state):
        outs = torch.cat([h(hidden_state) for h in self.heads], dim=-1)
        outs = self.fc_layer(outs)
        return outs

 

 

Feed Forward Networks 

Feed Forward Network는 우리가 알고 있는 가장 간단한 Fully Connected Layer, 즉, Linear 모델을 생각하면 된다. 

\[ \text{FFN}(x) = \text{max}(0, xW_1 + b_1) \cdot W_2 + b_2 \]

$\text{max}(\cdot)$의 구조는 Transformer에서는 $GELU$를 사용하기 때문이다. 

class FeedForward(nn.Module):
    def __init__(self, config):
        super(FeedForward, self).__init__()
        self.fc_layer1 = nn.Linear(config.hidden_size, config.intermediate_size) # 768 X 3072
        self.fc_layer2 = nn.Linear(config.intermediate_size, config.hidden_size) # 3072 X 768
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x):
        outs = self.fc_layer1(x)
        outs = self.gelu(outs)
        outs = self.fc_layer2(outs)
        outs = self.dropout(outs)
        return outs

 

Layer Normalization and Skip Connection

Transformer는 Batch Normalization이 아닌 Layer Normalization과 Skip Connection을 사용한다. LayerNorm과 Skip Connection은 간단하게 구현이 가능하다. Skip Connection에 대한 개념은 ResNet을 참고하면 된다. 

class TransformerEncoderLayer(nn.Module):
    def __init__(self, config):
        super(TransformerEncoderLayer, self).__init__()
        self.layer_norm1 = nn.LayerNorm(config.hidden_size)
        self.layer_norm2 = nn.LayerNorm(config.hidden_size)
        self.attention = MultiHeadAttention(config)
        self.ffn = FeedForward(config)

    def forward(self, x):
        hidden_state = self.layer_norm1(x)
        outs = x + self.attention(hidden_state)
        outs = outs + self.ffn(self.layer_norm2(outs))
        return outs

 

 

Token Embedding and Position Embedding

Token Embedding은 실제 입력 문장의 각 단어에 대한 Embedding을 의미한다. 위에서 사용한 Attention Network를 사용하면 마지막에 Weight-Sum을 하기 때문에 각 단어의 Embedding에 대한 위치 정보가 사라진다. 따라서, 위치 정보를 기억하기 위해 Position Embedding을 추가한다. Position Embedding은 문장의 최대 길이만큼 단어의 순서에 따라 1, 2, 3, .. 순으로 증가하는 단순한 함수를 이용한다. 입력 문장의 길이가 유동적인 경우 사인 함수나 코사인 함수와 같이 주기를 가지는 함수를 사용하기도 한다. 

class Embeddings(nn.Module):
    def __init__(self, config):
        super(Embeddings, self).__init__()
        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout()

    def forward(self, input_ids):
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)
        token_embeddings = self.token_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)

        embeddings = token_embeddings + position_embeddings 
        embeddings = self.layer_norm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

 

최종적으로 도출되는 Encoder는 아래의 구조와 같다.

class TransformerEncoder(nn.Module):
    def __init__(self, config):
        super(TransformerEncoder, self).__init__()
        self.embeddings = Embeddings(config)
        self.layers = nn.ModuleList([TransformerEncoderLayer(config) for _ in range(config.num_hidden_layers)])
    
    def forward(self, x):
        x = self.embeddings(x)
        for layer in self.layers:
            x = layer(x)

        return x

 

Decoder

Transformer의 Decoder는 Encoder와 매우 유사하기 때문에 차이가 존재하는 부분에 대해서만 다룬다. Encoder와의 차이는 바로 Maksed Multi-Head Attention이다. 입력 문장의 일부를 Masking함으로써 실제 정답을 모르게 처리를 해 학습의 성능을 개선시키려는 것이 목적이다.  Transformer는 처음 기계 번역을 위해 제안된 모델이기 때문에 Input Embedding은 특정 언어, Output Embedding은 번역하고자 하는 언어를 사용한다. 예를 들어, 한국어를 영어로 번역하고 싶은 경우에는 Input Embedding은 한국어로, Output Embedding은 영어로 사용하는 것이다. 

 

 

seq_len = inputs.input_ids.size(-1)
mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0) # tril -> lower triangular matrix 

scores.masked_fill(mask == 0, -float('inf'))


# tensor([[[ 2.9408e+01,        -inf,        -inf,        -inf,        -inf],
#          [-1.6171e+00,  2.8352e+01,        -inf,        -inf,        -inf],
#          [ 2.1406e+00, -5.4225e-01,  2.9296e+01,        -inf,        -inf],
#          [-2.3453e-01,  1.1241e+00,  6.5819e-03,  2.6097e+01,        -inf],
#          [ 6.5022e-01,  7.8864e-01, -2.1337e+00,  1.3018e+00,  2.6512e+01]]],
#        grad_fn=<MaskedFillBackward0>

위와 같이 Masking을 하고 싶은 값에 -inf로 설정하여 참조하지 않도록 설정하여 사용된다. 

 

def scaled_dot_product_attention(query, key, value, mask=None):
    dim_k = query.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    weights = F.softmax(scores, dim=-1)
    return weights.bmm(value)

따라서, Encoder의 scaled dot product attention 구조에서 mask 부분만 추가하게 되면 동일하게 사용이 가능하다. Decoder의 나머지 골조는 Encoder와 동일하다. Transformer는 기계 번역 뿐만 아니라 감성 분석, 질의 응답 등 다양한 NLP 분야에 사용되는 것뿐만 아니라 최근에는 비전 분야에도 사용되고 그래프 분야에서도 사용되는 매우 획기적인 기술이다. 비전 분야에서는 ViT 가 가장 대표적이며, 물체 감지 분야에 사용되는 Swin Transformer 등이 전부 Transformer 기반이다. 그래프 분야에서는 Graphormer가 가장 대표적이고, TokenGT, PatchGT 등의 기술이 있다. 

 

추천 시스템에서는 Sequential한 정보를 활용하는 BERT4Rec 등의 모델이 있다. 이렇게 다양한 분야에 응용되는 Transformer의 모델을 언급하는 이유는 매우 유용한 기술이기 때문에 한 번쯤은 뜯어서 보기를 바라는 마음으로 작성하였다.

 

이로써, Transformer가 어떻게 작동되는지에 대해서 알아보았다. 종합적인 코드를 확인하고 싶은 경우에는 여기를 참고하면 설명과 함께 작성해두었다.

 

본 글은 '트랜스포머를 활용한 자연어 처리' 책의 기반으로 코드를 작성하였으며, 설명을 추가하였습니다. [Original Code]