Kalman Filter Enhanced GRPO for Reinforcement Learning-Based Language Model Reasoning

TL;DR

This paper tackles high-variance policy-gradient updates in language-model reasoning by replacing GRPO's group-mean baseline with a lightweight Kalman filter that online estimates a latent reward baseline $\hat{x}_{i|i}$ and its uncertainty $P_{i|i}$. The resulting KRPO method integrates seamlessly with GRPO, adds no new learned parameters, and uses the Kalman-filtered baseline to compute adaptive advantages $A_i = \frac{r_i - \hat{x}_{i|i}}{\sqrt{P_{i|i}} + \varepsilon}$. Across arithmetic and OpenMath mathematical reasoning benchmarks, KRPO yields higher rewards and faster convergence than GRPO, with results robust to a range of hyperparameters and base models. This uncertainty-aware approach improves training stability and policy optimization for dynamic, noisy reward signals in LLM reasoning tasks, offering a practical alternative to standard baselines.

Abstract

The advantage function is a central concept in RL that helps reduce variance in policy gradient estimates. Recently, for language modeling, Group Relative Policy Optimization (GRPO) was proposed to compute the advantage for each output by subtracting the mean reward, as the baseline, for all outputs in the group. However, it can lead to high variance when the reward advantage is inaccurately estimated. In this work, we propose Kalman Filter Enhanced Group Relative Policy Optimization (KRPO) model, by using lightweight Kalman filtering to dynamically estimate the latent reward baseline and uncertainty. This filtering technique replaces the naive group mean, enabling more adaptive advantage normalization. Our method does not require additional learned parameters over GRPO. This approach offers a simple yet effective way to incorporate group-level uncertainty for advantage estimation, improving policy optimization in settings where highly dynamic reward signals are difficult to model for language models. Through the accuracies and rewards obtained from math question answering and reasoning, we show that using a more adaptive advantage estimation model, KRPO can improve the performance and show more stable return curves upon GRPO. The code is available at https://github.com/billhhh/KRPO_LLMs_RL.

Figures (8)

• Figure 1: The curves of returns/rewards within a batch for different difficulty levels of questions within the Arithmetic Dataset.
• Figure 2: The curves of returns/rewards within a batch for different difficulty levels of questions within OpenMath-Instruct Dataset.
• Figure 3: The training return curves of additional (a) base models and (b) datasets.
• Figure 4: The training return curves of different (a) KL divergence loss weights and (b) process and measure noises.
• Figure 5: The training return curves of Llama3.2-3B-Instruct model on normal level of Arithmetic data.