Resumo Técnico: Dropout-GRPO para Raciocínio Latente Contínuo

1. Definição do Problema

O artigo aborda uma incompatibilidade estrutural fundamental entre a Otimização de Política Relativa de Grupo (GRPO) e modelos de raciocínio latente contínuo (especificamente a arquitetura COCONUT).

• O Requisito da GRPO: A GRPO depende da diversidade de $K$ rollouts dentro de um grupo para calcular uma vantagem de média de grupo não nula ($A^{(k)} = r^{(k)} - \mu_r$). Se todos os rollouts produzirem trajetórias idênticas, a vantagem colapsa para zero, interrompendo a otimização.
• O Desafio do Raciocínio Latente: Ao contrário dos modelos padrão de Cadeia de Pensamento (CoT) que derivam diversidade através da amostragem estocástica de tokens, os modelos de raciocínio latente contínuo (como o COCONUT) operam via recorrência de estado oculto determinística. Dado um prompt e parâmetros fixos, múltiplos rollouts resultam em trajetórias latentes e respostas idênticas.
• A Consequência: Aplicar a GRPO padrão a esses modelos resulta em variância zero ($\sigma_r = 0$), fazendo com que o gradiente da política desapareça e impedindo a geração de qualquer sinal de aprendizado.

2. Metodologia: Dropout-GRPO

Os autores propõem o Dropout-GRPO, um método que injeta a estocasticidade necessária no processo de raciocínio latente através de um dropout estruturado, mantendo as garantias teóricas da inferência variacional.

Mecanismo Central: Dropout Variacional de Máscara Compartilhada

Em vez do dropout padrão ao nível de token, o método introduz estocasticidade ao nível de parâmetro:

1. Geração de Máscara: Para cada rollout $k$, uma única máscara de Bernoulli $\xi^{(k)}$ é extraída.
2. Consistência: Esta máscara é mantida constante ao longo de todos os $T$ passos de recorrência latente dentro daquele rollout específico.
3. Perturbação: A máscara induz uma perturbação estruturada $\tilde{\theta}(\xi)$ aos pesos da rede. A recorrência latente torna-se $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Replay de Máscara: Crucialmente, a semente do RNG para a máscara é armazenada durante a fase de rollout. Durante a atualização da política, a mesma máscara é regenerada e aplicada. Isso garante que a trajetória avaliada para a recompensa seja bit-idêntica à trajetória usada para o cálculo do gradiente quando $\theta = \theta_{old}$.

Interpretação Teórica

• Média de Modelo Bayesiano: Seguindo Gal e Ghahramani, a máscara compartilhada trata cada rollout como uma única amostra posterior de uma distribuição estruturada sobre os parâmetros. A política marginal $\bar{\pi}_\theta$ é interpretada como uma média de modelo bayesiano.
• Não Viesamento do Gradiente: Os autores provam que o gradiente substituto usando a vantagem de média única ($A^{(k)} = r^{(k)} - \mu_r$) é um estimador não viesado do gradiente da recompensa esperada da política marginal, escalonado por $(1 - 1/K)$.
• Redução de Variância: Ao reutilizar a mesma máscara (Common Random Numbers) durante a atualização como a utilizada durante o rollout, o método reduz a variância em comparação com a reamostragem de uma nova máscara no momento da atualização.

Refinamentos de Implementação

Para estabilizar o treinamento, quatro refinamentos específicos são introduzidos:

1. KL $k_3$ Huberizada: Substitui o clamping padrão no estimador de divergência KL por uma extensão do tipo Huber para evitar que os gradientes desapareçam para tokens com grandes log-razões, evitando "tokens escapados" que derivam sem regularização.
2. Anneamento de Região de Confiança: O coeficiente KL $\beta$ é suavizado proporcionalmente com a taxa de aprendizado $\eta_t$ para evitar que a política fique presa em um equilíbrio onde a tração da KL cancela o gradiente da política.
3. Filtro de Precisão ao Nível do Grupo: Grupos com precisão empírica $\mu_r$ fora de uma janela configurada (ex: $[0.20, 0.95]$) são descartados. Isso atua como um currículo emergente, focando o computo em prompts com variância não trivial.
4. Skip Coordenado por DDP: Garante a consistência do treinamento distribuído ao descartar grupos.

3. Principais Contribuições

1. Novo Paradigma para RL Latente: O artigo identifica a natureza determinística da recorrência latente como a causa da falha da GRPO e propõe o dropout variacional de máscara compartilhada como uma solução mínima e teoricamente limpa. É o primeiro método a aplicar com sucesso o RL relativo de grupo a modelos de raciocínio latente contínuo.
2. Justificativa Teórica:
   • Prova a definição do Jacobiano latente $\partial h_T / \partial \theta$ sob a regularidade padrão de Transformer.
   • Prova o não viés do gradiente substituto para a política marginal.
   • Demonstra que o replay de máscara fornece redução de variância por números aleatórios comuns.
3. Validação Empírica: Aplica com sucesso a GRPO ao modelo COCONUT no benchmark GSM8K, um regime onde a GRPO determinística falha completamente.
4. Implementação de Referência: Libera o código que realiza o replay de máscara bit-idêntico em toda a recorrência latente.

4. Resultados Experimentais

Os experimentos foram conduzidos utilizando o Qwen2.5-1.5B como modelo base com uma profundidade latente de $T=6$.

• Melhoria de Desempenho: O Dropout-GRPO melhorou a precisão GSM8K pass@1 de um baseline COCONUT de 27,29% para 29,01%.
• Recuperação de Capacidades: O método reverteu uma regressão de desempenho de 1,33 pontos no dataset SVAMP causada pelo treinamento de currículo, recuperando o desempenho para 44,00% (correspondendo à variância do modelo base).
• Estudos de Ablação:
  • GRPO Determinística: Falhou completamente, gerando uma trajetória de treinamento plana, idêntica ao SFT inicial devido à vantagem zero.
  • REINFORCE + Dropout (EMA): Falhou devido à alta variância; o gradiente de um único rollout era muito ruidoso para estabilizar o SFT inicial.
  • Dropout-GRPO: A única configuração que produziu um sinal de aprendizado viável, confirmando que as bases relativas de grupo são essenciais para a redução de variância neste cenário.

5. Significância e Alegações

O artigo posiciona este trabalho como uma abordagem prática e teoricamente fundamentada para o pós-treinamento de Large Language Models (LLMs) de raciocínio latente.

• Viabilidade: Demonstra que o pós-treinamento por RL é viável para modelos que raciocinam diretamente em estados ocultos, uma capacidade anteriormente bloqueada pela falta de diversidade de rollouts.
• Mecanismo: Reenquadra o dropout não como um regularizador (que é frequentemente desativado em LLMs modernos), mas como uma fonte de estocasticidade exógena estruturada necessária para a otimização de política variacional.
• Limitações: Os autores são modestos quanto à magnitude da melhoria (um ganho de 1,72 ponto no GSM8K), atribuindo isso à pequena escala do backbone (1.5B) e à escassez de datasets de raciocínio de passos concisos adequados. Eles observam que a política implantada (inferência sem dropout) difere da política marginal otimizada, introduzindo um viés dependente das constantes de Lipschitz.

Em resumo, o Dropout-GRPO preenche a lacavra entre o aprendizado por reforço relativo de grupo e o raciocínio latente contínuo, permitindo a otimização de modelos que anteriormente não podiam utilizar a GRPO devido às suas dinâmicas internas determinísticas.