[코드리뷰]BERT

이번 논문리뷰부터는 코드리뷰도 같이 진행하기로 하였다. 이번 코드리뷰는 huggingface의 transformers 패키지 내의 구현된 코드를 분석하는 것으로 하기로 하였다. 코드의 주소는 다음과 같다.

https://github.com/huggingface/transformers/blob/master/src/transformers/models/bert/modeling_bert.py

GitHub - huggingface/transformers: 🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

코드의 길이는 무려 1876줄이다.. 이렇게 긴 코드 분석은 처음인 것 같다. 코드를 어떻게 분석할까 고민하다가, 먼저 구역을 명확히 나누는 것을 골자로 분석해보기로 하였다.

1줄 ~ 16줄 : 라이센스 및 주석

19줄 ~ 57줄 : import문

60줄 ~ 90줄 : 전역변수 선언

93줄 ~ 163줄 : 함수 load_tf_weights_in_bert 정의

166줄 ~ 223줄 : 클래스 BertEmbeddings 정의

226줄 ~ 350줄 : 클래스 BertSelfAttention 정의

353줄 ~ 364줄 : 클래스 BertSelfOutput 정의

367줄 ~ 413줄 : 클래스 BertAttention 정의

416줄 ~ 428줄 : 클래스 BertIntermediate 정의

431줄 ~ 442줄 : 클래스 BertOutput 정의

445줄 ~ 527줄 : 클래스 BertLayer 정의

530줄 ~ 624줄 : 클래스 BertEncoder 정의

627줄 ~ 639줄 : 클래스 BertPooler 정의

642줄 ~ 656줄 : 클래스 BertPredictionHeadTransform 정의

659줄 ~ 676줄 : 클래스 BertLMPredictionHead 정의

679줄 ~ 686줄 : 클래스 BertOnlyMLMHead 정의

689줄 ~ 696줄 : 클래스 BertOnlyNSPHead 정의

699줄 ~ 708줄 : 클래스 BertPreTrainingHeads 정의

711줄 ~ 741줄 : 클래스 BertPreTrainedModel 정의

744줄 ~ 775줄 : 클래스 BertForPreTrainingOutput 정의

778줄 ~ 841줄 : 전역변수 DOCSTRING 선언

844줄 ~ 1021줄 : 클래스 BertModel 선언

1024줄 ~ 1129줄 : 클래스 BertForPreTraining 선언

1132줄 ~ 1282줄 : 클래스 BertLMHeadModel 선언

1285줄 ~ 1390줄 : 클래스 BertForMaskedLM 선언

1393줄 ~ 1491줄 : 클래스 BertForNextSentencePrediction 선언

1494줄 ~ 1593줄 : 클래스 BertForSequenceClassification 선언

1596줄 ~ 1688줄 : 클래스 BertForMultipleChoice 선언

1691줄 ~ 1773줄 : 클래스 BertForTokenClassification 선언

1776줄 ~ 1876줄 : 클래스 BertForQuestionAnswering 선언

클래스가 무려 25개이다. 비록 완벽하게 분석하지는 못할 것 같지만, 최대한 분석해보는 것을 목표로 하자..

 

1줄 ~ 16줄 : 라이센스 및 주석

라이센스에 관련된 주석부분이다. 이 파일을 사용하려면 라이센스를 잘 준수하라는 문구가 적혀있다. 그리고 이 파일은 PyTorch로 구현된 BERT model이라 적혀있다.

19줄 ~ 57줄 : import문

 

import warnings : 파이썬 경고문구와 관련된 라이브러리이다.

from dataclasses import dataclass : dataclass를 import한다. dataclass는 클래스를 마치 내장 자료 구조(리스트, 튜플 등)처럼 편리한 클래스로 만들어준다. 자세한 설명은 다음 링크를 참조한다. [출처] : https://www.daleseo.com/python-dataclasses/

[파이썬] 데이터 클래스 사용법 (dataclasses 모듈)

from typing import Optional, Tuple : 타입 힌트(type hint)를 지원하기 위한 라이브러리이다. 자세한 설명은 다음 링크를 참조한다. [출처] : https://www.daleseo.com/python-typing/

[파이썬] typing 모듈로 타입 표시하기

여기서 from ~ import 구문에서 .이 붙어있는데, 이는 상위 디렉토리 개수를 의미한다.(상대 경로) 정확히는 .이 하나면 현재 파일과 같은 디렉토리, .이 두개면 부모 디렉토리, 그리고 위로 올라갈수록 .이 하나씩 증가하는 것이다.

 

