Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Brain-SAM 기술 요약

문제 정의
자기공명영상(MRI)은 신경 영상 분야의 핵심이지만, 뇌 구조 및 병변의 정확한 분할(segmentation)은 여전히 어려운 과제로 남아 있다. Segment Anything Model(SAM) 및 그 후속 모델인 SAM2와 같은 파운데이션 모델들이 일반적인 자연 이미지에서 놀라운 일반화 성능을 보여주었음에도 불구하고, 이를 뇌 MRI에 적용하는 데에는 상당한 장애물이 존재한다. 기존의 의료용 적응 모델들(예: MedSAM, SAM-Med3D)은 심장, 간, 폐와 같은 구조를 포함하는 혼합 장기 데이터셋으로 학습되는 경향이 있다. 이러한 모델들은 복잡한 조직 구성, 불분명한 회백질-백질 경계, 매우 다양한 병변 형태와 같은 뇌 특유의 해부학적 특성에 적응하는 데 어려움을 겪는다. 또한, 대부분의 SAM 기반 적응 모델들은 3D 볼륨 스캔을 2D 슬라이스로 변환하여 처리함으로써 슬라이스 간의 공간적 맥락을 간과한다. 아울러, 수동 인터랙티브 프롬프트에 대한 의존도는 임상 현장에서 뇌전증 병변과 같이 미세한 병리 현상에 대해 정확한 프롬프트를 제공하기 어렵게 만들어 확장성을 제한한다. 마지막으로, 명확한 경계를 가진 2D 자연 이미지에 설계된 기존 SAM 마스크 디코더는 조직의 이질성과 흐릿한 경계가 특징인 뇌 MRI 데이터에서 최적의 성능을 내지 못하는 경우가 많다.

방법론
저자들은 3D 뇌 MRI 병변 분할에 특화된 모델인 Brain-SAM을 제안한다. 이 아키텍처는 세 가지 주요 수정을 통해 SAM2 프레임워크를 확장한다:

1. 이미지 인코더의 3D 적응: 기존의 2D Hiera 인코더(계층적 비전 트랜스포머)를 3D 볼륨 데이터를 직접 처리할 수 있도록 확장한다. 여기에는 네 가지 핵심 모듈의 확장이 포함된다:

   • 패치 임베딩(Patch Embedding): 2D 컨볼루션을 3D 컨볼루션으로 교체하여 볼륨 입력을 복셀 패치로 분할한다.
   • 위치 임베딩(Position Embedding): 깊이(depth), 높이(height), 너비(width) 전반의 공간적 관계를 포착하기 위해 2D에서 3D 표현으로 확장한다.
   • 토큰 시퀀스(Token Sequence): 3D 볼륨을 세 축 모두를 따라 평탄화(flattening)하여 볼륨 어텐션에 적합한 더 긴 토큰 시퀀스를 생성한다.
   • Q-풀링(Q-Pooling): MaxPool2d를 MaxPool3d로 교체하여 깊이 및 공간 차원 전반의 특징 다운샘플링을 가능하게 한다.

2. UNETR 기반 디코더: 미세한 국소 텍스처와 경계 세부 사항을 처리하는 데 있어 기존 SAM 디코더의 한계를 해결하기 위해, 저자들은 UNETR에서 영감을 받은 디코더를 도입한다. 이 구성 요소는 다중 스케일 특징 융합을 통해 트랜스포머의 전역 모델링 능력과 CNN의 국소적 미세 표현 능력을 통합한다. 해상도 변환을 요구하는 표준 UNETR 디코더와 달리, 이 설계는 Hiera 인코더의 고유한 계층적 특징을 활용하여 직접적인 융합 및 점진적 업샘플링을 수행한다. 아키텍처는 완전 자동 분할을 위한 헤드와 프롬프트 기반 분할을 위한 중간 특징 생성을 위한 헤드라는 듀얼 헤드 출력을 지원한다.

