Introducción de la búsqueda de vecino más cercano aproximado en Elasticsearch 8.0

Ha habido un aumento del interés en la búsqueda de vectores, gracias a una nueva generación de modelos de machine learning que pueden representar todo tipo de contenido como vectores, incluidos texto, imágenes, eventos y más. A menudo llamados “modelos de incrustación”, estas poderosas representaciones pueden capturar la similitud entre dos piezas de contenido de una manera que va más allá de sus características superficiales.

Los algoritmos de búsqueda de vecino más cercano k (kNN) encuentran los vectores en un sets de datos que son más similares a un vector de búsqueda. Junto con estas representaciones vectoriales, la búsqueda kNN abre posibilidades interesantes para la recuperación:

  • Encontrar pasajes que probablemente contengan la respuesta a una pregunta
  • Detección de imágenes casi duplicadas en un gran set de datos
  • Encontrar canciones que suenen similares a una canción dada

La búsqueda vectorial está lista para convertirse en un componente importante de la caja de herramientas de búsqueda, junto a técnicas tradicionales como la puntuación basada en términos.

Actualmente, Elasticsearch brinda soporte para almacenar vectores a través del tipo de campo dense_vector y usarlos para calcular puntuaciones de documentos. Esto permite a los usuarios realizar una búsqueda de kNN exacta mediante el escaneo de todos los documentos. Elasticsearch 8.0 se basa en esta funcionalidad para brindar soporte a la búsqueda de vecino más cercano aproximado (ANN) rápida. Esto representa un enfoque con mucha más escalabilidad, lo que permite ejecutar la búsqueda de vectores de manera eficiente en grandes sets de datos.

ANN en Elasticsearch

¿Qué es la búsqueda de vecino más cercano aproximado?

Existen estructuras de datos bien establecidas para kNN en vectores de baja dimensionalidad, como los árboles KD. De hecho, Elasticsearch incorpora árboles KD para brindar soporte a búsquedas en datos geoespaciales y numéricos. Pero los modelos de incrustación modernos para texto e imágenes suelen producir vectores de alta dimensionalidad de 100 - 1000 elementos, o incluso más. Estas representaciones vectoriales presentan un desafío único, ya que es muy difícil encontrar vecinos más cercanos de manera eficiente en altas dimensiones.

Ante esta dificultad, los algoritmos de vecino más cercano suelen sacrificar la precisión perfecta para mejorar su velocidad. Es posible que estos algoritmos de vecino más cercano aproximado (ANN) no siempre devuelvan los verdaderos k vectores más cercanos. Pero se ejecutan de manera eficiente, escalando a grandes sets de datos mientras mantienen un buen rendimiento.

Elegir un algoritmo de ANN

El diseño de algoritmos ANN es un área activa de investigación académica, y hay muchos algoritmos prometedores entre los cuales elegir. A menudo presentan diferentes concesiones en términos de velocidad de búsqueda, complejidad de implementación y costo de indexar. Afortunadamente, existe un gran proyecto de open source llamado ann-benchmarks que prueba los algoritmos líderes frente a varios sets de datos y publica comparaciones.

Elasticsearch 8.0 utiliza un algoritmo ANN llamado grafos de mundo pequeño jerárquico navegable (HNSW), que organiza los vectores en un grafo en función de su similitud entre sí. HNSW muestra un sólido rendimiento de búsqueda en una variedad de sets de datos de ann-benchmarks, y también obtuvo buenos resultados en nuestras propias pruebas. Otro beneficio de HNSW es que se utiliza ampliamente en la industria, habiéndose implementado en varios sistemas diferentes. Además del documento académico original, existen muchos recursos útiles para aprender sobre los detalles del algoritmo. Aunque el ANN de Elasticsearch se basa actualmente en HNSW, la característica está diseñada de forma flexible para permitirnos incorporar diferentes enfoques en el futuro.

¡Muéstrame el código!

Para indexar vectores para la búsqueda ANN, debemos establecer index: true y especificar la métrica de similitud que estamos usando para compararlos:

PUT index
{
 "mappings": {
   "properties": {
     "image-vector": {
       "type": "dense_vector",
       "dims": 128,
       "index": true,
       "similarity": "l2_norm"
     }
   }
 }
}

PUT index/_doc
{
 "image-vector": [0.12, 1.34, ...]
}
Luego, después de agregar vectores, podemos buscar los k vecinos más cercanos a un vector de búsqueda:
GET index/_knn_search
{
 "knn": {
   "field": "image-vector",
   "query_vector": [-0.5, 9.4, ...],
   "k": 10,
   "num_candidates": 100
 }
}
El nuevo endpoint _knn_search utiliza grafos HNSW para recuperar vectores similares de manera eficiente. A diferencia del kNN exacto, que realiza un escaneo completo de los datos, escala bien a sets de datos grandes. Aquí tienes un ejemplo que compara _knn_search con el enfoque exacto basado en búsquedas script_score en un sets de datos de 1 000 000 de vectores de imagen con 128 dimensiones, con un promedio de más de 10 000 búsquedas diferentes:
Approach         Queries Per Second    Recall (k=10)
script_score           5.257               1.000
_knn_search          849.286               0.945

