Kernelized Advantage Estimation: From Nonparametric Statistics to LLM Reasoning

Abstract: Recent advances in LLMs have increasingly relied on reinforcement learning (RL) to improve their reasoning capabilities. Three approaches have been widely adopted: (i) Proximal policy optimization and advantage actor-critic rely on a deep neural network to estimate the value function of the learning policy in order to reduce the variance of the policy gradient. However, estimating and maintaining such a value network incurs substantial computational and memory overhead. (ii) Group relative policy optimization (GRPO) avoids training a value network by approximating the value function using sample averages. However, GRPO samples a large number of reasoning traces per prompt to achieve accurate value function approximation, making it computationally expensive. (iii) REINFORCE-type algorithms sample only a single reasoning trajectory per prompt, which reduces computational cost but suffers from poor sample efficiency. In this work, we focus on a practical, resource-constrained setting in which only a small number of reasoning traces can be sampled per prompt, while low-variance gradient estimation remains essential for high-quality policy learning. To address this challenge, we bring classical nonparametric statistical methods, which are both computationally and statistically efficient, to LLM reasoning. We employ kernel smoothing as a concrete example for value function estimation and the subsequent policy optimization. Numerical and theoretical results demonstrate that our proposal achieves accurate value and gradient estimation, leading to improved policy optimization.

• The paper introduces KAE, a variance reduction method that leverages kernel smoothing to balance computational efficiency and gradient accuracy in policy optimization.
• It demonstrates substantial improvements with up to 70% reduction in value estimation MSE and significant gradient error reduction compared to GRPO and REINFORCE++.
• Empirical results on GSM8K and MATH benchmarks show that KAE stabilizes training and enhances test accuracy, matching oracle-level performance in resource-constrained settings.

Kernelized Advantage Estimation for LLM Reasoning: Nonparametric Statistical Approaches

Motivation and Background

This paper addresses the challenge of variance reduction in policy gradient reinforcement learning for LLM reasoning, particularly in resource-constrained environments. Traditional approaches such as PPO and A2C utilize deep neural networks to estimate value functions, significantly increasing computational and memory requirements. GRPO-type algorithms mitigate this by employing sample averages of rewards from multiple trajectories per prompt, but become costly as the required number of trajectories increases. REINFORCE variants reduce computational load by sampling a single reasoning trajectory per prompt but suffer from high variance and poor sample efficiency.

The motivation for the proposed Kernelized Advantage Estimation (KAE) stems from the need to balance computational efficiency with accurate, low-variance gradient estimation. By leveraging classical nonparametric statistical techniques—in particular, kernel smoothing—KAE aims to improve value function estimation and subsequent policy optimization without requiring heavy computational resources or extensive sampling.

Formalization and Algorithmic Framework

The RL-based LLM reasoning setup is formulated as a contextual bandit, with prompts $X$, completions $Y$, and verifiable rewards $Z = r(X, Y)$. Policy gradient algorithms estimate $\nabla_\theta J(\theta)$ by leveraging the score function of the autoregressive policy $\pi_\theta$.

Variance reduction is typically achieved via:

• A2C/PPO: Baseline subtraction using value network approximations, which introduces computational and memory overhead.
• GRPO: Group averaging across sampled completions, reducing dependence on value networks but requiring many trajectories.
• REINFORCE++: Baseline using the mean reward across prompts, further decreasing computational cost but at the expense of sample efficiency.

The generic algorithmic structure unifies these approaches, as illustrated in (Figure 1).

Figure 1: Unification of A2C, REINFORCE++, and GRPO within a baseline-driven advantage estimation framework.

KAE innovates by employing kernel smoothing across training iterations for each prompt, borrowing statistical strength from previous rewards while accounting for parameter drift. This results in more accurate value function estimates, even when group size or batch size is severely limited.

Theoretical Analysis

The theoretical contributions revolve around three axes:

