Назад к блогу

Как Pinterest ускорила поиск: квантование эмбеддингов и переход на SSD

Как Pinterest ускорила поиск: квантование эмбеддингов и переход на SSD

Pinterest перевела поиск по эмбеддингам с полностью резидентного в память HNSW на схему, где основная часть индекса живёт на SSD, а квантование сжимает векторы на десятки процентов без критичной потери качества. Разбор показывает, как устроены IVF-индексы и SPANN изнутри — от обучения грубого квантователя и назначения векторов по спискам до блочного добавления данных. Полезно тем, кто упирается в стоимость памяти при росте числа эмбеддингов и ищет рабочий компромисс между recall, задержкой и железом.

Поиск по векторам упирается в память: чем больше эмбеддингов, тем дороже держать индекс целиком в оперативной памяти. Pinterest прошла путь от стандартного HNSW через сжатые квантованные варианты к индексу, который держит в памяти лишь малую часть данных, а основную массу хранит на SSD. Разберём механику этого пути и устройство индексов, на которых он построен.

Путь платформы Manas: от HNSW к SPANN

Стандартный HNSW страдает от большого потребления памяти и сталкивается с растущими затратами, проблемами выделения оборудования и гибкости инфраструктуры. Первый шаг — сжатие векторов квантованием.

скалярное квантование уменьшает индексы HNSW на 59%, а IVF-индекс — на 75%, сохраняя recall выше 90% на разных нагрузках. Альтернатива — продуктовое квантование — сжимает сильнее: HNSW на 74%, IVF на 93%, но recall падает до 70–80%.

Для перехода на SSD Pinterest оценила DiskANN и SPANN. SPANN с PQ дала в 3 раза больше QPS, чем DiskANN, при втрое меньшей задержке и падении recall всего на 5%. В оценке рекомендаций Pin с индексированием более 5 млрд эмбеддингов SPANN сэкономил свыше 40% процессорного времени по сравнению с полностью резидентным HNSW.

Как устроен IVF-индекс: грубый квантователь и списки

IVF-индекс делит пространство на nlist областей. За распределение векторов по областям отвечает грубый квантователь. Обучение квантователя выполняет train_q1_impl: он проверяет, что квантователь задан и что размерность кодека совпадает с размерностью квантователя, после чего выбирает способ обучения по флагу quantizer_trains_alone. Этот флаг задаёт, кто обучает грубый квантователь: при значении 0 квантователь обучается на тех же данных, что и сам индекс, при значении 2 квантователь обучается отдельно, а значение 1 означает, что квантователь уже обучен и его обучение пропускается.

При quantizer_trains_alone == 0 квантователь сбрасывается. Если включён use_super_kmeans, создаётся SuperKMeans и его центроиды добавляются в квантователь; иначе создаётся Clustering, который обучается на данных, а в качестве индекса назначения получает assigner. При quantizer_trains_alone == 2 создаётся Clustering, и если clustering_index не задан, используется IndexFlatL2 как assigner, после чего центроиды добавляются в квантователь.

Назначение векторов в списки идёт через assigner: квантователь назначает каждому вектору номер списка, а затем для вектора вычисляется остаток, который и кодируется дальше.

train_q1 — обёртка над train_q1_ex с типом Float32, а тот вызывает train_q1_impl с codec=nullptr. train_q1_encoded отличается тем, что передаёт в train_q1_impl реальный кодек. Сферические центроиды поддержаны в ветке quantizer_trains_alone == 2: там разрешён METRIC_L2 или METRIC_INNER_PRODUCT при cp.spherical, потому что в этом случае L2 и inner product эквивалентны. При quantizer_trains_alone == 0 и use_super_kmeans требуется L2 или сферический inner-product.

Номер списка хранится компактно: coarse_code_size вычисляет число байт, начиная с nlist - 1 и сдвигая вправо на 8 бит, пока значение не обнулится. encode_listno пишет номер побайтово в little endian, а decode_listno читает его обратно и проверяет, что результат лежит в диапазоне [0, nlist).

Добавление и кодирование векторов

add_with_ids сначала проверяет, что квантователь и инвертированные списки заданы, выделяет массив под номера списков и просит квантователь назначить каждый вектор, после чего передаёт работу в add_core.

