[논문 리뷰] ProMamba: Prompt-Mamba for polyp segmentation

Jianhao Xie, Ruofan Liao외 4명|arXiv (Cornell University)|2024. 03. 20.

Natural Language Processing Techniques인용 수 8

한 줄 요약

ProMamba는 Vision-Mamba와 프롬프트 기술을 이용한 용종 분할로, 강한 교차 데이터셋 일반화와 여섯 개 데이터셋에서 평균 약 5%의 향상을 달성합니다.

ABSTRACT

Detecting polyps through colonoscopy is an important task in medical image segmentation, which provides significant assistance and reference value for clinical surgery. However, accurate segmentation of polyps is a challenging task due to two main reasons. Firstly, polyps exhibit various shapes and colors. Secondly, the boundaries between polyps and their normal surroundings are often unclear. Additionally, significant differences between different datasets lead to limited generalization capabilities of existing methods. To address these issues, we propose a segmentation model based on Prompt-Mamba, which incorporates the latest Vision-Mamba and prompt technologies. Compared to previous models trained on the same dataset, our model not only maintains high segmentation accuracy on the validation part of the same dataset but also demonstrates superior accuracy on unseen datasets, exhibiting excellent generalization capabilities. Notably, we are the first to apply the Vision-Mamba architecture to polyp segmentation and the first to utilize prompt technology in a polyp segmentation model. Our model efficiently accomplishes segmentation tasks, surpassing previous state-of-the-art methods by an average of 5% across six datasets. Furthermore, we have developed multiple versions of our model with scaled parameter counts, achieving better performance than previous models even with fewer parameters. Our code and trained weights will be released soon.

연구 동기 및 목표

• 대장내시경 기반 임상 의사결정을 지원하기 위해 정확한 용종 분할의 필요성을 제고한다.
• 데이터셋 간 용종 모양, 색상, 경계 불명확성의 변동성을 다룬다.
• 훈련 데이터 외의 보이지 않는 데이터셋에 대한 일반화를 향상시킨다.

제안 방법

• Vision-Mamba 아키텍처를 기반으로 한 용종 분할용 Prompt-Mamba를 도입한다.
• 경계 구분을 개선하기 위해 프롬프트 기술을 도입한다.
• 이전 방법과 비교해 보지 않은 데이터셋에서 더 높은 정확도를 보이도록, 확장 가능한 매개변수 수를 사용한다.
• 모델이 여섯 개 데이터셋에서 평균 5%의 상승으로 기존 최첨단을 상회한다는 것을 보고한다.

실험 결과

연구 질문

• RQ1 프롬프트를 갖춘 Vision-Mamba가 보이지 않는 데이터셋에서 용종 분할 정확도를 향상시킬 수 있는가?
• RQ2 Prompt-Mamba가 다양한 용종 데이터셋에서 이전 방법보다 일반화에 더 우수한가?
• RQ3 ProMamba의 서로 다른 매개변수 수가 성능과 효율성에 어떤 영향을 미치는가?
• RQ4 이것이 Vision-Mamba와 프롬프트를 용종 분할에 적용한 최초의 사례이며, 실제 이득은 무엇인가?

주요 결과

• 모델이 보지 않은 데이터셋에서 더 높은 분할 정확도를 달성하여 강한 일반화를 나타낸다.
• ProMamba는 여섯 개 데이터셋에서 평균 5%의 상승으로 이전 최첨단 방법을 능가한다.
• 매개변수 수를 늘린 여러 모델 변형이 더 많은 매개변수를 사용하지 않아도 더 나은 성능을 달성한다.
• 용종 분할에 대한 Vision-Mamba의 최초 적용 및 용종 분할 모델에서 프롬프트의 최초 사용.
• 코드와 학습 가중치의 공개가 예정되어 있다.