CoRAG: modelagem dinâmica de RAG encadeada usando MCTS (Monte Carlo Trees)

Resumo das principais contribuições da CORAG

CORAG (Restrição de custos) Recuperação Optimisation for Retrieval-Augmented Generation) é um sistema inovador de Retrieval Augmented Generation (RAG) projetado para lidar com os problemas existentes. RAG Principais desafios da metodologia. A seguir estão as principais contribuições do CORAG:

1. Consideração abrangente das associações entre blocos::
   • ponto de inovaçãoCORAG é a primeira estrutura a introduzir a otimização da ordem de combinação de blocos na tarefa RAG. Diferentemente das abordagens tradicionais, o CORAG não trata mais cada bloco de texto de forma independente ou considera os blocos apenas no nível de agrupamento, mas usa o Pesquisa em árvore de Monte Carlo (MCTS) para buscar sequencialmente a ordem ideal das combinações de blocos. Essa abordagem captura totalmente as associações complexas entre os blocos, evita informações redundantes e garante que as combinações de blocos selecionadas respondam totalmente à consulta do usuário.
   • de pontaA CORAG é capaz de gerar respostas mais precisas e relevantes, melhorando a qualidade da geração.
2. Solução do problema de não monotonicidade da utilidade do bloco::
   • ponto de inovaçãoCORAG irá restrição orçamentária integrado ao processo de otimização do portfólio de blocos, em vez de considerar o esgotamento do orçamento como uma condição de encerramento. Essa abordagem considera a utilidade do bloco de não monotonicidadeou seja, adicionar mais blocos nem sempre melhora a qualidade do resultado final.
   • de pontaO CORAG é capaz de evitar a inclusão excessiva de blocos irrelevantes ou redundantes, melhorando assim a precisão e a relevância das respostas geradas.
3. Adaptação a diferentes campos de pesquisa::
   • ponto de inovaçãoCORAG introduziu um Configuração do agenteO agente usa o Aprendizagem comparativa para adaptar dinamicamente a configuração do MCTS a diferentes domínios de consulta. O agente recomenda o melhor Modelo de reorganização responder cantando Configuração do MCTS.
   • de pontaEssa abordagem oferece ao CORAG a flexibilidade para lidar com uma ampla variedade de tipos de consultas, desde simples consultas por palavras-chave até problemas complexos de raciocínio, garantindo respostas de alta qualidade em diversos cenários de aplicativos.
4. Estratégias de pesquisa eficientes e dimensionáveis::
   • ponto de inovaçãoO CORAG emprega o MCTS para a pesquisa de nós com o técnica de expansão paralela Acelera o processo de pesquisa. Essa abordagem reduz o espaço de pesquisa exponencial para linear e equilibra efetivamente o explorações responder cantando Utilizar.
   • de pontaO CORAG é capaz de lidar com conjuntos de dados em grande escala, mantendo uma recuperação eficiente e alcançando um equilíbrio entre o custo computacional e a qualidade da recuperação.
5. Ganhos significativos de desempenho::
   • verificação experimentalOs resultados experimentais mostram que o CORAG supera os métodos de base existentes em vários conjuntos de dados de referência, melhorando as pontuações ROUGE em cerca de 30%. A CORAG também se destaca em termos de eficiência, fornecendo respostas de alta qualidade dentro de restrições de custo rigorosas.

Exemplo de fluxo de trabalho CORAG

Para ajudar o leitor a entender rapidamente como o CORAG funciona, apresentamos a seguir um exemplo completo de fluxo de trabalho que mostra como o CORAG processa uma consulta do usuário e gera a resposta final. Cada etapa contém entradas e saídas para representar o vai-e-vem do fluxo de trabalho.

Etapa 1: Entrada de consulta do usuário

• importaçãoUm usuário envia uma consulta em linguagem natural ao sistema CORAG, por exemplo:

  "Explique o processo de fotossíntese e liste os fatores que afetam sua eficiência."
  

• exportaçõesA consulta é passada para o Módulo de incorporação de consultas Processamento.

Etapa 2: Geração de incorporação de consulta

• importaçãoTexto de consulta do usuário.
• lidar comUso de treinamento prévio Modelos de incorporação(por exemplo, BGE-M3) converte o texto da consulta em uma representação vetorial.
• exportaçõesVetor de incorporação: Consulta o vetor de incorporação (por exemplo, um vetor de 1024 dimensões).

  Incorporação de consulta: [0,123, -0,456, 0,789, ... , -0.012]
  

Etapa 3: Configurar a previsão do agente

• importação: vetor de incorporação da consulta [0.123, -0.456, 0.789, ..., -0.012].
• lidar com::
  1. extração de recursosO vetor de incorporação é alimentado no Configuração do agente (usado em uma expressão nominal) rede de codificação A extração de recursos é realizada no
  2. Previsão de rearranjoPrevisão de saída da rede codificada da rede ideal Modelo de reorganizaçãoPor exemplo bge-reranker-large.
  3. Previsão de configuração do MCTSAo mesmo tempo, a rede de codificação prevê o MCTS ideal Parâmetros de configuraçãocomo o número de iterações, o fator de custo e o fator de exploração.
• exportações::
  • Modelo de reorganização ideal::bge-reranker-large
  • Parâmetros de configuração do MCTS::
    • Número de iterações:: 15
    • fator de custo: 0.2
    • Fator de exploração: 2.5

  Reranker ideal: bge-reranker-large
  Configuração do MCTS.
  - Iterações: 15
  - Coeficiente de custo: 0,2
  - Coeficiente de exploração: 2,5
  

Etapa 4: Recuperar blocos de texto em potencial

• importação::
  • Vetor de incorporação de consulta [0.123, -0.456, 0.789, ..., -0.012].
  • banco de dados de vetores(por exemplo, contendo blocos de texto pré-segmentados e seus vetores de incorporação).
• lidar comUse o vetor de incorporação da consulta para Pesquisa de similaridaderecupera os possíveis blocos de texto mais relevantes do banco de dados de vetores. A seguir, exemplos dos cinco primeiros blocos de texto recuperados:Bloco de texto 1::

  A fotossíntese é o processo pelo qual as plantas, as algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio.
  

  Bloco de texto 2::

  A fotossíntese ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em dois estágios: a reação da luz e a reação do carbono.
  

  Bloco de texto 3::

  Os fatores que afetam a eficiência da fotossíntese incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água.
  

  Bloco de texto 4::

  O efeito da intensidade da luz na fotossíntese está positivamente correlacionado, mas o excesso de luz leva ao fenômeno da fotoinibição, que reduz a eficiência da fotossíntese.
  

  Bloco de texto 5::

  O dióxido de carbono é uma das matérias-primas da fotossíntese e sua concentração afeta diretamente a taxa de fotossíntese.
  

• exportaçõesLista de blocos de texto em potencial e seus vetores de incorporação.

  Partes recuperadas.
  - Fragmento 1: "A fotossíntese é o processo pelo qual as plantas, as algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio."
  - Fragmento 2: "A fotossíntese ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em dois estágios: a reação da luz e a reação do carbono."
  - Fragmento 3: "Os fatores que afetam a eficiência da fotossíntese incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água."
  - Fragmento 4: "O efeito da intensidade da luz na fotossíntese está positivamente correlacionado, mas o excesso de luz leva ao fenômeno da fotoinibição, que reduz a eficiência da fotossíntese."
  - Fragmento 5: "O dióxido de carbono é uma das matérias-primas para a fotossíntese e sua concentração afeta diretamente a taxa de fotossíntese."
  

Etapa 5: Pesquisa em árvore do MCTS

• importação::
  • Lista de possíveis blocos de texto::
    • Parte 1: ...
    • Parte 2: ...
    • Parte 3: ...
    • Parte 4: ...
    • Parte 5: ...
  • Parâmetros de configuração do MCTS::
    • Número de iterações:: 15
    • fator de custo: 0.2
    • Fator de exploração: 2.5
  • Modelo de reorganização::bge-reranker-large
• lidar com::
  1. Inicialização do nó raizNó raiz: O nó raiz representa um estado vazio em que nenhum bloco de texto está selecionado.
  2. Expansão iterativa::
     • opção: Uso Algoritmo UCB Selecione o nó com a maior utilidade atual. Por exemplo, o pedaço 3 é escolhido para a primeira iteração porque tem a maior relevância para a consulta.
     • extensõesGerar todos os nós filhos possíveis (novas combinações de pedaços de texto). Por exemplo, os filhos do pedaço 3 podem ser o pedaço 1, o pedaço 2, o pedaço 4 e o pedaço 5.
     • avaliação: Uso Modelo de reorganização bge-reranker-large Avalie a utilidade de todas as novas combinações em paralelo. Por exemplo, avalie a utilidade de combinações como Pedaço 3 + Pedaço 1, Pedaço 3 + Pedaço 2 e assim por diante.
     • atualizaçãoAtualizar o utilitário do nó e as contagens de acesso e propagar as atualizações para cima.
  3. iterarRepetir as etapas de seleção, expansão, avaliação e atualização até que o número máximo de iterações (15) seja atingido.
  4. Condições de rescisãoNúmero máximo de iterações atingido: Número máximo de iterações atingido.