add_core разбивает добавление на блоки по 65536 векторов: если векторов больше, он рекурсивно вызывает себя для каждого поддиапазона, сдвигая указатели на данные, идентификаторы и номера списков. Это блокировка для избежания избыточных аллокаций. В остальной части add_core проверяет входные условия, считает число векторов без списка, выделяет буфер под коды, кодирует векторы и в параллельной секции OpenMP распределяет списки между потоками по условию list_no % nt == rank, добавляя записи в инвертированные списки и регистрируя их в прямой карте. Для вектора без списка при rank == 0 запись регистрируется с номером списка -1.

add_sa_codes работает с уже готовыми кодами: вычисляет размер грубого кода, в цикле извлекает из каждого кода номер списка, берёт идентификатор и добавляет запись в инвертированный список, после чего увеличивает ntotal.

IndexIVFPQ::encode_vectors при by_residual сначала вычисляет остатки: для вектора с отрицательным номером списка остаток заполняется нулями, иначе вычисляется как разность с центроидом. Затем остатки (или исходные векторы, если by_residual ложно) кодируются продуктовым квантователем. Если включён include_listnos, коды сдвигаются вправо на размер грубого кода, а в начало каждого кода записывается номер списка.

Распараллеливание через OpenMP применяется к независимым итерациям: в compute_residuals директива включается при n > 1000, в ProductQuantizer::compute_codes — для независимой обработки векторов при dsub < 16, в ProductQuantizer::decode — при n > 100.

Продуктовое квантование: обучение и кодирование

Обучение ProductQuantizer разбивает вектор на M подвекторов длиной dsub и для каждого подвектора обучает свой набор из ksub центроидов. Если train_type не Train_shared, для каждого подвектора создаётся Clustering, и в зависимости от final_train_type центроиды инициализируются: init_hypercube строит центроиды как среднее плюс-минус максимум по первым nbits координатам, а остальные координаты берёт равными среднему; init_hypercube_pca обучает PCAMatrix на данных и добавляет к среднему вклад собственных значений и собственных векторов. Затем Clustering обучается, и результат сохраняется. Если train_type == Train_shared, Clustering обучается сразу на n*M векторах, и одни и те же центроиды копируются во все M срезов.

Кодирование идёт по подвекторам: для каждого подвектора находится ближайший центроид и записывается его индекс. Реализация кодировщика выбирается по nbits: 8 бит — PQEncoder8, 16 — PQEncoder16, иначе PQEncoderGeneric. compute_codes обрабатывает данные блоками, чтобы не использовать слишком много памяти, а при dsub < 16 кодирует каждый вектор напрямую, иначе строит таблицы расстояний и по ним выбирает лучший центроид.

Декодирование для каждого подвектора извлекает индекс центроида и копирует соответствующий центроид в выходной вектор; при вызове с параметром n декодирование распараллеливается по векторам.

Таблицы расстояний и выбор реализации

compute_distance_table вычисляет для одного вектора таблицу расстояний до центроидов. Если транспонированные центроиды недоступны, используется регулярная версия; при их наличии — оптимизированная с транспонированными центроидами и длинами центроидов. compute_inner_prod_table для одного вектора всегда использует регулярную версию с inner products.

compute_distance_tables обрабатывает сразу много векторов: при dsub == 2, nbits >= 3 и nbits < 8 (только для AVX2, ARM NEON или RISCV_RVV) вызывается специализированная версия для двухмерных подвекторов, иначе при dsub < 16 параллельно по векторам вызывается compute_distance_table, а при большем dsub используется BLAS.

Выбор реализации в поиске по таблицам зависит от nbits: при 8 битах и ksub == 256 вызывается специализированный 8-битный поиск, иначе — общий вариант для 8 бит; при 16 битах — вариант для 16 бит; в остальных случаях — общий. Внутри варианта для 8/16 бит при M == 4 вызывается специализированная версия для четырёх подвекторов, при M % 4 == 0 — версия с развёрткой по четвёркам, иначе используется медленный общий цикл.

FastScan: SIMD-ускорение поиска

