트랜스포머(Transformer) (4)

1. Token Embedding + Positional Encoding

■ 트랜스포머는 다음 그림과 같이 input과 output은 각각 Input/Output Embedding + Positional Encoding을 거쳐 인코더와 디코더의 입력으로 들어간다.

■ 트랜스포머 모델에 사용할 파라미터는 다음과 같으며, 

# model parameter
batch_size = 128
max_len = 256
d_model = 512
n_layers = 6
n_heads = 8
d_ff = 2048
drop_prob = 0.1

■ 영어-독일어 단어 집합(vocabulary)의 크기와, <pad> 토큰의 정수 인덱스(또는 아이디)는 다음과 같다.

len(en_vocab), len(de_vocab)
```#결과#```
(5893, 7853)
````````````

en_vocab.get_stoi()['<pad>'], de_vocab.get_stoi()['<pad>']
```#결과#```
(1, 1)
````````````

---

1.1 Token Embedding

■ 단어 집합의 크기를 받아, d_model 차원으로 입력을 임베딩하는 것은 파이토치의 nn.Embedding 클래스를 상속받아서 다음과 같이 구현할 수 있다.

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, d_model):
        super().__init__(vocab_size, d_model, padding_idx=1)

■ 예를 들어, 독일어 시퀀스의 미니배치 (batch_size, seq_len)이 토큰 임베딩 클래스를 통과한다면, (batch_size, seq_len, embedding_dim=d_model=512)로 변환된다. 배치에 있는 데이터는 토큰의 정수 인덱스이므로 룩업 테이블 연산이 수행되는 것이다.

■ seq_len 차원을 가지는 batch_size 개의 벡터들이 임베딩 계층을 거쳐 batch_size 개의 (seq_len, embedding_dim=d_model) 형태로 변환된다.

---

1.2 Positional Encoding

■ 포지셔널 인코딩은 다음과 같이 구현할 수 있다.

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len, device):
        super().__init__()
        self.encoding = torch.zeros(max_len, d_model, device=device) # 입력과 더하기 위해 입력과 같은 크기로 생성
        self.encoding.requires_grad = False # 그래디언트 계산 불필요

        pos = torch.arange(0, max_len, device=device) # 0부터 max_len - 1까지의 숫자 생성
        # pos.shape: (max_len)
        pos = pos.float().unsqueeze(dim=1) # 단어의 위치를 나타내기 위해 1D -> 2D로 차원 확장
        # pos.shape: (max_len, 1)

        ## 2i 만들기
        # i는 d_model 차원의 인덱스 # i는 0부터 시작
        _2i = torch.arange(0, d_model, step=2, device=device).float() # 0부터 '차원 수(d_model)-1'까지 step=2의 숫자 생성

        # 사인 함수는 짝수 인덱스에 위치 정보 계산
        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        # 코사인 함수는 홀수 인덱스에 위치 정보 계산
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))

    def forward(self, x):
        # x.shape : (batch, seq_len) or (batch, seq_len, d_model)
        seq_len = x.size(1) # seq_len은 현재 배치의 시퀀스 길이(max_len보다 작거나 같은 값)
        return self.encoding[:seq_len, :] # 반환되는 텐서의 shape: (seq_len, d_model)

■ 여기서 입력으로 들어오는 x의 차원은 TokenEmbedding을 거쳤기 때문에 512이다.

■ Attention Is All You Need 논문에서는 다음과 같은 사인(sine)과 코사인(cosine)을 활용한 수식을 통해 512 차원 중 짝수 인덱스에는 사인, 홀수 인덱스에는 코사인 값을 넣기 때문에 입력으로 들어오는 임베딩의 차원이  d_model=512라면, max_len=256이 되어야 한다. 

 

■ pos = torch.arange(0, max_len)를 통해 0에서 255까지의 총 256개의 숫자를 생성하고,

_2i = torch.arange(0, d_model, step=2).float()는 0에서 255까지의 step=2의 숫자를 생성한다. 

■ 브로드캐스팅 연산을 하기 위해 pos의 차원을 dim=1 방향으로 확장하면, pos는 (256, 1)이 된다.

[[pos_0]
 [pos_1]
 ...
 [pos_255]]

■ (10000 ** (_2i / d_model))의 크기는 (256). 즉, 256개의 원소를 가진 1차원 배열이다.

■ pos / (10000 ** (_2i / d_model))에서 파이토치는 (10000 ** (_2i / d_model)) 항을 연산이 가능하도록 2차원 텐서 (1, 256)으로 간주하게 된다.

■ 이제 '/' 연산 과정에서 브로드캐스팅 연산이 진행되는데, (256, 1) 크기의 pos는 1개의 열이 256번 복제되어 (256, 256) 크기로 확장되고

[[pos_0, pos_0, ..., pos_0]       
 [pos_1, pos_1, ..., pos_1]
 ...
 [pos_255, pos_255, ..., pos_255]]

■ 분모(denominator)에 들어가는 (1, 256) 크기의  (10000 ** (_2i / d_model)) 항은 행이 256번 복제되어 (256, 256) 크기로 확장된다고 생각할 수 있다.

[[denom_0, denom_1, ..., denom_255]
 [denom_0, denom_1, ..., denom_255] 
 ...
 [denom_0, denom_1, ..., denom_255]]

■ 이제 두 텐서 모두 (256, 256) 크기를 가지며, 같은 위치에 있는 요소들끼리 나눗셈 연산이 수행되며, 결과 텐서의 크기는 (256, 256)이 된다.

■ 이러한 위치 인코딩 결과는 행이 max_len으로 시퀀스의 최대 길이, 열은 입력으로 들어오는 토큰 임베딩 벡터의 차원인 d_model=512로 이루어진 (max_len, d_model) 크기의 self.encoding에 다음과 같이 torch.sin(x) 256개, torch.cos(x) 256개의 값이 각각 짝수, 홀수 인덱스에 들어가게 된다.

