Consultando 3 bilhões de vetores

• Por causa do problema de consultar 3 bilhões de vetores do Jeff Dean, este é um registro de experimento técnico sobre o processo de implementar diretamente uma solução map-reduce otimizada
• Começando com uma implementação naive que calcula o dot product entre 3 bilhões de vetores float32 de 768 dimensões e 1.000 vetores de consulta, foi obtida uma melhora gradual de desempenho com vetorização em numpy e conversão para float32
• Com base em 3.000 vetores, o método naive levou cerca de 2 segundos; após a vetorização, caiu para 0,01 s e, com float32, para 0,0045 s
• Ao escalar para 3 bilhões, a matriz de resultados exige cerca de 8,6 TB de memória, causando OOM, o que exige otimizações adicionais como processamento em lote, mapeamento de memória, reescrita em Rust/C e uso de ferramentas como SimSIMD
• Enfatiza que, mais do que a solução técnica, o desafio central mais difícil é definir os requisitos (forma de retorno, especificações da máquina, se compressão é aceitável etc.)

---

O ponto de partida do problema

• Ao ver a troca sobre Jeff Dean e o problema de consultar 3 bilhões de vetores, tudo começou com a tentativa de implementar diretamente uma solução map-reduce otimizada
• Um vetor é um arranjo de números de ponto flutuante com n dimensões, e a busca vetorial é usada para encontrar palavras ou itens semanticamente semelhantes
• Um padrão comum no uso de embeddings em pesquisa, recomendação e aplicações de busca generativa como o Cursor

Implementação naive

• Suposição inicial: 3 bilhões de vetores de documentos a serem pesquisados e cerca de 1.000 vetores de consulta, todos armazenados em disco no formato .npy
• A dimensão dos vetores é 768, um tamanho comum usado em muitos modelos de consulta por embeddings baseados em similaridade
• Um método de laço duplo que compara cada vetor de consulta com todos os vetores de documentos para calcular o dot product e retornar o resultado
• Testando primeiro com 3.000 vetores, o resultado foi de cerca de 2 segundos em um MacBook M2 (escala 1 milhão de vezes menor que 3 bilhões)

Otimização por vetorização (Vectorized)

• Aplicação de uma abordagem vetorizada que substitui o laço duplo pela operação de multiplicação de matrizes do numpy (vectors_file @ query_vectors.T)
• Grande melhora para 0,0107 s com base em 3.000 vetores
• Ao expandir para 3 milhões de vetores, levou 12,85 s

Conversão para float32

• Otimização adicional ao converter os dados para np.float32
• Redução para 0,0045 s com base em 3.000 vetores
• Como em 3 milhões de vetores leva cerca de 13 s, estima-se que em 3 bilhões levaria 1.000 vezes mais, ou cerca de 3.216 minutos

Resumo da comparação de desempenho

• Método naive (3.000 vetores): 1,9877 s
• Método vetorizado (3.000 vetores): 0,0107 s
• Método vetorizado (3 milhões de vetores): 12,8491 s
• Método com float32 (3.000 vetores): 0,0045 s

Problema de OOM e direções de otimização adicional

• Memória necessária ao processar 3 bilhões de objetos em float32 (4 bytes): cerca de 8,6 TB, causando OOM na execução
• São necessárias otimizações adicionais na direção sugerida por Jeff Dean:
  • Alterar o código para usar geradores e operações de comparação em lote
  • Gravar no disco a cada certo volume de operações ou usar mapeamento de memória
  • Reescrever o código de otimização em nível de sistema em Rust ou C
  • Usar bibliotecas dedicadas à comparação de similaridade vetorial em larga escala, como SimSIMD

A importância de definir os requisitos

• Antes de otimizações adicionais, surgiram pontos pouco claros no próprio problema:
  • Se é para consultar todos os 3 bilhões de uma vez com um único vetor de consulta e retornar todos os resultados, ou se é uma busca ANN top-k
  • Se também é preciso retornar os vetores originais e se há necessidade de reranqueamento com base em similaridade
  • Se a ideia é consultar o conjunto inteiro de uma vez com 1.000 vetores de consulta
  • Se os vetores já estão em memória, se são lidos um a um do disco ou se o método é por streaming
  • Se a execução é local, quais são as especificações da máquina e se GPU pode ser usada
  • O impacto do tamanho do embedding na representação do resultado e no tamanho dos vetores de entrada e saída
  • Se a compressão de float64 para float32 é aceitável do ponto de vista de precisão
  • Se é um projeto pontual e quanto tempo há disponível para sua criação
• Mais do que a solução técnica em si, a tarefa mais difícil é definir corretamente os requisitos