Extreme Language Model Compression with Optimal Subwords and Shared Projections

Abstract

현재 사전 훈련된 모델들(ELMo GPT BERT XLNet)이 다양한 언어 이해 과제에서 좋은 성적을 거두고 있음. 그러나 모델들이 너무 커서 모바일이나 Edge 기기에서 사용하기가 어려움. 그래서 나온 knowledge distillation이 있지만 원래 Teacher 모델과는 다른 어휘로 Student 모델을 산출하는데에는 효과가 없음. 이 논문에서 훨씬 더 작은 어휘와 낮은 임베딩, hidden state dimensions을 가진 student 모델을 훈련시키기 위해 새로운 knowledge distillation을 도입함.

Teacher, student 모델을 동시에 훈련시켜 student 어휘를 위한 최적의 워드 임베딩을 얻는 dualtrain 메커니증을 채용. teacher 모델에서 student 모델로 layer-wise knowledge를 전달하는 공유된 투영 매트릭스와 결합함.

We combine this approach with learning shared projection matrices that transfer layer-wise knowledge from the teacher model to the student model.

이렇게 해서 BERT_base 모델을 60배이상 압축할 수 있어 7MB 미만의 언어 모델을 만들었음.

Introduction

Abstract에서 얘기 나왔다시피 모바일 및 Edge 디바이스에서 실시간 추론 등을 수행하기에는 자원이 제한되어 있어서 불가능합니다. 그래서 여기서의 목표는 사전 훈련된 교사 모델로부터 soft predicted label distribution과 같은 지식을 자용하여 student 모델을 훈련 시키는 것 입니다.

The goal is to train a small-footprint student model by borrowing knowledge, such as through a soft predicted label distribution, from a larger pre-trained teacher model.

이전에도 이러한 노력이 있었지만 간과된 중요한 병목 현상은 입력 어휘 크기와 해당 워드 임베딩 메트릭스인데 모든 모델 파라미터의 상당한 비율을 차지합니다. (30K WordPiece 토큰으로 구성된 BERT_base 모델의 임베딩 테이블은 모델 크기의 21% 차지)

기존 distillation기법은 teacher, student 모델이 동일한 어휘와 출력 공간을 공유해야해서 NLP 모델 워드 크기를 줄이는데 있어서 제한이 많았습니다.

개선하기 위해서 다음의 두 아이디어를 제시

• Dual Training : teacher, student 모델은 같은 순서에 따라 다른 어휘와 호환되지 않는 토큰화를 가짐. Distillation 동안 단일 sequence에서 teacher, student 모델의 어휘 토큰 단어를 혼합하여 teacher 모델에 제공합니다. 마스크된 언어 모델 테스크와 함께 teacher, student의 WordPiece 임베딩을 암시적으로 정렬할 수 있습니다. 즉, student 워드 임베딩이 컨텍스트로 사용되어 teacher word에 의한 토큰화 단어를 예측할 수 있고 반대도 가능합니다

Dual Training: Our teacher and student models have different vocabularies and incompatible tokenizations for the same sequence. To address this during distillation, we feed the teacher model a mix of teacher vocabulary-tokenized and student vocabulary-tokenized words within a single sequence. Coupled with the masked language modeling task, this encourages an implicit alignment of the teacher and student WordPiece embeddings, since the student vocabulary embedding may be used as context to predict a word tokenized by the teacher vocabulary and vice versa.

• Shared Variable Projections : Hidden state dimension을 줄이는 걸로 정보 손실을 최소화 하기 위해 teacher, student 모델의 학습 가능한 변수를 정렬를 위한 별도의 loss를 도입하였습니다. 이로써 layer-wise하게 지식을 student 모델에 보다 직접적으로 전달이 가능하게 됩니다.

To minimize the loss of information from reducing the hidden state dimension, we introduce a separate loss to align the teacher and student models’ trainable variables. This allows for more direct layer-wise transfer of knowledge to the student model.

위 두 가지 방법을 조합하여 12-layer 고압축 student BERT 모델을 훈련시켰고, teacher BERT_base 모델과 비교하여 61배(with 48 dimension size) 압축됨. 일반화 된 언어 모델링 관점과 다운 스트림 작업을 모두 측정하기 위한 실험을 수행합니다..

Related work

NLP 압축 작업은 크게 4가지 정도 입니다.

