[CVPR2018] Deep Mutual Learning

이번에 리뷰할 논문은 Mutual Learning 입니다. Mutual Learning에 대해서 들어보신 분들은 아마 잘 아실 것 같은데 저는 최근에 알게된 용어이긴 합니다. 리뷰를 시작하기 전에 간략하게 소개드리면 Mutual Learning과 가장 유사한 개념으로는 Knowledge Distillation이 있겠네요.

그럼 리뷰 시작합니다.

Introduction

컴퓨터비전 분야든 다른 분야든 상관없이, 인공지능은 상당히 놀라운 성과들을 보여주고 있습니다. 하지만 네트워크가 점점 더 깊어지거나, 넓어지거나, 또는 너무 방대한 양의 학습 파라미터들이 존재해야만 좋은 성능을 기대해볼 수 있습니다.

하지만 이렇게 방대해진 네트워크는 작은 메모리를 필요로하는 실제 application에서 적용하기 힘들뿐더러 실시간 추론도 매우 느린 속도로 진행되기 때문에 네트워크의 크기를 무작정 키울 수는 없습니다.

그래서 이전부터 네트워크의 크기는 최대한 줄이면서 그로 인해 발생할 수 있는 성능의 감소는 최소한으로 하는 방법론들이 여러가지 방면으로 제안되어 왔는데, 그 중 하나가 바로 Knowledge Distillation(KD)입니다.

KD의 경우 제가 알기론 제프리 힌튼님께서 14년도에 제안하신 걸로 알고 있는데, 아마 많은 분들이 알고 계실 것이고 저도 옛날 리뷰에서 한번 다뤘던걸로 기억합니다. 그럼에도 불구하고 혹시나 모르시는 분들을 위해 다시 한번 간략하게 설명하고 넘어가겠습니다.

그림1은 KD의 과정을 매우 잘 나타내는 그림입니다. 보시면 네트워크가 2개 존재하는데 남색깔 모델은 Teacher, 초록색 모델은 Student라고 적혀있는 모습이네요.

네 KD는 지식을 전이?하는 개념으로 저희가 가장 흔하게 사용하는 Ground Truth를 활용한 Supervised Learning 외에도 이미 잘 학습된 매우 크고 깊은 Teacher Model의 지식을 비교적 작고 가벼운 student model에게 전달해주는 것입니다.

그림1에서 아래 hard label이 ground turth를 의미하며 student의 결과와 직접적으로 loss를 계산하게 되며, 동시에 위에 Teacher model이 동일한 입력 x를 통과해 출력한 soft label과 student의 soft label을 비교함으로써 학습이 진행됩니다.

저기 위에 예시에서는 단순히 최종적인 class label에 대해서만 비교하여 loss를 계산하였지만, 중간 단계의 feature map 사이에서도 student의 feature가 teacher를 닮아가도록 비교하는 loss를 추가하는 방법들도 많이 존재합니다.

아무튼간에 핵심을 요약하자면 KD는 이미 사전에 잘 학습된 Teacher model이 존재한다 라는 것입니다. 즉 Teacher model은 사전에 이미 학습이 되었기 때문에 student model을 학습시켜야 할 때는 freeze 되어 어떠한 weight의 변화가 없어야만 합니다. 이러한 방식을 논문에서는 이를 one-way knowledge transfer 라고 정의합니다.

반면에, 저자의 논문 제목이자 새롭게 제안하는 방법론인 Mutual Learning는 KD랑 무엇이 다를까요? Mutual Learning은 직독직해 하자면 바로 상호학습이라는 뜻으로 그림2와 같이 teacher-student 관계가 아닌 student-student 관계를 의미하며 KD랑 달리 two-way knowledge transfer라고 정의합니다.

즉 서로가 동시에 학습을 함으로써 부족한게 있으면 보충해주는 협력관계라는 것이죠. 이게 얼핏보면 student-student model이 동시에 학습을 하게 되니 이거 경쟁하면서 학습하는건가?? 라고 생각하실 수 있는데, 보통 그런 경쟁적인 구도하면 GAN이 떠오르실 겁니다.

하지만 방금 전에도 얘기했다시피, 이 두 네트워크는 서로 경쟁이 아닌, 협력관계임이 중요합니다. 참고로 이러한 협력관계는 아무것도 학습이 안된 student-student model 뿐만 아닌, teacher-student 관계에서도 충분히 서로의 성능 향상에 도움을 준다고 합니다. 그리고 저자는 이러한 Mutal Learning이 KD보다 더 우수한 성능 향상을 기대할 수 있으며 그 외에도 여러 장점들이 있다고 합니다.

