arXiv CSopen access

Cost-Aware Learning

Cost-Aware Learning Clara Mohri*, §

Amir Globerson†, ‡

Tomer Koren†, ‡

Haim Kaplan†, ‡

Yishay Mansour†, ‡

arXiv:2604.28020v1 [cs.LG] 30 Apr 2026

May 1, 2026

Abstract We consider the problem of Cost-Aware Learning, where sampling different component functions of a finite-sum objective incurs different costs. The objective is to reach a target error while minimizing the total cost. First, we propose the Cost-Aware Stochastic Gradient Descent algorithm for convex functions, and derive its cost complexity to attain an error of ϵ. Furthermore, we establish a lower bound for this setting and provide a subset selection algorithm to further reduce the cost of training. We apply our theoretical insights to reinforcement learning with language models, where the computational cost of policy gradients varies with sequence length. To this end, we introduce Cost-Aware GRPO, an algorithm designed to reduce the cost of policy optimization while preserving performance. Empirical results on 1.5B and 8B LLMs demonstrate that our approach reduces the tokens used in policy optimization by up to about 30% while matching or exceeding baseline accuracy.

1

Introduction

Stochastic gradient descent (SGD) theory typically focuses on the number of steps required to attain a certain error tolerance. Implicitly, this analysis assumes that taking a gradient step on any data point incurs an identical computational cost. However, in modern foundation model training, the assumption of uniform cost is not always true. For example, in Reinforcement Learning from Human Feedback (RLHF) [Ouyang et al., 2022], computing the policy gradient on a long reasoning trace requires significantly more FLOPs than on a concise response. Motivated by the computational gap, we consider the problem of Cost-Aware Learning, where each training sample incurs a known, non-uniform computational cost. While our primary objective remains to learn an estimator with small error, this setting introduces a secondary objective: minimizing the cost of training. In particular, we consider a fixed training set on which the objective is to optimize a convex function. Each sample has an associated cost which is known, and incurred only when using that sample with the learning algorithm. We derive an algorithm called Cost-Aware SGD, which minimizes the cost of training to attain ϵ error. This algorithm leverages principles of importance sampling [Kloek and Van Dijk, 1978, Beygelzimer et al., 2009] to achieve small error while minimizing the associated cost. We further extend our analysis to perform dataset selection prior to training; this introduces a controlled bias into the estimator in exchange for a potentially significant reduction in overall training cost. Having established the optimal strategy in theory, we turn our attention to LLM post-training. In particular, we integrate principles of Cost-Aware Learning with GRPO [Shao et al., 2024] to introduce the Cost-Aware GRPO algorithm. In standard language model training, the number of FLOPs required to train ∗ Harvard University

† Tel Aviv University ‡ Google Research

§ Work done while Clara Mohri was a Student Researcher at Google Research. Correspondence to: [email protected].

1

Accuracy vs. Token Count (Qwen3-8B)

60

Accuracy (%)

55

Method

GRPO CA-GRPO

50 45

(1) 17% fewer tokens

40

(2) 51% fewer tokens

35 0

100M

200M

Cumulative Tokens

300M

400M

Figure 1: Qwen3-8B Base training using GRPO and Cost-Aware GRPO (CA-GRPO). We evaluate on AIME1983-2024 (pass@1/mean@32) and report cumulative tokens used in policy gradient computation. We plot every 20 training steps until peak accuracy is attained for both. CA-GRPO requires: (1) 17% fewer tokens to reach its absolute peak than GRPO requires to come within 2% of that performance, (2) 51% fewer tokens to match or exceed the peak accuracy of GRPO. on a sequence of tokens depends linearly on the number of tokens in that sequence. Hence, we consider the cost associated with each sample to be the number of tokens in a prompt and response. As such, Cost-Aware GRPO focuses on preserving the accuracy of the learned model while minimizing the FLOPs used in policy optimization. In Cost-Aware GRPO, we use the sampling strategy derived from our theoretical insights rather than standard sampling for policy optimization. At a high level, this sampling strategy aims to favor training on responses which have larger gradient norms and fewer tokens. We show that applying the optimal sampling distribution to the GRPO setting allows us to train a model with the same performance while reducing the FLOPs necessary to compute policy gradients throughout training. The contributions of this paper are as follows: 1. Formulation of Cost-Aware Learning. We formulate cost-aware learning in the convex and strongly convex setting. 2. Cost-Aware SGD. We propose the Cost-Aware SGD algorithm and provide analysis. We analyze the cost required to attain expected error ϵ, and compare our algorithm to two baselines. We provide synthetic validation showing the benefits of using Cost-Aware SGD. Further, we provide a lower bound and suggest a subset selection solution to further minimize the cost of training. 3. Cost-Aware GRPO. We apply our theoretical insights to GRPO and propose Cost-Aware GRPO. We run experiments with Qwen2.5-Math-1.5B-Instruct [Yang et al., 2024] and Qwen3-8B Base [Yang et al., 2025], and evaluate performance on various benchmarks to show that our method consistently improves token efficiency in policy optimization without compromising on performance.

2

2

Related Work

Importance sampling for SGD. A significant body of work focuses on accelerating SGD convergence via importance sampling, primarily through variance reduction [Zhao and Zhang, 2015, Needell et al., 2014]. A foundational result by Alain et al. [2015] established that the optimal sampling distribution for variance reduction assigns probabilities proportional to the per-sample gradient norms. In the context of deep learning, where exact gradient norms are expensive to compute, methods have been proposed to approximate these norms using upper bounds or loss values [Katharopoulos and Fleuret, 2018, Johnson and Guestrin, 2018]. However, these approaches operate under a standard optimization framework that treats iterations as the unit of resource consumption. They aim to maximize the information gain per step, implicitly assuming uniform computational cost across samples. Our work generalizes this by considering the cost-aware setting where the aim is to minimize training cost, leading to a sampling distribution that balances gradient magnitude against sample cost. In particular, the work of Chen et al. [2023b] shares similarities with our work by also studying a finite-sum setting with varying sample costs. However, our contributions differ significantly in analysis and scope: while their work focuses on asymptotic behavior, we provide finite-time guarantees complemented by a lower bound. Furthermore, we implement this cost-aware framework with language models via Cost-Aware GRPO. Training Data Pruning/ Selection. There is a large literature focused on pruning a training dataset. Typically, the objective is to obtain a better predictor or maintain performance with fewer samples. Related works include [Sorscher et al., 2022, Paul et al., 2021, Xie et al., 2023, Chen et al., 2023a]. Our work differs from these as it explicitly considers a variable cost per sample. Active learning. There is a significant body of work studying the problem of active learning, which considers when to request more information about training samples. We refer the reader to [Dasgupta, 2011] for an in-depth overview of the literature. Most related to our work, Beygelzimer et al. [2009], Cortes et al. [2020, 2019] use importance weighting to probabilistically select samples for which to query a label. Generally, these algorithms fix a hypothesis set a priori, in which hypotheses are eliminated as more data is acquired. Our work shares the similarity of deciding when to request more information, however, in contrast, our work requests the gradient of a sample. Furthermore, we do not use a current hypothesis set in order to make these decisions. RL for Language Models. While policy gradient methods [Sutton et al., 1999, Kakade, 2001, Schulman et al., 2015] have been used in many settings, the works of Ouyang et al. [2022], Dai et al. [2023] have garnered significant attention for their application to LLM post-training. Various methods such as DPO [Rafailov et al., 2023], PPO [Schulman et al., 2017, Zheng et al., 2023], and GRPO [Shao et al., 2024] have been shown to be extremely effective in improving accuracy on standard benchmarks. However, these methods incur substantial computational and memory costs. Several methods aim to improve data efficiency: [Wang et al., 2025, Fatemi et al., 2025, Li et al., 2025, Yu et al., 2025, Xu et al., 2025, Zheng et al., 2025]. Yu et al. [2025] introduces, among other contributions, a mechanism to continue resampling from a model until a sufficient amount of non-zero advantage responses have been generated, showing that this data helps the model achieve better performance. Xu et al. [2025] attempts to subsample the generations of a model in an effort to train on less data. Zheng et al. [2025] keeps track of the history of various prompts in an effort to remove low-signal prompts. Our work contrasts with these by using importance reweighting and taking into account the computational cost of sequences when making decisions.

3

Problem Formulation

Pn We consider minimizing a convex function f in the finite-sum setting where f (x) = n1 i=1 fi (x). We assume that f is minimized at x∗ , and that each component fi is a convex and Gi -Lipschitz function. In the cost-aware setting, each component i has an evaluation cost ci ≥ 0 for querying its gradient oracle. Namely, 3

Algorithm 1 Cost-Aware SGD 1: Input: Initial point x1 , iterations T , step size η, sampling distribution p∗ .

2: for t = 1 to T do 3: Sample index it ∼ p∗ 4: Query ∇fit (xt ), incur cost cit . 5: Compute importance-weighted stochastic gradient: 6: g̃t = np1 ∇fit (xt ) it

Update parameter: xt+1 = ΠX (xt − ηg̃t ) 8: end for PT 9: Output: x̄T = T1 t=1 xt (or suffix averaging solution for strongly convex case [Rakhlin et al., 2011]) 7:

upon querying a particular sample i at a point x, the algorithm obtains the exact gradient ∇fi (x) and incurs cost ci . We assume that the learning algorithm has access to Gi and ci for all samples i. For a given ϵ > 0, the objective of a (possibly randomized) cost-aware learning algorithm is to incur the minimum possible cost to obtain x̂ such that E[f (x̂) − f (x∗ )] ≤ ϵ.

4

Cost-Aware SGD

We address the problem presented in Section 3 via Algorithm 1, which combines importance weighting with projected SGD. We analyze the convergence behavior to derive the optimal sampling distribution p∗ that minimizes the total expected cost to reach a target accuracy. We compare this with other sampling strategies and validate the method synthetically. Furthermore, we provide a lower bound and consider the sub-optimality of using different proxies as a strategy for sampling. In Appendix C we extend the analysis to analyze the last iterate and non-convex settings.

4.1

Cost Analysis

Consider a fixed sampling PT distribution p ∈ ∆n . The total cost of the algorithm over T iterations is the random variable KT = t=1 cit . By the linearity of expectation, the expected total cost is: E[KT ] =

T X t=1

E [cit ] = T

it ∼p

n X

pi ci = T · C(p),

(1)

i=1

Pn where C(p) = i=1 pi ci represents the expected cost per iteration. Our goal is to minimize the total expected cost required to reach a target error ϵ. Let T (ϵ, p) denote the number of iterations required to achieve expected error ϵ using distribution p. We seek to find p which minimizes: K(ϵ, p) = T (ϵ, p) · C(p).

(2)

Case 1: General Convex Functions For general convex functions, standard √ convergence results for SGD [Nemirovsky et al., 1983, Bubeck et al., 2015] state that with step size η ∝ 1/ T , the expected suboptimality satisfies: p D S(p) ∗ √ E[f (x̄T ) − f (x )] ≤ , (3) T where S(p) = Ei∼p [∥g̃∥2 ] is the second moment of the gradient estimator, and D is the diameter of the parameter space. Expanding this second moment for importance sampling: S(p) =

n X i=1

pi

∇fi (x) npi 4

n

2

≤

1 X G2i , n2 i=1 pi

(4)

2

