A Comparative Benchmark of Federated Learning Strategies for Mortality Prediction on Heterogeneous and Imbalanced Clinical Data

기술 요약: 사망 예측을 위한 연합 학습 전략의 비교 벤치마크

문제 정의
전자 건강 기록(EHR)을 활용한 병원 내 사망 예측은 임상 의사결정을 개선할 수 있는 상당한 잠재력을 가지고 있다. 그러나 데이터 프라이버시 규제(예: GDPR, HIPAA)로 인해 다기관 데이터를 중앙 집중화하는 것이 차단되어 고성능 모델 개발에 어려움이 있다. 연합 학습(Federated Learning, FL)은 원시 데이터를 공유하지 않고도 협력적 학습을 가능하게 함으로써 이 문제를 해결한다. 그러나 의료 분야에서 FL의 실질적인 적용은 두 가지 결정적이고 복합적인 문제에 직면해 있다: 통계적 이질성(다양한 환자 인구 통계 및 임상 관행으로 인한 기관 간 비독립 동일 분포(non-IID) 데이터 분포)과 심각한 클래스 불균형(사망률은 일반적으로 낮으며, 흔히 2% 미만임). 표준 FL 알고리즘, 특히 연합 평균(FedAvg) 방식은 이러한 조건 하에서 성능 저하, 느린 수렴 또는 발산 문제를 겪는 경우가 많다. 대규모의 이질적이고 불균형한 임상 데이터셋에 대해 현대적인 FL 집계 전략을 구체적으로 비교한 포괄적인 경험적 벤치마크가 부족한 실정이다.

방법론
본 연구는 대규모 중환자 케어 데이터베이스인 MIMIC-IV(버전 3.1)를 사용하여 다섯 가지 서로 다른 FL 전략에 대한 엄격한 비교 분석을 수행한다.

• 데이터 시뮬레이션: 현실적인 non-IID 환경을 시뮬레이션하기 위해, 저자들은 466,351건의 입원 사례를 다섯 개의 별도 임상 케어 유닛(응급실, 퇴원 대기실, 분만실, 내과/심장내과, 내과)으로 분할하였다. 이러한 분할은 극심한 이질성을 생성하며, 사망률은 0.00%(분만실, 20,337건 중 사망 1건)에서 2.99%(내과)까지 다양하게 나타난다.
• 특징 공학(Feature Engineering): 모델은 인구 통계, 입원 수준의 진단/처치 횟수, 그리고 첫 24시간 동안 평균화된 11개 항목의 실험실 패널(예: 크레아티닌, 젖산)로 구성된 특징 세트를 활용한다. 결측치는 훈련 세트의 중앙값으로 대체되며, 결측 여부가 명시적으로 표시된다.
• 클래스 불균형 처리: 데이터를 중앙 집중화하지 않고 클라이언트 수준에서 불균형을 해결하기 위해, 각 클라이언트의 훈련 세트에 SMOTE-Tomek 리샘플링 기법을 로컬하게 적용한다. 단 하나의 양성 사례만 포함된 "분만실" 클라이언트를 위해서는 특정 폴백 메커니즘(클래스 가중 손실)이 사용된다.
• 모델 아키텍처: 모든 전략은 공유된 심층 신경망(DNN) 아키텍처(두 개의 은닉층(각 128, 64 뉴런)을 가진 다층 퍼셉트론)를 사용하여 성능 차이가 오직 집계 전략에 의해서만 기인하도록 보장한다.
• 평가된 전략:
  1. FedAvg: 표준 가중 평균.
  2. FedProx: 로컬 업데이트가 글로벌 모델에서 너무 멀리 벗어나지 않도록 제약하는 근사 항(proximal term)을 목적 함수에 추가함.
  3. FedAdagrad & FedAdam: 서버 측 적응형 옵티마이저로, 글로벌 모델을 위해 모멘텀과 적응형 학습률을 유지함.
  4. FedCluster: 방향성을 기준으로 클라이언트 업데이트를 클러스터링하여(DBSCAN 사용) 집계 전 이상치를 필터링함.
• 평가 지표: 불균형이 심한 환경에서 신뢰성이 높은 AUC-ROC와 AUC-PR(임계값 독립적)을 주요 지표로 선정하였다. 보조 지표로는 검증된 임계값에서의 F1-score, 정밀도(Precision), 재현율(Recall), 정확도(Accuracy)를 사용한다. 실험은 5개의 무작위 시드를 사용하여 50회의 통신 라운드 동안 실행된다.

주요 결과

