Resumo Técnico: LaDi-RL

Declaração do Problema

O aprendizado por reforço (RL) tornou-se um paradigma dominante para aprimorar as capacidades de raciocínio de Grandes Modelos de Linguagem (LLMs), particularmente através da otimização de trajetórias de raciocínio. No entanto, os métodos existentes otimizam predominantemente políticas sobre sequências discretas de tokens. Essa abordagem cria uma incompatibilidade fundamental: decisões críticas de raciocínio são frequentemente semânticas, globais e em nível de trajetória, enquanto a otimização em nível de token depende de perturbações locais.

Essa incompatibilidade leva ao colapso da diversidade, um fenômeno onde a política do modelo concentra a massa de probabilidade em um conjunto restrito de padrões de raciocínio, suprimindo estratégias alternativas. Embora a exploração em nível de token possa gerar variantes lexicais (por exemplo, substituir "portanto" por "então"), ela luta para descobrir estratégias semanticamente distintas (por exemplo, mudar do raciocínio algébrico para o geométrico). Consequentemente, modelos ajustados por RL frequentemente melhoram o pass@1 (precisão de amostra única), mas falham em manter ou melhorar o desempenho do pass@k em grandes $k$, indicando uma perda de trajetórias corretas diversas.

Além disso, mover a exploração para espaços latentes contínuos introduz um desafio de atribuição de crédito. No RL em espaço latente, a política atua em um espaço latente contínuo, mas as recompensas são observadas apenas após a trajetória latente ser decodificada em texto. Uma estratégia ingênua de "rollout" (simulação) entrelaça a qualidade do raciocínio latente com a qualidade da decodificação textual. Uma recompensa baixa pode decorrer de uma trajetória latente falha ou meramente de uma realização textual imperfeita de um latente sólido, criando um sinal de aprendizado ruidoso e enganoso.

Metodologia: LaDi-RL

Os autores propõem o Raciocínio com Difusão Latente e Aprendizado por Reforço (LaDi-RL), um framework que reformula a exploração em espaço latente contínuo usando modelos de difusão para prevenir o colapso de entropia.

1. Política de Difusão Latente

Em vez de gerar raciocínio token por token, o LaDi-RL codifica uma trajetória de raciocínio inteira em um bloco latente contínuo e compacto (por exemplo, 64 tokens) usando um Autoencoder Variacional (VAE). A política é parametrizada como um modelo de difusão sobre essas representações de trajetória.

• Mecanismo: Começando a partir de ruído, a política remove ruído iterativamente em direção a uma trajetória de raciocínio latente. Esse processo permite exploração estruturada e modelagem expressiva de distribuições, apoiando naturalmente estratégias de raciocínio multimodais.
• Orientação de Diversidade: Para prevenir ainda mais o colapso, os autores introduzem um mecanismo de orientação baseado em repulsão. Durante o processo de remoção de ruído, uma força repulsiva é aplicada a trajetórias latentes próximas, incentivando-as a se separarem no espaço latente. Isso atua como uma regularização geométrica local, preservando explicitamente a diversidade sem depender exclusivamente da entropia da política.

2. Rollouts Hierárquicos Latente-Texto

Para abordar o desafio de atribuição de crédito, o LaDi-RL emprega uma estratégia de rollout hierárquico $N \times M$:

• Amostragem Latente ($N$): Para uma dada consulta, a política de difusão amostra $N$ trajetórias de raciocínio latente distintas.
• Decodificação de Texto ($M$): Para cada trajetória latente amostrada, a política de texto gera $M$ textos de resposta candidatos.
• Recompensa Marginalizada pelo Decodificador: As recompensas das $M$ conclusões de texto são agregadas (médias) para estimar a utilidade da trajetória latente subjacente. Essa estimativa "marginalizada pelo decodificador" isola a qualidade do raciocínio latente da estocasticidade da decodificação de texto, fornecendo um sinal de recompensa mais limpo e de menor variância para otimizar a política de difusão.

