[논문 리뷰] Segment Any Cell: A SAM-based Auto-prompting Fine-tuning Framework for Nuclei Segmentation

Saiyang Na, Yuzhi Guo외 3명|arXiv (Cornell University)|2024. 01. 24.

Artificial Intelligence in Healthcare and Education인용 수 9

한 줄 요약

SAC는 Low-Rank Attention Adapter와 자동 프롬프트 생성기를 통합하여 샘 플링(/?)을 개선하고 MoNuSeg 및 DSB 데이터셋에서 자동 프롬프트를 사용한 최첨단 결과를 달성합니다.

ABSTRACT

In the rapidly evolving field of AI research, foundational models like BERT and GPT have significantly advanced language and vision tasks. The advent of pretrain-prompting models such as ChatGPT and Segmentation Anything Model (SAM) has further revolutionized image segmentation. However, their applications in specialized areas, particularly in nuclei segmentation within medical imaging, reveal a key challenge: the generation of high-quality, informative prompts is as crucial as applying state-of-the-art (SOTA) fine-tuning techniques on foundation models. To address this, we introduce Segment Any Cell (SAC), an innovative framework that enhances SAM specifically for nuclei segmentation. SAC integrates a Low-Rank Adaptation (LoRA) within the attention layer of the Transformer to improve the fine-tuning process, outperforming existing SOTA methods. It also introduces an innovative auto-prompt generator that produces effective prompts to guide segmentation, a critical factor in handling the complexities of nuclei segmentation in biomedical imaging. Our extensive experiments demonstrate the superiority of SAC in nuclei segmentation tasks, proving its effectiveness as a tool for pathologists and researchers. Our contributions include a novel prompt generation strategy, automated adaptability for diverse segmentation tasks, the innovative application of Low-Rank Attention Adaptation in SAM, and a versatile framework for semantic segmentation challenges.

연구 동기 및 목표

• SAM 기반 방법의 프롬프트 품질 병목현상을 해결하여 핵 분할 개선을 모티브로 제시한다.
• 효율적인 적응을 위한 SAM의 트랜스포머에 적용되는 LoRA 기반 파인튜닝 메커니즘을 개발한다.
• 핵 분할을 위한 고품질 양성/음성 프롬프트를 자동으로 생성하는 프롬프트 제너레이터를 만든다.
• 기존의 파인튜닝 및 적응 방법과 자동 프롬핑의 호환성을 입증한다.
• 핵 분할을 넘어서는 의미 분할 작업에 대한 SAC의 일반성을 보여준다.

제안 방법

• SAM의 ViT 어텐션 층 내에 Q와 V 행렬을 수정하고 K는 고정한 상태의 Low-Rank Attention Adapter(LoRA)를 적용한다.
• 이미지에서 BCE 목표를 사용해 고품질 프롬프트를 생성하는 Auto Prompt Generator(UNet 기반)를 도입한다.
• LoRA와 자동 프롬프트를 입력으로 받는 학습 가능한 SAM 마스크 디코더를 고정된 SAM 이미지 인코더와 함께 사용한다.
• 생성되거나 전문가 프롬프트를 처리하는 프롬프트 인코더(고정)를 사용하고 최종 마스크는 SAM 마스크 디코더에서 얻는다.
• 필요 시 분할을 돕기 위한 수동 프롬프트를 위한 선택적 추론 경로를 제공한다.

실험 결과

연구 질문

• RQ1 LoRA 기반 SAM 파인튜닝이 매개변수 수를 증가시키지 않으면서 핵 분할을 개선할 수 있는가?
• RQ2 자동 프롬프트 생성이 전문가 프롬프트만 사용하는 경우보다 분할 품질을 개선하는가?
• RQ3 중심점 기반 프롬프트 선택과 직접 확률 기반 선택이 분할 성능에 어떤 차이를 만드는가?
• RQ4 SAC가 핵 분할을 넘어서 다양한 의미 분할 작업에 적용 가능한가?
• RQ5 프롬프트의 양과 품질이 SAM 기반 핵 분할 성능에 어떤 영향을 미치는가?

주요 결과

• SAC는 MoNuSeg 및 DSB 데이터셋에서 SAM 기반 변형 및 다수의 CNN/트랜스포머 방법을 포함한 베이스라인보다 우수한 분할 점수(Dice, IoU, F1)를 달성한다.
• 제로샷 SAM(1-expert)은 파인튜닝된 접근법에 비해 성능이 낮은 반면, LoRA와 자동 프롬프트를 갖춘 SAC는 SAM-FT 및 MSA를 모든 지표에서 능가한다.
• LoRA 기반의 SAM 파인튜닝과 자동 프롬프트는 파라미터를 늘리지 않으면서 Dice 및 IoU를 MS보다 높게 달성한다.
• 점수 증가형 자동 프롬프트 사용(점수)과 소수의 전문가 프롬프트가 성능을 향상시키며, 중심점 기반 프롬프트 선택은 더 큰 포인트 수에서 직접 확률 기반 선택보다 우수하다.
• SAC는 SAM-FT 및 MSA보다 Dice 수렴 속도가 빠르고 전체적으로 더 같거나 더 적은 학습 가능 매개변수를 사용하여 효율성을 강조한다.