Uma breve introdução ao índice ScaNN

O ScaNN é a resposta da Google a um desafio familiar na pesquisa vetorial em grande escala: como aumentar o rendimento das consultas e reduzir a utilização de memória sem afetar de forma inaceitável a qualidade dos resultados. Conceptualmente, o ScaNN parte da receita clássica de IVF+PQ - clustering grosseiro mais quantização agressiva de produtos - mas acrescenta duas inovações importantes que alteram significativamente a fronteira do desempenho:

• Um objetivo de quantização sensível à pontuação que preserva melhor a ordem relativa dos verdadeiros vizinhos, melhorando a qualidade da classificação mesmo sob forte compressão.

• O FastScan é um caminho de pesquisa de PQ de 4 bits optimizado por SIMD que reduz o estrangulamento tradicional da carga de memória, mantendo mais trabalho nos registos da CPU.

Na prática, é uma boa escolha quando não há problema em trocar um pouco de recall por um QPS alto e um espaço de memória muito menor (geralmente compactando vetores para ~1/16 do tamanho original), como em cargas de trabalho de recomendação insensíveis ao recall.

Nesta postagem, revisitaremos o IVFPQ como linha de base, exploraremos as principais otimizações que o ScaNN introduz em cima dele e finalizaremos com resultados experimentais que fundamentam a discussão no desempenho medido.

Recapitulação do IVFPQ

O ScaNN foi proposto pelo Google em 2020, e o artigo relata uma melhoria de desempenho de 3 × em relação ao HNSW no conjunto de dados GloVe. Você pode consultar o artigo original e a implementação de código aberto para obter detalhes.

Antes de apresentar o ScaNN, vamos recapitular brevemente o IVFPQ, uma vez que o ScaNN é construído sobre a mesma estrutura geral.

IVFPQ significa Inverted File with Product Quantization (Arquivo invertido com quantificação de produto), um algoritmo usado para pesquisa eficiente e em larga escala do Approximate Nearest Neighbor (ANN) em bancos de dados vetoriais de alta dimensão. É uma abordagem híbrida que combina duas técnicas, o índice de ficheiro invertido (IVF) e a quantização de produtos (PQ), para equilibrar a velocidade de pesquisa, a utilização de memória e a precisão.

Como funciona o IVFPQ

O processo envolve duas etapas principais durante a indexação e a pesquisa:

• Camada IVF: os vetores são agrupados em nlist listas invertidas (clusters). No momento da consulta, apenas um subconjunto de clusters é visitado (nprobe) para compensar a recuperação e a latência/rendimento.

• Camada PQ: em cada agregado visitado, cada vetor de dimensão D é dividido em m subvectores, cada um de dimensão (D/m). Cada subvector é quantificado, atribuindo-o ao centróide mais próximo no seu sub-codificador. Se um sub-livro de códigos tiver 256 centróides, cada subvector pode ser representado por um código uint8 (um ID em [0, 255]).

O cálculo da distância pode então ser reescrito como a soma dos 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)

Aqui, L representa uma tabela de pesquisa. No momento da consulta, a tabela de pesquisa é construída, registando a distância entre a consulta e cada subvector quantificado. Todos os cálculos de distância subsequentes são convertidos em pesquisas na tabela seguidas de somatório.

Por exemplo, para vectores de 128 dimensões divididos em 32 subvectores de 4 dimensões cada, se cada subvector for codificado por um uint8 ID, o custo de armazenamento por vetor cai de (128 x 4) bytes para (32 x 1) bytes - uma redução de 1/16.

Optimizações ScaNN baseadas em IVFPQ

Em resumo, o ScaNN melhora o IVFPQ em dois aspectos:

1. Quantização: A ScaNN propõe um objetivo para além da simples substituição de cada subvector pelo seu centróide k-means mais próximo (ou seja, minimizando o erro de reconstrução).

2. Eficiência de pesquisa: O ScaNN acelera a pesquisa baseada em LUT, que é frequentemente limitada pela memória, através de um caminho FastScan amigável ao SIMD.

Perda de quantização sensível à pontuação

Uma vez que a análise se baseia na métrica IP, após o ScaNN decompor o erro de quantização em componentes paralelos e perpendiculares, apenas o componente paralelo afecta o resultado, pelo que deve ser aplicado um termo de penalização maior. Consequentemente, a função de perda pode ser reescrita da seguinte forma:

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

A figura abaixo mostra um exemplo bidimensional que ilustra que o erro causado pelo componente paralelo é maior e pode levar a resultados incorrectos do vizinho mais próximo, justificando assim uma penalização mais severa.

A figura da esquerda mostra uma quantização deficiente porque o desvio paralelo afecta o resultado final, enquanto a figura da direita mostra uma quantização melhor.

PQ FastScan de 4 bits

