PaperlayerView exact revision on arXiv

Source-linked AI summary

Online Causal Kalman Filtering for Stable and Effective Policy Optimization

Shuo He, Lang Feng, Xin Cheng, Lei Feng, Bo An

arXiv:2602.10609v2cs.CLcs.AI

TL;DR

Large-scale LLM reinforcement learning is destabilized by noisy token-level importance-sampling ratios, while common smoothing strategies overlook temporal structure within sequences. KPO models the desired ratio as an evolving latent state and estimates it online with a causal Kalman filter. Across challenging math-reasoning benchmarks, KPO reports superior performance and improved training stability, while preserving local structure and suppressing noise spikes.

  • Problem

    Token-level IS ratios can exhibit structurally inconsistent local off-policy deviation, potentially distorting adjacent-token policy-gradient updates and destabilizing optimization.

  • Method

    KPO models the desired token-wise IS ratio as a latent state evolving across tokens and applies an online causal Kalman filter using past and current observations.

  • Results

    KPO consistently improves performance over strong baselines on challenging math-reasoning benchmarks while enhancing training stability.

  • Takeaways & Limitations

    KPO preserves local structure-aware variation while smoothing noisy IS-ratio spikes for more stable and effective policy updates.

Abstract

from arXiv · show

Reinforcement learning for large language models suffers from high-variance token-level importance sampling (IS) ratios, which would destabilize policy optimization at scale. To improve stability, recent methods typically use a fixed sequence-level IS ratio for all tokens in a sequence or adjust each token's IS ratio separately, thereby neglecting temporal off-policy derivation across tokens in a sequence. In this paper, we first empirically identify that local off-policy deviation is structurally inconsistent at the token level, which may distort policy-gradient updates across adjacent tokens and lead to training collapse. To address the issue, we propose Online Causal Kalman Filtering for stable and effective Policy Optimization (KPO). Concretely, we model the desired IS ratio as a latent state that evolves across tokens and apply a Kalman filter to update this state online and autoregressively based on the states of past tokens, regardless of future tokens. The resulting filtered IS ratios preserve token-wise local structure-aware variation while strongly smoothing noise spikes, yielding more stable and effective policy updates. Experimentally, KPO achieves superior results on challenging math reasoning datasets compared with state-of-the-art counterparts.

1. Introduction

Large-scale RL for LLMs can become unstable because off-policy token-level IS ratios amplify gradient variance, while existing smoothing methods neglect temporal structure within sequences. KPO uses causal Kalman filtering to smooth local fluctuations while preserving coherent token-level structure, achieving stable and effective policy optimization on math-reasoning benchmarks.

  • Motivation: Large-scale off-policy RL can destabilize LLM policy optimization because token-level IS ratios amplify policy-gradient variance.Instability is also associated with entropy collapse and practical mismatches such as MoE routing discontinuities, train-inference differences, and numerical precision.
  • Motivation: Existing methods either apply one sequence-level IS ratio to every token or adjust token ratios independently, neglecting temporal structure across tokens.Sequence-level smoothing reduces variance but can obscure within-sequence heterogeneity.
  • Motivation: Off-policy deviation is expected to be globally heterogeneous but locally homogeneous within coherent reasoning paths.Adjacent tokens in a reasoning path share local semantics and should exhibit slowly varying deviation, whereas distinct paths can differ substantially.
  • Empirical analysis: The paper empirically examines token-level off-policy patterns using window-wise frequency, run lengths, and switching frequency during GRPO on Qwen3-4B.These statistics characterize how off-policy behavior changes across a response.
  • KPO: KPO applies causal Kalman filtering over token ratios to smooth local fluctuations while preserving within-sequence coherent structure.The filter uses past and current tokens online, making the estimated ratio autoregressive and structure-aware.
  • Results: KPO consistently improves math-reasoning benchmark performance over state-of-the-art baselines while enhancing training stability.The paper evaluates KPO on six challenging mathematical reasoning tasks.

2. Related Work

Related work addresses RL optimization through improved advantage estimation, normalization, group weighting, and importance-sampling ratio handling. Existing IS-ratio methods either smooth ratios at the sequence level or adapt them token by token.

  • Advanced variants of group-based RL: Recent group-based RL methods refine learning-signal estimation, normalization, and utilization to reduce biased or degenerate updates.Examples include Dr. GRPO, GVPO, BNPO, and GDPO, which target estimator collapse, group weighting, or normalization-related instability.
  • IS-ratio-oriented RL: GMPO and GSPO use fixed sequence-level IS ratios for all tokens within a sequence.Their approach averages token-wise ratios across the entire sequence.
  • IS-ratio-oriented RL: SAPO adaptively modulates token-wise ratios with soft gating functions based on ratio magnitudes.This provides continuous token-level adjustment rather than a single sequence-wide ratio.
  • IS-ratio-oriented RL: ASPO uses asymmetric ratio handling to tighten updates on risky high-ratio tokens while retaining learning signal.It represents another token-wise strategy for controlling importance-sampling effects.

3. Preliminaries

