DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

O artigo propõe o DyJR, um framework de regularização que utiliza um buffer dinâmico e a divergência Jensen-Shannon para preservar a diversidade em algoritmos de aprendizado por reforço on-policy, superando a ineficiência amostral do GRPO e evitando o colapso de modos sem comprometer a eficiência computacional.

O essencial

Imagine que você está ensinando um aluno muito inteligente (o Modelo de Linguagem) a resolver problemas de matemática complexos ou a escrever códigos de banco de dados. O método tradicional de ensino, chamado RL (Aprendizado por Reforço), funciona assim: o aluno tenta resolver um problema, acerta ou erra, recebe uma nota, e o professor ajusta a estratégia dele.

O problema é que, no método atual (chamado GRPO), assim que o aluno dá a resposta e recebe a nota, o professor joga a resposta no lixo. Ele esquece imediatamente o que foi feito e pede para o aluno tentar de novo do zero. Isso é muito desperdício de tempo e energia!

Alguns métodos anteriores tentaram salvar essas respostas antigas para reutilizá-las. Mas eles cometiam um erro grave: eles tratavam as respostas antigas como "verdades absolutas" e forçavam o aluno a decorar exatamente aquelas soluções. O resultado? O aluno perdia a criatividade, deixava de explorar novas ideias e ficava preso em apenas uma maneira de resolver as coisas (o que os cientistas chamam de "colapso de modo").

É aqui que entra o DyJR (o método proposto neste artigo). Vamos entender como ele funciona com uma analogia simples:

1. O Aluno e o Caderno de Anotações (O Buffer Dinâmico)

Imagine que o professor tem um caderno de anotações para guardar as tentativas do aluno.

• O problema dos métodos antigos: Eles guardavam tudo o que o aluno já fez, desde o primeiro dia de aula até hoje. Mas, como o aluno evolui rápido, as anotações de quando ele era iniciante (cheias de erros e tentativas variadas) não servem mais para o aluno avançado. Guardar tudo só ocupa espaço e confunde o aluno.
• A solução do DyJR: O professor usa um caderno inteligente com uma regra de "FIFO" (Primeiro a Entrar, Primeiro a Sair). Ele guarda apenas as tentativas mais recentes e relevantes.
  • A mágica do "Aquecimento": No início do curso, quando o aluno está descobrindo muitas formas diferentes de pensar (alta diversidade), o professor enche o caderno rapidamente para capturar essa criatividade.
  • A estabilização: Conforme o aluno fica mais experiente e o caderno fica cheio, ele começa a apagar as anotações mais antigas para dar lugar às novas. Isso garante que o aluno sempre tenha acesso a exemplos que fazem sentido para o nível atual dele, sem se perder no passado.

2. A Regra de Ouro: Diversidade vs. Precisão (A Regularização JS)

Aqui está o segredo principal.

• O erro comum: A maioria dos métodos diz: "Olhe para essa resposta antiga que estava certa. Agora, faça exatamente igual a ela." Isso faz o aluno ficar robótico e perder a capacidade de pensar fora da caixa.
• A abordagem do DyJR: O professor diz: "Olhe para todas as tentativas antigas que deram certo. Não tente copiar uma delas especificamente. Em vez disso, não se afaste muito do conjunto de todas essas boas tentativas."
  • Eles usam uma ferramenta matemática chamada Divergência de Jensen-Shannon (pense nela como um "medidor de distância" entre o que o aluno está fazendo agora e o que ele já fez de bom no passado).
  • O objetivo não é forçar o aluno a ser igual a um exemplo específico, mas garantir que ele continue explorando várias caminhos diferentes, mantendo a "diversidade" das soluções. É como dizer: "Não fique preso em um único caminho, mantenha várias portas abertas."

3. Por que isso é melhor? (Os Resultados)

O artigo mostra que, ao usar esse método:

