[CV 연구] Transformer Attention Visually Explained

[But what is a GPT? Visual intro to transformers Chapter 5, Deep Learning]

[Attention in transformers, visually explained Chapter 6, Deep Learning]

본 포스트는 3Blue1Brown 영상을 참고하여 개인적으로 정리한 글임을 밝힙니다.

Attention Pattern: Q, K dot product한 matrix

개념적이해: 문장에서 하나의 word token에 대한 initial embedding vector는 dot product연산으로 이전 word값들을 참조하여 가중치가 반영된 벡터들의 합으로 업데이트 되면서 high dimensional space 상에서 이전 word들과 연관성있게 벡터의 방향과 값의 크기가 업데이트 됩니다.

inital embedding vector는 Attention Pattern에서는 하나의 column vector에 해당합니다.

Column vector(initial embedding vector) 1개를 Tower라 하고 아래를 생각해봅시다.

“Tower”의 initial embedding vector와 “Eiffel”의 embedding vector와 dot product로 구한 가중치(Q, K 연산결과)만큼 “Tower”의 initial embedding vector를 “Eiffel”의 embedding vector 방향으로 high dimensional space상에서 이동합니다.

즉, 가중치가 반영된 벡터를 1개 더해서 “Tower”의 inital embedding vector를 “Eiffel”의 embedding vector 방향으로 업데이트가 됩니다.

–> 업데이트 되면서 “Tower”의 initial embedding vector가 high dimensional space 상에서 “Eiffel Tower” 벡터에 가까워지게 됩니다.

만약 문장이 “miniature Eiffel Tower”라면, “miniature”라는 context가 “Tower”에 주는 영향으로써 하나 더 추가되었다고 생각하면 됩니다. 즉, “Tower”의 initial embedding vector에 더해주는 context embedding vector가 하나 더 추가로 존재한다고 생각하시면 됩니다.

이 내용을 Attention Pattern에서 보면, 위와 같은 initial embedding vector에서 더해주는 context에 해당하는 embedding vector들이 여러 개 존재한다고 볼 수 있습니다. (위의 예시에서는 더해주는 벡터가 1개, 2개인 경우만 visualize해서 본 것으로 생각하시면 됩니다.)

이번에는 “miniature Eiffel Tower”보다 긴 문장인 “a fluffy blue creature roamed the verdant forest” 문장에서 Attention을 이해해봅시다.

creature의 initial embedding에 더해지는 임베딩 벡터들은, creature 앞에 해당하는 “a”, “fluffy”, “blue”에 대한 각각의 embedding vector이고,

이후 모든 임베딩 벡터들은 cheating을 방지하기 위해 0의 가중치를 가지도록 attention pattern을 계산했습니다.

즉, creature의 initial embedding vector는 “a”, “fluffy”, “blue”의 임베딩 벡터와의 유사도를 dot product로 계산하고, 그 값의 크기들만큼 high dimensional space 상에서 “creature”의 initial embedding vector에 더해져 “creature”의 initial embedding vector가 “a fluffy blue creature”의 임베딩 벡터 방향으로 업데이트 되게 됩니다.

이는 위의 간단한 miniature Eiffel Tower 예시와 동일하고, 단지 initial embedding vector에 더해지는 임베딩 벡터의 개수가 늘어, 더 많은 context를 고려하는 예시일 뿐이라고 이해할 수 있습니다. (문장이 길면 하나의 initial embedding vector에 더해지는 임베딩 벡터의 개수가 늘어난다는 의미로 해석가능합니다. 문장이 길면 전체 context를 고려하기 위해 더 많은 임베딩 벡터가 더해진다고도 해석할 수 있습니다.)

현재까지 모든 설명은 single head attention의 하나의 column vector(initial embedding vector)에 대한 것이었습니다.

[Single head attention]이 하나의 attention pattern을 가졌다면,

[Multi-headed attention]은 96개의 distinct한 attention pattern이 있는 것입니다.

“a fluffy blue creature roamed the verdant forest” 문장에서 Single head attention과 Multi-headed attention의 차이를 알아봅시다.

“creature”에 대한 1개의 initial embedding vector (Attention Pattern에서 column vector 1개)에 대해서만 고려해봅시다.

[Single head attention]

