Efficient Vector Search in the Wild: One Model for Multi-K Queries

O artigo apresenta o OMEGA, um método de busca aprendida que, ao treinar um modelo base apenas para K=1 e aplicar refinamentos dinâmicos, permite atender consultas de múltiplos valores de K com alta precisão e desempenho, reduzindo significativamente o tempo de latência e os custos de pré-processamento em comparação com métodos existentes.

O essencial

Imagine que você tem uma biblioteca gigante com milhões de livros (os vetores) e você precisa encontrar os mais parecidos com um livro que você acabou de pegar na mão (a consulta). O problema é que a biblioteca é tão grande que procurar um por um levaria uma eternidade.

Para resolver isso, os bibliotecários (os sistemas de banco de dados) organizam os livros em um mapa de conexões (um índice gráfico). Mas há um dilema:

• Se você pedir apenas 1 livro (K=1), o mapa é fácil de ler.
• Se você pedir 100 livros (K=100), você precisa explorar muito mais do mapa.

Até hoje, os "bibliotecários inteligentes" (modelos de IA) que ajudavam a parar a busca cedo demais eram treinados para apenas um número específico. Se você treinasse o robô para achar 1 livro, ele ficaria confuso e pararia muito cedo se você pedisse 100. Se você treinasse para 100, ele gastaria tempo demais procurando apenas 1.

Para ter um robô para cada situação, você teria que gastar uma fortuna e muito tempo treinando cada um deles antes de abrir a biblioteca. Isso é caro e lento.

A Solução: OMEGA (O "Mestre de Todos")

O artigo apresenta o OMEGA, um novo sistema que é como um chef de cozinha universal.

1. A Grande Ideia: "Um Modelo para Todos"

Em vez de treinar um robô para achar 1 livro, outro para 10, outro para 100, o OMEGA treina apenas um robô para achar o melhor livro de todos (o nº 1).

Como ele consegue achar os outros? Usando um truque de "máscara":

• Imagine que você achou o livro perfeito. Você o coloca em uma caixa e o esconde (mascara).
• Agora, você pergunta ao mesmo robô: "Qual é o novo melhor livro entre os que sobraram?"
• O robô, que já sabe achar o melhor, acha o segundo melhor com facilidade.
• Você esconde o segundo, e pede o terceiro. E assim por diante.

Isso transforma um problema difícil (achar 100 livros de uma vez) em 100 problemas fáceis (achar 1 livro de cada vez), usando o mesmo robô treinado.

2. O Problema do "Caminho" (Trajetória)

O robô precisa saber quando parar de procurar. Os robôs antigos olhavam para a distância exata do livro atual. Mas, quando você esconde (mascara) um livro, a distância muda e o robô fica confuso.

O OMEGA usa uma nova dica: a Trilha de Distância.
Em vez de olhar apenas para a distância atual, o robô olha para o padrão de como a distância está caindo enquanto ele anda pelo mapa. É como se ele olhasse para a inclinação da estrada. Se a estrada está ficando íngreme e descendo rápido, ele sabe que está chegando perto do tesouro. Esse padrão de "descida" funciona bem, mesmo quando você esconde livros no caminho.

3. O Truque Estatístico (Não chamar o robô toda hora)

Chamar o robô para cada livro (de 1 a 100) ainda pode ser lento. O OMEGA tem um segundo truque: Adivinhação Estatística.

O sistema aprendeu uma regra matemática: "Se eu já encontrei os 20 melhores livros, há 90% de chance de que o 50º melhor já esteja no meu grupo, sem eu precisar chamar o robô para verificar".
O sistema consulta uma "tabela de probabilidades" (como um oráculo rápido). Se a tabela diz que a chance de ter achado o que você quer é alta, ele para de chamar o robô e entrega o resultado. Isso economiza tempo e energia.

Por que isso é incrível? (Os Resultados)