음 그럼 Mutual Learning이 왜 KD보다 더 우수한 성능을 보여주는 것일까요? 상식적으로 봤을 때 KD의 경우 더 큰 모델로 사전학습된 모델이 support를 해주는 것인데, 어째서 작고 가벼운 모델 둘이 동시에 학습하는 것이 더 우수한 결과를 보이느냐 말입니다.

그것에 대해서 저자는 이론적으로 완벽하게 증명하지는 못한다고 합니다. 하지만 여러가지 추측들에 대해서는 논문에서 주장하고 있습니다. 저자의 기본적인 주장은 집단사고(groupthink)에 있습니다.

동일한 모델, 동일한 초기 조건, 동일한 GT, 동일한 loss를 사용하게 되면 당연히 두 모델의 feature representation이 동일할 수 밖에 없을 것입니다. 하지만 서로 다른 초기조건을 가진 체 동일한 학습을 하게 되면 서로 다른 표현방식을 배울 수 있다고 합니다.

또한 이 student 집단들이 다음으로 나올 가장 가능성이 높은 클래스에 대해서 동료들끼리 찾아내고 매칭함으로써 student model의 posterior entropy가 높아진다.. 고 하는데 posterior entropy에 대해서는 잘 알지 못하지만 결과적으로 이러한 이점은 모델들이 더욱 강력한 global minima에 수렴될 수 있도록 도와주며, testdata에 대해 더 일반화된 결과를 보여줄 수 있도록 해준다고 합니다.

요약하자면 Mutual Learning은 동료모델과의 협동적인 학습을 통해 단순하지만 매우 효율적으로 네트워크들의 일반화 능력을 향상시킨다고 볼 수 있으며, 저자는 해당 논문에서 CIFAR100과 instance recongnition problem인 Market1501데이터 셋에서 person re-dienticiation task를 수행함으로써 Mutual Learning의 효과를 실험적으로 증명하였습니다.

Formulation of Deep Mutual Learning(DML)

그러면 먼저 두 네트워크 집단에 대하여 DML에 대한 수식화를 설명하겠습니다.

M개의 클래스들로부터 N개의 샘플들을 X = \{ x_{i} \}^{N}_{i=1}라고 했을 때 이와 대응되는 label을 Y = \{ y_{i} \}^_{N}_{i=1} 라고 놓을 수 있습니다. 이 때 student network \Theta_{1} 이 제공하는 x_{i}에 대하여 class m의 확률 값은 아래 수식 1과 같습니다.

z^{m}는 \Theta_{1} 내에서 softmax layer를 타고 나온 결과 값을 의미합니다.

만약 multi-class classification을 수행한다면, \Theta_{1} 의 목적 함수는 GT label과 predicted value 사이에 cross entropy loss로 정의될 수 있습니다.

일반적인 supervised loss는 네트워크가 학습 데이터 셋에 대하여 올바른 label을 예측할 수 있도록 학습시킵니다. 여기서 \Theta_{1} 의 테스트 데이터에 대한 일반화 성능을 향상시키기 위해서는 또 다른 동료 네트워크 \Theta_{2} 를 학습 단계에서 추가합니다.

이 때 두 네트워크의 예측값 p_{1}, p_{2}에 대하여 정량적인 매칭을 하기 위해서 저자는 Kullback Leibler(KL) Divergence를 사용합니다. p_{1}, p_{2}에 대하여 KL distance를 계산하는 식은 아래와 같습니다.

KL Divergence에 대해서는 예전에 김형준 연구원님이 리뷰로 다룬 적이 있으므로 해당 글을 참고해주시면 좋을 것 같습니다.

아무튼 두 네트워크의 전체 loss function은 각각 식(5)와 (6)과 같습니다.

각각의 네트워크들은 모두 GT label에 대해서 올바르게 예측했는지를 학습(supervised loss L_{C})할 뿐만 아니라 동료 네트워크들이 추정한 확률값들이 서로 일치하는지에 대해서도 학습하게 됩니다.(KL mimicry loss)

