Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation
The Segment Anything Model (SAM) has recently gained popularity in the field of image segmentation due to its impressive capabilities in various segmentation tasks and its prompt-based interface. However, recent studies and individual experiments have show
arxiv.org
이전에 리뷰한 Segment Anything Model(SAM)을 이용하여 medical image segmentation에 적용한 논문이 있어 하는 리뷰
Abstract
Segment Anything Model(SAM)은 프롬프트 기반의 모델로 모든 segment작업에서 인상적인 성능을 보이고 있다. 하지만 최근의 연구들은 medical분야에서는 낮은 성능을 보인다. 본 논문에서는 SAM을 fine-tuning하여 성능을 향상시키기 보다는 효과적인 adaptation 기술을 통해 SAM에 의학 지식을 부여하는 Medical SAM Adapter(Med-SA)을 제안한다. 2D 기반의 SAM을 3D 의학 이미지에 적용하기 위해 Space-Depth Transpose(SD-Trans)와 프롬프트를 위한 Nyper-Prompting Adapter(HyP-Adpt)을 제안한다.
Introduction
SAM의 medical image에 대한 성능의 한계는 모델의 의학적 정보의 부족과 데이터셋에서 나타나는 객체간에 모호한 경계, 작은 마스킹 영역들, 저대비의 이미지(low contrast)가 원인이다. 최근의 연구 중 vanilla SAM을 이용해 의학 이미지에 특화하여 fully fine-tuning한 사례가 있지만 비용면에서 매우 비효율적이고 사전 훈련된 모델을 이용하는 것만으로도 매우 우수한 transferability(전이능력)을 지니고 있다고 한다. 본 논문에서는 NLP에서 사용되는 Adaption 모듈을 이용하여 약간의 파라미터만을 통해 모델을 학습시킨다. SD-Trans는 입력 임베딩의 spatial 차원을 depth 차원으로 바꾸어 self-attention 블록이 서로 다른 차원에 대한 정보를 처리할 수 있도록 한다. HyP-Adpt는 adaption embedding을 효율적으로 처리하고 visual prompt를 위한 가중치를 생성하여 넓고 깊은 prompt-adaption(?)을 생성한다.
Method
Med-SA architecture
우리는 medical task에서 SAM의 성능을 향상시키는 것이 목적이기 때문에 모든 파라미터를 학습시키는 대신, 모델은 frozen시킨 채 Adapter module만을 훈련시켰다. Adapter은 down-projection, ReLU, up-projection을 포함하며 모델에 bottleneck의 형태로 제공된다(Fig 1.b).
SAM의 인코더부분에서는 ViT블록마다 2개의 Adapter가 존재하는데 이는 각각 Multi-head attention 이후, residual connection 이전에 위치한다. 또한 2번째의 Adapter 이후에는 scale factor s를 이용해 임베딩 벡터의 크기를 조정한다.
디코더에서는 블록마다 3개의 Adapter을 사용하는데 첫번째로는 prompt embedding을 통합하기 위해 사용되는데 이 모듈을 Hyper-Prompting Adapter라고 부른다. 2번째는 인코더에서와 마찬가지로 MLP를 향상시키기 위한 것(scaling 부분에 있는 Adapter)이고 3번째는 image embedding-prompt간에 cross-attention을 하는 residual connection이후에 적용된다(Fig 1.d의 가장 오른쪽에 있는 Adapter).
SD-Trans
의료 이미지는 2D, 3D로 이루어져 있고 특히 3D는 각각의 슬라이스가 연관되어 있기 때문에 확실한 의사결정에 중요하다. SAM은 슬라이스 이미지에 대한 예측을 통합하여 3D 이미지 볼륨의 segmentation map을 획득할 수는 있으나 가까운 볼륨간에 관계를 놓칠 수 있다.
SD-Trans 는 각 블록내에 위치하는데 attention을 2개의 갈래로 나누어 space, depth에 대한 정보를 획득한다. 가중치를 공유하는 attention layer에 대해 깊이 \(D\), embedding \(N\), length \(L\)에 해당하는 \(D\times N\times L\)과 이를 transpose한 \(D\times N\times L\)을 입력한 후 둘을 더함으로써 깊이와 공간에 대한 정보를 모두 학습하게 된다.
HyP-Adpt
모델이 학습한 task와 downstream task는 다를 수 있기 때문에 visual prompt를 adapter에 적용하는 것은 매우 중요하다. 우리는 knowledge-condition을 위해 하나의 모델에서 다른 모델의 가중치를 생성하는 HyperNetwork로부터 영감을 받았다. HyP-Adpt는 projection과 reshaping만을 사용하며 adapter embedding에 행렬곱을 통해 직접 적용된다.
자세한 결과 및 implementation detail은 논문을 참조
