Transformer 구현하고 이해하기(2)

‍한지영·2022년 3월 14일
3

Transformer

목록 보기
2/2
post-thumbnail
  • 무려 작년 5월 1일 발제했던 10달묵은 트랜스포머 정리.
  • 🤚 tf 코드입니다.
  • 킹키독스 : https://wikidocs.net/31379
  • transformer 개념 정리, Positional encoding 함수 구현에 대해서는 이전 포스트를 참고하세요

2. 트랜스포머(코드)

2) multi-head Attention 구현

* Scaled dot product attention?

어텐션을 수행하기 위해서는 Query, Key, Value 벡터가 필요하다. Q K V 벡터는 어떻게 얻는가? 위 그림처럼, 각 단어 벡터들로부터의 연산으로 Q K V 벡터를 얻을 수 있다. 이 때 Q벡터 K벡터 V벡터의 차원은 어떻게 될까? 인코더의 초기 입력이자 트랜스포머 모든 출력층의 출력 차원인 "d_model"차원의 단어 벡터들보다 더 작은 차원을 가진다. 얼마나 더 작은 차원? Q,K,V의 차원은 트랜스포머의 하이퍼파라미터인 Num_heads 값에 의해 결정되는데, Q,K,V는 (d_model/num_heads)의 차원을 가지게 된다. 이는 차후 설명할 'Multi-head'와 관련 깊다.

다시 돌아가서, Q, K, V 벡터를 얻는 방법은 이 Student란 단어 벡터에 가중치 행렬을 곱하면서 더 작은 Q, K, V 벡터(그림의 Qstudent, Kstudent, Wstudent)를 얻을 수 있다. 모든 단어 벡터에 대해 위와 같이 서로 다른 3개의 가중치행렬을 곱하는 과정을 거치면 모든 단어 벡터들에 대한 Q K V 벡터를 얻을 수 있다.

Q, K, V 벡터를 얻었다면, Q K V 벡터를 이용해 어텐션 스코어를 구한다. 쿼리벡터는 모든 키 벡터에 대해서 어텐션 스코어를 구하게 된다. 트랜스포머에서 어텐션 스코어를 구하기 위해 사용하는 함수는 위 그림의 우상단 수식과 같다. 단순히 내적만을 사용하지 않고 특정 값으로 나눠주어 스케일링했다고 해서 scaled dot product attention이라고 한다.

attention을 수행하는 과정을 다시 정리해보자.
1. 단어 i에 대한 쿼리 벡터에 모든 키벡터를 곱한다.
2. 특정 값, 즉 루트 dk(dk= 키벡터의 차원)로 스케일링해준다.
✅ 스케일링까지 마쳐준 이 attention score는 쿼리벡터가 키벡터와 얼마나 연관되어 있는지를 보여주는 수치이다.
3. attention score에 softmax 함수를 사용하여 어텐션 분포를 구한다.
4. attention distribution과 value 벡터를 가중합을 해서 단어 i에 대한 최종 어텐션 값을 구한다.

그리고 이는 모든 쿼리 벡터에 반복된다.


이는 행렬 연산으로 손쉽게 처리될 수 있다. 다음은 행렬 연산으로 모든 어텐션값을 한번에 계산하는 모습이다.

어텐션 값 행렬을 구하기 위한 과정은 다음과 같다.
1. 단어 벡터가 아닌 문장 행렬에 가중치 행렬 wq wk wv를 곱해서 q,k,v 행렬을 구한다.
2. 쿼리 행렬과 전치한 키 행렬을 곱한다.
3. 루트 dk(dk= 키벡터의 차원)로 나눠 attention score 행렬을 구한다.
4. 소프트맥스 함수를 취해서 attention distribution을 구한다.
5. value 행렬을 곱해 최종적으로 어텐션 값 행렬을 구한다.


* Scaled dot product attention - 코드

이제 코드를 이해해보자.

간단한데, 우선 쿼리 행렬 Q와 키행렬 K를 곱해주고 있다.(matmul= 행렬곱 함수) Transpose_b=True를 해줌으로써 키행렬을 전치해서 곱해주고 있다.