• exportaçõesOrdem ideal de combinação de blocos de texto: ordem ideal de combinação de blocos de texto.

  Combinação ideal de pedaços.
  - Fragmento 3: "Os fatores que afetam a eficiência da fotossíntese incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água."
  - Fragmento 1: "A fotossíntese é o processo pelo qual plantas, algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio."
  - Fragmento 2: "A fotossíntese ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em dois estágios: a reação da luz e a reação do carbono."
  

  tomar nota deDurante o processo MCTS, o CORAG ajusta dinamicamente a ordem de combinação de acordo com a utilidade e, por fim, seleciona a ordem de combinação ideal.

Etapa 6: Gerar resposta final

• importação::
  • Ordem ideal de combinação de blocos de texto::
    • Parte 3: "Os fatores que afetam a eficiência fotossintética incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água."
    • Parte 1: "A fotossíntese é o processo pelo qual plantas, algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio."
    • Parte 2: "A fotossíntese ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em duas fases: a reação da luz e a reação do carbono."
  • consulta do usuárioDescrição: "Explique o processo de fotossíntese e liste os fatores que afetam sua eficiência."
• lidar com::
  1. Dicas de construçãoCombinando a ordem ideal de combinação de blocos de texto com consultas do usuário para construir o LLM do prompt de entrada. Exemplo:

     Consulta do usuário: Explique o processo de fotossíntese e liste os fatores que afetam sua eficiência.
     Informações relacionadas:
     1. Os fatores que afetam a eficiência da fotossíntese incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água.
     2. a fotossíntese é o processo pelo qual as plantas, as algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio.
     3. a fotossíntese ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em dois estágios: a reação da luz e a reação do carbono.
     

  2. Geração de uma resposta: Uso LLM(por exemplo, Llama3) Gere uma resposta final para a solicitação.
• exportaçõesResposta final em linguagem natural: A resposta final em linguagem natural.

  Resposta final.
  "A fotossíntese é o processo pelo qual as plantas, as algas e algumas bactérias usam a energia da luz para converter dióxido de carbono e água em glicose e oxigênio. Ela ocorre principalmente nos cloroplastos das folhas das plantas e é dividida em duas fases: a reação da luz e a reação do carbono. Os fatores que afetam a eficiência da fotossíntese incluem a intensidade da luz, a concentração de dióxido de carbono, a temperatura e a disponibilidade de água."
  

Etapa 7: Saída de resposta

• importaçãoResposta final em linguagem natural: A resposta final em linguagem natural.
• lidar comApresentação da resposta ao usuário: Apresente a resposta ao usuário.
• exportaçõesInterface do usuário: A interface do usuário exibe o resultado da resposta.

resumos

Com o exemplo de fluxo de trabalho acima, o princípio de funcionamento do CORAG pode ser resumido nas etapas a seguir:

1. Geração de incorporação de consultasConversão: converte uma consulta do usuário em uma representação vetorial.
2. Configuração da previsão do agenteMCTS: previsão do modelo de rearranjo ideal e dos parâmetros de configuração do MCTS.
3. Recuperar possíveis blocos de textoRecupera o bloco de texto mais relevante do banco de dados de vetores.
4. Pesquisa em árvore do MCTSUse o MCTS para procurar a ordem ideal de combinação de blocos de texto.
5. Gerar resposta finalUse o LLM para gerar a resposta final com base na combinação ideal de blocos de texto.
6. saída de respostaApresentação da resposta ao usuário: Apresente a resposta ao usuário.

Essas etapas demonstram como o CORAG combina efetivamente as três fases de recuperação, aprimoramento e geração para fornecer respostas em linguagem natural de alta qualidade. Com exemplos detalhados de dados, os leitores podem obter uma compreensão mais clara do processo de processamento de dados do CORAG e de como ele funciona.

texto original:: https://arxiv.org/pdf/2411.00744

legenda:: CORAG: um sistema de otimização de recuperação com restrições de custo para geração de aprimoramento de recuperação

autorZiting Wang, Haitao Yuan, Wei Dong (Universidade Tecnológica de Nanyang), Gao Cong (Universidade Tecnológica de Nanyang), Feifei Li (Grupo Alibaba)

resumos

Os modelos de linguagem grande (LLMs) demonstraram excelentes recursos em tarefas generativas, mas geralmente é difícil acessar informações atualizadas, o que pode levar à desilusão. A geração aprimorada por recuperação (RAG) aborda esse problema integrando o conhecimento de bancos de dados externos para obter respostas mais precisas e relevantes. Devido à limitação da janela de contexto dos LLMs, não é prático inserir diretamente todo o contexto do banco de dados externo no modelo. Em vez disso, apenas as informações mais relevantes, ou seja, os "pedaços" (chunks), são recuperados seletivamente. No entanto, a pesquisa atual sobre RAG enfrenta três desafios principais. Primeiro, as soluções existentes geralmente selecionam cada bloco de forma independente, ignorando as possíveis associações entre eles. Segundo, na prática, a utilidade dos blocos é "não monotônica", o que significa que adicionar mais blocos pode reduzir a utilidade geral. As abordagens tradicionais enfatizam a maximização do número de blocos incluídos, o que pode, inadvertidamente, prejudicar o desempenho. Terceiro, cada tipo de consulta de usuário tem características exclusivas que exigem um processamento personalizado, que as soluções atuais não levam em conta adequadamente.

Para superar esses desafios, propomos o CORAG, um sistema de otimização de recuperação com restrições de custo, para a geração de aprimoramento de recuperação. Adotamos uma estrutura de estratégia baseada em Monte Carlo Tree Search (MCTS) para encontrar as combinações ideais de blocos sequencialmente, considerando, assim, de forma abrangente, as associações entre os blocos. Além disso, em vez de considerar a exaustão do orçamento como condição de término, integramos as restrições orçamentárias ao processo de otimização das combinações de blocos, o que resolve com eficácia o problema da não monotonicidade da utilidade dos blocos. Além disso, ao projetar agentes de configuração, nosso sistema prevê a configuração ideal para cada tipo de consulta, melhorando a adaptabilidade e a eficiência. Os resultados experimentais mostram que nossa estrutura tem até 301 TP3T de melhoria de desempenho em relação ao modelo de linha de base, destacando a eficácia, o dimensionamento e a aplicabilidade da estrutura para aplicativos de contexto longo.

Formato de referência do PVLDB.

Ziting Wang, Haitao Yuan, Wei Dong, Gao Cong e Feifei Li. CORAG: um sistema de otimização de recuperação com restrição de custo para geração aumentada de recuperação . pvldb, 14(1): xxx-xxx, 2020.
doi:XX.XX/XXX.XX

Figura 1: Exemplo de ordem de combinação de blocos de dados.

1 Introdução

Embora os LLMs tenham demonstrado excelentes capacidades em tarefas generativas, eles frequentemente encontram dificuldades em adquirir informações atualizadas, o que pode levar a alucinações [10, 38]. Para enfrentar esses desafios, o RAG surgiu como uma solução importante. Ao integrar fontes de dados externas ao LLM, o RAG pode fornecer informações mais precisas, relevantes e atualizadas. Atualmente, a RAG tem sido amplamente estudada no contexto de LLMs, especialmente em tarefas que exigem conhecimento externo atualizado, como tarefas de perguntas e respostas [2, 22, 29], recuperação de informações médicas [1, 32] e análise de séries temporais [12, 26, 40]. Em geral, as fontes de dados externas são tão grandes que não é prático alimentá-las diretamente no LLM. Para resolver esse problema, os dados geralmente são particionados em partes não sobrepostas e armazenados em um banco de dados vetorial e, em seguida, o usuário consulta as partes mais úteis para criar dicas para o LLM. Portanto, projetar estruturas e algoritmos de pesquisa eficientes e precisos para encontrar os blocos mais relevantes tornou-se um tópico de pesquisa importante e foi amplamente estudado nas comunidades de banco de dados [39, 48] e de aprendizado de máquina [2, 35, 43].

No entanto, há três desafios principais com as abordagens existentes.

Desafio 1: ligações entre blocos. Atualmente, há duas abordagens principais para identificar os blocos mais relevantes. A primeira abordagem formula o problema como uma tarefa de AKNN (Approximate k Nearest Neighbours) [41, 45], em que cada bloco recebe uma pontuação e são selecionados os k blocos mais próximos classificados pela pontuação. A segunda abordagem agrupa os blocos e retorna todos os blocos no cluster mais relevante em resposta a uma consulta [22, 29]. Entretanto, ambos os métodos ignoram as possíveis associações entre os blocos: o primeiro método ignora completamente as associações, enquanto o segundo método considera apenas superficialmente todos os blocos em cada cluster, tratando-os como igualmente relevantes. Como resultado, esses métodos introduzem uma redundância substancial nos blocos selecionados quando vários blocos transmitem informações semelhantes ou redundantes.

