Zzong's Notes

Jina-ColBERT-v2

19 min read

Jina-ColBERT-v2: A General Purpose Multilingual Late Interaction Retriever

이 논문은 ColBERT와 같은 다중 벡터(multi-vector) 밀집 리트리버 모델의 효용성을 강조하며 출발합니다. ColBERT는 쿼리와 문서의 각 토큰 간 상호작용을 ‘지연 상호작용(late interaction)’ 방식으로 계산해, 검색 품질을 크게 높이면서도 기존 단일 벡터 기반 리트리버와 비슷한 추론 효율성을 달성합니다.

하지만 기존 ColBERTv2 모델은 영어 데이터인 MSMARCO로만 학습되어 다국어 검색에는 한계가 있었습니다. 이를 보완하려는 ColBERT-XM, BGE-M3 등 여러 시도들이 있었으나, 추가 파라미터가 필요하거나 임베딩 크기가 지나치게 커서 현실적인 활용에 제약이 있었습니다.

Jina-ColBERT-v2의 주요 특징은 다음과 같습니다.

  1. 다양한 데이터 기반 학습 대규모 대조 학습(contrastive tuning)과 소규모 지도 증류(supervised distillation)를 결합한 2단계 학습 방식을 도입했습니다.

  2. 뛰어난 범용 다국어 성능 고품질 및 저품질 언어 자원(기계 번역된 데이터 포함)을 폭넓게 활용해, 미학습 언어에 대해서도 우수한 검색 성능을 보입니다.

  3. 효율적인 추론 마트료시카 표현 손실(Matryoshka Representation Loss, MRL)을 활용하여 여러 크기의 임베딩을 동시에 학습할 수 있도록 했습니다. 이를 통해 사용자는 임베딩 차원을 필요에 따라 선택하여 저장 공간을 절약하면서도 성능 저하를 최소화할 수 있습니다.

아키텍처

  • 기반 언어 모델 개선 Jina-ColBERT-v2는 다국어 처리에 강점이 있는 XLM-RoBERTa를 기반으로 사용합니다. 여기에 플래시 어텐션(Flash Attention) 기법과 기존 절대적 위치 임베딩 대신 회전 위치 임베딩(Rotary Positional Embeddings, RoPE)을 적용해 효율성과 성능을 모두 높였습니다.

  • 다중 선형 헤드(Multiple Linear Heads) 원래 ColBERT에서는 언어 모델의 은닉 차원(예: 768)을 더 낮은 차원(예: 128)으로 줄이는 선형 헤드를 사용합니다. Jina-ColBERT-v2는 마트료시카 표현 손실(MRL)을 이용해 {64, 96, 128, 256, 512, 768} 등 다양한 차원의 선형 헤드를 동시에 학습합니다. 이로 인해 사용자는 저장 공간 상황이나 목적에 따라 최적의 임베딩 크기를 선택할 수 있으며, 예를 들어 차원을 128에서 64로 줄여도 nDCG@10 기준 성능 하락이 약 1.59%에 불과할 정도로 효율적입니다.

B) 훈련 방식

Jina-ColBERT-v2는 세 단계의 훈련 과정을 거칩니다.

  1. 쌍(pair) 데이터 학습 다양한 언어로 구성된 방대한 텍스트 쌍 데이터를 사용해 모델을 먼저 학습시킵니다. 총 약 4억5천만 개의 쌍으로 구성되며 이 중 절반은 영어이고 나머지는 주요 29개 언어입니다. 이 단계에서는 대조 손실 함수를 통해 의미적 유사성을 효과적으로 학습하게 됩니다.

  2. 트리플렛(triplet) 데이터 미세 조정 쿼리–긍정 문서–부정 문서(hard negative)의 트리플렛 구조 데이터를 활용해 추가 미세조정을 진행합니다. MSMARCO나 MIRACL처럼 고품질 인공 주석 데이터뿐 아니라 기계 번역 및 법률 분야 등 특정 도메인 합성 데이터도 포함됩니다. 이 과정에서는 뛰어난 다국어 교차 인코더(cross-encoder)를 교사 모델로 삼아 KL 발산 손실(KL divergence loss)을 통해 그 지식을 학생 모델(Jina-ColBERT-v2)에 증류합니다.

  3. 쿼리 증강 어텐션(Query Augmentation Attention) ColBERT 특유의 [MASK] 토큰 추가 방식으로 쿼리 표현력을 확장하는데, 본 논문에서는 어텐션 마스크를 조정함으로써 실제 쿼리 토큰들이 추가 [MASK] 토큰에도 집중할 수 있게 설계했습니다. 이를 통해 특히 다국어 검색(MIRACL)에서 더욱 뛰어난 성능 향상을 확인했습니다.

