💻 computer science

Hybrid Transformer-Mamba for Weakly Supervised Volumetric Medical Segmentation

이 논문은 평면 수준의 레이블로부터 3D 컨텍스트를 포착하기 위해 효율적인 교차 평면 모델링을 활용하는 하이브리드 Transformer-Mamba 아키텍처인 TranSamba를 소개하며, 이를 통해 약지도 학습 기반의 볼륨 의료 영상 분할에서 최첨단 성능을 달성한다.

원저자: Yiheng Lyu, Lian Xu, Coen Arrow, Mohammed Bennamoun, Farid Boussaid, Girish Dwivedi

게시일 2026-06-19

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yiheng Lyu, Lian Xu, Coen Arrow, Mohammed Bennamoun, Farid Boussaid, Girish Dwivedi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 컴퓨터에게 뇌의 종양이나 심장의 공동(cavity)처럼 3D 의료 스캔 내부의 특정 물체를 찾는 법을 가르치려 한다고 상상해 보세요. 보통 컴퓨터에게 이를 가르치려면, 사람이 모든 스캔 슬라이스마다 그 물체의 완벽한 윤곽선을 일일이 그려주어야 합니다. 이것은 마치 선생님에게 학생의 공책 모든 페이지를 하나하나 채점하라고 요구하는 것과 같습니다. 시간이 엄청나게 오래 걸리고 비용도 매우 많이 듭니다.

문제점: "2D 맹목성"
시간을 절약하기 위해 연구자들은 "약한 지도 학습(weak supervision)"을 사용합니다. 대신에 사람들은 컴퓨터에게 "이 슬라이스 뭉치 전체에 종양이 있다" 또는 "없다"라고만 말해줍니다. 이것은 선생님이 정확히 어디가 틀렸는지 지적하는 대신, 그냥 "이 단원 전체를 잘 해냈구나"라고 말하는 것과 같습니다.

문제는 대부분의 컴퓨터가 의료 스캔을 2D 사진의 묶음처럼 보도록 훈련된다는 점입니다. 컴퓨터는 한 슬라이스를 보고 추측을 한 뒤, 다음 슬라이스를 보고 또 다른 추측을 하며, 슬라이스끼리 서로 소통하는 법을 전혀 모릅니다. 이들은 인체가 3D 물체라는 사실을 놓칩니다. 종양은 단순히 한 페이지 위의 평면적인 원이 아닙니다. 종양은 여러 페이지에 걸쳐 뻗어 있는 3D 덩어리입니다.

해결책: TranSamba (전문가 팀)
저자들은 TranSamba라는 새로운 AI 모델을 구축했습니다. 이것을 아주 다른 두 가지 특성을 가진 전문가 팀이 퍼즐을 풀기 위해 협력하는 것으로 생각해보세요.

"지역 탐정" (Transformer): 이 부분은 스캔의 단일 슬라이스를 관찰하여 "헤이, 여기 목표물처럼 보이는 형태가 있어"라고 말하는 데 탁월합니다. 특정 사진 내부의 세부 사항에 집중하는 데 매우 능숙합니다.
"맥락 연결자" (Mamba): 이것이 새로운 주인공입니다. Mamba를 책의 페이지 사이를 빠르게 달리는 초고속 메신저라고 상상해 보세요. 단순히 한 페이지를 보는 대신, Mamba는 5페이지의 탐정에게 달려가 "이봐, 4페이지에 있는 것이 6페이지에 있는 것과 연결되어 있어"라고 빠르게 속삭여 줍니다. 이는 복잡한 수학 계산에 얽매이지 않으면서도 모델이 슬라이스 사이에서 물체가 어떻게 흐르는지 이해하도록 돕습니다.

그들이 협력하는 방식
과거의 방식은 모든 페이지를 한꺼번에 연결하려고 시도하는 것과 같았는데, 이는 마치 모든 사람이 동시에 소리를 질러서 대화가 되어도 될지 안 될지 모르는 혼란스러운 상황과 같습니다. 이는 계산적으로 매우 비효고적이고 느립니다.

TranSamba는 게임의 판도를 바꿉니다. 이 모델은 "맥락 연결자(Mamba)"를 사용하여 인접한 슬라이스 사이에서 정보를 직선 형태로 효율적으로 전달합니다. 이것은 마치 바톤을 다음 주자에게 매끄럽게 전달하는 계주와 같습니다. 이를 통해 "지역 탐정(Transformer)"은 바로 옆 슬라이스에서 무슨 일이 일어나고 있는지 알게 되어 훨씬 더 나은 추측을 할 수 있게 됩니다.

