[论文解读] Uncertainty-Aware Attention for Reliable Interpretation and Prediction

Jay Heo, Haebeom Lee+5 more|arXiv (Cornell University)|May 24, 2018

Machine Learning in Healthcare参考文献 29被引用 34

一句话总结

本文提出不确定性感知注意力（UA），一种基于变分推断的注意力机制，通过学习注意力权重中的实例特定方差来建模输入相关的不确定性。通过生成带有学习到的不确定性的注意力，UA在电子健康记录（EHR）任务中提升了预测的可靠性、校准性以及临床可解释性，在准确率、不确定性校准（ECE）以及通过“我不知道”决策实现的安全预测方面优于基线模型。

ABSTRACT

Attention mechanism is effective in both focusing the deep learning models on relevant features and interpreting them. However, attentions may be unreliable since the networks that generate them are often trained in a weakly-supervised manner. To overcome this limitation, we introduce the notion of input-dependent uncertainty to the attention mechanism, such that it generates attention for each feature with varying degrees of noise based on the given input, to learn larger variance on instances it is uncertain about. We learn this Uncertainty-aware Attention (UA) mechanism using variational inference, and validate it on various risk prediction tasks from electronic health records on which our model significantly outperforms existing attention models. The analysis of the learned attentions shows that our model generates attentions that comply with clinicians' interpretation, and provide richer interpretation via learned variance. Further evaluation of both the accuracy of the uncertainty calibration and the prediction performance with "I don't know" decision show that UA yields networks with high reliability as well.

研究动机与目标

• 为解决传统注意力机制在医疗等安全关键应用中的不可靠性问题，此类模型可能过度自信且校准性差。
• 使深度学习模型能够通过在注意力权重中建模实例特定的预测不确定性，识别自身的不确定性。
• 通过不确定性感知预测提升模型可靠性，包括在不确定时通过“我不知道”进行决策延迟的能力。
• 通过学习反映模型置信度的注意力方差，提供更丰富、与临床一致的解释。

提出的方法

• 将注意力权重形式化为具有输入相关均值和方差的高斯分布，实现异方差不确定性建模。
• 使用变分推断端到端训练不确定性感知注意力机制，联合优化注意力与不确定性。
• 在注意力机制中解耦均值与方差的预测，实现对注意力强度与不确定性的独立建模。
• 采用蒙特卡洛Dropout与随机高斯采样，估计预测不确定性以评估可靠性。
• 通过在不确定性超过阈值时延迟预测，引入“我不知道”（IDK）决策机制。
• 在多个EHR任务（包括脓毒症与死亡率预测）上验证该方法，使用期望校准误差（ECE）和IDK下的准确率进行评估。

实验结果

研究问题

• RQ1 能否通过在注意力权重中显式建模实例特定的不确定性，使注意力机制更加可靠？
• RQ2 与确定性或随机注意力基线相比，不确定性感知注意力在预测校准方面有何改进？
• RQ3 注意力中学习到的不确定性在多大程度上与临床对患者风险的判断一致？
• RQ4 当模型不确定时，通过“我不知道”选项延迟预测，能否提升安全性？
• RQ5 不确定性感知注意力是否能在真实世界的EHR风险预测任务中带来更好的性能与更可靠的预测？

主要发现

• 在八个EHR任务中，UA在预测准确率方面显著优于RETAIN-DA、RETAIN-SA和UA独立基线，平均ECE值分别为4.22 ± 0.82（PhysioNet死亡率）和1.78 ± 0.41（脓毒症），表现出更优的校准性。
• 在保持正确预测率相当的情况下，使用“我不知道”选项可使错误预测减少30%，证明了更高的可靠性。
• 临床医生验证的注意力解释显示，UA生成的注意力更符合临床推理，优于基线模型。
• 不确定性感知注意力通过学习有意义的方差模式，实现了更丰富的解释，这些模式反映了模型对每个输入的置信度。
• UA+（增加额外不确定性建模）进一步提升了校准性，尤其在胰腺癌预测等低频任务中表现更优。
• 该方法在多种EHR数据集（包括PhysioNet、MIMIC和Pancreatic）上表现出稳健性能，准确率与可靠性均持续提升。