Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning

O artigo apresenta o CAPO, um novo algoritmo de otimização de política que utiliza informações de curvatura para identificar e filtrar amostras instáveis durante o treinamento por reforço em LLMs, garantindo atualizações estáveis e alcançando uma eficiência de amostragem até 30 vezes superior à do GRPO padrão em tarefas de raciocínio matemático.

O essencial

Imagine que você está ensinando um gênio da matemática (um Modelo de Linguagem Grande, ou LLM) a resolver problemas complexos. Para isso, você usa um método chamado "Reinforcement Learning" (Aprendizado por Reforço). Pense nisso como um sistema de recompensas: o gênio tenta resolver um problema, e se acertar, ganha um ponto. Se errar, não ganha nada. O objetivo é fazer o gênio aprender a acertar cada vez mais, ajustando sua "mente" (os parâmetros do modelo) a cada tentativa.

O problema é que, às vezes, o gênio tenta mudar de ideia de forma tão brusca e descontrolada que ele esquece tudo o que sabia antes. É como se, ao tentar aprender a andar de bicicleta, ele decidisse pular da bicicleta e correr em direção a um muro porque achou que seria mais rápido. Isso é chamado de "colapso da política" (policy collapse). Para evitar isso, os pesquisadores atuais são muito cautelosos: eles dão passos minúsculos e usam muitos exemplos de treino, o que torna o processo lento e caro.

Este artigo, apresentado na conferência ICLR 2026, propõe uma solução inteligente chamada CAPO (Otimização de Política Consciente de Curvatura).

A Analogia do Carro e a Estrada Sinuosa

Para entender o CAPO, imagine que treinar esse modelo de IA é como dirigir um carro de Fórmula 1 em uma estrada de montanha cheia de curvas perigosas.

1. O Problema Atual (GRPO): Os métodos atuais olham apenas para a estrada logo à frente (o gradiente). Se o carro está indo rápido, eles tentam manter a velocidade, mas têm medo de fazer uma curva muito fechada. Para não capotar, eles dirigem muito devagar e com muito cuidado, gastando muita gasolina (dados de treino) para chegar ao destino.
2. O Perigo: Se o motorista tentar acelerar demais para ser mais eficiente (usar menos dados), ele pode entrar em uma curva muito fechada e capotar (o modelo "quebra" e para de aprender).
3. A Solução CAPO: O CAPO é como ter um GPS de alta tecnologia que não só vê a estrada à frente, mas também sente a curvatura do terreno. Ele sabe exatamente onde a estrada está ficando íngreme ou onde uma curva é perigosa antes mesmo de o carro chegar lá.

Como o CAPO Funciona (Passo a Passo)

O CAPO faz três coisas principais para manter o carro seguro e rápido:

1. Sentindo a Curvatura (Geometria de Segunda Ordem):
   Em vez de apenas olhar para a direção do carro, o CAPO analisa a "forma" da estrada. Ele calcula matematicamente se a próxima curva é suave ou se é um precipício. Ele usa uma técnica chamada "última camada" para fazer esse cálculo de forma rápida, sem precisar de supercomputadores gigantes. É como ter um sensor que diz: "Ei, essa curva aqui é perigosa, não entre nela em alta velocidade".

2. O Filtro de Segurança (Seleção de Dados):
   Durante o treino, o modelo gera milhares de tentativas de resposta. O CAPO olha para cada uma dessas tentativas antes de usá-las para ensinar o modelo.

   • Se uma tentativa parece que vai fazer o modelo dar um "pulo" perigoso (uma atualização instável), o CAPO diz: "Não, essa não!" e a descarta.
   • Se a tentativa parece segura e útil, ele diz: "Pode entrar!".
     É como um professor que, ao corrigir os exercícios de um aluno, decide não usar os exemplos que confundem o aluno, mas sim os que ajudam a entender o conceito sem causar confusão.

3. Aceleração Segura:
   Como o CAPO remove os exemplos perigosos, ele permite que o modelo use uma "velocidade" maior (uma taxa de aprendizado mais alta) sem medo de capotar. Isso significa que o modelo aprende muito mais rápido com menos exemplos.

Os Resultados Mágicos

Os autores testaram essa ideia em problemas de matemática (o conjunto de dados MATH) e descobriram coisas impressionantes:

