[23’] Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation
카테고리: Medical
🔍 Abstract
SAM(Segment Anything)을 medical domain에 fine-tuning시키기 위한 Med-SA(Medical SAM Adapter)를 제안한다. 논문에서는 흔히 사용되는 adaptation을 사용한 fine-tuning만으로는 성능 향상이 부족했기에 두 가지 테크닉을 도입하여 medical-domain specific fine-tuning을 진행하였다.
- SD-Trans(Space-Depth Transpose): 2D input만을 사용하는 SAM이 3D input의 spatial information을 잃어버리는 문제를 해결하기 위해, $xy$ 평면의 space 뿐만 아니라 $z$축으로의 depth 방향으로도 정보 교환이 이루어지도록 한다.
- Hyp-Adpt(Hyper-Prompting Adapter): prompt-conditioned adaptation을 통한 decoder fine-tuning을 진행한다.
1. Architecture
저자들은 SAM에 크게 세 가지 변형을 가했다. (1) Fine-tuning을 위한 adapter, (2) 3D correlation을 반영하기 위한 SD-Trans, 그리고 (3) prompt-conditioned adaptation을 위한 Hyp-Adpt이다. 이 중 핵심 contribution이라고 할 수 있는 SD-Trans와 Hyp-Adpt에 대해 자세히 살펴보자.
SD-Trans는 2D input slice 간의 correlation을 반영하기 위한 테크닉이라고 볼 수 있다. 여기서 slice는 z축과 수직으로 자른 단면들을 말한다. 3D input size가 $D \times H \times W \times C$라고 하면, 이후 transformer embedding을 통해 $D$개의 2D slice $H \times W \times C$는 각각 $N \times L$으로 변환된다. 이때 $N$은 embedding의 수, $L$은 embedding의 길이로 볼 수 있다. 따라서 일반적인 SAM을 사용했다면 $D$개의 2D slice embedding $N \times L$ 내부에서 interaction이 이루어진다. $D$개의 2D slice 간의 correlation은 반영되지 않는 것이다. 하지만 단순히 $D \times N \times L$을 transpose하여 $N \times D \times L$로 변환한 뒤 계산한다면, $D \times L$ 내부에서 interaction이 이루어지므로 $D$개의 2D slice 간의 correlation을 반영할 수 있다. 정리하면 다음과 같다.
- Spatial branch: Input $D \times N \times L \rightarrow$ intra-slice interaction
- Depth branch: Transposed input $N \times D \times L \rightarrow$ inter-slice interaction
Hyp-Adpt는 단순히 prompt를 multi-head attention에 cross-attention을 통해 추가하는 것을 넘어 prompt에 따라 model weight이 변하도록 하는 hyper-prompting 테크닉이다. Prompt embedding $e^ {\text{prompt}}$로부터 MLP, reshape를 거쳐 prompt weight $w^ {\text{prompt}}$를 얻는다. 그리고 adapter로부터 얻은 embedding $e^ {\text{down}}$에 아래와 같이 연속적으로 prompt weight를 적용한다.
\[e^ {\text{down}} _ {n+1} = \text{ReLU} (\text{Norm}(e^ {\text{down}} _ n \otimes w^ {\text{prompt}}))\]이처럼 model-weight 수준에서의 prompt conditioning을 진행하는 것의 이점은 다양한 input에 대해 flexible한 prompt를 적용할 수 있다는 점이다. Multi-modal input을 받아야 하는 medical domain에 적합한 adaptation이라는 것이다.
2. Experiments
저자들은 abdominal organ segmentation, multi-modality image segmentation 등으로 평가를 진행했는데 다른 SOTA 모델이라고 할 수 있는 MedSAM에 비해 주목할 만한 큰 성능 향상이 있었음을 알 수 있다. 경제적이면서 효율적인 fine-tuning이 잘 이루어졌다는 것을 알 수 있다.
Ablation study 결과에서는 단순 adapter만 사용한 경우보다 SD-Trans와 Hyp-Adpt를 추가하는 경우 더욱 성능이 향상되었음을 알 수 있다.
💡 Summary
Medical SAM Adapter는 SAM을 medical domain에 fine-tuning시키기 위한 방법론을 제안한다. 본 논문에서는 (1) fine-tuning을 위한 adapter, (2) 3D correlation을 반영하기 위한 SD-Trans, 그리고 (3) prompt에 의해 model weight이 변하는 hyper-prompting을 구현한 Hyp-Adpt를 사용하여 효율적이고 경제적인 성능 향상을 이룰 수 있었다. Medical domain 뿐만 아니라 3D input을 사용하거나 multi-modal input을 사용한다면 시도해볼 수 있는 방법이다.
댓글 남기기