Os modelos de embeddings multi-vetor, como o ColBERT, oferecem grande precisão na pesquisa, mas o número de vetores armazenados por documento faz rapidamente aumentar o custo do índice. Uma técnica de agrupamento, o token pooling, reduz este tamanho, e um treino adaptado permite, segundo a fonte, conservar melhor a qualidade dos resultados para um orçamento de armazenamento igual. Com pool_factor=4, um modelo treinado normalmente conserva 94,6 % do seu nDCG@10 não comprimido no NanoBEIR, contra 97,8 % para o modelo treinado considerando o pooling.
- Os modelos multi-vetor são muito precisos, mas armazenam muitos vetores por documento, o que encarece o índice.
- O token pooling agrupa vetores semelhantes e reduz o índice para cerca de um quarto.
- Um treino que considera o pooling conserva melhor a qualidade, com 97,8 % contra 94,6 % da pontuação não comprimida.
- A fonte sugere combinar quantização, token pooling e embeddings Matryoshka para reduzir ainda mais o índice.
Por que os modelos multi-vetor armazenam tantos vetores
Um modelo de embeddings clássico resume um documento inteiro num único vetor. Os modelos multi-vetor conservam, pelo contrário, um conjunto de vetores contextuais, um por token (uma unidade de texto, frequentemente uma palavra ou parte de uma palavra), o que permite que diferentes partes de uma consulta correspondam a diferentes partes de um documento. A fonte dá o exemplo de uma consulta sobre um computador portátil com longa autonomia, que pode encontrar uma passagem sobre o aparelho numa parte do documento e uma passagem sobre a autonomia da bateria noutra.
Este funcionamento tem um custo. Segundo o cálculo apresentado, com uma quantização a 4 bits por coordenada residual, 200 vetores de 128 dimensões representam 12,8 KB por documento, ou seja, 12,8 GB para um milhão de documentos. Reduzir cada documento para 50 vetores reduz este total para 3,2 GB, mantendo a dimensão e precisão idênticas.
O método baseia-se numa observação: muitos vetores de um mesmo documento podem ser próximos, pois vários tokens descrevem frequentemente noções semelhantes. A biblioteca HierarchicalTokenPooling, da Sentence Transformers, agrupa estes vetores por clustering de Ward, que funde em cada etapa os dois grupos que provocam o menor aumento da dispersão interna. Cada grupo é depois substituído pela sua média, e depois normalizado. Com pool_factor=4, um documento conserva aproximadamente um quarto dos seus vetores iniciais, de modo que um documento mais longo mantém proporcionalmente mais vetores.
Treinar o modelo para a compressão
Quanto maior for a compressão, mais fácil é perder informação útil. O treino que considera o pooling visa precisamente conservar a informação importante quando os vetores são fundidos. Cada lote de treino combina duas perdas: a pesquisa com os embeddings de documentos originais, e a pesquisa após compressão para cerca de um quarto dos vetores. Estas duas perdas são calculadas em média, e os dois objetivos atualizam o mesmo codificador.
Como o clustering toma decisões discretas, a fonte indica que o treino usa um estimador straight-through, que transmite os gradientes dos vetores agrupados para os tokens que os compõem, sem diferenciar a atribuição aos clusters. A classe MultiVectorTokenPoolingLoss encapsula uma perda de pesquisa existente, aqui CachedMultiVectorMultipleNegativesRankingLoss, com pesos de 0,5 para cada variante. O modelo inicial é answerdotai/ModernBERT-base, treinado em 120 000 triplos MS MARCO durante uma época.
Resultados no NanoBEIR
Com pool_factor=4, o nDCG@10 passa de 49,03 para 51,23 para o modelo treinado com pooling, reduzindo em 74,8 % o número de vetores de documentos em comparação com uma pesquisa não comprimida. O modelo não comprimido também melhora, de 51,84 para 52,39. Para uma média de 29,58 vetores por documento, o pooling adaptativo conserva 97,3 % da pontuação não comprimida do modelo treinado, contra 93,3 % para o modelo treinado normalmente.
Para reduzir o tamanho de um índice, a fonte sugere primeiro a quantização dos embeddings, depois o token pooling, eventualmente com um ajuste fino que considere o pooling, e finalmente os embeddings Matryoshka para reduzir a dimensão de cada vetor. Estas abordagens podem combinar-se. Os trabalhos em que o artigo se baseia são descritos no estudo Learn to Pool de Stefan Josef.
