Acelerando a fusão de gráficos HNSW

Da busca vetorial às poderosas REST APIs, o Elasticsearch oferece aos desenvolvedores o kit de ferramentas de busca mais abrangente. Mergulhe em notebooks de exemplo no GitHub para experimentar algo novo. Você também pode iniciar seu teste gratuito ou executar o Elasticsearch localmente hoje mesmo.

No passado, discutimos alguns dos desafios de ter que pesquisar em vários grafos HNSW e como conseguimos mitigá-los. Naquela ocasião, mencionamos algumas melhorias adicionais que tínhamos planejado. Este post é o culminar desse trabalho.

Você pode perguntar: por que usar vários gráficos? Este é um efeito colateral de uma escolha arquitetônica no Lucene: segmentos imutáveis. Como acontece com a maioria das escolhas arquitetônicas, há prós e contras. Por exemplo, recentemente disponibilizamos o Elasticsearch sem servidor. Nesse contexto, obtivemos benefícios muito significativos de segmentos imutáveis, incluindo replicação de índice eficiente e a capacidade de desacoplar índice e computação de consulta e dimensioná-los automaticamente de forma independente. Para quantização vetorial, as fusões de segmentos nos dão a oportunidade de atualizar parâmetros para adaptá-los às características dos dados. Nessa linha, acreditamos que há outras vantagens proporcionadas pela oportunidade de medir características de dados e revisitar opções de indexação.

Nesta postagem, discutiremos o trabalho que temos feito para reduzir significativamente a sobrecarga de construção de vários gráficos HNSW e, em particular, para reduzir o custo de mesclagem de gráficos.

Histórico

Para manter um número gerenciável de segmentos, o Lucene verifica periodicamente se deve mesclar segmentos. Isso equivale a verificar se a contagem de segmentos atual excede uma contagem de segmentos de destino, que é determinada pelo tamanho do segmento base e pela política de mesclagem. Se a contagem for excedida, o Lucene mescla grupos de segmentos enquanto a restrição for violada. Este processo foi descrito em detalhes em outro lugar.

Lucene opta por mesclar segmentos de tamanhos semelhantes porque isso atinge um crescimento logarítmico na amplificação de gravação. No caso de um índice vetorial, a amplificação de escrita é o número de vezes que um vetor será inserido em um gráfico. O Lucene tentará mesclar segmentos em grupos de aproximadamente 10. Consequentemente, os vetores são inseridos em um gráfico aproximadamente $1+\frac{9}{10}\log_{10}\left(\frac{n}{n_0}\right)$ vezes, onde $n$ é a contagem do vetor de índice e $n_0$ é a contagem do vetor de segmento base esperado. Devido ao crescimento logarítmico, a amplificação da gravação é de um dígito, mesmo para índices grandes. Entretanto, o tempo total gasto na fusão de gráficos é linearmente proporcional à amplificação da gravação.

Ao mesclar gráficos HNSW, já fazemos uma pequena otimização: mantendo o gráfico do maior segmento e inserindo vetores dos outros segmentos nele. Esta é a razão do fator 9/10 acima. Abaixo mostramos como podemos melhorar significativamente usando informações de todos os gráficos que estamos mesclando.

Fusão de gráficos HNSW

Anteriormente, mantínhamos o grafo maior e inseríamos vetores dos demais, ignorando os grafos que os continham. A principal conclusão que utilizamos a seguir é que cada grafo HNSW descartado contém informações importantes sobre a proximidade dos vetores que ele contém. Gostaríamos de usar essas informações para acelerar a inserção, pelo menos de alguns, dos vetores.

Nós nos concentramos no problema de inserir um gráfico menor $G_s=(V_s,E_s)$ em um gráfico maior $G_l=(V_l,E_l)$, já que esta é uma operação atômica que podemos usar para construir qualquer política de mesclagem.

A estratégia é encontrar um subconjunto de vértices de $J\subset V_s$ para inserir no gráfico grande. Em seguida, usamos a conectividade desses vértices no pequeno gráfico para acelerar a inserção dos vértices restantes $V_s \setminus J$. A seguir, usamos $N_s(u)$ e $N_l(u)$ para denotar os vizinhos de um vértice $u$ no gráfico pequeno e grande, respectivamente. Esquematicamente o processo é o seguinte.

