Transformer 간단하게 코드와 함께 살펴보기

본 글을 통해 오랜만에 다시 한번 Transformer 구조를 이해해보고, 코딩으로 봤을 때 어떻게 보면 좋을 지를 정리해보고자 한다.

아키텍처

 
일단 기본적인 Encoder-Decoder 아키텍처를 보면 다음과 같다.
기존 논문에서는 기계 번역 모델로 사용하였습니다.
아래처럼 영어 문장을 프랑스어 문장으로 번역하도록 하였습니다.

간단하게 인코더와 디코더의 역할을 보면 다음과 같습니다.

인코더	입력 문장에서 특징을 추출
디코더	특징을 사용하여 출력 문장을 생성

Encoder

여러 개의 Encoder 블록으로 구성됩니다.
입력 문장은 Encoder 블록을 거치며 마지막 인코더 블록의 출력이 디코더의 입력 특징이 됩니다.

간단하게 코드를 가져오면 다음과 같다.
EncoderLayer를 n개 만큼 만들고, 같은 값을 계속 상속받아서 사용하는 것을 알 수 있다.

# Implementing the Encoder
class Encoder(Layer):
    def __init__(self, vocab_size, sequence_length, h, d_k, d_v, d_model, d_ff, n, rate, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        self.pos_encoding = PositionEmbeddingFixedWeights(sequence_length, vocab_size, d_model)
        self.dropout = Dropout(rate)
        self.encoder_layer = [EncoderLayer(h, d_k, d_v, d_model, d_ff, rate) for _ in range(n)]
 
    def call(self, input_sentence, padding_mask, training):
        # Generate the positional encoding
        pos_encoding_output = self.pos_encoding(input_sentence)
        # Expected output shape = (batch_size, sequence_length, d_model)
 
        # Add in a dropout layer
        x = self.dropout(pos_encoding_output, training=training)
 
        # Pass on the positional encoded values to each encoder layer
        for i, layer in enumerate(self.encoder_layer):
            x = layer(x, padding_mask, training)
 
        return x

 

Decoder

디코더 역시 다양하게 블록으로 구성되어있습니다.

아래 코드를 보면, decoder layer에 encoder_output과 x(encoder_output)과 mask가 추가된 것을 볼 수 있다.

class Decoder(Layer):
    def __init__(self, vocab_size, sequence_length, h, d_k, d_v, d_model, d_ff, n, rate, **kwargs):
        super(Decoder, self).__init__(**kwargs)
        self.pos_encoding = PositionEmbeddingFixedWeights(sequence_length, vocab_size, d_model)
        self.dropout = Dropout(rate)
        self.decoder_layer = [DecoderLayer(h, d_k, d_v, d_model, d_ff, rate) for _ in range(n)
        
    def call(self, output_target, encoder_output, lookahead_mask, padding_mask, training):
        # Generate the positional encoding
        pos_encoding_output = self.pos_encoding(output_target)
        # Expected output shape = (number of sentences, sequence_length, d_model)

        # Add in a dropout layer
        x = self.dropout(pos_encoding_output, training=training)

        # Pass on the positional encoded values to each encoder layer
        for i, layer in enumerate(self.decoder_layer):
            x = layer(x, encoder_output, lookahead_mask, padding_mask, training)

        return x

 
각 디코더 블록은 인코더에서 피처 정보를 받아서 사용합니다.
인코더와 디코더를 수직으로 그리면 전체 그림이 종이의 다이어그램처럼 보입니다.
 

다이어그램

이런식으로 간단하게 그리면 다음과 같다. 

이제 좀 더 자세하게 어떻게 데이터가 들어가는 지를 알아보자.
 

Input Embedding and Positional Encoding

Input Embedding

 
일단 다들 아시겠지만, 토큰화작업이 필요하다. 
그리고 토큰화된 문장은 고정 길이 시퀀스입니다.
예를 들어, 최대 길이가 200인 경우 모든 문장에는 trailing 패딩이 있는 200개의 토큰이 있으며 직관적으로 표시하면 다음과 같습니다.
이 단어 토큰들은 사전에서 인덱스(숫자로) 바꿔줍니다. 
 
(‘Hello’, ‘world’, ‘!’, <pad>, <pad>, …, <pad>)
-> (8667, 1362, 106, 0, 0, …, 0)
 
숫자 8667은 이 예에서 토큰 Hello에 해당합니다. 또한 <EOS>(문장 끝 표시)와 같은 특수 문자를 포함할 수도 있습니다.
토크나이저 및 어휘 데이터 세트에 따라 다릅니다. Hugging Face의 모델을 사용하는 경우 이러한 세부 사항을 처리하는 모델에 대한 특정 토크나이저가 있습니다.
 
처음부터 모델을 구축하는 경우
문장을 토큰화하고, 단어장을 설정하고, 각 토큰에 인덱스를 할당하는 방법을 결정해야 합니다.
 
이러한 토큰을 신경망에 공급하기 위해 각 토큰을 신경 기계 번역 및 기타 자연 언어 모델의 일반적인 관행인 임베딩 벡터로 변환합니다.
과거 논문에서는 이러한 임베딩을 위해 512차원 벡터를 사용합니다.
따라서 문장의 최대 길이가 200이면 모든 문장의 모양은 (200, 512)가 됩니다.

Tokenizer

토크나이저는 단지 1가지 목적을 가지고 있습니다. 즉, 입력된 텍스트를 모델에서 처리할 수 있는 데이터로 변환하는 것입니다. 모델은 숫자만 처리할 수 있으므로, 토크나이저는 텍스트 입력을 숫자 데이터로 변환해야 합니다. 
 
여기에는 다양한 방법들이 존재합니다.
 
각각의 방법들은 아래 참고 문서에 있는 것을 보고 확인하면 좋을 것 같습니다.
 

1. 단어 기반 토큰화 (Word-based Tokenization)

2. 문자 기반 토큰화 (Character-based Tokenization)

3. 하위 단어 토큰화 (Subword Tokenization)

• Byte-level BPE (GPT-2에 사용됨)
• WordPiece (BERT에 사용됨)
• SentencePiece, Unigram (몇몇 다국어 모델에 사용됨)

 

 
 
 

Positional Encoding

 
기존에 자주 사용하던 LSTM과 큰 차이가 있는 부분이 positional encoding이라는 부분이 있습니다.
 
모델이 반복 없이 단어 위치를 알 수 있도록 각 임베딩에 위치 인코딩을 주입합니다.
위치 인코딩에 대한 자세한 내용은 이 문서를 참조하십시오.
https://kikaben.com/transformers-positional-encoding/
간단하게 설명하면 아래와 같은 포지션 인코딩 함수를 갖게 되고, seq_length 가 128이고 position dimension을 512로 하면
아래와 같은 각각의 포지션에 대한 정보가 다르게 맵핑되는 것을 알 수 있습니다.
이를 통해 문장에서 각 토큰들의 위치를 부여하는 방식입니다.
 

여기서 일단 인풋은 여기서 단어가 숫자로 변화된 상황이고, 그것을 임베딩하여 만든 시퀀스 임베딩입니다.
그래서 각각의 시퀀스의 임베딩은 토큰의 의미와 포지션의 의미가 같이 담겨져 있습니다. 
인코더는 이러한 벡터에 대해 선형 대수 연산을 수행하여 전체 문장에서 각 토큰에 대한 컨텍스트를 추출하고 여러 인코더 블록을 통해 대상 작업에 유용한 정보로 임베딩 벡터를 풍부하게 합니다.
 

예시

간단하게 배치가 1인 경우를 가정해보자.
max seq length를 4라고 해보자. 
그러면 (1,4)라는 데이터가 처음에 들어갈 것이다.
그때 그러면 각각의 인덱스에 대항 임베딩을 해줘야 한다. 
그때 3이라고 하면 다음과 같이 임베딩이 될 것이다.
임베딩되면 (1,4,3)이 된다.
그러면 각각의 단어들끼리도 포지션 정보를 줘서 (1,4,3)이다. 
이런 식으로 특정 시퀀스에 대해서 임베딩을 하고 임베딩을 한 것끼리 더하고 붙여주거나 하는 작업을 해준다.
(1,4) → (1,4,3)

Encoder Block

인코더 블록에 대해서 알아보자.
자세한 구조를 보면 이런식으로 self attention 이랑 layer norm으로 구조가 되어 있다. 

Self Attention

Attention is All You 라는 논문에서부터 시작된 개념이라 볼 수 있다.
 
아래와 같은 문장을 번역하고자 합시다.
 
”The animal didn't cross the street because it was too tired”
 
이 문장에서 "it"은 무엇을 가리킵니까? 거리(street)를 말하는 건가요 아니면 동물(animal)을 말하는 건가요?
사람에게는 간단한 질문이지만 알고리즘에게는 그렇게 간단하지 않습니다.
그렇지만 attention을 사용하게 되면 모델이 "it"이라는 단어를 처리할 때 self-attention을 통해 "it"을 "animal"과 연관시킬 수 있습니다.
모델이 각 단어(입력 시퀀스의 각 위치)를 처리할 때 Self Attention을 통해 입력 시퀀스의 다른 위치에서 이 단어를 더 잘 인코딩하는 데 도움이 되는 단서를 찾을 수 있습니다.
 
RNN에 익숙하다면 hidden state를 유지함으로써 RNN이 처리한 이전 단어/벡터의 표현을 현재 처리 중인 단어/벡터와 통합할 수 있는 방법을 생각해 보십시오.
Self-attention은 Transformer가 다른 관련 단어의 "이해"를 현재 처리 중인 단어로 만드는 방법입니다.

 
이제 그러면 위에 같이 특정 단어와 나머지 단어간의 중요도를 계산하는 방법을 알아보자.
 

1번째 단계

Self-attention 계산의 첫 번째 단계는 인코더의 각 입력 벡터(이 경우 각 단어의 임베딩)에서 세 개의 벡터를 만드는 것입니다.
따라서 각 단어에 대해 쿼리(query) 벡터, 키(key) 벡터 및 값(value) 벡터를 만듭니다.
이러한 벡터는 임베딩에 학습 과정 중에 훈련시킨 3개의 행렬을 곱하여 생성됩니다.
 
3개의 벡터는 attention를 계산하고 생각하는 데 유용한 추상화입니다. 
attention이 어떻게 계산되는지 아래에서 읽으면 이러한 각 벡터가 수행하는 역할에 대해 알아야 할 거의 모든 것을 알게 될 것입니다.

2번째 단계

Self-Attention 계산의 두 번째 단계는 점수를 계산하는 것입니다.
이 예에서 "Thinking"이라는 첫 단어에 대한 self-attention을 계산한다고 가정해 보겠습니다. 이 단어에 대해 입력 문장의 각 단어에 점수를 매겨야 합니다.
점수는 특정 위치에서 단어를 인코딩할 때 입력 문장의 다른 부분에 얼마나 많은 초점을 둘 것인지를 결정합니다.
 
점수는 점수를 매기는 각 단어의 키(key) 벡터와 쿼리(query) 벡터의 내적을 취하여 계산됩니다.
 
따라서 위치 #1에 있는 단어에 대한 셀프 어텐션을 처리하는 경우 첫 번째 점수는 q1과 k1의 내적이 됩니다. 두 번째 점수는 q1과 k2의 내적입니다.
즉 보면 각각의 단어와 다른 다언 같의 스코어를 다 계산합니다.

3,4번째 단계

세 번째와 네 번째 단계는 점수를 8로 나누는 것입니다(논문에서 사용된 키 벡터 차원의 제곱근 64).
이것은 더 안정적인 기울기를 갖게 합니다. 여기에 다른 가능한 값이 있을 수 있지만 이것이 기본값)
softmax 작업을 통해 결과를 전달합니다. Softmax는 점수를 정규화하여 모두 양수이고 합이 1이 되도록 합니다.
 

위에서 q1에 대해서 값을 구하고 확률값으로 만들어 줍니다.
 
이 softmax 점수는 각 단어가 이 위치에서 표현되는 양을 결정합니다. 분명히 이 위치에 있는 단어는 가장 높은 softmax 점수를 가지지만 때로는 현재 단어와 관련된 다른 단어에 주의를 기울이는 것이 유용합니다.

5번째 단계

다섯 번째 단계는 각 값(value) 벡터에 softmax 점수를 곱하는 것입니다(합산 준비).
- 여기서 직감은 집중하려는 단어의 값을 그대로 유지하고 관련 없는 단어를 제거하는 것입니다(예를 들어 0.001과 같은 작은 숫자를 곱하여).
여섯 번째 단계는 가중 값 벡터를 합산하는 것입니다. 이는 이 위치(첫 번째 단어에 대해)에서 self-attention 레이어의 출력을 생성합니다.

이것으로 self-attention 계산을 마칩니다. 결과 벡터는 피드포워드 신경망에 함께 보낼 수 있는 벡터입니다.
그러나 실제 구현에서 이 계산은 더 빠른 처리를 위해 행렬 형식으로 수행됩니다. 단어 수준에서 계산의 직관을 보았으니 이제 살펴보겠습니다.

Self Attention 행렬 계산

 

	Token	(Batch, Seq_Len)
	Token Embedding	(Batch, Seq_Len, Emb_Dim)
	Position Encoding	(Batch, Seq_Len) => (Batch, Seq_Len, Emb_Dim)
	Token Embedding + Position Encoding (SUM)	(Batch, Seq_Len, Emb_Dim)
Self Attention	Q(Wegiht (Emb_Dim, Emb_Dim2))	(Batch, Seq_Len, Emb_Dim2)
	K(Wegiht (Emb_Dim, Emb_Dim2))	(Batch, Seq_Len, Emb_Dim2)
	V(Wegiht (Emb_Dim, Emb_Dim2))	(Batch, Seq_Len, Emb_Dim2)
	QK^T	(Batch, Seq_Len, Seq_Len)
	softmax(QK^T)	(Batch, Seq_Len, Seq_Len)
	softmax(QK^T)V	(Batch, Seq_Len, Emb_Dim2)

 

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    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 = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

 

Multi-Head Self Attention 행렬 계산

	Token	(Batch, Seq_Len)
	Token Embedding	(Batch, Seq_Len, Emb_Dim)
	Position Encoding	(Batch, Seq_Len) => (Batch, Seq_Len, Emb_Dim)
	Token Embedding + Position Encoding (SUM)	(Batch, Seq_Len, Emb_Dim)
Multi Head Self Attention	Q(Wegiht (Emb_Dim, Emb_Dim2)) Emb_Dim2 (Head x D_Model)	(Batch, Seq_Len, Emb_Dim2) => (Batch, Head, Seq_Len, D_Model)
	K(Wegiht (Emb_Dim, Emb_Dim2)) Emb_Dim2 (Head x D_Model)	(Batch, Seq_Len, Emb_Dim2) => (Batch, Head, Seq_Len, D_Model)
	V(Wegiht (Emb_Dim, Emb_Dim2)) Emb_Dim2 (Head x D_Model)	(Batch, Seq_Len, Emb_Dim_v2) => (Batch, Head, Seq_Len, D_Model)
	QK^T	(Batch, Head, Seq_Len, Seq_Len)
	softmax(QK^T)	(Batch, Head, Seq_Len, Seq_Len)
	softmax(QK^T)V	(Batch, Head, Seq_Len, D_Model)
Transpose	softmax(QK^T)V => transpose(1,2)	(Batch, Seq_Len, Head, D_Model)
Reshape	softmax(QK^T)V => (batch, seq_len , -1)	(Batch, Seq_Len, Head x D_Model)
Fc	W (Head x D_Model , Emb_Dim3)	(Batch, Seq_Len, Emb_Dim3)

 
1) "인코더-디코더 주의" 계층에서 쿼리는 이전 디코더 계층에서 오고 메모리 키와 값은 인코더 출력에서 옵니다.
이를 통해 디코더의 모든 위치가 입력 시퀀스의 모든 위치에 참석할 수 있습니다.
이는 sequence-to-sequence 모델의 일반적인 인코더-디코더 주의 메커니즘을 모방합니다.
 
