[논문 리뷰] ALP-KD: Attention-Based Layer Projection for Knowledge Distillation

Peyman Passban, Yimeng Wu외 2명|arXiv (Cornell University)|2020. 12. 27.

Topic Modeling참고 문헌 29인용 수 8

한 줄 요약

이 논문은 BERT를 더 작은 모델로 압축할 때 중간 레이어를 건너뛰는 문제를 방지하기 위해 학습 가능한 어텐션 가중치를 사용해 다수의 교사 레이어를 통합하는 어텐션 기반 레이어 프로젝션 방법인 ALP-KD를 제안한다. 학생 모델이 최종 및 중간 표현을 모두 모방할 수 있도록 함으로써, GLUE 벤치마크에서 기존 방법들을 능가하는 우수한 성능을 달성한다.

ABSTRACT

Knowledge distillation is considered as a training and compression strategy in which two neural networks, namely a teacher and a student, are coupled together during training. The teacher network is supposed to be a trustworthy predictor and the student tries to mimic its predictions. Usually, a student with a lighter architecture is selected so we can achieve compression and yet deliver high-quality results. In such a setting, distillation only happens for final predictions whereas the student could also benefit from teacher's supervision for internal components. Motivated by this, we studied the problem of distillation for intermediate layers. Since there might not be a one-to-one alignment between student and teacher layers, existing techniques skip some teacher layers and only distill from a subset of them. This shortcoming directly impacts quality, so we instead propose a combinatorial technique which relies on attention. Our model fuses teacher-side information and takes each layer's significance into consideration, then performs distillation between combined teacher layers and those of the student. Using our technique, we distilled a 12-layer BERT (Devlin et al. 2019) into 6-, 4-, and 2-layer counterparts and evaluated them on GLUE tasks (Wang et al. 2018). Experimental results show that our combinatorial approach is able to outperform other existing techniques.

연구 동기 및 목표

• 기존 지식 증류 방법이 아키텍처 불일치로 인해 중간 교사 레이어를 건너뛰는 한계를 해결한다.
• 일对일 레이어 매칭에 의존하는 것 대신, 다수의 서로 다른 교사 레이어에서의 감독을 통해 학생 모델의 성능을 향상시킨다.
• 동적으로 관련 교사 레이어 표현을 가중하고 융합하기 위한 조합형 어텐션 기반 메커니즘을 개발한다.
• 어텐션을 통해 교사 레이어를 융합할 경우 학생 및 교사 모델 간 내부 표현 정렬이 향상됨을 입증한다.
• 깊은 교사 네트워크(예: 12층 BERT)를 더 작고 효율적인 학생 모델(6층, 4층, 2층 버전)로 압축할 때 최신 기술 수준의 성능을 달성한다.

제안 방법

• 다수의 교사 레이어 특징을 단일 가중 조합으로 융합하기 위한 새로운 어텐션 기반 레이어 프로젝션(ALP-KD) 메커니즘을 제안한다.
• 학생의 현재 레이어에 대한 관련성을 기반으로 각 교사 레이어의 중요도 점수를 학습 가능한 어텐션 메커니즘을 통해 할당한다.
• 다중 헤드 어텐션 모듈을 통해 계산된 가중치를 사용해, 다수의 교사 레이어에서 온 중간 특징의 가중합으로 통합된 교사 표현을 수식화한다.
• 학생의 중간 출력과 가중 조합된 교사 특징 간의 차이를 비교하는 증류 손실을 통해 융합된 교사 표현을 학생의 학습 과정에 통합한다.
• 레이어 간 일치가 필요 없이 비연속적이거나 아키텍처가 다른 레이어에 어텐션을 적용할 수 있도록 탄력적인 구성이 가능해져, 레이어별 정렬이 필요 없어진다.
• 하드 레이블에 대한 교차 엔트로피 손실과, 학생 출력과 어 attention-결합된 교사 특징 간의 증류 손실을 조합한 손실을 사용해 학생 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 어텐션을 통해 서로 다른 교사 레이어를 융합함으로써, 일대일 레이어 매칭을 넘어서 지식 증류 성능을 향상시킬 수 있는가?
• RQ2 교사 표현의 어텐션 기반 융합이 학생 및 교사 모델 간 중간 은닉 상태 정렬에 어떤 영향을 미치는가?
• RQ3 아키텍처 불일치로 인해 레이어를 버리지 않고 모든 교사 레이어의 정보를 활용함으로써 학생 모델 성능 향상이 측정 가능한가?
• RQ4 NLP 벤치마크에서 기존 방법인 PKD와 RKD에 비해 ALP-KD는 정확도와 표현 충실도 측면에서 얼마나 뛰어난가?
• RQ5 깊은 모델(예: 12층 BERT)을 매우 작은 아키텍처(예: 2층 BERT)로 압축할 때 ALP-KD는 성능 저하 없이 효과적으로 적용될 수 있는가?

주요 결과

• 12층 BERT를 6층, 4층, 2층 학생 모델로 압축할 때, ALP-KD는 PKD 및 RKD를 포함한 기존 지식 증류 방법들을 다수의 GLUE 벤치마크 작업에서 능가한다.
• 어텐션 기반 융합으로 인해 학생 및 교사 모델 간 중간 표현이 훨씬 유사해졌으며, 코사인 거리 시각화 결과 ALP-KD는 평균 거리 0.05를 기록했고, PKD는 0.20을 기록했다.
• PCA를 통한 은닉 상태 시각화 결과, ALP-KD 학생 모델이 기준 방법보다 교사 모델과 더 유사한 표현을 학습하는 것으로 확인되었다.
• 모든 교사 레이어의 정보를 활용함으로써 레이어 건너뛰기 문제를 효과적으로 완화했다.
• 깊은 교사 모델에서 매우 좁은 학생 모델로의 압축에서도 ALP-KD는 아키텍처 격차에 관계없이 뛰어난 성능 유지를 보였다.
• 제거 실험 결과, 어텐션 기반 조합이 필수적임을 입증했다. 단순 연결 또는 균일한 가중치는 경쟁력 있는 성능을 달성하기에 부족하다.