• Matrix approximation : 낮은 순위 근사법으로 모델의 matrix 매개변수를 압축. 즉, full-rank matrix 변수는 여러 low-rank matrix를 사용하여 근사되므로 유효 모델 매개 변수 수가 줄어듭니다.
• Parameter pruning/sharing : 모델 매개변수의 중복성을 탐색하고 다양한 신경망 아키텍처에서 중복 가중치와 뉴런을 제거합니다.
• Weight quantization : 모델 가중치를 낮은 정밀도의 정부 및 부동 소수점 숫자에 매핑하는데 중점둡니다. 효율적인 저 정밀 계산(low-precision calculations)을 지원하는 하드웨어에서 특히 효과적입니다.
• Knowledge distillation : 소규모 student 모델은 큰 teacher 모델과 다르게 매개 변수화되어 더 많은 모델링 자유를 제공합니다. 큰 모델의 소프트 출력 레이블 분포에 맞게 student 모델을 교육하면 기계 번역 및 언어 모델링과 같은 많은 작업에 적합합니다.

Teacher, student 모델 출력에 국한되지 않고 attention transfer를 통한 knowledge distillation이라던지 feature map 또는 중간 모델 출력(intermediate model output)을 수행합니다. 또한 transformer layer 수를 줄임으로써 모델 압축을 할 수 있습니다

다만, teacher, student 모델 vocab이 호환되지 않기 때문에 즉시 적용이 어렵고 임베딩 matrix size 또한 충분하게 focus 하지 못하였습니다.

• Student vocab이 적은 BERT에 대한 knowledge distillation
  • 왼쪽 : 기본 BERT 매개 변수가 있는 사전 훈련된 teacher BERT 모델 (30K vocab, 768 hidden state dimension)
  • 오른쪽 : 더 작은 voab(5K)과 hidden state dimension(48)으로 처음 훈련된 student BERT
  • Distrillation 하는 동안 teacher 모델은 각 입력 단어를 분할하기 위해 vocab을 임의로 선택
  • Transformer layer 옆의 빨간색/녹색은 학습 가능한 매개변수

Methodology

이 논문에서는 WordPiece 토큰 수를 줄이는 데에 중점을 둡니다. 이 섹션에서는 축소의 근본 원인과 그에 따른 문제, Dual Training, Shared Variable Projections 에 대해 논의 합니다.

Optimal subword embeddings via knowledge distillation

기본적인 distillation을 따릅니다.

• Teacher 모델(θt) : 30522 WordPiece token, 768 dimension, hidden state로 pre-trained 된 12-layer BERT_base
• Student 모델(θs) : 동일한 수의 transformer layer로 구성. 단, vocab이 적고 embedding/hidden dimension은 더 작습니다.
• BERT와 동일한 WordPiece 알로기즘 및 학습 코퍼스를 사용하여 student 모델은 4928 Wordpieces를 얻습니다.

Wordpiece token은 greedy segment 알고리즘을 훈련 코퍼스에 적용하여 얻은 sub-word units 입니다. 원하는 수의 wordpeice는 세그먼트화 된 말뭉치가 사용되는 wordpiece의 수가 최소가 되도록 선택됩니다. Teacher, student vocab을 보면 student의 vocab의 93.9% 정도가 teacher vocab에 있기 때문에 줄일 여지가 있습니다.

Google에서는 범용 student 언어 모델을 훈련 시키려고 노력하기 때문에 teacher 모델의 orginal training objective를 재사용하여 student 모델(마스크 언어 모델링, 다음 문장 예측)을 최적화합니다.

Student vocab은 teacher vocab의 완전한 하위 집합이 아니기 때문에 두 vocab은 동일한 단어를 다르게 토큰화 할 수 있습니다. 결과적으로 두 모델의 출력이 정렬되지 않을 수 있습니다. 따라서 해결 할 수 있는 두가지 대안을 설명합니다.

Dual training

Distaillation 과정에서 teacjer vocab에서 사용된 다른 세그먼트화된 토큰과 함께 student vocab을 사용한 시퀀스에서 세그먼트로 토큰(확률 pDT, 하이퍼 파라미터)을 무작위로 선택하여 teacher, student vocab을 혼합 할 것을 제안합니다.

During distillation, for a given training sequence input to the teacher model, we propose to mix the teacher and student vocabularies by randomly selecting (with a probability pDT , a hyperparameter) tokens from the sequence to segment using the student vocabulary, with the other tokens segmented using the teacher vocabulary.

그림에서 입력 컨텍스트 [ ‘I’, ‘like’, ‘machine’, ‘learning’]에서 I 및 machine이라는 단어는 teacher vocab을 사용하여 분할되며 like, learning은 student vocab에서 사용하여 분류됩니다.

Teacher, student vocab은 동일한 단어에 대한 표시 정렬을 권장합니다.

