멀티벡터 임베딩 모델용 Token Pooling-Aware Training 기법 공개

• Token Pooling-Aware Training은 ColBERT와 같은 멀티벡터 모델의 높은 저장 비용을 해결하기 위해 유사한 토큰 임베딩을 병합함. pool_factor=4 설정 시 NanoBEIR에서 문서 벡터를 74.8% 줄이면서 기존 nDCG@10 점수의 97.8%를 유지함.
• HierarchicalTokenPooling 모듈을 활용해 Ward 계층적 클러스터링으로 유사한 토큰 임베딩을 그룹화함. pool_factor=4 설정 시 문서당 평균 벡터 수가 193.22개에서 48.68개로 감소함.
• 학습 과정에서 원본 및 풀링된 표현의 검색 목적 함수를 평균화하는 로스 래퍼를 사용함. 직통 추정기(straight-through estimators)를 통해 이산 클러스터링 결정에 대한 그래디언트 역전파를 수행함.
• MaxSim 스코어링 및 Ward 클러스터링에 커널 가속을 적용해 SciPy 대비 30~150배 속도 향상을 달성함. 해당 접근법은 Sentence Transformers에 구현되어 있으며, answerdotai/ModernBERT-base 기반 체크포인트가 공개됨.

Implementation and Evaluation

• 120,000개의 MS MARCO 삼항쌍(triplets)과 128차원 토큰 투영으로 모델을 학습하고, NanoBEIR@13에서 평균 nDCG@10으로 평가함.
• 풀링 인식 학습은 모든 압축 수준에서 베이스라인 모델보다 우수한 성능을 보였으며, max_tokens=4에서 가장 큰 향상이 관찰되어 nDCG@10이 37.22에서 43.00으로 증가함.
• distance_threshold 보정을 통한 적응형 풀링은 코퍼스별 압축을 가능하게 하며, 저장 공간을 84.7% 줄이면서 풀링 전 성능의 97.3%를 유지함.

Storage and Compression

• 문서당 200개 벡터를 가진 100만 개 문서 코퍼스에서 pool_factor=4를 적용하고 4비트 잔차 양자화를 사용하면 총 저장 용량이 12.8 GB에서 3.2 GB로 감소함.
• 권장 압축 파이프라인 순서는 임베딩 양자화, 토큰 풀링, Matryoshka 토큰 임베딩 차원 축소임.
• 추론 시 재학습 없이 pool_factor를 동적으로 조정하여 다양한 저장 예산을 탐색할 수 있으며, 관련 문서의 랭킹 순위를 보존함.