* 논문 : https://arxiv.org/pdf/2004.02984.pdf
* 깃헙 : https://github.com/google-research/google-research/tree/master/mobilebert
개요
성능은 좋지만 무거운 BERT_large, 다이어트 시켜보자!
- Bottleneck 구조 도입
- Self attention과 FFN 사이의 밸런스를 맞춤
- 그 결과, 깊이는 BERT_large만큼 <깊지만>, 레이어는 <얇아진> 버전인 MobileBERT가 탄생
MobileBERT는 원래 모델과 마찬가지로 fine-tuning을 통해 어떠한 NLP 태스크에도 적용 가능하다.
성능 & 모델 사이즈 측면에서
- BERT_base보다 4.3배 작으면서 5.5배 빠른 모델을 확보
- GLUE 태스크에서 BERT_base 모델 대비 성능 하락은 0.6에 불과
- Pixel 4 모바일폰에서 63ms의 latency로 추론이 가능
- SQuAD에 있어서는 심지어 BERT_base보다 높은 성적인 EM=79.2/F1=90.0 을 달성하였다
모델 개요
- 먼저 teacher 모델로 inverted-bottleneck을 적용한 BERT_large 모델을 학습 (IB-BERT)
- 이후 IB-BERT의 지식을 얇고 긴 네트워크인 MobileBERT에 증류(Knowledge Distillation)
모델 아키텍처
[ Bottleneck & Inverted-Bottleneck ]
Bottleneck 구조
- 위의 표에서 보이는 것과 같이 MobileBERT는 BERT_large와 마찬가지로 24개의 레이어를 가지진다.
- 하지만 [MHA + FFN] 에서 사용하는 hidden 차원은 128차원에 불과하다.
- 여기서 [MHA + FFN]의 인풋과 아웃풋은 512차원을 가질 수 있도록 하기 위해 앞뒤로 Linear 레이어를 추가한다.
하지만 이와 같이 깊으면서 얇은 뉴럴넷을 학습시키는 것은 어려운 과제이다.
학습이 잘 이뤄질 수 있도록 하기 위해 논문에서는 선생님 모델을 학습시킨 후 knowledge distillation을 수행한다.
Inverted Bottleneck 구조
선생님 모델은 BERT_large에서 사용하는 [MHA + FFN]의 hidden 차원과 head 수를 가지지만,
MobileBERT에게 지식을 증류할 수 있도록 인풋 & 아웃풋을 512로 만드는 2개의 레이어를 추가한다.
Bottelneck과 반대로 인풋이 핵심 모듈의 차원보다 작기 때문에 이를 Inverted Botteleneck이라 부른다.
이러한 구조를 도입함으로써 IB-BERT는 BERT_large가 캐리할 수 있는 만큼의 성능을 유지하고,
이러한 모델의 body에서 나오는 아웃풋을 MobileBERT에게 학습시킴으로써 모델을 충분히 컴팩트하게 만들 수 있다.
[ Stacked Feed-Forward Networks ]
Bottleneck 구조의를 도입할 때 단점은 MHA 모듈과 FFN 모듈 사이의 균형이 무너진다는 점이다.
Transformer 구조에서 MHA와 FFN는 각각 다음과 같은 역할을 수행하고 있었다:
- MHA: 각기 다른 공간에서 오는 정보에 동시적으로 attend
- FFN : 모델의 비선형성을 증가시킴
BERT 모델에서 MHA와 FFN의 파라미터 개수의 비율은 항상 1:2로 유지되었다.
하지만 bottleneck 구조를 도입하게 되면, MHA의 인풋은 더 넓은 feature map(inter-block 크기)에서 오는 반면
FFN의 인풋은 더 얇은 bottleneck(intra-block 크기에 해당)에서 온다.
즉, 이 그림을 다시 보면
- MHA는 inter-block flow에 따라
. 512차원을 인풋으로 받아 4개의 attention head 적용
. 그 결과 128차원의 아웃풋 생성
- FFN은 intra-block flow에 따라
. MHA가 리턴한 128차원을 인풋으로 받아
. 512 차원의 intermediate hidden을 거쳐 128차원을 아웃풋
이렇게 되면, MHA 모듈이 상대적으로 더 많은 파라미터를 보유하게 된다.
이 이슈를 해결하기 위해 논문에서는 FFN을 여러 층으로 쌓는 stacked FFN을 도입한다.
실험에 따라, 각 MHA마다 4층의 FFN을 쌓았다.
[ Operational Optimization ]
모델의 latency를 분석한 결과, layer normalization과 gelu activation가 latency의 상당 부분을 차지하였다.
따라서, MobileBERT는 이 둘을 다음과 같은 새로운 연산으로 대체한다.
1. Layer Normalization 제거
n-channel의 hidden state h에 대한 레이어 정규화는 element-wise linear transformation으로 대체한다
이때 NoNorm은 추론시에도 LayerNorm과는 다른 성질을 가진다.
원래 제안된 레이어정규화는 벡터들의 배치에 대한 선형 연산은 아니기 때문이다.
2. ReLU activation 사용
gelu보다 간단한 연산인 relu를 사용한다.
[ Embedding Factorization ]
임베딩 부분은 모델 크기에서 상당 부분을 차지하기 때문에, MobileBERT에서는 임베딩 크기를 128로 설정한다.
이후 raw token embedding에 대해 kerner_size = 3인 1D convolution을 적용하여 512차원의 아웃풋을 생성한다.
Google research 깃허브 확인 >>
trigram_input=True -> trigram에 대한 convolution 결과를 인풋으로 사용하도록 하는 파라미터
## modeling.py line 218~
with tf.variable_scope("embeddings"):
# Perform embedding lookup on the word ids.
(self.word_embedding_output, self.embedding_table) = embedding_lookup(
input_ids=input_ids,
vocab_size=config.vocab_size,
embedding_size=config.embedding_size,
initializer_range=config.initializer_range,
word_embedding_name="word_embeddings",
use_one_hot_embeddings=use_one_hot_embeddings)
if config.trigram_input:
inputs = self.word_embedding_output
self.embedding_output = tf.concat(
[tf.pad(inputs[:, 1:], ((0, 0), (0, 1), (0, 0))),
inputs,
tf.pad(inputs[:, :-1], ((0, 0), (1, 0), (0, 0)))],
axis=2)
else:
self.embedding_output = self.word_embedding_output
if (config.trigram_input or
config.embedding_size != config.hidden_size):
self.embedding_output = dense_layer_2d(
self.embedding_output,
config.hidden_size,
create_initializer(config.initializer_range),
None,
name="embedding_transformation",
use_einsum=use_einsum)
# Add positional embeddings and token type embeddings, then layer
# normalize and perform dropout.
self.embedding_output = embedding_postprocessor(
input_tensor=self.embedding_output,
use_token_type=True,
token_type_ids=token_type_ids,
token_type_vocab_size=config.type_vocab_size,
token_type_embedding_name="token_type_embeddings",
use_position_embeddings=True,
position_embedding_name="position_embeddings",
initializer_range=config.initializer_range,
max_position_embeddings=config.max_position_embeddings,
dropout_prob=config.hidden_dropout_prob,
normalization_type=config.normalization_type)
학습 목적함수
MobileBERT 학습을 위해 2가지 KD 목적함수: 1) feature map 2) attention transfer 를 사용한다.
즉, l번째 레이어에서 레이어 간의 지식 증류 loss는 다음의 두 가지의 선형 결합으로 정의한다.
[ Feature Map Transfer , FMT ]
- BERT의 각 레이어는 이전 레이어의 아웃풋을 인풋으로 취한다.
- 따라서 각 레이어가 반환하는 feature map이 선생님 모델과 비슷한 것이 가장 중요하다.
- 이에 선생님 모델과 MobileBERT의 feature map 사이의 mean squared error loss를 KD 목적함수로 사용한다.
- 이때 손실함수를 정규화된 feature map 사이의 불일치와 feature map의 통계량에 대한 불일치에 대한 항으로 분해하여 학습하는 것이 안적적인 학습에 도움이 된다.
[ Attention Transfer , AT ]
- 각각의 attention head에서 선생님 모델과 MobileBERT의 self attention 분포간의 KL-divergence를 활용
- 각 레이어에서 attention transfer를 위해 선생님 모델 & 학생 모델 모두 attention head 개수를 4로 설정
[ Pre-training Distillation , PD ]
- 레이어간에 지식을 증류할 뿐만 아니라 사전학습을 수행할 때에도 distillation loss를 활용한다.
- 즉, 원래 MLM & NLP loss에 더해 새로운 MLM Knowledge distillation loss를 선형 결합하여 사용한다.
학습 전략
위의 loss들을 활용한다고 할 때, 학습 중 다양한 전략을 사용할 수 있다.
본 논문에서는 다음의 세 가지 전략을 고려한다.
[ Auxiliary Knowledge Transfer ]
- 중간에 있는 지식을 transfer 하는 것은 KD에 있어 보조적인(auxiliary) 태스크라고 간주한다.
- 이에 모든 레이어에서 정의되는 지식 transfer loss와 사전학습 distillation loss를 선형 결합한 단일 loss를 사용한다.
[ Joint Knowledge Transfer ]
- IB-BERT 선생님 모델의 중간 레이어 지식은 MobileBERT에 있어 최적의 해법이 아닐 수 있다.
- 이에 두 loss 항을 분리하여 먼저 layer 단의 지식을 전이학습한 후 사전학습 distillation을 수행한다.
[ Progressive Knowledge Transfer ]
- MobileBERT가 IB-BERT 선생님 모델을 완벽하게 따라하지 못할 수 있고, 특히 밑단에서 잘못된 결과는 상위 레이어의 지식 증류를 방해할 수 있다.
- 따라서 각 레이어를 점진적으로 학습하는 방법을 제안한다.
- 이 방법에서는 transfer를 L 단계로 나누어 진행한다. (L은 레이어 개수와 동일)
실험 결과
[ Model Setting ]
IB-BERT 선생님 모델과 MobileBERT 모델에 있어 최적의 셋팅을 찾기 위해 SQuAD 데이터셋을 활용한다.
2048 배치 사이즈고 125k 스탭동안 (원래 BERT 모델의 1/2 수준) 학습한 결과를 사용해 분석하였다.
IB-BERT 구조 탐색
선생님 모델은 정확도 하락이 없는 한도 내에서 inter-block hidden size(=feature map 크기)를 작게 유지하고자 한다.
- (a) - (e) 결과로 볼 때, inter-block 크기를 512까지 줄여도 정확도는 유지되었다. 따라서 IB-BERT_large의 inter-block 히든 사이즈는 512로 설정한다.
- (f) - (i) 결과를 볼 때, intra-block hidden size를 줄이면 성능이 급격하게 감소한다. 이 결과로 볼 때, non-linear 모듈이 가지는 representation power라고도 볼 수 있는 intra-block은 BERT에서 중요한 역할을 담당하며, 따라서 선생님 모델에서 이 차원 수를 줄일 수 없다고 결론지었다.
- (a)와 (f) 결과를 비교해볼 때, atention head 개수를 16개에서 4개로 줄이는 것은 성능을 저하시키지 않았다.
MobileBERT의 구조 탐색
- BERT_base보다 4배 작은 모델로 모델을 압축하기 위해, MobileBERT는 25M의 파라미터를 가지도록 한다.
- 이때 MHA와 FFN에 사용되는 파라미터의 비율을 조절하는 실험을 통해 최적의 구조를 탐색해 보았다.
- 그 결과, MHA와 FFN의 파라미터 비율은 0.4~0.6일 때 가장 좋았고, 이는 원래 Transformer 모델이 선택한 0.5라는 비율과 유사하다.
- 논문에서는 intra-block 크기 128에 FFN을 4층 쌓은 모델을 선택하였다. 이는 정확도와 학습 효율성을 고려한 구조이다.
- 또한 레이어간의 지식 증류를 위해 선생님 모델 & 학생 모델 모두 attention head 개수는 4로 설정한다.
[ Implementation Details ]
- 사전학습 데이터는BERT와 마찬가지로 BookCorpus, 영문 위키피디아 문서 사용
- IB-BERT_large 성능을 원래 BERT_large 성능과 동일하게 맞추기 위해 256 TPU v3 chips를 사용해 배치 사이즈 4,096으로 500k 스텝 학습 & LAMB optimizer 사용
- MobileBERT에 대한 사전학습 distillation 단계에서도 같은 학습 스케줄링 사용
- progressive knowledge transfer는 24개 레이어에 대해 240k 스텝씩 진행
- ablation study에서는 사전학습 스텝을 1/2만 진행함 (빠른 성능 측정을 위해)
- 다운스트림 태스크는 MobileBERT를 fine-tuning한 결과를 제시 (BERT를 fine-tuning하는 것처럼!)
- MobileBERT는 일반 BERT에 비해 큰 learning rate로 더 오래 fine-tuning해야 했음.
GLUE 태스크 결과
- OPT는 Operational Optimization 의미 (GeLU-> ReLU / NoNorm) : OPT하면 latency가 크게 감소하나, 약간의 성능 trade-off가 있음. OPT 하지 않으면 BERT_base에서 성능 하락이 없는 수준.
SQuAD 결과
- MobileBERT는 사이즈가 작으면서도 다른 모델의 성능을 크게 앞선 것을 볼 수 있음
Quantization
- TF Lite를 사용하여 post-training quantization을 진행해봄.
- 이를 통해 MobileBERT를 추가적으로 4배 압축하면서도 성능에 저하가 생기지 않음.
[ Ablation Studies ]
Operational Optimization
- Operational Optimization을 수행하면, FLOPS는 감소하지 않지만 latency를 매우 효과적으로 감소시킴.
- 이론적으로 계산하는 computation overhead(=FLOPS)와 실제 추론 속도 사이에는 괴리가 있다.
Training Strategies
- Auxiliary knowledge transfer / joint knowledge transfer / progressive knowledge transfer 뭐가 좋지?
- Progressive Knowledge Transfer 방법이 다른 두 방법에 비해 좋음
- AKT가 가장 안 좋았는데, 최적의 중간 레이어의 knowledge(어텐션 맵, feature map 등)는 선생님 모델과 학생 모델에서 다를 수 있고, 따라서 pre-training distillation 단계를 추가적으로 거치는 것이 좋은 성능을 얻는 데에 유효하다.
Training Objectives
- Attention Transfer / Feature Map Transfer / Pre-training Distillation / Operational Optimization 효과 비교
- Feature Map Transfer 방법이 MobileBERT 성능에 가장 유효한 요인이었다.
- Attention Transfer와 Pre-training Distillation 역시 긍정적으로 작용한다.
- IB-BERT_large 선생님 모델은 원래 BERT_large 모델만큼 강력한 모델이나, MobileBERT는 선생님에 비해서는 성능이 저조하다. 이 부분을 개선하기 위한 추가적인 연구가 필요하다.
2020/11/08 - [AI] - 모델 경량화 - BERT 경량화 / 추론 속도 향상 기법 정리
모델 경량화 - BERT 경량화 / 추론 속도 향상 기법 정리
BERT는 뛰어난 성능과 간단한 fine-tuning 기법에도 불구하고 - 거대한 모델 사이즈 (파라미터 개수) - 느린 추론 속도 - 복잡하고 비용이 많이 드는 사전학습 과정 으로 인해 그 사용성에 대해 제한이
littlefoxdiary.tistory.com
'AI' 카테고리의 다른 글
| 인공지능의 시대, 튜링 테스트를 넘어선 새로운 비전이 필요하다 (2) | 2021.01.09 |
|---|---|
| [논문리뷰] Are Sixteen Heads Really Better than One? (0) | 2020.12.29 |
| [논문리뷰] Tinybert: Distilling bert for natural language understanding (0) | 2020.11.30 |
| 모델 경량화 - BERT 경량화 / 추론 속도 향상 기법 정리 (0) | 2020.11.08 |
| [논문리뷰] Small Language Models Are Also Few-Shot Learners (2) | 2020.10.09 |
