ColBERT

ColBERT (BERT 기반 문맥화 지연 상호작용)는 쿼리와 각 문서를 토큰마다 하나의 임베딩으로 표현하고 지연 상호작용으로 문서 점수를 계산하는 신경 검색 모델입니다. 각 쿼리 토큰은 가장 유사한 문서 토큰을 찾고, MaxSim이라고 하는 이 최대 유사도들을 합산해 관련성 점수를 구합니다. Stanford University의 Omar Khattab과 Matei Zaharia는 SIGIR 2020 논문 ColBERT: 문맥화된 지연 상호작용을 통한 효율적이고 효과적인 구절 검색에서 ColBERT를 소개했습니다. 문서 토큰 임베딩은 오프라인에서 한 번만 계산되므로, 쿼리 시에는 후보 문서마다 전체 BERT 처리를 하는 대신 저렴한 유사도 계산만 수행합니다.

지연 상호작용의 위치#

신경 검색기는 쿼리와 문서가 언제 만나느냐에 따라 다릅니다. 이중 인코더는 각 텍스트를 단일 벡터로 압축하는데, 이는 벡터 검색의 표준 방식입니다. 문서 벡터를 미리 계산하므로 검색은 빠르지만, 단일 벡터는 개별 용어 간 차이를 흐릿하게 만듭니다. 교차 인코더는 쿼리와 문서를 함께 트랜스포머에 입력하므로 어텐션이 모든 토큰 쌍을 비교합니다. 정확도는 가장 높지만 쿼리와 문서의 각 쌍마다 한 번씩 실행해야 하므로 짧은 후보 목록의 재순위화에만 사용할 수 있습니다.

ColBERT는 이중 인코더처럼 쿼리와 문서를 따로 인코딩하지만, 토큰마다 128차원 벡터를 하나씩 유지하고 마지막 MaxSim 단계에서만 두 쪽이 상호작용하도록 합니다. 쿼리는 BERT의 마스크 토큰으로 고정 길이 32까지 패딩되며, 이를 통해 모델은 관련 용어로 쿼리를 확장할 추가 슬롯을 얻습니다. MS MARCO와 TREC CAR 구절 검색 벤치마크에서 ColBERT는 BERT 기반 순위 모델과 경쟁할 만한 성능을 보이면서도 실행 속도는 두 자릿수 배 빠르고 쿼리당 FLOPs는 네 자릿수 배 적었습니다.

MaxSim 점수 계산#

MaxSim 연산자는 몇 줄의 PyTorch 코드로 구현할 수 있습니다. 이 예제는 작은 BERT 계열 인코더의 문맥화된 토큰 임베딩에 연산을 적용합니다.

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
query = "who wrote the novel Dune"
docs = [
    "Frank Herbert published the science fiction novel Dune in 1965.",
    "A dune is a hill of sand shaped by the wind.",
]

# One L2-normalized contextual embedding per query token
q = F.normalize(encoder.encode(query, output_value="token_embeddings"), dim=-1)
for doc in docs:
    d = F.normalize(encoder.encode(doc, output_value="token_embeddings"), dim=-1)

    # MaxSim: each query token keeps its best-matching document token, and the maxima are summed
    score = (q @ d.T).max(dim=1).values.sum()
    print(f"{score.item():.2f}  {doc}")

소설 구절의 점수는 5.90이고 모래 언덕 구절의 점수는 4.33입니다. 두 구절 모두 "dune"을 포함하므로 두 경우 모두 강하게 일치하지만, 첫 번째 구절에서는 쿼리 토큰 "wrote"가 "published"와 코사인 유사도 0.72로 일치하고 "novel"은 자기 자신과 일치합니다. 반면 두 번째 구절에서 "wrote"의 최적 일치 점수는 0.29에 불과합니다.

ColBERT와 관련 용어 비교#

• BERT 자체: ColBERT가 검색용으로 미세 조정하는 사전 학습 인코더이며, 각 토큰을 128차원으로 투영하는 선형 레이어를 추가합니다.
• 지연 청킹: 텍스트를 분할하기 전에 문맥을 포함해 인코딩한 다음 토큰을 평균해 청크마다 벡터 하나를 만드는 방식입니다. ColBERT는 모든 토큰 벡터를 유지합니다.
• 시맨틱 검색: ColBERT가 수행하는 작업으로, 공유 키워드가 아니라 의미를 기준으로 구절의 순위를 매깁니다.
• 하이브리드 검색: 키워드 검색 결과와 단일 벡터 검색 결과를 결합해 정확한 용어 일치와 의미 기반 검색을 함께 수행합니다. ColBERT의 토큰 수준 매칭은 하나의 신경 모델 안에서 이러한 정확성의 일부를 포착합니다.

이미지에서의 지연 상호작용#

지연 상호작용은 이미지에도 적용됩니다. ColPali(2024)는 비전-언어 모델로 각 문서 페이지의 스크린샷을 패치 벡터 그리드로 임베딩하고, MaxSim으로 쿼리 토큰과 비교해 PDF 대상 멀티모달 RAG에 지연 상호작용을 적용합니다. 개방 어휘 탐지도 관련된 이유로 이미지 특징을 세밀하게 유지합니다. Ultralytics YOLOE는 "double-decker bus"와 같은 프롬프트 클래스마다 텍스트 임베딩을 저장하고, 이를 특징 맵의 모든 위치에 있는 영역 특징과 비교합니다. 따라서 짧은 프롬프트가 이미지 전체를 요약하는 단일 벡터 대신 작은 객체 하나와 일치할 수 있습니다.

ColBERTv2와 PLAID#

모든 토큰의 벡터를 저장하면 ColBERT 인덱스는 단일 벡터 인덱스보다 대략 한 자릿수 배 더 커집니다. 이 비용은 ColBERT가 검색 증강 생성(RAG) 파이프라인의 1단계 검색기나 재순위화기로 사용될 때 중요합니다. Stanford ColBERT 저장소는 두 가지 해결책을 구현합니다. PLAID는 전체 MaxSim을 실행하기 전에 클러스터 중심점을 기준으로 후보의 점수를 계산하는 검색 엔진입니다. ColBERTv2(NAACL 2022)는 각 토큰 벡터를 가장 가까운 중심점의 ID와 차원당 1비트 또는 2비트로 양자화한 잔차로 저장하고, 교차 인코더에서 증류한 잡음 제거 감독을 사용해 학습합니다. 잔차 압축으로 지연 상호작용의 저장 공간을 6~10배 줄였습니다.