От ANN к честному KNN на GPU: как мы пересобрали отбор кандидатов в рекомендациях Ozon

Привет! Мы команда рекомендательной системы Ozon, и сегодня мы хотим рассказать о нашем пути от приближённого поиска соседей (ANN) к точному KNN на GPU. Этот материал для тех, кто работает с рекомендациями, поиском или большими векторными пространствами и задумывается о том, можно ли выжать максимум из железа, не жертвуя качеством.

В индустрии уже есть примеры, когда команды рекомендаций уходят от готовых ANN‑индексов к более специализированным GPU‑решениям скоринга. Мы же опишем, как это выглядит в масштабах российского e‑commerce, и расскажем о результатах A/B‑тестов. Сразу оговоримся: это не «Hello, world» с парой тысяч векторов, а продакшен на десятки миллионов пользователей и сотни миллионов товаров, где каждый час пайплайна и каждый процент recall имеют цену.

Контекст: рекомендации в Ozon

Рекомендательная система Ozon ежедневно обслуживает десятки миллионов пользователей и сотни миллионов товаров. Если упростить, то отбор кандидатов устроен так: по поведению пользователя мы строим несколько векторных представлений — эмбеддингов (ALS, нейросетевые), а затем для каждого эмбеддинга ищем ближайших кандидатов в общем каталоге товаров. Близость векторов здесь — это прокси релевантности: чем ближе вектор товара к вектору пользователя (или другого товара), тем выше шанс, что товар «зайдёт».

Мы используем два основных источника кандидатов:

• user2item — подбираем товары под конкретного пользователя по векторам его интересов;

• item2item — подбираем похожие товары к заданному товару (карточка, корзина, избранное, заказы).

В текущем рассказе мы сосредоточимся на сценарии user2item как наиболее массовом и показательном. Item2item (поиск аналогов, аксессуаров) технически близок — те же GPU‑пайплайны с другими формулами близости, но мы оставим его за рамками этой статьи ради фокуса.

Важно понимать масштаб задачи. В user2item мы фактически решаем задачу «для каждого из десятков миллионов пользователей найти лучшие товары из каталога в сотни миллионов позиций». Это декартово произведение, которое в лоб посчитать было вычислительно сложно, поэтому исторически здесь и применялись приближённые методы.

Исторически отбор кандидатов в user2item мы строили на алгоритме приближённого поиска HNSW поверх Spark, а финальное ранжирование делала CatBoost‑модель. В упрощённом виде ежедневный пайплайн выглядел так.

1. Определяем, для кого и из чего считаем. Сначала отбираем список пользователей, для которых будем искать кандидатов, — берём верхние K миллионов по количеству совершённых действий. Параллельно формируем базовый список товаров, доступных для ранжирования. В него попадают товары, доступные к продаже и находящиеся в стоках.

2. Собираем кандидатов из двух источников. Для каждого пользователя нужно сформировать пары «user_client_id — item_id» под последующее ранжирование. Кандидаты приходят из двух мест:

   • ANN‑поиск — ищем по 100 наиболее релевантных товаров на каждую категорию первого уровня.

   • Item2item — собираем для пользователя его исторические клики, добавления в избранное и в корзину, заказы за последние N месяцев (назовём это ключами для item2item) и для каждого ключа добираем похожие товары. Сколько кандидатов добавлять на один ключ, рассчитываем исходя из заданного общего количества item2item‑кандидатов и числа ключей.
     Из обоих источников оставляем только тех кандидатов, которых пользователь не видел.

3. Готовим фичи и собираем пул. Параллельно с формированием пар собираем фичи: пользовательские счётчики и их агрегаты, фичи ценовых предпочтений, информацию о сегменте пользователя (socdem, RFM), конверсионность товаров. Затем джойним пары и фичи в единый пул — это один из самых тяжёлых стейджей, который работает до 9–10 часов.

4. Скорим и ранжируем. Пары скорит CatBoost‑модель, обученная на ранжирующий лосс. После этого формируем топ товаров для каждого пользователя: сортируем по скору и применяем правило разнообразия — ограничиваем количество товаров одного типа в топе через управляемый параметр. При необходимости можно навешивать и фильтры по колонкам‑счётчикам.

5. Выгружаем в хранилища сервисов. Финальные подборки складываем в Scylla, откуда их забирает онлайн‑слой рекомендаций.

Почему ANN перестал нас устраивать