Comecemos por rever o processo de cálculo do PQ: durante a consulta, as distâncias entre a consulta e os centros de agrupamento do subvector são pré-computadas para construir uma tabela de pesquisa. O cálculo da distância é então efectuado através de pesquisas na tabela para obter as distâncias dos segmentos e somá-las.

No entanto, as leituras frequentes de memória ainda se tornam um gargalo de desempenho. Se a tabela de pesquisa puder ser suficientemente pequena para caber nos registos, as operações de leitura da memória podem ser transformadas em instruções SIMD eficientes da CPU.

Primeiro, cada subvector é agrupado em 16 classes, pelo que um valor de 4 bits pode representar um centro de agrupamento - esta é a origem do nome "PQ de 4 bits". Em seguida, as distâncias tipicamente representadas como floats são convertidas para uint8 usando a Quantização Escalar (SQ). Desta forma, a tabela de pesquisa para um subvector pode ser armazenada num registo usando 16 × 8 = 128 bits.

Finalmente, examinemos a disposição de armazenamento do registo (utilizando o conjunto de instruções AVX2 como exemplo): os subvectores de 32 vectores são colocados num registo de 128 bits, combinados com a tabela de pesquisa. A operação "lookup" pode então ser eficientemente completada usando uma única instrução SIMD shuffle da CPU.

disposição dos registos

SIMD Shuffle para Lookup

Eis uma observação interessante: o artigo do ScaNN concentra-se inteiramente na primeira otimização, o que é razoável, uma vez que pode ser considerado um artigo sobre algoritmos que enfatiza as derivações matemáticas. No entanto, os resultados experimentais apresentados no documento são notavelmente impressionantes.

Os resultados experimentais apresentados no documento ScaNN.

Intuitivamente, a otimização da perda por si só não deveria produzir efeitos tão dramáticos. Outro blogue também salientou este facto - o que realmente faz a diferença é a parte do PQ FastScan de 4 bits.

Resultados experimentais

Usando a ferramenta de benchmark de banco de dados vetorial VectorDBBench, realizamos um teste simples. A vantagem de desempenho do ScaNN sobre os tradicionais IVFFLAT e IVF_PQ é bastante evidente. Após a integração no Milvus, no conjunto de dados Cohere1M, com a mesma taxa de recuperação, o QPS pode atingir 5x o do IVFFLAT e 6x o do IVF_PQ.

No entanto, o QPS é ligeiramente inferior ao de índices baseados em gráficos como o HNSW, pelo que não é a primeira escolha para casos de utilização de elevado QPS. No entanto, para cenários com menor recuperação, é aceitável (como em alguns sistemas de recomendação). A utilização do ScaNN sem carregar dados em bruto pode atingir um QPS impressionante com um espaço de memória extremamente reduzido (1/16 dos dados originais), o que o torna uma excelente opção de índice.

Tipo_de_índice	Caso	QPS	latência(p99)	recordar	memória
IVFFLAT	Desempenho1M	266	0.0173s	0.9544	3G
HNSW	Desempenho1M	1885	0.0054s	0.9438	3.24G
IVF_PQ	Desempenho1M	208	0.0292s	0.928	0.375G
ScaNN（com dados brutos: verdadeiro）	Desempenho1M	1215	0.0069s	0.9389	3.186G
ScaNN（com dados brutos: falso）	Desempenho1M	1265	0.0071s	0.7066	0.186G

Conclusão

O ScaNN baseia-se na estrutura familiar do IVFPQ, mas leva-o significativamente mais longe através de um profundo trabalho de engenharia tanto na quantização como na aceleração de pesquisa de baixo nível. Ao alinhar o objetivo da quantização com a qualidade da classificação e ao eliminar os estrangulamentos de memória no ciclo interno, o ScaNN combina a quantização consciente da pontuação com um caminho PQ FastScan de 4 bits que transforma um processo tradicionalmente limitado à memória num cálculo eficiente e compatível com SIMD.

Na prática, isso dá ao ScaNN um nicho claro e bem definido. Não se destina a substituir os índices baseados em grafos, como o HNSW, em contextos de alta recordação. Em vez disso, para cargas de trabalho insensíveis à recuperação com orçamentos de memória apertados, o ScaNN oferece alto rendimento com uma pegada muito pequena. Nas nossas experiências, após a integração no Milvus VectorDB, o ScaNN alcançou cerca de 5× o QPS do IVFFLAT no conjunto de dados Cohere1M, tornando-o uma forte escolha para a recuperação ANN de alto rendimento e baixa latência, onde a compressão e a eficiência são mais importantes do que a recuperação perfeita.

Se você estiver interessado em explorar mais o ScaNN ou discutir a seleção de índices em sistemas do mundo real, junte-se ao nosso Canal do Slack para conversar com nossos engenheiros e outros engenheiros de IA na comunidade. Também pode reservar uma sessão individual de 20 minutos para obter informações, orientação e respostas às suas perguntas através do Milvus Office Hours.