본문 바로가기
ML&DL/NLP(LLM)

[NLP][Code Review] Attention is all you need pytorch 코드 리뷰

by 거북이주인장 2024. 3. 30.

Paper motivation

  • RNN은 sequential한 데이터를 처리하는데 있어서 시간의 순서에 따른 hidden state인 $h_{t}$을 $h_{t-1}, t$의 함수로 만든다.
  • 즉, 현재의 $h_t$을 만들기 위해 이전 hidden state인 $h_{t-1}$이 있어야 하고 이러한 순차적이고 의존적인 특성이 병렬 학습을 하는데 큰 걸림돌이 된다. 앞 시점의 연산이 끝나지 않은 경우, 그 뒤의 연산을 수행할 수 없기 때문이다.
  • transformer는 이러한 recurrence 구조를 버리고 오직 attention 구조만 사용하여 input과 output의 global dependency를 학습하고자 한다.
  • 또한 recurrence 구조의 순차적인 특징이 없기 때문에 병렬 학습이 가능하고 이는 학습시간의 단축으로 이어진다.

My motivation

LLM이 요즘 매우 핫하다. openai의 chatgpt부터 해서 google의 gemini, claude sonnet, 오픈소스 모델인 llama, gemma, mistral, mixtral까지 엄청나게 다양한 LLM이 쏟아져 나오고 있다. 필자는 어쩌다보니 회사에서 LLM 관련된 개발을 하고 있는데, 빠르게 변화하고 진화하고 있는 LLM 필드에서 무언가 부족하다는 생각이 들었다. 그것은 LLM의 근간이 되는 transformer에 대한 이해.

NLP 분야에서 영향력이 가장 큰 논문 중 하나를 뽑으라면, attention is all you need 논문일 것이다. 이 연구에 대한 철저한 이해가 없이 LLM을 개발하고 있었으니.. 사실 어떤 LLM이 더 좋다고는 하는데 어떤 차이가 있어서 왜 더 좋은지 정확하게 설명할 수가 없었다. 이런 부족함을 채우기 위해 attention is all you need pytorch 구현 코드를 하나씩 살펴보고 논문의 내용과 연결해보는 과정이 필요하다고 생각했다.

transformer pytorch 구현체는 많이 있지만 그 중 가장 유명한 아래 깃헙 레포를 전적으로 참고했다.

https://github.com/jadore801120/attention-is-all-you-need-pytorch/tree/master

 

GitHub - jadore801120/attention-is-all-you-need-pytorch: A PyTorch implementation of the Transformer model in "Attention is All

A PyTorch implementation of the Transformer model in "Attention is All You Need". - jadore801120/attention-is-all-you-need-pytorch

github.com

또한 본문의 내용 및 그림은 아래 블로그 글에서 많이 참고했다. 여러 transformer 설명 글을 봤지만, 아래 블로그 글이 가장 설명을 잘해둔 곳이라고 생각한다.

https://cpm0722.github.io/pytorch-implementation/transformer

이제 본격적으로 들어가보자. 모델을 만들기 위한 단계는 크게 1) 데이터 준비 및 전처리 2) training 3) test로 구성된다.

환경 세팅 및 데이터 준비

먼저 환경설정을 해보자. 레포에서는 python 3.6에 torch을 1.3.1로 설치하라고 나온다. 근데 torch 버전이 많이 업그레이드 되었으므로 최근 버전에 맞춰서 가상 환경을 설정해본다. 필자는 python 3.8으로 세팅하였다.

torch==2.2.1
spacy==2.3.5
dill==0.3.3
torchtext==0.4.0

가상 환경을 설정하고 spacy 언어 모델을 다운받자.

> python -m spacy download en
> python -m spacy download de

그리고 preprocess.py을 실행하면 아래와 같이 오류가 발생한다.

> python preprocess.py -lang_src de -lang_trg en -share_vocab -save_data m30k_deen_shr.pkl
# requests.exceptions.SSLError: 
# HTTPSConnectionPool(host='www.quest.dcs.shef.ac.uk', port=443): 
# Max retries exceeded with url: /wmt16_files_mmt/training.tar.gz 
# (Caused by SSLError(CertificateError("hostname 'www.quest.dcs.shef.ac.uk' doesn't match either of '*.dcs.shef.ac.uk', '*.dcs.sheffield.ac.uk'")))

데이터를 다운 받는데 오류가 발생한 것으로 보인다. 이것저것 찾아보다가, 그냥 아래 깃헙에서 직접 다운받는게 빠르다고 판단하여, 압축 파일을 다운받고 .data/multi30k 디렉토리에 옮겼다.

 

dataset/data/task1/raw at master · multi30k/dataset

Multi30k Dataset. Contribute to multi30k/dataset development by creating an account on GitHub.

github.com

다운 받아야할 파일은 train.en, train.de, val.en, val.de, test2016.en, test2016.de이다.

이제 preprocess.py을 실행하면 m30k_deen_shr.pkl 파일이 만들어진 것을 확인할 수 있다.

> python preprocess.py -lang_src de -lang_trg en -share_vocab -save_data m30k_deen_shr.pkl
# /Users/user/personal/.venv/lib/python3.8/site-packages/urllib3/__init__.py:35: 
# NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+, currently the 'ssl' module is compiled with 'LibreSSL 2.8.3'. See: https://github.com/urllib3/urllib3/issues/3020 warnings.warn(
# Namespace(data_src=None, data_trg=None, keep_case=False, lang_src='de', lang_trg='en', max_len=100, min_word_count=3, save_data='m30k_deen_shr.pkl', share_vocab=True)
# [Info] Get source language vocabulary size: 5375
# [Info] Get target language vocabulary size: 4556
# [Info] Merging two vocabulary ...
# [Info] Get merged vocabulary size: 9521
# [Info] Dumping the processed data to pickle file m30k_deen_shr.pkl

Transformer의 큰 구조

transformer의 큰 구조이다. 입력 문장을 transformer라는 함수에 넣으면 출력 문장이 나온다. 이 경우에 해당하는 task가 여러개가 있을 것 같은데 논문에서는 기계 번역을 위해 transformer 구조를 사용했다. 즉, 불어를 영어로 번역하는 task에서는 input sentence가 불어로, output sentence가 영어로 나오는 것이다. 중간에 transformer라는 함수를 통과하면 영어 문장이 나온다.

논문에 나온, 조금 더 상세한 구조는 아래와 같다.

크게 encoder / decoder로 나누어서 본다. encoder는 입력 문장을 받고 decoder는 출력 문장뿐만 아니라 encoder에서 나온 결과물도 입력으로 받는다. 이를 쉽게 표현하면 아래와 같다.

encoder는 처음 문장을 받고 context를 생성하는데 이 과정을 encoding이라고 한다. context란 무엇일까? 문장을 고차원적으로 압축해서 담은 벡터라고 생각하면 된다. 예를 들어, 나는 오늘 점심에 도시락을 먹었는데 그게 가격 대비 상당히 맛있었어 라는 문장을 생각해보자. 이 문장에는 1) 나는 오늘 점심으로 도시락을 먹었다. 2) 그 도시락이 맛있었다. 3) 그 도시락이 가격도 괜찮았다. 라는 의미를 담고 있다. 인간은 이러한 문맥적인 의미를 자연어를 통해 바로 알 수 있지만 컴퓨터나 기계는 알 수 없다. 따라서 이 문장이 가지고 있는 의미를 context로 압축하여 위의 세가지 의미를 담고자 하는 것이다.

decoder는 encoder와 반대이다. decoder는 context을 받고 출력 문장을 최종적으로 만들어낸다. 조금 더 상세하게 들어가보면 context 뿐만 아니라 출력으로 생산하는 문장을 right shift한 문장도 함께 받기는한데 이는 이후에 조금 더 자세하게 살펴보자. 지금으로써는 decoder가 context와 문장을 입력으로 받고 출력 문장을 만들어내는 함수 정도로 이해하자.

 

이제 transformer가 어떻게 input을 받는지부터 output을 만드는 과정까지 end to end로 그 과정을 살펴보자.

Token preprocessing

transformer는 input으로 text sequence을 받는다. 근데 sequence의 길이가 엄청 길수도, 짧을수도 있다. 따라서 max_length을 정하여 그 길이보다 길다면 자르고, 짧다면 <PAD>을 통해서 빈 공간을 채워둔다. 또한 기존의 bag of words에 존재하지 않은 단어가 나온다면 그 단어를 <UNK>로 치환한다. 여기서 encoding/decoder 각각에 들어가는 sequence을 전처리하는 과정이 살짝 다르므로 나눠서 살펴보자.

Encoder sequence

구체적인 예시로 들어가보자. 아래의 세 문장이 있을 때, max_length=6 기준으로 encoder에 넣기 전에 아래의 전처리 과정을 거친다.

  • i am looking for happiness: $\{ x_1, x_2, x_3, x_4, x_5, <\text{/s}> \}$
  • i am bohyun:$\{ x_1, x_2, x_6, \text{</s>, <PAD>, <PAD>} \}$
  • i am @#: $\{ x_1, x_2, \text{<UNK>, </s>, <PAD>, <PAD> } \}$

문장의 끝에 </s>을 넣음으로써 문장이 종료됐다는 것을 표시한다. max_length=6 기준으로 그 길이가 짧다면 <PAD>을 채워 넣어준다.