C) 결과

  • 영어 기준 성능: BEIR 벤치마크에서 기존 ColBERTv2 대비 평균적으로 약 6.6% 높은 성능을 보였습니다.
  • 다국어 기준 성능: mMARCO 벤치마크에서 ColBERT-XM보다 모든 언어에서 우수했으며, 심지어 훈련되지 않은 이탈리아어나 베트남어·힌디어와 같은 언어에서도 강력한 결과를 나타냈습니다.
  • MIRACL 벤치마크에서는 BGE-M3보다는 약간 낮은 점수를 기록했지만 이는 BGE-M3가 Jina-ColBERT-v2(128차원)에 비해 토큰 임베딩 크기가 무려 여덟 배나 큰 점을 고려해야 합니다.

D) QnAs

D.1) 왜 Pair 데이터로 먼저 학습하고, 그 다음 Triplet 데이터로 학습했나요?

이는 “일반화에서 특수화(Generalization to Specialization)” 전략을 따르는 매우 효과적인 2단계 학습 방법입니다.

  • 1단계: Pair 데이터 학습 (대조 학습)
    • 목표: 모델에게 광범위하고 일반적인 의미론적 이해를 가르치는 단계입니다. 즉, “어떤 두 문장이 의미적으로 비슷한가?”를 배우게 됩니다.
    • 데이터: 4억 5천만 개의 방대한 텍스트 쌍(semantically related pairs)을 사용합니다. 이 데이터는 특정 “질의-응답” 형식이 아니라, 단순히 위키피디아의 인접 문단이나 유사한 주제의 글처럼 의미가 연결된 텍스트들입니다.
    • 효과: 이 단계를 통해 모델은 다양한 언어와 주제에 걸쳐 풍부한 단어와 문장의 표현(representation)을 학습합니다. 검색이라는 특정 작업에 얽매이지 않고, 언어 자체의 의미적 구조를 탄탄하게 만드는 사전 훈련(pre-training) 과정과 같습니다.
  • 2단계: Triplet 데이터 학습 (지도 증류)
    • 목표: 일반적인 의미 이해 능력을 바탕으로, 검색(retrieval)이라는 특정 작업에 모델을 미세 조정(fine-tuning)하는 단계입니다. 즉, “주어진 쿼리에 대해, 긍정 문서를 부정 문서보다 더 가깝게 배치하는 방법”을 집중적으로 배웁니다.
    • 데이터: (쿼리, 긍정 문서, 부정 문서)로 구성된 트리플렛(triplet)을 사용합니다. 이 데이터는 검색 작업에 매우 특화되어 있습니다.
    • 효과: 이 단계를 통해 모델은 단순히 ‘의미가 비슷하다’를 넘어, 쿼리의 의도에 맞는 문서를 정확히 순위 매기는(ranking) 능력을 갖추게 됩니다. 특히 ‘어려운 부정 예시(hard negatives)‘를 사용함으로써, 쿼리와 주제는 비슷하지만 정답이 아닌 까다로운 문서들을 구별하는 능력이 크게 향상됩니다.

요약: 거대한 Pair 데이터로 언어의 기초 체력을 기르고(일반화), 더 작지만 고품질인 Triplet 데이터로 검색 순위 결정이라는 전문 기술을 연마(특수화)하는 것입니다. 이 방식은 모델이 특정 데이터셋에 과적합(overfitting)되는 것을 방지하고, 범용 성능을 높이는 데 매우 효과적입니다.

D.2) KL Divergence Loss 학습은 구체적으로 어떻게 한 건가요? 그리고 교사 모델은 무엇을 썼나요?

