K-Prism: A Knowledge-Guided and Prompt Integrated Universal Medical Image Segmentation Model

기술 요약: K-PRISM

문제 정의

의료 영상 분할(segmentation)은 임상 워크플로(예: 종양 경계 설정, 장기 정량화)에서 매우 중요하지만, 기존의 딥러닝 방식은 파편화 문제를 겪고 있습니다. 기존 모델들은 일반적으로 다음과 같은 단일 지식 패러다임에 국한되어 있습니다:

1. 시맨틱 프라이어(Semantic Priors): 특정 작업이나 장기에 대해 대규모 레이블 데이터셋에 의존함.
2. 인-컨텍스트 지식(In-Context Knowledge): 일반화를 위해 소수의 퓨샷(few-shot) 참조 예시를 활용함.
3. 대화형 피드백(Interactive Feedback): 반복적인 정교화를 위해 사용자 입력(클릭, 스크리블)에 의존함.

이러한 파편화는 의료 기관이 수십 개의 작업별 모델을 유지하도록 강요하여 배포 복잡성을 높이고 성능의 불일치를 초래합니다. 더욱이, 실제 임상 현장에서는 이러한 다양한 지식 소스들을 동적으로 통합하여 사용합니다(예: 방사선 전문의가 해부학적 지식을 활용하고, 유사한 사례를 참조하며, 상호작용을 통해 경계를 정교화하는 과정). 현재의 AI 모델들은 이러한 유연성이 부족합니다. 기존의 어떤 프레임워크도 이 세 가지 유형의 지식을 단일 아키텍처 내에서 성공적으로 통합하지 못했습니다.

방법론

저자들은 시맨틱 프라이어, 인-컨텍스트 예시, 그리고 대화형 피드백을 단일 모델에 통합하도록 설계된 통합 분할 프레임워크인 K-Prism을 제안합니다. 핵심 혁신은 **이중 프롬프트 표현(dual-prompt representation)**과 MoE(Mixture-of-Experts) 디코더의 결합에 있습니다.

1. 이중 프롬프트 표현

K-Prism은 이질적인 지식 소스를 두 가지 상호 보완적인 프롬프트 유형으로 인코딩합니다:

• 1-D 희소 프롬프트(1-D Sparse Prompts): 무엇을 분할할 것인지 정의합니다. 이는 작업 수준 또는 인스턴스별 쿼리(예: 시맨틱 분할을 위한 클래스 임베딩, 인-컨텍스트를 위한 객체 쿼리, 또는 인터랙티브를 위한 클릭 좌표)를 인코딩합니다.
• 2-D 조밀 프롬프트(2-D Dense Prompts): 어디에 주의를 기울일지 나타냅니다. 이는 공간적 특징 맵을 변조하여 위치 단서(localization cues)를 주입하고 구조적 세부 사항을 정교화합니다(예: 참조 이미지-마스크 쌍에서 유도된 어피니티 행렬 또는 사용자 클릭 맵).

본 프레임워크는 세 가지 운영 모드를 지원합니다:

• 모드-1 (Semantic): 학습된 프라이어에 기반한 분할을 예측하기 위해 1-D 희소 프롬프트(학습 가능한 클래스 임베딩)만을 사용합니다.
• 모드-2 (In-Context): 1-D 희소 프롬프트(참조 마스크에서 유도된 객체 쿼리)와 2-D 조밀 프롬프트(참조 특징을 쿼리 이미지로 투영하는 어피니티 행렬)를 모두 사용하여 퓨샷 일반화를 가능하게 합니다.
• 모드-3 (Interactive): 이미지 특징과 융합된 2-D 조밀 프롬프트(양성/음성 클릭 및 이전 마스크 인코딩)를 사용합니다. 또한 사용자 클릭을 1-D 희소 쿼리로 포함함으로써 모드-1 또는 모드-2의 예측을 정교화할 수 있습니다.

2. Mixture-of-Experts (MoE) 디코더

서로 다른 지식 유형의 구별된 처리 요구 사항을 다루기 위해, K-Prism은 MoE 디코더를 채택합니다.

• 동적 라우팅(Dynamic Routing): 게이팅 네트워크가 쿼리별 전문가 가중치를 예측하여, 1-D 희소 프롬프트와 2-D 융합 특징 맵을 특화된 전문가들에게 동적으로 라우팅합니다.
• 양방향 상호작용(Bidirectional Interaction): 디코더는 1-D 프롬프트와 2-D 특징 맵이 상호작용하는 양방향 교차 주의(cross-attention) 레이어를 활용하여, 모델이 객체 의미론과 공간적 문맥을 공동으로 추론할 수 있게 합니다.
• 아키텍처: MoE는 교차 주의(CA)와 피드 포워드 네트워크(FFN) 레이어 모두에 적용되어, 패러다임 전환 시 아키텍처 수정 없이도 작업 인식 특화(task-aware specialization)를 가능하게 합니다.

