[논문 리뷰] Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty

Miguel Monteiro, Loïc Le Folgoc외 6명|arXiv (Cornell University)|2020. 06. 10.

Explainable Artificial Intelligence (XAI)참고 문헌 48인용 수 79

한 줄 요약

확률적 분할 네트워크(SSN)은 로짓 공간에서 저랭크 다변정규 분포를 사용해 분할 맵의 결합 분포를 모델링하며, 다수의 공간적으로 일관된 가설과 2D 및 3D 의학 영상에서의 효율적 추론을 가능하게 한다.

ABSTRACT

In image segmentation, there is often more than one plausible solution for a given input. In medical imaging, for example, experts will often disagree about the exact location of object boundaries. Estimating this inherent uncertainty and predicting multiple plausible hypotheses is of great interest in many applications, yet this ability is lacking in most current deep learning methods. In this paper, we introduce stochastic segmentation networks (SSNs), an efficient probabilistic method for modelling aleatoric uncertainty with any image segmentation network architecture. In contrast to approaches that produce pixel-wise estimates, SSNs model joint distributions over entire label maps and thus can generate multiple spatially coherent hypotheses for a single image. By using a low-rank multivariate normal distribution over the logit space to model the probability of the label map given the image, we obtain a spatially consistent probability distribution that can be efficiently computed by a neural network without any changes to the underlying architecture. We tested our method on the segmentation of real-world medical data, including lung nodules in 2D CT and brain tumours in 3D multimodal MRI scans. SSNs outperform state-of-the-art for modelling correlated uncertainty in ambiguous images while being much simpler, more flexible, and more efficient.

연구 동기 및 목표

• 의료 영상과 같이 여러 개의 가능한 경계가 존재하는 이미지 분할에서 고유한 불확실성을 포착해야 할 필요성을 동기 부여한다.
• 네트워크 아키텍처를 변경하지 않으면서 레이블 맵의 결합 분포를 모델링하는 경량의 확률적 프레임워크를 제안한다.
• 단일 순전파로 다수의 공간적으로 일관된 분할 가설을 생성하도록 한다.
• 2D CT 폐결절 및 3D 다중 모달 MRI 뇌종양 분할에의 적용 가능성을 보여준다.

제안 방법

• 로짓 맵을 저랭크 플러스 대각 다변정규로 모델링한다: η|x ~ N(μ(x), Σ(x)) 이고 Σ = P P^T + D.
• μ(x), P, D를 신경망의 출력으로 매개변수화하여 효율적인 계산과 확장성을 가능하게 한다.
• 로짓 공간에서 몬테카를로 샘플링으로 η에 대한 계산 불가능한 적분을 근사한다.
• 재매개화 트릭을 사용해 몬테카를로 샘플을 역전파하고 교차 엔트로피와 유사한 목적 함수를 적용한다.
• 추론은 단일 순전파에서 로짓 맵의 다수 샘플을 생성하며, 평균 로짓이 가장 가능성이 높은 분할을 제공한다.

실험 결과

연구 질문

• RQ1 저랭크 로짓 공간 분포를 통해 공간적으로 상관된 내재적 불확실성을 모델링하면 일관되고 다수의 분할 가설을 만들 수 있는가?
• RQ2 의료 영상 작업에서 SSN이 결정론적 및 기존 확률적 분할 방법에 비해 예측 성능과 불확실성 보정을 개선하는가?
• RQ3 이 접근법이 고차원 3D 데이터에 얼마나 확장되며 일반적인 분할 백본과 어떻게 통합되는가?

주요 결과

• SSN은 모호한 의학 영상에서 상관된 불확실성을 모델링하는 데 최신 방법보다 우수하다.
• 저랭크 공분산 모델은 폐결절 분할에서 대각 공분산 변형 및 결정론적 기준선보다 예측 성능과 불확실성 보정이 더 우수하게 나타난다.
• SSN은 최소한의 아키텍처 변경과 상대적으로 가벼운 계산 비용으로 2D 및 3D에서 다수의 그럴듯한 분할 샘플링을 가능하게 한다.
• BraTS 2017 3D 뇌종양 분할에서 확률 모델은 결정론적 모델과 유사한 Dice 스코어를 달성하는 동시에 분할의 분포를 제공한다.
• 프레임워크는 샘플의 추론 후 조작(예: 클래스 존재 여부 조정) 및 다양성의 온도 기반 제어를 가능하게 한다.