from ...activations import ACT2FN : activations.py에 있는 전역변수 딕셔너리 "ACT2FN"을 import하는 것이다. ACT2FN에는 activations.py에 선언되어 있는 활성화함수들이 담겨있는 딕셔너리이다. 파이썬의 함수는 일급 객체이므로 이렇게 모아놓는 것이 가능하다.

from ...file_utils import (ModelOutput, ...) : file_utils.py에서 여러 클래스를 import하는 것이다. 대략적으로 하는 역할들은 docstring 작성에 관련된 것들이며, ModelOutput은 모든 모델의 output을 dataclass로 변환하는 역할을 한다.

from ...modeling_outputs import (BaseModelOutputWithPastAndCrossAttentions, ...) : modeling_outputs.py에서 각종 모델에 관련된 output class들을 import한다.

from ...modeling_utils import (PreTrainedModel, ...) : modeling_utils.py에서 각 모델에 공통적으로 사용될 클래스들을 import한다.

from ...utils import logging : utils 디렉토리의 logging.py에서 로깅에 관련된 것들을 import한다.

from .configuration_bert import BertConfig : configuration_bert.py에서 BERT 모델의 configuration을 저장하는 클래스인 BertConfig를 import한다.

60줄 ~ 90줄 : 전역변수 선언

 

logger는 이 파일에서 logging을 하는 역할을 한다. _FOR_DOC 전역변수들은 코드에서 docstring에 들어갈 변수들이다. BERT_PRETRAINED_MODEL_ARCHIVE_LIST는 이 파일에서 따로 쓰이지는 않고 다른 파일에서 쓰이는 리스트인 것 같다.

93줄 ~ 163줄 : 함수 load_tf_weights_in_bert 정의

pytorch 모델에다가 tensorflow의 checkpoint를 load하는 함수이다. 자세한 내용은 패스함.

166줄 ~ 223줄 : 클래스 BertEmbeddings 정의

word, position로부터 embedding을 만들고, token_type embedding을 만드는 클래스이다. nn.Module을 상속받는다.

self.word_embeddings, self.position_embeddings, self.token_type_embeddings는 이름 그대로 word, position, token_type에 관한 embedding layer이다. nn.Embedding()은 embedding table이라 생각하면 되며, train시 같이 훈련되는 테이블이다.

1번째 인자는 num_embeddings로, embedding_dictionary의 size가 들어온다. (word라면 vocabulary의 크기)

2번째 인자는 embedding_dim으로, 임베딩 된 dense vector의 dimension을 의미한다.

3번째 인자(Optional)는 padding_idx로 훈련 시 gradient에 영향을 주지 않을 padding이 쓰인 index를 넣어주면 된다.

self.position_embedding_type에는 getattr()이 쓰였는데, 여기서 getattr(object, "name")은 object라는 오브젝트 내부의 "name"이라는 멤버를 반환하는 것을 의미한다. 그러므로, getattr(config, "position_embedding_type", "absolute")는 config.position_embedding_type이랑 같은 의미이며, "absolute"는 만약 해당 멤버가 존재하지 않을 경우에 사용되는 default값이다.

register_buffer는 훈련 시 훈련되지 않을 파라미터(layer)를 등록할 때 사용한다. "position_ids"라는 이름으로 [[0, 1, 2, ..., config.max_position_embeddings - 1]] 텐서를 등록한다. (논문을 살펴보면 position embedding 자리는 0부터 차례대로 숫자가 1씩 증가하는 형태였다)

그리고 torch의 version이 1.6.0보다 최신이면 위에서 만든 position_ids 모양의 텐서에 0으로 채운 텐서([[0, 0, 0, ..., 0]])을 "token_type_ids"라는 이름으로 등록한다.

먼저, input_ids는 (batch_size, sequence_length)의 shape을 가지며, vocabulary내 token의 index들이다. input_embeds는 (batch_size, sequence_length, hidden_size)의 shape을 가지며, input_ids 대신에 사용 가능하다. (이걸 사용 시 앞에서 정의한 word_embedding layer를 사용 안함)

position_ids는 (batch_size, sequence_length)의 shape을 가지며, position embedding에 해당한다. (이걸 사용 시 앞에서 정의한 self.position_ids[0, 1, 2, ..., config.max_embeddings - 1]에서 뽑아와서 사용한다.)

token_type_ids는 (batch_size, sequence_length)의 shape을 가지며, segment embedding에 해당한다. sentence A는 0, sentence B는 1로 나타낸다. (만약 인자로 받지 못했다면, 위에서 정의한 self.token_type_ids인 [0, 0, ..., 0]을 사용한다.)

