Federated Robustness Propagation: Sharing Robustness in Heterogeneous Federated Learning

Overview

• Federated Learning은 데이터 공유 없이 파라미터만 공유하여 학습하는 구조
• Adversarial training은 robustness를 향상시키지만 계산 비용이 커 일부 client만 수행 가능
• client 간 computing resource와 data distribution이 heterogeneous하여 성능 격차 발생
• 본 연구는 robustness를 client 간에 전달하는 Federated Robustness Propagation 방법 제안

Key Takeaways

Problem Setting

• federated learning 환경에서 client마다 데이터 분포와 자원이 다름
• adversarial training은 비용이 높아 일부 client만 적용 가능
• adversarial training을 하지 않는 client는 robustness가 낮음
• 제한된 환경에서 robustness를 효과적으로 공유하는 것이 핵심 문제

Main Idea

• BN statistics를 활용하여 adversarial robustness를 전달하는 구조 설계
• Dual BN Structure
  • BNc는 clean data에 대한 statistics
  • BNa는 adversarial data에 대한 statistics
  • affine parameter는 공유하고 BN statistics는 분리
• AT Client
  • clean branch와 adversarial branch를 함께 학습
  • PGD 기반 adversarial sample 생성
  • clean loss와 adversarial loss를 동시에 최적화
• ST Client
  • adversarial training 없이 학습 수행
  • server에서 전달받은 BNa를 활용하여 pseudo-noise calibration 수행
  • clean data만으로도 adversarial robustness를 일부 반영
• Robustness Propagation
  • server가 AT client의 BNa statistics를 수집
  • BNc 기반 similarity로 client 간 유사도 계산
  • weighted aggregation으로 적절한 BNa를 생성하여 전달

Result

• 일부 client만 adversarial training을 수행해도 전체 robustness 향상
• full adversarial training 대비 낮은 비용으로 유사한 성능 달성
• data heterogeneity와 hardware heterogeneity 환경 모두에서 안정적인 성능
• cosine similarity 기반 weighting이 uniform 방식보다 성능 우수
• 다양한 attack 설정에서도 robustness 유지
• BN statistics 품질과 domain 차이에 따라 성능이 영향을 받는 한계 존재