Self-Attentive Sentence Embedding(SASE) 구현해보기

모델 성명은 여기를 참고해주세요. 

pytorch version 1.8.0입니다. 

 

SASE모델을 크게 아래처럼 쪼개갰습니다. 

• embedder
• self-attention
• classifier
• loss function

먼저 들어가기 전에 라이브러리 import를 해주셔야합니다. 

import torch
import torch.nn as nn
import torch.nn.functional as F

1. Embedder

임베더는 말 그대로 초기 단어를 임베딩하고, bi-LSTM을 통해 문맥정보를 압축시켜주는 역할을 하는 부분입니다. 

저는 word2vec을 따로 학습시켜주었고요. 다음과 같이 구현해주었습니다.

모델 그림으로 보면 위의 빨간 네모박스 부분이라고 생각해주시면 됩니다. 

 

수식을 한번 볼까요?

윗 부분은 word2vec embedding을 불러오는 부분이라 생각하시면 됩니다. 

위의 수식은 bi-LSTM을 거쳐 나온 부분이라 생각하시면 됩니다. 

class SentenceEmbedder(nn.Module):
    def __init__(self, cfg, embeddings=None, vocab_size=None):
        super(SentenceEmbedder, self).__init__()
        self.hidden_layer = cfg['hidden_layer']
        self.att_hidden_layer = cfg['att_hidden_layer']
        self.gaussian = GaussianNoise(cfg['gaussian'])
        self.dropout = nn.Dropout(p=cfg['dropout'])
        
        ## 임베딩 불러오기
        if embeddings is not None:
            self.embedding_dim = cfg['embed_dim_word']
            self.embedding = nn.Embedding.from_pretrained(embeddings=embeddings,
                                                          freeze=True)
        else:
            self.embedding_dim = cfg['embed_dim_word']
            self.embedding = nn.Embedding(num_embeddings=vocab_size,
                                          embedding_dim=self.embedding_dim)
        ## LSTM 선언
        self.lstm = nn.LSTM(input_size=self.embedding_dim, hidden_size=self.hidden_layer, num_layers=2, batch_first=True, bidirectional=True)
        
    def forward(self, input):
        S = self.embedding(input)
        S = self.gaussian(S)
        S = self.dropout(S)
        H, _ = self.lstm(S)
        
        return H

저는 여기서 가우시안 노이즈를 사용했는데요. 필요하신 분들은 저처럼 작성하시면 됩니다. 

class GaussianNoise(nn.Module):
    def __init__(self, std):
        super().__init__()
        self.std = std

    def forward(self, x):
        if self.training:
            noise = torch.randn_like(x) * self.std
            return x + noise
        return x

2. Self-attention

이번에는 self-attention을 수행하는 부분입니다. 

그림에서보면 위의 빨간 부분에 해당합니다. 

결과로 나온 벡터를 가지고 attended vector인 M을 만드는 부분이죠.

먼저 Attention score를 구해줍니다. 

이후 기존 벡터에 곱해주어 attended vector를 완성합니다. 

class SelfAttention(nn.Module):
    def __init__(self, cfg):
        super(SelfAttention, self).__init__()
        
        self.hidden_layer = cfg['hidden_layer']
        self.att_hidden_layer = cfg['att_hidden_layer']
        self.hops = hops
        
        ## 어텐션 구하는 부분
        self.att_score = nn.Sequential(
            nn.Linear(self.hidden_layer*2, self.att_hidden_layer),
            nn.Tanh(),
            nn.Linear(self.att_hidden_layer, self.hops)
        )
        
        
    def forward(self,H):
        A = self.att_score(H)
        A = nn.functional.softmax(A, dim=1)
        A = weights.transpose(1, 2)
        M = torch.bmm(A, H)
        
        return M, A

3. Classifier&Loss function

드디어 마지막 부분입니다. 

여기서 loss function을 함께 붙인 이유는 frobenius norm계산을 클래스 내에서 할것이기 때문입니다. 

위의 빨간 박스 부분과 loss를 구하는 부분을 함께 구현해보겠습니다. 

classifier 부분은 단순히 fc layer만 거쳐주면됩니다. 

 

수식을 봐야할 부분은 loss function부분인데요. 

위와 같이 수식을 작성해 주시면 됩니다. 

class SelfAttentiveModel(nn.Module):
    def __init__(self, cfg, embeddings=None, char_vocab=None, vocab_size=None):
        super(SelfAttentiveModel, self).__init__()
        
        ## parameter setting
        self.hidden_layer = cfg['hidden_layer']
        self.att_hidden_layer = cfg['att_hidden_layer']
        self.classifier_hidden_layer = cfg['classifier_hidden_layer']
        self.hops = cfg['hops']
        self.dropout_f = nn.Dropout(p=cfg['dropout_dens'])
        
        ## embedder define
        self.embedder = SentenceEmbedder(cfg, embeddings=embeddings)
        
        ## attention define
        self.attention = SelfAttention(cfg)

        ## final layer
        self.fc = nn.Sequential(
            nn.Linear(self.hops*self.hidden_layer*4, self.classifier_hidden_layer),
            nn.ReLU(),
            nn.Linear(self.classifier_hidden_layer, cfg['target_num'])
        )

        
    def _frobenius(self, attention):
        batch_size = attention.shape[0]
        AAT = torch.bmm(attention, attention.transpose(1,2))
        I = torch.eye(self.hops).unsqueeze(0).repeat(batch_size, 1, 1).cuda()
        penalization_term = torch.norm(AAT - I) / batch_size
        
        return penalization_term
            
    def forward(self, input):
        ## embedding
        H = self.embedding(input)
        M, A = self.attention(H)
       
        M = M.view(-1,self.hops*self.hidden_layer*2)
   
        final = self.fc(M)
        loss = self._frobenius(A)

        return  final, loss

위에서 나온 loss를 모델에서 계산된 loss와 합산해주시면 됩니다.