where Gi is the Lipschitz constant for component fi . To guarantee error ϵ, we require T = D ϵS(p) . Substituting 2 this into the cost formula:  2  D S(p) K(ϵ, p) = C(p) ∝ S(p)C(p). (5) ϵ2 Case 2: Strongly Convex Functions For f that is µ-strongly convex, SGD with decaying step size ηt = 1/(µt) achieves a rate of O(S(p)/µT ) [Rakhlin et al., 2011]. The iterations required for error ϵ are T = 4S(p) µϵ . The expected total cost under distribution p to attain error ϵ becomes: K(ϵ, p) =



 4S(p) C(p) ∝ S(p)C(p). µϵ

(6)

Crucially, in both cases, minimizing the total cost is equivalent to minimizing the product J(p) = S(p)C(p).

4.2

Derivation of the Optimal Distribution

We now derive the sampling distribution that minimizes this cost objective. Theorem 4.1 (Optimal Cost-Weighted Distribution). The sampling distribution p∗ that minimizes the expected cost to reach error ϵ is: √ Gi / ci ∗ pi = Pn (7) √ . j=1 Gj / cj Under this optimal distribution, the minimum expected costs for general convex and µ-strongly convex functions are, respectively: D2 Kcvx (ϵ, p ) = 2 ϵ ∗

n

1X √ Gi ci n i=1

!2

4 Kstr-cvx (ϵ, p ) = µϵ ∗

,

n

1X √ Gi ci n i=1

!2 .

(8)

Proof. We minimize the product term J(p) = S(p)C(p). By the Cauchy-Schwarz inequality, for vectors √ √ u, v ∈ Rn defined by ui = pi ci and vi = Gi / pi : n X

! pi ci

i=1

n X G2 i

i=1

pi

! = ∥u∥ ∥v∥ ≥ ⟨u, v⟩ = 2

2

2

n X

√

Gi ci

!2 .

(9)

i=1

√ √ √ Equality holds if and only if u ∝ v, which implies pi ci ∝ Gi / pi ,P or pi ∝ Gi / ci . Substituting the √ minimum value of the product back into J(p), we obtain J(p∗ ) = n12 ( Gi ci )2 . Substituting J(p∗ ) into Eq. 5 and Eq. 6 yields the stated costs.

4.3

Cost-Improvement over Baselines

We compare the optimal sampling distribution to two baselines: uniform sampling and the “variance strategy,” which has been shown to optimally reduce variance in other work. 4.3.1

Uniform sampling

We consider the baseline of sampling uniformly. Under uniform sampling, the expected cost to achieve ϵ error is: ! n ! n X D2 X 2 Gi ci . (10) K(ϵ, punif ) = 2 2 ϵ n i=1 i=1

5

4.3.2

Variance strategy

We consider the baseline which adapts to Gi but not to costs. This baseline corresponds to the optimal variance-reduction strategy derived in classical importance sampling literature, which ignores computational costs. It has been established that minimizing the variance of the gradient estimator requires sampling proportionally to the gradient norms [Alain et al., 2015, Zhao and Zhang, 2015, Needell et al., 2014]. We call this the “variance” strategy, and refer to Appendix A.1 for its derivation. With (pvar )i = PGiG , the expected cost under sampling with pvar to achieve ϵ error is: ! n ! n X D2 X K(ϵ, pvar ) = 2 2 Gi Gi ci . ϵ n i=1 i=1

j

j

(11)

Theorem 4.2 (Cost Improvement). Let Kunif (ϵ), Kvar (ϵ), Kopt (ϵ) be the total cost to reach ϵ-accuracy using uniform, adaptive, and optimal sampling, respectively. Then Kopt ≤ Kunif , with the gap determined by the correlation between gradients and costs: P √ ( Gi ci )2 Kopt (ϵ) = P 2 P ≤ 1. (12) Kunif (ϵ) ( Gi )( ci ) In addition, Kopt ≤ Kvar , with the gap determined by the correlation between gradients and costs: P √ ( Gi ci )2 Kopt (ϵ) P = P ≤ 1. Kvar (ϵ) ( Gi )( i Gi ci )

(13)

The proof relies on an application of the Cauchy-Schwarz inequality, which we defer to Appendix A.2. We remark that, when all costs are uniform, there is still an improvement from Kopt over Kunif . This is because we leverage the second moments to minimize variance. With unit costs, Kopt = Kvar . We compare Kvar and Kunif carefully in Appendix A.2

4.4

Synthetic Validation

We evaluate the sampling strategies on a linear least squares task with N = 3000 data points generated randomly in dimension d = 50. The data vectors ai are scaled such that their norms are bounded by L = 10. This geometry ensures that the Lipschitz constants Gi (gradient upper bounds) are bounded but heterogeneous across the dataset. To obtain the gradient of a sample, we have ∇fi (x) = (a⊤ i x − yi )ai . Consequently, the gradient norm is the magnitude of the residual scaled by the data norm, ∥∇fi (x)∥ = 2 |a⊤ i x − yi |∥ai ∥ ≤ ∥ai ∥ · (2D). Evaluation costs ci are assigned randomly from a uniform distribution between 1 and 1000. We compare √ three sampling approaches: Uniform, Variance (pi ∝ Gi ), and Optimal (pi ∝ Gi / ci ). For the latter two, we evaluate both a static version (using global bounds Gi ) and a dynamic version (using instantaneous gradient norms ∥∇fi (xt )∥). The results are shown in Figure 2. • Cost Efficiency: For a fixed error target, the Optimal strategy incurs the least total cost. By weighting samples by the ratio of information to cost, it effectively identifies “computational bargains”—samples that provide significant variance reduction for a low price—while avoiding informative but overpriced data points. • Convergence Rate: The Variance strategy converges the fastest in terms of iteration count. This is expected, as it minimizes the gradient estimator’s variance without regard for computational budget. • Dynamic vs. Static: The dynamic strategies slightly outperform their static counterparts. This confirms that the static Lipschitz bounds Gi are approximations, and exploiting the local geometry of the loss landscape provides an even tighter sampling distribution. 6

Figure 2: Synthetic experiment in which we compare the error with the total training steps and the total cost incurred by the learning algorithm.

4.5

Proxy Sub-Optimality

In practice, it can be prohibitively costly to compute Gi . As such, one may not have access to the perfect sampling distribution p∗ . In this subsection, we consider the sub-optimality of using a proxy for p∗ . We recall that J(p) = S(p)C(p), and that K(ϵ, p) is the expected cost to attain error ϵ under distribution p. In Section 4.1, we showed that for any ϵ, K(ϵ, p) ∝ J(p). Hence, we characterize suboptimality of p′ by comparing J(p′ ) with J(p∗ ). In particular, J(p′ ) represents how cost scales with error for p′ . Hence, the quantity J(p′ )/J(p∗ ) represents the ratio between the cost incurred by p′ and that of p∗ . We recall the definitions for χ2 -divergence and Pearson correlation. The χ2 -divergence is defined as follows P P2 for two distributions P, Q: Dχ2 (P ||Q) = Q − 1. The Pearson correlation between random variables X and Y with expected values µX , µY , and standard deviations σX , σY , respectively, is defined as: ρX,Y =

Cov(X, Y ) . σX σY

In Theorem 4.3, we exactly characterize the sub-optimality due to using a different sampling distribution p′ in place of p∗ ; under a cost-biased projection, the sub-optimality is characterized by the χ2 -divergence of p′ and p∗ . In order to provide further intuition as to when one can use some proxy G′i in place of Gi to compute a proxy sampling distribution, we provide Theorem 4.4. The result implies that, when Gi and G′i have near-1 Pearson correlation, then E[J(p′ )] ≈ J(p∗ ). Both proofs are provided in Appendix A.5. Theorem 4.3 (Sub-optimality gap for distribution proxy). Let p∗ be the optimal cost-aware sampling distribution and p′ be any empirical sampling distribution. We define the cost-biased distribution p̃ corresponding to p as: pi ci p̃i = . C(p) The sub-optimality gap is exactly determined by the Pearson χ2 -divergence between the cost-biased optimal distribution and the cost-biased empirical distribution: J(p′ ) = 1 + Dχ2 (p̃∗ ||p̃′ ). J(p∗ ) Theorem 4.4 (Sub-optimality gap for Gi proxy). Let p′ be the empirical sampling distribution using estimated √ gradient norms G′i , such that p′i ∝ G′i / ci . Assume the estimates follow an additive noise model G′i = Gi + ϵi , where ϵi are independent, zero-mean noise terms with variance σϵ2 , independent of ci . Let ρ be the Pearson 7

correlation coefficient between the true norms Gi and the estimates G′i . Assuming the noise ϵi is relatively small compared to Gi , the expected optimality gap is approximately: E[J(p′ )] ≈1+ J(p∗ )



1 − ρ2 ρ2



2 Pn σG

√

ci i=1 Gi

Pn

i=1 Gi

√ , ci

2 where σG is the variance of the true gradient norms.

4.6

Lower Bound

Having established the upper bounds for Cost-Aware SGD, a natural question is whether it is possible to design a sampling strategy that achieves a lower total cost. In this section, we provide a lower Pnbound√for the class of convex functions. We show that even in this favorable setting, the dependency on ( i=1 Gi ci )2 is fundamental. Theorem 4.5. Fix ϵ, G > 0. For any number of components n ≥ G2 /ϵ2 , any set of nonnegative query costs {ci }ni=1 , and any (possibly randomized) algorithm with an expected error guarantee E[F (x̂) − min F (x)] ≤ ϵ, there exists a convex, G-Lipschitz (i.e., with Gi = G for all i) finite-sum problem instance over X = [−1, 1] such that the algorithm’s expected total query cost is at least:  !2  2 X√ G 1 Ω 2 ci  , ϵ n ∗ i∈S

where S ∗ ⊆ [n] is any set of components such that following cost-uniformity condition holds: √ maxi∈S ∗ ci G P . √ ≤ 1 40ϵ c ∗ j j∈S n

(14)

We defer the proof to Appendix B. The lower bound suggests that there is still room for improvement in cost-aware learning, and that the upper bound established in Theorem 4.1 may not be tight due to its dependence on J(p∗ ), despite matching in the n21ε2 dependence. In particular, it indicates that when the costs are non-uniform (in the sense that Eq. (14) only holds for a proper subset S ∗ of [n]) the algorithm may not necessarily incur some of the higher costs. This motivates a combination of importance-weighted sampling with subset selection, which we consider next.

5

Subset Selection

Pn √ The importance sampling strategy p∗ derived in Section 4 achieves a total cost proportional to ( i=1 Gi ci )2 . While optimal for unbiased estimation, this cost may still be prohibitive if high-cost samples have large gradients. One might therefore consider excluding costly samples entirely, introducing a controlled level of bias to reduce the training cost. We formally model this as aP subset selection problem, where we fix a subset π ⊆ {1, . . . , n} and minimize the proxy objective fπ (x) = n1 i∈π fi (x). Optimizing fπ introduces bias because the algorithm converges toward x∗π = argmin fπ (x) rather than ∗ x = argmin f (x). While this introduces a deterministic bias floor, it lowers the computational burden by restricting evaluations to the subset π. Pn Convergence of SGD on a Restricted Subset. Let f (x) = n1 i=1 fi (x) be convex with component gradients bounded by ∥∇fi (x)∥ ≤ Gi . Let π ⊆ {1, . . . , n} be a fixed subset, and let q be a sampling

