One Model Is Enough: Native Retrieval Embeddings from LLM Agent Hidden States

Este artigo propõe um método que capacita agentes de LLM a realizar recuperação nativa de conhecimento diretamente a partir de seus estados ocultos, eliminando a necessidade de um modelo de incorporação separado e mantendo 97% da qualidade de recuperação original.

O essencial

Imagine que você tem um assistente superinteligente (um LLM, ou Modelo de Linguagem) que precisa pesquisar informações na internet para responder às suas perguntas.

Até hoje, a forma como esse assistente funcionava era um pouco como se ele tivesse que escrever um bilhete para si mesmo antes de pedir ajuda. O processo era assim:

1. O assistente pensa na resposta e escreve uma pergunta em texto (ex: "Qual a capital da França?").
2. Depois, ele pega esse texto e o entrega a um segundo funcionário (um modelo de "embedding") que traduz essa frase em um código matemático complexo para encontrar o documento certo.

O problema: O assistente já tinha toda a informação na cabeça dele! Ele já "entendeu" o contexto da conversa. Escrever a pergunta e depois pedir para outro funcionário traduzir isso é como se você já soubesse a receita do bolo, mas ainda assim precisasse escrever a receita num papel e entregar para um vizinho traduzi-la para o cozinheiro. É redundante, lento e gasta energia extra.

A Solução: "Um Modelo é Suficiente"

O artigo propõe uma ideia genial: e se o assistente pudesse falar a língua do pesquisador diretamente, sem precisar escrever o bilhete?

Os autores criaram um "adaptador" (uma pequena peça de software chamada cabeça de projeção) que se conecta diretamente à mente do assistente.

A Analogia do Tradutor Instantâneo:
Pense no assistente como um doutor que tem um diagnóstico completo na cabeça (os "estados ocultos" ou hidden states), mas que só consegue falar em "medicoês" complexo.

• O jeito antigo: O doutor escreve o diagnóstico num papel em "medicoês", e um tradutor (o segundo modelo) lê o papel e o transforma em uma linguagem que o arquivo médico entende.
• O jeito novo: O doutor usa um óculos mágico (o adaptador). Assim que ele pensa no diagnóstico, o óculos traduz instantaneamente o pensamento dele para a linguagem do arquivo médico, sem precisar escrever nada no papel.

Como eles fizeram isso?

Eles não ensinaram o assistente a pesquisar do zero. Eles usaram uma técnica chamada Distilação de Conhecimento. É como se eles colocassem o assistente (aluno) e o tradutor especialista (professor) lado a lado.

O professor diz: "Olhe, quando eu vejo essa pergunta, eu crio este código específico. Você, assistente, olhe para o seu próprio pensamento e tente criar um código que seja o mais parecido possível com o meu."

Para garantir que o assistente aprendesse bem, eles usaram três regras de treino:

1. Alinhamento: "Seja parecido com o professor."
2. Contraste: "Não confunda perguntas diferentes. Se a pergunta A é diferente da B, seus códigos devem ser bem distintos."
3. Ranking (Classificação): "Se o professor acha que o documento X é melhor que o Y, você também deve achar."

Os Resultados

Eles testaram isso em um banco de dados de conversas complexas (QReCC) e os resultados foram impressionantes:

• Qualidade: O novo sistema manteve 97% da qualidade do sistema antigo. Ele quase não perdeu nada na precisão.
• Velocidade: O sistema ficou 21 vezes mais rápido. Como não precisa mais esperar o "segundo funcionário" traduzir o texto, a resposta chega instantaneamente.
• Simplicidade: Agora, o sistema só precisa de um modelo rodando, em vez de dois.

O Que Isso Significa para o Futuro?

Imagine que você está conversando com uma IA no seu celular. Antes, para ela pesquisar algo, ela tinha que "pensar alto" (gerar texto), "parar" para processar isso em outro lugar e só depois buscar. Agora, ela pode "pensar" e "buscar" ao mesmo tempo, usando a mesma energia mental.

Resumo da Ópera:
Os autores mostraram que não precisamos de dois robôs para fazer o trabalho de um. Com um pequeno ajuste (o "óculos mágico"), o robô principal já consegue fazer a pesquisa sozinho, tornando tudo mais rápido, mais barato e mais eficiente, sem perder a inteligência.

É como se, em vez de ter que desenhar um mapa para um guia turístico, o guia já nascesse sabendo ler o GPS diretamente.

Resumo técnico

Resumo Técnico: One Model Is Enough

1. O Problema Identificado

Atualmente, os agentes de LLM (Large Language Models) que utilizam Geração Aumentada por Recuperação (RAG) operam em um pipeline de dois modelos:

1. O LLM gera uma consulta de busca como texto natural.
2. Um modelo de embedding separado codifica esse texto em um vetor denso para recuperação contra um índice de documentos.

Os autores argumentam que essa arquitetura introduz uma redundância fundamental: o LLM já processou todo o contexto conversacional (intenção do usuário, histórico de diálogo e requisitos da tarefa) e codificou esse entendimento em seus estados ocultos (hidden states). A geração do texto da consulta é apenas uma projeção discreta e com perda de informação desse entendimento interno rico. O segundo modelo (o encoder de embedding) precisa reprocessar esse texto do zero para recuperar informações semânticas que o LLM já possuía, descartando o entendimento interno do agente e aproximando-o novamente. Isso aumenta a complexidade da infraestrutura e a latência.

