파이토치로 트랜스포머 구현하기

https://02vec.tistory.com/2

트랜스포머 논문 이해하기

 

파이토치로 트랜스포머를 구현해봅시다. 리뷰 할 코드는 독일어를 영어로 번역하는 트랜스포머를 구현합니다.

 

 

 
 

 
먼저 torchtext와 spaCy를 설치합니다. torchtext는 각종 자연어 처리 도구를 제공하며, spaCy는 자연어 문장을 자동으로 토큰화 해주는 기능 등을 제공합니다.
 
 
 
 
  

 
spaCy에서 영어 문장과 독일어 문장을 토큰화하는 tokenizer를 불러옵니다. "I am a graduate student."라는 영어 문장이 각각의 토큰으로 잘 분리된 것을 확인할 수 있습니다.
 
 

 
 
 
 

 

이를 이용하여 영어 문장과 독일어 문장 각각을 토큰화하여 반환하는 메서드를 정의합시다.
 
 

 
 
 
 

 
torchtext.data.Field를 통해 데이터셋에 적용할 전처리 사항을 명시할 수 있습니다. 소스 언어(독일어)에 대한 전처리를 SRC로, 타겟 언어(영어)에 대한 전처리를 TRG로 정의합시다. batch_first = True로 설정하여 모델 학습시 배치 크기가 첫 번째 차원에 위치하도록 합니다.
 
 
 

 
 
 

 
torchtext.datasets.Multi30k로 번역 데이터셋을 불러옵니다. exts 인자를 통해 소스 언어 타겟 언어를 명시합니다. ".de"는 독일어, ".en"은 영어를 의미합니다. 데이터셋 전처리를 명시한 SRC, TRG를 fields 인자로 지정합니다. 이렇게 구축된 데이터셋을 학습, 검증, 테스트 데이터로 분리했습니다.

 

 
 
 
 

 
학습, 검증, 테스트 데이터의 크기를 확인합니다.
 
 
 

 
 
 

 
학습 데이터 1개를 가져와서 소스 문장과 타겟 문장를 확인합니다. 소스 문장은 "ein mann, der mit einer tasse kaffee an einem urinal steht."이며, 타겟 문장은 "a man standing at a urinal with a coffee cup."입니다. 트랜스포머 모델은 소스 문장을 입력받았을 때 타겟 문장을 출력해야 합니다.
 
 
 

 

 

 
SRC, TRG의 build_vocab() 메서드를 호출하여 독일어와 영어 단어 사전을 생성합니다. 단어 사전은 모든 문장 중에서 2번 이상 등장한 단어로 구성됩니다. 독일어 사전(SRC.vocab)에는 7853개의 단어,가 영어 사전(TRG.vocab)에는 5893개의 단어가 있습니다.
 

 

 
 

 
단어 사전에서 각 단어의 인덱스를 확인할 수 있습니다. 단어 사전에 없는 단어의 인덱스는 0이며, 패딩 토큰의 인덱스는 1, <sos> 토큰의 인덱스는 2, <eos> 토큰의 인덱스는 3입니다. 영어 단어 사전에서 "hello"의 인덱스는 4112입니다.
 
 

 

 

 
 
 

 
길이가 비슷한 문장들을 묶어서 하나의 배치를 구성했습니다. 파이토치의 BucketIterator를 통해 길이가 비슷한 문장들로 배치를 구성할 수 있습니다. 배치 크기는 128입니다.
 
 
 

 
 

 
학습 데이터로더에서 첫 번째 배치를 가져온 후, 해당 배치에 담긴 가장 긴 독일어 문장의 길이를 확인해봅시다. 또한 배치의 첫 번째 문장을 출력해봅니다.
 
 
 

 
 
 

 
첫 번째 배치에 담긴 가장 긴 독일어 문장은 27개의 토큰으로 구성된 "ein mann, der einen rucksack trägt, schaut hinauf zu zinem wasserfall." 입니다. 배치에서 가장 긴 문장이 27개의 토큰을 가지므로, 배치의 나머지 문장들에 대해 남는 부분은 패딩 토큰으로 채워져 길이가 27로 맞춰집니다.
 

 

 

