Translate Policy to Language: Flow Matching Generated Rewards for LLM Explanations

1. O Problema

À medida que agentes inteligentes (baseados em RL e LLMs) se integram cada vez mais ao cotidiano humano, a capacidade de explicar suas políticas de decisão em linguagem natural torna-se crucial para a confiança, segurança e alinhamento. O objetivo é treinar um LLM de Explicação que, dado apenas o contexto de uma decisão (sem ver a decisão em si), gere uma explicação que permita a um terceiro inferir corretamente a decisão original do agente.

Os desafios principais identificados pelos autores são:

• Pluralismo e Probabilidade: As avaliações humanas sobre a qualidade de uma explicação são inerentemente plurais e probabilísticas (diferentes pessoas podem julgar a mesma explicação de formas diferentes).
• Limitações do Feedback de IA (RLAIF): Métodos existentes de Reinforcement Learning from AI Feedback (RLAIF) utilizam LLMs auxiliares (Proxy LLMs) para gerar recompensas substitutas (proxy rewards). No entanto, esses proxies introduzem erros e ruídos em relação à distribuição real de recompensas humanas. A maioria dos métodos atuais não gerencia explicitamente essa discrepância distribucional, tratando as recompensas como determinísticas ou falhando em capturar a incerteza inerente.
• Generalização: Muitas soluções de explicação são específicas de tarefas, exigindo engenharia pesada para cada cenário (ex: carros autônomos vs. robótica vs. QA).

2. Metodologia

Os autores propõem um framework geral que treina LLMs de explicação usando Reinforcement Learning from AI Feedback (RLAIF), mas com uma inovação central: o uso de Fluxos Normais Contínuos (CNFs) e, especificamente, Fluxo Retificado (Rectified Flow), para gerar recompensas distribucionais.

Arquitetura e Fluxo de Treinamento:

1. LLM de Explicação ($\pi_e$): Recebe o contexto da decisão (ex: histórico de estados/ações em RL ou a pergunta em QA) e gera uma explicação natural. A decisão real é oculta durante a geração para forçar o raciocínio baseado apenas no contexto.
2. Proxy LLMs: Atuam como avaliadores iniciais. Dada a explicação gerada e o contexto, eles tentam inferir a decisão. A probabilidade de inferir a decisão correta é usada como sinal de recompensa.
3. Modelo de Recompensa com Fluxo Retificado (Rectified Flow):
   • Em vez de usar diretamente a saída do Proxy LLM (que é ruidosa), os autores treinam um modelo generativo baseado em Fluxo Retificado ($\phi$).
   • Mapeamento de Ruído: O modelo aprende a mapear uma distribuição de ruído base (ex: Gaussiana) para a distribuição complexa de recompensas gerada pelos Proxy LLMs.
   • Atenção Cruzada: A arquitetura do modelo de fluxo é condicionada ao contexto e à explicação através de uma camada de atenção cruzada dentro de um Transformer. Isso permite que o modelo de recompensa selecione cues linguísticos e contextuais relevantes ao gerar a distribuição de recompensas.
   • Recompensas por Frase: O modelo gera recompensas densas (por frase) baseadas na mudança na probabilidade de inferir a decisão correta à medida que a explicação avança.
4. Treinamento Iterativo: O LLM de Explicação e o modelo de Fluxo Retificado são treinados alternadamente. O LLM é otimizado via PPO (Proximal Policy Optimization) para maximizar as recompensas geradas pelo Fluxo Retificado.

Fundamentação Teórica:

O artigo apresenta uma prova teórica (Teorema 1) mostrando que, sob certas suposições (o ruído do Proxy LLM segue a mesma forma funcional da distribuição base do fluxo), o modelo de Fluxo Retificado pode limitar provadamente o desvio da distribuição de recompensas aprendida em relação à verdadeira distribuição de recompensas humanas.

• O Fluxo Retificado atua como um processo de "desruído" (denoising), aprendendo a inversa do processo que o Proxy LLM aplica às recompensas humanas ideais.
• Isso permite que o modelo generalize para "falsos negativos" (explicações boas que o Proxy LLM classificou erroneamente como ruins), algo que métodos determinísticos falham em fazer.