MERGE-HNSW

Inputs $G_s$ and $G_l$

1$\;\;$Find $J\subset V_s$ to insert into $G_l$ using COMPUTE-JOIN-SET
2$\;\;$Insert each vertex $u\in J$ into $G_l$
3$\;\;$for $u\in V_s \setminus J$ do
4$\;\;\;\;J_u\leftarrow J\cap N_s(u)$
5$\;\;\;\;E_u\leftarrow \cup_{v\in J_u} N_l(u)$
6$\;\;\;\;W\leftarrow\,$FAST-SEARCH-LAYER$(J_u, E_u)$
7$\;\;\;\;neighbors \leftarrow\,$SELECT-NEIGHBORS-HEURISTIC$(u, W)$
8$\;\;\;\;J\leftarrow J \cup \{u\}$

Calculamos o conjunto $J$ usando um procedimento que discutiremos abaixo (linha 1). Em seguida, inserimos cada vértice em $J$ no gráfico grande usando o procedimento de inserção padrão HNSW (linha 2). Para cada vértice que não inserimos, encontramos seus vizinhos que inserimos e seus vizinhos no gráfico grande (linhas 4 e 5). Usamos um procedimento FAST-SEARCH-LAYER semeado com este conjunto (linha 6) para encontrar os candidatos para o SELECT-NEIGHBORS-HEURISTIC do artigo HNSW (linha 7). Na verdade, estamos substituindo SEARCH-LAYER para encontrar o conjunto de candidatos no método INSERT (Algoritmo 1 do artigo), que de outra forma não será alterado. Por fim, adicionamos o vértice que acabamos de inserir em $J$ (linha 8).

É claro que para que isso funcione, cada vértice em $V_s \setminus J$ deve ter pelo menos um vizinho em $J$. Na verdade, exigimos que para cada vértice em $u\in V_s \setminus J$ que $|J\cap N_s(u) |\geq k_u$ para algum $k_u<M$, a conectividade máxima da camada. Observamos que em gráficos HNSW reais vemos uma grande dispersão de graus de vértices. A figura abaixo mostra uma função de densidade cumulativa típica do grau do vértice para a camada inferior de um gráfico Lucene HNSW.

Exploramos o uso de um valor fixo para $k_u$ e também o tornamos uma função do grau do vértice. Esta segunda opção leva a maiores acelerações com impacto mínimo na qualidade do gráfico, então optei pela seguinte opção

Observe que $|N_s(u)$| é igual ao grau do vértice $u$ no gráfico pequeno por definição. Ter um limite inferior de dois significa que inseriremos todos os vértices cujo grau seja menor que dois.

Um argumento de contagem simples sugere que se escolhermos $J$ cuidadosamente, precisamos apenas inserir em torno de $\frac{1}{5}|V_s|$ em $G_l$ diretamente. Especificamente, colorimos uma aresta do gráfico se inserirmos exatamente um de seus vértices finais em $J$. Então sabemos que para cada vértice em $V_s \setminus J$ ter pelo menos $k_u$ vizinhos em $J,$ precisamos colorir pelo menos $\sum_{u\in V_s\setminus J} k_u$ arestas. Além disso, esperamos que

Aqui, $\mathbb{E}_U\left[N_s(U)|\right]$ é o grau médio do vértice no pequeno gráfico. Para cada vértice $u\in J,$ colorimos no máximo $|N_s(u)|$ arestas. Portanto, o número total de arestas que esperamos colorir é no máximo $|J|\, \mathbb{E}_U\left[|N_s(U)|\right]$. Esperamos que, ao escolher $J$ cuidadosamente, possamos colorir perto desse número de arestas e, portanto, para cobrir todos os vértices, $|J|$ precisa satisfazer

Isso implica que $|J|=\frac{1}{4}(|V_s|-|J|)=\frac{4}{5}\frac{1}{4}|V_s|=\frac{1}{5}|V_s|$.

