Resumo Técnico: LatentRAG

Declaração do Problema

Sistemas tradicionais de Geração Aumentada por Recuperação (RAG) frequentemente dependem de um processo de recuperação de etapa única, que luta com perguntas complexas que exigem raciocínio iterativo. Embora o RAG Agêntico aborde isso empregando Modelos de Linguagem de Grande Escala (LLMs) como agentes de busca que geram iterativamente pensamentos intermediários e subconsultas para recuperar informações, essa abordagem introduz latência significativa.

O principal gargalo nos métodos existentes de RAG Agêntico (por exemplo, Search-R1, AutoRefine) é a geração autoregressiva, token a token, de pensamentos e subconsultas em linguagem natural. Como demonstrado na análise do artigo, essas etapas de geração representam aproximadamente 90% da latência total de inferência, tornando o RAG Agêntico 15–20 vezes mais lento que o RAG de etapa única ingênuo. Além disso, a geração explícita de longas sequências de texto impede a paralelização eficiente e dificulta a otimização conjunta ponta a ponta do LLM e do modelo de recuperação devido à natureza discreta dos tokens de saída.

Metodologia: LatentRAG

Os autores propõem o LatentRAG, um framework que transfere tanto o raciocínio quanto a recuperação do espaço de linguagem discreto para o espaço latente contínuo. Em vez de gerar tokens explícitos de linguagem natural para pensamentos e subconsultas, o modelo opera sobre "tokens latentes"—os estados ocultos da última camada de tokens de espaço reservado especiais.

Componentes Principais

1. Geração Latente:

   • O LLM é solicitado com tokens especiais (por exemplo, <think1>...<thinkm> para pensamentos e <query1>...<queryn> para subconsultas).
   • Durante a etapa de pré-preenchimento, esses tokens especiais são processados em paralelo. O modelo extrai os últimos estados ocultos desses tokens para servir como tokens de pensamento latente ($H^\tau_t$) e subconsulta latente ($H^s_t$).
   • Isso permite que o modelo aloque computação interna sem a sobrecarga sequencial da geração autoregressiva de texto.

2. Recuperação Latente:

   • Como os tokens latentes residem no espaço de saída do LLM, um módulo projetor leve ($Proj_{ret}$) faz a ponte para o espaço de entrada do modelo de recuperação.
   • A subconsulta latente projetada é alimentada em um modelo de recuperação densa treinável para gerar embeddings.
   • Objetivo de Treinamento: Para alinhar os embeddings de subconsulta latente com o comportamento de recuperação de subconsultas em linguagem natural sem exigir anotações de documentos verdadeiros para cada etapa intermediária, os autores introduzem uma perda de divergência KL. Essa perda minimiza a divergência entre a distribuição de similaridade de documentos induzida pela subconsulta latente e aquela induzida por uma subconsulta de linguagem natural de referência (congelada). Isso permite a otimização conjunta totalmente diferenciável e ponta a ponta do LLM e do modelo de recuperação.

3. Decodificação Latente Paralela (Opcional):

   • Para abordar a natureza de "caixa preta" do raciocínio latente e garantir transparência, um mecanismo de decodificação paralelo é incorporado.
   • Módulos projetores mapeiam tokens latentes de volta ao espaço de entrada do LLM, permitindo que o modelo decodifique os pensamentos e subconsultas correspondentes em linguagem natural.
   • Crucialmente, como os tokens latentes são gerados em uma única passagem direta, a subsequente decodificação de pensamentos e subconsultas em diferentes etapas pode ser realizada em paralelo, reduzindo significativamente a latência de decodificação em comparação com a geração autoregressiva sequencial.

Contribuições Principais

1. Framework LatentRAG: Um sistema de RAG agêntico inovador que realiza raciocínio e recuperação inteiramente dentro do espaço latente, eliminando a sobrecarga de latência associada à geração explícita token a token de pensamentos e subconsultas.
2. Alinhamento no Espaço Latente: Um objetivo de treinamento baseado em divergência KL que alinha as representações de subconsulta latente do LLM com um modelo de recuperação densa pré-treinado. Isso facilita a otimização conjunta ponta a ponta e permite que tokens latentes sirvam como consultas de recuperação eficazes, apesar da escassez de dados em etapas intermediárias.
3. Mecanismo de Decodificação Paralela: Um método para traduzir tokens latentes de volta para linguagem natural para transparência. Esse mecanismo suporta decodificação paralela em múltiplas etapas de raciocínio, oferecendo um compromisso entre eficiência e interpretabilidade, mantendo uma vantagem significativa de velocidade sobre métodos explícitos.

Resultados Experimentais

Os autores avaliaram o LatentRAG em sete conjuntos de dados de referência (NQ, TriviaQA, PopQA, HotpotQA, 2Wiki, Musique, Bamboogle) usando vários modelos de recuperação e escalas de LLM.

• Desempenho: O LatentRAG alcança desempenho comparável aos métodos explícitos de RAG Agêntico mais avançados (por exemplo, Search-R1, AutoRefine), com diferenças de desempenho relativas geralmente dentro de 5%.
• Eficiência: O método reduz a latência de inferência em aproximadamente 90% em comparação com as bases de RAG Agêntico explícito.
  • Em média, o LatentRAG reduz a latência de ~5.000ms (Search-R1) para ~600ms.
  • Isso traz a latência do RAG Agêntico para níveis comparáveis ao RAG tradicional de etapa única (por exemplo, ~400–500ms).
• Escalabilidade: Os ganhos de eficiência mantêm-se em diferentes tamanhos de modelo (LLMs de 3B a 14B e modelos de recuperação de 0,6B a 8B). Aumentar o tamanho do LLM aumenta significativamente a latência para métodos explícitos devido a sequências de decodificação mais longas, enquanto o LatentRAG mantém a eficiência.
• Estudos de Ablação: O objetivo de recuperação de divergência KL mostrou-se superior à perda de cosseno ou à perda InfoNCE. A remoção do modelo de recuperação pré-treinado ou da perda de decodificação latente resultou em degradação de desempenho, confirmando a necessidade do viés indutivo do modelo de recuperação e da regularização fornecida pela decodificação.

Significado e Alegações

O artigo afirma que o LatentRAG conecta com sucesso a lacuna entre o alto desempenho do RAG Agêntico de múltiplas etapas e a eficiência do RAG de etapa única. Ao mover o raciocínio e a recuperação para o espaço latente, o framework evita as ineficiências computacionais da geração autoregressiva de texto, mantendo a capacidade de decompor perguntas complexas.

Os autores enfatizam que, embora o raciocínio latente tenha sido explorado para tarefas de geração exclusiva, sua aplicação ao RAG Agêntico introduz desafios únicos relacionados à escassez de dados e transparência. O LatentRAG aborda isso propondo uma estratégia de otimização conjunta e um mecanismo de decodificação paralelo opcional. O trabalho sugere uma mudança de motores de busca baseados em texto orientados a humanos para motores de busca baseados em embeddings orientados a agentes, oferecendo um paradigma mais eficiente para agentes de busca de informações na era dos modelos de linguagem de grande escala. Os autores observam modestamente que sua abordagem atual depende de ajuste fino supervisionado (SFT) usando trajetórias de agentes existentes, e trabalhos futuros poderiam explorar aprendizado por reforço para otimizar ainda mais as políticas de recuperação.