스터디/논문

트랜스포머: Attention is all you need 이해 및 구현

민서타 2024. 7. 22. 18:26

ABSTRACT

구글의 트랜스포머 모델은 자연어처리와 컴퓨터비전 분야의 교과서로 사용된다. 현재 AI논문 인용 횟수 1위(약 91,000회)를 달성할 정도로 지금까지도 주목받고 있는 모델이다. 다양한 트랜스포머에 대한 지식 정리글이 많지만, 이 포스팅은 초심자가 정확한 코드 구현보단 이해에 초점을 맞춰 작성해보려 한다

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


SUMMARY

핵심 아키텍쳐

-포지셔널 인코딩, 인코더, 디코더

-멀티 헤드 어텐션(인코더 셀프, 마스크드 디코더 셀프 어텐션, 인코더-디코더 어텐션)


0. 모델 구조

Figure 2 왼쪽: 인코더
Figure 3 오른쪽: 디코더

인코더는 모델에게 문장을 전달하고, 디코더는 이를 출력 문장으로 반환한다. 디코더는 인코더에서 나온 컨텍스트와 추가로 문장이 하나 더 들어간다.

자세한 내용은 인코더와 디코더 부분에서 다루겠다.


1.포지셔널 인코딩: 입력 시퀀스의 순서를 조절(Figure 1의 Input, output 임베딩 다음단계) 

트랜스포머는 기존 순방향 입력의 RNN 모델과 달리, 입력 시퀀스의 순서에 종속되지 않는다. 트랜스포머는 입력이 한 번에 들어간다. 여기서 단어 간의 순서로 인한 차이

"상추와 고기를 먹었어"와 "고기를 상추와 먹었어"라는 문장이 다른 것처럼

포지셔널 인코딩(위치에 따른 요소별 중요도 부여)을 사용하게 된 것이다. 여기서 사인과 코사인 함수를 사용한 이유는, 학습한 훈련 문장의 최대 길이가 테스트 문장의 길이의 최대값과 다를 때(훈련 문장 25, 테스트 40) 두 함수를 이용하면 -1과 1(짝수: sin, 홀수: cos)로 변환하여 모든값을 대변할 수 있게 된다.

class PositionalEncoding(nn.Module):

    def __init__(self, d_embed, max_len=256, device=torch.device("cpu")):
        super(PositionalEncoding, self).__init__()
        encoding = torch.zeros(max_len, d_embed)
        encoding.requires_grad = False
        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_embed, 2) * -(math.log(10000.0) / d_embed))
        encoding[:, 0::2] = torch.sin(position * div_term)
        encoding[:, 1::2] = torch.cos(position * div_term)
        self.encoding = encoding.unsqueeze(0).to(device)


    def forward(self, x):
        _, seq_len, _ = x.size()
        pos_embed = self.encoding[:, :seq_len, :]
        out = x + pos_embed
        return out

2. Encoder

-  Self Attention(같은 문장 내 두 token 사이의 어텐션 계산) * 어텐션: 넓은 범위의 전체 데이터에서 특정 부분에 집중하는 것
-  피드 포워드
-  정규화
- 멀티 헤드

인코더 구성요소: N개의 인코더 블록(논문, N=6)

* 인코더 블록: Multi-Head Attention Layer와 Position-wise Feed-Forward Layer로 구성

2-1. 셀프 어텐션(Scaled dot-pordict attention, 개별 토큰 사이의 전체 연관도 계산 N x N)

Query: 현 시점의 토큰, Key: 어텐션 대상 토큰, Value: 어텐션 대상 토큰)

어텐션은 문장 내에서 어떤 것에 집중할까를 의미하고, 같은 문장 내에서 어텐션을 계산하는 것을 셀프 어텐션이라고 한다(어텐션 계산을 수학적으로 scaled_dot_product 지칭하기도함) 

문장을 예시로 든다면 I ate chicken에서 첫 입력인 I가 Q가 되고 그 단어와의 연관도 ate가 Key와 Value가 되거나, chicken을 대상으로한 Key와 Value가 되어 관계(어텐션)를 계산한다.

논문 내에선 모두 같은 Shape을 가지고 있으며, 쿼리와 키값의 내적을 통한 스칼라값이 V에 곱해지며 기존의 Shape을 유지한 채 연산이 진행된다. 하지만 입력 토큰인 모든 문장 간의 사이즈가 균일하지 않기 때문에, 제로 패딩(0 또는 1 추가로 shape을 맞춤) 과정인 mask 과정이 사용되게 된다(0: 허수, 1: 실제 토큰) 