그리고 행렬곱을 키벡터의 차원을 뜻하는 dk의 루트값으로 나눠준다.(스케일링)

그리고 소프트맥스를 취해 어텐션 분포 행렬을 얻은 뒤, 마지막으로 value 행렬과 곱해줌으로써 최종 어텐션 값 매트릭스를 반환하고 있다.

그리고 빨간박스 부분이 마스킹에 해당하는 부분이다. 마스킹이란, 간단히 설명하면 어텐션 스코어 행렬의 마스킹할 위치에 매우 작은 음수값을 넣는 부분이다. 트랜스포머의 마스킹에는 패딩마스크와 룩어헤드 마스크가 있는데, 먼저 패딩 마스크 먼저 언급하겠다.

PADDING MASK?


쿼리 행렬과 키 행렬을 곱해준 위 attention score 행렬을 보자. 입력 문장에 pad 토큰이 들어가는 것을 알 수 있다.(패딩) pad 토큰에는 실질적인 의미가 없기 때문에, 키에 패드 토큰이 존재한다면 이에 대해서는 유사도를 구하지 않기 위해 "굉장히 작은 음수값"을 넣어주는데, 이것을 패딩 마스킹이라 한다.

아주 작은 음수값을 넣어주고, 소프트맥스 함수를 지나게 되면 마스킹된 이 위치의 값들은 0에 가까운 값이 되어 단어간 유사도를 구할 때 패드토큰이 반영되지 않게 된다. 즉, 키에 패드 토큰이 있는 경우에는 열 전체에 아주 작은 음수를 넣어줌으로써 어텐션에서 제외하는 것을 패딩 마스킹이라고 한다.

* Multi-head attention?

Scaled dot product attention에 대해서 이해했으니, 이제 멀티헤드 어텐션에 대해 이해해보자. 트랜스포머가 사용하는 3가지의 어텐션은 모두 멀티헤드 어텐션이다.

멀티 헤드 어텐션이란, 멀티 헤드의 뜻 그대로 머리가 여러개라는 뜻이며, 단어와 단어 사이의 연관성은 어떻게 보느냐에 따라 연관된 정도를 다르게 나타낼 수 있기 때문에 멀티헤드 어텐션을 이용한다. 한 번의 어텐션을 하는 것보다 여러 번의 어텐션을 병렬로, 즉 head의 개수만큼의 어텐션을 병렬로 하는 것이 더 효율적이기 때문이다.

앞서, q k v 차원은 num_heads라는 하이퍼 파라미터에 의해 결정되며, qkv는 d_model을 num_heads로 나눈 차원을 가진다고 언급했다.

이 의미는, 만약 num_heads가 8이라면 8개의 어텐션이 병렬로 이루어지며, 가중치 행렬값은 각각의 어텐션 해드마다 전부 상이하기 때문에 다른 시각으로 정보를 수집할 수 있다는 의미이다.

위 그림 좌측처럼 병렬 어텐션을 모두 수행하였다면, 우측처럼 모든 어텐션 헤드들을 연결하고, 어텐션헤드를 모두 연결한 후 가중치행렬을 곱하게 된다.

가중치행렬까지 곱한 결과물인 멀티헤드 어텐션 행렬은 인코더의 입력이었던 문장행렬의 크기와 동일한데, 이는 트랜스포머가 다수의 인코더를 쌓은 형태이기 때문에 인코더에서의 입력의 크기가 출력에서도 동일크기로 유지되어야 다음 인코더에서도 입력이 될 수 있기 때문이다.(몹시 중요하다)

* Multi-head attention - 코드

그럼 멀티헤드 어텐션 구현 코드를 보자.