The preliminaries define token-level importance sampling for PPO and describe GRPO’s shared sequence-level advantage without an explicit value model. They also motivate sequence-level ratio smoothing while noting its loss of within-sequence heterogeneity.

  • PPO: PPO defines each token’s IS ratio as the updated-policy probability divided by the old-policy probability.The ratio corrects distribution shift between the behavior policy and the updated policy.
  • PPO: PPO uses token-level advantages and clips the surrogate objective to cap overly large effective policy updates.The clipping range controls how far the ratio can influence the objective.
  • GRPO: GRPO avoids an explicit value model by sampling response groups, scoring them with a verifier, and constructing a sequence-level relative advantage shared across response tokens.The shared advantage is assigned to every token in the corresponding response.
  • Token-level policy gradients: In off-policy optimization, token-level IS ratios can vary within clipping bounds and contribute to instability through covariance effects involving high-probability, high-advantage tokens.On-policy ratios equal 1 for all tokens.
  • Sequence-level IS ratios: GSPO and GMPO reduce variance by replacing token-wise ratios with a fixed sequence-level ratio for every token.This fully smooths token-wise fluctuations into one scalar, regardless of within-sequence off-policy heterogeneity.

4. Method: Online Causal Kalman Filtering

KPO treats token-wise IS ratios as a structured but noisy time series whose off-policy degree varies across semantic segments yet remains locally coherent. It applies causal Kalman filtering to smooth ratios using past and current tokens, then uses the filtered values in policy optimization to retain token-level heterogeneity while reducing variance.

  • Motivation: KPO models per-token IS ratios as a structured, noisy time series with local coherence across adjacent tokens.The method targets latent off-policy dynamics that vary across semantic segments but should not fluctuate sharply token by token.
  • State-space model: The state-space model represents the desired IS ratio as a latent state and separates process noise Q from observation noise V.Q controls how quickly the latent ratio may drift, while V reflects the reliability of observed ratios and jointly determines the smoothing-tracking trade-off.
  • Filtering procedure: Causal filtering predicts each token’s state from past observations, computes an adaptive Kalman gain, and updates the estimate with the current observation.The gain increases when predicted uncertainty is large or measurement noise is small, and decreases when noisy observations should be down-weighted.
  • Filtering procedure: The filter updates mean and uncertainty online, allowing larger gains to track genuine changes while smaller gains suppress transient fluctuations and isolated spikes.This autoregressive design uses only past and current tokens, matching left-to-right token generation without future-token information.
  • Policy optimization objective: KPO exponentiates the filtered log-space estimate and replaces raw token-wise IS ratios in GRPO-style objectives with the resulting ratios.The objective retains per-token ratios rather than sequence-level aggregation, preserving token-wise heterogeneity while reducing variance through causal filtering.

5. Experiments

Experiments show that KPO delivers strong math-reasoning performance while stabilizing training and reshaping token-level IS-ratio dynamics toward coherent, low-variance segments. Stronger filtering generally improves rewards, whereas weaker filtering allows noise and performance degradation.

  • KPO consistently improves over GRPO, GMPO, and GSPO across six challenging math benchmarks.
  • KPO-clipped achieves the best avg@16 on AIME’24, AIME’25, AMC’23, MATH500, and Olympiad.The reported values are 37.91, 36.87, 87.50, 89.42, and 54.06, respectively.
  • Against GSPO, KPO-clipped improves AIME’24 avg@16/pass@16 from 32.70/60.00 to 37.91/63.33, while KPO-unclipped reaches 66.67 pass@16.On AIME’25, KPO-clipped improves GSPO from 29.16/50.00 to 36.87/60.00.
  • KPO remains stable and effective throughout training, while GRPO diverges after approximately 200 steps and GSPO and GMPO plateau.KPO finishes with the highest reward among the compared methods.
  • KPO maintains relatively high entropy with mild fluctuations, unlike GRPO’s early collapse to near-zero entropy.The reported dynamics suggest that KPO better preserves exploration during optimization.
  • Kalman filtering increases token run-lengths from short fragmented runs to long contiguous blocks and reduces switching frequency from 0.43 to 0.01.Up and Down run-lengths increase to 119.95 and 135.12, respectively, while On increases to 35.11.
  • The Q/V = 1e −6 setting yields the best final reward and stable improvement, whereas Q/V = 1e −2 produces lower rewards and mid-training degradation.The experiments identify a trade-off between stronger smoothing and responsiveness to short-term fluctuations.

6. Conclusion

KPO treats token-wise importance-sampling ratios as a temporally evolving process and applies causal Kalman filtering to smooth noise while preserving structural coherence. It improves stability and effectiveness, but its autoregressive filtering limits parallelization.

  • 6. Conclusion: KPO models token-wise importance-sampling ratios as an online state-space inference problem with a causal Kalman filter.The filter uses within-sequence temporal structure to estimate ratios autoregressively.
  • 6. Conclusion: KPO suppresses noisy importance-sampling spikes while retaining structural coherence, improving both policy-optimization stability and effectiveness.The paper reports superiority on challenging math-reasoning benchmarks and compatibility with existing pipelines.
  • 6. Conclusion: Autoregressive Kalman filtering is difficult to parallelize like token generation, motivating future work on parallelized filtering algorithms.This is the paper’s stated future direction.
  • 6. Conclusion: The training procedure computes token/action-wise log importance ratios as noisy observations and filters them in log space using process noise Q and measurement noise R.These noise parameters control the strength of temporal smoothing.