Por exemplo, conforme mostrado na Fig. 1, ao consultar a altura e o histórico da Torre Eiffel, se cada bloco for tratado de forma independente, o método guloso seleciona os blocos χ3 e χ1 porque eles têm as duas primeiras pontuações mais altas. Entretanto, esses dois blocos fornecem apenas informações sobre o histórico, o que não é suficiente para responder totalmente à consulta. Para responder melhor à consulta, é necessário incluir blocos que contenham o nome do construtor, por exemplo, χ4. Por outro lado, a abordagem de agrupamento retornará todos os blocos χ1, χ2, χ3 e χ4, o que gera redundância. A solução ideal escolherá χ3 e χ4, pois eles fornecem as informações necessárias sem redundância. Além disso, estudos [11, 19, 42] mostraram que a ordem dos blocos afeta o desempenho do LLM, um fator que também é ignorado pelos métodos existentes. No caso da Torre Eiffel, por exemplo, ao escolher os blocos χ3 e χ4, colocar o χ4 na primeira posição gera uma pontuação maior em comparação com a ordem oposta. Entretanto, determinar a ordem ideal das combinações de blocos é uma tarefa desafiadora, pois ambas exigem que o espaço de busca cresça exponencialmente com o número de blocos disponíveis. Neste documento, demonstramos ainda que esse problema é NP-difícil (consulte a Seção 2.1).

Desafio 2: não monotonicidade da utilidade. A solução atual baseia-se na suposição de que a inclusão de mais blocos sempre produz um resultado final melhor. Especificamente, na abordagem baseada em AKNN, exatamente k blocos são selecionados de forma determinística a cada vez. Na abordagem baseada em clustering, um limite de distância entre clusters e consultas é definido e todos os clusters dentro desse limite são retornados. Ambas retornam o maior número possível de blocos. Entretanto, na prática, a utilidade dos blocos não é monotônica. Mais especificamente, o excesso de blocos dilui as principais informações ao adicionar conteúdo relacionado à borda, produzindo ruído que reduz a clareza. Além disso, conflitos ou diferenças sutis entre blocos podem confundir o modelo e reduzir a qualidade da resposta. Por exemplo, conforme mostrado na Figura 1, adicionar o bloco χ1 reduz a utilidade quando χ3 e χ4 são selecionados, destacando o fato de que as pontuações de utilidade normalmente não são monotônicas na prática.

Desafio 3: diversidade de consultas. Há diferentes tipos de consultas de usuários, cada uma exigindo uma estratégia de classificação diferente com base em suas características exclusivas [47]. Nos sistemas RAG atuais, a pontuação de utilidade de um bloco geralmente é determinada pelo modelo de reclassificação atribuído. Até o momento, existem vários modelos de reclassificação, mas observamos que seu desempenho varia muito entre os tipos de consulta, e nenhum modelo de reclassificação fixo único supera consistentemente os outros em todas as variações de consulta (consulte os experimentos na Seção 6.3.4 para obter mais detalhes). As abordagens atuais [20, 46] geralmente dependem de modelos de reclassificação estáticos para classificação de blocos e não têm flexibilidade para se adaptar a diferentes contextos de consulta.

Declaração do problema: Existe um sistema RAG que leva em conta as associações entre blocos e a não monotonicidade dos utilitários, além de ser capaz de acomodar todos os tipos de consultas?

1.1 Nossa contribuição

Neste artigo, respondemos a essa pergunta de forma afirmativa, propondo uma nova estrutura de árvore de políticas baseada em MCTS para otimizar a recuperação de blocos em sistemas RAG. Em geral, nossa contribuição pode ser resumida da seguinte forma:

- Apresentamos a primeira estrutura RAG que considera a ordem das combinações de blocos em uma tarefa RAG. Em vez de considerar cada bloco de forma independente ou em um nível de agrupamento, usamos o MCTS para ajudar a pesquisar sequencialmente a ordem ideal de combinação de blocos. A ideia de alto nível é a seguinte: primeiro, inicializamos o nó raiz. Em seguida, durante as iterações, expandimos a árvore selecionando o nó com a maior utilidade e calculando a utilidade de seus nós de expansão. Após cada expansão, atualizamos a utilidade em toda a árvore de políticas. Nesse processo, a decisão de cada iteração depende dos blocos que foram selecionados, o que nos permite considerar totalmente as associações entre os blocos. Além disso, o MCTS reduz o espaço de pesquisa exponencial para linear, e aplicamos técnicas de expansão paralela para aumentar ainda mais a eficiência computacional. Com esse projeto, abordamos o Desafio 1.
- Diferentemente das estruturas RAG anteriores, que consideram a exaustão do orçamento como uma condição de encerramento, propomos uma nova formulação na qual as restrições orçamentárias são integradas ao processo de otimização das combinações de blocos para levar em conta a não monotonicidade da utilidade dos blocos, abordando, assim, o Desafio 2. Além disso, ao priorizar blocos altamente correlacionados e de baixo custo e levar em conta o comprimento dos tokens, reduzimos ainda mais o custo computacional.
- Propomos um agente baseado em aprendizado comparativo que adapta dinamicamente a configuração do MCTS à consulta, adaptando o modelo de reclassificação e a configuração ao domínio específico da consulta. Essa abordagem oferece flexibilidade e robustez para consultas dinâmicas e específicas do domínio, adaptando-se ao Desafio 3.
- Além disso, realizamos experimentos abrangentes comparando nossa estrutura com vários métodos de última geração. Os resultados validam a eficácia, a eficiência e a escalabilidade de nossa abordagem e melhoram o desempenho em 30% em comparação com a linha de base.

2 Conhecimento preparatório

Nesta seção, primeiro apresentamos as definições de alguns conceitos-chave, como ordem de blocos e combinação de blocos, na Seção 2.1. Em seguida, apresentamos uma prova NP-hard do problema de otimização da ordem dos blocos. Por fim, discutimos trabalhos relacionados na Seção 2.3.

2.1 Conceitos-chave

RAG & Chunks. O RAG é um método eficiente para melhorar o desempenho de modelos generativos, recuperando o contexto relevante de um corpus externo. Nessa abordagem, o corpus é primeiro particionado em unidades menores e mais gerenciáveis chamadas blocos, que são armazenadas em um banco de dados vetorial. Assim, podemos dar uma definição formal de um bloco da seguinte forma:

Definição 2.1 (Bloco). Deixe C representar um corpus de documentos e um bloco χ ser definido como um bloco contínuo de texto extraído de C. Formalmente, um bloco χ consiste em uma sequência de tokens (t1, t2, ... , tn), em que cada ti é um token em C e o tamanho n é definido pelo usuário.

Em um sistema RAG, cada bloco é incorporado em uma representação vetorial usando um modelo de incorporação que captura o significado semântico do bloco e permite a recuperação de conteúdo contextualmente semelhante. Quando uma nova consulta é recebida, o banco de dados vetorial executa uma pesquisa de similaridade para identificar os blocos semanticamente mais relevantes para a consulta. Esses blocos recuperados são então passados para um gerador (por exemplo, um modelo de linguagem grande) para gerar uma resposta final com base no conteúdo recuperado. Especificamente, quanto mais tokens um bloco contiver, maior será o custo incorrido pelo gerador. Assim, definimos o custo de um bloco como cost(χ) = |χ|, que é igual ao número de tokens no bloco.

Ordem de combinação de blocos. Em um sistema RAG, os resultados de uma pesquisa de banco de dados vetorial podem incluir vários blocos. Entretanto, devido às restrições de entrada do modelo generativo, não é prático usar todos esses blocos. Portanto, é necessário selecionar o subconjunto ideal de blocos, chamado de combinação de blocos, para atender a um determinado orçamento de custo. Além disso, a ordem dos blocos dentro da combinação tem um impacto significativo no desempenho do modelo generativo. O objetivo é identificar a ordem das combinações de blocos com a ordem ideal, formalmente definida da seguinte forma:

Definição 2.2 (Seleção ideal da ordem de combinação de blocos). Seja {χ1, χ2, ... , χk} seja um conjunto de blocos potenciais, ℛ seja um orçamento de custo, Φ = ⟨χφ1, ... χφm⟩ representa uma ordem de possíveis combinações de blocos, em que cada χφi é um bloco e o índice φi denota sua posição em Φ. Seja U(Φ) a pontuação de utilidade atribuída pelo modelo de reclassificação, que pode ser arbitrária ou composta. Nosso objetivo é encontrar a ordem de combinação dos blocos que maximiza as pontuações de utilidade sob a restrição de custo de alimentá-los em LLMs para gerar a resposta final, ou seja, procurar o

2.2 Provas de dificuldade NP

Para mostrar que a seleção sequencial da combinação de blocos é NP difícil, generalizamos o problema de hipergrupo máximo ponderado (MWHP) para ela. Como o MWHP é NP-difícil, mostramos que qualquer instância do MWHP pode ser convertida em uma instância de otimização de combinação de blocos em tempo polinomial.

2.2.1 Definição do problema da MWHP

Dado um hipergrafo ℋ = (V, E, w1, w2), em que V é o conjunto de vértices e E é o conjunto de hiper-redes, e cada hiper-redes contém um subconjunto de V. w1: v → ℝ e w2: e → ℝ são as funções de peso, que atribuem um peso a cada vértice e hiper-redes, respectivamente. Dado um subconjunto de vértices V' ⊆ V, dizemos que uma hipercategoria e pertence a V', ou seja, e ∈ V', se V' cobrir todos os vértices de e. O objetivo é encontrar k vértices e maximizar a soma dos pesos desses vértices e das hipercordas que eles cobrem:

