[論文レビュー] 3DSAM-adapter: Holistic adaptation of SAM from 2D to 3D for promptable tumor segmentation

Shizhan Gong, Yuan Zhong+6 more|arXiv (Cornell University)|Jun 23, 2023

Advanced Neural Network Applications被引用数 32

ひとこと要約

この論文は、体積医用腫瘍分割のためのSAMの2Dから3Dへの総合的でパラメータ効率の良い適応を提示し、最小限の調整可能パラメータとボリュームあたり1つのプロンプトで最先端の結果を達成します。

ABSTRACT

Despite that the segment anything model (SAM) achieved impressive results on general-purpose semantic segmentation with strong generalization ability on daily images, its demonstrated performance on medical image segmentation is less precise and not stable, especially when dealing with tumor segmentation tasks that involve objects of small sizes, irregular shapes, and low contrast. Notably, the original SAM architecture is designed for 2D natural images, therefore would not be able to extract the 3D spatial information from volumetric medical data effectively. In this paper, we propose a novel adaptation method for transferring SAM from 2D to 3D for promptable medical image segmentation. Through a holistically designed scheme for architecture modification, we transfer the SAM to support volumetric inputs while retaining the majority of its pre-trained parameters for reuse. The fine-tuning process is conducted in a parameter-efficient manner, wherein most of the pre-trained parameters remain frozen, and only a few lightweight spatial adapters are introduced and tuned. Regardless of the domain gap between natural and medical data and the disparity in the spatial arrangement between 2D and 3D, the transformer trained on natural images can effectively capture the spatial patterns present in volumetric medical images with only lightweight adaptations. We conduct experiments on four open-source tumor segmentation datasets, and with a single click prompt, our model can outperform domain state-of-the-art medical image segmentation models on 3 out of 4 tasks, specifically by 8.25%, 29.87%, and 10.11% for kidney tumor, pancreas tumor, colon cancer segmentation, and achieve similar performance for liver tumor segmentation. We also compare our adaptation method with existing popular adapters, and observed significant performance improvement on most datasets.

研究の動機と目的

• SAMの3D医用腫瘍分割におけるドメインと次元のギャップによるサブ最適で不安定な性能を解消する。
• 可能な限り事前学習済み重みを再利用しつつ、体積入力に対してSAMを総合的に適応させる。
• 2D事前学習を3D医療データへ橋渡しするためのパラメータ効率的なファインチューニングと軽量アダプターを開発する。
• 複数の腫瘍データセットにわたって、プロンプトへの頑健性を改善し高い分割精度を維持する。

提案手法

• 体積入力を受け付けるように画像エンコーダを最小限のパラメータ変更で改良し、3Dパッチ埋め込み、3D位置エンコーディング、メモリ効率的な戦略を用いた3Dアテンションブロックを導入する。
• 微調整用の軽量スペーシャルアダプターを導入し、ほとんどのSAM重みを凍結し、アダプターと正規化層のみを訓練する。
• プロンプトエンコーディングを、画像特徴マップから埋め込みを取得するビジュアルサンプラーに置換し、トークン爆発とプロンプトからのノイズを緩和するためにクロスアテンションを用いた小規模なグローバルクエリを使用する。
• マスクデコーダを軽量な3D CNNへ更新し、多層集約を備えて高解像度の3Dマスクを生成する。
• ボリュームあたり単一または少数のポイントプロンプトで訓練し、ノイズの多いプロンプトへの頑健性を向上させる背景前景サンプリングを含む。

Figure 1: Overview of our proposed method for 3DSAM-adapter. The original ViT is modified to support volumetric inputs. The prompt encoder is redesigned to support 3D point prompt, and the mask decoder is updated to 3D CNN with multi-layer aggregation to generate 3D segmentation.

実験結果

リサーチクエスチョン

• RQ1 SAMの総合的な2Dから3D適応は、事前学習済み知識を保持しつつ医療ボリュームの3D空間パターンを効果的にエンコードできるか。
• RQ2 視覚サンプラーに基づくプロンプトエンコーディングは、限られたプロンプト下で3D医療分節において従来の位置エンコーディングより優れているか。
• RQ3 3Dマスクデコーダの多層集約は、小さくコントラストの低い腫瘍の分割精度にどのような影響を与えるか。
• RQ4 ボリューム分割における調整可能パラメータ数と分割性能のトレードオフは、完全なファインチューニングおよび他のアダプターと比較してどうなるか。

主な発見

方法	腎臓腫瘍 Dice	腎臓腫瘍 NSD	膵臓腫瘍 Dice	膵臓腫瘍 NSD	肝臓腫瘍 Dice	肝臓腫瘍 NSD	結腸癌 Dice	結腸癌 NSD	チューニングされたパラメータ数
nnU-Net (Nat. Methods 2021)	73.07	77.47	41.65	62.54	60.10	75.41	43.91	52.52	30.76M
TransBTS (MICCAI 2021)	40.79	37.74	31.90	41.62	34.69	49.47	17.05	21.63	32.33M
nnFormer (arXiv 2021)	45.14	42.28	36.53	53.97	45.54	60.67	24.28	32.19	149.49M
Swin-UNETR (CVPR 2022)	65.54	72.04	40.57	60.05	50.26	64.32	35.21	42.94	62.19M
UNETR++ (arXiv 2022)	56.49	60.04	37.25	53.59	37.13	51.99	25.36	30.68	55.70M
3D UX-Net (ICLR 2023)	57.59	58.55	34.83	52.56	45.54	60.67	28.50	32.73	53.01M
SAM-B (1 pt/slice) [4]	36.30	29.86	24.01	26.74	6.71	7.63	28.83	33.63	–
Ours (1 pt/volume)	73.78	83.86	54.09	76.27	54.78	69.55	48.35	63.65	25.46M
SAM-B (3 pts/volume) [4]	74.91	84.35	54.92	77.57	56.30	70.02	49.43	65.02	25.46M
SAM-B (10 pts/slice) [4]	75.95	84.92	57.47	79.62	56.61	69.52	49.99	65.67	25.46M

• 1プロンプトあたりの我々の手法でのDiceスコアは腎臓73.78、NSD83.86；膵臓54.09、76.27；肝臓54.78、69.55；結腸48.35、63.65。
• ボリュームあたり3プロンプトを用いた場合Dice74.91、NSD84.35（腎75.??、膵臓54.92、肝77.57、結腸56.30、表に従う）、より多くのプロンプトで性能が向上。
• 最先端のベースライン（nnU-Net等）と比べ、提案手法は多くのデータセットで優位となることが多く、膵臓と結腸癌タスクで大幅な改善を示す。
• 既存のアダプターや完全なファインチューニングベースラインに対して、元のモデルの調整可能パラメータを16.96%のみで超える高いパラメータ効率を示す。
• ビジュアルサンプラーによるプロンプトエンコーディングは、位置エンコーディングに対して有意に優れており（KiTS21のアブレーションでDice改善約40%）。
• マスクデコーダの多層集約は、非集約型と比較してDiceを約15.75%改善。

Figure 2: Spatial adapter.