8

Figure 3: Synthetic validation of the greedy subset selection algorithm. Γ is a tunable parameter representing the allowable bias budget. As it is increased, the greedy algorithm chooses fewer (and cheaper) samples to train on. distribution supported on π. We define the bias upper bound βπ and the restricted second moment σπ2 (q) as: 1X Gj ≥ ∥∇f (x) − ∇fπ (x)∥, n j ∈π /   2 1 X G2i 2 1 ∇f (x) ≤ σπ (q) = E . i nq i i∼q n2 i∈π qi βπ =

(15) (16)

√ Assuming SGD is run with step size ηt ∝ 1/ T , standard results [Devolder et al., 2014] establish that the expected suboptimality of the averaged iterate satisfies: p   D σπ2 (q) ∗ √ E f (x̄T ) − f (x ) ≤ + Dβπ . (17) T Here, Dβπ represents the irreducible bias floor. The target error ϵ must strictly exceed this floor (ϵ > Dβπ ) 2 D 2 σπ (q) to be achievable. Solving for T , the number of steps required is T = (ϵ−Dβ 2 . Using the optimal importance π) P √ sampling strategy restricted to π (as in Theorem 4.1), the second moment becomes σπ2 (q∗ ) = n12 ( i∈π Gi ci )2 . Thus, when optimizing over subset π, the cost to obtain expected error ϵ is: !2 X √ D2 Gi ci . Kπ (ϵ) = 2 (18) n (ϵ − Dβπ )2 i∈π Optimal Subset Selection. We seek the subset π that minimizes this cost. We decouple this by fixing a maximum allowable bias budget Γ and finding the subset that minimizes the sampling cost subject to the constraint Dβπ ≤ Γ. Theorem 5.1 (Equivalence to Min-Knapsack Problem). Let Γ be the bias budget. Finding the optimal subset π ∗ that minimizes the per-step sampling cost subject to Dβπ ≤ Γ is equivalent to the Min-Cost Knapsack Problem (also known as the Covering Knapsack Problem). Let zi ∈ {0, 1} be the indicator variable where zi = 1 if i ∈ π. The problem is: min z

n X

z i ωi

subject to

i=1

n X

zi vi ≥ Vreq ,

(19)

i=1

√ where the is ωi = Gi ci , the “item value” is vi = Gi , and the required value coverage is Pn “item cost” Vreq = j=1 Gj − nΓ D . 9

Accuracy vs. Token Count (GRPO)

70

60

Accuracy (%)

60

Accuracy (%)

Accuracy vs. Token Count (GRPO+ZVF)

70

50 40

Method

50 40

Method

No sampling

p* * psmooth ( = 0.01)

30 20

No sampling

0.0

0.2

0.4

0.6

0.8

1.0

Cumulative Tokens

1.2

1.4

p* * psmooth ( = 0.01)

30 20

1.6 1e9

0.0

0.2

0.4

0.6

Cumulative Tokens

0.8

1.0

1e9

Figure 4: Cumulative tokens compared with AIME pass@1/mean@32 accuracy throughout training for both GRPO and GRPO+ZVF settings for the 1.5B model. We plot the accuracy on the y-axis and the cumulative number of tokens used in policy optimization on the x-axis, to compare the number of tokens used for a fixed accuracy. We evaluate every 100 steps. We defer the proof to Appendix A.3. A similar result for strongly convex objectives is provided in Appendix A.4. Greedy Approximation. The Min-Cost Knapsack problem is NP-hard [Garey and Johnson, 1979]. However, efficient approximations exist. A standard greedy strategy sorts items by their value-to-cost density ρi = vi /ωi . Substituting our parameters: ρi =

Gi √

Gi

1 =√ . ci ci

(20)

Remarkably, the optimal greedy strategy is simply to select samples with the lowest cost ci , regardless of their Lipschitz constant Gi . The greedy algorithm selects the cheapest samples until the coverage constraint Vreq is met. This provides a 2-approximation to the optimal integer solution [Csirik et al., 1991].

5.1

Synthetic Validation

Under a similar setting to Section 4.4, we implement the greedy approximation. Within a smaller bias budget, subset selection is able to attain similar error at a significantly reduced cost. As the selected subset becomes much smaller, the error increases due to bias.

6

Cost-Aware GRPO

We apply our theoretical analysis to the empirical setting of GRPO [Shao et al., 2024]. In particular, we focus on reducing the number of FLOPs used by policy optimization.

6.1

Datasets and Models

We perform experiments with the Qwen2.5-Math-1.5B-Instruct [Yang et al., 2024] and Qwen3-8B [Yang et al., 2025] models. We use the DAPO dataset [Yu et al., 2025] for training, and evaluate on the following benchmarks: MATH500 [Hendrycks et al., 2021], AMC [Art of Problem Solving, n.d.], and GSM8K [Cobbe et al., 2021], and AIME1983-2024 [Veeraboina, 2023]. All experiments are performed within the Verl framework 10

0.08

0.98

0.06

0.96

0.04

0.94

0.02 0.00

Correlation, Divergence vs. Training Step (1.5B, GRPO)

0.10

1.00

0.08

0.98

0.06

0.96

0.04

0.94

0.92

0.02

0.92

0

100

200

300

Training Step

400

500

0.90

0.00

2 Divergence

Pearson Correlation

1.00

Pearson Correlation

2 Divergence

Correlation, Divergence vs. Training Step (8B, GRPO)

0.10

200

400

600

800

1000

Training Step

1200

1400

0.90

Figure 5: Sub-optimality metrics for using |Ai | in place of Gi for GRPO experiments across two model sizes. The Pearson correlations are near 1. The cost-biased χ2 -divergence between the true p∗ and the distribution defined by this proxy is near 0. Results are obtained by running one step of training at each checkpoint and computing the sequence-level gradient contribution. [Sheng et al., 2025], and we make only slight modifications for the purpose of our algorithm. We report hyperparameters for all experiments in Appendix D.

6.2

Cost-Aware Learning with GRPO

Background on GRPO. In GRPO, training alternates between two stages. The first stage generates many responses for each prompt in order to build a dataset for the second stage. The second stage is where policy gradient updates are made. When training in an on-policy fashion, the gradient updates are made on the current set of prompts and responses, which can be considered as the current dataset. Typically in GRPO, the training dataset has “verifiable rewards,” such as a math dataset where the answers are known. The reward is 1 if the response is correct and 0 otherwise. Given M responses for a prompt, advantages are normalized: PM 1 ri − M k=1 rk Ai = . (21) std({rk }M k=1 ) Let P (Q) be the distribution over prompts and πθold (q) be the distribution over responses from policy πθold given prompt q. The objective contribution of a prompt q and response o is as follows:     |o| πθ (ot | q, o<t ) (t) πθ (ot | q, o<t ) 1 X min A , clip , 1 − ϵ, 1 + ϵ A(t) − βDKL (πθ ∥πref ). ℓ(o; θ, q) = |o| t=1 πθold (ot | q, o<t ) πθold (ot | q, o<t ) The GRPO objective for parameters θ is: " JGRPO (θ) = E

q∼P (Q) {oi }M i=1 ∼πθold (q)

# M 1 X ℓ(oi ; θ) . M g=1

(22)

Given n prompts and M rollouts per prompt, and letting oi,g be the g’th rollout from the i’th prompt, the empirical objective is: n M 1 XX JˆGRPO (θ) = ℓ(oi,g ; θ, qi ) . (23) nM i=1 g=1 In effect, the empirical GRPO objective is a finite sum objective over prompts and rollouts. If we sample according to another distribution p ∈ ∆nM , where we sample rollout pair (qi , oi,g ) with probability p(i,g) , the importance-weighted objective is: n

M

1 XX 1 (p) JˆGRPO (θ) = ℓ(oi,g ; θ, qi ) . nM i=1 g=1 p(i,g) 11

(24)

Algorithm 2 Cost-Aware GRPO Input: Initial policy πθ , Dataset D, Prompt batch size n, Rollouts per prompt M , Mini-batch size B, Iterations K for iteration k = 1 to K do // Phase 1: Data Collection Sample prompt batch X = {x1 , . . . , xn } ∼ D Generate M outputs per prompt: yi,j ∼ πθold (·|xi ) Form rollout dataset Dstep of size N = n × M Compute advantages for all samples in Dstep // Phase 2: Compute Sampling Distribution √ Estimate weights su ← Gu / cu for P all u ∈ Dstep Compute probabilities p∗u ← su / v∈Dstep sv // Phase 3: Importance Weighted Updates Set number of updates T ← ⌈N/B⌉ for update step t = 1 to T do // Sample mini-batch and optimize policy Sample mini-batch ∼ Multinomial(support = Dstep , probs = p∗ , size = B) (p∗ ) Update πθ by maximizing JˆGRPO (θ) over mini-batch end for end for Defining ci . The notion of cost fits seamlessly in this setting of GRPO because different prompt and response pairs have different lengths. The amount of FLOPs required in the policy gradient step depends linearly on the total number of tokens in the prompt and response. Hence, for each prompt and response pair, we define ci to be the sum of the length of the prompt and the length of the response. Proxy for Gi . It remains to find a strong proxy for Gi ; recall that Gi represents the norm of the gradient of a sample. In language model training, it would not be logical to compute the norm of the gradient of each sample in order to obtain an optimal sampling distribution, as this has a very high computational burden and defeats the purpose. Crucially, we estimate Gi with |Ai |, the magnitude of the advantage of sample i. We justify this choice by observing low cost-biased χ-squared divergence between the sampling proxy obtained this way and the true p∗ distribution (invoking Theorem 4.3), as well as high Pearson correlation between Gi and |Ai | (invoking Theorem 4.4). In Figure 5, we report the Pearson correlation between |Ai | and Gi across training steps for GRPO with the 8B and 1.5B models. In order to obtain Gi , we compute the sequence-level gradients associated with each sequence. Upon computing the true gradient norms, we compute the cost-biased χ2 divergence between the distributions obtained using |Ai | and Gi respectively. The complete results for all training settings can be found in Appendix D.7. Algorithm. Equipped with notions of cost and a gradient proxy, we combine cost-aware learning with GRPO in the policy-gradient stage of training. The algorithm is detailed in Algorithm 2. When we perform the policy gradient steps, we fill mini-batches according to the p∗ distribution rather than simply using all √ samples once. The p∗ distribution is obtained by letting p∗i ∝ Gi / ci as described in Section 4. Once the samples are chosen for each batch, we use importance sampling to mitigate bias. Aside from this, we make no changes to the GRPO pipeline.

6.3

Settings and Methods

We evaluate combinations of two training settings and three sampling methods. Training settings. (1) GRPO: Standard GRPO training as described in Section 6.2. (2) GRPO + ZVF: GRPO with zero-variance filtering (ZVF) [Yu et al., 2025], which removes zero-advantage prompts. To maintain a fixed batch size, rollouts are resampled until sufficient non-zero-advantage samples are obtained.

12

Setting

Method

AIME

AMC

MATH

GSM8K

Avg. Accuracy

GRPO GRPO GRPO

No sampling p∗ p∗smooth (α = 0.01)

61.3 65.6 65.3

64.1 71.2 68.1

73.2 73.2 72.3

86.2 86.0 86.0

71.2 74.0 72.9