2. Metodologia Proposta

O artigo propõe equipar o agente LLM com uma capacidade nativa de recuperação, eliminando a necessidade de um modelo de embedding separado durante a inferência.

• Abordagem: Acoplar uma cabeça de projeção leve (lightweight projection head) aos estados ocultos do LLM. Esta cabeça mapeia diretamente os estados ocultos gerados durante a inferência normal para o espaço vetorial de embeddings existente.
• Extração de Estados Ocultos: Durante o processo de geração autoregressiva do LLM, os estados ocultos da última camada ($h_i$) de cada token gerado são coletados. Isso adiciona sobrecarga negligenciável, pois os estados já são computados.
• Arquitetura da Cabeça de Projeção:
  1. Projeção de Entrada: Uma camada linear mapeia a dimensão oculta do LLM para uma dimensão interna.
  2. Codificador Transformer: Uma pilha de camadas de encoder com atenção auto-regressiva agrega informações através da sequência gerada.
  3. Agrupamento (Pooling): Uso de mean pooling sobre as posições válidas para comprimir a sequência em um único vetor.
  4. Projeção de Saída e Normalização: Uma camada linear final projeta para a dimensão alvo, seguida de normalização L2 para garantir que a similaridade por produto escalar seja equivalente à similaridade cosseno.
• Objetivos de Treinamento (Distilação de Conhecimento): A cabeça de projeção é treinada para imitar um modelo de embedding "professor" (teacher) usando uma combinação de três funções de perda:
  1. Perda de Alinhamento ($L_{align}$): Minimiza a distância angular entre o vetor projetado e o embedding do professor para cada consulta.
  2. Perda Contrastiva ($L_{contra}$): Usa InfoNCE para preservar a estrutura discriminativa relativa entre as consultas dentro de um batch.
  3. Perda de Distilação de Rank ($L_{rank}$): Usa divergência KL para transferir as preferências de classificação de documentos do professor para o aluno, alinhando as distribuições de pontuação dos documentos candidatos.

3. Contribuições Principais

• Identificação e Formalização da Redundância: O trabalho formaliza a ineficiência do pipeline de dois modelos e propõe a projeção de estados ocultos como uma alternativa viável.
• Novo Objetivo de Treinamento: Design de uma função de perda tripla (alinhamento, contraste e rank) específica para projetar estados ocultos de LLMs em espaços de embedding.
• Validação Empírica Rigorosa: Experimentos abrangentes com 12 configurações de ablação, intervalos de confiança bootstrap e testes de significância estatística, demonstrando qualidade de recuperação quase idêntica sem um modelo separado.

4. Resultados Experimentais

Os experimentos foram conduzidos no benchmark de busca conversacional QReCC, utilizando a família de modelos Qwen (LLM: Qwen3-8B como agente; Embedding: Qwen3-Embedding-8B como professor).

• Qualidade de Recuperação: O método proposto manteve 97% da qualidade de recuperação da linha de base (pipeline padrão).
  • Recall@10: 0.607 (vs. 0.637 da linha de base).
  • MRR@10: 0.293 (vs. 0.329 da linha de base).
  • A diferença é estatisticamente significativa, mas a sobreposição de acertos entre os métodos é de 84,2%.
• Latência: Redução drástica de 43.5 ms para 2.0 ms por consulta (aceleração de 21.8x), pois elimina a passagem de forward do modelo de embedding.
• Ablação de Perdas:
  • A perda de alinhamento foi o componente individual mais forte.
  • A distilação de rank sozinha falhou completamente (colapso), mas trouxe ganhos consistentes quando combinada com alinhamento.
  • A combinação de todas as três perdas foi superior a qualquer par.
• Hiperparâmetros: O treinamento estendido (80 épocas) com uma taxa de aprendizado mais baixa ($2 \times 10^{-4}$) foi crucial para o melhor desempenho, enquanto taxas mais altas causaram colapso do treinamento.

5. Significado e Limitações

• Significado: O trabalho demonstra que é possível simplificar drasticamente a arquitetura de sistemas RAG, removendo a dependência de modelos de embedding externos durante a inferência. Isso reduz custos computacionais, latência e complexidade de infraestrutura, permitindo que o LLM "busque com seus próprios pensamentos".
• Limitações:
  • Avaliado apenas em um único dataset (QReCC).
  • Utiliza configurações "mesma família" (Qwen), o que pode facilitar a transferência de representações; cenários de famílias cruzadas podem ser mais desafiadores.
  • Ainda há uma pequena lacuna de qualidade estatisticamente significativa em relação à linha de base.
  • O modelo de embedding ainda é necessário durante a fase de treinamento (para gerar os dados de ensino), embora seja dispensado na implantação.

Em resumo, o artigo propõe uma mudança de paradigma onde o agente de IA não precisa "falar" para "ouvir" (gerar texto para ser codificado), mas pode acessar diretamente seu próprio entendimento interno para realizar buscas, mantendo alta eficiência e qualidade.