Efficient Multivector Retrieval with Token-Aware Clustering and Hierarchical Indexing

1. Formulação do Problema

Modelos de recuperação multivetorial (por exemplo, ColBERT) alcançaram eficácia state-of-the-art em Recuperação de Informação (RI) ao codificar documentos como conjuntos de vetores em nível de token e utilizar "interação tardia" (MaxSim) para calcular a relevância. No entanto, sua implantação é dificultada por:

• Alto Custo Computacional e de Memória: Armazenar e pesquisar bilhões de vetores de token exige recursos massivos.
• Agrupamento Ineficiente: Soluções atuais dependem do agrupamento $k$-means padrão para gerar vetores centróides para compressão e coleta de candidatos.
  • Gargalo de Escalabilidade: O $k$-means escala mal com o tamanho do conjunto de dados e o número de clusters ($k$), frequentemente exigindo dezenas de horas ou GPUs caras para treinar até mesmo conjuntos moderadamente grandes (por exemplo, 262k centróides).
  • Desequilíbrio de Tokens: O $k$-means padrão trata todos os vetores igualmente. Consequentemente, tokens frequentes (stopwords) dominam a função de perda e consomem a maioria dos centróides, enquanto tokens raros e altamente discriminativos são sub-representados, degradando a qualidade da recuperação.

2. Metodologia

Os autores propõem o Tachiom, um sistema de recuperação que compreende duas inovações centrais: Agrupamento Consciente de Token (Tac) e um Índice Hierárquico para recuperação otimizada.

A. Agrupamento Consciente de Token (Tac)

O Tac aborda as limitações do $k$-means padrão explorando a propriedade estrutural da identidade do token em embeddings multivetoriais. Em vez de agrupar todos os vetores globalmente, o Tac decompõe o problema em subproblemas independentes por token.

• Estratégia de Alocação de Centróides: O Tac aloca o orçamento de centróides com base na frequência do token e na variância semântica, em vez de distribuição uniforme de vetores.
  • Tratamento de Cauda: Tokens extremamente raros ($n < \mu$) têm garantido pelo menos 1 centróide; tokens pequenos ($\mu \le n < \tau$) recebem 2.
  • Pontuação Amortecida: Para tokens ativos, um peso $w_j$ é calculado como $\sqrt{n_j} \cdot s_j$, onde $n_j$ é a frequência e $s_j$ é a variância semântica (variância componente a componente). A raiz quadrada amortecia a influência de tokens de alta frequência, impedindo que monopolizem o orçamento.
  • Limitação e Reconciliação: Pisos e tetos rígidos garantem representação mínima para tokens raros e previnem alocação excessiva. O orçamento final é reconciliado para corresponder ao alvo $k$.
• Paralelização: Ao resolver problemas independentes de $k_j$-means para cada tipo de token, o Tac permite paralelização massiva, reduzindo drasticamente a complexidade computacional em comparação com o $k$-means global.
• Compressão de Resíduos: Após o agrupamento, os resíduos (vetor - centróide) são normalizados e comprimidos usando Quantização de Produto (PQ).

B. Índice Hierárquico para Recuperação (Tachiom)

O Tachiom aproveita os centróides de alta resolução gerados pelo Tac para otimizar o pipeline de recuperação em duas fases:

1. Fase de Coleta (Pontuação Apenas com Centróides):
   • Um grafo Hnsw é construído sobre os centróides.
   • Cada centróide mantém uma lista invertida de IDs de documentos.
   • Para uma consulta, o sistema percorre o grafo Hnsw para encontrar os centróides mais próximos, recupera os documentos associados e calcula uma pontuação MaxSim aproximada usando apenas interações de centróides.
   • Isso contorna completamente os cálculos caros em nível de token durante a seleção inicial de candidatos.
2. Fase de Refinamento (MaxSim Otimizado por PQ):
   • Os principais candidatos são refinados usando tanto centróides quanto resíduos comprimidos por PQ.
   • Layout Otimizado para Cache: Os autores introduzem um layout de memória especializado para tabelas de distância PQ. Em vez de acessos de memória dispersos para cada token de consulta, as distâncias são organizadas em Macro-blocos (subespaços), Blocos de centróide e Micro-blocos de token de consulta. Isso permite recuperar distâncias para todos os tokens de consulta contra um centróide específico como um único bloco de memória contíguo, melhorando significativamente a localidade de cache e a utilização SIMD.

3. Principais Contribuições

1. Agrupamento Consciente de Token (Tac): Um algoritmo inovador que decompõe o agrupamento global em subproblemas por token. Ele garante teoricamente uma aceleração limitada inferior pela razão entre o peso total do vocabulário e o peso máximo de um único token.
2. Arquitetura Tachiom: Um sistema de recuperação que combina coleta de centróides baseada em Hnsw com um layout de PQ otimizado para cache, adaptado para interação tardia, permitindo pontuação completa eficiente.
3. Validação Teórica e Empírica: O artigo fornece um limite inferior teórico para a aceleração do Tac e demonstra que o método escala para milhões de centróides enquanto mantém ou melhora a eficácia da recuperação.

4. Resultados Experimentais

Os experimentos foram conduzidos em Ms Marco-v1 (8,8M passagens) e LoTTE (2,4M passagens) usando embeddings ColBERTv2.

• Eficiência de Agrupamento:
  • O Tac alcançou até 247$\times$ mais rápido no treinamento que o $k$-means padrão (Faiss) e 84$\times$ mais rápido que o Faiss com aceleração MKL.
  • O Tac agrupou com sucesso 600M vetores em 262k clusters em 8 minutos, enquanto concorrentes atingiram um tempo limite de 24 horas antes de alcançar 131k–524k clusters.
  • O Tac escala perfeitamente para mais de 4 milhões de centróides em ~100 minutos.
• Eficácia da Recuperação:
  • O Tac alcançou MRR@10 igual ou superior ao $k$-means padrão com orçamentos fixos de centróides, provando que priorizar tokens raros melhora a qualidade da aproximação.
• Velocidade de Recuperação:
  • O Tachiom alcançou até 9,8$\times$ mais rápido na recuperação ponta a ponta comparado a sistemas state-of-the-art (Emvb, Warp, Igp).
  • Ele igualou a eficácia máxima do Emvb (que usa variantes supervisionadas caras de PQ) enquanto usa PQ padrão e alcança um aceleramento de 5,5$\times$ sobre o Emvb.

5. Significado

Este trabalho muda fundamentalmente o paradigma para a eficiência da recuperação multivetorial:

• Escalabilidade: Remove o gargalo computacional do agrupamento, permitindo que sistemas utilizem milhões de centróides (granularidade mais fina), que anteriormente eram inviáveis de treinar.
• Compromisso Qualidade-Eficiência: Ao alinhar a alocação de centróides com objetivos de recuperação (favorecendo tokens raros discriminativos) em vez de erro de reconstrução puro, alcança melhor eficácia com menos recursos.
• Implantação Prática: A combinação de treinamento rápido e inferência ultra-rápida (via coleta Hnsw e PQ otimizado) torna a recuperação multivetorial de alto desempenho viável para ambientes de produção em grande escala.

Os autores lançaram sua implementação em Rust do Tac e do Tachiom no GitHub para garantir reprodutibilidade.