Decoder sequence

Encoder에 들어가는 입력은 번역하기 이전 문장이라면 decoder에 들어가는 입력은 번역 후의 문장이다. 구체적인 예시로 살펴보자.

  • 나는 행복을 찾고 있다: $\{ \text{<s>}, x_1, x_2, x_3, x_4, \text{</s>}\}$
  • 나는 보현이다.: $\{ \text{<s>}, x_1, x_5, \text{</s>}, \text{<PAD>}, \text{<PAD>} \}$
  • 나는 @#: $ \{ \text{<s>}, x_1, \text{<UNK>}, \text{</s>} ,\text{<PAD>}, \text{<PAD>} \} $

Decoder sequence에는 문장의 시작점에 을 넣고 끝나는 부분에 을 넣는다. 문장의 시작점에 <s>을 넣느냐의 여부만 encoder sequence의 경우와 다르고 나머지는 동일하게 처리한다.

Token embedding

이제 encoder/decoder에서 단어를 임베딩하는부분을 살펴보자. embed_size만큼의 각 단어별 임베딩 벡터 look-up table이 있고 거기서 각 단어에 맞는 임베딩 벡터를 찾는다. (이때, 주로 어떤 임베딩 벡터를 사용하는지 나중에 알아보자)

예를 들어, looking이라는 단어에 대한 임베딩 벡터를 찾고자 하면 vocab_size X embed_size 만큼의 임베딩 행렬에서 looking에 맞는 임베딩 벡터를 가져오는 것이다.

Positional encoding

positional encoding이란 무엇이며 왜 하는 것일까? 뒤에서 자세하게 살펴보겠지만, transformer의 근간이 되는 attention 연산에는 단어의 sequence에 대한 정보가 반영되어 있지 않다. 물론 decoder에서 self-attention하는 부분에는 subsequent masking을 통해 이를 어느정도 구현하는데, 그 이외의 attention에는 sequence에 대한 정보가 부족하다. RNN은 이전 단어에 대해 의존적인 대신에 이런 sequence 정보가 제대로 들어간다. transformer는 이러한 단점을 보완하기 위해 positional encoding을 사용한다.

 

positional encoding이란 문장 내에서 단어의 위치에 따라서 단어에 부여하는 일종의 서로 다른 값이다. 논문에서 아래와 같이 정의한다.

'pos'는 단어의 인덱스이고, '2i'는 임베딩 벡터의 인덱스이다. 즉, 그 단어가 문장 내에서 어떤 위치에 있는지에 따라서, 그리고 그 단어의 임베딩 벡터 인덱스에 따라서 삼각함수 값을 가지며, 이는 곧 주기적인 값이므로 일종의 '위치'에 대한 값을 의미할 수 있다.

 

max_length=512, embed_size=1024인 경우에 positional encoding에 대한 heatmap을 예시로 살펴보자. 구체적인 코드는 바로 다음에 살펴본다.

x축은 embed_size, y축은 max_length을 의미한다. y축에 따라서 sequence 내의 단어 위치가 달라지는 것이다. 즉, 수평으로 값들이 한 단어가 가지는 positional encoding 값인데 y축에 따라서 그 값이 달라지는 것을 알 수 있다. 이와 같이 sequence 내에서 단어가 어디에 위치하느냐에 따라 일종의 주기적인 값을 부여하는 것이다.

 

그러면 이를 코드로 어떻게 구현하는지 살펴보자.`init` 함수와 하위 메써드를 살펴보자.

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()

        # Not a parameter
        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))
   
    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ...

    def forward(self, x):
        ...
  • `d_hid`: 임베딩 벡터의 차원
  • `n_position`: max_length 값

먼저 `init` 함수이다. super을 통해서 PositionalEncoding 클래스가 상속받는 nn.Module의 init 함수를 실행한다.

중요한 것은 `self.register_buffer`의 역할이다. 모델을 학습할 때, sinusoid_table은 파라미터로 인식되면 안 되는데 이 역할을 `self.register_buffer`가 수행한다. (관련하여 나중에 따로 포스팅을 작성한다)

 

다음으로 `_get_sinusoid_encoding_table` 메써드를 살펴보자.

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        ...

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ''' Sinusoid position encoding table '''
        # TODO: make it with torch instead of numpy

        def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)

    def forward(self, x):
        ...

논문에서 홀수번째/짝수번째 positional encoding 값은 아래와 같이 정의된다.

삼각함수 안에 있는 것을 코딩한 것이 `[position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]` 이다. 

PositionalEncoding 클래스는 만들어질 때 `init` 함수에서 바로 `self.register_buffer`가 실행되면서 `self.pos_table`에 sinusoid_table 값이 저장된다. 그렇기 때문에 클래스를 만들자마자 바로  input을 넣을 수 있다.

`_get_sinusoid_encoding_table` 메써드를 통해 (1, 200, 512) 차원의 sinusoid_table이 만들어진다. 여기서 차원은 각각 (batch_size, max_length, embedding_size)이다. sequence의 길이가 최대 200이므로, 200차원까지 만들어두는 것이고 만약 그 길이가 200보다 작으면 잘라서 사용한다. 이는 `forward` 메써드에서 확인할 수 있다.

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        ...

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ...

    def forward(self, x):
        return x + self.pos_table[:, :x.size(1)].clone().detach()
  • `x`: (batch_size, sequence_length, embedding_size)인 입력 데이터

`forward` 메써드에서 sequence_length는 max_length인 200보다 작을 수 있다. 따라서 `x.size(1)`까지의 값을 self.pos_table에서 가져온다. 그리고 입력값에 더한다. 이로써 임베딩 벡터에 positional encoding을 더하는 과정이 완성된다.

Encoder

Encoder에 input으로 넣기 위한 전처리 단계와 positional encoding 처리에 대해 알아보았다. 그러면 이제 본격적으로 encoder을 알아보자. Encoder는 6개의 Encoder Block으로 구성되어있다. 논문에서는 레이어의 개수가 $N=6$이라고 나와있다. 첫번째 encoder block은 입력 문장을 받고 context을 출력한다. 두번째 encoder block은 첫번째 encoder block의 context을 받고 또 다른 context을 출력한다. 이런 과정이 6번 반복되는 것이다. 

 

그러면 왜 이런 과정을 6번이나 반복하는 것일까? 조금 추상적으로 들릴 수 있겠지만, 입력으로 들어오는 문장이 가지는, 더 높은 차원의 context을 담기 위한 것이다. 첫번째 encoder block은 입력 문장을 바로 받기 때문에 아무래도 낮은 수준의 context일텐데, encoder block을 쌓다보니까 context, context에 대한 context... 이런식으로 점차 높은 차원의 의미를 가지는 context가 만들어진다. 일종의 인간이 아니라 기계가 이해하기 쉽게 만드는 과정이라고 생각하면 된다.

 

그럼 하나의 encoder block은 어떻게 구성되어 있을까? 아래와 같이, Multi-head attention layer와 point-wise feed-forward layer로 구성되어 있다.

이 두 layer가 의미하는 바는 무엇일까? Multi-head attention을 먼저 살펴볼 것인데, 그 전에 먼저 attention 자체의 개념부터 살펴보자.

Attention

Multi-head attention은 무엇일까? scaled-dot product attention을 head의 개수만큼 (multi-head) 병렬적으로 수행하는 것을 의미한다. transformer의 핵심인 병렬 연산을, 바로 이 Multi-head attention (이하 MHA)이 가능하게 하는 것이다. 그렇다면 더 작게 들어가서 attention(scaled-dot product attention이랑 동치)이란 무엇일까? 유명한 아래 문장을 통해 attention의 개념을 잡아보자.

The animal didn't cross the street, because it was too tired.

 

여기서 'it'이 의미하는 것이 무엇일까? street일까? 인간은 자연어를 읽으면서 'it'을 'the animal'과 자연스럽게 연관시키지만 기계는 처음부터 그렇게 하지 못한다. 따라서 'it'이 문장의 다른 어떤 단어들과 연관이 있는지 자체적으로 학습을 통해 알아내야 한다. 'it'이 'the animal'과 관련있는 정도,  'it'이 'the street'와 관련있는 정도 등, 이 양적인 개념을 'attention'이라고 부른다. 다시 말해, 같은 문장 내에서 서로 다른 token끼리 연관있는 정도를 self-attention으로 정의한다. self가 붙은 이유는, 하나의 문장 내에서 각 token끼리의 연관성을 보기 때문이다. 이게 아니라 서로 다른 두 문장의 token이 각각에 가지는 연관성을 계산하는 것은 cross-attention이라고 한다. 말 그대로, 연관있는 정도를 '문장 간에' 보기 때문이다.

 

그러면 기존에 NLP에서 RNN을 많이 사용했는데, attention을 사용함으로써 얻는 이득은 무엇일까? 논문에서 나왔듯이, RNN은 순차적인 모델링을 하기 때문에 현재의 hidden state을 계산하기 위해서 반드시 이전의 hidden state 값이 필요하다. 즉, 이전 hidden state 값에 dependent하다는 것이다. 이는 병렬 연산을 하는데 큰 장애물이 된다. sequence가 길어지고 연산량이 많아짐에 따라 병렬 연산은 필수인데, RNN은 그것을 하지 못하니 한계를 가질 수밖에 없는 것이다. attention은 MHA을 통해 병렬 연산을 할 수 있다. 직관적으로 생각해보면 하나의 문장이 $n$개의 token으로 구성되어 있다고 하면, $n$개의 토큰 간의 attention을 계산하는데 $n^2$의 시간이 걸린다. 그러나 이는 어느 한 태스크에 의존적이지 않기 때문에 병렬적(multi-head)으로 할 수 있다. 따라서 RNN에 비해 더 빠르게 token들간의 관계를 파악할 수 있는 것이다.

