Prune Your Model Before Distill It

https://arxiv.org/abs/2109.14960

Prune Your Model Before Distill It

https://github.com/ososos888/prune-then-distill

GitHub - ososos888/prune-then-distill

 

Abstract

최근 연구에서는 student-friendly teacher를 knowledge distillation에 쓰는 것이 많이 연구되고 있다.

 

이 논문에서는 "prune, then distill"이라는 연구를 제안하면서, pruned teacher가 unpruned teacher보다 더 좋은 성능을 낼 수 있음을 보이며, distillation에서 regularizer 역할을 하는 것을 이론적으로 설명했다.

 

 

Introduction

On the efficacy of knowledge distillation ICCV(2019)라는 논문에서, student의 capacity가 작을 경우에는 less-trained teacher가 더 잘 가르친다는 것을 보였다.

 

최근에는 "잘 학습되고 큰 teacher network가 student에게 더 distillation을 잘 할 것이다." 라는 것과 반대로, student의 capacity가 작을 경우에는, student-friendly한 teacher가 더 distillation을 잘한다는 것을 제안하고 있다.

 

이 논문에서는 pruning을 이용해서 teacher를 student friendly한 형태로 만들고, 이를 distillation에 활용하는 것을 제안하는데,

진짜 pruning이 distillation 성능에 도움을 주는 것이 맞을까?를 대답하기 위해, 3가지 step으로 구성했다.

1. train the teacher network.

2. prune the teacher network.

3. distill the pruned network to the smaller network.

 

또한 3가지 실험으로 검증했다.

1. distill vgg19 to vgg11

2. self-distillation vgg19 and resnet18 

3. distill resnet18 to vgg16 and mobilenetV2

 

3가지 모두 unpruned teacher보다 pruned teacher가 더 좋음을 확인했다.

 

 

또한, KD는 label smoothing regularization의 관점에서도 볼 수 있는데, smoothing regularization으로 학습한 teacher는 original teacher보다, smoother label을 제공하는 것을 알아냈다고 한다. 그러니까, smoothing regularization으로 학습한 teacher를 distillation에 사용하면, LSR와 동일하다는 것을 알아냈다고 한다.

 

이와 마찬가지로 pruned 역시 sparsity-inducing regularizer로 볼 수 있으니까, pruned teacher는 distillation을 해줄때 sparsity-inducing regularizer효과를 줄 것이다라고 한다.

• 음... smoothing regularization처럼 pruning의 효과도 distillation에서 영향을 준다고 주장하는 건가?

무튼 이러한 pruned teacher가 더 좋은 kd를 줄 수 있다는 점에서, novel network compression scheme을 제안했다.

이 논문에서는 student architecture구성에 특이점이 있는데, teacher와 depth는 같지만, parameter가 적도록 구성했다.

왼쪽(기존KD), 오른쪽(논문이 제안한 pruned kd)

 

 

Prune, then Distill

 

pruned kd와 unpruned kd를 비교해보자!

vanila kd에서 loss alpha는 0.95, temperature는 10을 사용했다. (appendix 참조)

 

vgg19 to vgg11

대충 결과를 보면 Unpruned teacher보다 성능이 좋다. DBL은 layer별로 채널이 두배인 VGG19

 

self-distillation

same architecture에서의 Independent, KD, Pruned KD 성능비교.

마찬가지로, pruned teacher의 kd 성능이 좋았다. 

• 같은 모델인데도 성능이 더 좋네? 이유가 뭐지?

 

 

위의 실험들에서 주목할만한 점은, teacher의 acc가 높다고 무조건 distillation이 잘 되지 않는다는 점이다.

pruned teacher가 unpruned teacher에 비해서 acc가 낮아도, distillation이 잘 되어 student의 acc가 높은 것을 볼 수 있다.

 

 

Pruned Teacher as a Regularizer

"Weighted coin: Duality of adaptive dropout and regularization". NeurIPS (2021)

라는 논문에서, pruning은 sparsity-inducing regularization을 통해 empirical risk minimization problem을 해결했다고 한다.

그러므로, pruned teacher를 사용한 distillation도 이런 regularization효과가 distillation이 될 것이라고 주장한다.

 

 

Transferring Knowledge of Sparsity

이 논문에서는 unstructured pruned network에서 학습하는 network compression framework를 제안했는데,

제일 어려운 부분은 pruned teacher로부터 효율적으로 학습이 가능한 student architecture를 설계하는 것이라고 한다.

 

그림1처럼 teacher랑 같지만, 각 layer별 chennel이 줄어들어서, 남은 파라미터의 수가 대략 같도록 student를 구성한다.

대략적인 방법은 다음과 같다.

1. teacher network를 학습한다. 
2. 학습한 teacher network를 갖고 unstructured pruning을 진행한다.
3. pruned teacher network의 layer별 sparsity를 고려하여 student 아키텍쳐 설계한다.
4. pruned teacher network로 부터 student network에게 knowledge distillation을 한다.

 

위 framework에서 특별한 pruning method나 distillation 방법은 필요없다.

pruning을 진행할때는 LR rewinding을 했고, pruned teacher로 부터 vanila KD를 적용했다.

 

지금 위의 방식은 sparse network에서 채널 수가 더 적은 student network로 distillation하면서 채널 수를 더 줄이는 방법인데, 이는 residual distillation(residual network에서 불필요한 부분을 제거하는 것)과 유사하다고 한다.

 

distillation시 알맞은 student network를 찾으면서 채널 수를 줄이는 방법이니까 structured pruning의 관점으로도 볼 수 있다.

 

 

Experiments