• O aluno não fica entediado: Ele continua explorando novas ideias em vez de decorar uma única fórmula.
• Economia de energia: O sistema não precisa guardar milhões de exemplos antigos (o que exigiria computadores gigantescos). Ele guarda apenas o essencial e recente.
• Melhores notas: Em testes de matemática e criação de códigos (SQL), o método DyJR superou os métodos anteriores, alcançando resultados mais altos e consistentes.

Resumo em uma frase

O DyJR é como um professor sábio que não apenas guarda as respostas certas do aluno, mas organiza esse arquivo de forma inteligente (apenas o recente e relevante) e usa essas memórias para garantir que o aluno continue sendo criativo e não perca a capacidade de pensar em várias soluções diferentes ao mesmo tempo.

Em vez de forçar o aluno a ser uma "máquina de repetir", o DyJR o ajuda a ser um "explorador inteligente" que aprende com o passado sem ficar preso nele.

Resumo técnico

Resumo Técnico: DyJR

1. O Problema

O artigo aborda um gargalo crítico no Aprendizado por Reforço com Recompensas Verificáveis (RLVR) para Grandes Modelos de Linguagem (LLMs), especialmente em tarefas de raciocínio complexo (como matemática e SQL).

• Ineficiência de Amostragem: Algoritmos on-policy atuais, como o GRPO (Group Relative Policy Optimization), descartam os dados de "rollout" (trajetórias geradas) após uma única atualização. Isso desperdiça recursos computacionais valiosos e impede o aprendizado contínuo a partir de sucessos passados.
• Falhas dos Métodos de Replay Existentes: Métodos de Experience Replay (Replay de Experiência) anteriores tentam reutilizar dados históricos, mas frequentemente cometem dois erros fundamentais:
  1. Colapso de Modo (Mode Collapse): Ao atualizar diretamente a política com gradientes de dados históricos, o modelo tende a superajustar-se (overfitting) a caminhos de solução específicos, perdendo sua capacidade de exploração e diversidade.
  2. Custo Computacional Elevado: Manter grandes buffers de dados históricos consome muita memória de GPU e recursos de treinamento.
• Hipótese Central: Os autores argumentam que o valor principal dos dados históricos não é reforçar a precisão (já que o modelo atual é geralmente melhor), mas sim sustentar a diversidade de estratégias de raciocínio para evitar que o modelo converja prematuramente para um único caminho de solução.

2. Metodologia: DyJR

Os autores propõem o DyJR (Dynamic Jensen-Shannon Replay), um framework de regularização que redefine como os dados de replay são construídos e utilizados.

A. Construção de Dados: Buffer Dinâmico Sensível ao Tempo
Em vez de armazenar tudo indiscriminadamente, o DyJR utiliza uma estratégia não uniforme:

• Buffer Dinâmico (FIFO): O buffer mantém apenas amostras temporariamente próximas ao modelo atual (usando um protocolo First-In-First-Out com uma idade máxima $M$). Isso sincroniza o buffer com a evolução do modelo.
• Seleção Adaptativa de Dados:
  • Fase de Aquecimento (Early Stage): Durante os primeiros passos de treinamento (quando a entropia do modelo cai rapidamente), o buffer expande temporariamente sua capacidade para capturar padrões de alta entropia e evitar o colapso precoce.
  • Critério de Confiança: O sistema prioriza amostras com alta confiança (múltiplos caminhos corretos para a mesma pergunta), mas relaxa os critérios para tarefas difíceis, garantindo que soluções raras não sejam perdidas.

B. Utilização de Dados: Regularização por Divergência de Jensen-Shannon (JS)
O DyJR abandona a atualização direta de gradiente sobre os dados de replay (que causa overfitting) e introduz uma restrição de distribuição:

• Regularização JS: Em vez de maximizar a verossimilhança dos dados históricos, o algoritmo minimiza a Divergência de Jensen-Shannon ($D_{JS}$) entre a política atual e uma distribuição de referência construída a partir da mistura de políticas históricas no buffer.
• Vantagem da JS: Diferente da Divergência KL (que é assimétrica e pode forçar o modelo a cobrir modos de baixa probabilidade de forma agressiva), a $D_{JS}$ é simétrica e limitada. Isso atua como uma "âncora" flexível, impedindo que o modelo se afaste drasticamente de caminhos de sucesso diversos sem alterar agressivamente a direção de otimização.

