Revisiting Neural Retrieval on Accelerators [2023]

Довольно часто в рекомендательной системе есть кандидатогенератор, отдающий порядка тысячи кандидатов, а дальше в дело уже вступает ранжирующая модель.

Авторы данной статьи непрозрачно намекают, что использование dot product для генерации кандидатов - отстой. Матрица скалярных произведений всех на всех получается с помощью умножения 2 матриц вдоль размерности длины d - размера эмбеддинга, а значит ранг итоговой матрицы ограничен этим самым d, тогда как реальность сложнее.

Бывает, что в сервисе есть ещё одна, более лёгкая стадия ранжирования, которая может забрать больше документов из кандгена и проскорить получше их перед подачей в тяжёлую модель. Авторы этой работы решили добиться максимум возможного из этого подхода.

Начнём с самой первой стадии. Авторы почему-то ненавидят HNSW, и поэтому используют эмбеддинг небольшого размера (64), чтобы можно было проскорить ваще всю базу и получить 100к кандидатов. Но им не нужно ровно 100к, поэтому используется approximate top-k.

Допустим, мы хотим достать 100к из миллиарда. Их будет разделять порог по похожести x. Этот порог можно оценить, засэмплив миллион документов и посчитав похожесть 100-го. Затем можно, используя этот порог, пробежаться по всей базе и отсеять документы. Итоговая сложность алгоритма падает с O(N log k) до O(N + rN * log rk), где r - доля сабсэмплинга. В реальности получается в 2.5 раза быстрее.

Дальше к этим 100к кандидатам применяется так называемая Mixture of Logits (MoL) - модель, которая берёт на вход эмбеддинг пользователя и айтема и выдаёт их похожесть, которая по выразительности находится между MLP и Dot Product.

Оба входных вектора нарезаются на куски размером по d - будем называть их подвекторами. Далее мы берём ku подвекторов пользователя и kx подвекторов документа и считаем ku*kx скалярных произведений друг на друга. Далее этот вектор скалярно умножается на вектор весов pi(x, u), который считается как линейный слой + софтмакс от на основе конката тех самых скалярных произведений и ещё двух векторов юзерных и айтемных фичей.

Авторы уделяют много времени оптимизациям, в том числе дизайна своего GPU kernel, правда, к сожалению, я так и не увидел понятного описания того, во сколько раз в точно таких же условиях оно применяется медленнее, чем dot product. Но суммарно эти 2 стадии могут прожевать в 1.8 раз меньше, чем просто top-k dot product. По качеству оно его при этом сильно превосходит.

@knowledge_accumulator