Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning
Authors: Luckeciano Carvalho Melo
Alessandro Abate
Yarin Gal
Venue: ICLR 2026 Poster (2026)
Keywords: Reinforcement Learning
Policy Gradients
Large Language Models
Primary area: foundation or frontier models, including LLMs

Abstract: Reinforcement Learning, particularly through policy gradient methods, has played a central role in enabling reasoning capabilities of Large Language Models. However, the optimization stability of policy gradients in this setting remains understudied. As a result, existing implementations often resort to conservative hyperparameter choices to ensure stability, which requires more training samples and increases computational costs. Hence, developing models for reliably tracking the underlying optimization dynamics and leveraging them into training enables more sample-efficient regimes and further unleashes scalable post-training. We address this gap by formalizing the stochastic optimization problem of policy gradients with explicit consideration of second-order geometry. We propose a tractable computational framework that tracks and leverages curvature information during policy updates. We further employ this framework to design interventions in the optimization process through data selection. The resultant algorithm, Curvature-Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out. Theoretically, we establish monotonic improvement guarantees under realistic assumptions. On standard math reasoning benchmarks, we empirically show that CAPO ensures stable updates under aggressive learning regimes where baselines catastrophically fail. With minimal intervention (rejecting fewer than 8% of tokens), CAPO achieves up to 30$\times$ improvement in sample efficiency over standard GRPO for LLM reasoning.

Published as a conference paper at ICLR 2026
STABILIZING POLICY GRADIENTS FOR SAMPLE-EFFICIENT REINFORCEMENT LEARNING IN LLM REASONING
Luckeciano Melo Alessandro Abate Yarin Gal 1ATMI,, University of Oxford 2OXAAV, University of Oxford luckeciano carvalho melo@cs. ox. ac. uk
ABSTRACT
Reinforcement Learning, partieularly through policy gradient methods, has played a central role in enabling reasoning capabilities Large Language Models. How- ever, the optimization stability of policy gradients in this setting remains under- studied. As result, existing implementations often resort to conservative hyper- parameter choices to ensure stability, which requires more training samples and increases computational costs. Hence, developing models for reliably tracking the underlying optimization dynamics and leveraging them into training enables more sample-efficient regimes and further unleashes scalable post-training. We address this gap by formalizing the stochastic optimization problem policy gradients with explicit consideration of second-order geometry. We propose tractable computational framework that tracks and leverages curvature information during policy updates. We further employ this framework to design interventions in the optimization process through data selection. The resultant algorithm, Curvature- Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out. Theoretically, we establish monotonic improvement guarantees under realistic assumptions. On standard math reasoning benchmarks, we empirically show that CAPO ensures stable updates under aggressive learning regimes where baselines catastrophically fail. With minimal intervention (reject- ing fewer than 8% of tokens), CAPO achieves up to 30x improvement in sample efficiency over standard GRPO for LLM reasoning.
MATH Accuracy
INTRODUCTION
(A)
0.80
The emergence of reasoning capabilities in Large Language Models (LLMs) represents major shift in Al research. Beyond language understanding, reasoning has become core ingredient of widely deployed systems (Ope- nAl al.) 2024 Gemini 2025). enabling appli- cations such mathematical problem solving (Shao et al. 2024), code generation (Shojaee al. 2023). and agentic workflows (Yao al. 2023 This progress is primarily attributed to advances in scaling Reinforcement Learn- ing (RL) techniques for LLM post-training, particularly policy gradient methods such as PPO (Schulman et al.) (017). GRPO (Shao ct al. 2024). and variants (Yu et al. 2025 Liu al. 2025b). These methods enabled LLMss to develop behaviors for autonomous chain-of- thought reasoning (Gandhi et al., 2025) and to effectively scale test-time compute (Setlur
0.75
GRPO
0.70
0.65
0.60
0.0
0.2
0.4
0.6
0.8
1.0
Training Completions
x10*
Accuracy on MATH dataset from dir- ferent RL methods. CAPO (ours) achieves x0 greater sample efficiency under an aggressive (A) update regime (higher learning rate, smaller batch size), whereas GRPO suffers policy collapse.
Figure 1:
al. 2025

---

Published as a conference paper at ICLR 2026
Despite its success in LLM fine-tuning and other decision-makee makinn tasks (Bellemare et al. 2020 Mnih et al.] 2015), RL still faces fundamental challenges that limit its broader practicality and scalability. In particular, policy gradients suffer from optimization instabilities driven by the non- stationary nature of the RI. objective and the high variance of estimates (Castanyer et al. 2025). These problems are further compounded by the known pathologies of training deep networks (Pas-) canu et al., 20133; Pennington et al., 2017). These factors lead to several undesired consequences, such as catastrophic updates and policy collapse (Dohare et al. 2023), plasticity loss (Juliani & Ash 2024), sample inefficiency (Kaiser et al. 2020), and hyperparameter sensitivity (Henderson et al., 2018 As a result, the optimization dynamics of RL remain an active area of research from both theoretical and empirical standpoints (Mei et al., 2022 Lyle et al. 2022 Vaswani et al.) 2022)
Perhaps due to the recency of the topic, the optimization dynamics of RL in the context of L.Mss remains underexplored. These challenges persist in the LLM setting and may be even more pro- nounced, since training involves billion-parameter models with very deep architectures and sam- pling horizons that can extend arbitrarily. In practice, current implementations of RL for LLMs typically rely on conservative hyperparameters to ensure stability, such as low learning rates (c.g.. 3 x 10-6 less) and large batch sizes (e.g. thousands of generations per policy update) (Sheng al. 2024) Hugging Face 2025; Guo et al.] (025). These choices substantially increase the number of LLM generations required for learning, raising computational costs. Therefore, stabilizing these algorithms in sample-efficient regimes is crucial to further scale RL for LLM reasoning.
One promising direction is to design algorithms that explicitly model second-order geometry in the optimization landscape and incorporate this information into policy updates. In this work, we formalize the RL optimization problem accounting for curvature terms, namely the Hessian of the objective and the Fisher Information Matrix of the policy distribution. Building on this formulation, we introduce a computationally and numerically tractable model of optimization dynamics that ap- proximates this curvature information. This model enables continuous monitoring of gradient and curvature estimates during policy updates, scales to billion-parameter models and provides analyti- cal expressions for these quantities, which facilitate systematic analysis of the learning dynamics.
We further leverage this optimization model to plan the next policy gradient step It allows antic- ipating policy updates that potentially induce sudden shifts in the objective or policy distribution often associated with unstable optimization behavior - and intervening before taking the actual step identify particu-
in the LLM. We propose a simple data selection mechanism intervention: lar samples that heavily contribute to these abrupt shifts and mask them out of the policy gradient estimation. We refer to this method as Curvature-Aware Policy Optimization (CAPO).
We theoretically establish monotonic policy improvement guarantees under CAPO with practical assumptions. We then empirically validate CAPO on standard math reasoning benchmarks, show- ing that it yields stable optimization even in regimes with aggressive updates, where standard RL algorithms suffer catastrophic updates and policy collapse. As result, CAPO achieves up 30x improvement in sample efficiency compared to GRPO in the standard regime, as presented in Fig- ure [ Lastly, we show that its interventions are minimal, typically rejecting fewer than 8% of the tokens, with negligible computational overhead.
2 RELATED WORK
RL & LLMs. The use of RL techniques to optimize LLMs has been an active area of research in recent years. Early work focused on RL from Human Feedback (RLHF), which optimizes poli- cies toward modeled human preferences (Ziegler et al.) 2019: Stiennon et al. 2020 Ouyang et al.) 2022]. More recently, RI. for L.MM reasoning has gained significant attention for its effectiveness in enabling autonomous chain-of of-in-of-thought reasoning (Gandhi et al., 2025) and in scaling test-time com- pute (Setlur et al. 2025). This breakthrough was initially driven by the seminal works of the OpenAI o-series (OpenAl et al. 2024) and DeepScek-RI (Guo al.) 2025), which popularized GRPO (Shao) et al. 2024). Since then, the research community has studied several aspects of the training pipeline (Zhang et al. 2025), y including alternative objectives (Roux al. 2025 Hu al. 2025). , sampling
mechanisms
Yu et al. 2025), reward shapings (Yang et al. 2024). and different training configura-
In this work, "model" refers to the proposed computational model of curvatures and "policy" to the I.I.M. "Model gradients" computed under the former, while "policy gradients" denote the true L.M gradients.
2

---