2) 인코더에는 self-attention 레이어가 포함되어 있습니다. self-attention 계층에서 모든 키, 값 및 쿼리는 동일한 위치에서 옵니다. 이 경우 인코더에서 이전 레이어의 출력입니다. 인코더의 각 위치는 인코더의 이전 계층에 있는 모든 위치에 참석할 수 있습니다.
 
3) 유사하게, 디코더의 셀프 어텐션 레이어는 디코더의 각 위치가 해당 위치를 포함하여 디코더의 모든 위치에 참석할 수 있도록 합니다. auto-regressive 속성을 유지하기 위해 디코더에서 왼쪽 방향 정보 흐름을 방지해야 합니다.
우리는 마스킹 아웃(masking out)을 통해 scaled dot attention를 구현합니다
 
 
 

Add & Normalize

다음 단계는 크게 2가지 합쳐진 거라 보면 된다.
Add 는 residual connection 같은 개념이고 Normalize 는 Layer Norm을 사용하고 있다.

Add	Residual connection	solution for vanishing gradient problme
Normalize	Layer Norm	tricks to make life easier for training the model

1. Residual Connection (잔차 연결): 이것은 각 서브-레이어(예: multi-head attention, feed-forward neural network 등)의 입력을 그 출력에 추가하는 과정입니다. 잔차 연결은 심층 네트워크에서 그래디언트를 더 쉽게 역전파하도록 돕고, 모델이 더 깊은 층을 학습하는 데 도움이 됩니다. 이러한 잔차 연결은 "Add" 부분을 담당합니다.

