[论文解读] 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

一句话总结

该论文提出联邦不确定性感知学习算法（FUALA），一种新颖的联邦学习框架，用于分布式电子健康记录（EHR）数据，通过基于医院模型泛化性能动态加权其贡献，并在推理阶段实现基于集成的不确定性估计，从而提升模型鲁棒性。在分布外设置下的早产预测任务中，FUALA优于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家医院的87K份EHR记录的真实世界早产预测任务上评估所提方法。

提出的方法

• FUALA通过引入一个中央协调者（Federator），在FedAvg基础上扩展，向每家医院发送全局模型及一个随机分配的医院模型更新，用于泛化性能评估。
• 每家医院在本地训练全局模型E个周期，并在其自身数据上计算所分配模型的泛化度量（AU-ROC + PR-AUC）/2。
• Federator使用计算出的泛化得分作为动态采样权重进行模型聚合，降低泛化性能差的医院的贡献。
• 在推理阶段，最终预测为来自各医院模型最后几层的C个预测结果的集成，通过预测方差实现不确定性估计。
• 模型采用RETAIN作为基础架构，包含两条并行的RNN分支和一个共享的逻辑回归输出层，用于EHR事件序列建模。
• 不确定性通过集成预测的标准差量化，方差越大表示置信度越低。

实验结果

研究问题

• 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岁的产妇群体，早产预测的不确定性更高（预测方差更大），表明对高风险病例的置信度较低。
• 在误分类样本中，约35个集成模型中的17个对预测类别存在分歧，表明错误预测具有较高不确定性。
• 预测概率接近0或1时方差较低，表明置信度高；而中间概率值则表现出更高的不确定性。
• 用于加权的泛化性能度量在训练轮次中持续改善，因为不同医院被轮流分配用于评估，从而获得更优的聚合权重。
• 基于集成的不确定性估计表明，模型分歧与预测误差强相关，支持其在拒绝预测或预警机制中的应用。