2.2.2 Processo de normalização

Agora, construímos uma instância correspondente do problema de otimização de combinação de blocos a partir de uma determinada instância do MWHP. Para cada nó v ∈ V, criamos um bloco correspondente Xv. Definimos seu custo cost(Xv) ≡ 1. Uma ordem de combinação de blocos Φ corresponde a um subconjunto de vértices de V denotado por V(Φ) ⊆ V. Definimos sua utilidade como

Por fim, definimos B = k e buscamos

Usando Φdenota a solução de (4), então fica claro que V(Φ) é uma solução de (2) e a redução pode ser feita em tempo O(|V|-|E|).

Observe que essa indução pressupõe que, em nosso problema de otimização da combinação de blocos, permitimos que a reclassificação seja arbitrária, o que significa que as pontuações de utilidade também podem ser atribuídas arbitrariamente. A complexidade de encontrar a ordem ideal de combinação de blocos pode ser significativamente reduzida se forem feitas algumas suposições sobre o reclassificador. Por exemplo, se a reclassificação não levar em conta as associações e apenas somar linearmente as pontuações de utilidade de cada bloco, cada bloco poderá ser avaliado independentemente. No entanto, neste documento, lidamos com o caso mais geral e não fazemos nenhuma suposição sobre o modelo de reclassificação.

2.3 Trabalhos relacionados

2.3.1 Geração de aprimoramento de recuperação

O RAG [14, 20] é amplamente usado para lidar com tarefas de PNL com uso intensivo de conhecimento. Em um processo típico de RAG, um recuperador baseado em incorporação densa procura informações relevantes em bancos de dados externos, que são então usadas pelo LLM durante o processo de geração. Para aprimorar esse processo, vários estudos [5, 18, 22, 35] se concentraram em adaptar os recuperadores para atender melhor às necessidades de geração dos LLMs, desenvolvendo métodos de recuperação em várias etapas e filtrando informações irrelevantes. Embora existam muitos recuperadores de última geração [8, 9, 15, 16, 27, 34], a otimização de recuperadores e LLMs juntos em um processo de ponta a ponta [25, 31] é mais promissora. Por exemplo, a pesquisa [30] se concentra no treinamento conjunto do recuperador e do LLM simultaneamente ou em etapas. No entanto, isso exige a perda de agentes para otimização e complica o processo de treinamento, especialmente quando os bancos de dados incorporados precisam ser reindexados com frequência, o que acarreta altos custos computacionais. Portanto, abordagens como [5] decompõem consultas complexas de várias etapas em subintensidades menores para melhorar a abrangência da resposta sem reindexação frequente. No entanto, essas abordagens geralmente ignoram a função crucial da ordem de combinação de blocos, que pode afetar significativamente a qualidade geral da resposta do LLM. Até onde sabemos, este artigo é a primeira abordagem a considerar a ordem de combinação de blocos em uma tarefa RAG.

2.3.2 Classificação RAG

Os métodos de reclassificação são essenciais para melhorar o desempenho da recuperação no processo RAG [43, 44, 51]

Figura 2: Visão geral da arquitetura do sistema CORAG

Os métodos tradicionais de reclassificação [33, 50] geralmente dependem de modelos de linguagem de tamanho médio, como o BERT ou o T5, para classificar os contextos recuperados. No entanto, esses modelos geralmente têm dificuldade para capturar as relações semânticas entre consultas e contextos, especialmente em configurações de amostragem zero ou pequena. Portanto, pesquisas recentes [43] enfatizaram o potencial dos LLMs ajustados por comando para melhorar a reclassificação do contexto, mesmo na presença de informações ruidosas ou irrelevantes. Apesar desses avanços, os recursos de reclassificação do LLM em sistemas RAG permanecem subutilizados. Em particular, foi demonstrado que a classificação de blocos afeta o desempenho do LLM [19], enfatizando a necessidade de considerar a ordem das combinações de blocos nas tarefas de RAG. Entretanto, os modelos existentes não são adequados para situações em que uma sequência ou combinação específica de blocos é necessária para a recuperação ideal, em vez de blocos isolados. Portanto, pesquisas futuras precisam fazer melhor uso do LLM para sequenciar blocos de forma mais eficiente em resposta a consultas dentro da estrutura do RAG.

2.3.3 Aprendizado por reforço para modelos de idiomas grandes

Recentemente, o aprendizado por reforço (RL) tem sido cada vez mais usado em várias tarefas de gerenciamento de dados e RAG. As técnicas de RL podem permitir que grandes modelos de linguagem aprimorem seus recursos de geração explorando fontes de conhecimento externas, como mecanismos de pesquisa [13, 23]. Em particular, o feedback humano [4, 36, 37] pode ser integrado para ajudar os modelos a gerar respostas mais precisas e contextualmente relevantes por meio da estrutura de RL. Além disso, várias abordagens de otimização de consultas [17, 21, 49] aprimoram ainda mais o processo de recuperação, permitindo que o desempenho do modelo informe o ajuste da consulta e, por fim, melhore os resultados das tarefas posteriores. Neste trabalho, aplicamos uma técnica de RL leve, o MCTS, para otimizar o processo de pesquisa sequencial de combinação de blocos em sistemas RAG. Também apresentamos um agente de configuração para orientar o processo de pesquisa do MCTS. Até onde sabemos, esta é a primeira abordagem para esse problema específico.

3 Visão geral do sistema

Conforme mencionado anteriormente, as estruturas RAG existentes enfrentam três desafios principais: como levar em conta adequadamente as associações entre blocos e a não monotonicidade da utilidade sequencial das combinações de blocos e adaptar-se a diferentes domínios de consulta. Esses desafios levam à redução da relevância dos resultados. Para resolver esses problemas, apresentamos o CORAG, um sistema projetado para recuperar combinações de blocos ideais, levando em conta os domínios de consulta e os orçamentos dos usuários. Como o componente mais importante do nosso sistema, apresentamos o modelo de pesquisa de combinação de blocos ideal. O modelo emprega uma árvore de políticas baseada em MCTS para pesquisa sequencial de ordens de combinação de blocos, o que nos permite considerar totalmente as associações entre blocos (Desafio 1), bem como a não monotonicidade da utilidade da ordem de combinação de blocos (Desafio 2). Além disso, propomos um módulo de inferência de configuração que recomenda configurações ideais de MCTS e reclassificadores para vários domínios de consulta, abordando assim o Desafio 3. A seguir, apresentamos uma breve descrição desses dois módulos.

Pesquisa de combinação de blocos ideal: Uma abordagem direta para considerar as associações de blocos consiste em recuperar blocos potenciais de um banco de dados de vetores (conforme mostrado na Etapa 1) e, em seguida, explorar exaustivamente todas as combinações de blocos possíveis. No entanto, essa abordagem leva a uma latência e a um custo computacional significativos. Para atenuar esse problema, construímos uma árvore de estratégias (mostrada na Etapa 2) que reformula a busca da combinação ideal de blocos como um problema de busca no nó da árvore. Especificamente, o nó raiz da árvore de estratégias representa um estado inicial vazio e cada nó filho corresponde a uma combinação específica de blocos. Por exemplo, se o nó raiz tiver nós filhos que representam blocos χ1, um de seus filhos pode representar a combinação χ1+χ2, enquanto outro pode representar χ1+χ3

Projetamos um algoritmo de pesquisa baseado em MCTS para resolver esse problema. Diferentemente do MCTS tradicional, nossa abordagem expande o nó de maior utilidade em cada iteração enquanto avalia todos os nós filhos possíveis. Além disso, consideramos as restrições de custo e orçamento durante a pesquisa na árvore de estratégias. As utilidades dos nós são calculadas equilibrando a exploração com o controle de custos, otimizando a eficiência e a precisão.

Raciocínio de configuração: Uma solução simples para o ajuste da configuração seria enumerar todas as configurações ou reclassificações possíveis, computar os resultados em paralelo e, em seguida, escolher a melhor configuração. No entanto, isso levaria a um custo irreal do sistema RAG. Para otimizar as configurações (ou seja, o número de iterações, o fator de custo e o fator de exploração) do processo de pesquisa da árvore de políticas, introduzimos um agente de configuração que gera configurações dinamicamente com base no domínio da consulta. Para garantir a eficácia do modelo, empregamos uma abordagem de aprendizado comparativo usando pares de rótulos positivos e negativos: os rótulos positivos correspondem a embeddings de consulta do mesmo melhor ranqueador, enquanto os rótulos negativos vêm de um melhor ranqueador diferente. A função de perda conjunta é usada para otimizar simultaneamente a regressão (para ajuste de parâmetros) e o aprendizado de contraste (para aprimorar a discriminação de rótulos).

Resumo. O fluxo da nossa estrutura é mostrado na Figura 2. Primeiro, geramos embeddings para a consulta de entrada e depois a usamos para recuperar blocos potenciais do banco de dados de vetores. Esses embeddings de consulta também são alimentados em um agente de configuração que gera dinamicamente uma configuração MCTS ideal com base no campo de consulta. Usando essa configuração ideal, podemos pesquisar na árvore de políticas para determinar a melhor combinação de blocos e a melhor ordem entre os blocos potenciais recuperados. Por fim, essa combinação ideal de blocos é usada para criar as dicas finais para os LLMs.

