Resumo Técnico: DISA - Amostragem por Importância Offline para RL de LLM com Correspondência de Distribuição

1. Declaração do Problema

Agentes de raciocínio modernos são cada vez mais avaliados pela sua capacidade de gerar múltiplos caminhos de solução válidos (por exemplo, cadeias de pensamento diversas, planos ou rastros de uso de ferramentas) em vez de uma única resposta correta. Objetivos padrão de Aprendizado por Reforço (RL), como Otimização de Política Próxima (PPO) ou Otimização de Política Relativa por Grupo (GRPO), focam na maximização de recompensa. Embora eficazes para a precisão de resposta única, esses métodos tendem a sofrer de colapso de modo, onde a política concentra a massa de probabilidade em uma única trajetória de alta recompensa, eliminando a diversidade de soluções válidas necessária para uma avaliação robusta (por exemplo, métricas pass@k).

A RL com correspondência de distribuição oferece uma alternativa ao treinar a política para ajustar-se a uma distribuição-alvo onde a probabilidade da trajetória é proporcional à recompensa exponenciada:
$$\tilde{\pi}(o | q) = \frac{1}{Z(q)} \pi_{\text{ref}}(o | q) \exp(\beta \tilde{r}(q, o))$$
onde $Z(q)$ é a função de partição (uma constante de normalização dependente do prompt $q$).

O gargalo central nos métodos existentes de correspondência de distribuição (por exemplo, FlowRL, GFlowNets) é o acoplamento online entre a política ($\pi_\theta$) e o estimador da função de partição ($Z_\phi$). Esses métodos treinam conjuntamente ambos os componentes em rolagens online compartilhadas usando uma perda conjunta (por exemplo, Equilíbrio de Trajetória). Esse acoplamento cria dois problemas críticos:

1. Propagação de Erro: Erros de calibração no estimador da função de partição $Z_\phi$ fluem diretamente para o gradiente da política, distorcendo as atualizações.
2. Falta de Diagnósticos: Não há um sinal de qualidade independente para a função de partição; seu desempenho está emaranhado com a perda da política, tornando impossível diagnosticar se as falhas decorrem da política ou da estimação da função de partição.

2. Metodologia: Ancoragem por Amostragem por Importância Desacoplada (DISA)

Os autores propõem o DISA, um pipeline de três etapas que desacopla a estimação da função de partição do loop de otimização da política. A ideia central é que a função de partição $Z(q)$ depende apenas da política de referência $\pi_{\text{ref}}$, da função de recompensa e da temperatura inversa $\beta$, mas não da política treinável $\pi_\theta$. Portanto, ela pode ser estimada offline antes do início do treinamento de RL.

Etapa 1: Estimação por Amostragem por Importância (IS) Offline

Para um conjunto de prompts offline $\mathcal{D}_Z$, o método desenha $N$ trajetórias de proposta a partir de um modelo de proposta forte $p_T(\cdot|q)$ (por exemplo, um modelo professor maior).