En este ejemplo, la búsqueda ANN es órdenes de magnitud más rápida que el enfoque exacto. Su recuperación es de alrededor del 95 %, por lo que, en promedio, encuentra más de 9 de los 10 verdaderos vecinos más cercanos.

Puedes verificar el rendimiento de la búsqueda kNN en las evaluaciones comparativas cada noche de Elasticsearch. Estas evaluaciones comparativas funcionan con es-rally, una herramienta para la evaluación comparativa de Elasticsearch, específicamente el nuevo track de Rally dense_vector. Planeamos extender Rally para reportar la recuperación además de la latencia, ya que también es importante realizar un seguimiento de la precisión del algoritmo. Actualmente, estas evaluaciones comparativas prueban un set de datos de un par de millones de vectores, pero la búsqueda ANN ciertamente puede escalar más allá de esto con un mayor tiempo de índice o la adición de recursos de hardware.

Debido a que es un algoritmo aproximado, existen consideraciones especiales para ejecutar ANN en comparación con otros tipos de búsqueda. ANN tiene parámetros tanto al momento de la búsqueda como al momento de la indexación para controlar el equilibrio entre la latencia de búsqueda, la precisión de los resultados y el costo de indexar. Es importante medir la recuperación (recall) de la búsqueda de ANN en tu sets de datos para asegurarte de que la configuración funcione bien. Al comenzar con la búsqueda de kNN, la guía de referencia puede ser un lugar útil para empezar.

Con la tecnología de Apache Lucene

Muchas de las capacidades de búsqueda del núcleo de Elasticsearch funcionan gracias a la biblioteca Lucene, un proyecto open source gobernado por la Apache Software Foundation. Elasticsearch ANN no es la excepción, y está desarrollado sobre una nueva y emocionante característica de Lucene para almacenar y buscar vectores numéricos. Esta característica es el resultado de una gran colaboración que involucró a varios desarrolladores de diferentes organizaciones. Comenzando como una propuesta audaz, progresó rápidamente a una implementación funcional (y rápida). Luego vino el desafío de diseñar la API y completar la característica.

Desde entonces, la comunidad de Lucene ha continuado colaborando para impulsar la característica hacia adelante. Varios desarrolladores se interesaron e hicieron contribuciones, desde rediseñar nombres, hasta actualizaciones de algoritmos, mejoras de rendimiento, y más. Las capacidades de búsqueda vectorial de Lucene se están expandiendo rápidamente gracias a los esfuerzos de todos.

Además de la fructífera colaboración, desarrollar ANN en Lucene aporta otros beneficios importantes. La implementación de Lucene está diseñada a bajo nivel para integrarse correctamente con la funcionalidad existente, lo que permite que la búsqueda ANN interactúe sin problemas con otras características de Elasticsearch. Una integración tan profunda realmente no sería posible si dependiéramos de una biblioteca ANN externa. Por ejemplo, Lucene ANN gestiona de forma transparente los documentos eliminados omitiendo los 'tombstones' durante la búsqueda en el grafo. También respeta todas las garantías de compatibilidad de datos de Lucene, por lo que puedes estar seguro de que los datos vectoriales seguirán funcionando después de una actualización. Por último, la implementación escrita en Java al igual que Elasticsearch, lo que nos permite garantizar su seguridad y simplificar la gestión de la memoria.

¿Qué sigue?

En la versión 8.0, el endpoint _knn_search para la búsqueda eficiente de ANN se lanzará como una «vista previa técnica». La búsqueda de ANN es un tema relativamente nuevo no solo para Elastic, sino para la industria, y existen preguntas abiertas importantes sobre cómo debería comportarse. ¿Cuál es la mejor manera de combinar las puntuaciones de similitud vectorial con las puntuaciones de BM25 tradicionales? ¿Debería la búsqueda kNN admitir paginación? Desarrollar ANN como su propio endpoint experimental nos permitirá iterar rápidamente y probar su comportamiento. Planeamos integrar finalmente ANN en la API _search una vez que tengamos respuestas sólidas para estas preguntas. (Aunque _knn_search aún no está disponible para el público en general, el tipo de campo dense_vector se lanzó como GA en la versión 7.6 y sigue teniendo una API estable).

Algunas capacidades clave que planeamos admitir incluyen ANN con filtros, así como búsqueda "híbrida" donde los resultados de ANN se combinan con los de una búsqueda tradicional. También estamos trabajando para mejorar la velocidad de indexar, ya que crear grafos HNSW puede ser una operación exigente. Consideramos este lanzamiento solo como un comienzo y esperamos mejorar la búsqueda ANN en los próximos lanzamientos. Tus comentarios son realmente valiosos y ayudan a dar forma a la dirección de la característica. Nos encantaría recibir noticias tuyas en GitHub y en nuestros foros de debate (¡y en Lucene también)!

Prueba la búsqueda ANN en Elastic Cloud por logging en la consola de Elastic Cloud o registrándote para una prueba gratuita de 14 días.