Online Causal Kalman Filtering for Stable and Effective Policy Optimization

Título: Filtragem de Kalman Causal Online para Otimização de Políticas Estável e Eficiente (KPO)

1. O Problema

O artigo aborda um desafio crítico no aprendizado por reforço (RL) para Grandes Modelos de Linguagem (LLMs), especificamente em métodos de otimização de política em larga escala como o GRPO (Group Relative Policy Optimization).

• Alta Variância nas Razões de Amostragem por Importância (IS): Em configurações off-policy (onde a política atual difere da política de coleta de dados), a razão de amostragem por importância (IS) calculada token a token sofre de alta variância.
• Inconsistência Estrutural Local: A pesquisa empírica revela que, embora se espere que desvios off-policy sejam homogêneos localmente (dentro de um segmento de raciocínio), na prática, os tokens adjacentes exibem uma inconsistência estrutural. As razões IS variam bruscamente e frequentemente entre tokens vizinhos, criando "picos de ruído" e alternâncias rápidas (Up/Down/On).
• Consequências: Essa variabilidade local distorce as atualizações do gradiente da política entre tokens adjacentes, levando a uma instabilidade no treinamento, colapso de entropia e, em casos extremos, ao colapso total do treinamento (training collapse).
• Limitações das Soluções Atuais: Métodos recentes tentam estabilizar o treinamento usando uma razão IS fixa em nível de sequência (ignorando a estrutura local) ou ajustando tokens individualmente sem considerar a deriva temporal. Ambas as abordagens falham em capturar a coerência local necessária para uma otimização estável.

2. Metodologia: KPO (Kalman Policy Optimization)

Os autores propõem o KPO, um método que utiliza Filtragem de Kalman Causal Online para suavizar as razões IS token a token, preservando a estrutura temporal local.

• Perspectiva Temporal: O método trata a sequência de razões IS como uma série temporal observada com ruído. A ideia central é que o desvio off-policy deve evoluir gradualmente ao longo de tokens adjacentes (coerência local), e não saltar aleatoriamente.
• Modelo de Espaço de Estados:
  • Estado Latente ($\rho_t$): Representa a razão IS "desejada" e suavizada.
  • Observação ($z_t$): É o logaritmo da razão IS bruta calculada ($\log r_t$).
  • Equações: O modelo assume um processo de caminhada aleatória para o estado latente (com ruído de processo $Q$) e uma observação ruidosa (com ruído de observação $V$).
• Filtragem Causal (Autoregressiva):
  • O filtro de Kalman opera causalmente, utilizando apenas informações de tokens passados e do token atual para estimar o estado do token atual. Isso é crucial para manter a consistência com a geração autoregressiva de LLMs.
  • O filtro alterna entre Previsão (baseada no estado anterior) e Atualização (corrigindo a previsão com a nova observação usando o Ganho de Kalman, $K_t$).
  • O Ganho de Kalman adapta-se dinamicamente: se a incerteza do modelo é alta ou o ruído de observação é baixo, o filtro confia mais na nova medição; caso contrário, ele suaviza o ruído, mantendo a tendência temporal.
• Objetivo de Otimização: As razões IS filtradas ($\hat{r}_t = \exp(\hat{\rho}_t)$) substituem as razões brutas na função de perda do GRPO (com ou sem clipping), permitindo atualizações de política que são tanto numericamente suavizadas quanto estruturalmente coerentes.

3. Principais Contribuições

1. Descoberta Empírica de Padrões Locais: Os autores identificam e quantificam a inconsistência estrutural local nas razões IS token a token (alta frequência de troca de estados e curta duração de "runs"), demonstrando que isso é a causa raiz da instabilidade em otimizações off-policy.
2. Método KPO: Introdução de um novo método de otimização que aplica filtragem de Kalman causal para suavizar flutuações locais enquanto preserva a estrutura coerente dentro da sequência.
3. Resultados Robustos: Demonstração de que o KPO supera consistentemente os métodos state-of-the-art (como GRPO, GSPO e GMPO) em tarefas de raciocínio matemático complexo, melhorando tanto a estabilidade do treinamento quanto a performance final.

4. Resultados Experimentais

Os experimentos foram realizados em seis benchmarks desafiadores de raciocínio matemático (AIME'24, AIME'25, AMC'23, MATH500, Minerva, OlympiadBench) utilizando o modelo Qwen3-4B.

• Desempenho Superior: O KPO (especialmente a versão com clipping) alcançou os melhores resultados em avg@16 (precisão média) em 5 dos 6 benchmarks e em pass@16 (probabilidade de pelo menos uma resposta correta) em 3 benchmarks.
  • Exemplo: No AIME'24, o KPO-clipped atingiu 37.91 (avg@16) e 63.33 (pass@16), superando o GSPO (32.70 / 60.00).
• Estabilidade no Treinamento:
  • Recompensa: Enquanto o GRPO diverge após ~200 passos e o GSPO/GMPO estagnam cedo, o KPO mostra melhoria contínua e estável.
  • Entropia: O KPO previne o colapso de entropia (comportamento determinístico prematuro), mantendo a exploração.
  • Gradiente de Política: O KPO apresenta perda de gradiente com baixa variância, ao contrário das oscilações ruidosas do GRPO.
• Análise de Distribuição (Pré vs. Pós-Filtro):
  • Comprimento de Execução (Run-length): Aumentou drasticamente (de ~1.6 tokens para ~120-135 tokens para estados "Up/Down"), indicando que o filtro cria blocos contíguos e coerentes.
  • Frequência de Troca: Caiu de 0.43 (muito frequente) para 0.01 (muito raro), confirmando a remoção de alternâncias espúrias.
  • Variância: A variância global e local das razões IS foi reduzida para quase zero, eliminando o ruído de alta frequência.

5. Significado e Impacto

• Novo Paradigma Temporal: O trabalho introduz uma perspectiva temporal para a otimização de políticas em LLMs, tratando a estimativa de razões IS como um problema de inferência de estado em espaço de estados online, em vez de apenas uma agregação estatística.
• Equilíbrio entre Suavização e Estrutura: Diferente de métodos que suavizam excessivamente (perdendo nuances locais) ou não suavizam (mantendo ruído), o KPO encontra um equilíbrio ideal, preservando a heterogeneidade global entre segmentos de raciocínio enquanto garante coerência local.
• Aplicabilidade: O método é leve, compatível com pipelines existentes de RLHF/RL (como GRPO) e particulamente eficaz em cenários de raciocínio de longo horizonte (long-chain-of-thought), onde a estabilidade é crítica.
• Limitação Futura: O processo de filtragem é inerentemente sequencial (autoregressivo), o que dificulta a paralelização completa durante o treinamento, um ponto que os autores sugerem para pesquisas futuras.

Em resumo, o KPO resolve um problema fundamental de instabilidade no RL de LLMs ao aplicar princípios de controle de sinal (Filtro de Kalman) para "limpar" o ruído nas atualizações de política, resultando em modelos mais robustos e capazes de raciocínio complexo.