[논문구현] Attention is all you need - 코드를 만들어보자 - 1

Transformer

Transformer의 구조는 간략하게 설명해서 input sentence를 넣어서 output sentence를 생성해내는 model이다.
Trnasformer는 Encoder와 Decoder로 구성되어있다.
아래의 이미지는 전체적인 Transformer의 구조도이다.
왼쪽 부분을 Encoder, 오른쪽 부분을 Decoder라고한다. 전체적인 동작 과정은 논문 리뷰에 설명했으니 여기서는 코드를 위주로 풀어가려고 한다.

Encoder & Decoder

Encoder의 목표는 context를 제대로 생성(문장의 정보를 빠뜨리지 않고 압축)하는 것.

Decoder는 context를 input으로 받아 Sentence를 output으로 생성해낸다. context만 받는 것이 아니라 output으로 생성해내는 sentence를 right shift한 sentence도 함께 입력받는다. 따라서 Decoder는 sentence, context를 input으로 받아 sentence를 만들어내는 함수이다.

Encoder와 Decoder에 모두 context vector가 등장하는데, Encoder는 context를 생성해내고, Decoder는 context를 사용한다.
이러한 흐름으로 Encoder와 Decoder가 연결되어 전체 Transformer를 구성하는 것이다.

지금까지의 개념을 바탕으로 아주 간단한 Transformer의 model을 구현해보자.

import torch
import torch.nn as nn