“creature”의 initial embedding vector와의 “a”, “fluffy”, “blue”의 embedding vector와의 dot product로 구한 가중치(Q, K 연산결과)만큼 “creature”의 initial embedding vector를 “a’, “fluffy”, “blue”의 embedding vector 3개 방향으로 high dimensional space상에서 이동합니다.

–> 업데이트 되면서 high dimensional space 상에서 “creature”의 initial embedding vector가 이에 attend하는 “a”, “fluffy”, “blue”라는 context를 고려하는 “a fluffy blue creature”의 embedding vector에 가까워지게 됩니다.

[Multi head attention]

single head attention이 96개인 것입니다. 즉, Distinct한 Attention Pattern이 96개 있는 것입니다.

“creature”의 initial embedding vector와의 “a”, “fluffy”, “blue”의 embedding vector와의 dot product로 구한 가중치(Q, K 연산결과)만큼 “creature”의 initial embedding vector를 “a”, “fluffy”, “blue”의 embedding vector 3개 x 96개 방향으로 high dimensional space상에서 이동합니다.

–> 업데이트 되면서 high dimensional space 상에서 “creature”의 initial embedding vector가 이에 attend하는 “a”, “fluffy”, “blue”라는 context를 고려하는 “a fluffy blue creature”의 embedding vector에 가까워지게 됩니다.

Embedding Space, Attention Pattern, Softmax, Masking

Query, Key는 먼저 사전의 정의된 Embedding Matrix를 거쳐 initial Embedding Space 상으로 보냅니다. 그리고, 임베딩 벡터들은 initial Embedding Space 상에서 Attention Pattern에서 구하는 임베딩 벡터만큼 더해져 더 많은 context를 아우르는 임베딩 벡터로 이동하게 됩니다.

Query, Key에 대해 dot product를 취하여, dot product 결과의 크기를 회색원의 크기로 표현하면 아래와 같습니다.

Query, Key의 연산 결과를 Attention Pattern이라고 하며, 연산 자체의 결과는 Logits으로 표현됩니다.

우리는 Probability에 대해 다루고 싶으므로 Logits을 softmax 함수를 거쳐 확률값으로 변환해줍니다.

Softmax를 거친 Logits들의 합은 1로 변화합니다.

(이때, softmax의 temperature 값으로 가장 큰 probability의 값의 contribution 정도를 조절할 수 있습니다. t=0이면 제일 큰 softmax 값이 100%가 되고 나머지는 0%가 되도록 scaling이 조절되고, t가 커질수록 제일 큰 softmax 값의 기여도가 줄어들도록 확률값을 조절할 수 있습니다.)

Attention Pattern은 Attention is All you need 페이퍼의 수식과 대응하여 보면 아래와 같습니다.

이때 numerical stability를 위해 d_k를 나누는 작업이 추가되었습니다.

Attention Pattern을 계산할 때, 현재값 이후의 token을 Attend하는 것은 cheating에 해당하므로 이를 방지하기 위해, 해당 token들에 대해서는 masking을 하는 과정이 추가되어야 합니다. 이는 그 부분에 해당하는 값을 -∞로 바꾸고, softmax를 통과시켜 확률값을 0%로 바꾸는 것으로 해당 token부분에 대해서 masking을 수행할 수 있습니다.

결론적으로 Attention을 통해 각 token의 embedding vector는 더 많은 context를 고려한 벡터로 변화하게 됩니다.

Attention 예시 1)

또다른 예시로 “they crashed the car” 문장에서 “car”가 “they”, “crashed”, “the”의 context를 고려해서 바뀌는 것을 관찰할 수 있습니다.

좀 더 긴 문장에서 띄엄띄엄 존재하는 context에 대해서도 Attention의 효과를 관찰할 수 있습니다.

“wizard”, “Hogwarts”, “Hermione”의 context를 반영한 “Harry”

“Queen”, “Sussex”, “William”, “Harry”의 context를 반영한 “Harry”

Attention 예시 2)

“King”의 embedding vector는 context를 표현하는 “lived in Scotland”, “murdererd predecessor”, “in Shakespearean language”의 embedding vector들의 context를 반영하는 embedding vector로 업데이트 됩니다.

1) “King”의 embedding vector와 “lived in Scotland”, “murdered predecssor”, “in Shakespearean language”의 embedding vector들 간의 계산은 “King”을 Query vector로, 나머지는 Key vector로하여 dot product로 벡터간의 유사도를 계산한 Attention Pattern을 구하고,

