본문 바로가기

AI/논문리뷰

[논문리뷰] Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

논문 링크 : https://proceedings.neurips.cc/paper/2020/hash/a1d4c20b182ad7137ab3606f0e3fc8a4-Abstract.html

 

Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

Requests for name changes in the electronic proceedings will be accepted with no questions asked. However name changes may cause bibliographic tracking issues. Authors are asked to consider this carefully and discuss it with their co-authors prior to reque

proceedings.neurips.cc

해당 논문은 2020년 NeurIPS에 발표된 논문으로, Federated Learning(연합학습)과 Split Learning의 장점을 결합하여 FL환경에서 edge device의 computational resource를 최대한 활용할 수 있는 방법을 제안한 논문이다. 해당 논문을 읽은 이유는 연합학습 환경에서 feature vector(=embedding vector)를 서버에 전송하는 사례에 대해 알아보기 위해 읽게 되었다.

 

-Overall framework

 

-Group Knowledge Transfer

1. Reformulation

Figure 2(a)는 FL, (b)는 SL에 대한 간략한 그림을 나타내고 있다. FL에서는 client가 가지고 있는 local dataset을 학습시켜 global(server) model의 최적화 속도를 높임으로써 communication cost를 줄일 수 있다. 하지만, model의 size가 너무 크게 되면, communication cost에 대한 비용이 너무 크고, edge device의 컴퓨팅 자원의 한계 때문에 적합하지 않다. (b)의 SL같은 경우는 large model을 두개의 파트로 분리하여 edge device의 컴퓨팅 자원의 한계에 대한 문제를 해결하지만, mini batch 마다 model을 update하는 과정에서 communication cost에 대한 비용이 너무 많이 발생한다. 따라서 본 논문에서는 FL과 SL의 장점을 결합하여 (c)와 같은 형태로 학습을 수행하는 것을 목표로 한다.

 

따라서 본 논문에서는 server 모델만 optimization하는 기존의 FL과 다르게, 클라이언트 edge device 모델과 서버 모델을 모두 optimization하는 문제로 reformulation한다. 이는 Equation(2)와 (5)와 같이 나타내어진다.

 

2. Advantages and Challenges

해당 방법은 server는 모델의 큰 부분을, client(edge device)는 작은 부분으로 나뉘어서 학습을 수행하기 때문에, 컴퓨팅 자원을 효율적으로 사용할 수 있으며, model의 weight를 보내는 기존의 FL과 다르게 해당 방법은 client의 feature extractor output인 feature map을 전송하기 때문에, communication cost 측면에서도 FL에 비해서 효율적이다.

하지만, client는 소량의 local data만 가지고 있기 때문에, 최적화가 어렵고, client 모델이 잘 최적화 되지 않으면, quality가 낮은 feature map을 가지고 server 모델은 학습을 해야하기 때문에 server 모델 또한 최적화에 어려움을 겪을 수 있다.

 

3. Group Knowledge Transfer (FedGKT)

앞선 Challenge를 해결하기 위해, 해당 논문에서는 서버의 knowledge를 이용하여 edge device의 최적화를 돕는다. 서버의 knowledge를 transfer하기 위해서 distillation 기법을 활용한다. 위의 (2), (5) 식의 detail한 loss는 아래의 (6), (7)과 같다.

기존의 spervised learning에 사용하는 cross-entropy loss에 distillation loss를 추가했다. (6)은 서버에서 k개의 edge device의 logit에 대한 soft label과의 KL divergence Loss를 사용하며, (7)은 edge device에서 server의 logit에 해단 soft label과의 KL divergence loss를 사용한다. 이러한 loss를 통해서 서버는 edge device들의, edge device는 서버의 knowledge를 활용하여 학습을 수행한다. Knowledge distillation에 대한 개념은 아래의 논문을 통해 참고할 수 있다.

https://arxiv.org/abs/1503.02531

 

Distilling the Knowledge in a Neural Network

A very simple way to improve the performance of almost any machine learning algorithm is to train many different models on the same data and then to average their predictions. Unfortunately, making predictions using a whole ensemble of models is cumbersome

arxiv.org

논문에서 제안하는 FedGKT 알고리즘의 프레임워크를 다음과 같이 요약할 수 있다.

 

-Experiments

실험은 CIFAR-10, CIFAR-100, CINIC-10 데이터셋을 사용하여 진행하였으며, edge model은 ResNet-8, server-side model은 ResNet-55과 ResNet-109을 사용하여 평가되어진다. 또한 edge와 server는 같은 input과 output의 dimension을 가진다. 클라이언트는 총 16개를 사용하였다.

 

위의 그래프는 ResNet-56에서 3가지 데이터셋에 대한 FedAvg, FedGKT, Centralized learning 결과를 보여준다. 

 

테이블은 최종적인 결과에 대한 수치를 보여준다. 논문에서는 FedGKT가 FedAvg와 비슷하거나 약간 높은 성능을 보인다고 설명한다.

 

또한 Efficiency를 평가하기 위해서, edge device에서의 계산량을 측정하여 비교한다. 계산량은 FLOPs(부동소수점 연산)의 횟수로 정의하여 counting한다. FedGKT는 FedAvg와 다르게 edge device에서 아주 작은 크기의 ResNet-8을 사용한다. 따라서 Figure4 에서 보는것처럼 Baseline인 FedAvg로 ResNet-110과 56을 사용하였을 때보다 각각 17배 9배 작아진것을 확인할 수 있었다. (해당 논문의 가장 큰 contirubiton이라고 생각한다.)

 

Figure 5는 FL환경에서 모델을 교환하는 SL과 feature map만 교환하는 FedGKT의 communication cost를 비교한 것이고, communication cost 또한 효과적으로 줄였음을 실험적으로 보여주고 있다.

 

이상 논문리뷰였습니다.