[논문 리뷰] Attention is All You Need (NIPS, 2017)

Suhyeong·2024년 12월 15일

NLP

목록 보기

1/1

Overview

원문: Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems. 2017.

해당 연구는 Google Brain에서 2017년에 NeurIPS에 발표한 연구이며, 기존의 RNN 계열의 문제인 Long-Term dependency를 해결한 Transformer라 불리는 모델을 제안한다. 이는 BERT와 GPT와 같은 Pre-trained Language Model(PLM) 등등 현재 많은 LLM에 적용이 되고있는 아키텍처이다.

Introduction

RNN 계열의 모델 (LSTM, GRU)은 Language Modeling과 Machine Translation과 같은 Sequence Modeling과 Transduction의 가장 좋은 접근법으로 자리를 잡았다.
image (1) 당시의 최신의 연구들은 sequence의 긴 길이에서 메모리 제약을 없애기 위해 계산 효율성을 높이는 방법을 제시를 하였으나, 기본적인 sequence data의 제약은 많이 남아있다.
image (2) 이러한 문제점에 반해 Attention Mechanism은 입력이나 출력 sequence 내의 위치에 관계 없이 dependency를 모델링할 수 있게 된다. 기존의 RNN에 적용된 Attention Mechanism은 context와 함께 word vector에 softmax를 씌워 attention score를 계산하고 이를 통해 translation을 진행하는 방향으로 이루어졌었다.

이 연구는 위의 Attention Mechanism을 이용하여 각 sequence 내의 위치에 관계 없이 dependency를 모델링할 수 있도록 하여 global dependency를 생성하는데에 있어 Attention Mechanism만을 사용하는 새로운 아키텍처인 Transformer를 제시한다.

이 논문의 주요 Contribution은 다음과 같다.

본 연구에서는 순환 구조를 배제하고 전적으로 주의 메커니즘을 활용하여 입력과 출력 간의 전역 의존성을 파악하는 새로운 모델 아키텍처인 Transformer를 제안한다.

Transformer 모델은 기존 방식보다 훨씬 더 높은 수준의 병렬 처리를 가능하게 함으로써, 단 12시간의 훈련만으로도 번역 품질에서 최신 기술 수준(SOTA, State-Of-The-Art)을 달성한다.

이 아키텍처는 번역 품질을 혁신적으로 향상시킬 뿐만 아니라, 머신 러닝과 자연어 처리 분야에 새로운 방향을 제시한다.

Model Architecture

대부분의 seq2seq 모델들은 en-decoder의 구조를 가진다. encoder는 입력 sequence

(x_{1}, ..., x_{n})

를 연속적인 representation

z = (z_{1}, ..., z_{n})

의 sequence로 mapping한다.

z

를 통해 decoder는 symbol의 출력 sequence

(y_{1}, ..., y_{m})

를 한 번에 하나씩 생성한다. model은 각 step에서 auto-regressive 적이며, 다음 symbol을 생성할 때 이전에 생성된 symbol을 추가 입력으로 사용한다.

Encoder and Decoder Stacks

Encoder: Encoder의 경우 $N = 6$ 의 동일한 layer로 구성 되며, 각 층에는 2개의 sublayer가 있다. 첫 번째는 Multi-head self-attention layer이고, 두번째는 간단한 position-wise fully connected feed-forward network이다.

두 sublayer 모두 residual connection을 사용하며, 이후 layer norm을 진행한다. 이를 통해 알 수 있듯이 각 sublayer의 출력은 $\text{LayerNorm}(x + \text{Sublayer(x)})$ 이다. 이런 residual connection을 용이하게 하기 위하 모델의 모든 sublayer와 embedding layer는 $d_{\text{model}} = 512$ 의 차원을 출력한다.

Decoder: Deocoder 역시 Encoder와 동일하게 구성되며, Encoder와 다르게 self-attention layer를 수정하여 현재 예측하는 위치가 다음에 올 위치에 attention을 진행하지 않도록 masking을 진행하며, 위치 $i$ 에 대한 예측이 $i$ 보다 작은 위치 즉, 이 전의 정보에만 attention을 진행할 수 있도록 한다.

Attention

Attention은 query와 key-value 집합을 출력으로 mapping할 수 있다. 여기서 query, key, value는 모두 벡터이며, 출력 값들은 전부 weighted sum 연산이 진행된다.

