[LM] 그림으로 설명하는
트랜스포머

The Illustrated Transformer

by mashed moshirakano

Aug 25. 2023

출처: http://jalammar.github.io/illustrated-transformer/

트랜스포머의 텐서플로우 구현은 여기, 파이토치 구현은 여기서 찾아볼 수 있다.

위에서 조망하기 (A High-Level look)

어떤 문장이 주어졌을 때, 기계번역 블랙박스 모델은 해당 문장에 대한 번역결과를 산출한다

Encoding/Decoding component

인코딩 부분은 6개의 인코더가 쌓여진 형태로 구성된다. 디코딩 부분도 마찬가지로 6개의 디코더가 쌓여진 형태이다. (몇개의 인코더/디코더를 쌓을 지는 모델러의 재량에 달려있다)

1. Encoding component

이 때, 각 인코더들은 동일한 구조이며, weight를 공유하지 않는다. 하나의 인코더는 2개의 서브 레이어(self-attention, feed forward NN)로 분해할 수 있다.

self-attention

하나의 단어를 인코딩할 때, 입력된 문장의 다른 단어를 함께 고려할 수 있도록 돕는다.

feed forward NN

self-attention 레이어의 아웃풋은 feed forward NN의 인풋으로 피드된다. 동일한 feed forward NN이 각 위치에 독립적으로 적용된다

2. Decoding component

하나의 디코더안에도 인코더와 마찬가지로 self-attention, feed forward NN 레이어를 가지고 있지만, encoder-decoder attention 레이어가 중간에 자리잡고 있다. 이를 통해 입력된 문장에서 관련된 부분에 집중할 수 있게 한다 (seq2seq 모델에서의 어텐션의 역할과 비슷하다)

텐서를 그림으로 나타내기

트랜스포머의 주요 구조를 알아봤으니, 이제 다양한 벡터/텐서들이 트랜스포머 안에서 어떻게 흐르고 있는 지 살펴볼 것이다.

입력 문장 “나는 학교에 간다”가 주어졌을때, 각 단어들은 임베딩 알고리즘을 통해 벡터로 변환되었다.

각 단어는 512 차원의 벡터로 임베딩되었다. 본문에서는 4칸으로 단순화 시켜서 표현할 것이다

임베딩은 가장 인코더 중 첫번째 인코더(가장 바닥에 있는 인코더)에서만 수행된다. 나머지 5개의 인코더는 직전 인코더의 아웃풋을 입력으로 받게 된다. 보통 인코더는 사이즈 512의 벡터로 이루어진 리스트를 인풋으로 받게된다. 이 리스트의 사이즈(각 문장에서 최대 몇개의 단어를 고려할 것인지) 는 hyperparameter로 정할 수 있으나, 기본적으로는 학습 데이터에서 가장 긴 문장의 길이(문장 안의 단어의 갯수)와 동일하게 적용된다. 위의 예시에서는 (3, 512)가 인풋 행렬의 형태가 된다.

<임베딩 코드>

인풋 시퀀스의 단어들을 임베딩한 뒤, 각각의 단어들은 인코더안의 2개의 레이어를 각각 타고 흐른다. 여기서 우리는 트랜스포머의 가장 중요한 특성을 살펴볼 수 있는데, 이는 바로 각 단어들이 고유한 경로를 가지고 인코더를 통과한다는 것이다. self-attention 레이어에서는 이 경로들 간의 dependency가 존재하지만, feed-forward NN에서는 이러한 dependency가 존재하지 않는다. 따라서 feed forward 레이어를 통과할때, 다양한 경로들이 병렬적으로 처리될 수 있는 것이다.

self-attention 레이어에서 단어의 경로들 간에 dependency를 가진다는 것은,
주어진문장(sequence)안에 등장한 단어들끼리의 상대적인 중요도(weights to focus)를 학습한다는 것을 의미한다.

다음으로는, 좀 더 짧은 문장으로 바꿔서, 인코더의 서브 레이어 안에서 어떤 일이 벌어지는 지 더 자세하게 알아볼 것이다.

인코딩하기