Query, Key, Value

위에서 self-attention을 한 문장 내에서 토큰들간의 가지는 연관성이라고 정의했다. 그러면 구체적으로 이 attention 값을 어떻게 계산하는지 알아보자. 

The animal didn't cross the street, because it was too tired.

 

attention을 계산하기 위해 Query, Key, Value의 3가지 벡터가 사용된다. 이 세가지 벡터는 각기 다른 역할을 한다.

  • Query: 현재 시점의 token
  • Key: attention을 구하고자 하는 대상 token
  • Value: attention을 구하고자 하는 대상 token (key와 동일한 token)

'it'이 무엇을 의미하는지 attention을 계산하는 과정을 생각해보자. 이때, Query는 'it'으로 고정이다. 문장 내에서 'the animal', 'didn't'등의 token이 있는데, 이 token들이 attention을 계산할 대상 token이다. 즉, Key와 Value가 'the animal'일 때도 있고, 'the street'일 때도 있는 것이다. vanila attention은 모든 token과의 attention을 구하므로, 'it'과 가장 연관성이 높은 (attention이 가장 높은) token을 찾기 위해 문장 내에서 모든 token을 탐색한다. (나중에는 sparse attention 등의 알고리즘도 나온다.) 여기서 필자는 Key, Value는 동일한 대상 token을 의미하는데 왜 Key, Value로 나뉘어져 있는지 궁금했는데 지금은 attention 계산할 때 실제 값은 다르게 나오지만 같은 token을 의미한다고 이해하자. attention을 계산하는데 서로 다르게 사용되는 것이다.

 

그러면 Query, Key, Value 벡터가 어떻게 만들어질까? 세 벡터는 토큰별로 고정되어 있지 않고 transformer가 학습됨에 따라서 업데이트되는 fully connected layer에 token embedding vector을 넣어 만들어진다. 이때, Query, Key, Value와 관련된 서로 다른 세개의 FC가 존재하며 따라서 projection matrix을 공유하지 않고 별개로 존재한다. 즉, 위의 'it'과 다른 token의 attention을 계산하는 상황을 생각해보면 'it'에 대한 token embedding 벡터를 Query FC, Key FC, Value FC에 통과하여 Query, Key, Value 벡터를 만든다. 이때 만들어지는 Query, Key, Value 벡터의 shape는 모두 동일하다. 논문의 notation에 따라서 이 세벡터의 차원을 $d_{k}$로 정의한다. Query 벡터는 'it'에 대한 벡터이고, Key, Value 벡터는 예를 들어 'the street' 토큰에 대한 벡터이다. Key, Value 벡터는 의미적으로 동일한 token을 칭하지만, 서로 다른 FC를 통과했기 때문에 값은 다르다.

Scaled dot-product attention

Query, Key, Value 벡터를 정의했으니 이 벡터를 사용해서 attention을 어떻게 계산하는지 알아보자. 현재는 'it'에 대한 attention을 계산하는 과정이다. Query, Key, Value 벡터를 $Q, K, V$라고 정의하자. $Q$에 대한 attention은 아래와 같이 정의된다.

얼핏보면 복잡한데, 논문에서는 아래와 같이 그림으로 표현한다.

$Q,K$간에 행렬 곱을 하고 scaling을 한다. 경우에 따라서 masking을 하고 softmax을 취하여 일종의 가중치를 만든다. 최종적으로 $V$와 행렬 곱을 한다. encoder, decoder에 따라서 $Q,K,V$가 만들어지는 과정이 살짝 다른데 논의를 간단하기 위해 cross-attention이 아닌, encoder/decoder의 self-attention의 경우를 생각해보자. $Q,K,V$ 행렬은 아래와 같이 만들어진다.

위에서 $Q,K,V$을 만들기 위해 간단하게 FC layer을 통과시킨다고 언급한바 있다. 이를 그대로 적용하면 된다. FC layer이므로 원하는 차원인 $d_k$로 projection하는 projection matrix가 있다. Query, Key, Value는 서로 다른 FC layer을 가지기 때문에 $Q,K,V$는 서로 다른 projection matrix을 가진다. 다시 한번 말하지만, Query는 현재 시점의 단어이고 Key, Value는 attention을 구하고자 하는 대상 단어이다. $K,V$는 의미적으로는 같지만 서로 다른 FC layer을 가지므로 다른 값을 가진다. self-attention의 전체적인 도식을 살펴보자.

$Q$를 넣고 self-attention을 계산하면 $Q$와 차원이 동일한 attention 행렬이 나온다. 이제 행렬 전체로 보면 헷갈릴 수 있으니, 단어 하나하나를 살펴보자.

 

구체적으로 'it'과 'animal' 사이의 attention을 구하는 과정을 생각해보자. $d_k=3$이라고 하면, Query는 'it'이고 Key, Value는 'animal'이며 각각은 $1 \times 3$ 차원의 벡터이다.

scaled-dot product attention은 가장 먼저 $QK^T$을 한다. 'it'과 'animal'간의 $QK^T$을 하면 $1 \times 1$ scalar 값이 나오는데, 이것이 바로 'it'과 'animal'간의 attention score이다. 근데 $d_k=3$일 수도 있고 $d_k=512$일 수도 있다. 즉, 벡터의 차원이 높아지면 scalar 값의 커지므로 이를 scaling한다. 논문에서는 scaling을 통해 gradient vanishing도 방지할 수 있다고 한다.

위 예시에서는 'it'과 'animal'간의 attention score을 구하고 scaling을 했다. 근데 self-attention은 $Q$에 대한 attention을 구할 때, 대상 token을 나머지 모든 token으로 정한다고 했었다. 여기서 대상 token은 $K,V$을 의미한다고 했다. 위 예시에서는 $K,V$가 의미하는 token이 'animal'이었기 때문에 차원이 $1 \times 3$이었다. 이제 한 문장 안에 $n$개의 token이 있다고 하면 $K,V$의 차원은 $n \times 3$이 되는 것이다. 여기서 $d_k=3$으로 정의했음을 다시 한번 짚고 넘어가자.

$Q$는 여전히 'it'이고 $K,V$는 다른 모든 token이므로 $n \times 3$이다. 이제 $QK^T$을 통해 'it'과 나머지 token에 대한 attention score가 $1 \times n$차원으로 나올 것이다.

이렇게 구한 attention score는 Query 'it'과 나머지 token들간에 가지는 연관성을 의미한다. 근데, 지금 이 값의 scale은 Query에 따라서 서로 다르다. 일종의 가중치 개념이기 때문에 scale을 동일하게 해줄 필요가 있다. 이때 가장 많이 사용하는 것이 softmax이다. 합이 1이 되도록 scaling을 해줌으로써 'it'과 가장 많이 연관된 token은 가장 높은 가중치를 가지는 것이다.

 

이제 $V$와 attention score을 곱해야한다. 다시 한번 짚고 넘어가면, $V$는 Query인 'it'의 대상이 되는 token이며, attention score는 'it'과 다른 token이 가지는 연관성을 확률의 개념으로 표현한 것이다. 

softmax을 취했기 때문에 attention prob로 불러보자. attention prob와 $V$을 곱할 때, 어떤 요소들끼리 곱해지는지 살펴보자. $V$는 대상 token이고 attention prob는 각 요소가 합이 1인 가중치이다. attention prob의 'the'와 $V$의 첫번째 행인 'the'가 곱해진다. attention prob의 'animal'과 $V$의 두번째 행인 'animal'이 곱해진다. 즉, $V$에서 대상이 되는 token은 그 자신의 가중치를 attention prob으로부터 얻는다. 'it'이 Query라고 할 때, 대상이 되는 token이 $n$개가 있는데 그 중에 어떤 token이 더 중요한지를 attention prob을 통해 계산하는 과정이라고 생각하자.

 

주목할만한 점은 처음에 Query인 'it' 토큰의 차원이 $1 \times 3$ 이었는데 결과로 나오는 attention도 $1 \times 3$이라는 것이다.

 

지금까지는 'it'에 대한 attention을 구하는 과정이었다면, 이제 나머지 token의 attention을 구해야 한다. 이를 위해 아래와 같이 행렬로 확장하는 과정이 필요하다.

$Q$의 차원이 $n \times 3$으로 확장된 것을 볼 수 있다. 다만, $K,V$의 차원은 동일하다. 애초에 $Q$가 어떤 token이든 간에 다른 모든 token에 대한 attention을 계산하는 것은 동일하기 때문이다.

크게 다른 점은 없다. attention score 행렬이 $n \times n$이 되었다. 행은 Query을, 열은 Key을 의미한다. attention score의 각 행은 sum to 1이다. 즉, 각 행은 Query가 주어졌을 때 다른 token에 대한 가중치를 의미하고, $V$와 행렬곱을 한다.

 

위에서 살펴본 self-attention 도식을 다시한번 살펴보자.

