BERT는 어떻게 학습시킬까? (BERT Pretraining 시키기)

수많은 NLP Downstream Task에서 SOTA를 달성한 BERT에 대해 알아보자. 본 글에서는 모델의 구조와 성능에 대한 얘기가 아닌, BERT 학습의 전반적인 이야기를 해보고자 한다. 따라서 본 글은 BERT의 모델 구조에 대한 이해를 필요로 한다. 아직 BERT가 무엇인지 모른다면 아래의 여기를 참고하자. 

 

본 글은 NVIDIA/BERT 코드를 읽고 정리한 것이다. 전편과 이어지는 글이기에 본 글을 읽기 전 1편을 읽는 것을 강력히 추천드리는 바이다.

---

BERT Pretraining 시키기

 

 

2편에서는 1편에서 만든 Training Instance로 Pretraining 시키는 과정에 대해 알아볼 것이다. 1편의 마지막 내용을 살짝 떠올려보면, Training Instance를 만드는 작업이 아래와 같이 마무리되었다.

 

 

<입력 Sequence - 인코딩 됨>

 2      4    523  8 312   1   53 5234   4   323  123   3    21  33  22   21  123   3


<입력 Mask> (Padding을 마스킹하기 위해서)

1 ... 1 0 ... 0 (뒤에 0들은 padding 된 index들)


<마스킹 된 토큰의 라벨 - 인코딩 됨>

3444, 553


<마스킹 된 토큰의 위치>

1, 8


<Random Nex>

1 (True)


<Sequence ID> (해당 토큰이 Seq_A에 속하는지 Seq_B에 속하는지)

0 ... 0 1 ... 1

 

이제 위 Training Instance를 사용해서 BERT를 Pretraining 시킬 차례다.

 

 

1. Embedding

 

BERT를 Pretraining 시키기 전에 Encoding 한 토큰을 Embedding 해야 된다. 슬슬 헷갈리기 시작한다.. Encoding? Embedding?

 

한 토큰을 구분되는 Index로 1편에서 바꿔주었다. 하지만 각 단어가 고유한 Index를 가지기 때문에 (One-hot Vector과 비슷하게), 단어 사이에 의미의 유사성이 없다. 이는 사람이 문맥을 파악하는 방법과 상이하다. 사람은 단어 아버지와 어머니 사이에 모종의 연결고리가 있다는 것을 알고 있지 않는가? 이렇게 단어를 Encoding 만 하게 되면 모델로써는 두 단어가 어떤 관계에 있는지 알 도리가 없다. 이 문제를 해결하기 위해 나온 방법이 Embedding이다. 이외에도 Embedding을 하는 이유는 굉장히 많지만 본 글에서는 이를 설명하는 것이 목적이 아니기에 좀 더 자세한 설명은 여기를 참고하자.

 

BERT 는 세 종류의 Embedding을 거치는데 각각 아래와 같다. 이해를 돕기 위해 Batch Size = 8 , Sequence Length = 128 인 상황을 가정하겠다. 여기서 Sequence Length란 [CLS] Seq_A [SEP] Seq_B [SEP]에 있는 토큰들의 개수를 말한다. 또한 hidden_layer_size = embedding size = 768이라고 하겠다. 이는 한 토큰을 Embedding 할 때 길이가 768인 벡터로 Embedding 하겠다는 뜻이다.

 

1.1 Word Embedding  : Word Index Scalar> Embedding Vector

torch.Size([8, 128]) > torch.Size([8, 128, 768])  / Layer Size : (32200, 768)

 

토큰의 Index (BERT의 경우 0 이상 32199 이하의 정수)가 길이가 768인 Vector로 바뀌는 Embedding 이다. 따라서 스칼라가 벡터로 변환되었다. 물론 좀 더 자세히 보면, 스칼라가 One-hot Vecotr로 바뀌고, 그다음에 길이 768인 Vector로 전환되겠지만 말이다.

 

 

1.2 Position Embedding  : Position Index Scalar > Embedding Vector

