콘텐츠로 이동
Data Prep
상세

STM: 도메인 특화 Retriever를 위한 모듈형 프레임워크

Less Finetuning, Better Retrieval: Rethinking LLM Adaptation for Biomedical Retrievers

항목 내용
arXiv 2602.04731
저자 Sameh Khattab et al.
발표 2026-02-04
카테고리 cs.CL, cs.LG

1. 문제 정의

도메인 특화 Retriever의 과제

RAG(Retrieval-Augmented Generation)에서 retriever 성능이 핵심이지만, 특수 도메인(의료, 법률 등)에 적응시키는 데 어려움이 있다:

  1. 대규모 사전학습 필요: 도메인 코퍼스로 재학습 비용
  2. 일반 도메인 능력 손실: 파인튜닝 후 기존 능력 망각
  3. 데이터 부족: 특수 도메인의 레이블 데이터 희소

STM의 접근

Synthesize-Train-Merge 3단계 모듈형 프레임워크:

[Synthesize] 합성 Hard Negatives 생성
[Train] 도메인 특화 학습
[Merge] 모델 병합으로 일반 능력 보존

2. STM 프레임워크

2.1 Synthesize: 합성 Hard Negatives

Hard negative mining의 자동화:

쿼리: "COVID-19 백신의 부작용은?"

Easy Negative: "자동차 수리 방법"  ← 쉽게 구분
Hard Negative: "COVID-19 검사 방법" ← 관련 있지만 다름
    이것이 학습에 효과적

생성 방법:

  1. LLM으로 쿼리 분석
  2. 의미적으로 유사하지만 관련 없는 문서 생성
  3. 난이도 조절 가능
\[\text{hardness}(d, q) = \text{sim}(d, q) - \text{relevance}(d, q)\]

2.2 Train: Retrieval Prompt 최적화

Decoder-only LLM을 retriever로 활용:

[기존 방식]
입력: "쿼리 텍스트"
출력: 임베딩

[STM 방식]
입력: "[RETRIEVE] 의료 문헌에서 {쿼리} 관련 문서 검색"
출력: 최적화된 임베딩

프롬프트 최적화:

  • 태스크별 프롬프트 템플릿 학습
  • 도메인 키워드 자동 삽입
  • 검색 의도 명시화

2.3 Merge: Model Merging

파인튜닝 모델과 원본 모델의 가중치 병합:

\[\theta_{merged} = \alpha \cdot \theta_{finetuned} + (1-\alpha) \cdot \theta_{original}\]

병합 전략:

방법 설명 장점
Linear 단순 가중 평균 간단, 안정적
SLERP 구면 선형 보간 방향 보존
TIES Task-specific 가중치 선택 충돌 최소화
DARE 랜덤 드롭 후 병합 희소성 활용

3. 아키텍처 상세

전체 파이프라인

┌─────────────────────────────────────────────┐
│           1. Synthesize Phase               │
├─────────────────────────────────────────────┤
│  도메인 코퍼스 → LLM → Hard Negatives 생성   │
│        ↓                                    │
│  (쿼리, 정답, hard_neg) 트리플렛 구성        │
└─────────────────────────────────────────────┘
┌─────────────────────────────────────────────┐
│           2. Train Phase                    │
├─────────────────────────────────────────────┤
│  Base LLM + Retrieval Prompt                │
│        ↓                                    │
│  Contrastive Learning (InfoNCE)             │
│        ↓                                    │
│  Domain-specific Retriever                  │
└─────────────────────────────────────────────┘
┌─────────────────────────────────────────────┐
│           3. Merge Phase                    │
├─────────────────────────────────────────────┤
│  Finetuned Model + Original Model           │
│        ↓                                    │
│  Weight Interpolation (TIES/DARE)           │
│        ↓                                    │
│  Final Merged Retriever                     │
└─────────────────────────────────────────────┘

학습 목표

InfoNCE 손실 함수:

\[\mathcal{L} = -\log \frac{\exp(s(q, d^+)/\tau)}{\exp(s(q, d^+)/\tau) + \sum_{d^- \in \mathcal{N}} \exp(s(q, d^-)/\tau)}\]
  • \(s(q, d)\): 쿼리-문서 유사도
  • \(d^+\): 정답 문서
  • \(\mathcal{N}\): negative 문서 집합 (합성 hard negative 포함)
  • \(\tau\): temperature

4. 실험 결과

MTEB 벤치마크 성능

12개 의료 + 일반 태스크에서 평가:

모델 의료 태스크 일반 태스크 평균
BGE-base 0.512 0.687 0.600
E5-large 0.524 0.701 0.613
Task Expert 0.589 0.621 0.605
STM 0.632 0.695 0.664

향상 정도

비교 대상 향상폭
Task-specific Expert 대비 최대 +23.5%
평균 향상 +7.5%

핵심 발견

  1. 일반 능력 보존: 병합 후에도 일반 도메인 성능 유지
  2. 대규모 사전학습 불필요: 합성 데이터 + 병합으로 충분
  3. 모듈형 설계: 각 단계 독립적 개선 가능

5. Ablation Study

각 컴포넌트 기여도

설정 의료 성능 일반 성능
Base only 0.512 0.687
+ Synthesize 0.573 0.672
+ Train 0.608 0.654
+ Merge 0.632 0.695

Hard Negative 수의 영향

Hard Neg 수 성능 비고
0 (random only) 0.541 baseline
3 0.598 +10.5%
7 0.621 +14.8%
15 0.632 +16.8%
31 0.629 포화

6. 구현 가이드

최소 요구사항

# 필요 라이브러리
transformers >= 4.35
sentence-transformers >= 2.2
mergekit  # 모델 병합용

# 하드웨어
# 7B 모델 기준: GPU 24GB+ (A100/H100 권장)

핵심 하이퍼파라미터

파라미터 권장값 설명
learning_rate 2e-5 도메인 학습 LR
batch_size 32 대조 학습 배치
hard_neg_ratio 0.7 hard negative 비율
merge_alpha 0.3~0.5 병합 시 도메인 가중치
temperature 0.02 InfoNCE temperature

7. 참고 자료

  • arXiv 원문
  • 관련 연구: Dense Retrieval, Model Merging, Hard Negative Mining
  • 참고 도구: mergekit, sentence-transformers

정리일: 2026-03-01