# 사인 함수는 짝수 인덱스에 위치 정보 계산
self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
# 코사인 함수는 홀수 인덱스에 위치 정보 계산
self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))

■ self.encoding[:, 0::2]은 모든 행(max_len)의 짝수 인덱스 열(2i)을 의미한다. 0::2는 0번 인덱스부터, 2씩 건너뛰며 열을 선택한다는 뜻이다. 즉, [0, 2, 4, 6, ...]같은 짝수 번째 열들을 의미한다.  

■ self.encoding[:, 1::2]은 모든 행(max_len)의 홀수 인덱스 열(2i+1)을 의미한다. [1, 3, 5, 7, ...]같은 홀수 번째 열들을 의미한다.

■ 그러므로 (max_len, d_model=512) 크기의 self.encoding에 대하여,

- self.encoding[:, 0::2]에는 256개의 사인 값들이 짝수 인덱스의 열에 들어가고

- self.encoding[:, 1::2]에는 256개의 코사인 값들이 홀수 인덱스의 열에 들어간다.

■ 최종적으로 필요한 위치 인코딩 결과는, 배치 단위로 들어오는 (seq_len, d_model=512) 크기의 토큰 임베딩과 동일한 크기 (seq_len, d_model)이 필요하다. 그러므로 forward() 과정에서 다음과 같이 현재 배치의 시퀀스 길이 seq_len 만큼의 행만 가져와서 결과를 반환한다.

    def forward(self, x):
        # x.shape : (batch, seq_len, d_model)
        seq_len = x.size(1) # seq_len은 현재 배치의 시퀀스 길이(max_len보다 작거나 같은 값)
        return self.encoding[:seq_len, :] # 반환되는 텐서의 shape: (seq_len, d_model)

■ seq_len=30이라고 가정하자. 그렇다면, 토큰 임베딩의 결과는 (batch_size, seq_len=30, d_model=512)이고, 포지셔널 인코딩의 결과는 (seq_len=30, d_model=512)이다. 마찬가지로 두 결과를 더할 때, 브로드캐스팅 연산이 진행된다.

■ 참고로, 위치 인코딩 결과는 학습 과정에서 업데이트될 필요가 없다. 그러므로 self.encoding에 requires_grad=False를 설정하였다.

---

1.3 최종 모델 입력을 위한 임베딩

■ 텍스트를 모델에 입력하기 위한 최종 임베딩은 다음과 같이 토큰 임베딩과 포지셔널 임베딩의 덧셈(+)으로 구현할 수 있다.

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, dropout, device):
        super().__init__()
        self.tok_embedding = TokenEmbedding(vocab_size=vocab_size, d_model=d_model)
        self.pos_embedding = PositionalEncoding(d_model=d_model, max_len=max_len, device=device)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        tok_emb = self.tok_embedding(x)
        pos_emb = self.pos_embedding(x)
        return self.dropout(tok_emb + pos_emb)

■ 순전파 과정(forward())에서 TransformerEmbedding의 입력으로 들어오는 x는 (batch_size, seq_len=30)의 입력 시퀀스이다. 

■ (batch_size, seq_len=30) 크기의 x는 TokenEmbedding과 PositionalEncoding을 각각 거쳐 (batch_size, seq_len=30, embedding_dim=d_model=512) 크기의 토큰 임베딩과 (seq_len=30, d_model=512) 크기의 위치 인코딩으로 변환된다.

■ 트랜스포머 모델에 최종적으로 사용하는 임베딩은 토큰 임베딩(tok_emb) + 위치 인코딩(pos_emb)이다. 이 덧셈 과정에서 브로드캐스팅 연산이 수행된다. 

■ 파이토치는 크기가 달라도 사칙연산을 수행할 수 있도록 자동으로 크기를 맞춰 연산을 수행하게 해주는 브로드캐스팅 연산을 지원한다. 

■ tok_emb + pos_emb에서 tok_emb과 pos_emb의 두 번째, 세 번째 차원은 동일하지만, 첫 번째 차원(batch_size)은 pos_emb에 존재하지 않지만, pos_emb은 tok_emb의 모든 계층에 브로드캐스트된다.

■ (seq_len=30, d_model=512)의 pos_emb 텐서가 batch_size 개수만큼 복사되어 각각의 배치마다 pos_emb가 element-wise sum을 한다고 볼 수 있다.

■ 그러므로 TransformerEmbedding의 결과로 반환되는 텐서의 형상은 (batch_size, seq_len, d_model)을 유지하게 된다.

---

---

2. Transformer Encoder

■ 하나의 인코더 계층(인코더 블록)에는 멀티 헤드 (셀프) 어텐션 하위 계층(sub layer)과 피드 포워드 하위 계층, 그리고 안정적인 학습이 가능하도록 도와주는 층 정규화와 잔차 연결로 구성된다. 

- 논문에서는 아래의 트랜스포머 모델 아키텍처 그림처럼 층 정규화 방식으로 사후 정규화를 사용했지만, 사전 정규화를 적용해 보았다.

■ 그리고 이 인코더 블록을 N번 반복해서 인코더 블록을 N개 쌓는다. 

 

■ 위의 그림과 같은 트랜스포머 인코더는 다음과 같이 구현할 수 있다. 