HNSW — отличный алгоритм, и долгое время он нас полностью устраивал. Но по мере роста каталога и усложнения продуктовых требований накопился ряд ограничений.

Приближённость выдачи. ANN по определению не гарантирует, что найдены действительно ближайшие соседи. На «голове» распределения это почти незаметно, но на «хвостах» — у редких товаров и пользователей с нетипичным поведением — recall проседает. А именно там часто и прячется самый ценный, неочевидный кандидат, ради которого рекомендации и существуют.

Сложность встраивания бизнес‑логики. Нам мало просто найти ближайшие векторы. Нужно тут же отфильтровать недоступные для пользователя товары, применить штрафы за дизлайки и уже купленное, учесть рекламные маски, обеспечить разнообразие выдачи. В классическом ANN‑индексе вся эта логика живёт снаружи: сначала достаём «сырых» кандидатов с запасом, потом отдельными Spark‑джобами джойним фичи, фильтруем и пересортировываем. Именно отсюда и берётся тот самый тяжёлый постпроцессинг с сотнями миллиардов строк — мы вынуждены поднимать огромный промежуточный пул просто потому, что не можем отсечь лишнее на этапе поиска.

Дороговизна и негибкость пересборки индекса. Индекс HNSW нужно строить и поддерживать, а изменение базы товаров или геометрии пространства требует его пересборки. Это сковывает эксперименты: быстро проверить гипотезу «давайте иначе считать близость» или «давайте подмешаем ещё один эмбеддинг» без дорогой переиндексации не получается.

В какой‑то момент мы задали себе еретический по меркам ANN вопрос: а что, если посчитать всё честно? Современная GPU способна перемножать матрицы с колоссальной пропускной способностью, и полный перебор «каждый запрос против всей базы» — это, по сути, одно большое матричное умножение, под которое железо и заточено. Если мы всё равно упираемся в GPU и постпроцессинг, то прямой перебор может оказаться и точнее, и в сумме дешевле — особенно если встроить бизнес‑логику прямо в вычисление.

Решение: честный KNN на GPU

Мы решили попробовать честный KNN на GPU — полный перебор всех векторов с отбором топ‑K на каждый ключ. Звучит как «в лоб», но именно прямой перебор на современных картах оказывается и быстрым, и куда более гибким, чем поддержка ANN‑индекса.

Первая версия: ежедневный офлайн‑расчёт на GPU

Мы начали с простой идеи: загрузить все эмбеддинги товаров (назовём их базой, base) в память GPU и для каждого ключа (query — напомним, что ключом может быть пользователь или товар) за один проход посчитать расстояния и взять ближайших соседей.

Концептуально всё сводится к перемножению матрицы запросов на матрицу базы и взятию топ‑K по строкам. Но для сотен миллионов товаров и десятков миллионов пользователей наивная реализация не помещается ни в память, ни в разумное время: матрица скоров query × base — это сотни триллионов чисел. Поэтому данные нужно грамотно шардировать (резать базу) и реплицировать (резать запросы), а сам топ считать потоково, не сохраняя полную матрицу расстояний.

Пайплайн в коде. Мы реализовали пайплайн на PySpark + PyTorch, который ежедневно запускается в Hadoop‑кластере и использует выделенные GPU‑ноды. Его упрощённая схема:

1. Подготовка данных на PySpark и сохранение их в HDFS.

2. Распаковка данных на GPU‑нодах.

3. Расчёт скоров и отбор топ‑K поверх них, сохранение кандидатов в HDFS.

4. Выгрузка подборок в рантайм в key‑value‑хранилище.

Алгоритмическая эволюция движка

Самое интересное началось дальше. Раз весь проход по базе мы и так контролируем сами, то всю бизнес‑логику, которая раньше жила в тяжёлом постпроцессинге, можно встроить прямо внутрь вычисления топа — там же, где считаются расстояния, и без выгрузки промежуточных миллиардов строк обратно в Spark.

Булевы маски и фильтрация. Первым делом добавили фильтры по булевым маскам. Например, в base есть колонка рекламности товара. За один проход мы можем посчитать сразу два топа: общий и только по рекламным товарам, без повторного пересчёта расстояний. Аналогично выделяем отдельные подборки для продуктов Ozon Селект и Ozon Fresh, просто фильтруя по соответствующей маске. Получается, что одна и та же дорогая операция расчёта скоров переиспользуется под несколько продуктовых нужд.

