Nova técnica torna sistemas RAG muito melhores na recuperação dos documentos corretos

9 de outubro de 2024 16:17

A recuperação aumentada por geração (RAG) tornou-se um método popular para vincular grandes modelos de linguagem (LLMs) ao conhecimento externo. Os sistemas RAG normalmente usam um modelo de incorporação para codificar documentos em um corpus de conhecimento e selecionar aqueles que são mais relevantes para a consulta do usuário.

No entanto, os métodos de recuperação padrão muitas vezes não levam em conta detalhes contextuais que podem fazer uma grande diferença em conjuntos de dados específicos de aplicação. Em um novo artigo, pesquisadores da Universidade de Cornell introduzem “incorporações de documentos contextuais”, uma técnica que melhora o desempenho dos modelos de incorporação ao torná-los cientes do contexto em que os documentos são recuperados.

As limitações dos bi-encoders

A abordagem mais comum para recuperação de documentos em RAG é usar “bi-encoders”, onde um modelo de incorporação cria uma representação fixa de cada documento e a armazena em um banco de dados vetorial. Durante a inferência, a incorporação da consulta é calculada e comparada às incorporações armazenadas para encontrar os documentos mais relevantes.

Os bi-encoders tornaram-se uma escolha popular para recuperação de documentos em sistemas RAG devido à sua eficiência e escalabilidade. No entanto, os bi-encoders frequentemente têm dificuldades com conjuntos de dados especializados e contextuais porque são treinados em dados genéricos. Na verdade, quando se trata de corporações de conhecimento especializadas, eles podem ficar aquém de métodos estatísticos clássicos, como BM25, em determinadas tarefas.

“Nosso projeto começou com o estudo do BM25, um algoritmo tradicional para recuperação de texto”, disse John (Jack) Morris, estudante de doutorado na Cornell Tech e coautor do artigo, ao VentureBeat. “Fizemos uma pequena análise e vimos que quanto mais fora do domínio o conjunto de dados, mais o BM25 supera as redes neurais.”

BM25 alcança sua flexibilidade calculando o peso de cada palavra no contexto do corpus que está indexando. Por exemplo, se uma palavra aparece em muitos documentos no corpus de conhecimento, seu peso será reduzido, mesmo que seja uma palavra-chave importante em outros contextos. Isso permite que o BM25 se adapte às características específicas de diferentes conjuntos de dados.

“Modelos de recuperação densos baseados em redes neurais tradicionais não podem fazer isso porque definem pesos uma vez, com base nos dados de treinamento”, disse Morris. “Tentamos projetar uma abordagem que pudesse corrigir isso.”

Incorporações de documentos contextuais

Os pesquisadores de Cornell propõem dois métodos complementares para melhorar o desempenho dos bi-encoders, adicionando a noção de contexto às incorporações de documentos.

“Se você pensar na recuperação como uma ‘competição’ entre documentos para ver qual é o mais relevante para uma determinada consulta de pesquisa, usamos ‘contexto’ para informar o codificador sobre os outros documentos que estarão na competição”, disse Morris.

O primeiro método modifica o processo de treinamento do modelo de incorporação. Os pesquisadores usam uma técnica que agrupa documentos similares antes de treinar o modelo de incorporação. Em seguida, usam o aprendizado contrastivo para treinar o codificador a distinguir documentos dentro de cada cluster.

O aprendizado contrastivo é uma técnica não supervisionada na qual o modelo é treinado para diferenciar entre exemplos positivos e negativos. Ao ser forçado a distinguir entre documentos similares, o modelo se torna mais sensível a diferenças sutis que são importantes em contextos específicos.

O segundo método modifica a arquitetura do bi-encoder. Os pesquisadores aumentam o codificador com um mecanismo que lhe dá acesso ao corpus durante o processo de incorporação. Isso permite que o codificador leve em conta o contexto do documento ao gerar sua incorporação.

A arquitetura aumentada funciona em duas etapas. Primeiro, calcula uma incorporação compartilhada para o cluster ao qual o documento pertence. Em seguida, combina essa incorporação compartilhada com as características únicas do documento para criar uma incorporação contextualizada.

Essa abordagem permite que o modelo capture tanto o contexto geral do cluster do documento quanto os detalhes específicos que o tornam único. A saída ainda é uma incorporação do mesmo tamanho que um bi-encoder normal, portanto, não requer alterações no processo de recuperação.

O impacto das incorporações de documentos contextuais

Os pesquisadores avaliaram seu método em vários benchmarks e descobriram que ele superou consistentemente os bi-encoders padrão de tamanhos semelhantes, especialmente em configurações fora do domínio, onde os conjuntos de dados de treinamento e teste são significativamente diferentes.

“Nosso modelo deve ser útil para qualquer domínio que seja materialmente diferente dos dados de treinamento e pode ser pensado como uma substituição barata para o ajuste fino de modelos de incorporação específicos de domínio”, disse Morris.

As incorporações contextuais podem ser usadas para melhorar o desempenho dos sistemas RAG em diferentes domínios. Por exemplo, se todos os seus documentos compartilham uma estrutura ou contexto, um modelo de incorporação normal desperdiçaria espaço em suas incorporações ao armazenar essa estrutura ou informação redundante.

“As incorporações contextuais, por outro lado, podem ver a partir do contexto ao redor que essa informação compartilhada não é útil e descartá-la antes de decidir exatamente o que armazenar na incorporação”, disse Morris.

Os pesquisadores lançaram uma versão pequena de seu modelo de incorporação de documentos contextuais (cde-small-v1). Ele pode ser usado como uma substituição direta para ferramentas populares de código aberto, como HuggingFace e SentenceTransformers, para criar incorporações personalizadas para diferentes aplicações.

Morris afirma que as incorporações contextuais não estão limitadas a modelos baseados em texto e podem ser estendidas a outras modalidades, como arquiteturas de texto para imagem. Há também espaço para melhorá-las com algoritmos de clustering mais avançados e avaliar a eficácia da técnica em escalas maiores.

Fonte