[논문 리뷰] Attention is All You Need

 출처 : https://arxiv.org/abs/1706.03762

Attention Is All You Need

 구글 연구팀에서 2017년 발표한 논문입니다. 구글의 기존 Seq2Seq 모델에 대응한 '페이스북의 ConvS2S'가 발표된 지 얼마되지 않아 Transformer라는 기가 막힌 아키텍처를 선보인 것이 바로 이 논문입니다. 사실 너무 많이 읽어서 적당히 쓰려고 했는데, 제가 좋아하는 논문이기도 하고 Transformer 자체는 현재도 SOTA라고 할만한 아키텍처기 때문에 꽤 공들여 써보려고 합니다.

 

 1. Introduction

 구글은 2014년에 Seq2Seq 아키텍처를 발표하였습니다. 다음 링크에서 자세한 원본을 확인해보실 수 있습니다. https://arxiv.org/abs/1409.3215

Sequence to Sequence Learning with Neural Networks

 물론 Seq2Seq의 성능은 뛰어납니다. 그러나 Seq2Seq 모델은 분명한 단점이 존재합니다. 첫 번째, 훈련 시간이 오래 걸립니다. 이는 이전 타임스탭의 출력을 입력으로 받아야하는 (디코더의) 특성상 병렬처리가 안 되기 때문입니다. 논문에는 나오지 않으나, 두 번째는 Teacher Forcing, 즉 추론 결과와 정답 간의 괴리입니다. 사실 디코더의 특정 타임스탭은 다음 타임스탭에 그대로 입력으로 주어지지 않습니다. 기본적으로는 이전 타임스탭의 정답을 입력으로 주고, 성능 향상을 위해 추론 결과를 합치는 정도이죠. 

 결국 이 모든 단점은 Seq2Seq의 recurrent한 특성 때문에 발생합니다. 따라서 구글은 여기에서 기존의 Recurrent 아키텍처를 버리고, 전적으로 attention에만 의존한 모델을 선보이려 합니다. 

 

3. Model Architecture

 모델의 기본적인 목표는 기존과 동일합니다. 어떤 문장(symbol representation)을 입력으로 받아, 실수 좌표상 벡터(contiuous representation)로 인코딩하고, 디코더는 이를 다시 타켓 언어의 문장으로 펼칩니다. 그러나 구조는 기존의 것과는 사뭇 다른 것이 보입니다.

 인코더와 디코더는 각각의 블록을 스택한 구조입니다. 논문에서는 둘 다 6개 스택했다고 나와있습니다. 이전 GNMT 논문에서도 모델이 작은 변화를 탐색할만큼 깊은 구조를 쌓아야 한다고 했고, 그래디언트 소멸을 방지하기 위해서 Residual connection을 사용했었는데, 이번에도 이는 같음을 알 수 있습니다. 다만 내부는 오직 Attention과 FC layer로만 이루어져 있습니다.

 또한 512차원의 hidden state를 출력하도록 세팅했으며, illegal connection을 방지하기 위한 Masked Multi-Head Attention도 눈에 띕니다. 자세한 내용은 차차 살펴보도록 하겠습니다.

 

3.2 Attention

 Attention은 기본적으로 Key에서 Query하여 적절한 Value를 가져오는 과정입니다. 편의상 Query를 Q, Key를 K, Value를 V라고 하겠습니다. GNMT에서 확인했듯이 Q는 디코더의 가장 하단 레이어의 linear transformation입니다. 이렇게 생성된 질의와 K, 즉 인코더의 모든 타임스탭 간 유사도를 얻습니다. V는 사실상 K와 같으며, Q와 K의 유사도에 softmax를 취한 Attention weight와 곱해져 유용한 정보가 됩니다.

 결국 Attention을 학습한다는 것은 적절한 Q를 생성할 수 있는 linear formation의 가중치를 학습한다는 말입니다. 더 효율적인 Q를 만들 수 있다면 더 좋은 성능을 낼 수 있다는 것이죠. Transformer는 이를 위해서 Multi-Head Attention이라는 구조를 사용합니다. 간단히 말하면 질의를 여러 개 생성하겠다는 거죠. 이 개수를 h(head)라고 하겠습니다. h만큼 Q, K, V가 head 차원으로 축소 및 분할되고, 이는 각각의 Scaled Dot-Product Attention을 거쳐 정보를 출력합니다. 출력된 h개의 정보들은 Concatenation되고, 출력 차원으로 확장됩니다.. 어렵게 말한듯 하지만, 결국 여러 번 질문하여 마지막엔 정보를 하나로 합친 것입니다. 논문에서는 8개의 head를 사용했다고 합니다.