• Eficiência Extrema: O CAPO foi capaz de aprender 30 vezes mais rápido do que os métodos tradicionais quando usados em modo "agressivo". É como se o carro de Fórmula 1 chegasse ao destino em 30 minutos em vez de 15 horas.
• Estabilidade: Enquanto os outros métodos capotavam (perdiam a capacidade de aprender) quando tentavam ir rápido, o CAPO manteve a estabilidade o tempo todo.
• Intervenção Mínima: O CAPO é tão inteligente que rejeita menos de 8% das tentativas de resposta. Ou seja, ele não precisa jogar fora quase nada; apenas remove os poucos exemplos que realmente poderiam causar problemas.

Resumo em uma Frase

O CAPO é como um copiloto de segurança para a inteligência artificial: ele olha para o futuro, identifica onde o caminho pode ficar perigoso, remove os obstáculos antes que o carro entre neles e permite que o carro corra mais rápido e mais longe, gastando menos combustível.

Isso é crucial porque treinar essas IAs hoje em dia é extremamente caro e demorado. Se conseguirmos torná-las mais eficientes, poderemos ter IAs mais inteligentes e capazes de resolver problemas complexos (como medicina ou engenharia) de forma mais acessível.

Resumo técnico

Título: Estabilização de Gradientes de Política para Aprendizado por Reforço Eficiente em Amostra no Raciocínio de LLMs

1. O Problema

O Aprendizado por Reforço (RL), particularmente através de métodos de gradiente de política (como PPO e GRPO), tem sido fundamental para habilitar capacidades de raciocínio em Modelos de Linguagem de Grande Escala (LLMs). No entanto, a otimização desses modelos enfrenta desafios críticos:

• Instabilidade de Otimização: A natureza não estacionária do objetivo do RL e a alta variância das estimativas de gradiente levam a atualizações catastróficas e ao colapso da política (onde o modelo perde a capacidade de aprender e degrada seu desempenho).
• Ineficiência de Amostra: Para mitigar a instabilidade, as implementações atuais recorrem a hiperparâmetros conservadores (taxas de aprendizado muito baixas e tamanhos de lote grandes). Isso aumenta drasticamente o número de gerações de LLM necessárias para o treinamento, elevando os custos computacionais.
• Falta de Análise de Segunda Ordem: A dinâmica de otimização em LLMs é pouco estudada sob a ótica da geometria de segunda ordem (curvatura), dificultando a previsão de comportamentos instáveis antes que ocorram.

O objetivo do trabalho é desenvolver um método que permita regimes de treinamento mais agressivos (mais eficientes em amostras) sem sacrificar a estabilidade, identificando e mitigando as causas raiz das atualizações instáveis.

2. Metodologia: CAPO (Curvature-Aware Policy Optimization)

Os autores propõem o CAPO, um framework computacional que modela e utiliza informações de curvatura (geometria de segunda ordem) para estabilizar o treinamento.

A. Modelagem da Paisagem de Otimização (Geometria de Segunda Ordem)
Em vez de calcular Hessianos ou Matrizes de Informação de Fisher (FIM) completos (que são intratáveis para modelos com bilhões de parâmetros), o CAPO utiliza uma aproximação baseada em duas ideias principais:

1. Modelo de Última Camada (Last-Layer Model): O modelo assume que a curvatura relevante para a estabilidade reside principalmente na camada de saída (logits) do LLM. Isso reduz a complexidade de modelar a curvatura de todo o modelo para apenas os pesos da última camada.
2. Aproximação de Deslocamentos Direcionais: O método não materializa as matrizes completas. Em vez disso, estima os deslocamentos direcionais no objetivo ($m_H$) e na distribuição da política ($m_F$) usando produtos internos esparsos.
   • $m_H(\Delta\theta)$: Estima a mudança no objetivo (Hessiano).
   • $m_F(\Delta\theta)$: Estima a mudança na distribuição da política (FIM), correlacionada com a divergência KL.

B. Mecanismo de Intervenção via Seleção de Dados
O CAPO atua como um mecanismo de "rejeição de amostras" (rejection sampling) em nível de token:

