[논문 리뷰] FedBABU: Towards Enhanced Representation for Federated Image Classification

Jaehoon Oh, Sangmook Kim외 1명|arXiv (Cornell University)|2021. 06. 04.

Privacy-Preserving Technologies in Data참고 문헌 48인용 수 48

한 줄 요약

FedBABU는 연합 학습 중 모델을 학습 가능한 바디와 고정된 헤드로 분리하여 추론 시 헤드를 미세 조정해 개인화를 효율적으로 가능하게 하며, 이질적인 데이터 설정에서 FedAvg보다 표현력과 개인화를 향상시킨다.

ABSTRACT

Federated learning has evolved to improve a single global model under data heterogeneity (as a curse) or to develop multiple personalized models using data heterogeneity (as a blessing). However, little research has considered both directions simultaneously. In this paper, we first investigate the relationship between them by analyzing Federated Averaging at the client level and determine that a better federated global model performance does not constantly improve personalization. To elucidate the cause of this personalization performance degradation problem, we decompose the entire network into the body (extractor), which is related to universality, and the head (classifier), which is related to personalization. We then point out that this problem stems from training the head. Based on this observation, we propose a novel federated learning algorithm, coined FedBABU, which only updates the body of the model during federated training (i.e., the head is randomly initialized and never updated), and the head is fine-tuned for personalization during the evaluation process. Extensive experiments show consistent performance improvements and an efficient personalization of FedBABU. The code is available at https://github.com/jhoon-oh/FedBABU.

연구 동기 및 목표

• 데이터 이질성 하에서 단일 글로벌 모델이 개인화된 모델과 어떻게 상호작용하는지 이해를 자극한다.
• 신경망을 보편적 바디와 개인화된 헤드로 분리하여 개인화 저하의 원인을 식별한다.
• 연합 학습 동안 헤드를 고정하고 바디만 업데이트하는 FedBABU를 제안하고 검증한다.
• 클라이언트 간 고정된 헤드가 강한 표현력과 로컬 디바이스에서의 개인화를 가속화한다는 점을 보인다.

제안 방법

• 모델 매개변수를 바디(추출기)와 헤드(분류기)로 분리한다.
• FedBABU 동안 각 클라이언트에서 바디만 업데이트하고 헤드는 고정된 무작위로 초기화된 글로벌 헤드를 사용한다.
• 연합 학습 동안 헤드를 업데이트하거나 집계하지 않고 바디 매개변수만 클라이언트 간에 집계한다.
• 초기 글로벌 모델의 성능과 클라이언트별 미세 조정 성능을 비교해 개인화를 평가한다.
• 고정된 헤드가 공동으로 학습된 헤드와 비슷한 중앙 집중적 성능을 제공할 수 있음을 보인다.
• 필요에 따라 학습된 헤드를 템플릿으로 대체해 헤드 없이도 표현력 파악을 평가한다.

실험 결과

연구 질문

• RQ1 헤드와 바디를 분리하는 것이 이질적인 데이터에서 연합 학습의 개인화 저하를 완화할 수 있는가?
• RQ2 연합 학습 중 헤드를 냉동시키면 이후 개인화 미세 조정에 대한 보편적 표현력의 질이 향상되는가?
• RQ3 다양한 데이터 이질성 하에서 FedBABU가 초기 정확도와 개인화 정확도 측면에서 FedAvg 및 다른 개인화 방법들과 어떻게 비교되는가?
• RQ4 헤드 없이 클래스 템플릿을 사용해 FedBABU의 글로벌 모델의 표현력이 견고한가?
• RQ5 FedBABU가 FedProx와 같은 정규화된 FL 프레임워크에 일반화될 수 있는가?

주요 결과

• FedBABU의 바디만 학습하는 경우 FedAvg에 비해 또는 그보다 강력한 표현력을 보이며, 특히 이질성이 큰 경우에 그렇다.
• 클라이언트 간 고정된 헤드는 표현 학습에 강력한 지도를 제공하여 개인화 효율성을 개선한다.
• FedBABU는 빠른 개인화를 가능하게 하며, 때로는 몇 개의 미세 조정 에폭만으로도 FedAvg를 능가한다.
• 헤드가 학습되지 않을 때도 FedBABU는 초기 정확도가 더 높거나 경쟁력 있으며 개인화를 희생하지 않았다.
• 헤드를 미세 조정하거나 전체 모델을 미세 조정하는 것이 FedAvg의 개인화에 더 좋은 성능을 보이나, FedBABU는 헤드만 미세 조정해도 강한 결과를 달성할 수 있다.