앞서 말한 것 처럼, 인코더는 (단어)벡터들의 리스트를 인풋으로 받는다. 이를 self-attention 레이어 → feed-forward NN에 통과시키고 나온 결과를 다음 인코더에 넘겨준다.

1. 위에서 조망한 self-attention

다음의 문장을 번역한다고 하자.

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

위의 문장에서 "it"은 어떤 것을 의미하는가? street? animal? 만약 모델이 "it"을 처리하고 있었다면, self-attention은 "it"과 "animal"을 연결하게 된다. 즉, 한 문장에서 특정 위치에 있는 단어를 처리할 때, self-attention 은 해당 문장에서 다른 위치에 있는 단어를 함께 볼 수 있게 한다.

RNN을 떠올려보면, hidden state를 통해 이전에 처리한 단어벡터와 현재 처리된 단어벡터를 통합했었다. self-attention은 트랜스포머로 하여금 관련있는 다른 단어를 현재 처리하고 있는 단어에 함께 녹일 수 있도록 한다.

2. self-attention 을 자세히 뜯어보기

벡터를 사용하여 self-attention 을 어떻게 계산하는 지 살펴보고, 어떻게 실제로 구현되는지 알아보자.

x1을 가중치 행렬 WQ와 곱하면 해당 단어의 query vector q1을 얻을 수 있다. 이러한 방식으로 인풋 시퀀스안의 각 단어에 대한 query, key, value projection 을 수행하게 된다.

첫번째 스텝: query-key-value vector 생성
인코더의 인풋 벡터들(문장 안의 단어 벡터들) 각각에 대해 3개의 벡터를 만들어내는 것이다. 즉, 각 단어에 대해 query, key, value 벡터를 만들어낸다. 이 벡터들은 학습 과정에서 학습된 3개의 행렬에 임베딩을 곱해서 얻어진다.

예를 들어, “thinking” 라는 단어의 임베딩 벡터 x1; shape=(1, 4) 에 대해, query weight 행렬 WQ;shape=(4, 3)을 곱하면 query vector q1;shape=(1,3) 이 도출된다. 이때의 WQ는 학습 과정에서 업데이트된다.

두번째 스텝: dot product 계산

“thinking” 라는 단어의 self-attention을 계산해보자. 이를 위해서는 “thinking”와 입력 문장(”thinking machines”)안의 모든 단어간에 스코어를 계산해야한다. 이 스코어의 역할은, 특정 위치의 단어를 인코딩할 때 입력 문장안의 각 단어별로 얼만큼의 focus를 줄지를 결정한다. 현재 처리하고 있는 단어의 query 벡터와 문장 안의 모든 단어의 key 벡터를 dot product 하여 스코어를 얻을 수 있다.

만약 우리가 첫번째 위치에 있는 단어의 self-attention을 구한다면, 첫번째 스코어는 q1("thinking"의 query 벡터)과 k1("thinking"의 key 벡터)의 내적값이 되고, 두번째 스코어는 q1과 k2("machines"의 key 벡터)의 내적값이 될 것이다.

세번째 스텝: 차원수를 고려한 스케일링

스코어를 8로 나눠주는 것이다 (이때 8이라는 숫자는 논문에서 쓰인 key 벡터의 차원수 64를 루트 취한 값이다. 이를 통해 더 안정적인 gradient를 얻을 수 있다고 한다)

네번째 스텝: softmax 연산

softmax를 통해 스코어들을 normalize하여 모든 값이 양수이고 합이 1이 될 수 있게 한다. 이 softmax 스코어를 통해, 첫번째 위치에서 문장안의 각 단어들이 얼만큼의 표현력을 가지는 지 알 수 있다.

다섯번째 스텝: sofmax score X value vector → weighted value vector

각 value 벡터와 softmax 스코어를 곱해준다. 직관적으로 이해해보자면, 우리가 focus하고 싶은(즉, 현재의 단어와 관련이 높은) 단어들은 온전한 벡터 값을 보존하되, 관련없는 단어들의 벡터 값은 희석되게 하는 것이다.