Similar to cross-lingual training in Lample & Conneau (2019), this encourages alignment of the representations for the same word as per the teacher and student vocabularies.

이것은 마스크 언어 모델링 작업을 통해 수행됩니다. 이제 모델은 teacher vocab를 사용하여 세그먼트화된 컨텍스트 단어를 사용하여 student vocab에서 단어를 예측하는 법을 배우고 그 반대로도 학습해야합니다.

모델 매개변수 θt뿐만 아니라 teacher embedding에서도 이러한 방식으로 student embedding을 효과적으로 학습 할 수 있습니다.

Teacher 모델 입력에 대해서만 이중 학습을 수행합니다. Student 모델은 student vocab을 사용하여 독점적으로 세그먼트화된 단어를 받습니다. 또한, 마스크된 언어 모델링 중에 모델은 질문의 단어를 분류하는데 사용된 teacher, student vocab에 서로 다른 softmax 레이어를 사용합니다.

Shared projections

Student 모델을 훈련하기 위해 teacher 모델 출력에만 의존하는 것은 일반화가 잘 되지 않을 수 있습니다. 따라서, 일부 접근법은 student 모델의 주간 예측을 선생의 예측과 일치시키고 조정하는 시도를 합니다. 그러나, 출력 공간이 동일하지 않기 때문에 정렬이 어렵습니다.

대신에, 더 작은 치수를 갖는 teacher 모델 파라미터(θt)로부터 student 파라미터(θs)로의 정보 손실을 최소화 하려고 합니다. 정렬을 하기위해 모델 매개 변수를 동일한 공간에 투영함으로써 해결할 수 있습니다. 그림에서 θt의 각 변수를 θs의 해당 변수와 같은 모양으로 투용합니다.

예를 들어 모양이 768 X 768인 모든 학습 가능한 변수 θt에 대해 두개의 투영(projection) 행렬 U, V를 학습하여 student 모델 θ't에 투영합니다. d는 student 모델의 hidden dimension입니다.

U, V는 해당 차원의 모든 BERT 모델 매개 변수에 공통이며 distillation 후 미세 조정 또는 추론에 필요하지 않습니다.

Student 변수와 teacher 변수의 투용을 정렬하기 위해 별도의 평균 제곱 오차 손실을 소개합니다.

↓ 표시는 하향 투영을 나타냅니다.

$$L^↓p = \sum{\theta'_t,\theta'_s\subset\theta_t,\theta_s} ||U\theta'_tV - \theta'_s||^2$$

위 손실 함수는 student 공간에서 학습 가능한 변수를 정렬합니다. θt에서와 같은 모양으로 훈련 가능한 변수를 θs로 투영할 수 있습니다. 아래 방정식은 teacher 공간에서 학습 가능한 변수를 비교할 수 있습니다.

$$L^↑p = \sum{\theta'_t,\theta'_s\subset\theta_t,\theta_s} ||\theta'_t - V\theta'_sU||^2$$

Optimization objective

최종 손실 함수에는 teacher 모델은 dual-vocab입력으로 훈련되었고 정적이지 않기 때문에 추가적인 투영 손실 외에도 teacher 모델 뿐만 아니라 student 모델를 위한 마스크된 언어 모델링 교차 엔트로피 손실이 포함되었습니다.

P(yi = c|θs)과 P(yi = c|θt)은 각각 클래스 c에 대한 teacher, student 모델 예측 확률을 나타내고 1은 지표함수를 나타냅니다. 아래 방정식은 최종 손실 L_final을 정의하며, e는 하이퍼 파라미터 입니다.

$$L_{ce} = \sum_i(\sum_{c\subset C}[1_{[y_i]=c}logP(y_i=c|\theta_s)] + \sum_{c \subset C}[1_{[y_i]=c}logP(y_i=c|\theta_t)])$$

$$L_{final} = L_p + \epsilon ×L_{ce}$$

Experiments

Knowledge distillation을 평가하기 위해 두가지 클래스의 실험을 설계합니다.

1. 명시적 평가를 위해 마스크 단어 예측 작업을 사용하여 student 언어 모델을 평가
2. 다운 스트림 문장 및 문장 쌍 분류 작업에 student 언어 모델 출력 위에 작업 별 affine 레이어를 추가하여 미세 조정(습득한 표현의 품질을 암시적으로 평가)

Language Model Training

Student BERT 언어 모델을 훈련시키기 위해 teacher BERT 모델을 distillation하는 동안 BooksCorpus와 영어 위키를 훈련시키는데 사용된 것과 같은 vocab을 사용합니다.