4 Pesquisa de combinação de blocos

Conforme mencionado anteriormente, a ordem das combinações de blocos tem um impacto significativo na eficiência da construção de dicas para LLMs. Devido ao grande número de combinações possíveis, especialmente quando um grande número de blocos está envolvido, não é viável enumerar todas as ordens possíveis de combinação de blocos. Nesta seção, propomos uma nova abordagem que alcança um bom equilíbrio entre eficiência e precisão no problema de busca da ordem ideal de combinação de blocos. Na Seção 4.1, modelamos o problema como a busca de nós ideais em uma árvore de políticas (Seção 4.1). Em seguida, propomos um algoritmo baseado em MCTS para resolver esse problema de busca de nós (Seção 4.2).

4.1 Modelagem de pesquisa de árvore de estratégia

Para encontrar a ordem combinatória ideal, a primeira etapa é encontrar uma estrutura de dados que possa enumerar com eficiência todas as ordens combinatórias possíveis. Uma opção natural é uma árvore, na qual, percorrendo os nós da raiz até as folhas, podemos explorar todas as possíveis respostas.

Árvore de estratégias. Conforme mostrado na Fig. 3, construímos uma árvore de estratégias para representar a ordem de todas as combinações possíveis de blocos extraídas do banco de dados de vetores. Especificamente, o nó raiz simboliza o estado inicial sem blocos, e cada nó subsequente representa um bloco selecionado entre os blocos potenciais. Assim, os nós filhos são gerados a partir de seus nós pais, selecionando o próximo bloco disponível da fila de blocos potenciais e mesclando-o à sequência estabelecida pelo nó ancestral. Por exemplo, se um nó representa a sequência de combinação de blocos {χ1}, então um nó filho pode conter sequências de combinação subsequentes, como {χ1, χ2}, {χ1, χ3} ou {χ1, χ4}. Portanto, definimos formalmente a árvore de estratégias da seguinte forma:

Definição 4.1 (Árvore de estratégias). Dada uma consulta q e um conjunto de blocos potenciais {χ1, χ2, ... , χn}, construímos uma árvore de políticas T. χn}, construímos uma árvore de políticas T. O nó raiz de T representa o estado inicial sem nenhum bloco. Cada nó não raiz subsequente contém um conjunto de blocos, mesclando os blocos recém-selecionados dos blocos potenciais restantes na sequência de seus nós pais. Esse processo constrói sequencialmente uma combinação ordenada de blocos em cada nó não raiz, e nosso objetivo é encontrar o nó com a maior pontuação de utilidade.

Em uma árvore de estratégias, nosso objetivo é selecionar um nó que contenha blocos ordenados que ofereçam o benefício máximo com o custo mínimo. Para isso, precisamos projetar uma função de cálculo de utilidade para avaliar a compensação entre benefícios e custos. Essa função é quantificada pelo que definimos como "utilidade do nó", conforme descrito abaixo.

Utilitário nodal. A métrica de utilidade consiste em dois componentes: o benefício derivado da seleção de uma combinação de blocos e o custo de usar os blocos como prompts para LLMs. Especificamente, os benefícios são quantificados pelos LLMs, que medem a similaridade entre os blocos selecionados e a consulta. Em particular, nós os denotamos como valores de nó V. Em seguida, usamos o algoritmo UCB (Upper Confidence Bound) [3] para equilibrar o equilíbrio entre a exploração (valores de nó V(vi)) e a exploração (contagens de pesquisa N(vi)) dos valores de nó V(vi) e contagens de pesquisa N(vi). Para um determinado nó vi, com relação ao custo, consideramos o custo de rotulagem definido na Seção 2 e medido pela proporção do custo da combinação de blocos atual em relação ao orçamento total alocado B. Assim, a utilidade do nó é definida da seguinte forma:

Definição 4.2 (Utilidade nodal). Considerando uma árvore de estratégias e um orçamento de custos B, a utilidade de um nó não raiz é definida como:

em que V(vi) é o valor estimado do benefício da combinação de blocos no nó vi, determinado pelo modelo de treinamento, N(vi) é o número de visitas ao nó vi, o que facilita a exploração de nós menos visitados, e N é o número total de visitas a todos os nós na árvore de políticas para garantir um equilíbrio entre exploração e aproveitamento. Além disso, cost(vi) denota o custo de rotulagem do nó vi, B é o orçamento total de rotulagem, c regula a compensação de exploração e λ serve como um fator de penalidade de custo para melhorar a eficiência do custo.

Modelagem de seleção de nó ideal. Com base nas utilidades dos nós definidos, a tarefa de selecionar a ordem ideal de combinação de blocos, conforme descrito na Seção 2, é reformulada como a seleção do nó ideal na árvore de políticas T . Sob uma determinada restrição orçamentária B, o objetivo é identificar os nós vi ⊆ T para maximizar a utilidade U(vi) e, ao mesmo tempo, garantir que o custo total associado a vi não exceda B. Formalmente, isso é expresso como:

em que V(vi) é o benefício estimado da combinação de blocos no nó vi e cost(vi) representa seu custo associado. Essa fórmula permite escolher o bloco que maximiza a utilidade dentro de um determinado orçamento.

Figura 3: Fluxo de trabalho de otimização de blocos baseado em MCTS

Algoritmo 1 Pesquisa em árvore de estratégias baseada em MCTS

4.2 Pesquisa de árvore de políticas baseada em MCTS

Motivação. A enumeração de todos os nós da árvore de políticas encontrará o nó ideal, mas resultará em um alto custo computacional. Para resolver esse problema, uma abordagem direta é aplicar uma estratégia gulosa que navega pela árvore iterativamente, começando pelo nó raiz. Em cada iteração, selecione o nó filho com o maior benefício e continue até que o orçamento se esgote. Entretanto, é provável que essa abordagem leve a resultados abaixo do ideal. Por exemplo, χ1 pode ter um benefício um pouco maior do que χ2, mas χ2 + χ3 pode ter um benefício significativamente maior do que χ1 + χ3. Nesse caso, uma abordagem gananciosa pode levar a resultados abaixo do ideal. Portanto, é necessário revisitar o nó pai de alta eficiência. Ao mesmo tempo, precisamos reduzir a exploração de nós de baixo benefício.

Para atingir nosso objetivo, propomos um método de pesquisa de árvore de estratégia baseado em MCTS, projetado para selecionar e classificar eficientemente as combinações de blocos. Essa abordagem explora iterativamente o espaço de possíveis sequências de blocos, otimizando uma determinada consulta dentro de uma restrição orçamentária especificada.

Visão geral. As políticas baseadas em MCTS são descritas no Algoritmo 1. Primeiro, inicializamos o nó raiz da árvore de políticas usando consultas de entrada. Quando o orçamento computacional não se esgota, executamos iterativamente duas etapas principais: seleção de nós e atualização de utilidades. Quando o limite de iteração ou o orçamento é atingido, interrompemos o processo e pesquisamos recursivamente na árvore para encontrar o nó com a utilidade mais alta. Diferentemente das estratégias MCTS tradicionais, que geralmente se concentram apenas no nó raiz, nossa abordagem também considera nós de nível intermediário promissores para maximizar a utilidade da combinação de blocos.

Explicação detalhada dos principais recursos. Explicamos melhor essas duas funções principais a seguir:

1. Seleção de nós (Algoritmo 2). Selecionamos recursivamente o nó com o maior valor de utilidade que tem maior probabilidade de levar à combinação ideal de blocos. Especificamente:

Algoritmo 2 NodeSelection

Algoritmo 3 UtilityUpdate

• Selecione. Identificamos o nó com o valor máximo de utilidade. Se 𝑣 ainda não foi expandido, geramos todos os nós descendentes possíveis e os incluímos na árvore de estratégias. Se 𝑣 tiver sido expandido, selecionamos o nó descendente com a utilidade mais alta para exploração adicional.
• Extensões. Depois de selecionar o nó com a maior utilidade, nós o estendemos gerando todos os possíveis nós filhos. Cada nó filho representa uma nova ordem de possíveis combinações de blocos. Nossa abordagem usa a expansão paralela, que calcula e avalia vários subnós simultaneamente. Esse paralelismo explora a capacidade da rede de valores de lidar com várias combinações a um custo computacional semelhante ao de um único nó, aumentando a eficiência da pesquisa.
• Utilidade calculada. Calculamos o valor de utilidade para cada novo nó filho usando a fórmula de utilidade. O modelo de reclassificação R processa várias combinações de blocos em paralelo para gerar.

5 Configuração do agente

Depois de abordar a eficiência do processamento das associações de blocos dentro do orçamento do usuário, a tarefa restante é projetar um sistema que se adapte ao domínio de cada consulta. O processo do MCTS envolve várias configurações importantes, inclusive a seleção do ranqueador, o número de iterações, o fator de exploração e o fator de custo. A otimização dessas configurações em diferentes tipos de consulta é particularmente desafiadora. Para resolver esse problema, propomos uma estrutura de agente de configuração que prevê o ranqueador e a configuração ideais para cada consulta. Nesta seção, primeiro apresentamos a estrutura do agente na Seção 5.1 e, em seguida, descrevemos o processo de aprendizagem do modelo na Seção 5.2.

