[논문 리뷰] FedSoup: Improving Generalization and Personalization in Federated Learning via Selective Model Interpolation

Ming‐Hui Chen, Meirui Jiang외 3명|arXiv (Cornell University)|2023. 07. 20.

AI in cancer detectionComputer Science인용 수 3

한 줄 요약

FedSoup는 연간 모델 선택과 피드백 기반 모델 패치를 통해 일반화와 개인화를 동시에 향상시키는 선택적 모델 보간 방법을 제안한다. 이는 클라이언트별로 별도의 전역 모델 풀을 유지하고 국소 모델과 전역 모델 간의 보간을 통해 이루어지며, 시간에 따른 모델 선택과 피드백 기반 모델 패치를 통해 더 평탄한 최소값을 달성하여 추론 비용 증가 없이도 분포 외 일반화 성능을 크게 향상시킨다.

ABSTRACT

Cross-silo federated learning (FL) enables the development of machine learning models on datasets distributed across data centers such as hospitals and clinical research laboratories. However, recent research has found that current FL algorithms face a trade-off between local and global performance when confronted with distribution shifts. Specifically, personalized FL methods have a tendency to overfit to local data, leading to a sharp valley in the local model and inhibiting its ability to generalize to out-of-distribution data. In this paper, we propose a novel federated model soup method (i.e., selective interpolation of model parameters) to optimize the trade-off between local and global performance. Specifically, during the federated training phase, each client maintains its own global model pool by monitoring the performance of the interpolated model between the local and global models. This allows us to alleviate overfitting and seek flat minima, which can significantly improve the model's generalization performance. We evaluate our method on retinal and pathological image classification tasks, and our proposed method achieves significant improvements for out-of-distribution generalization. Our code is available at https://github.com/ubc-tea/FedSoup.

연구 동기 및 목표

• 데이터 분포 이질성 하에서 크로스실 피드백 기반 학습에서 국소 개인화와 전역 일반화 간의 상충 관계를 해결한다.
• 과적합으로 인한 날카로운 최소값과 낮은 분포 외 성능을 유발하는 개인화된 피드백 기반 학습 모델의 문제를 해결한다.
• 재학습이 필요 없고 통신 비용이 낮은 단일 학습 세션 내에서 모델 수프를 피드백 기반 환경에 적응시킬 수 있는 방법을 개발한다.
• 이전 전역 모델과 국소 업데이트의 보간을 통해 더 평탄한 최소값을 탐색하여 모델의 강건성을 향상시킨다.
• 추론 또는 메모리 오버헤드 증가 없이도 내부 분포 및 외부 분포 데이터에서 모두 높은 성능을 유지한다.

제안 방법

• 피드백 기반 학습 중에 국소 검증 성능 기반으로 클라이언트별로 이전 전역 모델 풀을 유지한다.
• 탐욕적인 시간적 모델 선택 전략을 적용하여 보간에 적합한 고성능 전역 모델을 식별하며, 손실 곡면의 다른 기초 영역에 속한 모델은 회피한다.
• 선택된 전역 모델과 국소 모델을 보간하여 개인화되고 일반화된 모델인 '수프'를 구성한다.
• 피드백 기반 모델 패치 도입: 국소 모델을 전역 모델 수프와 보간하여 개인화와 전역 성능 간의 균형을 맞춘다.
• 평균 주도 헤시안 고유값을 추정하기 위해 파wer 반복법을 사용하여 최소값의 날카움을 측정한다.
• 다양한 학습 에포크에서의 모델을 사용하여 피드백 기반 학습 환경에 적합하게 모델 수프 원칙(고성능 모델의 가중 평균)을 적용한다.

실험 결과

연구 질문

• RQ1 데이터 이질성 하에서 피드백 기반 학습에서 국소 개인화와 전역 일반화 간의 성능 상충 관계는 무엇이 원인인가?
• RQ2 어떻게 모델 보간을 통해 피드백 기반 환경에서 더 평탄한 최소값을 달성하고 일반화 성능을 향상시킬 수 있는가?
• RQ3 다른 학습 에포크에서 온 모델의 선택적 보간은 개인화와 분포 외 일반화 성능 향상에 기여하는가?
• RQ4 다중 재학습 런 또는 높은 통신 비용 없이도 모델 수프 기법을 피드백 기반 학습에 어떻게 적응시킬 수 있는가?
• RQ5 FedSoup는 제한된 국소 데이터로 인한 과적합을 어느 정도 완화하면서도 전역 성능를 유지하는가?

주요 결과

• Camelyon17 병리학 데이터셋에서 FedSoup는 FedAvg 대비 새로운 도메인 일반화 성능에서 AUC 2.87 포인트 향상했다.
• 망막 혈관도 데이터셋에서 FedSoup는 분포 외 일반화 성능에서 86.24%의 정확도를 기록하여 FedBABU(85.09%)와 FedProx(85.18%)를 1.05점 이상 앞섰다.
• 손실 곡면의 날카움이 감소함에 따라 낮은 중앙값 주도 헤시안 고유값을 통해 더 평탄한 최소값을 확인할 수 있었다.
• FedSoup는 높은 국소 성능(85.71% 정확도)을 유지하면서도 뛰어난 전역 성능(96.00% 정확도)을 달성하여 기준 모델 대비 더 나은 국소-전역 상충 관계를 보였다.
• 실험 간 변동성이 낮아 더 높은 안정성을 보였으며, 예를 들어 Retina 데이터셋에서 전역 AUC의 분산은 FedAvg 대비 0.43(0.83 대비)로 낮았다.
• 작은 데이터셋인 망막 혈관도에서 성능 향상이 가장 두드러졌으며, 이는 국소 데이터가 제한된 상황에서 과적합을 효과적으로 완화함을 시사한다.