Se SEARCH-LAYER dominar o tempo de execução, isso sugere que poderíamos atingir uma velocidade de até $5 vezes$ maior no tempo de mesclagem. Dado o crescimento logarítmico da amplificação de gravação, isso significa que mesmo para índices muito grandes, normalmente dobraríamos apenas o tempo de construção em comparação à construção de um gráfico.

O risco dessa estratégia é que prejudicamos a qualidade do gráfico. Inicialmente tentamos com um no-op FAST-SEARCH-LAYER. Descobrimos que isso degrada a qualidade do gráfico a ponto de a recuperação como função da latência ser impactada, principalmente ao mesclar em um único segmento. Em seguida, exploramos várias alternativas usando uma busca limitada do gráfico. No final, a escolha mais eficaz foi a mais simples. Use SEARCH-LAYER mas com um ef_construction baixo. Com essa parametrização conseguimos obter gráficos de excelente qualidade e ainda diminuir o tempo de mesclagem em pouco mais de 30% em média.

Calculando o conjunto de junção

Encontrar um bom conjunto de junção pode ser formulado como um problema de cobertura de grafos HNSW. Uma heurística gulosa é uma heurística simples e eficaz para aproximar coberturas ótimas de grafos. A abordagem que adotamos seleciona os vértices um de cada vez para adicionar a $J$ em ordem decrescente de ganho. O ganho é definido da seguinte forma:

Aqui, $c(v)$ denota a contagem de vizinhos de um vetor $v$ em $J$ e $1\{\cdot\}$ é a função indicadora. O ganho inclui a mudança na contagem do vértice que adicionamos a $J$, ou seja, $\max(k_v-c(v),0)$, pois nos aproximamos do nosso objetivo adicionando um vértice menos coberto. O cálculo do ganho é ilustrado na figura abaixo para o vértice central laranja.

Mantemos o seguinte estado para cada vértice $v$:

1. Seja velho,
2. Seu ganho $Gain(v)$,
3. A contagem de vértices adjacentes em $J$ denotada por $c(v)$,
4. Um número aleatório no intervalo [0,1] que é usado para desempate.

O pseudocódigo para calcular o conjunto de junções é o seguinte.

COMPUTE-JOIN-SET

Inputs $G_s$

1$\;\;C\leftarrow\emptyset$
2$\;\;Gain_{exit}\leftarrow 0$
3$\;\;$for $u\in G_s$ do
4$\;\;\;\;C\leftarrow C \cup \{(\text{false}, k_u+deg(u), 0, \text{rand in }[0,1])\}$
5$\;\;\;\;Gain_{exit}\leftarrow Gain_{exit}+k_u$
6$\;\;Gain_{tot}\leftarrow 0$
7$\;\;$while $Gain_{tot}<Gain_{exit}$ do
8$\;\;\;\;v^*\leftarrow\,$ maximum gain vertex in $C$
9$\;\;\;\;$Remove the state for $v^*$ from $C$
10$\;\;\;$if $v^*$ is not stale then
11$\;\;\;\;\;J\leftarrow J\cup\{v^*\}$
12$\;\;\;\;\;Gain_{tot}\leftarrow Gain_{tot}+Gain(v^*)$
13$\;\;\;\;\;$for $u \in N_s(v^*)$ do
14$\;\;\;\;\;\;\;$mark $u$ as stale if $c(v^*)<k_{v^*}$
15$\;\;\;\;\;\;\;$mark neighbors of $u$ stale if $c(u)=k_u-1$
16$\;\;\;\;\;\;\;c(u)\leftarrow c(u)+1$
17$\;\;\;$else
18$\;\;\;\;\;Gain(v^*)\leftarrow \max(k_v-c(v^*),0)+\sum_{u\in N_s(v^*)\setminus J} 1\left\{c(u)<k_u\right\}$
19$\;\;\;\;\;$if $Gain(v^*)>0$ then
20$\;\;\;\;\;\;\;C\leftarrow C \cup \{(\text{false},Gain(v^*),c(v^*),\text{copy rand})\}$
21$\;\;$return $J$

Primeiro inicializamos o estado nas linhas 1-5.