2) Attention Pattern의 가중치만큼 “lived in Scotland”, “murdered predecessor”, “in Shakespearean language”를 Value_down으로 인코딩한 각각 value embedding vector에 가중치를 주고,

3) “King”의 embedding vector에 더하여, “King”의 embedding vector가 context인 “lived in Scotland”, “murdered predecessor”, “in Shakespearean langauge”를 반영한 embedding vector로 업데이트 됩니다.

Attention 예시 3)

Attention 예시 2)와 동일한 과정으로 이전에 들었던 “a fluffy blue creature roamed the verdant forest”에서 “creature”의 embedding vector에 context로 주어지는 “a”, “fluffy”, “blue”의 embedding vector로 인해 업데이트 되는 과정을 Attention Pattern에서 표현하면 아래와 같습니다.

이 과정을 high dimensional space 상에서 “creature”의 embedding vector에 “fluffy”의 Value embedding vector와 “blue”의 Value embedding vector가 더해져 “creature”의 embedding vector가 “blue fluffy creature”의 embedding vector로 업데이트되는 과정을 시각화 해보겠습니다.

1) “fluffy”의 Value가 “creature”의 embedding vector에 더해지는 과정 시각화

2) “blue”의 Value가 “fluffy creature”의 embedding vector에 더해지는 과정 시각화

시각화를 위해 1), 2)로 나눠 순차적으로 Value가 더해지는 과정을 보았지만, 실제 Attention Pattern에서는 아래와 같이 “creature”의 embedding vector에 “fluffy”, “blue”의 Value가 더해지는 과정이 한번에 이루어집니다.

Value Matrix 연산 과정

Value matrix를 12,288 x 12,288로 사용할 수는 있지만 파라미터 수가 너무 많습니다.

우리는 Value matrix의 파라미터 수를 Query와 Key matrix의 파라미터 수의 합만큼으로 나타내고 싶습니다.

=> 이는 코드 구현상에서 Query, Key, Value_down, Value_up의 파라미터를 같도록 만들어주어, Multi-head attention에서 Query, Key, Value의 head dim을 하나의 hyper parameter로 조절할 수 있게 해줍니다.

따라서 Query, Key matrix가 128 x 12,288 matrix의 파라미터를 가지므로, Value matrix를 128 x 12,288 matrix 2개를 연산한 것으로 쪼갭니다.

이렇게 쪼개진 matrix 2개는 각각 linear map을 수행한다고 생각하시면 됩니다.

앞서 high dimensional space상에 Value matrix를 통해 Value를 구했던 과정에서 Value matrix는 linear map을 수행한다고 생각하면 됩니다.

12,288 dims을 128 dims로 linear map하는 matrix를 Value_down matrix라 해봅시다.

128 dims을 12,288 dims로 linear map하는 matrix를 Value_up matrix라 해봅시다.

결론적으로 여기서 수행하는 Linear map은 “Low rank” transformation과 같습니다.

이로써 Query, Key, Value (Value_up, Value_down)의 파라미터 수가 같아지도록 만들었습니다.

즉, transformer의 multi-headed attention 상에서 Query, Key, Value의 head 수를 128로 동일하게 맞춰줄 수 있게, Value의 파라미터 수를 조절해주었다고 생각해도 됩니다. 즉, Multi-headed attention에서 Query, Key, Value의 head_dim을 하나의 hyper parameter로써 64, 128, 256 등으로 바꿔줄 수 있게 되었습니다.

Value matrix를 linear map(Low rank transformation)으로 취급하여 나눴던 Value_up, Value_down matrix 중에서 Value_up만 따로 행렬로 묶습니다.

우리가 실제 코드 Implementation에서는 Query, Key, Value에서 Value 부분은 Value_down matrix로 코딩이 되어 있고, 이 matrix가 128 x 12,288 의 dimension을 가지게 됩니다.

그리고 Value_up matrix들은 Output Matrix로써 따로 분리되어 있습니다.

ML 사람들은 여러 개의 matrix를 하나의 single matrix multiplication으로 compress하는 것을 좋아하기 때문에, 이처럼 Value_up matrix와 Value_down matrix를 분리하여, Value_up 행렬만 모아서 Output matrix로 만듭니다.