#논문공부 #Attention is all you need #NLP

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing systems (pp. 5998-6008).

- Attention이라는 개념은 Neural machine translation by jointly learning to align and translated의 논문에서 처음 제시된 아이디어라고 한다. align and translate jointly로 표현되어 있는데, context vector로 구체적으로 표현된다. 위 논문에서는 context vector는 특정 시점의 소스 언어(단어 시점)와 전 시점의 예측한 단어와 가장 관련이 있는 정보를 함축하고 있는 벡터로 attention is all you need의 attention의 아이디어를 구축하는데 기초가 된다. 

- 이 논문에서 중요하게 다루고 개념인 "attention"은 "어떤 부분에 더 주의를 기울일 것인가?"로 요약할 수 있다. input단어 정보들을 받아서 다음 번의 단어를 추론하는 데 있어서 input 정보의 모든 토큰에 동일한 비율로 참고하는 것이 아니라 특히 현 시점에 있어서 중요하게 참고 될 수 있는 토큰 정보를 활용하겠다는 것이다. 

"나는 바다에 가고 싶어" 라는 문장을 영어로 번역하는데 있어서 현재 생성한 문장이, "I would like to go" 라고 하면 다음에 생성해야 할 단어가 '싶어'를 번역한 "the sea"부분이다. 이 때, input의 sequence에서 "가고"라는 토큰 정보에 특히 주의를 깊게 기울이겠다는 뜻이다. 결과적으로, '싶어'를 encoding하는데 있어서 '가고'의 토큰 정보를 반영하면서 encoding하게 된다. input에서 어떤 세부 input 토큰에 주의를 기울일지 나타내는 것이 attention이다. 또한 이러한 과정은 구체적으로 self-attention layer에서 이뤄진다. 

이 부분은 Self-attention의 가장 중요한 부분이라고 할 수 있다. 논문마다 attention을 표현하는 함수는 다양하게 표현될 수 있다. 본 논문에서는 "Scaled dot-product attention"으로 부르는 함수를 사용하는데 이는 mapping된 query(Q)와 key(K)-value(V)를 활용한다. 이것은 데이터베이스에서 친숙한 개념인 key-value의 개념과 유사하다. query는 처리하고 있는 단어로 해석할 때, query의 단어가 key 단어 리스트와 얼마나 유사한지를 내적으로 계산하여 score값에 반영한 뒤, 이를 각 단어의 value 벡터와 곱하여 최종적으로 query단어를 주변 토큰 맥락정보를 반영하여 인코딩하게 된다. 모든 계산은 메트릭스 형태로 이뤄지기 때문에 계산이 빠르게 이뤄진다. 'scaled'라는 수식어가 붙는 이유는 Query, Key를 dot-product한 후에, 이것의 차원을 축소하기 위해서 루트 key 차원의 수로 나눠지기 때문이다. 이런 트릭를 거쳐야 학습이 조금 더 안정적으로 이뤄진다고 한다. 

다음 문장을 표현하는 dictionary가 있다고 생각해보자. 

dic = { "나는" : 1, "바다에":2, "가고":3, "싶어":4}    {"key" : "value"} 

(실제로는 value값이 위의 정수값이 아니라 벡터값으로 표현된다.)

이 때 "떠나고" 라는 단어가 query에 있다고 보자. 이 때, "떠나고"라는 단어와 모든 key 값들의 유사도를 구하고 이것들을 weight값으로 표현하여 다음의 출력을 예측하는데 있어서 사용하게 된다. 앞의 설명에서는 key, value, query가 숫자이거나 인간이 이해할 수 있는 언어로 표현되어 있지만 실제 코드에서는 이것이 적절하게 embedding상태인 벡터값으로 표현되어 있다. 

Attention에서의 Query는 현재 time-stamp의 디코더의 출력, Keys는 각 time-step별 인코더의 출력값, Values는 각 time-step별 인코더의 출력값을 의미한다. attention함수는 아래와 같이 표현될 수 있으며, 

이 때, input에 들어가는 Query, key의 차원은 dk, value의 차원은 dv이다. 현 시점 decoder의 출력값인 query를 받아, 이 시점의 key값과 내적(dot-product)한다. 이후 루트 dk로 나눈뒤에, 각각의 value들의 weight를 구한다. 이 때, 정규화를 위해서 softmax 함수를 적용하게 된다. 여기에서, 내적하는 방식을 적용한 이유로 더 빠르고, 실제로 공간 효율성이 있기 때문에 사용했다고 한다. 또한 이는 안정적인 gradient를 만들어서 학습시 유용하다고 한다. 

