서론
연구배경 및 필요성
image segmentation 문제에 자주 사용되는 CNN(Convolutional Neural Network)는 intrinsic locality라는 단점이 있어 국소적으로 feature map을 추출하게 된다는 단점이 있다.
Transformer의 경우 global self-attention 을 사용하여 이미지 전체와의 종속성을 파악할 수 있지만 low-level detail에 취약하다는 단점이 있다.
의료 영상의 경우 low-level detail(작은 장기, 혈관 등) 뿐만 아니라 넓은 범위에서의 종속성(큰 장기 등)도 segmentation에 꼭 필요하기 때문에, 각 구조의 단점을 보완할 필요성이 있다.
논문의 주요 목표와 기여
저자가 제안한 TransUNet은 U-Net과 Transformer의 강점을 결합하여 의료 영상 분할의 성능을 개선하는 데 중점을 준 모델이다. 기존 U-Net이 가지고 있던 Intrinsic locality 라는 한계점을 Transformer의 Self-attention이라는 global context 처리 능력으로 보완하고, Transformer의 단점은 CNN의 로컬 정보로 보완하고자 했다.
이론적 배경 탐구
1) UNet
Semantic Segmentation을 위해 고안된 구조. 인코더 - 디코더가 대칭적인 모델이며, 블럭 내가 아닌 인코더 -> 디코더 전역적 skip-connection 있음.
2) Transformer
Self-Attention을 이용해 입력 간의 전역 의존성을 효율적으로 학습하고 병렬 처리를 통해 빠르고 유연하게 동작하는 모델.
3) ViT(Vision Transformer)
ViT는 이미지를 패치(patch)로 나누고 Transformer 구조를 적용해 전역 문맥을 학습하는 모델. CNN 없이도 대규모 데이터에서 뛰어난 성능을 보임.
Summary and Key points
CNN-Transformer Hybrid as Encoder
Patch Embedding 전에 CNN으로 feature map을 추출하는 방식을 사용한다. 이렇게 하면 1) decoding path에서 high-resolution CNN feature map을 사용 가능하고(skip connection을 통해) 2) 따라서 더 나은 성능을 기대할 수 있다.
Cascaded Upsampler(CUP)
순차적 업샘플러라는 뜻으로, 디코딩 과정에서 여러 단계의 upsampling으로 구성되어 사용되며 final segmentation mask를 산출하기 위해 hidden feature를 decoding한다. 2x upsampling, 3x3 conv, ReLU 등을 포함한다.
Skip-Connection of UNet
인코딩 과정에서 1/2, 1/4, 1/8 resolution 상의 feature를 디코딩 과정으로 skip-connection으로 넘겨줘 concat하면 transformer에서 소실된 local detail 정보를 추가해 줄 수 있다.
결과
MICCAI 2015에서 제공된 30개의 CT scan 데이터를 사용하여 abdominal multi-organ segmentation(8 of target organ)을 진행했고, Automated Cardiac Diagnosis Challenge(ACDC)에서 제공된 100개의 MRI scan을 사용하여 left ventricle, right ventricle, myocardium segmentation을 진행했다.
V-Net, DARR, U-Net, AttnUNet과 비교했을 때도 성능이 개선되었다.
+) DSC(Dice Similarity Coefficient)
두 집합간의 겹치는 정도를 측정하는 지표로 정확도 평가에 사용된다.
++)HD(Hausdorff Distance)
한쪽의 점 집합에서 다른쪽 점 집합 거리의 최대값으로 계산한다. 이 거리가 좁을수록 target과 경계가 유사하다.
ViT-None 구조와 ViT-CUP 비교 -> CUP의 유의미한 성능 개선 능력 확인
ViT-CUP 구조와 R50-ViT-CUP 비교 -> Hybrid encoder의 유의미한 성능 개선 능력 확인
R50-ViT-CUP 구조와 TransUNet 비교 -> Skip-connection 구조 통한 local+global context 융합 성능이 좋음
Ablation Study
1) Number of skip-connections
skip connection의 개수를 0개, 1개, 3개로 바꿔가며 성능을 관찰한다.
0개 -> R50-ViT-CUP
1개 -> 1/4 resolution에서만 skip
3개 -> 1/2, 1/4, 1/8 resolution에서 skip
그 결과 skip의 개수가 더 많을수록 더 나은 성능을 보였다.
+) Attention Gate를 통과하는 것은 transformer를 통과하는 것과 어떻게 다를까?
transformer를 통과하면 모든 patch 간의 global 관계를 학습하며 새로운 feature map을 형성한다. 반면 AG는 전체의 관계를 보는 대신 기존의 local detail을 보존하면서 필요없는 정보를 억제하는 역할만 한다.
2) Input resolution
Input image의 resolution을 224x224에서 512x512로 바꾸고 patch 크기는
같게 하여 진행했을 때, Seqence의 길이는 5배정도 길어지며 성능이 개선되었다. 즉, 넣어주는 영상의 resolution이 높을수록 segmentation 성능이 좋아진다.
3) Sequence length & patch size
Patch의 크기를 줄이면 2)와 같이 sequence의 길이가 증가하며 transformer 가 더 복잡한 관계를 학습 가능하므로 성능이 개선된다.
시각화
저자가 제시한 (b)TransUNet는 순수 CNN 기반 모델인 (d) 및 (e)와 비교했을 때 더 정확한 분할을 보여줬다. 예를 들어, 첫 번째 슬라이스에서 (d)와 (e)는 right kidney를 거의 분할하지 못한 반면 TransUNet에서는 target과 거의 흡사한 양상을 보였다. 이는 transformer를 통해 원본 위치 상 동떨어져 있는 patch와의 관계를 학습했기 때문으로 생각된다.
(c)는 첫 번째 슬라이스에서는 양호한 분할을 보여주는 듯 했으나 세 번째 슬라이스에서 left kidney의 정확한 모양이 미분할되었다. 이는 인코더-디코더 간의 skip connection이 없어 low-level detail이 전달되지 않았기 때문으로 예상된다.
결론
시사점
본 논문에서 제안한 TransUNet의 CNN-Transformer Hybrid 구조는 skip connection을 통해 local detail과 global context라는 장점들을 융합할 수 있었다. 이는 의료 영상 분할에서의 성능을 크게 향상시켰으며, CT와 MRI 데이터를 통한 실험에서 유의미한 성과를 거두었다. 더 나아가, 저자는 이 모델의 확장성을 증명하기 위해 지난 2023년 이를 3D 영상 데이터에 적용한 후속 연구를 발표했다.
개선 방안 및 연구 가능성
skip connection의 ablation study 부분에서 저자는 [[Attention Gate(AG)|Attention Gate]]를 1 개의 skip connection의 경로에 시범삼아 사용했을 때 성능이 개선되었지만 메모리 상의 이유로 이를 사용하지 않았다고 했다. 이 방식을 채택하는 대신 기존의 트랜스포머 구조보다 경량화된 Swin transformer를 사용하여 인코더의 윈도우 크기와 계층 수를 조절한다면 성능을 개선할 수 있을 것이다.
Reference
https://philipshrimp.github.io/posts/neural_networks/cnn_basics/
https://goldenrabbit.co.kr/2022/08/03/
https://github.com/Beckschen/TransUNet
Chen, J., Qi, X., Dou, Q., et al. (2021). TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation. arXiv preprint arXiv:2102.04306.
Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. MICCAI, 234-241.
Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention is All You Need. NeurIPS, 5998-6008.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., et al. (2020). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv preprint arXiv:2010.11929.
Comments
1) intrinsic locality 가 무엇인가요? CNN에서만 발생되는 문제인가요?
논문에서 intrinsic locality라고 표현한 것은 CNN 구조에서 주로 문제점으로 지적되는 국부성(전체 영상에 비해 좁은 receptive field 때문에 전체적인 관계나 멀리 떨어진 픽셀 간의 관계를 학습하지 못함) 입니다. 이를 방지하기 위해 convolution 레이어를 여러 번 통과하는 등의 방식을 통해 receptive field를 넓히며 위치별 패턴을 찾으려 합니다. CNN에서만 발생되는 문제라고 단정지을 수는 없지만(국부적으로, 패치 내에서만 관계성을 학습하는 모든 구조에서 발생한다고 볼 수 있음) CNN의 대표적인 단점인 것으로 알고 있습니다.
2) 국소적으로 보는것이 단점이 되는 이유가 무엇인가요?
CNN 구조에서는 임의의 위치에 있는 두 픽셀을 관련짓기 위해 필요한 연산량이 거리에 따라 증가합니다. 논문의 experiment에서 사용한 medical image의 경우(multi-organ segmentation) 다른 장기 간의 공간적 관계가 멀리 떨어져 있을 수 있으므로 global context 없이 국소적으로 학습하는 것이 단점이 됩니다. visualization했을 때, left kidney를 global context 없이 학습한 U-Net만이 segmentation하지 못한 것을 바탕으로 유추할 수 있습니다.
3) segmentaiton 말고 다른 medical domain task에서도 넓은 범위의 종속성이 중요한가요?
넓은 범위의 종속성은 segmentation에서 좋은 성능을 보였기 때문에 다른 task에도 필요하다고 생각합니다. 병변과 그 주변 조직과의 관계를 학습하여 classification하거나, 정확한 공간적 위치를 파악해 두 이미지를 mapping하는 registration task 등에서도 전역적인 관계를 학습한다면 더 정확한 위치관계를 알 수 있게 되고, 결과적으로는 성능을 향상시킬 것으로 보입니다.
4) self-attention 말고 다른 attention 기법이 이미 많이 나와있음에도 CNN + transformer 형태의 architecture를 구태여 섞어야 하는 이유가 설득력 있어 보이지 않습니다 😕 왜 섞는게 좋은지 설득할 수 있을만큼 명확한 이유가 있으면 좋겠는데 저자 생각 + 리뷰어 생각은 어떠신지요?
저자가 제시한 아이디어는 CNN으로 local detail을 보고, transformer로 global 정보를 보는 것이었습니다. 이 과정에서 단순 self-attention으로도 충분한 전역 학습이 가능하다고 판단한 것 같습니다. 또한, skip connection ablation study에서 연산 효율성을 위해 attention gate에 대한 실험을 더 하지 않은 것으로 보아 연산량 측면에서 부담이 생긴 것으로 보입니다. attention의 인사이트는 중요한 것에만 집중한다 에 있다고 생각하는데, 이로 인한 저수준 정보의 손실을 방지하기 위해, 또 계산효용성을 위해 굳이 하이브리드 인코더 구조를 사용했다고 생각됩니다. 이를 대체할 수 있는 attention 기법에 대해서는 더 연구해 보겠습니다.
5) 인코더로 CNN+Transformer를 사용하고, 디코더로 Unet의 디코더 부분을 사용한다고 이해했습니다. fig3의 unet 구조의 디코더 부분을 보면, unet 또한 단계별로 upsampling이 되는 것 같은데 본 논문에서 제안하는 Cascaded Upsampler와는 어떤 차이가 있는지 궁금합니다.
UNet의 디코더와 같은 연산 블럭을 가지고 있지만 transformer에서 나온 토큰을 재배열하여 CNN 인코더의 skip connection과 결합시키는 과정이 추가되어 있습니다. 기존의 UNet 디코더는 업샘플링을 사용해 복원하고 인코더에서 받아온 피처맵을 concat한 다음 conv를 적용합니다. TransUNet에서는 각 upsampling 단계에서 transformer의 출력 토큰들을 다시 공간에 reorganization해야 하고, 이에 skip connection의 CNN 피처를 projection하는 등의 채널, 해상도 매칭 과정을 거친다는 점이 추가되었습니다.
논문에서 CUP의 성능을 증명하기 위해 사용한 대조군은 naive upsampling baseline으로, 최종 토큰을 skip connection 없이 그저 bilinearly upsampling하는 것입니다. 이름(CUP)은 저자가 붙여줬지만... UNet의 구조를 거의 그대로 가져온 것이므로, 왜 굳이 "UNet 구조"를 transformer와 결합하려 했는가의 이유가 된다고 생각합니다.
6) Skip connection이 무엇인지, 어떤 역할을 하는지 간단히 설명 부탁드립니다! 또한, skip connection이 low-level detail을 전달한다고 했는데, high-level semantic feature와의 균형은 어떻게 유지되는지 궁금합니다.
TransUNet에서 차용한 UNet에서의 skip connection은 encoder와 decoder 사이의 정보를 전달하기 위한 구조입니다. CNN 인코더에서 이미지를 다운샘플링하면 세부 구조가 손실될 수 있는데, 이를 방지하기 위해 encoder의 중간 출력을 decoder에 직접 연결해 전달하는 것이 skip connection 입니다.
각 단계(1/4, 1/8, 1/16 resolution)에서 얻어진 feature map을 decoder에 전달하는 과정에서 다시 공간적 mapping을 해 더해줍니다. 이 과정에서 가중치의 동적 조절(둘 사이의 균형 유지에 대해) 등은 밝혀지지 않아, 원래의 skip-connection처럼 그저 원 정보를 더해주는 것 자체에 의의가 있는 것 같습니다.
