논문 리뷰 Similarity-Preserving Knowledge Distillation

ICCV 2019 에 발표된 논문 Similarity-Preserving 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 를 달성 할 수 없기 때문인 것 같습니다.