Em cada iteração do loop principal, extraímos inicialmente o vértice de ganho máximo (linha 8), desfazendo empates aleatoriamente. Antes de fazer qualquer alteração, precisamos verificar se o ganho do vértice está obsoleto. Em particular, cada vez que adicionamos um vértice em $J$ , afetamos o ganho de outros vértices:

1. Como todos os seus vizinhos têm um vizinho adicional em $J,$ seus ganhos podem mudar (linha 14)
2. Se algum dos seus vizinhos estiver agora totalmente coberto, os ganhos de todos os seus vizinhos podem mudar (linhas 14-16)

Nós recalculamos os ganhos de forma lenta, então só recalculamos o ganho de um vértice se quisermos inseri-lo em $J$ (linhas 18-20). Como os ganhos sempre diminuem, nunca podemos perder um vértice que devemos inserir.

Observe que precisamos apenas monitorar o ganho total de vértices que adicionamos a $J$ para determinar quando sair. Além disso, enquanto $Gain_{tot}<Gain_{exit} ,$ pelo menos um vértice terá ganho diferente de zero, então sempre progredimos.

Resultados

Realizamos experimentos em quatro conjuntos de dados que juntos abrangem nossas três métricas de distância suportadas (euclidiana, cosseno e produto interno):

1. quora-E5-small: 522931 documentos, 384 dimensões e usa similaridade de cosseno,
2. cohere-wikipedia-v2: 1 milhão de documentos, 768 dimensões e usa similaridade de cosseno,
3. gist: 1M documentos, 960 dimensões e utiliza distância euclidiana e
4. cohere-wikipedia-v3: 1 milhão de documentos, 1.024 dimensões e usa o produto interno máximo.

Para cada conjunto de dados, avaliamos dois níveis de quantização:

1. int8 – que usa um inteiro de 1 byte por dimensão e
2. Churrasco – que usa um único bit por dimensão.

Por fim, para cada experimento, avaliamos a qualidade da pesquisa em duas profundidades de recuperação e examinamos depois da construção do índice e depois da fusão forçada em um único segmento.

Em resumo, alcançamos acelerações substanciais e consistentes na indexação e mesclagem, mantendo a qualidade do gráfico e, consequentemente, o desempenho da pesquisa em todos os casos.

Experimento 1: quantização int8

As acelerações médias da linha de base até o candidato, as mudanças propostas, são:

Aceleração do tempo de índice: 1,28$vezes$

Forçar aceleração de fusão: 1,72$vezes$

Isso corresponde à seguinte repartição nos tempos de execução

Para completar, os horários exatos são

Abaixo, mostramos os gráficos de recall versus latência que comparam o candidato (linhas tracejadas) à linha de base em duas profundidades de recuperação: recall@10 e recall@100 para índices com vários segmentos (o resultado final da nossa estratégia de mesclagem padrão após indexar todos os vetores) e após a mesclagem forçada para um único segmento. Uma curva mais alta e mais à esquerda é melhor, o que significa maior recuperação com menor latência.

Como você pode ver, para índices de múltiplos segmentos, o candidato é melhor para o conjunto de dados Cohere v3 e um pouco pior, mas quase comparável, para todos os outros conjuntos de dados. Após a fusão em um único segmento, as curvas de recall são quase idênticas para todos os casos.

Experimento 2: quantização BBQ

As acelerações médias da linha de base até o candidato são:

Aceleração do tempo de índice: 1,33$vezes$

Forçar aceleração de fusão: 1,34$vezes$

Isso corresponde à seguinte repartição nos tempos de execução

Para completar, os horários exatos são

Para índices de múltiplos segmentos, o candidato é melhor para quase todos os conjuntos de dados, exceto o Cohere v2, onde a linha de base é ligeiramente melhor. Para os índices de segmento único, as curvas de recall são quase idênticas para todos os casos.

Conclusão

O algoritmo discutido neste blog estará disponível no próximo Lucene 10.2 e na versão do Elasticsearch baseada nele. Os usuários poderão aproveitar o desempenho de mesclagem aprimorado e o tempo reduzido de criação de índice nessas novas versões. Essa mudança faz parte do nosso esforço contínuo para tornar o Lucene e o Elasticsearch rápidos e eficientes para pesquisas vetoriais e híbridas.