[논문 리뷰] Attention Is All You Need

RNN 없이 시퀀스를 처리할 수 있을까

2017년까지만 해도 번역이든 요약이든 시퀀스를 다루는 모델은 거의 다 RNN 계열이었다. LSTM이든 GRU든 결국 t번째 단어를 처리하려면 t-1번째까지의 계산이 끝나야 한다는 게 핵심 전제였고, 이 순차성 때문에 병렬화가 안 됐다. 문장이 길어지면 메모리 제약 때문에 배치 크기도 줄여야 했고, 학습 속도는 더 느려졌다.

Attention은 이미 RNN과 결합된 형태로 쓰이고 있었다. Bahdanau attention처럼 인코더-디코더 사이에 attention을 끼워 넣어서 멀리 떨어진 단어 사이의 의존성도 잘 잡아내는 식이었다. 근데 이 논문 저자들은 한 발 더 나갔다. "그럼 RNN은 왜 필요한 거지? Attention만으로 인코더와 디코더를 다 만들면 안 되나?"라는 질문을 던진 거다. 그렇게 나온 게 Transformer고, recurrence와 convolution을 아예 빼버리고 attention만으로 시퀀스를 처리한다.

결과는 꽤 인상적이었다. WMT 2014 영어-독일어 번역에서 28.4 BLEU로 기존 최고 기록(앙상블 포함)을 2점 이상 앞섰고, 영어-프랑스어에서는 41.8 BLEU로 단일 모델 기준 새로운 state-of-the-art를 찍었다. 그것도 8개 GPU로 3.5일 학습한 결과였다. 이전 최고 모델들이 훨씬 오래 걸렸던 걸 생각하면 학습 비용 대비 효율이 압도적이었다.

전체 구조 - 인코더-디코더는 그대로, 내부만 다 바뀌었다

Transformer도 큰 틀에서는 인코더가 입력 시퀀스를 받아 연속적인 표현으로 바꾸고, 디코더가 그걸 받아서 출력을 한 토큰씩 생성하는 인코더-디코더 구조를 따른다. 다만 그 안에서 순환 구조 대신 self-attention과 position-wise feed-forward network를 쌓아 올렸다는 게 다르다.

Figure 1. Transformer 전체 아키텍처 (인코더-디코더 구조)

왼쪽이 인코더, 오른쪽이 디코더다. 둘 다 동일한 레이어를 N=6번 쌓은 구조인데, 인코더 레이어는 multi-head self-attention과 feed-forward network 두 개의 서브레이어로 이루어져 있고, 디코더는 여기에 인코더의 출력을 보는 attention 레이어가 하나 더 끼어 있다. 각 서브레이어마다 residual connection을 걸고 그 뒤에 layer normalization을 적용한다. 즉 서브레이어의 출력은 항상 LayerNorm(x + Sublayer(x)) 형태다. 이 residual connection을 살리려고 모든 서브레이어와 임베딩 레이어의 출력 차원을 d_model=512로 통일했다.

디코더 쪽에서 한 가지 더 신경 쓴 게 있다. 디코더는 출력을 순차적으로 생성해야 하는데, self-attention은 기본적으로 모든 위치를 다 보기 때문에 자기보다 뒤에 나올 단어까지 미리 봐버리면 안 된다. 그래서 masking을 걸어서 i번째 위치는 자기 자신을 포함해서 그 이전 위치까지만 attention을 줄 수 있게 막는다. 여기에 출력 임베딩을 한 칸씩 밀어주는 것까지 더해서, i번째 위치의 예측이 오직 i보다 이전의 알려진 출력에만 의존하도록 만든 거다.

Attention의 핵심 - Query, Key, Value

Attention 함수를 한 문장으로 정의하면, 쿼리(query)와 키-값(key-value) 쌍의 집합을 받아서 출력을 만드는 매핑이다. 출력은 value들의 가중합인데, 이때 각 value에 곱해지는 가중치는 query와 그에 대응하는 key 사이의 호환성(compatibility)으로 정해진다.

