[2025-1] 박서형 - Swin-Unet: Unet-like Pure Transformer forMedical Image Segmentation

https://arxiv.org/abs/2105.05537

 

Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

 

 

 

1. introduction

medical 분야에서 segmentation 작업이 중요해졌는데 여기서 U-Net 같은 FCNN 기반의 방법들이 뛰어난 성능을 보였다.그러나 CNN의 locality 한계로 long range dependency 학습에 어려움이 생겼고 이를 transformer로 보완하고자 하는 시도가 있었다. Swin-Unet은 Swin Transformer의 성공에 영감을 받아 제안된 구조로 transformer 기반의 U 자형 구조로 Encoder, Bottleneck, Decoder, Skip Connection으로 구성된다. 이때 patch expanding layer를 통해 높은 정확도와 generalization 성능을 보여주었다. 이 연구는 transformer를 활용한 medical image segementation의 새로운 가능성을 제시한 점에서 의의가 있다. 

 

 

 

2. Related work

Swin-Unet은 좋은 성능을 보여준 Swin Transformer 블록을 기본 단위로 medical image segmentation을 위한 U자형 encoder-decoder 구조를 구축했다. 그리고 최근 연구에서 transformer에 cnn 구조를 보완적으로 도입하고자 하는 시도가 많았던 것과 달리 Swin-Unet은 순수 transformer만을 이용해 작업을 수행한다. 

 

 

 

3. Method

1) architecture overview

 

 

Swin-Unet은 Encoder, Bottleneck, Decoder, Skip Connections로 구성되어 있고 기본 단위는 swin transformer block이다. 

 

Encoder : 우선 encoder는 input image를 sequence embedding으로 변환하기 위해 4x4 patch로 분할하고 각 패치의 차원은 4x4x3=48이 된다.  그 후 linear embedding layer를 통과시켜 패치를 원하는 차원으로 변환시켜준다. 그리고 patch merging layer를 거쳐 downsampling과 차원 증가를 한 후 swin transformer block을 통과하여 feature를 학습한다. 

 

Decoder : U-Net의 대칭 구조에 영감을 받아 설계되었다. swin transforemer block과 patch expanding layer로 이루어져 있으며 인코더의 swin transformer의 output이 context feature와 skip connection을 통해 융합되어 downsampling에 의한 공간 정보 소실을 방지한다. 그리고 decoder는 encoder와 달리 patch expanding layer를 통해 upsampling을 수행하여 input image의 resolution을 복원한다. 이를 통해 나온 feature map을 linear projection layer에 적용해 픽셀 단위 segmentation 예측을 출력한다. 

 

 

 

2) Swin transformer block

 

전통적인 multi-head self-attention과 달리 shifted window를 기반으로 구성된다. swin transformer는 layernorm, multi-head self attention, residual connection, 2 layer mlp로 구성되어 있다. 

 

- W-MSA : window 기반의 multi-head self-attention

- SW - MSA :  shifted window 기반의 multi-head self-attention

 

swin transformer block에서 input이 통과하면서 거치는 수식은 다음과 같다. 

 

이때 self-attention 공식은 다음과 같다. 

 

 

 

3) patching merging layer

patch merging layer는 input image를 4개의 부분으로 나누고 concatenate해서 resolution은 2배 downsampling 시키고 feature dimension은 4배로 증가시킨다. 이후 linear layer를 통과하면 featur dimension을 두 배로 키운다.

 

 

 

4)  bottleneck

두 개의 swin transformer만 사용하여 구성되어 feature dimension과 resolution은 변하지 않고 더 깊은 feature representation을 학습할 수 있게 한다. 

 

 

 

5) patch expanding layer

resolution을 upsampling하면서 feature dimension은 반으로 줄여 효율적으로 공간 정보를 복구한다. 

 

- linear layer : feature dimension 2배로 증가 ( W/32 × H/32 × 8C  →  W/32 × H/32 × 16C )

-rearrange operation : resolution 2배 증가, feature dimension 1/4배 감소 ( W/32 × H/32 × 16C → W/16 × H/16 × 4C )

 

 

 

6) skip connection

 

U-Net과 유사하게 encoder에서 추출된 featur( shallow feature )와 decoder의 deep feature를 concat하여 downsampling에 의한 공간 정보 소실을 방지한다. 이때 concat은 linear layer를 거쳐 upsampling되어 차원이 동일한 feature끼리 진해ㅐㅇ된다. 

 

 

 

 

4. Experiment

1) segmentation accuracy

Swin-Unet은 Synapse multi-organ CT 데이터셋에서 가장 뛰어난 성능을 보였다.

 

2) generalization

ACDC 데이터셋에서도 Swin-Unet은 뛰어난 성능을 보이며 다양한 종류의 medical image에서 우수한 성능을 보일 수 있음을 증명했다.

 

3) upsampling & number of skip connection

upsampling을 도입한 Swin-Unet이 도입하지 않은 다른 모델보다 좋은 성능을 보이며 upsampling의 중요성을 보여줬다. 또한 skip connection의 수가 3일 때 가장 높은 성능을 보였다. 이로써 upsampling과 skip connection의 수 또한 성능을 높이는 중요한 요소임을 보여줬다.