class TransformerEncoder(nn.Module):
    def __init__(self, src_vocab_size, max_len, d_model, d_ff, n_head, n_layers, dropout, device):
        super().__init__()
        self.emb = TransformerEmbedding(d_model=d_model,
                                        max_len=max_len,
                                        vocab_size=src_vocab_size,
                                        dropout=dropout,
                                        device=device)

        self.layers = nn.ModuleList([EncoderLayer(d_model=d_model,
                                                  d_ff=d_ff,
                                                  n_head=n_head,
                                                  dropout=dropout)
                                         for _ in range(n_layers)])

    def forward(self, src, src_mask):
        output = self.emb(src)

        for mod in self.layers:
            output = mod(output, src_mask)

        return output

■ nn.ModuleList는 파이토치에서 사용되는 모듈들을 리스트 형태로 관리하는 클래스이다. 이를 통해 동적으로 모듈들을 추가 또는 삭제할 수 있다. 
■ 이렇게 정의된 nn.ModuleList는 자동으로 모델의 파라미터들과 함께 관리되며, 모델의 forward 연산에서 호출될 수 있다. 

■ 예를 들어, nn.ModuleList를 사용하여 여러 개의 레이어(layer)를 동적으로 관리하는 모델을 정의할 수 있다.

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linears = nn.ModuleList()
        for i in range(5):
            self.linears.append(nn.Linear(10, 20))

    def forward(self, x):
        for layer in self.linears:
            x = layer(x)
        return x

- 이 예에서는 nn.ModuleList를 사용하여 nn.Linear() 모듈을 5개 추가한다.

- forward 함수에서 nn.ModuleList에 있는 모듈들을 순차적으로 꺼내 연산을 수행하게 된다.

■ 여기서 구현한 트랜스포머의 인코더 블록을 N개씩 쌓는 것도 같은 방식에 해당한다.

■ TransformerEncoder 클래스의 forwrad() 함수를 보면, 입력(src)를 토큰 임베딩 + 위치 인코딩에 통과시켜 얻은 결과를, 첫 번째 인코딩 블록에 패딩 마스크와 같이 입력으로 넣는 것을 볼 수 있다. 이 과정을 통해 첫 번째 인코더 블록 결과로 나온 output은 다시 두 번째 인코딩 블록에 패딩 마스크와 같이 입력으로 들어가게 된다. 이 과정을 N(n_layers)번 반복해서 인코더의 최종 결과를 반환한다.

---

2.1 Multi-Head Attention