5.1 Estrutura de modelagem

Motivação. Para enfrentar o Desafio 3, que exige a adaptação ao domínio de cada consulta e a recomendação de configurações ideais, uma solução simples é usar um classificador MLP para atribuir cada consulta ao seu melhor ranqueador. No entanto, os experimentos iniciais mostram que a classificação MLP tem um desempenho ruim. Após uma análise mais aprofundada, observamos que tipos semelhantes de consultas tendem a compartilhar os mesmos melhores ranqueadores e configurações. Portanto, usar a rede siamesa com aprendizado de contraste para aproximar as consultas da mesma categoria e, ao mesmo tempo, afastar as consultas de categorias diferentes é uma abordagem mais viável.

A Figura 4 apresenta uma visão geral do nosso agente configurado, que consiste em dois módulos principais responsáveis pela transformação da entrada em um gráfico de recursos. Primeiro, o módulo de incorporação de entrada gera incorporações da consulta de entrada. Posteriormente, a rede de codificação processa esses embeddings para produzir mapas de recursos, que são usados para derivar várias configurações das definições do MCTS.

As seções a seguir descrevem cada componente em detalhes e explicam a lógica do projeto.

(1) Insira a consulta incorporada em. Para capturar com eficiência os fatores de várias consultas, dada a diversidade de tipos de consultas, como fatos explícitos, fatos implícitos, racionalidade interpretável e racionalidade oculta [47], usamos o modelo de incorporação BGE-M3 [6] para gerar incorporações para cada consulta. Essas incorporações aprimoram a estrutura de aprendizado ao mapear tipos semelhantes de consultas para a mesma categoria de reclassificação. Representados em um espaço de 1024 dimensões, os embeddings capturam os recursos semânticos subjacentes, permitindo que a rede de codificação os compare e classifique com eficiência. Essa etapa ajuda a melhorar a relevância da recuperação de diferentes tipos de consulta. Além disso, o uso do mesmo modelo de incorporação também gera incorporações anotadas de reclassificação ideais, incluindo seus recursos exclusivos e metadados associados, permitindo que o modelo alinhe as consultas à reclassificação ideal.

(2) Geração de mapas de recursos usando redes codificadas: o Para otimizar a tarefa de seleção de reclassificadores e recomendar o melhor reclassificador para cada consulta em diferentes tipos de consulta, usamos redes de codificação para aprender eficientemente representações que são úteis para a classificação e a previsão de configuração. Usamos uma rede siamesa que consiste em três camadas totalmente conectadas para essa tarefa. Ela processa os embeddings de consulta de entrada de dimensão d = 1024 e aprende os resultados da classificação e as previsões de configuração do MCTS (ou seja, número de iterações e λ). As ramificações da rede de codificação compartilham pesos, e cada ramificação aplica uma transformação linear seguida pela ativação RELU. Sequencialmente, a primeira camada oculta reduz a dimensionalidade para 512, a segunda para 256 e a terceira para 128. A camada de saída final fornece previsões de classificação que especificam o melhor ranqueador para cada consulta, bem como saídas de regressão que são usadas para prever as configurações MCTS mais eficientes para orientar o processo de pesquisa. A saída de classificação identifica o melhor ranqueador para cada consulta, enquanto a saída de regressão determina as melhores definições de configuração do MCTS.

5.2 Treinamento conjunto

Nesta seção, descrevemos o processo de treinamento para o desenvolvimento de agentes de configuração. Conforme mostrado na Fig. 4, implementamos três tarefas de treinamento conjunto para melhorar a eficiência do treinamento do modelo. As duas primeiras tarefas envolvem classificação e regressão para selecionar o melhor ranqueador e prever os melhores valores dos hiperparâmetros do MCTS, respectivamente. Além disso, incorporamos uma abordagem de aprendizado comparativo para aprimorar ainda mais o processo de aprendizado.

5.2.1 Perdas de classificação e regressão.

Dado o rótulo de reclassificação previsto Y(pred) e sua reclassificação ideal real correspondente Y(true), a perda de classificação L(cla) é calculada da seguinte forma.

em que F(cla) representa a perda de entropia cruzada entre os rótulos de reclassificação previstos e reais. Essa função de perda ajuda a classificar com precisão o ranqueador ideal para cada consulta. Da mesma forma, a perda de regressão L(reg) é definida como.

em que 𝐹reg é o erro quadrático médio (MSE) entre o parâmetro previsto do MCTS 𝑝pred e o parâmetro real do MCTS 𝑝true. Essa métrica garante uma previsão precisa da configuração do MCTS em termos do número de iterações e 𝜆.

Figura 4: Visão geral do Configuration Agent

5.2.2 Aprendizagem contrastiva.

Para distinguir com eficiência entre diferentes domínios de consulta e recomendar a melhor configuração para cada consulta, usamos o aprendizado comparativo para aproximar as consultas no mesmo domínio e, ao mesmo tempo, afastar as incorporações de diferentes classes de reclassificação.

Preparação da comparação. Para preparar o conjunto de dados de treinamento, precisamos determinar o ranqueador e a configuração ideais para cada consulta. Neste estudo, os ranqueadores ideais e as configurações correspondentes para cada consulta foram determinados por meio de experimentos extensivos com várias configurações. Posteriormente, os pares de consultas são gerados com base nessas anotações de reclassificação ideais. Os pares positivos consistem em consultas que compartilham o mesmo reclassificador ideal, facilitando a minimização de sua incorporação no espaço de recursos. Por outro lado, os pares negativos consistem em consultas com diferentes reclassificadores com o objetivo de maximizar a distância de incorporação entre eles. Como alguns reclassificadores têm desempenho semelhante em determinadas consultas, selecionamos apenas os casos em que o ROUGE-L difere em mais de 10% para formar nosso conjunto de dados de treinamento.

Perda de contraste. Conforme mostrado na Fig. 4, para um determinado par positivo (𝑥𝑖, 𝑥+𝑖) e o par negativo (𝑥𝑗, 𝑥-𝑗), primeiro usamos o modelo de codificação para gerar seus mapas de recursos correspondentes. Esses mapas de recursos são então usados para calcular a perda de contraste 𝐿con. Especificamente, esse processo pode ser representado da seguinte forma:

Entre eles.𝑓𝜃(𝑥) denota a função de incorporação.𝐹con é uma função de similaridade aplicada a dois tipos de pares: pares positivos (com solicitantes semelhantes) e pares negativos (com solicitantes diferentes). Essa função de perda foi projetada para garantir que as consultas com o mesmo solicitante estejam mais próximas no espaço de incorporação e que as consultas com solicitantes diferentes estejam mais distantes.

5.2.3 Todo o processo de treinamento.

Por fim, a função de perda total L(total) é a soma das perdas de comparação, classificação e regressão, como segue.

Em particular, a perda de contraste Lcon(θ) incentiva as incorporações de consultas com a mesma reclassificação ideal a ficarem mais próximas, enquanto afasta as incorporações de consultas com reclassificações diferentes. A perda de classificação Lcla(θ) ajuda o modelo a identificar corretamente os reclassificadores usando a entropia cruzada, enquanto a perda de regressão Lreg(θ) minimiza o erro na previsão da configuração ideal do MCTS.

Observações. Depois que a perda total Ltotal é calculada, os parâmetros de rede θ são atualizados usando o gradiente descendente em uma taxa de aprendizado η. Esse processo de otimização é repetido em vários ciclos E e lotes, garantindo que tanto o classificador quanto as previsões de parâmetros sejam aprimorados ao longo do tempo.

6 experimentos

O estudo experimental tem como objetivo responder às seguintes perguntas.
- RQ1 Como o nosso CORAG se compara a outros métodos em pipelines de RAG com restrições de custo?
- Qual é a eficiência do RQ2 CORAG com diferentes tamanhos de bloco?
- RQ3 Quais são os gargalos atuais no RAG?
- Quais são os gargalos do RAG?
- RQ4 Qual é a escalabilidade do CORAG com diferentes tamanhos de conjuntos de dados?
- RQ5 Qual é a eficácia de cada projeto no CORAG?

6.1 Configuração experimental

Ambiente. Integramos nosso sistema com a popular estrutura RAG LlamaIndex. Os experimentos foram executados em um servidor Linux equipado com uma CPU Intel Core i7-13700K (12 núcleos, 24 threads, 5,3 GHz), 64 GB de RAM e SSD NVMe de 1 TiB. O módulo Configuration Agent foi implementado em uma GPU NVIDIA RTX 4090 equipada com 24 GB de VRAM usando o PyTorch 2.0.

Tabela 1: Dados estatísticos usados no experimento.

Conjunto de dados.Para avaliar o desempenho do CORAG em diferentes cenários, realizamos experimentos em dois conjuntos de dados diferentes com diferentes focos de tarefas: (1) WikiPassageQA [7] é um benchmark de perguntas e respostas que contém 4.165 perguntas e mais de 100.000 blocos de texto destinados a avaliar a recuperação em nível de parágrafo. (2) MARCO [24] é um conjunto de dados abrangente dedicado a tarefas de processamento de linguagem natural com ênfase principal em perguntas e respostas e recuperação de informações. Conforme mostrado na Tabela 1, tanto o WikiPassageQA quanto o MARCO fornecem perguntas e anotações de parágrafo de alta qualidade, o que os torna adequados para avaliar a eficácia da recuperação. Em nossos experimentos, solicitamos aos LLMs que gerassem respostas verdadeiras para cada conjunto de dados. Por exemplo, se usarmos o Llama3 para avaliar o desempenho do CORAG, também usaremos o Llama3 para gerar verdades básicas na mesma configuração experimental para manter a imparcialidade e o alinhamento com as características dos LLMs.

