본 논문은 Microsoft Research에서 2020년에 작성한 논문이다.
논문 링크: https://arxiv.org/pdf/2002.10957
1. Introduction
본 논문은 수억 개의 파라미터를 갖는 일반적인 NLP 모델을 경량화한 MiniLM 모델을 제안한다.
- 보통의 Language Model (이하 LM)들은 대용량의 데이터셋으로 Pre-training 되고, 그다음 연구 목적에 맞는 downstream 작업에 맞춰 Fine-tuning 되는 것이 일반적이다. → 모델의 크기가 큰 경우가 많아 실사용에 무리가 있다.
- 2015년 Knowledge Distillation 기법, 한국말로는 지식 증류 기법이라고 불리는 방법론이 등장한다. 간단히 설명을 하자면, Fine-tuning 후에 높은 성능을 가지는 모델을 Teacher로 삼아서 작은 모델 Student가 큰 모델을 모방하며 학습하는 방법이다. 아마 뒤에서 나오는 설명을 보면 이해가 갈 것이다. → 모델이 큰 경우 Fine-tuning이 어렵다고 한다.
- 기존 연구들에서는 DistillBERT, MobileBERT, TinyBERT 경량 모델을 설계하였다. 최종 결괏값에 softmax 함수를 취한 soft target 확률이나 Teacher 모델 특정 layer의 hidden state 값 등 다양한 요소들을 모방하도록 설계한 모델들이다. → 아키텍처 설계 방면에서 제약이 있는 점을 문제점으로 꼽고 있다.
- 본 논문에서는 이러한 아이디어에 영감을 받아 Transformer 기반 task-agnostic 한 deep self-attention distillation framework를 제안한다. 메인 contribution은 다음과 같다.
- Teacher 모델의 마지막 레이어에 대해서만 지식 증류 기법을 적용하여 아키텍처 설계를 조금 더 유연하게 가져간다.
- Deep self-attention knowledge를 강화하기 위해 Value 값에 대해서도 지식 증류 기법을 적용한다.
- Student와 Teacher 모델 사이에 Asisistant 모델을 추가하여 성능을 향상시킨다.
아마 여기까지만 보고도 본 논문에 대한 감이 오는 분들이 있을 수도 있겠다. 최고..❣️
더 깊이 알고 싶은 사람들을 위해 지금부터 논문 내용을 정리해 보겠다.
2. Preliminary
아마 이 논문을 읽고 계신 여러분들은 Transformer와 BERT에 대한 지식을 가지고 있을 거라고 생각한다.
그럼에도 불구하고 혹시 나는 설명이 더 필요하다! 싶으신 분들은 주저 말고 댓글을 달아주길 바란다.
지금 모르는 내용을 평생 모르는 것보단, 오늘부터라도 알아서 내일은 또 다른 지식을 쌓는 편이 훨씬 낫다.
(1) Input Representation
텍스트를 어떻게 임베딩했는지에 대한 내용이다.
BERT 모델과 동일하게 WordPiece를 사용해서 임베딩 했다고 한다.
설명을 덧붙여보자면, BERT에선 "forcasted" 같은 word는 "forcast", "##ed" 같은 subword로 나뉘고, 입력 텍스트의 시작 부분은 [CLS] 토큰으로 시작되고, 각기 다른 문장은 [SEP] 토큰으로 구분한다.
(2) Backbone Network: Transformer
Transformer의 각 레이어에서는 길이가 x인 vector를 받아 다음 상태의 hidden vector를 생성한다.
그래서 hidden vector H는 입력 x에 대한 contextualized 표현을 갖게 된다고 불린다.
현재 l-1번째 레이어를 보고 있다고 하자.
그다음 l번째 레이어에 대한 Query, Key, Value를 생성하기 위해, l번째 레이어에 저장된 Weight들과 l-1번째 레이어에서 만들어진 hidden vector를 사용한다.
l번째 레이어의 Query, Key로 scaled dot-product 연산을 수행한 결과에 softmax를 취해주면 드디어 attention distribution이 계산된다.
이걸 왜 굳이 굳이 설명을 하느냐? 하실 수도 있다.
본 논문에서는 이 attention distribution에 대해 지식 증류 기법을 사용하기 때문에 이 부분을 짚고 넘어가면 뒤에서 이해가 더 잘 될 것이라 준비해 봤다. 😋
(3) Transfer Distillation
지식 증류 기법은 다음 논문에서 모델 경량화를 위해 소개된 학습 방법론이다.
Distilling the Knowledge in a Neural Network (2015) https://arxiv.org/pdf/1503.02531
수식은 별로 그렇게 어렵지 않다.
D는 학습 데이터를 의미하고, f는 각 모델에서 추출된 feature를 의미한다.
그리고 L은 loss 함수를 의미하는데, 보통 MSE나 KL-divergence를 사용하는 것이 일반적이다.
쉽게 설명하자면, Stduent와 Teacher 모델에서 만들어낸 결과가 최대한 유사해지도록!! 학습하는 것이 이것의 목적이다.
앞에서 언급한 것처럼 DistillBERT, MobileBERT, TinyBERT과 같은 다른 경량 모델들은 다양한 요소들에 대해 Distillation 기법을 적용한 것들이다.
다른 모델들에 대해 더 탐구해서 전체 모델을 비교해 보는 것도 좋은 공부가 될 것 같다. ☀️
3. Deep Self-Attention Distillation
이제부터 위 사진에서 1,2로 표시한 부분에 대해 먼저 설명하고 그다음에 Assistant 모델에 대해 설명을 진행해 보겠다.
(1) Self-Attention Distribution Transfer
Teacher와 Student 마지막 레이어의 attention distribution의 KL-divergence가 최소가 되도록 수식이 짜여있다.
여기서 Ah는 헤드 개수, x는 임베딩된 입력 토큰의 개수이고, L과 M은 두 모델 내부의 레이어 개수이다.
L과 M은 이것이 두 모델의 마지막 레이어라는 점을 알려주고 있다.
마지막 레이어의 값만 사용하기 때문에 Student 모델의 레이어 개수를 유연하게 설정할 수 있으며, 저자들은 그로 인해 아키텍처 설계에 드는 노력을 줄일 수 있었다는 내용을 언급한다.
(2) Self-Attention Value-Relation Transfer
Value-relation에 대한 수식은 (1)과 같은 형태를 띠고 있으며, VR 값은 아래와 같이 계산된다.
Value값과 그것의 전치 행렬에 대해 scaled dot-product 연산을 수행하고 softmax 값을 취하면 VR 값을 얻을 수 있다.
이 두 VR의 차원은 동일하게 '어텐션 헤드 개수 x 임배딩 차원 x 임배딩 차원' 만큼의 크기를 갖게 된다.
논문 사진을 보면 양쪽 모델 Values 아래에 적힌 차원이 다른 것을 확인할 수 있다.
하지만 scaled dot-product 연산을 거친 후에는 차원이 통일되기 때문에, 이 또한 아키텍처 설계에 도움이 되었을 뿐만 아니라 별도의 파라미터를 추가할 필요가 없어 모델을 더 작게 만들 수 있었다고 한다.
(3) Teacher Assistant
Teacher 모델에 비해 레이어 개수와 hidden state 사이즈가 2배 이상 작은 Student 모델을 학습시킬 때는 중간에 Assistant 모델을 추가해서 학습하는 게 더 성능이 좋았다고 한다.
방법은 뭐 크게 어려울 건 없다.
먼저 Teacher를 이용해 Assistant를 학습시켜 주고, 그다음 Assistant 모델로 Student 모델을 학습시킨다.
이 방법론은 Improved Knowledge Distillation via Teacher Assistant (2019)라는 논문에서 영감을 받았다고 한다.
4. Experiments
자.. 이제 한숨 돌리고 결과를 찬찬히 살펴보자.
(1) Distillation Setup
- Teacher 모델로 BERT(base) 사용
| 레이어 | hidden state 크기 | 어텐션 헤드 | 파라미터 |
| 12개 | 768 | 12개 (Student에서도 동일해야 함) | 1억 900만 개 |
- 데이터셋: English Wikipedia & BookCorpus
- 전처리: WordPiece 토큰화 방식 (어휘 크기: 30,522, 최대 시퀀스 길이: 512)
- 기타: Adam 설정값, learning rate, epochs 등등 자세한 내용은 논문에 나와 있으니 참고 바랍니다!
이 외에도 multilingual MiniLM 학습 내용과 훈련 방법에 대한 내용이 나와 있었다.
(2) Downstream Tasks
1. Extractive Question Answering (EQA)
EQA는 주어진 문단 P와 질문 Q가 주어졌을 때, 그에 대한 정답을 문단 P안에서 선택하는 것이다.
이때 정답은 문단 P안에서 끊기지 않고, 연속적이어야 한다. (하나의 뭉탱이만을 골라야 한다!)
입력은 다음과 같은 형태로 구성이 된다.
[CLS] Q [SEP] P [SEP]
해당 태스크에선 각 토큰이 답변의 시작인지 끝인지를 예측하기 위해 2개의 FC를 썼다고 한다.
평가는 QA 벤치마크에서 주로 사용되는 SQuAD 2.0 이란데이터셋에서 평가를 진행했다고 한다.
2. GLUE (General Language Understanding Evaluation)
GLUE는 아래와 같은 9개의 분류 task로 구성되어 있다고 한다.
CoLA, SST, MRPC, STS, QQP, MNLI, QNLI, RTE, WNLI
해당 태스크에선 CLS 토큰 다음에 FC 레이어를 추가해서 레이블을 예측했다고 한다.
(3) Main Results
SQuAD2.0에선 F1 스코어를, 그 외에 task에선 모두 정확도를 평가지표로 사용했다.
혹시 F1 스코어를 모른다면, 구글에 검색해서 수식을 이해해 보길 바란다. 하나도 안 어렵다. 💛
- BERT와의 비교: 속도가 2배 더 빠름에도 불구하고, SQuAD2.0 or MNLI 등의 작업에서 comparable 한 성능을 보인다.
- TinyBERT와의 비교: SQuAD2.0에서 3.0 이상의 F1 스코어, CoLA에선 5.0 이상의 정확도 향상을 보인다.
저자들은 Student 모델의 크기를 달리하며 비교 실험을 진행했다.
여기서 잠깐~!
MLM-KD는 soft target 확률에 대해 지식 증류기법을 수행한 버전을 의미하고,
w/TA는 앞에서 언급한 (Teacher) Assistnat 모델을 추가해서 학습한 실험을 의미한다.
- 모든 항목에서 TinyBERT 보다 성능이 좋다는 것을 체크하자.
- Assistant 모델을 추가하는 것, 특히 작은 모델에 대해서는 그 효과가 더 크게 발휘하는 것을 볼 수 있다.
(4) Ablation Study
Value-Rel는 3-2 Self-Attention Value-relation Transfer에서 소개했던 그 지식 증류 기법이다.
이 결과를 통해 모듈의 필요성이 입증되었다!
본 논문에선 KL-divergence가 최소화되는 방식으로 loss가 짜여 있다.
Value-MSE는 value 행렬에 대해 MSE를 계산하는 방식으로 loss 수식을 만든 버전이다.
표에서 확인할 수 있듯이, 논문에서 제시한 방법이 더 좋다!
마지막으로, 모든 레이어 별로 증류하는 방법과 마지막 레이어에 대해서만 증류하는 방법의 성능을 비교해 보자.
실험 결과, MiniLM이 더 좋은 성능을 보였다. (개인적으로 이 부분은 조금 의외였다...!)
당연하게도 저자들은 하나의 레이어에 대해서만 증류를 진행하니까 훈련 속도도 더 빨랐고, 레이어 매핑 어려움도 비교적 덜 했다고 언급한다.
5. Disscusion
여기까지 오느라 다들 수고 많으셨습니다 ~ 😀
이제부터 또 다른 다양한 실험 결과를 보도록 하지요 !_!
(1) Better Teacher Better Student
a와 b는 각각 12-layer와 6-layer 모델을 의미한다.
크기가 더 큰 MiniLM a 모델의 경우 b 모델보다 더 뛰어났으며, SQuAD2.0과 GLUE 몇몇 태스크에서 BERT보다 더 좋은 성능을 내는 것을 볼 수 있다.
(2) MINI LM for NLG Tasks
두 가지 Natural Language Generation 태스크에 대한 결과를 살펴봅시다.
1. Question Generation
이 작업은 주어진 입력과 '답변'이 있을 때, 모델이 그에 맞는 질문을 생성하도록 요구한다.
크게 나뉜 두 행은 각각 다른 데이터셋으로 구분이 되어있다.
평가지표는 BLEU@4, METEOR, ROUGE-L을 사용했다.
표를 보면, UNILM 보다는 좋지 않지만, 나머지 모델에 비해 더 좋은 성능을 보이고 있음을 알 수 있다.
2. Abstractive Summarization
Abstractove Summarization은 주어진 문서를 핵심 내용만 갖도록 압축하는 작업이다.
파라미터 수가 현저히 적음에도 불구하고, 큰 모델들에 대적할만한 성능을 가지는 것을 볼 수 있다! 😁
(3) Multilingual MiniLM
저자들은 또한 다국어 pre-trained 모델 XLM-R(base)을 Teacher로 설정하여 학습을 진행했다.
원래 다국어 모델의 경우 더 많은 계산을 필요로 하기 때문에 훈련이 오래 걸린다는 단점이 있는데, MiniLM에선 마지막 레이어의 값들만 뽑아 쓰기 때문에 훈련 속도가 훨씬 더 빨랐다고 한다.
1. XNLI
15개 언어에 대한 분류 작업을 수행한 결과이다.
이 결과에서도 또한 크기가 큰 모델들과 비슷한 성능을 보인다!
2. MLQA
이 데이터셋은 영어 중심이던 기존 데이터셋들관 다르게, 다국어 간 교차를 통한 자연어 추론 문제로 구성되어 있다.
위의 표는 SQuAD1.1 데이터로 훈련하고, MLQA 데이터셋으로 테스트한 결과이다.
평가 지표는 각각 F1 스코어와 Exact Match 스코어다.
(EM: 모든 요소들이 정확하게 일치하는지 측정하는 지표인데, 사실 이 task에서 어떻게 계산됐는지 감이 잘 안 잡힌다. 😂)
여기서 먼저 전체적으로 제일 좋은 성능을 보여주는 fine-tuned 버전인 XLM-R(base) ‡ 모델의 결과를 보자.
MiniLM a 모델은 그보다는 뒤지지만, 어느 정도 괜찮은 성능을 가지고 있다.
6. Conclusion
본 논문에서는 경량 NLP 모델인 MiniLM을 제안하였다.
메인 contribution은 바로 트랜스포머의 마지막 레이어에 대해서만 지식 증류 기법을 적용했다는 점.
그 지식 증류 기법은 attention distribution과 value-relation 값에 적용했으며, Teacher 모델보다 크기가 훨씬 작은 Student 모델을 학습시킬 땐 중간에 Assistant 모델을 두어서 더 좋은 모델을 설계했다.
크기가 더 작음에도 불구하고 기본 모델들보다 더 좋은 성능을 보이기도 했으며, 실험에서 BERT가 아닌 다른 모델에 대해서도 지식 증류 기법을 사용할 수 있음을 보였다.
원래 지식 증류 기법에 대해 어느 정도는 알고 있었지만, 이 논문을 읽고 난 후에야 완벽하게(?) 체화할 수 있었다.
아무래도 2020년 논문이다 보니 현재는 더더더더 엄~~청 다양한 압축 기법이 소개되어 있을 것이다!
온디바이스AI 강의 수강 중에 edge device에 관한 논문을 읽다가 찾게 된 논문이었는데, 개인적으로 엄청 재밌게 읽었다.
나의 글이 여러분께 도움이 되길 바라며 이만 물러나겠습니다~ 👋🏻
궁금한 점이 있다면 언제든 댓글 남겨주세요. 😍
'On-device AI' 카테고리의 다른 글
| [논문 리뷰] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter (2019) (2) | 2024.07.12 |
|---|---|
| [논문 리뷰] MiniVLM: A Smaller and Faster Vision-Language Model (2020) (0) | 2024.06.06 |