[논문 리뷰] Joint Learning of Brain Lesion and Anatomy Segmentation from Heterogeneous Datasets

Nicolas Roulet, Diego Fernández Slezak외 1명|arXiv (Cornell University)|2019. 03. 08.

Advanced Neural Network Applications참고 문헌 15인용 수 5

한 줄 요약

이 논문은 이질적인 데이터셋에서 분리된 레이블을 가진 데이터를 사용하여 뇌 병변 및 해부학적 구조 분할을 동시에 학습하기 위해 적응형 교차 엔트로피(ACE) 손실 함수를 제안한다. ACE 손실은 클래스별 교차 엔트로피 손실의 동적 가중치를 통해 병변 및 조직 레이블 간의 모순을 해결하여, 표준 교차 엔트로피 손실 및 딱스 손실이 실패하는 상황에서도 효과적인 공동 학습을 가능하게 한다. 다중 네트워크 접근 방식과 경쟁 가능한 성능을 달성하며, 특히 뇌 조직 분할에서 뛰어난 성능을 보이며 다양한 병변 유형에 걸쳐 뛰어난 안정성을 유지한다.

ABSTRACT

Brain lesion and anatomy segmentation in magnetic resonance images are fundamental tasks in neuroimaging research and clinical practice. Given enough training data, convolutional neuronal networks (CNN) proved to outperform all existent techniques in both tasks independently. However, to date, little work has been done regarding simultaneous learning of brain lesion and anatomy segmentation from disjoint datasets. In this work we focus on training a single CNN model to predict brain tissue and lesion segmentations using heterogeneous datasets labeled independently, according to only one of these tasks (a common scenario when using publicly available datasets). We show that label contradiction issues can arise in this case, and propose a novel adaptive cross entropy (ACE) loss function that makes such training possible. We provide quantitative evaluation in two different scenarios, benchmarking the proposed method in comparison with a multi-network approach. Our experiments suggest that ACE loss enables training of single models when standard cross entropy and Dice loss functions tend to fail. Moreover, we show that it is possible to achieve competitive results when comparing with multiple networks trained for independent tasks.

연구 동기 및 목표

• 이질적인 데이터셋에서 분리된 레이블을 가진 데이터를 사용하여 단일 딥 러닝 모델을 동시에 뇌 병변 및 해부학적 분할을 수행하도록 효과적으로 학습시키는 데 도전하는 것.
• 병변 및 조직 레이블 간의 충돌이 발생하는 상황, 특히 한 데이터셋에서는 배경으로 표시되는 것이 다른 데이터셋에서는 조직으로 표시되는 경우와 같은 레이블 모순 문제를 식별하고 해결하는 것.
• 이러한 충돌하는 레이블에도 불구하고 효과적인 공동 학습이 가능한 손실 함수를 개발하는 것.
• 다중 네트워크 및 표준 손실 기반 접근 방식과의 성능 비교를 통해 다양한 병변 유형에 걸쳐 분할 정확도 및 일반화 능력을 평가하는 것.
• 공동 학습 시 ACE 손실을 사용할 경우, 입력 모odal이 제한된 경우에도 전용 다중 네트워크 모델과 유사한 성능을 달성할 수 있음을 보여주는 것.

제안 방법

• 레이블의 빈도와 일관성에 기반해 클래스 가중치를 동적으로 조정함으로써 레이블 모순을 해결하는 적응형 교차 엔트로피(ACE) 손실 함수를 제안한다.
• ACE 손실은 클래스별 교차 엔트로피 항목을 계산하고, 특히 배경 및 조직 클래스에서 충돌하는 레이블의 영향을 최소화하기 위해 적응형 가중치를 적용한다.
• 각 데이터셋이 병변 또는 해부학적 분할 중 하나의 작업만을 위한 레이블을 가진 이질적인 데이터셋을 통합하여 ACE 손실을 사용해 단일 U-Net 모델을 엔드 투 엔드로 학습한다.
• 동일한 네트워크가 동시에 병변 및 조직 분할 예측을 수행하는 다중 작업 학습 프레임워크를 활용하며, 작업 간 특징 표현을 공유한다.
• 공개된 데이터셋인 WMH(백질 경화성 병변), BRATS(종양), T1-강조 MRI를 사용하며, 공간적 및 강도 도메인을 일치시키기 위해 철저한 전처리를 수행한다.
• ACE 손실을 표준 다중 분류 교차 엔트로피 손실 및 딱스 손실과 비교하는 분석 실험을 수행하며, 각 작업에 대해 별도의 네트워크를 사용하는 Multi-UNet 기준선과의 성능 평가를 실시한다.

실험 결과

연구 질문

• RQ1 이질적인 데이터셋에서 분리된 레이블을 가진 데이터를 사용할 때, 단일 딥 네트워크가 뇌 병변 및 해부학적 분할을 효과적으로 학습시킬 수 있는가?
• RQ2 표준 손실 함수인 교차 엔트로피 또는 딱스 손실을 사용할 경우 병변 및 조직 레이블 간의 레이블 모순이 학습 과정에 어떤 영향을 미치는가?
• RQ3 제안된 ACE 손실 함수가 레이블 모순 문제를 어느 정도 완화하고 공동 학습 시나리오에서 분할 성능을 향상시키는가?
• RQ4 ACE로 학습된 단일 모델의 성능은 병변 및 해부학적 분할을 위한 별도의 모델을 사용하는 다중 네트워크 접근 방식과 비교해 어떻게 되는가?
• RQ5 입력 모odal이 병변 및 해부학적 데이터셋 모두에 공통으로 존재하는 것들로 제한된 경우 ACE를 사용한 공동 학습의 한계는 무엇인가?

주요 결과

• ACE 손실은 표준 교차 엔트로피 손실 및 딱스 손실이 레이블 모순으로 인해 실패하는 상황에서도 단일 U-Net 모델을 공동 병변 및 해부학적 분할을 위해 성공적으로 학습시킬 수 있다.
• WMH 분할에서 ACE 단일 모델은 Multi-UNet 기준선과 유사한 성능을 보이며, 윌콕슨 서열-순위 검정에 의해 통계적으로 유의미한 차이가 없음을 확인했다.
• 뇌 조직 분할에서는 ACE 모델이 여러 구조에서 Multi-UNet을 능가하거나 동등한 성능을 보이며, 공동 최적화를 통한 향상된 특징 학습을 나타낸다.
• 부종 및 종양 분할에서는 ACE 모델이 성능 저하를 보였는데, 주로 단일 T1 이미지만 사용된 제한된 입력 모달 때문이며, Multi-UNet은 T1, T1g, T2, FLAIR 등 여러 시퀀스를 활용했다.
• 이러한 결과는 계층적 다중 작업 학습의 가능성을 입증하며, 병변 및 조직 클래스의 공동 예측이 레이블 상호작용을 포착함으로써 전체 분할 정확도를 향상시킬 수 있음을 보여준다.
• 저자들은 ACE 손실이 클래스 불균형과 레이블 갈등을 효과적으로 처리함을 관찰하며, 이는 다중 작업 의료 영상 분할 분야에 광범위하게 적용될 잠재력을 지닌다는 것을 시사한다.