ZVF ZVF ZVF

No sampling p∗ p∗smooth (α = 0.01)

58.1 57.6 62.9

64.7 63.0 68.0

73.2 72.9 73.3

86.1 85.9 85.8

70.6 69.8 72.5

Table 1: Comparative best performance (Math-1.5B-Instruct, pass@1/mean@32). For each setting and method, we report the best averaged checkpoint accuracy across four benchmarks. On easier benchmarks, all methods perform similarly. On harder benchmarks, Cost-Aware GRPO variants outperform in particular in the GRPO setting and the smoothed setting with ZVF. Complete results are in Table 6. Sampling methods. (1) No sampling: Standard training on all generated rollouts without resampling or importance weighting. (2) p∗ sampling: We sample rollouts according to p∗ (Section 6.2, Algorithm 2) and apply importance weighting to correct for bias. (3) Smoothed p∗ sampling: We also consider p∗smooth (α) = (1 − α)p∗ + αU, where U is the uniform distribution. For the 1.5B model, we evaluate all combinations of settings and sampling methods. For the 8B model, we restrict our attention to the GRPO setting, and evaluate both p∗ and p∗smooth sampling. Additional implementation and hyperparameter details are provided in Appendix D.

6.4

Results

We evaluate all methods across two primary axes: benchmark accuracy and the cumulative tokens used in policy optimization. Because token count in this step is directly proportional to computational FLOPs for policy optimization, this evaluation demonstrates how cost-aware learning minimizes training costs without sacrificing performance. All reported accuracies denote pass@1/mean@32. Qwen2.5-Math-1.5B-Instruct Results. Figure 4 plots the cumulative policy optimization tokens against AIME accuracy for the Qwen2.5-Math-1.5B-Instruct model. Under the standard GRPO setting, both p∗ and p∗smooth (α = 0.01) demonstrate significant efficiency gains, matching the baseline’s peak accuracy using 28% and 32% fewer tokens respectively. Furthermore, both strategies ultimately surpass the baseline’s maximum performance by 5%. We also plot the steps immediately following the peak to illustrate that the baseline suffers from performance degradation, a known instability in GRPO [Yu et al., 2025, Khatri et al., 2025, Gao et al., 2025, Qi et al., 2025], whereas our method does not. Under GRPO+ZVF, p∗smooth (α = 0.01) maintains an advantage, achieving the baseline’s peak accuracy using 13% fewer tokens before ultimately outperforming it by 5%. While p∗ sampling also demonstrates initial token efficiency in this setting, its performance eventually degrades. We report corresponding metrics under the AMC benchmark in Figure 11: (1) p∗ sampling with GPRO requires 34% fewer than the baseline to achieve the same accuracy and improves peak accuracy by 7%, and (2) p∗smooth (0.01) sampling with GRPO+ZVF requires 5% fewer tokens than the baseline to achieve the same accuracy and improves peak accuracy by 3%. Finally, Table 1 details the peak checkpoint accuracy across all four benchmarks. Cost-aware sampling proves highly consistent across the evaluation sets, yielding a robust improvement over the baseline in overall average accuracy. Qwen3-8B Results. Figure 1 illustrates the accuracy of Qwen3-8B on the AIME benchmark up to peak accuracy. We observe that p∗ sampling achieves the baseline’s exact peak accuracy using 51% fewer tokens. Because the baseline’s late-stage improvement is marginal, we also evaluate an asymmetric comparison:

13

Setting

Method

AIME

AMC

MATH

GSM8K

Avg. Accuracy

GRPO GRPO GRPO

No sampling p∗ p∗smooth (α = 0.01)

58.3 58.5 58.5

72.8 74.7 74.0

87.0 86.7 86.7

96.0 95.2 94.9

78.5 78.8 78.5

Table 2: Comparative best performance (8B-Base, pass@1/mean@32). For each setting and method, we report the best averaged checkpoint accuracy across four benchmarks. We find that all methods have very similar peak performance. Complete results are in Table 7. matching p∗ when it reaches 99% of its peak accuracy. Even under this relaxed threshold, p∗ requires 17% fewer tokens than the baseline. This efficiency trend extends to the AMC benchmark (Figure 13), where p∗ sampling reaches the baseline’s maximum accuracy using 47% fewer tokens. We also implement p∗smooth (α = 0.01) and p∗smooth (α = 0.05) sampling for the 8B model, as reported in Figure 12. While the smoothed variants improve upon the baseline in terms of token efficiency as well, neither performs as well as p∗ sampling. We believe this may be due to the high fidelity of |Ai | as a proxy for Gi in this setting, as shown in Figure 5. Finally, Table 2 reports the peak average checkpoint accuracy across benchmarks remains highly consistent, with p∗ yielding a slight overall performance increase alongside efficiency gains. We observe the smoothed variants attain similar accuracy and report additional metrics in Appendix D.9. Summary. Overall, these results empirically demonstrate the ability of Cost-Aware GRPO to preserve accuracy while significantly reducing the FLOPs required during policy optimization. As language model responses become longer, training on their responses becomes costlier. Cost-Aware GRPO effectively mitigates this cost, reducing the training steps required for convergence and potentially raising the ceiling for final downstream performance.

6.5

Ablations

We consider various ablations to Cost-Aware GRPO in order to assess its robustness and generality.

Accuracy vs. Token Count (CISPO)

Accuracy vs. Token Count ( 2 noise for p * -NOISE- 2 ) 60 55

50

Accuracy (%)

Accuracy (%)

60

Method

p*

p*-NOISE-0.01 p*-NOISE-0.1 p*-NOISE-1 p*-NOISE-5 No sampling

40

30

0.0

0.2

0.4

0.6

0.8

1.0

Cumulative Tokens

1.2

1.4

45 40

Method

35

No sampling

p* * psmooth ( = 0.01)

30 25

1.6 1e9

(a) Results on Qwen2.5-Math-1.5B-Instruct when adding zero-mean, σ 2 -variance noise to the proxy |Ai | prior to computing the sampling distribution under GRPO.

50

0.0

0.5

1.0

1.5

Cumulative Tokens

2.0

2.5

1e9

(b) AIME accuracy (pass@1/mean@32) for training Qwen2.5-Math-1.5B-Instruct with CISPO. Cost-aware learning is robust to training objectives.

Training objective. We consider a variant of the GRPO objective, CISPO, initially proposed by Chen et al. [2025]. This variant applies a stop-gradient to the importance weight of the objective. For this setting, we base our hyperparameters on findings from Khatri et al. [2025]. Results, reported in Figure 6b, show 14

that our method is not specific to the GRPO objective; rather, it generalizes across training objectives. p∗ sampling and p∗smooth (0.01) sampling require 31% and 34% fewer tokens respectively to attain the baseline’s peak accuracy. This finding is further supported by strong sub-optimality metrics, performance on additional benchmarks, and smoothed variants in Appendix D.6. Proxy sub-optimality. We ablate the noise level associated with our proxy |Ai |. In particular, we add zero-mean noise with variance σ 2 to |Ai | prior to computing the sampling distribution. Results are shown in Figure 6a. We observe that Cost-Aware GRPO is robust to noise up until σ 2 = 0.1, but as noise approaches σ 2 = 1.0, it deteriorates with respect to the baseline. Choice of α and proxy. In Appendix D.5, we ablate our choice of α and report results for p∗smooth with α = 0.05 and α = 0.1 as well. We also ablate our choice of approximating Gi with |Ai | and consider using √ p∗i = 1/ ci , which corresponds to substituting Gi with 1 for all samples. We observe that it does not perform as well, which implies a need for correctly approximating the gradient norm of a sample. In Appendix D.9, we report smoothed variants for the 8B model as well.

7

Conclusion

In this paper, we formalize the problem of cost-aware learning and propose Cost-Aware SGD (Algorithm 1). Our analysis derives the optimal reweighting distribution for convex and strongly convex objectives, providing theoretical guarantees on the expected cost to attain an error of ϵ. Furthermore, we establish a lower bound for this setting and introduce a subset selection algorithm that further reduces training costs at the expense of introducing controllable bias. We then apply these theoretical insights to the policy optimization step of GRPO, where the FLOPs required to train on a sequence scale directly with its length. To mitigate this computational cost, we introduce Cost-Aware GRPO. Our experiments with Qwen2.5-Math-1.5B-Instruct and Qwen3-8B demonstrate the broad effectiveness of this approach across standard GRPO, CISPO, and GRPO+Zero-Variance Filtering (ZVF) settings. Ultimately, Cost-Aware GRPO significantly improves token efficiency, reducing the tokens used in policy optimization by up to 30%, while maintaining or exceeding the final reasoning performance of standard baselines. Future Directions. In Cost-Aware GRPO, we primarily rely on the magnitude of the advantage as a proxy for the gradient norm; identifying and evaluating alternative proxies remains an interesting direction for future research. Furthermore, having demonstrated that cost-aware sampling generalizes successfully from standard GRPO to the CISPO objective, a natural next step is to explore its integration into a broader class of RLVR and RLHF algorithms. Finally, on the theoretical front, extending our finite-sum SGD analysis to a more general stochastic optimization regime would be meaningful.

15

References Guillaume Alain, Alex Lamb, Chinnadhurai Sankar, Aaron Courville, and Yoshua Bengio. Variance reduction in sgd by distributed importance sampling. arXiv preprint arXiv:1511.06481, 2015. Art of Problem Solving. Amc problems and solutions. https://artofproblemsolving.com/wiki/index. php/AMC_Problems_and_Solutions, n.d. Accessed: 2026-01-28. Alina Beygelzimer, Sanjoy Dasgupta, and John Langford. Importance weighted active learning. In Proceedings of the 26th annual international conference on machine learning, pages 49–56, 2009. Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning, 8(3-4):231–357, 2015. Aili Chen, Aonian Li, Bangwei Gong, Binyang Jiang, Bo Fei, Bo Yang, Boji Shan, Changqing Yu, Chao Wang, Cheng Zhu, et al. Minimax-m1: Scaling test-time compute efficiently with lightning attention. arXiv preprint arXiv:2506.13585, 2025. Lichang Chen, Shiyang Li, Jun Yan, Hai Wang, Kalpa Gunaratna, Vikas Yadav, Zheng Tang, Vijay Srinivasan, Tianyi Zhou, Heng Huang, et al. Alpagasus: Training a better alpaca with fewer data. arXiv preprint arXiv:2307.08701, 2023a. Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941, 2018. Ziang Chen, Jianfeng Lu, Huajie Qian, Xinshang Wang, and Wotao Yin. Hetersgd: Tackling heterogeneous sampling costs via optimal reweighted stochastic gradient descent. In Francisco Ruiz, Jennifer Dy, and JanWillem van de Meent, editors, Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pages 10732–10781. PMLR, 25–27 Apr 2023b. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. Corinna Cortes, Giulia DeSalvo, Claudio Gentile, Mehryar Mohri, and Ningshan Zhang. Region-based active learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2801–2809. PMLR, 2019. Corinna Cortes, Giulia Desalvo, Claudio Gentile, Mehryar Mohri, and Ningshan Zhang. Adaptive region-based active learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 2144–2153. PMLR, 13–18 Jul 2020. URL https://proceedings.mlr.press/v119/cortes20a.html. János Csirik, Johannes Bartholomeus Gerardus Frenk, Martine Labbé, and Shuzhong Zhang. Heuristics for the 0-1 min-knapsack problem. Acta Cybernetica, 10(1-2):15–20, 1991. Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rlhf: Safe reinforcement learning from human feedback. arXiv preprint arXiv:2310.12773, 2023. Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. Sanjoy Dasgupta. Two faces of active learning. Theoretical computer science, 412(19):1767–1781, 2011.