현재 기준이 되는 단어인 $Q$와 대상이 되는 단어인 $K,V$을 조합하여 $Q$의 attention 행렬을 계산했다. 그러면 여기서 Query의 attention 행렬이 의미하는 바는 무엇일까? attention 행렬의 'it'에 대응되는 행을 살펴보자. 'it'에 해당하는 attention 벡터이다. 이 벡터에는 query인 'it'가 주어졌을 때, 문장 내에서 다른 단어와의 연관성이 압축되었다고 볼 수 있다. 즉, 문장 내에서 'it'은 'animal' 또는 'tired'와도 연관될 수 있다. 어느 단어가 더 연관됐는지는 알 수 없겠지만 그래도 이 연관된 정도가 'it'에 대응되는 attention 벡터에 들어갔다고 기대할 수 있는 것이다.

Masking

MHA을 설명하면서 한가지 빼먹은 부분이 있는데 바로 Masking이라고 적혀있는 단계이다. masking은 무엇이며 왜 하는 것일까? masking이란 0과 1이라는 값으로 행렬을 구성하는 것을 의미한다. 어떤 값을 0으로 구성하는 것일까? 서두에서, max_length에 따라서 padding을 하는 과정을 살펴보았다. 'the animal didn't cross the street, because it was too tired' 라는 문장을 살펴보자. 여기서 만약에 max_length을 14로 정한다면, 문장의 길이 (현재는 편의상 `split`으로 tokenize 한다고 가정한다.)가 11인 위 문장은 빈공간이 남게 된다. 이를 <PAD> 단어를 통해 채운다. 즉, 데이터를 아래와 같이 전처리한다.

  • { the, animal, didn't, cross, the, street, because, it, was, too, tired, <PAD>, <PAD>, </s> }

attention을 구하는 과정을 생각해보면 이 sequence에 대한 $QK^T$ 행렬은 아래와 같이 만들어진다.

MHA 계산하는 과정에서 살펴본 것처럼 여기서 행별로 softmax을 한다. 그런데 <PAD> Query에 대응되는 행을 살펴보자. 사실 <PAD>는 의미가 없는 단어이다. 그냥 max_length을 채우기 위한 역할을 할 뿐이다. 따라서 <PAD> Query에 대응되는 attention 값을 계산할 필요가 없다. 따라서 아래와 같이 <PAD>, </s>에 대응되는 부분을 모두 0로 masking 한다.

또, 일반적인 단어가 Query로 주어졌을 때 <PAD>, </s>와 같은 단어에 대한 attention을 계산할 필요가 없다. 따라서 Key에서 <PAD>, </s>에 대응되는 부분도 모두 0으로 masking한다.

이렇게 세팅해두면 'it'에 대한 attention을 계산하고 softmax을 취할 때, <PAD>, </s>에 대응되는 부분은 모두 0으로 매핑된다. 더 정확히는 $n \times n$의 Query, Key 행렬에서 0으로 masking된 부분을 -1e9와 같이 매우 큰 음수값으로 채워넣음으로써 softmax을 했을 때, 0에 가까운 값을 가지도록 한다. 이를 통해 무의미한 단어들이 가중치를 갖는 것을 방지하는 것이다.

scaled-dot product attention을 코드로 구현해보자.

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn
  • `temperature`: scaling할 dimension이다. 보통, $\sqrt{w_dim}$ 으로 정의된다.
  • `q`: query 행렬. (batch_size, n_heads, seq_length, dk)
  • `k`: key 행렬. (batch_size, n_heads, seq_length, dk)
  • `v`: value 행렬. (batch_size, n_heads, seq_length, dv)
  • `mask`: 0,1로 구분되는 masking 행렬. (batch_size, n_heads, seq_length, dv)

Scaled-dot product attention의 핵심 연산은 $softmax(QK^T / \sqrt{d}) V^T$이다. 이는 11번째 줄에 나와있는데, 여기서 key 행렬이 transpose 되어있다. 이 부분을 이해하기 위해 $Q,K,V$ 행렬의 차원을 살펴보자.

q 행렬의 차원은 (batch_size, n_heads, seq_length, dk) 이다. 여기서 n_heads * dk을 하면 임베딩 벡터의 차원이 나온다. $QK^T$은 sequence에 있는 단어들끼리 내적을 구하는 것이므로, batch, head 차원에서는 똑같은 연산만 해주면 된다. 즉, `q[0][0]` 행렬과 `k[0][0]` 행렬 곱이 결국 $QK^T$가 되는 것인데, 이를 위해서 key 행렬을 transpose 해줘야 한다. 따라서 `k.transpose(2,3)`을 해주는 것이다.

이제 결과로 나오는 `attn` 값은 (batch_size, n_heads, seq_length, dk) 행렬과 (batch_Size, n_heads, dk, seq_length) 행렬의 곱이므로 (batch_size, n_heads, seq_length, seq_length) 행렬이 된다.

 

이제 masking 하는 부분인 13~14번째 줄을 살펴보자. encoder 단계이든, decoder 단계이든, attention하는 과정에서 masking이 들어간다. 이 masking은 <PAD>가 샘플링되지 않게 하거나, decoder에서는 subsequent_masking 기술을 사용하므로 masking의 형태가 약간씩 다르다. 어찌됐든, 위 클래스에서는 `mask`가 들어오면 `attn.masked_fill(mask == 0, -1e9)`을 통해 0으로 마스킹된 부분을 절대값이 매우 큰 음수로 치환한다.

파이참에서 디버깅을 돌려보니, `mask` 행렬의 차원이 (batch_size, 1, 1, seq_length)로 나온다. `attn`의 차원은 (batch_size, n_heads, seq_length, seq_length)이라서, 살짝 헷갈렸는데 아마도 masking을 `attn`의 (seq_length, seq_length)에서 행별로 동일하게 적용하는 것 같다. (파이토치가 참 편리한 기능이 많은 것 같다)

 

16번째 줄에서 `F.softmax(attn, dim=-1)`을 통해 (batch_size, n_heads, seq_length, seq_length) 차원에서 (seq_length, seq_length)의 행별로 softmax을 구한다. (`F.softmax(attn, dim=-1)[0][0][0,:].sum()`을 해보면 1이 나온다.)

 

17번째 줄에서 (batch_size, n_heads, seq_length, seq_length)인 `attn`와 (batch_size, n_heads, seq_length, dv)인 `v`을 행렬곱하여 self-attention을 하고 (batch_size, n_heads, seq_length, dv)인 `output`을 얻는다.

Multi-head attention

여태까지 MHA를 구성하는 scaled-dot product attention에 대해서 상세하게 살펴보았다. MHA는 scaled-dot attention을 병렬적으로 $h$회만큼 수행하고 최종적으로 결과를 concat한다. 그러면 굳이 한번에 행렬 연산을 하지 않고 쪼개서 병렬로 scaled-dot product attention을 하는 이유는 무엇일까?

논문에서는 다른 위치에 있는 상이한 representation subspaces을 joint하게 모델이 학습할 수 있다고 한다. 'it'에 대한 attention을 계산한다고 생각해보자. 이때, single attention으로 하는 것보다 multi-head attention으로 하는 것이 'the animal'뿐만 아니라 'tired' 등과 같이 다른 단어와의 연관성도 잘 담을 것이기 때문에 'it'이 'the animal'이고 'tired'하다고까지 모델이 알 수 있는 것이다. 조금 추상적이지만, single attention에 비해서 더 입체적인 정보를 담을 수 있다고 이해하자.

 

multi-head attention을 살펴보기 전에 single attention인 scaled-dot product attention을 다시 한번 살펴보자.

Query의 attention 행렬은 $n \times d_k$인데 이 연산을 $h$번만큼 병렬적으로 수행한다고 생각하면 된다. single attention일 때, $Q,K,V$가 서로 다른 FC를 가졌으므로 multi-head으로 확장한다면 총  $3 \times h$개만큼의 FC를 가지게 된다. 각 attention의 output은 $n \times d_k$이고 이를 $h$번만큼 반복하므로 concat하면 총 $n \times (d_k \times h)$ 차원의 attention 행렬이 생기는 것이다. 여기서 $d_k \times h$을 $d_{model}$이라고 정의하자. 보통 $d_{model}$은 encoder의 input의 차원인 $d_{embed}$와 동일한 값을 사용한다. 이제 그림으로 다시 살펴보자.

$n \times d_k$ 차원의 single attention 결과가 $h$개만큼 반복되어 concat 되었고 총 $n \times d_{model}$ 차원이 되었다. 근데 잘 생각해보면 각 head 당 서로 다른 3개의 $Q,K,V$의 FC를 가질 필요가 있을까? $Q,K,V$ 행렬을 만들 때, $d_k$가 아니라 concat한 차원인 $d_{model}$으로 projection을 해보자. 그러면 multi-head을 위해 $3 \times h$개만큼의 FC를 가질 필요가 없고 기존처럼 3개의 FC만 가지면 된다. $Q,K,V$가 $n \times d_{model}$ 차원으로 만들어졌을 때, $d_{model} = d_k \times h$만큼 multi-head로 쪼개서 각 head별로 attention을 계산한다. 그리고 최종 concat하는 것이다. 그림으로 표현하면 아래와 같다.

흥미로운 것은 multi-head attention을 단지 여러개의 single head attention을 concat함으로써 구현한다는 것이다. 논문에서는 서로 다른 문맥을 담기 위해 multi-head attention을 한다고 했는데 생각보다 간단하게 구현되는 것 같아서 신기했다.

 

encoder의 입력값의 차원은 $n \times d_{embed}$이었음을 기억해보자. 현재 차원은 $n \times d_{model}$ 이다. 보통 $d_{embed} = d_{embed}$이긴 하지만, 다를 수도 있으므로 차원을 FC을 통해 맞춰준다.

함수 도식으로 생각해보자. $n \times d_{embed}$ 입력이 들어왔는데, multi-head attention을 통해서 $n \times d_{embed}$ 차원의 행렬을 만들었다.

 

이제 코드로 구현해보자.

class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)


    def forward(self, q, k, v, mask=None):

        ...
  • `n_head`: head의 수
  • `d_model`: 임베딩 차원
  • `d_k`: $K$의 하나의 head의 차원. `d_v`와 동일함.
  • `d_v`: $V$의 하나의 head의 차원. `d_k`와 동일함

