[Transformer 이해하기 2] Neural Machine Translation by Jointly Learning to Align and Translate, Attention 설명 및 코드리뷰

다음 포스팅을 참고하여 작성하였음을 밝힙니다.

https://wikidocs.net/22893

1) 어텐션 메커니즘 (Attention Mechanism)

 

 

들어가기에 앞서 Seq2Seq 모델을 살펴봅시다.

 

Seq2Seq with 2 LSTM layers

 

Seq2Seq는 번역 과정에서 입력문장을 벡터 $z$로 임베딩하게 됩니다. 이러한 Seq2Seq는 두 가지 문제점이 있는데,

* 고정 벡터에 입력문장 정보를 압축하려고 하다보니 정보 손실이 생기고,

* 인코더 / 디코더에서 사용하는 RNN이 고질적으로 그래디언트 손실을 유발한다는 것입니다.

 

Attention은 이러한 단점을 극복하기 위해, 디코더에서 문장을 번역하는 매 step마다 전체 입력 문장을 참고합니다. 

디코더에서, 해당 step의 번역할 부분과 관련이 높은 입력문장 부분을 집중적으로 참고하는 모델이라 하여 Attention이라는 이름을 가집니다.

 

"Attention은 번역 과정에서 전체 입력문장을 참고한다" 는 기본 내용을 인지한 상태로 pytorch 코드와 Attention 메커니즘을 함께 살펴봅시다.

 

 

논문 링크: https://arxiv.org/abs/1409.0473

깃허브 링크: https://github.com/bentrevett/pytorch-seq2seq/blob/master/3%20-%20Neural%20Machine%20Translation%20by%20Jointly%20Learning%20to%20Align%20and%20Translate.ipynb

GitHub - bentrevett/pytorch-seq2seq: Tutorials on implementing a few sequence-to-sequence (seq2seq) models with PyTorch and Torc

 

 

 

 

 

 Introduction

 

Attention이 계산되는 과정을 먼저 살펴봅시다.

 

 

* Attention 함수는 "주어진 쿼리(Query)에 대해서 모든 키(Key)와의 유사도를 각각 구해 키와 매핑되어 있는 값(Value)에 반영" 합니다. key, value는 딕셔너리를 구성하는 요소입니다. 

 

* 여기서 주어진 쿼리는 $s_{t}^{T}$ 입니다. $s_{t}^{T}$는 디코딩 부분의 LSTM을 통과하여 나온 임베딩 벡터고, 이 벡터가 인코더의 임베딩 벡터들인 $h_1, \cdots, h_4$와 얼마나 유사도 (Attention score) 를 가지고 있는지 보는 것입니다. (인코더의 임베딩 벡터가 Key, Value가 됩니다)

 

* 위 이미지에서는 벡터의 내적으로 유사도를 산출했지만, 유사도를 산출하는 방법은 여러가지가 있습니다.

 

* 문장 길이 (토큰의 개수) 만큼의 Attention score들을 softmax 함수를 통과시켜서 Attention Distribution을 구합니다.

 

* Attention Distribution vector $\alpha = [\alpha_1, \alpha_2, \alpha_3, \alpha_4]$와 인코더의 임베딩 행렬인 $h = [h_1, h_2, h_3, h_4]$ 를 행렬곱하여 Context vector (Attention value) 를 구합니다. $$\text{Attention}(Q, K, V) = \text{Attention value}$$

 

(Seq2Seq 에서는 인코더를 통과하고 나온 단일 임베딩 벡터를 Context vector 로 부르는 차이가 있습니다)

 

* 디코더의 임베딩 벡터와 Context vector를 단순히 이어붙이고 (concatenate),

크기가 (concatenate 벡터 길이) x (임베딩 벡터 길이) 인 Dense Layer를 통과시켜 최종적으로 당초의 임베딩 벡터와 같은 크기의 벡터를 생성하고, Softmax 함수를 통과시켜서 다음 토큰을 예측합니다.

 

 

이제 코드를 살펴봅시다!

 

 

 

 

 Building the Seq2Seq Model

 