• Pesos de Importância: Para cada trajetória $o_i$, um peso de importância é calculado: $w(q, o_i) = \frac{\pi_{\text{ref}}(o_i|q)}{p_T(o_i|q)} \exp(\beta \tilde{r}(q, o_i))$.
• Agregação no Espaço Logarítmico: Para garantir estabilidade numérica e reduzir viés, a estimativa da função de partição para o prompt $q$ é computada no espaço logarítmico usando um agregador logsumexp:
  $$\log \hat{Z}_{\text{IS}}(q) = \text{logsumexp}_i \left( \log \pi_{\text{ref}}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
  Isso fornece uma estimativa não enviesada de $Z(q)$ na escala linear, com um viés controlado de $O(1/N)$ na escala logarítmica.

Etapa 2: Amortização via Regressão

As estimativas por prompt $\log \hat{Z}_{\text{IS}}(q)$ servem como rótulos para treinar um regressor condicionado ao prompt leve $g_\psi(q)$ (por exemplo, uma pequena MLP).

• O regressor é treinado via regressão de mínimos quadrados para minimizar o erro quadrático entre sua previsão e os rótulos de IS offline.
• Uma vez treinado, os parâmetros do regressor $\psi^*$ são congelados. Isso cria uma estimativa reutilizável e offline da função de partição que generaliza para prompts fora do conjunto de estimação inicial.

Etapa 3: Aprendizado de Política Online Ancorado

Durante a fase online de RL, o regressor congelado $g_{\psi^*}(q)$ substitui a função de partição aprendível $Z_\phi(q)$ na perda de Equilíbrio de Trajetória (TB).

• Objetivo: A política $\pi_\theta$ é otimizada para minimizar o resíduo:
  $$\mathcal{L}_{\text{DISA}}(\theta) = \mathbb{E} \left[ \sum_i \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{\text{ref}}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
• Desacoplamento: Crucialmente, os gradientes são bloqueados de fluir através de $g_{\psi^*}$. A função de partição atua como uma constante fixa, por prompt (uma "âncora"), durante as atualizações da política.

3. Contribuições Principais

1. Identificação do Acoplamento como Obstáculo Central: O artigo identifica o emaranhamento entre a estimação da função de partição e o aprendizado da política como a principal fonte de erros de calibração e cegueira diagnóstica na RL com correspondência de distribuição.
2. O Pipeline DISA: Uma abordagem inovadora de três etapas que move a estimação da função de partição para offline via amostragem por importância, amortiza-a via regressão e a congela para RL online.
3. Garantias Teóricas:
   • Preservação Exata: Se a função de partição offline for exata, o minimizador global do objetivo DISA corresponde à distribuição-alvo teórica inclinada pela recompensa.
   • Estacionariedade com Viés: Mesmo com um viés apenas de prompt (onde o erro depende apenas de $q$ e não da trajetória), a distribuição-alvo permanece um ponto estacionário do fluxo de gradiente on-policy.
   • Perturbação Limitada: Erros residuais de estimação entram no objetivo da política apenas através de um envelope de perturbação limitado (derivado via Cauchy-Schwarz), garantindo estabilidade.
4. Validação Empírica: Avaliação abrangente em benchmarks de matemática e código demonstrando que o DISA preserva os benefícios da correspondência de distribuição enquanto evita as armadilhas do acoplamento online.

4. Resultados Experimentais

Os autores avaliaram o DISA em dois backbones de pesos abertos (Qwen2.5-7B e Qwen3-4B-Base) em seis benchmarks de matemática (incluindo AIME 2024/2025, AMC 2023, MATH-500) e três benchmarks de código (LiveCodeBench, CodeForces, HumanEval+).

• Raciocínio Matemático:
  • O DISA igualou ou superou ligeiramente a linha de base acoplada online FlowRL (por exemplo, 65,5 vs. 64,8 Mean@8 no Qwen3-4B-Base).
  • Superou significativamente as linhas de base de maximização de recompensa GRPO e GSPO (por exemplo, +16,1 pontos sobre o GRPO no Qwen3-4B-Base).
• Geração de Código:
  • No pass@16 (uma métrica sensível à diversidade de soluções), o DISA alcançou as pontuações mais altas entre todos os métodos (por exemplo, 51,5 no Qwen3-4B-Base), superando FlowRL, GSPO e GRPO.
  • No pass@1, o DISA permaneceu competitivo, mostrando nenhuma regressão significativa em comparação com métodos de maximização de recompensa.
• Análise de Diversidade:
  • Uma avaliação com LLM como juiz sobre a diversidade de estratégias de solução (usando uma escala Likert de 1 a 5) mostrou que o DISA reteve a maior fração da diversidade original do backbone (3,72/5,0), comparado ao FlowRL (3,40) e significativamente melhor que GRPO/GSPO (~3,0).
• Estudos de Ablação:
  • Força da Proposta: Usar um modelo de proposta mais fraco (Qwen3-4B vs. 235B) aumentou a variância dos rótulos offline, levando a uma queda de desempenho, mas o regressor ainda convergiu, confirmando o trade-off viés-variância previsto pela teoria.
  • Temperatura Inversa ($\beta$): Estudos de sensibilidade confirmaram que sub-inclinação ($\beta=10$) leva a uma perda leve de sinal, enquanto super-inclinação ($\beta=20$) causa um "abismo de variância" e colapso de modo, validando a análise teórica dos pesos de amostragem por importância.
• Destilação vs. RL: Quando comparado ao ajuste fino (SFT) nas mesmas trajetórias de alta recompensa usadas para a estimação offline do DISA, o DISA forneceu um ganho adicional de até 13,8 pontos Mean@8, demonstrando que ajustar a distribuição suave inclinada pela recompensa recupera informações da paisagem descartadas por alvos SFT rígidos.

5. Significado e Alegações

O artigo alega que o DISA representa uma mudança arquitetônica significativa no pós-treinamento de LLMs para raciocínio. Ao desacoplar a estimação da função de partição do loop da política, o DISA:

• Elimina Erros de Calibração: Impede que erros de estimação da função de partição distorçam os gradientes da política.
• Permite Diagnósticos Independentes: A qualidade da função de partição pode ser monitorada separadamente do treinamento da política através da perda de validação do regressor.
• Preserva a Diversidade: Mantém com sucesso a estrutura multimodal do espaço de soluções, o que é crítico para tarefas de raciocínio complexo onde múltiplos caminhos válidos existem.
• Solidez Teórica: O método fornece garantias formais de que o processo de estimação offline não altera o ponto fixo teórico do objetivo de correspondência de distribuição, desde que o erro de estimação seja limitado.

Os autores concluem que o DISA oferece uma alternativa robusta e teoricamente fundamentada à correspondência de distribuição acoplada online, alcançando desempenho state-of-the-art tanto em precisão quanto em diversidade, sem a instabilidade associada ao treinamento conjunto de funções de partição e políticas.