← 소식 목록

AURA: 학습형 희소 검색에 의미 표현과 단어 정보량을 함께 담다

연구 목적

검색 엔진은 질의와 관련된 문서를 빠르게 찾아야 합니다. BM25 같은 전통적인 희소 검색은 역색인(inverted index) 덕분에 빠르지만, 단어가 겹쳐야만 관련성을 알아챕니다. 학습형 희소 검색(Learned Sparse Retrieval, LSR)은 사전학습 언어모델로 단어 가중치를 학습해 이 한계를 줄였고, 대표 모델이 SPLADE입니다.

하지만 SPLADE도 여전히 희소한 단어 표현으로만 관련성을 판단합니다. 그래서 문서 전체의 의미를 담는 데 한계가 있습니다. 반대로 dense retrieval은 전체 의미에는 강하지만 단어 수준의 일치에는 약합니다. 또 SPLADE의 희소성 정규화는 모든 토큰을 비슷하게 누르기 때문에, 정보가 많은 단어와 적은 단어를 구분하지 않습니다.

핵심 아이디어

AURA(Alignment of Unified-semantic and Rare-lexical Attributes)는 SPLADE 기반 틀 위에 세 가지를 더합니다.

  • Dense retrieval loss: BERT의 [CLS] 임베딩으로 dense 표현을 함께 학습합니다. 검색할 때는 희소 점수와 dense 점수를 섞어 씁니다. (S_total = α·S_sparse + (1−α)·S_dense, α = 0.75)
  • IDF-FLOPS loss: FLOPS 정규화에 IDF 가중치를 곱합니다. 불용어처럼 흔한 토큰은 더 강하게 억제하고, 문서를 구별하는 드문 토큰은 살립니다. IDF는 MS MARCO에서 비정상적인 subword 토큰을 걸러낸 뒤 계산했고, 숫자·알파벳 토큰은 남겼습니다.
  • 스케줄러: 주 손실(sparse retrieval loss)은 처음부터 그대로 쓰고, 두 보조 손실의 가중치는 (t/T)² 꼴로 천천히 키웁니다. T = 50,000 step에서 1이 되고 이후 고정됩니다.

실험 설정

  • 학습·in-domain 평가: MS MARCO Passage Ranking, TREC DL 2019
  • zero-shot 평가: BEIR 12개 데이터셋 (ArguAna, Climate-FEVER, DBPedia, FEVER, FiQA, HotpotQA, NFCorpus, NQ, Quora, SCIDOCS, SciFact, TREC-COVID)
  • 지표: MS MARCO는 MRR@10, Recall@1000 / TREC DL과 BEIR는 nDCG@10
  • 비교 대상: 희소 검색(BM25, DeepCT, doc2query-T5, COIL-tok, DeepImpact, SPLADE, SPLADE-max)과 dense 검색(ANCE, TCT-ColBERT, TAS-B, RocketQA 등)
  • 정규화 계수는 SPLADE-max와 같게 λ_q = 3×10⁻⁴, λ_d = 1×10⁻⁴를 썼습니다.

결과

MS MARCO와 TREC DL 2019 결과입니다(일부 행만 발췌).

ModelMS MARCO dev MRR@10MS MARCO dev R@1000TREC DL 2019 nDCG@10TREC DL 2019 R@1000
BM250.1840.8530.5060.745
TAS-B0.3470.9780.7170.843
SPLADE0.3220.9550.6650.813
SPLADE-max0.3440.9680.7040.812
AURA0.3440.9690.7240.843

MS MARCO에서는 SPLADE-max와 비슷한 MRR@10을 유지했고, TREC DL 2019에서는 nDCG@10 0.724, R@1000 0.843을 기록했습니다.

BEIR nDCG@10 결과입니다(일부 행만 발췌, "Avg. all"은 MS MARCO 포함 평균).

CorpusBM25ContrieverSPLADE-maxAURA
FEVER0.7530.7580.7190.778
HotpotQA0.6030.6380.6380.656
TREC-COVID0.6560.5960.6630.735
SciFact0.6650.6770.6090.656
Avg. all0.4260.4890.4650.489
Avg. zero-shot0.4430.4960.4700.496

AURA는 zero-shot 평균을 SPLADE-max의 0.470에서 0.496으로 높였고, 사전학습 모델인 Contriever와 같은 평균을 냈습니다. 특히 개체명이나 의학 용어처럼 정확한 단어 일치가 중요한 HotpotQA, TREC-COVID에서 향상이 컸습니다.

구성 요소를 하나씩 뺀 실험에서는 dense retrieval loss를 뺄 때 zero-shot 평균이 0.465로 가장 크게 떨어졌습니다. IDF 가중치를 빼면 0.481, 스케줄러를 빼면 0.490이었습니다.

추가로, 평균 문서 길이보다 긴 질의에서는 활성화된 질의 토큰 중 상위 30%만 쓰면 ArguAna 점수가 0.396에서 0.426으로 오르고, zero-shot 평균도 0.498로 조금 올랐습니다.

한계와 다음 단계

  • dense 표현을 함께 쓰기 때문에, SPLADE처럼 희소 표현만 쓰는 방법보다 색인·저장·추론 비용이 더 들 수 있습니다.
  • 두 보조 손실에 같은 형태의 스케줄러를 쓰고 있어, 항상 최적이라고 보긴 어렵습니다. 손실별 스케줄러나 적응형 스케줄링이 다음 과제입니다.
  • dense 학습이 들어가면서 희소 표현의 활성 토큰 수가 늘어납니다(예시 질의에서 SPLADE-max 39개, AURA 75개). 논문은 이것이 SciFact에서 성능이 조금 낮은 이유일 수 있다고 봅니다.
  • 앞으로 더 다양한 검색 환경에서 일반성을 확인하고, 표현 제어와 스케줄링을 더 정교하고 자동화된 방식으로 발전시킬 계획입니다.