MHA에서 query / key / value에 대한 projection 행렬은 고정된 것이 아니라 학습하는 파라미터이다. 따라서 `nn.Linear`을 통해 이를 정의해준다. 위에서 언급했듯이 $Q,K,V$의 output_dim은 $d_{model}$이 되도록 하고, 후에 multi-head로 쪼갠다.

`self.fc`을 통해 입력값의 차원과 동일해지도록 차원을 조정해준다. $Q,K,V$의 projection matrix는 $d_{embed}$에서 $d_{model}$로 벡터를 projection한다. 보통은 $d_{embed} = d_{model}$이지만 그렇지 않은 경우가 있을 수도 있으므로 이를 원래 차원으로 돌려 놓는 과정을 FC을 통해서 수행한다. 더 정확하게는, `self.fc = nn.Linear(n_head * d_v, d_model, bias=False)`가 아니라 `self.fc = nn.Linear(n_head * d_v, d_embed, bias=False)`이어야 할 것 같다. 기존에 MHA을 하기 전에 $n \times d_{embed}$ 차원을 가졌는데, 이와 동일하게 만들어주는 과정이기 때문이다. 다만, 대부분 $d_{embed} = d_{model}$으로 세팅해두기 때문에 위 코드에서 `d_{embed}`가 아닌, `d_{model}`로 둔 것 같다.

`forward` 메써드를 살펴보자.

class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        ...


    def forward(self, q, k, v, mask=None):

        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        q, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(self.fc(q))
        q += residual

        q = self.layer_norm(q)

        return q, attn

17~19번째 줄이 좀 헷갈리니 차원을 살펴보자.

`self.w_qs(q)`는 query linear layer에 query 행렬을 통과시킨다. 이는 곧 (batch_size, seq_length, embed_size)의 행렬을 계산한다. 여기서 embed_size을 n_head만큼 나눠서 Multi-head을 구현하는 것이다. `n_head * d_k = embed_size`으로 설정하여 `self.w_qs(q)`을 (batch_size, seq_length, n_head, d_k)의 4차원 텐서로 변환한다. 똑같은 작업을 q,k,v에 한다.

22번째 줄에서 transpose는 굳이 왜 하는 것일까? (batch_size, seq_length, n_head, d_k) 차원인 텐서는 서로 행렬 곱을 하기에 적당하지 않다. batch_size, n_head 별로 행렬 곱을 해야하기 때문이다. 따라서 이 차원을 (batch_size, n_head, seq_length, d_k)로 바꿔서 scaled-dot product attention을 진행한다.

24~25번째 줄에서, mask는 현재 (batch_size, 1, seq_length) 차원이다. 현재 $Q,K,V$의 차원이 (batch_size, n_head, seq_length, d_k)인 4차원이므로 이에 맞게 unsqueeze을 해준다.

31번째 줄은 무엇을 하는 것일까? `self.attention`을 통해서 MHA을 하고, 결과로 나오는 `q`의 차원은 (batch_size, n_head, seq_length, d_k)이다. 이제 MHA을 끝냈으니 더이상 n_head만큼 차원을 더 가지고갈 필요가 없는 것이다. 따라서 transpose을 통해 우선 (batch_size, seq_length, n_head, d_k)로 바꾸고, `view`을 통해 (batch_size, seq_length, n_head * d_k)로 차원을 바꿔준다.

32번째 줄은 FC하는 부분인데 위에서 언급했듯이 기존의 차원으로 되돌리는 과정이다. 그 이후는 add & norm하는 부분이다. 최종적으로 `q`의 차원은 (batch_size, seq_length, n_head * d_k)이다.

Position-wise feed-forward networks

FC > Relu > FC로 구성된 layer이다. 입력으로는 MHA을 통과한 $n \times d_{embed}$ 차원의 행렬이 들어온다. 이를 다시 중간 차원인 $d_{ff}$을 정의하고 $n \times d_{ff}, d_{ff} \times d_{embed}$ 차원을 가지는 FC을 통과시킨다. 논문에서 수식은 아래와 같이 나온다.

코드로 구현해보자.

class PositionwiseFeedForward(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid) # position-wise
        self.w_2 = nn.Linear(d_hid, d_in) # position-wise
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        residual = x

        x = self.w_2(F.relu(self.w_1(x)))
        x = self.dropout(x)
        x += residual

        x = self.layer_norm(x)

        return x
  • `d_in`에서 `d_in`으로 가는 FFN을 정의한다.
  • `d_in`: 입력 차원이자, 출력 차원
  • `d_hid`: hidden layer의 차원

논문에서의 수식과 동일하다. FC > Relu > FC을 통과한다. dropout을 하고 add & normalization을 한 뒤에 최종 값을 반환한다.

 

encoder의 구조를 다시 한번 살펴보자.

encoder는 MHA layer, FFN layer로 구성되어 있다. MHA을 통과하면 $n \times d_{embed}$ 차원의 입력값은 동일하게 $n \times d_{embed}$이 된다. FFN도 마찬가지로 $n \times d_{embed}$의 출력값을 만들어낸다.

 

이제 하나의 encoder block을 코드로 살펴보자. MHA와 FFN으로 구성되어 있을 것이다.

