Breve introducción al Índice ScaNN

ScaNN es la respuesta de Google a un reto conocido en la búsqueda vectorial a gran escala: cómo aumentar el rendimiento de las consultas y reducir el uso de memoria sin que la calidad de los resultados se vea afectada de forma inaceptable. Conceptualmente, ScaNN parte de la receta clásica de IVF+PQ (agrupación gruesa y cuantificación agresiva del producto), pero incorpora dos innovaciones importantes que cambian significativamente la frontera del rendimiento:

• Un objetivo de cuantificación basado en la puntuación que conserva mejor el orden relativo de los vecinos verdaderos, mejorando la calidad de la clasificación incluso en condiciones de compresión intensa.

• FastScan es una ruta de búsqueda PQ de 4 bits optimizada para SIMD que reduce el tradicional cuello de botella de la carga de memoria al mantener más trabajo en los registros de la CPU.

En la práctica, es una buena opción cuando se está de acuerdo en cambiar algo de memoria por un alto QPS y una huella de memoria mucho más pequeña (a menudo comprimiendo vectores a ~1/16 del tamaño original), como en cargas de trabajo de recomendación insensibles a la memoria.

En este artículo, volveremos a IVFPQ como punto de partida, exploraremos las optimizaciones clave que ScaNN introduce sobre él y terminaremos con resultados experimentales que fundamentan el debate en el rendimiento medido.

Recapitulación de IVFPQ

ScaNN fue propuesto por Google en 2020, y en el documento se informa de una mejora del rendimiento 3 veces superior a HNSW en el conjunto de datos GloVe. Puede consultar el artículo original y la implementación de código abierto para obtener más información.

Antes de presentar ScaNN, repasaremos brevemente IVFPQ, ya que ScaNN se basa en el mismo marco general.

IVFPQ son las siglas de Inverted File with Product Quantization (Archivo invertido con cuantificación de producto), un algoritmo utilizado para la búsqueda eficiente y a gran escala de Vecinos Cercanos Aproximados (RNA) en bases de datos vectoriales de alta dimensión. Se trata de un enfoque híbrido que combina dos técnicas, el índice de archivo invertido (IVF) y la cuantificación de productos (PQ), para equilibrar la velocidad de búsqueda, el uso de memoria y la precisión.

Funcionamiento del IVFPQ

El proceso consta de dos pasos principales durante la indexación y la búsqueda:

• Capa IVF: los vectores se agrupan en nlist listas invertidas (clusters). En el momento de la consulta, sólo se visita un subconjunto de clústeres (nprobe) para compensar la recuperación y la latencia/rendimiento.

• Capa PQ: dentro de cada clúster visitado, cada vector de dimensión D se divide en m subvectores, cada uno de dimensión (D/m). Cada subvector se cuantifica asignándolo al centroide más cercano de su sub-libro de códigos. Si un subcodificador tiene 256 centroides, cada subvector puede representarse mediante un código uint8 (un ID en [0, 255]).

El cálculo de la distancia puede reescribirse como la suma de los subvectores:

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)

Aquí, L representa una tabla de búsqueda. En el momento de la consulta, se construye la tabla de consulta y se registra la distancia entre la consulta y cada subvector cuantificado. Todos los cálculos de distancia posteriores se convierten en búsquedas en la tabla seguidas de una suma.

Por ejemplo, para vectores de 128 dimensiones divididos en 32 subvectores de 4 dimensiones cada uno, si cada subvector se codifica mediante un ID uint8, el coste de almacenamiento por vector se reduce de (128 x 4) bytes a (32 x 1) bytes, es decir, 1/16 menos.

Optimizaciones de ScaNN basadas en IVFPQ

En resumen, ScaNN mejora IVFPQ en dos aspectos:

1. Cuantización: ScaNN propone un objetivo que va más allá de la simple sustitución de cada subvector por su centroide k-means más cercano (es decir, minimizar el error de reconstrucción).

2. Eficiencia de búsqueda: ScaNN acelera la búsqueda basada en LUT, que a menudo está limitada por la memoria, mediante una ruta FastScan compatible con SIMD.

Pérdida de cuantificación en función de la puntuación

Dado que el análisis se basa en la métrica IP, después de que ScaNN descomponga el error de cuantificación en componentes paralelos y perpendiculares, sólo el componente paralelo afecta al resultado, por lo que debe aplicarse un término de penalización mayor. En consecuencia, la función de pérdida puede reescribirse como sigue:

