💻 computer science

Primus: Enforcing Attention Usage for 3D Medical Image Segmentation

본 논문은 하이브리드 모델의 한계를 어텐션 활용을 극대화함으로써 극복하고, 9 개의 공개 데이터셋에서 기존 CNN 기반 방법론을 능가하거나 이에 필적하는 최첨단 성능을 달성하는 3D 의료 영상 분할을 위한 최초의 경쟁력 있는 트랜스포머 중심 아키텍처인 Primus 와 PrimusV2 를 소개합니다.

원저자: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

게시일 2026-05-01

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Tassilo Wald, Saikat Roy, Fabian Isensee, Constantin Ulrich, Sebastian Ziegler, Dasha Trofimova, Raphael Stock, Michael Baumgartner, Gregor Köhler, Klaus Maier-Hein

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

"Primus: 3D 의료 영상 분할을 위한 어텐션 사용 강제"라는 논문에 대한 설명을 창의적인 비유를 곁들여 쉬운 언어로 번역한 것입니다.

큰 문제: "가짜" 트랜스포머

3D 의료 스캔 (예: 신장의 CT 스캔) 을 보고 종양 주위에 완벽한 윤곽선을 그릴 수 있는 로봇을 만들려고 상상해 보세요.

오랫동안 이 일을 수행하는 데 가장 뛰어난 로봇은 CNN(합성곱 신경망)이었습니다. CNN 을 매우 유능한 지역 탐정으로 생각하세요. 이 탐정은 작은 이웃 구역의 픽셀들을 살펴보고 단서를 모은 다음 다음 구역으로 이동합니다. 국소적인 세부 사항을 파악하는 데는 뛰어나지만, 때로는 장기 전체의 '큰 그림'을 놓치기도 합니다.

그런 다음 트랜스포머가 등장했습니다. 이들은 전체 이미지를 한 번에 바라보고 신장의 윗부분이 아랫부분과 어떻게 관련되는지 이해할 수 있는 '초과관찰자'와 같습니다. 이들은 언어 처리 (예: 챗봇) 와 자연 사진 분야에서 유명해졌습니다. 사람들은 "트랜스포머가 책 전체를 읽거나 풍경 전체를 볼 수 있다면, 의료 스캔에도 완벽할 것이다!"라고 생각했습니다.

하지만 함정이 있었습니다: 연구자들이 3D 의료 스캔에 트랜스포머를 적용해 보았을 때, 그 성능은 그다지 좋지 않았습니다. 그들은 종종 구식 CNN 탐정보다 느리고 정확도도 낮았습니다.

수사: 누가 일을 하고 있는가?

이 논문의 저자들은 왜 이러한 트랜스포머들이 실패하는지 조사하기로 결정했습니다. 그들은 CNN 과 트랜스포머를 모두 사용하려는 9 개의 인기 있는 '하이브리드' 모델들을 살펴보았습니다.

그들은 충격적인 비밀을 발견했습니다: 트랜스포머는 대부분 잠들어 있었습니다.

비유: 유명한 고가의 건축가 (트랜스포머) 를 고용해 집을 설계하게 하고, 동시에 100 명의 일반 벽돌공 (CNN) 을 고용했다고 상상해 보세요. 집이 완성되었을 때, 저자들은 건축가가 실제로 설계도를 그리지 않았음을 깨달았습니다. 벽돌공들이 집 전체를 직접 지었고, 건축가는 그저 고개만 끄덕이며 서 있었을 뿐입니다.
증거: 연구자들이 이러한 모델에서 '건축가' (트랜스포머 블록) 를 제거했을 때, 모델의 성능은 거의 동일하게 유지되었습니다. 어떤 경우에는 트랜스포머를 제거함으로써 모델이 쓸모없는 구성 요소에 에너지를 낭비하지 않게 되어 오히려 더 빨라지고 약간 더 좋아지기도 했습니다.

이 논문은 이를"UNet 지수"라고 부릅니다. 대부분의 기존 모델들은 높은 지수를 보였는데, 이는 실제로는 무거운 쓸모없는 트랜스포머 배낭을 멘 CNN 일 뿐이라는 뜻입니다.

해결책: Primus 와 PrimusV2

저자들은 질문했습니다: "트랜스포머가 반드시 일을 하도록 만든다면 어떨까?" 그들은 Primus(라틴어로 '첫 번째')와 PrimusV2라는 두 가지 새로운 아키텍처를 개발했습니다.