2. Layer Normalization (층 정규화): 잔차 연결 후에는 층 정규화가 적용됩니다. 층 정규화는 각 샘플의 각 특징을 독립적으로 정규화합니다. 이는 학습 과정을 안정화하고, 모델의 학습 속도를 향상시키며, 일반적으로 더 나은 성능을 달성하는 데 도움이 됩니다. 이것이 "Normalize" 부분입니다.

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = 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.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))

 

Position-wise Feed-Forward Networks

각각의 attention sub layer에는 위 레이어, 인코더 및 디코더의 각 레이어에는 fully connected feed-forward 네트워크가 포함되어 있으며 각 위치에 개별적으로 동일하게 적용됩니다.
이것은 사이에 ReLU 활성화가 있는 두 개의 선형 변환으로 구성됩니다.
 
FFN(x)=max(0,xW1+b1)W2+b2

class PositionWiseFeedForwardLayer(nn.Module):

    def __init__(self, fc1, fc2):
        super(PositionWiseFeedForwardLayer, self).__init__()
        self.fc1 = fc1   # (d_embed, d_ff)
        self.relu = nn.ReLU()
        self.fc2 = fc2 # (d_ff, d_embed)


    def forward(self, x):
        out = x
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)
        return out

역할
FFNN는 Transformer의 중요한 구성 요소로서, 각 단어나 시퀀스 위치의 특징을 학습하는 데 중요한 역할을 합니다.

