ColBERT: BERT 기반 지연 상호작용 검색

ColBERT: BERT 기반 지연 상호작용 검색

질문과 문서를 BERT로 따로 인코딩하되, 문장 벡터 하나로 뭉개지 않고 토큰마다 벡터를 남긴 뒤 마지막에 값싼 'MaxSim' 계산으로 맞춰 보는 '지연 상호작용(late interaction)' 검색 모델입니다. BERT 재순위 모델에 가까운 정확도를 내면서 170배 넘게 빨랐습니다.

들어가기 전에

이 해설은 아래 내용을 안다고 보고 씁니다. 처음이라면 먼저 읽고 오면 수식과 손계산을 훨씬 쉽게 따라올 수 있습니다.

• 벡터 검색과 근사 최근접 이웃: 평탄 · IVF · PQ · HNSW — 벡터 인덱스와 근사 검색을 먼저 정리하는 수업
• 오픈 도메인 질의응답을 위한 밀집 문단 검색 (DPR) — 질문 · 문단을 BERT 두 개로 벡터화해 내적으로 찾는 검색기로 BM25를 앞선 논문

한눈에 보기

• 검색 모델의 세 가지 설계 — 표현 중심(벡터 하나), 상호작용 중심(교차 인코더), 그리고 둘을 섞은 지연 상호작용 — 의 차이를 알게 됩니다.
• 질문 토큰마다 "문서에서 가장 닮은 토큰"을 찾아 더하는 MaxSim 점수를 손으로 계산해 봅니다.
• 질문을 [mask]로 채우는 질문 증강(query augmentation) 이 왜 필요한지 이해합니다.
• 속도 · 정확도 · 저장 공간 사이의 균형을 논문 수치로 판단할 수 있게 됩니다.

초록 (번역)

최근 자연어 이해(NLU)의 발전이 정보 검색(IR)의 빠른 진보를 이끌고 있으며, 이는 주로 문서 순위 매기기를 위해 깊은 언어 모델(LM)을 미세조정한 덕분입니다. 이런 LM 기반 순위 모델은 매우 효과적이지만, 이전 방법들보다 계산 비용을 몇 자릿수나 늘립니다. 특히 관련도 점수 하나를 계산하려고 질문-문서 쌍마다 거대한 신경망을 통과시켜야 하기 때문입니다. 이 문제를 풀기 위해 우리는 깊은 LM(특히 BERT)을 효율적인 검색에 맞게 고친 새로운 순위 모델 ColBERT를 제안합니다. ColBERT는 질문과 문서를 BERT로 각각 따로 인코딩한 뒤, 둘 사이의 세밀한 유사성을 모델링하는 값싸면서도 강력한 상호작용 단계를 쓰는 '지연 상호작용' 구조를 도입합니다. 이 세밀한 상호작용을 늦추되 그대로 유지함으로써, ColBERT는 깊은 LM의 표현력을 활용하면서 동시에 문서 표현을 미리 오프라인으로 계산해 둘 수 있게 되어 질문 처리 속도를 크게 높입니다. 전통적인 모델이 검색한 문서를 재순위하는 비용을 줄이는 것을 넘어, ColBERT의 가지치기 친화적 상호작용 방식은 벡터 유사도 색인을 활용해 대규모 문서 모음에서 직접 끝까지 검색하는 것도 가능하게 합니다. 최근의 두 문단 검색 데이터셋으로 ColBERT를 폭넓게 평가했습니다. 그 결과 ColBERT의 효과는 기존 BERT 기반 모델과 대등하면서(그리고 BERT를 쓰지 않는 모든 기준 모델보다 좋으면서), 실행 속도는 두 자릿수 배 빠르고 질문당 FLOPs는 네 자릿수 배 적었습니다.

무엇을 풀었나

BERT로 검색 결과의 순위를 다시 매기면 정확도가 크게 올랐습니다. 논문 그림 기준으로 이전 최고 방법보다 MRR@10이 거의 7%p 높았습니다. 대신 질문 하나에 문서 1,000개를 재순위하면 BERT를 1,000번 돌려야 했고, 고성능 GPU에서도 지연 시간이 수만 밀리초까지 늘었습니다. 응답이 100ms만 늦어져도 사용자 경험과 매출에 영향을 준다는 점에서 실서비스에 쓰기 어려웠습니다.