비유하자면 일종의 검색이다. query는 "지금 나는 무엇을 찾고 있는가"이고, key는 "각 위치는 어떤 정보를 갖고 있는가"를 나타내는 인덱스이고, value는 실제로 가져올 내용물이다. query와 key가 잘 맞아떨어질수록(내적값이 클수록) 그 위치의 value를 더 많이 가져오는 식이다. 파이썬 딕셔너리와 다른 점은, 딕셔너리는 key가 정확히 일치해야 값을 가져오지만 attention은 유사도만큼 비례해서 여러 value를 섞어 가져온다는 점이다.

Scaled Dot-Product Attention

논문에서 쓴 attention 함수는 이렇게 정의된다.

Attention(Q, K, V) = softmax( QK^T / √d_k ) V

Q, K, V는 각각 query, key, value를 행렬로 쌓은 것이다. Q와 K를 내적해서 query와 각 key 사이의 유사도(attention score)를 구하고, 이걸 √d_k로 나눠준 다음 softmax를 씌워서 합이 1이 되는 가중치로 만든다. 이 가중치를 V에 곱해서 최종 출력을 얻는다.

여기서 흥미로운 부분은 √d_k로 나누는 스케일링이다. 흔히 쓰이는 attention 방식은 두 가지인데, 하나는 feed-forward network로 호환성을 계산하는 additive attention이고 다른 하나는 내적을 쓰는 dot-product attention이다. 이론적 복잡도는 둘이 비슷하지만, dot-product attention은 행렬곱 연산으로 최적화하기 좋아서 실제로는 훨씬 빠르고 메모리 효율적이다. 문제는 d_k가 커질수록 내적값의 분산도 커진다는 점이다. q와 k의 각 성분이 평균 0, 분산 1인 독립 확률변수라고 가정하면, 내적 q·k는 평균 0, 분산 d_k를 갖는다. 즉 d_k가 클수록 내적값이 커지고, 이게 softmax를 그래디언트가 거의 0인 영역으로 밀어넣어 버린다. 그래서 단순 dot-product attention은 d_k가 큰 경우 additive attention보다 오히려 성능이 떨어졌다. 저자들은 이 문제를 해결하려고 내적값을 √d_k로 나눠서 분산을 다시 1 근처로 맞췄다.

Figure 2. Scaled Dot-Product Attention (좌)과 Multi-Head Attention (우)

Multi-Head Attention

d_model 차원의 Q, K, V로 attention을 한 번만 수행하는 대신, 논문은 Q, K, V를 각각 다른 학습된 linear projection으로 h번 투영한 다음 각각에 대해 attention을 병렬로 수행하는 방식을 택했다. 이렇게 나온 h개의 출력을 이어붙이고(concat) 다시 한 번 linear projection을 거쳐 최종 출력을 만든다.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)

왜 굳이 head를 여러 개로 쪼갰을까. 하나의 attention만 쓰면 평균을 내는 과정에서 서로 다른 표현 부분공간(representation subspace)의 정보가 뭉개진다. head를 여러 개 두면 각 head가 서로 다른 관점에서 정보를 뽑아낼 수 있다. 어떤 head는 인접한 단어 사이의 국소적 관계를 보고, 어떤 head는 문장 전체에 걸친 장거리 의존성을 보는 식으로 역할이 나뉜다. 같은 문서를 여러 독자가 각자의 관점으로 동시에 읽는 것과 비슷하다.

Transformer 안에서 Attention이 쓰이는 세 가지 방식

Multi-head attention은 모델 안에서 역할이 조금씩 다른 세 군데에 쓰인다.

첫 번째는 인코더-디코더 attention이다. query는 디코더의 이전 레이어에서 오고, key와 value는 인코더의 최종 출력에서 온다. 디코더의 모든 위치가 입력 시퀀스의 모든 위치를 참고할 수 있게 해주는 부분으로, 기존 seq2seq 모델의 attention 메커니즘과 본질적으로 같은 역할이다.

두 번째는 인코더의 self-attention이다. 여기서는 query, key, value가 모두 같은 곳, 즉 인코더 이전 레이어의 출력에서 나온다. 인코더의 각 위치는 이전 레이어의 모든 위치를 참고할 수 있다.

