Parameter-efficient Prompt Tuning and Hierarchical Textual Guidance for Few-shot Whole Slide Image Classification

1. 연구 배경 및 문제 정의 (Problem)

Whole Slide Images (WSI) 는 기가픽셀 (giga-pixel) 규모의 초고해상도 디지털 병리 이미지로, 컴퓨터 병리학 (CPath) 분야에서 중요한 역할을 합니다. 그러나 WSI 는 다음과 같은 고유한 도전 과제를 가지고 있습니다.

• 데이터 부족 및 주석 비용: WSI 는 진단적으로 중요한 악성 세포가 매우 작은 영역에 국한되어 있어, 정확한 인스턴스 (패치) 수준의 주석을 얻는 데 막대한 전문가 시간과 비용이 소요됩니다. 이로 인해 슬라이드 수준의 약한 지도 학습 (Weakly Supervised Learning) 이 주로 사용되지만, 데이터가 부족한 Few-shot 환경에서는 모델 학습이 어렵습니다.
• 기존 VLM 기반 방법의 한계: 최근 Vision-Language Models (VLM, 예: CLIP, CONCH) 이 Few-shot WSI 분류에 적용되고 있으나, 다음과 같은 문제점이 있습니다.
  1. 과도한 파라미터 및 추론 오버헤드: 기존 프롬프트 튜닝 (Prompt Tuning) 방법 (예: CoOp) 은 크로스-어텐션 (Cross-attention) 등을 사용하여 학습 가능한 파라미터 수를 급격히 늘리고 추론 비용을 증가시킵니다.
  2. 하드 필터링 (Hard Filtering) 의 정보 손실: VLM 의 텍스트 임베딩과 정렬도가 낮은 인스턴스를 강제로 제거하는 방식은, 텍스트로 표현되지 않았을지라도 진단적으로 중요한 정보를 가진 영역을 실수로 삭제하여 정보 손실을 초래할 수 있습니다.
  3. 계층적 구조 무시: WSI 는 세포 수준의 미세 구조부터 조직 수준의 거시 구조까지 계층적 공간 구조를 가지고 있으나, 기존 방법들은 이를 충분히 반영하지 못합니다.

2. 제안 방법론: HIPSS

저자들은 HIPSS (Hierarchical textual guidance based WSI representation learning and Prompt tuning with Scaling and Shifting features) 라는 새로운 프레임워크를 제안합니다. 이는 두 가지 핵심 기여로 구성됩니다.

A. 텍스트 인코더 내 스케일링 및 쉬프팅 (Scaling and Shifting Features, SSF) 기반 프롬프트 튜닝

• 개념: 기존 VLM 의 텍스트 인코더를 고정하고, 특정 레이어에 학습 가능한 스케일 ($\gamma$) 과 쉬프트 ($\beta$) 파라미터만 추가하여 특징을 변형합니다.
• 방식: $y = \gamma \cdot x + \beta$ 공식을 사용하여 사전 학습된 특징을 목표 분포에 맞춥니다.
• 장점:
  • 파라미터 효율성: 전체 가중치를 학습하거나 크로스-어텐션을 추가하는 대신, 레이어당 2 개의 파라미터만 학습하여 학습 가능한 파라미터 수를 획기적으로 줄입니다.
  • 추론 효율성: 재파라미터화 (Re-parameterization) 기법을 통해 학습된 선형 변환을 모델에 병합할 수 있어, 추론 시 추가적인 계산 오버헤드가 발생하지 않습니다.

B. 소프트 계층적 텍스트 안내 (Soft Hierarchical Textual Guidance) 기반 WSI 표현 학습

• 계층적 구조 활용: WSI 를 먼저 $4096 \times 4096$ 크기의 영역 (Region) 으로 나누고, 이를 다시 $256 \times 256$ 크기의 인스턴스 (Instance) 로 분할합니다.
• 소프트 안내 전략:
  • 하드 필터링 제거: 텍스트 임베딩과의 정렬도가 낮은 인스턴스를 버리는 대신, 어텐션 가중치 (Attention Weights) 를 텍스트 임베딩을 기반으로 부드럽게 조정합니다.
  • 가중치 정제 (Refinement): 인스턴스 (또는 영역) 임베딩과 텍스트 임베딩 간의 코사인 유사도를 계산하여, 유사도가 낮은 경우 가중치를 줄이고 높은 경우 증폭하는 3 단계 정제 메커니즘을 적용합니다. 이는 불확실하지만 유용한 정보를 가진 인스턴스를 보존하면서도 신뢰할 수 있는 매칭을 강조합니다.