• Economia de Tempo (Pré-processamento): Para treinar o OMEGA, você gasta o tempo de treinar apenas um robô (o do K=1). Os sistemas antigos precisavam de 2x, 3x ou até 6x mais tempo de treinamento para cobrir diferentes números de resultados.
• Velocidade na Hora da Busca: OMEGA é mais rápido porque não faz buscas desnecessárias. Ele para exatamente quando precisa, seja para 1 livro ou 100.
• Custo Total: Mesmo com um pouco mais de inteligência na hora da busca, o tempo total gasto (treinamento + busca) é menor. É como comprar um carro que gasta um pouco mais de gasolina, mas que é tão rápido que você chega ao trabalho 2 horas antes, economizando seu tempo de vida.

Resumo em uma Analogia Final

Imagine que você está procurando um amigo em uma multidão.

• Sistemas Antigos: Você contrata um guarda-costas para procurar 1 pessoa, outro para 10, outro para 50. É caro contratar todos e demora para treiná-los.
• OMEGA: Você contrata um guarda-costas super esperto que sabe achar a pessoa mais próxima.
  • Se você quer 1 amigo, ele acha.
  • Se você quer 10, ele acha o primeiro, você diz "esconde ele", e ele acha o segundo mais próximo, e assim por diante.
  • Além disso, ele tem um "olho de águia" (a estatística) que diz: "Já achei 20, a chance de ter achado os 50 que você quer é tão alta que não preciso mais chamar o chefe, posso entregar a lista agora".

O OMEGA é a primeira ferramenta que faz tudo isso com um único modelo, economizando tempo de preparação e entregando resultados mais rápidos e precisos, seja você um usuário pedindo 1 resultado ou 100.

Resumo técnico

1. O Problema

A busca por vizinhos mais próximos aproximados (ANNS - Approximate Nearest Neighbor Search) é fundamental para serviços de banco de dados vetoriais modernos, como motores de recomendação e RAG (Retrieval-Augmented Generation). No entanto, existem desafios críticos na implantação real ("no wild"):

• Consultas Multi-K: Diferentes aplicações exigem diferentes valores de $K$ (número de resultados retornados) para a mesma base de dados. O valor de $K$ varia dinamicamente dependendo do contexto da aplicação.
• Falha de Generalização dos Modelos Atuais: Métodos de busca aprendida (learned search) de última geração (como DARTH e LAET) são treinados para um valor específico de $K$. Quando aplicados a outros valores de $K$:
  • Para $K > K_{treino}$: O modelo tende a parar a busca muito cedo (sub-busca), resultando em baixa precisão (recall).
  • Para $K < K_{treino}$: O modelo continua buscando além do necessário (sobre-busca), resultando em alta latência.
• Custo de Pré-processamento Proibitivo: Treinar modelos separados para cada $K$ comum ou um único modelo que generalize todos os $K$ exige um tempo de pré-processamento (treinamento) massivo. Em ambientes de produção (como o da Alibaba), o custo de pré-processamento já consome ~22% dos recursos de computação. Aumentar esse tempo para treinar múltiplos modelos é economicamente inviável.

2. Metodologia: OMEGA

O artigo apresenta o OMEGA (One-Model Efficient Generalized ANNS), uma abordagem que utiliza um único modelo treinado para $K=1$ para atender consultas de qualquer $K$ com alta precisão e baixa latência.

Ideia Central

A busca para o top-$K$ pode ser decomposta em $K$ buscas sequenciais para o top-1.

1. O modelo prevê se o conjunto de busca atual contém o vizinho mais próximo (top-1).
2. Ao encontrar o top-1, ele é "mascarado" (excluído) do conjunto de busca.
3. O mesmo modelo é invocado novamente para encontrar o próximo top-1 (que se torna o top-2 global) no conjunto restante.
4. Este processo itera até encontrar os $K$ resultados.

Desafios e Soluções Técnicas

Desafio 1: Como treinar um modelo base eficaz para generalização com mascaramento?

• Problema: Características usadas em modelos existentes (como a distância mínima absoluta) falham quando vetores são mascarados, pois a distribuição de distâncias muda drasticamente.
• Solução (Trajetória de Distância): O OMEGA utiliza a trajetória de distância (o padrão de redução de distância durante a exploração do grafo) como característica principal.
  • A trajetória mostra uma tendência de queda acentuada à medida que se aproxima do alvo.
  • Essa tendência é robusta ao mascaramento: mesmo após remover o top-1, a trajetória para encontrar o próximo top-1 (o antigo top-2) mantém o mesmo padrão de queda.
  • O modelo é treinado apenas com consultas top-1, usando uma janela deslizante para capturar estatísticas (média, variância, percentis) das últimas distâncias visitadas.

