The fully convolutional transformer for medical image segmentation 리뷰

Tragakis, Athanasios, Chaitanya Kaul, Roderick Murray-Smith, and Dirk Husmeier. "The fully convolutional transformer for medical image segmentation." In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 3660-3669. 2023. 

WACV 2023 Open Access Repository

 

Abstract

트랜스포머는 medical image analysis에 아직 초기단계이고 fine-grained segmentation을 위한 능력을 보유하고 있지 않다. Fully Convolutional Transformer(FCT)는 이미지 표현을 학습하기 위한 CNN의 장점과 long-term dependencies를 포착하기 위한 트랜스포머의 장점을 결합하였다. FCT는 medical imaging에서 첫 fully convolutional transformer이며 long-range semantic을 학습하기 위한 1단계와 계층적으로 전역정보를 학습하기 위한 2단계로 나누어 진다. FCT는 사전훈련없이도 우월한 성능을 보였다.

 

The Fully Convolutional Transformer

이미지 데이터 \(X\)와 그에 해당하는 semantic/binary segmentation maps \(Y\)로 이루어진 데이터셋 \(\{X, Y\}\)를 가정한다. 각각의 이미지는 \(x_i \in R^{H\times W\times C}\), \(C=\{3,...,N\}\)이고 \(y_i\in R^{H\times W\times K}\), \(K\in\{1,...,D\}\) 이다. FCT의 입력으로는 3D이미지의 조각들인 2D 패치이다. FCT는 기존 방식들과 다르게 over-lapping 패치로부터 특성을 추출하고 임베딩 후 Multi-head attention을 수행한다. 이후 Wide-Focus 모듈을 이용하여 이미지 내 세부 정보를 추출하여 projection한다.

FCT Layer

FCT는 가장 먼저 LayerNorm-Conv-Conv-MaxPool 의 순으로 시행된다(Convolutional Attention블록이 수행되기 전). Conv는 연속된 3x3 kernel size로 이미지의 특징을 추출하며 Gelu 활성화 함수가 뒤따른다. 다른 proposed block들과의 첫번째 차이점은 Convolutional Attention의 적용이다.

 

MaxPool이후에는 Depthwise-Convolution으로 토큰맵을 생성하는데 3x3 kernel size, \(s\times s\) stride, valid padding으로 입력과 출력의 크기를 고정시키고 다시 LayerNorm이 뒤따른다(그림에는 존재하지 않지만 코드로 확인해보면 Convolutional Projection, MHSA에서 중간중간 계속 LayerNorm이 존재함).

 

MHSA에서 기존 블록들과 또 다른 차이점이 나오는데 우리는 Linear projection이 아닌 Depthwise-Convolution을 사용하였다(위에랑 다른 Depthwise임). 이는 계산비용을 줄이면서 더 나은 spatial context를 추출할 수 있다.

 

마지막으로 Wide-Focus로 불리는 모듈은 medical image 속 fine-grained information을 추출하기 위한 모듈이다. 그림과 같이 총 3갈래의 conv layer로 나눠진 후 합쳐지는데 이는 2개의 dilated conv(dilation=2, 3)와 일반 conv의 출력을 합친 후 다시 추가적인 conv layer를 통해 특징을 병합한다.

 

Encoder

인코더는 총 4개의 FCT layer로 구성되어 있으며 각기 다른 크기에서 각 클래스들의 ROI feature를 획득하기 위해 피라미드 스타일의 multiple 이미지 input을 목표로 하지만 이같은 동작이 없어도 뛰어난 성능을 보였다.

 

Decoder

디코더는 bottleneck의 출력을 입력받아 binary/segmentation maps을 resample하도록 학습된다. 인코더와 똑같은 구조를 가지고 있으며 encoder의 출력과 skip-connection으로 연결되어 더해진다. (28x28)크기가 되는 가장 낮은 구간에서는 supervision을 사용하지 않았는데, 이는 ROI가 segmentation되기 너무 작고 모델의 성능을 저하시키기 때문이다.

 

Implementation Details

loss function : equally weighted combination of cross-entropy and dice loss

optimizer : Adam

learning rate : 1e-3

scheduler : plateau (thorugh monitoring the validation loss)

epochs : warm up 50 + training 250

data augmentation

- random rotation(0 ~ 360)

- zoom range (max 0.2)

- shear range (max 0.1)

- horizontal/vertical shift (max 0.3)

- horizontal/vertical flip

모든 모델은 사전훈련 가중치를 사용하지 않고 random initialzed 되었다.

 

자세한 실험결과는 논문을 참고하길 바람