• 성능 계층: FedProx가 가장 견고한 전략으로 나타났으며, 가장 높은 평균 AUC-ROC(0.897)와 AUC-PR(0.230)을 달-성했다. 쌍체 t-검정 결과, FedProx의 AUC-ROC 우위는 다른 모든 전략에 대해 통계적으로 유의미했다.
• 전략별 특이점:
  • FedAdagrad는 AUC-ROC에서 꾸준한 개선(0.890 도달)을 보였으나, 확률 교정(calibration)이 불량하여 가장 낮은 정밀도(0.149)와 F1(0.232)을 기록했다.
  • FedAdam은 높은 불안정성을 보였으며, 10~15 라운드 부근에서 급격한 손실 스파이크와 AUC-ROC 붕괴를 경험한 후 회복되었다.
  • FedCluster는 가장 높은 F1-score(0.280)와 정밀도(0.252)를 달성했으나, AUC-ROC(0.873)에서는 뒤처졌는데, 이는 정보가 있으나 노이즈가 섞인 업데이트를 필터링했음을 시사한다.
  • FedAvg는 FL 전략 중 최악의 성능(AUC-ROC 0.857)을 보였으며, F1 점수는 개선됨에도 불구하고 순위 품질(ranking quality)은 시간이 지남에 따라 저하되었다.
• 중앙 집중형 vs. 연합형: 중앙 집중식 베이스라인(명시적 불균형 처리가 없는 단순 모델)은 모든 연합 전략보다 유의미하게 높은 성능(AUC-ROC 0.929, AUC-PR 0.312)을 보였다. 이는 FL이 프라이버시를 보존하지만, 동일한 데이터에 대한 중앙 집중식 훈련과 비교했을 때 다소 실질적인 예측 성능 비용이 발생함을 나타낸다.
• 클라이언트별 이질성: 글로벌 연합 모델이 모든 클라이언트에게 균등하게 작용하지 않았다. AUC-ROC는 0.902(응급실)에서 0.809(내과)까지 다양하게 나타났으며, 이는 대형 클라이언트에 의해 주도된 집계가 작고 임상적으로 구별되는 인구 집단에서는 성능이 저하될 수 있음을 보여준다.
• 차분 프라이버시(Differential Privacy, DP): FedProx에 사례별 DP-SGD를 적용했을 때 AUC-ROC는 유지(0.896)되었으나, AUC-PR과 F1은 크게 감소(약 15-22%)하였고 훈련 시간은 4.1배 증가하였다.

주요 기여

1. 포괄적 벤치마크: 대규모, non-IID, 불균형한 임상 사망 예측 작업에 대해 베이스라인, 정규화, 적응형 및 클러스터링 기반 FL 전략을 최초로 직접 비교하였다.
2. 방법론적 엄격성: 매우 불균형한 임상 환경에서는 고정된 임계값 기반 지표(F1)보다 임계값 독립적 지표(AUC-ROC/AUC-PR)를 사용하는 것이 중요함을 입증하였다. 이를 통해 특정 전략이 순위 능력(ranking ability)은 향상시키면서도 교정된 의사결정 능력(calibrated decision-making)은 저하시킬 수 있음을 밝혔다.
3. 클라이언트별 분석: 집계 지표가 서로 다른 케어 유닛 간의 상당한 성능 격차를 가릴 수 있음을 강조하며, FL 시스템을 평가할 때 개별 클라이언트 분포에 대한 평가가 필요함을 역설하였다.
4. 프라이버시-유용성 트레이드오프: 이 문맥에서 차분 프라이버시 구현의 구체적인 비용을 정량화하여, 순위 능력은 보존되지만 운영 지표와 계산 효율성은 저하됨을 보여주었다.

의의 및 주장
본 논문은 FedProx가 클라이언트 드리프트(client drift)를 관리할 수 있는 능력 덕분에 이질적이고 불균형한 임상 예측 작업에 있어 평가된 집계 전략 중 가장 견고한 선택이라고 주장한다. 그러나 저자들은 연합을 결정하는 기준은 예측 성능의 우월함이 아니라 주로 프라이버시 및 데이터 거버넌스 요구사항에 의해 결정되어야 한다고 겸허히 결론짓는다. 왜냐하면 중앙 집중식 훈련이 측정 가능한 우위를 유지하기 때문이다.

본 연구는 방법론적인 경고를 제공한다. 즉, 자연스러운 클래스 분포를 보존하는 평가 세트를 구축하고, 적절한 지표를 선택하는 것이 의료 분야에서 FL 전략을 정확하게 평가하는 데 필수적이라는 것이다. 이 연구는 FL이 프라이버시를 보존하는 협업을 위한 유효한 솔루션이지만, 정교한 알고리즘 선택과 현실적인 기대치(중앙 집중식 모델과의 성능 격차 등)를 요구하는 특정 과제들(이질성, 불균형, 집계 비용)을 수반한다는 점을 강조한다.