기술 요약: ENSAM – 대화형 3D 의료 영상 분할을 위한 효율적인 파운데이션 모델

문제 정의

3D 의료 영상의 정확한 분할은 진단 및 치료 계획과 같은 임상 작업에 매우 중요합니다. 딥러닝 모델이 이러한 작업을 자동화해 왔지만, 자연 이미지를 위해 설계된 최근의 파운데이션 모델인 SAM(Segment Anything Model) 및 SAM2는 상당한 적응 과정 없이는 의료 모달리티에 일반화하는 데 어려움을 겪습니다. 기존의 의료 특화 모델들은 트레이드오프(trade-off) 문제에 직면하곤 합니다. 일부는 반복적인 정교화(iterative refinement) 지원이 부족하고(예: MedSAM), 다른 모델들은 막대한 계산 자원과 사전 학습된 가중치를 필요로 하며, 많은 모델이 다양한 모달리티(CT, MRI, PET, 초음파, 현미경 이미지)에 걸쳐 일반화하는 데 실패합니다. 또한, 표준화된 숨겨진 테스트 세트에서 이러한 모델들을 비교하기 위한 통합 벤치마크가 이전에 부족했습니다.

본 연구에서 다루는 구체적인 과제는 CV-PR 2025 Foundation Models for Interactive 3D Biomedical Image Segmentation Challenge입니다. 목표는 시뮬레이션된 대화형 환경 내에서 사용자 프롬프트(바운딩 박스 또는 클릭)가 주어졌을 때 모든 해부학적 또는 병리적 구조를 분할할 수 있는 보편적이고 프롬프트 가능한 모델을 개발하는 것입니다. "코어셋 트랙(coreset track)"의 핵심 제약 조건은 **전체 훈련 데이터의 10%**로 제한된다는 점과, 외부의 사전 학습된 가중치를 활용하지 않고 처음부터(from scratch) 훈련해야 하며, 엄격한 계산 예산(단일 32 GB GPU) 하에서 작동해야 한다는 점입니다.

방법론: ENSAM 아키텍처

저자들은 U-Net 스타일의 아키텍처를 기반으로 구축된 경량화되고 프롬프트 가능한 모델인 ENSAM(Equivariant, Normalized, Segment Anything Model)을 제안합니다. 이는 이미지 인코더, 프롬프트 인코더, 마스크 디코더의 세 가지 주요 구성 요소로 이루어져 있습니다.

1. 이미지 인코더

특정 분할 벤치마크에서 CNN보다 성능이 떨어지는 경우가 많은 비전 트랜스포머(Vision Transformer) 대신, ENSAM은 SegResNet 기반 인코더를 활용합니다. 이 아키텍처는 다운샘플링 레이어(StridedConv3D)와 교차 배치된 잔차 블록(residual blocks)의 캐스케이드로 구성됩니다. 이는 단일 채널 입력을 다채널 특징 맵으로 변환하며 볼륨 일관성을 유지합니다.

2. 프롬프트 인코더 및 상호작용 메커니즘

프롬프트 인코더는 3D 바운딩 박스, 전경 클릭, 배경 클릭을 처리합니다. 주요 혁신 사항은 다음과 같습니다:

• Lie 회전 위치 인코딩 (LieRE): 등변성(equivariance)을 보존하고 효율성을 높이기 위해, 모델은 절대 위치 인코딩 대신 상대 위치 인코딩을 사용합니다. 이는 키(key) 및 쿼리(query) 벡터에 위치 의존적 회전 행렬을 적용함으로써 달성됩니다. 회전 행렬 $R(p)$는 Lie 대수 $so(n)$로부터 유도되며, 이를 통해 어텐션 점수가 상대적 위치$p_j - p_i$에만 의존하도록 보장합니다.
• 정규화된 어텐션 (Normalized Attention): 최근의 정규화된 트랜스포머 연구에서 영감을 받아, ENSAM은 각 최적화 단계 이후 모든 가중치 행렬과 활성화 함수에 $\ell_2$ 정규화를 적용하여 표준 LayerNorm 및 가중치 감쇠(weight decay)를 대체합니다. 이는 출력을 단위 하이퍼스피어(unit hypersphere) 내로 제한하며, 다음과 같은 수정된 잔차 업데이트 전략을 사용합니다:
  $$x \leftarrow \text{Norm}(\text{Norm}(x) + \lambda(\text{Norm}(\text{nBlock}(x)) - \text{Norm}(x)))$$
  여기서 $\lambda$는 학습 가능한 고유 학습률(eigen learning rate)입니다. 이는 수렴을 가속화하고 수치적 안정성을 향상시키는 것을 목표로 합니다.
• 상호작용 흐름: 사용자 프롬프트와 이전 분할 로짓(strided convolution을 통해 다운샘플링됨)은 네 단계의 과정을 통해 이미지 임베딩과 상호작용합니다: (1) 프롬프트에 대한 정규화된 셀프 어텐션, (2) 프롬프트-이미지 어텐션, (3) MLP 처리, (4) 이미지-프롬프트 어텐션. 이 사이클은 두 번 반복됩니다.

3. 마스크 디코더

디코더는 인코더 구조를 반영하지만, 업샘플링 레이어당 하나의 잔차 블록을 사용합니다. 스킵 연결(skip connections)을 결را(concatenate)하고, 이를 ResBlock3D를 통해 처리한 후, 삼선형 업샘플링(trilinear upsampling)을 적용하여 분할 로짓을 출력합니다.

