[논문 리뷰] Test-Time Robust Personalization for Federated Learning

Liangze Jiang, Tao Lin|arXiv (Cornell University)|2022. 05. 22.

Privacy-Preserving Technologies in Data인용 수 7

한 줄 요약

이 논문은 테스트 시기 정규화를 통한 개인화 방법을 제안하며, 테스트 시기 미세조정과 학습 가능한 앙상블 가중치를 사용하여 모델 헤드를 동적으로 적응시킨다. 테스트 샘플의 역사를 바탕으로 국소 및 글로벌 헤드 파라미터를 최적화하고, 특징 표현을 활용함으로써, 특히 분포 이탈 및 비독립 동일 분포(Non-IID) 데이터 조건에서 기준 모델 대비 뚜렷한 정확도 향상을 달성한다.

ABSTRACT

Federated Learning (FL) is a machine learning paradigm where many clients collaboratively learn a shared global model with decentralized training data. Personalized FL additionally adapts the global model to different clients, achieving promising results on consistent local training and test distributions. However, for real-world personalized FL applications, it is crucial to go one step further: robustifying FL models under the evolving local test set during deployment, where various distribution shifts can arise. In this work, we identify the pitfalls of existing works under test-time distribution shifts and propose Federated Test-time Head Ensemble plus tuning(FedTHE+), which personalizes FL models with robustness to various test-time distribution shifts. We illustrate the advancement of FedTHE+ (and its computationally efficient variant FedTHE) over strong competitors, by training various neural architectures (CNN, ResNet, and Transformer) on CIFAR10 andImageNet with various test distributions. Along with this, we build a benchmark for assessing the performance and robustness of personalized FL methods during deployment. Code: https://github.com/LINs-lab/FedTHE.

연구 동기 및 목표

• 데이터 분포 이탈과 비독립 동일 분포(Non-IID) 클라이언트 데이터 조건 하에서 연합 학습의 모델 개인화 과제를 해결하기 위해.
• 글로벌 모델 재학습 없이도 클라이언트 데이터에 접근하지 않고도 테스트 시기 성능을 향상시키기 위해.
• 로컬 테스트 통계를 활용하여 추론 중에 개인화된 모델 헤드를 동적으로 적응시키는 방법을 개발하기 위해.
• 추론 중 글로벌 및 로컬 모델 구성 요소를 모두 활용하여 분포 이탈에 대한 강건성을 향상시키기 위해.
• 테스트 시기 동안 글로벌 및 로컬 헤드 간 앙상블 가중치를 학습 가능한 방식으로 적응적으로 최적화하기 위해.

제안 방법

• 이 방법은 글로벌 모델 헤드 파라미터 $\mathbf{w}^g$ 와 로컬 모델 헤드 파라미터 $\mathbf{w}^l$ 를 사용하며, 둘 다 글로벌 모델에서 초기화된다.
• 테스트 시기 미세조정은 들어오는 테스트 샘플 $\mathbf{x}_n$ 에 대해 $\tau$ 개의 최적화 단계를 거쳐 로컬 헤드 $\mathbf{w}^l$ 를 미세조정하며, 학습 가능한 스텝 크기 $\eta$ 를 사용한다.
• 학습 가능한 가중치 $e \in [0,1]$ 를 사용하여 글로벌 및 로컬 헤드의 앙상블을 형성하며, 이는 테스트 시기 적응 중 기울기 기반 최적화를 통해 업데이트된다.
• 적응형 헤드 가중치 및 특징 표현 적응을 위한 테스트 기록 $\mathbf{h}^{\text{history}}$ 를 유지한다.
• 특징 표현은 글로벌 $\mathbf{h}^g$ 와 로컬 $\mathbf{h}^l$ 구성 요소로 분할되며, 후자는 테스트 샘플에서의 미세조정을 통해 업데이트된다.
• 초기화 파rameter $\alpha = 0.1$ 와 $\beta = 0.3$ 는 앙상블 가중치 및 헤드 업데이트의 최적화 목표에서 손실 항목 간 균형을 조절한다.

실험 결과

연구 질문

• RQ1 추론 중 분포 이탈에 대해 어떻게 연합 학습의 모델 개인화를 강건하게 만들 수 있는가?
• RQ2 글로벌 모델을 재학습하지 않고도 테스트 시기 동적 적응을 통해 성능을 향상시킬 수 있는가?
• RQ3 학습 가능한 글로벌 및 로컬 헤드 앙상블을 사용할 경우 테스트 시기 정확도에 어떤 영향을 미치는가?
• RQ4 로컬 헤드의 테스트 시기 미세조정은 비독립 동일 분포(Non-IID) 데이터 조건 하에서 일반화에 어떤 영향을 미치는가?
• RQ5 테스트 기록을 유지함으로써 개인화된 예측의 강건성이 어느 정도 향상되는가?

주요 결과

• 제안된 방법은 분포 이탈 조건 하에서 표준 개인화 기준 모델 대비 뚜렷한 정확도 향상을 달성한다.
• 적응형 앙상블 가중치를 사용한 테스트 시기 미세조정은 정적 헤드 앙상블 대비 오류율을 최대 25% 감소시킨다.
• 클라이언트 데이터 분포가 학습 데이터에서 벗어나도 높은 성능 유지를 유지하며, 강건성을 입증한다.
• 학습 가능한 앙상블 가중치 $e$ 의 사용은 글로벌 및 로컬 헤드 간 자동 균형 조정을 가능하게 하여 일반화 성능을 향상시킨다.
• 비독립 동일 분포(Non-IID) 데이터 설정에서도 효과적이며, 다양한 클라이언트 데이터 분포 간 일관된 성능 향상을 보인다.
• 제거 실험을 통해 테스트 시기 미세조정과 기록 기반 적응 모두 성능 향상에 기여한다는 것이 확인되었다.