- Attention 함수를 한 번 사용하는 것보다는, 각각 attention함수를 h번 사용하여 여러개의 attention을 만드는 것이 더 이롭다고 한다. 여러개의 attention을 사용하여 ("Multi-Head attention") 하나로 합치는 과정을 통해서 최종적으로 하나의 attention을 만들게 된다. 본 논문에서는 8개의 attention head를 사용하였다고 한다.

- 한편, 단순히 query와 key값을 내적하기에는 무리가 있다. 한국어와 영어의 단어들의 분포가 단어의 의미에 따라서 비슷하게 분포해 있을 것이지만, 언어가 다르기에 완전히 벡터 공간에서 "바다"의 벡터값이 "sea"의 벡터값과 일치하지 않을 것이다. 따라서 이들 언어를 연결해줄 수 있는 적절한 선형 변환이 필요하다. 선형 변환을 위한 파라미터값들은 신경망 가중치 파라미터로써, 역전파를 통해서 학습된다. 

 

이와 관련한 코드 부분

 

class sae(nn.Module): 
        def __init__(self, vocab_size, embed_size = 512):
            super().__init__()
            dim = embed_size
            num_layers = 1

            # architecture
            self.embed = nn.Embedding(vocab_size, dim, padding_idx = PAD_IDX)
            self.pe = self.pos_encoding(dim)
            self.layers = nn.ModuleList([self.layer(dim) for _ in range(num_layers)])

        def forward(self, x):
            mask = x.eq(PAD_IDX).view(x.size(0), 1, 1, -1)
            x = self.embed(x)
            h = x + self.pe[:x.size(1)]
            for layer in self.layers:
                h = layer(h, mask)
            return h

        @staticmethod
        def pos_encoding(dim, maxlen = 1000): # positional encoding
            pe = Tensor(maxlen, dim)
            pos = torch.arange(0, maxlen, 1.).unsqueeze(1)
            k = torch.exp(-np.log(10000) * torch.arange(0, dim, 2.) / dim)
            pe[:, 0::2] = torch.sin(pos * k)
            pe[:, 1::2] = torch.cos(pos * k)
            return pe

        class layer(nn.Module): # encoder layer
            def __init__(self, dim):
                super().__init__()

                # architecture
                self.attn = embed.sae.attn_mh(dim)
                self.ffn = embed.sae.ffn(dim)

            def forward(self, x, mask):
                z = self.attn(x, x, x, mask)
                z = self.ffn(z)
                return z

        class attn_mh(nn.Module): # multi-head attention
            def __init__(self, dim):
                super().__init__()
                self.D = dim # dimension of model
                self.H = 8 # number of heads
                self.Dk = self.D // self.H # dimension of key
                self.Dv = self.D // self.H # dimension of value

                # architecture
                self.Wq = nn.Linear(self.D, self.H * self.Dk) # query
                self.Wk = nn.Linear(self.D, self.H * self.Dk) # key for attention distribution
                self.Wv = nn.Linear(self.D, self.H * self.Dv) # value for context representation
                self.Wo = nn.Linear(self.H * self.Dv, self.D)
                self.dropout = nn.Dropout(DROPOUT)
                self.norm = nn.LayerNorm(self.D)

            def attn_sdp(self, q, k, v, mask): # scaled dot-product attention
                c = np.sqrt(self.Dk) # scale factor
                a = torch.matmul(q, k.transpose(2, 3)) / c # compatibility function
                a = a.masked_fill(mask, -10000) # masking in log space
                a = F.softmax(a, -1)
                a = torch.matmul(a, v)
                return a # attention weights

            def forward(self, q, k, v, mask):
                b = q.size(0) # batch_size (B)
                x = q # identity
                q = self.Wq(q).view(b, -1, self.H, self.Dk).transpose(1, 2)
                k = self.Wk(k).view(b, -1, self.H, self.Dk).transpose(1, 2)
                v = self.Wv(v).view(b, -1, self.H, self.Dv).transpose(1, 2)
                z = self.attn_sdp(q, k, v, mask)
                z = z.transpose(1, 2).contiguous().view(b, -1, self.H * self.Dv)
                z = self.Wo(z)
                z = self.norm(x + self.dropout(z)) # residual connection and dropout
                return z

        class ffn(nn.Module): # position-wise feed-forward networks
            def __init__(self, dim):
                super().__init__()
                dim_ffn = 2048

                # architecture
                self.layers = nn.Sequential(
                    nn.Linear(dim, dim_ffn),
                    nn.ReLU(),
                    nn.Dropout(DROPOUT),
                    nn.Linear(dim_ffn, dim)
                )
                self.norm = nn.LayerNorm(dim)

            def forward(self, x):
                z = x + self.layers(x) # residual connection
                z = self.norm(z) # layer normalization
                return z