FastScan требует жёстких условий: bbs должен быть кратен 32, nbits равен 4, а размерность fine-квантователя совпадать с размерностью вектора. При инициализации вычисляются ksub, округлённое до чётного M2 и размер кода, и проверяется, что размер кода совпадает с размером кода fine-квантователя. Если индекс владеет своими инвертированными списками, они заменяются на block-packed списки с упаковщиком кодов.

Если инвертированные списки не block-packed (например, индекс не владеет ими), упаковщик кодов не создаётся — вызывающий сам управляет списками. Упаковщик создаётся только когда приведение к block-packed спискам удаётся, и тогда старый упаковщик удаляется. Шаг упаковки кодов по умолчанию нулевой, что означает стандартный шаг в M байт.

Реализации поиска отличаются способом обработки запросов. search_implem_1 обрабатывает по одному запросу: строит таблицу подстановок, перебирает списки, пропуская отрицательные номера и пустые списки, и для каждого непустого списка оценивает расстояния. search_implem_2 использует 8-битные таблицы и хранит промежуточные расстояния в 16-битных целых, преобразуя их во float в конце.

search_implem_10 принимает обработчик результатов и сканер кодов, применяет early-stop из параметров поиска: лимит числа просмотренных кодов, лимит числа списков, требование заполнить top-k и лимит пустых бакетов. Проверка лимитов по кодам и спискам выполняется перед началом следующего списка, а для range-search пустые или невалидные списки увеличивают счётчик пустых бакетов. search_implem_12 группирует запросы по общему номеру списка и обрабатывает пачки запросов. search_implem_14 тоже группирует запросы по списку, но использует локальные кучи на поток.

Early-stop опции по пустым бакетам и по k-NN требуют per-query реализаций, потому что счётчики просмотренных кодов, посещённых списков и пустых бакетов ведутся внутри цикла по каждому запросу.

Таблицы подстановок и выбор реализации в FastScan

compute_LUT_uint8 сначала строит float-таблицы, а затем в параллельном цикле преобразует их в 8-битные таблицы с нормализаторами. Оценка по таблицам читает коды через битовый поток: для каждого кодового слова накапливается смещение, а для каждого подвектора читается код и прибавляется значение из таблицы. При наличии масштаба для оставшихся подвекторов значения прибавляются с этим масштабом.

Наличие AVX2 или ARM NEON влияет на построение таблиц: при dsub == 2, nbits >= 3 и nbits < 8 вызывается специализированная версия, иначе при dsub < 16 используется параллельный цикл по векторам, а при большем dsub — BLAS.

Поиск: preassigned, range_search и сканеры

search_preassigned проверяет входные условия, вычисляет число просматриваемых списков как минимум из nlist и nprobe. Режим параллелизации задаётся параметром parallel_mode: при нулевом значении потоки OpenMP распределяют запросы по срезам — каждый поток обрабатывает свой диапазон запросов, а исключения в параллельном регионе перехватываются и перебрасываются после него; при ненулевом значении внешняя параллелизация не применяется, и функция поиска вызывается для всех запросов сразу. Нулевой режим — основная ветка, рассчитанная на 95% случаев использования; ненулевые режимы добавляют сложность за счёт иных способов распараллеливания, обработки прерываний и ошибок, сбора статистики и сбора минимумов и максимумов.

sorted-range shortcut применяется, когда селектор — это диапазон с флагом assume_sorted и список не использует итератор: тогда селектор обнуляется, чтобы применить специальную обработку. Иначе селектор диапазона обнуляется и используется общая обработка. Это ограничивает сканирование секцией array-backed списка, потому что для итерируемых списков селектор и фильтр должны проверяться по каждой записи.

range_search сначала квантует запросы, получая ближайшие списки и расстояния до них, затем делает prefetch списков и передаёт всё в range_search_preassigned. Там вычисляются число просматриваемых списков, лимит кодов и бюджет ранней остановки по пустым бакетам: значение 0 отключает остановку по пустым бакетам. Этот бюджет поддерживается только при нулевом режиме параллелизации. Для каждого потока создаётся частичный результат и сканер; в функции сканирования списка номер берётся из массива ключей, отрицательные ключи пропускаются, а пустые списки отсеиваются.