FFNN의 주요 역할은:

1. 복잡도 추가: FFNN은 Transformer에 비선형성을 추가하여 모델이 복잡한 패턴을 학습할 수 있게 합니다. FFNN은 선형 변환 후에 활성화 함수(예: ReLU 또는 GELU)를 적용하여 비선형성을 추가합니다.

2. 특징 학습: FFNN는 각 위치의 입력에 독립적으로 적용되므로, 각 위치에서 서로 다른 특징을 추출하고 학습할 수 있습니다. 이것은 각 단어나 시퀀스 위치가 다른 위치와 독립적으로 처리될 수 있음을 의미합니다.

3. 모델 용량 증가: FFNN는 Transformer의 모델 용량을 증가시킵니다. 이는 모델이 더 큰 입력 공간과 출력 공간을 다룰 수 있게 해 주며, 따라서 더 다양한 패턴과 상호작용을 학습할 수 있습니다.

따라서, Position-wise Feed-Forward Neural Network는 Transformer 모델의 핵심 부분으로, 각 위치에서 복잡한 특징을 학습하고 추출하는 데 중요한 역할을 합니다.
 

Encoder Block 대략적인 전체 코드

import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    # MultiHeadAttention 구현을 위한 공간. 이 부분을 본인의 코드로 채워주세요.
    pass