세 번째는 디코더의 self-attention이다. 인코더와 비슷하게 디코더의 각 위치가 디코더 내의 모든 위치를 참고할 수 있게 해주는데, 다만 auto-regressive한 특성을 지키기 위해 자기 자신을 포함해 그 이전 위치까지만 보도록 막아야 한다. 구현 방식은 간단한데, scaled dot-product attention 안에서 허용되지 않는 연결에 해당하는 softmax 입력값을 -∞로 설정해서 softmax를 거치면 그 위치의 가중치가 0이 되게 만드는 식이다.

Position-wise Feed-Forward Network

attention 서브레이어 다음에는 각 위치마다 독립적으로 동일하게 적용되는 fully connected feed-forward network가 붙는다. ReLU를 사이에 둔 두 번의 선형 변환이다.

FFN(x) = max(0, xW_1 + b_1) W_2 + b_2

같은 레이어 안에서는 위치마다 동일한 파라미터를 쓰지만, 레이어가 바뀌면 다른 파라미터를 쓴다. 커널 크기 1짜리 convolution을 두 번 적용하는 것과 같다고 볼 수도 있다. 입력과 출력의 차원은 d_model=512이고, 중간 레이어의 차원은 d_ff=2048로 더 넓다.

순서 정보는 어떻게 넣을까 - Positional Encoding

Transformer에는 recurrence도 convolution도 없다. 그러다 보니 같은 단어 집합이라도 순서가 뒤바뀌면 self-attention 입장에서는 사실상 구분이 안 된다. attention 연산 자체가 순서에 무관(permutation-invariant)하기 때문이다. 그래서 입력 임베딩에 위치 정보를 더해주는 positional encoding이 필요하다.

저자들은 서로 다른 주파수를 가진 sin, cos 함수를 사용했다.

PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

pos는 시퀀스 내 위치, i는 임베딩 차원의 인덱스다. 차원마다 서로 다른 파장의 sin/cos 곡선이 할당되는데, 파장은 2π부터 10000·2π까지 기하급수적으로 늘어난다. 이 함수를 고른 이유는 임의의 고정된 offset k에 대해 PE(pos+k)가 PE(pos)의 선형 함수로 표현될 수 있어서, 모델이 상대적 위치 관계를 더 쉽게 학습할 수 있을 거라고 본 것이다.

차원 인덱스가 낮을수록(빨간색) 파장이 짧아서 빠르게 진동하고, 차원 인덱스가 높을수록(파란색) 파장이 길어서 천천히 변한다. 이 여러 주파수의 조합이 합쳐지면서 각 위치마다 고유한 패턴이 만들어진다. 학습 가능한 positional embedding을 대신 써봤을 때도 거의 같은 성능이 나왔는데(뒤에 나오는 ablation 실험의 (E) 행 참고), sin/cos 방식을 최종으로 택한 이유는 학습 때 본 시퀀스보다 더 긴 시퀀스에도 외삽(extrapolate)할 수 있을 거라는 기대 때문이었다.

왜 Self-Attention인가

저자들은 self-attention을 recurrent 레이어, convolutional 레이어와 비교하면서 세 가지 기준을 세웠다. 레이어당 전체 연산 복잡도, 병렬화 가능한 연산량(순차적으로 필요한 최소 연산 수), 그리고 장거리 의존성 사이의 경로 길이다. 마지막 기준이 특히 중요한데, 입력과 출력의 어느 두 위치 사이든 신호가 오가는 경로가 짧을수록 먼 거리의 의존성을 학습하기 쉬워지기 때문이다.

Table 1. 레이어 타입별 연산 복잡도, 병렬화 가능성, 최대 경로 길이 비교

표를 보면 self-attention은 모든 위치를 단 한 번의 연산으로 연결하기 때문에 최대 경로 길이가 O(1)이다. 반면 recurrent 레이어는 순차적으로 처리해야 하니 O(n)만큼의 연산이 필요하고 경로 길이도 O(n)이다. convolution은 커널 크기가 k일 때 한 번에 인접한 영역만 연결할 수 있어서, 모든 위치 쌍을 연결하려면 O(n/k)개의 레이어를 쌓거나 dilated convolution을 쓰면 O(log_k(n))이 필요하다.