16

Olivier Devolder, Francois Glineur, and Yurii Nesterov. First-order methods of smooth convex optimization with inexact oracle. Mathematical Programming, 146(1):37–75, 2014. Mehdi Fatemi, Banafsheh Rafiee, Mingjie Tang, and Kartik Talamadupula. Concise reasoning via reinforcement learning. arXiv preprint arXiv:2504.05185, 2025. Chang Gao, Chujie Zheng, Xiong-Hui Chen, Kai Dang, Shixuan Liu, Bowen Yu, An Yang, Shuai Bai, Jingren Zhou, and Junyang Lin. Soft adaptive policy optimization. arXiv preprint arXiv:2511.20347, 2025. Michael R. Garey and David S. Johnson. Computers and Intractability: A Guide to the Theory of NPCompleteness. W. H. Freeman and Company, New York, 1979. Series of Books in the Mathematical Sciences. Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4):2341–2368, 2013. Nicholas J. A. Harvey, Christopher Liaw, Yaniv Plan, and Sikander Randhawa. Tight analyses for non-smooth stochastic gradient descent. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1579–1613. PMLR, 25–28 Jun 2019. URL https://proceedings.mlr.press/v99/harvey19a.html. 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. arXiv preprint arXiv:2103.03874, 2021. Prateek Jain, Dheeraj Nagaraj, and Praneeth Netrapalli. Making the last iterate of sgd information theoretically optimal. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1752–1755. PMLR, 25–28 Jun 2019. URL https://proceedings.mlr.press/v99/jain19a.html. Tyler B Johnson and Carlos Guestrin. Training deep models faster with robust, approximate importance sampling. Advances in Neural Information Processing Systems, 31, 2018. Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001. Angelos Katharopoulos and Francois Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pages 2525–2534. PMLR, 2018. Devvrit Khatri, Lovish Madaan, Rishabh Tiwari, Rachit Bansal, Sai Surya Duvvuri, Manzil Zaheer, Inderjit S Dhillon, David Brandfonbrener, and Rishabh Agarwal. The art of scaling reinforcement learning compute for llms. arXiv preprint arXiv:2510.13786, 2025. Teun Kloek and Herman K Van Dijk. Bayesian estimates of equation system parameters: an application of integration by monte carlo. Econometrica: Journal of the Econometric Society, pages 1–19, 1978. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023. Xuefeng Li, Haoyang Zou, and Pengfei Liu. Limr: Less is more for rl scaling. arXiv preprint arXiv:2502.11886, 2025. Deanna Needell, Nathan Srebro, and Rachel Ward. Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm. Advances in neural information processing systems, 27, 2014. AS Nemirovsky, DB Yudin, and E Dawson. Wiley-interscience series in discrete mathematics, 1983. 17

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022. Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Finding important examples early in training. Advances in neural information processing systems, 34:20596–20607, 2021. Penghui Qi, Zichen Liu, Xiangxin Zhou, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. Defeating the training-inference mismatch via fp16. arXiv preprint arXiv:2510.26788, 2025. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728–53741, 2023. Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. arXiv preprint arXiv:1109.5647, 2011. John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889–1897. PMLR, 2015. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning, pages 71–79. PMLR, 2013. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. 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. In Proceedings of the Twentieth European Conference on Computer Systems, pages 1279–1297, 2025. Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35: 19523–19536, 2022. Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In S. Solla, T. Leen, and K. Müller, editors, Advances in Neural Information Processing Systems, volume 12. MIT Press, 1999. URL https://proceedings. neurips.cc/paper_files/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf. Hemish Veeraboina. Aime problem set 1983-2024, 2023. hemishveeraboina/aime-problem-set-1983-2024.

URL https://www.kaggle.com/datasets/

Yiping Wang, Qing Yang, Zhiyuan Zeng, Liliang Ren, Liyuan Liu, Baolin Peng, Hao Cheng, Xuehai He, Kuan Wang, Jianfeng Gao, et al. Reinforcement learning for reasoning in large language models with one training example. arXiv preprint arXiv:2504.20571, 2025. Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy S Liang. Data selection for language models via importance resampling. Advances in Neural Information Processing Systems, 36:34201–34227, 2023. Yixuan Even Xu, Yash Savani, Fei Fang, and Zico Kolter. Not all rollouts are useful: Down-sampling rollouts in llm reinforcement learning. arXiv preprint arXiv:2504.13818, 2025. 18

An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409.12122, 2024. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476, 2025. Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity, 2020. URL https://arxiv.org/abs/1905.11881. Peilin Zhao and Tong Zhang. Stochastic optimization with importance sampling for regularized loss minimization. In international conference on machine learning, pages 1–9. PMLR, 2015. Haizhong Zheng, Yang Zhou, Brian R Bartoldson, Bhavya Kailkhura, Fan Lai, Jiawei Zhao, and Beidi Chen. Act only when it pays: Efficient reinforcement learning for llm reasoning via selective rollouts. arXiv preprint arXiv:2506.02177, 2025. Rui Zheng, Shihan Dou, Songyang Gao, Yuan Hua, Wei Shen, Binghai Wang, Yan Liu, Senjie Jin, Qin Liu, Yuhao Zhou, et al. Secrets of rlhf in large language models part i: Ppo. arXiv preprint arXiv:2307.04964, 2023.

19

Contents 1 Introduction

1

2 Related Work

3

3 Problem Formulation

3

4 Cost-Aware SGD 4.1 Cost Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.2 Derivation of the Optimal Distribution . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.3 Cost-Improvement over Baselines . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.4 Synthetic Validation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.5 Proxy Sub-Optimality . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.6 Lower Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

4 4 5 5 6 7 8

5 Subset Selection 5.1 Synthetic Validation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

8 10

6 Cost-Aware GRPO 10 6.1 Datasets and Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10 6.2 Cost-Aware Learning with GRPO . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11 6.3 Settings and Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 6.4 Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 6.5 Ablations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14 7 Conclusion

15

A Proofs 21 A.1 Variance reduction strategy derivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 A.2 Cost-Improvement Theorem & Comparison Analysis . . . . . . . . . . . . . . . . . . . . . . . 21 A.3 Min Cost Knapsack . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 A.4 Strongly Convex Subset Selection . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 A.5 Sub-Optimality Proofs . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 B Lower bound

27

C Extensions and Generality of the Theoretical Analysis 31 C.1 Average Iterate (x̄T ) vs. Last Iterate (xT ) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 C.2 Generality of the Cost Objective in Non-Convex Settings . . . . . . . . . . . . . . . . . . . . . 31 C.3 Differences between Theory and Practice . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 D Experimental Details and Ablations 32 D.1 Experimental Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 D.2 Implementation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 D.3 Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 D.4 Computational Requirements . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 D.5 Ablation on variants of p∗ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 D.6 Full CISPO results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36 D.7 Full sub-optimality results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 37 D.8 Additional results for Qwen2.5-Math-1.5B-Instruct . . . . . . . . . . . . . . . . . . . . . . . . 37 D.9 Additional results for Qwen3-8B . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 38

20

A

Proofs

In this section we provide the proofs and derivations for statements made in the paper.

A.1

Variance reduction strategy derivation

In this case, we have:



 ! n n X X G2j 1   J(p) = S(p)C(p) = pi , n2 j=1 pj i=1

where p ∈ ∆n . From Cauchy-Schwarz, we can lower bound J(p): !2 n 1 X Gi . J(p) ≥ 2 n i=1 √ √ This lower bound is attained when Gi / pi ∝ pi , or equivalently, pi ∝ Gi . Hence, when costs are uniform, we focus on minimizing the second moment. This gives: Gi (pvar )i = P . j Gj

(25)

When we apply pvar to the non-uniform cost setting, the total expected cost to attain error ϵ becomes: !2 n ! n X D2 X Gi Kvar (ϵ) = 2 2 (pvar )i ci ϵ n i=1 i=1 !2 n ! n X Gi D2 X P Gi ci = 2 2 ϵ n i=1 j Gj i=1 ! n ! n X D2 X Gi Gi ci . = 2 2 ϵ n i=1 i=1

A.2

Cost-Improvement Theorem & Comparison Analysis

Theorem A.1 (Cost Improvement). Let Kunif (ϵ), Kvar (ϵ), Kopt (ϵ) be the total cost to reach ϵ-accuracy using uniform, adaptive, and optimal sampling, respectively. Then Kopt ≤ Kunif , with the gap determined by the correlation between gradients and costs: P √ ( Gi ci )2 Kopt (ϵ) = P 2 P ≤ 1. Kunif (ϵ) ( Gi )( ci ) In addition, Kopt ≤ Kvar , with the gap determined by the correlation between gradients and costs: P √ ( Gi ci )2 Kopt (ϵ) P = P ≤ 1. Kvar (ϵ) ( Gi )( i Gi ci ) Proof. Equation (12) follows simply from Cauchy-Schwarz and by definition of uniform sampling. We √ √ also apply Cauchy-Schwarz to obtain the result in Equation (13). In particular, substituting ui = Gi , vi = Gi ci P P P 2 in Cauchy-Schwarz, ( i ui vi ) ≤ ( i u2i )( i vi2 ) we have X

 √ 2 X X Gi ci ≤ Gi Gi ci .

21

(26)

Theorem A.2 (Comparison of Adaptive vs. Uniform). Unlike the optimal strategy, the adaptive strategy pvar is not guaranteed to outperform uniform sampling. The ratio of their costs is given by: Pn Pn ( i=1 Gi )( i=1 Gi ci ) Kvar (ϵ) Pn . (27) = Pn Kunif (ϵ) ( i=1 G2i )( i=1 ci ) Let E[·] denote the expectation taken over a uniform distribution of indices i ∼ Uniform(1, . . . , n). Then Kvar ≤ Kunif holds if and only if: E[G] Cov(G, c) ≤ E[c] Var(G). (28) Consequently, if gradients and costs are negatively correlated or uncorrelated, adaptive sampling is superior. However, if gradients and costs are strongly positively correlated (i.e., samples with large gradients are disproportionately expensive), uniform sampling may yield a lower total cost. P P P P Proof. Recall that Kunif ∝ ( G2i )( ci ) and Kvar ∝ ( Gi )( GP i ci ). We can rewrite these sums in terms of expectations over uniformly sampled indices. For any vector x, xi = n E[x]. Thus: (n E[G])(n E[Gc]) E[G] E[Gc] Kvar = = . Kunif (n E[G2 ])(n E[c]) E[G2 ] E[c] We seek the condition under which this ratio is ≤ 1, which implies E[G] E[Gc] ≤ E[G2 ] E[c]. We apply the definitions E[Gc] = Cov(G, c) + E[G] E[c] and E[G2 ] = Var(G) + E[G]2 . Substituting these into the inequality:
E[G](Cov(G, c) + E[G] E[c]) ≤ Var(G) + E[G]2 E[c]. Expanding both sides: E[G] Cov(G, c) + E[G]2 E[c] ≤ E[c] Var(G) + E[G]2 E[c]. Subtracting the common term E[G]2 E[c] from both sides yields the final condition: E[G] Cov(G, c) ≤ E[c] Var(G). Since E[G], E[c], and Var(G) are non-negative, this inequality always holds if Cov(G, c) ≤ 0 (negative or zero correlation). It is violated only if Cov(G, c) is sufficiently large and positive.