class PositionWiseFeedForward(nn.Module):
    # PositionWiseFeedForward 구현을 위한 공간. 이 부분을 본인의 코드로 채워주세요.
    pass

class EncoderBlock(nn.Module):
    def __init__(self, d_model, d_ff, n_heads):
        super().__init__()
        
        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = PositionWiseFeedForward(d_model, d_ff)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.dropout = nn.Dropout(0.1)

    def forward(self, x, mask):
        # x shape: (batch_size, seq_length, d_model)
        # mask shape: (batch_size, seq_length)
        
        # Multi-Head Self-Attention with Residual Connection and Layer Normalization
        attn_output = self.self_attention(x, x, x, mask)  
        # attn_output shape: (batch_size, seq_length, d_model)
        
        x = self.norm1(x + self.dropout(attn_output))  
        # x shape after norm1 and dropout: (batch_size, seq_length, d_model)
        
        # Position-wise Feed-Forward Network with Residual Connection and Layer Normalization
        ffn_output = self.ffn(x)  
        # ffn_output shape: (batch_size, seq_length, d_model)
        
        x = self.norm2(x + self.dropout(ffn_output))
        # x shape after norm2 and dropout: (batch_size, seq_length, d_model)

        return x

 

MASK (SRC MASK)

 
위에 코드의 forward를 보면 이제까지 이야기 하지 않은 mask라는 것이 있다.
그리고 그 마스크는 attention 함수에서 사용하는 것을 알 수 있다.
 
(패딩)마스킹은 주로 특정 위치의 값을 무시하게 하기 위해 사용됩니다.
이는 주로 시퀀스의 끝에 패딩이 추가되었을 때 유용하며, 패딩된 위치는 모델의 학습 또는 추론에 영향을 주지 않아야 합니다.
 
즉 이게 필요한 이유는 seq_len라는 것이 정해져있기 때문에, 각 문장마다 길이는 달라지고, 그때 사용하지 않는 토큰들에 대해서는 보통 <pad> 처리를 하는데, 그것을 attention에서는 반영이 안되게 하는 것이라고 볼 수 있다.
 
