Resumo Técnico: Além do Colapso de Modos: Correspondência de Distribuição para Raciocínio Diverso

1. Declaração do Problema

O artigo identifica o colapso de modos como uma limitação fundamental em métodos de aprendizado por reforço (RL) on-policy, particularmente aqueles utilizados para treinar Modelos de Raciocínio de Grande Escala (LRMs) via Aprendizado por Reforço com Recompensas Verificáveis (RLVR). Embora métodos como a Otimização de Política Relativa em Grupo (GRPO) tenham alcançado sucesso na resolução de problemas matemáticos e geração de código, eles exibem uma tendência a concentrar a massa de probabilidade em uma única trajetória de alta recompensa assim que esta é descoberta. Esse comportamento suprime estratégias de solução alternativas, interrompe a exploração e leva a um raciocínio frágil que falha em generalizar.

Os autores argumentam que esse fenômeno não é meramente um artefato de implementação, mas uma consequência direta do objetivo de otimização: o RL on-policy padrão minimiza implicitamente a divergência de Kullback-Leibler (KL) reversa ($D_{KL}(\pi_\theta \| p^*)$). A KL reversa é inerentemente "buscadora de modos", amplificando a primeira trajetória bem-sucedida encontrada em vez de manter uma distribuição sobre múltiplas soluções de alta recompensa diversas. Em contraste, a divergência de KL direta ($D_{KL}(p^* \| \pi_\theta)$) exibe um comportamento "cobridor de modos", incentivando a política a corresponder ao suporte completo da distribuição alvo, mas é tipicamente intratável porque requer amostragem a partir da distribuição global ponderada por recompensa.

2. Metodologia: DMPO

Para abordar o colapso de modos sem exigir amostragem global intratável, os autores propõem a Otimização de Política de Correspondência de Distribuição (DMPO).

Insight Central

O DMPO aproxima a minimização da KL direta no nível de grupo em vez de globalmente. Em vez de amostrar de uma distribuição alvo global, o algoritmo constrói uma distribuição alvo sobre um grupo específico de trajetórias amostradas para um determinado prompt.

Formulação Algorítmica

1. Distribuição Alvo no Nível de Grupo: Para um grupo de $G$ trajetórias $\mathcal{O} = \{o_1, \dots, o_G\}$ amostradas da política atual $\pi_{\theta_{old}}$, uma distribuição alvo $p(o_i | \mathcal{O})$ é construída proporcional ao exponencial de suas recompensas (uma distribuição de Boltzmann):
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   Isso é computacionalmente tratável, pois requer apenas a função de partição no nível de grupo.

2. Distribuição de Política no Nível de Grupo: Para alinhar a política a esse alvo, o método define uma distribuição de política $q_\theta(o_i | \mathcal{O})$ baseada em log-verossimilhanças normalizadas por comprimento para contabilizar o viés do comprimento da trajetória:
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   onde $\varphi(o_i)$ é a log-probabilidade média por token.

3. Função de Perda de Correspondência de Distribuição: A política é alinhada ao alvo minimizando a divergência entre $q_\theta$ e $p$. Os autores utilizam o Erro Quadrático Médio (MSE) em vez da divergência KL para estabilidade numérica, prevenindo a explosão do gradiente quando a probabilidade da política se aproxima de zero:
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. Objetivo Unificado: O objetivo final combina o termo padrão de maximização de recompensa do GRPO com a regularização de correspondência de distribuição:
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   Aqui, $\mathcal{L}_{GRPO}$ impulsiona a exploração (maximização de recompensa), enquanto $\mathcal{L}_{DM}$ impõe a exploração ao impedir que a política colapse para um único modo dentro da região de alta recompensa.

3. Contribuições Principais

A. Algoritmo DMPO

O artigo apresenta o DMPO como uma solução principiante e prática para o colapso de modos. Ao aproximar a KL direta no nível de grupo, ele alcança um comportamento cobridor de modos sem exigir uma função de partição aprendida ou amostragem global. O método é compatível com os frameworks GRPO existentes, exigindo apenas a adição de um único termo de regularização.