예를 들어, softmax score가 0.0001과 같이 작은 수가 나왔을 때, value 벡터와 곱해지면 그 단어 벡터의 값은 매우 작아지게 된다.

여섯번째 스텝: value vector 의 가중합

위에서 softmax score와 value vector들이 곱해진 결과인 weighted value vector들을 모두 더해준다. 이것이 (첫번째 단어에 대한) self-attention 레이어의 아웃풋이 된다.

요약하면,

self-attention 레이어는 각 단어에 얼만큼의 가중치를 줄지 계산하는 레이어이다.

이를 위해서는 각 단어의 query/key/value 벡터와 가중치 행렬이 필요하다.

1~4번째 스텝은 가중치 행렬(softmax score)을 계산하는 과정이고,

5~6 번째 스텝은 각 단어 벡터에 가중치 행렬을 곱한 뒤 더해서 하나의 벡터 값이 산출되게 하는 과정이다.

"나는 학교에 간다" 라는 문장이 주어졌다면, 가중치 행렬을 구하는 방식은 다음과 같이 해석할 수 있다.

(1) q_v(학교에) 와 k_v(모든 단어 in 문장)간의 유사도를 구하고 = dot product

(2) 이 유사도(내적 값)를 scaling = (divide by root(dim) & softmax)

self-attention 레이어에 통과시켜 얻은 아웃풋은 다음과 같이 해석할 수 있다.

v(학교에) = 0.01 x v(나는) + 0.9 x v(학교에) + 0.09 x v(간다) = 각 단어 벡터별로 focus 한 결과의 합

scaled dot product attention 라고도 하며, 연산과정을 위와 같이 도식화 할 수 있다.

머리가 여러 개 달린 괴물이 등장하다

논문에서는 self-attention 레이어에 multi-headed attention을 붙여서 더욱 고도화시켰다. multi-headed attention 은 기존의 attention 레이어를 두개의 관점에서 발전시켰다.

(1) 다른 위치에 focus할 수 있는 모델의 능력을 확장시켰다. 물론, 기존의 self-attention이 다른 단어의 인코딩을 조금씩 담고 있었지만, 그래도 처리하고 있는 단어 그자체의 인코딩에 가장 큰 영향을 받는다

(2) attention 레이어가 여러 개의 representation subspace를 가질 수 있게 하였다. 이제부터 알아보겠지만, multi-headed attention을 통해 여러개의 query-key-value 가중치 행렬 집합을 얻을 수 있다(트랜스포머는 8개의 attention head를 가지고 있기 때문에, 각 인코더와 디코더는 각각 8개의 qkv가중치 행렬 집합을 가지고 있다) 각각의 qkv 가중치 행렬 집합은 랜덤으로 초기화된다. 그리고 학습이 끝나면 각 qkv 가중치 행렬을 이용하여, 인풋 임베딩들(혹은 이전 인코더/디코더에서 얻은 벡터들)을 다른 representation subspace로 projection 시키게 된다.

위에서 설명한 것처럼 self-attention 계산을 하게 되면, 8개의 서로 다른 가중치 행렬을 곱해서 8개의 서로 다른 결과 행렬 Z가 나오게 된다.

이제 우리는 약간 곤란한 상황에 직면하게 되었다. feed-forward 레이어는 8개의 행렬이 들어올 줄은 모르고 있기 때문이다 (그저 하나의 단어벡터에 대한 하나의 self-attention 결과 행렬을 기다리고 있을 뿐)

그래서 우리는 이 8개의 행렬을 하나의 행렬로 압축시킬 것이다. 이 8개의 행렬을 이어붙인 뒤(concatenate), 추가적인 가중치 행렬 WO와 곱해주는 것이다.

요약하면, 다음과 같다

1) 인풋 행렬을 받는다

2) 각 단어를 임베딩한다

3) (첫번째 인코더라면 임베딩 벡터에 대해, 그렇지 않으면 이전 인코더의 결과 벡터에 대해) 8개의 Wq, Wk, Wv 가중치 행렬을 만들고 입력벡터와 곱해준다