4. 훈련 전략

• 최적화 도구 (Optimizer): 저자들은 표준 경사 하강법 대신, SVD(Singular Value Decomposition)의 효율적인 근사를 사용하여 2D 가중치 행렬에 작동하는 Muon optimizer를 채택했습니다. 1D 파라미터의 경우 AdamW를 사용합니다.
• 데이터 처리: 훈련은 시뮬레이션된 상호작용을 포함합니다. 초기 바운딩 박스는 정답(ground truth)으로부터 무작위 오프셋을 주어 생성되며, 이후 가장 큰 오류 영역(과소 분할의 경우 전경, 과다 분할의 경우 배경)에 다섯 번의 반복 클릭이 배치됩니다.
• 손실 함수: Generalized Dice Loss와 Cross-Entropy(가중치 1:2)의 복합 손실 함수가 사용됩니다.
• 코어셋 선택: 코어셋 트랙의 경우, 훈련 세트는 모달리티 간의 균등한 표현을 보장하기 위해 선정된 약 4,471개의 샘플(전체의 10%)로 제한되었으며, 정규화 또는 주석 문제가 있는 데이터셋은 제외되었습니다.

주요 기여

1. 효율적인 파운데이션 모델: ENSAM은 단 하나의 32 GB GPU에서 단 6시간 만에, 5,000개 미만의 볼륨을 사용하여 처음부터 훈련되었음에도 불구하고 범용 3D 의료 분할에서 경쟁력 있는 성능을 달acia합니다.
2. 아키텍처 혁신: Lie 회전 위치 인코딩과 정규화된 어텐션을 SegResNet 백본에 통합함으로써, 절대 위치 인코딩 및 표준 어텐션에 비해 훈련 효율성을 크게 높이고 분할 품질을 개선했습니다.
3. 최적화 도구 적용: 본 논문은 3D 의료 분할 작업에서 Muon optimizer의 효능을 입증하며, 이것이 AdamW보다 훈련 속도 및 최종 손실 측면에서 우수함을 보여줍니다.
4. 벤치마킹: 본 연구는 숨겨진 다중 모달리티 테스트 세트에 대한 엄격한 평가를 제공하여, 제한된 데이터 및 계산 조건 하에서 최첨단 베이스라인(VISTA3D, SAM-Med3D, SegVol)과 공정한 비교를 수행합니다.

결과

정량적 성능

숨겨진 테스트 세트(코어셋 트랙)에서 ENSAM은 다음을 달성했습니다:

• DSC AUC: 2.404
• NSD AUC: 2.266
• 최종 DSC: 0.627
• 최종 NSD: 0.597

비교:

• ENSAM은 네 가지 지표 모두에서 VISTA3D와 SAM-Med3D를 능가했습니다.
• SegVol의 DSC AUC 및 NSD AUC와는 대등했으나, 최종 DSC(0.627 vs 0.622)에서는 SegVol을 앞질렀습니다. 다만, 최종 NSD에서는 약간 뒤처졌습니다.
• 전체 챌린지 순위에서 ENSAM은 10개 팀 중 5위를 차지했으며, 사전 학습된 가중치를 사용하지 않은 모델 중 가장 우수한 성능을 보였습니다.

절제 연구 (Ablation Studies)

실험을 통해 절대 위치 인코딩에서 상대 위치 인코딩(LieRE)으로 전환하는 것이 훈련 효율성을 높인다는 점을 확인했습니다. 또한, LieRE를 Muon optimizer와 결려했을 때 기준 설정(Absolute PE + Adam)보다 가장 빠른 수렴과 최상의 최종 손실을 얻었습니다.

정성적 분석

• 성공 사례: ENSAM은 큰 구조물(예: CT에서의 간, 비장, 신장; MRI에서의 척추체)을 정확하게 묘사하였고, 2D+t 초음파 시퀀스를 효과적으로 처리했습니다.
• 실패 사례: 현미경 이미지의 미세한 디테일(가는 혈관)에서 어려움을 겪었으며, 전처리 과정 중 강도 클리핑(intensity clipping)으로 인해 PET 이미지에서 과다 분할을 보이는 경우가 있었습니다. 또한 움직임 아티팩트로 인해 초음파에서 거친 가장자리(jagged edges)를 보이기도 했습니다.

의의 및 주장

저자들은 ENSAM이 의료 영상을 위한 고성능 파운데이션 모델을 막대한 계산 자원이나 외부 사전 학습된 가중치에 의존하지 않고도 효율적으로 처음부터 훈련할 수 있음을 입증한다고 주장합니다. 논문은 상대 위치 인코딩과 Muon optimizer의 결합이 이러한 효율성의 핵심 동력임을 강조합니다.

본 연구는 현재의 아키텍처가 반드시 최적은 아니며, 향후 더 깊은 네트워크 탐색, 물리적 좌표를 통한 이방성 간격(anisotropic spacing) 처리, 그리고 더 다양한 프롬프트 유형(예: scribble, lasso)의 통합이 도움이 될 수 있음을 겸허히 언급하며 마무리됩니다.

결론적으로, ENSAM은 의료용 파운데이션 모델에 대한 접근성을 민주화하는 중요한 단계이며, 엄격한 제약 조건 하에서도 효과적인 대화형 분할이 가능하다는 것을 증명합니다.