Apresentamos a busca de vizinho mais próximo aproximado no Elasticsearch 8.0

Houve um aumento no interesse pela busca vetorial, graças a uma nova geração de modelos de machine learning que podem representar todos os tipos de conteúdo como vetores, incluindo texto, imagens, eventos e muito mais. Frequentemente chamadas de “modelos de embedding”, essas representações poderosas podem capturar a similaridade entre dois conteúdos de uma maneira que vai além de suas características superficiais.

Os algoritmos de busca de k vizinhos mais próximos (kNN) encontram os vetores em um conjunto de dados que são mais semelhantes a um vetor de consulta. Combinada com essas representações vetoriais, a busca kNN abre possibilidades interessantes para a recuperação:

• Encontrar passagens que provavelmente contêm a resposta para uma pergunta
• Detectando imagens quase duplicadas em um grande conjunto de dados
• Encontrar músicas que soam semelhantes a uma determinada música

A busca vetorial está pronta para se tornar um componente importante da caixa de ferramentas de busca, ao lado de técnicas tradicionais como a pontuação baseada em termos.

Atualmente, o Elasticsearch oferece suporte ao armazenamento de vetores por meio do tipo de campo dense_vector e ao uso deles para calcular pontuações de documentos. Isso permite que os usuários realizem uma busca kNN exata ao escanear todos os documentos. O Elasticsearch 8.0 baseia-se nessa funcionalidade para oferecer suporte a buscar rápida de vizinho mais próximo aproximado (ANN). Isso representa uma abordagem com muito mais escalabilidade, permitindo que a busca vetorial seja executada de forma eficiente em grandes conjuntos de dados.

ANN no Elasticsearch

O que é buscar vizinho mais próximo aproximado?

Existem estruturas de dados bem estabelecidas para kNN em vetores de baixa dimensão, como árvores KD. Na verdade, o Elasticsearch incorpora árvores KD para dar suporte a buscar em dados geoespaciais e numéricos. Mas os modelos modernos de embedding para texto e imagens normalmente produzem vetores de alta dimensão de 100 – 1000 elementos, ou até mais. Essas representações vetoriais apresentam um desafio único, pois é muito difícil encontrar vizinhos mais próximos de forma eficiente em altas dimensões.

Diante dessa dificuldade, os algoritmos de vizinho mais próximo geralmente sacrificam a precisão perfeita para melhorar sua velocidade. Esses algoritmos de vizinho mais próximo aproximado (ANN) podem nem sempre retornar os verdadeiros k vetores mais próximos. Mas eles são executados de forma eficiente, com redimensionamento para grandes conjuntos de dados enquanto mantêm um bom desempenho.

Escolhendo um algoritmo ANN

O design de algoritmos de ANN é uma área ativa de pesquisa acadêmica, e existem muitos algoritmos promissores para escolher. Eles geralmente apresentam diferentes compensações em termos de velocidade de buscar, complexidade de implementação e custo de indexação. Felizmente, existe um ótimo projeto open source chamado ann-benchmarks que testa os principais algoritmos em relação a vários conjuntos de dados e publica comparações.

O Elasticsearch 8.0 usa um algoritmo ANN chamado Hierarchical Navigable Small World gráficos (HNSW), que organiza vetores em um grafo com base na similaridade entre eles. O HNSW apresenta um forte desempenho de buscar em uma variedade de conjuntos de dados ann-benchmarks e também teve um bom desempenho em nossos próprios testes. Outro benefício do HNSW é que ele é amplamente utilizado no setor, tendo sido implementado em vários sistemas diferentes. Além do artigo acadêmico original, existem muitos recursos úteis para aprender sobre os detalhes do algoritmo. Embora o ANN do Elasticsearch seja atualmente baseado em HNSW, o recurso foi projetado de forma flexível para nos permitir incorporar diferentes abordagens no futuro.

Mostre-me o código!

Para indexar vetores para buscar ANN, precisamos definir index: true e especificar a métrica de similaridade que estamos usando para compará-los:

Neste exemplo, buscar ANN é ordens de magnitude mais rápida do que a abordagem exata. Seu recall é de cerca de 95%, portanto, em média, ele encontra mais de 9 dos 10 vizinhos mais próximos verdadeiros.

