Embedding-based retrieval na busca do Facebook

01 · CONTEXTO

Por que a busca social pede embeddings

No momento do paper, a busca do Facebook ainda se baseava sobretudo em um modelo de correspondência Booleana, embora a recuperação por embeddings já fosse usada em buscadores web havia anos. Ali, a intenção de busca não depende só do texto: é influenciada também por quem faz a busca e pelo contexto em que ela acontece. Por isso, para os autores, a recuperação por embeddings no Facebook não é um problema de embedding de texto.

Uma busca tem um estágio de recuperação, que procura um conjunto de documentos relevantes com baixa latência e baixo custo computacional, e um estágio de ranking, que coloca os documentos mais desejados no topo com modelos mais complexos. A diferença de escala é grande: o ranking costuma considerar centenas de documentos por sessão, enquanto a recuperação precisa processar bilhões ou trilhões de documentos do índice.

A recuperação por embeddings (EBR) usa embeddings para representar a consulta e os documentos e converte o problema de recuperação em uma busca de vizinhos mais próximos no espaço de embeddings.

02 · MODELO E TREINO

Embedding unificado de dois lados

Os autores formulam a recuperação como um problema de otimização de recall. O conjunto de resultados-alvo de uma consulta pode vir de cliques ou de documentos considerados relevantes por avaliação humana. Como métrica offline, propõem rodar uma busca KNN no índice inteiro e usar o recall@K.

O modelo unificado tem dois lados: de um, a requisição de busca, com texto da consulta, pessoa que busca e contexto; do outro, o documento. Os dois codificadores são, por padrão, redes separadas, com a opção de compartilhar parte dos parâmetros. Os vetores densos são comparados por similaridade de cosseno, e a perda triplet tem a intuição de separar o par positivo do par negativo por uma margem de distância.

A margem exige ajuste: o valor ótimo varia bastante entre tarefas de treino, e margens diferentes resultaram em uma variação de 5-10% no recall de KNN.

Entradas além do texto fazem diferença. Ao trocar embeddings de texto por embeddings unificados, o recall melhorou +18% na busca de eventos e +16% na de grupos. Entre os atributos adicionados estão a localização, vantajosa em buscas por negócios, grupos e eventos locais, e atributos sociais, vindos de um modelo de embedding separado treinado sobre o grafo social.

Comparados à correspondência Booleana por termos, embeddings treinados só com atributos de texto se mostraram bons em dois cenários: a correspondência aproximada de texto e a opcionalização de termos, como descartar um termo da consulta e ainda recuperar o grupo esperado.

Os autores atribuem a regressão a um viés: os negativos sem clique tendem a ser casos difíceis, enquanto a maioria dos documentos do índice são casos fáceis, que não casam com a consulta. Treinar só com negativos assim muda a representatividade dos dados de treino em relação à tarefa real de recuperação.

03 · SERVING

Recuperação híbrida sobre o índice invertido

A forma direta de servir seria juntar os candidatos da recuperação existente com os de uma busca KNN de embeddings. Os autores consideraram isso subótimo por três motivos: alto custo de desempenho no experimento inicial, alto custo de manutenção por haver dois índices e grande sobreposição possível entre os dois conjuntos de candidatos.

A solução foi uma recuperação híbrida, que integra KNN de embeddings e correspondência Booleana para pontuar documentos. Para isso usaram a biblioteca Faiss na quantização dos vetores e a integraram à recuperação por índice invertido. Os autores citam duas vantagens da busca aproximada baseada em índice invertido: custo de armazenamento menor, pela quantização dos vetores, e integração mais fácil ao sistema de recuperação existente.

A integração foi feita no Unicorn, o motor de recuperação por trás da maioria dos produtos de busca do Facebook. Os autores estenderam a representação do documento para incluir embeddings e criaram o operador de consulta (nn), que casa os documentos cujo embedding está dentro de um raio do embedding da consulta. Na indexação, o embedding de cada documento é quantizado e vira um termo, para seu cluster grosso, e um payload, para o resíduo quantizado. Na consulta, o operador é reescrito internamente como um OR dos termos dos clusters mais próximos do embedding da consulta.

O modelo de dois lados foi decomposto em um modelo de consulta e um de documento, servidos separadamente. A consulta passa por um serviço online de inferência, em tempo real. Os documentos são processados em lote offline com Spark, e os embeddings gerados são publicados, junto de outros metadados, no índice direto; com a quantização adicional (grossa e PQ), também no índice invertido.

1. A consulta chega com o texto, a pessoa que busca e o contexto, e vai ao serviço online que calcula o embedding da requisição em tempo real.

2. O embedding da consulta segue para o motor de recuperação híbrida.

3. O motor reescreve o operador de vizinhos mais próximos em um OR dos termos dos clusters mais próximos e o combina com os termos Booleanos da expressão.

4. Fora do caminho da consulta, a inferência em lote publica no índice os embeddings dos documentos, já quantizados.

5. Os candidatos recuperados seguem para os estágios de ranking e filtragem.

Como o suporte a vizinhos mais próximos foi construído sobre primitivas existentes, e não como um sistema separado, o motor herdou atualizações em tempo real, planejamento e execução eficientes de consultas e consultas multi-hop. Entre os modos de busca, o de raio deu um melhor equilíbrio entre desempenho do sistema e qualidade do que o top-K, e é o usado na produção descrita no paper.