Published as a conference paper at ICLR 2026
tions (Liu et al. 2025b Team et al. 2023). Our work fits within this line of research by investigating RL for LLMs from an optimization dynamics perspective, proposing a model of the optimization landscape and using it to design stable policy gradient updates.
Optimization Dynamics in RL. The non-convex and non-stationary nature of RL training has mo- tivated large body of work understanding and stabilizing optimization dynamics in RL agents. In the context of policy gradients, prior research has investigated the role of baselines (Mei et al.) 2022), variance reduction techniques (Greensmith et al. 2001), and emergent pathologies such as plasticity or capacity loss (Sokar et al. 2023 Klein et al., 2024 and policy collapse (Dohare et al., 2023]. Beyond these analyses. past literature has also developed conservative policy optimization methods (Schulman et al. 2015) [217]] Achiam et al.] (017). While this line of work is exten- sive and evolving, we primarily highlight the recent contribution of Castanyer et al., which, like ours, examines the stabilization of policy gradients through curvature-informed interventions. Their methodology, however, differs: they apply natural gradients with K-FAC (Eschenhagen et al. al.] (023) in general deep RL environments, whereas our work develops new approximation of curvature that is tractable at the scale of LLMs and is incorporates it into optimization through data selection.
Improving RL for LLM Reasoning. In the context of LLM research, nascent but growing litera- ture explores improvements to RL training for reasoning. These works typically propose heuristics that target specific problems observed during training- for example, noisy gradient estimates, lim- ited output diversity, or large policy updates. Common approaches include rethinking advantage estimation (Liu et al. 2025a AAhmadian al. (024), controlling policy entropy (Yu et al. [025] Cui et al., 2025), and bounding advantage estimates or log-likelihoods Yang et al. (2025a b). In contrast, our work takes more principled approach. Rather than introducing heuristics address isolated issues, we develop a framework based on second-order stochastic optimization that funda- mentally explains these instabilities and addresses them in a unified manner.
3 PRELIMINARIES
Problem Statement. We formulate the problem of next-token generation as a Markov Decision
Process (MDP), defined by the tuple M (8, A, P, R. po. 7, T),, in which S is a state space, A is an action space, P S x A + A(S) a transition function, R S x A + [-Tbound: +oound a bounded reward function, Pa S - A(S) an initial state distribution, Y E [o. 1] a discount factor, and T the length of the horizon. In the LLM setting, let V be a token vocabulary and L E N a maximum sequence length, including both prompt and generated tokens. S = of all finite sequences, with each state & E S representing the concatenation of the prompt and the tokens generated up to time with total length most L. AA is the space spanned by V: at each step, the policy selects a token Or E V. P is governed by autoregressive sampling and takes the form of trivial deterministic function 8r+11 8, a,, where o denotes concatenation. The initial state distribution Po specifies a distribution over prompts. During policy optimization, one typically optimizes a parameterized LLM T S x + A(A), with the objective of maximizing the expected cumulative reward over the generated sequence:
V" is the set
J(0) = EErwTo Σn E' 22'(sszaa))
where T denotes a trajectory, 50 ~ po(so), as ~ (a:) I st). and 8c-1 = 8t ag.
Policy Gradient (PG) methods optimize a stochastic policy by differentiating (() with respect to the policy parameters (Williams) 1992) and can be written as (Sutton et al. 1999;
T
VoJ(O) ETATE
Σ
Vo log To(a; st) R(st,
This expectation can be estimated via Monte Carlo sampling under the current policy mæ. However, such estimates often have high variance. A standard remedy is to subtract a baseline b(s;) which leaves the gradient unbiased while reducing variance. In practice, this is typically done by replacing the reward with an estimate of the advantage function A(s;, a.). For the rest of this work, we will assume the advantage version of this objective.
3

---

Published as a conference paper at ICLR 2026
Group Relative Policy Optimization
(Shoo et al. 2024) is widely used method for RL in LLMs. it optimizes a surrogate objective that employs off-policy
Akin to PPO (Schulman et al. correction Kakade & Langford
2017).
2002)
with clipping strategy to prevent large deviations:
JGRPO (0) = B
1
min (ro(st, cllp(rrr az), c, 1 +0) AGRPO (st, 0z)
Tr
BKR.. (TT( se) Thase( (- 8c))
where (st 0z) = xe(le))) and T is the sampling policy. The KL divergence term acts as a reg- T (az [s))
ularizer that penalizes deviation from Thmse the initial LLM. In contrast to standard PG methods, GRPO draws samples in groups: for each prompt sy po(ss), it generates group of trajectories (Ti) œg. Contributions from all state-action pairs of a trajectory are averaged (rather than dis- counted), which effectively assume ? = with per-trajectory normalization. Finally, the advantage estimator is defined as:
A at)
R(T) - R
R
G L
R(T)),
1 G
(((T)) R) (4)
tE
G
where i((T) is the return for trajectory and é is small constant for numerical stability.
4 MODELING THE OPTIMIZATION LANDSCAPE WITH SECOND-ORDER GEOMETRY
In this section, we develop model of the optimization landscape. We formulate the reinforcement learning (RL) optimization problem with policy gradients by explicitly incorporating second-orde geometric information. Building on this formulation, introduce tractable computational model that approximates the role of curvature during learning. Our hypothesis is that by providing a simple but effective approximation of second-order gradients, one could track sudden shifts in the objective or policy and anticipate potentially unstable updates.
The Higher-Order Objective. Consider the objective function J(0)) in Equation I After an update step A0, the new objective (@ + A0) is given by the following Taylor expansion:
J(O + A0) = ((0) + (oJ A6 + 1 AA H()) AO +0(|A0)))
(5)
2
m AA0)
where H0) denotes the Hessian of the objective. Equation holds under Lipschitz continuous
Hessian (see Assumption A.I). with detailed proof in Appendix A As the cubic term may be negative, we can establish guaranteed lower bound J(Ø + A0) > ((Ø) mn m=()))))))) 0(1 A0 A0 In practice, the cubic term is often negligible, and we approximate the objective change by m (A0). Crucially, standard gradient ascent ignores the Hessian contribution, which can lead to decrease in the objective for non-convex problems (such as RL) when this contribution is sufficiently negative.
The Fisher Information Matrix. The Hessian captures the local curvature of the objective func- tion. In RL, however, the objective is non-stationary, and what ultimately matters is how updates change the policy distribution. For instance, update may produce only small change in the objective while inducing a large shift in the policy. This alters how future trajectories are sampled and may destabilize learning. Therefore, it is necessary track the geometry of the policy distri- bution directly, which what the Fisher Information Matrix (FIM) enables. One can show that the directional curvature under the Fisher geometry approximates the average KL divergence between policy and before and after small step A0:
DKL T0+ = AO F(0) +(()))+(()))))))
(6)
mF(A0)
where DKL (To || To+20)
Egnds [K 1)) To420 b))]. and F(0)
Egwder (-))) Va log (a Ve log (a(a | 8) is the FIM. The proof is in Appendix
4

---

Published as a conference paper at ICLR 2026
Similarly to the Hessian case, the cubic term is often negligible and focus m One can further show that enforcing a trust region DKL (To Te+10) S during policy updates leads to monotonic improvement of the true objective, given sufficiently small 6 (Schulman et al. 2015).
Ultimately, we aim to design a model that approximates m (20) and m (A0) without explicitly computing gradients curvature terms in the high-dimensional parameter space of the LLM. This approach can be viewed as a form of model-based RL, but from a different perspective: whereas prior work typically models components of the MDP, such the dynamics or reward function, we instead model the optimization process itself, which allows us to plan gradient estimates.
4.1 COMPUTATIONAL MODEL
For an LLM with d parameters, both Hessian and FIM are d x d matrices, which is intractable for billion-size parameter spaces. Even approximations such K-FAC (Eschenhagen et al. 2023) would incur unfeasible memory cost. Therefore, we need to devise a computational model that scalable and effectively provides curvature information stabilize policy gradients. Next, we describe our methodology.
Last-Layer Model. Since modeling the full Hessian or Fisher Information Matrix (FIM) in- feasible, we restrict attention to curvature in parameter subspace. To this end, we adopt a sim- ple last-layer approach. An LLM
softmax policy over the token vocabulary Te(a | 8) = E. exp(fots.aa) where fe(s, a) E R are the logits produced by the network. Letting fe (st) denote
exp(fols.a))
the full logits vector, with 0 = (0, 1b). we represent the pre-softmax layer as fo(*;) = (h)(
where W R Kxd, is the last-layer weight matrix, K = dim(V), and ha(ss) Rdd, We then define
&b vec(W) E kk.d; In Appendix C we show that the last-layer model gradient O(사b) of the
objective in Equation
T Σ
T( = EYoNo
y A'L (c) - =T()s)) @ ho(st)
where @ denotes a Kronecker product, Ca, E V is the one-hot action vector ea. = 1aa at).
and To (st)) ) the policy distribution vector. We use the vectorization operation vec(.) only for conve- nience and it does not introduce new assumptions. In this work, we use tilde superscript to denote bosel-based gradients and curvatures, in contrast to the actual policy gradient g(0) == V J().
Under the last-layer model, the Hessian of the objective takes the following form:
T Σ +' ((s,,)((ee, ((a,- (T0(s:))
(+bb)= = ETTNTT
F(si)) -
ha(()))a((s))
(8)
where F(s:) is the FIM for state st. In Lemma C.1
we show that this expression can be estimated via samples. Similarly, the last-layer approximation of the FIM is:
((b) Eyano [((ea) - (T(s)))(ea -oo)())___)()))))))))))))) T) @ (o())))))()))))
Computing Directional Curvatures. Even with the approximated model, the curvature matrices have dimension Kd; x Kd,, For current LLMs, where K > 10s and d; > 103, fully materializing these matrices computationally infeasible. Fortunately, our goal is to approximate the shifts in the objective and policy, m (A0) and m (A00)) Thus, we only need to approximate the directional curvatures AO C(@) AO, without explicitly materializing the full Hessian or FIM. In Appendix D we present mechanism that enables this computation without constructing large tensors. Our method requires storing only (Kdd) tensors per state-action sample, instead of the ((Kd,))) entries of the full curvature matrices.
Exploiting Gradient Sparsity. We further reduce complexity by exploiting the structure of gradi- ents arising from LLM generation. Standard LLM decoding relies on selective sampling methods (e.g., top-k, nucleus sampling) Wolf et al. (2020) to improve generation quality, as most of the prob- small subset l of the vocabulary (Fan et al. 2018 Holtzman al.) 2020), typically with & < 100. Consequently, only & tokens have son-zero probability at each gener- ation step, which implies that only the k de parameters of the last-layer weight matrix W associated
ability mass is concentrated
5

---

