[논문 리뷰] Extract then Distill: Efficient and Effective Task-Agnostic BERT Distillation

Cheng Chen, Yichun Yin외 5명|arXiv (Cornell University)|2021. 04. 24.

Topic Modeling참고 문헌 30인용 수 4

한 줄 요약

이 논문은 사전에 학습된 교사 모델 파라미터를 재사용하여 더 작은 학생 모델를 초기화하는 유연하고 효율적인 방법인 Extract Then Distill (ETD)을 제안한다. 중요도 점수 또는 무작위 선택을 통해 가장 중요한 뉴런(FFN, 어텐션 헤드, 히든 상태)을 추출하고 재사용함으로써, ETD는 훈련 비용을 70% 감소시키면서도 기준 distillation 방법보다 성능을 향상시킨다.

ABSTRACT

Task-agnostic knowledge distillation, a teacher-student framework, has been proved effective for BERT compression. Although achieving promising results on NLP tasks, it requires enormous computational resources. In this paper, we propose Extract Then Distill (ETD), a generic and flexible strategy to reuse the teacher's parameters for efficient and effective task-agnostic distillation, which can be applied to students of any size. Specifically, we introduce two variants of ETD, ETD-Rand and ETD-Impt, which extract the teacher's parameters in a random manner and by following an importance metric respectively. In this way, the student has already acquired some knowledge at the beginning of the distillation process, which makes the distillation process converge faster. We demonstrate the effectiveness of ETD on the GLUE benchmark and SQuAD. The experimental results show that: (1) compared with the baseline without an ETD strategy, ETD can save 70\% of computation cost. Moreover, it achieves better results than the baseline when using the same computing resource. (2) ETD is generic and has been proven effective for different distillation methods (e.g., TinyBERT and MiniLM) and students of different sizes. The source code will be publicly available upon publication.

연구 동기 및 목표

• 대규모 비지도 학습 코퍼스에서 광범위하게 훈련이 필요한 태스크 무관 BERT distillation의 높은 계산 비용을 해결하기 위해.
• 교사 모델과 다른 너비와 깊이를 가진 학생 모델을 포함하여, 어떤 크기의 학생 모델이라도 효율적이고 효과적으로 distillation을 수행할 수 있도록 하기 위해.
• 교사 모델의 파라미터를 재사용하는 일반적인 초기화 전략을 개발하여 수렴 속도를 가속화하고 성능을 향상시키기 위해.
• 절단 기반 추출을 통한 파라미터 재사용이 무작위 초기화나 단순 레이어 복사보다 더 효과적임을 입증하기 위해.

제안 방법

• ETD는 너비 방향 파라미터 재사용을 위한 두 가지 변형, ETD-Rand(무작위 뉴런 추출)와 ETD-Impt(중요도 기반 뉴런 추출)를 도입한다.
• 지식 전이의 광범위한 커버리지 확보를 위해 교사 모델의 Transformer 레이어들에서 균일한 레이어 선택을 수행한다.
• FFN 뉴런, 어텐션 헤드, 히든 상태에 대한 중요도 점수를 계산하여 선택적 파라미터 추출을 안내한다.
• 추출된 파라미터는 미세조정 이전에 학생 모델을 초기화하는 데 사용되어 광범위한 distillation 훈련의 필요성을 줄인다.
• TinyBERT 및 MiniLM과 같은 기존 distillation 프레임워크와 호환된다.
• 추출 과정은 효율적이며 거의 무시할 수 있는 계산 오버헤드를 유발한다.

실험 결과

연구 질문

• RQ1 구조적 추출을 통해 사전에 학습된 교사 모델 파라미터를 재사용하면 태스크 무관 BERT distillation의 효율성이 향상되는가?
• RQ2 중요도 기반 파라미터 선택이 무작위 또는 균일한 선택보다 학생 모델 초기화에서 성능을 더 잘 끌어올리는가?
• RQ3 ETD는 교사 모델과 동일한 너비와 깊이를 가지지 않은 학생 모델에도 적용 가능한가?
• RQ4 ETD는 성능을 유지하거나 향상시키면서 계산 비용을 얼마나 줄일 수 있는가?
• RQ5 레이어 선택 전략(균일, 상위, 하위)의 선택이 distillation 수렴과 최종 성능에 어떤 영향을 미치는가?

주요 결과

• ETD는 기준 방법 대비 distillation의 계산 비용을 70% 감소시키면서도 성능에 손상 없이 수행한다.
• 동일한 계산 자원을 사용할 경우, ETD는 GLUE 및 SQuAD 벤치마크 양쪽 모두에서 기준 distillation 방법보다 더 우수한 성능을 달성한다.
• 중요도 점수를 사용해 파라미터를 선택하는 ETD-Impt는 ETD-Rand 및 무작위 초기화보다 일관되게 뛰어난 성능을 보인다.
• 제거 실험 결과, 가장 중요도가 낮은 뉴런조차도 긍정적인 기여를 하며, 중요도 순서를 뒤집는 것은 성능 저하를 초래한다.
• 균일한 레이어 선택 전략이 상위 또는 하위 레이어 선택 전략보다 우수하며, 얕은 레이어에서 깊은 레이어로의 균형 잡힌 지식 전이를 보장한다.
• 이 방법은 TinyBERT 및 MiniLM과 같은 다양한 distillation 프레임워크에 일반적이고 효과적이며, 다양한 크기의 학생 모델과도 잘 작동한다.