Love Me, Love My Label: Rethinking the Role of Labels in Prompt Retrieval for Visual In-Context Learning

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

• 배경: 시각적 문맥 학습 (Visual In-Context Learning, VICL) 은 증강된 예시 (demonstrative prompts) 를 통해 시각 기반 모델이 다양한 작업을 수행하도록 유도하는 기술입니다. VICL 의 성능은 선택된 프롬프트의 질에 크게 의존합니다.
• 기존 방법의 한계: 기존 연구 (SupPR, Partial2Global 등) 는 프롬프트 이미지 간의 시각적 유사성 (image similarity) 에만 집중하여 프롬프트를 검색하거나 재순위화 (reranking) 했습니다.
• 핵심 문제:
  • 레이블 불일치 (Label Inconsistency): 시각적으로 유사하지만 레이블 (예: 객체 클래스, 세그멘테이션 마스크 등) 이 쿼리와 일치하지 않는 프롬프트가 검색될 수 있습니다.
  • 성능 저하: 실험 결과, 시각적 유사성은 높지만 레이블이 일치하지 않는 프롬프트는 VICL 성능을 저하시키는 것으로 확인되었습니다. 반면, 쿼리와 프롬프트 간의 레이블 일관성 (Label Consistency) 이 높을수록 성능이 향상됩니다.
  • 정보 손실: 기존 검색 방식에서는 테스트 시 쿼리 레이블을 알 수 없다는 이유로 프롬프트 측의 레이블 정보를 무시하거나 배제하는 경향이 있었습니다.

2. 제안 방법: LaPR (Label-aware Prompt Retrieval)

저자들은 레이블 정보를 프롬프트 검색 과정에 명시적으로 통합하기 위해 LaPR이라는 새로운 프레임워크를 제안했습니다.

2.1. 핵심 아키텍처

1. 레이블 인식 임베딩 (Label-Aware Representation):

   • 프롬프트 측: 프롬프트 이미지 ($I^p$) 와 해당 레이블 ($L^p$) 을 결합하여 Joint Representation을 생성합니다. 이를 통해 레이블 정보를 명시적으로 임베딩에 주입합니다.
   • 쿼리 측: 테스트 시 쿼리 레이블 ($y_q$) 을 알 수 없다는 문제를 해결하기 위해 혼합 전문가 (Mixture-of-Experts, MoE) 메커니즘을 도입합니다.
     • 전문가 (Experts): 각 전문가 ($E_k$) 는 서로 다른 모드 (예: 긴 꼬리, 뿔, 날카로운 부리 등 세부 특징) 를 포착하도록 설계됩니다.
     • 라우터 (Router): 쿼리 이미지의 특징을 기반으로 각 모드의 가중치 ($\pi_q$) 를 추론합니다. 이를 통해 쿼리의 암시적 (unknown) 레이블을 추정하고, 쿼리에 가장 적합한 프롬프트 정보를 추출합니다.

2. 쿼리 관련 프롬프트 추출:

   • 라우터가 추론한 가중치 ($\pi_q$) 를 사용하여 프롬프트 측의 모드별 임베딩을 가중 합산하여, 쿼리에 가장 관련성이 높은 프롬프트 임베딩을 생성합니다.

2.2. 최적화 전략 (Alternating Optimization)

전문가 (Experts) 와 라우터 (Router) 의 역할을 분리하여 안정적으로 학습시키기 위해 교대 최적화 (Alternating Optimization) 방식을 사용합니다. 미니배치당 두 단계의 업데이트를 수행합니다.

• 단계 1: 전문가 학습 (Expert Step)
  • 목표: VICL 작업 수행 성능을 극대화합니다.
  • 방법: 라우터 가중치를 고정하고, 성능 기반 대비 손실 (VICL performance-guided contrastive loss, $\mathcal{L}_{PG}$) 을 사용하여 각 전문가가 특정 모드를 효과적으로 인코딩하도록 학습시킵니다. (실제 VICL 모델의 예측 결과를 기반으로 양/음성 샘플을 선정)
