[논문 리뷰] Knowledge Distillation from Internal Representations

Gustavo Aguilar, Ling Yuan외 4명|arXiv (Cornell University)|2019. 10. 08.

Topic Modeling참고 문헌 20인용 수 23

한 줄 요약

이 논문은 BERT 유사 모델의 내부 표현에서 지식 distillation을 제안하며, 큰 교사 모델의 소프트 레이블과 은닉층 활성화 값을 따라가도록 학습하는 학생 네트워크를 설계한다. 어텐션 맵과 은닉 상태에 대한 KL 발산을 사용해 여러 내부 레이어에서 distillation을 구현함으로써 일반화 능력이 크게 향상되고, GLUE 벤치마크 데이터셋 전반에서 표준 지식 distillation보다 뛰어난 성능을 보인다.

ABSTRACT

Knowledge distillation is typically conducted by training a small model (the student) to mimic a large and cumbersome model (the teacher). The idea is to compress the knowledge from the teacher by using its output probabilities as soft-labels to optimize the student. However, when the teacher is considerably large, there is no guarantee that the internal knowledge of the teacher will be transferred into the student; even if the student closely matches the soft-labels, its internal representations may be considerably different. This internal mismatch can undermine the generalization capabilities originally intended to be transferred from the teacher to the student. In this paper, we propose to distill the internal representations of a large model such as BERT into a simplified version of it. We formulate two ways to distill such representations and various algorithms to conduct the distillation. We experiment with datasets from the GLUE benchmark and consistently show that adding knowledge distillation from internal representations is a more powerful method than only using soft-label distillation.

연구 동기 및 목표

• 표준 지식 distillation은 출력 확률만 전이하므로 교사의 내부 언어학적 지식을 유지하지 못하는 한계를 해결하기 위해.
• 큰 교사 모델에서 소프트 레이블뿐 아니라 내부 표현(예: 어텐션 맵과 은닉 상태)까지 전이함으로써 모델 압축을 향상시키기 위해.
• 학생 모델이 출력 분포뿐만 아니라 교사의 내부 행동과도 유사하게 일반화되도록 하여, 교사와 유사한 의사결정 경계를 따르도록 보장하기 위해.
• 소프트 레이블 distillation만으로는 달성할 수 없는 성능 향상을 달성하기 위해, 특히 자원이 제한된 또는 제한된 추론 환경에서 효과적인 방법을 개발하기 위해.

제안 방법

• 이 방법은 최종 출력 확률 외에도 교사 모델의 여러 레이어에 걸친 중간 은닉 표현에서도 지식 distillation을 수행한다.
• 각 레이어에서 교사 및 학생 모델의 어텐션 확률을 일치시키기 위해 KL 발산을 사용하여 고차원적 언어학적 추상화를 포착한다.
• 선택된 레이어에서 교사 및 학생 모델의 은닉 상태 간에 별도의 KL 발산 손실을 적용하여 내부 표현의 구조적 일치를 가능하게 한다.
• 진행형 및 스택형 지식 distillation 전략을 모두 지원하여, 여러 교사 레이어의 지식을 단일 학생 레이어로 압축하는 데에灵活性를 제공한다.
• 학생은 하드 레이블에 대한 교차 엔트로피 손실, 소프트 레이블에 대한 distillation 손실, 어텐션 맵과 은닉 상태에 대한 추가 distillation 손실을 조합하여 학습된다.
• 기존의 FitNet과 같은 접근 방식과 달리, 표현 일치를 위한 추가 파rameter를 도입하지 않고도 동일한 공간 내에서 표현을 직접 일치시킴으로써, 표현 일치를 위한 추가 파rameter를 생략한다.

실험 결과

연구 질문

• RQ1 큰 교사 모델의 내부 표현(예: 어텐션 맵과 은닉 상태)을 distillation함으로써, 표준 지식 distillation을 넘어서 학생 모델의 일반화 능력과 성능을 향상시킬 수 있는가?
• RQ2 출력 확률 외에도 학생의 내부 행동을 교사와 일치시킴으로써, 특히 모호하거나 도전적인 예제에서 더 견고하고 정확한 예측을 이끌 수 있는가?
• RQ3 내부 표현 distillation은 표준 KD와 비교해 오류 패tern과 교사의 예측과의 일관성 측면에서 어떻게 다를까?
• RQ4 추가 파rameter를 도입하거나 추론 복잡도를 증가시키지 않고도 깊은 트랜스포머 레이어에 저장된 언어학적 지식을 효과적으로 압축할 수 있는가?
• RQ5 다른 레이어(예: 초기 레이어 대비 후기 레이어)에서 distillation을 수행할 경우, 최종 학생 성능과 일반화 능력에 어떤 영향을 미치는가?

주요 결과

• 내부 표현 distillation을 적용한 학생 모델은 QQP 개발 세트에서 교사 모델의 정확 예측의 97.9%를 달성했으며(36,191/36,967건), 표준 KD의 경우 95.7% (35,401건)였다.
• 내부 표현 distillation을 적용한 학생 모델은 776건의 잘못된 예측을 하였고, 표준 KD 학생 모델의 경우 1,566건이었으며, 이는 교사의 의사결정 경계와 더 잘 일치함을 시사한다.
• QQP 데이터셋에서, 교사가 정확하게 예측했지만 표준 KD 학생 모델이 실패한 경우(예: 표본 3 및 4, 표 4 참조)에 내부 표현 distillation을 적용한 학생 모델은 100%로 정확히 식별하였다.
• 교사가 정확한 경우, 내부 표현 distillation을 적용한 방법은 표준 KD 대비 잘못된 예측 수를 50% 이상 감소시켜, 더 뛰어난 일관성과 일반화 능력을 보였다.
• 교사가 오류를 범한 경우(예: 표 4의 표본 1), 내부 표현 distillation을 적용한 학생 모델은 더 경계를 두었으며(확률 0.4221 대비 표준 KD의 0.9999), 과신을 줄이고 보다 잘 校정된 예측을 보였다.
• 이 방법은 모든 네 개의 GLUE 벤치마크 데이터셋에서 표준 KD를 일관되게 뛰어넘었으며, 정확도 향상과 교사의 행동과의 일치성 향상 측면에서 견고한 성과를 보였다.