멀티 헤드 어텐션의 구현은 크게 다섯 가지 파트로 구성된다.

  1. WQ, WK, WV(Q, K, V 행렬을 만들기 위한 가중치행렬)에 해당하는 밀집층(Dense layer)을 지나게한다. 앞서 설명한 가중치행렬을 곱하는 과정을 구현하면 입력을 밀집층을 지나게 하는 형태가 된다.
  2. 지정된 헤드 수(num_heads)만큼 쿼리 키 밸류를 나눈다(split).
  3. 스케일드 닷 프로덕트 어텐션을 num_heads만큼 병렬적으로 수행하고,
  4. 나눠졌던 헤드들을 연결(concatenatetion)한다.
  5. 마지막으로 WO에 해당하는 밀집층을 지나게 한다. (WO는 어텐션 헤드들을 연결 후에 곱해주는 행렬.)

이제 코드를 보자.

먼저 멀티헤드 어텐션 클래스에서는 assert 함수를 이용해서 d_model 을 num_heads로 나눈 값, 즉 q,k,v의 차원이 나누어 떨어지도록 조건을 걸고 있다.(assert는 뒤의 조건이 True가 아니면 AssertError를 발생시키는 함수)

그리고 WQ, WK, WV, WO에 해당하는 밀집층을 정의하고 있다. Dense layer 정의에서 유닛을 모든층의 출력차원인 d_model로 지정한 이유는, WQ, WK, WV의 경우 헤드의 개수대로 나누기 위해서 여전히 d_model의 차원이어야 하기 때문이고, WO의 경우 마지막 덴스 레이어를 지나고 나오는 결과값이 멀티헤드 어텐션의 출력인데, 서브레이어들의 결과값, 즉 모든 출력층의 출력 차원은 d_model의 차원이어야 하기 때문이다.

또한 num_heads의 개수만큼 q k v를 스플릿하는 함수 역시 정의하고 있다.

그리고 call 함수를 보면, 실제로 받은 input을 각각 밀집층을 지나게 하고, 헤드를 나누고, 헤드의 수만큼 스케일드 닷 프로덕트 어텐션을 수행하고, 나눠졌던 헤드들을 연결한 후 마지막으로 덴스레이어를 지나게 해서 아웃풋을 산출하는 것을 확인할 수 있다.


3) Transformer Encoder 구현

* 코드

위에서 구현했던 multi-head attention 클래스를 바탕으로 인코더를 구현해보자.

위 encoder_layer 함수가 받고 있는 파라미터는 dff, d_model, num_heads, dropout이다. 이 중 앞서 설명하지 않은 것은 dff가 유일한데, dff란 트랜스포머 내부의 피드 포워드 신경망의 은닉층의 크기를 의미한다.

코드를 보면, 먼저 input 함수 -> 입력층의 shape을 정의해준다. 그리고 인코더에서는 입력 문장에 패딩이 있을 수 있으므로 어텐션 시 패딩 토큰을 제외하도록 패딩 마스크를 사용한다.이 패딩 마스크는 MultiHeadAttention 함수의 mask의 인자값으로 들어가고 있다.

그리고 트랜스포머 인코더를 이루는 두 개의 서브층인 멀티 헤드 어텐션과 FFNN의 구현이 이루어 지고 있다.

첫번째 서브층인 멀티헤드 인코더 셀프 어텐션에는 셀프 어텐션인 만큼 q k v 모두 동일하게 인코더의 인풋을 넣어주고 있다. 그리고 패딩마스크를 사용해주고 있다.

multihead attention층과 FFNN층 사이에는 드롭아웃, 잔차연결, 층정규화가 이루어지고 있다. 먼저, dropout rate를 받아서 과적합을 방지하기 위해 dropout을 중간에 진행하, 이후 잔차연결과 층정규화 과정도 거친다.

잔차연결?

이 때 잔차연결이란, 서브층의 입력과 출력을 더하는 것을 말한다. 일반적인 신경망 모델 학습 시 모델의 층이 깊어질수록 학습 결과가 좋을 수 있다고 알려져 있다. 하지만 층을 너무 깊이 쌓거나 노드 수를 너무 크게 증가시키면 입력 정보가 여러 층을 거치면서 이전 층에 대한 정보 손실이 발생할 수 있. 따라서 이전 층의 정보를 이용하기 위해 이전 층의 정보를 연결하는 잔차 연결(skip connection)을 적용한다.