class EncoderLayer(nn.Module):
    ''' Compose with two layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, enc_input, slf_attn_mask=None):
        enc_output, enc_slf_attn = self.slf_attn(
            enc_input, enc_input, enc_input, mask=slf_attn_mask)
        enc_output = self.pos_ffn(enc_output)
        return enc_output, enc_slf_attn
  • `d_moel`: 임베딩 모델의 차원
  • `d_inner`: FFN할 때 hidden layer의 차원

위에서 열심히 코드를 봤다면 아주 간단한 클래스이다. EncoderLayer는 기본적으로 MHA, FFN으로 구성되어 있다. 따라서 MHA을 통과시긴 enc_output을 FFN에 통과시킨다. 이러한 EncoderLayer는 Encoder 구조에서 총 n_layers만큼 반복된다. self-attention이기 때문에 `self.slf_attn`의 입력으로 모두 동일한 `enc_input`이 들어간다. 이는 이후에 decoder의 cross-attention과는 다른 점이므로 다시 한번 주의 깊게 살펴보자.

 

이제 encoder block을 쌓아올린 encoder을 코드로 살펴보자.

class Encoder(nn.Module):
    ''' A encoder model with self attention mechanism. '''

    def __init__(
            self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, dropout=0.1, n_position=200, scale_emb=False):

        super().__init__()

        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, src_seq, src_mask, return_attns=False):

        ...

 

  • `n_src_vocab`: input sequence의 vocab size
  • `d_word_vec`: 단어 임베딩 벡터의 차원
  • `n_layers`: Encoder을 반복할 횟수
  • `n_head`: MHA을 할 때, head의 수
  • `d_k`: MHA할 때, 하나의 head에서 key의 차원
  • `d_v`: MHA할 때, 하나의 head에서 value의 차원
  • `d_model`: FFN할 때 output 차원?
  • `d_inner`: Linear > Relu > Linear할 때, 중간 차원
  • `pad_idx`: padding index
  • `dropout`: dropou하는 비율
  • `n_position`: max length
  • `scale_emb`: embedding 값을 스케일링할지 여부

transformer에서 encoder는 PE 및 MHA와 FFN으로 이루어지는데, 이를 위한 여러 클래스를 정의한다. 하나 주의 깊게 살펴봐야할 것은 `self.layer_stack`이다.

self.layer_stack = nn.ModuleList([
    EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
    for _ in range(n_layers)])

 

sublayer로 MHA와 FFN을 담고있는 `EncoderLayer`을 정의한다. 이를 `n_layers`만큼 반복하는데, 논문에서 나온 것처럼 보통 6으로 지정하여 총 6번 MHA, FFN을 반복하도록 설정한다.

다음으로 `forward` 메써드를 살펴보자.

class Encoder(nn.Module):
    ''' A encoder model with self attention mechanism. '''

    def __init__(
            self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, dropout=0.1, n_position=200, scale_emb=False):

        ...

    def forward(self, src_seq, src_mask, return_attns=False):

        enc_slf_attn_list = []

        # -- Forward
        enc_output = self.src_word_emb(src_seq)
        if self.scale_emb:
            enc_output *= self.d_model ** 0.5
        enc_output = self.dropout(self.position_enc(enc_output))
        enc_output = self.layer_norm(enc_output)

        for enc_layer in self.layer_stack:
            enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask)
            enc_slf_attn_list += [enc_slf_attn] if return_attns else []

        if return_attns:
            return enc_output, enc_slf_attn_list
        return enc_output,

 

  • `src_seq`: input sequence로, (batch_size, sequence_length)이다.
  • `src_mask`: input sequence에 대한 mask로, (batch_size, 1, sequence_length)이다. (중간에 1이 들어가있는 이유는?)

15~19번째 줄은 `src_seq`을 임베딩하고 PE을 더한 후, dropout, layer_norm까지 하는 부분이다.

그리고 21~23번째 줄에서는 self.n_layers만큼 EncoderLayer에서 MHA, FFN을 진행한다.

`return_attns`에 따라서 attention 값까지 반환할지, 아니면 enc_output만 반환할지 결정한다. 어떤 경우에 attention 값까지 반환하는지는.. 아직 파악하지 못했다.

Decoder

Encoder에 대해서 상세하게 살펴봤으니 이제 Decoder을 살펴볼 차례이다.

앞서 encoder는 입력 문장이 들어가면 문맥이 압축된 context을 만들어낸다고 했다. decoder는 encoder가 만든 context와 또 다른 문장이 들어간다. 즉, 두 종류의 입력을 받는 것이다. 위의 모델 구조를 보면 decoder의 최종 output이 transformer 모델의 최종 output이다. 즉, 영한 번역 task을 생각해보면 output sentence는 한글 문장이 되는 것이다.

 

그렇다면 여기서 의문이 생길 수 있다. decoder가 입력으로 받는 또 다른 문장이 무엇을 의미할까? 위 모델 구조에서는 shifted right라고 나와있는데 이것이 의미하는 바는 무엇일까? 얼핏 생각해보면 outputs을 decoder의 입력으로 넣는다는 것 자체가 잘 이해가 가지 않는다. 학습 데이터의 label을 학습하는데 넣는다는 것일까? ML에서 학습을 할 때 ground rule은 모델이 label을 보지 않아야한다는 것 아닌가? 이를 이해하기 위해 아래의 개념에 대해 살펴보자.

Teacher Forcing

결론부터 말하면, decoder에 outputs(shifted right)을 넣는다는 것은 맨 마지막 token인 <EOS>을 제외한 문장을 decoder의 입력으로 넣는다는 것이다. 이렇게 조금 변형된 label을 decoder에 입력으로 넣는 것을 teacher forcing이라고 한다.

 

sequence 기반의 번역 model이 있다고 생각해보자. 입력 문장이 들어왔을 때 새로운 문장을 만들어내야한다. 랜덤하게 파라미터를 초기화한 학습 초반을 생각해보자. 열심히 출력 문장을 만들고 있지만 초기화한 파라미터 때문에 중간에 만들어지는 token이 엉터리일 가능성이 크다. RNN으로 생각해보면, 이전에 나온 token을 다음 token을 만들어내는데 쓰이는데, 이전에 나온 token이 엉터리이니, 이 엉터리 token에 기반한 다음 token도 엉터리일 가능성이 크다. 그러면 점점 결과가 좋아지기는 힘들 것이다. 이런 현상을 방지하기 위해 teacher forcing을 사용한다. teacher forcing이란, 지도 학습에서 label data를 input으로 활용하는 것이다. 학습 과정에서 model이 생성해낸 token을 다음 token 생성 때 사용하는 것이 아니라, 실제 label data의 token을 사용하게 되는 것이다. 우선 정확도 100%을 달성하는 이상적인 model을 생각해보자.

정상적인 output token을 만들어내기 때문에 그 다음에 생성하는 token도 정상적인 token임을 알 수 있다. 하지만 실제로는 이렇지 않다.

rnn의 token이 뒤죽박죽이고 이 token이 다음 token 생성시에 입력으로 들어가기 때문에 출력으로 나오는 token도 정답과 멀어지는 것이다. 이러면 제대로 된 학습이 진행되기 힘들기 때문에 label data을 input으로 사용한다.

정확히는 label data을 그대로 사용하는게 아니라 마지막 token인 <EOS>을 제외하고 rnn의 input으로 사용한다.

 

하지만 잘 생각해보면, 이런 label을 가지고 있는 경우는 학습할때에 한정된다. test을 하거나 실제 모델이 배포된 상황에서는 label data을 가지고 있지 않다. 따라서 이러한 경우에는 model이 생성해낸 이전 token을 사용하게 된다.

 

이와 같이 학습 과정과 실제 사용에서 상황의 차이가 발생하지만, model의 학습 성능을 향상시킬 수 있다는 점에서 많은 encoder-decoder 구조의 모델에서 사용하는 기법이다.

Tearch forcing in transformer (subsequent masking)

위의 예시는 RNN 모델의 경우이다. 즉, RNN과 같이 이전 token을 현재 token을 만들어내는데 사용하는 모델 구조라면 위와 같이 ground truth[:-1] (맨 마지막 token을 제외)을 차례대로 넣어도 상관이 없다. 두번째 RNN Cell을 살펴보자. 해당 cell은 '그동안'이라는 단어를 만드는데 있어서 '어머니'라는 단어를 입력으로 받는다. 이때, '그동안'이라는 단어는 이 cell에게 주어지지 않는다. 즉, 최소한 해당 cell이 만들어내야하는 output을 이 cell에게 알려줘서는 안되는 것이다.

 

이 구조를 그대로 transformer에 사용해도 될까? 결론부터 말하면, 그렇게 하면 안 된다. transformer는 병렬 연산을 위해 임베딩 행렬을 그대로 받는다. 이는 encoder든, decoder든 동일하다. decoder에서 MHA을 하는 과정을 생각해보자. encoder와 마찬가지로 self-attention을 한다고 하면 '어머니'가 query token일 때, 나머지 token인 '그동안', '잘', '지내셨어요?'와의 attention을 계산하게 된다. 즉, teacher forcing을 하기 위해서 label data을 그대로 넣었는데 '어머니' token에 대해 다음에 올 token인 '그동안' token을 모델에 그대로 알려주게 되는 것이다. 

 

그러면 teacher forcing을 transformer에 적절하게 적용하기 위해서는 어떤 기법이 필요할까? '어머니' token이 현재 query일 때, 그 이후에 나오는 token들과 attention을 계산하지 못하도록 설정하면 되는 것이 아닌가? 이를 위해 transformer는 subsequent masking을 구현한다.

def get_subsequent_mask(seq):
    ''' For masking out the subsequent info. '''
    sz_b, len_s = seq.size()
    subsequent_mask = (1 - torch.triu(
        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()
    return subsequent_mask

`torch.triu()` 함수는 상삼각행렬을 만드는 함수이다. 1에서 이 텐서를 빼면 하삼각행렬이 만들어진다. '어머니 그동안 잘 지내셨어요?' 문장에 대해서 subsequent mask 행렬을 만들면 아래와 같다.

[[1,0,0,0],
 [1,1,0,0],
 [1,1,1,0],
 [1,1,1,1]]

이렇게 masking을 하고 attention을 진행한다. 코드 상으로는 $QK^T / \sqrt{d_k}$을 하고 masking을 진행하는데, softmax 이전에만 하면 상관 없다. 위의 subsequent masking을 자세하게 살펴보자. 첫번째 행을 살펴보자. 첫번째 token에 대해 attention을 할 때, 그 자신과의 attention만 구하고 다른 단어에 대해서는 구하지 않겠다는 뜻이다. 두번째 행을 살펴보자. 두번째 token에 대해 attention을 할 때, 자신을 포함한 그 이전 token과의 attention만 계산하겠다는 뜻이다. 이와 같이 decoder에서는 label data가 입력으로 들어가기 때문에 현재 token 기준으로 그 다음 token에 대해 알지 못하도록, 0으로 masking을 하는 것이다.

 

masking하는 함수는 위의 `get_subsequent_mask`로 구현되어 있고, masking이 모델에서 어떻게 녹아들어갔는지 살펴보자.

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn

다시 scaled-dot product attention 코드이다. 중간에 보면 `attn.masked_fill` 라인이 있다. mask 값이 0이라면 매우 큰 음수 값으로 채워 넣어서 softmax 계산 시에 0에 가까운 가중치를 가지게 만든다. 여기서 mask는, encoder에서 <PAD>을 masking 하는 것, encoder-decoder에서 cross-attention시에 하는 masking, 과 decoder에서 self-attention시에 subsequent mask을 하는 것 모두를 의미한다.

Decoder Block

encoder와 비슷하게 decoder역시 $N$개의 decoder block을 쌓은 구조이다. 이때, encoder와 다른 점은 각 decoder block으로 encoder의 context가 추가적인 입력으로 주어진다는 것이다.

encoder block은 MHA, FFN으로 구성되었음을 기억해보자. decoder block은 어떻게 구성이 될까?

크게 보면 self-attention, cross-attention, FFN으로 구성된다. 잘 살펴보면, decoder block이 encoder로부터 받는 context 입력은 cross-attention에서만 쓰임을 알 수 있다.

 

Self multi-head attention layer

self-attention layer는 right shifted outputs 문장들간에 attention을 계산하는 과정이다. 이때, 위에서 살펴보았듯이 subsequent masking을 사용하여 현재 token이 다음 token을 모르도록 설정한다. 논문에서는 masked multi-head attention이라고 표시되어 있다. 핵심은 encoder의 self-attention과는 다르게 subsequent masking을 한다는 점이다. 이 차이점은 encoder와 decoder의 입력값의 차이에서 비롯됨을 꼭 기억하고 넘어가자.

 

Cross multi-head attention layer

decoder block의 가장 핵심적인 부분이다. encoder의 context와 decoder의 self-attention에서 나온 값을 입력으로 받는다. 이 두 입력은 완전히 다르게 사용된다. decoder의 self MHA에서 나온 값은 cross MHA의 Query로, encoder의 context는 cross MHA의 Key, Value로 사용된다. 이 점을 반드시 기억하고 넘어가자. 이전의 self-attention과는 형태가 달라졌다. self-attention에서는 Query, Key, Value 모두 하나의 문장 내에 존재하는 token을 의미했다. cross-attention에서는 서로 다른 문장의 token 간의 attention을 계산하는 것이다.

 

그러면 왜 decoder의 self MHA에서 나온 값은 cross MHA의 Query로, encoder의 context는 cross MHA의 Key, Value로 들어가는 것일까? $Q,K,V$의 역할을 다시 한번 생각해보자. $Q$는 현재 시점의 token, $K,V$는 attention을 계산하고자 하는, 대상이 되는 token이다. decoder에서 teacher forcing을 통해 들어온 label data는 encoder의 입력과 어떤 관계를 가지고 있는지, 이를 attention으로 계산해야 한다. 따라서 기준이 되는 token이 decoder의 label data가 되는 것이고 대상이 되는 token이 encoder의 입력값이 되는 것이다.

 

Position-wise feed-forward layer

encoder의 FFN과 동일하다. FC > Relu > FC의 과정을 거친다.

 

이제 하나의 decoder block을 코드로 구현해보자.

class DecoderLayer(nn.Module):
    ''' Compose with three layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(
            self, dec_input, enc_output,
            slf_attn_mask=None, dec_enc_attn_mask=None):
        dec_output, dec_slf_attn = self.slf_attn(
            dec_input, dec_input, dec_input, mask=slf_attn_mask)
        dec_output, dec_enc_attn = self.enc_attn(
            dec_output, enc_output, enc_output, mask=dec_enc_attn_mask)
        dec_output = self.pos_ffn(dec_output)
        return dec_output, dec_slf_attn, dec_enc_attn