왜 중요한가

속도와 효율성: Mamba 부분이 매우 효율적이기 때문에, 스캔에 슬라이스가 수백 개가 있더라도 모델이 느려지지 않습니다. 모델은 빠르고 방대한 양의 컴퓨터 메모리를 필요로 하지 않습니다.
더 나은 결과: 저자들은 이 모델을 뇌 종양, 신장 종양, 심장 공동 등 세 가지 다른 유형의 의료 스캔에 대해 테스트했습니다. 모든 경우에서, TranSamba는 이러한 "약한" 레이블을 사용하는 다른 어떤 방법보다 물체를 더 잘 찾아냈습니다. 그것은 마치 표지만 보고 추측하는 대신 마침내 책 전체를 읽는 법을 알아낸 팀과 같았습니다.

한계점
논문은 이 모델이 물체가 어디에 있는지 찾는 데는 뛰어나지만, 아주 작은 물체(예: 거대한 바위와 비교했을 때의 먼지 한 점)를 찾는 데는 여전히 약간 어려움을 겪는다고 언급합니다. 또한, "약한" 레이블로 훈련되었기 때문에, 현재는 연구용 도구이며 의사가 병원에서 직접 사용하기 전에는 몇 가지 추가 단계가 필요할 수 있습니다.

요약하자면
TranSamba는 날카로운 눈을 가진 지역 관찰자와 슬라이스 사이의 점들을 연결하는 빠르고 효율적인 메신저를 결려 결합함으로써, 컴퓨터가 3D 의료 스캔을 이해하도록 가르치는 새로운 AI 아키텍처입니다. 이를 통해 컴퓨터는 단순하고 저렴한 레이블로부터 학습하면서도, 수학 계산을 위해 슈퍼컴퓨터를 필요로 하지 않고 높은 정확도로 의료적 문제를 찾아낼 수 있습니다.

기술 요약: 약지도 학습 기반 볼륨형 의료 영상 분할을 위한 하이브리드 트랜스포머-맘바(Transformer-Mamba)

문제 정의
볼륨형 의료 영상 분할(예: CT, MRI)은 일반적으로 노동 집약적인 복셀 단위 주석(voxel-level annotation)을 필요로 한다. 약지도 분할(Weakly supervised segmentation)은 국소화 정보가 부족한 슬라이스 수준(plane-level) 레이블만을 사용하여 모델을 학습시키는 해결책을 제공한다. 클래스 활성화 매핑(Class Activation Mapping, CAM)은 분류 모델로부터 거친 수준의 국소화 정보를 추출할 수 있지만, 기존 방식들은 의료 데이터의 볼륨 특성으로 인해 상당한 문제에 직면한다. 대부분의 현재 접근 방식은 개별 슬라이스를 독립적으로 처리하는 2D 인코더에 의존하여 3D 컨텍스트를 간과한다. 또한, 표준적인 교차 평면 자기 주의 집중(cross-plane self-attention, SA)을 통해 볼륨 컨텍스트를 포착하려고 시도하는 것은 평면의 수( $N$ )에 따라 이차 함수적으로 증가하는 계산 비용 때문에 매우 비효리적이다.

방법론: TranSamba
저자들은 이러한 한계를 해결하기 위해 비전 트랜스포머(ViT)와 맘바(Mamba, 구조적 상태 공간 모델) 아키텍처를 결합한 하이브리드 아키텍처인 TranSamba를 제안한다. 핵심 설계 철학은 3D 컨텍스트를 효율적으로 활용하기 위해 Mamba의 효율성을 활용하는 동시에, 평면 내 객체 국소화를 위한 트랜스포머의 효능을 유지하는 것이다.