검색 모델 설계는 크게 두 갈래였습니다.

• 표현 중심(representation-focused): 질문과 문서를 각각 벡터 하나 로 만들고 비교합니다(SBERT · DPR 방식). 문서 벡터를 미리 계산할 수 있어 빠르지만, 문서 전체를 점 하나로 눌러 담아 세부 단어 일치 정보가 사라집니다.
• 상호작용 중심(interaction-focused): 질문과 문서의 단어끼리 직접 맞대 봅니다. BERT에 둘을 함께 넣는 교차 인코더가 가장 강력한 형태입니다. 정확하지만 문서를 미리 계산해 둘 수 없습니다.

ColBERT의 질문은 이것입니다. 세밀한 단어 단위 비교는 살리면서, 문서 쪽 계산은 미리 해 둘 수 없을까? 답은 상호작용을 "BERT 안"이 아니라 "BERT를 다 지난 뒤"로 미루는 것입니다.

핵심 아이디어

123456flowchart LR
  Q["질문 [Q] + [mask] 채움"] --> BQ[BERT] --> LQ[선형층 128차원 + 정규화] --> EQ[질문 토큰 벡터 32개]
  D["문서 [D]"] --> BD[BERT] --> LD[선형층 128차원 + 정규화 + 문장부호 제거] --> ED[(문서 토큰 벡터들 · 미리 저장)]
  EQ --> M[MaxSim: 질문 토큰마다 가장 닮은 문서 토큰]
  ED --> M
  M --> S[합계 = 관련도 점수]

1. 지연 상호작용과 MaxSim

직관. 채용 담당자가 '필요 역량 목록'의 항목마다 지원서에서 가장 잘 맞는 문장 하나를 찾아 점수를 매기고, 그 점수를 모두 더하는 방식입니다. 지원서 문장들은 미리 정리해 둘 수 있고, 비교는 마지막에 한 번만 합니다.

단계별 설명.

1. 질문을 BERT에 넣어 토큰마다 벡터를 얻습니다. 이것이 질문 벡터 묶음 $E_q$입니다.
2. 문서도 BERT에 넣어 토큰마다 벡터를 얻습니다. 이것이 문서 벡터 묶음 $E_d$이고, 미리 계산해 디스크에 저장해 둡니다.
3. 질문 토큰 하나를 골라, 문서의 모든 토큰 벡터와 유사도를 재고 가장 큰 값 하나만 남깁니다(MaxSim).
4. 질문의 모든 토큰에 대해 3을 반복하고 그 값들을 더합니다. 이 합이 관련도 점수입니다.

수식.

숫자로 해 보기 (설명용 예시, 논문 수치 아님). 질문 토큰 2개, 문서 토큰 3개, 2차원 단위 벡터라고 합시다.

• 질문: $q_1 = (1, 0)$, $q_2 = (0, 1)$
• 문서: $d_1 = (0.8, 0.6)$, $d_2 = (0.6, 0.8)$, $d_3 = (1, 0)$

$q_1$과의 내적은 $(0.8, 0.6, 1.0)$이라 최댓값 1.0, $q_2$와의 내적은 $(0.6, 0.8, 0)$이라 최댓값 0.8입니다. 점수는 $1.0 + 0.8 = 1.8$입니다. 만약 문서 전체를 평균 벡터 하나로 줄였다면 $(0.8, 0.47)$ 정도가 되어, 질문의 두 측면이 각각 문서의 어느 부분과 정확히 맞는지가 흐려집니다.

왜 '평균'이 아니라 '최대'인가. 논문의 절제 실험에서 MaxSim 대신 평균 유사도를 쓴 모델은 MRR@10이 눈에 띄게 낮았습니다. 질문의 각 단어가 문서의 특정 단어에 '집중'해 맞춰지는 것이 중요하다는 해석입니다. 또 벡터 하나([CLS], 4,096차원)로 비교하는 모델도 ColBERT보다 크게 낮아, 토큰 단위 비교 자체가 핵심임을 보였습니다.

왜 빠른가. 일반 BERT 재순위는 문서 $k$개마다 '질문+문서' 긴 입력을 BERT에 넣어야 하고, 어텐션 비용은 길이의 제곱으로 늡니다. ColBERT는 짧은 질문만 BERT에 한 번 넣고, 나머지는 미리 저장된 벡터와의 행렬 곱 · 최댓값 · 합계뿐입니다. 상호작용 단계에는 학습되는 파라미터도 없습니다.

