[논문 리뷰] ViM-UNet: Vision Mamba for Biomedical Segmentation

Anwai Archit, Constantin Pape|arXiv (Cornell University)|2024. 04. 11.

Image Processing Techniques and Applications인용 수 4

한 줄 요약

이 논문은 비전 링크드 모델(Vision Mamba) 아키텍처를 기반으로 한 새로운 생물의학 영상 분할 모델인 ViM-UNet을 소개한다. 이는 기존의 UNet과 UNETR 네트워크를 대체한다. ViM-UNet은 LIVECell 및 CREMI와 같은 전자현미경 영상 분할 작업에서 경쟁력 있거나 더 뛰어난 성능을 기록하며, UNETR보다 훨씬 더 파rameter 효율적이고 빠른 속도를 보이며, 맥락이 풍부한 생물의학 작업에 대해 링크드 모델의 잠재력을 입증한다.

ABSTRACT

CNNs, most notably the UNet, are the default architecture for biomedical segmentation. Transformer-based approaches, such as UNETR, have been proposed to replace them, benefiting from a global field of view, but suffering from larger runtimes and higher parameter counts. The recent Vision Mamba architecture offers a compelling alternative to transformers, also providing a global field of view, but at higher efficiency. Here, we introduce ViM-UNet, a novel segmentation architecture based on it and compare it to UNet and UNETR for two challenging microscopy instance segmentation tasks. We find that it performs similarly or better than UNet, depending on the task, and outperforms UNETR while being more efficient. Our code is open source and documented at https://github.com/constantinpape/torch-em/blob/main/vimunet.md.

연구 동기 및 목표

• CNN 기반 UNet과 트랜스포머 기반 UNETR의 생물의학 분할에서의 한계, 특히 높은 파라미터 수와 계산 비용을 해결하기 위해.
• 최근 제안된 전역 수용 영역과 효율적인 시퀀스 모델링 기능을 가진 비전 링크드(Vision Mamba, ViM) 아키텍처가 생물의학 영상 분할에서 트랜스포머의 우월한 대안이 될 수 있는지 탐색하기 위해.
• 신경선 분할과 같은 대규모 맥락적 이해가 필요한 어려운 전자현미경 영상 데이터셋에서 ViM-UNet의 성능을 평가하기 위해.
• 다양한 데이터셋에서 UNet, UNETR, nnUNet, U-Mamba와의 정확도, 추론 속도, 메모리 효율성 측면에서 ViM-UNet의 성능을 비교하기 위해.
• 파라미터 수와 추론 시간을 분석하여, 특히 데이터가 적은 환경과 3D 분할 시나리오에서 모델의 일반화 능력과 효율성을 검증하기 위해.

제안 방법

• UNETR의 ViT 백본을 대체하기 위해 비전 링크드(ViM) 인코더를 양방향 상태공간모델(SSM)로 조정하여 전역 맥락 모델링을 가능하게 하며, 계산 비용을 줄인다.
• 비전 트랜스포머와 유사하게 16×16 패치 기반 처리를 사용하지만, 다중 헤드 어텐션 대신 SSM을 사용하여 장거리 의존성을 효율적으로 구현한다.
• 표준 UNet 스타일 디코더를 사용하며, 각 헤드당 두 개의 합성곱층과 역합성곱을 사용한다. 모든 모델 간 공정한 비교를 위해 동일한 아키텍처를 공유한다.
• 모든 모델을 Adam 옵timizer를 사용해 학습하며, 기초 학습률 1e-4로 시작하고, 100k 반복 동안 손실 감소가 정체될 경우 학습률 감소 전략을 적용한다.
• 두 가지 인스턴스 분할 헤드를 구현한다: 하나는 전경 및 경계 확률 지도(워터셰드 후처리 필요), 다른 하나는 중심점 및 거리 예측.
• LIVECell(단면형광 현미경)과 CREMI(볼륨 전자현미경, 2D 슬라이스)에서 평가하며, 주요 평가 지표로 평균 분할 정확도를 사용한다.

Figure 1: Example images with segmentation from ViM-UNet Small .

실험 결과

연구 질문

• RQ1 비전 링크드 아키텍처(ViM-UNet)가 어려운 생물의학적 인스턴스 분할 작업에서 UNet 및 UNETR와 경쟁하거나 더 뛰어난 성능을 낼 수 있는가?
• RQ2 ViM-UNet의 전역 시야가 전자현미경 영상에서 신경선 분할과 같이 대규모 맥락 이해가 필요한 작업에서 UNet보다 성능 우위를 제공하는가?
• RQ3 UNETR의 높은 파라미터 수와 계산 요구량을 감안할 때, ViM-UNet의 추론 속도와 메모리 효율성은 어떻게 비교되는가?
• RQ4 ViM-UNet이 광범위한 사전학습 없이도 UNETR를 능가할 수 있으며, 다양한 형태학적 특성을 가진 다양한 데이터셋에서 잘 일반화되는가?
• RQ5 맥락이 핵심적인 3D 분할 및 세포 추적 작업에서 ViM-UNet이 트랜스포머 기반 모델의 실용적이고 효율적인 대안이 될 수 있는가?

주요 결과

• LIVECell 데이터셋에서 ViM-UNet는 거리 예측 기반 UNet과 유사한 성능(0.05초 추론 시간)을 기록했으며, 더 많은 파라미터를 가진 UNETR(0.54초)를 뛰어넘었다.
• LIVECell에서 경계 분할 성능은 UNet에 이어 두 번째로 높았으며, UNETR는 가장 열악한 성능을 보였다.
• CREMI 데이터셋에서 ViM-UNet는 가장 높은 평균 분할 정확도를 기록했으며, UNet과 UNETR를 모두 앞섰다. 이는 대규모 구조 작업에서 전역 맥락 모델링의 이점이 뚜렷하게 나타남을 시사한다.
• ViM-UNet는 파라미터 수가 Tiny 기준 18M, Small 기준 39M로 UNETR의 113M(Base), 334M(Large), 665M(Huge)보다 훨씬 적었으며, 경쟁 가능한 성능를 유지했다.
• ViM-UNet의 추론 시간은 Tiny와 Small 모두 0.05초로 매우 낮았고, UNETR의 0.15초에서 0.54초에 비해 뚜렷하게 빠르며, 학습 시 최대 48GB VRAM이 필요한 UNETR와 달리 최대 10GB VRAM만으로도 가능했다.
• 결과적으로 ViM-UNet는 추론 메커니즘의 오버헤드 없이 효율적인 전역 모델링을 통해 3D 분할 및 세포 추적과 같이 맥락이 중요한 작업에서 특히 효과적임을 시사한다.

Figure 2: Results for our and external methods; circles highlight the three best methods.