$$\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}(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,

La figura siguiente muestra un ejemplo bidimensional que ilustra que el error causado por el componente paralelo es mayor y puede conducir a resultados incorrectos del vecino más cercano, lo que justifica una penalización más severa.

La figura de la izquierda muestra una mala cuantificación porque el desplazamiento paralelo afecta al resultado final, mientras que la figura de la derecha muestra una mejor cuantificación.

PQ FastScan de 4 bits

Revisemos primero el proceso de cálculo de PQ: durante la consulta, las distancias entre la consulta y los centros de cluster del subvector se calculan previamente para construir una tabla de búsqueda. A continuación, el cálculo de las distancias se realiza mediante búsquedas en la tabla para obtener las distancias entre segmentos y sumarlas.

Sin embargo, las frecuentes lecturas de memoria siguen siendo un cuello de botella para el rendimiento. Si la tabla de consulta puede hacerse lo suficientemente pequeña como para caber en registros, las operaciones de lectura de memoria pueden transformarse en instrucciones SIMD eficientes para la CPU.

En primer lugar, cada subvector se agrupa en 16 clases, por lo que un valor de 4 bits puede representar un centro de clúster -este es el origen del nombre "PQ de 4 bits". A continuación, las distancias representadas normalmente como valores flotantes se convierten a uint8 mediante cuantificación escalar (SQ). De este modo, la tabla de consulta de un subvector puede almacenarse en un registro que utiliza 16 × 8 = 128 bits.

Por último, examinemos la disposición de almacenamiento del registro (utilizando como ejemplo el conjunto de instrucciones AVX2): los subvectores de 32 vectores se colocan en un registro de 128 bits, combinados con la tabla de consulta. De este modo, la operación de "búsqueda" puede completarse eficientemente utilizando una única instrucción de CPU SIMD shuffle.

disposición de registros

SIMD Shuffle para la búsqueda

He aquí una observación interesante: el artículo de ScaNN se centra por completo en la primera optimización, lo cual es razonable ya que puede considerarse un artículo de algoritmos que hace hincapié en las derivaciones matemáticas. Sin embargo, los resultados experimentales presentados en el artículo son extraordinariamente impresionantes.

Los resultados experimentales presentados en el artículo de ScaNN.

Intuitivamente, la optimización de la pérdida por sí sola no debería producir efectos tan espectaculares. Otro blog también lo ha señalado: lo que realmente marca la diferencia es la parte de PQ FastScan de 4 bits.

Resultados experimentales

Hemos realizado una sencilla prueba con la herramienta de referencia de bases de datos vectoriales VectorDBBench. La ventaja de rendimiento de ScaNN sobre los tradicionales IVFFLAT e IVF_PQ es bastante evidente. Tras la integración en Milvus, en el conjunto de datos Cohere1M y con la misma tasa de recuperación, QPS puede alcanzar 5 veces el de IVFFLAT y 6 veces el de IVF_PQ.

Sin embargo, QPS es ligeramente inferior al de índices basados en gráficos como HNSW, por lo que no es la primera opción para casos de uso de alto QPS. Pero para escenarios con menor recall, es aceptable (como en algunos sistemas de recomendación), el uso de ScaNN sin cargar datos sin procesar puede lograr impresionantes QPS con una huella de memoria extremadamente baja (1/16 de los datos originales), lo que lo convierte en una excelente elección de índice.

Tipo_índice	Caso	QPS	latencia(p99)	recuperación	memoria
IVFFLAT	Rendimiento1M	266	0.0173s	0.9544	3G
HNSW	Rendimiento1M	1885	0.0054s	0.9438	3.24G
IVF_PQ	Rendimiento1M	208	0.0292s	0.928	0.375G
ScaNN（con_datos_sin procesar: true）	Rendimiento1M	1215	0.0069s	0.9389	3.186G
ScaNN（con_datos_sin procesar: falso）	Rendimiento1M	1265	0.0071s	0.7066	0.186G

Conclusión

ScaNN se basa en el conocido marco IVFPQ, pero lo lleva mucho más lejos gracias a un profundo trabajo de ingeniería tanto en la cuantificación como en la aceleración de la búsqueda de bajo nivel. Alineando el objetivo de cuantificación con la calidad de la clasificación y eliminando los cuellos de botella de memoria en el bucle interno, ScaNN combina la cuantificación basada en la puntuación con una ruta PQ FastScan de 4 bits que convierte un proceso tradicionalmente limitado por la memoria en un cálculo eficiente y SIMD.

En la práctica, esto da a ScaNN un nicho claro y bien definido. No está pensado para sustituir a los índices basados en grafos, como HNSW, en entornos de alta recuperación. En cambio, para cargas de trabajo insensibles a las llamadas con presupuestos de memoria ajustados, ScaNN ofrece un alto rendimiento con un tamaño muy reducido. En nuestros experimentos, tras la integración en Milvus VectorDB, ScaNN alcanzó aproximadamente 5 veces el QPS de IVFFLAT en el conjunto de datos Cohere1M, lo que lo convierte en una opción sólida para la recuperación de RNA de alto rendimiento y baja latencia en la que la compresión y la eficiencia importan más que la recuperación perfecta.

Si está interesado en profundizar en ScaNN o en debatir sobre la selección de índices en sistemas del mundo real, únase a nuestro canal de Slack para charlar con nuestros ingenieros y otros ingenieros de IA de la comunidad. También puede reservar una sesión individual de 20 minutos para obtener información, orientación y respuestas a sus preguntas a través de Milvus Office Hours.