[BoostCamp AI Tech] RNN & LSTM & Attention

RNN

RNN이라고 하는 것은 순환 신경망을 의미한다. 

이것은 자기 자신을 순회하며 학습을 한다. 위의 그림과 비슷하게 이해하면 될 것이다. 

RNN의 이러한 성질 때문에, 이전 입력의 정보를 사용할 수 있다. 정보가 매번 들어오고 동일한 모델을 순회하기 때문이다. 

RNN은 이전 정보를 사용할 수 있다는 장점 때문에, 시계열 데이터를 활용하기 유리했다. 

 

우리는 개발자니까, 좀 더 구체적으로 보자면 다음과 같다. 아래와 같은 구조를 n번 반복하는 것이다. 

동일한 구조를 계속해서 사용하기 때문에 입력이 많아져도 모델의 크기가 증가하지 않는다. 

입력으로는 문장이 들어올 수 있다. 각 입력 단어들이 x 값으로 들어간다. 

 

하지만, RNN의 모양을 생각해보면 GPU의 장점을 사용할 수 없다. 앞에서 나온 출력 결과가 있어야지만, 뒤에서 이를 사용할 수 있기 때문이다. 또한 RNN의 심각한 문제중 하나는 기울기 소실과 장거리 의존성을 설계하기 어렵다는 것이다. 

 

기울기 소실이 왜 생길까? 

이는 위의 계산 그래프를 통해 알 수 있다. 계산 그래프 상에서 역전파를 생각하면 tanh와 Whh 값이 계속해서 곱해지는데, tanh의 값은 항상 1보다 작기에 기울기 소실이 생길 가능성이 높아진다. 또한 Whh 값에 따라서 기울기 폭발이라는 문제가 생길 수도 있다. 이때, 기울기 클리핑을 사용하면 폭발 문제는 해결할 수 있다. 하지만, 여전히 기울기 소실이 문제다.

LSTM

 

이러한 기울기 소실을 해결하기 위해서 나온 모델이 LSTM이다. 아래와 같이 Cell state를 추가한다. 이를 highway라고 표현하기도 한다. 

기울기 소실이 일어나는 부분은 셀 상태에서 곱하기 부분에 해당하는 Forget Gate이다. 이 값에 따라서 역전파를 할 때 얼마나 정보가 지워지는지 알 수 있고, 만약 1이라면 정보가 지워지지 않는 것이다. Input Gate도 곱셈 노드가 있는 부분이고, 새로운 정보를 얼마나 넣을까 하는 이야기다. input gate 값이 작을수록 덜 학습된다. output gate도 마찬가지이다. 

 

seq2seq

이러한 모델들을 응용하면 번역 AI를 만들 수 있다. 이때 중요한 것은 입력으로 무엇을 넣고, 출력이 무언인가이다. 

한국어 문장을 번역하면 영어 문장이 나온다. 때문에 입력도 문장, 출력도 문장이 되어야 한다. 

 

위에서 본 RNN 계열 모델의 개념을 생각하면 매 시점마다 정보가 기억된다.

이것을 출력으로 이어주지 않고, 정보만 따로 저장하면 그것이 언어모델의 인코더가 된다. 

 

RNN에서는 정보를 활용하여 출력을 얻었지만, 아래 구조에서는 정보를 기억만 하고 있다. 

이렇게 저장한 정보는 뒤에 나오는 디코더에서 활용한다. 

(초기 seq2seq 모델은 각 시점을 기억하지 않고, 마지막 출력만 기억했다.)

 

전체적인 구조는 오른쪽 그림과 같아지는데, 이 때 dot 연산을 통해서 가장 유사한 정보를 활용한다. 

 

먼저 attention score를 계산한다. 현재 decoder에서 연산 중인 값과 내적을 한다. 그러면 여러 벡터들이 나올 것인데, 이것이 attention score 들이다.  

(이러한 것이 가능한 이유는 벡터가 정보를 충분히 담고 있다고 가정하는 것이다.)

 

인코더를 코드 상으로 구현하면 다음과 같다. 

 

# Encoder
class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hidden_dim, n_layers, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hidden_dim, n_layers, dropout=dropout, batch_first=True)
    
    def forward(self, src):
        embedded = self.embedding(src)
        outputs, (hidden, cell) = self.rnn(embedded)
        return outputs, (hidden, cell)

 

forward에서 return으로 출력 값과 Hidden 값을 넘겨준다. 

 

이 Attantion score 값을 다시 softmax를 해서, 유사도에 대한 확률 값으로 만들어준다. 그리고 이것을 다시~ h 벡터와 곱해준다. 

 

앞에서 본 encoder에서 사용한 정보를 활용해서 인자로 hidden과 encoder_outputs를 받는다.

# Attention Mechanism
class Attention(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.Wa = nn.Linear(hidden_dim, hidden_dim)
        self.Ua = nn.Linear(hidden_dim, hidden_dim)
        self.Va = nn.Linear(hidden_dim, 1)
    
    def forward(self, hidden, encoder_outputs):
        batch_size = encoder_outputs.size(0)
        seq_len = encoder_outputs.size(1)
        
        hidden = hidden.unsqueeze(1).repeat(1, seq_len, 1)  # Repeat hidden state for each encoder output
        energy = torch.tanh(self.Wa(hidden) + self.Ua(encoder_outputs))
        attention_weights = torch.softmax(self.Va(energy).squeeze(-1), dim=1)
        
        context_vector = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        return context_vector, attention_weights

 

아까도 보았듯이 Attention 구조는 encoder의 출력과 hidden 벡터 사이의 연산이다. 

 

# Decoder
class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hidden_dim, n_layers, dropout=0.5):
        super().__init__()
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim + hidden_dim, hidden_dim, n_layers, dropout=dropout, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, input, hidden, cell, encoder_outputs, attention):
        embedded = self.embedding(input)
        context_vector, _ = attention(hidden[-1], encoder_outputs)
        rnn_input = torch.cat((embedded, context_vector.unsqueeze(1)), dim=-1)
        output, (hidden, cell) = self.rnn(rnn_input, (hidden, cell))
        prediction = self.fc(output.squeeze(1))
        return prediction, hidden, cell

 

decoder에서는 이 정보를 활용하여 예측한다. 

 

최종적으로 합치면 이렇게 나온다. 

 

# Full Seq2Seq Model with Attention
class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, attention):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.attention = attention
    
    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        batch_size = src.size(0)
        trg_len = trg.size(1)
        vocab_size = self.decoder.embedding.num_embeddings
        outputs = torch.zeros(batch_size, trg_len, vocab_size).to(src.device)
        
        encoder_outputs, (hidden, cell) = self.encoder(src)
        
        input = trg[:, 0]
        
        for t in range(1, trg_len):
            output, hidden, cell = self.decoder(input, hidden, cell, encoder_outputs, self.attention)
            outputs[:, t, :] = output
            
            teacher_force = torch.rand(1).item() < teacher_forcing_ratio
            top1 = output.argmax(1) 
            input = trg[:, t] if teacher_force else top1
        
        return outputs