query, key value가 무엇인지는 아래에 간단히 정리해 두었다.

query: 현재 단어와 다른 단어들 사이의 관계를 평가하는 지표 vector

key: 비교 대상이 되는 다른 단어들의 속성 vector

value: query와 key의 매칭 정도에 따라 가중치가 부여된 정보 vector

Scaled Dot-Product Attention

image (4)

Transformer의 attention은 scaled dot-product attention이라 불리는 attention mecahnism을 진행한다.

입력은 $d_k$ dimension의 queries와 keys, $d_v$ dimension의 values로 구성되며, queries와 모든 keys의 dot-product를 계산하고, 이를 $\sqrt{d_k}$ 로 scaling을 진행한다. 이후 softmax를 적용하여 value에 대한 attention weight를 얻게 된다. 이를 수식으로 표현하면 다음과 같다.

\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^{T}} {\sqrt{d_{k}}})V \quad\quad\quad\quad(1)

이때 까지 가장 일반적으로 사용되는 attention mechanism은 additive attention과 dot-product attention이 있다. dot-product의 경우 해당 연구와 비슷 하지만 Transformer의 경우 scaling을 진행하기 때문에 다르다는 것을 알 수 있다.

Transformer의 경우 큰 $d_k$ 값에 대해서 dor-product attention의 크기가 커져서 softmax 함수의 출력이 매우 작은 확률 분포로 극단적으로 치우쳐 몇몇 항목은 매우 높은 가중치를 받고 대부분의 항목들이 거의 0에 가까운 가중치를 받게 되기 때문에 scaling을 진행해준다.

Multi-Head Attention

$d_{\text{model}}$ -dimension의 keys, values 그리고 queries를 가진 단일 attention mechanism을 적용하는 대신, 해당 연구에서는 queries, keys, values를 $d_{k}, d_{k}, d_{v}$ dimension으로 서로 다르게 학습된 $h$ 번의 linear projection을 진행하는 것이 더 좋은 성능을 낼 수 있다고 말하고 있다.

이렇게 projection된 queries, keys, values에 대하여 attention function을 병렬로 수행하여 $d_v$ -dimension의 출력 값을 얻게 된다. 이 출력 값들은 concatenate 되고, 다시 projection되여 최종 값으로 Figure 2와 같이 나타난다.

Multi-Head Attention은 모델이 다른 representation space의 다른 위치에서 정보에 공동으로 주목할 수 있게 해준다. 그러나 단일 attention의 경우 평균화가 이를 방해하게된다. 위에서 설명했던 Multi-Head Attention은 아래와 같이 표현될 수 있다.

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^{O}\\ \quad\quad\quad\text{where} \quad \text{head}_i = \text{Attention}(QW_{i}^{Q}, KW_{i}^{K}, VW_{i}^{V})

Position-wise Feed-Forward Networks

encoder와 decoder의 각 layer는 FFn (Feed-Forward networks)를 가지고 있으며 ReLU 활성화 함수를 가지고 있으며 총 2개의 linear layer를 가지고 있다.

\text{FFN}(x) = \text{max}(0, xW_{1} + b_{1})W_{2} + b_{2}\quad\quad\quad\quad(2)

FFn 내부에서 입력의 차원은 $d_{\text{model}} = 512$ 이고, $d_{\text{ff}} = 2048$ 까지 dimension을 증가시켰다가 다시 $d_{\text{model}} = 512$ 으로 감소 시키며 좀 더 세부적인 representation을 얻을 수 있도록 연산되어진다.

Positional Encoding

Transformer의 경우 문장의 단어들이 어느 위치에 오는지 알 수 없다. 이를 해결하기 위해 positional encoding을 도입한다. model에 순서 정보를 입력하기 위해 각 위치에 대한 사인과 코사인 함수 기반 encoding을 input embedding에 추가한다. $pos$ 와 dimension $i$ 에 따라 다음과 같이 계산된다.

P E_{(pos, 2i)} = sin(pos / 10000^{2i / d_{\text{model}}})

PE_{(pos, 2i)} = cos(pos / 10000^{2i / d_{\text{model}}})

Why Self-Attention

image (5)

