Token pooling no treinamento reduz o custo de armazenamento em embeddings multi-vector

O problema de fundo é de armazenamento. Modelos multi-vector como o ColBERT guardam um vector contextual por token, e o MaxSim soma, para cada vector da consulta, a melhor correspondência no documento. Mesmo com a quantization residual de 2 ou 4 bits do Fast-PLAID, 200 vectors de 128 dimensões ocupam 12,8 KB por documento, ou 12,8 GB para um milhão de documentos. Com 50 vectors por documento, o total cai para 3,2 GB.

O HierarchicalTokenPooling mescla vectors parecidos com clustering hierárquico de Ward, representa cada grupo pela média e normaliza de novo para o MaxSim. Com pool_factor=4, sobra cerca de um quarto dos vectors. O primeiro token é protegido por padrão, e só os documentos são comprimidos, não as consultas. O treinamento ciente de pooling usa um loss wrapper no Sentence Transformers: metade do peso vai para o retrieval sem pooling e metade para o retrieval após pooling. Como o Ward é discreto, o gradiente passa por um straight-through estimator. No experimento com ModernBERT-base e 120 mil triplets do MS MARCO, a retenção do nDCG@10 do NanoBEIR em pool_factor=4 subiu de 94,6% para 97,8%, com 49,03 contra 51,23 de nDCG@10. O autor não fez busca de hiperparâmetros e só o loss difere entre as duas execuções.

Para quem opera um índice multi-vector em produção, o recado é que o custo de armazenamento não precisa vir só de quantization. A leitura aqui é que o pooling pode ser tratado como parte do treino, e não só como etapa de indexação. O próprio post sugere ordem de ataque: quantization primeiro, depois pooling, depois Matryoshka. Vale validar com o seu corpus e o seu motor de busca, porque os números vêm de um único experimento.