Does Your ViT Still Need U-Net for Segmentation?

기술 요약: 당신의 ViT는 여전히 U-Net을 필요로 하는가?

문제 정의
지난 10년 동안 의료 영상 분할(segmentation)은 U-Net 스타일의 인코더-디코더 구조에 의해 지배되어 왔습니다. 이러한 설계는 스킵 연결(skip connections)과 대칭적 디코딩을 통해 공간적 세부 사항의 손실을 성공적으로 해결하지만, 제한된 수용 영역(receptive field)을 가진 합성곱 연산에 의존하기 때문에 전역적 문맥(global context) 모델링이 어렵다는 제약이 있습니다. Vision Transformer(ViT)는 장거리 의존성(long-range dependencies)을 포착하기 위한 셀프 어텐션 메커니-즘을 제공하지만, 초기 ViT 기반 분할 방법들(예: TransUNet, SwinUNet)은 사전 학습된 ViT 표현력이 정밀한 밀집 예측(dense prediction)을 수행하기에는 불충분하다고 간주되었기 때문에 여전히 U-Net 디코더를 유지했습니다.

최근 대규모 사전 학습(예: DINOv2, SigLIP)의 발전은 ViT의 표현 능력을 근본적으로 재정의하였으며, 방대한 미라벨 데이터셋으로부터 강력한 시각적 특징을 학습할 수 있게 했습니다. 이러한 진화는 다음과 같은 중요한 질문을 던집니다: 현대의 강력하게 사전 학습된 ViT 백본이 가용할 때, 의료 영상 분할에서 U-Net 패러다임이 여전히 필요한가? 만약 그렇지 않다면, 디코더 없이 효과적으로 분할 마스크를 예측하기 위해 인코더 전용(encoder-only) 프레임워크를 어떻게 설계해야 하는가?

방법론: EoSeg
저자들은 현대적인 사전 학습된 ViT 백본을 기반으로 한 쿼리 기반 설계를 채택하여 U-Net 디코더를 제거한 프레임워크인 **EoSeg (Encoder-only Segmentation)**를 제안합니다. 이 방법론은 핵심적인 아키텍처 선택들에 대한 체계적인 탐색을 통해 개발되었습니다:

1. 백본 선택: 본 연구는 인코더 전용 설정에서 다양한 사전 학습된 ViT 백본(Vi-T-L/14, SigLIP, MTP, DINOv2, DINOv3 포함)을 평가합니다. DINOv2가 최적의 백본으로 선택되었으며, 디코더 없이도 가장 높은 베이스라인 성능(Synapse 데이터셋에서 81.86% mDice)을 달성했습니다.
2. 마스크 생성:
   • 디컨볼루션(Deconvolution) vs. 쿼리 기반(Query-based): 초기에는 저해상도 ViT 특징을 업샘플링하기 위해 디컨볼루션 모듈을 테스트했습니다. 그러나 저자들은 쿼리 기반 마스크 예측 방식이 더 우수하다는 것을 발견했습니다. 픽셀 단위의 디코딩 대신, 학습 가능한 쿼리 임베딩이 전역적 시맨틱 프로브(global semantic probes) 역할을 수행합니다.
   • 메커니즘: 쿼리 벡터는 마스크 임베딩 공간으로 투영되며, 업샘플링된 전역 특징 맵과 채널 단위 도트 곱(channel-wise dot product)을 통해 상호작용하여 직접 이진 마스크를 생성합니다.
