[논문 리뷰] FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning

Hong-You Chen, Wei‐Lun Chao|arXiv (Cornell University)|2020. 09. 04.

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

한 줄 요약

FedBE는 비독립 동일 분포(non-i.i.i.d.) 데이터와 딥 네ural 네트워크 환경에서 강건성을 향상시키기 위해 베이지안 모델 앙상블 기반의 집약 방법을 제안한다. 국소 모델에 적합된 분포에서 고품질의 전역 모델을 샘플링하고, 스위치드 웨이트 평균화(stochastic weight averaging)를 통한 지식 정복을 통해 FedAvg보다 뛰어난 정확도를 달성한다. 특히 깊은 네트워크와 이질적인 환경에서 뛰어난 성능을 보인다.

ABSTRACT

Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact. Our code is publicly available at https://github.com/hongyouc/FedBE.

연구 동기 및 목표

• 모델 드리프트와 낮은 일반화 성능이 발생하는 비독립 동일 분포 데이터 설정에서 FedAvg의 성능 저하 문제를 해결한다.
• 단순한 모델 평균화의 한계를 극복하기 위해 불확실성을 포착하고 예측의 강건성을 향상시키기 위해 베이지안 모델 앙상블을 활용한다.
• 클라이언트가 재사용할 수 있는 단일 전역 모델로 앙상블 예측을 정복하여 다중 라운드 페더레이티드 러닝을 효과적으로 구현한다.
• 클라이언트의 훈련 절차를 수정하지 않고도 기존 페더레이티드 러닝 프레임워크와의 호환성을 확보하기 위해 집약 단계만 대체한다.
• 서버의 레이블이 없는 데이터와 테스트 데이터 간의 분포 이탈에 대해 강건성을 확보하여 실용적 적용성을 향상시킨다.

제안 방법

• 국소 클라이언트 모델에 기반해 가우시안 또는 딜레트 분포를 피팅하여 모델 분포를 구축하고, 다양한 전역 모델 후보를 샘플링한다.
• 다양한 샘플된 전역 모델의 예측을 집약하여 앙상블 예측을 수행함으로써 강건성과 정확도를 향상시킨다.
• 앙상블 예측에 대해 레이블이 없는 서버 데이터를 활용해 의사 레이블을 생성하고, 이를 단일 전역 모델로의 지식 정복에 활용한다.
• 노이즈가 있는 앙상블 예측에 과적합되지 않도록 정복 과정에서 스위치드 웨이트 평균화(SWA)를 적용하여 일반화 성능을 향상시킨다.
• 앙상블 예측을 교사 신호로 사용하여 학생 전역 모델을 정복 손실을 기반으로 훈련시킨다.
• 클라이언트 측 훈련 및 동적 모멘타 메커니즘을 유지하면서 기존 FL 프레임워크에 FedBE를 집약하기 위해 집약 단계만 대체한다.

실험 결과

연구 질문

• RQ1 FedAvg에 비해 비독립 동일 분포 조건에서 베이지안 모델 앙상블이 페더레이티드 러닝의 모델 집약에 성능 향상을 이룰 수 있는가?
• RQ2 ResNet과 같이 FedAvg가 일반적으로 성능이 저하되는 깊은 신경망(예: ResNet20) 환경에서 FedBE는 어떻게 성능을 내는가?
• RQ3 서버의 레이블이 없는 데이터를 정복에 활용할 경우, 특히 테스트 데이터와 분포가 다를 때 그 영향은 무엇인가?
• RQ4 부분적 클라이언트 참여 및 시스템 이질성과 같은 실용적 시스템 제약 조건 하에서도 FedBE는 성능을 유지할 수 있는가?
• RQ5 기존 정규화 기법(예: FedProx)과 결합했을 때 FedBE의 전체 성능에 어떤 영향을 미치는가?

주요 결과

• 비독립 동일 분포 Tiny-ImageNet에서 ResNet20를 사용할 경우 FedBE는 FedAvg보다 5.9% 높은 정확도(32.4% 대 35.4%)를 달성하고, MobileNetV2에서는 3.4% 높은 정확도(25.5% 대 27.8%)를 기록한다.
• 비독립 동일 분포 CIFAR-10에서 FedBE는 ResNet20를 사용해 77.1%의 정확도를 달성했으며, 이는 FedAvg(69.9%)와 FedProx(69.4%)를 크게 앞서는 성능이다.
• 서버의 레이블이 없는 데이터가 테스트 데이터와 다른 분포를 가질 경우에도 FedBE는 높은 정확도를 유지하여 도메인 이탈에 강건함을 입증한다.
• 부분적 참여(라운드당 100명 중 10명) 조건에서 FedBE는 Tiny-ImageNet(ResNet20)에서 35.4%의 정확도를 기록했으며, 이는 FedAvg(32.4%)를 초월한다.
• 지속적인 로컬 훈련 에포크 수가 다른 시스템 이질성 조건에서 FedBE는 CIFAR-10(ResNet20)에서 77.1%의 정확도를 달성했으며, 이는 FedAvg(69.9%)와 FedProx(69.4%)를 모두 앞선다.
• FedProx와 조합했을 경우 FedBE는 추가로 성능 향상을 보였으며, CIFAR-10(ResNet20)에서 77.5%의 정확도를 기록하여 존재하는 정규화 기법과의 호환성과 상호보완성을 입증한다.