Фильтрация по истории пользователя. Просмотренные товары исключаем из выдачи: на вход передаётся список id элементов из base, и через sparse tensor мы убираем этих кандидатов из подборки ещё до отбора топа. Это дешевле, чем фильтровать готовый топ постфактум и рисковать тем, что после фильтрации в нём останется слишком мало позиций.

Фильтрация по категориям. Дизлайки мы тоже учитываем на уровне категорий. Здесь пришлось придумать специальный механизм хранения индекса: для каждого пользователя держим массив категорий, значения в котором отсортированы в том же порядке, что и товары в base: по категориям первого и последнего уровня. Благодаря согласованному порядку мы можем на GPU быстро «наложить» категориальные штрафы на вектор скоров без дорогих джойнов и пересортировок.

Управление разнообразием (diversity). Чтобы в топе не было перекоса в одну категорию, внедрили возможность применения эвристик: не больше N товаров из одной категории. Алгоритм: сперва набираем расширенный топ с учётом лимитов на категории, а затем финально сортируем.

Кастомные вычисления скоров близости. Для простоты расчёта в качестве скора часто используют dot‑произведение или косинусную близость, но в некоторых сценариях релевантность — это более сложная функция от пары эмбеддингов. Мы обобщили механизм вычисления скора по паре векторов: вместо готовой метрики близости вызывается небольшая нейросеть, которая принимает на вход эмбеддинги пользователя и товара и вычисляет скор. Это позволяет, в частности, комбинировать несколько эмбеддингов пользователя — например, отражающих интересы пользователя по разным категориям, — нелинейным образом прямо во время GPU‑прохода.

KNN only: несколько эмбеддингов в одном проходе

Современные рекомендательные модели часто порождают несколько векторных представлений одного объекта, например отдельные нейросетевые модели и модели ALS, отдельные эмбеддинги для интересов, для ценовой чувствительности и так далее. Раньше каждое такое представление означало отдельный источник кандидатов и отдельную джобу. Мы же добавили возможность считать KNN сразу по нескольким векторам и комбинировать результаты с помощью нейросетевой модели.

Финальный скор кандидата считается через применение нейросети поверх скоров по отдельным эмбеддингам. Более того, в формуле можно смешивать скор модели с различными бизнес‑требованиями, например с рекламными ставками. Это и есть механизм сквозного ранжирования товаров на первой стадии ранжирования: качество (близость эмбеддингов) и бизнес‑сигналы (ставки) живут в одной формуле и настраиваются согласованно.

Техническая эволюция: быстрее и ближе к real‑time

Шардирование и реплики

Практически сразу встал вопрос, как масштабировать систему одновременно по двум осям: по размеру базы и по числу запросов. Для этого мы адаптировали код под механизмы шардирования и реплицирования, которые позволяют параллелить вычисления: шардирование — по базе, реплицирование — по запросам. Всё это раскладывается на несколько GPU.

Логика такая: при загрузке данных мы разбиваем базу товаров на M шардов, а все запросы — на N реплик. Каждая пара «шард — реплика» занимает одну GPU, при этом на одну реплику может приходиться несколько шардов. Запросы с одной реплики идут на каждый шард этой реплики, результаты расчёта собираются со всех шардов и переранжируются на CPU с учётом разнообразия. Реплики при этом работают полностью независимо друг от друга.

Что это даёт на практике:

• Реплицирование параллелит запросы и кратно ускоряет расчёт — чем больше реплик, тем быстрее обрабатываются все пользователи.

• Шардирование распределяет базу по нескольким GPU, благодаря чему мы можем её утяжелять, например добавлять товары и расширять эмбеддинги, не упираясь в память одной карты.

Ускорение через CUDA‑ядра

Стандартные torch‑операции (torch.mm, torch.topk) хороши для прототипа, но bottleneck’ом оказался именно отбор top‑K — после расчёта матрицы скоров нужно для каждого пользователя пройтись по каждой категории, применить квоты и diversity. Этот шаг мы переписали на CUDA‑ядра: heap‑based и segmented gatherer’ы, которые собирают top‑K прямо на GPU с учётом иерархических квот по категориям разного уровня и булевых масок, например рекламности. Это позволило ускорить этап отбора top‑K в 2,5–3 раза по сравнению с torch.topk и постобработкой на CPU.

Выход в near real‑time (NRT) и Triton

Джобы, которые обновляют подборки раз в сутки, — это хорошо для офлайн‑сценариев, но хочется реагировать на действия пользователя здесь и сейчас, поэтому мы пошли в NRT‑расчёты подборок пользователя и обернули всё в сервис на базе Nvidia Triton.