3. 다층 쿼리 모델링 (Multi-level Query Modeling): U-Net의 스킵 연결에서 영감을 받아, 이 프레임워크는 단순히 마지막 블록만이 아니라 마지막 3개의 트랜스포머 블록으로부터의 표현을 활용합니다. 학습 가능한 쿼리가 이 블록들 이전에 주입되고 이미지 토큰과의 셀프 어텐션을 통해 정제됨으로써, 모델이 여러 단계로부터 시맨틱 정보를 집계할 수 있도록 합니다.
4. 학습 가능한 블록 퓨전 (Learnable Block Fusion): 세 개의 중간 블록으로부터의 예측을 결합하기 위해 학습 가능한 퓨전 모듈이 도입됩니다. 퓨전 게이트는 쿼리 표현을 사용하여 각 블록의 출력에 대한 적응형 가중치를 예측하며, 이를 통해 모델이 최종 예측에 대한 각 레벨의 기여도를 동적으로 결정할 수 있게 합니다.
5. 쿼리 구성: 프레임워크는 각 쿼리가 마스크와 클래스 레이블을 동시에 예측하는 쿼리 기반 분류 패러다임을 채택합니다. 실험 결과, 카테고리당 두 개의 쿼리를 할당하는 것이 과도한 파라미터화 없이 클래스 내 변이(intra-class variations)를 포착하는 데 최적의 성능을 보였습니다.
6. 지도 학습 전략 (Supervision Strategy): 의료 분할의 일반적인 관행과는 달리, 이 특정 인코더 전용 설정에서는 **심층 지도 학습(deep supervision, 중간 출력에 손실을 적용하는 방식)**이 성능 향상에 도움이 되지 않는다는 것이 밝혀졌으며, 최종 설계에서 제외되었습니다.

주요 기여

• 아키텍처 재평가: 본 논문은 현대적인 대규모 사전 학습된 ViT 백본을 활용할 경우, 고품질 의료 영상 분할을 위해 U-Net 스타일의 디코더가 더 이상 필수 조건이 아님을 보여주는 광범위한 평가를 제공합니다.
• EoSeg 프레임워크: 저자들은 인코더 특징으로부터 직접 마스크 예측을 가능하게 하기 위해 다층 쿼리 모델링과 학습 가능한 블록 퓨전을 통합한 쿼리 기반 인코더 전용 프레임워크인 EoSeg를 제시합니다.
• 포괄적 검증: 제안된 설계는 CT (Synapse), MRI (ACDC), 조직병리 (GlaS, MoNuSeg), 내시경 (Kvasir-Seg), 피부 병변 (ISIC-2016, ISIC-2017) 등 5가지 서로 다른 의료 영상 양상을 아우르는 7개의 벤치마크 데이터셋을 통해 검증되었습니다.

실험 결과
EoSeg는 테스트된 모든 양상에서 기존의 CNN 기반, 트랜스포머 기반 및 하이브리드 인코더-디코더 방식들을 일관되게 능가했습니다:

• Synapse (CT): 85.50% mean Dice (mDice)를 달성하여, 이전 최고 성능인 EViT-UNet(80.87%)보다 4.63 포인트 앞섰습니다. 또한 대동맥(Aorta, 90.80%) 및 비장(Spleen, 94.13%)과 같은 특정 장기에서도 최고 점수를 기록했습니다.
• ACDC (MRI): 91.73% mDice를 달성하여, 이전 최고 성능을 0.24 포인트 초과했습니다.
• GlaS (조직병리): 93.27% mDice 및 87.79% mIoU를 달성하여, 이전 최고 성능보다 각각 0.83 및 1.29 포인트 앞섰습니다.
• MoNuSeg (조직병리): 80.51% mDice 및 67.45% mIoU를 달성했습니다.
• Kvasir-Seg & ISIC: 폴립 및 병변 분할에서 우수한 성능을 입증하였으며, 매끄러운 경계와 일관된 구조를 가진 높은 mIoU 및 mDice 점수를 기록했습니다.

정성적 분석 결과, EoSeg는 경쟁 방법들과 비교했을 때 더 매끄러운 경계를 생성하고 아티팩트가 적으며, 디코더 기반의 특징 복원 없이도 미세한 구조적 세부 사항을 효과적으로 보존함을 확인했습니다.

의의 및 주장
본 논문은 초기 ViT 백본의 한계로 인해 의료 영상 분할에서 U-Net 패러다임이 지배적이었던 시대는 끝났다고 주장합니다. 대신, 대규모 사전 학습을 통해 학습된 강력한 표현력은 인코더 전용 설계가 정확한 분할 마스크를 직접 예측할 수 있게 합니다.

저자들은 EoSeg의 성공이 단순히 디코더를 제거했기 때문이 아니라, 적절한 쿼리 기반 예측 패러다임을 다층 모델링 및 적응형 퓨전과 결과적으로 결합했기 때문임을 강조합니다. 이러한 발견은 대규모 사전 학습된 비전 모델을 구축하는 미래의 의료 영상 분할 프레임워크가 복잡한 디코더 아키텍처에서 벗어나, 현대적 ViT의 본질적인 표현 능력에 의존할 수 있음을 시사합니다.