그들이 어떻게 트랜스포머를 깨워 일을 하게 했는지 살펴보겠습니다:

1. 세부 사항을 누르지 않기 (토크나이저)

일반적인 트랜스포머는 3D 이미지를 데이터 토큰으로 변환하기 위해 거대한 덩어리 (거대한 두꺼운 케이크 조각을 자르는 것처럼) 로 잘라내곤 합니다. 이는 작은 종양이나 얇은 혈관과 같은 작지만 중요한 세부 사항을 버리게 만듭니다.

해결책: Primus 는 '고해상도 토크나이저'를 사용합니다. 거대한 조각 대신 8x8x8 보크셀 (voxel) 크기의 더 작고 정교한 조각을 사용합니다.
비유: 모든 거리가 흐릿한 선으로만 표시된 도시 지도를 보는 대신, Primus 는 개별 주택까지 볼 수 있는 지도를 봅니다. 이는 트랜스포머가 작업할 수 있는 더 많은 세부 정보를 제공합니다.

2. '반복적' 접근법 (PrimusV2)

조각을 더 작게 잘랐음에도 불구하고, 트랜스포머는 때로 복잡한 장기들의 3D 형태를 바로 이해하는 데 어려움을 겪었습니다.

해결책: PrimusV2 는 '반복적 패치 임베딩 (Iterative Patch Embedding)'을 사용합니다. 한 번에 전체 덩어리를 이해하려 시도하는 대신, 양파를 껍질 벗기거나 스케치를 다듬는 것처럼 이미지를 단계별로 처리합니다. 트랜스포머가 데이터를 보기 전에 몇 가지 작고 스마트한 합성곱 단계를 사용하여 데이터를 준비합니다.
비유: 복잡한 퍼즐을 풀려고 한다고 상상해 보세요. Primus 는 모든 조각을 그냥 테이블 위에 쏟아붓지 않습니다. 먼저 색과 가장자리 조각으로 분류 (반복적 단계) 한 다음, '초과관찰자' (트랜스포머) 가 최종 그림을 훨씬 쉽게 볼 수 있도록 합니다.

3. GPS 제공 (3D RoPE)

트랜스포머는 본질적으로 공간에 대해 '맹목'입니다. '왼쪽'과 '오른쪽'이 다르다는 것을 알려주지 않으면 알지 못합니다.

해결책: 저자들은 특별한 3D '회전 위치 임베딩 (Rotary Position Embedding, RoPE)'을 추가했습니다.
비유: 이는 트랜스포머에 깊이, 너비, 높이를 동시에 이해하는 GPS 시스템을 부여하는 것과 같습니다. 이는 종양이 장기 나머지 부분에 비해 3D 공간에서 정확히 어디에 위치하는지 알고 있습니다.

결과: 트랜스포머가 마침내 승리하다

이러한 변경 사항 이후, 결과는 인상적이었습니다:

Primus는 표준 '골드 스탠다드'인 CNN(nnU-Net) 과 경쟁할 수 있는 최초의 트랜스포머 기반 모델이 되었습니다.
PrimusV2는 단순히 경쟁하는 것을 넘어, 대부분의 테스트에서 표준 CNN 을 능가했으며 오늘날 이용 가능한 가장 복잡하고 최상급의 CNN 들과도 대등한 성능을 보였습니다.

핵심 요약:
이 논문은 트랜스포머가 3D 의료 영상 처리를 위한 최고의 도구가 될 수 있음을 증명합니다. 단, CNN 의 조수처럼 대우하지 않아야 합니다. 트랜스포머가 주인공이 되도록 시스템을 설계하고, 고해상도 데이터를 공급하며 3D 공간을 이해할 수 있는 적절한 도구를 제공해야 합니다.

한 문장으로 요약한 내용

저자들은 Primus라는 새로운 AI 모델을 개발하여 마침내 '초과관찰자'인 트랜스포머가 의료 영상 처리에서 중추적인 역할을 하도록 강제함으로써, 올바르게 설계될 경우 트랜스포머가 전통적인 '지역 탐정'인 CNN 들을 능가할 수 있음을 증명했습니다.

"Primus: 3D 의료 이미지 분할을 위한 어텐션 사용 강제" 논문의 상세 기술 요약입니다.

1. 문제 제기