Você pode verificar o desempenho de buscar kNN nos benchmarks noturnos do Elasticsearch. Estes benchmarks são impulsionados pelo es-rally, uma ferramenta para benchmarking do Elasticsearch, especificamente a nova track dense_vector do Rally. Planejamos estender o Rally para fazer um relatório sobre o recall além da latência, pois também é importante rastrear a precisão do algoritmo. Atualmente, estes benchmarks testam um conjunto de dados de alguns milhões de vetores, mas a busca ANN certamente pode redimensionar além disso com um tempo de índice maior ou a adição de recursos de hardware.

Desenvolvido com Apache Lucene

Muitas das capacidades de buscar do núcleo do Elasticsearch são impulsionadas pela biblioteca Lucene, um projeto open source governado pela Apache Software Foundation. O ANN do Elasticsearch não é exceção e foi desenvolvido com base em um novo e empolgante recurso do Lucene para armazenar e buscar vetores numéricos. Este recurso é o resultado de uma grande colaboração envolvendo vários desenvolvedores de diferentes organizações. Começando como uma proposta ousada, progrediu rapidamente para uma implementação funcional (e rápida). Depois veio o desafio de projetar a API e completar o recurso.

Desde então, a comunidade Lucene continuou a colaborar para impulsionar o recurso. Vários desenvolvedores se interessaram e fizeram contribuições, desde redesign de nomes, até atualizações de algoritmos, melhorias de desempenho e muito mais. As capacidades de busca vetorial do Lucene estão se expandindo rapidamente graças aos esforços de todos.

Além da colaboração frutífera, desenvolver ANN no Lucene traz outros grandes benefícios. A implementação do Lucene foi projetada em baixo nível para integrar-se corretamente à funcionalidade existente, o que permite que a busca ANN interaja perfeitamente com outros recursos do Elasticsearch. Uma integração tão profunda não seria realmente possível se dependêssemos de uma biblioteca ANN externa. Por exemplo, o ANN do Lucene trata documentos excluídos de forma transparente ao ignorar 'tombstones' durante a busca no gráfico. Ele também respeita todas as garantias de compatibilidade de dados do Lucene, para que você possa ter certeza de que os dados vetoriais ainda funcionarão após uma atualização. Por fim, a implementação escrita em Java, assim como o Elasticsearch, o que nos permite garantir sua segurança e simplificar o gerenciamento de memória.

O que vem a seguir?

Na versão 8.0, o endpoint _knn_search para busca eficiente de ANN será lançado como uma "prévia técnica". A busca de ANN é um tópico relativamente novo não apenas para a Elastic, mas para o setor, e existem questões em aberto significativas sobre como ela deve se comportar. Qual é a melhor maneira de combinar pontuações de similaridade vetorial com as tradicionais pontuações BM25? Buscar kNN deve oferecer suporte à paginação? Desenvolver o ANN como seu próprio endpoint experimental nos permitirá iterar rapidamente e testar seu comportamento. Planejamos integrar o ANN ao endpoint da API _search assim que tivermos respostas sólidas para essas perguntas. (Embora o _knn_search ainda não esteja com disponibilidade geral, o tipo de campo dense_vector tornou-se GA na versão 7.6 e continua a ter uma API estável.)

Algumas capacidades principais que planejamos dar suporte incluem ANN com filtros, bem como buscar “híbrida”, onde os resultados de ANN são combinados com aqueles de uma consulta tradicional. Também estamos trabalhando para melhorar a velocidade de indexação, já que a construção de gráficos HNSW pode ser uma operação cara. Consideramos este lançamento apenas um começo e estamos ansiosos para melhorar a capacidade de buscar ANN nos próximos lançamentos. Seu feedback é muito valioso e ajuda a moldar a direção do recurso. Adoraríamos ouvir você no GitHub e em nossos fóruns de discussão (e no Lucene também)!

Experimente buscar ANN no Elastic Cloud fazendo logging no console Elastic Cloud ou inscrevendo-se para uma avaliação gratuita de 14 dias.