A.3

Min Cost Knapsack

Theorem A.3 (Equivalence to Min-Knapsack Problem). Let Γ be the maximum allowable bias budget such that convergence is possible. Finding the optimal subset π ∗ that minimizes the expected total cost subject to this budget is equivalent to solving the Min-Cost Knapsack Problem. Specifically, let zi ∈ {0, 1} be the indicator variable where zi = 1 if i ∈ π. The problem is: min z

n X

z i ωi

subject to

i=1

n X

zi vi ≥ Vreq ,

(29)

i=1

√ where the is ωi = Gi ci , the “item value” is vi = Gi , and the required value coverage is Pn “item cost” Vreq = j=1 Gj − nΓ D . Proof. We aim the cost complexity factor derived for the optimal distribution restricted to P to minimize √ π: J(π) = ( i∈π Gi ci )2 . The constraint is on the bias error contribution Dβπ ≤ Γ < ϵ. Recalling the definition of bias:   X X 1 1 βπ ≤ Gj =⇒ D Gj  ≤ Γ. n n j ∈π /

j ∈π /

22

Minimizing the sum of gradients excluded from π is equivalent to maximizing the sum of gradients included in π. We rewrite the constraint: X j ∈π /

Gj ≤

n n X X X X nΓ nΓ nΓ =⇒ =⇒ Gi ≥ . Gj − Gi ≤ Gj − D D D i∈π j=1 i∈π j=1

Pn √ Therefore, if we let Vreq = j=1 Gj − nΓ D , let ωi = Gi ci as the penalty for including item i, and let vi = Gi as the gain toward satisfying the constraint, we recover the formulation in Equation (19). Using this algorithm, the worst case E[K(ϵ)], where ϵ > Γ, for Dβπ ≤ Γ is: D2 J(π) . {π:βπ ≤Γ/D} n2 (ϵ − Γ)2 min

P The bias constraint Dβπ ≤ Γ determines the required gradient coverage Vreq = i Gi − nΓ/D, and thus defines a specific Min-Cost Knapsack instance as in Theorem 5.1. By varying Γ over its feasible range and solving the associated knapsack problem for each choice, we obtain an optimal subset π(Γ) minimizing the total cost subject to that bias budget. Since the total expected complexity E[Kπ (ϵ)] =

D2 J(π) n2 (ϵ − Dβ

π)

2

depends on both the numerator (the sampling cost) and the denominator (the bias floor), enumerating over Γ allows us to evaluate the achievable error–cost tradeoff for every subset. Consequently, we can compute the optimal target accuracy ϵ by selecting the bias level Γ that yields the smallest overall expected cost.

A.4

Strongly Convex Subset Selection

Theorem A.4 (Subset Selection for Strongly Convex Functions). Assume f is µ-strongly convex. Let zi ∈ {0, 1} P be the indicator variable for including data point i in subset π. We seek to minimize the cost √ complexity zi Gi ci subject to a budget Γ on the function value suboptimality f (x∗π ) − f (x∗ ) ≤ Γ. This problem is structurally equivalent to the Min-Cost Knapsack formulation in Theorem 5.1, with a modified requirement Vreq . The optimal solution to the LP relaxation is obtained by selecting indices i strictly based on increasing cost ci . Proof. Using the Polyak-Lojasiewicz (PL) inequality and the fact that ∇fS (x∗S ) = 0, the gradient at the proxy minimizer is determined solely by the excluded gradients: ∇f (x∗π ) =

1X ∇fj (x∗π ). n j ∈π /

We bound the norm of this gradient using the triangle inequality (∥ bound definition (∥∇fj (x)∥ ≤ Gj ): ∥∇f (x∗π )∥ =

1X ∇fj (x∗π ) n j ∈π /

1X ∥∇fj (x∗π )∥ ≤ n j ∈π /

1X ≤ Gj . n j ∈π /

23

P

v∥ ≤

P

∥v∥) and the gradient upper

1 Substituting this back into the PL inequality (f (x∗π ) − f (x∗ ) ≤ 2µ ∥∇f (x∗π )∥2 ) yields:

 2 !2 n X X 1 1  ∗ ∗ Gj  = (1 − zi )Gi . f (xπ ) − f (x ) ≤ 2µn2 2µn2 i=1 j ∈π /

We enforce the constraint f (x∗π ) − f (x∗ ) ≤ Γ. Because the gradient norms Gi are non-negative, we can take the square root of the inequality to linearize the constraint without altering the optimization region: n X

p (1 − zi )Gi ≤ n 2µΓ.

i=1

√ This is now a linear constraint on the excluded mass. Let B = n 2µΓ be the effective budget for the excluded gradients. We rewrite this in terms of the included variables zi : n X j=1

Gj −

n X

zi Gi ≤ B =⇒

i=1

n X

zi Gi ≥

i=1

n X

Gj − B.

j=1

P P √ The optimization problem is to minimize zi (Gi ci ) subject to zi Gi ≥ Vreq . This is identical to Equation 19. This problem is structurally equivalent to the Min-Cost Knapsack formulation in Theorem 5.1, with a modified requirement Vreq . Its LP relaxation is a fractional knapsack problem whose optimal solution is obtained by ordering indices by increasing evaluation cost ci (equivalently, decreasing density vi /ωi ).

A.5

Sub-Optimality Proofs

Theorem A.5 (Sub-optimality Gap via Cost-Biased Divergence). Let p∗ be the optimal cost-aware sampling ′ distribution and Pnp be any empirical sampling distribution. Let the expected cost per iteration for a distribution p be C(p) = i=1 pi ci .We define the cost-biased distribution p̃ corresponding to p as: p̃i =

pi ci . C(p)

The optimality gap is exactly determined by the Pearson χ2 -divergence between the cost-biased optimal distribution and the cost-biased empirical distribution: J(p′ ) = 1 + Dχ2 (p̃∗ ||p̃′ ). J(p∗ ) Proof. Recall that minimizing the total expected cost is equivalent to minimizing the product J(p) = S(p)C(p). √ Pn G G / c The optimal distribution is defined as p∗i = Zi den i , where Zden = j=1 √cjj . We can rearrange this to express the true gradient norms in terms of p∗i : √ Gi = Zden p∗i ci . Pn G2 Substituting this expression for Gi into the variance term S(p′ ) = n12 i=1 p′i yields the objective for the i empirical distribution: ! n X Z2 (p∗i )2 ci ′ J(p′ ) = den C(p ) . 2 n p′i i=1 Evaluating this same expression for the optimal distribution p∗ (where p′i = p∗i ) gives J(p∗ ) = Taking the ratio of the two costs isolates the gap: Pn (p∗ )2 C(p′ ) i=1 ci pi′ J(p′ ) i = J(p∗ ) C(p∗ )2 24

2 Zden ∗ 2 n2 C(p ) .