3. Otimização Conjunta

O framework otimiza duas políticas heterogêneas simultaneamente usando Otimização de Política Relativa de Grupo (GRPO):

• Política Latente: Otimizada usando a recompensa média das $M$ textos associados a cada bloco latente. O objetivo incentiva o modelo de difusão a explorar modos de raciocínio diversos e de alta qualidade.
• Política de Texto: Otimizada usando as diferenças relativas entre as $M$ textos gerados a partir do mesmo bloco latente. Isso força a política de texto a encontrar a realização textual ótima para uma dada estratégia de raciocínio, desacoplando-a da exploração de novas estratégias.
• Função de Perda: Uma soma ponderada das perdas latente e textual garante que a política latente explore cadeias de pensamento diversas enquanto a política de texto permanece um gerador fiel e eficaz.

Resultados Principais

Os autores avaliaram o LaDi-RL em benchmarks de geração de código e raciocínio matemático, comparando-o com baselines de RL em nível de token (por exemplo, GRPO Vanilla, GRPO com vantagens de entropia) e outros métodos de raciocínio latente.

• Ganhos de Desempenho: O LaDi-RL alcançou melhorias absolutas de pass@1 de +9,4% na geração de código e +5,7% no raciocínio matemático em comparação com as baselines de RL em nível de token.
• Preservação da Diversidade: Diferentemente do GRPO vanilla, que exibe colapso de diversidade (onde o pass@k cai abaixo do modelo base em grandes $k$), o LaDi-RL consistentemente superou o desempenho pass@k do modelo base em todos os benchmarks. Por exemplo, no AIME 2025, alcançou uma melhoria absoluta de +13,4 no pass@128.
• Eficiência: Ao comprimir o raciocínio em um bloco latente (64 tokens) antes da decodificação, o LaDi-RL reduziu significativamente o número de tokens gerados (por exemplo, de ~3.447 tokens no GRPO padrão para ~422 tokens no LaDi-RL para tarefas de código), resultando em uma redução de 33% no tempo de inferência ponta a ponta.
• Estudos de Ablação:
  • Trajetória vs. Nível de Token: A exploração no espaço de trajetória latente produziu diversidade semântica significativamente maior (0,52 vs. 0,18–0,22 para nível de token) e desempenho mais forte.
  • Difusão vs. VAE: A política baseada em difusão superou um amostrador VAE de passo único, demonstrando a importância da remoção de ruído iterativa para aprender distribuições complexas de raciocínio latente.
  • Rollouts Hierárquicos: A estrutura $N \times M$ foi crítica; aumentar $N$ (diversidade latente) impulsionou a diversidade semântica, enquanto aumentar $M$ (média de texto) forneceu sinais de recompensa mais limpos.

Significado e Alegações

O artigo afirma que o LaDi-RL fornece uma alternativa principiada ao RL em nível de token para melhorar o raciocínio de LLMs. Seu significado principal reside em:

1. Mitigar o Colapso da Diversidade: Ao deslocar a exploração de perturbações locais de tokens para difusão em nível de trajetória global, o método mantém um conjunto diverso de estratégias de solução, impedindo que a política convirja para mínimos locais subótimos.
2. Resolver a Atribuição de Crédito: O rollout hierárquico e a estimativa de recompensa marginalizada pelo decodificador efetivamente desentrelaçam a qualidade do raciocínio do ruído de decodificação, permitindo a otimização estável de políticas latentes.
3. Eficiência e Qualidade: A abordagem alcança desempenho de última geração tanto em precisão (pass@1) quanto em diversidade (pass@k), enquanto reduz significativamente a sobrecarga computacional através da compressão latente.

Os autores concluem que a difusão latente oferece um mecanismo robusto para exploração estruturada, permitindo que modelos descubram estratégias de raciocínio semanticamente distintas que métodos em nível de token perdem, desbloqueando assim limites de raciocínio mais altos sem sacrificar a diversidade de soluções.