Published as a conference paper at ICLR 2026
with these logits yield non-zero gradients. We therefore store and operate these gradients in sparse form. This sparsity also applies to the computation of directional curvatures in Equations 58 and 63 as these reduce dot products involving sparse vectors (e.g., (ee, (({)))) and the model-based update step AQ). Naturally, as we estimate gradients with more samples, the representation expands to cover all tokens generated, but typically << K. For instance, our experiments presented < 10 Overall, the memory and dot product complexity reduce to O(k - -d).
Modeling the Step AQ. A final design choice concerns how model the planned update steps, AO. Under the last-layer model, these steps take the form Aub. This choice essentially determines how we represent the optimizer. A simple option is model the update as stochastic gradient descent (SGD) step, A ag, where a is the learning rate. Alternativel,, we can match the LLM optimizer, which in our case is Adam (Kingma & Ba 2015), i.e., Aúb a a "Vrte" where Pa and Ộ the bias-corrected first and second moment estimates of the gradient.
5 CURVATURE-AWARE POLICY OPTIMIZATION
We may now compute the objective and policy shifts under our model as: m "(bb) A&b A (*b), m=(bb)
+ F(b))
(10)
and estimate m and m via samples following the methodology described in the subsection 4.]
We now design an algorithm that intervenes the optimization of the underlying LLM policy using the model-based updates. Since our objective is to stabilize policy gradients in sample-efficien regimes, natural choice is to construct an algorithm that follows the principles of trust-region methods (Murphy) 2022). We implement this idea through rejection sampling mechanism.
Given a batch of collected trajectories, we partition into disjoint subsets b; C subset, we compute a proposed step Aab; and evaluate the shifts defined in Equation TO
B.
For each We then
accept subset if it satisfies the (local) trust-region constraints 8p, δн. and 6 shigh.
5H S. ( <
shigh
mp(A))) S ÓF-
(11)
The accepted subsets are subsequently used to compute the gradient update of the LLM policy. Con- ceptually, this mechanism is analogous to token masking. Overall, this data selection mechanism is simple, computationally inexpensive, and flexible, as it can be applied at different granularities, in- cluding tokens, sentences, groups, full batches. The formal pseudocode is provided in Algorithm Next, we establish theoretical results for monotonic policy improvement under CAPO.
Theorem 5.1 (Monotonic improvement under CAPO). Fix thresholds λμ O and 6y >> O. Let B be a batch of sampled trajectories. Split B into disjoint N subsets by C B. and propose candidate subset updates { [.0)..N. Retain those satisfying:
(H(A0,) 2 6µ = w Mr
(A.)))) S 8F,
(12)
with > O and M, defined in Assumption E.
Let B denote the superset of the B accepted EieB A0, Then, for two policies To and
subsets, and define the aggregated update: A
To+20. with (A* a)l < e, we obtain:
27
J(To+10)= - J(TO) 2 w w - O
C= (1 -= 71= e V2.
(13)
Thus choosing w 2 Ov guarantees monotonic improvement: (TT+ 2 ((Te).
The proof is provided in Appendix E Observe that shigh is not required to establish monotonic improvement. Nonetheless, it serves a safeguard against overly aggressive steps. In practice, introducing this upper cap reduces the observed M and which allows the use of smaller òh- Finally, we note that Theorem relies on the true objective and policy shifts, whereas in practice these quantities are approximated using model.
6 EXPERIMENTS AND DISCUSSION
In this section, we evaluate (i) how the proposed computational model captures the optimization landscape, and (ii) how this information can be used to stabilize RL. optimization dynamics through

---

Published as a conference paper at ICLR 2026
MATH Accuracy
TEST Accuracy
GRPO (A)
CAPO DrGRPO
REINFORCE
0.42
GRPO (A)
CAPO
REINFORCE
0.80
0.40
0.75
GRPO at 6 10
0.38
GRPO at 60
0.70
0.36
0.65
0.34
0.60
0.32
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Training Completions x10
Training Completions x10*
Comparison with baseline methods on policy gradient stability. While the setup with more aggressive updates makes all methods more sample-efficient, it also leads the baselines to pol- icy collapse. In contrast, CAPO prevents collapse and achieves up 30x greater sample efficiency than GRPO under aggressive updates.
Figure 2:
CAPO. Our central hypothesis is that an inexpensive yet effective approximation of second-order geometry can track unstable shifts in the objective and policy, and that this information can in turn be used to stabilize aggressive update regimes, leading to more sample-efficient RL in LLMs.
Experimental Setup. We consider standard RL setup for finetuning LLMs on reasoning tasks. Our implementation builds on the Open-R1 open-source project (Hugging Face 2025 and we maximize an iccuracy-based reward. Following prior work, we fine-tune a Qwen2.5 Matth- LLM (Qwen et al. 2025) on mathematical reasoning questions. Our primary evaluation metric is accuracy, but we also track optimization-relatee quantities such as gradient and curvature statistics and token rejection rates. Since goal is to evaluate sample efficiency, report all metrics as a function of the number of training completions (i.e., LLM trajectories generated). Appendix G provides additional details regarding implementation, hyperparameters, and compute resource
Datasets Benchmarks. We train our policies on the MATH dataset (Hendrycks et al. 2021)
For evaluation, we consider eight benchmarks: GSM8K (Cobbe al. 2021). MATH500 (Light-
man et al. 2023). OlympiadBench (He et al. (024). MinervaMath (Lewkowycz et al.] 2022). GPQA:Diamond (Rein al. 2023). AMC23, AIME24, and AIME25. Most of these benchmarks contain mathematical questions at varying levels (high school, graduate. and olympiad), while GPQA focuses on general STEM-related problems. For simplicity, we report the average perfor- mance across all eight benchmarks, which we refer to as "TEST" in the results.
Comparison Methods. We evaluate our approach against two GRPO variants. The first corresponds to the standard "conservative" update regime implemented in the Open-R1 codebase. The second, which we denote "GRPO (A)," adopts a more aggressive regime intended to improve sample eff- ciency, with learning rate 5x higher and a batch size 12x smaller. This matches the configuration used by CAPO. We also evaluate Dr.G (Liu et al. 2025a) and REINFORCE (William) 1992). both under the same aggressive regime.
CAPO operationalization. CAPO optimizes the same objective as GRPO, but leverages the data selection mechanism introduced in Section For fair comparison, use the same hyperparam- eters as GRPO (A), We implement CAPO with token-level selection, i.c., proposing steps A&h; and rejecting samples per-token basis. Finally, model optimization steps using Adam.
6.1 EXPERIMENTS
We highlight and analyze the following questions to evaluate our hypothesis and proposed method:
We release our code httpa///tpp///ttttoooo/oo/////er

---

Published as a conference paper at ICLR 2026
1.00
1e-2
(Token)
le
(Global)
1
(Token)
(Giobal)
GRPO (A)
0.75
0.75
DrCiPO
REINFORCE
0.50
1.0
0.50
0.25
0:25
0.00
6 Training Completions x10?
4
2
Training Completions x107
Training Completions
Training Completions
Figure 3: Evaluation of policy and objective shifts estimates from the proposed computational model during training. Unstable methods exhibit large and abrupt directional curvatures, while stable ones maintain much smaller and smoother shifts. CAPO, by applying token-level bounds, also ensures
shifts at the global (batch) level, supporting the rationale of Theorem 5.1
Does CAPO prevent instability in LLM policy gradients? Does it lead to better sample ef- ficiency? Figure 2 reports accuracy for all methods on MATH and on the TEST benchmark set. First, we observe that the more aggressive setup does lead to more sample-efficient learning than the conservative across all methods. However, for the baselines, this improvement comes the cost of stability. Under the aggressive regime, all baseline methods suffer from policy collapse, with performance dropping well below that of the base model and therefore losing the ability learn fur- ther. In contrast, CAPO maintains stable performance throughout training, remaining effective long after all other methods have collapsed. This demonstrates that CAPO effectively prevents instability under aggressive updates. As a result, CAPO requires 30x fewer completions on MATH and 9x fewer completions on TEST compared standard conservative GRPO.
What does the proposed computational model reveal about the optimization landscape? To analyze this question. we examine the policy shift m and the objective shift m at both the token level and the global (batch) level over the course of training, presented in Figure For my, we find that unstable methods (GRPO (A), DrGRPO, REINFORCE) exhibit very high global directional curvatures during training, whereas stable methods (CAPO, standard GRPO) maintain much smaller shifts. In particular, the global mp correlates closely with the instability observed in Figure D showing that the model, despite its simplicity, remains informative about optimization dynamics.
For my, we observe similar trends: unstable methods show abrupt shifts, while stable ones produce smoother, better- behaved curves. Note that, while higher me directly signals instability since it tracks policy shifts, a higher m does not necessarily directly imply instability. This is because m depends the adopted advantage function (Equation 33)) and the normalization strategy of each method. Still, sharp peaks in the my curves also correlate with training instabilities. Lastly, we highlight that CAPO, by applying local bound per token, also ensures well-behaved shifts the global level, which supports the rationale of Theorem 5. Overall, these results highlight that the computational model provides meaningful information about the optimization landscape, and that CAPO effectively leverages this information to stabilize training.
Can we extend curvature-aware selection to other RL methods? To test this, we extend Dr.GRPO and REIN-
Token Rejection Rate
FORCE by incorporating our proposed curvature-aware selection, resulting in Dr.CAPO and ReinCAPO, respec- tively. Figure 4 reports the evaluation results for these
0.08
0.06
methods. In all cases, incorporating the selection strat-
0.04
egy improves upon the base method and prevents policy collapse. These findings suggest that the proposed com- putational model and intervention mechanism are broadly applicable across different policy optimization objectives.
0.02
0.00
2
4
8
Training Completions x103 Figure 5: Token rejection rate under CAPO. It maintains low rejection rate over training, stabilizing learning with minimal intervention.
How aggressive is CAPO's intervention to ensure sta- bility? We analyze the extent of token rejection required by CAPO to maintain stable gradients, measured by the token rejection rate during training (Figure 5) The re- jection rate peaks at about 8% in the carly stages of opti-
8

---

Published as a conference paper at ICLR 2026
MATH Accuracy
TEST Accuracy
DrCAPO
CAPO
REINFORCE
0.42
DICAPO (A)
CAPO
REINFORCE
0.80
0.40
0.75
0.38
0.70
0.36
0.65
0.34
0.60
0.33
0.32
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Training Completions x x10*
Training Completions
x10*
Figure 4: Evaluation of extended versions of RL methods with curvature-aware selection. In- corporating curvature-aware selection consistently improves the base methods, preventing policy collapse and demonstrating the broader applicability our approach across different policy opti- mization objectives.
mization, when higher learning rates produce more aggressive updates, but quickly decreases and remains below 2% for the remainder of training. Overall, this shows that CAPO guides optimiza- tion toward stable curvature regions while keeping its intervention minimal, allowing the LLM to continue leveraging the vast majority of samples.
Additional Experiments. We provide computational cost analysis of CAPO in Appendix H where we show that the additional components incur minor overhead. Additionally, we present further experiments in Appendix [ including ablation study the optimizer model and de- tailed evaluation of other heuristics traditionally used to ensure stability (c.g., PPO clipping and KL regularization), highlighting their limitations in the LLM setup.
7 FINAL REMARKS
In this work, propose computational framework that models curvature information and inte- grates it into policy updates through CAPO. We provide theoretical guarantees for CAPO and show that it is effective at identifying samples that contribute unstable updates, preventing policy col- lapse in aggressive training regimes where standard RL methods for LLM reasoning fail. As a result, CAPO achieves up 30x improvement in sample efficiency compared widely used training setups, while requiring only minimal intervention and computational overhead. Overall, it enables more sample-efficien learning regimes, supporting further scalability post-training scalability.
Limitations. Despite the encouraging results, we acknowledge some limitations of our work. First, due to compute budget constraints, focused on experiments smaller, academic scale. While we demonstrated the effectiveness of CAPO against commonly used RI. methods, future work could extend these results distinct problem settings and longer training schedules. Second, the choice of CAPO thresholds depends the problem setting (MDP, objective function, base policy) and may require tuning across different scenarios. Nonetheless, this is not a major concern, as the thresholds can be tuned solely the training distribution.
Future Work. Beyond scalability, future research may explore different parametrizations of the computational model (for instance, by extending it to deeper layers) and investigate their impact computational tractability and curvature estimates. In addition, future work may evaluate CAPO extensions to other intervention mechanisms, such soft masking regularization methods.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank Shreshth Malik, Clare Lyle, and Yoav Gelberg for the insightful discussions in the early stages of this project. Luckeciano C. Melo acknowledges funding from the Air Force Office of Scientific Research (AFOSR) European Office of Aerospace Research & Development (BOARD) under grant number FF8855-2211-7000
Yarin Gal is supported by Turing AI Fellowship financed by the UK government's s Office for Artificial Intelligence, through UK Research and Innovation (grant reference EP/V030302/I and delivered by the Alan Turing Institute. We gratefully acknowl- edge compute resources provided by the Alan Turing Institute. This work was also funded under the Horizon Europe grant 101213369 DVPS.
REFERENCES
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 22-31, PMLR, 06-11 Aug 2017.
htc70///chiio http::ppss/pprceed..ngg mmlr. wwwps/ewweeseeeeeeeoo ach am Tal.aaatttttt html
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back to basics: Revisiting REINFORCE-style optimiza- tion for learning from human feedback in LLMs. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pp. 12248- 12267, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v 1/2/5///1//2222..llooooo...//// URL https: awt//acaaaa http///aaaannthooooy.. org/ 2024 a124...c--oongg 6627
S. Amari, A. Cichocki, and H. H. Yang. A new learning algorithm for blind signal separation. In Proceedings of the 9th International Conference on Neural Information Processing Systems, N1PS'95, pp. 757-763, Cambridge, MA, USA, 1995. MIT Press.
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251- 276, 1998. doi: 10.1 02216600009966993007744444
Marc Bellemare, Salvatore Candido, Pablo Castro, Jun Gong, Marlos Machado, Subhodeep Moitra, Sameera Ponda, and Ziyu Wang- Autonomous navigation of stratospheric balloons using rein- forcement learning. Nature, 588:77-82, 12 2020. doi: 10.1038/s41586-000 01988//4158-00202999999999
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https: / //www. wandb. com/ Software available from wandb.com.
Roger Creus Castanyer, Johan Obando-Ceron, Lu Li, Pierre-Luc Bacon, Glen Berseth, Aaron Courville, and Pablo Samuel Castro. Stable gradients for stable learning at scale in deep rein- forcement learning, 2025. JRRpptttts///ax.v..oom
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2 10. 14168, 2021.
Ganqu Cui, Yuchen Zhang, Jiacheng Chen, Lifan Yuan, Zhi Wang. Yuxin Zuo, Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, Zhiyuan Liu, Hao Peng. Lei Bai, Wanli Ouyang, Yu Cheng, Bowen Zhou, and Ning Ding. The entropy mechanism of reinforcement learning for reasoning language models, 2025. URL /ttt/ttss////xiiii /ttp///axvvv/..////////e///..avr wttg httppss//aaccvv.orr//aa////// http//aab/22555.211777
Shibhansh Dohare, Qingfeng Lan, and A. Rupam Mahmood. Overcoming policy collapse in deep reinforcement learning. In Sixteenth European Workshop on Reinforcement Learning, 2023, URL
htttps://openzeviiet
Runa Eschenhagen, Alexander Immer, Richard Turner, Frank Schneider, and Philipp Hennig. cronecker-factored approximate curvature for modern neural network architectures. In A. Oh, T. Naumann, A. Globerson, K. Saenko. M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 33624-33655. Curran Associates, Inc.,
10

---

Published as a conference paper at ICLR 2026
2023. URL https://ppoeee....s.sarasss.ccis http:a/perriilleeeapa/////0 File/
Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In Iryna Gurevych and Yusuke Miyao (eds.). Proceedings of the 56th Annual Meeting of the Associa- tion for Computational Linguistics (Volume 1: Long Papers), pp. 889-898, Melbourne, Aus- tralia, July 2018. Association for Computational Linguistics. doi: 1111666/11//188000658 URL
Kanishk Gandhi, Ayush Chakravarthy, Anikait Singh, Nathan Lile, and Noah D. Goodman. Cogni- tive behaviors that enable self-improving reasoners, or, four habits of highly effective stars, 2025.
URL https: /at:://axxxii /ttxxvxxvvvv//////// http://rrgbbbbbbb.bbb.. /ttpp//////.a..o./////////30 01307
Gemini. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality long context, and
next generation agentic capabilities, 2025. URL
Evan Greensmith, Peter L. Bartlett, and Jonathan Baxter. Variance reduction techniques for gra- dient estimates in reinforcement learning. In Proceedings of the 15th International Conference on Neural Information Processing Systems: Natural and Synthetic, NIPS'01, pp. 1507-1514, Cambridge, MA, USA, 2001. MIT Press.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, 7. F. Wu, Zhibin Gou, Zhi- hong Shao, Zhuoshu Li, Ziyi Gao. Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng. Chengda Lu, Chenggang Zhao, Chengqi Deng, Chong Ruan, Damai Dai, Deli Chen. Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Hanwei Xu, Honghui Ding, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jingchang Chen, Jingyang Yuan, Jinhao Tu, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaichao You, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingxu Zhou, Meng Li, Miaojun Wang. Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang. Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. S. Li, Shuang Zhou, Shaoqing Wu, Tao Yun, Tian Pei, Tianyu Sun, T. Wang. Wangding Zeng, Wen Liu, Wenfeng Liang, Wenjun Gao. Wenqin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang. Xiaokang Chen, Xiaotao Nie. Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, X. Q. Li. Xiangyue Jin, Xiaojin Shen, Xiaosha Chen, Xiaowen Sun, Xiaoxiang Wang. Xinnan Song, Xinyi Zhou, Xianzu Wang, Xinxia Shan, Y. K. Li, Y. Q. Wang. Y. X. Wei, Yang Zhang. Yanhong Xu, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Wang, Yi Yu, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma. Yiyuan Liu, Yongqiang Guo, Yuan Ou, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Y. X. Zhu, Yanping Huang, Yaohui Li, Yi Zheng, Yuchen Zhu, Yunxian Ma. Ying Tang. Yukun Zha, Yuting Yan, Z. Z. Ren, Zehui Ren, Zhangli Sha. Zhe Fu, Zhean Xu, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan, Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepseek-rl incentivizes reasoning in llms through reinforcement learn- ing. Nature, 645(808 ():8830))::333666 2025. ISSN 1476-4687. doi: 10. URL https: http::::/doo...g////// 10t.///11..33311186888 025 09422-2
Chaoqun He, Renjie Luo. Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, Jie Liu, Lei Qi, Zhiyuan Liu, and Maosong Sun. Olympiadbench: A challenging benchmark for promoting agi with olympiad-leve bilingual mul- timodal scientific problems, 2024.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. Proceedings of the AAAI Conference on Artificial Intelligence, 32(1), Apr. 2018. doi: 13009//aai.v2ii.. 11691.11664, URL atttps://// aaai. brc/innexx.php/AAAA wtw.c//vveeee///////////
11

---

Published as a conference paper at ICLR 2026
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious of neural text degeneration. In International Conference on Learning Representations, 2020. URL https:
//openzeview. net http://rwww..eeeeeoooeeeaaaaaooe oow.:
Jian Hu, Jason Klein Liu, Haotian Xu, and Wei Shen. Reinforce++: An efficient rlhí algorithm with robustness to both prompt and reward models, 2025. URL htt.ttts//aaxxi..rrg/.. abs/
7551103262
Hugging Face. Open rl: A fully open reproduction of deepseek-rl, January 2025. URL https:
open- rl
Arthur Juliani and Jordan T. Ash. A study of plasticity loss in on-policy deep inforcement learning. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Pa-
quet, J. Tomezak, and C. Zhang (eds.). Advances in Neural Information Process- ing Systems, volume 37, pp. 113884-113910. Curran Associates, Inc., 2024. URL
https://proceediinn newripps....uu....ooo..........m www.ineurip.... cc wwtpccppppppppeeeeeeeeceeeee.eee httppppprr...llspppppp////////// ce 19999966636592111888338
cd6f www.cccffff-pp-ppee Conference.pdd.
Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohi- uddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020. URL
https::/cpenreview net f
Sham M. Kakade and John Langford. Approximately optimal approximate reinforcement learn- ing. In International Conference on Machine Learning, 2002. URL https://api.
Gemanticscholar. wrt/Corppstiii
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015. URL htt.:/tp://aariiv.orr///s//////s
6980
Timo Klein, Lukas Miklautz, Kevin Sidak, Claudia Plant, and Sebastian Tschiatschek. Plastccity loss in deep reinforcement learning: A survey, 2024. URL /ttpp://aaxvvv.or////b//a
04832
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with lan- guage models, 2022. http::tttt///ttttee wwt.:oogg///aaoaooooooo... abs 2206 14858
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.200500 2023.
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi. Tianyu Pang. Chao Du, Wee Sun Lee, and Min Lin. Understanding rl ll-zero-like training: A critical perspective. In 2nd Al for Math Workshop @ ICML 2025, 2025a. URL battps:://ppereete "ttpeerrrrieeeenet//rruuuum j LpC1
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. Understanding l--zeo-like training: A critical perspective, 2025b. URL https:
//axxiv.orr//abs//5000 20783
Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dynam- ics and generalization in deep reinforcement learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th Inter- national Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 14560 14581. PMLR, 17-23 Jul 2022. URL hattp:://ppp httpss////ccceeengg...te mlr. preas/vl 62 e22a html
12

---

Published as a conference paper at ICLR 2026
Jincheng Mei, Wesley Chung, Valentin Thomas, Bo Dai, Csaba Szepesvari, and Dale Schu- urmans. The role of baselines in policy gradient optimization. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 17800. Curran Associates, Inc., 2022. URL httpps://proceedd ngs neurips cc cttp:c/appepeeeee///ppppppeee fitps//aap/se/a//2// FLLe/18022166636666eccccc3333333
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep rein- forcement learning. Nature, 518(7540):529 533, February 2015. ISSN 00280836. URL htt.://////nuuue12222
Kevin P. Murphy. Probabilistic Machine Learning: An introduction. MIT Press, 2022. URL
http:::/probm..giiee
OpenAI, :, Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, Alex Iftimic, Alex Karpenko, Alex Tachard Passos, Alexander Neitz, Alexander Prokofiev, Alexander Wei, Allison Tam, Ally Bennett, Ananya Kumar, Andre Saraiva, Andrea Vallone, Andrew Duberstein, Andrew Kondrich, Andrey Mishchenko, Andy Applebaum, Angela Jiang, Ashvin Nair, Barret Zoph. Behrooz Ghor- bani, Ben Rossen, Benjamin Sokolowsky, Boaz Barak, Bob McGrew, Borys Minaiev, Botao Hao, Bowen Baker, Brandon Houghton, Brandon McKinzie, Brydon Eastman, Camillo Lugaresi, Cary Bassin, Cary Hudson, Chak Ming Li, Charles de Bourcy, Chelsea Voss, Chen Shen, Chong Zhang, Chris Koch, Chris Orsinger, Christopher Hesse, Claudia Fischer, Clive Chan, Dan Roberts, Daniel Kappler, Daniel Levy, Daniel Selsam, David Dohan, David Farhi, David Mely, David Robinson, Dimitris Tsipras, Doug Li, Dragos Oprica, Eben Freeman, Eddie Zhang, Edmund Wong, Eliz- abeth Proehl, Enoch Cheung, Eric Mitchell, Eric Wallace, Erik Ritter, Evan Mays, Fan Wang, Felipe Petroski Such, Filippo Raso, Florencia Leoni, Foivos Tsimpourlas. Francis Song, Fred von Lohmann, Freddie Sulit, Geoff Salmon, Giambattista Parascandolo, Gildas Chabot, Grace Zhao, Greg Brockman, Guillaume Leclere, Hadi Salman, Haiming Bao, Hao Sheng, Hart An- drin, Hessam Bagherinezhad, Hongyu Ren, Hunter Lightman. Hyung Won Chung, Ian Kivlichan, Ian O' O'Connell, Ian Osband, Ignasi Clavera Gilaberte, Ilge Akkaya, Ilya Kostrikov, Ilya Sutskever, Irina Kofman, Jakub Pachocki, James Lennon, Jason Wei, Jean Harb, Jerry Twore, Jiacheng Feng, Jiahui Yu, Jiayi Weng, Jie Tang, Jieqi Yu, Joaquin Quiñonero Candela, Joe Palermo, Joel Parish, Johannes Heidecke, John Hallman, John Rizzo, Jonathan Gordon, Jonathan Uesato, Jonathan Ward, Joost Huizinga, Julie Wang, Kai Chen, Kai Xiao, Karan Singhal, Karina Nguyen, Karl Cobbe, Katy Shi, Kayla Wood, Kendra Rimbach, Keren Gu-Lemberg, Kevin Liu, Kevin Lu, Kevin Stone, Kevin Yu, Lama Ahmad, Lauren Yang, Leo Liu, Leon Maksin, Leyton Ho, Liam Fedus, Lilian Weng, Linden Li, Lindsay McCallum, Lindsey Held, Lorenz Kuhn, Lukas Kon- draciuk, Lukasz Kaiser, Luke Metz, Madelaine Boyd, Maja Trebacz, Manas Joglekar, Mark Chen, Marko Tintor, Mason Meyer, Matt Jones, Matt Kaufer, Max Schwarzer, Meghan Shah, Mehmet Yatbaz, Melody Y. Guan, Mengyuan Xu, Mengyuan Yan, Mia Glaese Mianna Chen, Michael Lampe, Michael Malek, Michele Wang, Michelle Fradin, Mike McClay, Mikhail Pavlov, Miles Wang, Mingxuan Wang, Mira Murati, Mo Bavarian, Mostafa Rohaninejad, Nat McAleese, Neil Chowdhury, Neil Chowdhury, Nick Ryder, Nikolas Tezak, Noam Brown, Ofir Nachum, Oleg Boiko, Oleg Murk, Olivia Watkins, Patrick Chao, Paul Ashbourne, Pavel Izmailov, Peter Zhokhov, Rachel Dias, Rahul Arora, Randall Lin, Rapha Gontijo Lopes, Raz Gaon, Reah Miyara, Reimar Leike, Renny Hwang, Rhythm Garg, Robin Brown, Roshan James, Rui Shu, Ryan Cheu, Ryan Greene, Saachi Jain, Sam Altman. Sam Toizer, Sam Toyer, Samuel Miserendino. Sandhini Agar- wal, Santiago Hernandez. Sasha Baker, Scott McKinney, Scottie Yan, Shengjia Zhao, Shengli Hu, Shibani Santurkar, Shraman Ray Chaudhuri, Shuyuan Zhang. Siyuan Fu, Spencer Papay, Steph Lin, Suchir Balaji, Suvansh Sanjeev, Szymon Sidor, Tal Broda, Aidan Clark, Tao Wang, Tay- lor Gordon, Ted Sanders, Tejal Patwardhan, Thibault Sottiaux, Thomas Degry, Thomas Dimson, Tianhao Zheng, Timur Garipov, Tom Stasi, Trapit Bansal, Trevor Creech, Troy Peterson, Tyna Eloundou, Valerie Qi, Vineet Kosaraju, Vinnie Monaco, Vitchyr Pong. Vlad Fomenko, Weiyi Zheng, Wenda Zhou, Wes McCabe, Wojeiech Zaremba, Yann Dubois, Yinghai Lu, Yining Chen, Young Cha, Yu Bai, Yuchen He, Yuchen Zhang, Yunyun Wang. Zheng Shao, and Zhuohan Li. Openai ol system card 2024. URL http///axvvv... http:///wwwwv..o/aaa/////aaa 16720
13

---

Published as a conference paper at ICLR 2026
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal. Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Pro- ceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 978171385 171787188888.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning. volume 28 of Proceedings of Machine Learning Research, pp. 1310-1318, Atlanta, Georgia, USA, 17-19 Jun 2013, PMLR. URL www.tttpss//cc /ttttp:://proceeingg.s mlr. wttppp//////sssanaa.... 3.tml
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In I. Guyon, U. Von Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
nwwrips. cttpp/ssrccpppppeireeaaaeappp wtt.ffflsspppprr//00// HTL1//FE000a67773000111111120100
pdff
Qwen, :, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025.
URL httph:tt/////. //apx::////xxx/ iv. org/ abs/ /tabs//4122 2412. 15115
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate-level google-proof q&a benchmark, 2023. URL httph: wttpvv/aaxvvvvvva org abs 2311.
Nicolas Le Roux, Mare G. Bellemare, Jonathan Lebensold, Arnaud Bergeron, Joshua Greaves, Alex Fréchette, Carolyne Pelletier, Eric Thibodcau-Laufee, Sándor Toth, and Sam Work. Tapered off- policy reinforce: Stable and efficient reinforcement learning for llms, 2025. URL https::/ arxiv.org/abs. httpgg//bs/5533.... 2503. 4286
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Francis Bach and David Blei (eds.), Proceedings the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889 1897, Lille, France, 07-09 Jul 2015. PMLR. URL
mmlr.
html
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. URL /ttp://te//aax....rg////////// 06347
Amrith Setlur, Nived Rajaraman, Sergey Levine, and Aviral Kumar. Scaling test time compute with- out verification or suboptimal, 2025.
12118
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathe- matical reasoning in open language models, 2024. URLhttps://aaxt hww.voorg/aaa org/abs/2422
03300
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. arXiv preprint arXiv: 2409. 19256, 2024.
14

---

Published as a conference paper at ICLR 2026
Parshin Shojace, Aneesh Jain, Sindhu Tipimeni, and Chandan K. Reddy. Exccution-basee code gen- eration using deep reinforcement learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
net. fowp://wwwwweeeeoeoteuuuaaam fo
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon in deep reinforcement learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 32145-32168 PMLR, 23-29 Jul 2023. URL pktts::/pprceeet mageocedding..... mmlr.
press/v202 http///ooar22aammmm
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feed- back. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 3008-3021. Curran Associates, Inc., 2020. URL
http::/appppffeeeeeeeeeaeee files/paper/20200 Ei F1c12208881112208866666633333333
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Proceedings of the 13th International Conference on Neural Information P Processing Systems, NIPS'99, pp. 1057-1063, Cambridge, MA, USA, 1999. MIT Press.
Yunhao Tang and Rémi Munos. On few pitfalls in kl divergence gradient estimation for rl, 2025. URL [tRLhhtps:://// http://xxivv..gaab//////// http:/bs/250000 09477
Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, Chuming Tang, Congcong Wang, Dehao Zhang, Enming Yuan, Enzhe Lu, Fengxiang Tang, Flood Sung, Guangda Wei, Guokun Lai, Haiqing Guo, Han Zhu, Hao Ding, Hao Hu, Hao Yang, Hao Zhang, Haotian Yao, Haotian Zhao, Haoyu Lu, Haoze Li, Haozhen Yu, Hongcheng Gao, Huabin Zheng, Huan Yuan, Jia Chen, Jianhang Guo, Jianlin Su, Jianzhou Wang. Jie Zhao, Jin Zhang, Jingyuan Liu, Junjie Yan, Junyan Wu, Lidong Shi, Ling Ye. Longhui Yu, Mengnan Dong, Neo Zhang. Ningchen Ma, Qiwei Pan, Qucheng Gong, Shaowei Liu, Shengling Ma, Shupeng Wei, Sihan Cao, Siying Huang. Tao Jiang, Weihao Gao, Weimin Xiong, Weiran He, Weixiao Huang, Weixin Xu, Wenhao Wu, Wenyang He, Xianghui Wei, Xianqing Jia, Xingzhe Wu, Xinran Xu, Xinxing Zu, Xinyu Zhou, Xuehai Pan, Y. Charles, Yang Li, Yangyang Hu, Yangyang Liu, Yanru Chen, Yejie Wang, Yibo Liu, Yidao Qin, Yifeng Liu, Ying Yang. Yiping Bao, Yulun Du, Yuxin Wu, Yuzhi Wang. Zaida Zhou, Zhaoji Wang, Zhaowei Li, Zhen Zhu, Zheng Zhang. Zhexu Wang, Zhilin Yang. Zhiqi Huang. Zihao Huang, Ziyao Xu, Zonghan Yang. and Zongyu Lin. Kimi k1.5: Scaling reinforcement learning with llms, 2025.
www.///////.//////////.//eeeee// abs 2200125111111 12599
Sharan Vaswani, Olivier Bachem, Simone Totaro, Robert Müller, Shivam Garg, Matthieu Geist, Marlos C. Machado, Pablo Samuel Castro, and Nicolas Le Roux. A general class of surro- gate functions for stable and efficient reinforcement learning. In Gustau Camps- Francisco J. R. Ruiz, and Isabel Valera (eds.), Proceedings of The 25th International Conference Arti- ficial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 8649. PMLR, 28-30 Mar 2022. URL httpt:tt/////m vaswani22a. html
mlr. press/v151/
R. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8:229-256, 1992.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State-of- the-art natural language processing, 2020. URL httttp::/aaxxi..oot www.w/wwwww....//./..e 03771]
An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, Keming Lu, Mingfeng Xue, Runji Lin, Tianyu Liu,
15

---

Published as a conference paper at ICLR 2026
Xingzhang Ren, and Zhenru Zhang. Qwen2.5-mate technical report: Toward mathematical ex- pert model via self celf-improeemett, 2024. URL httpp::/aaxxy....
abs 2249./2222222
Shihui Yang, Chengfeng Dou, Peidong Guo, Kai Lu, Qiang Ju. Fei Deng, and Rihui Xin. Depo:
Dynamic clipping policy optimization, 2025a. URL fttps::/arx.v..rgg/ wtt..or//aabb///gg
02333
Zhihe Yang, Xufang Luo, Zilong Wang, Donggi Han, Zhiyuan He, Dongsheng Li, and Yunjian Xu. Do not let low probabbiiyl tokens over-dominate in for llms. 2025b. URL https:
//azxi..rr//ab///3333 12999
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. ReAct: Synergizing reasoning and acting in language models. In International Conference on Learning Representations (ICLR), 2023.
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan. Xiaochen Zuo Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, Haibin Lin, Zhiqi Lin, Bole Ma. Guang- ming Sheng, Yuxuan Tong, Chi Zhang, Mofan Zhang, Wang Zhang. Hang Zhu, Jinhua Zhu, Jiaze Chen, Jiangjie Chen, Chengyi Wang, Hongli Yu, Yuxuan Song, Xiangpeng Wei, Hao Zhou, Jingjing Liu, Wei-Ying Ma, Ya-Qin Zhang, Lin Yan, Mu Qiao, Yonghui Wu, and Mingx- uan Wang. Dapo: An open-source IIm reinforcement learning system at scale, 2025. URL https: /tapps://arxiv.. www.e/orgraaa... httgaabs/553..22
Kaiyan Zhang, Yuxin Zuo, Bingxiang He, Youbang Sun, Runze Liu, Che Jiang, Yuchen Fan, Kai Tian, Guoli Jia, Pengfei Li, Yu Fu, Xingtai Lv, Yuchen Zhang, Sihang Zeng, Shang Qu, Haozhan Li, Shijie Wang, Yuru Wang, Xinwei Long, Fangfu Liu, Xiang Xu, Jiaze Ma, Xuekai Zhu, Ermo Hua, Yihao Liu, Zonglin Li. Huayu Chen, Xiaoye Qu, Yafu Li, Weize Chen, Zhenzhao Yuan, Junqi Gao. Dong Li, Zhiyuan Ma, Ganqu Cui, Zhiyuan Liu, Biqing Qi, Ning Ding, and Bowen Zhou. A survey of reinforcement learning for large reasoning models, 2025. URL https:
08827
Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint ar www:::19900089900 2019. htts: : // //arxvv..../..////// org/ abs 1909. 08593
16

---

Published as a conference paper at ICLR 2026
A DERIVATION OF THE SECOND-ORDER OPTIMIZATION OBJECTIVE
In this section, we formally derive the higher-order expansion of the objective function around given parameter vector, and present conditions for monotonic improvement. We start by highlighting a smoothness assumption required for our analysis.
Assumption A.1 (Lipschitz continuity of the Hessian), There exists a constant L2 2 0 such that, for all 7 E [0, 1] and all AO E R.
V =20) V JO))l) s L2T ||00|.
(14)
Assumption A.I is standard in the analysis of trust-region and cubic-regularized methods, and holds locally for smooth policy parameterizations.
Proposition A.1 (Second-order expansion with integral remainder). Let J Rd R be three times differentiable, and denote & (() and H (0) For any update direction A R the objective value at the perturbed parameter 0 + 26 admits the expansion
A H / (1- (1-T) (V-)(0+TA0)-H)00 -
J(0+20)= = J 2e+
dr. (15)
Under Assumption
A.I
the following lower-bound holds
J(e + A0) 2 J(0) + T0 A H
A0
L2 ||0|
(16)
Proof. Let ((T) J(0 +TO0) for E [0, Then ('(0) -g A0 and d² A HA0. The
(one-dimensional) Taylor formula with integral remainder gives
((1) = 600) + +00) + /.' (1 (d") dr.
(17)
Since (7) - -0 H)00, we obtain equation 15
Assumption
A.1
implies + H) ||00|² J[ + ~A)
H op
Lar |
Solving the integral gives (1 L2T dr = L2/6. Since this term can negative, a worst-case bound yields the inequality 16
17

---

Published as a conference paper at ICLR 2026
B DERIVATION OF THE POLICY DIVERGENCE QUADRATIC APPROXIMATION
In this section, we formally derive the higher-order expansion of the KL term around small step AO. Throughout this derivation, we assume standard regularity assumptions hold (e.g., parameter- independent support, differentiability of log Te. and dominated convergence so that differentiation may pass under the expectation). The state averaging distribution d, fixed.
Assumption
B.1 (Lipschitz
continuity of the Fisher curvature).
Let F(0)
Egwder (-1a)
V log (a s) Va log (a {
There exists constant LF 2 0 such that, for all E [0, 1J and all A0 e ER4
F 7 F(0)]l.g S LF |A0||-
(18)
Assumption B.I is analogous to the Assumption A.1 applied to the Fisher geometry.
Lemma B.1 (The grad-log- prob identity). Under regularity assumptions, the following identity holds:
Bgnder am Vo looooolla s)] = O.
(19)
Proof. Fix By normalization,
Σ. Ta 8) 1. Differentiating. Ea VoTo(a 8) = Since
Vene no Ve log me, obtain
Σ Ta 8) Ve log To s) = 0,
(20)
iec. Eu Va log (a 8]] = Averaging over 8 ~ dx preserves zero.
Lemma B.2 (Fisher identity). Under regularity assumptions, the following identity holds:
E E Tog
s)] E V log Ta
Vo log
[] =: F(0).
(21)
Proof. Fix s. Twice differentiating normalization gives V E. To(a 8) = Σa VāTo(a s) = 0.
Using va = no VV log To + Vo log no Vo log ग T we obtain
0 = Σ Ta(a 8) Vi log ( (s) + Σ Te (a 8) V log (a 8) V log (a 8) (22)
Recognizing expectations over a ~ To(- and multiplying by -1 yields
-En (1)) V log Te(a s)] En (c)) Vo log (a s) V log (a s)
(23)
Averaging over 5 dx gives the result.
Proposition B.1 (Second-order expansion with integral remainder). Define the average forward KL
as
DRI To+20) Egnde [KL [KL(TT(- |ss To+ s))]
(24)
Then, for any update A0.
F(O)A0+ (0-)) F(0)) AO dr.
DKL To+20)=
(25)
And, under Assumption B.1 the following holds:
(RL T0+ =0TT F(0)20 0(1001)).
(26)
Proof. Let ((T) = DKL T0+720). By the Taylor expansion with integral remainder,
Φ(1) o(0) 6' (0) (0) + S (1 ()) o²())) dr.
(27)
18