self-attention은 모든 입력 위치를 일정한 연산 수로 연결하는 반면, RNN은 sequence 길이에 비례하여 연산수가 증가한다. 이러한 특성 덕분에 self-attention은 Complexity와 Maximum Path를 줄여주어, long-term dependecies를 효과적으로 학습한다. CNN과 비교했을 때도 self-attention은 더 적은 연산을 사용하며, 더 짧은 path를 제공한다. 언급한 자세한 수치는 Table 1에서 확인할 수 있다.

Training

Training Data and Batching

WMT 2014 English-German dataset에서 약 450만의 문장 pair를 훈련하였으며, source-target vocab은 약 37000개 token으로 구성되었다. 또한 English-French에서는 더 큰 WMT 2014 dataset을 사용하였으며, 3600만 문장과 32000 word-piece vocab으로 token을 나누었다.

Hardware and Schedule

8개의 NVIDIA P100 GPU를 탑재한 한 대의 기계에서 모델을 훈련했다. 논문에 걸쳐 설명된 hyper-parameter를 사용하여, 각 훈련 step에서는 약 0.4초가 소요되었다. base model은 총 100,000단계 또는 12시간 동안 훈련했다. large model은 300,000단계(3.5일) 동안 훈련했다.

Optimizer

Adam optimizer를 사용하였으며, $\beta_{1} = 0.9$ , $beta_{2} = 0.98$ 그리고 $\epsilon = 10^{-9}$ 를 설정하였으며, learning rate는 다음 공식에 따라 변화된다.

lrate = d_{\text{model}}^{-0.5} \cdot \text{min}(step\_num^{-0.5}, (step\_num \cdot warmup\_steps^{-1.5})\quad\quad\quad(3)

이는 Noam Scheduler라고 불리며, 처음 $warmup\_ steps$ 학습 단계 동안 학습률을 선형적으로 증가 시키고, 이후 step에 역제곱근에 비례하여 감소 시키는 것을 의미하며, 해당 연구에서는 $warmup\_ steps$ 는 4000으로 설정하였다.

import torch
from torch.optim.lr_scheduler import _LRScheduler

class NoamLR(_LRScheduler):
    """The LR scheduler proposed by Noam

    This is simplified to use only torch components, removing external dependencies.

    Ref:
        "Attention Is All You Need", https://arxiv.org/pdf/1706.03762.pdf
    """

    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        model_size: int = 320,
        warmup_steps: int = 25000,
        last_epoch: int = -1,
    ):
        self.model_size = model_size
        self.warmup_steps = warmup_steps
        super(NoamLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        """Calculate learning rate using Noam scheme."""
        step_num = self.last_epoch + 1
        return [
            lr * self.model_size**(-0.5) * 
            min(step_num**(-0.5), step_num * self.warmup_steps**(-1.5))
            for lr in self.base_lrs
        ]

    def __repr__(self):
        return f"{self.__class__.__name__}(model_size={self.model_size}, warmup_steps={self.warmup_steps})"

위의 코드는 Transformer의 Noam Schduler를 Pytorch를 이용하여 구현한 코드이다.

Results

Machine Translation

image (6)

WMT 2014 English-to-German translation task에서 Transformer big 모델이 다른 Ensemble model들을 포함한 다른 모델들 보다 BLEU score가 2.0 이상 높았으며, 28.4라는 SOTA (State-Of-The-Art)를 달성하였다.

또한 WMT 2014 English-to-French에서도 단일 모델들에서 모두 앞서는 성능을 냈으며, BLEU score 41.0을 달성하였다.

Model Variations

image (7)

각각의 Transformer의 설정을 바꿔가며 English-to-German dataset의 성능 변화를 보는 것으로 $d_k$ 를 줄이면 model의 품질이 떨어지는 것을 확인할 수 있다.

Conclusion

이 연구에서는 Attention 기반 model들의 미래에 매우 기대하고 있으며, 다른 task에도 이를 적용할 계획이며, 텍스트 이외의 입출력 형식, 예를 들어 이미지, 오디오 및 비디오를 포함하는 문제에 Transformer를 확장할 예정이라고 밝혔다.

여기까지 요즘 매우 뜨고있는 LLM과 이미지 쪽에서도 현재 많이 쓰이는 ViT, Swin Transformer와 같은 많은 모델들의 토대로 쓰이는 Transformer 리뷰를 마치도록 하겠다.

Suhyeong

딥린이