Linha de base.Comparamos o desempenho do CORAG com duas linhas de base típicas do RAG:

- RAPTOR [29]: o RAPTOR constrói uma árvore hierárquica de resumo de documentos incorporando, agrupando e resumindo recursivamente blocos de texto para obter uma abstração em vários níveis. Essa abordagem é consistente com a abordagem baseada em agrupamento discutida na Seção 1. Concluímos a construção da árvore dentro de restrições orçamentárias aproximadas.
- NaiveRAG: esse é um método básico para recuperar blocos relevantes. Primeiramente, os blocos candidatos são recuperados do banco de dados de vetores com base na pesquisa de similaridade de vetores e, em seguida, são classificados usando o modelo de reclassificação. Esse método é o método AKNN mencionado na Seção 1. Para satisfazer a restrição de custo, usamos uma estratégia de alocação de orçamento gulosa para recuperar blocos até que o orçamento se esgote completamente.

Além disso, removemos o agente de configuração em nossa abordagem como linha de base para avaliar seu impacto no desempenho do CORAG, referindo-nos a essa versão como CORAG sem agente. Por fim, implementamos uma abordagem chamada CORAG Upper, que estabelece um limite superior explorando todas as combinações possíveis de blocos e selecionando a melhor ordem. Devido ao grande número de combinações possíveis, no caso do CORAG Upper, limitamos a exploração a combinações de menos de seis blocos.

Observações.Outras abordagens, como o GraphRAG [22], dependem muito de chamadas frequentes a LLMs para resumir blocos e construir índices, incorrendo em custos enormes (por exemplo, bilhões de tokens) que excedem nossas restrições rigorosas de custo. Portanto, essas abordagens não são viáveis para resolver nosso problema. Para uma comparação justa, excluímos esses tipos de abordagens RAG de nossos experimentos.

Configurações de hiperparâmetros: os hiperparâmetros do CORAG são determinados automaticamente pelo agente de configuração, enquanto o NaiveRAG não exige nenhum hiperparâmetro. Para os outros métodos de linha de base, garantimos a consistência usando os mesmos hiperparâmetros para comparações justas. Especificamente, definimos o fator de exploração como 2,4, o número de iterações como 10 e o fator de custo λ como 0,1. Experimentos preliminares mostram que essa configuração otimiza o desempenho da linha de base. Também realizaremos mais estudos de ablação para validar essas configurações.

Configurações de parâmetros de aprendizado. Em nossa abordagem, o agente de configuração é treinado usando o aprendizado de contraste. Os hiperparâmetros usados nesse processo incluem perda de contraste (margem=1,0), taxa de aprendizado (lr=0,001), tamanho do lote (32), número de ciclos (num_epochs=60) e modelo de incorporação (ou seja, BAAI/bge-m3 [6]).

Métricas de avaliação. Avaliamos a eficácia comparando as pontuações do Rouge de respostas de verdade e respostas geradas, usando Rouge-1, Rouge-2 e Rouge-L como métricas de avaliação. Para avaliar a eficiência, medimos a latência necessária para responder às consultas usando métodos diferentes.

Figura 5: Comparação de eficiência

6.2 Comparação de desempenho

6.2.1 RQ1: Comparação de roteiros.

Conforme mostrado na Fig. 2, comparamos o desempenho do CORAG com várias linhas de base em diferentes conjuntos de dados, principalmente usando o WikiPassageQA e o MARCO. A avaliação é realizada em três tamanhos de bloco diferentes, usando as métricas Rouge-1, Rouge-2 e Rouge-L para avaliar a melhoria das respostas geradas pelo LLM devido ao nosso método de recuperação. Os métodos RAG convencionais (por exemplo, NaiveRAG e RAPTOR) apresentam uma melhoria significativa de cerca de 251 TP3 T. Não é de surpreender que o CORAG não exceda o limite superior, o que representa um caso extremo em que todas as ordens de combinação possíveis são enumeradas exaustivamente, o que é claramente ineficiente e impraticável. Concluindo, o CORAG supera o desempenho da linha de base, aumentando a relevância da pesquisa e, ao mesmo tempo, reduzindo efetivamente o espaço de pesquisa.

6.2.2 RQ2: Avaliação da eficiência.

Conforme mostrado na Fig. 5, como o CORAG se baseia em um algoritmo de pesquisa em árvore, o agente ajuda a prever o melhor ranqueador e os melhores parâmetros para uma determinada consulta. Portanto, é fundamental avaliar o impacto de diferentes tamanhos de blocos e conjuntos de dados sobre a eficiência da tarefa de otimização da recuperação. Testamos a eficiência usando diferentes conjuntos de dados e tamanhos de blocos e observamos que o NaiveRAG, usando métodos tradicionais de recuperação, atinge tempos de recuperação mais curtos, mas pontuações de Rouge mais baixas. O CORAG superior tem bom desempenho em termos de Rouge, mas é significativamente menos eficiente, pois explora todo o espaço de pesquisa. Da mesma forma, o RAPTOR utiliza um LLM externo para sumarização e apresenta baixa eficiência. Por outro lado, nossa abordagem CORAG consegue um equilíbrio eficaz entre eficiência e relevância da pesquisa.

6.2.3 RQ3: Decomposição do desempenho.

Mostramos uma decomposição do desempenho da linha de base NaiveRAG para destacar os gargalos dos sistemas RAG atuais. Para enfrentar o desafio de pesquisar a ordem ideal de combinações de blocos, a implementação usando o NaiveRAG requer as seguintes etapas: a) obter a incorporação da consulta, b) recuperar as possíveis combinações de blocos, c) reclassificar as possíveis combinações de blocos e d) refinamento de dicas. A Figura 6 informa a latência média de cada etapa, na configuração experimental anterior.

Tabela 2: Comparação do Rouge nos conjuntos de dados WikiPassage QA e MARCO

6.2.4 RQ4: Avaliação da escalabilidade

O CORAG apresenta excelente escalabilidade, especialmente ao trabalhar com grandes conjuntos de dados (como o WikiPassageQA e o MARCO) que contêm um grande número de parágrafos. O número de blocos pode ser facilmente ampliado para 100 mil ou mais, dividindo cada parágrafo em 256 blocos rotulados. Apesar do aumento significativo no volume de dados, nosso tempo de recuperação aumentou em apenas 101 TP3T em relação à abordagem tradicional, demonstrando a eficiência do nosso sistema no gerenciamento de tarefas de recuperação em grande escala. Notavelmente, nosso sistema supera o método CORAG Upper, exigindo apenas um décimo do tempo de recuperação e ainda fornecendo pontuações ROUGE competitivas. Esse equilíbrio eficaz entre o desempenho e a sobrecarga computacional destaca a capacidade do sistema de podar com eficiência o espaço de pesquisa, garantindo uma recuperação rápida mesmo em conjuntos de dados enormes. Portanto, nossa abordagem é adequada para cenários que exigem processamento de dados em grande escala e alta precisão de recuperação.

Figura 6: Decomposição de desempenho

Tabela 3: Comparação do desempenho com diferentes orçamentos

6.3 RQ5: Estudos de ablação

6.3.1 Estudos de ablação com diferentes orçamentos

Conforme mostrado na Tabela 3, avaliamos o CORAG como um sistema com restrições de custo e investigamos o impacto de diferentes orçamentos no desempenho geral. Usando o conjunto de dados MARCO, definimos os limites de orçamento para 1024, 2048 e 8192 tokens e avaliamos os resultados usando o ROUGE.CORAG supera consistentemente todas as linhas de base nesses níveis de orçamento. Notavelmente, o custo médio de marcação do CORAG permanece abaixo de cada limite de orçamento, sugerindo que a utilidade dos blocos não é monotônica, conforme destacado no Desafio 2. À medida que o orçamento aumenta, o CORAG se beneficia de um espaço de pesquisa expandido e consegue incluir informações mais relevantes, em vez de apenas aumentar o número de blocos.

6.3.2 Estudos de ablação com diferentes coeficientes de exploração

Conforme mostrado na Fig. 7, realizamos um estudo de ablação para avaliar o impacto de diferentes coeficientes de exploração sobre o desempenho do sistema, testando especificamente valores de C de 0, 1, 2 e 3. Os resultados mostram que um coeficiente de exploração de aproximadamente 2 proporciona o melhor desempenho, alcançando um equilíbrio ideal entre exploração e aproveitamento durante o processo de busca. Esse equilíbrio permite que o sistema descubra com eficiência informações relevantes e, ao mesmo tempo, mantenha o foco em blocos de alto potencial, o que, em última análise, leva a uma melhor resposta do RAG. Em contrapartida, coeficientes de exploração mais baixos levaram a resultados abaixo do ideal devido à subexploração, enquanto coeficientes de exploração mais altos levaram a resultados abaixo do ideal devido ao desvio excessivo do foco. Essas descobertas enfatizam a função essencial do coeficiente de exploração no desempenho do processo de pesquisa do CORAG e destacam a importância do ajuste cuidadoso dos parâmetros.

