Cost-Aware Learning
For practitioners training large language models, this work reduces computational cost of policy optimization without sacrificing performance, though the theoretical contributions are incremental.
This paper introduces cost-aware learning where sampling different components of a finite-sum objective incurs different costs, proposing Cost-Aware SGD for convex functions with cost complexity analysis and a lower bound. Applied to RL for LLMs, Cost-Aware GRPO reduces tokens used in policy optimization by up to 30% while matching or exceeding baseline accuracy on 1.5B and 8B models.
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.