[논문 리뷰] Class-Aware Adversarial Transformers for Medical Image Segmentation

Chenyu You, Ruihan Zhao외 6명|PubMed|2022. 01. 26.

Radiomics and Machine Learning in Medical Imaging참고 문헌 24인용 수 78

한 줄 요약

CASTformer는 2D 의학 영상 분할을 위한 GAN 기반 트랜스포머 프레임워크로, 계층적 다중 스케일 생성기와 클래스 인식 트랜스포머 모듈, 그리고 트랜스포머 기반 판별기를 사용하여 분할 정확도를 높인다. Synapse와 LiTS 데이터셋에서 Dice와 Jaccard 점수의 현 상태 최우수 성능을 달성하며 주목할 만한 개선을 보인다.

ABSTRACT

Transformers have made remarkable progress towards modeling long-range dependencies within the medical image analysis domain. However, current transformer-based models suffer from several disadvantages: (1) existing methods fail to capture the important features of the images due to the naive tokenization scheme; (2) the models suffer from information loss because they only consider single-scale feature representations; and (3) the segmentation label maps generated by the models are not accurate enough without considering rich semantic contexts and anatomical textures. In this work, we present CASTformer, a novel type of adversarial transformers, for 2D medical image segmentation. First, we take advantage of the pyramid structure to construct multi-scale representations and handle multi-scale variations. We then design a novel class-aware transformer module to better learn the discriminative regions of objects with semantic structures. Lastly, we utilize an adversarial training strategy that boosts segmentation accuracy and correspondingly allows a transformer-based discriminator to capture high-level semantically correlated contents and low-level anatomical features. Our experiments demonstrate that CASTformer dramatically outperforms previous state-of-the-art transformer-based approaches on three benchmarks, obtaining 2.54%-5.88% absolute improvements in Dice over previous models. Further qualitative experiments provide a more detailed picture of the model's inner workings, shed light on the challenges in improved transparency, and demonstrate that transfer learning can greatly improve performance and reduce the size of medical image datasets in training, making CASTformer a strong starting point for downstream medical image analysis tasks.

연구 동기 및 목표

• 향상된 분할의 동기 부여를 위해 기존 트랜스포머 기반 의학 분할 모델의 다중 스케일, 의미론적, 맥락 한계를 다루고 이를 개선하려 한다.
• 피라미드 구조의 생성기, 클래스 인식 트랜스포머 모듈, 그리고 GAN 기반 학습으로 글로벌 및 로컬 특징 학습을 강화하는 CASTformer를 제안한다.
• 다양한 의학 영상 벤치마크에서의 성능 향상을 입증하고 전이 학습 및 구성 요소 기여를 분석한다.

제안 방법

• 트랜스포머 기반 생성기(CATformer)와 Vision Transformers에서 사전 학습된 판별기를 갖춘 GAN 프레임워크를 도입한다.
• 분할을 위한 다중 스케일 표현 학습을 위한 기능 피라미드를 Incorporate 한다.
• 해부학적으로 의미 있는 영역을 반복적으로 샘플링하는 클래스 인식 트랜스포머(CAT) 모듈을 개발한다.
• 장거리 맥락을 포착하는 트랜스포머 인코더 모듈(TEM)을 활용한다.
• 효율적인 다중 스케일 융합 및 마스크 예측을 위한 경량 All-MLP 디코더를 사용한다.
• 실감도와 정확도 간의 균형을 맞추기 위해 WGAN-GP 목적과 분할 손실(Dice 및 교차 엔트로피)을 함께 학습한다.

실험 결과

연구 질문

• RQ1 피라미드 구조의 트랜스포머 기반 생성기가 단일 스케일 접근법보다 다중 스케일 의학 영상 분할을 개선할 수 있는가?
• RQ2 트랜스포머 내의 클래스 인식 샘플링 전략이 해부학적으로 의미 있는 영역의 로컬라이제이션을 향상시키는가?
• RQ3 적대적 학습과 트랜스포머 기반 판별기가 의학 이미지의 분할 정밀도와 의미적 일관성을 향상시키는가?
• RQ4 전이 학습과 사전 학습 CV 백본이 제한된 의학 데이터셋에서 성능에 어떤 영향을 미치는가?
• RQ5 개별 CASTformer 구성 요소(CAT 모듈, TEM, GAN 학습)의 전체 성능에 대한 기여는 무엇인가?

주요 결과

프레임워크	평균	대동맥	담낭	신장 (좌)	신장 (우)	간	췌장	비장	위	DSC	Jaccard	95HD	ASD
CASTformer (ours)	82.55	74.69	22.73	5.81	89.05	67.48	86.05	82.17	95.61	67.49	91.00	81.55
CATformer (ours)	82.17	73.22	16.20	4.28	88.98	67.16	85.72	81.69	95.34	66.53	90.74	81.20

• CASTformer는 Synapse에서 Dice 82.55 및 Jaccard 74.69로 최첨단 성능을 달성한다( CASTformer 행).
• LiTS에서 CASTformer는 Dice 73.82% 및 Jaccard 64.91%를 달성하여 TransUNet보다 Dice에서 5.88%p, Jaccard에서 4.66%p 향상시켰다.
• CATformer (GAN 없이)도 이전 방법을 능가하며 Synapse에서 Dice 82.17 및 Jaccard 73.22를 달성했다.
• CV 프리트레인 백본으로의 전이 학습은 특히 소형 데이터셋에서 성능을 크게 향상시킨다.
• 절삭(Ablation) 실험은 클래스 인식 트랜스포머와 TEM이 모두 유의미하게 기여함을 보여주며, 둘 중 하나를 제거하면 Dice 이득이 감소한다.