주요 기여

1. 통합 프레임워크: K-Prism은 시맨틱, 인-컨텍스트, 인터랙티브 지식 유형을 단일 아키텍처 내에 통합한 최초의 의료 영상 분할 모델입니다.
2. 이중 프롬프트 설계: 1-D 희소 프롬프트와 2-D 조밀 프롬프트의 참신한 조합을 통해 다양한 지식 소스를 유연하게 인코딩할 수 있습니다.
3. MoE 디코더: MoE 디코더의 도입은 동적 라우팅과 다양한 작업 및 모달리티 간의 공동 학습을 가능하게 하여, 상충하는 최적화 목표로 인한 성능 저하를 방지합니다.
4. 포괄적 평가: 다양한 모달리티(CT, MRI, X-ray, 병리, 초음파)와 분할 패러다임을 아우르는 18개의 공개 데이터셋에서 모델의 성능을 검증했습니다.

실험 결과

실험 결과, K-Prism은 세 가지 패러다임 모두에서 SOTA(State-of-the-Art) 성능을 달성했습니다:

• 시맨틱 분할(Semantic Segmentation): 12개의 인-디스트리뷰션(in-distribution) 데이터셋에서 K-Prism은 평균 Dice 점수 **86.21%**를 기록하며, 범용 모델인 Hermes(85.02%)와 태스크 특화 모델인 nnU-Net의 크로스 데이터셋 일반화 성능(83.45%)을 능가했습니다.
• 인-컨텍스트 분할(In-Context Segmentation): K-KPrism은 인-디스트리뷰션 데이터셋에서 평균 Dice 84.82%, 외부 데이터셋에서 **82.49%**를 달성하여 기존의 인-컨텍스트 방식(예: Iris, UniverSeg)을 크게 앞질렀습니다. 또한 미학습 클래스에 대한 강력한 적응력을 보여주었으며, 차순위 방법인 Iris의 26.07% 대비 31.91%의 평균 Dice를 기록했습니다.
• 인터랙티브 분할(Interactive Segmentation): K-Prism은 인-디스트리뷰션 데이터셋에서 목표 정확도에 도달하는 데 가장 적은 클릭 수를 요구했습니다(NoC90: 1.95, NoC95: 3.51). 5번의 클릭 후 **95.50%**의 Dice 점수를 달성하여 SegNext(93.80%) 및 MedSAM(91.21%)과 같은 베이스라인을 넘어섰습니다. 결정적으로, 미학습 클래스 데이터셋에서도 높은 효율성과 정확도를 유지했습니다.

의의 및 주장

본 논문은 K-Prism이 파편화된 AI 모델과 유연하고 다각적인 임상 실무 사이의 결정적인 간극을 해결한다고 주장합니다. 세 가지 지식 패러다임을 통합함으로써 K-Prism은 다음과 같은 효과를 제공합니다:

• 배포 복잡성 감소: 의료 기관이 서로 다른 작업이나 지식 유형을 위해 별도의 모델을 유지할 필요성을 제거합니다.
• 임상 워크플로 모사: 시맨틱, 퓨샷, 인터랙티브 모드 간의 원활한 전환을 가능하게 하여, 임상의가 기본 모델을 변경하지 않고도 가용한 지식(예: 일반적인 해부학적 프라이어에서 특정 참조 사례 또는 인터랙티브 정교화로 전환)을 활용할 수 있게 합니다.
• 기반 기술로서의 역할: 단순한 분할 도구를 넘어, 임상의가 프라이어나 인-컨텍스트 학습을 통해 초기 마스크를 생성하고 이를 최소한의 상호작용으로 정교화할 수 있는 의료 파운데이션 모델 및 대규모 어노테이션 파이프라인을 위한 실질적인 인프라로서 자리매김합니다.

저자들은 인-컨텍스트 성능의 참조 품질 의존성, MoE 디코더의 계산 비용, 심각한 도메인 변화(미학습 클래스)에 대한 과제 등 한계점을 인정합니다. 그러나 통합된 설계를 통해 인터랙티브 정교화가 다른 패러다임의 실패를 보완할 수 있으므로, 실제 현장 배포를 위한 견고한 솔책을 제공한다고 주장합니다.