(Preparing Data 부분은 이전 [Transformer 이해하기 1]과 동일합니다)

[Transformer 이해하기 1] 에서는 인코더와 디코더에 단방향 LSTM을 사용했으나, 이번엔 양방향 GRU를 사용해볼 것입니다. 양방향 RNN은 다음 포스팅과 아래 이미지를 참고하시기 바랍니다. 

 

양방향 RNN 모식도

 

인코더 부분 코드를 살펴봅시다.

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
        super().__init__()
        
        self.embedding = nn.Embedding(input_dim, emb_dim)
        
        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
        
        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src):
        
        #src = [src len, batch size]
        
        embedded = self.dropout(self.embedding(src))
        
        #embedded = [src len, batch size, emb dim]
        
        outputs, hidden = self.rnn(embedded)
                
        #outputs = [src len, batch size, hid dim * num directions]
        #hidden = [n layers * num directions, batch size, hid dim]
        
        #hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]
        #outputs are always from the last layer
        
        #hidden [-2, :, : ] is the last of the forwards RNN 
        #hidden [-1, :, : ] is the last of the backwards RNN
        
        #initial decoder hidden is final hidden state of the forwards and backwards 
        #  encoder RNNs fed through a linear layer
        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
        
        #outputs = [src len, batch size, enc hid dim * 2]
        #hidden = [batch size, dec hid dim]
        
        return outputs, hidden

 

* 토큰화된 문장을 input으로 받고, 임베딩 차원으로 변환합니다. (emb_dim)

* fc (fully-connected) 는 GRU가 양방향이기 때문에 enc_hid_dim * 2 를 input 크기로 받습니다.

* outputs는 산출된 hidden layer 이므로, [문장 길이, 배치사이즈, 인코더 임베딩] 크기를 갖습니다.

* hidden은 디코더로 진입하는 벡터이므로, fc로 크기가 조절되어 [배치사이즈, 디코더 임베딩] 크기를 갖습니다.

 

 

이제 디코더로 진입하여 매 스텝마다 계산될 Attention을 코드로 살펴봅시다.

class Attention(nn.Module):
    def __init__(self, enc_hid_dim, dec_hid_dim):
        super().__init__()
        
        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
        self.v = nn.Linear(dec_hid_dim, 1, bias = False)
        
    def forward(self, hidden, encoder_outputs):
        
        #hidden = [batch size, dec hid dim]
        #encoder_outputs = [src len, batch size, enc hid dim * 2]
        
        batch_size = encoder_outputs.shape[1]
        src_len = encoder_outputs.shape[0]
        
        #repeat decoder hidden state src_len times
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
        
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        
        #hidden = [batch size, src len, dec hid dim]
        #encoder_outputs = [batch size, src len, enc hid dim * 2]
        
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim = 2))) 
        
        #energy = [batch size, src len, dec hid dim]

        attention = self.v(energy).squeeze(2)
        
        #attention= [batch size, src len]
        
        return F.softmax(attention, dim=1)

 

* $\text{Attention}(Q, K, V) = \text{Attention value}$

* Q에 해당하는 디코더의 임베딩 값을 hidden 으로 할당합니다.

* 디코더의 hidden 과 인코더의 encoder_outputs로 Attention Distribution을 산출합니다.

* 위의 예시 이미지에서는 인코더/디코더의 hidden layer 크기를 동일하게 두고 내적하는 방법을 사용했지만, 위 코드에서는 인코더/디코더의 hidden layer를 concatenate 하고 Dense layer (attn, v) 를 통해서 길이를 input 문장 길이와 맞게 조절해주었습니다.

 

 