자연어 처리 및 자연 이미지 컴퓨터 비전에서 트랜스포머의 성공에도 불구하고, 3D 의료 이미지 분할에 미치는 영향은 제한적이었습니다. 현재 벤치마크는 여전히 컨볼루션 신경망 (CNN), 특히 nnU-Net 프레임워크가 지배하고 있습니다. 저자들은 의료 영상 분야의 기존 트랜스포머 기반 접근법에서 두 가지 치명적인 단점을 확인했습니다.

컨볼루션에 대한 과도한 의존: 대부분의 "트랜스포머" 아키텍처는 실제로 하이브리드 모델로, 대부분의 매개변수와 계산 용량이 컨볼루션 블록 (인코더/디코더) 에 존재하여 트랜스포머 구성 요소는 부수적인 역할에 그칩니다.
비효율적인 어텐션: 많은 하이브리드 모델에서 트랜스포머 블록을 제거해도 성능 저하가 미미하여, 모델이 CNN 만으로 표현을 학습하고 있음을 시사합니다. 또한, 기존 순수 또는 준순수 트랜스포머 모델은 작은 희소 라벨링 의료 데이터셋에서 처음부터 훈련하는 트랜스포머의 어려움으로 인해 강력한 CNN 베이스라인보다 종종 낮은 성능을 보입니다.

2. 방법론

저자들은 기존 모델의 한계를 입증하기 위한 엄격한 해체 분석과 새로운 진정한 트랜스포머 중심 아키텍처 설계라는 두 가지 접근법을 제안합니다.

A. 기존 아키텍처 분석

저자들은 UNet Index 라는 지표를 사용하여 TransUNet, SwinUNETR, UNETR 등 인기 있는 9 가지 트랜스포머 기반 모델을 분석했습니다.

UNet Index 정의: 표준 U-Net 대비 트랜스포머 블록 외부의 컨볼루션 매개변수 양을 측정하는 비율입니다.
결과: 대부분의 모델은 높은 UNet Index(무거운 CNN 의존성) 를 보였습니다. 트랜스포머 블록을 항등 매핑 (identity mappings) 으로 대체했을 때, 9 개 모델 중 6 개가 Dice Similarity Coefficient(DSC) 에서 3% 미만의 감소만 보여, 트랜스포머가 성능의 주된 동인이 아님을 입증했습니다.

B. Primus 및 PrimusV2 아키텍처

이러한 문제를 해결하기 위해 저자들은 3D 의료 영상을 위한 최초의 경쟁력 있는 트랜스포머 중심 아키텍처인 Primus와 PrimusV2를 설계했습니다. 이러한 모델은 컨볼루션 매개변수를 최소화하여 네트워크가 어텐션 메커니즘에 의존하도록 강제합니다.

주요 설계 혁신:

고해상도 토큰화:
- 국소 세부 정보를 잃는 큰 $16\times16\times16$ 패치를 사용하는 표준 ViT와 달리, Primus는 미세한 해부학적 구조를 보존하기 위해 $8\times8\times8$ 패치 크기를 사용합니다.
- PrimusV2는 반복적 패치 임베딩을 도입합니다. 단일 투영 대신 잔여 컨볼루션 블록의 체인을 사용하여 토큰화 전에 반복적으로 다운샘플링하고 의미론적 특징을 추출합니다. 이는 무거운 CNN 백본에 의존하지 않고 더 풍부한 토큰을 생성합니다.
3D 회전 위치 임베딩 (RoPE):
- 표준 절대 위치 임베딩은 3D 의료 데이터에서 흔히 발생하는 공간적 이동에 어려움을 겪습니다. 저자들은 **축 방향 회전 위치 임베딩 (RoPE)**을 3D 에 적용하여 상대적 거리와 방향을 어텐션 메커니즘에 직접 인코딩함으로써 체적적 공간 상관관계를 더 잘 포착하도록 했습니다.
최신 트랜스포머 구성 요소:
- MLP 블록: 자연 이미지 트랜스포머 (예: EVA-02) 에서 효과가 입증된 SwiGLU(Swish 활성화 게이트형 선형 유닛) 와 LayerNorm 채택.
- 안정화 메커니즘: 대규모 모델 (Primus-L) 의 훈련을 안정화하고 발산을 방지하기 위해 LayerScale과 Post-Attention Normalization 사용.
- 경량 디코더: 토큰을 독립적으로 업샘플링하는 전치 컨볼루션으로 구성된 디코더로, 컨볼루션 영향을 최소화하고 UNet Index 를 낮게 유지합니다.