def scaled_dot_product_attention(query, key, value):
  dim_k = query.size(-1)
  scoers = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
  weights = F.softmax(scores, dim=-1)
  return torch.bmm(weights, value)
  
  # 실제 어텐션에선 mask 추가
  def calculate_attention(self, query, key, value, mask):
    # query, key, value: (n_batch, h, seq_len, d_k)
    # mask: (n_batch, 1, seq_len, seq_len)
    d_k = key.shape[-1]
    attention_score = torch.matmul(query, key.transpose(-2, -1)) # Q x K^T, (n_batch, h, seq_len, seq_len)
    attention_score = attention_score / math.sqrt(d_k)
    if mask is not None:
        attention_score = attention_score.masked_fill(mask==0, -1e9)
    attention_prob = F.softmax(attention_score, dim=-1) # (n_batch, h, seq_len, seq_len)
    out = torch.matmul(attention_prob, value) # (n_batch, h, seq_len, d_k)
    return out

 

멀티헤드 어텐션: 셀프 어텐션을 n번 병렬처리 시도 및 종합

이러한 셀프 어텐션을 여러 문장 또는 구조에서 반복해서 쌓은 것을 멀티헤드 어텐션이라 한다.


2-2. Position-wise Feed-Forward Layer

개념적으로는 역전파 과정중 생길 수 있는 기울기 소실 방지를 위해 output에 input값을 추가해주는 것이다.

class ResidualConnectionLayer(nn.Module):

    def __init__(self):
        super(ResidualConnectionLayer, self).__init__()


    def forward(self, x, sub_layer):
        out = x
        out = sub_layer(out)
        out = out + x
        return out

Final encoder

class Encoder(nn.Module):
    def __init__(self, n_input_vocab, d_model, head, d_ff, max_len, padding_idx, dropout, n_layers, device):
        super().__init__()

        # Embedding
        self.input_emb = nn.Embedding(n_input_vocab, d_model, padding_idx=padding_idx)
        self.pos_encoding = PositionalEncoding(d_model, max_len, device)
        self.dropout = nn.Dropout(p=dropout)

        # n개의 encoder layer를 list에 추가
        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model=d_model, 
                                                         head=head, 
                                                         d_ff=d_ff, 
                                                         dropout=dropout)
                                             for _ in range(n_layers)])

    def forward(self, x, padding_mask):
        # 1. 입력에 대한 input embedding, positional encoding 생성
        input_emb = self.input_emb(x)
        pos_encoding = self.pos_encoding(x)

        # 2. add & dropout
        x = self.dropout(input_emb + pos_encoding)

        # 3. n번 EncoderLayer 반복
        for encoder_layer in self.encoder_layers:
            x, attention_score = encoder_layer(x, padding_mask)

        return x

3. Decoder

디코더 구성요소: N개의 디코더 블록(논문, N=6)

* 디코더 블록: Self Multi-Head Attention Layer와 Cross Multi-Head Attention Layer,

   Position-wise Feed-Forward Layer로 구성(이때 input인 Context는 모든 디코더 블록에 각각 적용)

- 교사 강요(Teach Forcing)
- 마스크드 멀티 헤드 셀프 어텐션(셀프 멀티 헤드 어텐션)
- 인코더-디코더 어텐션(크로스 멀티 헤드 어텐션)

3-1. 교사 강요

창작의 영역인 생성에선, 우리가 원하는 100% 정확도의 문장을 만들 수 없다. 그래서 지도학습의 정답 데이터를 가져와 학습시키는 교사 강요 개념이 파생되었다. 하지만 순차 입력인 RNN과 달리 트랜스포머는 병렬 처리의 입력이라, 모든 정답을 알게되는 현상이 생기는데. 이 때 예측해야할 토큰은 보이지 않도록 처리하는 것을 subsequent masking 기법이라 한다.

def make_subsequent_mask(query, key):
    # query: (n_batch, query_seq_len)
    # key: (n_batch, key_seq_len)
    query_seq_len, key_seq_len = query.size(1), key.size(1)

    tril = np.tril(np.ones((query_seq_len, key_seq_len)), k=0).astype('uint8') # lower triangle without diagonal
    mask = torch.tensor(tril, dtype=torch.bool, requires_grad=False, device=query.device)
    return mask
    # 하삼각행렬 생성
    [[1,0,0,0,0,0,0],
     [1,1,0,0,0,0,0],
     [1,1,1,0,0,0,0],...]]

i번째 토큰은 자기자신인 i만 볼 수 있게 된다. 하지만 인코더의 입력과 마찬가지로

패드 마스킹을 적용해야 한다.

def make_tgt_mask(self, tgt):
    pad_mask = self.make_pad_mask(tgt, tgt)
    seq_mask = self.make_subsequent_mask(tgt, tgt)
    mask = pad_mask & seq_mask
    return pad_mask & seq_mask
