Краткое введение в индекс ScaNN

ScaNN - это ответ Google на привычную проблему крупномасштабного векторного поиска: как увеличить пропускную способность запросов и снизить потребление памяти без неприемлемого ущерба для качества результатов. Концептуально ScaNN основывается на классическом рецепте IVF+PQ - грубая кластеризация плюс агрессивное квантование по продуктам - но включает два важных новшества, которые существенно сдвигают границу производительности:

• Задача квантования с учетом баллов, которая лучше сохраняет относительное упорядочивание истинных соседей, улучшая качество ранжирования даже при сильном сжатии.

• FastScan - это оптимизированный для SIMD 4-битный путь поиска PQ, который уменьшает традиционное узкое место в памяти, сохраняя больше работы в регистрах процессора.

На практике он является сильным выбором, когда вы готовы пожертвовать некоторым количеством отзыва ради высокого QPS и гораздо меньшего объема памяти (часто сжимая векторы до ~1/16 от исходного размера), например, в рекомендательных нагрузках, не чувствительных к отзыву.

В этом посте мы вновь рассмотрим IVFPQ в качестве базового уровня, изучим ключевые оптимизации, которые ScaNN вводит поверх него, и завершим обсуждение экспериментальными результатами, основанными на измеренной производительности.

Краткое описание IVFPQ

ScaNN была предложена компанией Google в 2020 году, и в статье сообщается о повышении производительности на 3× по сравнению с HNSW на наборе данных GloVe. За подробностями вы можете обратиться к оригинальной статье и реализации с открытым исходным кодом.

Прежде чем представить ScaNN, мы кратко расскажем о IVFPQ, поскольку ScaNN построен на основе того же общего фреймворка.

IVFPQ расшифровывается как Inverted File with Product Quantization- алгоритм, используемый для эффективного и масштабного поиска ближайших соседей (ANN) в высокоразмерных векторных базах данных. Это гибридный подход, сочетающий две техники - инвертированный индекс файла (IVF) и квантование продукта (PQ)- для баланса скорости поиска, использования памяти и точности.

Принцип работы IVFPQ

Процесс включает в себя два основных этапа индексирования и поиска:

• Уровень ЭКО: векторы кластеризуются в инвертированные списки (кластеры) nlist. Во время запроса посещается только подмножество кластеров (nprobe), чтобы найти компромисс между отзывом и задержкой/пропускной способностью.

• Уровень PQ: внутри каждого посещенного кластера каждый D-мерный вектор разбивается на m подвекторов, каждый из которых имеет размерность (D/m). Каждый субвектор квантуется путем присвоения ему ближайшего центроида в его субкодексе. Если в подкодексе 256 центроидов, то каждый подвектор может быть представлен кодом uint8 (идентификатор в диапазоне [0, 255]).

Вычисление расстояния может быть переписано как сумма по подвекторам:

D(q, X) = D(q, u0) + D(q, u1) + D(q, u2) + ... + D(q, un)

= L(q, id1) + L(q, id2) + L(q, id3) + ... + L(q, idn)

Здесь L представляет собой таблицу поиска. Во время запроса строится таблица поиска, в которой записывается расстояние между запросом и каждым квантованным подвектором. Все последующие вычисления расстояний преобразуются в поиск по таблице с последующим суммированием.

Например, для 128-мерных векторов, разбитых на 32 субвектора по 4 измерения каждый, если каждый субвектор кодируется идентификатором uint8, стоимость хранения одного вектора снижается с (128 x 4) байт до (32 x 1) байт - сокращение на 1/16.

Оптимизации ScaNN на основе IVFPQ

В целом, ScaNN улучшает IVFPQ в двух аспектах:

1. Квантование: ScaNN предлагает цель, выходящую за рамки простой замены каждого субвектора на ближайший к нему центроид k-средних (т. е. минимизация ошибки реконструкции).

2. Эффективность поиска: ScaNN ускоряет поиск на основе LUT, который часто ограничен памятью, с помощью SIMD-дружественного пути FastScan.

Потери при квантовании с учетом баллов

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