3. 주요 기여

실증적 해체: UNet Index 와 제거 실험 (ablation studies) 을 통해 현재 "트랜스포머" 의료 모델이 사실은 약간의 어텐션 추가가 포함된 CNN 임을 정량적 증거로 제시했습니다.
최초의 경쟁력 있는 트랜스포머 중심 모델: 대부분의 매개변수와 FLOPs 를 트랜스포머 블록 내에 유지하면서 최첨단 성능을 달성하는 Primus와 PrimusV2를 도입했습니다.
아키텍처 혁신: 고해상도 토큰화 ( $8\times8\times8$ ), 반복적 패치 임베딩, 3D RoPE 를 결합하면 트랜스포머가 의료 영상의 데이터 부족 및 국소 특징 요구 사항을 극복할 수 있음을 입증했습니다.
자기지도 학습의 효율성: CNN 중심 아키텍처에 비해 토큰을 마스킹할 때 VRAM 을 현저히 적게 (40% 감소) 사용하여 트랜스포머 중심 설계가 Masked Image Modeling(MIM) 사전 훈련에 더 효율적임을 강조했습니다.

4. 실험 결과

이 모델들은 CT, MRI 등 다양한 모달리티와 장기를 아우르는 9 개의 공개 데이터셋(AMOS22, KiTS23, TotalSegmentator-BTCV, SBM 등 포함) 에서 평가되었습니다.

베이스라인 대비 성능:
- Primus: 현재 트랜스포머 기반 방법을 능가하고 기본 nnU-Net과 직접 경쟁하여 9 개 데이터셋 중 4 개에서 이를 초과했습니다.
- PrimusV2: UNETR++, SwinUNETRv2 를 포함한 모든 다른 트랜스포머 베이스라인을 크게 능가했으며, ResEnc-L 및 MedNeXt와 같은 최첨단 CNN 과 동등한 성능에 도달했습니다.
- 격차 축소: PrimusV2 는 최상위 CNN 과의 성능 격차를 0.05~~0.29 DSC 포인트로 줄였으며, 이는 이전 트랜스포머 시도들이 8~~11 DSC 포인트 뒤처졌던 것에 비해 엄청난 개선입니다.
제거 실험 통찰:
- PrimusV2 에서 트랜스포머를 제거하면 약 5.0 DSC 감소가 발생하여, 모델이 어텐션 메커니즘에 크게 의존하고 있음을 확인했습니다 (과거 하이브리드 모델에서는 감소폭이 0 에 가까웠던 것과 대조적).
- PrimusV2 의 반복적 토크나이저는 표준 투영 대비 약 2 DSC 향상을 제공했으며, 특히 작은 병변 (예: SBM 데이터셋의 뇌 전이) 에 대한 성능을 개선했습니다.
확장성: 모델 크기를 S 에서 L 로 확장하는 것이 매우 작은 데이터셋에서는 과적합으로 인해 체감 수익이 감소했지만, PrimusV2-L 은 더 큰 데이터셋에서 전체적으로 가장 높은 점수를 달성했습니다.

5. 의의

이 연구는 3D 의료 이미지 분할 분야에서 패러다임 전환을 의미합니다.

순수 트랜스포머의 실현 가능성: 트랜스포머가 CNN 에 대한 추가 기능이 아니라 의료 영상에서 성능의 주된 동인이 될 수 있음을 입증했습니다.
미래 대비: 강력한 트랜스포머 중심 베이스라인을 확립함으로써, 트랜스포머는 뛰어나지만 CNN 은 어려움을 겪는 자기지도 학습 (SSL) 및 다중 모달 통합(예: 이미지와 텍스트 보고서 결합) 과 같은 분야에서 의료 AI 의 채택을 가능하게 합니다.
임상 적용 가능성: 추론 시간 분석에 따르면, 가장 큰 PrimusV2 모델조차도 뇌졸중 감지 등 임상적 의사결정 시간 내에 적합하여 실제 배포가 가능합니다.

결론적으로, Primus와 PrimusV2는 고해상도 토큰, 3D RoPE, 반복적 임베딩과 같은 올바른 아키텍처 적응을 통해 트랜스포머가 3D 의료 분할에서 최상위 CNN 과 맞먹거나 능가할 수 있음을 보여주었으며, 차세대 의료 AI 파이프라인의 문을 열었습니다.