3. 학습 및 미세 조정 전략:

   • 사전 학습(Pretraining): 모델은 SAM2 가중치로 초기화되며, 견고한 의료 표현을 구축하기 위해 BraTS2021 데이터셋(1,000 케이스, 4개 모달리티)만을 대상으로 완전 자동 분할을 위한 사전 학습 단계를 거친다.
   • 어댑터 기반 미세 조정(Adapter-Based Fine-Tuning): 여러 뇌 병변 데이터셋(종양, 뇌졸중, 뇌전증)에 적응하면서 치명적 망각(catastrophic forgetting)을 방정하기 위해, 각 Hiera 블록 앞에 경량 어댑터 모듈(3D 컨볼루션을 사용하는 bottleneck 구조)을 삽ante한다. 미세 조정 중에는 백본을 동결하고 공유된 어댑터 파라미터만 업데이트한다.
   • 프롬프트 메커니즘: 포인트 프롬프트가 주요 상호작용 방식으로 사용되며, 이를 3D로 확장한다. 마스크 생성 전략은 내적(dot-product) 연산에서 업데이트된 키를 사용하는 대신, UNETR 기반 디코더의 중간 특징 표현과 처리된 쿼리를 융합하도록 설계되었다.
   • 정교화(Refinement): 사후 처리 정교화 메커니즘은 사용자 클릭을 사용하여 3D 가우시안 커널을 통해 마스크 확률을 국소적으로 조정한다.

주요 기여

1. SAM2의 3D 확장: 본 연구는 SAM2 프레임워크를 2D 자연 이미지에서 3D 의료 영상 분석으로 성공적으로 확장하여, 슬라이스별 근사 방식 없이 직접적인 볼륨 처리를 가능하게 했다.
2. 하이브리드 디코더 아키텍처: UNETR에서 영감을 받은 디코더의 도입을 통해, 모델이 단일 아키텍처 내에서 완전 자동 및 인터랙티브 분할을 모두 지원할 수 있게 하여 국소적인 미세 세부 사항과 전역적인 문맥 의존성을 효과적으로 포착한다.
3. 포괄적인 데이터셋 큐레이션: 저자들은 뇌 종양, 뇌졸중, 뇌전증을 아우르는 대규모 데이터셋을 큐레이션하여 모델이 광범위한 임상 응용 분야를 지원할 수 있도록 했다.

결과
6개의 MRI 모달리티를 포함하는 7개의 공개 데이터셋(BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI)에 대해 광범위한 실험이 수행되었다. Brain-SAM은 작업별 특화 모델(nnU-Net, UNETR, Swin UNETR) 및 파운데이션 모델(SAM-Med3D, VISTA3D)과 비교되었다.

• 전반적인 성능: Brain-SAM은 모든 베이스라인 모델보다 우수한 Dice 점수와 낮은 95% 하우스도르프 거리(HD95)를 달성했다. 구체적으로, nnU-Net과 비교했을 때 뇌전증 병변에서 22%, 뇌 전이암에서 9%, 수막종에서 6%의 Dice 개선을 달성했다.
• 작은 병변 분할: 모델은 작은 병변(예: 뇌전증)을 분할하는 데 있어 상당한 강점을 보였으며, 강력한 베이스라인 모델들보다 15%~18% 높은 Dice 점수를 기록했다. 또한 다른 방법들과 비교하여 작은 타겟(<5,000 복셀)에 대해 재현율(recall)과 정밀도(precision) 사이의 더 나은 균형을 유지했다.
• 절제 연구(Ablation Studies):
  • 3D-ViT 백본을 3D-Hiera로 교체했을 때, 작은 병변에 대해 Dice 점수가 최대 14% 향상되었고 HD95가 유의미하게 감소했다.
  • UNETR 기반 디코더를 통합했을 때, 디코더가 없는 버전보다 평균 Dice 점수가 8% 증가했다.
  • 어댑터 기반 미세 조정 전략은 전체 미세 조정과 대등한 성능을 달성하면서도 매우 효율적인 파라미터 사용량(전체 미세 조정의 72.02M 대비 39.4M 개의 튜닝 가능한 파라미터)을 입증했다.

의의 및 주장
본 논문은 Brain-SAM이 다운스트림 뇌 MRI 분석 작업에 대해 견고한 사전 학습 모델을 제공한다고 주장한다. 불분명한 경계나 작은 병변 크기와 같은 뇌 조직의 특정 해부학적 및 구조적 문제를 해결함으로써, 기존의 범용 파운데이션 모델보다 더 신뢰할 수 있고 정확한 분할을 제공한다. 저자들은 이러한 능력이 특히 뇌전증 및 뇌종양과 같은 질환의 조기 진단 및 치료 계획 수립에 있어 더 정보에 기반한 임상적 의사결정을 지원할 수 있다고 상정한다. 자동 및 프롬프트 가이드 모드 모두에서 기능할 수 있는 모델의 능력은 높은 정확도를 유지하면서 수동 주석 작업의 부담을 줄이는 것을 목표로 한다. 저자들은 현재 연구가 단일 모달리티 데이터에 국한되어 있음을 언급하며, 향후 연구에서 병변 특성화를 더욱 강화하기 위해 다중 모달리티 통합을 탐구할 수 있다고 제언한다.