트랜스포머에서도 멀티헤드어텐션의 input과 출력인 attention을 더해주고 있다. 트랜스포머에서 입력과 출력의 차원은 동일하기 때문에, 덧셈 연산이 가능하다. 위 인코더의 구조에서 인코더 인풋(positional encoding을 거친 embedding)으로부터 add & norm으로 향하는 화살표가 바로 잔차연결을 의미한다.


다시 코드로 돌아오면, 어텐션의 결과는 잔차연결을 거친 후 keras가 제공하는 함수로 층정규화되고 있다. 층 정규화는 텐서의 마지막 차원에 대해서 평균과 분산을 구하고, 이를 가지고 값을 정규화하여 학습을 돕는 것으로, 이 잔차연결과 층정규화, 또 드롭아웃의 과정은 트랜스포머에서 각 서브층 이후마다 수행된다.


그 다음은 두번째 서브층인 피드 포워드 신경망이다. 피드포워드 신경망 식을 그림으로 표현하면 위와 같은데, 여기서 중요한 것은 두번째 서브층을 지난 인코더의 최종 출력 역시 인코더의 입력 크기였던 (문장의 길이, d_model)의 크기를 유지한다는 점이다.

코드를 보면, 피드 포워드 신경망은 앞선 멀티헤드 어텐션의 결과로 나온 행렬을 인풋값으로 받고, 첫번째 덴스 레이어의 유닛은 피드 포워드 신경망의 은닉층의 크기인 dff로 지정하고 있음을 알 수 있다. 또한 트랜스포머의 모든 층의 출력차원은 d_model이어야 하기에 두번째 서브층, 피드 포워드 신경망의 최종 출력을 결정지을 units를 d_model로 지정해 둔 것을 확인할 수 있다.

피드포워드 신경망 서브층 이후에도 드롭 아웃, 잔차 연결과 층 정규화가 수행된다. 이렇게 하나의 인코더층을 지난 아웃풋은 다음 인코더 층으로 들어가 동일한 연산을 반복한다.(논문은 6개의 인코더, 6개의 디코더를 쌓고 있다)


transformer encoder 쌓기

하나의 인코더 층을 구현했으니, 이제 여러 인코더 층을 쌓는 코드이다.
먼저 인풋값을 임베딩하고, 포지셔널 인코딩까지 진행한후에 드롭아웃 해주고, num_layers 수만큼 아까 구현한 인코더 층을 쌓고 있다.(encoder layer def) output이 반복문을 돌면서 다음 레이어로계속 들어가는 것을 볼 수 있다. 그리고 마지막 인코더 층에서 얻는 행렬을 디코더로 보내줌으로써 트랜스포머 인코더의 인코딩 연산이 끝나게 된다.


4) Transformer Decoder 구현

* Transformer Decoder?


트랜스포머의 디코더는 지금껏 말했듯 세 개의 서브층으로 구성된다.

  1. Masked Decoder Self-Attention
  2. Encoder-Decoder Attention
  3. FFNN

1.디코더 셀프 어텐션은 q k v 값이 모두 디코더의 것인 반면에 인코더 디코더 어텐션의 경우 q값은 디코더에 존재하지만 k v값은 인코더의 값이라는 차이가 있다. 그리고 디코더 셀프 어텐션의 또 하나의 특징은 인코더 셀프 어텐션과 달리 마스크 매트릭스를 이용해서 뒤의 단어의 정보를 이용하지 않는다는 차이점이 있었다.

* 코드


역시나 인풋의 차원을 d_model로 지정해주고 있고, 인코더 디코더 어텐션에서 인코더 아웃풋을 받아와서 써야하기 때문에 인코더 아웃풋 역시 차원을 d_model로 지정해주고 있다.

그리고 앞서 인코더 레이어에도 있었던 패딩 마스크 말고도, 룩어헤드 마스크를 선언해주고 있다. 인코더와 달리 디코더에서는 룩어헤드 마스크 행렬을 이용하는 이유? 디코더에서는 단어를 전체적으로 이용하면 안되기 때문이다.

