[논문 리뷰] ASPS: Augmented Segment Anything Model for Polyp Segmentation

Huiqian Li, Dingwen Zhang외 4명|arXiv (Cornell University)|2024. 06. 30.

Vehicle License Plate Recognition인용 수 4

한 줄 요약

이 논문은 도메인 갭과 분포 외 일반화 성능이 열 劣한 문제를 해결하기 위해 Segment Anything Model(SAM)의 새로운 변형인 ASPS를 제안한다. SAM의 고정된 ViT와 훈련 가능한 CNN 인코더를 크로스브랜치 특징 증강(Cross-branch Feature Augmentation, CFA)을 통해 통합하고, 예측 신뢰도 캘리브레이션을 향상시키기 위해 불확실성 유도 예측 정규화(Uncertainty-guided Prediction Regularization, UPR)를 적용함으로써, 프롬프트가 필요 없이 다섯 개인 폴립 데이터셋에서 최신 기준 성능을 달성한다.

ABSTRACT

Polyp segmentation plays a pivotal role in colorectal cancer diagnosis. Recently, the emergence of the Segment Anything Model (SAM) has introduced unprecedented potential for polyp segmentation, leveraging its powerful pre-training capability on large-scale datasets. However, due to the domain gap between natural and endoscopy images, SAM encounters two limitations in achieving effective performance in polyp segmentation. Firstly, its Transformer-based structure prioritizes global and low-frequency information, potentially overlooking local details, and introducing bias into the learned features. Secondly, when applied to endoscopy images, its poor out-of-distribution (OOD) performance results in substandard predictions and biased confidence output. To tackle these challenges, we introduce a novel approach named Augmented SAM for Polyp Segmentation (ASPS), equipped with two modules: Cross-branch Feature Augmentation (CFA) and Uncertainty-guided Prediction Regularization (UPR). CFA integrates a trainable CNN encoder branch with a frozen ViT encoder, enabling the integration of domain-specific knowledge while enhancing local features and high-frequency details. Moreover, UPR ingeniously leverages SAM's IoU score to mitigate uncertainty during the training procedure, thereby improving OOD performance and domain generalization. Extensive experimental results demonstrate the effectiveness and utility of the proposed method in improving SAM's performance in polyp segmentation. Our code is available at https://github.com/HuiqianLi/ASPS.

연구 동기 및 목표

• SAM의 사전 훈련에서 사용된 자연 이미지와 내 endoscopy 폴립 이미지 사이의 도메인 갭을 해결하기 위해.
• SAM이 전반적인 낮은 주파수 패턴에 집중함에 따라 자주 간과되는 폴립 이미지의 국소적이고 고주파 특징을 포착할 수 있는 능력을 향상시키기 위해.
• 분포 외(end-of-distribution, OOD) 엔도스코피 이미지에 대한 예측 불확실성을 줄이고 신뢰도 캘리브레이션을 향상시키기 위해.
• 임상 적용에서 프롬프트 의존성을 제거하기 위해 엔드 투 엔드, 프롬프트 없는 훈련을 가능하게 하기 위해.
• 경량적이고 효율적인 적응 프레임워크를 통해 도메인 일반화 능력과 강건성을 향상시키기 위해.

제안 방법

• 가장자리 특징 증강(Cross-branch Feature Augmentation, CFA)을 도입하여, 훈련 가능한 CNN 인코더의 특징과 고정된 비전 트랜스포머(Vision Transformer, ViT) 인코더의 특징을 융합함으로써 다중 척도 및 고주파 특징 학습 능력을 향상시킨다.
• 불확실성 유도 예측 정규화(Uncertainty-guided Prediction Regularization, UPR)를 적용하여, 지도 데이터를 힌트로 사용해 훈련 중에 신뢰도 점수를 캘리브레이션하고 불확실성을 줄인다.
• 목 부분과 트랜스포머 블록의 정규화 레이어를 적응시켜 엔도스코피 도메인 분포에 더 잘 적응하도록 한다.
• 푸리에 분석을 통해 CNN 브랜치가 ViT 기준보다 더 많은 고주파 성분을 포착하고 있음을 검증한다.
• 프롬프트 없이 엔드 투 엔드로 훈련하여 엔도스코피 이미지에 직접 추론이 가능하게 한다.
• UPR에서 SAM의 IoU 예측을 지도 신호로 활용하여 불확실성 감소를 유도하고 분포 외 일반화 성능을 향상시킨다.

실험 결과

연구 질문

• RQ1 순수한 ViT 기반 SAM과 비교해 하이브리드 CNN-ViT 인코더 아키텍처가 폴립 세그멘테이션의 특징 표현 능력을 향상시키는가?
• RQ2 훈련 중에 불확실성 정규화를 도입함으로써 분포 외 일반화 능력과 신뢰도 캘리브레이션 성능이 향상되는가?
• RQ3 기존 SAM 가중치를 미세조정하거나 프롬프트에 의존하지 않고도 도메인 특화된 특징 증강을 달성할 수 있는가?
• RQ4 다중 수준 특징 통합과 위치 임베딩 교체가 세그멘테이션 정확도와 강건성에 어떤 영향을 미치는가?
• RQ5 제안된 방법이 다양한 폴립 데이터셋에서 기존의 SAM 기반 모델보다 얼마나 뛰어나게 성능을 발휘하는가?

주요 결과

• ASPS는 다섯 개인 폴립 데이터셋에서 평균 Dice 스코어 0.914와 평균 IoU 0.843을 기록하며, 기준 SAM 및 MedSAM, SurgicalSAM과 같은 이전 방법들을 초월한다.
• CFA 모듈만으로도 ViT-B 기준 대비 평균 Dice는 31.7% 향상되고 평균 IoU는 41.4% 향상된다.
• CFA와 UPR의 조합이 가장 뛰어난 성능을 보이며, CVC-ClinicDB 데이터셋에서 Dice 스코어 0.950과 IoU 0.905를 기록한다.
• 목 레이어 정규화(neighborhood normalization, NN)와 힌트 기반 정규화를 사용한 UPR가 가장 높은 성능을 기록하여, 불확실성 감소에 효과적임을 입증한다.
• 푸리에 분석 결과, CNN 브랜치가 ViT 인코더보다 유의미하게 더 많은 고주파 성분을 포착하고 있음이 확인되어 국소 세부 정보 학습 능력 향상에 기여함을 뒷받침한다.
• 정성적 결과에서는 ASPS 예측이 특히 폴립의 미세한 경계와 질감 세부 정보를 더 정확히 포착함을 보여준다.