여기서 KL divergence 기반의 mimicry loss는 비대칭적이며, 그러므로 두 네트워크에 대한 해당 loss는 서로 다른 값을 가지게 됩니다. 그래서 KL divergence를 대신하여 대칭적인 특성을 지니는 Jensen-Shannon Divergence loss를 저자는 사용하기도 했었는데, 저자는 결국 실험적으로 대칭이냐 비대칭이냐는 유의미한 실험적 차이를 나타내지 못했다고 합니다.

Optimisation

DML은 반복적으로 얘기하는 것이지만, 두 모델이 서로 동시에, 협동적으로 최적화된다고 하였습니다. 일반적인 네트워크가 그러하듯, mini-batch 단위로 모델이 업데이트 되며, 이 때 저자는 특정 모델만 먼저 업데이트 되는 것이 아닌 두 모델이 동시에 parameter를 업데이트하였다고 합니다.

또한 동일한 mini-batch를 입력으로 하여 학습을 하며, 전체적인 학습 과정은 아래 알고리즘을 통해 확인하실 수 있습니다. 참고로 GPU가 2개 존재하게 되는 경우, 아래 알고리즘에서 step 1,2는 첫번째 GPU가, step 3,4,는 두 번째 GPU가 병렬적으로 수행이 가능하다고 합니다.

Extension to Larger Student Cohorts

저자는 student cohort 내에서 1개가 아닌 더 많은 수의 네트워크들을 구성하여 학습을 시켰다고 합니다. 즉 위에서 제가 설명드린 내용들은 \Theta_{1} 과 \Theta_{2} 가 서로 mutual learning을 하는 것이라면, 이번에는 \Theta_{1} ... \Theta_{K} (K \geq 2) 가 되는 것입니다. 그리고 이 경우의 목적 함수는 다음과 같습니다.

수식 10은 K개의 네트워크가 있을 때 사용되는 목적함수로, 각각의 student network는 나머지 K-1개의 네트워크들을 각각 하나의 teacher network로 보고 mimicry loss를 계산하게 됩니다. 아까 위에서 설명한 수식 5,6은 K가 2일 때의 예시라고 보시면 되겠습니다.

아 그리고 저자가 KL divergence를 통해 계산한 mimicry loss에 대하여 K-1로 나누어준 것에 대해서는 일단 모델이 Ground Truth label에 대해 더 비중을 높인 체로 학습되기를 원해서라고 합니다.

아무튼 정리하자면 수식 10의 경우에는 K번째 네트워크가 K-1개의 남은 모든 네트워크들과 각각 KL distance를 계산하여 평균을 취한 것으로 볼 수 있습니다. 반면에 저자는 K-1개의 네트워크 전체를 하나의 single teacher로 볼 수도 있다고 합니다.

수식 11을 보시면 아시겠지만, k번째 네트워크의 확률값 p_{k}가 평균확률값 p_{avg}과 KL distance를 한번만 계산해주는 모습입니다. 그리고 이 평균확률값은 말 그대로 K-1개의 네트워크들의 확률값을 다 더해 평균을 계산한 것이죠.

이렇게 하게 될 경우 K-1개의 네트워크들은 일종의 하나의 앙상블 모델로 볼 수 있게 되는 것입니다. 저자는 이렇게 수식 10을 통한 K-1개의 분리된 teacher 방식과 수식11을 통한 개의 앙상블 teacher 방식 중에 무엇이 좋은지에 대하여 실험을 진행하였는데, 결과부터 말씀드리면 K-1개의 분리된 teacher를 사용하는 방식이 더 좋은 성능을 보였다고 합니다.

이것에 대하여 저자는 수식11 과정 중 확률값의 평균을 구하는 model averaging step이 teacher model의 posterior probabilities를 true class에 더욱 강조(peak)한다고 합니다.

그 결과 전반적인 class에 대하여 posterior entropy가 감소하게 되는데, 이는 posterior entorpy 값을 높게 만듦으로써 더 강인한 결과를 생성해내려는 DML의 목적과는 모순적인 과정이기 때문에 성능이 감소한다고 합니다.

posterior probabilities? posterior entropy? 와 같은 단어들은 저에게도 너무 생소하여 이해를 돕기 위한 설명을 추가하기에는 힘드네요. 개인적으로 공부하여 보충해보도록 하겠습니다.

아무튼 요약하자면, K개의 student model을 추가하여 mutual learning을 할 때는 K-1개의 student model을 하나의 앙상블 teacher로 적용하는 것보다는 K-1개의 분리된 teacher model로 사용하는 것이 entropy 상승 관점에서는 적절하여 성능도 더 향상된다고 보시면 되겠습니다.