위 트랜스포머 layer 구성을 보면 디코더도 인코더와 동일하게 임베딩 과정과 포지셔널 인코딩 과정을 거치는데, rnn을 사용하지 않는 트랜스포머의 특성 상 트랜스포머의 디코더는 인코더처럼 문장 행렬 전체를 한 번에 입력 받는다. 그리고 이 디코더는 입력받은 이 문장 행렬로부터 단어를 예측하도록 훈련이 된다.

seq2seq와 비교를 해보면, seq2seq의 디코더에 쓰이는 rnn 계열 신경망은 순차적으로 시퀀셜하게 인풋을 입력받으므로 단어 예측에 현재 시점 이전의 입력된 단어들만 참고할 수 있지만, 트랜스포머는 문장 행렬 전체를 한번에 인풋으로 받으므로 단어를 예측하는데에 이보다 더 후시점의 단어까지도 참고할 수 있는 현상이 발생한다.

이를 극복하기 위해 트랜스포머 디코더에서는 룩어헤드 마스크를 도입한 것이다.

디코더의 첫번째 서브층인 “마스크드 멀티 헤드 셀프 어텐션”에서는 기본적으로 인코더의 멀티 헤드 셀프 어텐션과 동일한 연산을 수행한다.

제일 먼저 쿼리 행렬과 전치한 키 행렬을 곱하는데, 이로인해 어텐션 스코어 매트릭스가 만들어진다. 그리고 이 스코어 행렬에, 자기 자신보다 미래에 있는 단어들은 참고하지 못하도록 위와 같은 마스킹을 합니다.

이를 구현하려면 어떻게 해야 할까?

look_ahead_mask를 구현하는 코드는 위와 같다.


만들어진 룩 어헤드 마스크는 패딩 마스크와 마찬가지로 앞서 구현했던 멀티 헤드 어텐션 안에 있는 닷 프로덕트 어텐션 함수(위 코드)에 mask라는 인자로 전달된다. 패딩 마스킹을 써야 하는 경우에는 앞서 인코더를 구현할 때 보았던 것처럼 인자로 패딩마스크를 전달하고, 룩 어헤드 마스킹을 써야 할 때는 룩-어헤드 마스크를 전달하게 된다. 그러면 해당 마스크에 아주 작은 음수 값을 곱해줌으로써, 이 값이 소프트맥스 함수를 지나고나면 0에 가까워져 해당 토큰, 즉 패드 토큰이나 디코더에서 단어를 예측할 때 사용할 수 없는, 미래의 단어들을 어텐션에 반영하지 않게 되는 것이다.


간단히 이 scaled dot product attention 안의 마스킹 코드에 대해 정리를 해보면,

인코더의 셀프 어텐션 서브층에서는 인자로 패딩 마스크를,

디코더의 마스크드 셀프 어텐션 서브층에서는 인자로 룩어헤드 마스크를,

디코더의 인코더-디코더 어텐션에서는 패딩 마스크를 전달한다고 보면 된다.

다만 룩어헤드 마스크를 한다고해서 패딩 마스크가 불필요한 것은 아니기 때문에 룩어헤드 마스크는 패딩 마스크를 포함하도록 구현하고 있다.(create_look_ahead_mask def 참고)


다시 디코더 레이어 구현 코드로 돌아가보자.

첫번째 서브층, 즉 마스크드 셀프 어텐션에서 쓸 룩어헤드 마스크와 두번째 서브층인 인코더 디코더 어텐션에서 쓸 패딩 마스크를 정의해주고 있다.

그리고 첫번째 어텐션 서브층을 쌓아주는데 이는 마스크드 셀프 어텐션으로, 넘겨주는 Q K V값이 모두 디코더의 인풋으로 같음을 확인할 수 있다. 그리고 마스크는 룩어헤드 마스크를 넘겨준다.

그리고 디코더에서도 인코더에서와 같이 서브층 이후에 잔차 연결과 층 정규화가 수행되고 있다. 잔차 연결을 위해 attention1의 결과값과 인풋값을 더해주고 있고, 그위에 층정규화 함수를 씌워주고 있다.

