기술 요약: 도메인 적응형 의미 분할을 위한 비전 트랜스포머 조건부 UNet

문제 제기

비전 트랜스포머 (ViT) 는 자연 이미지 도메인에서 최첨단 (SOTA) 분할 모델의 핵심 구조로 자리 잡았으나, 생체 의학 이미징에 적용하는 것은 여전히 난제입니다. 생체 의학 데이터셋은 종종 주석이 제한되어 있으며, 불규칙하고 희소하며 얇거나 대비가 낮은 형태를 가진 특수 구조를 특징으로 합니다. 저자들은 특정 성능 격차를 식별했습니다: 프롬프트 가능 ViT 모델에서 일반적으로 사용되는 경량 픽셀 디코더 (예: 단순 선형 헤더 또는 업샘플링 마스크) 는 고정밀 생체 의학 마스크에 필요한 국소 귀납적 편향 (local inductive bias) 이 종종 부족합니다. 또한, 생체 의학 작업을 위해 대규모 ViT 기반 모델을 미세 조정하는 것은 계산 비용이 많이 들고 데이터 집약적이며, 기존 하이브리드 접근 방식은 엄격한 검증 하에서 잘 조정된 UNet(예: nnU-Net) 보다 일관되게 우수한 성능을 발휘하지 못합니다.

방법론: ViTC-UNet

본 논문은 ViT 의 전역적 시각적 사전 지식과 CNN(합성곱 신경망) 의 국소적 귀납적 편향 사이의 격차를 해소하기 위해 설계된 ViTC-UNet 프레임워크를 제안합니다. 이 아키텍처는 교차 도메인 설정에서도 엔드 투 엔드 ViT 미세 조정을 피하고, 프리트레인된 ViT 표현을 고정 (frozen) 한 상태에서 UNet 을 조건부로 적용함으로써 이를 달성합니다.

핵심 구성 요소

1. 고정된 백본 인코더: 모델은 대규모 자연 이미지로 프리트레인된 고정된 ViT 백본을 활용합니다 ( DINOv2 및 SAM 2로 평가됨). 이는 Transformer 미세 조정의 데이터 요구 사항을 우회하면서도 고품질의 범용 특징을 유지합니다.
2. 조건부 디코더: 이 모듈은 자연 이미지 임베딩과 생체 의학 공간 표현 사이의 격차를 연결합니다.
   • 다중 레이어 퍼셉트론 (MLP) 을 followed by SAM 에서 수정된 양방향 어텐션 블록 (two-way attention blocks) 의 일련의 구조를 사용합니다.
   • 구조 토큰 (Structure Tokens): 고정된 출력 채널 대신, 모델은 공유되는 공간 위치 임베딩과 짝을 이루는 학습 가능한 "구조 토큰"(대상 클래스당 하나) 을 사용합니다.
   • 메커니즘: 이미지 임베딩과 구조 토큰은 서로에게 어텐션합니다. 구조 토큰은 이미지별 컨텍스트를 집계하는 반면, 시각적 표현은 요청된 대상에 의해 조절됩니다.
   • 위상 일관성: 저자들은 ViT 백본이 위상동형사상 (homeomorphism, 연속적이고 전단사적인 매핑) 으로 작용하므로 디코더가 이러한 위상적 특성을 유지해야 한다고 주장합니다. 양방향 어텐션 블록은 위상동형 변환으로 간주되어 견고한 하류 디코딩을 용이하게 하는 잘 정립된 표현 공간을 보장합니다.
3. 다단계 및 멀티스케일 잠재 조건부:
   • 단일 병목 표현에 의존하는 대신, 디코더는 연속적인 어텐션 블록을 통해 일련의 구조 조건부 잠재 상태를 생성합니다.
   • 이러한 상태는 다단계 조건부(블록 간 이미지 - 토큰 상호작용 정제) 와 멀티스케일 조건부(UNet 디코딩 경로의 해당 수준과 융합) 를 제공합니다.
   • 더 깊은 단계는 대상별 전역 컨텍스트를 수신하는 반면, 더 얕은 단계는 정밀한 경계와 희소 대상의 복원을 지원합니다.
4. 구조 조건부 UNet 픽셀 디코더:
   • 저자들은 픽셀 디코더로 nnU-Net을 활용합니다.
   • $n$ 번째 어텐션 블록의 출력은 $n$ 번째 가장 얕은 합성곱 블록의 스킵 연결과 연결됩니다.
   • 단일 채널 출력: 핵심적으로, 픽셀 디코더는 대상 구조의 수와 관계없이 단일 출력 채널을 사용합니다. 의미적 정체성은 출력 차원성과 분리되며, 모델은 입력 구조 토큰에 조건부로 이진 마스크를 생성합니다. 다중 클래스 분할은 각 대상 토큰에 대해 추론 과정을 반복함으로써 달성됩니다.