Extension to Semi-supervised Learning

논문에서 제안하는 DML은 semi-supervised learning에서도 손쉽게 적용이 가능합니다. 방식은 매우 간단한데, 사용가능한 labelled data(GT)에 대해서는 cross-entropy loss를 활용하여 supervised loss로 학습을 하며, label이 존재하지 않는 데이터들을 포함한 모든 학습 데이터에 대해서는 KL distance를 기반으로 한 mimicry loss를 사용하는 것입니다.

KL distance를 활용한 loss는 label에 상관없이 사용할 수 있기 때문에 semi-supervised learning에서 아주 손쉽게 적용할 수 있게 되는 것이죠. 이러한 semi-supervised learning에서의 목적함수는 수식 12로 나타낼 수 있습니다.

여기서 \mathcal{L}, \mathcal{X}는 각각 labelled data와 total data(labelled + unlabelled data)의 집합을 의미합니다. 즉 label이 있는 데이터는 cross-entropy, 그 외에는 KL distnace loss를 사용하는 것입니다.

Experiments

저자는 총 4개의 데이터 셋에 대해서 실험을 진행하였는데, 대략적으로 ImageNet, CIFAR-10, CIFAR-100, Market-1501 데이터 셋이 이에 해당합니다.

또한 논문에서는 실험을 위해 총 4개의 네트워크를 사용하였는데 각각의 학습파라미터 크기는 위에 표에 나와있습니다. 좌측 2개가 가벼운 모델, 우측 2개가 무거운 모델로 아마 무거운 모델들은 teacher network 역할로 사용할 것 같네요.

표2는 CIFAR에 대한 결과입니다. 가벼운 모델의 조합(ResNet-32 + ResNet-32)이나 일반적인 Knowledge Distillation 조합(WRN-28-10 + ResNet-32) 또는 둘다 무거운 조합(WRN-28-10 + WRN-28-10)에 대해서 모두 mutual learning이 효과가 있었다는 것을 확인할 수 있습니다.

참고로 각각 따로 학습시킨 결과를 지칭하는 Independent와 mutal learning한 결과인 DML의 성능 차이를 뺸 것이 DML-Ind이므로 DML-Ind의 차이값이 크다는 것은 mutual learning의 효과를 잘 보고 있다는 것입니다.

CIFAR-10과 달리 CIFAR-100의 경우 작은 모델 둘이서 mutual learning을 하더라도 큰 성능 개선이 발생하는 모습이 보이며 큰 모델과 작은 모델끼리의 조합일 때 작은 모델의 성능 향상이 가장 많은 상승을 보이는 듯 합니다.

Comparision with Model Distillation

다음은 Mutual learning이 Knowledge Distillation보다 좋은가?에 대한 실험 결과입니다.

표3을 보시면 1 distills 2 라는 항목이 보이실텐데, 저 부분이 바로 사전학습된 Net1을 가지고 Net2를 distillation한 결과로 볼 수 있습니다.

Distillation을 수행하게 되면 Net2의 성능이 혼자 학습하였을 때 보다 당연히 향상된 것을 확인하실 수 있으시겠지만, Mutual Learning을 통해 함께 학습하게 되는 경우 더 큰 폭의 성능 향상을 얻을 수 있습니다.(Market-1501 데이터 셋 기준 DML은 50.15에서 58.47로 8.32 mAP가 증가)

결론

Knowledge Distillation이라는 컨셉은 정말로 단순하지만 성능은 제법 크게 향상시킬 수 있다는 점에서 메리트가 있는 방법론이었습니다. 하지만 미리 사전학습을 시켜야한다는 점과 고성능의 teacher 네트워크를 준비해야한다는 점에서는 아무래도 명확한 단점이 존재하였죠.

반면에 Mutual Learning은 작은 네트워크들끼리 뭉쳐서 학습을 하더라도 성능 향상을 기대해볼 수 있으며, 가장 좋은 점은 아무래도 사전학습을 미리 시킬필요가 없다는 것이겠네요. 게다가 KD보다 더 큰 성능도 기대해볼 수 있다는 점에서 매우 좋은 컨셉인 것 같습니다.

CVPR 리버틀 기간이라 논문 속 실험 결과에 대해서는 다 다루지 못하였습니다. 해당 논문은 mutual learnig의 근본이면서 동시에 내용도 상당히 간단하여서 직접 읽어보시는 것을 추천합니다.