---

Published as a conference paper at ICLR 2026
Then 6(0) = 0, and (r(r) V log me me+rso |s]] A0, by Lemma
B.1 d'(0) 0.
Differentiating again and applying Lemma B.2
d"(r) = 0 F(O + ~00)
s²(0) = AOT F(0)20.
(28)
Substituting the evaluated terms yields the expansion.
Finally, Assumption B. implies
A
- A0| S LFT A
(29)
Integrating (1 - ()r dr 1/6, so the remainder term is (1)
19

---

Published as a conference paper at ICLR 2026
C DERIVATION OF GRADIENTS AND CURVATURES UNDER LAST-LAYER MODEL
In this section, we formally derive the gradient and curvature expressions assuming the last-layer model.
Proposition C.1 (Gradient w.r.t.last-layer model softmax policy). Let us consider sofimax
exp(fe((.))
policy (a 8) En Expexx(ol(aaa) Let us also denote the pre-softmax layer by fo(s)) =
(h)( W RKxx, hø(st) E Rdd Define sb == vec(W) RKd, with 0 = (0. ab),
K dim(V). Then the policy gradient with respect to ab of PG objective:
T
J(0) = Ey TEYNTT
Σ
7' (z) aog 81)
(30)
given by:
T ((()) = Ex Σ A(s, (ca- ro) @ m(ss))
(31)
where ea R.K, K dim(V), denotes the one-hot vector of the realized action as time (i.e., ea = Ea,h TT( E RR is the vector of action probabilities at 8t, and & denotes the Kronecker product.
Proof. Starting from the advantage version of Equation the policy gradient with respect ab is given by
Ø(*b) = Eyyno
Σ
of A(st, az) V log (a sx)
(32)
===
With logits f(s:) Who(s:), the Jacobian of the log-softmax with respect to f(s)) is:
i log ne(a 8) Ca To(si) R
(33)
0f(st)
Vectorizing W gives:
0f(s)) Ozb
Ix 2 ho(se) E RxKd
(34)
By the chain rule,
V logmo(a = (ca-so(ss)) (Ix h)(s:))) (ea -o((s)) ho(ss).
where we used standard Kronecker product identities to obtain a vector in Kd Plugging the ex- pression for V log To
8t) into Equation 32 yields
T r A(s ar) (ca- 1 =o(si)) & h(s))
T( = Erwre
(35)
The Hessian of the Objective. For the Hessian, we start by extending the PG Theorem for Hessians: Lemma C.1 (Hessian the Policy Gradient). Let To a s) be a differentiable stochastic policy and consider the discounted policy gradient objective
J(0) = ETT
Σ ot A(se, ar) log (ot sx))
(36)
1=
where A(st,a.) is the advantage function at time Then, the Hessian of (() is given by
20

---

Published as a conference paper at ICLR 2026
T Σ
(OJ(0) = Br
yf A(s an) (Vo To
V log s + Vi log To sx)))
(37)
Proof. Taking the first derivative of J(0), we obtain
VoJ
Σ t=
7 A(si ai) Vo Volog 8t)
(38)
Differentiating more yields
T
(T)( = Vo RT
Σ
of A(t,, ar) Vo log To(a: 8t)
(39)
Expanding the expectation explicitly over state action pairs weighted by the discounted state distri- bution d. (x)) gives
(O)( = dd (55)))eo Vol To(a s) A(s,a) Vo log ro 1
(40)
Applying the product rule, we obtain
V8J(0) Σ
d (st) Σ (a s)A(s, a) Ve log Te(a V) log (a(a) 8) +V2 log (a
s))
(41)
Rewriting in expectation form gives the final result:
(OJ(() ETT-TO
Σ 7 A(si, ar) (Vo log ma 8t) log To(a; (8) V2 log (a 8t))
(42)
Now, we can state the Hessian form under the last-layer model:
Proposition C.2 (Hessian under Last-Layer Model). Let us consider a softmax policy so(a
exp (fo(s.a))
E. Exrexp (fo(esi))) Let us also denote the pre-sofimax layer by f((t) Who(s:), W
RKxd ho(()) E Rdd Define ub == vec(W) E Rd with = (0,b)) = dim(V). Then, the Hessian of the discounted policy gradient objective
T Σ
J(0) = Errae
rd A(s,, az) Tog st)
(43)
is given by
T
((bb) (4b)==2 ETTTT
Σ 7 A(s,a) A(ssaa (lea (o( [aa
T T & ho(s)))o(()))
(44)
K is the one-hot vector of action a, T (8t) E R is the vector of action probabilities, and F(s+) = diag( (aagg()((ss))) mo(s))=o(s)) is the Fisher information matrix state St-
where ea E KR
Proof. From Proposition
db log (a 81) (ea no (s:)) @ hg(st).
(45)
21

---

Published as a conference paper at ICLR 2026
Hence, the outer product is
V log st) V log (a(a st)
(46)
((e, TT((s)) @ ho(ss)) ((ea mo @ ho(ss))
(ea - (o(s)))e - =o(s))) @ ho(st))a(ss))
where we applied the identity (u @ () & v)T = (nu1) @ (vv Next, we compute the second derivative. Since V log To aa st) (ea =0(ss)) (s)):). follows that
w a s2) = V ()) @ hs()))
(47)
Using (st) (diag((T($s))) To(s))TTOs:)] @ ho(s)), we obtain
V *
log (a $2)
(48)
(diag(no(<:)) - To (T)) @ (o())))((s)) (st) @ (hg()))
(49)
Finally, substituting both terms into the general Hessian expression from Lemma C.I
(+b) E, Es,orx [A(s, a) (V log To (a: st) V (a st) + V log T st))]
yields:
T Σ
(+bb) [2)))= = ETTWTT
A(s,,)) ((oo(s))))))))) no(ss)) F(se)) s (o(()))a((s)
1=0
(50)
Proposition C.3 (Fisher Information under the Last-Layer Model), Let consider softmax policy
exp(Fo(s,a)
(a 5) Exp (Fossaaa) Let also denote the pre-softmax layer by f(sc) Who(s:),
W E Rxd (x(x) .hg(st) Rd.. Define vb vec(W) E Rdd, with O (0, 4)) K dim(V). Then, the Fisher information matrix with respect to 2b is
F(bb) Ex [(Car (e(ss)) T) & ho(s))ho(ss)]]
(51)
where ta, RRK the one-hot vector of the realized action as and TO(SE) E RK is the vector of action probabilities at state 8t-
Proof. From Proposition
Ves log (a: (s)) = (Co, - To(s))) 8 ho(st).
(52)
Therefore,
V log (al st) V log (a | st) ((ca,- To(ss)) ( (53)
where the last step follows from the Kronecker identity (u @ zv @ x) T (wv @ (ax Substituting this into the definition of the discounted Fisher information matrix yields the result. 0
22

---

Published as a conference paper at ICLR 2026
D DIRECTIONAL CURVATURES COMPUTATION
In this section, we present our mechanisms to compute Hessian and Fisher directional curvatures.
D.1 DIRECTIONAL FISHER CURVATURE
For the last-layer parameters vo vec(W)) with W R K dim(V), denote by U ()ve E IR Kxd. the corresponding matrix form of the direction. We aim to compute the curvature of the Fisher information matrix along direction Arh in parameter space. Recall the Fisher information matrix under the Last-Layer Model (Equation 9):
F(b) En [(ubu)) @ (heh,))
(54)
where := ca. Te(ss) E RK is the policy enror vector and h == ho(ss) E Rd4 is the feature vector. Using the Kronecker Vector identity vec(X) T (A @ B) vec(X) Tr(AXBX")
F( = E, [vec(U)] (wru @ h4h, vec(U]]
(55)
= =E. [T(uuu U (ahTUT)
(56)
Let U/ == Uh, E RK Then Tr(ure, )mee) (u, vz)2. And we obtain:
Asp (bb) Er~wo [(u un v²)
(57)
We can estimate the Equation above with samples. Given batch of N state-action- time samples ((8,, az, 4)) an estimator of the curvature is:
N Σ (uT ((Th))))
(58)
with u4 = - and h4 = (s;)) practice, A itself is typically estimated from data (e.g., stochastic gradient direction), hence not strictly deterministic. Therefore, estimating Equation 58 introduces mild bias as u, and ht are statistically dependent.
Cost Analysis. The computation requires only vector and matrix-vector operations. Per sample, we compute Uh, at cost (Kd)) and the dot product u, v at cost O(K). followed by a sealar square. In memory, only store U (Kd parameters) and the per-sample vectors u, and h4. This is dramatically cheaper than materializing the full Fisher matrix F e kdxkd, which would require (Kd)² entries.
D.2 DIRECTIONAL HESSIAN CURVATURE
We now consider the curvature of the Hessian along a direction Asb. We also assume the same notation as subsection Recall the Hessian under the Last-Layer model (Equation 8)
H(b) Ex A(s, a) uuu. - F(s)) ha(s)hihh
(59)
where F(s) = diag( To (s)) To (s)To(ss) T is the Fisher matrix at state = Ca. (T()s) ) E policy error vector and hi = hg((t) ) E IRd, is the feature vector.
The directional curvature along Agp is
Σ r A(st, vec(U) (mwuy) F(st)) @ heh vec(T)
A&h ATTH((AAA EE,
(60)
1=0
23

---

Published as a conference paper at ICLR 2026
Using the Kronecker- Fronecker--Vectrr Vector identity vec(X)) @ voc(X) Tr(AXXX)) obtain:
(b E,
Uhete Σ 7e A(s,, as) Uhath U (61)
Let U2 Ubr. Then the traces simplify via Tr(u (e,, = (wt 21)2
(r(F(st) mn = v/ F(st) vt:
where the first equality uses un Tu To (u v) wv and Tr(ab") bT a, and the second uses Tr( xx = x Ar. Hence,
ALN()) = BrNsr
Σ+ A(st, ar) ((v) vz)² vf F(s))
(62)
t=
We can estimate the Equation above via samples, noting the same remarks as in subsection D.I The sample-based baseed estimator
N Σ 2h A A(s,, (Iu) (1) (F)))) U, ca: -T (s:), 62 (hg(s.).
AU H
(63)
Cost Analysis. The computation again only involves vectors and matrix-vector operations. Per sample, compute v Uh, cost O(Kd), then (u, u2)2 at cost O(K). The second term requires an analogous computation to the Fisher case in subsection [D.]] Hence, the complexity remains O(Kd) per sample, and the memory cost is linear in K and d, avoiding materialization of the full Hessian Hl kd
24

---

Published as a conference paper at ICLR 2026
E MONOTONIC POLICY IMPROVEMENT UNDER CAPO
In this section, we formalize the conditions of monotonic improvement under CAPO.
Assumption E.1 (Bounded curvature and step norms). Let be differentiable policy with objective J(e). Write g(0) V)J((), (())
(z) and F(O)
Exwde a (-1a) Ve log (a|s)) Ve log ( [. For AO E g define the quadratic diagnostics
mH mH (A)) = g(0) Ao AO (O))20, m=(A0) = Ao T F(0)20. (64)
Assume:
(i) (Hessian operator norm bound) | S M for some finite M > [H)))) == 240
where
[H(0)]]
T=]
(ii) (Per-candidate step bound) Every candidate update considered by the algorithm satisfies
A S for some
0.
Remarks. The step norm bound standard in practice, since learning rates, clipping, or trust-region constraints ensure ||0|| <. The Hessian bound (0) S M is more restrictive globally, but
over any compact region parameter space visited by the algorithm, continuity of H (Ø) implies a finite M.
Lemma E.1 (Surrogate-tre true performance gap). For any policies ι and n', with DKr. (T((T)) the averaaee forward KL under dr,
2y
J(T') Z La(T') C (VDR(T C=
71- v/2,
(65)
where (s) S with finite, and (a(') == ((x) + Egwde, amt (-10) (s, a)}. Moreover, writing = and N² 20 for a parameter step A0.
L (To ( g AO + z A H()) Ao + o(113)13)
(66)
Proof. The proof of equation 65 is in Achiam et al. 2017). For Equation 66 we define ((0)) =
Lug ). Note that (0) ) Now compute the gradient of v at 0 0:
V ((0) 0'=0 Bards ove (- (4 a)]
Eander army [Ax (8, (8,a) a) Ver log no (als)))
Exwder awe (A a) Vo log To(a [s] =: g().
(67)
where g()) is precisely the policy gradient. Differentiate once more:
V 11 (00)) - = Euwday amarr (-14) ATo (s, a) Vz Tor (a |s)
0==0
Egnd ( (115) Ano (s, a) Vo log (a Vo log To s)
0-=0"
(()).
By the second-order Taylor expansion, ( + A0) = (0) + g(0) AO + (() AO + o(1A0)1))
(68)
which is exactly equation 66
Theorem E.1 (Monotonic improvement under CAPO, restated). Fix thresholds δн > 0 and 6p > 0. Let B be a batch of sampled trajectories. Split B into disjoint N subsets b; C B. and propose candidate subset updates {A0,).. Js.N. Retain those satisfying:
m (20,) 2 бн = w Mr
mp(10; S 8F,
(69)
with > and M, defined in Assumption subsets, and define the aggregated update: Ao EieBace Ad; Then, for two policies no and A0, we obtain:
Let Baee denote the superset of the B accepted
(TT+ - ((T)) 2 C
(70)
Thus choosing 2 Ovdr guarantees monotonic improvement: I (TT+ J(mo)-
25

---

Published as a conference paper at ICLR 2026
Proof. We first establish bounds in the global Fisher and Hessian directional curvatures.
Fisher global bound. Since F = 0, the quadratic form ((u) == u
Fu is convex. Thus:
A FA#
Σ A8:
Σ
A8, S
1 Σ A0, FAR.
(71)
B
B
B
iEBucc
iEBacc
xEBare
The inequality above follows from:
1
Σ
AB.)
F
Σ
A0.)
(72)
B
B
IEBoco
iEBocc
(A0: 20J) F(A0: A02) 2 0,
(73)
2B2
i.jEBace
because F 2 0 Oimpliss each summand is nonnegative. Hence:
AO FAO K
Σ POT
S 3 Σ 2mF(A0,) S 26p.
(74)
B
B
iEBare
iEBoco
Hessian global bound. Expanding mH (O0):
mH( (10) = g(0) AO + 2
A HAO
= g
Σ
A0:
+ (
Σ
A0,
H
E A0,)
2 B
iEBacc
iEBacc
JEBuce
(75)
1 E g(0)
AO,
Σ
HAO
B
2B
iEB.co
i.jEBaco
We can decompose the quadratic form:
AO,
HA0
AO HAO + Σ
AO. HAO,.
(76)
J.E
iEBuco
i.EBB (gf
Substituting equation i6 equation 75 and grouping yields
m (0)
mn (A0,)
B-1
A0. HAO:
Σ
A0 HD0,- (77)
B
2B²
2B
BEBuuc
SEBeen
f.
iデj
By the operator norm bound || llpp S -M and Cauchy-Schware
AA HA0,1 M || ||0,
Hence, using ||0,|| S for all i,
Σ
2 HAO, < M
Σ B.E $73
HOTHAO; 2 -MBB-11-
(78)
iEB..o
Substituting into equation 77
mH (A0)
JIF
1
Σ mH mH(A0.) M M (1- -)
(79)
B
iEB..c
If each accepted subset satisfies mH (AO)) I
Mr then averaging gives
B Eies,,. m (AO,) 2 w Mr Plugging into equation
79 yields
Mr
my A w + M Mr (1 (1- 1 = +
w.
(80)
B
From Equations 65 and 66 of Lemmal E.1 have that:
J(mo+s0) J(TO) (()) A0 + 2 A0T HO)) AO ~||)0)13 C CV DDL. PRII(T ( 10+
my mx(1e)
mr (r (011))))))
(81)
Then, using mp(A0) < 8F, mH m=(A0) w, and assuming the cubic terms negligible. J(TO+s0) - J(TO) 2 w -C55p.
(82)
Thus choosing w 2 CVDF guarantees monotonic improvement: (T0+A0) 2 (TTo).
26

