Se você faz busca nos próprios documentos com um modelo no estilo ColBERT, o índice provavelmente é maior do que gostaria. Esses modelos multivetoriais guardam um vetor para cada token, e não um por documento, o que os torna precisos e caros de armazenar. Um novo post no blog da Hugging Face mostra como o Sentence Transformers agora consegue treinar um modelo para continuar preciso depois que esses vetores são comprimidos, recuperando a maior parte da qualidade que costuma se perder.
O problema de armazenamento é concreto. Com 4 bits por coordenada, 200 vetores de 128 dimensões ocupam 12,8 KB por documento, ou 12,8 GB para um milhão de documentos. Reduzir cada documento a 50 vetores leva isso a 3,2 GB. O HierarchicalTokenPooling da biblioteca junta vetores de tokens parecidos com agrupamento de Ward, seja por um fator fixo, seja por um limite de distância que pode ser calibrado para uma média-alvo, e deixa os vetores da consulta intactos. O agrupamento funciona com qualquer modelo multivetorial, mas um modelo que nunca foi treinado assim pode perder distinções importantes.
A novidade é um invólucro de função de perda, o MultiVectorTokenPoolingLoss, que treina com duas versões de cada lote: os vetores completos do documento e uma versão reduzida a cerca de um quarto. A ideia lembra os embeddings Matryoshka, que treinam prefixos menores para continuarem úteis. No experimento do autor, dois modelos ModernBERT-base foram treinados por uma época com as mesmas 120.000 triplas do MS MARCO, diferindo apenas na perda. Com fator de agrupamento 4, o modelo comum manteve 94,6% da nota NanoBEIR sem agrupamento; o treinado para isso manteve 97,8%, com o nDCG@10 subindo de 49,03 para 51,23 e 74,8% menos vetores de documento. A qualidade sem agrupamento também melhorou um pouco. Um kernel de GPU para o agrupamento seria 30-150x mais rápido que o SciPy.
Para equipes que usam o PyLate, a biblioteca da LightOn para busca no estilo ColBERT construída sobre o Sentence Transformers, isso é diretamente relevante: o trabalho parte do estudo da própria LightOn sobre regularização do agrupamento, e a receita de treino roda numa única A10G pelo Hugging Face Jobs em menos de quatro horas. O conselho prático é combinar quantização, agrupamento e truncamento Matryoshka, e medir tamanho do índice e qualidade no seu próprio corpus.
As ressalvas são do próprio autor. É um único experimento com um único modelo base e sem busca de hiperparâmetros, avaliado no NanoBEIR, um conjunto pequeno de testes. O ganho é mais claro no fator de agrupamento para o qual o modelo foi treinado, ainda que se transfira para outros ajustes. Os resultados com documentos longos ou de domínio específico podem ser diferentes.