[논문 리뷰] Uncertainty-Aware Attention for Reliable Interpretation and Prediction

Jay Heo, Haebeom Lee외 5명|arXiv (Cornell University)|2018. 05. 24.

Machine Learning in Healthcare참고 문헌 29인용 수 34

한 줄 요약

이 논문은 입력에 따라 변하는 불확실성을 모델링하기 위해 개별 예측에 특화된 분산을 학습하는 변분 추론 기반의 주의 메커니즘인 불확실성 인식 주의(Uncertainty-Aware Attention, 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.

연구 동기 및 목표

• 의료와 같은 안전이 중요한 응용 분야에서 기존 주의 메커니즘이 과신하고 캘리브레이션 되어 있지 않은 점을 해결하기 위해.
• 주의 가중치에 개별 예측에 특화된 예측 불확실성을 모델링하여 딥 러닝 모델이 자신의 불확실성을 인식할 수 있도록 하기 위해.
• 불확실성 인식 예측을 통해 모델의 신뢰성을 향상시키고, 불확실할 경우 '모르겠다'라는 결정을 유보함으로써 가능하게 하기 위해.
• 모델의 자신감을 반영하는 주의 가중치의 분산을 학습함으로써 더 풍부하고 임상적으로 일치하는 해석을 제공하기 위해.

제안 방법

• 입력에 따라 변화하는 평균과 분산을 갖는 가우시안 분포로 주의 가중치를 설정하여 이방성 불확실성 모델링을 가능하게 한다.
• 변분 추론을 사용하여 불확실성 인식 주의 메커니즘을 엔드 투 엔드로 훈련시키며, 주의와 불확실성 모두를 동시에 최적화한다.
• 주의 메커니즘에서 평균과 분산 예측을 분리하여 주의 강도와 불확실성 각각을 별도로 모델링할 수 있도록 한다.
• 예측 불확실성을 평가하기 위해 몬테카를로 드롭아웃과 스토케스틱 가우시안 샘플링을 활용한다.
• 불확실성이 임계값을 초과할 경우 예측을 유보하는 '모르겠다(IDK)' 결정 메커니즘을 도입한다.
• 예상 캘리브레이션 오차(ECE)와 IDK 조건 하에서의 정확도를 사용하여, 패혈증 및 사망률 예측과 같은 다양한 EHR 작업에서 방법을 검증한다.

실험 결과

연구 질문

• RQ1 주의 가중치에 개별 예측에 특화된 불확실성을 명시적으로 모델링함으로써 주의 메커니즘이 더 신뢰성 있게 개선될 수 있는가?
• RQ2 불확실성 인식 주의는 결정론적 또는 스토케스틱 주의 기준 모델과 비교해 예측 캘리브레이션을 어떻게 향상시키는가?
• RQ3 학습된 주의 가중치의 불확실성은 환자 위험에 대한 임상적 해석과 얼마나 일치하는가?
• RQ4 불확실할 경우 '모르겠다' 옵션을 통해 예측을 유보함으로써 모델이 안전성을 향상시킬 수 있는가?
• RQ5 불확실성 인식 주의는 실제 세계의 EHR 위험 예측 작업에서 더 높은 성능과 더 신뢰할 수 있는 예측을 이끌어낼 수 있는가?

주요 결과

• UA는 패혈증 예측(1.78 ± 0.41)과 PhysioNet 사망률 예측(4.22 ± 0.82)에서 평균 ECE 값이 각각 4.22 ± 0.82와 1.78 ± 0.41로 기준 모델인 RETAIN-DA, RETAIN-SA 및 UA-독립 기준 모델보다 예측 정확도에서 뚜렷한 승리를 거두었으며, 캘리브레이션 성능이 뛰어나다.
• '모르겠다' 옵션을 사용할 경우 동일한 정확한 예측 비율에서 잘못된 예측이 30% 감소하여 더 높은 신뢰성을 입증하였다.
• 임상의가 검증한 주의 해석 결과, UA가 생성한 주의 가중치는 기준 모델보다 임상적 추론과 더 잘 일치하였다.
• 불확실성 인식 주의는 입력에 따라 모델 자신감을 반영하는 의미 있는 분산 패턴을 학습함으로써 더 풍부한 해석을 제공하였다.
• 추가적인 불확실성 모델링을 적용한 UA+는 특히 희귀 질환인 췌장암 예측과 같은 저빈도 작업에서 캘리브레이션 성능을 더욱 향상시켰다.
• PhysioNet, MIMIC, 췌장암 데이터셋을 포함한 다양한 EHR 데이터셋에서 우수한 성능을 보이며 정확도와 신뢰성 측면에서 일관된 향상을 보였다.