다중 벡터 임베딩 모델은 각 쿼리 토큰이 문서 내에서 일치하는 항목을 찾을 기회를 갖게 합니다. 이는 검색에 매우 표현력이 뛰어나지만, 이러한 문서 벡터들이 빠르게 쌓입니다. 수백만 개의 문서를 가진 코퍼스는 수억 개의 벡터를 포함할 수 있습니다!
토큰 풀링(Token pooling)은 유사한 토큰 벡터들을 병합하여 이러한 표현을 더 작게 만듭니다. 하지만 제가 압축하는 정도가 강해질수록, 유용한 무언가를 잃어버리기 쉬워집니다. 풀링 인식 학습(Pooling-aware training)은 인코더가 자신의 벡터들이 병합될 때 중요한 정보를 보존하도록 가르치며, 동일한 저장 예산으로 더 나은 검색을 수행할 수 있도록 돕습니다.
본 게시물에서는 풀링이 어떻게 작동하는지 설명하고, Sentence Transformers를 사용하여 압축된 검색 모델을 학습시키며, 그 결과로 나타나는 품질과 저장 공간의 트레이드오프(tradeoff)를 탐구하겠습니다. pool_factor=4를 사용했을 때, 제가 일반적으로 훈련한 모델은 풀링되지 않은 NanoBEIR nDCG@10의 94.6%를 유지했지만, 토큰 풀링 인식 학습을 적용하자 동일한 저장 예산에서 그 유지율이 97.8%로 높아졌습니다! 어떻게 되었는지 살펴보겠습니다.
본 게시물은 Stefan Josef의 Learn to Pool 및 LightOn의 풀링 정규화(pooling regularization) 연구에서 다루어진 풀링 인식 학습을 기반으로 하며, 여기서는 타겟팅된 학습과 적응형 예산(adaptive budgets)도 탐구합니다. 저는 여기서 방법론 이해에 초점을 맞추고 짧고 재현 가능한 Sentence Transformers 예시를 통해 이를 실제로 적용해 보겠습니다.
풀링 인식 학습은 동일한 벡터 예산에서 검색 품질을 향상시킵니다. 자세한 내용은 아래를 참고하세요.
- 다중 벡터 검색이 많은 벡터를 필요로 하는 이유
- 토큰 풀링(Token Pooling)
- 풀링 인식 학습(Token Pooling-Aware Training)
- 훈련된 모델 사용하기
- 평가 결과
- 압축 기법 결합
- 참고 문헌
- 코드 및 모델
단일 벡터 임베딩 모델은 전체 문서를 하나의 벡터로 압축합니다. 반면, ColBERT와 같은 다중 벡터 모델은 일련의 문맥적 토큰 벡터들을 유지하여, 쿼리의 서로 다른 부분이 문서의 서로 다른 부분과 일치할 수 있도록 합니다.
예를 들어, 배터리 수명이 긴 노트북에 대한 질의는 문서의 한 부분에서 노트북에 관한 증거를 찾고 다른 부분에서 배터리 내구성을 찾는 식일 수 있습니다. MaxSim 스코어링을 사용하면 각 쿼리 벡터가 가장 잘 일치하는 문서 벡터를 선택하고, 이 점수들을 합산합니다:
Fast-PLAID가 지원하는 2비트 또는 4비트 잔차 양자화(residual quantization)를 사용하더라도, 문서당 수백 개의 벡터를 저장하는 것은 많은 용량을 차지합니다. 잔차 좌표당 4비트를 기준으로 할 때, 128 차원의 200개 벡터는 문서당 12.8 KB의 양자화된 잔차(quantized residuals)가 필요하며, 백만 개의 문서를 처리하면 총 12.8 GB에 달합니다. 각 문서를 50개의 벡터로 풀링(pooling)하면 차원과 정밀도를 유지하면서 이를 3.2 GB로 줄일 수 있습니다.
유용한 관찰은 많은 문서 벡터들이 유사할 수 있다는 것입니다. 여러 토큰이 밀접하게 관련된 개념을 설명할 수 있으며, 이는 이들을 함께 표현할 기회를 만듭니다.
HierarchicalTokenPooling은 Ward 계층적 클러스터링(Ward hierarchical clustering)을 사용하여 유사한 토큰 임베딩(token embeddings)들을 그룹화합니다. 각 토큰 벡터당 하나의 그룹에서 시작하여, 이웃하는 두 그룹이 자신들의 그룹 내에서 얼마나 퍼져 있는지에 가장 작은 증가를 일으키는 방식으로 반복적으로 병합합니다. 이러한 퍼짐 정도(spread)는 각 벡터와 해당 그룹의 평균 간 제곱 거리 합으로 측정되므로, Ward 방법은 유사한 벡터들로 구성된 응집력 있는 그룹을 선호합니다.
이러한 병합 과정은 점진적으로 더 큰 그룹들의 계층 구조를 만듭니다. 이를 통해 우리는 예산에 맞는 압축 수준을 선택할 수 있습니다. 그런 다음 각 그룹을 그 평균값(mean)으로 표현하고, MaxSim 스코어링 준비가 되도록 결과 벡터들을 다시 정규화합니다.
압축률을 문서 길이에 비례하도록 설정할 수 있습니다:
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
pool_factor=4를 사용하면 문서가 원래 벡터의 약 4분의 1을 유지하게 되므로, 더 긴 문서는 비례적으로 더 많은 용량을 가질 수 있습니다. 기본 설정으로는 첫 번째 토큰은 병합으로부터 보호되며, 풀링(pooling)은 오직 문서 임베딩(document embeddings)에만 적용되고 쿼리 임베딩(query embeddings)은 변경되지 않습니다.
계층적 토큰 풀링(HierarchicalTokenPooling)
세 가지 대안적인 제어 기능을 제공합니다:
| 설정 | 표현하는 내용 |
|---|
max_tokens=16 | 보호 토큰을 포함하여 문서당 최대 16개의 벡터를 사용합니다. |
pool_factor=4 | 각 문서의 원래 벡터 중 약 4분의 1에 해당합니다. |
distance_threshold=... | 공유된 Ward 병합 거리 임계값으로, 임베딩 기하학에 따라 벡터 개수가 달라지도록 허용합니다. |
이러한 차이점은 문서 길이가 가변적일 때 중요합니다. pool_factor=4를 사용하면 더 긴 문서가 약 4분의 1의 벡터를 유지할 수 있어, 내용에 따라 표현(representation)을 확장할 수 있습니다. distance_threshold를 사용하면 더 많은 고유한 임베딩을 보존해야 할 때 벡터 개수가 증가할 수 있습니다.
Adaptive pooling은 공유된 거리 임계값을 사용하여 유사한 임베딩을 더 공격적으로 병합하는 동시에, 내용이 뚜렷한 문서의 벡터는 더 많이 보존합니다. 이는 풀링 인식 목적 함수(pooling-aware objective) 없이 학습된 모델을 포함하여 모든 다중 벡터 모델과 함께 사용할 수 있습니다.
임계값을 수동으로 선택하기보다는, 목표 평균 벡터 개수에 맞춰 이를 보정(calibrate)할 수 있습니다. calibration_documents에 있는 코퍼스 대표 샘플이 주어지면:
from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
model = MultiVectorEncoder("lightonai/LateOn")
...
보정 과정에는 문서 임베딩만 필요하며, 인덱싱 시 재사용할 수 있는 임계값을 반환합니다. 목표는 평균이므로 개별 문서는 16개보다 많거나 적은 벡터를 받을 수 있습니다. 달성된 평균값은 다른 코퍼스에서 달라질 수도 있습니다.
풀링은 일반적인 다중 벡터 모델에서도 이미 유용할 수 있습니다. 어려움은 문서가 검색에 중요한 차이점을 포함하고 있지만, 이 벡터들을 병합하면 그 차이점을 복구하기 어렵게 만드는 경우에 발생합니다. 원래의 학습 목적 함수는 모델에게 이러한 압축된 표현(compressed representations)으로 작동하도록 요구한 적이 없습니다. 하지만, 모델에게 정확히 그렇게 하도록 가르칠 수 있습니다.
Matryoshka 임베딩을 사용해 본 적이 있다면, 훈련 원리가 익숙하게 느껴질 것입니다. Matryoshka는 여러 임베딩 너비에서 검색 목표(retrieval objective)를 적용하여 작은 접두사들이 유용하게 유지되도록 가르칩니다. 풀링 인식 학습(Pooling-aware training)은 더 적은 토큰 벡터를 가진 표현에 이 목표를 적용합니다.
제 예시의 경우, 각 훈련 배치는 두 가지 손실(losses)을 기여합니다:
- 원본 문서 임베딩을 사용한 검색.
- 문서를 원래 벡터의 약 4분의 1로 압축한 후의 검색.
저는 이들을 평균화합니다:
첫 번째 항은 모델이 전체 표현(full representations)으로 검색하도록 가르칩니다. 두 번째 항은 유용한 일치 항목들이 풀링을 겪고도 살아남도록 토큰 임베딩을 배열하는 방법을 가르칩니다. 두 목표 모두 동일한 인코더를 업데이트하여, 압축 전후의 검색에 대한 피드백을 제공합니다.
Ward 클러스터링은 어떤 토큰들이 함께 속하는지에 대해 이산적인 결정을 내립니다. 훈련 중, HierarchicalTokenPooling은 클러스터 할당을 미분하지 않으면서 풀링된 벡터로부터 그 멤버 토큰들로 기울기(gradients)를 전달하기 위해 straight-through estimator를 사용합니다. 그러면 인코더는 풀링된 표현으로 발생한 검색 오류로부터 학습할 수 있습니다.
MultiVectorTokenPoolingLoss는 기존의 다중 벡터 검색 손실(multi-vector retrieval loss)을 감싸서 사용자가 선택한 풀러(poolers)를 사용하여 적용합니다:
from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import (
CachedMultiVectorMultipleNegativesRankingLoss,
...
CachedMultiVectorMultipleNegativesRankingLoss는 검색 목표를 제공합니다. 이는 각 쿼리가 배치 내의 부정적인 문서들보다 자신의 관련 문서를 더 높은 점수로 평가하도록 가르칩니다. 기본적으로 모델의 유사성 설정(이 예시에서는 MaxSim)과 커널 설정을 사용하여 model.similarity를 사용합니다. similarity_fn_name="meanmaxsim"으로 구성된 모델은 훈련 중 MeanMaxSim을 사용할 것입니다.
그래디언트 캐싱(Gradient caching)은 인코더가 더 큰 규모의 배치 내 네거티브 샘플을 유지하면서도 작은 미니배치로 처리할 수 있게 합니다. 제 학습 스크립트는 훈련 배치 크기 256에 대해 미니배치 크기 64를 사용합니다.
MultiVectorTokenPoolingLoss는 압축 목적 함수(compression objective)를 추가합니다. 여기서 None은 한 변형이 원본 임베딩을 사용하는 것을 의미하며, pooling은 factor-4 변형을 제공합니다. 두 가중치는 손실 값을 평균화합니다. 풀링 변형들은 계산된 토큰 임베딩을 재사용하므로, 추가적인 풀러는 해당 변형에 대해 별도의 인코더 패스가 필요하지 않으면서 클러스터링 및 스코어링 작업을 추가합니다.
풀러는 또한 입력 임베딩을 통해 그래디언트를 보존하는 사용자 정의 BaseTokenPooling 구현체일 수 있습니다. 원하는 배포 표현(representation)에 맞게 훈련 목적 함수를 조정할 수 있습니다.
factor-4 풀러를 선택한 배포 구성으로 교체하거나, 여러 풀러를 포함하고 각 손실에 대한 가중치를 지정할 수 있습니다.
완전한 훈련 스크립트에는 데이터셋, 트레이너(trainer), 검증 손실, 그리고 훈련 중 소규모 검색 평가가 포함되어 있습니다. 저장소에서 체크아웃하여 다음 명령어로 두 변형을 로컬에서 실행합니다:
uv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py --baseline
uv run examples/multi_vector_encoder/training/token_pooling/training_cached_contrastive.py
이 스크립트는 훈련 및 커널 종속성을 선언하므로, uv가 이를 자동으로 설치합니다. 호환 가능한 커널은 자동으로 활성화됩니다.
Hugging Face Jobs를 사용하여 클라우드 GPU에서도 훈련할 수 있습니다. CLI를 설치하고 Job을 실행하며 모델을 업로드할 수 있는 토큰으로 로그인하세요:
pip install -U huggingface_hub
hf auth login
hf jobs uv run \
...
이 명령어는 A10G GPU에서 4시간의 시간 제한으로 풀링 인식 변형을 실행합니다. HF_TOKEN을 시크릿(secret)으로 전달하면, 훈련이 완료될 때 스크립트가 학습된 모델을 사용자의 Hub 계정에 업로드할 수 있습니다. 일반 손실 기준선(ordinary-loss baseline)도 실행하려면, 스크립트 URL 뒤에 --baseline을 붙여 명령어를 반복합니다.
커널(kernels)을 사용하면 MaxSim 스코어링, Ward 클러스터링 및 지원되는 백본(backbone) 연산을 가속화할 수 있습니다. 훈련 및 커널 종속성을 설치하세요:
pip install -U "sentence-transformers[train,kernels]"
위의 훈련 예시에서 손실을 구성하기 전에 모델과 풀러(pooler) 초기화를 교체하십시오:
model = MultiVectorEncoder(
"answerdotai/ModernBERT-base",
device="cuda",
...
Hierarchical poolers는 기본적으로 모델의 커널 설정을 상속받으므로, 모델에만 커널을 활성화하면 됩니다.
Sentence Transformers는 사용자의 환경에 맞는 호환 가능한 커널을 선택하며, 대체 구현이 없는 곳에서는 기존 구현을 사용합니다. 예를 들어, 기본 MaxSim 커널은 전체 토큰 유사성 행렬을 메모리에 구축하지 않고 쿼리를 문서와 비교하여 스코어링하지만, CUDA와 Triton이 필요합니다.
기본 Ward 클러스터링 커널은 CPU 대신 CUDA에서 실행되어 계층적 클러스터링 속도를 상당히 높여줍니다. 설정에 따라 SciPy보다 30~150배 빠르다고 할 수 있습니다.
풀링 인식 모델을 로드하고 압축된 문서 임베딩으로 검색해 보겠습니다. 커널 설정은 모델과 함께 저장되지 않으므로, 추론(inference)에서도 사용하려면 로드할 때 use_kernels=True를 추가하십시오:
from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling")
...
풀러는 인코더의 순방향 파이프라인(forward pipeline) 이후에 문서 임베딩을 압축하여, 총 벡터 수를 58개에서 16개로 줄이는 동시에 쿼리 임베딩은 변경하지 않습니다. Mars가 여전히 가장 높은 유사성 점수를 받고, 그 뒤를 Jupiter와 Venus가 따르므로, 훨씬 적은 수의 문서 벡터를 저장했음에도 불구하고 순위는 유지됩니다.
훈련된 모델과 동일한 저장 예산을 탐색하기 위해 추론 시 pool_factor를 변경할 수 있습니다. 반복 사용을 위해서는 token_pooling=pooling도 전달할 수 있습니다.
MultiVectorEncoder를 초기화하여 런타임 기본값으로 설정할 때, 풀러의 tasks는 여전히 어떤 입력이 압축되는지 제어합니다(예: 문서만). 이를 조사하기 위해, 저는 answerdotai/ModernBERT-base에서 두 개의 모델을 훈련시켰으며, 각 모델은 새로운 128차원 토큰 투영(token projection)을 가집니다:
- tomaarsen/mve-ModernBERT-base-token-pooling-baseline (캐시된 대조 손실(cached contrastive loss)만 사용).
- tomaarsen/mve-ModernBERT-base-token-pooling (위에서 언급된 unpooled 및 factor-4 목표(objectives)로 감싼 캐시된 대조 손실을 사용).
