논문 리뷰 Improved Knowledge Distillation via Teacher Assistant

Improved Knowledge Distillation via Teacher Assistant 은 AAAI 2020에 제출된 논문입니다. 저자 중에는 딥마인드 소속도 있습니다.

key Idea

이 논문의 핵심 아이디어를 잘 설명하고 있는 그림입니다. 저자는 기존의 Knowledge Distillation 의 Teacher, Student(이하 T, S) 관계 사이에 중간 크기의 모델인 Teacher Assistant(이하 TA) 모델을 사용하여 여러번 Distillation 을 반복하면 S 의 성능을 더 높일 수 있음을 보여주고 있습니다.

Model Structure

CNN 과 ResNet 을 사용하여 실험을 하는대, ResNet 은 ResNet 논문에서 제시한 형태 그대로를 사용하고, CNN 은 아래와 같이 데이터셋 별로 2개의 모델 구조를 가지고 실험합니다.

Intuition

저자는 S 의 layer 수를 2개로 고정하고 T 의 layer 수를 4, 6, 8, 10 으로 늘려가며 S 의 성능을 관찰했을때, T 가 일정 수준 이상으로 커지면 T 의 성능은 좋아지지만 S 의 성능은 오히려 떨어지는 모습을 보여주며 T 의 크기가 큰 것이 S 의 좋은 성능을 보장하지 않음을 보여주고 있습니다. 이러한 관찰로부터

1. 충분히 큰 T 를 사용하여 S 의 성능을 높이고 싶으나,

2. T 와 S 의 차이가 클 경우 S 가 T 를 잘 mimicking 하기엔 capacity gap 이 너무 벌어지고

3. T 가 더 정확해질수록 T 의 soft label output 이 점점 hard label output 처럼 변하는 문제를 생각할 수 있습니다.

1 을 만족하면서 2, 3 을 해소하기 위해 T 와 S 사이에 T 보다 작은 TA 를 넣어 T → TA 로 Distillation 하여 잘 학습된 TA 를 만들고 다시 TA → S 로 Distillation 하여 최종적으로는 좋은 성능의 S 를 만들고자 합니다.

Experiment

결과적으로 성능이 좋으니 논문이 accept 되겠죠? 자잘한 실험결과는 넘겨두고, 가장 큰 contribution 이라고 생각되는 실험만 정리하겠습니다.

저자는 Multiple Teacher Assistant Knowledge Distillation 을 제안하고 있습니다. 위 Figure 의 각 열은 layer 수를 의미합니다. 첫번째 열의 네번째 행에 있는 oval(56.19) 에서 두번째 열의 네번째 행에 있는 oval(56.75) 로 향하는 화살표

oval(56.19) → oval(56.75)

는 layer 10 인 T 에서 layer 8 인 S 로 Distillation 했을 때의 성능(정확도)이 56.75 라고 해석하면 됩니다.

결론적으로 TA Distillation 을 하는것이 하지 않는것보다 성능이 좋고 같은 layer 수를 가지고 있다면 여러번 하는것이 그렇지 않은것보다 더 좋음을 보여주고 있습니다.

이 외에도 저자는 왜 TA Distillation 이 먹히는지 수학적으로 증명하고, 실험적으로도 보여주고 있는데 이에 대한 내용은 요청이 있을시 추가작성 하겠습니다.