[論文レビュー] Federated Uncertainty-Aware Learning for Distributed Hospital EHR Data

Sabri Boughorbel, Fethi Jarray+4 more|arXiv (Cornell University)|Oct 27, 2019

Privacy-Preserving Technologies in Data参考文献 16被引用数 14

ひとこと要約

本稿では、分散型電子歴史記録（EHR）データを対象とした新規のフェデレーテッドラーニングフレームワーク、Federated Uncertainty-Aware Learning Algorithm（FUALA）を提案する。FUALAは、病院のモデル汎化性能に基づき動的に寄与度を重み付けすることでモデルのロバスト性を向上させるとともに、推論時にアンサンブルベースの不確実性推定を可能にする。FUALAは、分布外設定下での早産予測において、FedAvg や weighted-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データを用いた機械学習モデルの学習において、非i.i.d.データとプライバシー制約の課題に対処すること。
• フェデレーテッドトレーニングプロセスに不確実性推定を組み込むことで、医療分野におけるフェデレーテッドラーニングの汎化性能と信頼性を向上させること。
• 一般化性能を動的重み付けメカニズムとして用いることで、低性能または過信したモデルの影響を軽減すること。
• 各病院のモデルの最終層を用いたアンサンブルによる不確実性に配慮した推論を可能とすること。
• 本手法を、42の病院からなる87,000件のEHR記録を用いた実世界の早産予測タスクで評価すること。

提案手法

• FUALAは、一般化評価のための病院モデル更新をランダムに割り当てられる中央のFederatorを導入することで、FedAvgを拡張する。
• 各病院は、グローバルモデルをEエポック分ローカルで訓練し、割り当てられたモデルの自身のデータにおける一般化指標（AU-ROC + PR-AUC）/2 を算出する。
• Federatorは、算出された一般化スコアを動的サンプリング重みとして用い、一般化性能が低い病院からの寄与度を低減する。
• 推論時、最終予測は各病院のモデルの最終層からのC個の予測のアンサンブルであり、予測分散を用いた不確実性推定が可能になる。
• モデルは、2本の並列RNNブランチと共有されたロジスティック出力層を備えたRETAINをベースアーキテクチャとして用いる。
• 不確実性は、アンサンブル予測の標準偏差で定量化され、高い分散は低い信頼性を示す。

実験結果

リサーチクエスチョン

• RQ1 一般化性能に基づく動的モデル重み付けが、非i.i.d. EHRデータにおけるフェデレーテッドラーニング性能を向上させ得るか？
• RQ2 アンサンブルベースの不確実性推定は、分布シフト下での早産予測におけるモデル信頼性をどのように向上させるか？
• RQ3 FUALAは、分布外テストデータにおけるAU-ROCおよびAU-PRにおいて、FedAvg や weighted-FedAvg より優れた性能を示すか？
• RQ4 予測不確実性は、異なる母体の年齢層や分類結果ごとにどのように変化するか？
• RQ5 不確実性指標を用いて、臨床予測モデルにおける高リスクまたは誤分類ケースを特定できるか？

主な発見

• FUALAは、テストセットでAU-ROC 67.8 ± 2.7、AU-PR 19.6 ± 4.3を達成し、FedAvg（AU-ROC 62.2 ± 2.2、AU-PR 18.6 ± 2.9）および weighted-FedAvg（AU-ROC 63.5 ± 2.7、AU-PR 14.7 ± 3.0）を上回った。
• 41歳以上の母体年齢層における早産予測では、より高い不確実性（高い予測分散）を示し、高リスクケースでの信頼性が低いことを示した。
• 誤分類されたサンプルでは、アンサンブル内の約35個中17個のモデルが予測クラスに合意しなかったため、誤った予測における高い不確実性が確認された。
• 0または1に近い確率の予測は低分散を示し、高い信頼性を示したのに対し、中間確率ではより高い不確実性が観察された。
• 異なる病院が順次評価対象に割り当てられるに従い、重み付けに用いる一般化性能指標が時間経過とともに向上し、より良い集約重みが得られた。
• アンサンブルベースの不確実性推定により、モデルの不一致が予測誤差と強く相関していることが明らかとなり、拒否またはアラート機構への応用が有効であることが裏付けられた。