torch.Size([8, 128]) > torch.Size([8, 128, 768]) / Layer Size : (128, 768)

 

BERT의 Input에는 해당 토큰의 위치가 어디인지에 대한 정보가 없다. 하지만 사람이 문장을 해석할 때 단어 위치는 굉장히 중요한 지표이므로, 이 정보를 추가해주어야 한다. 위 예시에서 Position Index Scalara 은 0~127 사이의 스칼라 값을 가질 것이다.

 

 

1.3 Token Type Embedding  : Token Type Scalar > Embedding Vector

torch.Size([8, 128]) > torch.Size([8, 128, 768]) / Layer Size : (2, 768)

 

여기서 토큰 타입이랑 해당 토큰이 Seq_A 에 속하는지, Seq_B에 속하는지 에 대한 정보이다. A에 속하면 0, B에 속하면 1의 값을 가질 것이다. 따라서 이 또한 스칼라에서 벡터로의 Embedding이다.

 

 

여기서 중요한 것은 세 Embedding 모두 Layer이라는 것이다. 즉, Model이 학습할 때 그 대상이라는 것이다. PyTorch에서는 torch.nn.Embedding 모듈을 사용하며 아래 코드에서 확인할 수 있다. torch.nn.Embedding의 모델 구조는 FFNN (Fully Connected) 이다. 이를 확인했던 방법은, 공식 문서를 봐도 되지만 그냥 Paramemter 개수를 확인했더니.. 그랬다. Embedding 의 전체 코드는 아래와 같다.

 

import torch.nn as nn