Multi-Head Attention

 내부의 Scaled Dot-Product Attention은 간단합니다. 내적을 이용하여 유사도를 구하는 Attention을 Dot-Product Attention이라고 할 뿐입니다. 이는 마치 코사인 유사도를 구하는 것과 비슷합니다. 단 유사도를 그냥 softmax를 통과시키는 것이 아니라, head의 차원으로 나누어줍니다. 예를 들어 hidden state가 512차원이고, head가 8개였다고 한다면 512/8 = 64의 제곱근으로 유사도를 나누어줍니다. 이는 Q와 K의 내적이 head의 차원에 비례하여 커지기 때문에, 안정적인 학습을 위해 스케일링해주는 것입니다.

Dot-Product Attention

 Transformer 모델은 이 Attention-based Architecture를 통해 인코딩과 디코딩까지 해내고 있습니다. 특정 단어의 의미를 파악할 때 다른 위치에 있는 토큰들까지 확인하는 방식인데, 이를 Self-Attention이라고 합니다. 다만 디코더를 확인해보면 Masked Multi-Head Attention이라는 것이 있는데, 이는 디코더의 auto-regressive property 때문입니다. 디코더에 인코더와 같은 Self-Attention을 적용하게 되면 미래를 참고하는 일이 발생할 수 있습니다. 예를 들어 모델은 2번째 단어를 생성 중인데, 5번째 나올 단어를 이미 알고 있다면 이치와 맞지 않습니다. 이것을 Illegal Connection이라고 합니다. 이를 방지하게 위해서 Attention에서 softmax를 통과하기 전, 현재 타임스탭 이후의 자리에 -∞로 마스킹해줍니다. 이 경우 softmax를 통과하면 자연스럽게 값은 0이 되어, 어떠한 가중치도 없게 됩니다.

 

 각 블록에서 최종 출력은 FFN을 통과합니다. 활성화 함수로 ReLU를 사용하고 있고, 두 레이어를 거친다는 것을 하단의 식에서 알 수 있습니다.

3.5 Positional Encoding

 

 Seq2Seq 기반 인코더-디코더 구조가 태생적으로 순차성을 보장하는 것과는 달리, Transformer는 Self-Attention을 배경으로 하기 때문에 어순에 대한 정보가 없습니다. 따라서 각 블록은 심볼 입력을 실수 차원으로 변환한 뒤, 즉 임베딩한 뒤 Positional Encoding이라는 과정을 거칩니다. 이는 따로 학습하는 구조가 아니며, 특정한 식에 의해서 깊이별로, 그리고 타임스탭별로 다르게 책정된 상수를 더해주는 것 뿐입니다. 여기서 i는 차원이며, pos는 타임스탭입니다. 그리고 홀수 차원과 짝수 차원에 대해서 각각 사인함수와 코사인함수를 사용하여 구분하였습니다.

Positional Encoding

 역시 구글인지라 Positional Encoding을 학습하는 방식도 사용해봤지만, Sinusoidal한 방식이 더 잘 먹혀서 이렇게 했다고 합니다.

 

4. Why Self-Attention

 왜 Seq2Seq이나 ConvS2S가 아닌 Self-Attention을 사용해야 하는지에 대해서 설명한 대목입니다. 세 가지 측면에서 접근하고 있는데 하나는 계산 복잡도(Complexity per Layer)입니다. n이 문장의 길이, d가 모델의 차원이라고 할 경우 대부분의 경우 n < d이므로 Recurrent나 Convolutional Layer에 비해 Self-Attention이 가장 낮은 복잡도를 보임을 알 수 있습니다.

 두 번째는 순차 연산(Sequential Operation)입니다. 반대로 말하면 병렬 가능성이라고도 할 수 있겠습니다. Recurrent가 문장에 대해 n번의 연산을 수행해야 하는 반면 Self-Attention은 illegal connection만 방지해준다면 레이어 단위의 병렬 처리가 가능합니다. Convolutional Layer와는 동점이네요.

 마지막은 최장경로입니다. 즉 얼마나 장거리 의존에 대해 잘 파악하는가인데, 한 토큰의 해석에 다른 모든 토큰을 확인하는 Self-Attention의 특성상 최장 거리는 1이며 가장 장거리 의존을 잘 파악한다고 볼 수 있겠습니다.

 이외에 Attention weight를 확인하여 결과에 대한 해석이 가능하다는 부가적인 장점도 존재합니다! (사실 딥러닝의 가장 큰 단점이 블랙박스라고 생각하는 입장에서, 이게 핵심이 아닐까라고도 생각합니다.)

 

