[논문 리뷰] Few-shot brain segmentation from weakly labeled data with deep heteroscedastic multi-task networks

Richard McKinley, Michael Rebsamen외 4명|arXiv (Cornell University)|2019. 04. 04.

Brain Tumor Detection and Classification참고 문헌 17인용 수 6

한 줄 요약

이 논문은 ABCD 연구에서 제공한 T1-강조 MRI 스캔 9장의 약한 레이블만을 사용하여 소수의 뇌 분할 방법을 제안한다. 깊이 있는 이종성 다중작업 네트워크를 활용하여 난류 불확실성(aleatoric uncertainty)을 모델링한다. 새로운 레이블 뒤집기 불확실성 네트워크는 각 볼록체에서 잘못 레이블링될 가능성을 학습하며, 평균 딱스 점수 0.87±0.11을 달성한다. 이는 표준 네트워크와 이전의 이종성 방법보다 유의미하게 뛰어나며, 학습된 레이블 스무딩이 최소한의 노이지 데이터에서도 성능 향상에 기여함을 보여준다.

ABSTRACT

In applications of supervised learning applied to medical image segmentation, the need for large amounts of labeled data typically goes unquestioned. In particular, in the case of brain anatomy segmentation, hundreds or thousands of weakly-labeled volumes are often used as training data. In this paper, we first observe that for many brain structures, a small number of training examples, (n=9), weakly labeled using Freesurfer 6.0, plus simple data augmentation, suffice as training data to achieve high performance, achieving an overall mean Dice coefficient of $0.84 \pm 0.12$ compared to Freesurfer over 28 brain structures in T1-weighted images of $\approx 4000$ 9-10 year-olds from the Adolescent Brain Cognitive Development study. We then examine two varieties of heteroscedastic network as a method for improving classification results. An existing proposal by Kendall and Gal, which uses Monte-Carlo inference to learn to predict the variance of each prediction, yields an overall mean Dice of $0.85 \pm 0.14$ and showed statistically significant improvements over 25 brain structures. Meanwhile a novel heteroscedastic network which directly learns the probability that an example has been mislabeled yielded an overall mean Dice of $0.87 \pm 0.11$ and showed statistically significant improvements over all but one of the brain structures considered. The loss function associated to this network can be interpreted as performing a form of learned label smoothing, where labels are only smoothed where they are judged to be uncertain.

연구 동기 및 목표

• 소수의 약한 레이블이 부여된 MRI 스캔만으로 높은 성능의 뇌 분할이 가능할 수 있는지 조사한다. 이는 대규모 레이블링 데이터셋이 필수적이라는 전제를 도전한다.
• 훈련 데이터가 부족하고 레이블이 노이지한 상황에서 이종성 불확실성 모델링이 분할 성능에 어떤 영향을 미치는지 평가한다.
• 각 볼록체에서 레이블 오류 발생 확률을 추정하는 새로운 이종성 네트워크를 제안하고, 이를 통해 학습된 레이블 스무딩을 가능하게 한다.
• 소수의 레이블, 약한 지도 학습 환경에서 표준 네트워크, 예측 분산 네트워크, 그리고 제안된 레이블 뒤집기 불확실성 네트워크의 성능을 비교한다.

제안 방법

• ABCD 연구의 9개의 약한 레이블이 부여된 T1-강조 MRI 볼륨을 기반으로, 보조 레이블로 Freesurfer 6.0 분할 결과를 사용하여 깊이 있는 다중작업 U-Net 아키텍처를 훈련한다.
• 이종성 두 가지 변형을 도입한다: 하나는 몬테카를로 드롭아웃을 사용해 예측 분산을 추정하고, 다른 하나는 볼록체당 레이블 불일치 확률(레이블 뒤집기 확률)을 예측하는 새로운 방법이다.
• 레이블 뒤집기 불확실성 네트워크는 고도로 예측된 잘못된 레이블링 확률을 가진 볼록체의 손실을 약화시켜, 불확실하거나 노이지한 레이블을 효과적으로 가중치를 낮춘다.
• 손실 함수는 심각한 클래스 불균형을 다루기 위해 포칼 손실(γ=2)과 불확실성 인식 손실 항목을 조합하며, 이는 예측된 불확실성에 따라 기울기 흐름을 조절한다.
• 데이터 증강에는 중앙선을 기준으로 반사, 무작위 회전, 강도 재조정이 포함되어 최소한의 데이터로 일반화 성능을 향상시킨다.
• 모델 추론은 Axial, Sagittal, Coronal 시점에서의 3D 앙상블 평균을 사용하며, 최종 분할 결과를 도출하기 위해 로짓을 평균화한다.

실험 결과

연구 질문

• RQ1 수동 레이블링 없이 단지 9개의 약한 레이블이 부여된 훈련 케이스만으로도 고정확도의 뇌 분할이 가능할 수 있는가?
• RQ2 소수의 레이블, 약한 지도 학습 환경에서 이종성 네트워크를 통한 난류 불확실성 모델링이 분할 성능 향상에 기여하는가?
• RQ3 레이블 오류 발생 확률(레이블 뒤집기 확률)을 예측하는 새로운 불확실성 헤드가 표준 예측 분산 추정보다 더 높은 성능을 내는가?
• RQ4 불확실성이 높은 영역에서만 레이블 스무딩을 수행하는 학습된 레이블 스무딩이, 균일한 스무딩이나 스무딩 없이 노이즈 유도 오류를 줄이는 데 더 효과적인가?

주요 결과

• 일반 네트워크는 28개의 뇌 구조에서 평균 딱스 계수 0.84±0.12를 기록했으며, 160,000개의 훈련 샘플 이후 과적합으로 인해 성능이 저하되었다.
• 예측 분산 이종성 네트워크는 평균 딱스 점수 0.85±0.14를 기록했으며, 28개 뇌 구조 중 25개에서 표준 네트워크보다 통계적으로 유의미하게 향상된 성능(𝑝<10⁻⁶)을 보였다.
• 제안된 레이블 뒤집기 이종성 네트워크는 평균 딱스 점수 0.87±0.11로 가장 높은 성능을 기록했으며, 28개 중 1개를 제외한 모든 뇌 구조에서 표준 모델 및 예측 분산 모델보다 유의미하게 뛰어났다(𝑝<10⁻⁶).
• 레이블 뒤집기 모델은 높은 훈련 안정성을 보였으며, 어려운 분할 영역인 악순환체 및 제3뇌실과 같은 영역에서도 초기 성능 붕괴 없이 높은 성능을 유지했다.
• 레이블 뒤집기 지도는 잘못된 레이블링 가능성이 높은 볼록체를 식별하며, 분할 출력의 품질 측정 지표로 잠재력을 보였다.
• 초기 결과에 따르면, 신중하게 선택된 케이스를 활용하면 일시적 훈련도 가능할 것으로 보이며, 이는 소수의 레이블 학습 기술이 의료 영상 분야에 더욱 넓게 적용될 수 있음을 시사한다.