1. Para um lote de trajetórias, o algoritmo calcula o passo de atualização proposto ($\Delta\psi$) baseado no modelo de última camada.
2. Avalia os deslocamentos $m_H$ e $m_F$ para subconjuntos de dados (tokens).
3. Critério de Aceitação: Um subconjunto é aceito apenas se satisfizer restrições de "região de confiança" (trust-region):
   • $m_H$ deve estar dentro de um intervalo aceitável (garantindo melhoria no objetivo).
   • $m_F$ deve ser pequeno (garantindo que a mudança na política não seja abrupta).
4. Mascaramento: Tokens ou subconjuntos que violam essas condições (indicando curvatura extrema ou mudança brusca na política) são mascarados (rejeitados) e não contribuem para o cálculo do gradiente real do LLM.

C. Fundamentação Teórica
Os autores provam teoreticamente que, sob suposições realistas (limitação da norma do Hessiano e do passo), o CAPO garante melhoria monótona da política. Ao rejeitar amostras que causam grandes deslocamentos na distribuição da política, o algoritmo evita o colapso, mesmo com taxas de aprendizado altas.

3. Resultados Experimentais

Os experimentos foram conduzidos no modelo Qwen2.5-Math-7B em tarefas de raciocínio matemático (dataset MATH e benchmarks como GSM8K, GPQA, etc.).

• Estabilidade sob Regimes Agressivos: Enquanto métodos baselines (GRPO padrão, GRPO agressivo, REINFORCE) sofreram colapso de política quando submetidos a taxas de aprendizado 5x maiores e tamanhos de lote 12x menores, o CAPO manteve a estabilidade e continuou a melhorar.
• Eficiência de Amostra: O CAPO alcançou uma melhoria de até 30x na eficiência de amostra em comparação com o GRPO padrão (conservador). Isso significa que o CAPO atingiu o mesmo nível de precisão com 30 vezes menos gerações de trajetórias.
• Intervenção Mínima: O mecanismo de rejeição é altamente seletivo. Em média, menos de 8% dos tokens são rejeitados durante o treinamento, com a taxa caindo para abaixo de 2% após as fases iniciais.
• Custo Computacional: A sobrecarga computacional do CAPO é mínima (< 3% do tempo total de iteração de aprendizado), pois evita o cálculo de Hessianos completos e opera em representações esparsas.
• Generalização: A técnica foi aplicada com sucesso a outros algoritmos (DrGRPO, REINFORCE), demonstrando que o mecanismo de seleção baseado em curvatura é agnóstico ao algoritmo de otimização subjacente.

4. Contribuições Principais

1. Formalização do Problema: Estabelecimento de uma formulação explícita do problema de otimização de RL em LLMs considerando termos de segunda ordem (Hessiano e FIM).
2. Framework Computacional Tractável: Desenvolvimento de um modelo de última camada que permite estimar curvaturas e deslocamentos de política de forma escalável para modelos de bilhões de parâmetros, sem custos proibitivos de memória.
3. Algoritmo CAPO: Proposta de um método de seleção de dados (mascaramento de tokens) que intervém proativamente no processo de otimização, rejeitando amostras que levariam a atualizações instáveis.
4. Garantias Teóricas e Empíricas: Prova de melhoria monótona sob CAPO e validação empírica de que o método permite treinar LLMs de raciocínio com hiperparâmetros agressivos, superando o trade-off tradicional entre estabilidade e eficiência.

5. Significado e Impacto

Este trabalho é significativo porque aborda um gargalo fundamental no treinamento de LLMs: a ineficiência causada pela necessidade de hiperparâmetros conservadores. Ao demonstrar que é possível estabilizar o treinamento através da compreensão da geometria da otimização (curvatura) e não apenas através de heurísticas de clipping ou regularização KL, o CAPO abre caminho para:

• Treinamento mais rápido e barato: Redução drástica no custo computacional de pós-treinamento (post-training) de LLMs.
• Escalabilidade: Possibilidade de treinar modelos maiores e mais complexos em regimes de aprendizado que antes eram considerados instáveis.
• Mudança de Paradigma: Sugere que a estabilidade em RL para LLMs pode ser alcançada através de seleção inteligente de dados baseada em modelos de curvatura, em vez de apenas restringir o tamanho do passo de atualização globalmente.

Em resumo, o CAPO oferece uma solução elegante e eficiente para o problema de instabilidade em RL para LLMs, permitindo que a comunidade explore regimes de treinamento mais agressivos e eficientes sem sacrificar a robustez do modelo.