2. 질문 · 문서 인코더의 세부 설계

직관. 같은 번역가(BERT 하나)가 질문과 문서를 모두 읽되, 앞에 붙은 꼬리표([Q] / [D])로 지금 무엇을 읽는지 구분합니다.

단계별 설명.

1. 질문 · 문서 인코더는 BERT 하나를 공유 합니다. 질문 앞에는 [Q], 문서 앞에는 [D] 특수 토큰을 [CLS] 바로 뒤에 붙입니다.
2. BERT 출력 위에 활성화 함수 없는 선형층을 두어 차원을 $m = 128$로 줄입니다(BERT 은닉 차원보다 훨씬 작게). 저장 공간과 GPU 전송 시간이 줄어듭니다.
3. 출력 벡터를 길이(L2 노름) 1로 정규화합니다. 그래서 내적이 곧 코사인 유사도입니다.
4. 문서 쪽에서는 문장부호에 해당하는 벡터를 걸러 내 저장량을 줄입니다.

수식 (논문 식을 이 글의 용어로 정리).

3. 질문 증강: [mask]로 빈칸을 채운다

직관. 짧은 질문 "서울 맛집"만으로는 정보가 부족합니다. 빈칸 몇 개를 덧붙여 두면, BERT가 문맥을 보고 그 빈칸에 '추천', '가격', '위치' 같은 관련 의미를 스스로 채워 넣습니다.

단계별 설명.

1. 질문이 32토큰보다 짧으면 [mask] 토큰으로 32개까지 채웁니다(길면 32개에서 자름).
2. BERT는 [mask] 자리에도 문맥에 맞는 벡터를 만들어 냅니다.
3. 이 벡터들도 MaxSim에 참여하므로, 질문에 없던 관련 단어를 찾거나 기존 단어의 비중을 조절하는 효과가 납니다.

왜 이렇게 했나. 논문은 이것을 '부드러운 질문 확장'으로 설명하고, 질문 증강을 뺀 모델은 MRR@10이 눈에 띄게 낮았다고 보고합니다.

4. 학습

질문 $q$, 관련 문서 $d^+$, 비관련 문서 $d^-$의 삼중항을 씁니다. 두 문서의 점수 $S_{q,d^+}$, $S_{q,d^-}$에 소프트맥스를 씌워 $d^+$를 고르도록 교차 엔트로피로 학습합니다(쌍 단위 소프트맥스 교차 엔트로피).

(위 식은 논문의 '쌍 단위 소프트맥스 교차 엔트로피'를 풀어 쓴 것입니다.)

• 학습률 $3 \times 10^{-6}$, 배치 32, Adam, $N_q = 32$, $m = 128$
• MS MARCO는 공식 BERT-base로 시작해 20만 번 반복 학습

숫자로 해 보기 (설명용). $S_{q,d^+} = 1.8$, $S_{q,d^-} = 1.2$라면 $d^+$를 고를 확률은 $e^{1.8} / (e^{1.8} + e^{1.2}) \approx 6.05 / (6.05 + 3.32) \approx 0.65$, 손실은 약 0.43입니다.

5. 두 가지 쓰는 법: 재순위와 끝까지 검색

• 재순위(re-rank): BM25가 뽑은 상위 1,000개 문서의 저장된 벡터를 GPU로 옮겨 한꺼번에 MaxSim 점수를 매깁니다. 논문은 이때 비용의 대부분이 '벡터를 모아 GPU로 옮기는 일'이라고 밝힙니다.
• 끝까지 검색(end-to-end): 문서 모음 전체의 토큰 벡터를 FAISS 색인(IVFPQ: 공간을 칸 $P$개로 나누고 벡터를 바이트 단위로 압축)에 넣습니다. ① 질문 토큰 벡터 32개 각각으로 가까운 문서 토큰을 찾아 그 문서들을 후보로 모으고, ② 후보만 정확한 MaxSim으로 다시 순위를 매기는 2단계입니다. MaxSim이 '토큰 하나의 최근접 이웃'으로 쪼개지는 성질 덕분에 가능한 방식입니다.

결과

MS MARCO 재순위 (BM25 상위 1,000개를 다시 순위 매김)

