Zzong's Notes

Rethinking the Role of Token Retrieval in Multi-Vector Retrieval

12 min read

구글 딥마인드 연구원들이 발표한 “다중 벡터 검색에서 토큰 검색의 역할에 대한 재고찰(Rethinking the Role of Token Retrieval in Multi-Vector Retrieval)“이라는 논문은 기존 정보 검색 기술의 효율성과 정확성을 한 단계 끌어올린 새로운 모델 ‘XTR(ConteXtualized Token Retriever)‘을 제시합니다.

기존 모델의 한계

ColBERT와 같은 기존의 다중 벡터 검색 모델은 질의(query)와 문서에 포함된 각 토큰 간의 상호작용을 계산하여 높은 정확도를 보여주었습니다. 하지만 이는 복잡하고 느린 3단계 추론 과정(토큰 검색 → 정보 수집 → 점수 계산)을 필요로 했습니다. 특히 후보 문서의 ‘모든’ 토큰 벡터를 불러와야 하는 ‘정보 수집(gathering)’ 단계는 상당한 계산 비용과 시간을 요구하여 대규모 서비스에 적용하기 어려운 한계가 있었습니다. 또한, 모델 학습은 주로 마지막 ‘점수 계산’ 단계에 맞춰져 있어, 정작 후보군을 결정하는 첫 ‘토큰 검색’ 단계의 성능이 최적화되지 않는 문제도 있었습니다.

A.1.1.1) XTR의 혁신적인 접근법

XTR은 이러한 문제점을 ‘토큰 검색’ 단계의 역할을 재정의함으로써 해결합니다. XTR의 핵심 아이디어는 다음과 같습니다.

  • 효율적인 2단계 프로세스: XTR은 가장 중요하고 유의미한 문서 토큰을 먼저 검색하도록 모델을 훈련시킵니다. 이렇게 선별된 핵심 토큰들만으로도 충분히 정확한 점수 계산이 가능해져, 비용이 많이 드는 ‘정보 수집’ 단계를 완전히 제거하고 검색 과정을 단순화했습니다.
  • 새로운 학습 목표: XTR은 학습 과정에서 토큰 검색 단계를 시뮬레이션하는 새로운 목적 함수를 도입했습니다. 이를 통해 학습 목표와 실제 추론 과정 사이의 불일치를 해소하고 토큰 검색 자체의 성능을 극대화했습니다.
  • 비용 효율적인 점수 계산: 후보 문서 순위를 매기는 점수 계산(scoring) 단계의 연산 비용(FLOPs)을 ColBERT 대비 4,000배 이상 줄였습니다. 이는 검색된 토큰 정보만을 사용하기 때문에 가능하며, 이로 인해 전체적인 검색 효율이 비약적으로 향상되었습니다.

A.1.1.2) 주요 성과

XTR은 정보 검색 분야의 권위 있는 벤치마크인 BEIR에서 별도의 증류(distillation) 기술 없이 기존 최고 성능(SOTA)을 2.8 nDCG@10 만큼 뛰어넘는 새로운 기록을 달성했습니다. 또한, LoTTE, EntityQuestions, 그리고 18개 언어를 포함하는 다국어 검색 벤치마크 MIRACL에서도 뛰어난 성능을 입증했습니다.

분석 결과, XTR은 기존 ColBERT 모델에 비해 실제 정답 문서에 포함된 ‘골든 토큰(gold tokens)‘을 찾아내는 능력이 월등히 뛰어났습니다. 이는 단순히 단어의 표면적 일치보다 문맥적 연관성을 파악하는 능력이 향상되었음을 의미합니다.

결론적으로, XTR은 다중 벡터 검색의 복잡한 과정을 단순화하여 효율성을 크게 높이면서도, 오히려 더 뛰어난 정확도를 달성한 혁신적인 모델이라고 할 수 있습니다. 이는 향후 검색 엔진, 추천 시스템 등 다양한 분야에 실질적으로 기여할 수 있는 중요한 연구 성과로 평가됩니다.

네, 아주 정확한 질문입니다. 사용자의 궁금증을 단계별로 명확하게 설명해 드리겠습니다.

결론부터 말씀드리면, XTR은 ColBERT와 같은 다중 벡터 모델 아키텍처를 기반으로 하지만, 전혀 다른 방식으로 학습시켰습니다. 기존 ColBERT를 그대로 다시 학습시킨 것이 아니라, 학습 목표 자체를 새롭게 설계한 것입니다. 그리고 핵심 연산은 MaxSim(Maximum Inner Product Search)을 변형하여 사용합니다.

A.1.2) ColBERT 모델을 다시 학습시킨 것인가?