• 단계 2: 라우터 학습 (Router Step)
  • 목표: 추론된 모드 혼합 비율이 실제 쿼리 레이블과 일치하도록 조정합니다.
  • 방법: 전문가를 고정하고, 레이블 기반 대비 손실 (Label-guided contrastive loss, $\mathcal{L}_{LG}$) 과 부하 균형 손실 (Load-balancing loss, $\mathcal{L}_{LB}$) 을 사용하여 학습합니다.
    • $\mathcal{L}_{LG}$: 레이블 일관성이 높은 프롬프트를 양 (Positive) 으로 간주하여 라우터가 올바른 모드를 선택하도록 유도합니다.
    • $\mathcal{L}_{LB}$: 모든 전문가가 고르게 사용되도록 하여 특정 전문가만 편중되는 것을 방지합니다.

3. 주요 기여 (Key Contributions)

1. VICL 에서의 첫 번째 레이블 인식 프롬프트 검색: 프롬프트 검색 시 이미지 유사성뿐만 아니라 레이블 불일치 문제를 명시적으로 해결하고, 레이블의 중요성을 개념적으로 재정의했습니다.
2. 새로운 아키텍처 설계: 프롬프트 측에는 이미지 - 레이블 결합 임베딩을, 쿼리 측에는 MoE 기반의 라우팅 메커니즘을 도입하여 쿼리의 암시적 레이블을 추정하고 관련 정보를 추출합니다.
3. 역할 분리 및 교대 최적화: 전문가와 라우터의 학습 목표를 분리하여, 전문가는 모드별 표현력을 강화하고 라우터는 레이블 일관성을 맞추도록 설계했습니다.
4. 광범위한 실험 검증: 다양한 백본 (Feature Extractor) 과 교차 폴드 (Cross-fold) 시나리오에서 뛰어난 일반화 능력을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 작업: 전경 분할 (Foreground Segmentation, Pascal-5i), 단일 객체 감지 (Single-object Detection, Pascal VOC), 이미지 컬러화 (Colorization, ImageNet) 작업에서 평가되었습니다.
• 성능:
  • 분할 (Segmentation): 기존 최상위 방법 (RH-Partial2Global 등) 대비 평균 mIoU 에서 6.00%p 향상.
  • 감지 (Detection): 3.46%p 향상.
  • 볼팅 (Voting) 적용 시: 재순위화 기반 방법과 동급 이상의 성능을 보이면서도 검색 기반 방법으로는 최상위 성능을 기록했습니다.
• 교차 폴드 전이 (Cross-Fold Transferability): 한 폴드에서 학습된 모델을 다른 폴드에 적용했을 때, 기존 방법 (SupPR, Partial2Global) 대비 약 14% 이상 성능이 우수했습니다. 이는 MoE 가 카테고리 정보보다 더 세분화된 모드를 포착하고 라우터가 이를 적응적으로 선택하기 때문입니다.
• 백본 일반화: CLIP 인코더뿐만 아니라 DINOv2 와 같은 다른 특징 추출기에서도 일관된 성능 향상을 보여주었습니다.

5. 의의 및 결론 (Significance)

• 레이블의 중요성 재발견: VICL 프롬프트 검색에서 레이블 정보는 단순한 부가 정보가 아니라, 성능을 결정하는 핵심 요소임을 입증했습니다.
• 새로운 패러다임: "이미지만 보고 검색"하는 기존 방식에서 "이미지와 레이블의 일관성을 고려하여 검색"하는 방식으로 패러다임을 전환했습니다.
• 실용성: 추론 시 쿼리 레이블이 없더라도 라우팅 메커니즘을 통해 효과적으로 레이블 정보를 활용함으로써, 실제 응용 환경에서의 VICL 성능을 크게 향상시킬 수 있는 기반을 마련했습니다.

이 논문은 시각적 문맥 학습의 성능을 극대화하기 위해 레이블 정보의 체계적인 활용이 필수적임을 강조하며, 이를 위한 새로운 기술적 프레임워크를 제시했다는 점에서 의의가 큽니다.