멀티 벡터 임베딩 모델의 토큰 풀링 대응 학습을 해설, Sentence Transformers로 구현
Token Pooling-Aware Training for Multi-Vector Embedding Models
Hugging Face 블로그는 멀티벡터 임베딩 모델에 토큰 풀링 인식 학습을 도입하는 방법을 Sentence Transformers 코드 예시와 함께 설명했다. pool_factor=4 설정에서 일반 학습 모델의 NanoBEIR nDCG@10 유지율은 94.6%, 풀링 인식 학습 모델은 97.8%였으며, 두 모델 모두 문서 벡터가 4분의 1로 줄어든다.
멀티벡터 검색은 문서마다 많은 토큰 벡터를 보관하기 때문에 저장 용량이 크게 불어난다. 예를 들어 128차원, 4비트 잔차의 벡터 200개는 문서 하나당 12.8KB이며, 문서가 100만 개라면 12.8GB가 된다. 이를 50개 벡터로 풀링하면 3.2GB에 들어갈 수 있다. HierarchicalTokenPooling은 Ward 방식의 계층적 클러스터링으로 비슷한 벡터들을 통합하고, 각 그룹의 평균을 다시 정규화한다. 용량 지정 방식은 max_tokens, pool_factor, distance_threshold의 세 가지이며, 임계값은 목표 평균 벡터 수에 맞춰 보정할 수 있다.
학습에서는 원래 문서 임베딩으로 계산한 검색 손실과, 약 4분의 1로 압축한 뒤의 검색 손실을 각각 0.5의 가중치로 평균한다. 클러스터 할당은 미분하지 않고, straight-through estimator로 기울기를 구성 토큰에 전달한다. 구현은 MultiVectorTokenPoolingLoss가 기존 손실을 감싸는 형태이며, 풀링 버전은 추가 인코더 통과가 필요하지 않다. CUDA 커널을 켜면 조건에 따라 Ward 클러스터링이 SciPy보다 30~150배 빨라진다고 설명되어 있다.
평가에서는 ModernBERT-base를 기반으로, MS MARCO의 12만 트리플릿으로 1에폭 학습한 두 모델을 NanoBEIR@13에서 비교했다. pool_factor=4일 때 nDCG@10은 49.03에서 51.23으로 올랐고, 벡터 수는 압축하지 않은 경우보다 74.8% 적다. 압축하지 않은 경우에도 51.84에서 52.39로 개선되었다. 고정 상한과 적응형 임계값에서도 효과가 나타났으며, 평균 29.58개 벡터에서는 대응 모델이 비압축 점수의 97.3%를 유지했고, 일반 모델은 93.3%였다.
압축 기법의 조합에 대해서는 양자화, 토큰 풀링, Matryoshka에 의한 차원 축소 순으로 검토할 것을 권장한다. MultiVectorMatryoshkaLoss와 MultiVectorTokenPoolingLoss를 함께 사용하면 차원 수와 벡터 수를 모두 학습할 수 있다. 학습 스크립트는 uv나 Hugging Face Jobs로 실행할 수 있으며, 학습된 체크포인트도 공개되어 있다.
출처: Hugging Face Blog · huggingface.co