Онлайн‑сервис на Triton Inference Server принимает на вход эмбеддинги пользователя, параметры подборки и отдаёт готовый топ‑K. При этом эмбеддинги товаров закешированы прямо в GPU‑памяти, чтобы не платить за их загрузку на каждый запрос. Эмбеддинги базы товаров обновляются раз в день: мы подгружаем в Triton новую версию модели, в которую эти эмбеддинги уже «зашиты». Таким образом, актуализация базы — это, по сути, выкатка новой версии модели. Приятный побочный эффект — единый механизм версионирования и отката: откатить базу можно так же, как откатывают любую модель. Тяжёлые эмбеддинги пользователя, использующиеся для отбора кандидатов, на текущий момент пересчитываются ежедневно, но даже в такой схеме получилось показать улучшения по метрикам за счёт учёта RT истории показанных товаров. С возможностью пересчитывать подборку по мере изменения эмбеддинга пользователя мы будем экспериментировать в следующих итерациях. Батчирование запросов мы отдаём самому Triton — он собирает входящие запросы в группы, что позволяет эффективнее утилизировать GPU без ручного управления очередями на стороне сервиса.

Результаты

Переход от ANN‑поиска с ранжированием CatBoost моделью на Spark‑кластере к единой модели retrieval as ranking на GPU дал эффект сразу в двух плоскостях — инфраструктурной и продуктовой.

Инфраструктура и TTM. Главный технический выигрыш — мы заметно упростили и ускорили пайплайн:

• Time‑to‑market подборки сократился с ≈30 часов для всех пользователей до нескольких минут для самых активных пользователей и ≈12 часов для медианного пользователя. Конкретно в user2item‑источнике в несколько раз ускорили доставку свежих кандидатов пользователю и научились реагировать на изменения поведения пользователя заметно оперативнее.

• Освободили несколько десятков машин Spark‑ресурсов. Это стало возможным потому, что мы убрали тяжёлые джобы сбора и объединения фич с кандидатами и часть другой предподготовки. Освободившиеся ресурсы кластера тут же направили на другие задачи команды.

• Получили возможность скорить всю товарную базу с фильтрами и разнообразием с честным расчётом топа.

Продуктовые метрики (A/B относительно старого пайплайна). На глобальном уровне мы увидели рост ключевых метрик покупательской активности: вырастили как глобальное количество покупок на пользователя, так и рекламную выручку.

Ожидаемо получилось улучшить на единицы процентов метрики в самих рекомендациях:

• Количество заказанных товаров на пользователя с рекомендаций.

• Доля пользователей с заказом из рекомендаций.

Выводы и планы

Честный KNN на GPU оказался для нас тем случаем, когда решение в лоб выигрывает у умного приближённого. Мы получили точную выдачу без потерь recall на «хвостах», встроили всю бизнес‑логику прямо в проход по векторам, ускорили TTM почти вдвое и освободили существенный объём CPU‑ и RAM‑ресурсов — и всё это сопровождается ростом продуктовых метрик в A/B‑тесте.

Несколько уроков, которые мы вынесли:

• Не бойтесь полного перебора, если позволяют размерности и железо. Современные GPU отлично справляются с матричными операциями, и честный перебор часто оказывается и точнее, и в сумме дешевле в поддержке, чем приближённый индекс.

• Собственная имплементация даёт свободу. Возможность встроить фильтры, штрафы, diversity и комбинацию эмбеддингов прямо в вычисление топа важнее, чем кажется. Именно она убирает тяжёлый постпроцессинг и открывает простор для экспериментов.

• Точный KNN полезен даже как эталон. Имея ground truth, вы наконец можете честно измерить, сколько теряет ваш приближённый поиск.

Куда хотим двигаться дальше:

• Увеличение реактивности сбора подборок — ускорить расчёт retrieval эмбеддингов, экспериментировать с частотой пересчёта user2item‑подборок пользователю.

• Использование статистик взаимодействия с документами в RT — актуализировать товарные статистики и эмбеддинги внутри GPU с минимальной задержкой.

• Развитие инструментов сбора подборок — добавлять учёт новых сигналов и новые логики разнообразия прямо в проход на GPU, не возвращаясь к тяжёлому постпроцессингу.

• Интеграция с генеративным ретривалом — выжимать максимальную пользу из комбинации подходов с обработкой документной базы на GPU и SOTA‑подходов по моделированию пользователя.