• 구조:
  1. Region Encoder: 각 영역 내 인스턴스들을 텍스트 안내 어텐션으로 집계하여 영역 임베딩 생성.
  2. WSI Encoder: 생성된 영역 임베딩들을 다시 텍스트 안내 어텐션으로 집계하여 전체 슬라이드 임베딩 생성.

C. 학습 전략

• LLM 기반 프롬프트 생성: LLM 을 사용하여 WSI 전체와 특정 영역에 대한 시각적 특징을 설명하는 텍스트 프롬프트를 생성합니다.
• 대조 학습 (Contrastive Learning): 생성된 텍스트 임베딩과 이미지 임베딩 간의 대조 손실 (Contrastive Loss) 을 최소화하여 이미지 - 텍스트 공간에서의 정렬을 강화합니다.

3. 주요 기여 (Key Contributions)

1. 파라미터 효율적 프롬프트 튜닝: 텍스트 인코더의 특징에 스케일링 및 쉬프팅을 적용하여 학습 파라미터를 대폭 줄이고 추론 비용을 낮추는 새로운 방법을 제안했습니다.
2. 소프트 계층적 텍스트 안내: WSI 의 고유한 계층적 구조를 반영하고, 하드 필터링 없이 텍스트 안내를 통해 정보 손실을 방지하는 표현 학습 방법을 도입했습니다.
3. 성능 입증: 유방암, 폐암, 난소암을 포함한 3 개의 주요 병리 데이터셋에서 Few-shot 설정 (1~16 샷) 하에 기존 최첨단 (SOTA) 방법들을 일관되게 능가하는 성능을 보였습니다.
4. 약지도 종양 국소화: 분류 작업 외에도 학습된 특징 표현의 질을 입증하기 위해, 별도의 국소화 튜닝 없이도 강력한 약지도 종양 국소화 (Weakly-supervised Tumor Localization) 성능을 달성했습니다.

4. 실험 결과 (Results)

데이터셋: Camelyon16 (유방암), TCGA-Lung (폐암), UBC-OCEAN (난소암).
평가 지표: AUC (Area Under the Curve).

• 분류 성능:

  • Camelyon16: 1 샷 설정에서 SOTA 대비 7.2% 향상, 4 샷에서 10.9% 향상.
  • TCGA-Lung: 4 샷 설정에서 7.8% 향상.
  • UBC-OCEAN: 1 샷 설정에서 13.8% 향상 (가장 큰 개선 폭).
  • 전반적으로 모든 샷 설정에서 기존 방법 (TOP, FOCUS, ViLa-MIL 등) 을 상회했습니다.

• 효율성 및 파라미터:

  • 학습 가능한 파라미터 수를 유방/폐암 데이터셋에서 18.1%, 난소암 데이터셋에서 5.8% 감소시켰습니다.
  • 추론 시간은 슬라이드당 12.6% 단축되었습니다.
  • LoRA 기반 프롬프트 튜닝 대비 파라미터는 59.7% 적게 사용하면서 성능은 16.4% 더 높았습니다.

• 국소화 성능 (Dice Coefficient):

  • Camelyon16 데이터셋에서 기존 Few-shot 방법들 (TOP: 0.203) 보다 훨씬 높은 0.732의 Dice 점수를 기록하여, 모델이 종양 위치를 정확하게 파악하고 있음을 입증했습니다.

• 애블레이션 연구 (Ablation Study):

  • SSF 기반 프롬프트 튜닝과 텍스트 안내 어텐션 정제 전략을 모두 적용했을 때 가장 높은 성능을 보였으며, 각 구성 요소가 성능 향상에 기여함을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 Few-shot WSI 분류의 핵심 난제인 '데이터 부족'과 '계산 효율성'을 동시에 해결하는 효과적인 프레임워크를 제시합니다.

• 임상적 가치: 적은 양의 주석 데이터로도 높은 정확도의 진단과 종양 위치 파악이 가능해져, 희귀 질환이나 데이터가 부족한 임상 환경에서의 AI 적용 가능성을 높였습니다.
• 기술적 혁신: VLM 의 사전 학습 지식을 활용하면서도, WSI 의 계층적 구조를 텍스트 안내를 통해 자연스럽게 통합하고, 하드 필터링 없이 정보를 보존하는 '소프트'한 접근법은 향후 멀티모달 의료 영상 분석에 중요한 방향성을 제시합니다.
• 확장성: 제안된 방법은 계산 비용이 낮고 효율적이므로, 실제 임상 시스템에 배포하기에 적합하며, 종양 분할 (Segmentation) 등 다른 다운스트림 작업으로도 확장 가능성이 큽니다.

결론적으로 HIPSS 는 파라미터 효율성과 표현 학습의 정밀함을 결합하여, 의료 영상 분석 분야에서 Few-shot 학습의 새로운 표준을 제시하는 중요한 연구로 평가됩니다.