Эффективный поиск сходства векторов в рабочих процессах рекомендательных систем с использованием Milvus и NVIDIA Merlin
Этот пост был впервые опубликован на Medium-канале NVIDIA Merlin, отредактирован и переопубликован здесь с разрешения. Он был совместно написан Burcin Bozkaya и William Hicks из NVIDIA, а также Filip Haltmayer и Li Liu из Zilliz.
Введение
Современные рекомендательные системы (Recsys) состоят из конвейеров обучения/инференса, включающих несколько этапов: прием данных, предварительную обработку данных, обучение моделей и настройку гиперпараметров для поиска, фильтрации, ранжирования и оценки релевантных элементов. Важнейшим компонентом конвейера рекомендательной системы является извлечение или обнаружение объектов, наиболее релевантных пользователю, особенно при наличии больших каталогов элементов. Этот этап обычно включает поиск приближенных ближайших соседей (ANN) по индексированной базе данных низкоразмерных векторных представлений (т. е. эмбеддингов) атрибутов продуктов и пользователей, созданных с помощью моделей глубокого обучения, которые обучаются на взаимодействиях между пользователями и продуктами/сервисами.
NVIDIA Merlin, open-source фреймворк, разработанный для обучения сквозных моделей с целью предоставления рекомендаций в любом масштабе, интегрируется с эффективным индексом векторной базы данных и фреймворком поиска. Одним из таких фреймворков, недавно привлекших большое внимание, является Milvus, open-source векторная база данных, созданная Zilliz. Она предлагает быстрые возможности индексирования и запросов. Milvus недавно добавил поддержку ускорения на GPU, которая использует GPU NVIDIA для поддержки AI-рабочих процессов. Поддержка ускорения на GPU — отличная новость, потому что ускоренная библиотека векторного поиска делает возможными быстрые параллельные запросы, положительно влияя на требования к задержке в современных рекомендательных системах, где разработчики ожидают множество одновременных запросов. У Milvus более 5 млн загрузок Docker, ~23 тыс. звезд на GitHub (по состоянию на сентябрь 2023 года), более 5 000 корпоративных клиентов, и он является ключевым компонентом многих приложений (см. варианты использования cases).
В этом блоге демонстрируется, как Milvus работает с фреймворком Merlin Recsys во время обучения и инференса. Мы показываем, как Milvus дополняет Merlin на этапе извлечения элементов благодаря высокоэффективному top-k поиску по векторным эмбеддингам и как его можно использовать с NVIDIA Triton Inference Server (TIS) во время инференса (см. рисунок 1). Наши результаты бенчмарка показывают впечатляющее ускорение от 37x до 91x с GPU-ускоренным Milvus, который использует NVIDIA RAFT с векторными эмбеддингами, сгенерированными Merlin Models. Код, который мы используем для демонстрации интеграции Merlin-Milvus, и подробные результаты бенчмарка, а также библиотека, которая помогла провести наше бенчмарк-исследование, доступны здесь.
Рисунок 1. Многоэтапная рекомендательная система с фреймворком Milvus, участвующим в этапе извлечения. Источник оригинальной многоэтапной иллюстрации: этот пост в блоге.
Проблемы, с которыми сталкиваются рекомендательные системы
Учитывая многоэтапную природу рекомендательных систем и доступность различных компонентов и библиотек, которые они интегрируют, существенной задачей является бесшовная интеграция всех компонентов в end-to-end pipeline. Мы стремимся показать в наших примерных ноутбуках, что интеграцию можно выполнить с меньшими усилиями.
Еще одна задача рабочих процессов рекомендательных систем — ускорение определенных частей pipeline. Хотя GPU, как известно, играют огромную роль в обучении больших нейронных сетей, векторные базы данных и ANN search начали использовать GPU лишь недавно. С ростом размеров товарных каталогов e-commerce или баз данных стриминговых медиа, а также количества пользователей, использующих эти сервисы, CPU должны обеспечивать необходимую производительность для обслуживания миллионов пользователей в эффективных Recsys workflows. GPU-ускорение в других частях pipeline стало необходимым для решения этой задачи. Решение в этом блоге отвечает на этот вызов, показывая, что ANN search эффективен при использовании GPU.
Технологические стеки для решения
Давайте начнем с обзора некоторых основных принципов, необходимых для выполнения нашей работы.
NVIDIA Merlin: open-source библиотека с высокоуровневыми API, ускоряющая рекомендательные системы на NVIDIA GPUs.
NVTabular: для предварительной обработки входных табличных данных и feature engineering.
Merlin Models: для обучения моделей deep learning и, в данном случае, изучения embedding-векторов пользователей и объектов на основе данных о взаимодействиях пользователей.
Merlin Systems: для объединения рекомендательной модели на базе TensorFlow с другими элементами (например, feature store, ANN search с Milvus) для обслуживания через TIS.
Triton Inference Server: для этапа inference, на котором передается вектор признаков пользователя и генерируются рекомендации продуктов.
Containerization: все вышеперечисленное доступно через container(s), предоставляемые NVIDIA в NGC catalog. Мы использовали Merlin TensorFlow 23.06 container.
Milvus 2.3: для выполнения GPU-ускоренного векторного индексирования и запросов.
Milvus 2.2.11: то же, что и выше, но на CPU.
Pymilvus SDK: для подключения к серверу Milvus, создания индексов векторной базы данных и выполнения запросов через Python-интерфейс.
Feast: для сохранения и извлечения атрибутов пользователей и объектов в (open source) feature store как части нашего end-to-end RecSys pipeline.
Несколько базовых библиотек и фреймворков также используются под капотом. Например, Merlin опирается на другие библиотеки NVIDIA, такие как cuDF и Dask, обе доступны в рамках RAPIDS cuDF. Аналогично, Milvus использует NVIDIA RAFT для примитивов GPU-ускорения и модифицированные библиотеки, такие как HNSW и FAISS, для поиска.
Понимание векторных баз данных и Milvus
Приблизительный поиск ближайших соседей (ANN) — это функциональность, с которой реляционные базы данных не справляются. Реляционные БД предназначены для работы с табличными данными с предопределенными структурами и напрямую сопоставимыми значениями. Индексы реляционных баз данных опираются на это, чтобы сравнивать данные и создавать структуры, которые используют знание о том, меньше или больше каждое значение другого. Векторы эмбеддингов нельзя напрямую сравнивать друг с другом таким образом, поскольку нам нужно знать, что представляет каждое значение в векторе. Они не могут сказать, обязательно ли один вектор меньше другого. Единственное, что мы можем сделать, — вычислить расстояние между двумя векторами. Если расстояние между двумя векторами мало, мы можем предположить, что признаки, которые они представляют, похожи, а если оно велико, мы можем предположить, что данные, которые они представляют, сильнее различаются. Однако эти эффективные индексы имеют свою цену; вычисление расстояния между двумя векторами требует больших вычислительных затрат, а векторные индексы не являются легко адаптируемыми и иногда не поддаются изменению. Из-за этих двух ограничений интеграция таких индексов в реляционные базы данных более сложна, поэтому необходимы специализированные векторные базы данных.
Milvus был создан для решения проблем, с которыми реляционные базы данных сталкиваются при работе с векторами, и был изначально спроектирован для обработки этих векторов эмбеддингов и их индексов в большом масштабе. Чтобы соответствовать статусу cloud-native, Milvus разделяет вычисления и хранение, а также различные вычислительные задачи — запросы, подготовку данных и индексирование. Пользователи могут масштабировать каждую часть базы данных для обработки разных сценариев использования, будь то интенсивная вставка данных или интенсивный поиск. Если происходит большой приток запросов на вставку, пользователь может временно масштабировать индексные узлы горизонтально и вертикально, чтобы справиться с загрузкой данных. Аналогично, если данные не загружаются, но выполняется много поисковых запросов, пользователь может сократить количество индексных узлов и вместо этого масштабировать узлы запросов для повышения пропускной способности. Такая архитектура системы (см. рисунок 2) потребовала от нас мыслить в парадигме параллельных вычислений, что привело к созданию системы, оптимизированной для вычислений, с множеством возможностей для дальнейшей оптимизации.
Рисунок 2. Архитектура системы Milvus
Milvus также использует множество передовых библиотек индексирования, чтобы предоставить пользователям как можно больше возможностей настройки их системы. Он улучшает их, добавляя возможность обработки CRUD-операций, потоковых данных и фильтрации. Далее мы обсудим, чем отличаются эти индексы и каковы плюсы и минусы каждого из них.
Пример решения: интеграция Milvus и Merlin
Пример решения, который мы представляем здесь, демонстрирует интеграцию Milvus с Merlin на этапе извлечения элементов (когда k наиболее релевантных элементов извлекаются с помощью ANN-поиска). Мы используем реальный набор данных из соревнования RecSys, описанный ниже. Мы обучаем deep learning модель Two-Tower, которая изучает векторные эмбеддинги для пользователей и элементов. В этом разделе также представлен план нашей работы по бенчмаркингу, включая метрики, которые мы собираем, и диапазон параметров, которые мы используем.
Наш подход включает:
загрузку и предварительную обработку данных
обучение deep learning модели Two-Tower
построение индекса Milvus
поиск по сходству в Milvus
Мы кратко описываем каждый шаг и отсылаем читателя к нашим notebooks за подробностями.
Набор данных
YOOCHOOSE GmbH предоставляет набор данных, который мы используем в этом интеграционном и бенчмарк-исследовании для RecSys 2015 challenge, и он доступен на Kaggle. Он содержит события кликов/покупок пользователей из европейского онлайн-ритейлера с такими атрибутами, как ID сессии, временная метка, ID товара, связанный с кликом/покупкой, и категория товара, доступные в файле yoochoose-clicks.dat. Сессии независимы, и нет признаков возвращающихся пользователей, поэтому мы считаем каждую сессию принадлежащей отдельному пользователю. Набор данных содержит 9,249,729 уникальных сессий (пользователей) и 52,739 уникальных товаров.
Загрузка и предварительная обработка данных
Инструмент, который мы используем для предварительной обработки данных, — NVTabular, ускоренный на GPU, высокомасштабируемый компонент Merlin для feature engineering и предварительной обработки. Мы используем NVTabular для считывания данных в память GPU, перестановки признаков по мере необходимости, экспорта в parquet-файлы и создания train-validation split для обучения. В результате получается 7,305,761 уникальный пользователь и 49,008 уникальных товаров для обучения. Мы также категоризируем каждый столбец и его значения в целочисленные значения. Набор данных теперь готов для обучения с моделью Two-Tower.
Обучение модели
Мы используем модель глубокого обучения Two-Tower для обучения и генерации эмбеддингов пользователей и товаров, которые позже используются для векторного индексирования и запросов. После обучения модели мы можем извлечь выученные эмбеддинги пользователей и товаров.
Следующие два шага являются необязательными: модель DLRM, обученная ранжировать извлеченные товары для рекомендации, и используемое хранилище признаков (в данном случае Feast) для хранения и извлечения признаков пользователей и товаров. Мы включаем их для полноты многоэтапного рабочего процесса.
Наконец, мы экспортируем эмбеддинги пользователей и товаров в parquet-файлы, которые позже можно повторно загрузить для создания векторного индекса Milvus.
Создание индекса Milvus и выполнение запросов к нему
Milvus облегчает векторное индексирование и поиск по сходству через «сервер», запущенный на машине инференса. В нашем notebook #2 мы настраиваем это, устанавливая через pip сервер Milvus и Pymilvus, а затем запускаем сервер с его стандартным портом прослушивания. Далее мы демонстрируем создание простого индекса (IVF_FLAT) и выполнение запросов к нему с использованием функций setup_milvus и query_milvus соответственно.
Бенчмаркинг
Мы разработали два бенчмарка, чтобы продемонстрировать аргументы в пользу использования быстрой и эффективной библиотеки векторного индексирования/поиска, такой как Milvus.
Использование Milvus для создания векторных индексов с двумя наборами эмбеддингов, которые мы сгенерировали: 1) эмбеддинги пользователей для 7.3M уникальных пользователей, разделенные как 85% train set (для индексирования) и 15% test set (для запросов), и 2) эмбеддинги товаров для 49K продуктов (с разделением train-test 50–50). Этот бенчмарк выполняется независимо для каждого векторного набора данных, а результаты приводятся отдельно.
Использование Milvus для создания векторного индекса для набора данных эмбеддингов 49K товаров и выполнения запросов 7.3M уникальных пользователей к этому индексу для поиска по сходству.
В этих бенчмарках мы использовали алгоритмы индексирования IVFPQ и HNSW, выполняемые на GPU и CPU, а также различные комбинации параметров. Подробности доступны на нашей странице GitHub.
Компромисс между качеством поиска и пропускной способностью является важным фактором производительности, особенно в производственной среде. Milvus позволяет полностью контролировать параметры индексирования, чтобы исследовать этот компромисс для конкретного сценария использования и добиться лучших результатов поиска с использованием эталонной истины. Это может означать увеличение вычислительных затрат в виде снижения пропускной способности или количества запросов в секунду (QPS). Мы измеряем качество ANN-поиска с помощью метрики полноты и предоставляем кривые QPS-полнота, демонстрирующие этот компромисс. Затем можно выбрать приемлемый уровень качества поиска с учетом вычислительных ресурсов или требований бизнес-сценария к задержке/пропускной способности.
Также обратите внимание на размер пакета запросов (nq), используемый в наших бенчмарках. Это полезно в рабочих процессах, где несколько одновременных запросов отправляются на инференс (например, офлайн-рекомендации запрашиваются и отправляются списку получателей электронной почты или онлайн-рекомендации создаются путем объединения поступающих параллельных запросов и их одновременной обработки). В зависимости от сценария использования TIS также может помочь обрабатывать эти запросы пакетами.
Результаты
Теперь мы приводим результаты для трех наборов бенчмарков как на CPU, так и на GPU, используя типы индексов HNSW (только CPU) и IVF_PQ (CPU и GPU), реализованные в Milvus.
Векторный поиск сходства: товары vs. товары
С этим самым маленьким набором данных каждый запуск для заданной комбинации параметров берет 50% векторов товаров в качестве векторов запросов и запрашивает топ-100 похожих векторов из оставшейся части. HNSW и IVF_PQ обеспечивают высокую полноту при протестированных настройках параметров, в диапазоне 0,958–1,0 и 0,665–0,997 соответственно. Этот результат показывает, что HNSW демонстрирует лучшую полноту, но IVF_PQ с малыми настройками nlist дает весьма сопоставимую полноту. Также следует отметить, что значения полноты могут сильно варьироваться в зависимости от параметров индексирования и запросов. Значения, которые мы приводим, были получены после предварительных экспериментов с общими диапазонами параметров и дальнейшего приближения к выбранному подмножеству.
Общее время выполнения всех запросов на CPU с HNSW для заданной комбинации параметров находится в диапазоне от 5,22 до 5,33 сек. (быстрее при увеличении m, относительно без изменений при ef), а с IVF_PQ — от 13,67 до 14,67 сек. (медленнее при увеличении nlist и nprobe). Ускорение на GPU действительно оказывает заметный эффект, как видно на Рисунке 3.
На Рисунке 3 показан компромисс между полнотой и пропускной способностью по всем запускам, выполненным на CPU и GPU с этим небольшим набором данных с использованием IVF_PQ. Мы обнаружили, что GPU обеспечивает ускорение от 4x до 15x по всем протестированным комбинациям параметров (большее ускорение при увеличении nprobe). Это рассчитывается как отношение QPS в запусках на GPU к QPS в запусках на CPU для каждой комбинации параметров. В целом этот набор представляет небольшую сложность для CPU или GPU и демонстрирует перспективы дальнейшего ускорения на более крупных наборах данных, как обсуждается ниже.
Рисунок 3. Ускорение на GPU с алгоритмом Milvus IVF_PQ, выполняющимся на GPU NVIDIA A100 (поиск сходства товар-товар)
Векторный поиск сходства: пользователи vs. пользователи
Для гораздо более крупного второго набора данных (7,3 млн пользователей) мы выделяем 85% (~6,2 млн) векторов как “train” (набор векторов для индексирования), а оставшиеся 15% (~1,1 млн) — как “test” или набор векторов запросов. HNSW и IVF_PQ в этом случае работают исключительно хорошо, со значениями полноты 0,884–1,0 и 0,922–0,999 соответственно. Однако они требуют значительно больше вычислительных ресурсов, особенно IVF_PQ на CPU. Общее время выполнения всех запросов на CPU с HNSW составляет от 279,89 до 295,56 сек., а с IVF_PQ — от 3082,67 до 10932,33 сек. Обратите внимание, что это время запросов является суммарным для 1,1 млн запрошенных векторов, поэтому можно сказать, что один запрос к индексу все еще выполняется очень быстро.
Однако запросы на базе CPU могут оказаться нежизнеспособными, если сервер инференса ожидает многие тысячи одновременных запросов для выполнения поиска по инвентарю из миллионов товаров.
Графический процессор A100 обеспечивает впечатляющее ускорение от 37x до 91x (в среднем 76,1x) по всем комбинациям параметров с IVF_PQ с точки зрения пропускной способности (QPS), как показано на рисунке 4. Это согласуется с тем, что мы наблюдали на небольшом наборе данных, что говорит о том, что производительность GPU достаточно хорошо масштабируется при использовании Milvus с миллионами векторов эмбеддингов.
Рисунок 4. Ускорение GPU с алгоритмом Milvus IVF_PQ, работающим на GPU NVIDIA A100 (поиск сходства пользователь-пользователь)
На следующем подробном рисунке 5 показан компромисс между recall и QPS для всех комбинаций параметров, протестированных на CPU и GPU с IVF_PQ. Каждый набор точек (верхний для GPU, нижний для CPU) на этой диаграмме отображает компромисс, возникающий при изменении параметров индексации/запросов векторов для достижения более высокого recall за счет снижения пропускной способности. Обратите внимание на значительную потерю QPS в случае GPU при попытке достичь более высоких уровней recall.
Рисунок 5. Компромисс Recall-Throughput для всех комбинаций параметров, протестированных на CPU и GPU с IVF_PQ (пользователи против пользователей)
Поиск векторного сходства: пользователи против объектов
Наконец, рассмотрим еще один реалистичный сценарий использования, где векторы пользователей запрашиваются по векторам объектов (как показано в Notebook 01 выше). В этом случае индексируются 49 тыс. векторов объектов, а для каждого из 7,3 млн векторов пользователей выполняется запрос топ-100 наиболее похожих объектов.
Именно здесь становится интересно, поскольку выполнение запросов для 7,3 млн векторов пакетами по 1000 к индексу из 49 тыс. объектов оказывается трудоемким на CPU как для HNSW, так и для IVF_PQ. GPU, похоже, справляется с этим случаем лучше (см. рисунок 6). Самые высокие уровни точности IVF_PQ на CPU при nlist = 100 вычисляются в среднем примерно за 86 минут, но значительно варьируются по мере увеличения значения nprobe (51 мин. при nprobe = 5 против 128 мин. при nprobe = 20). GPU NVIDIA A100 значительно ускоряет производительность в 4x–17x раз (более высокие ускорения при увеличении nprobe). Помните, что алгоритм IVF_PQ благодаря своей технике квантования также уменьшает объем потребляемой памяти и предоставляет вычислительно жизнеспособное решение для ANN-поиска в сочетании с ускорением на GPU.
Рисунок 6. Ускорение GPU с алгоритмом Milvus IVF_PQ, работающим на GPU NVIDIA A100 (поиск сходства пользователь-объект)
Подобно рисунку 5, компромисс между recall и throughput показан на рисунке 7 для всех комбинаций параметров, протестированных с IVF_PQ. Здесь все еще можно увидеть, что может потребоваться немного пожертвовать точностью ANN-поиска в пользу увеличения пропускной способности, хотя различия гораздо менее заметны, особенно в случае запусков на GPU. Это говорит о том, что можно ожидать относительно стабильно высоких уровней вычислительной производительности с GPU, при этом сохраняя высокий recall.
Рисунок 7. Компромисс Recall-Throughput для всех комбинаций параметров, протестированных на CPU и GPU с IVF_PQ (пользователи против объектов)
Заключение
Мы с радостью поделимся несколькими заключительными замечаниями, если вы дочитали до этого места. Мы хотим напомнить вам, что сложность и многоэтапность современных Recsys требуют производительности и эффективности на каждом этапе. Надеемся, этот блог дал вам убедительные причины рассмотреть использование двух важных возможностей в ваших конвейерах RecSys:
Библиотека Merlin Systems от NVIDIA Merlin позволяет легко подключить Milvus, эффективный движок векторного поиска с ускорением на GPU.
Используйте GPU для ускорения вычислений при индексировании векторных баз данных и ANN-поиске с технологией, такой как RAPIDS RAFT.
Эти результаты показывают, что представленная интеграция Merlin-Milvus обладает высокой производительностью и гораздо менее сложна, чем другие варианты для обучения и инференса. Кроме того, оба фреймворка активно развиваются, и в каждом релизе добавляется множество новых функций (например, новые GPU-ускоренные индексы векторных баз данных от Milvus). Тот факт, что поиск по сходству векторов является важнейшим компонентом различных рабочих процессов, таких как компьютерное зрение, большие языковые модели и рекомендательные системы, делает эти усилия еще более ценными.
В заключение мы хотели бы поблагодарить всех из Zilliz/Milvus, Merlin и команд RAFT, кто внес вклад в работу над этим проектом и публикацией в блоге. Будем рады услышать от вас, если у вас будет возможность внедрить Merlin и Milvus в ваши recsys или другие рабочие процессы.
Читать далее
The Great AI Agent Protocol Race: Function Calling vs. MCP vs. A2A
Compare Function Calling, MCP, and A2A protocols for AI agents. Learn which standard best fits your development needs and future-proof your applications.
Vector Databases vs. Graph Databases
Use a vector database for AI-powered similarity search; use a graph database for complex relationship-based queries and network analysis.
Zilliz Cloud BYOC Upgrades: Bring Enterprise-Grade Security, Networking Isolation, and More
Discover how Zilliz Cloud BYOC brings enterprise-grade security, networking isolation, and infrastructure automation to vector database deployments in AWS