아닙니다. 더 정확히는, ColBERT의 “학습 방법”을 근본적으로 바꾼 것입니다.

  • ColBERT의 문제점: ColBERT는 학습 시에 (미니배치 내의) 정답 문서와 오답 문서 간의 최종 점수 차이를 벌리는 데에만 집중합니다. 이 최종 점수는 문서 내 모든 토큰을 대상으로 MaxSim 연산을 수행한 결과입니다. 하지만 실제 서비스(추론)에서는 비용 문제로 모든 문서의 모든 토큰을 볼 수 없고, 1단계인 ‘토큰 검색’을 통해 후보 문서들을 추려냅니다. 즉, 학습 방식(최종 점수 최적화)과 실제 사용 방식(토큰 검색이 중요) 사이에 괴리가 있었습니다.
  • XTR의 접근법: XTR은 이 문제를 해결하기 위해, “학습 과정 자체에서 ‘토큰 검색’ 단계를 흉내 내자”는 아이디어를 도입했습니다. 모델이 최종 점수만 잘 내는 것이 아니라, 애초에 올바른 토큰을 잘 검색하도록 직접적으로 훈련시키는 것입니다.

A.1.3) 구체적으로 어떻게 학습시켰는가? (MaxSim 연산이 아닌가?)

네, 여기서 MaxSim 연산에 대한 중요한 차이가 발생합니다. XTR의 학습 방식은 “In-Batch Token Retrieval” 이라고 불립니다.

A.1.3.1) 기존 ColBERT의 학습 (문서 단위 경쟁)

  1. 하나의 쿼리(Query)가 주어집니다.
  2. 미니배치(mini-batch) 안에 있는 정답 문서 D+와 여러 오답 문서 D-들이 있습니다.
  3. 각 문서(D+, D-)에 대해 독립적으로 MaxSim 점수를 계산합니다.
    • f_colbert(Q, D) = Σ (각 쿼리 토큰과 D 문서 내 가장 유사한 토큰 간의 유사도)
  4. 이 점수를 바탕으로 D+의 점수는 높이고, D-들의 점수는 낮추도록 학습합니다.
    • 문제: 오답 문서 D- 안에 쿼리 토큰과 아주 유사한 토큰이 딱 하나만 있어도, 1단계 ‘토큰 검색’에서는 이 오답 문서가 상위에 노출될 수 있습니다. 하지만 학습 시에는 문서 전체 점수가 낮기 때문에 이 문제가 간과될 수 있습니다.

A.1.3.2) 새로운 XTR의 학습 (토큰 단위 경쟁)

  1. 하나의 쿼리(Query)가 주어집니다.
  2. 미니배치 안에 있는 모든 문서(D+, D-들)의 모든 토큰들을 한곳에 모아 거대한 토큰 풀(pool)을 만듭니다.
  3. 쿼리의 각 토큰(q_i)에 대해, 이 거대한 토큰 풀 전체에서 가장 유사한 상위 k개(ktrain)의 문서 토큰을 찾습니다. 이것이 바로 학습 중에 ‘토큰 검색’을 시뮬레이션하는 과정입니다.
  4. 이제 특정 문서 D의 점수를 계산할 때, 기존처럼 D 안의 모든 토큰을 보는 것이 아닙니다. 오직 3번의 경쟁에서 살아남아 상위 k개 안에 든 D 소속 토큰들만을 대상으로 MaxSim 점수를 계산합니다.
  5. 만약 정답 문서 D+의 토큰 중 단 하나도 3번의 경쟁에서 이기지 못했다면, D+의 점수는 0이 됩니다. 이는 모델에게 엄청난 패널티(loss)를 주게 되고, 모델은 “아, 정답 문서의 토큰들이 경쟁에서 이길 수 있도록 토큰 임베딩을 만들어야겠구나” 라고 학습하게 됩니다.

A.1.4) 결론 및 비유

  • ColBERT의 MaxSim: 쿼리 토큰이 각 문서에 들어가서 “이 문서 안에서 나랑 가장 닮은 토큰이 누구지?”를 찾는 ‘지역 탐색’에 가깝습니다.
  • XTR의 MaxSim: 쿼리 토큰이 “이 훈련용 미니배치 전체에서 나랑 가장 닮은 토큰 k개를 뽑아라!” 라고 명령한 뒤, “그 우승자들 중에 이 문서 출신이 몇이나 되지?”를 확인하는 ‘전국 대회’ 방식에 가깝습니다.

이러한 “In-batch” 경쟁을 통해 XTR은 단순히 최종 문서 점수를 높이는 것을 넘어, 토큰 검색 단계 자체의 성능(리콜, recall)을 극적으로 향상시켰습니다. 그 결과, 비용이 많이 드는 ‘정보 수집(gathering)’ 단계를 생략하고도 더 빠르고 정확한 검색이 가능해진 것입니다.