트랜스포머 모델을 학습하기 위한 데이터의 전처리를 끝냈습니다. 이제 모델을 구현해봅시다.

 

 

 

 

 
트랜스포머에서 매우 중요한 컴포넌트인 Multi-head Attention을 구현해봅시다. query, key, value 벡터의 차원은 모두 $d_{model}$로 지정했습니다.
 
 
 

 

 
$d_{model}\mod h = 0$을 만족하는 어텐션 헤드의 개수 $h$로 Multi-head attention을 수행합니다. multi-head attention을 구현할 때 각 head $i$에 곱해지는 가중치 $QW_{i}^Q, KW_{i}^K, VW_{i}^V$를 하나하나 변수로 정의하기보다는, $(d_{model},  d_{k})$ 크기의 행렬 $W_{Query}, W_{Key}, W_{Value}$에 모든 head에 대한 가중치를 담아 처리하는 것이 구현 상 간편합니다. 코드에서 fc_q, fc_k, fc_v가 이러한 가중치 행렬에 해당합니다.
 
 
 

 

 
모든 head의 어텐션 결과를 concat하고 선형 변환을 취한 값이 multi-head attention의 결과입니다. 코드에서 선형 변환은 $(d_{model}, d_{model})$차원의 fc_o입니다. scale은 head의 차원값인 어텐션 스케일링 값입니다.
 
 

 

 

 
Multi-head attention의 forward 메서드를 살펴봅시다. 배치에 들어있는 128개 학습 데이터 각각에 대해 multi-head attention을 수행합니다. 메서드의 인자로 전달되는 query, key, valu 변수는 사실 임베딩 벡터이므로, 각 변수에 fc_q, fc_k, fc_v로 선형 변환을 취하여 실제 query, key, value에 해당하는 값인 Q, K, V를 얻습니다.


어텐션은 각 head에서 단독으로 수행됩니다. 이를 위해 Q, K, V의 차원을 변환해야 합니다. 각 head에 따라 행렬이 분리 되도록 $d_{model}$ 차원을 $(h, d_{model}/h)$ 차원으로 분리합니다. view() 메서드로 차원을 변환할 수 있습니다. view() 메서드의 두 번째 인자인 -1은 query/key/value의 개수입니다. 이러한 개수는 가변적이므로 -1로 지정해야 합니다. 그 후 permute 메서드로 차원의 위치를 변경합니다.
 

 

$Attention(Q, K, V) = {softmax({QK^{T} \over \sqrt{d_{k}}})V}$

 
 
차원의 위치를 바꿔주는 이유는 $QK^{T}$를 계산하기 위해서입니다. 이 연산은 코드 상으로 (query_len, head_dim) $\times$ (head_dim, key_len)이므로, $QK^{T}$를 계산하기 위해선 -1로 지정한 두 번째 차원을 세 번째 차원인 head_dim과 교체해야 합니다.
 
 

 

 

 

Key의 전치 행렬을 구하기 위해 K.permute(0, 1, 3, 2)를 수행합니다. $QK^{T}$ 값을 head의 차원으로 나누어 스케일링 합니다. 마스킹 된 부분은 $-10^{10}$으로 대체하고, softmax를 취하여 어텐션 스코어를 구합니다. 어텐션 스코어에 value를 곱하고, 각 head의 결과를 concat합니다. 코드에선 (n_heads, head_dim) 차원을 다시 hidden_dim 차원으로 변환함으로써 concat이 수행됩니다. 결과적으로 (batch_size, query_len, hidden_dim) 차원의 행렬을 얻을 수 있고, 여기에 fc_o 선형 변환을 취하여 multi-head attention의 결괏값을 얻습니다. 추후 어텐션 스코어를 시각화할 수 있도록 attention 변수도 함께 반환합니다.

 

 

인코더 레이어의 multi-head attention은 소스 문장의 전체적인 맥락을 학습합니다. 소스 문장을 배치 단위로 묶어 한 번에 처리하고 있는데, 배치 안에 있는 각 문장의 길이는 서로 다를 수 있습니다. 그래서 배치의 가장 긴 문장보다 짧은 문장에는 패딩 토큰을 추가한다고 했습니다. 패딩 토큰은 문장 안에서 의미를 갖지 않으므로, 어텐션에 관여하면 안 됩니다.

 