다음 문장 예측 손실이 성능에 약간 영향을 미치기 때문에 overall distillation loss를 계산하기 위해 마스크 언어 모델링 작업만 사용합니다.

Student 모델을 사용하여 teacher 모델 입력 vocab을 분할 할 확률인 pDT를 사용하여 teacher 모델 입력에 대해 이중 학습이 가능합니다. shared projection을 포함한 실험의 경우 투영 행렬 U와 V는 Xavier 초기화를 사용했습니다. 손실 중량 계수(loss weight coefficient)는 튜닝 후 1로 설정됩니다.

기존과 달리 다운 스트림 작업에서 미세 조정되지 않은 teacher BERT 언어 모델을 직접 distillation하여 작업에 구애받지 않는 student language model을 얻습니다. 그리고나서 다운 스트림 작업을 위해 미세 조정합니다. 4X4 pod(32 TPU)으로 구성된 cloud tpu를 수행합니다. Learning rate는 0.00125이고 배치 사이즈는 4096인 LAK를 250K 단계에 대해 loss를 최적화 했습니다. student 모델 차원에 따라 2~4일 소요

Models and Baselines

Distillation된 student 모델의 3 가지 평가합니다.

• Teacher, student vocab에 대한 dual train
• Teacher 모델의 하향 투영(dual-train + shared Projdown) dual train 배개변수
• Teacher 모델의상향 투영(dual-train + shared Projup) dual train 배개변수

5K WordPiece를 사용하여 9개의 변형에 대해 임베딩 및 hidden dimensions는 48, 96 및 192로 student 모델을 train 합니다.

언어 모델링 평가를 위해 NoKD (No Knowledge Distillation)없이 baseline을 평가합니다. 모델은 distillation된 student 모델과 동일하게 매개 변수화 되었지만 teacher 모델 objective에 대해 직접 학습 했습니다. 다운 스트림 같은 경우 NoKD 및 PKD(Patiend Knowledge Distillation)과 비교합니다. Teacher 모델의 hidden state를 사용하여 12-layer BERT_base 모델을 3-layer, 6-layer BERT 모델로 distillation 합니다.

For the language modeling evaluation, we also evaluate a baseline without knowledge distillation (termed NoKD), with a model parameterized identically to the distilled student models but trained directly on the teacher model objective from scratch. For downstream tasks, we compare with NoKD as well as Patient Knowledge Distillation (PKD) from Sun et al. (2019), who distill the 12-layer BERTBASE model into 3 and 6-layer BERT models by using the teacher model’s hidden states

• BERT_base와 비교한 student 모델의 크기 요약
• #Params: student 모델의 매개 변수 수
• Model Size: MB
• FLOPS(Floating Point Operations Per Second) rate: 모델에 대한 추론에 필요한 부동 소수점 연산의 상대 배율을 측정

Evalution tasks and Datasets

언어 관점을 명시적으로 평거하기 위해 Reddit 데이터 set를 사용

다운 스트림 language understanding task에 대한 암시적 평가를 위해 GLUE 벤치 마크의 3가지 task에 대해 미세조정하고 평가합니다.

• 64K training 인스턴스가 포함된 two-way sentence sentiment classification task인 SST-2 (Stanford Sentiment Treebank)
• 3.7K training 인스턴스를 사용하여 paraphrases를 식별하는 two-way sentence pair classification task인 MRPC(Microsoft Research Paraphrase Corpus)
• premise-hypothesis relations를 식별하기 위한 hree-way sentence pair classification task with 393K training instances인 MNLI(Multi-Genre Natural Language Inference) 장르 일치 및 장르 불일치 에 대한 별도의 개발 및 테스트 세트가 있습니다. (따로 조정)

모든 다운 스트림 작업 평가에 대한 learning rate는 0.0002이고 batch size는 32인 LAMB을 사용하여 10 epoch를 미세 조정합니다. 최대 시퀀스 길이는 128 token으로 훈련됩니다.

Distillation된 student 모델에 대한 마스크 언어 모델링 작업 정확도 및 스크래치에서 미세 조정 baseline.

Results

위 표에서 다양한 모델과 NoKD 기준에 대한 마스크 단어 예측 정확도 수치가 포함되어 있습니다. dualtrain이 모든 모델 차원에 대한 기준보다 현저히 향상되고, dualtrain에 추가된 shared projection losses이 단어 예측 정확도를 더 향상 시킨다는 것을 관찰합니다. Teacher 공간으로 투영되는 sharedProjUp은 모든 모델 차원에서 SharedProjDown보다 성능이 훨씬 뛰어납니다. 예상대로 hidden state가 작아질수록 성능 저하가 있습니다.

