Transformer 정리 (2) : Attention Is All You Need (attention의 의미, transformer는 왜 병렬화가 쉬운지)

Transformer 정리 두 번째 순서로 Transformer 개념을 처음 제안한 논문 [1]을 바탕으로 Transformer 개념에 대하여 정리하고자 한다. Machine translation task에서 RNN 개념을 완전히 제거하고 오로지 attention mechanism(=fully connected layer 개념)만을 사용해서 더 향상된 결과를 보여줄 수 있다는 것을 보이는 논문이다. Fully connected layer의 가능성(all-to-all 분석 가능 & 매우 쉬운 병렬화)을 위주로 분석하고자 한다.

 

 

이전글

2022.05.18 - [Deep Learning/개념] - Transformer 정리 (1) : Attention (Neural machine translation by jointly learning to align and translate)

 

 

 

Attention만을 통한 네트워크 구현 : 기존과의 비교

Fig 1. 기존 Attention mechanism을 도입한 RNN 기반 번역 네트워크(왼쪽)과 attention mechanism만으로 구성된 번역 네트워크(오른쪽) [1][2]

 

Fig 1.은 RNN을 이용한 encoder - decoder 구조(왼쪽)와 transformer를 이용한 encoder - decoder 구조(오른쪽)를 보여주는 그림이다. 

 

왼쪽 RNN-based encoder - decoder 구조는 transformer가 제안되기 이전, machine translation 분야에서 가장 주류로 사용되던 네트워크 구조이다. RNN을 통해서 영어(번역하고자 하는 문장)을 하나의 encoded vector로 만들고, decoder로 encoded vector를 읽어서(=decode) 프랑스어(번역 결과)를 출력한다. 이때 문장을 하나의 벡터로 encode하는데, 문장이 길어질수록 몇몇 단어가 잊혀질수 있을 것이다.

 

이를 해결하기 위해 attention mechanism이 제안되었는데, 왼쪽 그림과 같이 "Je"를 추론할 때, 영어의 모든 단어를 직접 고려할 수 있도록 connection을 만들어주는 것이다. 즉 Fig 2.와 같이 all(영어 단어들) -> one(vector) -> all(프랑스 단어들) 관계가 all(영어 단어들) -> all(프랑스 단어들), fully connected 관계가 된 것이다.

 

Fig 2. 단어들간의 관계

 

Fig 1.의 오른쪽은 transformer-based encoder - decoder 구조를 나타낸 것이다. RNN은 단어의 하나씩 순차적으로 분석해야 하는 단점이 있다. 만약 문장이 10개의 단어를 포함하고 있다면 10개의 단어 하나 하나씩을 순차적으로 분석하여야 한다. 따라서 병렬화가 쉽지 않다. 또한 10번째 단어를 기억하는 과정에서 첫 번째 단어를 어느정도 잊을 가능성도 있을 것이다.

 

Transformer는 이러한 시간적 dependency가 무조건 있을 수 밖에 없는 RNN 구조를 버리고 dense (=fully connected) layer만으로 구성되었다. 즉 모든 단어를 한번에 분석하여 encode하겠다는 의미이다. (다만 단어의 순서도 중요하기에 positional encoding 과정 = 단어에 순번을 부여하는 과정을 추가하였다.)

 

예를 들어 RNN은 1초에 "I", 2초에 "am", 3초에 "a", 4초에 "student"를 분석한다면, transformer는 1초에 "I(0)", "am(1)", "a(2)", "student(3)" 4개 단어를 동시에 분석할 수 있다.

 

이는 matmul로 쉽게 구현할 수 있는데 (C=A*B 연산에서 A=단어*단어길이, B=단어길이*encode길이 => C=단어*encode길이), matmul은 CPU나 GPU와 같은 temporal, ISA-based 하드웨어에서 가장 효율적으로 최적화할 수 있는 연산 중 하나이다. 특히 A의 height와 weight 모두가 적당히 클수록 더 효율적인 연산이 가능해진다 (즉 단어를 동시에 많이 분석하는게 이득이라는 의미).

 

 

 

Multi-head Attention의 특징 = 매우 쉬운 병렬화

Fig 3. Self-attention 기법 구현[1]. Multi-Head Attention을 사용하여 병렬화가 쉬움

 

Transfomer의 한 모듈은 multi-head attention과 그것들을 결과물을 다시 한번 globally 분석하는 dense (=matmul=fc) layer로 구성된다.

 

Attention은 Query, Key, Value의 3가지 요소들 사이의 상관 관계를 확인하는 것으로, query가 들어왔을 때(영향을 받는 단어), 그것에 대한 key(영향을 주는 단어)와의 상관관계를 구하고, 그 결과(영향력)을 Value로 평가한다. Query와 Key와의 상관관계가 Value와 비슷하다면 높은 attention 점수를 얻게되며, 비슷한 의미를 가진 단어들끼리 attention 점수가 높아지도록 학습한다.

 

Eq 1. Attention equation

 

Eq 1.을 보면 알겠지만 attention도 매우 효율적인 matmul로 쉽게 계산할 수 있다.

 

이때 논문[1]은 Multi-Head attention을 도입하는데, 이는 위와 같은 attention을 독립적으로 여러번 하겠다(=병렬로 구현하겠다)는 의미이다. Fig 2.의 오른쪽 그림이 multi-head attention을 잘 나타낸다. n개의 head가 있으며, 각각의 head는 Q, K, V를 구하기 위해 총 3번의 독립적인 linear(=FC=dense) 연산을 한다. 그리고 n개의 head의 output을 하나의 vector로 concatenate 하면 최종 output이 나온다. 미리 각 attention마다 결과값을 저장할 메모리 offset을 구해서 그 위치에 저장만 하면 concatenation도 병렬로 쉽게 구현할 수 있다. 즉 n개의 head는 병렬로 매우 쉽게 처리할 수 있다.

 

Multi-head attention은 8개의 독립적인 attention을 시도하는데, 하나의 multi-head attention 내의 attention끼리는 서로간에 의존성이 없기에 매우 쉽게 병렬화할 수 있다. 논문에서도 8개의 attention 헤드들을 각각 다른 GPU에서 실행시키며 병렬화했을 것이고, 그렇게 쉽지만 매우 효율적으로 계산 속도 향상을 달성했을 것이다.

 

최종 output에 linear 연산을 한번 더 도입하였는데, 이는 n개의 head의 결과들끼리의 연관관계를 분석하는 과정일 것이다. n개의 head의 결과는 독립적으로 계산되었고, 그 결과도 단순 concat만 하기에, 그들 사이의 관계를 분석할 linear 단계가 추가적으로 필요했을 것으로 예상된다.

 

논문[1]에 따르면, multi-head attention을 사용하는 이유는 한번의 큰 attention을 하는 것보다 여러 번의 작은 attention을 하는 게 더 좋기 때문이라고 한다. 즉 병렬화가 쉽다는 거지, 병렬화가 쉬워서 사용한 것이 아닌 것이다. 사실 큰 attention 한번 하는 것은 여러 번의 작은 attention 하는것과 동일한 연산량을 가지며, 오히려 작은 attention을 여러번 할 경우, 추가적인 linear operation이 필요하기에 더 많은 연산량이 필요할 것이다.

 

Fig4. Single Attention과 Multi-head Attention 비교

 

 

 

Reference

[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 6000–6010.

[2] Bahdanau, D., Cho, K. H., & Bengio, Y. (2015). Neural machine translation by jointly learning to align and translate. Paper presented at 3rd International Conference on Learning Representations, ICLR 2015, San Diego, United States.