1. Value Function Estimation: Under standard smoothness and kernel regularity assumptions, KAE's value estimator achieves bias and variance bounds akin to classical kernel regression. With the optimal bandwidth, it attains Stone's minimax optimal rate: $O(N^{-2p/(2p+1)})$ for $p$-times differentiable value functions ($N$ being the effective sample size). GRPO and REINFORCE++ estimators are shown to be inconsistent in the resource-constrained regime.
2. Gradient Estimation: The mean squared error of the KAE gradient estimator differs from an oracle estimator (using the true value function) by a term proportional to the MSE of the value estimator. As the kernel bandwidth is refined, KAE's gradient estimation achieves asymptotic equivalence with the oracle estimator and strictly outperforms GRPO and REINFORCE++.
3. Policy Optimization: Provided Polyak-Lojasiewicz-type conditions and appropriate learning rate schedules, the gap between the optimal and learned policy's expected reward depends directly on the gradient estimator's MSE. KAE thus achieves policy optimization performance matching that of the oracle, and exceeding both GRPO and REINFORCE++ baselines in the limit.

Empirical Evaluation

Value Estimation

Strong empirical results validate theoretical claims. On GSM8K and MATH benchmarks, KAE consistently achieves 60–70% reduction in MSE over GRPO, and over 90% reduction versus REINFORCE++. Its performance is robust with respect to kernel bandwidth and kernel function choices, as shown in sensitivity analyses.

Figure 2: MSE of KAE’s value estimator under triangular and exponential kernels; robustly outperforming GRPO and REINFORCE++ across bandwidth settings.

Gradient Evaluation

Improvements in value estimation translate into more accurate gradient estimates. KAE achieves 5–9% reduction in gradient MSE over GRPO and 32–65% over REINFORCE++.

Policy Optimization

KAE's superior value and gradient estimation substantially impact downstream learning. On mathematical reasoning benchmarks (e.g., MATH, DAPO), KAE outperforms GRPO-type baselines in average accuracy by up to 11.8%, with maximum gains as high as 79.9% in certain settings.

In resource-constrained single-stream settings, KAE stabilizes learning and avoids the performance degradation observed with REINFORCE, improving test accuracy by 14.9% on GSM8K and 6.6% on MATH.

Figure 3: Test accuracy during training; KAE sustains performance while REINFORCE and variants decline after reaching peak accuracy.

Ablation studies confirm that prompt sampling schedule alone is insufficient for stabilization; the gains stem directly from the use of kernel smoothing for advantage estimation.

Oracle Property Demonstrated

In the one-shot regime, KAE achieves performance nearly indistinguishable from an oracle algorithm with access to the true value function, validating the theoretical oracle property in practical settings.

Figure 4: Expected rewards for one-shot GRPO, oracle, and KAE; KAE successfully matches oracle-level performance in the one-shot regime.

Implications and Future Directions

Practically, KAE enables statistically efficient policy optimization for LLM reasoning in settings where training a deep value network or extensive sampling is infeasible—e.g., small labs or public sector deployments. Theoretically, it demonstrates how classical nonparametric statistics can address variance reduction in RL for generative models, achieving optimal convergence rates and robust performance.

The framework naturally extends to other nonparametric estimators beyond kernel smoothing (e.g., local polynomials, sieves), providing a broader pathway for integrating statistical techniques with deep RL. Further work may incorporate adaptive bandwidth selection, leverage more diverse kernels, or explore multi-step smoothing strategies for even greater efficiency.

These findings inform not only best practices for RL-based LLM training but also foundational understanding of policy gradient variance reduction, opening additional avenues for principled algorithm design in deep reasoning models.

Conclusion

Kernelized Advantage Estimation introduces a novel variance-reduction strategy for policy gradient RL in LLM reasoning by leveraging kernel smoothing over training iterations. Achieving optimal convergence rates and demonstrating empirically and theoretically validated superiority over GRPO and REINFORCE-type baselines, KAE provides a practical, scalable solution for high-quality policy optimization under computational constraints. Its robustness and oracle property underscore the value of statistical methodologies in modern RL-based generative AI.