Google ScaNN

01 · CONTEXTO

Quando buscar por significado

Este caso parte de um post da Google Research, de julho de 2020, que apresenta o ScaNN, uma biblioteca de busca por similaridade de vetores, e o artigo "Accelerating Large-Scale Inference with Anisotropic Vector Quantization", que descreve a técnica de compressão por trás dela.

A busca por embeddings serve a consultas que dependem de significado e não de campos fáceis de indexar. Modelos de ML transformam textos e imagens em vetores de alta dimensão em que entradas parecidas ficam próximas, e esse é o ponto de partida dos sistemas de busca por vetores.

No modelo de duas torres, uma rede neural mapeia a consulta e outra mapeia os itens do banco para o mesmo espaço. Para responder, o sistema mapeia a consulta e então procura, entre todos os embeddings do banco, os mais próximos.

1. A aplicação envia a consulta, por exemplo um texto, à API de busca.

2. A API passa a consulta à torre de consulta, que a mapeia para o espaço de embeddings.

3. A torre de itens mapeia os itens do banco para o mesmo espaço, e os vetores resultantes ficam guardados.

4. O índice recebe o embedding da consulta e procura, entre os embeddings do banco, os mais próximos.

02 · O PROBLEMA

Achar vizinhos em escala

Uma forma comum de medir a semelhança entre consulta e itens é o produto interno, e a busca do vizinho mais próximo nesse critério é chamada de maximum inner-product search, ou MIPS. O conjunto de embeddings costuma ser grande demais para busca exaustiva, e a alta dimensionalidade dificulta podar candidatos.

Com bancos de milhões ou bilhões de itens, a MIPS costuma ser o gargalo de velocidade da inferência e a busca exaustiva é impraticável. Isso leva a algoritmos aproximados, que trocam um pouco de acurácia por um ganho grande de velocidade sobre a força bruta.

03 · COMPRESSÃO

Quantização que mira a MIPS

Várias soluções de ponta para MIPS comprimem os itens do banco para calcular uma aproximação do produto interno em uma fração do tempo da força bruta. Na quantização aprendida, um codebook de vetores é treinado a partir do banco e usado para representar aproximadamente seus elementos.

Esquemas anteriores quantizavam para minimizar a distância média entre cada vetor e sua forma quantizada. O post observa que isso não equivale a otimizar a acurácia da busca, e que codificações com distância média maior podem dar acurácia de MIPS superior.

A razão é a direção do erro. O erro paralelo ao vetor é mais prejudicial porque afeta desproporcionalmente os produtos internos altos, os que a MIPS quer estimar bem. Por isso a técnica, chamada de quantização anisotrópica, penaliza mais esse erro e aceita errar mais nos produtos internos baixos.

1. Os vetores do banco alimentam o treino de um codebook, o conjunto de vetores usado para representá-los de forma aproximada.

2. A perda usada no treino penaliza mais o erro paralelo ao vetor original.

3. O quantizador representa cada vetor do banco por um centro do codebook, e o índice comprimido mantém essas representações.

4. Na consulta, o estimador calcula produtos internos aproximados sobre as representações comprimidas, em uma fração do tempo da força bruta.

5. As estimativas alimentam a seleção dos melhores itens; a quantização anisotrópica melhora justamente a estimativa dos produtos internos com chance de estar entre os melhores.

04 · RESULTADO

Velocidade contra acurácia

O post mede o resultado em um benchmark do ann-benchmarks, o glove-100-angular: o ScaNN superou onze outras bibliotecas de busca por similaridade de vetores, atendendo cerca do dobro de consultas por segundo, para uma dada acurácia, que a segunda mais rápida.

A acurácia é medida por recall@k, uma métrica comum de acurácia da busca de vizinhos mais próximos. O post mostra o ScaNN com desempenho superior de forma consistente em vários pontos do compromisso entre velocidade e acurácia.

05 · LIÇÕES

O que levar para a entrevista

Os autores veem a quantização anisotrópica como um exemplo de ganho obtido ao otimizar o algoritmo para o objetivo final, a acurácia da busca, e não para uma meta intermediária como a distorção de compressão.

Em um desenho de busca vetorial, vale dizer qual recall o produto exige, qual velocidade ele precisa e como a compressão dos vetores muda esse compromisso. Escolher a métrica de erro da compressão em função da tarefa de busca é a lição central deste caso.