이어서 디코더 코드를 살펴봅시다.

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
        super().__init__()

        self.output_dim = output_dim
        self.attention = attention
        
        self.embedding = nn.Embedding(output_dim, emb_dim)
        
        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
        
        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, input, hidden, encoder_outputs):
             
        #input = [batch size]
        #hidden = [batch size, dec hid dim]
        #encoder_outputs = [src len, batch size, enc hid dim * 2]
        
        input = input.unsqueeze(0)
        
        #input = [1, batch size]
        
        embedded = self.dropout(self.embedding(input))
        
        #embedded = [1, batch size, emb dim]
        
        a = self.attention(hidden, encoder_outputs)
                
        #a = [batch size, src len]
        
        a = a.unsqueeze(1)
        
        #a = [batch size, 1, src len]
        
        encoder_outputs = encoder_outputs.permute(1, 0, 2)
        
        #encoder_outputs = [batch size, src len, enc hid dim * 2]
        
        weighted = torch.bmm(a, encoder_outputs)
        
        #weighted = [batch size, 1, enc hid dim * 2]
        
        weighted = weighted.permute(1, 0, 2)
        
        #weighted = [1, batch size, enc hid dim * 2]
        
        rnn_input = torch.cat((embedded, weighted), dim = 2)
        
        #rnn_input = [1, batch size, (enc hid dim * 2) + emb dim]
            
        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
        
        #output = [seq len, batch size, dec hid dim * n directions]
        #hidden = [n layers * n directions, batch size, dec hid dim]
        
        #seq len, n layers and n directions will always be 1 in this decoder, therefore:
        #output = [1, batch size, dec hid dim]
        #hidden = [1, batch size, dec hid dim]
        #this also means that output == hidden
        assert (output == hidden).all()
        
        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        weighted = weighted.squeeze(0)
        
        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))
        
        #prediction = [batch size, output dim]
        
        return prediction, hidden.squeeze(0)

 

* 디코더는 hidden layer, 각 스텝의 output을 input으로 받습니다.

* 각 스텝의 output을 디코더 임베딩 크기에 맞춰주고, hidden layer와 encoder_outputs로 Attention을 계산합니다.
(위 코드의 weighted 에 해당합니다)

* 계산한 Attention을 디코더 임베딩 벡터와 concatenate 하고, 이것과 hidden layer를 디코더의 GRU에 진입시킵니다. 

* 토큰으로 output을 가져오기 위해 fc_out (Dense) 으로 맞춰줍니다.

 

 

이제 전체 Seq2Seq 모델에서 Attention이 어떻게 작동하는지 살펴봅시다.

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
        
    def forward(self, src, trg, teacher_forcing_ratio = 0.5):
        
        #src = [src len, batch size]
        #trg = [trg len, batch size]
        #teacher_forcing_ratio is probability to use teacher forcing
        #e.g. if teacher_forcing_ratio is 0.75 we use teacher forcing 75% of the time
        
        batch_size = src.shape[1]
        trg_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim
        
        #tensor to store decoder outputs
        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
        
        #encoder_outputs is all hidden states of the input sequence, back and forwards
        #hidden is the final forward and backward hidden states, passed through a linear layer
        encoder_outputs, hidden = self.encoder(src)
                
        #first input to the decoder is the <sos> tokens
        input = trg[0,:]
        
        for t in range(1, trg_len):
            
            #insert input token embedding, previous hidden state and all encoder hidden states
            #receive output tensor (predictions) and new hidden state
            output, hidden = self.decoder(input, hidden, encoder_outputs)
            
            #place predictions in a tensor holding predictions for each token
            outputs[t] = output
            
            #decide if we are going to use teacher forcing or not
            teacher_force = random.random() < teacher_forcing_ratio
            
            #get the highest predicted token from our predictions
            top1 = output.argmax(1) 
            
            #if teacher forcing, use actual next token as next input
            #if not, use predicted token
            input = trg[t] if teacher_force else top1

        return outputs

 

* 인코더 부분에서 Attention 산출을 위한 encoder_outputs을 계산하는 것 외에는 기존 Seq2Seq 모델과 동일합니다.

 

---

 

Attention이 어떤 방식으로 구현되는지 살펴보았습니다. 디코더의 매 스텝마다 입력 문장의 정보들을 참고하여, 유사도가 높은 토큰들에 집중하여 성능을 개선한 것을 볼 수 있습니다.

 

이제 Attention이 무엇인지 이해했으니, 다음 포스팅에서는 "Attention is all you need" 논문에 실린 Transformer를 살펴보겠습니다. 

 

감사합니다!