연산 복잡도 측면에서는 self-attention이 O(n²·d)이고 recurrent가 O(n·d²)인데, 시퀀스 길이 n이 표현 차원 d보다 작은 경우(기계번역에서 흔히 쓰는 word-piece, byte-pair 표현이 그렇다)에는 self-attention이 더 빠르다. 다만 n이 아주 길어지는 경우를 대비해서, 출력 위치 주변의 이웃 r개만 보는 restricted self-attention도 언급하는데 이건 후속 연구로 남겨뒀다.

부가적인 이점도 있다. self-attention은 어텐션 분포를 직접 들여다볼 수 있어서 모델 해석이 더 쉬워진다. 논문 부록에 실린 attention 시각화를 보면, 개별 head가 단순히 무작위로 흩어진 패턴이 아니라 문장의 통사적·의미적 구조와 관련된 행동을 학습한 걸 확인할 수 있다. 예를 들어 동사 'making'과 멀리 떨어진 'more difficult'를 잇는 장거리 의존성을 특정 head가 정확히 따라가거나, its 같은 대명사가 가리키는 대상을 찾는 anaphora resolution에 특화된 head가 따로 나타나기도 한다.

학습은 어떻게 시켰나

학습 데이터는 WMT 2014 영어-독일어 약 450만 문장쌍을 byte-pair encoding으로 인코딩해서 3.7만 개 토큰의 공유 vocabulary를 썼고, 영어-프랑스어는 3,600만 문장으로 더 큰 데이터셋에 3.2만 word-piece vocabulary를 썼다. 배치는 길이가 비슷한 문장끼리 묶어서 한 배치에 소스/타겟 각각 약 2.5만 토큰이 들어가도록 구성했다.

학습은 8개의 P100 GPU 한 대로 진행했다. base 모델은 step당 약 0.4초가 걸려서 10만 step, 12시간 만에 학습이 끝났다. big 모델은 step당 1초가 걸렸고 30만 step, 3.5일을 학습했다.

옵티마이저는 Adam을 썼는데(β1=0.9, β2=0.98, ε=10⁻⁹), 학습률 스케줄이 특이하다.

lrate = d_model^(-0.5) · min(step_num^(-0.5), step_num · warmup_steps^(-1.5))

처음 warmup_steps(4000)동안은 학습률을 선형으로 증가시키다가, 그 이후로는 step 번호의 역제곱근에 비례해서 서서히 감소시킨다. 학습 초반에 너무 큰 학습률로 불안정해지는 걸 막으면서도, 이후에는 충분히 빠르게 수렴하도록 설계한 스케줄이다.

정규화는 두 가지를 썼다. 하나는 residual dropout인데, 각 서브레이어의 출력이 residual로 더해지기 전과 정규화되기 전에 dropout을 적용했고, 임베딩과 positional encoding을 더한 값에도 dropout을 걸었다(base 모델 기준 P_drop=0.1). 다른 하나는 label smoothing(ε_ls=0.1)으로, 정답 레이블에 대한 확신을 일부러 낮춰서 모델이 너무 확신에 차서 예측하지 않도록 했다. 이건 perplexity 자체는 약간 나빠지지만 정확도와 BLEU 점수는 오히려 올라가는 효과가 있었다.

실험 결과

기계 번역

WMT 2014 영어-독일어 번역에서 big 모델은 28.4 BLEU로 기존 모든 모델(앙상블 포함)을 2점 이상 앞서며 새로운 state-of-the-art를 세웠다. 심지어 base 모델조차 더 적은 학습 비용으로 이전의 모든 단일 모델과 앙상블을 넘어섰다. 영어-프랑스어에서는 41.8 BLEU를 기록해서, 이전 최고 모델 대비 학습 비용은 1/4도 안 되는 수준으로 단일 모델 기준 최고 점수를 갱신했다.