---

Published as a conference paper at ICLR 2026
F PSEUDOCODE OF CAPO
In this Appendix, we present CAPO's algorithm.
Algorithm 1: Curvature-Aware Policy Optimization (CAPO)
Input Policy #a; batch B of sampled trajectories; thresholds (6p, бн, ohiahh:;
optimizer for the last-layer model (e.g., SGD or Adam).
Output: Updated policy parameters
while done do
// Collect data with the current policy
Sample batch B = {r}; of trajectories, ne.
Partition B into disjoint subsets (b.),=
for 1,............. I....,N N in parallel do
// Buile 1d last-layer meta-model stats subset b; Estimate model-based gradient 0(*b) using Equation 7
Propose Adt; with the optimizer model (e.g., Ae; a for SGD, Adam's 's rule)
Compute directional curvatures A H1)) Ath, A T (()) Au as in Appendix D
Compute objective and policy shifts under the last-layer model:
ma (Av) < A + 2 T ( Arh, mr(Aw) + 2 AØ F(b) T Arh.
// Local cuastrreiinn acceptance test
6H S m z 0 ghigh
and mF(Ad)) s òf then
Mark subset b; as ACCEPT; add to Bace else
REJECT b;.
11 Compute the actual policy update on accepted data nee + then Estimate the objective on accepted samples (e.g., GRPO/PPO surrogate):
J(Ø) ogctbjectivel US, b.).
// Poliic icy Gradient and parameter update
0- VoJ
return 0
27