By the definition of the cost-biased distribution, we have pi = p̃i C(p) ci . Substituting this expression for both p∗ and p′ into the summation term: 2  ∗ ) n n p̃∗i C(p X X ci (p∗i )2 ci ′ = ci . ′) pi p̃′i C(p i=1 i=1 ci Simplifying the terms inside the sum: =

n X i=1

∗ 2

ci

(p̃∗i )2 C(pc2 ) p̃′i

n

=

i C(p′ )

ci

C(p∗ )2 X (p̃∗i )2 . C(p′ ) i=1 p̃′i

Substituting this simplified summation back into our original cost ratio: ! n n X J(p′ ) (p̃∗i )2 C(p′ ) C(p∗ )2 X (p̃∗i )2 = . = J(p∗ ) C(p∗ )2 C(p′ ) i=1 p̃′i p̃′i i=1 Using the definition of the Pearson χ2 -divergence, Dχ2 (P ||Q) =

P P2 Q

− 1, we arrive at the identity.

Theorem A.6. Let p∗ be the optimal cost-aware sampling distribution using the true gradient norms Gi , and √ let p′ be the empirical sampling distribution using estimated gradient norms G′i , such that p′i ∝ G′i / ci .Assume the estimates follow an additive noise model G′i = Gi + ϵi , where ϵi are independent, zero-mean noise terms with variance σϵ2 , independent of ci . Let ρ be the Pearson correlation coefficient between the true norms Gi and the estimates G′i . Assuming the noise ϵi is relatively small compared to Gi , the expected optimality gap is approximately bounded by: Pn √ci   E[J(p′ )] 1 − ρ2 i=1 Gi 2 ≈1+ σG Pn √ J(p∗ ) ρ2 i=1 Gi ci 2 where σG is the variance of the true gradient norms.

Proof. The √cost objective for a distribution p is J(p) = S(p)C(p). For the estimated distribution p′ , we have G′ / c p′i = Pn i ′ i√ . Substituting this into the cost objective yields: j=1

Gj / cj

1 J(p ) = 2 n ′

√ n X G2 ci i=1

!

i

n X

G′i

i=1

√ G′i ci

! .

Under the model G′i = Gi + ϵi , the covariance between G and G′ is strictly the variance of G, namely 2 2 Cov(G, G′ ) = σG . The variance of G′ is V ar(G′ ) = σG + σϵ2 . The Pearson correlation coefficient ρ is defined as: Cov(G, G′ ) σ2 σG p G ρ= = =p 2 2 + σ2 2 σG σG′ σG σG σ ϵ G + σϵ Solving for the noise variance σϵ2 in terms of ρ: 2 σϵ2 = σG



1 − ρ2 ρ2

 .

Because G′i appears in the denominator of the second moment, we apply a second-order Taylor expansion of 1/G′i centered around Gi : 1 1 ϵi ϵ2 ≈ − 2 + i3 Gi + ϵi Gi Gi Gi Substituting this into the first term of J(p′ ) and taking the expectation with respect to the noise ϵ: 25

 E

  √  √ G2i ci √ 1 E[ϵi ] E[ϵ2i ] 2√ 2 ci ≈ G + = G c − c + σ . i i i i ϵ G′i Gi G2i G3i Gi

Taking the expectation of the full cost J(p′ ), noting that the two sums are weakly dependent for large n and E[G′i ] = Gi : ! √ ! X n  n √ √ 1 X ′ 2 ci Gi ci + σϵ Gi ci E[J(p )] ≈ 2 n Gi i=1 i=1 n X

1 E[J(p )] ≈ 2 n ′

Recognizing that J(p∗ ) = n12

Pn

i=1

i=1 Gi

√

Gi ci

!2

σ2 + ϵ2 n

n √ X ci i=1

!

Gi

√
2 ci , we divide by J(p∗ ):

Pn √ci E[J(p′ )] i=1 Gi ≈ 1 + σϵ2 Pn √ . J(p∗ ) i=1 Gi ci 2 Finally, substituting σϵ2 = σG



1−ρ2 ρ2



completes the proof.

26

n X i=1

√

Gi ci

!

B

Lower bound

We consider the minimization of a finite-sum objective: n

min F (x) = x∈X

1X fi (x). n i=1

where each fi : X → R is convex and G-Lipschitz continuous. We assume access to an exact first-order oracle: when queried at index i and point x, the oracle returns the exact gradient ∇fi (x). Crucially, we assume there is a variable cost ci > 0 to querying the i-th component (at any point x). The goal is to minimize the total expected accumulated cost required to output a point x̂ such that E[F (x̂) − F (x∗ )] ≤ ϵ, where the expectation is taken over any randomization in the algorithm producing x̂. Theorem B.1. Fix ϵ, G > 0. For any number of components n ≥ G2 /ϵ2 , any set of nonnegative query costs {ci }ni=1 , and any (possibly randomized) algorithm with an expected error guarantee E[F (x̂) − min F (x)] ≤ ϵ, there exists a convex, G-Lipschitz finite-sum problem instance over X = [−1, 1] such that the algorithm’s expected total query cost is at least:  !2  2 X √ 1 G Ω 2 ci  , ϵ n ∗ i∈S

where S ∗ ⊆ [n] is any set of components such that following cost-uniformity condition holds: √ maxi∈S ∗ ci G P . √ ≤ 1 40ϵ c ∗ j j∈S n

(30)

Remark. If we let the costs be sorted c1 ≤ c2 ≤ . . . ≤ cn , then the set S ∗ that maximizes the lower bound (if it is nonempty) must be of the form S = [k] for some 1 ≤ k ≤ n, namely, the set of k cheapest costs. Indeed, if j ∈ S ∗ then also i ∈ S ∗ for any i < j, since otherwise including i into S ∗ increases the lower bound, while still satisfying Eq. (14): it can only increase the denominator on the LHS, while not affecting the numerator. For reasonably uniform costs, say perfectly uniform up to a small (constant) fraction of arbitrarily large outliers, the condition in Eq. (14) is satisfied by fixing S ∗ to contain all components except for the outliers, as long as G/ϵ is larger than a constant. It could be, however, that the maximal S ∗ is empty (and the lower bound is zero) if the costs are highly non-uniform. For example, if the costs are exponentially increasing, say √ ci = 4i for all i, then the ratio between the maximum and the sum of ci over any set of the form S ∗ = [k] is G between 12 and 1, and so the cost-uniformity condition implies that ϵ ≤ 20n . Thus, for any larger value of ϵ there is no set S ∗ for which the condition is satisfied. Proof of Theorem 4.5. The proof proceeds by constructing a hard distribution of instances based on linear functions and applying information-theoretic lower bounds. We break the argument into four steps. Step 1: Definition of hard instance. Consider the 1-dimensional domain X = [−1, 1]. We define the component functions as linear functions: ∀ i ∈ [n] :

fi (x) = −bi x,

Pn where bi ∈ R is a slope parameter. The global objective is F (x) = −x · b̄, where b̄ = n1 i=1 bi . The minimizer x∗ depends entirely on the sign of the average slope b̄; namely, x∗ = sign(b̄). The value of the minimum is F (x∗ ) = −|b̄|, and for any candidate x̂ ∈ [−1, 1], the optimality gap is F (x̂) − F (x∗ ) = b̄(x∗ − x̂). In particular, for any algorithm producing x̂, E[F (x̂) − F (x∗ )] ≥ b̄ · P[sign(x̂) ̸= sign(b̄)]. 27

(31)

We define the hard distribution D as the uniform mixture of two priors, D0 and D1 , parameterized by ∆1 , . . . , ∆n ≥ 0: • Under D0 : bi = clip(zi ), where zi ∼ N (+∆i , σ 2 ) are i.i.d., σ = G/2; • Under D1 : bi = clip(zi ), where zi ∼ N (−∆i , σ 2 ) are i.i.d., σ = G/2. Here, clip(zi ) = max(−G, min(G, z)) is a truncation to [−G, G] that strictly enforces the Lipschitz condition (deterministically). The bi are then distributed according to a truncated Gaussian distribution with variance σ 2 strictly smaller than G2 . For reasons that will become apparent in later steps, we will set the parameters ∆i so that the following conditions are met: n

1X ∆i = 20ϵ, n i=1 Step 2: Reduction to hypothesis testing.

and

∆i ≤

G 2

for all i ∈ [n].

(32)

Define the hypothesis testing error probability as

perr = 12 p0err + 12 p1err , where p0err = PD0 (x̂ < 0) and p1err = PD1 (x̂ > 0). We will next show that the expected error (Eq. (31)) is lower bounded in terms of perr . Conditioning on D0 , a sufficient condition for the algorithm to fail is if x̂ ≤ 0 and b̄ > 0, in which case the optimization error is at least |b̄|; thus ED0 [F (x̂) − F (x∗ )] ≥ ED0 [|b̄| · I(x̂ ≤ 0) · I(b̄ > 0)].

(33)

We next bound the expected value (under D0 ) of the mean of truncate variables. Let g(δ) = E[clip(Z + δ, −G, G)] where Z ∼ N (0, σ 2 ). The mean of the clipped variable is E[bi ] = g(∆i ). By the mean value theorem, g(∆i ) = g(0) + g ′ (ξ)∆i = g ′ (ξ)∆i for some ξ ∈ (0, ∆i ). The derivative is g ′ (t) = P(−G ≤ Z + t ≤ G). With σ = G/2 and ∆i ≤ G/2, the interval [−G − ∆i , G − ∆i ] always contains [−G/2, G/2] = [−σ, σ]. The Gaussian Pnmass in this interval is > 0.5, thus E[bi ] ≥ 0.5∆i . The aggregate mean then satisfies E[b̄] ≥ 0.5( n1 i=1 ∆i ) = 10ϵ. The truncated variables bi are independent and bounded in [−G, G]. By Hoeffding’s inequality, b̄ concentrates around its expectation: defining the event A = {b̄ ≥ 5ϵ}, the probability of the complement δ ′ = PD0 (Ac ) is bounded by    
3nϵ2 2n2 (5ϵ)2 δ ′ ≤ PD0 b̄ ≤ ED0 [b̄] − 5ϵ ≤ exp − ≤ exp − ≤ 0.05, n(2L)2 G2 where the final inequality is true since we assume n ≥ G2 /ϵ2 . Conditioning on the event A (where b̄ ≥ 5ϵ > 0) and recalling Eq. (33), we have ED0 [F (x̂) − F (x∗ )] ≥ ED0 [|b̄| · I(x̂ ≤ 0) · I(A)] ≥ 5ϵ PD0 ({x̂ ≤ 0} ∩ A) ≥ 5ϵ(PD0 (x̂ ≤ 0) − PD0 (Ac )) = 5ϵ(p0err − δ ′ ). Similarly, ED1 [F (x̂) − F (x∗ )] ≥ 5ϵ(p1err − δ ′ ), thus ED [F (x̂) − F (x∗ )] ≥ 5ϵ(perr − δ ′ ). The guarantee E[F (x̂) − F (x∗ )] ≤ ϵ implies 5ϵ(perr − 0.05) ≤ ϵ, which leads to perr ≤ 0.2 + 0.05 = 0.25. Thus, the algorithm must identify the correct hypothesis with probability at least 0.75.

28

Step 3: Information-theoretic lower bound. Let P and Q denote the probability distributions of the entire sequence of oracle responses observed by the algorithm when the instance is drawn from D0 and D1 , respectively. Let P ′ and Q′ denote the distributions of the underlying unclipped Gaussian vectors z = (z1 , . . . , zn ) under D0 and D1 . Since the observed gradients −bi are deterministic functions of zi (specifically, bi = clip(zi )), the data processing inequality for KL divergences implies: DKL (P ∥ Q) ≤ DKL (P ′ ∥ Q′ ). The variables zi are independent, and zi ∼ N (∆i , σ 2 ) under D0 , while zi ∼ N (−∆i , σ 2 ) under D1 . The KL divergence between two univariate Gaussians with means µ1 , µ2 and common variance σ 2 is 12 (µ1 − µ2 )2 /σ 2 . Here, the separation is 2∆i and σ 2 = G2 /4, so the KL divergence for zi is upper bounded by 8∆2i /G2 . Since the gradient oracles are exact and the objectives are linear, repeated queries to the same component yield identical results. Thus, information is gained only from the set of distinct components queried. Let qi be the probability (under D0 ) that component i is queried at least once by the algorithm. By the chain rule for KL divergence (and linearity of expectation), the total divergence is bounded by: DKL (P ∥ Q) ≤ DKL (P ′ ∥ Q′ ) ≤

n X i=1

qi

8∆2i . G2

(34)

To relate the divergence to the error probability perr , we use Le Cam’s argument, which states that for any hypothesis testing problem with uniform prior, perr ≥ 12 (1 − T V (P, Q)) , where T V (P, Q) is the total variation distance. Given our requirement that perr ≤ 41 , we must have: 1 1 4 ≥ 2 (1 − T V (P, Q))

By Pinsker’s inequality, T V (P, Q) ≤ divergence: q 1 2 ≤

q

=⇒ T V (P, Q) ≥ 12 .

1 2 DKL (P ∥ Q), which yields the required lower bound on the KL

1 2 DKL (P ∥ Q)

=⇒ DKL (P ∥ Q) ≥ 12 .

Substituting the upper bound from Eq. (34) yields the condition: n X

qi ∆2i ≥

i=1

G2 . 16

(35)

Step 4: Optimizing the lower Pn bound. From Eq. (35), any algorithm solving the problem must satisfy the information constraint i=1 qi ∆2i ≥ G2 /16. Since the total expected query cost C (under mixture Pthe n distribution D) is lower bounded by half times the expected cost under D0 , we also have C ≥ 12 i=1 qi ci . For a fixed set of parameters ∆i , we can lower bound the algorithm’s expected cost in terms of the “information density” ρi = ∆2i /ci of the components; indeed, ! n n n X X X G2 G2 2 ≤ qi ∆i = (qi ci )ρi ≤ qi ci max ρi =⇒ C ≥ . i 16 32 maxi ρi i=1 i=1 i=1 Pn To maximize this lower bound, the adversary minimizes max ρi subject to the constraints n1 i=1 ∆i = 20ϵ and ∆i ≤ G/2 for all i (recall Eq. (32)). The minimum is achieved when the ratios ρi are equalized across √ the active components. Thus, we set ∆i = α ci for components i in an subset S ∗ ⊆ [n], and ∆i = 0 P“active” n otherwise. The scaling factor α is determined using the constraint n1 i=1 ∆i = 20ϵ: 1 X √ 20nϵ (α ci ) = 20ϵ =⇒ α = P √ . n ci i∈S ∗ ∗ i∈S

29

(36)

To certify the additional condition ∆i ≤ G/2 for all i, we require for all i ∈ S ∗ : √ √ 20nϵ ci maxi∈S ∗ ci √ G G G α ci ≤ ⇐⇒ P ⇐⇒ 1 P . √ ≤ √ ≤ 2 cj 2 40ϵ cj j∈S ∗ j∈S ∗ n This is the precisely the uniformity condition given in the theorem statement. The subset S ∗ is defined as the largest set of cheapest components satisfying this condition. Finally, substituting α back into the cost lower bound, we conclude: G2 C≥ 32

P

√ 2 ci G2 = 20nϵ 12800ϵ2

i∈S ∗

30

1 X√ ci n ∗ i∈S

!2 .

C

Extensions and Generality of the Theoretical Analysis

C.1

Average Iterate (x̄T ) vs. Last Iterate (xT )

While Algorithm 1 returns the average iterate x̄T following standard conventions in the SGD literature, the properties of the last iterate xT have also been extensively studied [Shamir and Zhang, 2013, Jain et al., 2019, Harvey et al., 2019]. As shown by Harvey et al. [2019], for a fixed step size, the last-iterate of SGD is a factor of log(T ) worse than the average. While the average iterate generally achieves superior performance—sometimes strictly so—specific step-size modification schemes [Jain et al., 2019] allow the expected suboptimality of the last iterate to enjoy the same convergence guarantees. Consequently, sampling with the optimal distribution p∗ derived in Theorem 4.1 provides similar theoretical guarantees for the last iterate as we establish for the averaged iterate.

C.2

Generality of the Cost Objective in Non-Convex Settings

The mathematical optimality of p∗ is not strictly limited to the convex setting. Its derivation depends solely on the relationship between iteration complexity (T ) and estimator variance (S(p)). Across various smoothness conditions, the number of iterations T required to reach an ϵ-stationary point is fundamentally bottlenecked by the stochastic gradient variance. Specifically: • Non-Convex SGD: Reaching an ϵ-stationary point requires an iteration complexity of T ∝ S(p) ϵ2 [Ghadimi and Lan, 2013]. • Adaptive Optimizers (Adam): √  Recent theoretical analyses of Adam-type algorithms demonstrate S(p) √ a convergence rate of Õ [Chen et al., 2018], implying T ∝ S(p) ϵ2 (modulo sub-polynomial T logarithmic factors). • Generalized Smoothness: LLM loss landscapes exhibit generalized smoothness, where local curvature grows proportionally with the gradient norm. Optimization theory establishes that gradient clipping and adaptive methods guarantee convergence in this regime [Zhang et al., 2020]. Crucially, this theory  establishes that the dominant stochastic term in the iteration complexity explicitly retains the O

S(p) ϵ2

dependence on the variance S(p). Because the total expected training cost K is the product of iterations and the expected per-step cost (T · C(p)), the objective to minimize reduces to: J(p) = S(p)C(p) Since this objective remains structurally identical across these diverse optimization settings, minimizing the variance-cost product via p∗ continues to be the precise optimal strategy. Thus, the optimal distribution √ p∗i ∝ Gi / ci derived in Theorem 4.1 remains justified in non-convex regimes.

C.3

Differences between Theory and Practice

In the GRPO experiments, our training objective is not convex. We assume in the main paper that the objective is convex for the simplicity of the analysis, although the previous subsection explains why one might expect the result to hold in non-convex settings as well. Furthermore, in GRPO, the objective changes every time a new set of rollouts is collected. As such, one can think of our implementation as many rounds of cost-aware learning. With every new set of rollouts, we define a new sampling distribution over the newly generated pool. The policy is then optimized by iterating over mini-batches sampled from this reweighted pool, ensuring the cost-aware distribution actively shapes the variance of the gradient updates at each step. 31

D

Experimental Details and Ablations

D.1

Experimental Details

To ensure the reproducibility of our experiments, we provide a comprehensive overview of the hyperparameter configurations and the computational framework employed. Our implementation is built upon the Verl framework, utilizing vLLM [Kwon et al., 2023] for efficient inference and Flash-Attention [Dao et al., 2022] for optimized memory throughput.

D.2

Implementation

The primary modifications to the Verl codebase involve the construction of mini-batches for the policy gradient update and the integration of importance sampling. To maintain training stability, we re-center the importance weights around 1. This normalization allows us to maintain a consistent gradient norm across experiments. Crucially, this simply multiplies the training objective by a constant.

32

D.3

Hyperparameters

Table 3 summarizes the key hyperparameters used for both the GRPO baseline and our GRPO + ZVF method for the 1.5B experiments. For the Qwen3-8B experiments, we detail hyperparameters in the Table 4. We report Scale-RL hyperparameters in Table 5, largely influenced by the findings of [Khatri et al., 2025]. Hyperparameter Base Model Algorithm Learning Rate Weight Decay Train Batch Size Rollout Samples per Prompt (n) Rollout Temperature PPO Mini-Batch Size KL Loss Coefficient Clip Ratio Low Clip Ratio High

Value Qwen/Qwen2.5-Math-1.5B-Instruct GRPO 1 × 10−6 0.1 256 16 1.0 64 0.001 0.2 0.28

Table 3: Key Training Hyperparameters for GRPO and GRPO + ZVF.

Hyperparameter Base Model Algorithm Learning Rate Total Epochs Train Batch Size PPO Mini-Batch Size Rollout Samples per Prompt (n) Rollout Temperature Max Prompt Length Max Response Length KL Loss Coefficient

Value Qwen/Qwen3-8B GRPO 1 × 10−6 15 512 256 5 1.0 1024 4096 0.001

Table 4: Key Training Hyperparameters for Qwen3-8B Training.

33

Hyperparameter Base Model Advantage Estimator Policy Loss Type Learning Rate LR Warmup Steps Weight Decay Train Batch Size PPO Mini-Batch Size Gradient Clipping KL Loss Coefficient Entropy Coefficient Clip Ratio Low Clip Ratio High Clip Ratio C LM Head Precision

Value Qwen/Qwen2.5-Math-1.5B-Instruct GRPO CISPO 5 × 10−7 100 0.01 512 128 1.0 0.001 0 0.0 5.0 3.0 FP32

Table 5: Key Training Hyperparameters for ScaleRL/CISPO Optimized Script

D.4

Computational Requirements

All experiments were conducted using the hardware configurations standard to the Verl framework. The training durations are as follows: • GRPO: Approximately 24 hours (1 day). • GRPO + ZVF: Approximately 36 hours (1.5 days). • Scale RL: Approximately 36 hours (1.5 days). For the Qwen3-8B experiments, we used 8 NVIDIA H100 GPUs per training run. The training runs required around 5 days each.

34

D.5

Ablation on variants of p∗

In this section we ablate the different variants of p∗ . In particular, we also report p∗smooth for α = 0.05, α = 0.1. Furthermore, we consider p∗LEN , in which (p∗LEN )i = √1ci . This sampling distribution only considers the length of sequences in its distribution. We report these variants both for GRPO and GRPO+ZVF. In particular with the GRPO setting, we see that p∗LEN does not attain very strong performance. We hypothesize that this is because many shorter responses have 0 advantage, and hence there is little training signal. As in the main paper, we also report final performance for all methods.

Figure 7: Cumulative tokens compared with AIME pass@1/mean@32 accuracy throughout training for both GRPO and GRPO+ZVF methods.

Setting

Method

AIME

AMC

MATH500

GSM8K

Avg. Accuracy

GRPO GRPO GRPO GRPO GRPO GRPO

No sampling p∗ p∗smooth (α = 0.01) p∗smooth (α = 0.05) p∗smooth (α = 0.1) p∗ -LEN

61.3 65.6 65.3 65.6 65.4 60.3

64.1 71.2 68.1 68.3 66.0 64.8

73.2 73.2 72.3 72.8 72.6 73.0

86.2 86.0 86.0 86.1 85.8 86.0

71.2 74.0 72.9 73.2 72.5 71.0

ZVF ZVF ZVF ZVF ZVF ZVF

No sampling p∗ p∗smooth (α = 0.01) p∗smooth (α = 0.05) p∗smooth (α = 0.1) p∗ -LEN

58.1 57.6 62.9 60.5 61.4 58.0

64.7 63.0 68.0 65.9 68.3 63.3

73.2 72.9 73.3 73.0 72.9 73.9

86.1 85.9 85.8 86.4 86.0 86.2

70.6 69.8 72.5 71.4 72.2 70.3

Table 6: (Math-1.5B-Instruct, pass@1/mean@32): Complete results.

35

Setting

Method

AIME

AMC

MATH500

GSM8K

Avg. Accuracy

GRPO GRPO GRPO GRPO

No sampling p∗ p∗smooth (α = 0.01) p∗smooth (α = 0.05)

58.3 58.5 58.5 58.3

72.8 74.7 74.0 73.6

87.0 86.7 86.7 86.6

96.0 95.2 94.9 95.6

78.5 78.8 78.5 78.5

Table 7: (8B-Base, pass@1/mean@32): Complete results.

D.6

Full CISPO results

We report performance across benchmarks. The CISPO curves are reported in Figure 8 and the best checkpoint performance is reported in Table 8.

Figure 8: Full CISPO objective results for Qwen2.5-Math-1.5B-Instruct on AIME. We see robustness to smoothed cost-aware learning.

Setting

Method

AIME

AMC

MATH500

GSM8K

Avg. Accuracy

GRPO GRPO GRPO GRPO GRPO

No sampling p∗ p∗smooth (α = 0.01) p∗smooth (α = 0.05) p∗smooth (α = 0.1)

57.2 61.3 61.7 60.5 61.0

61.8 65.7 65.4 66.0 65.2

73.0 73.6 73.5 73.0 73.6

85.9 86.2 86.0 86.1 86.0

69.5 71.7 71.6 71.4 71.5

Table 8: (Math-1.5B-Instruct, pass@1/mean@32): Complete results for CISPO.

36

D.7

Full sub-optimality results

In the main paper, we report the sub-optimality metrics for the 8B and 1.5B models under the GRPO setting. Here, we also include the sub-optimality metrics for the GRPO+ZVF setting and the CISPO settings in Figures 9 and 10 respectively.

Figure 9: Sub-optimality metrics for 1.5B GRPO+ZVF training run.

D.8

Figure 10: Sub-optimality metrics for 1.5B CISPO training objective run.

Additional results for Qwen2.5-Math-1.5B-Instruct

We report AMC performance as a function of cumulative tokens in Figure 11.

Accuracy vs. Token Count (GRPO)

Accuracy vs. Token Count (GRPO+ZVF)

70 65

Accuracy (%)

Accuracy (%)

65 60 55

Method

No sampling

50

p* * psmooth ( = 0.01)

45 0.0

0.2

0.4

0.6

0.8

1.0

Cumulative Tokens

1.2

1.4

60 55

Method

50

No sampling

p* * psmooth ( = 0.01)

45

1.6 1e9

0.0

0.2

0.4

0.6

Cumulative Tokens

0.8

1.0

1e9

Figure 11: Cumulative tokens compared with AMC pass@1/mean@32 accuracy throughout training for both GRPO and GRPO+ZVF settings for the 1.5B model.

37

D.9

Additional results for Qwen3-8B

We present the results for the smoothed variants of the 8B model on AIME in Figure 12. Additionally, we plot AMC performance as a function of cumulative tokens in Figure 13. In Table 7, we report full results on best checkpoint accuracy for smoothed variants on four benchmarks.

Accuracy vs. Token Count (Qwen3-8B)

60

Accuracy (%)

55

Method

50

No sampling

p* * psmooth ( = 0.01) * psmooth ( = 0.05)

45 40 35 0

100M

200M

Cumulative Tokens

300M

400M

Figure 12: Qwen3-8B AIME results for all variants.

Accuracy vs. Token Count (Qwen3-8B, AMC) 75

Method

Accuracy (%)

70

GRPO CA-GRPO

65

(1) 30% fewer tokens

60

(2) 47% fewer tokens

55 0

100M

200M

300M

Cumulative Tokens

400M

Figure 13: AMC results for Qwen3-8B Base.

38

500M