주요 기여

• 새로운 아키텍처: 학습 가능한 토큰과 양방향 어텐션 디코더를 통해 고정된 ViT 표현에 UNet 을 조건부로 적용하여 전역적 사전 지식과 국소적 귀납적 편향을 결합한 ViTC-UNet 도입.
• 효율성과 적응성: 엔드 투 엔드 ViT 미세 조정 없이 SOTA 성능을 달성하여 데이터가 부족한 생체 의학 도메인에 적합함.
• 레이블 공간 유연성: 의미적 정체성을 출력 계층에서 분리함으로써, 네트워크 아키텍처를 수정하거나 전체 모델을 재학습하지 않고도 학습 가능한 토큰을 통해 새로운 대상 구조를 점진적으로 추가할 수 있음.
• 이론적 통찰: 위상동형사상에 기반한 디코더 설계에 대한 이론적 근거를 제공하며, ViT 임베딩 공간의 위상적 속성을 보존하는 것이 분할 충실도에 필수적임을 주장함.

결과

모델은 BiomedParse 가 큐레이션한 CT 및 MRI 데이터셋 (Abdomen1K, Aorta, AirwayTree, AMOS, BraTS 등 포함) 의 컴파일레이션에서 평가되었습니다.

• 성능: DINOv2 기반 ViTC-UNet 은 테스트된 데이터셋 전반에 걸쳐 평균 전경 mIoU 0.90을 달성하여 nnU-Net 베이스라인 (0.79) 과 BiomedParse 집계 결과 (0.88) 를 능가했습니다. 15 개 벤치마크 중 14 개에서 nnU-Net 과 동등하거나 더 나은 성능을 보였습니다.
• 하이브리드와의 비교: 이전 하이브리드 접근 방식인 ViT 임베딩을 단순히 UNet 병목에 연결하는 것만으로는 부족함이 입증되었습니다. 희소하고 미세 구조화된 대상을 위한 해부학적 형태를 정확하게 구획하기 위해서는 제안된 다단계 및 멀티스케일 조건부가 필수적입니다.
• 점진적 학습: 2 단계 훈련 실험에서 모델은 이전에 학습된 구조의 성능을 저하시키지 않고 새로운 클래스를 성공적으로 통합하여 레이블 공간 확장 능력을 확인했습니다.
• 백본 민감도: 성능은 백본의 잠재 공간 기하학에 크게 의존했습니다. 자기지도 학습인 DINOv2 가 지도 학습인 SAM 2 보다 우수한 성능을 보였는데, 이는 디코더에서 모델링된 위상동형 변환에 DINOv2 의 직교 선형 부분 공간이 더 유리함을 시사합니다.

중요성 및 주장

본 논문은 ViTC-UNet 이 대규모 시각적 사전 지식을 고복잡도 생체 의학 분할에 적응시키기 위한 계산 효율적이고 샘플 효율적인 경로를 제공한다고 주장합니다. UNet 을 프롬프트 가능하고 구조 조건부 픽셀 디코더로 취급함으로써, 범용 시각적 표현이 강력한 생체 의학 분할 파이프라인을 효과적으로 보완할 수 있음을 보여줍니다.

저자들은 이 접근 방식이 대규모 비전 모델과 특수 임상 응용 분야 사이의 격차, 특히 낮은 신호 대 잡음비와 레이블 부족이 있는 시나리오를 연결한다고 주장합니다. 그들은 대규모 트랜스포머의 미세 조정이라는 막대한 계산 비용 없이도 고충실도 디코딩을 가능하게 하면서도, 아키텍처 변경이 아닌 토큰 기반 조건부를 통해 새로운 해부학적 구조로 확장할 수 있는 능력을 유지한다고 강조합니다.

한계

저자들은 몇 가지 한계를 인정합니다:

• 2D 처리: ViT 인코더가 2D 자연 이미지로 프리트레인되었기 때문에 파이프라인은 데이터를 슬라이스 단위로 처리하여 3D 볼륨 컨텍스트를 완전히 활용하지 못합니다.
• 이산 토큰: 이산적이고 유한한 구조 토큰의 사용은 오픈 보카불러리 일반화 가능성을 제한합니다.
• 잠재 기하학: 모델의 성능은 인코더 잠재 공간의 직교성에 의존하며, 다른 기하학을 가진 공역 (codomain) 에 적응하려면 추가 조사가 필요합니다.
• 프롬프트 유형: 현재 작업은 의미적 구조 토큰에 초점을 맞추고 있으며, 상호작용적 시각적 프롬프트 (예: 점 또는 바운딩 박스) 는 아직 통합되지 않았습니다.