---

Published as a conference paper at ICLR 2026
G REPRODUCIBILITY STATEMENT
Code Release. To ensure the reproducibility of our research findings, we release our code at
https: /ttp:://ittin /ht.hhb....uugggg
1m Our implementation based on
PyTorch (Paszke et al. 2017) and HuggingFacc (Wolf et al. 2020). All baselines are available in the released code. We also plan to publish all the experiments logs in WandB (Biewald 2020).
Reproducibility. We detail our methodology in Sections 4.1 and and our experimental setup in Section 6 We provide all hyperparameters used in this work in Appendix For all experiments in this paper, report the results over five seeds with standard errors. For the MATH benchmark, we report in-training performance every step, while for the TEST benchmark set we evaluate check- points every 10 learning steps. For better visualization, we applied smoothing with exponential moving average on the curves. All datasets are open-source and available online for academic use.
Compute Resources. We execute all RL experiments using 4 NVIDIA H100 GPUs. Each seed in regime with aggressive updates takes approximately hours, while the standard regime takes approximately one day. Evaluation is done separately in the same hardware, taking approximately 90 minutes per seed.
LLM Usage Details. We use LLMs for paper writing to improve grammar, enhance clarity and writing flow, and assist with code and mathematical iterations. All outputs generated by the LLMs were thoroughly reviewed and verified by the authors to ensure factual accuracy and correctness.
28

---

Published as a conference paper at ICLR 2026
H COMPUTATIONAL COST ANALYSIS
Execution Time. T repoors breakdown of CAPO's execution time, including both the model estimations and the masking process. The table shows the average time (in seconds) of each oper- ation, averaged over all learning iterations, measured on our NVIDIA 4xH100 hardware. The total learning iteration time include LLM generations and forward and backward passes. We find that CAPO contributes less than 3% of the total step time in a learning iteration, resulting in minimal training overhead. Most of the cost arises from computing the Adam gradient and updating its mo- ments, since this also requires computing batch gradients on sparse representations. Lastly, the cost of computing the mask is minimal, below 0.01 seconds.
Memory cost. CAPO uses only volatile GPU memory, since all operations are transient and ten- sors are discarded after the masking generation. The main memory usage comes from maintaining token-level gradient tensors, which have shape (N, T, K, D), corresponding to batch size, com- pletion length, top-K K probabilities, and the number parameters experiments, with N 24, T 1024, K 50, and D 896, this amounts to volatile memory footprint of approximately GB, which minimal given the scale of LLM training. For compar- ison, this is significantly less expensive than performing KI. regularization, which requires storing an additional copy of the LLM in memory for the reference policy.
the last-layer model. our
Step
Avg. Time (s)
% of Total
Learning Iteration (Total)
135.84
100.00%
LLM Generations
55.50
40.85%
Total CAPO time
3.99
2.94%
Compute token-level gradients
0.04
0.03%
Compute Adam token gradients
0.31
0.38%
Compute & log mn
0.09
0.07%
Compute & log m
0.01
0.01%
Update Adam Moments
3.34
2.46%
Compute Hessian Mask Compute Fisher Mask
0.00
0.00%
0.00
0.00%
Table 1: Breakdown of the execution time of CAPO. CAPO contributes less than 3% of the total step time, resulting in minimal overhead relative to standard training.
29

