nnFormer: Volumetric Medical Image Segmentation via a 3D Transformer
Transformer, the model of choice for natural language processing, has drawn scant attention from the medical imaging community. Given the ability to exploit long-term dependencies, transformers are promising to help atypical convolutional neural networks t
ieeexplore.ieee.org
Abstract
자연어 처리를 위해 설계된 트랜스포머는 medical analysis에서도 주목받고 있다. 하지만 최근의 접근법들은(2022년초에 아카이브 등록된 논문이니 참고 바랍니다.) 트랜스포머를 단지 convolutional representations을 보조하기위해 전역 정보를 얻는 모듈로만 사용하고 있다. volumetric medical image segmentation을 위해 설계된 3D 트랜스포머 기반의 nnFormer는 conv+self-attention뿐만 아니라 locan/global volume-based self-attention을 통해 volume 표현을 학습한다. 또한 기존 U-Net에서 사용되던 skip-connection의 연산을 concat/summation이 아닌 skip attention의 방법을 제안한다.
Introduction
transformer와 convnet을 조합하거나 transformer만을 이용한 U-Net 형태의 여러 모델들이 제안되었지만 좋은 성능에도 불구하고 self-attention과 convolution의 조화가 적절하지 않다. nnFormer은 하이브리드 형태의 구조를 통해 convolution과 self-attention간에 학습 성능을 향상시켰다.
Method
Encoder
nnFormer은 \(X\in R^{H\times W\times D}\)을 만족하는 3D 패치 이미지 이다. H, W, D는 각 높이, 너비, 입력 scan의 깊이이다. ViT와 Swin Transformer의 embedding block에서 사용된 큰 크기의 kernel보다는 작은 kernel size가 초기 단계의 임베딩에서 더욱 좋은 성능을 보였다. 이는 패치단위의 위치 인코딩을 하는 트랜스포머보다 pixel-level을 통해 공간 정보를 더욱 정확하게 주입하였다. 또한 작은 사이즈의 kernel은 계산비용을 줄이면서 동일한 크기의 receptive field를 제공한다. Fig 2. b에 표기된 것처럼 embedding block은 kernel size=3인 4개의 conv layer로 구성되어 있다. 또한 각각의 conv이후에는 GELU + LayerNorm이 수행된다.
Local Volume-based Multi-head Self-attention(LV-MSA)
임베딩 이후에는 고차원의 텐서 \(X_e\)를 트랜스포머 블록으로 입력한다. 중요한 점은 다양한 크기로 down-sampling된 객체들과 고차원에서의 공간 정보를 엮어 long-term dependencies를 확보하는 것이다. Swin-Transformers와 다르게 3D local volume을 이용해 self-attention을 계산한다. 우리는 3D local volume을 적용한 Fig 3a와 같은 블록을 사용하였다.
수식으로 표현하면 다음과 같으며 \(l\)은 layer index, MLP = multi-layer perceptron을 뜻한다. 또한 SLV는 Shifted Local Volume을 뜻한다.
Down-sampling layer
우리는 patch merging 을 대체한 straightforward strided convolution로 성능을 향상시켰다. conv를 이용한 다운샘플링은 다양한 시각에서 객체 개념을 계층적으로 모델에 주입하기 좋다.
Skip Attention
conv에서의 skip connection은 concat이나 summation 연산을 통해 featuer정보를 합산한다. 우리는 \(l\)번째 인코더의 출력에 대해 linear projection을 수행하여 \(K^l, V^l\)의 key, value를 획득한다. 이는 디코더의 업샘플링 중 획득한 \(Q^l\)와 어텐션 연산을 수행한다.
Experiments and Results
자세한 결과는 논문을 참조하고 바란다.
initial lr : 0.01, 'poly' decay strategy
optimizer : SGD, momentum:0.99, weight_decay : 3e-5
loss : cross entropy + dice loss
epochs : 1000
augmentation : rotation, scaling, gaussian noise, gaussian blur, brightness and contrast, simulation of low resolution, gamma augmentation, mirroring (이렇게나 많이..?)
Deep supervision
- 마스크를 디코더의 각 스테이지의 출력에 맞는 크기로 변형시켜 단계적인 출력에 대해 loss를 계산함
ex) 디코더의 스테이지별 출력이 64, 128, 256 이라면 마스크 또한 다운샘플링을 통해 64, 128, 256의 크기로 변형한 후 손실 계산