class Transformer(nn.Module):
    def __init__(self, encoder, decoder):
        super(Transformer, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def encode(self, x):
        out = self.encoder(x)
        return out

    def decode(self, z, c):
        out = self.decoder(z, c)
        return out

    def forward(self, x, z):
        c = self.encode(x)
        y = self.decode(z, c)

        return y

Encoder 구조

Encoder는 위와 같은 구조로 이루어져 있다.
N개로 Encoder Block이 쌓여진 형태이다. 논문에서는 N=6로 설정했다. Encoder Block은 input과 output의 형태가 동일하다. 어떤 matrix를 input으로 받는다고 했을 때, Encoder Blcok이 도출해내는 output은 input과 완전히 동일한 shape을 갖는 matrix가 된다.


Encoder Block N개가 쌓여 Encoder를 이룬다고 했을 때, 첫번째 Encoder Block의 input은 전체 Encoder의 input으로 들어오는 문장 embedding이 된다. 첫번째 block이 output을 생성해내면 이를 두 번째 block이 input으로 사용하고, 또 그 output을 세 번째 block이 사용하는 식으로 연결되며, 가장 마지막 N번째 block의 output이 전체 Encoder의 output, 즉, context가 된다.


이러한 방식으로 block들이 연결되기 때문에, Encoder Block의 input과 output의 shape는 필연적으로 반드시 동일해야만 한다. 여기서 주목해야 하는 지점은 위에서 계속 언급했던 context 역시 Encoder의 input sentence와 동일한 shape를 가진다는 것이다. 즉, Encoder Block 뿐만 아니라 Encoder 전체도 shape에 대해 멱등(Idempotent)하다.


각 Encoder Block은 input으로 들어오는 vector에 대해 더 높은 차원(넓은 관점)에서의 context를 담는다. 높은 차원에서의 context라는 것은 더 추상적인 정보라는 의미이다. Encoder Block은 내부적으로 어떠한 Mechanism을 사용해 context를 담아내는데, Encoder Block이 겹겹이 쌓이다 보니 처음에는 원본 문장에 대한 낮은 수준의 context였겠지만 이후 context에 대한 context, context의 context에 대한 context … 와 같은 식으로 점차 높은 차원의 context가 저장되게 된다. Encoder Block의 내부적인 작동 방식은 곧 살펴볼 것이기에, 여기서는 직관적으로 Encoder Block의 역할, Encoder 내부의 전체적인 구조만 이해하고 넘어가자.

class Encoder(nn.Module):

    def __init__(self, encoder_block, n_layers):
        super(Encoder, self).__init__()
        self.layers = []

        for i in range(n_layers):
            self.layers.append(copy.deepcopy(encoder_block))

    def forward(self, x):

        out = x 

        for layer in self.layers:
            out = layers(out)
        return out 

forward()를 주목해보자. Encoder Block들을 순서대로 실행하면서, 이전 block의 output을 이후 block의 input으로 넣는다. 첫 block의 Encoder 전체의 input인 x가 된다. 이후 가장 마지막 block의 output, 즉 context를 return한다.

Encoder Block은 크게 Multi-Head Attention Layer, Position-wise Feed-Forward Layer로 구성된다. 각각의 layer에 대한 자세한 설명은 아래에서 살펴보도록 하고, 우선은 Encoder Block의 큰 구조만을 사용해 간단하게 구현해보자.

class EncoderBlock(nn.Module):

    def __init__(self, self_attention, position_ff):
        super(EncoderBlock, self).__init__()
        self.self_attention = self_attention
        self.position_ff = position_ff

    def forward(self, x):
        out = x
        out = self.self_attention(out)
        out = self.position_ff(out)

        return out

Attention

Multi-Head Attention은 Scaled Dot-Proudct-Attention을 병렬적으로 여러 개 수행하는 layer이다. 때문에 Multi-Head Attention을 이해하기 위해서는 Scaled Dot-Product Attention에 대해 먼저 알아야만 한다. Attention이라는 것은 넓은 범위의 전체 data에서 특정한 부분에 집중한다는 의미이다. Scaled Dot-Product Attention 자체를 줄여서 Attention으로 부르기도 한다. 다음의 문장을 통해 Attention의 개념을 이해해보자.

• The animal didn’t cross the street, because it was too tired.

위 문장에서 ‘it’은 무엇을 지칭하는 것일까? 사람이라면 직관적으로 ‘animal’과 연결지을 수 있지만, 컴퓨터는 ‘it’이 ‘animal’을 가리키는지, ‘street’를 가리키는지 알지 못한다. Attention은 이러한 문제를 해결하기 위해 두 token 사이의 연관 정도를 계산해내는 방법론이다.
위의 경우에는 같은 문장 내의 두 token 사이의 Attention을 계산하는 것이므로, Self-Attention이라고 부른다. 반면, 서로 다른 두 문장에 각각 존재하는 두 token 사이의 Attention을 계산하는 것을 Cross-Attention이라고 부른다.

Query, Key, Value

구체적으로 어떤 방식으로 행렬 곱셈을 사용해 Attention이 수행되는지 알아보자. 우선은 문제를 단순화하기 위해 Cross-Attention이 아닌 Self-Attention의 경우를 보겠다. 위의 예시 문장을 다시 가져와보자.

• The animal didn’t cross the street, because it was too tired.

Attention 계산에는 Query, Key, Value라는 3가지 vector가 사용된다. 각 vector의 역할을 정리하면 다음과 같다.

Query: 현재 시점의 token을 의미
Key: attention을 구하고자 하는 대상 token을 의미
Value: attention을 구하고자 하는 대상 token을 의미 (Key와 동일한 token)

위 문장에서 ‘it’이 어느 것을 지칭하는지 알아내고자 하는 상황이다. 그렇다면 ‘it’ token과 문장 내 다른 모든 token들에 대해 attention을 구해야 한다. 이 경우에는 Query는 ‘it’으로 고정이다. Key, Value는 서로 완전히 같은 token을 가리키는데, 문장의 시작부터 끝까지 모든 token들 중 하나가 될 것이다. Key와 Value가 ‘The’를 가리킬 경우 ‘it’과 ‘The’ 사이의 attention을 구하는 것이고, Key와 Value가 마지막 ‘tired’를 가리킬 경우 ‘it’과 ‘tired’ 사이의 attention을 구하는 것이 된다. 즉, Key와 Value는 문장의 처음부터 끝까지 탐색한다고 이해하면 된다. Query는 고정되어 하나의 token을 가리키고, Query와 가장 부합하는(Attention이 가장 높은) token을 찾기 위해서 Key, Value를 문장의 처음부터 끝까지 탐색시키는 것이다. 각각의 의미는 이해했으나, Key와 Value가 완전히 같은 token을 가리킨다면 왜 두 개가 따로 존재하는지 의문이 들 수 있다. 이는 이후에 다룰 것이나, 결론부터 말하자면 Key와 Value의 실제 값은 다르지만 의미적으로는 여전히 같은 token을 의미한다. Key와 Value는 이후 Attention 계산 과정에서 별개로 사용하게 된다.

Query, Key, Value가 각각 어떤 token을 가리키는지는 이해가 됐을 것이다. 하지만, 그래서 Query, Key, Value라는 세 vector의 구체적인 값은 어떻게 만들어지는지는 우리는 아직 알지 못한다. 정말 간단하게도, input으로 들어오는 token embedding vector를 fully connected layer에 넣어 세 vector를 생성해낸다. 세 vector를 생성해내는 FC layer는 모두 다르기 때문에, 결국 self-attention에서는 Query, Key, Value를 구하기 위해 3개의 서로 다른 FC layer가 존재한다. 이 FC layer들은 모두 같은 input shape, output shape를 갖는다. input shape가 같은 이유는 당연하게도 모두 다 동일한 token embedding vector를 input으로 받기 때문이다. 한편, 세 FC layer의 output shape가 같다는 것을 통해 각각 별개의 FC layer로 구해진 Query, Key, Value가 구체적인 값은 다를지언정 같은 shape를 갖는 vector가 된다는 것을 알 수 있다. 정리하자면, Query, Key, Value의 shape는 모두 동일하다. 앞으로 이 세 vector의 dimension을 dk로 명명한다. 여기서 k는 Key를 의미하는데, 굳이 Query, Key, Value 중 Key를 이름으로 채택한 이유는 특별히 있지 않고, 단지 논문의 notation에서 이를 채택했기 때문이다. 이제 위에서 얘기했던 Key, Value가 다른 값을 갖는 이유를 이해할 수 있다. input은 같은 token embedding vector였을지라도 서로 다른 FC layer를 통해서 각각 Key, Value가 구해지기 때문에 같은 token을 가리키면서 다른 값을 갖는 것이다.

Scaled Dot-Product Attention

이제 Query, Key, Value를 활용해 Attention을 계산해보자. Attention이라고 한다면 어떤 것에 대한 Attention인지 불명확하다. 구체적으로, Query에 대한 Attention이다. 이 점을 꼭 인지하고 넘어가자. 이후부터는 Query, Key, Value를 각각 Q, K, V로 축약해 부른다. Query의 Attention은 다음과 같은 수식으로 계산된다.

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

아래는 그림으로 계산의 흐름을 표현한 것이다.

출처: Attention is All You Need [https://arxiv.org/pdf/1706.03762.pdf]

key-query-value 정리