이 코드를 요약하자면, 사용자가 인자로 input_ids(word embedding), position_ids(position embedding), token_type_ids(segment embedding)을 넘겨서 사용할 수 있고, 넘기지 않았다면 config에서 구성한 값으로 사용된다. 그리고 각각 이 텐서들은 embedding layer를 거치며, embeddings라는 변수에 더해지게 된다. 이때, embeddings의 shape은 [batch_size, sequence_length, config.hidden_size]가 될것이다.

226줄 ~ 350줄 : 클래스 BertSelfAttention 정의

Bert 내에서 self-attention을 하는 클래스이다. nn.Module을 상속받는다.

먼저, config의 hidden_size가 num_attention_heads(head 숫자)로 나누어 떨어지는 지 확인한다. 이게 나누어 떨어지지 않으면, 후에 Multi-Headed Attention 구조를 만들 수 없기 때문에 여기서 검사한다.

self.num_attention_heads는 attention_head의 개수이다.

self.attention_head_size는 1개 head의 hidden_size를 의미한다.

self.all_head_size는 분리하기 전 원래 hidden_size를 의미한다.

self.query, self.key, self.value는 논문에 나온 Q, K, V를 의미한다. shape은 [config.hidden_size, self.all_head_size]이다.

self.position_embedding_type은 position_embedding의 방식이다. 각각 'Absolute', 'relative_key', 'relative_key_query' 세 종류가 있는데, 이 중 BERT 논문에서 나온 방식은 'Absolute'이다.

self.is_decoder는 이 attention을 decoder에서처럼 동작하게 할 것(단방향)인지의 여부이다. False이면 encoder처럼 동작한다. (양방향)

이 함수는 score 계산을 위해서 전치를 하는 함수이다.

mixed_query_layer에는 hidden_states (batch_size, sequence_length, config.hidden_size)가 들어오게 되며, 결과 output의 shape은 (batch_size, sequence_length, self.all_head_size)가 된다.

is_cross_attention은 cross-attention(attention의 종류 중 하나인듯 함. 잘 모르겠음)인지 아닌지 여부이다.

결과적으로, query_layer, key_layer, value_layer의 shape은 [batch_size, self.num_attention_heads, sequence_length, self.attention_head_size]가 된다.

논문의 공식대로 attention_score를 QK^T로 계산하였다. shape은 [batch_size, self.num_attention_heads, sequence_length, sequence_length]이다.

만약, attention_mask가 있다면(decoder에서 사용) attention_scores에 더함으로써 마스킹을 한다.

optional하게는 head에도 mask를 할 수 있다.

그 다음에는 attention_probs에 V를 matmul 함으로써, Attention(Q, K, V)인 context_layer를 구해내게 된다.

그 다음에는 context_layer의 shape을 다시 [batch_size, sequence_length, self.num_attention_heads,  self.attention_head_size]로 바꿔준다. 그리고 multi-head를 concat함으로써 [batch_size, sequence_length, self.all_head_size(config.hidden_size)]로 바꿔주게 된다.

만약, output_attentions가 True이면 attention_probs도 같이 반환해주고, 아니라면 context_layer만 반환한다. 만약, decoder라면 past_key_value도 같이 return해준다.

353줄 ~ 364줄 : 클래스 BertSelfOutput 정의

Bert내에서 self-attention의 output을 받아 추가적인 layer를 거친 후에 output을 내주는 부분이다.

dense layer, layer normalization, dropout layer 정도를 거치게 된다.

367줄 ~ 413줄 : 클래스 BertAttention 정의

앞에서 정의한 BertSelfAttention과 BertSelfOutput을 통해 전체적인 Attention을 구현한 클래스이다.

 

여기서 self.pruned_heads라는 것이 있는데, 실제 Transformer의 attention의 multi-head 구조에서 일부 head를 제거(pruning)해도 성능에 지장이 없거나 오히려 상승하였다는 논문이 있다. 그것을 구현한 듯 하다.

역시 pruning에 대한 것으로, linear_layer에 대해서도 pruning을 진행하는 것 같다. pruning에 관한 것은 배운 것이 없어서, 생략한다.

결국 이 클래스는 원래 구현한 SelfAttention에 pruning을 추가시킨 것뿐이다.

416줄 ~ 428줄 : 클래스 BertIntermediate 정의

Attention Output 이후에 있는 중간 layer이다.

ACT2FN은 미리 activations을 정의해놓은 dictionary이다. config.hidden_act는 hidden layer에 사용할 activation의 이름이다.

431줄 ~ 442줄 : 클래스 BertOutput 정의

중간 layer 이후에 오는 output layer이다.

여기 hidden_states는 input_tensor와 더하고 Layer Normalization을 한다. Add & Norm을 구현하였다.

445줄 ~ 527줄 : 클래스 BertLayer 정의

여기서 decoder에서 사용하는 것이 아니면, cross_attention은 사용할 수 없다.