UNETR 논문리뷰

ABSTRACT

medical image segmentation분야에서 encoder와 decoder를 FCNNs통해 구현한 UNET구조가 성능이 좋고, 유망하다.

encoder는 global, local features를 얻는다. decoder는 encoder의 정보를 기반으로 segmentation을 진행한다.

FCNNs를 통해 구현한 모델이 많은 성공을 이뤘지만, convolutional layer를 통해 구성된 FCNNs는 locality틍성을 가진다.( long-range dependencies : 장거리 공간정보를 얻지 못함 )

long-range sequence learning을 하는 트랜스포머의 영감을 받아 volumetric mdeical image segmentation문제를 seq-to-seq prediction 문제로 변경하였다. 이 고유한 모델 구조인 UNETR(UNEt Transformers)을 소개

 

UNETR은 Unet(encoder,decoder)구조를 가지면서 transformer를 encoder로 사용해 input volume에서 seqence 정보를 학습해 효과적으로 global multi-sclae 정보를 얻는다. 뿐만 아니라 transformer encoder는 skip connections을 통해 decoder로 연결된다.

 

INTRODUCTION

 

- FCNNs와 encoder-decoder구조를 가지는 UNET은 다양한 medical semantic segmentation task에서 sota달성

> UNET의 encoder에서 feature들을 downsampling을 통해 global contextual representations을 학습, decoder는 추출된 representations에서 upsampling을 진행하면서 semantic prediction진행, 추가로 encoder와 decoder를 skip connection으로 연결해 downsampling중 손실된 공간정보를 decoder에 전달.

> 이런 FCNN-based 접근은 좋은 성능을 가지지만 locality특성을 가지기 때문에 long-range dependencies 문제를 가짐 

 

- 자연어처리 분야에서 좋은 성능을 기록중인 Transformer를 computer vision분야로 접목시킨 ViT가 등장

> long-range dependencies를 modeling하고 global context를 얻는 능력이 뛰어난 Transformer를 back-bone encoder로 사용

 

- 3D segmentation문제를 1D sequcnce-to-sequnce예측 문제로 변경한 UNETR에 대해 소개

> global context를 잘 얻는 Transfomer를 encoder로 사용, 공간정보를 잘 얻는 CNN을 decoder로 사용, skip connection을 통해 encoder, decoder연결 

 

RELATED WORK

CNN, 3D CNN, ViT

 

METHODOLOGY

- ARCHITECTURE

Overview of UNETR architecture

UNETR은 Transformer를 쌓아만든 encoder(decoder에 skip connection을 통해 연결되는)를 가지는 contracting-expanding 패턴을가짐.

일반적으로 NLP에서 Transformer는 1D sequence를 input embedding으로 사용한다. UNETR도 유사하게 3D volume input x를 non-overlabping patches인 xv생성 

xv는 다시 linear layer를 거처 k dimensional embedding space로 project된다. 추가로 추출된 patch의 위치정보를 보존하기 위해 E_pos라는 embedding값을 더해준다. 그 후 트랜스포머 Encoder에 넣는다.

Unet처럼 Encoder로 부터 얻은 representation Z_i는  Deconvolutional layer를 커처 Convolutional layer인 Decoder에 skip connection으로 연결되어 

 

- Loss Function

soft dice loss와 cross entropy loss를 혼합해서 사용

 

- Dataset

BTCV : 13개의 장기에 대한 CT scan dataset, 512X512 pixel의 사진이 80~225slice로 구성됨, 

MSD : 뇌종양 데이터셋 

 

- Evalutation Metrics

Dice score

95% Hausdorff Distance