Melhorando a Recuperação e RAG com Ajuste Fino de Modelos de Embedding

Ajuste Fino de Modelos de Embedding para Melhor Recuperação e RAG

Resumo: Ajustar finamente um modelo de embedding em dados do domínio específico pode melhorar significativamente a busca vetorial e a precisão da geração aumentada por recuperação (RAG). Com o Databricks, é fácil ajustar finamente, implantar e avaliar modelos de embedding para otimizar a recuperação para seu caso de uso específico — aproveitando dados sintéticos sem rotulagem manual.

Por que é Importante: Se seu sistema de busca vetorial ou RAG não está recuperando os melhores resultados, ajustar finamente um modelo de embedding é uma maneira simples, porém poderosa, de aumentar o desempenho. Seja lidando com documentos financeiros, bases de conhecimento ou documentação de código interna, o ajuste fino pode fornecer resultados de busca mais relevantes e melhores respostas de LLM subsequentes.

O que Descobrimos: Ajustamos finamente e testamos dois modelos de embedding em três conjuntos de dados corporativos e vimos melhorias significativas nas métricas de recuperação (Recall@10) e no desempenho de RAG subsequente. Isso significa que o ajuste fino pode ser um divisor de águas para a precisão, sem a necessidade de rotulagem manual, utilizando apenas seus dados existentes.

Quer experimentar o ajuste fino de embeddings? Fornecemos uma solução de referência para ajudar você a começar. O Databricks torna a busca vetorial, RAG, reclassificação e ajuste fino de embeddings fáceis. Entre em contato com seu Gerente de Contas Databricks ou Arquiteto de Soluções para mais informações.

Por que Ajustar Finamente Embeddings?

Modelos de embedding potencializam os sistemas modernos de busca vetorial e RAG. Um modelo de embedding transforma texto em vetores, permitindo encontrar conteúdo relevante com base no significado, em vez de apenas em palavras-chave. No entanto, modelos prontos para uso nem sempre são otimizados para seu domínio específico — é aí que entra o ajuste fino.

Ajustar finamente um modelo de embedding em dados específicos do domínio ajuda de várias maneiras:

• Aumenta a precisão da recuperação: Embeddings personalizados melhoram os resultados de busca ao se alinharem com seus dados.
• Melhora o desempenho de RAG: Uma melhor recuperação reduz alucinações e permite respostas de IA generativa mais fundamentadas.
• Otimiza custo e latência: Um modelo ajustado finamente menor pode, às vezes, superar alternativas maiores e mais caras.

Neste post, demonstramos que ajustar finamente um modelo de embedding é uma maneira eficaz de melhorar o desempenho de recuperação e RAG para casos de uso corporativos específicos de tarefas.

Resultados: O Ajuste Fino Funciona

Ajustamos finamente dois modelos de embedding (gte-large-en-v1.5 e e5-mistral-7b-instruct) em dados sintéticos e os avaliamos em três conjuntos de dados de nossa Suite de Benchmarking de Inteligência de Domínio (DIBS) (FinanceBench, ManufactQA e Databricks DocsQA). Em seguida, comparamos-os com o text-embedding-3-large da OpenAI.

Principais Conclusões:

• O ajuste fino melhorou a precisão da recuperação em todos os conjuntos de dados, superando frequentemente os modelos base de forma significativa.
• Embeddings ajustados finamente tiveram desempenho igual ou superior à reclassificação em muitos casos, mostrando que podem ser uma solução autônoma robusta.
• Melhor recuperação levou a um melhor desempenho de RAG no FinanceBench, demonstrando benefícios de ponta a ponta.

Desempenho de Recuperação

