Year: 2026
Type: preprint
Abstract: We consider the problem of Cost-Aware Learning, where sampling different components of a finite-sum objective incurs different costs. The objective is to reach a target error while minimizing the total cost. We propose Cost-Aware SGD, which uses a distribution based on gradient norms and costs to sample components. We provide a thorough analysis of this algorithm, including cost-improvement bounds over baselines, a characterization of distribution proxy sub-optimality, and a lower bound. We apply our theoretical insights to reinforcement learning with language models, where the computational cost of sequence-level policy gradients varies with length. We find that the advantage magnitude serves as a high-fidelity proxy for gradient norms, and use this to introduce Cost-Aware GRPO. Empirical results on 1.5B, 4B, and 8B LLMs demonstrate that this algorithm significantly reduces the tokens used in policy optimization while matching or exceeding baseline accuracy.
Source: arXiv (Cornell University)
Institutions: Harvard University, Tel Aviv University, Google (United States)
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W7159687670
Cites:
Cited by:
Related to: