[논문 리뷰] Federated Uncertainty-Aware Learning for Distributed Hospital EHR Data

Sabri Boughorbel, Fethi Jarray외 4명|arXiv (Cornell University)|2019. 10. 27.

Privacy-Preserving Technologies in Data참고 문헌 16인용 수 14

한 줄 요약

이 논문은 분산 전자 의료 기록(EHR) 데이터를 위한 새로운 피어드 학습 프레임워크인 연합 불확실성 인식 학습 알고리즘(FUALA)을 제안한다. FUALA는 병원의 모델 일반화 성능에 기반해 기여도를 동적으로 가중하고, 추론 시 앙상블 기반 불확실성 추정을 가능하게 하여 모델의 강건성을 향상시킨다. FUALA는 비독립 동일 분포(out-of-distribution) 설정에서 전기 출생 예측 작업에서 FedAvg와 가중치 기반 FedAvg를 능가하며, 불확실성 지표는 고위험 연령군과 잘못 분류된 케이스에서 더 높은 예측 변동성을 드러낸다.

ABSTRACT

Recent works have shown that applying Machine Learning to Electronic Health Records (EHR) can strongly accelerate precision medicine. This requires developing models based on diverse EHR sources. Federated Learning (FL) has enabled predictive modeling using distributed training which lifted the need of sharing data and compromising privacy. Since models are distributed in FL, it is attractive to devise ensembles of Deep Neural Networks that also assess model uncertainty. We propose a new FL model called Federated Uncertainty-Aware Learning Algorithm (FUALA) that improves on Federated Averaging (FedAvg) in the context of EHR. FUALA embeds uncertainty information in two ways: It reduces the contribution of models with high uncertainty in the aggregated model. It also introduces model ensembling at prediction time by keeping the last layers of each hospital from the final round. In FUALA, the Federator (central node) sends at each round the average model to all hospitals as well as a randomly assigned hospital model update to estimate its generalization on that hospital own data. Each hospital sends back its model update as well a generalization estimation of the assigned model. At prediction time, the model outputs C predictions for each sample where C is the number of hospital models. The experimental analysis conducted on a cohort of 87K deliveries for the task of preterm-birth prediction showed that the proposed approach outperforms FedAvg when evaluated on out-of-distribution data. We illustrated how uncertainty could be measured using the proposed approach.

연구 동기 및 목표

• 병원 간 분산된 EHR 데이터에서 비독립 동일 분포(non-i.i.d.) 데이터와 개인정보 보호 제약 조건을 고려한 기계 학습 모델 훈련의 과제를 해결하기 위해.
• 연합 학습 과정에 불확실성 추정을 통합하여 의료 분야의 연합 학습에서 모델 일반화 능력과 신뢰성을 향상시키기 위해.
• 일반화 성능를 동적 가중치 기반 메커니즘으로 사용하여 성능이 낮거나 과도하게 자신감 있는 모델의 영향을 줄이기 위해.
• 각 병원의 모델 최종 레이어를 활용한 앙상블을 통해 추론 시 불확실성 인식 예측을 가능하게 하기 위해.
• 실제 전기 출생 예측 작업을 42개 병원의 87,000건의 EHR 기록을 사용하여 평가하기 위해.

제안 방법

• FUALA는 FedAvg를 확장하여, 글로벌 모델과 무작위로 할당된 병원 모델 업데이트를 각 병원에 전달하는 중앙 집중형 'Federator'를 도입한다. 이는 일반화 성능 평가를 위해 사용된다.
• 각 병원은 E 에포크 동안 로컬로 글로벌 모델을 훈련하고, 할당된 모델에 대해 자체 데이터 기반 일반화 지표(AU-ROC + PR-AUC)/2를 계산한다.
• Federator는 계산된 일반화 점수를 동적 샘플링 가중치로 사용하여 모델 집합을 수행하며, 일반화 성능가 낮은 병원의 기여도를 감소시킨다.
• 추론 시 최종 예측은 각 병원 모델의 마지막 레이어에서 유도된 C개의 예측 앙상블이며, 이는 예측 분산을 통해 불확실성 추정이 가능하다.
• 모델은 EHR 이벤트의 시퀀스 모델링을 위해 두 개의 병렬 RNN 브랜치와 공유된 로지스틱 출력 레이어를 갖는 RETAIN을 기본 아키텍처로 사용한다.
• 불확실성은 앙상블 예측의 표준편차를 사용하여 정량화되며, 높은 분산은 낮은 신뢰도를 의미한다.

실험 결과

연구 질문

• RQ1 일반화 성능에 기반한 동적 모델 가중치 기법이 비독립 동일 분포 EHR 데이터에서 연합 학습 성능을 향상시킬 수 있는가?
• RQ2 앙성 기반 불확실성 추정은 분포 이탈 상황에서 전기 출생 예측의 모델 신뢰성에 어떻게 기여하는가?
• RQ3 FUALA가 비독립 동일 분포 테스트 데이터에서 AU-ROC 및 AU-PR 측정치에서 FedAvg와 가중치 기반 FedAvg를 능가하는가?
• RQ4 예측 불확실성은 다양한 모성 연령대와 분류 결과에 따라 어떻게 변화하는가?
• RQ5 불확실성 지표는 임상 예측 모델에서 고위험 또는 잘못 분류된 케이스를 식별하는 데 사용될 수 있는가?

주요 결과

• FUALA는 테스트 세트에서 AU-ROC 67.8 ± 2.7과 AU-PR 19.6 ± 4.3을 기록하여, FedAvg(62.2 ± 2.2 AU-ROC, 18.6 ± 2.9 AU-PR)와 가중치 기반 FedAvg(63.5 ± 2.7 AU-ROC, 14.7 ± 3.0 AU-PR)를 모두 능가했다.
• 41세 초과 연령군에서 전기 출생 예측에 대해 높은 불확실성(높은 예측 분산)을 보였으며, 이는 고위험 케이스에서 낮은 신뢰도를 의미한다.
• 잘못 분류된 샘플의 경우, 앙성 내 약 17개 모델이 예측 클래스에 대해 의견을 다르게 하여 잘못된 예측에서 높은 불확실성을 보였다.
• 0 또는 1에 가까운 확률 예측은 낮은 분산을 보였으며, 이는 높은 신뢰도를 의미한다. 반면 중간 확률 값은 더 높은 불확실성을 보였다.
• 다양한 병원가 평가 대상으로 할당되면서 일반화 성능 지표가 라운드가 진행되면서 향상되어 보다 우수한 집합 가중치가 유도되었다.
• 앙성 기반 불확실성 추정은 모델 간 이견이 예측 오류와 강하게 상관됨을 드러내어, 거부 또는 경고 메커니즘 구현에 활용 가능함을 뒷받침했다.