04 · AJUSTE DO ANN

Quantização e parâmetros de busca aproximada

A quantização dos embeddings tem dois componentes: a quantização grossa, que agrupa os vetores em clusters, tipicamente por K-means, e a quantização de produto, mais fina, que permite calcular distâncias de forma eficiente. Entre os parâmetros a ajustar estão o número de clusters grossos, que afeta desempenho e recall, o número de bytes do PQ e o nprobe, que define quantos clusters serão varridos e também afeta desempenho e recall.

Os autores também recomendam implantar várias configurações de ANN online para entender o impacto real em desempenho da recuperação por embeddings no sistema, o que ajuda a decidir o orçamento de capacidade e a reduzir o escopo da busca offline de parâmetros.

05 · HARD MINING

Negativos e positivos difíceis

No modelo de busca de pessoas, os K primeiros resultados por embedding tinham, em geral, o mesmo nome, e o modelo nem sempre colocava o alvo acima dos demais, mesmo com atributos sociais. Os autores achavam muito provável que os negativos de treino fossem fáceis demais, por serem amostras aleatórias, em geral com nomes diferentes. Para o modelo diferenciar melhor resultados parecidos, propuseram usar como negativos difíceis amostras mais próximas dos positivos no espaço de embeddings.

No HNM online, os negativos difíceis são selecionados dinamicamente em cada lote de treino: os documentos de maior similaridade viram os negativos mais difíceis. Habilitá-lo melhorou o recall em +8.38% para pessoas, +7% para grupos e +5.33% para eventos. O ótimo foi de no máximo dois negativos difíceis por positivo, e usar mais do que isso passa a piorar o modelo. A limitação é que a chance de haver um negativo difícil entre amostras aleatórias pode ser baixa.

1. Os logs de busca alimentam a mineração de exemplos, que no ponto de partida usa cliques como positivos e documentos aleatórios como negativos.

2. O modelo de dois lados é treinado com a perda triplet, que separa o par positivo do negativo por uma margem de distância.

3. No HNM offline, o modelo gera os melhores resultados de cada consulta, e a busca aproximada reduz bastante o tempo de computação dessa geração.

4. Os negativos difíceis são amostrados de uma faixa intermediária do ranking, e não dos exemplos mais difíceis.

5. O modelo é retreinado misturando negativos fáceis e difíceis, aumentando a proporção de fáceis em relação aos difíceis.

6. Depois do treino, o modelo é separado em um codificador de consulta e outro de documento, servidos separadamente.

O HNM offline gera os K melhores resultados de cada consulta, seleciona negativos difíceis por uma estratégia e retreina o modelo com os novos triplets. Um resultado que parece contraintuitivo: modelos treinados só com negativos difíceis não superaram os treinados com negativos aleatórios. Usar os exemplos mais difíceis também não foi a melhor estratégia, e amostrar entre as posições 101-500 do ranking deu o melhor recall do modelo.

Misturar negativos aleatórios e difíceis ajudou: aumentar a proporção de fáceis continuou melhorando o recall, saturando em 100:1. A transferência de aprendizado do modelo difícil para o fácil trouxe ganho adicional. Para gerar os melhores resultados sem uma busca exaustiva, a busca aproximada reduz bastante o tempo de computação, e rodá-la em um shard aleatório bastou para gerar negativos eficazes.

Do lado dos positivos, os autores buscaram resultados novos, que a produção ainda não recuperava com sucesso, minerando alvos potenciais em sessões de busca que falharam. Um modelo treinado só com esses positivos difíceis alcançou recall semelhante ao dos dados de cliques com apenas 4% do volume de dados.

06 · DEPOIS DA RECUPERAÇÃO

Seleção, ranking, ensemble e resultado

Para melhorar eficiência e qualidade, os autores fizeram seleção de consultas e de índice. O EBR não é acionado para certas consultas em que seria fraco e sem valor extra. No índice, escolheram, por exemplo, apenas usuários ativos no mês, eventos recentes e páginas e grupos populares.

Como os estágios de ranking existentes foram desenhados para os cenários de recuperação anteriores, os autores levaram a similaridade de embedding para o ranking como feature, e a similaridade de cosseno foi a que teve melhor desempenho entre as opções testadas. Como o EBR pode ter menor precisão que a correspondência por termos, também montaram um laço de feedback com avaliação humana para retreinar o modelo de relevância e filtrar os resultados irrelevantes.

A análise do HNM levou à ideia do ensemble: são necessários exemplos difíceis para melhorar a precisão, mas os fáceis também importam para representar o espaço de recuperação. Daí a abordagem multiestágio, em que o primeiro modelo foca em recall e o segundo se especializa em diferenciar resultados mais parecidos. Entre as opções testadas para o segundo estágio, o modelo treinado com negativos sem clique teve o melhor recall de KNN, mas sofreu muito mais perda de precisão ao quantizar.

Outra composição em cascata usou o embedding de texto para pré-selecionar candidatos que casam com o texto e depois o modelo unificado para reordenar e devolver os melhores candidatos. Os autores relatam ganhos significativos de métricas em testes A/B online nas verticais em que avaliaram o EBR.