$$\begin{array}{r}{\displaystyle \ell ({\mathrm{xi}}_{},{\mathrm{x~i}}_{},w)}\end{array}$$$$∥2\begin{aligned} \ell\left(x_i, \tilde{x}_i, w\right) &=h_{\|}\left(w,\left\|x_i\right\|\right)\left\|r_{\|}\left(x_i, \tilde{x}_i\right)\right\|^2 +h_{\perp}\left(w,\left\|x_i\right\|\right)\left\|r_{\perp}\left(x_i,$$ \tilde{x}_i\right$$)\right\|^2 \end{aligned}$$ $$\begin{array}{r}\ell (x\end{array}$$i, x~i,w) =h∥ $$\begin{array}{l}(w,\end{array}$$∥xi$$\begin{array}{l}\Vert )\end{array}$$ $$\begin{array}{l}r\end{array}$$∥ $$\begin{array}{l}(x\end{array}$$i, x~i) 2+h⊥ $$\begin{array}{l}(w,\Vert x\end{array}$$i$$\begin{array}{l}\Vert )\Vert r\end{array}$$⊥ $$\begin{array}{l}(x\end{array}$$i,

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

На левом рисунке показано плохое квантование, поскольку параллельное смещение влияет на конечный результат, в то время как на правом рисунке показано лучшее квантование.

4-битный PQ FastScan

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

Однако частые обращения к памяти все равно становятся узким местом в производительности. Если таблицу поиска сделать достаточно маленькой, чтобы она помещалась в регистры, операции чтения памяти можно преобразовать в эффективные SIMD-инструкции процессора.

Сначала каждый субвектор разбивается на 16 классов, так что 4-битное значение может представлять центр кластера - отсюда и название "4-битный PQ". Затем расстояния, обычно представляемые в виде плавающих чисел, преобразуются в uint8 с помощью скалярной квантизации (SQ). Таким образом, таблица поиска для одного субвектора может храниться в регистре, используя 16 × 8 = 128 бит.

Наконец, рассмотрим схему хранения регистров (на примере набора инструкций AVX2): субвекторы 32 векторов помещаются в 128-битный регистр, объединенный с таблицей поиска. Операция "поиск" может быть эффективно выполнена с помощью одной SIMD-инструкции процессора.

расположение регистров

SIMD Shuffle для поиска

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

Экспериментальные результаты, представленные в статье ScaNN.

Интуитивно понятно, что оптимизация только потерь не должна приводить к таким впечатляющим результатам. В другом блоге на это также обратили внимание - на самом деле разница заключается в 4-битной части PQ FastScan.

Экспериментальные результаты

Используя бенчмарк векторных баз данных VectorDBBench, мы провели простой тест. Преимущество производительности ScaNN над традиционными IVFFLAT и IVF_PQ вполне очевидно. После интеграции в Milvus на наборе данных Cohere1M при одинаковом коэффициенте отзыва QPS достигает 5 раз по сравнению с IVFFLAT и 6 раз по сравнению с IVF_PQ.

Однако QPS немного ниже, чем у индексов на основе графов, таких как HNSW, поэтому он не является первым выбором для сценариев с высоким коэффициентом запоминания. Но для сценариев с меньшим количеством отзывов это приемлемо (например, в некоторых рекомендательных системах), использование ScaNN без загрузки исходных данных позволяет достичь впечатляющего QPS при крайне малом объеме памяти (1/16 от исходных данных), что делает его отличным выбором индекса.

Тип индекса	Случай	QPS	задержка(p99)	отзыв	память
IVFFLAT	Производительность1M	266	0.0173s	0.9544	3G
HNSW	Производительность1M	1885	0.0054s	0.9438	3.24G
ЭКО_PQ	Производительность1M	208	0.0292s	0.928	0.375G
ScaNN（with_raw_data: true）	Производительность1M	1215	0.0069s	0.9389	3.186G
ScaNN（with_raw_data: false）	Производительность1M	1265	0.0071s	0.7066	0.186G

Заключение

ScaNN опирается на знакомый фреймворк IVFPQ, но значительно расширяет его за счет глубокой инженерной работы в области квантования и ускорения низкоуровневого поиска. Согласовывая цель квантования с качеством ранжирования и устраняя узкие места памяти во внутреннем цикле, ScaNN сочетает квантование с учетом баллов с 4-битным PQ FastScan-путем, который превращает традиционно ограниченный памятью процесс в эффективное, SIMD-дружественное вычисление.

На практике это дает ScaNN четкую и хорошо определенную нишу. Он не предназначен для замены индексов на основе графов, таких как HNSW, в условиях большого количества обращений. Напротив, для нечувствительных к отзывам рабочих нагрузок с ограниченным бюджетом памяти ScaNN обеспечивает высокую пропускную способность при очень малой занимаемой площади. В наших экспериментах, после интеграции в Milvus VectorDB, ScaNN достигла примерно 5× QPS по сравнению с IVFFLAT на наборе данных Cohere1M, что делает ее отличным выбором для высокопроизводительного поиска ANN с низкой задержкой, где сжатие и эффективность важнее идеального запоминания.

Если вы заинтересованы в дальнейшем изучении ScaNN или обсуждении выбора индекса в реальных системах, присоединяйтесь к нашему Slack-каналу, чтобы пообщаться с нашими инженерами и другими инженерами по ИИ в сообществе. Вы также можете заказать 20-минутную индивидуальную сессию, чтобы получить знания, рекомендации и ответы на свои вопросы через Milvus Office Hours.