Busca eficiente de similaridade vetorial em fluxos de trabalho de recomendação usando Milvus com NVIDIA Merlin
Este post foi publicado originalmente no canal Medium da NVIDIA Merlin e editado e republicado aqui com permissão. Foi escrito em conjunto por Burcin Bozkaya e William Hicks, da NVIDIA, e Filip Haltmayer e Li Liu, da Zilliz.
Introdução
Sistemas de recomendação modernos (Recsys) consistem em pipelines de treinamento/inferência envolvendo múltiplos estágios de ingestão de dados, pré-processamento de dados, treinamento de modelos e ajuste de hiperparâmetros para recuperação, filtragem, ranqueamento e pontuação de itens relevantes. Um componente essencial de um pipeline de sistema de recomendação é a recuperação ou descoberta de coisas que são mais relevantes para um usuário, particularmente na presença de grandes catálogos de itens. Essa etapa normalmente envolve uma busca de vizinhos mais próximos aproximados (ANN) em um banco de dados indexado de representações vetoriais de baixa dimensão (ou seja, embeddings) de atributos de produtos e usuários criadas a partir de modelos de aprendizado profundo que treinam com base em interações entre usuários e produtos/serviços.
NVIDIA Merlin, um framework de código aberto desenvolvido para treinar modelos ponta a ponta para fazer recomendações em qualquer escala, integra-se a um índice de banco de dados vetorial eficiente e a um framework de busca. Um desses frameworks que ganhou muita atenção recentemente é o Milvus, um banco de dados vetorial de código aberto criado pela Zilliz. Ele oferece recursos rápidos de indexação e consulta. O Milvus adicionou recentemente suporte a aceleração por GPU, que usa GPUs NVIDIA para sustentar fluxos de trabalho de IA. O suporte a aceleração por GPU é uma ótima notícia porque uma biblioteca de busca vetorial acelerada torna possíveis consultas simultâneas rápidas, impactando positivamente os requisitos de latência nos sistemas de recomendação atuais, nos quais os desenvolvedores esperam muitas solicitações simultâneas. O Milvus tem mais de 5M pulls no docker, ~23k estrelas no GitHub (em setembro de 2023), mais de 5.000 clientes Enterprise e é um componente central de muitas aplicações (veja casos de uso).
Este blog demonstra como o Milvus funciona com o framework Merlin Recsys em tempo de treinamento e inferência. Mostramos como o Milvus complementa o Merlin na etapa de recuperação de itens com uma busca altamente eficiente de embeddings vetoriais top-k e como ele pode ser usado com o NVIDIA Triton Inference Server (TIS) em tempo de inferência (veja a Figura 1). Nossos resultados de benchmark mostram uma impressionante aceleração de 37x a 91x com o Milvus acelerado por GPU, que usa NVIDIA RAFT com os embeddings vetoriais gerados pelos Merlin Models. O código que usamos para mostrar a integração Merlin-Milvus e os resultados detalhados de benchmark, junto com a biblioteca que facilitou nosso estudo de benchmark, estão disponíveis aqui.
Figura 1. Sistema de recomendação em múltiplos estágios com o framework Milvus contribuindo para a etapa de recuperação. Fonte da figura original de múltiplos estágios: este post de blog.
Os desafios enfrentados pelos sistemas de recomendação
Dada a natureza em múltiplos estágios dos sistemas de recomendação e a disponibilidade de vários componentes e bibliotecas que eles integram, um desafio significativo é integrar todos os componentes perfeitamente em um pipeline de ponta a ponta. Nosso objetivo é mostrar que a integração pode ser feita com menos esforço em nossos notebooks de exemplo.
Outro desafio dos fluxos de trabalho de recomendadores é acelerar certas partes do pipeline. Embora se saiba que desempenham um papel enorme no treinamento de grandes redes neurais, as GPUs são adições recentes aos bancos de dados vetoriais e à busca ANN. Com o aumento do tamanho dos inventários de produtos de e-commerce ou dos bancos de dados de mídia de streaming e do número de usuários que utilizam esses serviços, as CPUs devem fornecer o desempenho necessário para atender milhões de usuários em fluxos de trabalho de Recsys performáticos. A aceleração por GPU em outras partes do pipeline tornou-se necessária para enfrentar esse desafio. A solução neste blog aborda esse desafio mostrando que a busca ANN é eficiente ao usar GPUs.
Pilhas de tecnologia para a solução
Vamos começar revisando alguns dos fundamentos necessários para realizar nosso trabalho.
NVIDIA Merlin: uma biblioteca de código aberto com APIs de alto nível que aceleram recomendadores em GPUs NVIDIA.
NVTabular: para pré-processamento dos dados tabulares de entrada e engenharia de atributos.
Merlin Models: para treinar modelos de deep learning e para aprender, neste caso, vetores de embedding de usuários e itens a partir de dados de interação de usuários.
Merlin Systems: para combinar um modelo de recomendação baseado em TensorFlow com outros elementos (por exemplo, feature store, busca ANN com Milvus) a serem servidos com TIS.
Triton Inference Server: para a etapa de inferência, em que um vetor de atributos do usuário é passado e recomendações de produtos são geradas.
Conteinerização: tudo o que foi mencionado acima está disponível por meio de container(s) que a NVIDIA fornece no NGC catalog. Usamos o Merlin TensorFlow 23.06 container.
Milvus 2.3: para realizar indexação e consultas vetoriais aceleradas por GPU.
Milvus 2.2.11: igual ao anterior, mas para fazê-lo na CPU.
Pymilvus SDK: para conectar-se ao servidor Milvus, criar índices de bancos de dados vetoriais e executar consultas por meio de uma interface Python.
Feast: para salvar e recuperar atributos de usuários e itens em uma feature store (open source) como parte do nosso pipeline RecSys de ponta a ponta.
Várias bibliotecas e frameworks subjacentes também são usados nos bastidores. Por exemplo, o Merlin depende de outras bibliotecas NVIDIA, como cuDF e Dask, ambas disponíveis em RAPIDS cuDF. Da mesma forma, o Milvus depende do NVIDIA RAFT para primitivas de aceleração por GPU e de bibliotecas modificadas, como HNSW e FAISS, para busca.
Entendendo bancos de dados vetoriais e Milvus
Vizinho mais próximo aproximado (ANN) é uma funcionalidade com a qual bancos de dados relacionais não conseguem lidar. Bancos de dados relacionais são projetados para lidar com dados tabulares com estruturas predefinidas e valores diretamente comparáveis. Índices de bancos de dados relacionais dependem disso para comparar dados e criar estruturas que aproveitam o fato de saber se cada valor é menor ou maior que o outro. Vetores de embedding não podem ser comparados diretamente entre si dessa forma, pois precisamos saber o que cada valor no vetor representa. Eles não conseguem dizer se um vetor é necessariamente menor que o outro. A única coisa que podemos fazer é calcular a distância entre os dois vetores. Se a distância entre dois vetores for pequena, podemos assumir que os atributos que eles representam são semelhantes e, se for grande, podemos assumir que os dados que eles representam são mais diferentes. No entanto, esses índices eficientes têm um custo; calcular a distância entre dois vetores é computacionalmente caro, e índices vetoriais não são facilmente adaptáveis e, às vezes, não são modificáveis. Devido a essas duas limitações, integrar esses índices é mais complexo em bancos de dados relacionais, e é por isso que bancos de dados vetoriais criados para esse fim são necessários.
Milvus foi criado para resolver os problemas que bancos de dados relacionais encontram com vetores e foi projetado desde o início para lidar com esses vetores de embedding e seus índices em grande escala. Para cumprir o selo cloud-native, o Milvus separa computação e armazenamento e diferentes tarefas de computação — consultas, preparação de dados e indexação. Os usuários podem escalar cada parte do banco de dados para lidar com outros casos de uso, sejam eles intensivos em inserção de dados ou intensivos em busca. Se houver um grande fluxo de solicitações de inserção, o usuário pode escalar temporariamente os nós de índice horizontal e verticalmente para lidar com a ingestão. Da mesma forma, se nenhum dado estiver sendo ingerido, mas houver muitas buscas, o usuário pode reduzir os nós de índice e, em vez disso, escalar os nós de consulta para obter maior throughput. Esse design de sistema (veja a Figura 2) exigiu que pensássemos com uma mentalidade de computação paralela, resultando em um sistema otimizado para computação com muitas portas abertas para otimizações adicionais.
Figura 2. Design do sistema Milvus
O Milvus também usa muitas bibliotecas de indexação de última geração para oferecer aos usuários o máximo de personalização possível para seu sistema. Ele as aprimora adicionando a capacidade de lidar com operações CRUD, dados em streaming e filtragem. Mais adiante, discutiremos como esses índices diferem e quais são os prós e contras de cada um.
Exemplo de solução: integração do Milvus e do Merlin
A solução de exemplo que apresentamos aqui demonstra a integração do Milvus com o Merlin na etapa de recuperação de itens (quando os k itens mais relevantes são recuperados por meio de uma busca ANN). Usamos um conjunto de dados real de um desafio RecSys, descrito abaixo. Treinamos um modelo de deep learning Two-Tower que aprende embeddings vetoriais para usuários e itens. Esta seção também fornece o blueprint do nosso trabalho de benchmarking, incluindo as métricas que coletamos e o intervalo de parâmetros que usamos.
Nossa abordagem envolve:
Ingestão e pré-processamento de dados
Treinamento do modelo de deep learning Two-Tower
Construção de índice do Milvus
Busca por similaridade no Milvus
Descrevemos brevemente cada etapa e encaminhamos o leitor aos nossos notebooks para detalhes.
Conjunto de dados
A YOOCHOOSE GmbH fornece o conjunto de dados que usamos neste estudo de integração e benchmark para o desafio RecSys 2015 e está disponível no Kaggle. Ele contém eventos de clique/compra de usuários de um varejista online europeu com atributos como ID de sessão, carimbo de data/hora, ID do item associado ao clique/compra e categoria do item, disponíveis no arquivo yoochoose-clicks.dat. As sessões são independentes, e não há indício de usuários recorrentes, então tratamos cada sessão como pertencente a um usuário distinto. O conjunto de dados tem 9,249,729 sessões únicas (usuários) e 52,739 itens únicos.
Ingestão e pré-processamento de dados
A ferramenta que usamos para pré-processamento de dados é NVTabular, um componente de engenharia de atributos e pré-processamento do Merlin acelerado por GPU e altamente escalável. Usamos o NVTabular para ler dados na memória da GPU, reorganizar atributos conforme necessário, exportar para arquivos parquet e criar uma divisão treino-validação para o treinamento. Isso resulta em 7,305,761 usuários únicos e 49,008 itens únicos para treinar. Também categorizamos cada coluna e seus valores em valores inteiros. O conjunto de dados agora está pronto para treinamento com o modelo Two-Tower.
Treinamento do modelo
Usamos o modelo de aprendizado profundo Two-Tower para treinar e gerar embeddings de usuários e itens, posteriormente usados em indexação e consulta vetorial. Após treinar o modelo, podemos extrair os embeddings aprendidos de usuários e itens.
As duas etapas a seguir são opcionais: um modelo DLRM treinado para classificar os itens recuperados para recomendação e um feature store usado (neste caso, Feast) para armazenar e recuperar atributos de usuários e itens. Nós as incluímos para a completude do fluxo de trabalho em múltiplos estágios.
Por fim, exportamos os embeddings de usuários e itens para arquivos parquet, que podem posteriormente ser recarregados para criar um índice vetorial Milvus.
Construção e consulta do índice Milvus
O Milvus facilita a indexação vetorial e a busca por similaridade por meio de um “servidor” iniciado na máquina de inferência. Em nosso notebook nº 2, configuramos isso instalando via pip o servidor Milvus e o Pymilvus e, em seguida, iniciando o servidor com sua porta de escuta padrão. Em seguida, demonstramos a construção de um índice simples (IVF_FLAT) e a consulta nele usando as funções setup_milvus e query_milvus, respectivamente.
Benchmarking
Projetamos dois benchmarks para demonstrar o caso de uso de uma biblioteca rápida e eficiente de indexação/busca vetorial, como o Milvus.
Usando o Milvus para construir índices vetoriais com os dois conjuntos de embeddings que geramos: 1) embeddings de usuários para 7.3M usuários únicos, divididos em 85% conjunto de treino (para indexação) e 15% conjunto de teste (para consulta), e 2) embeddings de itens para 49K produtos (com uma divisão treino-teste de 50–50). Este benchmark é feito independentemente para cada conjunto de dados vetoriais, e os resultados são relatados separadamente.
Usando o Milvus para construir um índice vetorial para o conjunto de dados de embeddings de itens de 49K e consultando os 7.3M usuários únicos nesse índice para busca por similaridade.
Nestes benchmarks, usamos algoritmos de indexação IVFPQ e HNSW executados em GPU e CPU, juntamente com várias combinações de parâmetros. Detalhes estão disponíveis em nossa página do GitHub.
A compensação entre qualidade da busca e throughput é uma consideração importante de desempenho, especialmente em um ambiente de produção. O Milvus permite controle completo sobre os parâmetros de indexação para explorar essa compensação para um determinado caso de uso, a fim de alcançar melhores resultados de busca com a verdade fundamental. Isso pode significar maior custo computacional na forma de taxa de throughput reduzida ou consultas por segundo (QPS). Medimos a qualidade da busca ANN com uma métrica de recall e fornecemos curvas QPS-recall que demonstram a compensação. Então, pode-se decidir sobre um nível aceitável de qualidade de busca considerando os recursos computacionais ou os requisitos de latência/throughput do caso de negócio.
Além disso, observe o tamanho do lote de consultas (nq) usado em nossos benchmarks. Isso é útil em fluxos de trabalho nos quais várias solicitações simultâneas são enviadas para inferência (por exemplo, recomendações offline solicitadas e enviadas para uma lista de destinatários de e-mail ou recomendações online criadas pelo agrupamento de solicitações concorrentes que chegam e pelo processamento de todas elas de uma só vez). Dependendo do caso de uso, o TIS também pode ajudar a processar essas solicitações em lotes.
Resultados
Agora relatamos os resultados para os três conjuntos de benchmarks tanto em CPU quanto em GPU, usando os tipos de índice HNSW (somente CPU) e IVF_PQ (CPU e GPU) implementados pelo Milvus.
Busca por similaridade vetorial de Itens vs. Itens
Com este menor conjunto de dados, cada execução para uma determinada combinação de parâmetros usa 50% dos vetores de itens como vetores de consulta e consulta os 100 principais vetores similares do restante. HNSW e IVF_PQ produzem alto recall com as configurações de parâmetros testadas, nos intervalos de 0,958–1,0 e 0,665–0,997, respectivamente. Esse resultado sugere que HNSW tem melhor desempenho em relação ao recall, mas IVF_PQ com configurações pequenas de nlist produz recall altamente comparável. Também devemos observar que os valores de recall podem variar muito dependendo dos parâmetros de indexação e consulta. Os valores que relatamos foram obtidos após experimentação preliminar com intervalos gerais de parâmetros e aprofundamento em um subconjunto selecionado.
O tempo total para executar todas as consultas em CPU com HNSW para uma determinada combinação de parâmetros varia entre 5,22 e 5,33 sec.s (mais rápido à medida que m aumenta, relativamente inalterado com ef) e com IVF_PQ entre 13,67 e 14,67 sec.s (mais lento à medida que nlist e nprobe aumentam). A aceleração por GPU tem, de fato, um efeito perceptível, como visto na Figura 3.
A Figura 3 mostra a compensação recall-throughput em todas as execuções concluídas em CPU e GPU com este pequeno conjunto de dados usando IVF_PQ. Constatamos que a GPU fornece uma aceleração de 4x a 15x em todas as combinações de parâmetros testadas (maior aceleração à medida que nprobe aumenta). Isso é calculado tomando a razão do QPS das execuções em GPU sobre o QPS das execuções em CPU para cada combinação de parâmetros. No geral, este conjunto apresenta pouco desafio para CPU ou GPU e mostra perspectivas de aceleração adicional com os conjuntos de dados maiores, conforme discutido abaixo.
Figura 3. Aceleração por GPU com o algoritmo Milvus IVF_PQ executado em GPU NVIDIA A100 (busca por similaridade item-item)
Busca por similaridade vetorial de Usuários vs. Usuários
Com o segundo conjunto de dados muito maior (7,3M usuários), reservamos 85% (~6,2M) dos vetores como “train” (o conjunto de vetores a ser indexado), e os 15% restantes (~1,1M) como conjunto de vetores “test” ou de consulta. HNSW e IVF_PQ apresentam desempenho excepcionalmente bom nesse caso, com valores de recall de 0,884–1,0 e 0,922–0,999, respectivamente. Eles são, no entanto, computacionalmente muito mais exigentes, especialmente com IVF_PQ na CPU. O tempo total para executar todas as consultas em CPU com HNSW varia de 279,89 a 295,56 sec.s e com IVF_PQ de 3082,67 a 10932,33 sec.s. Observe que esses tempos de consulta são cumulativos para 1,1M vetores consultados, então pode-se dizer que uma única consulta contra o índice ainda é muito rápida.
No entanto, a consulta baseada em CPU pode não ser viável se o servidor de inferência espera muitos milhares de solicitações concorrentes para executar consultas contra um inventário de milhões de itens.
A GPU A100 oferece uma aceleração impressionante de 37x a 91x (média de 76,1x) em todas as combinações de parâmetros com IVF_PQ em termos de throughput (QPS), mostrada na Figura 4. Isso é consistente com o que observamos com o conjunto de dados pequeno, o que sugere que o desempenho da GPU escala razoavelmente bem usando o Milvus com milhões de vetores de embedding.
Figura 4. Aceleração da GPU com o algoritmo Milvus IVF_PQ executado na GPU NVIDIA A100 (busca de similaridade usuário-usuário)
A Figura 5 detalhada a seguir mostra o trade-off recall-QPS para todas as combinações de parâmetros testadas em CPU e GPU com IVF_PQ. Cada conjunto de pontos (superior para GPU, inferior para CPU) neste gráfico representa o trade-off enfrentado ao alterar parâmetros de indexação/consulta de vetores para alcançar maior recall às custas de menor throughput. Observe a perda considerável de QPS no caso da GPU à medida que se tenta alcançar níveis mais altos de recall.
Figura 5. Trade-off Recall-Throughput para todas as combinações de parâmetros testadas em CPU e GPU com IVF_PQ (usuários vs. usuários)
Busca de similaridade vetorial de usuários vs. itens
Por fim, consideramos outro caso de uso realista em que vetores de usuários são consultados em relação a vetores de itens (conforme demonstrado no Notebook 01 acima). Nesse caso, 49K vetores de itens são indexados, e 7,3M vetores de usuários são consultados individualmente para os 100 itens mais semelhantes.
É aqui que as coisas ficam interessantes, porque consultar 7,3M em lotes de 1000 contra um índice de 49K itens parece demorado na CPU tanto para HNSW quanto para IVF_PQ. A GPU parece lidar melhor com esse caso (veja a Figura 6). Os níveis mais altos de acurácia do IVF_PQ na CPU quando nlist = 100 são computados em cerca de 86 minutos em média, mas variam significativamente à medida que o valor de nprobe aumenta (51 min. quando nprobe = 5 vs. 128 min. quando nprobe = 20). A GPU NVIDIA A100 acelera consideravelmente o desempenho por um fator de 4x a 17x (acelerações maiores à medida que nprobe fica maior). Lembre-se de que o algoritmo IVF_PQ, por meio de sua técnica de quantização, também reduz a pegada de memória e fornece uma solução de busca ANN computacionalmente viável combinada com a aceleração por GPU.
Figura 6. Aceleração da GPU com o algoritmo Milvus IVF_PQ executado na GPU NVIDIA A100 (busca de similaridade usuário-item)
Semelhante à Figura 5, o trade-off recall-throughput é mostrado na Figura 7 para todas as combinações de parâmetros testadas com IVF_PQ. Aqui, ainda é possível ver como talvez seja necessário abrir mão ligeiramente de alguma acurácia na busca ANN em favor de maior throughput, embora as diferenças sejam muito menos perceptíveis, especialmente no caso das execuções em GPU. Isso sugere que se pode esperar níveis relativamente consistentemente altos de desempenho computacional com a GPU, ainda alcançando alto recall.
Figura 7. Trade-off Recall-Throughput para todas as combinações de parâmetros testadas em CPU e GPU com IVF_PQ (usuários vs. itens)
Conclusão
Teremos prazer em compartilhar algumas considerações finais se você chegou até aqui. Queremos lembrá-lo de que a complexidade e a natureza de múltiplos estágios dos Recsys modernos exigem desempenho e eficiência em cada etapa. Esperamos que este blog tenha dado a você motivos convincentes para considerar o uso de dois recursos críticos nos seus pipelines de RecSys:
A biblioteca Merlin Systems da NVIDIA Merlin permite que você integre facilmente o Milvus, um mecanismo eficiente de busca vetorial acelerado por GPU.
Use GPU para acelerar computações de indexação de bancos de dados vetoriais e busca ANN com tecnologia como RAPIDS RAFT.
Essas descobertas sugerem que a integração Merlin-Milvus apresentada é altamente performática e muito menos complexa do que outras opções para treinamento e inferência. Além disso, ambos os frameworks são desenvolvidos ativamente, e muitos novos recursos (por exemplo, novos índices de banco de dados vetorial acelerados por GPU pelo Milvus) são adicionados a cada lançamento. O fato de a busca por similaridade vetorial ser um componente crucial em vários fluxos de trabalho, como visão computacional, modelagem de grandes linguagens e sistemas de recomendação, torna esse esforço ainda mais valioso.
Para concluir, gostaríamos de agradecer a todos da Zilliz/Milvus e da Merlin e às equipes da RAFT que contribuíram para o esforço de produzir este trabalho e a postagem no blog. Esperamos receber notícias suas, caso você tenha a oportunidade de implementar Merlin e Milvus em seu recsys ou em outros fluxos de trabalho.
Continue lendo
Why and How to Migrate from Self-Hosted Milvus to Zilliz Cloud
A simple, step-by-step guide to migrating from Milvus to Zilliz Cloud. Learn both endpoint and backup methods for a smooth, scalable vector database migration.
The Real Bottlenecks in Autonomous Driving — And How AI Infrastructure Can Solve Them
Autonomous driving faces a data bottleneck. Learn how AI-native vector databases like Zilliz solve scale, cost, and insight challenges across AV pipelines.
Building RAG Pipelines for Real-Time Data with Cloudera and Milvus
explore how Cloudera can be integrated with Milvus to effectively implement some of the key functionalities of RAG pipelines.