search1 работает без параллельной секции: он сам квантует один запрос, делает prefetch списков, создаёт сканер, устанавливает запрос, а затем в цикле пропускает отрицательные ключи и пустые списки, для остальных устанавливает список и сканирует коды.

Восстановление и проверка целостности

get_InvertedListScanner в базовом IVF-индексе не реализован и всегда выбрасывает исключение. В IndexIVFPQ он возвращает сканер, создаваемый через фабрику сканеров, а в IndexIVFPQFastScan — объект сканера FastScan.

reconstruct по идентификатору получает из прямой карты номер списка и смещение, проверяет, что номер списка меньше nlist, а смещение меньше размера списка, и вызывает восстановление по смещению. reconstruct_n перебирает все списки и все смещения, берёт идентификатор записи и, если он попадает в запрошенный диапазон, восстанавливает вектор в нужную позицию выходного массива.

check_ids_sorted проверяет, что внутри каждого списка идентификаторы не убывают: для каждой пары соседних идентификаторов при нарушении порядка увеличивается счётчик, и функция возвращает признак отсутствия нарушений.

decode_vectors в базовом IVF-индексе не реализован и выбрасывает исключение. В IndexIVFFastScan восстановление распаковывает код, декодирует номер списка, декодирует fine-квантователем и при by_residual добавляет восстановленный центроид.

Обслуживание индекса: удаление, обновление, слияние

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

update_vectors ведёт себя по-разному в зависимости от типа прямой карты. При хеш-таблице применяется подход «сначала удалить, потом добавить»: создаётся селектор по новым идентификаторам, вызывается удаление, проверяется, что удалены все записи, и только затем векторы добавляются заново. При array-карте подход другой: чтобы не создавать дыры в непрерывном диапазоне идентификаторов, векторы кодируются заново и коды обновляются напрямую через прямую карту.

train_encoder_on_encoded_rows обучает энкодер IVF на строках, которые не являются fp32: определяет лимит числа векторов для обучения, при необходимости субдискретизирует через перестановку с фиксированным seed, декодирует выбранные строки в fp32 и затем при by_residual назначает их квантователю, вычисляет остатки и обучает энкодер на остатках, иначе — на декодированных данных.

check_compatible_for_merge проверяет, что другой индекс — это IVF-индекс, что у него совпадают размерность, число списков, размер кода, что у квантователей одинаковое число векторов, что типы индексов совпадают и что у обоих нет прямой карты. Если включён дорогой режим проверки, дополнительно для каждого списка восстанавливаются векторы из квантователей обоих индексов и проверяется их равенство.

merge_from сначала проверяет совместимость, затем сливает инвертированные списки, увеличивает ntotal на число векторов другого индекса и обнуляет ntotal другого. replace_invlists удаляет старые списки, если индекс владел ими, затем проверяет, что у нового списка совпадают число списков и размер кода (или размер кода помечен как невалидный), после чего присваивает новые списки. copy_subset_to увеличивает ntotal другого индекса на число скопированных записей.

Пограничные случаи и компромиссы

В исходниках зафиксированы конкретные компромиссы. В поиске IVF прямо сказано, что концептуально простая функция становится очень сложной из-за нескольких способов параллелизации, обработки прерываний и ошибок, сбора статистики и сбора минимумов и максимумов, а основная ветка кода рассчитана на нулевой режим параллелизации. В кодировании PQ объясняется, почему выделяется буфер под расстояния: это кажется бессмысленным, но делается ради того, как компилятор генерирует код, и базовая версия оказывается быстрее из-за векторизации, на которую, вероятно, влияет и предсказатель переходов. Для случая, когда M не равно 4 и не кратно 4, используется относительно медленный вариант.

FastScan жёстко требует, чтобы bbs был кратен 32, а nbits равен 4, иначе выбрасывается исключение. В одной из реализаций поиска дополнительно требуется bbs == 32. В range-search явно запрещены k-NN-настройки: лимит числа списков, требование заполнить top-k и лимит кодов не соблюдаются при range-search в FastScan.

Обработка потерь, одновременных кандидатов и неполных данных

При восстановлении по смещению в FastScan код извлекается из упакованного списка: создаётся буфер, в него пишется номер списка, затем для каждого подвектора читается упакованный элемент и записывается в битовый поток, после чего вектор восстанавливается. В IndexIVFPQ восстановление проще: берётся код по смещению, декодируется, и при by_residual к результату прибавляется центроид.