4) Q, K, V 행렬을 통해 attention을 계산한다

5) 결과 행렬 Z들을 이어 붙인뒤, 가중치행렬 WO에 곱해서 레이어의 아웃풋을 만든다

positional encoding으로 sequence의 순서를 표현하다

인풋 시퀀스의 단어의 순서를 어떻게 고려할 지에 대해서 설명하겠다. 트랜스포머는 각 인풋 임베딩에 벡터를 더해준다. 이 벡터들은 각 단어의 위치를 학습할 수 있도록 특정한 패턴을 띄게 된다. 직관적으로 풀어보면,

임베딩에 positional encoding 벡터를 더해준 뒤에, Q/K/V 벡터로 projection 되고 내적 어텐션을 시킨 뒤 임베딩 간의 유의미한 거리를 가지게 된다. 만약 임베딩의 차원이 4라면, 실제 positional encoding은 다음과 같다

다음의 figure를 통해 positonal encoding을 좀 더 자세히 알아 보자.

각 행은 각 벡터의 positional encoding에 대응된다. 즉, 첫번째 행은 인풋 시퀀스(입력 문장)의 첫번째 단어임베딩에 더해지는 벡터인 것이다. 각 행은 512개의 값을 가지고 (더해지는 단어임베딩의 차원과 동일하다) 각 값은 -1과 1사이에 있다. 또한, 중간에서 갈라지는 패턴을 찾아 볼 수 있다. 이는, 왼쪽 반절의 값들이 사인 함수로부터 생성되고, 오른쪽 반절의 값들이 코사인 함수로 부터 생성되었기 때문이다. 이렇게 생성된 두 부분을 이어서 하나의 positional encoding 벡터를 만들게 된다.

논문에서는, positional encoding을 생성하기 위해서 get_timing_signal_1d() 이라는 함수가 쓰였다. 이를 통해, 보지 못한 길이의 시퀀스에도 scale할 수 있게 된다 (예를 들어, 학습된 모델이 학습 데이터셋에 등장한 그 어떤 문장보다 더 긴 문장을 번역해야되는 경우)

2020년 7월 업데이트: 이 논문에서는 두 함수가 생성한 신호를 이어붙이지 않고, 엮어서(interweave) 하여 positional encoding을 만든다.

The Residuals

마지막으로 인코더의 구조에서 한가지 짚고 넘어가야 될 것이 있다. 서브레이어 (self-attention, ffnn) 각각은 residual connection을 주변에 두르고 있다. 그리고 layer-normalization을 하게 된다. 즉, self-attention → residual connection + layer norm → feed-forward → residual connection + layer norm 인 것이다.

벡터와 self-attention관련 layer-norm 연산을 시각화하면 다음과 같다.

Decoding

인코딩에서는 인풋 시퀀스를 처리하는 것으로 시작했다. 디코딩에서는, 가장 윗부분의 인코더의 아웃풋이 어텐션 벡터 K와 V 집합으로 변환된다. 이 벡터들은 각 디코더안의 encoder-decoder attention 레이어에서 디코더가 인풋 시퀀스에 적절한 장소에 focus할 수 있도록 사용된다.

계속 같은 작업을 반복하다가, special symbol에 도달하게 되면 종료한다 (special symbol이란 트랜스포머의 디코더가 최종 아웃풋을 냈다는 지표로 쓰인다) 다음 타임스텝에서, 각 스텝의 아웃풋은 가장 밑바닥(처음)의 디코더에 주입된다. 디코더는 자신의 디코딩 결과를 인코더가 했던 것처럼 bubble up?한다. 인코더의 인풋처럼, 디코더의 인풋을 임베딩하고 positional encoding을 더해서 각 단어의 위치를 표시할 수 있게 한다.

디코더에서의 self-attention 은 인코더와 약간 다르다. 디코더에서는 self-attention 레이어가 (아웃풋 시퀀스에서) 이전 위치만 참고하도록 허용된다. self-attention을 계산 시, softmax 스텝 전에, 다음에 오는 위치를 마스킹 처리(-inf로 세팅)하여 이를 수행한다.