첫번째 MHA는 output sequence간의 self-attention이고 두번째 MHA는 enc_output과 dec_output간의 cross-attention임을 다시 한번 코드로 확인하자. 동일한 `MultiHeadAttention` 클래스를 사용했지만, `self.slf_attn`과 `self.enc_attn`은 입력으로 들어가는 $Q,K,V$가 다르다. 즉, 첫번째 MHA는 query, key, value에 모두 dec_input이 들어간다. dec_input간의 self-attention을 보는 것이기 때문이다. 이와 다르게, 두번째 MHA는 query에 dec_input이, key, value에는 enc_output이 들어간다. enc_output과 dec_output간의 attention을 보는 것이기 때문이다. 최종적으로 나오는 dec_output은 다시 FFN을 통과하여 하나의 decoder block을 완성한다.

 

그러면 이 decoder block을 $N$개만큼 쌓은 모델인 decoder을 코드로 살펴보자.

class Decoder(nn.Module):
    ''' A decoder model with self attention mechanism. '''

    def __init__(
            self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):

        super().__init__()

        self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):

        ...
  • `n_trg_vocab`: decoder vocab의 수
  • 그 외에는 Encoder 클래스와 동일한 argument

임베딩 모델, PE 클래스, dropout 등을 정의한다. Encoder 클래스와 마찬가지로 `n_layers`만큼 DecoderLayer을 쌓는다.

`forward` 메써드를 살펴보자.

class Decoder(nn.Module):
    ''' A decoder model with self attention mechanism. '''

    def __init__(
            self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):

        ...

    def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):

        dec_slf_attn_list, dec_enc_attn_list = [], []

        # -- Forward
        dec_output = self.trg_word_emb(trg_seq)
        if self.scale_emb:
            dec_output *= self.d_model ** 0.5
        dec_output = self.dropout(self.position_enc(dec_output))
        dec_output = self.layer_norm(dec_output)

        for dec_layer in self.layer_stack:
            dec_output, dec_slf_attn, dec_enc_attn = dec_layer(
                dec_output, enc_output, slf_attn_mask=trg_mask, dec_enc_attn_mask=src_mask)
            dec_slf_attn_list += [dec_slf_attn] if return_attns else []
            dec_enc_attn_list += [dec_enc_attn] if return_attns else []

        if return_attns:
            return dec_output, dec_slf_attn_list, dec_enc_attn_list
        return dec_output,
  • `trg_seq`: output sequence로, 임베딩을 하기 전에는 (batch_size, sequence_length) 이다.
  • `trg_mask`: decoder mask로, (batch_size, sequence_length, sequence_length) 이다. decoder에서는 autoregressive하게 masking을 하기 때문에 size[1]에 1이 아닌 sequence_length가 들어가있는 것 같다.
  • `enc_output`: encoder 단계에서 나온 행렬이다. (batch_size, sequence_length, embedding_size) 이다.
  • `src_mask`: encoder의 output와 decoder의 output간의 attention을 할 때 사용되는 mask이다.

15~19번째 줄은 output sequence을 임베딩 벡터로 만들고, PE을 더한 후에 dropout, layer_norm을 진행한다.

21~25번째 줄은 self.n_layers만큼 decoder layer을 반복한다. 이때, 내부적으로는 output sequence간의 self-MHA, enc_output과 dec_output간의 cross-MHA, 그리고 FFN이 진행된다.

 

Transformer

길고 긴 설명이 거의 끝이 났다. encoder, decoder을 코드로 모두 살펴보았으니 이제 전체적인 모델의 구조를 코드로 살펴볼 차례이다. 앞서 정의한 Encoder, Decoder 클래스를 사용하고, Encoder 클래스의 출력값을 Decoder 클래스로 잘 전달하면 된다.

class Transformer(nn.Module):
    ''' A sequence to sequence model with attention mechanism. '''

    def __init__(
            self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx,
            d_word_vec=512, d_model=512, d_inner=2048,
            n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
            trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True,
            scale_emb_or_prj='prj'):

        super().__init__()

        self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx

        # In section 3.4 of paper "Attention Is All You Need", there is such detail:
        # "In our model, we share the same weight matrix between the two
        # embedding layers and the pre-softmax linear transformation...
        # In the embedding layers, we multiply those weights by \sqrt{d_model}".
        #
        # Options here:
        #   'emb': multiply \sqrt{d_model} to embedding output
        #   'prj': multiply (\sqrt{d_model} ^ -1) to linear projection output
        #   'none': no multiplication

        assert scale_emb_or_prj in ['emb', 'prj', 'none']
        scale_emb = (scale_emb_or_prj == 'emb') if trg_emb_prj_weight_sharing else False
        self.scale_prj = (scale_emb_or_prj == 'prj') if trg_emb_prj_weight_sharing else False
        self.d_model = d_model

        self.encoder = Encoder(
            n_src_vocab=n_src_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=src_pad_idx, dropout=dropout, scale_emb=scale_emb)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=trg_pad_idx, dropout=dropout, scale_emb=scale_emb)

        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p) 

        assert d_model == d_word_vec, \
        'To facilitate the residual connections, \
         the dimensions of all module outputs shall be the same.'

        if trg_emb_prj_weight_sharing:
            # Share the weight between target word embedding & last dense layer
            self.trg_word_prj.weight = self.decoder.trg_word_emb.weight

        if emb_src_trg_weight_sharing:
            self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight


    def forward(self, src_seq, trg_seq):

        ...
  • `trg_emb_prj_weight_sharing`: last dense layer의 weight matrix와 target word embdding의 weight matrix을 같게 설정할지에 대한 옵션. 논문에서는 동일하게 설정한다고 나온다.
  • `emb_src_trg_weight_sharing`: source / target word weight matrix을 동일하게 설정할지에 대한 옵션

`self.encoder`, `self.decoder`는 위에서 살펴본, transformer을 구성하는 큰 두 줄기이다. `self.trg_word_prj`은 decoder에서 나온 output (batch_size, sequence_length, embedding_dim)을 target vocab의 수를 output_dim으로 projection하는 Linear layer이다. 이는 곧, 최종적으로 target vocab에 대한 softmax 확률을 구하여 높은 확률을 가지는 단어를 sampling하기 위한 전 단계라고 보면 된다.

`forward` 메써드를 살펴보자.

class Transformer(nn.Module):
    ''' A sequence to sequence model with attention mechanism. '''

    def __init__(
            self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx,
            d_word_vec=512, d_model=512, d_inner=2048,
            n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
            trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True,
            scale_emb_or_prj='prj'):

        ...


    def forward(self, src_seq, trg_seq):

        src_mask = get_pad_mask(src_seq, self.src_pad_idx)
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

        enc_output, *_ = self.encoder(src_seq, src_mask)
        dec_output, *_ = self.decoder(trg_seq, trg_mask, enc_output, src_mask)
        seq_logit = self.trg_word_prj(dec_output)
        if self.scale_prj:
            seq_logit *= self.d_model ** -0.5

        return seq_logit.view(-1, seq_logit.size(2))
  • `src_seq`: 번역 이전의 sequence. (batch_size, sequence_length)이다.
  • `trg_seq`: 번역 이후의 sequence. (batch_size, sequence_length)이다.