---

Published as a conference paper at ICLR 2026
I ADDITIONAL EXPERIMENTS
Ablation of the Optimizer Model. We conducted an ablation study on the impact of the optimizer representation. This choice reflects a trade-off between step accuracy and computational cost: SGD is cheaper, but the LLM policy is optimized with Adam. Figure 6 shows the results on the MATH dataset. For CAPO, representing the optimizer with either SGD or Adam yields similar perfor- mance. However, for Dr.CAP and ReinCAPO, the SGD variant is insufficient to prevent policy collapse. This suggests that matching the optimizer representation provides more robust choice across different setups.
CAPO
DICAPO
Rein CAPO
0.80
0.80
0.80
0.75
MATH Accuracy
0.75
0.70
0.65
0.70
0.60
0.65
0.65
0.60
0.60
0.50
0.0
0.2
0.6
0.2
0.4
0.8
0.4
0.8 1.0
ning Completions
x10
Training Completions K10*
Training Completions x10*
Figure 6: Ablation study of the optimizer model. For CAPO, both representations yield similar performance, whereas for Dr.CAPO and ReinCAPO, only the Adam-based representation prevents policy collapse, indicating that matching the optimizer provides more robust choice across setups.
Is PPO clipping enough to ensure stability? PPO clipping (Schulman al. 2017) is heuristic designed to prevent large updates by clipping the probability ratio between the current policy and the old policy that collected the data. This raises the question of whether clipping alone is sufficient to avoid policy collapse in our LLM setup. We note that clipping is primarily intended to facilitate off-policy updates, whereas experiments with on-policy data already reveal instability in current RL methods. Nevertheless, we conducted additional experiments using off off-policy data reused for iterations under different clipping ratios. Figure 7 shows results for two setups:
2 (minimal
off-policy shift) and I = 5 (moderate shift). We find that the standard clipping ratio (e = 0.2) does prevent collapse. More aggressive ratios alleviate instabilities but reduce performance, likely due to the strong bias introduced in the gradients. This trade-off becomes more pronounced as t increases.
KL Clipping,
KL Clipping, t
0.80
0.80
MATH Accuracy
0.75
0.75
0.65
0.65
0.60
0.60
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Gradient Steps
x10²
Gradient Steps
x10²
GRPO
CAPO
0.0.2
=
0.002
Figure/ Effect "PPO clipping" on GRPO stability. Standard clipping 0.2) fails to prevent collapse, while more aggressive ratios improve stability but reduce overall performance, with the trade-off becoming more severe as increases.
30

---

