[2022] BERT Learns to Teach: Knowledge Distillation with Meta Learning

 

[2022] BERT Learns to Teach: Knowledge Distillation with Meta Learning

Wangchunshu Zhou1, Canwen Xu, Julian McAuley

 

본문의 논문 ACL 2022로, 링크를 확인해 주세요.

 

Abstract

 

 training 중에 teacher 모델이 고정되는 전통적 지식 증류(KD) 방법에 대한 간단하면서도 효과적인 대안인 메타 학습을 통한 지식 증류(MetaDistil)를 제안한다. teacher network가 meta-learning framework에서 증류된 student network의 성능에서 피드백을 얻어 student network에  더 나은 지식 전달을 학습할 수 있다. 다시 말해, 이 부분이 learning to teach이다. 

 그리고 향상된 inner-learner에 중점을 둔 meta learning algorithms에서 inner-learner과 meta-learner 간의 alignment을 개선하기 위해 pilot update를 같이 소개하고 있다.

 간단한 knowledge distillation에 관한 내용은 여기를 확인하면 쉽게 이해할 수 있다.

 

 

 

 

 summary

 

 저자들은 knowlege distillation의 단점이 있음을 지적한다.

 

 ① teacher model은 student의 capacity를 인지하지 않는다.

 

 교육학에서 최근 연구들은 student 중심의 학습을 제안하는데, 이는 student의 성향과 학습 능력을 고려하여 student의 성능 향상을 보여왔다. 하지만, conventional KD에서는, student는 student model의 학습 능력과 성능과는 상관없이 teacher에서 수동적으로 지식을 수용한다. 그리하여 그 후 연구에서 student를 인식하여 teacher과 함께 학습하는 distillation과 task 특화 목적의 student를 소개하였다.

 

 ② teacher model은 distillation에 최적화되지 않는다.

 

이전까지는 teacher이 자신의 추론 성능을 최적화하는데 훈련되었지만, 자신의 지식을 student에 전달할 필요를 인지하지 않았다. 간혹 차선책으로 사용하기도 했다. 마치 박사과정의 학생이 문제 해결 능력을 갖추었지만, 교수님이 더 좋은 질의 training을 요구하는 것과 같다.

 

 위 두 가지의 단점을 보완하기 위해 저자들은 MetaDistil을 제안하며, 이는 distillation 과정을 통해 teacher의 지식을 전달하는 성능을 향상하기 위해 student의 학습과정에 대한 피드백을 활용하는 방법이다. 그리고 meta-learning의 최적화를 위해 bi-level learners의 학습을 align하는 pilot update를 제안한다. 

 

 

 

 저자들의 방법은 student 중심 학습론을 기반으로 하여, 기존 training set에서 분할된 별도의 data set인 "quiz set"에서, student model의 성능을 기반으로 한 output을 teacher가 조정하는 것이다.

 training 시작에 앞서, 각 training 단계에서 student $S$에서 $S'$을 복사하고 common knowledge distillation loss로 $S'$을 업데이트한다. 이 과정을 "teaching experiment" 라고 한다. 이 단계에서 우리는 quizzed한 experimental student인 $S'$를 얻게 된다. 그 후, quiz set에서 샘플을 얻고, 이 샘플들에서 $S'$의 loss를 계산한다. 저자들은 이 loss를 피드백 신호로 사용하여 2차 도함수를 계산하고 gradient descent를 수행하여 teacher를 meta-update한다. 마지막으로 $S'$을 버리고 같은 teacher를 $S$로 training batches에 distill하기 위해 업데이트한다. 

 meta learning을 사용한 것은 teacher model이 완전히 미분 가능한 방법으로 student에서 피드백을 받기 위함이다. 

 

 

 

 Experiment

 

 

 

 

 

Conclusion

 

 저자들의 큰 핵심은 본래의 방식에서 역으로 생각하였다.

 기본적으로 많이 아는 teacher-student model은 student 모델을 계속해서 업데이트를 하지만, 이 논문은 teacher model이 S의 복제인 S'에 맞추어 업데이트가 되고, 업데이트가 된 T는 S에게 KD를 수행하는 것이다. 꽤나 괜찮은 발상이어서 논문 자체의 내용도 재미있었다.