그러므로 $QK^T$에서 key가 패딩 토큰인 부분을 $-10^{10}$과 같은 매우 작은 값으로 대체합니다. 이 값에 softmax가 취해지면 어텐션 스코어 값이 0에 가까워지므로 패딩 토큰이 어텐션 과정에 미치는 영향이 거의 없다고 해석할 수 있습니다. 중요한 점은 패딩 토큰에 해당하는 query는 마스킹 하면 안된다는 점입니다. 이는 이해가 되지 않았던 부분이며, 인터넷에 있는 대부분의 트랜스포머 구현 글에서도 이유를 설명하지 않고 있습니다. 이에 대한 설명은 여기에서 다루고 있습니다.

 

 

 

 

 

 

해당 글에선 패딩 토큰인 query는 마스킹하면 안 된다고 합니다. 이유를 설명하기 전에, 먼저 패팅 토큰인 query를 위와 같이 마스킹 했다고 합시다. 여기에 softmax를 취하면 query가 패딩 토큰인 각 열(column)은 모두 동일한 값을 가지게 됩니다. 위 예제에선 그 값이 $1/6$입니다.

 

 

이는 패딩 토큰이 가지고 있던 정보를 모두 제거하는 행위입니다. 마스킹의 목적은 어텐션에서 패딩 토큰이 다른 토큰에 영향을 끼치지 않도록 하는 것이지, 패딩 토큰의 정보를 완전히 제거하는 것이 아닙니다. 모델은 패딩 토큰의 존재를 인지할 수 있어야 합니다. 이치럼 패팅 토큰인 query를 마스킹하고 얻은 어텐션은 모델 학습을 위한 형편 없는 표현(representation)이 됩니다. 그러므로, 어텐션에서 패딩 토큰이 다른 토큰에 영향을 끼치지 못하게 하려면 key가 패딩인 부분만 마스킹해야 합니다.

 

 

 

 

 

 

 

Multi-head attention 뒤에 오는 fully connected layer를 정의합니다. (query_len, hidden_dim) $\times$ (hidden_dim, pf_dim) $\times$ (pf_dim, hidden_dim) = (query_len, hidden_dim) 순으로 연산이 진행됩니다.
 
 

$FFN(x) = \max (0, xW_{1}+b_{1})W_{2}+b_{2}$
 
 

논문의 수식을 그대로 구현하되, 첫 번째 FC 레이어와 relu activation 뒤에 dropout을 적용했습니다.
 
 
 
 

 

 

 

multi-head attention과 feed forward layer 컴포넌트를 완성했으며, 이를 기반으로 인코더 레이어를 정의할 수 있습니다.

 

 

 
 
 

 
src는 (batch_size, src_len, hidden_dim) 차원으로 인코더 레이어에 입력됩니다. 논문에 나와있는 인코더 레이어의 구조 그대로, src에 multi-head attention(self attention)을 취해줍니다. skip connection과 layer norm을 취해주고 feed forward layer를 취해줍니다. 다시 skip connection과 layer norm을 취해줍니다. skip connection을 진행할 때 multi-head attention 결괏값에 dropout을 취해주고 있습니다

 

 

 
 
 

 

여러 개의 인코더 레이어를 쌓아 인코더 아키텍처를 정의합니다.
 
 

 

 

 

 

 

인코더는 (batch_size, src_len) 차원의 입력값을 받습니다. src_len은 배치에 담긴 가장 긴 문장의 길이입니다. 먼저 각 단어를 임베딩 벡터로 변환해야 하며, 여기에 각 단어의 위치 정보를 삽입하기 위해 포지셔널 인코딩 벡터가 더해져야 합니다. 논문에선 sin, cosine 함수를 이용한 포지셔널 인코딩을 수행하지만, 본 코드에선 파이토치에서 기본적으로 제공하는 단어 임베딩 모듈을 이용합니다. input_dim은 단어를 표현한 one-hot 벡터의 차원입니다. 논문에서 설명한 그대로, tok_embedding + pos_embdding의 결괏값을 $\sqrt d_{model}$로 스케일링하고 dropout을 적용합니다. 인코더는 총 n_layers 개수의 인코더 레이어로 구성됩니다.
 
 

 
 
 

 
파이토치의 nn.Embedding() 메서드에 pos 변수를 전달함으로써 포지셔널 인코딩 벡터를 만듭니다. pos 변수가 만들어지는 과정을 이해하기 위해 간단한 예제를 만들었습니다. 인코더의 출력 차원은 입력 차원과 동일한 (batch_size, src_len, hidden_dim)입니다.

 

 

 

 

 

 

 

 

