Búsqueda eficiente de similitud vectorial en flujos de trabajo de recomendación usando Milvus con NVIDIA Merlin
Esta publicación se publicó por primera vez en el canal de Medium de NVIDIA Merlin y se editó y volvió a publicar aquí con permiso. Fue escrita conjuntamente por Burcin Bozkaya y William Hicks de NVIDIA y Filip Haltmayer y Li Liu de Zilliz.
Introducción
Los sistemas de recomendación modernos (Recsys) consisten en canalizaciones de entrenamiento/inferencia que implican múltiples etapas de ingesta de datos, preprocesamiento de datos, entrenamiento de modelos y ajuste de hiperparámetros para recuperar, filtrar, clasificar y puntuar elementos relevantes. Un componente esencial de una canalización de un sistema de recomendación es la recuperación o el descubrimiento de elementos que son más relevantes para un usuario, particularmente en presencia de grandes catálogos de artículos. Este paso normalmente implica una búsqueda de vecinos más cercanos aproximados (ANN) sobre una base de datos indexada de representaciones vectoriales de baja dimensión (es decir, embeddings) de atributos de productos y usuarios creadas a partir de modelos de aprendizaje profundo que se entrenan con interacciones entre usuarios y productos/servicios.
NVIDIA Merlin, un framework de código abierto desarrollado para entrenar modelos de extremo a extremo con el fin de hacer recomendaciones a cualquier escala, se integra con un índice de base de datos vectorial y un framework de búsqueda eficientes. Uno de esos frameworks que ha recibido mucha atención recientemente es Milvus, una base de datos vectorial de código abierto creada por Zilliz. Ofrece capacidades rápidas de indexación y consulta. Milvus añadió recientemente soporte de aceleración por GPU que utiliza GPU de NVIDIA para sostener flujos de trabajo de IA. El soporte de aceleración por GPU es una gran noticia porque una biblioteca de búsqueda vectorial acelerada hace posibles consultas concurrentes rápidas, lo que impacta positivamente en los requisitos de latencia de los sistemas de recomendación actuales, donde los desarrolladores esperan muchas solicitudes concurrentes. Milvus tiene más de 5 millones de descargas de Docker, ~23 mil estrellas en GitHub (a septiembre de 2023), más de 5.000 clientes empresariales y es un componente central de muchas aplicaciones (véanse los casos de uso).
Este blog demuestra cómo Milvus funciona con el framework Recsys de Merlin durante el entrenamiento y la inferencia. Mostramos cómo Milvus complementa a Merlin en la etapa de recuperación de elementos con una búsqueda top-k altamente eficiente de embeddings vectoriales y cómo puede usarse con NVIDIA Triton Inference Server (TIS) durante la inferencia (véase la Figura 1). Nuestros resultados de benchmark muestran una impresionante aceleración de 37x a 91x con Milvus acelerado por GPU que utiliza NVIDIA RAFT con los embeddings vectoriales generados por Merlin Models. El código que utilizamos para mostrar la integración Merlin-Milvus y los resultados detallados del benchmark, junto con la biblioteca que facilitó nuestro estudio de benchmark, están disponibles aquí.
Figura 1. Sistema de recomendación multietapa con el framework Milvus contribuyendo a la etapa de recuperación. Fuente de la figura multietapa original: esta publicación de blog.
Los desafíos que enfrentan los recomendadores
Dada la naturaleza multietapa de los recomendadores y la disponibilidad de varios componentes y bibliotecas que integran, un desafío significativo es integrar todos los componentes sin problemas en un pipeline de extremo a extremo. Nuestro objetivo es mostrar que la integración puede realizarse con menos esfuerzo en nuestros notebooks de ejemplo.
Otro desafío de los flujos de trabajo de recomendadores es acelerar ciertas partes del pipeline. Aunque se sabe que desempeñan un papel enorme en el entrenamiento de grandes redes neuronales, las GPU son incorporaciones recientes a las bases de datos vectoriales y a la búsqueda ANN. Con el aumento del tamaño de los inventarios de productos de comercio electrónico o de las bases de datos de medios en streaming y del número de usuarios que utilizan estos servicios, las CPU deben proporcionar el rendimiento necesario para servir a millones de usuarios en flujos de trabajo Recsys eficientes. La aceleración por GPU en otras partes del pipeline se ha vuelto necesaria para abordar este desafío. La solución de este blog aborda este desafío mostrando que la búsqueda ANN es eficiente cuando se usan GPU.
Stacks tecnológicos para la solución
Comencemos revisando primero algunos de los fundamentos necesarios para llevar a cabo nuestro trabajo.
NVIDIA Merlin: una biblioteca de código abierto con API de alto nivel que aceleran los recomendadores en GPU NVIDIA.
NVTabular: para el preprocesamiento de los datos tabulares de entrada y la ingeniería de características.
Merlin Models: para entrenar modelos de aprendizaje profundo y para aprender, en este caso, vectores de embeddings de usuarios e ítems a partir de datos de interacción de usuarios.
Merlin Systems: para combinar un modelo de recomendación basado en TensorFlow con otros elementos (p. ej., feature store, búsqueda ANN con Milvus) para servirlo con TIS.
Triton Inference Server: para la etapa de inferencia, donde se pasa un vector de características de usuario y se generan recomendaciones de productos.
Containerización: todo lo anterior está disponible mediante contenedor(es) que NVIDIA proporciona en el catálogo NGC. Usamos el contenedor Merlin TensorFlow 23.06.
Milvus 2.3: para realizar indexación y consultas vectoriales aceleradas por GPU.
Milvus 2.2.11: igual que lo anterior, pero para hacerlo en CPU.
Pymilvus SDK: para conectarse al servidor Milvus, crear índices de bases de datos vectoriales y ejecutar consultas mediante una interfaz de Python.
Feast: para guardar y recuperar atributos de usuarios e ítems en una feature store (de código abierto) como parte de nuestro pipeline RecSys de extremo a extremo.
También se utilizan varias bibliotecas y frameworks subyacentes bajo el capó. Por ejemplo, Merlin depende de otras bibliotecas de NVIDIA, como cuDF y Dask, ambas disponibles en RAPIDS cuDF. Del mismo modo, Milvus depende de NVIDIA RAFT para primitivas de aceleración por GPU y de bibliotecas modificadas como HNSW y FAISS para la búsqueda.
Comprender las bases de datos vectoriales y Milvus
La búsqueda aproximada de vecinos más cercanos (ANN) es una funcionalidad que las bases de datos relacionales no pueden manejar. Las BD relacionales están diseñadas para manejar datos tabulares con estructuras predefinidas y valores directamente comparables. Los índices de bases de datos relacionales dependen de esto para comparar datos y crear estructuras que aprovechan el conocimiento de si cada valor es menor o mayor que otro. Los vectores de embedding no pueden compararse directamente entre sí de esta manera, ya que necesitamos saber qué representa cada valor del vector. No pueden decir si un vector es necesariamente menor que otro. Lo único que podemos hacer es calcular la distancia entre los dos vectores. Si la distancia entre dos vectores es pequeña, podemos asumir que las características que representan son similares, y si es grande, podemos asumir que los datos que representan son más diferentes. Sin embargo, estos índices eficientes tienen un coste; calcular la distancia entre dos vectores es costoso desde el punto de vista computacional, y los índices vectoriales no son fácilmente adaptables y a veces no son modificables. Debido a estas dos limitaciones, integrar estos índices es más complejo en las bases de datos relacionales, por lo que se necesitan bases de datos vectoriales diseñadas específicamente.
Milvus se creó para resolver los problemas que las bases de datos relacionales encuentran con los vectores y fue diseñado desde cero para manejar estos vectores de embedding y sus índices a gran escala. Para cumplir con la etiqueta cloud-native, Milvus separa la computación y el almacenamiento, así como diferentes tareas de computación: consultas, preparación de datos e indexación. Los usuarios pueden escalar cada parte de la base de datos para manejar otros casos de uso, ya sean con muchas inserciones de datos o con muchas búsquedas. Si hay una gran afluencia de solicitudes de inserción, el usuario puede escalar temporalmente los nodos de índice de forma horizontal y vertical para manejar la ingesta. Del mismo modo, si no se están ingiriendo datos, pero hay muchas búsquedas, el usuario puede reducir los nodos de índice y, en su lugar, escalar los nodos de consulta para obtener mayor rendimiento. Este diseño del sistema (véase la Figura 2) nos exigió pensar con una mentalidad de computación paralela, lo que dio como resultado un sistema optimizado para computación con muchas puertas abiertas a futuras optimizaciones.
Figura 2. Diseño del sistema de Milvus
Milvus también utiliza muchas bibliotecas de indexación de última generación para ofrecer a los usuarios tanta personalización de su sistema como sea posible. Las mejora añadiendo la capacidad de manejar operaciones CRUD, datos transmitidos en streaming y filtrado. Más adelante, analizaremos en qué se diferencian estos índices y cuáles son las ventajas y desventajas de cada uno.
Solución de ejemplo: integración de Milvus y Merlin
La solución de ejemplo que presentamos aquí demuestra la integración de Milvus con Merlin en la etapa de recuperación de ítems (cuando los k ítems más relevantes se recuperan mediante una búsqueda ANN). Utilizamos un conjunto de datos real de un desafío RecSys, descrito a continuación. Entrenamos un modelo de aprendizaje profundo Two-Tower que aprende embeddings vectoriales para usuarios e ítems. Esta sección también proporciona el esquema de nuestro trabajo de benchmarking, incluidas las métricas que recopilamos y el rango de parámetros que utilizamos.
Nuestro enfoque implica:
Ingesta y preprocesamiento de datos
Entrenamiento del modelo de aprendizaje profundo Two-Tower
Construcción del índice de Milvus
Búsqueda de similitud en Milvus
Describimos brevemente cada paso y remitimos al lector a nuestros notebooks para obtener detalles.
Conjunto de datos
YOOCHOOSE GmbH proporciona el conjunto de datos que usamos en este estudio de integración y benchmark para el desafío RecSys 2015 y está disponible en Kaggle. Contiene eventos de clic/compra de usuarios de un minorista online europeo con atributos como un ID de sesión, marca de tiempo, ID de artículo asociado con el clic/compra y categoría del artículo, disponibles en el archivo yoochoose-clicks.dat. Las sesiones son independientes y no hay indicios de usuarios recurrentes, por lo que tratamos cada sesión como perteneciente a un usuario distinto. El conjunto de datos tiene 9,249,729 sesiones únicas (usuarios) y 52,739 artículos únicos.
Ingesta y preprocesamiento de datos
La herramienta que usamos para el preprocesamiento de datos es NVTabular, un componente de ingeniería de características y preprocesamiento de Merlin acelerado por GPU y altamente escalable. Usamos NVTabular para leer datos en la memoria de la GPU, reorganizar características según sea necesario, exportar a archivos parquet y crear una división de entrenamiento-validación para el entrenamiento. Esto da como resultado 7,305,761 usuarios únicos y 49,008 artículos únicos para entrenar. También categorizamos cada columna y sus valores en valores enteros. El conjunto de datos ahora está listo para el entrenamiento con el modelo Two-Tower.
Entrenamiento del modelo
Usamos el modelo de aprendizaje profundo Two-Tower para entrenar y generar embeddings de usuarios y artículos, que luego se utilizan en la indexación y consulta vectorial. Después de entrenar el modelo, podemos extraer los embeddings de usuarios y artículos.
Los siguientes dos pasos son opcionales: un modelo DLRM entrenado para clasificar los artículos recuperados para recomendación y un almacén de características usado (en este caso, Feast) para almacenar y recuperar características de usuarios y artículos. Los incluimos para la integridad del flujo de trabajo multietapa.
Finalmente, exportamos los embeddings de usuarios y artículos a archivos parquet, que luego pueden recargarse para crear un índice vectorial de Milvus.
Creación y consulta del índice de Milvus
Milvus facilita la indexación vectorial y la búsqueda de similitud mediante un “servidor” lanzado en la máquina de inferencia. En nuestro notebook #2, configuramos esto instalando mediante pip el servidor de Milvus y Pymilvus, y luego iniciando el servidor con su puerto de escucha predeterminado. A continuación, demostramos cómo construir un índice simple (IVF_FLAT) y consultarlo utilizando las funciones setup_milvus y query_milvus, respectivamente.
Benchmarking
Hemos diseñado dos benchmarks para demostrar el caso de uso de una biblioteca rápida y eficiente de indexación/búsqueda vectorial como Milvus.
Usar Milvus para crear índices vectoriales con los dos conjuntos de embeddings que generamos: 1) embeddings de usuarios para 7.3M usuarios únicos, divididos en 85% conjunto de entrenamiento (para indexación) y 15% conjunto de prueba (para consulta), y 2) embeddings de artículos para 49K productos (con una división entrenamiento-prueba 50–50). Este benchmark se realiza de forma independiente para cada conjunto de datos vectoriales, y los resultados se informan por separado.
Usar Milvus para crear un índice vectorial para el conjunto de datos de embeddings de artículos de 49K y consultar los 7.3M usuarios únicos contra este índice para búsqueda de similitud.
En estos benchmarks, usamos algoritmos de indexación IVFPQ y HNSW ejecutados en GPU y CPU, junto con varias combinaciones de parámetros. Los detalles están disponibles en nuestra página de GitHub.
La compensación entre calidad de búsqueda y rendimiento es una consideración de rendimiento importante, especialmente en un entorno de producción. Milvus permite un control completo sobre los parámetros de indexación para explorar esta compensación para un caso de uso determinado con el fin de lograr mejores resultados de búsqueda con la verdad fundamental. Esto puede significar un mayor coste computacional en forma de una tasa de rendimiento reducida o consultas por segundo (QPS). Medimos la calidad de la búsqueda ANN con una métrica de recuperación y proporcionamos curvas QPS-recuperación que demuestran la compensación. Luego se puede decidir un nivel aceptable de calidad de búsqueda dados los recursos de cómputo o los requisitos de latencia/rendimiento del caso de negocio.
Además, tenga en cuenta el tamaño de lote de consulta (nq) utilizado en nuestros benchmarks. Esto es útil en flujos de trabajo donde se envían múltiples solicitudes simultáneas a inferencia (por ejemplo, recomendaciones offline solicitadas y enviadas a una lista de destinatarios de correo electrónico o recomendaciones online creadas agrupando solicitudes concurrentes que llegan y procesándolas todas a la vez). Dependiendo del caso de uso, TIS también puede ayudar a procesar estas solicitudes en lotes.
Resultados
Ahora informamos los resultados para los tres conjuntos de benchmarks tanto en CPU como en GPU, utilizando los tipos de índice HNSW (solo CPU) e IVF_PQ (CPU y GPU) implementados por Milvus.
Búsqueda de similitud vectorial de ítems vs. ítems
Con este conjunto de datos más pequeño, cada ejecución para una combinación de parámetros determinada toma el 50% de los vectores de ítems como vectores de consulta y consulta los 100 vectores similares principales del resto. HNSW e IVF_PQ producen una alta recuperación con la configuración de parámetros probada, en el rango de 0,958–1,0 y 0,665–0,997, respectivamente. Este resultado sugiere que HNSW funciona mejor con respecto a la recuperación, pero IVF_PQ con configuraciones pequeñas de nlist produce una recuperación muy comparable. También debemos señalar que los valores de recuperación pueden variar mucho según los parámetros de indexación y consulta. Los valores que informamos se han obtenido después de una experimentación preliminar con rangos generales de parámetros y profundizando más en un subconjunto seleccionado.
El tiempo total para ejecutar todas las consultas en CPU con HNSW para una combinación de parámetros determinada oscila entre 5,22 y 5,33 seg.s (más rápido a medida que m aumenta, relativamente sin cambios con ef) y con IVF_PQ entre 13,67 y 14,67 seg.s (más lento a medida que nlist y nprobe aumentan). La aceleración por GPU sí tiene un efecto notable, como se ve en la Figura 3.
La Figura 3 muestra la compensación recuperación-rendimiento en todas las ejecuciones completadas en CPU y GPU con este pequeño conjunto de datos usando IVF_PQ. Encontramos que GPU proporciona una aceleración de 4x a 15x en todas las combinaciones de parámetros probadas (mayor aceleración a medida que nprobe aumenta). Esto se calcula tomando la relación de QPS de las ejecuciones en GPU sobre QPS de las ejecuciones en CPU para cada combinación de parámetros. En general, este conjunto presenta poco desafío para CPU o GPU y muestra perspectivas de una mayor aceleración con los conjuntos de datos más grandes, como se discute a continuación.
Figura 3. Aceleración por GPU con el algoritmo IVF_PQ de Milvus ejecutándose en GPU NVIDIA A100 (búsqueda de similitud ítem-ítem)
Búsqueda de similitud vectorial de usuarios vs. usuarios
Con el segundo conjunto de datos mucho más grande (7,3M usuarios), reservamos el 85% (~6,2M) de los vectores como “entrenamiento” (el conjunto de vectores a indexar), y el 15% restante (~1,1M) como conjunto de vectores de “prueba” o consulta. HNSW e IVF_PQ funcionan excepcionalmente bien en este caso, con valores de recuperación de 0,884–1,0 y 0,922–0,999, respectivamente. Sin embargo, son mucho más exigentes computacionalmente, especialmente con IVF_PQ en la CPU. El tiempo total para ejecutar todas las consultas en CPU con HNSW oscila entre 279,89 y 295,56 seg.s y con IVF_PQ entre 3082,67 y 10932,33 seg.s. Tenga en cuenta que estos tiempos de consulta son acumulativos para 1,1M vectores consultados, por lo que se puede decir que una sola consulta contra el índice sigue siendo muy rápida.
Sin embargo, la consulta basada en CPU puede no ser viable si el servidor de inferencia espera que muchos miles de solicitudes concurrentes ejecuten consultas contra un inventario de millones de ítems.
La GPU A100 ofrece una aceleración impresionante de 37x a 91x (con un promedio de 76,1x) en todas las combinaciones de parámetros con IVF_PQ en términos de rendimiento (QPS), como se muestra en la Figura 4. Esto es coherente con lo que observamos con el conjunto de datos pequeño, lo que sugiere que el rendimiento de la GPU escala razonablemente bien usando Milvus con millones de vectores de embeddings.
Figura 4. Aceleración de la GPU con el algoritmo IVF_PQ de Milvus ejecutándose en una GPU NVIDIA A100 (búsqueda de similitud usuario-usuario)
La siguiente Figura 5 detallada muestra la compensación entre recall y QPS para todas las combinaciones de parámetros probadas en CPU y GPU con IVF_PQ. Cada conjunto de puntos (arriba para GPU, abajo para CPU) en este gráfico representa la compensación que se afronta al cambiar los parámetros de indexación/consulta de vectores para lograr un recall más alto a costa de un menor rendimiento. Tenga en cuenta la considerable pérdida de QPS en el caso de la GPU al intentar alcanzar niveles de recall más altos.
Figura 5. Compensación entre recall y rendimiento para todas las combinaciones de parámetros probadas en CPU y GPU con IVF_PQ (usuarios frente a usuarios)
Búsqueda de similitud vectorial usuarios frente a ítems
Finalmente, consideramos otro caso de uso realista en el que los vectores de usuario se consultan contra vectores de ítems (como se demostró en el Notebook 01 anterior). En este caso, se indexan 49K vectores de ítems, y 7,3M vectores de usuario se consultan cada uno para obtener los 100 ítems más similares.
Aquí es donde las cosas se ponen interesantes porque consultar 7,3M en lotes de 1000 contra un índice de 49K ítems parece llevar mucho tiempo en la CPU tanto para HNSW como para IVF_PQ. La GPU parece manejar mejor este caso (véase la Figura 6). Los niveles de precisión más altos de IVF_PQ en CPU cuando nlist = 100 se calculan en unos 86 minutos de promedio, pero varían significativamente a medida que aumenta el valor de nprobe (51 min. cuando nprobe = 5 frente a 128 min. cuando nprobe = 20). La GPU NVIDIA A100 acelera considerablemente el rendimiento por un factor de 4x a 17x (mayores aceleraciones a medida que nprobe aumenta). Recuerde que el algoritmo IVF_PQ, mediante su técnica de cuantización, también reduce la huella de memoria y proporciona una solución de búsqueda ANN computacionalmente viable combinada con la aceleración de la GPU.
Figura 6. Aceleración de la GPU con el algoritmo IVF_PQ de Milvus ejecutándose en una GPU NVIDIA A100 (búsqueda de similitud usuario-ítem)
De manera similar a la Figura 5, la compensación entre recall y rendimiento se muestra en la Figura 7 para todas las combinaciones de parámetros probadas con IVF_PQ. Aquí, aún se puede ver cómo puede ser necesario renunciar ligeramente a algo de precisión en la búsqueda ANN a favor de un mayor rendimiento, aunque las diferencias son mucho menos notables, especialmente en el caso de las ejecuciones en GPU. Esto sugiere que se pueden esperar niveles relativamente constantes y altos de rendimiento computacional con la GPU, sin dejar de lograr un recall alto.
Figura 7. Compensación entre recall y rendimiento para todas las combinaciones de parámetros probadas en CPU y GPU con IVF_PQ (usuarios frente a ítems)
Conclusión
Nos complace compartir algunas observaciones finales si ha llegado hasta aquí. Queremos recordarle que la complejidad y la naturaleza multietapa de los RecSys modernos requieren rendimiento y eficiencia en cada paso. Esperamos que este blog le haya dado razones convincentes para considerar el uso de dos funciones críticas en sus pipelines de RecSys:
La biblioteca Merlin Systems de NVIDIA Merlin le permite integrar fácilmente Milvus, un motor eficiente de búsqueda vectorial acelerado por GPU.
Use GPU para acelerar los cálculos de indexación de bases de datos vectoriales y la búsqueda ANN con tecnología como RAPIDS RAFT.
Estos hallazgos sugieren que la integración Merlin-Milvus presentada es muy eficiente y mucho menos compleja que otras opciones para el entrenamiento y la inferencia. Además, ambos frameworks se desarrollan activamente, y se añaden muchas funciones nuevas (por ejemplo, nuevos índices de bases de datos vectoriales acelerados por GPU de Milvus) en cada lanzamiento. El hecho de que la búsqueda de similitud vectorial sea un componente crucial en diversos flujos de trabajo, como visión por computadora, modelado de grandes lenguajes y sistemas de recomendación, hace que este esfuerzo valga aún más la pena.
Para finalizar, nos gustaría agradecer a todos los de Zilliz/Milvus y Merlin, y a los equipos de RAFT que contribuyeron al esfuerzo de producir este trabajo y la publicación del blog. Esperamos tener noticias suyas, si tienen la oportunidad de implementar Merlin y Milvus en su recsys u otros flujos de trabajo.
Sigue leyendo
Zilliz Cloud Enterprise Vector Search Powers High-Performance AI on AWS
Zilliz Cloud on AWS powers secure, scalable, ultra-fast vector search for enterprise AI apps, with BYOC, sub-10ms latency, and zero-DevOps simplicity.
Build for the Boom: Why AI Agent Startups Should Build Scalable Infrastructure Early
Explore strategies for developing AI agents that can handle rapid growth. Don't let inadequate systems undermine your success during critical breakthrough moments.
Milvus WebUI: A Visual Management Tool for Your Vector Database
Explore Milvus WebUI to monitor, manage, and optimize your vector database with real-time insights, performance tracking, and system health monitoring.