[논문 리뷰] FedDAR: Federated Domain-Aware Representation Learning

Aoxiao Zhong, Hao He외 3명|arXiv (Cornell University)|2022. 09. 08.

Machine Learning in Healthcare인용 수 4

한 줄 요약

FedDAR는 각 클라이언트의 데이터가 사전 정의된 도메인의 혼합으로 이루어진 도메인 혼합 비i.i.d. 데이터를 위한 새로운 피어드 학습 프레임워크를 제안한다. 이는 도메인 간 공유되는 인코더를 사용하면서도 도메인별로 특화된 헤드를 사용함으로써 표현 학습과 예측을 분리하여, 번갈아가며 최적화하고 이차적 집계를 통해 도메인 인식형 개인화를 실현하며, 선형 수렴성과 우수한 성능을 달성한다. 이는 합성 및 실제 의료 데이터셋에서 입증된다.

ABSTRACT

Cross-silo Federated learning (FL) has become a promising tool in machine learning applications for healthcare. It allows hospitals/institutions to train models with sufficient data while the data is kept private. To make sure the FL model is robust when facing heterogeneous data among FL clients, most efforts focus on personalizing models for clients. However, the latent relationships between clients' data are ignored. In this work, we focus on a special non-iid FL problem, called Domain-mixed FL, where each client's data distribution is assumed to be a mixture of several predefined domains. Recognizing the diversity of domains and the similarity within domains, we propose a novel method, FedDAR, which learns a domain shared representation and domain-wise personalized prediction heads in a decoupled manner. For simplified linear regression settings, we have theoretically proved that FedDAR enjoys a linear convergence rate. For general settings, we have performed intensive empirical studies on both synthetic and real-world medical datasets which demonstrate its superiority over prior FL methods.

연구 동기 및 목표

• 비i.i.d. 의료 데이터 설정에서 데이터 이질성을 다루기 위해 클라이언트의 데이터를 사전 정의된 도메인의 혼합으로 모델링하는 것, i.i.d. 또는 클라이언트별 분포를 가정하는 것과는 다름.
• 기존 피어드 학습 방법이 클라이언트별 개인화는 하지만 잠재된 도메인 구조를 忽시하는 한계를 극복하기 위한 것, 특히 인종이나 스캐너 프로토콜과 같은 도메인이 분포 이탈을 유발하는 의료 분야에서 중요하다.
• 공유 표현을 학습하면서도 도메인별 예측 헤드를 가능하게 하여 모델의 강건성과 공정성을 향상시키는 방법 개발.
• 선형 설정에서 제안된 방법의 수렴 속도 등 이론적 및 실증적 검증을 통해 도메인 기반 개인화가 도메인 혼합 피어드 학습 환경에서 클라이언트 기반 개인화보다 우수한 성능을 내는지 확인.

제안 방법

• FedDAR는 공유 인코더와 도메인별 예측 헤드로 모델을 분해하여 표현 학습과 예측을 분리함으로써 일반화 성능 향상.
• 로컬 헤드를 여러 로컬 에포크 동안 업데이트하고, 샘플 재가중을 통해 도메인 공정성을 확보하면서 인코더 가중치를 가중 평균으로 집계.
• 헤드 집계를 위해 이차 최적화 전략을 도입하여 글로벌 헤드의 수렴성과 최적성 향상.
• 이중 단계 학습 프로세스를 적용: FedAvg 웜업 이후 인코더와 헤드의 번갈아 가며 로컬 업데이트 및 글로벌 집계.
• 학습 중에 부족한 도메인의 기여도를 균형 잡기 위해 도메인 인식 손실 재가중 적용.
• 선형 및 비선형 모델 모두 지원하며, 실제 의료 응용에서 CNN과 ResNet-34에서 평가.

실험 결과

연구 질문

• RQ1 명시적으로 도메인 혼합을 모델링하는 피어드 학습 방법이 비i.i.i.d. 의료 데이터 설정에서 표준 및 개인화된 피어드 학습 베이스라인을 능가하는가?
• RQ2 표현 학습과 도메인별 예측을 분리하면 공동 최적화 대비 수렴성과 성능이 향상되는가?
• RQ3 예측 헤드의 도메인 인식 집계가 다양한 클라이언트 데이터 구성에서 모델의 공정성과 일반화에 어떤 영향을 미치는가?
• RQ4 선형 설정에서 제안된 방법에 대해 수렴 속도 등의 이론적 보장 조건을 확립할 수 있는가?
• RQ5 실제 의료 데이터셋에서 FedDAR는 부족한 도메인의 편향을 얼마나 줄이고 성능을 향상시키는가?

주요 결과

• FedDAR는 선형 회귀 설정에서 선형 수렴성을 확보하여 강력한 이론적 수렴 성질을 입증.
• 합성 데이터에서 FedAvg, FedProx, FedPer, FedRep보다 모든 도메인에서 테스트 정확도 측면에서 우월한 성능 확보.
• FairFace 데이터셋에서 최신 기술 대비 모든 인종 그룹의 평균 정확도가 높으며, 공정성 향상과 편향 감소 효과 확인.
• 실제 EXAM 데이터셋에서 로컬 검증 세트에서 평균 AUC 0.89 달성하여 FedProx, FedRep, FedMinMax를 크게 능가.
• 아블레이션 스터디를 통해 이차 헤드 집계 및 도메인 인식 재가중이 불균형한 도메인에서 성능 향상에 핵심적임을 확인.
• 한 클라이언트가 도메인당 100개의 샘플 뿐이어도 FedDAR는 강건한 성능 유지, 저자료 도메인에 대한 내구성 입증.