디코더 레이어를 정의합시다. 디코더 레이어에선 타겟 문장에 대한 셀프 어텐션, 인코더-디코더 어텐션이 수행됨을 유의합시다.

 

 

 

 

 

 

 

디코더 레이어의 forward() 메서드 또한 논문의 구조를 그대로 따랐기 때문에 연산 부분에 대한 추가적인 설명은 필요하지 않을 것 같습니다. 그러나 타겟 문장을 셀프 어텐션할 때 사용할 마스킹 행렬을 유의해야 합니다.

 

 

 

 

 

 

논문 리뷰에서 설명했듯이, 타겟 문장에 대한 셀프 어텐션에서 현재 $i$번째 단어를 출력해야 하는 상황이라면, 출력 문장의 $1∼i−1$번째 단어들만 가지고 어텐션을 수행해야 합니다. 이를 구현하기 위해 마스킹 행렬에서 $i$번재 단어 뒤에 오는 단어 토큰들이 key로 사용되지 않게끔 마스킹해줍니다. 이를 look ahead masking이라고 부릅니다. 물론 타겟 문장에도 패딩 토큰이 존재하기 때문에 패딩 마스킹과 look ahead 마스킹에서 한 쪽이라도 마스킹 된 부분을 모두 마스킹 처리하여 실제로 사용할 마스킹 행렬을 만듭니다.

 

 

인코더-디코더 어텐션에서 query로 타겟 문장(셀프 어텐션 값)이 사용되고, key로 소스 문장(인코더 출력값)이 사용되는 것을 확인할 수 있습니다. 그러므로 key의 패딩 토큰 부분을 마스킹하면 됩니다.

 

 

 

 

 

 

디코더 레이어가 정의되었으니, 디코더 또한 정의합시다. 트랜스포머를 학습할 때는 디코더가 한 번만 실행(feed-forwarding) 되지만, 추론(테스트) 때에는 auto-regressive하게 동작하므로 디코더를 한 번 실행할 때마다 출력 문장의 단어가 순서대로 하나씩 출력됨을 유의합시다.

 

 

 

 

 

 

 

디코더의 forward() 메서드도 논문을 그대로 따르고 있어서 추가적으로 설명할 부분은 없을 것 같습니다. 인코더에서 소스 문장을 임베딩 벡터로 변환한 뒤, 포지셔널 인코딩 벡터와 더해줍니다. 디코더의 출력 차원은 입력 차원과 동일한 (batch_size, tar_len, hidden_dim)입니다.

 

 

논문에선 디코더의 마지막 부분에 softmax 함수를 넣어주는 반면, 현재 디코더의 구현에선 softmax 함수가 없습니다. 그 이유는 모델 학습 시 loss function으로 torch.nn.CrossEntropyLoss() 메서드를 사용할 것인데, 이 메서드에서 내부적으로 디코더의 output에 softmax를 적용해주기 때문입니다. 

 

 

 

 

 

 

인코더와 디코더로 전체 트랜스포머 모델을 구성해봅시다. forward() 메서드는 매우 간단합니다. 소스 문장을 인코더에 넣어 결괏값을 받고, 그 값을 이용하여 디코더 연산을 수행합니다. Transformer 클래스에서 가장 중요한 점은 마스킹 행렬을 어떻게 생성하고 있는지 입니다.

 

 

 

 

 

 

 

먼저, 소스 문장에 대한 마스킹 행렬을 만드는 make_src_mask() 메서드를 봅시다. src는 (batch_size, src_len) 차원입니다. 배치를 구성하는 각 문장에 대해 패딩 토큰인 부분을 체크하여 마스킹 행렬을 만들어줍니다. 그러면 마스킹 행렬의 크기도 (batch_size, src_len)가 됩니다. 그런데 마스킹 행렬은 multi-head attention에서 $QK^{T}$에 적용됩니다. 이 텐서의 크기는 (batch_size, n_heads, query_len, key_len)이므로, 마스킹 행렬을 (batch_size, 1, 1, src_len) 크기로 변환해줘야 합니다.

 

 

 

 

 

만들어진 소스 문장 마스킹 행렬은 위와 같습니다. 각 문장의 토큰은 원으로 표시 했으며, 주황색 원은 패딩 토큰을 나타냅니다. batch_size개의 문장은 서로 다른 패딩 토큰의 개수를 가질 수 있습니다.

 

 

 

 

 

 

 