Desafio 2: Como reduzir a sobrecarga de múltiplas invocações do modelo?

• Problema: Chamar o modelo $K$ vezes (uma para cada resultado) pode anular os ganhos de latência, especialmente para $K$ grandes.
• Solução (Previsão Estatística): O sistema explora uma propriedade estatística: dada a descoberta de $N$ vizinhos mais próximos, a probabilidade de o $r$-ésimo vizinho (onde $r > N$) já estar no conjunto de busca segue uma distribuição previsível.
  • Um tabela de lookup pré-perfilada (2D) mapeia pares $(N, r)$ para probabilidades de recall.
  • Antes de invocar o modelo para o próximo passo, o sistema calcula o recall esperado estatisticamente. Se o recall previsto atingir a meta, a busca é interrompida imediatamente sem invocar o modelo.

3. Arquitetura do Sistema

• Fase Offline (Pré-processamento): Após a compactação do índice vetorial, o sistema treina apenas um modelo (GBDT - Gradient Boosting Decision Tree) usando consultas top-1. O tempo de treinamento é minimizado com early stopping dinâmico.
• Fase Online (Serving):
  1. Executa a busca no grafo (ex: HNSW).
  2. Usa o modelo treinado para prever se o top-1 atual foi encontrado.
  3. Se encontrado, mascara o vetor e continua.
  4. Periodicamente, consulta a tabela estatística para prever se a meta de recall para o $K$ total já foi atingida, permitindo parar a busca antecipadamente.

4. Resultados Principais

O OMEGA foi avaliado em três conjuntos de dados públicos (BIGANN, DEEP, GIST) e três bases de dados de produção da Alibaba, sob cargas de trabalho reais multi-K.

• Desempenho de Latência:
  • Com o mesmo orçamento de pré-processamento, o OMEGA reduz a latência média em 6% a 33% comparado aos métodos de busca aprendida de última geração (DARTH, LAET).
  • Para atingir a mesma latência ótima que os baselines, o OMEGA requer apenas 16% a 30% do tempo de pré-processamento.
• Precisão (Recall):
  • Mantém o mesmo alvo de recall (geralmente 0.95 ou 0.99) que os métodos de referência em todas as consultas multi-K.
  • Evita a degradação de precisão para $K$ grandes e a perda de desempenho para $K$ pequenos.
• Custo Computacional Total:
  • Considerando o custo de pré-processamento + tempo de serviço (serving), o OMEGA reduz o uso total de CPU em 4% a 24% comparado a métodos fixos e 1% a 24% comparado a métodos aprendidos existentes.
• Eficiência de Treinamento:
  • Elimina a necessidade de treinar múltiplos modelos. O tempo de pré-processamento é fixo e mínimo (equivalente a treinar um único modelo top-1).

5. Significado e Contribuições

• Primeira Solução Generalizável K: O OMEGA é o primeiro método de busca aprendida que generaliza efetivamente para múltiplos valores de $K$ usando um único modelo treinado, resolvendo o dilema entre precisão e latência em cenários reais.
• Viabilidade Econômica: Ao reduzir drasticamente o tempo de pré-processamento (que é um custo oculto não cobrado dos usuários em muitos serviços de nuvem), o OMEGA torna a adoção de métodos de busca aprendida viável para provedores de banco de dados vetoriais em escala.
• Inovação em Features: A introdução da "trajetória de distância" como característica robusta a mascaramento é uma contribuição fundamental para o aprendizado de máquina em estruturas de grafos dinâmicos.
• Otimização Híbrida: A combinação de aprendizado de máquina (modelo top-1) com estatística preditiva (tabela de lookup) cria um sistema híbrido que minimiza a sobrecarga de inferência.

O código do OMEGA foi aberto como open-source e está sendo integrado ao banco de dados vetorial de código aberto da Alibaba, Zvec.