class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings.
    """
    
    def __init__(self, config):
        super(BertEmbeddings, self).__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)  # (32200, 768)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)  # (128, 768)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)  # (2, 768)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)  # 0.1


    def forward(self, input_ids, token_type_ids):
        seq_length = input_ids.size(1)  # 128
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) # torch.Size([128])
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids) # torch.Size([8, 128])

        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = words_embeddings + position_embeddings + token_type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings  # (batchSize, sequenceLength, hidden_size)

 

요약하면 아래와 같다.

Embedding Vector = Dropout(Normalize(Word Embedding + Position Embedding + Token Type Embedding))

 

 

 

2. BERT Layer

 

 

아시다시피 BERT는 Transformer의 Encoder 를 직렬로 이은 구조를 갖는다. 앞서 Embedding에서 봤던 것이 빨간색 테두리 아래의 부분이고, 이제는 BERT 모델의 핵심인 빨간색 테두리 부분을 보자. 빨간 테두리 부분을 BERT Block (Encoder Block), 혹은 BERT Layer 이라고 하자. 한 Bert Layer은 세 가지 Layer로 구성되어 있다.

 

 

2.1 BERT Attention

 

    2.1.1 Bert Self Attention (hidden_state, attention_mask)

 

class BertSelfAttention(nn.Module):

    def __init__(self, config):
        super(BertSelfAttention, self).__init__()
        # hidden_size = 768, num_attention_heads = 12
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
        self.num_attention_heads = config.num_attention_heads  # 12
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)  # 64
        self.all_head_size = self.num_attention_heads * self.attention_head_size  # 768

        self.query = nn.Linear(config.hidden_size, self.all_head_size)  # (768, 768)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)  # (768, 768)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)  # (768, 768)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)  # 0.1


    def transpose_for_scores(self, x):
        #  (batch_size, seq_length, num_attention_heads, attention_head_size)
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        # Multi-Head Self Attention
        # (batch_size, seq_length, all_head_size) -> (batch_size, seq_length, num_attention_heads, attention_head_size)
        x = torch.reshape(x, new_x_shape)
        
        return x.permute(0, 2, 1, 3)  # (batch_size, 12, 128, 64)


    def transpose_key_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = torch.reshape(x, new_x_shape)
        
        return x.permute(0, 2, 3, 1)   # (batch_size, 12, 64, 128)


    def forward(self, hidden_states, attention_mask):
        mixed_query_layer = self.query(hidden_states)  # (batchSize, 128, 768)
        mixed_key_layer = self.key(hidden_states)  # (batchSize, 128, 768)
        mixed_value_layer = self.value(hidden_states)  # (batchSize, 128, 768)

        query_layer = self.transpose_for_scores(mixed_query_layer)  # (batch_size, 12, 128, 64)
        key_layer = self.transpose_key_for_scores(mixed_key_layer)  # (batch_size, 12, 64, 128)
        value_layer = self.transpose_for_scores(mixed_value_layer)  # (batch_size, 12, 128, 64)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer)  # (batch_size, 12, 128, 128)
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        attention_scores = attention_scores + attention_mask 

        # Normalize the attention scores to probabilities.
        attention_probs = F.softmax(attention_scores, dim=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs, value_layer)  # (batch_size, 12, 128, 64)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()  # (batch_size, 128, 12, 64)
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = torch.reshape(context_layer, new_context_layer_shape)  # (batch_size, 128, 768)
        
        return context_layer

 

전체 코드는 위와 같다. 각각의 토큰에 대해서 Query, Key, Value를 만들고 Self Attention을 수행한다. Multi Head Attentoin이기 때문에 768길이 Query, Key, Value를 12개의 64 길이로 나누는 것을 볼 수 있다.

 

 

 

    2.1.2 Bert Self Output (context, hidden_state)

 

class BertSelfOutput(nn.Module):

    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)  # (768, 768)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)


    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)  # (batch_size, 128, 768)
        hidden_states = self.dropout(hidden_states)  # (batch_size, 128, 768)
        hidden_states = self.LayerNorm(hidden_states + input_tensor) # Skip Connection
		
        return hidden_states

 

전체 코드는 위와 같다. Self Attention Layer을 거쳐 나온 Vecotr에 FFNN을 한 번 거치고, Skip Connection을 적용하는 것을 볼 수 있다.

 

 

 

2.2 BERT Intermediate

class BertIntermediate(nn.Module):

    def __init__(self, config):
        super(BertIntermediate, self).__init__()
        # (8, 128, 768) -> (8, 128, 3072)
        self.dense_act = LinearActivation(config.hidden_size, config.intermediate_size, act=config.hidden_act)

    def forward(self, hidden_states):
        hidden_states = self.dense_act(hidden_states)
        
        return hidden_states

 

Embedding Size 가 768이었던 Vecotr을 3072 (4배)로 늘려주는 FFNN을 거친다. Activation Function으로는 gelu를 사용하였다.

 

 

 

2.3 BERT Output

class BertOutput(nn.Module):

    def __init__(self, config):
        super(BertOutput, self).__init__()
        # (8, 128, 3072) -> (8, 128, 768)
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        
        return hidden_states  # (8, 128, 768)

 

이전에 3072로 늘어났던 Size를 다시 768로 바꾸고, Dropout & Normalize를 적용해준다. 여기서 나온 Output이 다음 Bert Layer의 Input으로 들어가게 된다. (다음 블록의 BERT Self Attention)

 

 

 

 

3. BERT Pooler

 

BERT 모델의 최종 Output으로 2가지를 내놓는다.

 

1) tanh(마지막 hidden_state의 첫 토큰) -> -1~1 사이의 Scala 값 / Result : (8, 2)

2) 마지막 hidden_state / Result : (8, 128, 32200)

 

1) 번으로는 A와 B가 연속된 Sequence 인지 (Random Next) 판단하고,

2) 번으로는 마스킹된 단어를 예측하게 된다.

 

2) 번에서 길이 768 Vector을 32200 Vector로 바꾸게 되는데, 이때 layer으로는 Word Embedding Layer을 쓴다. 절묘하게 맞아떨어지지 않는가? 

 

---

BERT 소개를 마치며

 

2편이 1편에 비해서 상당히 저퀄이라 좀 아쉽긴 하지만, 글로 전달한다는 것이 참 어렵다는 것을 다시 한번 느꼈다. 기회가 된다면 누군가에게 세미나를 해주고 싶다. 꼭 공유하고 싶은 내용이다.