타깃 문장에 대한 마스킹 행렬을 만드는 make_trg_mask() 메서드를 봅시다. 문장의 패딩 토큰을 마스킹하는 행렬을 만드는 방법은 이미 살펴본 바와 같습니다. 여기에 look ahead 마스크를 추가해야 합니다. look ahead 마스크는 trg_sub_mask라는 변수로 정의합니다. look ahead 마스크는 그저 (trg_len, trg_len) 크기의 삼각행렬입니다. torch.tril()과 torch.ones() 메서드로 쉽게 삼각행렬을 만들 수 있습니다. 두 마스킹 행렬에 &연산을 적용하면 타깃 문장 마스킹 행렬이 완성됩니다. 행렬의 크기는 (batch_size, 1, tag_len, trg_len)입니다.

 

 

 

 

 

 

이제 모델을 학습해봅시다. 학습에 필요한 하이퍼 파라미터를 정의합니다.

 

 

 

 

 

 

 

인코더와 디코더, 트랜스포머 객체를 선언합니다.

 

 

 

 

 

 

 

 

트랜스포머 모델은 약 900만 개의 파라미터를 가집니다.

 

 

 

 

 

 

 

모델의 가중치는 xavier uniform initiailzation으로 초기화 해줍니다.

 

 

 

 

 

 

 

adam optimizer로 학습을 합니다. 또한 모델이 예측한 타겟 문장과 정답 타겟 문장 간의 CrossEntropyLoss를 계산할 때 패딩인 토큰은 loss에 영향을 끼치지 않도록 설정합니다.

 

 

 

 

 

 

 

CrossEntropyLoss() 메서드에서 전달할 input과 target은 위와 같은 shape를 가지면 됩니다. 타겟 문장의 각 단어도 인덱스로만 표현해주면 됩니다.

 

 

 

 

 

 

 

 

 

트랜스포머를 학습시키는 메서드입니다. 트랜스포머 인코더의 입력값으로 src를 전달하고 있으며, 디코더의 입력값으로 trg[:, -1]을 전달하고 있습니다. 학습을 통해 모델이 출력값은 trg[:, 1]에 가까워져야 합니다. 정답이나 모델의 출력의 패팅 토큰의 경우 loss를 계산하는 데 반영되지 않으므로, 모델이 정답이 <pad>인 위치에 <good>을 출력하든 말든 loss에 반영되지 않습니다.

 

 

 

 

 

모델을 검증할 때 사용할 메서드입니다.

 

 

 

 

 

 

 

10 epoch로 모델을 학습해봅시다. 각 epoch 마다 걸리는 시간을 측정하기 위해 epoch_time 메서드를 정의했습니다.

 

 

 

 

 

 

 

 

학습 중 validation loss가 가장 적은 epoch의 모델을 저장합니다.

 

 

 

 

 

 

 

저장된 모델을 불러온 뒤, 테스트 데이터에 대한 loss과 PPL 값을 출력합니다.

 

 

 

 

 

 

 

직접 임의의 독일어 문장을 모델에 넣어 영어로 번역된 문장을 얻을 수 있습니다. sentence 객체가 실제 문자열(문장)인 경우, 이를 독일어 tokenizer로 토큰화하고 소문자로 변환합니다. 반면 이미 토큰화된 단어 묶음이라면 소문자로 변환만 해줍니다. 시작에 <sos> 토큰을, 끝에 <eos> 토큰을 붙여줍니다. 그 후 각 단어 토큰을 인덱스로 변환합니다. 

 

 

 

 

 

 

 

이제 트랜스포머 모델에 src, trg를 입력하여 번역 결과를 얻을 수 있습니다. trg는 <sos> 토큰 하나만 가지고 시작하며, 모델의 출력 단어 토큰(trg의 가장 마지막 단어에 대한 모델의 출력 단어)을 하나씩 붙여가면서 번역된 단어를 하나씩 얻습니다. 모델이 <eos> 토큰을 출력하면 번역을 완료했다는 뜻입니다.

 

 

 

 

 

 

 

우리가 얻은 값은 번역된 단어 토큰의 인덱스이므로, 이를 다시 실제 단어로 변환한 뒤 <sos> 토큰을 제외하고 출력 문장을 반환합니다.

 

 

 

 

 

 

 

이처럼 원하는 독일어 문장을 영어 문장으로 번역할 수 있습니다.