GLUE 서버에서 얻은 다운 스트림 언어 이해 작업 테스트 set에 대한 distillation 모델, teacher 모델 및 baseline과 teacher BERT_base와 비교한 각 모델의 크기 매개변수 및 압축 비율 결과.

MNLI-m 및 MNLI-mm은 MNLI에 대한 장르 일치 및 불일치 테스트 세트를 나타냄

Teacher, student 모델 어휘가 다르기 때문에 teacher BERT_base 모델의 마스크 단어 예측 정확도는 student 모델과 직접 비교할 수 없습니다.

위 표는 BERT_base teacher 모델 및 PKD, NoKD baseline에 대한 다운 스트림 언어 이해 task와 모델 크기에 대한 결과입니다.

제안된 접근 방식으로 훈련된 모델은 동일하게 매개변수화된 NoKD baseline에서 강력하고 일관되게 개선되며, dualtrain 및 shared projection기술이 BERT_base teacher 모델에 대해 적은 손실로 효과적입니다.

PKD 기준과 비교할 때 192 dimension 모델은 PKD 모델 중 하나보다 높은 압축률을 달성하고 3-layer PKD 기준보다 성능이 우수하며 거의 5배정도로 작고 작업 정확도에서 6-layer 기준과 비슷합니다.

192차원에서 96차원 모델로의 성능 저하가 최소입니다.

MRPC task의 경우 dual train으로 훈련된 96차원 모델은 80.5%의 정확도를 달성하며 매개변수가 12배가 많은 6-layer PKD 보다 높습니다.

압축률이 높은 48 차원 모델은 메모리 공간이 25배나 더 큰 3-layer PKD 모델과 유사한 성능을 나타냅니다.

Discussion

• Shared projections and model performance : 다운 스트림 작업 성능의 경우 모든 실험에서 지속적인 개선이 되었습니다. 다만, MRPC와 48 차원 모델(가장 작은 데이터 세트와 모델)에 대해서만 일관된 개선이 되지 않았습니다. 이는 teacher 모델 내부로부터 train signal를 제공하는 것보다 직접적인 방법이 낮은 데이터 또는 작은 모델 시나리오에 더 중요할 수 있습니다. 그러나 더 큰 모델과 더 많은 데이터의 경우 linear projection은 teacher와 student 매개 변수 공간을 정렬하는 상당히 간단한 기능이기 때문에 모델에서 사용할 수 있는 자유도를 감소 시킬 수 있습니다. 관심있는 관련 비교는 up-projection and down-projection of the model variables 입니다. 언어 모델링 작업에서는 상향 투영이 다운 스트림 작업에서는 약간 더 우수합니다. 잘 훈련된 teacher 모델의 매개변수는 teacher space에서 고품질의 로컬 최소값을 나타내며 투사 중 쉽게 검색이 가능합니다.
• Vocabulary size tradeoffs : 입력 vocab 크기와 관련된 문제는 자연어 처리 문제와 관련이 있습니다. NLP 모델의 vocab을 줄이는 데 대한 한가지 우려는 평균 token화 된 시퀀스 길이를 늘려 모델 학습을 어렵게 하는 것입니다. 그래서 이 작업에서는 기계 번역과 같은 입력 시퀀스 길이의 영향을 받지 않는 짧은 텍스트에 대한 분류 작업을 고려합니다. 많은 실제 응용 프로그램은 짧은 텍스트 입력을 중심으로 하고 있기 때문에 vocab 크기와 시퀀스 길이 사이의 더 나은 균형이 필요합니다.
• Order of distillation and fine-tuning : 기존에는 다운 스트림 작업에서 teacher 모델을 미세조정 한 다음에 distaillation을 합니다. 그러나 이 논문에서는 목표가 단어의 크기가 큰 wordpiece vocab에 의해 결정되는 정도를 탐구하는 것이며, 미세 조정된 teacher 모델을 distillation하여 실험을 하였고 다운 스트림 작업에서 더 나은 결과를 얻을 수 있도록 합니다.

Conclusion

BERT에 대한 knowledge distillation의 효과를 향상시키기 위해 2가지 새로운 아이디어를 제안했으며 작은 embedding, hidden dimension 뿐만 아니라 훨씬 작은 vocab을 사용하는데 중점을 두었습니다. 이 두가지 기술을 결합하여 일련의 압축된 12-layer student BERT 모델을 train 했습니다. 일반화된 언어 관점과 4개의 표준화된 다운 스트림 task를 모두 평가하기 위해 모델에 대한 실험을 통해 모델의 정확도와 압축 효율에 미치는 효과를 보여줍니다.