[논문 리뷰] Trustworthy Medical Segmentation with Uncertainty Estimation.

Giuseppina Carannante, Dimah Dera외 3명|arXiv (Cornell University)|2021. 11. 10.

Adversarial Robustness in Machine Learning참고 문헌 44인용 수 7

한 줄 요약

이 논문은 MRI 및 CT 데이터셋에서 노이즈와 적대적 공격에 대해 뛰어난 내성성을 보이며, 모델 파라미터의 평균과 공분산을 전파하기 위해 일阶 테일러 근사법을 사용하여 불확실성을 추정하는 베이지안 딥러닝 프레임워크를 제안한다. 이 방법은 증거 하한 bound(ELBO)를 최대화하여, 분할 맵과 불확실성 맵을 동시에 생성한다. 고도로 불확실성 값은 손상되거나 분포 외의 입력에서 모델의 신뢰도 저하를 나타낸다.

ABSTRACT

Deep Learning (DL) holds great promise in reshaping the healthcare systems given its precision, efficiency, and objectivity. However, the brittleness of DL models to noisy and out-of-distribution inputs is ailing their deployment in the clinic. Most systems produce point estimates without further information about model uncertainty or confidence. This paper introduces a new Bayesian deep learning framework for uncertainty quantification in segmentation neural networks, specifically encoder-decoder architectures. The proposed framework uses the first-order Taylor series approximation to propagate and learn the first two moments (mean and covariance) of the distribution of the model parameters given the training data by maximizing the evidence lower bound. The output consists of two maps: the segmented image and the uncertainty map of the segmentation. The uncertainty in the segmentation decisions is captured by the covariance matrix of the predictive distribution. We evaluate the proposed framework on medical image segmentation data from Magnetic Resonances Imaging and Computed Tomography scans. Our experiments on multiple benchmark datasets demonstrate that the proposed framework is more robust to noise and adversarial attacks as compared to state-of-the-art segmentation models. Moreover, the uncertainty map of the proposed framework associates low confidence (or equivalently high uncertainty) to patches in the test input images that are corrupted with noise, artifacts or adversarial attacks. Thus, the model can self-assess its segmentation decisions when it makes an erroneous prediction or misses part of the segmentation structures, e.g., tumor, by presenting higher values in the uncertainty map.

연구 동기 및 목표

• 의료 영상에서 노이즈 또는 분포 외의 이미지에 민감한 딥러닝 모델의 취약성을 해결하기 위해.
• 기존 분할 모델이 불확실성 정량화 없이 단지 점 추정치만 제공하는 한계를 극복하기 위해.
• 모델 파라미터 분포의 모멘트를 통해 예측 불확실성을 추정하여 분할 결정의 자기 평가 기능을 구현하기 위해.
• 의료 영상에서 딥러닝 모델의 신뢰성을 높이기 위해, 고도의 불확실성 값을 오류 예측이나 종양과 같은 누락된 구조물과 연관지기 위해.

제안 방법

• 베이지안 신경망 원리를 적용한 인코더-디코더 아키텍처를 사용하여 의료 영상 분할을 수행한다.
• 모델 파라미터 분포의 첫 번째 두 모멘트(평균과 공분산)를 전파하기 위해 일계 테일러 급수 근사를 적용한다.
• 학습 데이터를 바탕으로 모델 파라미터의 사후분포를 학습하기 위해 증거 하한 bound(ELBO)를 최대화한다.
• 예측 분할 맵과 예측 분포의 공분산 행렬에서 유도된 해당 불확실성 맵을 포함한 두 가지 출력을 생성한다.
• 예측 분포의 공분산 행렬을 활용하여 공간적 위치 간의 분할 결정 불확실성을 정량화한다.
• 노이즈 및 적대적 변형에 대한 내성성을 평가하기 위해 MRI 및 CT 데이터셋에서 프레임워크를 훈련 및 평가한다.

실험 결과

연구 질문

• RQ1 베이지안 딥러닝 프레임워크는 노이즈 또는 손상된 입력에서 의료 영상 분할 모델의 내성성을 향상시킬 수 있는가?
• RQ2 모델 파라미터 분포의 모멘트를 통한 불확실성 추정은 예측 분할의 낮은 신뢰 영역을 효과적으로 식별하는가?
• RQ3 기존 모델과 비교해 볼 때, 제안된 불확실성 맵은 적대적 공격이나 영상 아티팩트를 얼마나 잘 탐지하는가?
• RQ4 종양과 같은 구조물이 누락되었을 경우, 프레임워크는 예측의 불확실성 값을 높게 표시함으로써 스스로 예측을 평가할 수 있는가?
• RQ5 불확실성 정량화가 임상 적용을 위한 딥러닝 모델에 대한 신뢰도를 얼마나 향상시키는가?

주요 결과

• 제안된 프레임워크는 벤치마크 MRI 및 CT 데이터셋에서 최신 기술 대비 더 높은 내성성을 보였다.
• 불확실성 맵은 노이즈, 아티팩트 또는 적대적 변형에 의해 손상된 영상 패치에 대해 일관되게 높은 불확실성 값을 할당하였다.
• 모델이 종양과 같은 구조물을 분할하지 못할 경우, 해당 영역에서 불확실성 맵이 높은 값을 나타내어 낮은 신뢰도를 자가 평가하는 것으로 나타났다.
• 예측 분포의 공분산 행렬은 다양한 공간적 위치에서의 분할 결정 불확실성을 효과적으로 포착하였다.
• 프레임워크의 불확실성 추정은 분포 외 입력을 신뢰성 있게 탐지할 수 있도록 지원하여 더 안전한 임상적 배포를 가능하게 하였다.
• 추가적인 추론 시간이나 아키텍처 복잡성 없이도, 베이지안 파라미터 추정을 통해 분할의 신뢰성을 향상시킬 수 있었다.