[논문 리뷰] Multi-Compound Transformer for Accurate Biomedical Image Segmentation

Yuanfeng Ji, Ruimao Zhang외 5명|arXiv (Cornell University)|2021. 06. 28.

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

한 줄 요약

이 논문은 다중 복합 트랜스포머(MCTrans)를 제안하며, 자기 및 교차 어텐션 메커니즘을 통해 크로스스케일 특징 의존성과 의미 관계를 모델링함으로써 생물의학 영상 분할을 향상시키는 새로운 비전 트랜스포머 아키텍처이다. MCTrans는 UNet 유사 네트워크와 통합되어 여섯 가지 벤치마크에서 최신 기술(SOTA) 성능을 달성하며, Dice 스코어에서 UNet을 최대 4.34% 향상시켰다.

ABSTRACT

The recent vision transformer(i.e.for image classification) learns non-local attentive interaction of different patch tokens. However, prior arts miss learning the cross-scale dependencies of different pixels, the semantic correspondence of different labels, and the consistency of the feature representations and semantic embeddings, which are critical for biomedical segmentation. In this paper, we tackle the above issues by proposing a unified transformer network, termed Multi-Compound Transformer (MCTrans), which incorporates rich feature learning and semantic structure mining into a unified framework. Specifically, MCTrans embeds the multi-scale convolutional features as a sequence of tokens and performs intra- and inter-scale self-attention, rather than single-scale attention in previous works. In addition, a learnable proxy embedding is also introduced to model semantic relationship and feature enhancement by using self-attention and cross-attention, respectively. MCTrans can be easily plugged into a UNet-like network and attains a significant improvement over the state-of-the-art methods in biomedical image segmentation in six standard benchmarks. For example, MCTrans outperforms UNet by 3.64%, 3.71%, 4.34%, 2.8%, 1.88%, 1.57% in Pannuke, CVC-Clinic, CVC-Colon, Etis, Kavirs, ISIC2018 dataset, respectively. Code is available at https://github.com/JiYuanFeng/MCTrans.

연구 동기 및 목표

• 기존 비전 트랜스포머가 크로스스케일 의존성과 레이블 간 의미 관계를 모델링하지 못하는 데 기인한 한계를 해결한다.
• 의료 영상 내 다양한 해부학적 구조에서 특징 표현의 일관성과 분류 가능성을 향상시킨다.
• 단일 스케일 어텐션을 넘어서 장거리 문맥 모델링을 향상시켜 크기가 다양한 병변을 더 잘 포착할 수 있도록 한다.
• 다중 스케일 특징 학습과 의미 구조 탐색을 통합한 유일한 프레임워크를 구현하여 정확한 분할을 달성한다.
• 최소한의 계산 비용으로 다양한 생물의학 영상 분할 벤치마크에서 제안된 방법의 유효성을 입증한다.

제안 방법

• MCTrans는 다중 특징 스케일 간 픽셀 수준의 관계를 모델링하기 위해 트랜스포머-자기어텐션(TSA) 모듈을 활용하여 포괄적인 크로스스케일 문맥 인코딩을 가능하게 한다.
• 의미적 카테고리 표현을 위해 학습 가능한 프록시 임베딩을 도입하였으며, 이는 자기어텐션을 통해 카테고리별 의존성을 캡처하도록 정교화된다.
• 트랜스포머-교차어텐션(TCA) 모듈을 통해 특징 표현과 프록시 임베딩 간의 상호작용을 가능하게 하여 특징의 분류 능력과 일관성을 향상시킨다.
• 업데이트된 프록시 임베딩에 보조 손실을 적용하여 동일 카테고리 내 특징 상관관계를 명시적으로 향상시키고, 서로 다른 카테고리 간 분리 가능성을 개선한다.
• 엔코더-디코더 아키텍처를 기반으로 한 UNet 유사 구조에 MCTrans 모듈을 원활하게 통합하여 엔드 투 엔드 학습 및 추론을 수행한다.
• 다중 스케일 컨볼루션 특징은 토큰으로 임bedding되며, TSA 및 TCA 모듈을 거쳐 업샘플링되고 병합되어 분할 출력을 생성한다.

실험 결과

연구 질문

• RQ1 단일 스케일 자기어텐션에 비해 크로스스케일 어텐션 메커니즘이 생물의학 영상 분할에서 특징 표현을 향상시키는가?
• RQ2 다른 레이블 간 의미 관계를 학습하는 것은 특징의 분류 가능성과 분할 정확도에 어떤 영향을 미치는가?
• RQ3 보조 감독을 통한 학습 가능한 프록시 임베딩 도입이 소규모 및 복잡한 데이터셋에서 모델 성능을 얼마나 향상시키는가?
• RQ4 기본 UNet과 비교해 계산 비용이 약간 증가하는 것으로도 MCTrans 프레임워크는 높은 성능을 유지하는가?
• RQ5 MCTrans는 병변 및 해부학적 구조 분할을 포함한 다양한 생물의학 분할 작업에 일반화 가능한가?

주요 결과

• Pannuke 데이터셋에서 MCTrans는 UNet 대비 3.64% 향상된 Dice 스코어를 기록하였으며, CVC-Clinic, CVC-Colon, ETIS, Kavirs, ISIC2018 데이터셋에서는 각각 3.71%, 4.34%, 2.8%, 1.88%, 1.57% 향상되었다.
• Pannuke 데이터셋에서 ResNet-34 백본을 사용할 경우 MCTrans는 Dice 스코어 68.40%를 기록하여 유사한 파라미터 수를 가진 UNet++(66.08%) 및 CENet(66.50%)를 능가했다.
• 절단 실험 결과, 프록시 임베딩에 대한 보조 손실을 제거할 경우 성능이 68.40%에서 67.87%로 떨어지며, 이는 보조 손실이 특징 상관관계 및 분류 가능성을 향상시키는 데 기여함을 확인한다.
• TSA 및 TCA 모듈의 최적 수는 4로 확인되었으며, 이 이상으로 증가할 경우 성능 저하가 발생하여 모델 용량과 데이터셋 크기 간 균형이 중요함을 시사한다.
• MCTrans는 FLOPs가 18.065G(UNet의 14.037G 대비)로 다소 증가함에도 불구하고 높은 정확도 향상과 함께 효율성을 유지함을 입증했다.
• 그림 4의 정성적 결과에서는 MCTrans가 특히 작은 또는 비정형적인 병변에 대해 더 정밀하고 일관성 있는 분할 마스크를 생성하는 것으로 나타났다.