첫 번째로 리뷰할 논문은 Google AI에서 제안한 Decentralized Learning with Multi-Headed Distillation이라는 논문입니다. 제목에서 느껴지듯이 Multi head를 통한 distillation이라는 점과 Google에서 제안한 논문이라는 점에 궁금해서 읽게 되었습니다. 간단하게 리뷰해도록 하겠습니다!
리뷰한 내용의 부족한 점, 잘못된 점에 대한 피드백은 언제든 환영합니다!
Abstract
본 눈문에서는 개인 데이터를 가진 에이전트들이 서로의 데이터, 가중치를 공유하지 않고도 학습 가능한 새로운 Decentalized learning을 제안합니다. 특히, 그들은 개인 에이전트들의 multiple auxiliary head들이 Knowledge distillation을 통해 성능 향상을 가져올 수 있다고 주장합니다.
Introduction
기존의 decentralized learning의 대표적인 방법으로는 연합 학습 (Federated Learning)이 있습니다. 연합 학습은 각 에이전트들의 모델들이 협력하여 글로벌 모델을 학습시키는 학습 방법입니다. 대부분의 연합 학습은 FedAvg(https://arxiv.org/pdf/1602.05629.pdf)방법에 기반하고 있는데 이 방법에는 2가지의 단점이 있습니다. 그들은 gradient aggregation을 위해서 호환 가능한 weight space를 가져야 합니다. 그리고 그들은 일반적으로 중앙 집중식 방식으로 학습을 하다 보니 각 에이전트들이 오랜 기간 isolation 될 수 없습니다. 이와 같이 기존의 방법들은 여러 가지 제약을 가진다고 주장합니다.
또한 그들은 distillation 방법은 오직 특정 sample에 대한 활성화를 표현하므로 communication에 있어서 정보의 bit가 감소할 수 있고 서로 다른 아키텍처에서의 output과 representation 간의 knowledge를 교환할 수 있다고 이야기합니다. 따라서 distillation 방법은 decentralized learning에서 기존의 방법보다 더 적은 제약을 가지고 적합합니다.
Method
Intro에서 간단하게 설명했듯이 본 논문에서는 Auxiliary head를 사용합니다. 위의 Figure 2에서는 총 3개의 head를 사용하고 있습니다. Main은 기존의 Client가 가지는 head이며 Aux1, Aux2는 새로 추가된 head입니다. 그림과 같이 서로의 Auxiliary head가 knowledge distillation을 수행하고 있음을 보여줍니다. 아래에서는 이 distillation을 위해 어떤 loss를 사용했는지 보여줍니다.
첫 번째 loss term은 Embedding distillation loss로 $\Psi$는 local에서 계산된 결과이며 $\Phi$는 같은 sample $x$의 입력과 연결된 다른 clients에서 계산된 결과입니다. 따라서 해당 loss term은 같은 입력 sample $x$로부터 각 client들의 embedding 결과의 차이를 최소화하여 같은 sample에 대한 각 client들의 결과가 compatibility 하도록 강요합니다.
두 번째 loss term은 Prediction distillation loss입니다. 해당 loss term은 기존의 KL diversece loss(우항)과 유사하게 구성되어 Auxiliary head들의 prediction과 main head의 KL diversence입니다. 또한 여기에 $Q$는 가장 높은 confidence를 가지는 client에 대해서 1 값을 가지도록 선택하는 항입니다. 이를 통해 더욱 명백한 선택을 할 수 있다고 이야기합니다.
마지막으로는 Self-ditillation with multiple auxiliary head loss로써 각 auxiliary head들 간의 KL diversence로 위의 식과 유사합니다. 따라서 최종적인 loss term은 아래와 같이 위의 loss를 모두 합쳐서 표현한 식입니다.
Experiments
실험에 앞서 본 논문도 Federated Learning에 기반한 세팅을 가지고 있습니다. 위의 그림은 일반적인 Federated Learning의 알고리즘의 학습 방법으로 각 Clients들은 자신들 만의 데이터를 가지고 있다고 가정하며, 이들은 Server에서 aggregation 되어 public sample로 업데이트됩니다. 본 논문에서는 public sample로써 전체 데이터의 10%를 사용하며 나머지 데이터는 각 Client들의 private data가 됩니다.
Label assignment
Dataset에서 각 client들은 subset $l_{i}$에 할당됩니다. 이때, skwness parameter $s$가 등장하는데 $s$는 각 client에 배정되는 label을 조정하기 위한 값입니다. 그래서 primary label $l$을 가지는 client는 다른 client들에 비해 해당 label $l$이 할당될 확률이$s+1$배 더 높습니다. 만약 $s= 0$이면, 모든 client에게 모든 샘플들이 동등하게 할당되어 있으며 $s\rightarrow \infty$의 경우에는 label $l$을 primary label로 가진 client에게만 label $l$이 할당됩니다.
위의 실험에서 가장 눈에 띄는 결과는 $v_{aux}$의 유무가 즉, prediciotn distillation이 embedding distillation보다 더 많은 영향을 준다는 것을 알 수 있습니다. 결국 Auxiliary head가 성능 향상에 중요한 역할임을 보여줍니다.
또한 Auxiliary head에 따른 성능 그래프에서도 $s=0$인 경우에 3 heads 까지는 성능이 향상했으며 $s=100$의 경우에도 shared dataset에 대해서 성능이 향상됨을 보여줍니다.
본 논문에서는 Client의 Auxiliary head 간 distillation이 궁극적으로 shared model의 성능을 올릴 수 있다고 주장합니다. 매우 간단한 모델 세팅과 기존의 Knowledge distillation의 특성을 잘 살려 Federated Learning에 적용한 논문이 아닐까 싶습니다. 개인적인 생각으로는 각 Clinet들이 서로 각자의 private dataset으로 학습하면서 이들 Auxiliary head들 간의 다양한 Knowledge가 공유된 게 성능 향상에 중요한 역할을 하지 않았나 싶습니다.
비슷하게 기존의 classification에서 Multiple head를 통해 성능 향상을 가져온 Collaborative Learning for Deep Neural Networks 논문도 참고해 보시면 좋을 것 같습니다.