Align Your Query: Representation Alignment for Multimodality Medical Object Detection

1. 문제 정의 (Problem Definition)

• 배경: 의료 객체 탐지 (Medical Object Detection) 는 CXR(흉부 X 선), CT, MRI, 병리 등 다양한 영상 모달리티 (Modality) 에서 해부학적 구조나 이상 소견을 식별하고 위치를 파악하는 중요한 작업입니다.
• 핵심 문제: 기존 탐지 모델은 단일 모달리티 데이터로 학습될 때 좋은 성능을 보이지만, 서로 다른 통계적 특성과 시각적 특징을 가진 혼합된 다중 모달리티 (Mixed Multimodality) 데이터셋으로 학습할 경우 성능이 급격히 저하됩니다.
• 원인: 각 모달리티 (예: X 선 vs MRI) 는 고유의 분포를 가지며, 이로 인해 모델의 표현 공간 (Representation Space) 이 단절되고 이질적이게 되어, 하나의 모델이 모든 모달리티에서 견고하게 일반화하기 어렵습니다.
• 기존 방법의 한계: 기존 멀티모달 학습 기법들은 주로 이미지 인코더를 사전 학습하거나 크로스 어텐션 (Cross-attention) 을 통해 텍스트와 이미지를 결합하는 방식인데, 이는 DETR(Detection Transformer) 스타일의 탐지기 구조를 변경하거나 오버헤드를 증가시키는 경우가 많습니다.

2. 제안 방법 (Methodology)

저자들은 **객체 쿼리 (Object Queries)**의 표현을 모달리티 컨텍스트와 정렬 (Alignment) 하는 새로운 프레임워크를 제안합니다. 이 프레임워크는 두 가지 핵심 구성 요소로 이루어져 있습니다.

A. 모달리티 토큰 (Modality Tokens)

• 정의: 각 영상 모달리티 (예: CXR, CT) 와 대상 클래스 (예: 폐암, 뇌종양) 를 결합한 텍스트 기반의 컴팩트한 임베딩입니다.
  • 예시: "Aortic enlargement in CXR", "Brain tumor in MRI"
• 특징:
  • 사전 학습된 텍스트 인코더 (CLIP, BiomedCLIP 등) 를 사용하여 생성됩니다.
  • 추가적인 주석 (Annotation) 이 필요 없으며, 경량화되어 있습니다.
  • 모델의 백본 (Backbone) 이나 헤드 (Head) 구조를 변경하지 않고도 모달리티 정보를 명시적으로 제공합니다.

B. Multimodality Context Attention (MoCA)

• 목적: 탐지 과정에서 객체 쿼리에 모달리티 컨텍스트를 주입하기 위한 메커니즘입니다.
• 작동 원리:
  1. 해당 이미지의 모달리티 토큰을 객체 쿼리 세트 (Object Query Set) 에 추가합니다.
  2. 디코더의 셀프 어텐션 (Self-attention) 레이어에서 이 토큰과 객체 쿼리들을 혼합합니다.
  3. 각 객체 쿼리가 모달리티 토큰에 주의를 기울여 (Attend), 모달리티별 컨텍스트를 반영하도록 유도합니다.
• 장점: DETR 스타일 아키텍처를 유지하며, 추가적인 크로스 어텐션 스트림을 도입하지 않아 지연 시간 (Latency) 증가가 미미합니다.

C. QueryREPA (Query Representation Alignment)

• 목적: 탐지 학습 전에 객체 쿼리 표현을 모달리티 토큰과 정렬하는 사전 학습 (Pretraining) 단계입니다.
• 손실 함수: InfoNCE 기반의 **대조적 손실 (Contrastive Loss)**을 사용합니다.
  • 같은 모달리티의 쿼리 평균과 해당 모달리티 토큰을 긍정 쌍 (Positive Pair) 으로, 다른 모달리티 토큰을 부정 쌍 (Negative Pair) 으로 설정합니다.
• 모달리티 배치 샘플링 (Modality Batch Sampling): 각 미니배치에 서로 다른 모달리티의 이미지를 균일하게 포함시켜, 특정 모달리티가 지배하는 것을 방지하고 모달리티 간 분리 (Disentanglement) 를 극대화합니다.
• 효과: 탐지 헤드를 학습하기 전에 쿼리 공간이 모달리티를 인식하고 클래스에 충실한 (Class-faithful) 형태로 재구성됩니다.