디코더의 두번째 서브층은 인코더 디코더 어텐션이다. 주의할 것은 인코더 디코더 어텐션의 경우 계속 언급해왔듯이 q와 kv 값이 같지 않다는 사실이다. 즉, 셀프어텐션이 아니라는 사실이다.


빨간 박스를 잘 보면, 한개의 화살표는 잔차연결과 층정규화를 거친 마스크드 셀프 어텐션의 출력값이 인풋으로 들어오고 있고, 두 개의 화살표는 인코더 아웃풋에서 오는 것을 확인할 수 있다. 즉, 인코더 디코더 어텐션에서는 쿼리로는 디코더 행렬을, 키와 밸류로는 인코더 행렬을 받아서 위 그림 하단과 같은 어텐션 스코어를 구하는 행렬곱을 해주신다고 보면 된다.

이는 트랜스포머는 모든 서브레이어들의 결과값, 즉 모든 출력층의 출력 차원이 d_model의 차원이기 때문에 가능한 결과이다.

그리고 다시 코드를 보면 디코더 두번째 서브층의 마스크로는 패딩 마스크를 넘겨주는 것을 확인할 수 있다. 서브레이어 다음에는 늘 그렇듯 드롭아웃, 잔차연결, 층 정규화를 해주고, 세번째 서브층인 피드포워드신경망에는 잔차연결과 층정규화까지 완료된 인코더 디코더 어텐션의 출력값을 인풋으로 넣게되고, 인코더의 피드포워드 신경망과 마찬가지로 첫번째 덴스 레이어의 유닛은 피드 포워드 신경망의 은닉층의 크기인 dff로 지정하고, 트랜스포머의 모든 층의 출력차원은 d_model이어야 하기에 최종 출력을 결정지을 units은 d_model로 지정해 둔 것을 확인할 수 있다. 그리고 정말 마지막으로 잔차연결 후 층정규화 함수를 씌우고, 그 값을 리턴하고 있다.


transformer decoder 쌓기

디코더 레이어를 쌓는 코드는 인코더 레이어를 쌓는 코드와 거의 같다.

인풋값을 임베딩하고, 포지셔널 인코딩까지 진행한후에 드롭아웃 해주고, num_layers 수만큼 아까 구현한 디코더 층을 쌓는다. output이 반복문을 돌면서 계속 인풋으로 들어가는 것을 볼 수 있다. 그리고 디코더는 아웃풋과 패딩 마스크 이외에도 룩어헤드 마스크, 인코더 아웃풋도 사용하기 때문에 이도 같이 넘겨주는 것을 볼 수 있다.

5) Transformer 구현


마지막으로 지금까지 구현한 인코더와 디코더 함수를 끼워맞춰 트랜스포머를 조립한다.

코드를 보며 인지해 두어야 될 부분은 아래 정도가 있다.

  1. 인코더 함수의 입력(트랜스포머 인코더에서 필요한 것)은 입력 문장과 패딩 마스크
  2. 디코더의 함수의 입력(트랜스포머 디코더에서 필요한 것)은 디코더 인풋, 인코더 아웃풋, 룩어헤드 마스크과 패딩 마스크
  3. 인코더의 출력은 디코더에서 인코더 디코더 어텐션에서 사용되기 위해 디코더로 전달되고, 디코더의 출력은 예측을 위한 출력층으로 전달됨

여기서 예측을 위한 출력층이란 디코더의 끝단에 있는 부분, 즉 다중 클래스 분류 문제를 풀 수 있도록, unit이 vocab_size 인 신경망을 추가해준 return 전 마지막줄을 의미한다.


이상 10달묵은 트랜스포머 구조 설명/코드 구현(TF)을 마친다!🤗
다시 보면서 트랜스포머 리마인드한김에, 시간나면 깃헙에 이 TF 코드 토치로 다 바꾸는 연습해보고 트랜스포머 포스팅 하나 더 하도록 하겠다.(언제가 될진 모르겠지만👊👊)

profile
NLP 전공 잡식성 문헌정보 석사생

0개의 댓글