3. Principais Contribuições

1. Framework de Recompensas Distribucionais: Introdução de um método que modela explicitamente a distribuição probabilística de recompensas para tarefas de explicação, capturando o pluralismo das avaliações humanas.
2. Uso de Fluxo Retificado para RLAIF: Aplicação inovadora de Fluxos Normais Contínuos (especificamente Rectified Flow) como modelos de recompensa, oferecendo uma garantia teórica de erro reduzido em relação a métodos baseados apenas em LLMs proxy.
3. Arquitetura com Atenção Cruzada: Design de um modelo de recompensa que integra informações contextuais e linguísticas via atenção cruzada, permitindo uma avaliação mais precisa e generalizável.
4. Generalização para Agentes Diversos: O método é aplicado com sucesso tanto em tarefas de RL (SMAC, AI2-THOR) quanto em tarefas de LLM (MMLU, MathQA), demonstrando ser agnóstico à tarefa e ao agente.

4. Resultados Experimentais

Os autores avaliaram o método em quatro benchmarks: SMAC (jogo de estratégia), AI2-THOR (robótica/ambiente físico), MMLU (conhecimento profissional) e MathQA (problemas matemáticos).

• Desempenho Geral: O método proposto superou consistentemente todas as linhas de base, incluindo:
  • SFT (Supervised Fine-Tuning).
  • RLHF (PPO, DPO, KTO) usando recompensas determinísticas.
  • RLAIF padrão (usando recompensas diretas do Proxy LLM).
  • Frameworks de Raciocínio (CoT+SFT, ReFT).
  • Modelos de ponta como Skywork, DeepSeek-R1 e o3-mini.
• Métricas de Melhoria:
  • Em tarefas de RL (SMAC), houve um aumento de 12,4% na acurácia em comparação ao uso direto de recompensas do Proxy LLM.
  • Em tarefas de LLM (MMLU, MathQA), as melhorias variaram de 6,9% a 11,0%.
  • No benchmark complexo AI2-THOR, o método superou o KTO e o ReFT em 7,4% e 6,0% na taxa de sucesso, respectivamente.
• Avaliação Humana: Participantes humanos avaliaram as explicações geradas, encontrando-as superiores em:
  • Acurácia de Inferência: Terceiros conseguiram prever a decisão real com maior precisão.
  • Raciocínio Lógico: Maior coerência lógica.
  • Ação (Actionability): Maior facilidade em tomar uma decisão baseada na explicação.
  • Carga Cognitiva: Menor esforço mental necessário para entender a explicação.
• Ablação e Robustez:
  • A remoção da camada de atenção cruzada degradou significativamente o desempenho, especialmente na generalização para amostras negativas.
  • O método demonstrou robustez a recompensas ruidosas (ruído artificial nas recompensas finais) e funcionou bem com diferentes backbones de LLM (Llama, Qwen, Gemma).

5. Significado e Impacto

Este trabalho representa um avanço significativo na interseção entre IA Explicável (XAI) e Aprendizado por Reforço.

• Superação do Ruído de Proxy: Ao modelar a distribuição de recompensas em vez de apenas um valor pontual, o método mitiga o problema fundamental de que LLMs proxy não são perfeitos, permitindo recuperar sinais de recompensa que seriam perdidos em abordagens determinísticas.
• Escalabilidade e Generalização: A abordagem oferece um caminho para treinar agentes explicáveis de forma geral, sem necessidade de engenharia específica para cada domínio, facilitando a implantação de sistemas de IA seguros e transparentes em ambientes críticos (saúde, finanças, robótica).
• Fundamentação Teórica: A provisão de limites de erro teóricos para a recuperação da distribuição humana a partir de proxies ruidosos adiciona rigor matemático ao campo de RLAIF, sugerindo que modelos generativos baseados em ODEs (como Fluxos Retificados e Difusão) são ferramentas superiores para alinhamento de preferências complexas.

Em resumo, o artigo demonstra que modelar a incerteza e a pluralidade das recompensas através de Fluxo Retificado é uma estratégia superior para treinar LLMs a gerarem explicações que refletem fielmente a política de decisão de agentes complexos.