Figura 7: Comparação de ROUGE entre diferentes valores de C

6.3.3 Estudos de ablação com diferentes fatores de custo

Conforme mostrado na Fig. 8, foi realizado um estudo de ablação para avaliar o impacto de diferentes coeficientes de custo sobre o desempenho do sistema, testando especificamente os valores 0, 0,1, 0,2 e 0,3. Os resultados mostram que a introdução de coeficientes de custo no utilitário leva a uma ligeira diminuição na pontuação ROUGE. Essa diminuição se deve ao fato de que, na ausência de uma restrição de custo, a CORAG tende a produzir produtos mais longos, embora ao custo da eficiência de custo. No entanto, apesar da ligeira diminuição nas pontuações ROUGE, a diminuição permanece dentro de 51 TP3T, o que é aceitável. Esses resultados destacam a importância de ajustar eficientemente os coeficientes de custo para equilibrar a riqueza de resultados e as restrições de custo, enfatizando ainda mais a função do nosso agente de configuração para permitir o ajuste eficiente da configuração para otimizar o desempenho do CORAG.

Figura 8: Comparação de ROUGE para diferentes valores de lambda

6.3.4 Estudos de ablação com diferentes rearranjadores

Para avaliar o impacto de diferentes rearranjadores no desempenho da recuperação, realizamos um estudo de ablação usando seis modelos de rearranjadores amplamente reconhecidos: jina-reranker-v1-turbo-en, jina-reranker-v2-base-multilingual, bge-reranker-v2-m3, bge- reranker-large, bge-reranker-base e gte-multilingual-reranker-base. Esses rearranjadores foram avaliados no conjunto de dados MARCO usando o modelo llama3-8B configurado com um fator de custo fixo de 0,1, um fator de exploração de 2,4 e uma restrição orçamentária de 1024.

Tabela 4: Comparação de desempenho com diferentes reordenadores

Os resultados da Tabela 4 mostram que há diferenças de desempenho entre os diferentes reordenadores, destacando a importância de selecionar cuidadosamente um reordenador para otimizar o desempenho do sistema RAG sob restrições operacionais específicas. Entre os reorganizadores, o gte-multilingual-reranker-base e o bge-reranker-large apresentam desempenho consistentemente forte na tarefa de GQ, sugerindo que esses modelos de reorganização são altamente eficientes na captura de informações relevantes para diferentes consultas de GQ. Observamos que, à medida que o tamanho do bloco aumenta no estudo de ablação, cada reposicionador individual apresenta desempenho inferior ao do agente em termos de reposicionadores recomendados para diferentes consultas. Isso sugere que o agente de configuração explora de forma eficaz a diversidade de reposicionadores, ajustando dinamicamente a configuração para melhorar os resultados da recuperação. A capacidade do agente de configuração de recomendar melhores opções de reorganização e configurações de parâmetros destaca sua função de maximizar o desempenho do sistema RAG, especialmente sob restrições como orçamento limitado.

6.4 Estudos de caso

A Figura 9 mostra três exemplos que comparam a qualidade de recuperação do CORAG e da abordagem NaiveRAG tradicional, destacando por que nossa abordagem supera a abordagem de linha de base. Devido à sua simples pesquisa e reordenação top-k, o NaiveRAG muitas vezes perde informações importantes e relevantes para a intenção da consulta, pois geralmente recupera blocos com base na correspondência de palavras-chave em vez de relevância para a consulta. Por exemplo, para a consulta "Are leafy flowers shrubs?" (As flores com folhas são arbustos?) o NaiveRAG recupera conteúdo que contém palavras-chave correspondentes, mas não fornece a classificação real da flor folhosa. Por outro lado, a estratégia de combinação de blocos do CORAG recupera o contexto que inclui a categoria da flor da folha, permitindo que o LLM forneça uma resposta mais precisa. Em outro caso, o NaiveRAG recuperou termos e cláusulas legais que incluíam "oxyfluorfen", mas não compreendia a intenção da consulta, enquanto o CORAG fornecia conteúdo que vinculava o oxyfluorfen a casos de uso em algodão, o que exigia uma pesquisa que o NaiveRAG não conseguia capturar. A pesquisa de similaridade vetorial do NaiveRAG não conseguiu capturar as relações lógicas entre os blocos. Por fim, para a consulta "De onde vêm as bactérias?" o NaiveRAG recupera blocos que contêm a palavra-chave "bactérias", mas não aborda sua origem, enquanto o CORAG fornece uma resposta mais completa, incluindo a origem das bactérias e suas condições de reprodução. Esses casos ilustram que o CORAG é excelente na recuperação de informações logicamente conectadas, o que o torna mais adequado do que o NaiveRAG para consultas que exigem mais do que a simples correspondência de palavras-chave.

Figura 9: Estudo de caso da CORAG

7 Insights sobre o RAG e opções futuras de design

7.1 Deficiências do RAG atual

Apresentamos uma análise do sistema RAG atual, revelando os desafios de desempenho enfrentados durante as fases de recuperação (S1), aprimoramento (S2) e geração (S3).

S1: Despesas gerais de pesquisa Os sistemas RAG atuais geralmente utilizam LLMs para estruturas de resumo e indexação, ignorando os custos computacionais associados a LLMs externos, aumentando assim a sobrecarga computacional. Os reordenadores baseados em modelos, embora melhorem a relevância durante a recuperação, introduzem uma latência significativa, o que pode prejudicar a eficiência contextual sensível à latência. A construção de índices econômicos e as otimizações de reordenação são essenciais para equilibrar a eficiência e o desempenho.

S2: Custos indiretos aprimorados As técnicas pós-recuperação, como a ordenação otimizada da combinação de blocos, aumentam a relevância contextual, mas exigem computação adicional. As estratégias de poda que minimizam o espaço de pesquisa e otimizam a ordem de combinação são essenciais para equilibrar o custo computacional e aumentar a relevância contextual. A otimização eficiente da combinação de blocos com ênfase na ordem e na coerência é essencial para reduzir o custo e melhorar o desempenho da recuperação.

S3: Custos indiretos de geração A engenharia eficiente de dicas para combinações ideais de blocos requer recursos computacionais significativos. O refinamento e a compactação de dicas específicas da consulta são essenciais para reduzir as despesas gerais e, ao mesmo tempo, manter a relevância e a concisão da entrada. Estratégias adaptáveis para lidar com diferentes tipos de consulta e requisitos específicos de domínio garantem a eficiência da solicitação sem comprometer a qualidade do resultado.

7.2 Opções de design

Para enfrentar os desafios acima, as opções de design a seguir têm como objetivo otimizar o desempenho do sistema RAG:

P1: Co-projeto de processos de recuperação e reordenação No CORAG, a extensão paralela na pesquisa em árvore acelera o processamento da consulta e reduz significativamente a latência, permitindo a recuperação e a reordenação simultâneas. Otimizações futuras podem resolver o gargalo eliminando atrasos específicos do estágio e melhorando ainda mais a eficiência da classificação. Essa abordagem de co-projeto gerencia com eficiência a ordem de combinação de blocos para melhorar o processo de classificação e a pontuação associada.

P2: Otimização de estruturas de árvore e iterações de pesquisa Os resultados mostram que as alturas mais curtas da árvore da estratégia melhoram a eficiência da pesquisa, reduzindo as despesas gerais de computação, o que é particularmente benéfico para grandes conjuntos de dados. A minimização da altura da árvore na pesquisa baseada em árvore melhora a velocidade da pesquisa de blocos contextualmente relevantes, reduzindo significativamente a latência e o custo computacional. Essa abordagem de otimização melhora o desempenho do sistema RAG em conjuntos de dados de grande escala.

P3: Projeto de ponta dinâmica A seleção do rearranjador com base no tipo de consulta e o uso de modelos de dicas adaptáveis podem melhorar a relevância da recuperação do LLM. As estruturas de dicas dinâmicas alinhadas à intenção da consulta e ao contexto específico do domínio podem manter a qualidade da saída dentro das restrições de recursos. Essa abordagem de engenharia de dicas adaptativas alcança um equilíbrio eficaz entre eficiência e qualidade de recuperação, abordando a natureza dinâmica das consultas em sistemas RAG.

8 Conclusão

Considerando a não monotonicidade da utilidade dos blocos, a correlação entre os blocos e a diversidade de diferentes domínios de consulta, propomos o CORAG, um sistema de otimização de recuperação baseado em aprendizado. Inicialmente, modelamos a ordem das combinações de blocos como uma árvore de políticas e exploramos essa árvore usando o MCTS, com o objetivo de encontrar a ordem ideal das combinações de blocos. Introduzimos um agente de configuração que prevê com precisão a configuração e o rearranjador ideais para uma determinada consulta. Além disso, elaboramos uma estratégia de expansão de consulta paralela para expandir vários nós em cada iteração. Os resultados experimentais mostram que nossa abordagem supera significativamente as abordagens de última geração dentro das restrições de custo e também se destaca em termos de eficiência.