그래서 multi head attention에 구현을 보면 다음과 같이 되어 있다.
아마 이렇게 하는 이유를 생각해보면, batch 단위로 데이터를 넣다 보니, batch 에서 각각의 데이터마다 길이가 다를 수 있다.
만약 batch가 1 이라면 데이터 한개를 가지고 하는 거니까 상관이 없겠지만, 여러개를 병렬적으로 처리하기 위해서는 pad mask이라는 개념이 필요해보이긴 한다. 

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()

        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        self.qkv_proj = nn.Linear(d_model, d_model * 3)  # for Q, K, V projections
        self.fc_out = nn.Linear(d_model, d_model)  # final output projection

        self.scale = self.head_dim ** -0.5  # scale factor for dot product attention

    def split_heads(self, x, batch_size):
        x = x.view(batch_size, -1, self.n_heads, self.head_dim)
        return x.transpose(1, 2)  # shape: (batch_size, n_heads, seq_length, head_dim)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        QKV = self.qkv_proj(torch.cat((query, key, value), dim=-1))  # Q, K, V projections
        Q, K, V = torch.split(QKV, self.d_model, dim=-1)

        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)

        energy = torch.matmul(Q, K.transpose(-2, -1)) * self.scale

        if mask is not None:
            # original mask shape: (batch_size, seq_length)
            mask = mask.unsqueeze(1).unsqueeze(2)
            # mask shape after unsqueeze: (batch_size, 1, 1, seq_length)
            energy = energy.masked_fill(mask, float("-inf"))


        attention = torch.softmax(energy, dim=-1)

        out = torch.matmul(attention, V)
        out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        out = self.fc_out(out)

        return out

 
 

Encoder 전체 코드 (대략적인 모습)

class TransformerEncoder(nn.Module):
    def __init__(self, d_model, d_ff, n_heads, num_blocks):
        super().__init__()
        
        self.encoder_blocks = nn.ModuleList([
            EncoderBlock(d_model, d_ff, n_heads) for _ in range(num_blocks)
        ])

    def generate_padding_mask(self, input_seq):
        # 패딩 토큰의 위치에서는 1, 그렇지 않은 경우는 0
        pad_mask = (input_seq == 0)
        return pad_mask

    def forward(self, x):
        # x shape: (batch_size, seq_length)
        
        # Generate padding mask
        pad_mask = self.generate_padding_mask(x)
        # pad_mask shape: (batch_size, seq_length)
        
        for block in self.encoder_blocks:
            x = block(x, pad_mask)
            # x shape after each block: (batch_size, seq_length, d_model)

        return x


# Initialize the Transformer Encoder with 3 blocks
encoder = TransformerEncoder(d_model=512, d_ff=2048, n_heads=8, num_blocks=3)

 
여기까지 한 번 정리를 하면, transformer encoder에 태우기 위해서 필요한 것은데이터를 토큰화하는 작업이 필요하고
embedding lookup table을 구축하여, 각각의 토큰을 임베딩해야 하고, position 에 대해서 인코딩을 할 수 있어야 한다.
그리고 특정 토큰에 대해서 padding을 할 게 있으면 padding 처리를 통해 학습에 관여하지 않게 한다.
 
해당 데이터를 encoder block에 넣으면 다음과 같은 과정을 거치게 된다.
1. multi head attention
2. add & normalize
3. feed forward (Position-wise Feed Forward Network)
4. add & normalize
 
그 다음에 다음 encoder block에 또 mask랑 encoder block에서 나온 것을 그대로 넣어주면 된다. 
여기까지 하면 encoder에서 하는 것에 대한 결과는 최종적으로 다음과 같이 나온다.
 
(batch_size, seq_length, d_model)
즉 들어가는 것과 나오는 것의 대한 차원은 같게 나온다.
 

Decoder SIDE

decoder에서는 다음과 같이 진행되게 된다.

즉 encoder output에 나온 k,v의 값을 decoder에서 사용하게 된다.

크게 Decoder Layer 2가지로 나눌 수 있다.

1	"masked" self-attention layer	현재 및 이전 위치의 정보만을 참조
2	encoder-decoder attention layer	encoder의 출력을 참조

 

mask

특정 위치의 토큰을 숨기기 위해 사용하는데, 이를 통해 Transformer 모델이 특정 정보를 무시

src_mask	소스 시퀀스에서 패딩 토큰의 위치를 표시합니다. 일반적으로, 입력 시퀀스는 여러 개의 시퀀스를 함께 처리하기 위해 패딩 토큰으로 채워집니다. src_mask는 이러한 패딩 토큰이 어텐션 계산에 포함되지 않도록 하는 역할을 합니다. src_mask는 소스 시퀀스와 동일한 길이를 가지며, 패딩 토큰의 위치에서는 True (또는 1), 그 외의 위치에서는 False (또는 0)를 가집니다.
tgt_mask	Transformer Decoder에서 현재 및 이전 위치의 토큰만 참조하도록 합니다. 이를 통해 Decoder가 각 위치에서 출력을 생성할 때 "미래" 토큰을 "보지 못하게" 합니다. 이는 보통 "look-ahead" 마스크 또는 "causal" 마스크라고 부릅니다.

 