• ColBERT는 BERT-base와 대등한 정확도에 지연 시간은 170배 넘게, FLOPs는 13,900배 적었습니다. 지연 61ms 중 질문 인코딩과 상호작용은 13ms뿐이고 나머지는 벡터를 모아 옮기는 시간입니다.
• 재순위할 문서 수 $k$가 늘수록 격차가 커졌습니다. BERT의 FLOPs가 $k = 10$일 때 약 180배, $k = 1000$일 때 13,900배, $k = 2000$일 때 23,000배였습니다.

MS MARCO 끝까지 검색 (880만 문단 전체에서 상위 1,000개)

• 끝까지 검색한 ColBERT의 Recall@50이 공식 BM25의 Recall@1000(81.4)보다 높았습니다. BM25가 아예 놓친 문서를 찾아내 MRR@10도 재순위보다 올랐습니다.

TREC CAR (공식 지표 MAP): BM25 15.3, BM25 + BERT-base 31.0, BM25 + ColBERT 31.3, BM25 + BERT-large 33.5.

저장 공간 (MS MARCO, 재순위 기준)

• 가장 작게 줄인 설정도 MRR@10이 1%p 남짓만 낮았습니다. 색인은 GPU 4개 서버 한 대로 약 3시간이면 만들 수 있었습니다.

왜 중요한가 · 이후 영향

• '다중 벡터(multi-vector) 검색'이라는 갈래를 열었습니다. 벡터 하나(DPR)와 교차 인코더 사이에 실용적인 중간 지점이 있다는 것을 보였습니다.
• RAG 검색 품질 향상의 실전 도구가 됐습니다. 후속작 ColBERTv2가 저장 공간 문제를 줄이면서, RAG 파이프라인의 검색기나 재순위기로 쓰이게 됐습니다.
• 재순위와 끝까지 검색을 한 모델로 할 수 있다는 점도 실무에서 매력적입니다. 같은 색인으로 BM25 결과를 다시 매길 수도, 전체 모음에서 바로 찾을 수도 있습니다.

한계 · 주의할 점

• 저장 공간이 큽니다. 문서마다 토큰 수만큼 벡터를 저장하므로, 기본 설정에서 MS MARCO에 286GiB가 필요했습니다. 벡터 하나짜리 검색보다 한 자릿수 이상 큽니다. 이 문제를 정면으로 다룬 것이 ColBERTv2입니다.
• 끝까지 검색은 여전히 무겁습니다. 458ms로 BM25(62ms)보다 느리고, 색인 구축에 GPU가 필요합니다.
• 최고 정확도는 아닙니다. BERT-large 재순위(36.5)보다는 조금 낮습니다. 비용 대비 정확도가 강점입니다.
• 평가는 영어 웹 문단(MS MARCO)과 위키백과(TREC CAR)에 한정됩니다.

스스로 점검

1. 질문 토큰 $q = (1, 0)$, 문서 토큰 $(0.6, 0.8)$과 $(0, 1)$뿐이라면 이 질문 토큰의 MaxSim 값은?

   내적이 각각 0.6과 0이므로 최댓값 0.6입니다.

2. ColBERT가 일반 BERT 재순위보다 빠른 근본 이유는?

   문서 쪽 BERT 계산을 미리 해 저장해 두고, 질문 시점에는 짧은 질문만 BERT에 한 번 넣은 뒤 값싼 내적 · 최댓값 · 합계만 하기 때문입니다.

3. 질문을 [mask]로 32토큰까지 채우는 이유는?

   BERT가 빈칸 자리에 문맥에 맞는 의미를 채워 넣어 질문을 부드럽게 확장하고 단어 비중을 조절하게 하려는 것입니다. 이 장치를 빼면 정확도가 눈에 띄게 떨어졌습니다.

용어 정리

이어서 읽을 것

• 앞: Sentence-BERT · 밀집 문단 검색 (DPR)
• 뒤: ColBERTv2 · Contriever · RAG 총정리 (서베이)

다음에 볼 것

이 글을 이해했다면 이어서 보면 좋은 항목입니다.

• ColBERTv2: 가볍고 효과적인 지연 상호작용 검색 — 교차 인코더 증류와 중심점+잔차 압축으로 ColBERT를 더 정확하고 6~10배 가볍게 만든 검색기

출처

• 원문: https://arxiv.org/abs/2004.12832 (2020 · Omar Khattab 외)