ICCV 2019 에 발표된 논문 Similarity-Preserving Knowledge Distillation 을 리뷰합니다.
이 글의 내용을 이해하기 위해서는 Knowledge distillation 이 무엇인지 이해가 필요합니다. 아래 링크를 참조해주세요.
Knowledge Distillation 이란?
Similarity-Preserving Knowledge Distillation
논문의 Abstract 에서 KD 의 목적, 이 논문의 Contribution, Key Idea, 효과 혹은 결과를 파악할 수 있습니다. 앞에서 KD 의 목적은 이미 설명했습니다. 이 논문의 Contribution 은 Key Idea 에 대한 설명 이후에 자세하게 설명하는 것이 더 이해를 쉽게 할 수 있습니다.
Key Idea
이 논문의 핵심 아이디어는, 의미론적으로 유사한 이미지에 대한 Activation map 은 비슷하더라는 관찰에서 시작합니다. 이를 예를 들어 설명하면, '개' 라는 클래스의 이미지인 골든 리트리버와 래브라도 리트리버 이미지를 모델에 입력했을때, 중간 계산과정의 결과인 Activation map 을 시각화해보면 비슷하다! 라는 느낌입니다.
골든 리트리버와 래브라도 리트리버
아무래도 두 견종은 거의 비슷하게 생겼기 때문에, Convolution filter 를 통과했을 때 생성되는 Feature map 과 Activation function 을 통과했을 때의 Activation map 이 유사할것이라는 것은 직관적으로 이해할 수 있습니다.
이를 행렬을 통해 시각화해서 표현해보면,
위 그림에서 빨간색 행렬로 표현된 서로 같은 이미지의 Activation map "Alpha" 를 Alpha * Alpha^T 를 하여 나오는 행렬의 Element sum 은
서로 다른 이미지의 Activation map "Alpha" 와 "Beta" 의 같은 계산 결과보다 Element sum 이 더 작을것이라는 것을 알 수 있습니다. 당연하죠, 같은 행렬은 0의 위치가 같기 때문입니다.
이와 같은 원리로 같은 클래스의 서로 다른 이미지는 Activation map 이 비슷할 것이기 때문에 상대적으로 서로 다른 이미지의 Activation map 보다 위 계산결과처럼 값이 클 것이다라고 생각할 수 있습니다.
위 Figure 에서 각 열은 개별 이미지를 의미합니다. CIFAR-10 을 사용했기 때문에 1번부터 1000까지는 비행기 클래스의 이미지라고 생각할 수 있죠. 행은 모델의 각각의 채널들을 의미하는데, 각 채널이 같은 클래스의 이미지에 대해서 Activation map 이 비슷한 것들을 볼 수 있습니다.
이 논문은 Teacher network 의 Activation map 이 비슷한 성질을 Student network 에 전달(혹은 학습)시킬 수 있다면 Student network 의 성능을 올릴 수 있지 않을까? 라는 핵심 아이디어를 가지고 출발합니다.
Contribution
저자는 새로운 KD loss function 을 제안한 것이 이 논문의 Contribution 이라고 주장합니다. 이 논문의 Loss fucntion 은 다음과 같은데,
왼쪽 항은 이전 KD Loss function 에서도 사용되는 Loss term 이니, 오른쪽 Loss term 만 이해하면 이 논문의 Loss function 을 이해한다고 할 수 있습니다.
오른쪽 Loss term 의 유도과정을 살펴보겠습니다. l 번째 레이어의 Activation map 은
로 표현할 수 있습니다. Teacher 와 Student 의 모델 크기가 다르기 때문에 Channel, height, width 는 서로 다르지만, batch size 는 동일합니다.
Activation map 을 Reshaping 하면
로 새롭게 표현할 수 있고, 새로운 Q 를 Transepose 해서 Q 와 곱한 뒤 각 열을 L2 normalize 하면
각 모델의 Activation similiarity 를 비교할 수 있는 행렬인 를 구할 수 있습니다. 이를 시각화해보면 아래와 같습니다.
잘 학습이 된 Student network 라면 Teacher network 와 Activation similarity matrix G 가 위 그림처럼 유사하게 나올 것입니다.
이제 계산된 G 를 이해했기 때문에 새로운 Loss term 을 설명할 수 있습니다.
새로운 Loss term 은 위와 같은데, G 간의 차이를 Frobenius norm 을 통해 계산하고 있습니다. 앞에서 b^2 로 나누어 주는 것은 Batch size 에 상관 없이 Loss 를 계산하기 위해 1 Batch 와 같도록 만들어주기 위함입니다.
Teacher 와 Student 의 모델 Depth 가 다른 경우는 Block 단위로 모델의 구조를 생각하고 같은 Depth Block 의 마지막 Layer 의 출력을 사용한다고 합니다. 어떤 Depth 에서 Activation similiarity 를 비교할 것인가는 사용자의 몫으로 남겨두고 있지만, 입력에 가까운쪽보다는 출력에 가까운쪽의 Layer 를 사용하는 것이 더 좋을 것이라고 말하고 있습니다. 입력에 가까운 하부 Layer 는 Edge 같은 모든 클래스가 가지고 있는 특징에 Activation 됬다고 생각할 수 있기 때문입니다.
Experiment and Conclusion
이 논문에서 실행한 실험에 대해서 살펴보겠습니다.(논문과 같은 순서는 아닙니다.)
저자는 새로운 Loss term 이 증가할 수록 Student model 의 Error 가 증가하는 경향성이 있음을 통해 Loss term 이 잘 설계되었다고 주장하고 있습니다.
저자는 Depth 와 Width 가 서로 같거나 다른 경우에 학습이 잘 되는지 확인하였는데, 전반적으로 자신들이 제안한 Loss function 이 우수하다고 보여주고 있으며, 이 Table 에서 특기할만한 점은 초록색, 파란색 원으로 표시한 부분을 보면 두 모델이 같은 모델임에도 KD 를 통해 더 낮은 error 를 달성했음을 알 수 있습니다.
다음 실험에서는 ImageNet 으로 학습한 뒤 DTD 로의 Transfer learning 을 한 결과를 보여주고 있는데, 이 경우도 자신들의 Loss function 이 전반적으로 좋은 성능을 보여주고 있다고 합니다. 이전 실험과 다르게, KD 를 적용한 모델이 적용하지 않은 모델보다 성능이 떨어짐을 색깔로 표시한 영역에서 확인할 수 있지만 이는 DTD 가 5740개 이미지, 47개 클래스로 이루어진 데이터셋이기 때문에, 데이터셋 자체가 짧은 학습으로 낮은 Error 를 달성 할 수 없기 때문인 것 같습니다.
Hyper parameter 인 γ 의 변화에 대해서도 실험을 했는데, γ 에 대해서는 Robust 함을 알 수 있습니다.
기존 KD 방법론으로 알려진 Hinton 교수님의 KD 와 AT 에 자신들의 Loss term 을 결합한 경우에 대해서도 보여주고 있고,
Teacher 와 Student 가 서로 다른 모델일때도 잘 적용됨을 보여주고 있습니다.