16~17번째 줄에 encoder, decoder에 맞는 masking을 구한다. decoder에서 output sequence끼리 MHA을 하는 부분에서는 subsequent_mask을 사용하여 이후 단어들에 대한 attention을 계산하지 않고, 샘플링에서도 나오지 않게 된다.

19~23번째 줄에는 encoder을 통과시겨 enc_output을 얻고, 이를 다시 decoder에 통과시켜 dec_output을 얻는다. Linear projection을 통해 logit을 얻고 최종적으로 이를 반환한다.

여기서 차원을 다시 한번 살펴보면 아래와 같다.

  • `enc_output`: (batch_size, sequence_length, embedding_size)
  • `dec_output`: (batch_size, sequence_length, embedding_size)
  • `seq_logit`: (batch_size, sequence_length, decoder_vocab_size)
  • `seq_logit.view(-1, seq_logit.size(2))`: (batch_size*sequence_length, decoder_vocab_size)

즉, 코드 상에서는 softmax을 적용하여 확률까지 계산하지는 않고 logit까지만 계산한다. 어차피 학습 과정이기 때문에 굳이 여기서 확률을 계산할 필요가 없다. inference할 때, softmax을 사용하여 확률을 계산한다. 학습할 때는 logits (score)과 gold_label을 활용하여 cross-entropy loss을 구한다.

Training model

이제 본격적으로 모델을 학습해보자.

python train.py -data_pkl m30k_deen_shr.pkl \
	-log m30k_deen_shr \
	-embs_share_weight \
	-proj_share_weight \
	-label_smoothing \
	-output_dir output \
	-b 256 \
	-warmup 128000 \
	-epoch 400

train.py 파일의 train 함수를 살펴보자.

def train(model, training_data, validation_data, optimizer, device, opt):
    ...
    valid_losses = []
    for epoch_i in range(opt.epoch):
        print('[ Epoch', epoch_i, ']')

        start = time.time()
        train_loss, train_accu = train_epoch(
            model, training_data, optimizer, opt, device, smoothing=opt.label_smoothing)
        train_ppl = math.exp(min(train_loss, 100))
        # Current learning rate
        lr = optimizer._optimizer.param_groups[0]['lr']
        print_performances('Training', train_ppl, train_accu, start, lr)

        start = time.time()
        valid_loss, valid_accu = eval_epoch(model, validation_data, device, opt)
        valid_ppl = math.exp(min(valid_loss, 100))
        print_performances('Validation', valid_ppl, valid_accu, start, lr)

        valid_losses += [valid_loss]

        checkpoint = {'epoch': epoch_i, 'settings': opt, 'model': model.state_dict()}

        ...

앞, 뒤에 로그 등과 관련된 코드는 생략하였다.

  • epoch만큼 전체 데이터에 대해 학습을 한다.
  • 먼저 `train_epoch` 함수를 통해 training_data에 대한 학습을 시작한다.
  • 다음으로, `eval_epoch` 함수를 통해 validation_data에 대한 loss와 accuracy을 계산한다.

`train_epoch` 함수를 살펴보자.

def train_epoch(model, training_data, optimizer, opt, device, smoothing):
    ''' Epoch operation in training phase'''

    model.train()
    total_loss, n_word_total, n_word_correct = 0, 0, 0 

    desc = '  - (Training)   '
    for batch in tqdm(training_data, mininterval=2, desc=desc, leave=False):

        # prepare data
        src_seq = patch_src(batch.src, opt.src_pad_idx).to(device)
        trg_seq, gold = map(lambda x: x.to(device), patch_trg(batch.trg, opt.trg_pad_idx))

        # forward
        optimizer.zero_grad()
        pred = model(src_seq, trg_seq)

        # backward and update parameters
        loss, n_correct, n_word = cal_performance(
            pred, gold, opt.trg_pad_idx, smoothing=smoothing) 
        loss.backward()
        optimizer.step_and_update_lr()

        # note keeping
        n_word_total += n_word
        n_word_correct += n_correct
        total_loss += loss.item()

    loss_per_word = total_loss/n_word_total
    accuracy = n_word_correct/n_word_total
    return loss_per_word, accuracy
  • `model.train()`: 모델을 학습 모드로 전환한다.
  • `optimizer.zero_grad()`: gradient을 0으로 초기화한다.
  • `model(src_seq, trg_seq)`: src_seq 데이터와 trg_seq을 넣고 점수를 계산한다. 이때, `pred`의 차원은 (output_seq_length, dec_vocab_size)이다.
  • `cal_performance()`: pred와 gold의 차이를 loss 관점에서 계산한다. 이때, gold는 단어의 인덱스들, 즉 label이고 pred는 (output_seq_length, dec_vocab_size)이다. 따라서, 각 단어들마다 일종의 점수가 dec_vocab_size 길이만큼 나오는 것이고, cross-entropy loss을 계산하는데 사용한다.
  • `loss.backward()`: gradient을 계산한다.
  • `optimizer.step_and_update_lr()`: lr을 업데이트하고 gradient만큼 파라미터를 업데이트한다.

다음으로 `eval_epoch` 함수를 살펴보자.

def eval_epoch(model, validation_data, device, opt):
    ''' Epoch operation in evaluation phase '''

    model.eval()
    total_loss, n_word_total, n_word_correct = 0, 0, 0

    desc = '  - (Validation) '
    with torch.no_grad():
        for batch in tqdm(validation_data, mininterval=2, desc=desc, leave=False):

            # prepare data
            src_seq = patch_src(batch.src, opt.src_pad_idx).to(device)
            trg_seq, gold = map(lambda x: x.to(device), patch_trg(batch.trg, opt.trg_pad_idx))

            # forward
            pred = model(src_seq, trg_seq)
            loss, n_correct, n_word = cal_performance(
                pred, gold, opt.trg_pad_idx, smoothing=False)

            # note keeping
            n_word_total += n_word
            n_word_correct += n_correct
            total_loss += loss.item()

    loss_per_word = total_loss/n_word_total
    accuracy = n_word_correct/n_word_total
    return loss_per_word, accuracy
  • `model.eval()`: 모델을 validation 모드로 전환한다. 이때는 gradient을 계산하지 않고 파라미터를 업데이트하지도 않는다.
  • `model(src_seq, trg_seq)`: 한 epoch을 돌고나서 업데이트된 모델의 파라미터를 사용하여 pred을 계산한다.
  • `cal_performance`: pred와 gold의 차이를 cross-entropy loss 관점에서 계산한다.
  • `accuracy`: 학습할 때 사용하지 않은 validation data을 얼마나 잘 맞추는지에 대한 정확도이다.

학습이 제대로 진행되면 아래와 같이 epoch별로 batch 진행 상황이 뜬다. (아직 한참 남았네..)

Namespace(batch_size=256, cuda=True, d_inner_hid=2048, d_k=64, d_model=512, d_v=64, d_word_vec=512, data_pkl='m30k_deen_shr.pkl', dropout=0.1, embs_share_weight=True, epoch=400, label_smoothing=True, lr_mul=2.0, max_token_seq_len=100, n_head=8, n_layers=6, n_warmup_steps=128000, no_cuda=False, output_dir='output', proj_share_weight=True, save_mode='best', scale_emb_or_prj='prj', seed=None, src_pad_idx=1, src_vocab_size=9521, train_path=None, trg_pad_idx=1, trg_vocab_size=9521, use_tb=False, val_path=None)
[Info] Training performance will be written to file: output/train.log and output/valid.log
[ Epoch 0 ]
  - (Training)   :  15%|█▍        | 17/114 [03:36<17:30, 10.83s/it]

마무리

이렇게까지 코드를 자세하게 뜯어본 적은.. 대학원 때 scikit-learn 코드를 본 이후로 처음인 것 같다. 다행스럽게도 레포의 저자가 코드를 아주 깔끔하게 작성해서 이해하는데 크게 어렵지는 않았다. 특히, 이해가 안 가는 부분이 있으면 파이참에서 breakpoint을 걸어서 텐서의 차원을 확인해봄으로써 이해도를 크게 높일 수 있었던 것 같다. 아직 pytorch에 완전하게 익숙하지 않아서 더 많이 공부해야할 것 같다. 다음에는 Bert 코드와 karparthy의 nanoGPT 코드를 리뷰하고자 한다. karparthy는 AI 분야에서 혁신적인 분인데.. 그분은 얼마나 또 깔끔하게 코드를 작성했을지.. 벌써부터 기대가 된다!

'ML&DL > NLP(LLM)' 카테고리의 다른 글

[Prompt Engineering] Automatic Prompt Optimization with "Gradient Descent" and Beam Search Paper Review  (0) 2024.05.24
[NLP][Code Review] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding pytorch 코드 리뷰  (0) 2024.05.06
[NLP][Prompt Engineering] What is Prompt Optimization  (0) 2024.05.03
[Paper review][Inference Optimization] Efficient Memory Management for Large Language Model Serving with PagedAttention  (0) 2023.09.23

관련글

댓글

티스토리툴바