[論文レビュー] How to Efficiently Adapt Large Segmentation Model(SAM) to Medical Images

Xinrong Hu, Xiaowei Xu+1 more|arXiv (Cornell University)|Jun 23, 2023

Advanced Neural Network Applications被引用数 29

ひとこと要約

SAMエンコーダのみを微調整し、軽量で非プロンプト可能な予測ヘッド（ViT AutoSAM、CNN、または Linear）を追加してSAMを医療画像分割へ適応させ、プロンプトなしの少数ショット学習と多クラスマスクを実現する。

ABSTRACT

The emerging scale segmentation model, Segment Anything (SAM), exhibits impressive capabilities in zero-shot segmentation for natural images. However, when applied to medical images, SAM suffers from noticeable performance drop. To make SAM a real ``foundation model" for the computer vision community, it is critical to find an efficient way to customize SAM for medical image dataset. In this work, we propose to freeze SAM encoder and finetune a lightweight task-specific prediction head, as most of weights in SAM are contributed by the encoder. In addition, SAM is a promptable model, while prompt is not necessarily available in all application cases, and precise prompts for multiple class segmentation are also time-consuming. Therefore, we explore three types of prompt-free prediction heads in this work, include ViT, CNN, and linear layers. For ViT head, we remove the prompt tokens in the mask decoder of SAM, which is named AutoSAM. AutoSAM can also generate masks for different classes with one single inference after modification. To evaluate the label-efficiency of our finetuning method, we compare the results of these three prediction heads on a public medical image segmentation dataset with limited labeled data. Experiments demonstrate that finetuning SAM significantly improves its performance on medical image dataset, even with just one labeled volume. Moreover, AutoSAM and CNN prediction head also has better segmentation accuracy than training from scratch and self-supervised learning approaches when there is a shortage of annotations.

研究の動機と目的

• 医療画像領域へSAM（自然画像基盤モデル）の適応の必要性を動機づける。
• SAMエンコーダを凍結し、非プロンプト可能な予測ヘッドを追加する軽量微調整戦略を提案する。
• 少量のラベルデータの下で三つのヘッドアーキテクチャを評価する。
• 公開医療画像データセットで、ゼロ-shot/自己教師ありのベースラインと比較してラベル効率の改善を示す。）
• method([
• SAMエンコーダの重みを凍結し、軽量なタスク特化ヘッドを付加して微調整する。
• SAMマスクデコーダを非プロンプトヘッドに置換; AutoSAMでクラスごとに埋め込みを複製して多クラスマスクを可能に。
• 三つのヘッドアーキテクチャを評価: ViTベースのAutoSAM, CNNベースのヘッド（UNet風デコーダ）, Linearヘッド。
• 少量のラベル付きボリューム（1つまたは5つ）をCross-EntropyとDice損失の混合で訓練。
• UNet from scratchとSimCLRベースの自己教師付き事前訓練、そしてプロンプト付きゼロショットSAM（ボックス）と比較。

提案手法

• Freeze SAM encoder weights and attach a lightweight, task-specific head for finetuning.
• Replace the SAM mask decoder with a non-promptable head; enable multi-class masks by duplicating embeddings per class in AutoSAM.
• Evaluate three head architectures: ViT-based AutoSAM, CNN-based head (UNet-like decoder), and Linear head.
• Train with a small number of labeled volumes (1 or 5) using a mix of Cross-Entropy and Dice losses.
• Compare against UNet trained from scratch and SimCLR-based self-supervised pretraining, as well as the original zero-shot SAM with box prompts.

実験結果

リサーチクエスチョン

• RQ1 SAMのエンコーダを凍結し、軽量な非プロンプトヘッドを追加することで、限られたアノテーションで医用画像のセグメンテーションで競争力を持てるか？
• RQ2 少数ショット設定で最良の性能を発揮するヘッドアーキテクチャはどれか（AutoSAM ViT、CNN、Linear）？
• RQ3 AutoSAMは医療データセット全体でプロンプトなしの効率的な多クラス分割を可能にするか？

主な発見

手法	Dice%	ASSD	RV	Myo	LV	平均
UNET	13.45 ± 1.89	16.24 ± 4.14	22.95 ± 0.47	17.55 ± 2.05	51.55 ± 6.42
UNET + SimCLR	14.25 ± 6.52	19.40 ± 6.36	27.54 ± 9.80	20.40 ± 3.95	33.14 ± 4.39
Encoder + LN	0.00 ± 0.00	20.42 ± 13.20	48.40 ± 22.50	22.94 ± 12.32	49.38 ± 12.32
Encoder + CNN	30.66 ± 14.28	39.96 ± 8.14	50.55 ± 13.56	40.39 ± 11.90	38.13 ± 16.42
AutoSAM (ft all)	17.10 ± 9.76	30.05 ± 7.77	43.82 ± 13.91	30.32 ± 10.05	25.93 ± 1.94
AutoSAM	31.66 ± 13.26	33.49 ± 9.23	52.83 ± 16.49	39.32 ± 12.82	23.59 ± 2.07
sup w/ UNET	40.36 ± 2.36	52.23 ± 3.80	62.91 ± 5.58	51.83 ± 3.41	32.28 ± 1.40
5 volumes / UNET + SimCLR	45.48 ± 4.65	58.20 ± 6.12	68.95 ± 3.88	57.18 ± 3.20	28.98 ± 7.13
5 volumes / Encoder + LN	22.07 ± 11.2	37.38 ± 11.56	33.69 ± 27.63	31.05 ± 16.14	-
5 volumes / Encoder + CNN	59.87 ± 1.86	62.81 ± 2.82	78.96 ± 2.79	67.21 ± 1.32	25.46 ± 11.14
5 volumes / AutoSAM (ft all)	22.43 ± 18.03	37.08 ± 13.49	53.75 ± 15.08	37.76 ± 15.22	24.44 ± 9.92
5 volumes / AutoSAM	58.48 ± 3.90	62.18 ± 2.97	80.58 ± 1.42	67.08 ± 2.56	17.54 ± 3.65
5 volumes / unsup SAM (box)	53.57 ± 0.86	39.60 ± 0.65	0.00 ± 0.00	31.06 ± 0.41	7.83 ± 0.67

• 軽量ヘッドでSAMエンコーダを微調整するだけで、1つのラベル付きボリュームで医療分割を大幅に改善。
• AutoSAMおよびCNNヘッドはゼロ-shotやSimCLRの自己教師あり前学習よりも低データ域で優れている; Linearヘッドは過学習のため劣る。
• AutoSAM（ViTベースのヘッド）とCNNヘッドは他のベースラインより高いDiceスコアを達成し、AutoSAMはしばしばASSDでより良い。
• 大きなSAMエンコーダサイズ（vit-h）は一般に結果を改善するが、Encoder + CNNよりAutoSAMはエンコーダサイズの感度が低い。
• ラベルデータが増えると（5ボリューム）、特にDiceスコアでAutoSAMとCNNヘッドの性能差が広がる。
• AutoSAMはプロンプトなしで1回の推論で複数クラスのマスクを生成でき、クラスごとに埋め込みを複製する。