Table 2. WMT 2014 영어-독일어/영어-프랑스어 BLEU 점수 비교

base 모델은 마지막 5개 체크포인트(10분 간격으로 저장)를 평균낸 단일 모델을 썼고, big 모델은 마지막 20개를 평균냈다. 추론 시에는 beam size 4의 beam search와 length penalty α=0.6을 적용했고, 최대 출력 길이는 입력 길이+50으로 제한했다.

어떤 요소가 성능에 영향을 주는가

저자들은 base 모델을 여러 방식으로 바꿔가며 English-to-German 개발셋에서 성능 변화를 측정했다.

Table 3. Transformer 아키텍처 변형에 따른 성능 변화

몇 가지 눈에 띄는 결과가 있다. head 수를 바꿔본 (A) 행을 보면, head가 1개일 때(single-head attention)는 최적 설정보다 0.9 BLEU 낮았고, 반대로 head를 너무 많이 늘려도 성능이 떨어졌다. head 수에는 적당한 지점이 있다는 뜻이다. (B) 행에서는 attention key의 차원 d_k를 줄이면 모델 품질이 나빠지는 걸 확인했는데, 이건 query와 key의 호환성을 판단하는 게 생각보다 단순한 문제가 아니며 dot product보다 더 정교한 호환성 함수가 도움이 될 수도 있다는 걸 시사한다. (C), (D) 행에서는 예상대로 모델이 클수록 성능이 좋아지고, dropout이 과적합 방지에 확실히 도움이 된다는 걸 보여줬다. 그리고 앞서 언급한 (E) 행, 즉 학습된 positional embedding으로 바꿔도 sin/cos 방식과 거의 같은 결과가 나왔다.

다른 task에도 통할까 - 영어 구문 분석

Transformer가 번역 외의 task에도 일반화되는지 확인하려고 영어 constituency parsing(구문 분석) 실험도 진행했다. 이 task는 출력이 강한 구조적 제약을 받고 입력보다 훨씬 길다는 점에서 번역과는 다른 도전 과제였고, 기존 RNN seq2seq 모델은 데이터가 적은 환경에서 state-of-the-art를 내지 못했던 영역이었다.

Table 4. 영어 구문 분석(WSJ Section 23) 결과 비교

Penn Treebank의 WSJ 코퍼스(약 4만 문장)만으로 학습한 4-layer Transformer는 task별 튜닝을 거의 하지 않았는데도 91.3 F1을 기록해서, Recurrent Neural Network Grammar를 제외한 기존에 보고된 모든 모델을 능가했다. 1,700만 문장 규모의 준지도학습 환경에서는 92.7 F1까지 올라갔다. RNN 기반 seq2seq 모델과 달리, Transformer는 4만 문장짜리 WSJ 데이터만으로 학습해도 BerkeleyParser를 능가했다는 점이 특히 의미 있다.

마무리

이 논문이 흥미로운 건 RNN과 CNN을 완전히 들어내고도 번역 성능과 학습 속도 양쪽에서 다 이겼다는 점이다. self-attention 하나로 인코더와 디코더를 구성할 수 있다는 걸 실증적으로 보여줬고, 그 구조가 번역뿐 아니라 구문 분석처럼 성격이 다른 task에도 잘 들어맞는다는 것까지 확인했다.

다만 논문 자체에서도 한계를 인정한 부분이 있다. self-attention의 연산량은 시퀀스 길이의 제곱에 비례하기 때문에, 아주 긴 시퀀스를 다루려면 이웃만 보는 restricted attention 같은 추가 장치가 필요하다고 언급했다. 저자들은 이걸 후속 연구 과제로 남겨뒀고, 텍스트가 아닌 이미지, 오디오, 비디오 같은 다른 모달리티로 확장하는 것도 향후 방향으로 제시했다. 실제로 이후 몇 년 사이에 이 두 방향 모두 활발하게 연구가 진행됐다.

#Transformer #AttentionIsAllYouNeed #SelfAttention #MultiHeadAttention #NLP #논문리뷰