Published as a conference paper at ICLR 2026
MATH Accuracy
Max Gradient Norm
0.80
10a
Max Gradient Norm
0.75
364
10a
0.70
10
0.65
100
0.60
0.0
0.2
0.4
0.6
0.8
1.0
GRPO
CAPO
8=0.01
6=0.1 1
1,0
Gradient Steps
x10
GRPO
CAPO
0=0.01
B=0.1
B=1.0
Figure 8: Effect of KL regularization on GRPO stability. (Left) Accuracy on the MATH dataset under different levels of KL regularization. Stronger regularization (3 1.0) reduces instability but degrades performance. (Right) Maximum gradient norms (before clipping), averaged across seeds. KL regularization produces unbounded gradients that may drive the optimization into unstable re- gions.
Is KL regularization enough to ensure stability? Another common strategy to mitigate instabili- ties is to add KL regularizer that penalizes deviations from the base policy (see Equation 33. The rationale is that keeping the policy close to the base model may prevent large distributional shifts, such as those associated with policy collapse. In Figure 8 (left), we test different levels of regular- ization. We observe trend similar to clipping: only stronger regularization (3 1.0) helps prevent catastrophic updates, but at the cost of performance.
A more fundamental limitation of KL regularization becomes evident when examining its gradient:
Ta 8) buse (a 8) + +1]
V DKDKL (Ta || Tbsse) E Eand', ammty Vg log To(a s) log
(83)
Differentiating through the KL term introduces multiplicative log factor, which can produce un- bounded gradients. More concretely, as Tasee s) + O, the gradient magnitude diverges, effec- tively "exploding" the LLM policy gradient. We observe this empirically in Figure & (right), which shows the maximum gradient norms (before gradient clipping) over training, averaged across seeds. While gradient clipping reduce the gradient's 's magnitude, it does not alter its direction, which may still drive the optimization into unstable regions.
Finally, there are also practical drawbacks to KL regularization. First, it requires storing full copy of the base model in memory, which has led prior work to abandon the technique (Liu et al. (002bb). Second, differentiating KL estimates as loss functions typically yields biased approximations of the true KL gradient (Tang & Munos 2025)
31

---

Published as a conference paper at ICLR 2026
J HYPERPARAMETERS
In this section, we present the hyperparameters used our experiments. Table lists the hyperpa- rameters common to all training configurations and algorithms. Table 3 specifies the learning rate and batch size for the conservative and aggressive setups. Finally, Table reports the hyperparame- ters specific to curvature-awar masking, along with their values for each method. Due to compute budget constraints, we performed manual hyperparameter tuning, primarily searching across differ- ent orders of magnitude of both 8F and dr. For simplicity, we implemented single symmetric threshold for the Hessian, i.e., rejecting samples outside the interval -n < < бн-
Hyperparameter
Value
LLM Generation
Max Prompt Length
512
Max Completion Length
1024
Num Generations per Prompt
8
Temperature
0.9
Training
Gradient Steps
100
Warmup Ratio
0.1
Iterations per Batch
Optimizer LR Scheduler
Adam
Cosine
KI. 8
0.0
Table 2: Training Hyperparameters.
Hyperparameter
Standard Setup
Aggressive Setup
Learning Rate
3 x 10
1.5 x 10
Total Batch Size
1152
96
Table 3: Hyperparameters for the standard (conservative) and aggressive regimes.
Hyperparameter
CAPO 10 2
Dr.CAPO
ReinCAPO
Hessian 6н
5 10
10
Fisher 62
10
10-3
10
Table 4: Curvature-aware masking thresholds for CAPO, Dr.CAPO and ReinCAPO.
32

---

Published as a conference paper at ICLR 2026
K MONOTONIC POLICY IMPROVEMENT UNDER CAPO IN THE UNDISCOUNTED, FINITE-HORIZON SETTING
Appendix E. formalizes the conditions under which CAPO guarantees monotonic improvement in the standard discounted, infinite-horizon setting. Although this formulation is general and aligned with prior RL literature, this section extends the analysis the undiscounted, finite-horizon setting, which better reflects the LLM reasoning setup and is more consistent with the assumptions underlying practical algorithms such as GRPO.
For this analysis, we consider mnite-horizon Markov decision process (MDP) with horizon T E N, state space S. action space A. transition kernel P(s' ss a). reward function R S x A + E, and initial state distribution Po. A (stochastic) policy и is conditional distribution x(a s) over actions given states, The return of policy x is given by:
Σ R(st, at)
J(x) = E
(84)
Furthermore, define the advantage function as A a) (s, a) (s). For second
policy x', we also define the n' a'aveeagedd advantage of π at state s: A (s) == (a [A A (s. a)).
Lemma K.1 (Performance Difference Lemma, Finite Horizon, 7 = 1). Let π and x be two policies. Then
T-1 Σ E Es [A5 (s)]
- ((T)
(85)
where d ds,c(s) = (s s) denotes the time-t state-marginal under
Proof.
We start from the identity Q. a) r(., a) + E Be Rearranging.
r(s, a) Qx(s, a) E. [V A= + V*(s) E, (v+(V=(s'))-
(86)
Consider a trajectory (80, 87
generated by policy n'. Then
T-1 Σ
T-1
r(s a:)
Σ t=
(z), + (s.) E[V((s s/ ail.
(87)
t=
Taking expectation under nt' and using the law of total expectation,
T-1
J( = E,p
Σ (x an)
- Ex
Σ
V=((s) - va(oeta)
(88)
1=0
where V= = 0 by definition. The second sum telescopes:
Vi(+s) - V=(+++1)= V= - V=(3t) V=(so).
(89)
0=0
Thus,
J(TT) = E,"
Σ An (s),
+E Esprps (V (so))
(90)
s(x)
Therefore,
T-1
- - = Σ
E (Ax
(91)
1=
We can rewrite each term as
Ex - (At a+) Bender (a =
=B Bswdgr,1A5 (8))
(92)
which proves the claimed identity.
33

---

Published as a conference paper at ICLR 2026
We now bound the difference between the state marginals d and dz, in terms of how different the policies are. For 2 O, we first define the policy-induced transition kernels:
Pa(s' 6) Σ M s)P(s 8,a),
(a (8) Σ n' sa|s))(( s) 18,a).
(93)
Then o Ps and of "'41 PP..
Lemma K.2 State-Distributicn Shift Bound, Finite Horizon). Let IT, nt be two policies with the same initial state distribution d=.0 dx Po- Then, for all t O..... T 1,
t-1
Σ |d dz S Banden
x'(. s))
(94)
Proof. Define the difference vector òū = dx',t dz Then:
de+1 = dr',t11 d
d d
= (drr, dz,1) Pr' + dr -T) Pr)
b Px. (Pr P7).
(95)
Since Px is row-stochastic, Pr Pr. 1|2 1 |1 Next, we bound the term dr.i(P - Px). Let
w == (P. - P,), so w(s') = E. dx (s) Px (s' 8) Ps s)) Then:
|w|1
Σ (s'')| Σ Σ d (s) s) Pa 1s))
ΣΣ dr. Pr s) Pp(s' [s)
Ed.((s) Σ (ar(s') (8) - - |s]
E (Px ( Pa(,1
(96)
For each fixed s, using (s 8) - Ps(s' 8) = (x s) x(a ss) s, a) and the fact
that Ex, P(s' s, a) = 1, we obtain:
[ (-18) Pa( (6)1))111 =
Σ (a 8) - m(a s)) (s,a)
| [s) a(a P)) |s,a)
Σ It((a| 8) - x(a s)
2Drv(a s') ()).
(97)
Hence ||||| 2 ds.r(s)DDv(tt (), ( s)). Combining these two bounds and using the triangle inequality,
||8
||8.Pr. + dx. Pa)ll
< || + (P.. - P=)
< ||ōel1 + +04--
(98)
By definition, d dr so d = and ||ó011 I = O. Unrolling the recursion:
1
|| 2
7'(- (1)))] Eandu.k Drv(T(: 1),
(99)
k=
34

---

Published as a conference paper at ICLR 2026
We now define surrogate objective based the reference policy æ and the state distributions d Lemma K.3 (Surrogate-True Performance Gap, Finite Horizon), For any policies π and T', with DRL ( (') the average forward KL under da.
1)((2T 1))
J(T') 2 [x(x) - C. CVDKL(TTT))
C== TV
(100)
3
== where (AT(s, [)) S E with e finite, and L (x' == ((T) Ei= Egdda. [AT
Proof. By Lemma K.1 J( - [ Et= Londy , [AA (s)]. Subtracting the surrogate:
T-1
J(z') (T²(-²)
Σ
Ba A7 (s) Bende, A (s)
T-1
Σ
Σ (d): - da, (x)) A (s).
(101)
0
Taking absolute values and using (s)l I c and applying Lemma K.2
(s)
[J(7') - Lx(TT) $
Σ el eelldwr - dn,ell S
Eandayx
(.),(()))
T-1
2c
Σ
Egnda Drv( 18,6)))))((( ((s())))))
[s))]
Σ
k=0
T-1
2e
Σ
(T - 1 - k) Bonde, (TV(T( (s), s))}
k=0
(102)
For the KL-based bound, Pinsker's inequality and Jensen's 's inequality. For each t:
Egndd, Drv( s), [s) < Bandr, V (KL( ()) || IT( (s))
1 } Bandnn (DRL(T( (8) [[ 7( {)]]
(103)
For conciseness, we define D/ := DKL( ( 8) || x'(. (s). Then:
T-1 Σ -T- k)v 1 V2c Σ DA DR.
T-1
[J(T') - L LT S 2c
(104)
k=
where we have set by ;= T- k. By SScwaaySchwarz,
T-1
T-1 Σ
T-1
Σ by D2 -
DD
D
(105)
0=0
k=
k=
We note that
T-1
T-1
TT 1)T(2T -
T-1
Σ
-
b
-
Σ km
Dj = T DKL.
(106)
6
km
j=0
Therefore
IT - -1)((22 - 1)
[J(T') [T7( S v2eg
VT
6
IT -1( 1)
e / DKL
T
(107)
3
n
35

---

Published as a conference paper at ICLR 2026
The proof of Theorem 5.1 for the fmite-horizon setting follows exactly the one in Appendix but applying Lemma K.3 instead of Lemmal E.1
Infinite-Horizon vs. Finite-Horizon bounds. We highlight that, in both settings, the final guarantee takes the same form J (TT - (mT)
CVOO, where C'
27 [I~ 2 for the infinite-
horizon case, and C TV [T- (T- e for the finite-horizon case. In both cases, the constant C scales as (HH))) where Hell denotes the effective horizon: Heπ = I in the finite-horizon setting, and Helf = T- in the infinitehhorizon setting. Practically, this implies that both bounds are equally tight within their respective regimes.
36

---

Published as a conference paper at ICLR 2026
L A CLOSER LOOK AT MODEL ESTIMATES mh AND THE KL POLICY SHIFT
In this section, we analyze the relationship between the model's estimate of directional Fisher cur- vature, f. and the actual policy shift induced by an update, measured by DKL ( To+ Our goals are two-fold: (i) to clarify what CAPO requires from the underlying model in order to approx- imate a trust-region and to assess how well this approximation holds, and (ii) to examine the impact of CAPO's s updates on the true change in policy.
Does CAPO require fully calibrated model? Although well-calibrated estimates are sufficient condition for CAPO's data-selection mechanism to function effectively, they are not necessary. To illustrate this, consider a simple case where the estimated directional Fisher curvature satisfies my = > 0, where a >> 1 or << 1. Such model is clearly miscalibrated, yet
DK [Tell it preserves a strong correlation with the true policy shift. In CAPO, if we aim to enforce the trust- region condition DKI. (To, TT+A0) < 6, we can simply set the Fisher-threshold to 0; = ad, which recovers the desired constraint. More generally, CAPO only requires that the estimates be monotonically correlated with the true policy change, that large prospective shifts (those most likely to trigger instability or collapse) are reliably identified.
A natural way evaluate the quality of the model's estimates is measure their correlation with the true policy changes. Although we do not have direct access to this quantity, we can estimate it via samples. In particular, the KL divergence can be reliably estimated using a standard Monte Carlo estimator, which has manageable variance and leverages token-level information. We therefore compute these estimates and report the resulting Spearman correlations in the Table 5 where the is evaluated under both GRPO and CAPO updates at both token and global level. We find that the model estimates exhibit moderately strong correlation with the actual policy change, indicating consistent monotonic relationship. Notably, this correlation remains high under both GRPO and CAPO, suggesting that the estimates are meaningful even when they are not used intervene in the update.
Estimate
(GRPO)
p (CAPO)
mp (Token)
0.622
0.459
mr (Global)
0.596
0.498
Table 5: Spearman correlations p between Fisher directional curvature estimates mh and the estimated policy change DKL. (To (Too|TT We report correlations for both GRPO and CAPO updates. The results indicate that the estimates rh maintain a consistent monotonic relationship with the true policy shift across algorithms, reliable identifying the scale of the policy shifts
Ultimately, does CAPO induce a bound on the true DKL (Te ? In Figure 9 we present the policy shifts over the course of training for both algorithms. GRPO frequently presents peaked shifts, which are often associ- ated with unstable or overly aggressive updates. In contrast, CAPO generally maintains stable, small shifts, suggesting that it is effective in practically implementing a trust-region behav- ior throughout training.
0.10
GRPO
CAPO
0.08
Da o0)
0.06
0.04
0.02
0.00
0
20
40
60
80
100
Gradient Steps
Figure 9: Estimated policy KI. shifts dur- ing training. GRPO exhibits frequent sharp spikes in policy divergence, indicative of unstable updates, whereas CAPO maintains consistently small shifts, reflecting its ability to enforce trust- region-like behavior throughout training.
37

---

Published as a conference paper at ICLR 2026
M FURTHER QUESTIONS
This Appendix presents additional clarification questions aimed at improving the understanding of the proposed method and experiments. These questions were raised during the peer-review process, and we refer to the OpenReview page for the full discussion.
What is the effect of token selection in the sample efficiency evaluation? In Figure TO plot
the accuracy curves (analogous to Figs. 1 and 2) as a function of the accepted tokens. We observe that these curves closely resemble those obtained when accuracy is plotted against the number of completions. This suggests that the effect of masking on the total number of generated (and ac- cepted) tokens is small, consistent with the rejection rates reported in Figure 5. also indicates that the learned policies behave similarly in terms of token generation, showing that CAPO improves training sample efficiency without incurring additional inference-ting time costs.
MATH Accuracy
TEST Accuracy
0.42
CAPO
REINFORCE
GRPO GRPO (A)
REINFORCE
GRPD (A)
DrGRPO
DrGRPO
0.80
0.40
0.75
0.38
GRPO 3,2 x00
GRPO at = 3.2 10'
0.70
0.36
0.65
0.34
0.60
0.32
0
1
3 Accepted Tokens
2
5 x10
Accepted Tokens
x10
Figure 10: Sample efficiency curves as function of the number of accepted tokens. The trends closely match those obtained when using the number of completions, indicating that masking has minimal impact on token generation and that CAPO improves sample efficiency without added inference cost.
What the similarities between CAPO and TRPO? What are the differences? In terms of similarities, both CAPO and TRPO share the same motivation: devise conservative optimization procedure that implements a safe optimization region, typically expressed as a KL ball constraint. This idea predates TRPO, with its roots in natural gradient methods from optimization literature Amari 1998); Amari et al. (1995). What both CAPO and TRPO do is devise practical instantia- tions of the natural gradient that is suitable for their respective problem settings.
Methodologically, TRPO incorporates only the Fisher matrix in its updates, relying on a first-order approximation of the objective. In contrast, CAPO additionally leverages second-order curvature information of the objective through its Hessian, as shown in Equation 5 and further incorporated in the theoretical development in Equation 68 The main difference, however, lies in the implementa- tion, which crucially leads to different scalability properties.
Senjjugat-GGradieee (CG) algorithm to approx- imate the natural gradient step without fully materializing the Fisher matrix. Then, TRPO employs line search algorithm to solve the constrained optimization problem. The CG algorithm involves maintaining five vectors of size d (the gradient, current iterate, the residual, the search direction, and the -ector buffer), where is the number of parameters in the policy. While this memory cost is feasible for small deep networks (as usual in traditional Deep RL research), is prohibitive for LLM scale, where d is in the billions.
TRPO incorporates the Fisher matrix by employing
Furthermore, the CG algorithm is iterative and each iteration costs roughly the same as a backward pass, unless you sacrifice your Fisher matrix estimation by subsampling data. TRPO uses ten itera-
38

---

Published as a conference paper at ICLR 2026
tions. Considering the execution time in our setup (Appendix H), this overhead is also prohibitive. Lastly, the line search algorithm requires M additional forward passes in the whole batch (M the number of search trials), which is also substantial cost in setup (also illustrated in Appendix H). Overall, TRPO's memory and execution costs are prohibitive to LLM scale. CAPO, in contrast, leverages the last layer model and the optimizations described in Section 4.1, resulting in much lower costs, as evaluated in Table of Appendix H
In summary, while TRPO and CAPO share the same motivation and draw from the same semi- nal work on natural gradients, CAPO offers a formulation that scales to the memory and compute demands of LLM policies.
39


Reviews:
[{"review_id": "PjXrcTq90a", "reviewer": "Reviewer_NcAd", "summary": "The paper proposes a last-layer computational model, it tracks directional curvature for both the objective (Hessian) and the policy (Fisher). And uses these diagnostics to reject unstable data subsets/tokens before updating a large language model with policy gradients. The algorithm Curvature-Aware Policy Optimization (CAPO), claims monotonic improvement under local trust-region tests and reports up to $30\\times$ fewer completions to reach target accuracy on math-reasoning benchmarks. While rejecting $<8\\%$ tokens and incurring $<3\\%$ runtime overhead. Experiments are on MATH and a pooled “TEST” suite; GRPO and variants are the primary baselines, with both conservative and “aggressive” update regimes.", "questions": "1. Provide a $\\gamma=1$ analogue of Theorem 5.1 (or a trust-region argument) with finite constants, alternatively, empirically show that your $m_F$ threshold enforces a measured average $D_{\\mathrm{KL}}$ bound per update across training. \n\n2. Calibrate the surrogate $m_F(\\Delta\\psi)$ and $m_H(\\Delta\\psi)$ against the true forward KL and objective change measured on the LLM before/after updates; report correlation and calibration plots. \n3. Report sample efficiency vs. accepted tokens and vs. effective gradient norm/variance, not just vs. completions. To isolate the effect of masking.  \n4. Include sensitivity sweeps for $(\\delta_F,\\delta_H,\\delta_H^{\\text{high}})$; explain how choices interact with optimizer modeling (SGD vs. Adam) and aggressive vs. conservative regimes.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "1. Targets a high-impact problem: stabilizing policy-gradient training for LLM reasoning under aggressive hyperparameters. \n2. Provides a tractable, analytic surrogate for curvature using a last-layer model, with clear formulas and pseudocode (Algorithm 1).", "weaknesses": "1. The monotonic-improvement guarantee uses a constant $C=\\frac{2\\gamma}{(1-\\gamma)^2}\\,\\epsilon/\\sqrt{2}$ (Eq. (13)), the paper also states that GRPO’s per-trajectory normalization “effectively assumes $\\gamma=1$.” With $\\gamma \\to 1$, $C$ blows up, making the bound vacuous precisely for the main regime evaluated. No $\\gamma=1$ reformulation or trust-region alternative is provided, and there is no calibration showing that the chosen $m_F$ thresholds enforce a true KL trust region.  \n\n2. The acceptance test relies on thresholds $(\\delta_H,\\delta_F,\\delta_H^{\\text{high}})$ and bounded step norms/curvatures ($M,r,\\epsilon$), but the paper manually tunes thresholds and never verifies that the true (not surrogate) preconditions of Theorem 5.1 hold in practice.  \n\n4. The headline “$30\\times$ fewer completions” is measured in a regime where baselines catastrophically collapse, robust efficiency claims should hold against *stable* baselines at their best hyperparameters (e.g., larger batch/smaller LR, PPO-style clipping with tuned schedules, or natural-gradient/K-FAC/TRPO-style trust regions). Such baselines and sweeps are absent.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PjOMm60SeS", "reviewer": "Reviewer_1Let", "summary": "This paper introduces Curvature-Aware Policy Optimization CAPO, a reinforcement learning algorithm designed to stabilize policy gradient training for reasoning-focused LLMs. The authors first propose a computationally efficient approximation of second-order geometry using a last-layer curvature model, enabling estimation of directional Hessian and Fisher Information Matrix contributions without computing full curvature. They then use these estimates to detect unstable samples and mask them during gradient computation, enforcing trust-region-like constraints. Theoretically, they prove monotonic policy improvement under bounded curvature conditions. Empirically, CAPO is applied to Qwen2.5-Math-7B, showing up to 30× improvement in sample efficiency under aggressive update regimes where GRPO and other baselines collapse.", "questions": "Please refer to weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a curvature-aware method that is theoretically grounded.\n- It has a clear trust-region interpretation and monotonic improvement guarantee.\n- Empirically, it has substantial gains in sample efficiency with minimal token masking.\n- It is very relevant for real LLM training pipelines as it can be generalized to other methods including Dr.GRPO and REINFORCE.\n\nOverall, I think this paper proposes an interesting and novel algorithm that shows potential in small scale experiments. The method is well theoretically-grounded.", "weaknesses": "- Evaluation scale is limited to a single model family (Qwen2.5-Math-7B) and trained on a single dataset. Larger scale experiments could be  helpful to showcase the benefits of CAPO.\n- The algorithm adds complexity to already complex RL frameworks with added components and hyper-parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "djPnspMDir", "reviewer": "Reviewer_Fbr1", "summary": "This paper proposes an algorithm named CAPO which uses the second-order information of the policy objective for LLM reasoning to enhance the training stability.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ This problem studied in this paper is well-motivated and important in RL training.\n\n+ The presentation is clear, straight-forward and easy to follow.\n\n+ The method proposed in the paper seems reasonable to me, and the empirical results support the approach.", "weaknesses": "+ The main contribution of CAPO is that it incorporates two additional constraints given in Equation 10 which are based on the approximated Hessian and the Fisher matrix, to the original GRPO/PPO surrogate objective. This idea, however, has been proposed in the TRPO paper (https://arxiv.org/pdf/1502.05477, Section 6 and Appendix C). While I recognize that CAPO only estimates those second-order matrices with respect to the last softmax layer, which is a meaningful improvement compared to the original TRPO method that approximates the Hessian of all the model parameters, it would be much more convincing if it can show that CAPO performs well without other constraints such as ratio clipping and KL penalty.\n\nTherefore, I am completely neutral on this work despite my rating as there is no such an option at ICLR. \n\n\nMinor issue: I am unsure if Figure 7 is really about the comparison of CAPO and PPO as the legend says GRPO. The very difference between GRPO and PPO is not the use of ratio clipping, but whether advantages are estimated using a critic model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)