Após comparar em três conjuntos de dados, descobrimos que o ajuste fino de embeddings melhora a precisão em dois desses conjuntos. A Figura 1 mostra que, para FinanceBench e ManufactQA, os embeddings ajustados finamente superaram suas versões base, às vezes superando até mesmo o modelo de API da OpenAI (cinza claro). Para o Databricks DocsQA, no entanto, a precisão do text-embedding-3-large da OpenAI supera todos os modelos ajustados finamente. É possível que isso ocorra porque o modelo foi treinado na documentação pública do Databricks. Isso mostra que, embora o ajuste fino possa ser eficaz, ele depende fortemente do conjunto de dados de treinamento e da tarefa de avaliação.

Ajuste Fino vs. Reclassificação

Em seguida, comparamos os resultados acima com a reclassificação baseada em API usando voyageai/rerank-1 (Figura 2). Um reclassificador normalmente pega os k melhores resultados recuperados por um modelo de embedding, reclassifica esses resultados por relevância para a consulta de busca e, em seguida, retorna os k melhores reclassificados (em nosso caso, k=30 seguido por k=10). Isso funciona porque os reclassificadores geralmente são modelos maiores e mais poderosos do que os modelos de embedding, e também modelam a interação entre a consulta e o documento de uma forma mais expressiva.

O que descobrimos foi:

• O ajuste fino do gte-large-en-v1.5 superou a reclassificação no FinanceBench e ManufactQA.
• O text-embedding-3-large da OpenAI se beneficiou da reclassificação, mas as melhorias foram marginais em alguns conjuntos de dados.
• Para o Databricks DocsQA, a reclassificação teve um impacto menor, mas o ajuste fino ainda trouxe melhorias, mostrando a natureza dependente do conjunto de dados desses métodos.

Reclassificadores geralmente incorrem em latência e custo de inferência adicionais por consulta em relação aos modelos de embedding. No entanto, eles podem ser usados com bancos de dados vetoriais existentes e, em alguns casos, podem ser mais econômicos do que re-embedding dados com um modelo de embedding mais recente. A escolha de usar um reclassificador depende do seu domínio e dos seus requisitos de latência/custo.

O Ajuste Fino Ajuda no Desempenho de RAG

Para o FinanceBench, uma melhor recuperação se traduziu diretamente em melhor precisão de RAG quando combinada com o GPT-4o (ver Apêndice). No entanto, em domínios onde a recuperação já era forte, como no Databricks DocsQA, o ajuste fino não agregou muito — destacando que o ajuste fino funciona melhor quando a recuperação é um gargalo claro.

Como Ajustamos Finamente e Avaliamos Modelos de Embedding

Aqui estão alguns dos detalhes mais técnicos de nossa geração de dados sintéticos, ajuste fino e avaliação.

Modelos de Embedding

Ajustamos finamente dois modelos de embedding de código aberto:

• gte-large-en-v1.5 é um modelo de embedding popular baseado no BERT Large (434M parâmetros, 1,75 GB). Escolhemos executar experimentos neste modelo devido ao seu tamanho modesto e licença aberta. Este modelo de embedding também é suportado atualmente na API de Modelos Fundamentais do Databricks.
• e5-mistral-7b-instruct pertence a uma classe mais recente de modelos de embedding construídos sobre LLMs robustos (neste caso, Mistral-7b-instruct-v0.1). Embora o e5-mistral-7b-instruct seja melhor nos benchmarks de embedding padrão, como o MTEB, e seja capaz de lidar com prompts mais longos e sutis, ele é muito maior que o gte-large-en-v1.5 (pois tem 7 bilhões de parâmetros) e é um pouco mais lento e mais caro para servir.

Em seguida, comparamos com o text-embedding-3-large da OpenAI.

Conjuntos de Dados de Avaliação

Avaliamos todos os modelos nos seguintes conjuntos de dados de nossa Suite de Benchmarking de Inteligência de Domínio (DIBS): FinanceBench, ManufactQA e Databricks DocsQA.

Relatamos recall@10 como nossa principal métrica de recuperação; isso mede se o documento correto está entre os 10 principais documentos recuperados.