def create_src_mask(src):
    src_mask = (src != PAD_ID).unsqueeze(1).unsqueeze(2)
    # src shape: (batch_size, seq_length)
    # src_mask shape: (batch_size, 1, 1, seq_length)
    return src_mask
    
def create_tgt_mask(tgt):
    tgt_pad_mask = (tgt != PAD_ID).unsqueeze(1).unsqueeze(3)
    # tgt_pad_mask shape: (batch_size, 1, seq_length, 1)
    
    tgt_len = tgt.shape[1]
    # torch.tril 함수는 주어진 텐서의 하삼각부분만을 유지하고 나머지를 0으로 만드는 함수입니다. 
    # 이를 이용하여 각 위치에서 자신과 이전 토큰만 "볼 수 있는" 마스크를 생성합니다. 
    
    tgt_sub_mask = torch.tril(torch.ones((tgt_len, tgt_len))).bool()
    # tgt_sub_mask shape: (seq_length, seq_length)
    
    # 이 두 마스크를 합치면, 패딩 토큰과 미래 토큰 모두에 대한 마스크를 생성할 수 있습니다.
    tgt_mask = tgt_pad_mask & tgt_sub_mask
    # tgt_mask shape: (batch_size, 1, seq_length, seq_length)
    return tgt_mask

tgt_sub_mask는 이런 느낌으로 생성됩니다.

Decoder Block

decoder의 입력(즉, 타겟 시퀀스)에 대한 쿼리(Query)를 사용하고, encoder의 출력(즉, 소스 시퀀스의 처리된 표현)에 대한 키(Key) 및 값(Value)을 사용하여 작동합니다.
 
encoder_decoder_attention 계층의 입력으로 x(decoder의 입력), encoder_out(encoder의 출력), encoder_out(다시 한번), src_mask(소스 마스크)를 전달하고 있습니다. 
첫번째 encoder_out은 key로, 두번째 encoder_out은 value로 사용되며, x는 query로 사용됩니다. 
src_mask는 소스 시퀀스의 패딩 토큰을 마스킹하기 위해 사용됩니다. 
 
이렇게 계산된 결과는 잔여 연결(residual connection)을 통해 원래의 x에 더해지고, 그 결과는 layer normalization을 통과하여 최종 출력을 생성합니다.
 
 

class DecoderBlock(nn.Module):
    def __init__(self, d_model, d_ff, n_heads):
        super().__init__()

        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)

        self.encoder_decoder_attention = MultiHeadAttention(d_model, n_heads)
        self.norm2 = nn.LayerNorm(d_model)

        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm3 = nn.LayerNorm(d_model)

    def forward(self, x, encoder_out, src_mask=None, tgt_mask=None):
        # Self Attention
        residual = x
        x = self.norm1(x + self.self_attention(x, x, x, tgt_mask))
        # x shape: (batch_size, target_seq_length, d_model)

        # Encoder-Decoder Attention
        residual2 = x
        x = self.norm2(x + self.encoder_decoder_attention(x, encoder_out, encoder_out, src_mask))
        # x shape: (batch_size, target_seq_length, d_model)

        # Position-Wise Feed Forward
        x = self.norm3(x + self.feed_forward(x))
        # x shape: (batch_size, target_seq_length, d_model)

        return x

 
 

Decoder 전체 코드

class Decoder(nn.Module):
    def __init__(self, d_model, d_ff, n_heads, n_layers):
        super().__init__()

        self.layers = nn.ModuleList([DecoderBlock(d_model, d_ff, n_heads) for _ in range(n_layers)])

    def forward(self, x, encoder_out, src_mask=None, tgt_mask=None):  # x, encoder_out shape: (batch_size, seq_length, d_model)
        for layer in self.layers:
            x = layer(x, encoder_out, src_mask, tgt_mask)
        return x  # output shape: (batch_size, seq_length, d_model)