이는 지식 증류(Knowledge Distillation) 기법을 사용한 것입니다.

  • 교사 모델 (Teacher Model): 논문에서는 “독점적인 다국어 교차 인코더(a proprietary multilingual cross-encoder)” 를 사용했다고 밝혔습니다. 논문에서는 구체적으로 어떤 모델을 교사 모델로 사용했는지 밝히지 않았습니다.
  • 교차 인코더(Cross-encoder)는 쿼리와 문서를 한 번에 입력받아 관련성을 매우 정확하게 판단하지만, 속도가 느려 실제 검색 시스템에는 바로 쓰기 어려운 모델입니다. 이 강력하지만 느린 모델이 ‘교사’ 역할을 합니다.
  • KL Divergence를 이용한 학습 과정:
    1. 훈련 데이터 트리플렛(Query, Positive Doc, Negative Doc)을 준비합니다.
    2. 교사 모델이 (Query, Positive Doc) 쌍의 관련성 점수와 (Query, Negative Doc) 쌍의 관련성 점수를 계산합니다. 이 두 점수에 Softmax 함수를 적용하여 **확률 분포 (Teacher’s Distribution)**를 만듭니다. 예를 들어, [0.95, 0.05] 와 같은 분포가 나옵니다.
    3. 학생 모델 (Jina-ColBERT-v2) 역시 자체적인 MaxSim 방식으로 두 쌍의 관련성 점수를 계산하고, 똑같이 Softmax를 적용해 **자체 확률 분포 (Student’s Distribution)**를 만듭니다. 예를 들어, 처음에는 [0.6, 0.4] 와 같은 미숙한 분포를 생성할 수 있습니다.
    4. KL Divergence Loss는 이 두 확률 분포, 즉 교사의 분포([0.95, 0.05])와 학생의 분포([0.6, 0.4]) 사이의 차이를 계산합니다.
    5. 모델은 이 Loss를 최소화하는 방향으로 자신의 가중치를 업데이트합니다. 즉, 학생 모델(Jina-ColBERT-v2)이 교사 모델의 정교한 판단 방식을 모방하도록 학습하는 것입니다.

결론적으로, Jina-ColBERT-v2는 속도는 빠르면서도, 느리지만 정확한 교사 모델의 ‘판단력’을 그대로 전수받게 됩니다.

D.3) Negative 샘플은 어떻게 채택한 건가요?

논문에서는 **“어려운 부정 샘플(hard negatives)“**을 사용했다고 강조합니다. 이는 검색 성능에 매우 중요합니다. 어려운 부정 샘플이란, 쿼리와 키워드는 겹치지만 실제로는 정답이 아닌 문서를 의미합니다.

이러한 샘플은 주로 다음과 같은 방식으로 수집됩니다.

  1. 기존 벤치마크 데이터셋 활용: MSMARCO, MIRACL과 같은 데이터셋들은 이미 잘 구축된 어려운 부정 샘플을 포함하고 있습니다. 보통 BM25와 같은 검색 시스템으로 상위 K개의 문서를 가져온 뒤, 그중에서 정답이 아닌 문서들을 부정 샘플로 활용하는 방식입니다.
  2. 모델 기반 마이닝: 다른 검색 모델(주로 교차 인코더)을 사용해 쿼리와 의미적으로 매우 유사하지만 정답이 아닌 문서들을 직접 채굴하여 부정 샘플로 사용하기도 합니다.

Jina-ColBERT-v2는 이렇게 다양한 출처에서 얻은, 품질 좋은 어려운 부정 샘플을 포함한 트리플렛 데이터로 학습하여 까다로운 검색 환경에서의 변별력을 높였습니다.

D.4) 쿼리 증강 어텐션(Query Augmentation Attention) 마스크는 어떻게 수정한 건가요?