class Decoder(nn.Module):
    def __init__(self, n_output_vocab, d_model, head, d_ff, max_len, padding_idx, dropout, n_layers, device):
        super().__init__()

        # output embbeding 
        self.output_emb = nn.Embedding(n_output_vocab, d_model, padding_idx=padding_idx)
        self.pos_encoding = PositionalEncoding(d_model, max_len, device)
        self.dropout = nn.Dropout(p=dropout)

        # n 개의 decoder layer
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model=d_model, 
                                                         head=head, 
                                                         d_ff=d_ff, 
                                                         dropout=dropout)
                                             for _ in range(n_layers)])

    def forward(self, x, memory, look_ahead_mask, padding_mask):
        # 1. 입력에 대한 output embedding, positional encoding 생성
        # (batch_size, seq_len, d_model)
        output_emb = self.output_emb(x) 
        # (seq_len, d_model)
        pos_encoding = self.pos_encoding(x) 

        # 2. add & dropout
        # (batch_size, seq_len, d_model)
        x = self.dropout(output_emb + pos_encoding) 

        # 3. n번 DecoderLayer 반복
        for decoder_layer in self.decoder_layers:
            x = decoder_layer(x, memory, look_ahead_mask, padding_mask)

        return x

3-2. 마스크드 셀프 멀티 헤드 어텐션

인코더의 셀프 멀티 헤드 어텐션과 동일하지만 추가로 mask로 들어오는 인자에

패드 마스킹 + subsequent 마스킹 포함 


3-3. 인코더-디코더 어텐션

이전 마스크드 셀프 멀티 어텐션의 출력물과 인코더의 Context를 입력으로 받는다.

디코더의 출력물은 쿼리로, 인코더의 출력물은 Key와 Value로 사용한 크로스 멀티 헤드 어텐션을 계산

즉, 교사 강요로 얻어진 문장에 최대한 유사한 예측 문장을 도출하고자 하는 것이다.


이후 FC layer를 통해 Shape을 맞추고, Softmax를 통해 각 단어에 대한

확률 값으로  output을 도출한다.


Final Transformer

class Transformer(nn.Module):
    def __init__(self, n_input_vocab, n_output_vocab, d_model, head, d_ff, max_len, padding_idx, dropout, n_layers, device):
        super().__init__()
        self.padding_idx = padding_idx
        self.device = device

        # Encoder
        self.encoder = Encoder(n_input_vocab=n_input_vocab, 
                               d_model=d_model, 
                               head=head, 
                               d_ff=d_ff, 
                               max_len=max_len, 
                               padding_idx=padding_idx, 
                               dropout=dropout, 
                               n_layers=n_layers,
                               device=device)

        # Decoder
        self.decoder = Decoder(n_output_vocab=n_output_vocab, 
                               d_model=d_model, 
                               head=head, 
                               d_ff=d_ff, 
                               max_len=max_len, 
                               padding_idx=padding_idx, 
                               dropout=dropout, 
                               n_layers=n_layers,
                               device=device)

        # linear layer 
        # (batch_size, seq_len, n_output_vocab)
        self.linear = nn.Linear(d_model, n_output_vocab)

    def forward(self, src, tgt):
        # 1. 입력에 따른 mask 생성
        padding_mask = self.make_padding_mask(src, src)
        enc_dec_padding_mask = self.make_padding_mask(tgt, src)
        look_ahead_mask = self.make_padding_mask(tgt, tgt) * self.make_look_ahead_mask(tgt)

        # 2. encoder
        memory = self.encoder(src, padding_mask)

        # 3. decoder
        output = self.decoder(tgt, memory, look_ahead_mask, enc_dec_padding_mask)

        # 4. linear layer
        output = self.linear(output)

        return output

    def make_padding_mask(self, q, k):
        # q,k의 size = (batch_size, seq_len)
        _, q_seq_len = q.size()
        _, k_seq_len = k.size()

        q = q.ne(self.padding_idx)  # padding token을 0, 나머지를 1로 만들어줌
        q = q.unsqueeze(1).unsqueeze(3) # (batch_size, 1, q_seq_len, 1)
        q = q.repeat(1,1,1,k_seq_len)   # (batch_size, 1, q_seq_len, k_seq_len)

        k = k.ne(self.padding_idx)
        k = k.unsqueeze(1).unsqueeze(2) # (batch_size, 1, 1, k_seq_len)
        k = k.repeat(1,1,q_seq_len,1)   # (batch_size, 1, q_seq_len, k_seq_len)

        # and 연산
        # (batch_size, 1, q_seq_len, k_seq_len)
        mask = q & k

        return mask

    def make_look_ahead_mask(self, tgt):
        _, seq_len = tgt.size()

        # torch.tril 마스킹
        # (seq_len, seq_len)
        mask = torch.tril(torch.ones(seq_len,seq_len)).type(torch.BoolTensor).to(self.device)

        return mask

REFERENCE

도서: 트랜스포머를 활용한 자연어 처리, 한빛 미디어

파이썬 코드인용: https://cpm0722.github.io/pytorch-implementation/transformer

 

[NLP 논문 구현] pytorch로 구현하는 Transformer (Attention is All You Need)

Paper Link

cpm0722.github.io

https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Transformers

 

GitHub - sgrvinod/a-PyTorch-Tutorial-to-Transformers: Attention Is All You Need | a PyTorch Tutorial to Transformers

Attention Is All You Need | a PyTorch Tutorial to Transformers - sgrvinod/a-PyTorch-Tutorial-to-Transformers

github.com

 

반응형