5. Training

 

 8대의 P100 GPU를 사용하여 3.5일 동안 훈련했다고 합니다. Adam을 사용하되 파라미터는 일명 국룰(?)인 (0.9, 0.98)을 사용한 모습이네요. 출력층에는 Dropout을 0.1의 확률을 주어 달았습니다. 또한 단어에 대해서 확률을 0.1씩 더해주는 스무딩을 해주었습니다.

 base model같은 경우는 영-프 번역에서 성적이 밀리는 모습을 보이는 듯 싶으나, 주목할 만한 것은 Training Cost 대비 성능 향상입니다. 특히 Big model의 경우에는 모든 점수에서 이전까지의 모델을 이겼음에도 불구하고 거의 최하 수준의 리소스를 소비하고 있습니다. Base model은 더더욱 적은 리소스를 사용하고, 그만큼 성능을 포기한 거라고 보면 되겠죠. 

 

정리 

 Transformer는 현재 BERT, GPT 시리즈의 근본 아키텍처는 물론, 자연어처리가 아닌 다른 영역에서도 활발하게 사용되고 있습니다. 이것은 분명 뛰어난 발명입니다. 모델을 만들 때는 Inductive Bias가 영향을 끼치기 마련인데, 그럼에도 불구하고 RNN-based model이 아닌 Attention을 기반으로 번역기가 탄생할 수 있다는 것은 놀라운 일입니다.

 

 그러나 좋은 말만 써있는 논문의 특성상 문제점들도 분명 고려해봐야 할 것입니다. 첫 번째로는 메모리 문제입니다. Transformer의 학습 시간은 분명 체감이 될 정도로 빠릅니다. 그러나 메모리 사용량도 체감이 될 정도로 큽니다. 이는 병렬처리를 위한 큰 공간이 필요한 것, 그리고 Positional Encoding을 위한 저장 공간이 필요한 것 등 다양한 DP 요소가 영향을 미칩니다. 실제로 코랩 환경에서 배치 사이즈를 128 정도로만 가져가도 다운되버리는 현상이 발생했습니다. 충분한 크기의 작업공간이 있지 않는한 배치 사이즈는 줄어들게 될 것이고, 이는 gradient accumulation 등을 사용한다고 하더라도 결과적으로 훈련 시간을 다시 길게 만들 것입니다. 그래서 최근에는 Transformer 기반 인코더와 Sequential한 디코더를 방향도 제시되고 있다고 합니다. 

 

 두 번째는 훈련의 문제입니다. 논문에 제시된 아키텍처를 그대로 구현하다보면 훈련이 제대로 되지 않는다는 것을 느끼게 됩니다. 이는 Transformer가 초기 학습한 sample이 많이 없을 무렵, 깊은 Layer로 인해서 기울기가 불안정하기 때문입니다. 따라서 Adam과 같은 Adaptive Learning rate 방법론은 잘못된 방향으로 적응하게 됩니다. 이것을 해결하기 위해서 Learning rate decay, RAdam 등 다양한 방법론이 제시되었으나, 저같은 경우 pre-LN을 적용했을 때 가장 좋은 효과를 보았습니다. (참고로 TabNet도 비슷한 현상이 있었으나 이는 RAdam으로 해결되었습니다.)

 Transformer 모델의 훈련에 관련된 논문은 다음에 기회가 있다면 리뷰하도록 하겠습니다.

 

아래는 관련 소스코드를 구현해 놓은 레포지토리입니다.

링크: https://github.com/jeongchanyoung-1234/NMT

jeongchanyoung-1234/NMT