■ 트랜스포머의 멀티 헤드 어텐션은 크게 ① 셀프 어텐션과 ② 크로스 어텐션을 고려하여 구현해야 한다.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.n_head = n_head
        self.attention = ScaleDotProductAttention()

        self.w_q = nn.Linear(d_model, d_model) # Q 생성을 위한 선형 (변환) 층
        self.w_k = nn.Linear(d_model, d_model) # K 생성을 위한 선형 층
        self.w_v = nn.Linear(d_model, d_model) # V 생성을 위한 선형 층

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

    def forward(self, querys, keys, values, mask=None):
        # querys.shape: (batch_size, query_length, d_model)
        # keys.shape: (batch_size, key_length, d_model)
        # values.shape: (batch_size, value_length, d_model)
        batch_size, d_model = querys.shape[0], querys.shape[-1]
        T_q, T_k = querys.size(1), keys.size(1)
 
        querys = self.w_q(querys).view(batch_size, T_q, self.n_head, d_model // self.n_head).transpose(1, 2)
        keys = self.w_k(keys).view(batch_size, T_k, self.n_head, d_model // self.n_head).transpose(1, 2)
        values = self.w_v(values).view(batch_size, T_k, self.n_head, d_model // self.n_head).transpose(1, 2)
        # querys.shape: (batch_size, n_head, T_q, d_models // n_head)
        # keys.shape: (batch_size, n_head, T_k, d_models // n_head)
        # values.shape: (batch_size, n_head, T_k, d_models // n_head)

        output = self.attention(querys, keys, values, mask=mask) # scale dot product attention
        # output.shape: (batch_size, n_head, T_q, d_model // n_head)
        # output = attention weight @ V
		
        B, H, T, C = output.shape	
        output = output.transpose(1, 2).contiguous().view(B, T, H*C)
        # output.shape: (B, T, d_model)
        output = self.w_concat(output)

        return output

■ 그 이유는 위와 같이 순전파 과정(forward())에서 인코더의 셀프 어텐션에 사용하는 Q, K, V는 모두 동일한 입력 시퀀스로부터 계산되지만, 디코더의 크로스 어텐션에서 사용하는 Q는 디코더의 마스크드 멀티 헤드 셀프 어텐션 결과이고, K와 V는 인코더의 결과이기 때문이다.

■ forward() 함수를 보면, (batch_size, query_length, d_model) 크기의 Q와 (batch_size, key_length, d_model) 크기의 K, (batch_size, value_length, d_model) 크기의 V가 스케일드 닷-프로덕트 어텐션 연산을 수행하기 위해 입력으로 들어온다.

■ 그다음, 멀티 헤드 어텐션을 수행하기 위해 Q, K, V를 각각 (batch_size, n_head, T_q, d_models // n_head), (batch_size, n_head, T_k, d_models // n_head), (batch_size, n_head, T_k, d_models // n_head) 크기의 4차원 텐서로 변환한다.

■ 여기서 인코더의 셀프 어텐션만 고려한다면 T_q = T_k = T_v가 성립하지만, 디코더의 크로스 어텐션까지 고려한다면 T_q와 T_k = T_v의 길이가 다를 수 있다는 것을 생각해야 한다.

■ 쿼리, 키, 값의 초기 형상은 (batch_size, seq_len, d_model=512)이지만, 이를 어텐션 헤드별로 분리하면 각 텐서는 (batch_size, self.n_head=8, seq_len, d_model // n_head = 64)의 4차원 형상을 갖게 된다. 

■ 이를 그림으로 표현하자면, 다음과 같이 (seq_len, d_model=512) 행렬을 n_head 개의 (seq_len, d_model // n_head = 64) 행렬로 분리한 것으로 볼 수 있다.

- h = n_head

■ 쿼리, 키, 값 텐서가 위와 같은 구조를 가지므로, 스케일드 닷 프로덕트 어텐션에서 각 어텐션 헤드는 어텐션(셀프 어텐션, 크로스 어텐션) 메커니즘의  병렬 연산이 가능하다. 

■ 그리고 "병렬 어텐션" 연산이 가능하므로, 입력 정보의 서로 다른 특징을 학습할 수 있다는 점을 알 수 있다.  

■ 스케일드 닷-프로덕트 어텐션을 거쳐 나온 어텐션 결과의 크기는 (batch_size, n_head=8, T, d_model // n_head=64)가 된다.

■ 다음 단계는 위의 오른쪽 그림처럼 병렬 어텐션 연산 결과를 하나로 연결(concat)하여 선형 층에 통과시키는 것이다. 이를 위해 어텐션 결과 텐서의 형상을 (batch_size, T, d_model)로 변환한다.

■ 각각 64차원을 가지는 것 8개를 연결하므로, 연결 결과 마지막 차원은 512=d_model가 된다. 초기 쿼리, 키, 값 텐서의 마지막 차원으로 되돌리는 것이다.

■ 이를 통해 입/출력 차원을 유지되므로 트랜스포머 인코더, 디코더 블록을 쉽게 쌓을 수 있다.

---

2.2 Scaled Dot-Product Attention

■ 스케일드 닷 프로덕트 어텐션은 다음과 같이 구현할 수 있다.

import math

class ScaleDotProductAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, querys, keys, values, mask=None):
        # querys.shape == keys.shape == values.shape: (B, n_head, T, d_model // n_head)
        d_k = querys.size(-1)

        # Q와 K^T의 dot product로 어텐션 스코어(유사도)를 계산해야 하므로 transpose(2, 3)
        scores = querys @ keys.transpose(2, 3) / math.sqrt(d_k)

        if mask is not None:
          scores = scores.masked_fill(mask == False, -1e9)

        weights = self.softmax(scores)

        return weights @ values

■ 스케일드 닷 프로덕트 어텐션에서 멀티 헤드 어텐션을 수행하므로, forward() 함수의 입력으로 들어오는 Q, K, V는 (batch_size, n_head, T, d_model // n_head) 형태의 텐서이다.

- Q의 T는 T_q, K와 V의 T는 T_k = T_v이다.

■ 그러므로, 쿼리와 키 간의 유사도(scores)를 계산하기 위해 키의 형상을 (batch_size, n_head, d_model // n_head, T_k) 형태로 변경해준다.

■ (batch_size, n_head, T_q, d_model // n_head=64)와 (batch_size, n_head, d_model // n_haed=64, T_k)의 행렬 곱 연산은 행렬 곱의 조건(첫 번째 행렬의 열과 두 번째 행렬의 행의 크기가 같아야 한다.)을 갖추면 된다. 

■ T_q = 30, t_k = 27이라고 한다면, 어텐션 스코어 행렬(scores)의 형상은 (batch_size, n_head=8, T_q=30, T_k=27)이 된다. 이때, 각 값에는 \( \dfrac{1}{\sqrt{d_k}} \)로 스케일링이 적용된다. 

■ (batch_size, n_head, T_q, d_model // n_head) @ (batch_size, n_head, d_model // n_haed, T_k)의 결과인 (batch_size, n_head, T_q, T_k)는 각 배치 내의 각 헤드에 대해 행렬 곱 연산이 일어난 것이다. 

■ 각 배치 내의 각 헤드별 (T_q, T_k) 행렬은 각 토큰이 다른 모든 토큰에 대해 가지는 어텐션 스코어(유사도)를 나타낸다. 

■ (T_q, T_k) 행렬의 \( (i, j) \) 위치의 값은 \( i \)번째 쿼리 토큰이 \( j \)번째 키 토큰에 대한 유사도를 나타낸다.

- key_1, key_2, ..., key_k는 키의 토큰들, query_1, query_2, ... , query_q는 쿼리의 토큰들

- 이때, 셀프 어텐션 연산을 한다면 Q, K, V는 모두 동일한 시퀀스로부터 얻어졌기 때문에 T_q = T_k가 성립한다.

- 반면, 크로스 어텐션을 한다면 Q는 디코더 K, V는 인코더로부터 얻어진 것이므로, T_q = T_k가 보장되지 않는다. 

■ 이렇게 계산한 어텐션 스코어 행렬에 대하여, forwrad에 전달된 mask가 패딩 마스크이면 <pad> 위치에 마스킹이 적용된다. ■ 정확하게는 <pad> 위치에 음수 값을 넣어서 소프트맥스 함수에 통과시켜 <pad>의 어텐션 가중치를 0을 만들어서 <pad> 토큰의 영향을 받지 않도록 하는 것이다.

■ 어텐션 스코어 행렬에 소프트맥스 함수를 적용할 대상은 키(key)이다. 각 쿼리 토큰은 모든 키 토큰들에 대해 어느 정도의 가중치를 부여할지 결정해야 하기 때문이다. 

■ 즉, 특정 쿼리 토큰 query_i에 대해, 모든 키 토큰 key_1, key_2, ... ,key_k 와의 스코어들을 정규화하여 합이 1이 되는 확률 분포로 만들어야 한다. 이를 통해 각각의 쿼리가 각 키에 얼머나 "집중"할지에 대한 확률을 얻게 된다. 

■ 그러므로 scores 텐서의 (batch_size, n_head, T_q, T_k) 형태에서, 마지막 차원인 T_k를 따라 소프트맥스가 적용되어야 한다. 

■ 그다음 scores 텐서와 V 텐서(값, value)의 행렬 곱 연산을 통해 어텐션 결과를 계산하고, 계산 결과를 반환한다.

■ 이때, scores 텐서는 (batch_size, n_head, T_q, T_k), V 텐서는 (batch_size, n_head, T_v = T_k, d_model / n_head)이므로 어텐션 결과 텐서의 형상은 (batch_size, n_head, T_q, d_model // n_head)가 된다. 

---

2.3 Feed-Forward

■ 피드 포워드 계층은 첫 번째 선형 변환 결과를 활성화 함수에 통과시키고, 활성화 함수를 거쳐 나온 결과에 다시 선형 변환을 하는 방식으로 다음과 같이 구현할 수 있다.

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1): # d_model = 512, d_ff = 2048
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = nn.GELU() # 활성화 함수

    def forward(self, x):
        x = self.linear2(self.dropout1(self.activation(self.linear1(x)))) # 잔차 연결
        x = self.dropout2(x)
        return x

---

2.4 Encoder Layer

■ 인코더 블록(인코더 계층)은 다음과 같이 구현할 수 있다. 

class EncoderLayer(nn.Module):
    def __init__(self, d_model, d_ff, n_head, dropout):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model, device=device)
        self.attention = MultiHeadAttention(d_model=d_model, n_head=n_head)
        self.dropout1 = nn.Dropout(dropout)

        self.norm2 = nn.LayerNorm(d_model, device=device)
        self.feed_forward = PositionwiseFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, src, src_mask):
        _x = src
        norm_x1 = self.norm1(src)
        attn_output = self.attention(querys=norm_x1, keys=norm_x1, values=norm_x1, mask=src_mask)
        x = self.dropout1(attn_output)

        # add
        x = x + _x

        _x = x
        norm_x2 = self.norm2(x)
        x = self.feed_forward(norm_x2)
        x = self.dropout2(x)

        # add
        x = x + _x

        return x

■ forward() 함수를 보면 다음 그림과 같은 사전 정규화 방식으로 진행되며, 인코딩 블록은 입력(src)과 패딩 마스크(src_mask)를 입력으로 받는 것을 볼 수 있다.  잔차 연결은 화살표 모양 그대로 입력을 다시 더해주는 형태로 구현한다.

- 이때의 입력은 토큰 임베딩 + 위치 인코딩을 통과한 임베딩이며, (batch_size, src_seq_len=30, d_model=512)의 형상을 가진다.

■ 인코더 블록의 첫 번째 과정을 보면, 층 정규화를 적용한 입력 시퀀스(norm_x1)를 querys, keys, values로 사용하고 패딩 마스크를 적용한 상태에서 어텐션 연산을 수행하는 것을 볼 수 있다. 

■ 먼저 인코더 블록의 첫 번째 sub layer인 멀티 헤드 어텐션이 사전 정규화를 통해 진행되는 과정을 보면,

- _x = src에서 _x는 (batch_size, src_seq_len, d_model=512),

- LayerNorm을 통과한 norm_x1도 (batch_size, src_seq_len, d_model=512)이다.

- 그리고 멀티 헤드 어텐션 결과인 attn_output도 (batch_size, T_q = src_seq_len, d_model=512)이므로

- 잔차 연결(x = x + _x)의 결과 텐서도 (batch_size, src_seq_len, d_model=512)이 된다.

■ 첫 번째 sub layer 결과 x는 다시 다음 그림처럼 피드 포워드 계층의 결과와 잔차 연결로 연결된다.

- 입력으로 들어오는 x는 (batch_size, src_seq_len, d_model=512), _x = x에서 _x는 (batch_size, src_seq_len, d_model=512),

- LayerNorm을 통과한 norm_x2도 (batch_size, src_seq_len, d_model=512)이다.

- FFN에서 (batch_size, src_seq_len, d_model=512) \( \rightarrow \) (batch_size, src_seq_len, d_ff=2048) \( \rightarrow \) (batch_size, src_seq_len, d_model=512)가 된다.

- 그러므로, 잔차 연결(x = x + _x)의 결과 텐서도 (batch_size, src_seq_len, d_model=512)이 된다.

■ 여기 까지의 과정이 TransformerEncoder 클래스의 forward() 에서 for 문을 1번 진행한 결과이다. 즉, 하나의 인코더 블록(인코더 계층)이 수행된 결과이다.

    def forward(self, src, src_mask):
        output = self.emb(src)

        for mod in self.layers:
            output = mod(output, src_mask)

■ 이 과정을 n_layers만큼 반복한다. 

---

---

3. Transformer Decoder

■ 트랜스포머의 디코더 블록의 첫 번째 sub layer에서도, 인코더 블록(Encoder Layer)처럼 멀티 헤드 셀프 어텐션을 수행한다. 

■ 단, 이때 룩-어헤드 마스크를 통해 마스크드 멀티 헤드 셀프 어텐션을 수행한다. 이것이 인코더 블록과의 첫 번째 차이점이다.

■ 두 번째 차이점은 크로스 어텐션이다. 크로스 어텐션에서 Q는 trg(위 그림에서는 Outputs)의 토큰 임베딩과 위치 인코딩을 더한 후, 마스크드 멀티헤드 셀프 어텐션 계층을 통과한 결과이다. 

■ 그리고 K와 V는 트랜스포머 인코더의 최종 출력 결과이다. 

■ 이후 과정은 인코더 블록과 동일하게 피드 포워드 계층을 통과한다. 마찬가지로 디코더 블록을 N개 쌓을 경우, 이 과정을 N(n_layers)번 반복한다.

class TransformerDecoder(nn.Module):
    def __init__(self, trg_vocab_size, max_len, d_model, d_ff, n_head, n_layers, dropout, device):
        super().__init__()
        self.emb = TransformerEmbedding(d_model=d_model,
                                        max_len=max_len,
                                        vocab_size=trg_vocab_size,
                                        dropout=dropout,
                                        device=device)

        self.layers = nn.ModuleList([DecoderLayer(d_model=d_model,
                                                  d_ff=d_ff,
                                                  n_head=n_head,
                                                  dropout=dropout)
                                         for _ in range(n_layers)])

        self.linear = nn.Linear(d_model, trg_vocab_size)

    def forward(self, trg, src, trg_mask, src_mask):
        trg = self.emb(trg)

        for mod in self.layers:
            trg = mod(trg, src, trg_mask, src_mask)
        output = self.linear(trg)

        return output

■ 이때, 디코더의 마스크드 멀티 헤드 어텐션을 진행할 때, 룩-어헤드 마스크와 패딩 마스크가 같이 적용되어야 한다. 룩-어헤드 마스크를 적용했다고 해서, <pad> 토큰의 영향을 무시하기 위해 사용하는 패딩 마스크가 필요 없는 것이 아니기 때문이다.

■ TransformerDecoder의 forward에서 나온 최종 결과인 output은 선형 층과 Softmax를 거쳐 최종 예측을 계산하게 된다. 이때 코드에서는 선형 층만 구현하고 Softmax 함수는 구현하지 않았는데, 이는 손실 함수로 CrossEntropyLoss를 사용할 것이기 때문이다.

- CrossEntropyLoss는 내부적으로 LogSoftmax와 NLLLoss를 함께 처리하므로, 따로 Softmax를 적용할 필요가 없다. 

---

3.1 Decoder Layer

■ 디코더 블록 구현은 다음과 같다.

class DecoderLayer(nn.Module):
    def __init__(self, d_model, d_ff, n_head, dropout):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model, device=device)
        self.attention = MultiHeadAttention(d_model=d_model, n_head=n_head)
        self.dropout1 = nn.Dropout(dropout)

        self.norm2 = nn.LayerNorm(d_model, device=device)
        self.cross_attention = MultiHeadAttention(d_model=d_model, n_head=n_head)
        self.dropout2 = nn.Dropout(dropout)

        self.norm3 = nn.LayerNorm(d_model, device=device)
        self.feed_forward = PositionwiseFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, trg, src, trg_mask, src_mask):
        _x = trg
        norm_x1 = self.norm1(trg)
        attn_output = self.attention(querys=norm_x1, keys=norm_x1, values=norm_x1, mask=trg_mask)
        x = self.dropout1(attn_output)
        # add
        x = x + _x

        _x = x
        norm_x2 = self.norm2(x)
        cross_attn_output = self.cross_attention(querys=norm_x2, keys=src, values=src, mask=src_mask)
        x = self.dropout2(cross_attn_output)
        # add
        x = x + _x

        _x = x
        norm_x3 = self.norm3(x)
        x = self.feed_forward(norm_x3)
        x = self.dropout3(x)
        # add
        x = x + _x

        return x

■ 인코더 블록과 마찬가지로 사전 정규화 방식을 적용했으며, 마스크드 멀티 헤드 어텐션을 수행할 때 룩-어헤드 마스크가 필요하므로 해당 마스크를 전달하고, 셀프 어텐션을 수행할 때 룩-어헤드 마스크는 필요 없으므로 패딩 마스크만 전달한다.

■ 또한 인코더 블록과 동일하게, 입/출력 차원이 d_model로 유지된다.

---

---

4. Transformer

■ 아래의 코드는 Seq2Seq처럼 트랜스포머 인코더와 디코더를 연결한 것이다.

class Transformer(nn.Module):
    def __init__(self, src_pad_idx, trg_pad_idx, trg_sos_idx, src_vocab_size, trg_vocab_size, d_model, n_head, max_len,
                 d_ff, n_layers, dropout, device):
        super().__init__()
        self.src_pad_idx = src_pad_idx
        self.trg_pad_idx = trg_pad_idx
        self.trg_sos_idx = trg_sos_idx
        self.device = device
        
        self.encoder = TransformerEncoder(d_model=d_model,
                               n_head=n_head,
                               max_len=max_len,
                               d_ff=d_ff,
                               src_vocab_size=src_vocab_size,
                               dropout=dropout,
                               n_layers=n_layers,
                               device=device)

        self.decoder = TransformerDecoder(d_model=d_model,
                               n_head=n_head,
                               max_len=max_len,
                               d_ff=d_ff,
                               trg_vocab_size=trg_vocab_size,
                               dropout=dropout,
                               n_layers=n_layers,
                               device=device)

    def forward(self, src, trg):
        src_mask = self.make_src_mask(src)
        trg_mask = self.make_trg_mask(trg)
        encoder_output = self.encoder(src, src_mask)
        output = self.decoder(trg, encoder_output, trg_mask, src_mask)
        return output

    def make_src_mask(self, src):
        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
        return src_mask

    def make_trg_mask(self, trg):
        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(3)
        trg_len = trg.shape[1]
        trg_sub_mask = torch.ones((trg_len, trg_len), dtype = torch.bool).tril().to(self.device)
        trg_mask = trg_pad_mask & trg_sub_mask
        return trg_mask

■ forward()를 보면 트랜스포머 인코더에 입력 시퀀스(src)와 패딩 마스크를 전달해서 인코더의 최종 결과인 encoder_output을 계산한다. 

■ 그다음, 디코더에 룩-어헤드 마스크와 패딩 마스크 그리고 타겟 시퀀스(trg)와 encoder_output를 전달한다. encoder_output을 전달하는 이유는 크로스 어텐션을 계산하기 위해서이다. 

■ 패딩 마스크를 만드는 함수는 make_src_mask이며, 룩-어헤드 마스크를 만드는 함수는 make_trg_mask이다.

■ 먼저, 패딩 마스크를 만드는 방법은 입력으로 들어오는 batch_size 개수의 src_len 차원을 가지는 각각의 텐서에 대해, 텐서의 각 요소가 self.src_pad_idx와 같지 않은지 위와 같이 비교하는 것이다.

■ 이를 통해 패딩 토큰이 아닌 위치는 True가 되고, 패딩 토큰인 위치는 False가 된다.

■ 예를 들어, 배치 크기가 2이고, src 텐서가 다음과 src_len=4의 차원을 갖는다고 하자. 패딩 토큰의 정수 인덱스(또는 아이디)가 0이라면, 

src_pad_idx = 0

src = torch.tensor([
    [1, 2, 3, 0],  # 첫 번째 시퀀스: 마지막 토큰이 패딩
    [4, 5, 0, 0]   # 두 번째 시퀀스: 마지막 두 토큰이 패딩
])

src.shape, src
```#결과#```
(torch.Size([2, 4]),
 tensor([[1, 2, 3, 0],
         [4, 5, 0, 0]]))
````````````

■ 다음과 같이 src 텐서의 각 요소 중 패딩 토큰(정수 0)인 요소는 False로 마스킹된다. 

(src != src_pad_idx)
```#결과#```
tensor([[ True,  True,  True, False],
        [ True,  True, False, False]])
````````````

■ 이때 인코더의 패딩 마스크(src_mask)를 만들 때, unsqueeze(1), unsqueeze(2)를 해서, (batch_size, 1, 1, src_len)으로 차원을 맞춰준다.

- 인코더의 패딩 마스크는 셀프 어텐션에 사용할 패딩 마스크를 의미한다.

■ 이 마스크를 멀티 헤드 어텐션 내부에서 사용할 때, 어텐션 스코어 텐서 (batch_size, n_head, T_q=query_len, T_k=key_len)와 브로드캐스팅을 통해 연산될 수 있도록 하기 위해서이다.

■ src_mask의 shape을 (batch_size, 1, 1, src_len)으로 만들면, masked_fill() 연산 시, num_head 차원과 T_q(=query_len) 차원에 대해 브로드캐스팅되어 마스킹이 적용된다. 즉, 모든 어텐션 헤드의 모든 행(어텐션 스코어 벡터)에 대해 동일한 마스킹이 적용되는 것이다.

■ 예를 들어, 어테션 스코어 텐서가 scores = torch.rand(2, 3, 2, 4)이라고 하자. 그렇다면, 2개의 배치에 각각 3개의 어텐션 헤드가 있으며, 이 어텐션 헤드는 (T_q, T_k) = (2, 4) 행렬이다.  

■ 만약, 첫 번째 배치의 <pad> 토큰 위치는 네 번째 원소, 두 번째 배치의 <pad> 토큰 위치는 세 번째와 네 번째 원소라고 한다면, 차원을 확장한 src_mask는 다음과 같을 것이다. 

(torch.Size([2, 1, 1, 4]),
 tensor([[[[ True,  True,  True, False]]],
 
 
         [[[ True,  True, False, False]]]]))

■ 그러므로, scores.masked_fill(src_mask==False, -1e9)을 적용했을 때,
- 첫 번째 배치의 모든 헤드(3개 헤드)의 (T_q, T_k) 행렬에서 <pad> 토큰의 위치(네 번째 원소)는 -1e9로 마스킹되며,
- 두 번째 배치의 모든 헤드(3개 헤드)의 (T_q, T_k) 행렬에서 <pad> 토큰의 위치(세 번째와 네 번째 원소)는  -1e9로 마스킹된다.

■ 즉, 어텐션 스코어 벡터에서 어떤 쿼리 위치(어텐션 스코어 벡터의 행)에서든, 키의 특정 위치(어텐션 스코어 벡터의 열) j가 패딩이라면, 그 위치 j는 어텐션 대상에서 제외시키기 위해 unsqueeze(1), unsqueeze(2)로 src_mask의 차원을 확장시키는 것이다. 이를 통해, 쿼리가 무엇이든 상관없이 키의 패딩된 부분은 마스킹이 적용된다.

- 어텐션 스코어 행렬에서 패딩된 부분은 음수 값으로 변경

- \( \rightarrow \) 어텐션 스코어 행렬이 소프트맥스를 통과하면 음수 값이 있는 위치는 0으로(또는 0에 가까운 값) 변환. 즉, 해당 위치의 어텐션 가중치는 0

- \( \rightarrow \) 어텐션 가중치와 값(value)의 가중합을 수행. 결과적으로 패딩된 부분은 어텐션 연산에서 제외

scores.masked_fill(src_mask==False, -1e9)
```#결과#```
tensor(
		## 첫 번째 배치	
        # 어텐션 헤드 1
		[[[[ 4.9098e-01,  4.1584e-01,  7.9349e-01, -1.0000e+09],
          [ 2.8226e-01,  5.8958e-01,  3.8476e-01, -1.0000e+09]],
          
		# 어텐션 헤드 2
         [[ 8.8747e-01,  1.7320e-02,  8.4008e-01, -1.0000e+09],
          [ 4.8746e-01,  1.0147e-01,  9.1344e-01, -1.0000e+09]],
		# 어텐션 헤드 3
         [[ 5.9229e-01,  3.8054e-01,  5.3053e-01, -1.0000e+09],
          [ 1.4220e-01,  6.2324e-01,  7.6716e-01, -1.0000e+09]]],
		
		## 두 번째 배치
        # 어텐션 헤드 1
        [[[ 7.4712e-01,  4.1168e-01, -1.0000e+09, -1.0000e+09],
          [ 9.6008e-01,  2.6578e-01, -1.0000e+09, -1.0000e+09]],
		# 어텐션 헤드 2
         [[ 3.8476e-01,  5.9740e-01, -1.0000e+09, -1.0000e+09],
          [ 9.5313e-01,  8.0833e-01, -1.0000e+09, -1.0000e+09]],
		# 어텐션 헤드 3
         [[ 8.1780e-01,  8.1319e-01, -1.0000e+09, -1.0000e+09],
          [ 5.1216e-01,  9.1919e-01, -1.0000e+09, -1.0000e+09]]]])
````````````

 

■ 다음은 룩-어헤드 마스크를 만드는 방법이다.

■ 디코더의 마스크드 셀프 어텐션에서 룩-어헤드 마스크의 기능은 패딩 토큰을 마스크하는 것이 아니다. 그러므로 마스크드 멀티 헤드 셀프 어텐션 계층에 룩-어헤드 마스크와 패딩 마스크를 같이 전달해야 한다. 이는 룩-어헤드 마스크에 패딩 마스크 기능을 포함시키는 방법으로 구현할 수 있다. 

■ 마찬가지로 타겟 시퀀스에서 패딩이 아닌 위치는 True, 패딩인 위치는 False인 타겟 시퀀스의 패딩 마스크(trg_pad_mask)를 만든다.

■ 룩-어헤드 마스크의 크기는 trg_len x trg_len이다. 룩-어헤드 마스크를 어텐션 스코어 행렬에 적용할 텐데, 어텐션 스코어 행렬의 행을 쿼리, 열을 키로 본다면, 셀프 어텐션이므로 동일한 쿼리와 키를 보는 것이다. 그러므로 룩-어헤드 마스크의 크기는 trg_len x trg_len이다.

- trg_len은 타겟 시퀀스(trg)의 길이

■ 그래서 룩-어헤드 마스크는 torch.ones((trg_len, trg_len), dtype = torch.bool).tril()을 통해 (trg_len x trg_len) 크기를 가지며, 주대각선 위쪽이 모두 False, 주대각선부터 주대각선 아래쪽이 모두 True인 마스크 행렬을 만들어야 한다.

torch.ones((4, 4), dtype = torch.bool).tril()
```#결과#```
tensor([[ True, False, False, False],
        [ True,  True, False, False],
        [ True,  True,  True, False],
        [ True,  True,  True,  True]])
````````````

■ 최종적으로 trg_pad_mask &(and 연산) trg_sub_mask을 통해 패딩 마스크가 적용된 룩-어헤드 마스크를 반환한다.

■ 타겟 시퀀스의 패딩 마스크(trg_pad_mask)는 unsqueeze(1), unsqueeze(3)로 (batch_size, 1, trg_len, 1) 차원이 확장되며, (trg_len, trg_len) 크기의 룩-어헤드 마스크(trg_sub_mask)와 &(and) 연산을 하게 된다. 

■ 마찬가지로 브로드캐스팅 연산이 진행된다. 브로드캐스팅 규칙은 각 텐서는 최소한 1차원 이상을 가지고 있어야 하며(빈 텐서는 사용할 수 없다.) 두 텐서의 각 차원 크기가 다음 조건을 만족하는지 비교한다.

- 각 차원이 서로 동일합니다, 또는

- 각 차원중의 하나의 크기가 반드시 1입니다, 또는

- tensor들 중 하나의 차원이 존재하지 않습니다.

- 이때, "차원을 비교하는 순서는 맨 뒤(마지막 차원)에서부터 맨 앞(첫 번째 차원)으로"이다.

■ 현재 trg_pad_mask와 trg_sub_mask의 차원은 다음과 같다.

trg_pad_mask.shape:  (batch_size, 1, trg_len, 1)
trg_sub_mask.shape: 		  (trg_len, trg_len)

■ 생략된 부분은 자동으로 차원을 맞춰준다.

trg_pad_mask.shape:  (batch_size, 1, trg_len, 1)
trg_sub_mask.shape:	 	(1, 1, trg_len, trg_len)

■ 그리고 & 연산 시, trg_pad_mask는 (batch_size, 1, trg_len, 1 -> trg_len)으로 마지막 차원이 확장되고,

trg_sub_mask는 (1 -> batch_size, 1, trg_len, trg_len)으로 확장된다. 그러므로 & 연산 결과 텐서는 (batch_size, 1, trg_len, trg_len) 형태가 된다.

■ 이때, 결과 텐서 (batch_size, 1, trg_len, trg_len)에서 (trg_len, trg_len) 행렬은 패딩 마스크와 룩-어헤드 마스크가 모두 적용된 상태이다.

■ 이제 이 결과 텐서(trg_mask)는 마스크드 멀티 헤드 '셀프 어텐션'에서 어텐션 스코어 텐서 (batch_size, n_head, T_q, T_k)에 적용된다.

■ 이때, 모든 어텐션 헤드에 trg_mask가 적용되는데, "셀프 어텐션"이므로 T_q = T_k = trg_len이다. 그러므로 (trg_len, trg_len) 크기의 마스크(trg_mask)를  어텐션 스코어 텐서 (batch_size, n_head, T_q, T_k)에 적용할 수 있는 것이다.

---

참고) Tensor는 서로 다른 ndim에 대해서 어떻게 연산할까? (Broadcasting Semantics)

Tensor는 서로 다른 ndim에 대해서 어떻게 연산할까? (Broadcasting Semantics)

참고) https://tutorials.pytorch.kr/beginner/introyt/tensors_deeper_tutorial.html

Pytorch Tensor 소개