C. Objetivo de Otimização
A função de perda total combina o objetivo padrão do GRPO (on-policy) com o termo de regularização JS:
$$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$$
Onde $\alpha_{JS}$ é um hiperparâmetro que controla a força da regularização.

3. Contribuições Principais

1. Redefinição do Paradigma de Replay: Mudança de foco da "otimização de precisão" para a "regularização de diversidade". O replay serve para manter a capacidade exploratória, não apenas para repetir acertos.
2. Estratégia de Construção de Dados Dinâmica: Introdução de um buffer que expande e contrai baseado na fase de treinamento e na proximidade temporal, otimizando o uso de memória e capturando padrões de alta entropia iniciais.
3. Regularização Eficiente via JS: Substituição de atualizações de gradiente diretas por uma restrição de divergência, prevenindo o colapso de modo com um custo computacional mínimo.
4. Análise Granular: Fornecimento de uma análise detalhada da evolução da probabilidade de tokens (Rank-k), demonstrando como o DyJR mantém a diversidade de tokens de alta e baixa probabilidade durante o treinamento.

4. Resultados Experimentais

Os experimentos foram conduzidos em tarefas de raciocínio matemático (usando Qwen3-4B) e geração de SQL (usando Llama-3.1-8B).

• Desempenho em Matemática: O DyJR alcançou uma precisão média de 34.1% em 6 benchmarks matemáticos, superando significativamente o GRPO base (29.8%) e outros métodos de replay como RLEP, Ex-GRPO e DPH-RL.
  • Melhorias notáveis em benchmarks difíceis (ex: +7.4% no AMC23 e +2.9% no HMMT25).
• Desempenho em SQL: No dataset BIRD, o DyJR superou o GRPO em +3.3% (Pass@1) e +4.5% (Pass@16). No Spider, melhorou +5.0% (Pass@1) e +7.0% (Pass@16).
• Eficiência de Memória: O método requer armazenar apenas cerca de 2k pares de dados (vs. 28k em métodos como RLEP), com sobrecarga de memória de GPU insignificante (<1GB).
• Análise de Diversidade: Gráficos de evolução mostram que, enquanto o GRPO colapsa rapidamente para uma probabilidade de Rank-1 próxima a 1.0 (perdendo diversidade), o DyJR mantém uma distribuição mais equilibrada entre Rank-1, Rank-2 e Rank-3, permitindo exploração contínua.
• Escalabilidade: O DyJR demonstra melhor escalabilidade em métricas Pass@k (até k=1024), indicando que o modelo aprende a gerar múltiplas soluções corretas, não apenas uma.

5. Significado e Impacto

O artigo oferece uma solução prática e eficiente para um dos maiores desafios do RL em LLMs: o equilíbrio entre exploração e exploração (exploration-exploitation trade-off).

• Viabilidade de Escala: Ao reduzir drasticamente a necessidade de armazenamento de dados históricos e evitar o colapso de modo, o DyJR torna o RLVR mais escalável para tarefas complexas de raciocínio.
• Mudança de Perspectiva: O trabalho sugere que a "memória" em RL para LLMs não deve ser um repositório estático de acertos, mas um mecanismo dinâmico para preservar a diversidade de pensamento, especialmente nas fases iniciais de treinamento.
• Aplicabilidade Geral: A abordagem é validada em diferentes arquiteturas (Qwen, Llama) e domínios (Matemática, SQL), sugerindo que é uma técnica robusta para melhorar a generalização de modelos de linguagem.

Em suma, o DyJR demonstra que a diversidade de dados históricos é mais valiosa do que a precisão bruta para o treinamento de RL, e que essa diversidade pode ser preservada de forma eficiente através de uma regularização baseada em divergência de Jensen-Shannon e buffers dinâmicos.