B. MM-NP-Bench

Os autores introduzem o MM-NP-Bench, um benchmark multimodal projetado para avaliar rigorosamente a exploração em modelos de raciocínio.

• Escopo: Estende o NP-Bench baseado em texto para incluir representações visuais de 10 tarefas de otimização combinatória NP-difíceis (por exemplo, TSP, Coloração de Grafos, Cobertura de Vértices).
• Avaliação com Métricas Duplas: Para tornar o colapso de modos observável, o benchmark utiliza duas métricas:
  • Taxa de Sucesso (SR): Mede a satisfação de restrições (viabilidade).
  • Razão de Qualidade (QR): Mede a qualidade da otimização relativa a uma linha de base heurística.
  • Significado: O colapso de modos manifesta-se como alta SR mas baixa QR (encontrando soluções válidas mas subótimas).
• Infraestrutura: Inclui geradores paramétricos, verificadores baseados em regras e solucionadores heurísticos para fornecer linhas de base quase ótimas.

C. Validação Empírica

O artigo fornece experimentos extensivos demonstrando que o DMPO supera bases fortes (GRPO, GSPO, GPG, FlowRL, etc.) em múltiplos domínios.

4. Resultados

Desempenho em Tarefas de Otimização

• MM-NP-Bench (Multimodal): O DMPO alcançou uma Razão de Qualidade de 43,1%, uma melhoria relativa de 12% sobre o GRPO (38,4%). Também alcançou a maior Taxa de Sucesso (61,9% vs. 55,7%).
• NP-Bench (Apenas texto): O DMPO alcançou uma Razão de Qualidade de 43,9%, uma melhoria relativa de 9% sobre o GRPO (40,1%).
• Ablação: O método mostrou robustez às escolhas de hiperparâmetros ($\lambda \in [0,3, 0,7]$ e $\alpha \in [0,8, 1,2]$).

Generalização e Transferência

• Raciocínio Matemático: O treinamento em tarefas de otimização com DMPO melhorou o desempenho em benchmarks de raciocínio matemático (AIME, AMC, etc.) em 2,0% em comparação com o GRPO.
• Tarefas Fora de Domínio: O DMPO demonstrou uma melhoria de 2,3% em benchmarks de raciocínio visual fora de domínio (LogicVista, MathVista, etc.), sugerindo que o treinamento que preserva a diversidade aprimora as capacidades de raciocínio geral.

Dinâmica de Treinamento

A análise dos logs de treinamento revelou que, enquanto o GRPO sofre de "colapso de comprimento" (degenerando em respostas curtas e triviais) e platôs de recompensa prematuros, o DMPO sustenta a melhoria contínua e preserva cadeias de raciocínio robustas (~400 tokens). O DMPO mantém uma entropia menor que o GRPO, indicando "diversidade estruturada" (execução confiante de diferentes estratégias lógicas) em vez de confusão estocástica.

5. Significado e Alegações

O artigo afirma estabelecer a correspondência de distribuição como uma abordagem prática e principiante para prevenir o colapso de modos no RL on-policy. Os autores argumentam que:

1. Identificação da Causa Raiz: O colapso de modos no RLVR é fundamentalmente causado pela natureza buscadora de modos da minimização da KL reversa.
2. Solução Viável: A minimização da KL direta, tradicionalmente considerada intratável, pode ser efetivamente aproximada no nível de grupo para preservar a diversidade de soluções.
3. Benefício Geral: Preservar a diversidade durante o treinamento em tarefas específicas de otimização transfere-se para capacidades de raciocínio geral aprimoradas, sugerindo que a exploração é um componente crítico do raciocínio robusto em Modelos de Raciocínio de Grande Escala.

O trabalho posiciona o DMPO não meramente como um regularizador de diversidade, mas como um mecanismo que permite a descoberta de soluções de maior qualidade que objetivos buscadores de modos perdem, aprimorando assim tanto o desempenho de otimização quanto a robustez do raciocínio através de modalidades.