При неполных данных в поиске с восстановлением для отрицательного ключа восстановленный вектор заполняется значением -1, а в поиске с возвратом кодов для отрицательного ключа код заполняется 0xff. Перед восстановлением по смещению проверяются границы: номер списка меньше nlist, смещение меньше размера списка.

reconstruct_orig_invlists требует наличия исходных инвертированных списков и пустого нулевого списка, после чего параллельно по спискам распаковывает коды и добавляет записи в исходные списки. Слияние результатов по потокам выполняется через OpenMP: сначала инициализируются результаты как куча, затем в критической секции для каждого запроса добавляются локальные результаты, после чего выполняется переупорядочивание.

Как сравнивали и что получилось

Сравнение шло по трём метрикам: размер индекса, Recall@100 и QPS. Участники — HNSW плюс SQ, базовый IVF, IVF плюс PQ и IVF плюс SQ. Для SSD-обслуживания отдельно сравнивались SPANN с PQ и DiskANN.

КонфигурацияРазмер индексаRecall@100QPS
HNSW + SQ50 ГБ92.92%305.2
Базовый IVF97 ГБ91.69%1659.8
IVF + PQ6.8 ГБ76.00%1747.9
IVF + SQ25 ГБ95.71%1588.8

HNSW плюс SQ даёт индекс 50 ГБ с Recall@100 92.92% при 305.2 QPS. Базовый IVF-индекс — 97 ГБ с Recall@100 91.69% при 1659.8 QPS. IVF плюс PQ сжимает индекс до 6.8 ГБ, но Recall@100 падает до 76.00% при 1747.9 QPS. IVF плюс SQ даёт 25 ГБ и лучший Recall@100 95.71% при 1588.8 QPS.

Для SSD-обслуживания SPANN с PQ показал в 3 раза больше QPS, чем DiskANN, при втрое меньшей задержке и падении recall всего на 5%. В оценке рекомендаций Pin с индексированием более 5 млрд эмбеддингов SPANN сэкономил свыше 40% процессорного времени по сравнению с полностью резидентным HNSW.

Размер датасета, размерность, nlist, nprobe, M и nbits не указаны, и процедура повторения замера не описана.

Что из этого следует на практике

Квантование — это обмен памяти на recall, и характер обмена зависит от метода. Скалярное квантование сжимает умеренно и сохраняет recall выше 90%: индексы HNSW уменьшаются на 59%, IVF — на 75%. Продуктовое квантование сжимает сильнее (HNSW на 74%, IVF на 93%), но recall падает до 70–80%. Если recall критичен, а память не бесконечна, скалярное квантование даёт лучший баланс; если нужно максимальное сжатие и допустимо потерять часть качества — продуктовое.

Перенос основной массы данных на SSD меняет профиль затрат. Архитектура SPANN держит в памяти лишь небольшой центроидный индекс, а большие posting lists хранит на SSD, оптимизируя IOPS. Это даёт в 3 раза больше QPS, чем DiskANN, при втрое меньшей задержке и падении recall всего на 5%. На масштабе более 5 млрд эмбеддингов такая схема экономит свыше 40% процессорного времени по сравнению с полностью резидентным HNSW.

Жёсткие требования FastScan ограничивают применимость: bbs кратен 32, nbits равен 4, а в одной из реализаций поиска bbs строго равен 32. Если параметры не подходят, индекс просто не инициализируется. Early-stop по пустым бакетам и по k-NN работает только в per-query реализациях, потому что счётчики ведутся внутри цикла по каждому запросу; в range-search FastScan k-NN-настройки не соблюдаются вовсе.

Обновление векторов устроено по-разному в зависимости от типа прямой карты: при хеш-таблице запись удаляется и добавляется заново, а при array-карте кодируется заново и обновляется на месте, чтобы не создавать дыры в непрерывном диапазоне идентификаторов. Слияние индексов возможно только при совпадении размерности, числа списков, размера кода, числа векторов квантователя и типа индекса, и только если у обоих нет прямой карты.

Где смотреть в коде

Источники

Похожее