[논문 리뷰] Transformer-based out-of-distribution detection for clinically safe segmentation

Mark S. Graham, Petru-Daniel Tudosiu외 10명|arXiv (Cornell University)|2022. 05. 21.

Machine Learning in Healthcare인용 수 9

한 줄 요약

이 논문은 VQ-GAN을 사용해 두개의 뇌 CT 스캔을 압축하고, 변환기(transformer)를 통해 데이터의 가능도를 추정함으로써 임상 영상 분할을 위한 3D 변환기 기반의 분포 외(OOD) 탐지 방법을 제안한다. 이는 먼 OOD 및 가까운 OOD 스캔을 성공적으로 식별하며, 영상 가능도가 분할 성능과 강하게 상관관계가 있음을 보여주어 분할 이전에 문제 있는 입력을 신뢰성 있게 걸러낼 수 있다.

ABSTRACT

In a clinical setting it is essential that deployed image processing systems are robust to the full range of inputs they might encounter and, in particular, do not make confidently wrong predictions. The most popular approach to safe processing is to train networks that can provide a measure of their uncertainty, but these tend to fail for inputs that are far outside the training data distribution. Recently, generative modelling approaches have been proposed as an alternative; these can quantify the likelihood of a data sample explicitly, filtering out any out-of-distribution (OOD) samples before further processing is performed. In this work, we focus on image segmentation and evaluate several approaches to network uncertainty in the far-OOD and near-OOD cases for the task of segmenting haemorrhages in head CTs. We find all of these approaches are unsuitable for safe segmentation as they provide confidently wrong predictions when operating OOD. We propose performing full 3D OOD detection using a VQ-GAN to provide a compressed latent representation of the image and a transformer to estimate the data likelihood. Our approach successfully identifies images in both the far- and near-OOD cases. We find a strong relationship between image likelihood and the quality of a model's segmentation, making this approach viable for filtering images unsuitable for segmentation. To our knowledge, this is the first time transformers have been applied to perform OOD detection on 3D image data. Code is available at github.com/marksgraham/transformer-ood.

연구 동기 및 목표

• 깊이 학습 모델이 분포 외(OOD) 의료 영상에서 확신 있게 잘못된 예측을 내릴 수 있는 임상적 위험을 해결하기 위해.
• 3D CT 분할에서 먼 OOD 및 가까운 OOD 시나리오에 대해 기존의 불확실성 추정 방법(예: 드롭아웃, 앙상블)의 성능을 평가하기 위해.
• 처리 이전에 분할 실패를 일으킬 가능성이 높은 데이터를 걸러내는 강력한 OOD 탐지 시스템을 개발하기 위해.
• VQ-GAN 잠재공간에서 변환기를 사용해 추정한 영상 가능도가 분할 네트워크 성능과 강하게 상관관계가 있는지 조사하기 위해.
• 비두개 CT, MRI, 손상된 스캔을 포함한 다양한 OOD 데이터에서 접근 방법을 검증하기 위해.

제안 방법

• 3D 두개 CT 볼륨을 VQ-GAN을 사용해 이산 잠재공간으로 압축하여 효율적인 모델링을 위한 차원 감소를 수행한다.
• 잠재 토큰의 순차적 자동회귀 모델을 학습시켜 입력 영상의 정확한 가능도 추정이 가능하도록 한다.
• 각 입력 볼륨의 로그 가능도를 계산하여 분포 적합성의 척도로 사용한다.
• 가능도 점수를 OOD 위험의 대체 지표로 사용하며, 낮은 가능도가 데이터 분포 이탈의 가능성을 시사한다.
• 가능도 임계값을 적용해 분류 네트워크에 입력하기 전에 사전 필터링으로 통합한다.
• Monte Carlo 드롭아웃을 사용해 분할 네트워크를 학습하고 평가하여 가능도 점수와의 성능 상관관계를 분석한다.

실험 결과

연구 질문

• RQ1 기존의 표준 불확실성 추정 방법(예: 드롭아웃, 앙상블)은 3D 분할에서 먼 OOD 및 가까운 OOD 의료 영상에 적용했을 때 실패하는가?
• RQ2 VQ-GAN 잠재공간에서의 변환기 기반 가능도 추정기가 먼 OOD 및 가까운 OOD 3D 의료 영상 모두를 효과적으로 탐지할 수 있는가?
• RQ3 영상 가능도와 분할 네트워크 성능(예: 거짓 양성 비율) 간에 강한 상관관계가 존재하는가?
• RQ4 제안된 OOD 탐지 방법은 이상 입력에서 분할 네트워크가 확신 있게 잘못된 예측을 내리는 것을 방지하는가?
• RQ5 VQ-GAN에 인지적 손실 및 적대적 손실을 포함할 경우 OOD 탐지 성능에 어떤 영향을 미치는가?

주요 결과

• 모든 표준 불확실성 추정 방법(예: 드롭아웃, 앙상블)은 OOD 데이터에서 실패했으며, 먼 OOD 및 가까운 OOD 모두에서 확신 있게 잘못된 예측을 내놓았다.
• 제안된 VQ-GAN + 변환기 접근법은 비두개 CT, MRI, 손상된 스캔 포함 10개의 다른 OOD 클래스와의 분포 외 두개 CT를 식별하는 데 있어 100% AUC를 달성했다.
• 영상 로그 가능도는 분할 네트워크 성능과 강하게 상관관계가 있었다: 가능도가 -7000 이하인 입력은 유의미하게 높은 거짓 양성 비율을 보였다.
• 손상된 스캔 중 분포 외 데이터와 유사한 가능도를 가진 경우(예: 노이즈 σ=0.01, AUC=0.49)는 분할 네트워크에서 더 잘 처리되었으며, 미세한 손상이 있음을 시사했다.
• 이전의 생성 기반 OOD 탐지 방법과 비교해 극도로 OOD인 샘플에 높은 가능도를 부여하는 일반적인 실패 유형을 피함으로써 더 뛰어난 강건성을 보였다.
• 공간 가능도 맵은 손상 원인(예: 뇌골 제거, 노이즈)을 국소화하여 문제 영역을 설명 가능한 방식으로 탐지할 수 있었다.