class DecoderBlock(nn.Module):
    def __init__(self, d_model, d_ff, n_heads):
        super().__init__()

        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)

        self.encoder_decoder_attention = MultiHeadAttention(d_model, n_heads)
        self.norm2 = nn.LayerNorm(d_model)

        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm3 = nn.LayerNorm(d_model)

    def forward(self, x, encoder_out, src_mask=None, tgt_mask=None):  # x, encoder_out shape: (batch_size, seq_length, d_model)
        # Self Attention
        residual = x
        x = self.norm1(residual + self.self_attention(x, x, x, tgt_mask))  # x shape after self attention: (batch_size, seq_length, d_model)

        # Encoder-Decoder Attention
        residual = x
        x = self.norm2(residual + self.encoder_decoder_attention(x, encoder_out, encoder_out, src_mask))  # x shape after encoder-decoder attention: (batch_size, seq_length, d_model)

        # Position-Wise Feed Forward
        residual = x
        x = self.norm3(residual + self.feed_forward(x))  # x shape after feed forward: (batch_size, seq_length, d_model)

        return x  # output shape: (batch_size, seq_length, d_model)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()

        assert d_model % n_heads == 0

        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)

        self.fc = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None):  # query, key, value shape: (batch_size, seq_length, d_model)
        batch_size = query.shape[0]

        # Linear layers
        Q = self.W_q(query)  # Q shape: (batch_size, seq_length, d_model)
        K = self.W_k(key)    # K shape: (batch_size, seq_length, d_model)
        V = self.W_v(value)  # V shape: (batch_size, seq_length, d_model)

        # Split the last dimension into (n_heads, d_k)
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k)  # Q shape: (batch_size, seq_length, n_heads, d_k)
        K = K.view(batch_size, -1, self.n_heads, self.d_k)  # K shape: (batch_size, seq_length, n_heads, d_k)
        V = V.view(batch_size, -1, self.n_heads, self.d_k)  # V shape: (batch_size, seq_length, n_heads, d_k)

        # Transpose seq_length and n_heads dimensions
        Q = Q.transpose(1, 2)  # Q shape: (batch_size, n_heads, seq_length, d_k)
        K = K.transpose(1, 2)  # K shape: (batch_size, n_heads, seq_length, d_k)
        V = V.transpose(1, 2)  # V shape: (batch_size, n_heads, seq_length, d_k)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # scores shape: (batch_size, n_heads, seq_length, seq_length)

        if mask is not None:  # mask shape: (batch_size, 1, 1, seq_length)
            scores = scores.masked_fill(mask == 0, float("-inf"))

        attention_weights = torch.softmax(scores, dim=-1)  # attention_weights shape: (batch_size, n_heads, seq_length, seq_length)

        out = torch.matmul(attention_weights, V)  # out shape after attention: (batch_size, n_heads, seq_length, d_k)

        # Transpose n_heads and seq_length dimensions
        out = out.transpose(1, 2)  # out shape: (batch_size, seq_length, n_heads, d_k)

        # Combine the last two dimensions
        out = out.contiguous().view(batch_size, -1, self.n_heads * self.d_k)  # out shape: (batch_size, seq_length, d_model)

        out = self.fc(out)  # out shape: (batch_size, seq_length, d_model)

        return out  # output shape: (batch_size, seq_length, d_model)

 

특정 테스크

번역

예를 들어 번역이라는 테스크에 적용한다고 하면, output은 vocab_size만큼 될 것이다.

class Transformer(nn.Module):
    def __init__(self, d_model, d_ff, n_heads, n_layers, vocab_size):
        super().__init__()

        self.encoder = Encoder(d_model, d_ff, n_heads, n_layers)
        self.decoder = Decoder(d_model, d_ff, n_heads, n_layers)

        # This linear layer transforms the output of the decoder to the size of the vocabulary
        self.out = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        encoder_out = self.encoder(src, src_mask)
        decoder_out = self.decoder(tgt, encoder_out, src_mask, tgt_mask)

        # Apply the output linear layer and softmax to get probabilities
        output_probabilities = F.softmax(self.out(decoder_out), dim=-1)

        return output_probabilities

 
 
 
 
 
 
 
 
 

Softmax and Output Probabilities

디코더는 인코더의 입력 기능을 사용하여 출력 문장을 생성합니다.
인풋 feature는 풍부한 임베딩 벡터에 불과합니다.
 

 
 
 
 

참고

http://nlp.seas.harvard.edu/2018/04/03/attention.html
https://cpm0722.github.io/pytorch-implementation/transformer
http://jalammar.github.io/illustrated-transformer/
https://peterbloem.nl/blog/transformers
https://kikaben.com/transformers-positional-encoding/
https://wikidocs.net/166796
https://machinelearningmastery.com/implementing-the-transformer-encoder-from-scratch-in-tensorflow-and-keras/
https://machinelearningmastery.com/implementing-the-transformer-decoder-from-scratch-in-tensorflow-and-keras/
https://kikaben.com/transformers-encoder-decoder/

Transformer’s Encoder-Decoder: Let’s Understand The Model Architecture - KiKaBeN