[논문 리뷰] BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability

Shangqi Gao, Hangqi Zhou외 2명|arXiv (Cornell University)|2023. 03. 03.

Domain Adaptation and Few-Shot Learning인용 수 4

한 줄 요약

BayeSeg는 계층적 베이지안 사전을 사용하여 의료 영상을 도메인 불변 형상과 도메인 특화 외관 성분으로 분해하는 딥 베이지안 세그멘테이션 프레임워크를 제안한다. 이는 해석 가능하고 일반화 능력이 뛰어난 세그멘테이션을 가능하게 한다. 변분 추론을 통해 레이블 통계를 모델링함으로써, 교차 시퀀스, 교차 모odal, 교차 사이트 분포 이탈 상황에서 전립선 및 심장 MRI 세그멘테이션 분야에서 최신 기술 수준의 일반화 성능을 달성한다.

ABSTRACT

Due to the cross-domain distribution shift aroused from diverse medical imaging systems, many deep learning segmentation methods fail to perform well on unseen data, which limits their real-world applicability. Recent works have shown the benefits of extracting domain-invariant representations on domain generalization. However, the interpretability of domain-invariant features remains a great challenge. To address this problem, we propose an interpretable Bayesian framework (BayeSeg) through Bayesian modeling of image and label statistics to enhance model generalizability for medical image segmentation. Specifically, we first decompose an image into a spatial-correlated variable and a spatial-variant variable, assigning hierarchical Bayesian priors to explicitly force them to model the domain-stable shape and domain-specific appearance information respectively. Then, we model the segmentation as a locally smooth variable only related to the shape. Finally, we develop a variational Bayesian framework to infer the posterior distributions of these explainable variables. The framework is implemented with neural networks, and thus is referred to as deep Bayesian segmentation. Quantitative and qualitative experimental results on prostate segmentation and cardiac segmentation tasks have shown the effectiveness of our proposed method. Moreover, we investigated the interpretability of BayeSeg by explaining the posteriors and analyzed certain factors that affect the generalization ability through further ablation studies. Our code will be released via https://zmiclab.github.io/projects.html, once the manuscript is accepted for publication.

연구 동기 및 목표

• 다양한 영상 시퀀스, 모달리티, 사이트 간 분포 이탈 문제를 해결하기 위한 의료 영상 세그멘테이션 과제에 대응한다.
• 도메인 특화 데이터 증강이나 대규모 레이블링 데이터셋에 의존하지 않고도 새로운 도메인에서의 모델 일반화 능력을 향상시킨다.
• 통계적 사전을 통해 해부학적 형상을 도메인 불변 요인으로 명시적으로 모델링함으로써 해석 가능성을 향상시킨다.
• 불확실성 추정과 인과 기반 표현 학습을 지원하는 통합된 영상 및 레이블 통계 모델링 프레임워크를 개발한다.
• 사후 해석을 통한 설명 가능한 세그멘테이션 예측을 제공함으로써 임상적 신뢰도 및 규제 준수를 가능하게 한다.

제안 방법

• 각 의료 영상을 두 개의 잠재 변수로 분해한다: 공간적으로 상관된 형상 변수(도메인 불변)와 공간적으로 변동성이 있는 외관 변수(도메인 특화).
• 형상 및 외관 성분에 계층적 베이지안 사전를 부여하여, 형상이 도메인 간 안정적인 해부학적 구조를 포착하도록 강제한다.
• 세그멘테이션 마스크를 형상 변수의 국소적으로 스무스한 함수로 모델링하여 공간 일관성과 해석 가능성을 보장한다.
• 이미지, 외관, 형상, 세그멘테이션 레이블 간의 공동 확률 모델을 구축하기 위해 가우시안 마르코프 무작위 필드 기반의 생성 과정을 사용한다.
• 신경망을 활용한 변분 베이지안 프레임워크를 통해 사후 추론을 구현하여 엔드 투 엔드 학습과 불확실성 정량화를 가능하게 한다.
• t-SNE 시각화와 아블레이션 연구를 통해 학습된 형상 표현의 도메인 불변성과 일반화에 미치는 영향을 검증한다.

실험 결과

연구 질문

• RQ1 형상을 도메인 불변 요인으로 명시적으로 통계 모델링하면 교차 도메인 의료 영상 세그멘테이션에서 일반화 성능 향상에 기여하는가?
• RQ2 이미지 및 레이블 통계의 계층적 베이지안 사전 모델링이 해석 가능성과 분포 이탈에 대한 강건성 향상에 어떻게 기여하는가?
• RQ3 외관과 형상 성분을 분리함으로써 기존 딥 러닝 베이스라인 대비 일반화 성능 향상이 어느 정도 이루어지는가?
• RQ4 형상 및 외관 성분의 사후 분포를 의미적으로 해석하여 모델 예측을 설명할 수 있는가?
• RQ5 해부학적으로 관련된 변화와 해부학적으로 관련 없는 변화 등의 요인이 형상 기반 세그멘테이션의 일반화 성능에 어떤 영향을 미치는가?

주요 결과

• BayeSeg는 원천 도메인인 LGE 시퀀스에서 80.9 ± 6.8%의 Dice 스코어를 기록했고, 타겟 도메인인 T2 시퀀스에서는 79.4 ± 10.2%를 기록하여 교차 시퀀스 일반화에서 모든 아블레이션 베이스라인을 초월했다.
• T2 시퀀스에서의 성능은 아블레이션 #6(79.1 ± 10.9%)보다 뛰어나, 통합된 이미지-레이블 모델링과 명시적 형상 사전의 유용성을 확인했다.
• t-SNE 시각화 결과, BayeSeg의 형상 표현은 LGE와 T2 시퀀스 간에 뚜렷한 도메인 이탈이 감소한 것으로 나타났다.
• 아블레이션 연구를 통해 이미지 및 레이블 통계의 통합 모델링이 일반화 성능 향상에 기여했으며, 형상 사전나 레이블 스무딩을 생략한 변형 대비 전체 모델이 더 높은 성능을 보였다.
• 사후 분포 해석 결과, 형상 성분은 도메인 간에 안정적이었고, 외관 성분은 모달리티 특화 변화를 포착했다.
• 강한 형상 불변성에도 불구하고, T2에서의 최악의 성능은 높은 해부학적 관련 이탈 탓이었으며, 이는 현재의 형상 모델링이 해부학적 변동성을 완전히 제거하지 못하고 있음을 시사한다.