아키텍처 개요: 인코더는 $L$ 개의 하이브리드 레이어로 구성된다. 각 레이어는 교차 평면 맘바(Cross-Plane Mamba, CPM) 블록과 그 뒤를 잇는 트랜스포머 블록을 포함한다.
교차 평면 모델링 (CPM): 3D 컨텍스트를 효율적으로 포착하기 위해, $N$ 개의 인접한 평면에서 추출된 패치 토큰들을 하나의 시퀀스로 쌓고 $M \times N$ 길이(여기서 $M$ 은 평면당 패치 수)로 재구성한다. CPM 블록은 Mamba 모듈을 사용하여 이 토큰들 사이의 상호작용을 모델링한다. 결정적으로, 전역 평면 표현을 보존하기 위해 클래스 토큰을 분리하여 Mamba 블크에서 제외한다. 이를 통해 모델은 선형 시간 내에 평면 간 의존성을 학습할 수 있다.
평면 내 모델링 (Transformer): CPM 블록의 출력은 다시 $N$ 개의 병렬 시퀀스로 재구성된다. 이들은 Class-to-Patch(C2P) 자기 주의 집중을 사용하는 표준 트랜스포머 블록에 의해 처리된다. 교차 평면 단계에서 풍부해진 토큰 표현은 평면 내 주의 집중 맵을 강화하여 객체 국소화를 개선한다.
학습 및 추론: 모델은 평면 수준 레이블을 사용하여 학습된다. 손실 함수는 클래스 토큰의 전역 평균 풀링(GAP)과 패치 토큰의 전역 가중 순위 풀링(GWRP)에서 유도된 교차 엔트로피 항을 결합한다. 추론 시에는 모든 트랜스포머 블록으로부터 C2P 주의 집중 가중치를 집계하여 주의 집중 맵을 생성한다.

주요 기여

하이브리드 아키텍처: 평면 내 자기 주의 집중의 강점을 희생하지 않으면서 효율적인 교차 평면 모델링을 가능하게 하는 CPM 블록과 ViT 백본을 통합한 TranSamba를 도입하였다.
복잡도 효율성: 이 아키텍처는 평면의 수( $N$ )에 대해 선형 시간 복잡도를 달성하며, 표준 교차 평면 자기 주의 집중의 이차 함수적 스케일링과 달리 배치 처리에 대한 상수 공간 복잡도를 유지한다.
최첨단 성능(SOTA): 본 방법론은 교차 평면 모델링이 볼륨형 의료 영상 분할을 위한 일반화 가능한 전략임을 입증하였으며, 기존 연구에서 자주 사용되는 복잡한 다중 스케일 집계나 자기 지도 학습 규제 전략에 의존하지 않고도 다양한 모달리티와 병리 현상에서 우수한 결과를 달성하였다.

실험 결과
TranSamba는 세 가지 서로 다른 데이터셋에 대해 평가되었다:

BraTS: 뇌 종양 분할 (T2 FLAIR).
KiTS: 신장 종양 분할.
LASC: 좌심방 공동(Left atrium cavity) 분할.
성능: TranSamba는 세 데이터셋 모두에서 최첨단(SOTA) 결과를 달eric하였으며, 두 번째로 우수한 방법들보다 평균적으로 Dice 유사 계수(DSC)에서 5.3%, IoU에서 2.4% 높은 격차를 보였다.
절제 연구(Ablation Studies):
- 교차 평면 vs. 평면 내: CPM 블록을 평면 내 Mamba 블록(V2)으로 대체했을 때 성능이 현저히 낮아졌으며, 이는 성능 향상이 단순히 Mamba의 사용 때문이 아니라 구체적으로 교차 평면 모델링에서 기인함을 확인시켜 준다.
- 모델링 범위: 16개 평면에 걸쳐 모델링했을 때 가장 좋은 결과를 얻었으며, 이는 연속적인 구조물에 대해서는 장거리 모델링보다 국소적인 교차 평면 상호작용이 더 유익하다는 것을 시사한다.
- 레이어 설계: CPM 블록을 트랜스포머 블록 앞에 배치하는 것(CrossIn 디자인)이 최적임이 밝혀졌으며, 이는 주의 집중 추출이 풍부한 교차 평면 컨텍스트로부터 이득을 얻도록 보장한다.
객체 크기 분석: 성능은 객체 크기와 양의 상관관계를 보였다. TranSamba는 크거나 거대한 객체에는 탁월한 성능을 보였으나, 작은 객체(의료 분할의 알려진 문제)에 대해서는 어려움을 겪었으며, 특히 객체가 작은 경우가 많았던 KiTS 데이터셋에서 그러했다.

의의 및 주장
본 논문은 TranSamba가 3D 컨텍스트의 필요성과 약지도 학습의 계산적 제약 사이의 간극을 효과적으로 메운다고 주장한다. 저자들은 교차 평면 모델링과 평면 내 모델링을 분리함으로써, 효율적인 선형 시간 교차 평면 모델링이 볼륨형 의료 영상 분할에서 객체 국소화를 개선하기 위한 데 결정적이고 일반화 가능한 구성 요소임을 입증하였다. 이 연구는 향후 임상 적용을 위해 후처리 또는 다단계 접근 방식이 필요할 수 있음을 시사하지만, 핵심적인 아키텍처 혁신은 다양한 의료 영상 시나리오의 약지도 작업에 대한 견고한 토대를 제공한다.

유사한 논문