O padrão ouro para a qualidade do modelo de embedding é o benchmark MTEB, que incorpora tarefas de recuperação como BEIR, bem como muitas outras tarefas não de recuperação. Embora modelos como gte-large-en-v1.5 e e5-mistral-7b-instruct tenham um bom desempenho no MTEB, ficamos curiosos para ver como eles se saíram em nossas tarefas corporativas internas.

Dados de Treinamento

Treinamos modelos separados com dados sintéticos adaptados para cada um dos benchmarks acima:

Para gerar o conjunto de treinamento para cada domínio, pegamos documentos existentes e geramos consultas de amostra baseadas no conteúdo de cada documento usando LLMs como Llama 3 405B. As consultas sintéticas foram então filtradas por qualidade por um LLM-as-a-judge (GPT4o). As consultas filtradas e seus documentos associados foram então usados como pares contrastivos para finetuning. Usamos negativos in-batch para treinamento contrastivo, mas adicionar negativos difíceis poderia melhorar ainda mais o desempenho (ver Apêndice).

Ajuste de Hiperparâmetros

Executamos varreduras em:

• Taxa de aprendizado, tamanho do lote, temperatura softmax
• Contagem de épocas (1-3 épocas testadas)
• Variações de prompt de consulta (por exemplo, "Query:" vs. prompts baseados em instrução)
• Estratégia de pooling (pooling médio vs. pooling do último token)

Todo o finetuning foi feito usando as bibliotecas de código aberto mosaicml/composer, mosaicml/llm-foundry e mosaicml/streaming na plataforma Databricks.

Como Melhorar a Busca Vetorial e RAG no Databricks

O finetuning é apenas uma abordagem para melhorar o desempenho da busca vetorial e RAG; listamos algumas abordagens adicionais abaixo.

Para Melhor Recuperação:

• Use um modelo de embedding melhor: Muitos usuários trabalham inadvertidamente com embeddings desatualizados. Simplesmente trocar por um modelo de desempenho superior pode gerar ganhos imediatos. Verifique o ranking MTEB para os melhores modelos.
• Experimente a busca híbrida: Combine embeddings densos com busca baseada em palavras-chave para melhorar a precisão. O Databricks Vector Search facilita isso com uma solução de um clique.
• Use um reranker: Um reranker pode refinar os resultados reordenando-os com base na relevância. O Databricks oferece isso como um recurso integrado (atualmente em Private Preview). Entre em contato com seu Gerente de Contas para experimentar.

Para Melhor RAG:

• Otimize seus prompts: Pequenos ajustes nos prompts de LLM podem melhorar drasticamente as respostas. O DSPy pode ajudar a automatizar esse processo (veja Construa aplicativos genAI usando DSPy no Databricks).
• Atualize seu LLM: Se a recuperação for forte, mas as respostas forem fracas, considere usar um modelo generativo melhor.
• Finetune um LLM: Se o seu domínio for único e você tiver dados suficientes, o finetuning de um modelo como o Llama 3 pode aumentar ainda mais a qualidade do RAG. Veja Treinamento de Modelos Databricks: Faça o Finetune do Seu LLM no Databricks para Tarefas e Conhecimento Específicos para mais detalhes.

Comece com Finetuning no Databricks

O finetuning de embeddings pode ser uma vitória fácil para melhorar a recuperação e o RAG em seus sistemas de IA. No Databricks, você pode:

• Finetune e sirva modelos de embedding em infraestrutura escalável.
• Use ferramentas integradas para busca vetorial, reranking e RAG.
• Teste rapidamente diferentes modelos para encontrar o que funciona melhor para seu caso de uso.

Pronto para experimentar? Criamos uma solução de referência para facilitar o finetuning — entre em contato com seu Gerente de Contas Databricks ou Arquiteto de Soluções para obter acesso.

Apêndice

Tabela 1: Comparação de gte-large-en-v1.5, e5-mistral-7b-instruct e text-embedding-3-large. Mesmos dados da Figura 1.

Geração de Dados Sintéticos para Treinamento

