다중 벡터 임베딩 모델을 위한 토큰 풀링 인식 학습

다중 벡터 검색에서 토큰 풀링 이해하기

ColBERT와 같은 다중 벡터 임베딩 모델은 문서를 단일 벡터가 아닌 문맥적 토큰 벡터의 집합으로 표현합니다. 이 접근 방식은 쿼리 토큰을 가장 관련성 높은 문서 토큰과 매칭하는 MaxSim 점수 계산을 통해 높은 표현력의 검색을 가능하게 합니다. 그러나 이러한 세분성은 수백만 개의 문서를 포함하는 말뭉치가 수억 개의 벡터를 생성할 수 있어 상당한 저장 공간 요구 사항을 초래합니다.

토큰 풀링은 유사한 토큰 벡터를 더 적은 수의 대표 벡터로 병합하여 이를 완화합니다. 표준 풀링은 모든 모델에 적용할 수 있지만, 인코더가 원래 압축된 표현을 사용하여 검색을 수행하도록 훈련되지 않았기 때문에 종종 정보 손실로 이어집니다. 풀링 인식 학습(Pooling-aware training)은 모델이 벡터를 병합할 때 검색 정확도를 유지하도록 명시적으로 학습시킴으로써 이 문제를 해결합니다.

풀링 인식 학습의 메커니즘

풀링 인식 학습은 전체 표현과 압축된 표현 모두에 검색 목표를 동시에 적용함으로써 Matryoshka 임베딩과 유사하게 작동합니다. 학습 중에 각 배치는 두 가지 손실을 계산합니다:

  1. 비풀링 손실(Unpooled Loss): 원래의 전체 문서 임베딩 세트를 사용한 검색.
  2. 풀링 손실(Pooled Loss): 문서를 목표 벡터 수로 압축한 후의 검색 (예: pool_factor를 4로 사용).

총 손실은 이 두 구성 요소의 평균입니다. 이는 인코더가 중요한 정보가 클러스터링 과정에서 살아남을 수 있도록 토큰 임베딩을 구조화하도록 강제합니다. HierarchicalTokenPooling 모듈은 Ward 계층적 클러스터링을 사용하여 벡터를 그룹화하며, 학습 과정은 클러스터 할당의 이산적인 특성에도 불구하고 그래디언트를 원래 토큰 임베딩으로 다시 전달하기 위해 straight-through estimator를 사용합니다.

구현 및 성능

개발자는 Sentence Transformers 라이브러리 내의 MultiVectorTokenPoolingLoss 래퍼를 사용하여 이를 구현할 수 있습니다. 기존 검색 손실을 래핑함으로써 이 접근 방식은 다양한 풀링 변형에 대해 별도의 인코더 패스를 요구하지 않고 효율적인 학습을 가능하게 합니다.

주요 성능 향상

NanoBEIR 벤치마크에서 ModernBERT-base를 사용한 평가에 따르면, 풀링 인식 학습은 동일한 저장 공간 예산에서 표준 모델보다 훨씬 뛰어난 성능을 보입니다:

  • 유지율: pool_factor=4(벡터 4배 감소)에서 풀링 인식 학습은 비풀링 nDCG@10 점수의 97.8%를 유지하며, 일반적인 학습 모델의 94.6%와 비교됩니다.
  • 유연성: 이 학습의 이점은 학습 중에 사용된 특정 pool_factor를 넘어 확장되며, 다양한 고정 캡 및 적응형 풀링 임계값 전반에서 일관된 개선을 보여줍니다.
  • 적응형 보정: 사용자는 거리 임계값을 보정하여 목표 평균 벡터 수를 달성할 수 있으며, 이를 통해 모델이 정보 밀도가 높은 문서에 더 많은 벡터를 동적으로 할당할 수 있습니다.

권장 압축 전략

인덱스 크기가 주요 제약 사항인 애플리케이션의 경우, 다음의 계층적 접근 방식을 권장합니다:

  1. 양자화(Quantization): Fast-PLAID(2비트 또는 4비트 잔차 양자화)와 같은 기술을 적용하여 좌표당 저장 공간을 줄입니다.
  2. 토큰 풀링(Token Pooling): 풀링 인식 미세 조정을 사용하여 총 문서 벡터 수를 줄입니다.
  3. Matryoshka 임베딩: Matryoshka 토큰 임베딩을 활용하여 개별 벡터의 차원을 줄입니다.

이러한 기술은 상호 보완적이며 MultiVectorMatryoshkaLoss 및 MultiVectorTokenPoolingLoss와 같은 여러 손실 함수를 래핑하여 벡터 수와 차원을 동시에 최적화할 수 있습니다.

Sources