[論文レビュー] 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）を提案する。学習された不確実性を用いて注意を生成することで、電子歴史記録（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.

研究の動機と目的

• 医療などの安全が重要な応用分野において、従来の注意メカニズムが過信的でキャリブレーションが不十分である問題に対処すること。
• 注意重みにおけるインスタンス固有の予測不確実性をモデリングすることで、深層学習モデルが自らの不確実性を認識できるようにすること。
• 不確実性認識予測によりモデルの信頼性を向上させることで、不確実性が高い場合に「知らない」という意思決定を実行できるようにすること。
• モデルの自信を反映する注意重みの分散を学習することで、より豊富で臨床的に整合性のある解釈を可能にすること。

提案手法

• 入力に依存する平均と分散を持つガウス分布として注意重みを定式化し、非定常不確実性モデリングを可能にする。
• 変分推論を用いて、注意と不確実性を同時に最適化するエンドツーエンドの学習を実現する。
• 注意メカニズムにおける平均と分散の予測を分離することで、注意の強度と不確実性を独立してモデリング可能にする。
• 予測の不確実性を信頼性の評価に用いるために、モンテカルロドロップアウトと確率的ガウスサンプリングを採用する。
• 不確実性がしきい値を超えた場合に予測を延期する「知らない」（IDK）意思決定メカニズムを導入する。
• Expected Calibration Error（ECE）とIDK下での精度を用いて、sepsisや死亡予測を含む複数のEHRタスクで手法を検証する。

実験結果

リサーチクエスチョン

• RQ1 注意重みにおけるインスタンス固有の不確実性を明示的にモデリングすることで、注意メカニズムの信頼性を向上させることができるか？
• RQ2 不確実性認識注意は、決定論的または確率的注意のベースラインと比較して、予測のキャリブレーションをどのように改善するか？
• RQ3 学習された注意重みの不確実性は、患者リスクに関する臨床的解釈とどの程度整合するか？
• RQ4 不確実性が高い場合に「知らない」という選択肢で予測を延期することで、モデルの安全性を向上させることができるか？
• RQ5 不確実性認識注意は、実世界のEHRリスク予測タスクにおいて、より優れたパフォーマンスとより信頼性の高い予測をもたらすか？

主な発見

• UAは、PhysioNet死亡予測（ECE: 4.22 ± 0.82）とsepsis予測（ECE: 1.78 ± 0.41）を含む8つのEHRタスクにおいて、RETAIN-DA、RETAIN-SA、およびUA独立ベースラインを顕著に上回り、平均的なECE値が優れたキャリブレーションを示した。
• 「知らない」オプションを用いることで、正しい予測率が同等の状態で誤った予測が30％削減された。これは、より高い信頼性を示している。
• 臨床医による検証済みの注意解釈では、UAが生成した注意重みが、ベースラインモデルと比較して臨床的推論とよりよく一致していることが確認された。
• 不確実性認識注意は、入力ごとのモデルの自信を反映する意味のある分散パターンを学習することで、より豊かな解釈を可能にした。
• 追加の不確実性モデリングを施したUA+は、膵臓がん予測のような低頻度タスクにおいても、特にキャリブレーションが向上した。
• PhysioNet、MIMIC、Pancreaticを含む多様なEHRデータセットにおいて、本手法は一貫した精度と信頼性の向上を示し、優れたロバストネスを示した。