B.1. The experimental setting for Figure 1

Figure 1’s experimental analysis records token-level importance-sampling ratios from off-policy GRPO training and measures their temporal patterns. The analysis uses window-wise frequency, contiguous run lengths, and switching frequency over valid tokens.

  • B.1. The experimental setting for Figure 1: The analysis records token-level importance-sampling ratios for 960 samples at training step 100 under the original GRPO setup.Sequences use a maximum length of 4096 tokens; batch size is 32 and mini-batch size is 8.
  • B.1. The experimental setting for Figure 1: The experiment uses windows of 50 tokens and analyzes valid sequence positions while excluding padding through a mask.The maximum of 80 windows per sample follows from the sequence-length limit.
  • B.1. The experimental setting for Figure 1: Window-wise off-policy frequency counts up-policy and down-policy tokens within each 50-token window and aggregates frequencies across samples containing that window.Up and down are determined by the sign of the log importance ratio.
  • B.1. The experimental setting for Figure 1: Run lengths measure maximal contiguous sequences of the same token type after scanning each valid sample from left to right.Token types are up, down, or on, assigned from log-ratio signs or predefined ratio bins.
  • B.1. The experimental setting for Figure 1: Switch frequency measures adjacent token-type changes within non-overlapping windows and averages the resulting window-level rates.For sequences shorter than 50 valid tokens, the full valid sequence is used; sequences of length at most one receive rate 0.

B.2. The details of experiments

The Math experiments use uniform training and evaluation settings across methods, including fixed response lengths, learning rate, rollout groups, reward definition, and batch sizes.

  • B.2. The details of experiments: All Math methods use maximum response lengths of 4096 tokens and an actor learning rate of 1 × 10^-6.These are uniform hyperparameters across the compared methods.
  • B.2. The details of experiments: Training uses group-based rollouts with group size 8 and a binary rule-based reward assigning 1 for success and 0 for failure.The reward function is explicitly binary.
  • B.2. The details of experiments: Training and evaluation batch sizes are 32 and 64, respectively, with loss aggregation set to sequence-mean-token-mean.The passage also states that evaluation uses nucleus sampling with top-…; the value is truncated in the supplied text.

B.3. The meanings of training metrics

The paper interprets training metrics as indicators of update stability, exploration, and learning effectiveness. It also defines post-filtering token classes with a tolerance interval around ratio 1.

  • B.3. The meanings of training metrics: Policy Gradient Loss is the primary optimization signal: gradual decreases suggest stable learning, while sharp spikes or large magnitudes indicate potentially aggressive updates.The metric is used to monitor optimization stability.
  • B.3. The meanings of training metrics: Policy Gradient Clip Fraction measures the fraction of updates affected by clipping, with persistently high values implying frequent extreme updates.Moderate clipping levels are described as controlled steps.
  • B.3. The meanings of training metrics: Entropy tracks policy stochasticity: adequate entropy supports exploration, whereas rapid decay indicates mode collapse or over-confident generation.Excessively high entropy can slow convergence.
  • B.3. The meanings of training metrics: Mean Reward is the average return per rollout, and a smooth upward trend reflects effective learning.Abrupt drops typically signal instability or reward hacking.
  • B.3. The meanings of training metrics: After Kalman filtering, tokens are on-policy when r ∈ [1 − 0.0003, 1 + 0.0004], with values outside the interval classified as down-policy or up-policy.The tolerance is used because exact ratio equality becomes rare after smoothing.

C.2. The frequency-domain analysis

The frequency-domain analysis centers the token-level IS-ratio sequence, removes its mean, and uses the DFT spectrum to quantify low-frequency energy. Kalman filtering shifts the sequence toward predominantly low-frequency variation, with LFR rising from 0.12 to 0.98.

  • The analysis mean-centers the token-level ratio sequence before computing spectral energy, preventing the DC component from dominating.The centered signal is transformed with a discrete Fourier transform, and spectral power is defined from the transform magnitude.
  • The DFT spectrum defines low-frequency bins as symmetric frequency indices around zero up to cutoff kc.The symmetric-bin construction consistently counts low-frequency energy for real-valued sequences.
  • LFR measures the fraction of total signal energy below cutoff frequency ωc = 2πkc/T, distinguishing rapid token-to-token fluctuations from longer-span variation.Small LFR indicates high-frequency concentration, while large LFR indicates predominantly low-frequency variation.
  • 0.12 is the unfiltered sequence's LFR, consistent with a spectrum dominated by high-frequency energy.
  • 0.98 is the Kalman-filtered sequence's LFR, implying that it is almost entirely explained by low-frequency variation.
Loading 2602.10609v2…