3. 주요 기여 (Key Contributions)

1. 검출기 무관한 (Detector-agnostic) 모달리티 토큰: 추가 주석 없이 텍스트 기반 토큰을 생성하여 모달리티와 클래스 컨텍스트를 명시화하는 경량 프레임워크 제안.
2. MoCA (Multimodality Context Attention): 아키텍처 수정 없이 디코더의 셀프 어텐션 레이어에 토큰을 추가하여 모달리티 정보를 주입하는 드롭인 (Drop-in) 메커니즘.
3. QueryREPA: 대조적 손실을 이용한 쿼리 정렬 사전 학습 단계로, 레이블이 부족한 대규모 의료 데이터셋을 활용할 수 있게 하여 하류 작업 (Downstream task) 의 성능을 향상시킴.
4. 검증: 다양한 의료 모달리티 (CXR, CT, MRI, Colonoscopy, Pathology) 가 혼합된 데이터셋에서 기존 SOTA 모델들을 압도하는 성능 달성.

4. 실험 결과 (Results)

• 데이터셋: VinBigData(CXR), COVID-19 CT, LIDC-IDRI, NeoPolyp, BR35H, ACDC, MoNuSeg 등 7 개의 공개 데이터셋을 통합한 혼합 데이터셋 사용.
• 성능 (Quantitative):
  • DINO 베이스라인에 MoCA 와 QueryREPA 를 적용한 모델은 AP 41.3, AP50 65.5를 기록하여, 기존 SOTA 모델 (Deformable DETR, Sparse R-CNN, GLIP, Grounding DINO 등) 보다 모든 지표에서 우수한 성능을 보였습니다.
  • 특히 PubMedCLIP 인코더를 사용할 때 가장 높은 성능을 기록했습니다.
  • 다양한 객체 크기 (Small, Medium, Large) 에서 모두 일관된 성능 향상을 보였습니다.
• 비교 분석:
  • 텍스트 인코더: CLIP, BiomedCLIP, PubMedCLIP 등 다양한 인코더에서 일관된 성능 향상을 보여 방법론의 강건성을 입증했습니다.
  • 레이턴시: MoCA 는 디코더 헤드의 실행 시간을 약 0.2ms 만 증가시켜 (예: DINO 9.37ms → 9.59ms) 거의 무시할 수 있는 오버헤드를 가집니다.
• 비정형 데이터 활용 (Unlabeled Data):
  • QueryREPA 는 바운딩 박스 주석이 없는 이미지 (예: NIH Chest X-ray 의 11 만 장 이상 이미지) 에 대해서도 클래스 레이블만 있으면 사전 학습이 가능합니다. 이를 통해 주석 부족 문제를 해결하고 탐지 성능을 추가로 향상시킬 수 있음을 입증했습니다.
• 시각화 (Qualitative):
  • UMAP 시각화를 통해 QueryREPA 적용 전에는 모달리티 간 쿼리가 섞여 있었으나, 적용 후에는 모달리티별로 명확하게 분리된 클러스터를 형성함을 확인했습니다.
  • 정성적 결과에서 기존 모델들이 놓치거나 부정확하게 탐지했던 병변을 MoCA+QueryREPA 모델이 정확하게 찾아냈습니다.

5. 의의 및 결론 (Significance)

이 논문은 의료 영상 탐지의 다중 모달리티 학습이라는 난제를 해결하기 위해, 객체 쿼리 (Object Query) 레벨의 표현 정렬이라는 새로운 관점을 제시했습니다.

• 실용성: 복잡한 아키텍처 변경 없이 기존 DETR 기반 모델에 쉽게 적용 가능하며, 계산 비용이 매우 낮습니다.
• 확장성: 레이블이 부족한 의료 데이터 환경에서도 대량의 데이터를 활용할 수 있는 경로를 열어주어, 실제 임상 환경에서의 강건한 AI 모델 개발에 기여합니다.
• 일반성: DETR 스타일뿐만 아니라 Sparse R-CNN, DiffusionDet 등 다양한 현대적 탐지 모델에도 적용 가능한 범용적인 프레임워크입니다.

결론적으로, **"Align Your Query"**는 이질적인 의료 영상 데이터의 통계적 편차를 극복하고, 텍스트 기반 컨텍스트를 활용하여 단일 모델로 다양한 모달리티에서 높은 정확도의 객체 탐지를 가능하게 하는 획기적인 접근법입니다.