이는 ColBERT의 고유한 특징인 쿼리 증강(Query Augmentation) 을 개선한 것입니다.

  • 기존 방식:
    • 원래 쿼리 뒤에 [MASK] 토큰을 여러 개 추가합니다. (예: "what is colbert" [MASK] [MASK] ...)
    • 이 [MASK] 토큰들은 쿼리의 전반적인 의도나 확장된 의미를 학습하는 역할을 합니다.
    • 하지만 표준 어텐션 메커니즘에서는, 앞에 있는 실제 쿼리 토큰들(what, is, colbert)이 뒤에 있는 [MASK] 토큰들을 볼 수 없었습니다(attend to 할 수 없었습니다). 오직 [MASK] 토큰들만 앞의 쿼리 토큰들을 볼 수 있었죠.
  • Jina-ColBERT-v2의 수정 방식:
    • 어텐션 마스크(Attention Mask)를 수정하여, 실제 쿼리 토큰들(what, is, colbert)이 뒤따라오는 [MASK] 토큰들을 볼 수 있도록 허용했습니다.
    • 이유: 이렇게 함으로써 [MASK] 토큰이 학습한 ‘추상적이고 확장된 쿼리 의미’가 실제 쿼리 토큰의 표현(representation)에 직접적으로 영향을 미칠 수 있게 됩니다. 즉, what이라는 단어의 임베딩조차도 [MASK] 토큰의 정보를 받아 더 풍부하고 문맥에 맞는 의미를 갖게 됩니다.

D.4.1) 원래 BERT 류는 Bi-directional이라 모든 토큰끼리 상호작용하지 않나요?

질문자님께서 “BERT는 원래 양방향”이라고 이해하신 것은 정확합니다. 하지만 기존 ColBERT는 쿼리 증강이라는 특수 목적을 위해 이 양방향성에 의도적인 제약을 걸어두었습니다. Jina-ColBERT-v2는 이 제약을 풀어 쿼리 토큰과 증강용 [MASK] 토큰 간의 상호작용을 더욱 활발하게 만들어 성능을 개선한 것입니다.

BERT 계열 모델은 원래 모든 토큰이 서로를 참조할 수 있는 완전한 양방향 어텐션 구조를 가지고 있습니다. 예를 들어, [CLS] what is colbert [SEP]와 같은 입력에서는 what과 colbert가 서로의 정보를 자유롭게 주고받을 수 있습니다.

하지만 ColBERT 모델에서는 쿼리 증강(Query Augmentation)을 위해 특별한 어텐션 마스크를 적용합니다. 기존 ColBERT(v1/v2)는 다음과 같은 방식으로 동작했습니다.

  • 입력 시퀀스 예시: [Q] what is colbert [Q_END] [MASK] [MASK]...
  • [MASK] 토큰은 실제 쿼리 토큰들의 문맥을 파악하여 확장된 의미를 학습하는 역할을 합니다.
  • 이때 실제 쿼리 토큰들은 [MASK] 토큰을 참조하지 못하도록 제한되어 있었습니다. 즉, [MASK] → (what, is, colbert) 방향으로는 정보 흐름이 가능하지만, (what, is, colbert) → [MASK] 방향은 차단되어 있었던 것입니다. 이는 아직 학습되지 않은 [MASK]의 정보가 쿼리 본연의 표현에 영향을 주지 않게 하려는 의도로 볼 수 있습니다.

Jina-ColBERT-v2에서는 이 제한이 해제되었습니다. 어텐션 마스크 설정을 수정하여 실제 쿼리 토큰들이 [MASK] 토큰도 자유롭게 참조할 수 있게 만들었습니다. 이제 [MASK]가 학습한 추상적이고 확장된 개념을 다시 실제 쿼리 토큰들이 받아들여 자신의 표현에 반영할 수 있게 되어, 보다 풍부하고 강력한 상호작용이 가능해졌습니다.

정리하자면, 기본 BERT의 양방향성은 맞는 설명입니다. 다만 ColBERT에서는 특수한 목적(쿼리 증강)을 위해 일시적으로 이 양방향성에 제한을 두었고, Jina-ColBERT-v2에서 그 제약이 해소되어 더욱 강력한 표현력을 얻게 된 것입니다.

링크된 언급

1
Retrieval

...teraction을 기준점으로 두고, serving 최적화는 PLAID, WARP An Efficient Engine for Multi Vector Retrieval, MUVERA, Jina ColBERT v2를 함께 보면 좋다. D) Ranking And Applications Retrieval 이후에는 Cascade Ranking System 안에서 pre ranking, ranking...