[논문 리뷰] Medical Transformer: Gated Axial-Attention for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Poojan Oza외 2명|arXiv (Cornell University)|2021. 02. 21.

Radiomics and Machine Learning in Medical Imaging인용 수 6

한 줄 요약

이 논문은 의료 영상 분할을 위한 게이트형 축방향 어텐션 네트워크인 Medical Transformer (MedT)를 제안한다. 이는 축방향 자기어텐션에 학습 가능한 게이트 메커니즘을 통합하여 장거리 특징 모델링을 향상시킨다. 사전 훈련 없이도 세 가지 의료 영상 데이터셋에서 최신 기술 수준의 성능을 달성하며, 기존의 컨볼루션 및 트랜스포머 기반 모델을 능가한다. 이는 새로운 국소-전역 훈련 전략(LoGo) 덕분이다.

ABSTRACT

Over the past decade, Deep Convolutional Neural Networks have been widely adopted for medical image segmentation and shown to achieve adequate performance. However, due to the inherent inductive biases present in the convolutional architectures, they lack understanding of long-range dependencies in the image. Recently proposed Transformer-based architectures that leverage self-attention mechanism encode long-range dependencies and learn representations that are highly expressive. This motivates us to explore Transformer-based solutions and study the feasibility of using Transformer-based network architectures for medical image segmentation tasks. Majority of existing Transformer-based network architectures proposed for vision applications require large-scale datasets to train properly. However, compared to the datasets for vision applications, for medical imaging the number of data samples is relatively low, making it difficult to efficiently train transformers for medical applications. To this end, we propose a Gated Axial-Attention model which extends the existing architectures by introducing an additional control mechanism in the self-attention module. Furthermore, to train the model effectively on medical images, we propose a Local-Global training strategy (LoGo) which further improves the performance. Specifically, we operate on the whole image and patches to learn global and local features, respectively. The proposed Medical Transformer (MedT) is evaluated on three different medical image segmentation datasets and it is shown that it achieves better performance than the convolutional and other related transformer-based architectures. Code: https://github.com/jeya-maria-jose/Medical-Transformer

연구 동기 및 목표

• 의료 영상에서 장거리 상관관계를 모델링하는 데에 한계가 있는 컨볼루션 네트워크의 문제를 해결한다.
• 표준 트랜스포머의 효과적 훈련을 방해하는 의료 영상 분야의 데이터 부족 문제를 극복한다.
• 전역적 맥락과 국소적 세부 정보를 명시적으로 모델링함으로써 분할 정확도를 향상시키는 트랜스포머 기반 아키텍처를 개발한다.
• 전체 이미지와 패치 수준의 표현을 효과적으로 활용하여, 소규모 의료 영상 데이터셋에서 성능을 향상시키는 훈련 전략을 설계한다.

제안 방법

• 2D 자기어텐션에서 주목사의 가중치를 제어하기 위해 학습 가능한 게이트를 추가하는 게이트형 축방향 어텐션 메커니즘을 제안한다. 이는 관련된 공간적 맥락에만 집중함으로써 특징 표현을 향상시킨다.
• 전체 이미지에서 훈련하는 전역 브랜치와 이미지 패치에서 훈련하는 국소 브랜치를 동시에 훈련하는 국소-전역(LoGo) 훈련 전략을 도입한다. 이는 모델이 전역 맥락과 세밀한 국소 특징을 모두 학습할 수 있도록 한다.
• 축방향 어텐션을 사용하여 2D 자기어텐션을 두 개의 1D 연산으로 분해함으로써 계산 비용을 감소시키면서도 공간 모델링 능력을 유지한다.
• 기본 컨볼루션 레이어 대신 축방향 어텐션 블록을 사용하는 U-Net 유사 인코더-디코더 아키텍처를 사용한다.
• 이중 단계 훈련 프로토콜을 적용한다. 처음 10 에포크 동안 게이트 업데이트 없이 훈련하여 특징을 안정화한 후, 기준 학습률 0.001로 모든 구성 요소를 함께 훈련한다.
• 전역 브랜치를 오직 두 개의 인코더/디코더 레이어로 제한하고 국소 브랜치에서 패치 기반 처리를 사용함으로써 모델 파라미터를 줄여, 다중 브랜치 설계임에도 불구하고 효율성을 유지한다.

실험 결과

연구 질문

• RQ1 사전 훈련 없이도 축방향 자기어텐션을 갖춘 트랜스포머 기반 아키텍처가 의료 영상 분할에서 장거리 상관관계를 효과적으로 모델링할 수 있는가?
• RQ2 축방향 어텐션 메커니즘에 학습 가능한 게이트를 도입함으로써 소규모 의료 영상 데이터셋에서 성능 향상이 어떻게 이루어지는가?
• RQ3 국소-전역(LoGo) 훈련 전략은 전체 이미지 또는 패치 단독으로 훈련하는 것과 비교해 특징 학습을 어느 정도 향상시키는가?
• RQ4 제안된 의료 트랜스포머(MedT)가 제한된 의료 데이터에서 기존의 컨볼루션 및 트랜스포머 기반 모델보다 분할 정확도에서 뛰어난 성능을 보일 수 있는가?

주요 결과

• MedT는 뇌 초음파 영상 데이터셋에서 F1 스코어 88.84를 기록하여, U-Net(85.37), Res-UNet(87.5), Axial-UNet(88.39)을 모두 앞선다.
• 제거 실험 결과 각 구성 요소—게이트형 축방향 어텐션, LoGo 훈련, 다중 브랜치 설계—가 성능 향상에 기여하는 것으로 확인된다.
• 전역 인코더 레이어를 오직 두 개로 제한한 LoGo 전략이 강력한 성능을 보이며, 국소 패치 수준의 학습과 결합할 경우 최소한의 전역 모델링으로도 충분함을 시사한다.
• 140만 개의 파라미터만을 가진 MedT가 원래 U-Net(1,250만 개)과 Res-UNet(532만 개)보다도 훨씬 많은 파라미터를 가진 모델들보다 뛰어난 성능을 보이며, 파라미터 효율성을 입증한다.
• 정성적 결과에서는 MedT가 U-Net과 Res-UNet이 잘못 분류하는 복잡한 배경 무늬를 가진 초음파 스캔에서 배경 영역을 정확히 분할하는 것으로 나타났다.
• ImageNet 사전 훈련이 필요 없음에도 불구하고, TransUNet 및 TransFuse와 같이 대규모 사전 훈련에 의존하는 유사한 연구들과는 달리, MedT는 최신 기술 수준의 성능을 달성한다.