Para todos os conjuntos de dados, as consultas no conjunto de treinamento não foram as mesmas que as consultas no conjunto de teste. No entanto, no caso do Databricks DocsQA (mas não do FinanceBench ou ManufactQA), os documentos usados para gerar consultas sintéticas foram os mesmos documentos usados no conjunto de avaliação. O foco do nosso estudo é melhorar a recuperação em tarefas e domínios específicos (em oposição a um modelo de embedding generalizável de zero-shot); portanto, vemos isso como uma abordagem válida para certos casos de uso de produção. Para FinanceBench e ManufactQA, os documentos usados para gerar dados sintéticos não tiveram sobreposição com o corpus usado para avaliação.

Existem várias maneiras de selecionar passagens negativas para treinamento contrastivo. Elas podem ser selecionadas aleatoriamente ou podem ser pré-definidas. No primeiro caso, as passagens negativas são selecionadas dentro do lote de treinamento; elas são frequentemente referidas como "in-batch negatives" ou “soft negatives”. No segundo caso, o usuário pré-seleciona exemplos de texto que são semanticamente difíceis, ou seja, potencialmente relacionados à consulta, mas ligeiramente incorretos ou irrelevantes. Este segundo caso é às vezes chamado de "hard negatives". Neste trabalho, usamos simplesmente in-batch negatives; a literatura indica que o uso de hard negatives provavelmente levaria a resultados ainda melhores.

Detalhes de Finetuning

Para todos os experimentos de finetuning, o comprimento máximo da sequência é definido como 2048. Em seguida, avaliamos todos os checkpoints. Para toda a criação de benchmarks, os documentos do corpus foram truncados para 2048 tokens (não divididos em blocos), o que foi uma restrição razoável para nossos conjuntos de dados específicos. Escolhemos os baselines mais fortes em cada benchmark após varrer prompts de consulta e estratégia de pooling. 

Melhorando o Desempenho do RAG

Um sistema RAG consiste em um recuperador e um modelo generativo. O recuperador seleciona um conjunto de documentos relevantes para uma consulta específica e, em seguida, os alimenta para o modelo generativo. Selecionamos os melhores modelos gte-large-en-v1.5 com finetuning e os usamos para o primeiro estágio de recuperação de um sistema RAG simples (seguindo a abordagem geral descrita em Long Context RAG Performance of LLMs e The Long Context RAG Capabilities of OpenAI o1 and Google Gemini). Em particular, recuperamos k=10 documentos, cada um com um comprimento máximo de 512 tokens, e usamos GPT4o como o LLM generativo. A precisão final foi avaliada usando um LLM-as-a-judge (GPT4o).

No FinanceBench, a Figura 3 mostra que o uso de um modelo de embedding com finetuning leva a uma melhoria na precisão do RAG downstream. Além disso, é competitivo com o text-embedding-3-large. Isso é esperado, pois o finetuning do gte levou a uma grande melhoria no Recall@10 em relação ao gte baseline (Figura 1). Este exemplo destaca a eficácia do finetuning do modelo de embedding em domínios e conjuntos de dados específicos.

No conjunto de dados Databricks DocsQA, não encontramos melhorias ao usar o modelo gte com finetuning acima do gte baseline. Isso é um tanto esperado, pois as margens entre os modelos baseline e com finetuning nas Figuras 1 e 2 são pequenas. Curiosamente, mesmo que o text-embedding-3-large tenha um Recall@10 (ligeiramente) maior do que qualquer um dos modelos gte, ele não leva a uma precisão de RAG downstream maior. Como mostrado na Figura 1, todos os modelos de embedding tiveram um Recall@10 relativamente alto no conjunto de dados Databricks DocsQA; isso indica que a recuperação provavelmente não é o gargalo para o RAG, e que o finetuning de um modelo de embedding neste conjunto de dados não é necessariamente a abordagem mais frutífera.

Gostaríamos de agradecer a Quinn Leng e Matei Zaharia pelo feedback neste post do blog.

(Esta publicação no blog foi traduzida utilizando ferramentas baseadas em inteligência artificial) Publicação original